Add SQL planner support for calling round function with two argumen…

…ts (#2503) * test * specify correct signature for ROUND function
apache · May 10, 2022 · 8a29ed5 · 8a29ed5
1 parent 06fa7cc
commit 8a29ed5
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/datafusion/core/src/sql/planner.rs b/datafusion/core/src/sql/planner.rs
@@ -4355,6 +4355,10 @@ mod tests {
                     Field::new("j3_id", DataType::Int32, false),
                     Field::new("j3_string", DataType::Utf8, false),
                 ])),
+                "test_decimal" => Ok(Schema::new(vec![
+                    Field::new("id", DataType::Int32, false),
+                    Field::new("price", DataType::Decimal(10, 2), false),
+                ])),
                 "person" => Ok(Schema::new(vec![
                     Field::new("id", DataType::UInt32, false),
                     Field::new("first_name", DataType::Utf8, false),
@@ -4700,6 +4704,14 @@ mod tests {
         quick_test(sql, expected);
     }
 
+    #[tokio::test]
+    async fn round_decimal() {
+        let sql = "SELECT round(price/3, 2) FROM test_decimal";
+        let expected = "Projection: round(#test_decimal.price / Int64(3), Int64(2))\
+        \n  TableScan: test_decimal projection=None";
+        quick_test(sql, expected);
+    }
+
     #[ignore] // see https://github.com/apache/arrow-datafusion/issues/2469
     #[tokio::test]
     async fn aggregate_with_grouping_sets() {

diff --git a/datafusion/expr/src/function.rs b/datafusion/expr/src/function.rs
@@ -525,6 +525,15 @@ pub fn signature(fun: &BuiltinScalarFunction) -> Signature {
             ],
             fun.volatility(),
         ),
+        BuiltinScalarFunction::Round => Signature::one_of(
+            vec![
+                TypeSignature::Exact(vec![DataType::Float64, DataType::Int64]),
+                TypeSignature::Exact(vec![DataType::Float32, DataType::Int64]),
+                TypeSignature::Exact(vec![DataType::Float64]),
+                TypeSignature::Exact(vec![DataType::Float32]),
+            ],
+            fun.volatility(),
+        ),
         // math expressions expect 1 argument of type f64 or f32
         // priority is given to f64 because e.g. `sqrt(1i32)` is in IR (real numbers) and thus we
         // return the best approximation for it (in f64).