mirrors/datafusion-sqlparse
1
0
Fork 0
mirror of https://github.com/apache/datafusion-sqlparser-rs.git synced 2025-08-04 06:18:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

adding support for scale in CEIL and FLOOR functions (#1377)

Browse source
This commit is contained in:
Seve Martinez 2024-08-14 06:11:40 -07:00 • committed by GitHub
parent b072ce2589
commit c2f46ae07b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 113 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

48
src/ast/mod.rs
View file

@ -500,6 +500,24 @@ pub enum ExtractSyntax {
Comma,
}
/// The syntax used in a CEIL or FLOOR expression.
///
/// The `CEIL/FLOOR(<datetime value expression> TO <time unit>)` is an Amazon Kinesis Data Analytics extension.
/// See <https://docs.aws.amazon.com/kinesisanalytics/latest/sqlref/sql-reference-ceil.html> for
/// details.
///
/// Other dialects either support `CEIL/FLOOR( <expr> [, <scale>])` format or just
/// `CEIL/FLOOR(<expr>)`.
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
pub enum CeilFloorKind {
/// `CEIL( <expr> TO <DateTimeField>)`
DateTimeField(DateTimeField),
/// `CEIL( <expr> [, <scale>])`
Scale(Value),
}
/// An SQL expression of any type.
///
/// The parser does not distinguish between expressions of different types
@ -674,16 +692,22 @@ pub enum Expr {
/// ```sql
/// CEIL(<expr> [TO DateTimeField])
/// ```
/// ```sql
/// CEIL( <input_expr> [, <scale_expr> ] )
/// ```
Ceil {
expr: Box<Expr>,
field: DateTimeField,
field: CeilFloorKind,
},
/// ```sql
/// FLOOR(<expr> [TO DateTimeField])
/// ```
/// ```sql
/// FLOOR( <input_expr> [, <scale_expr> ] )
///
Floor {
expr: Box<Expr>,
field: DateTimeField,
field: CeilFloorKind,
},
/// ```sql
/// POSITION(<expr> in <expr>)
@ -1230,20 +1254,20 @@ impl fmt::Display for Expr {
ExtractSyntax::From => write!(f, "EXTRACT({field} FROM {expr})"),
ExtractSyntax::Comma => write!(f, "EXTRACT({field}, {expr})"),
},
Expr::Ceil { expr, field } => {
if field == &DateTimeField::NoDateTime {
Expr::Ceil { expr, field } => match field {
CeilFloorKind::DateTimeField(DateTimeField::NoDateTime) => {
write!(f, "CEIL({expr})")
} else {
write!(f, "CEIL({expr} TO {field})")
}
}
Expr::Floor { expr, field } => {
if field == &DateTimeField::NoDateTime {
CeilFloorKind::DateTimeField(dt_field) => write!(f, "CEIL({expr} TO {dt_field})"),
CeilFloorKind::Scale(s) => write!(f, "CEIL({expr}, {s})"),
},
Expr::Floor { expr, field } => match field {
CeilFloorKind::DateTimeField(DateTimeField::NoDateTime) => {
write!(f, "FLOOR({expr})")
} else {
write!(f, "FLOOR({expr} TO {field})")
}
}
CeilFloorKind::DateTimeField(dt_field) => write!(f, "FLOOR({expr} TO {dt_field})"),
CeilFloorKind::Scale(s) => write!(f, "FLOOR({expr}, {s})"),
},
Expr::Position { expr, r#in } => write!(f, "POSITION({expr} IN {in})"),
Expr::Collate { expr, collation } => write!(f, "{expr} COLLATE {collation}"),
Expr::Nested(ast) => write!(f, "({ast})"),

20
src/parser/mod.rs
View file

@ -1708,12 +1708,22 @@ impl<'a> Parser<'a> {
self.expect_token(&Token::LParen)?;
let expr = self.parse_expr()?;
// Parse `CEIL/FLOOR(expr)`
let mut field = DateTimeField::NoDateTime;
let keyword_to = self.parse_keyword(Keyword::TO);
if keyword_to {
let field = if self.parse_keyword(Keyword::TO) {
// Parse `CEIL/FLOOR(expr TO DateTimeField)`
field = self.parse_date_time_field()?;
}
CeilFloorKind::DateTimeField(self.parse_date_time_field()?)
} else if self.consume_token(&Token::Comma) {
// Parse `CEIL/FLOOR(expr, scale)`
match self.parse_value()? {
Value::Number(n, s) => CeilFloorKind::Scale(Value::Number(n, s)),
_ => {
return Err(ParserError::ParserError(
"Scale field can only be of number type".to_string(),
))
}
}
} else {
CeilFloorKind::DateTimeField(DateTimeField::NoDateTime)
};
self.expect_token(&Token::RParen)?;
if is_ceil {
Ok(Expr::Ceil {

64
tests/sqlparser_common.rs
View file

@ -2494,6 +2494,66 @@ fn parse_floor_number() {
verified_stmt("SELECT FLOOR(float_column) FROM my_table");
}
#[test]
fn parse_ceil_number_scale() {
verified_stmt("SELECT CEIL(1.5, 1)");
verified_stmt("SELECT CEIL(float_column, 3) FROM my_table");
}
#[test]
fn parse_floor_number_scale() {
verified_stmt("SELECT FLOOR(1.5, 1)");
verified_stmt("SELECT FLOOR(float_column, 3) FROM my_table");
}
#[test]
fn parse_ceil_scale() {
let sql = "SELECT CEIL(d, 2)";
let select = verified_only_select(sql);
#[cfg(feature = "bigdecimal")]
assert_eq!(
&Expr::Ceil {
expr: Box::new(Expr::Identifier(Ident::new("d"))),
field: CeilFloorKind::Scale(Value::Number(bigdecimal::BigDecimal::from(2), false)),
},
expr_from_projection(only(&select.projection)),
);
#[cfg(not(feature = "bigdecimal"))]
assert_eq!(
&Expr::Ceil {
expr: Box::new(Expr::Identifier(Ident::new("d"))),
field: CeilFloorKind::Scale(Value::Number(2.to_string(), false)),
},
expr_from_projection(only(&select.projection)),
);
}
#[test]
fn parse_floor_scale() {
let sql = "SELECT FLOOR(d, 2)";
let select = verified_only_select(sql);
#[cfg(feature = "bigdecimal")]
assert_eq!(
&Expr::Floor {
expr: Box::new(Expr::Identifier(Ident::new("d"))),
field: CeilFloorKind::Scale(Value::Number(bigdecimal::BigDecimal::from(2), false)),
},
expr_from_projection(only(&select.projection)),
);
#[cfg(not(feature = "bigdecimal"))]
assert_eq!(
&Expr::Floor {
expr: Box::new(Expr::Identifier(Ident::new("d"))),
field: CeilFloorKind::Scale(Value::Number(2.to_string(), false)),
},
expr_from_projection(only(&select.projection)),
);
}
#[test]
fn parse_ceil_datetime() {
let sql = "SELECT CEIL(d TO DAY)";
@ -2501,7 +2561,7 @@ fn parse_ceil_datetime() {
assert_eq!(
&Expr::Ceil {
expr: Box::new(Expr::Identifier(Ident::new("d"))),
field: DateTimeField::Day,
field: CeilFloorKind::DateTimeField(DateTimeField::Day),
},
expr_from_projection(only(&select.projection)),
);
@ -2528,7 +2588,7 @@ fn parse_floor_datetime() {
assert_eq!(
&Expr::Floor {
expr: Box::new(Expr::Identifier(Ident::new("d"))),
field: DateTimeField::Day,
field: CeilFloorKind::DateTimeField(DateTimeField::Day),
},
expr_from_projection(only(&select.projection)),
);