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

Author: seve-martinez

src/ast/mod.rs

@@ -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
@@ -683,16 +701,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>)
@@ -1251,20 +1275,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})"),

src/parser/mod.rs

@@ -1712,12 +1712,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 {

tests/sqlparser_common.rs

@@ -2495,14 +2495,74 @@ 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)";
     let select = verified_only_select(sql);
     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)),
     );
@@ -2529,7 +2589,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)),
     );