support min,max for decimal data type

Author: liukun4515

datafusion-examples/examples/csv_decimal_sql.rs

@@ -0,0 +1,54 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::error::Result;
+use datafusion::prelude::*;
+use std::sync::Arc;
+
+/// This example demonstrates executing a simple query against an Arrow data source (CSV) and
+/// fetching results
+#[tokio::main]
+async fn main() -> Result<()> {
+    // create local execution context
+    let mut ctx = ExecutionContext::new();
+
+    let testdata = datafusion::test_util::arrow_test_data();
+
+    // schema with decimal type
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("c1", DataType::Decimal(10, 6), false),
+        Field::new("c2", DataType::Float64, false),
+        Field::new("c3", DataType::Boolean, false),
+    ]));
+
+    // register csv file with the execution context
+    ctx.register_csv(
+        "aggregate_simple",
+        &format!("{}/csv/aggregate_simple.csv", testdata),
+        CsvReadOptions::new().schema(&schema),
+    )
+    .await?;
+
+    // execute the query
+    let df = ctx.sql("select c1 from aggregate_simple").await?;
+
+    // print the results
+    df.show().await?;
+
+    Ok(())
+}

datafusion/Cargo.toml

@@ -52,8 +52,10 @@ avro = ["avro-rs", "num-traits"]
 [dependencies]
 ahash = "0.7"
 hashbrown = { version = "0.11", features = ["raw"] }
-arrow = { version = "6.2.0", features = ["prettyprint"] }
-parquet = { version = "6.2.0", features = ["arrow"] }
+arrow = { path = "/Users/kliu3/Documents/github/arrow-rs/arrow", features = ["prettyprint"] }
+#arrow = { version = "6.2.0", features = ["prettyprint"] }
+#parquet = { version = "6.2.0", features = ["arrow"] }
+parquet = { path = "/Users/kliu3/Documents/github/arrow-rs/parquet", features = ["arrow"] }
 sqlparser = "0.12"
 paste = "^1.0"
 num_cpus = "1.13.0"

datafusion/src/execution/context.rs

@@ -3895,6 +3895,31 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn aggregate_decimal() -> Result<()> {
+        let mut ctx = ExecutionContext::new();
+        // schema with data
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("c1", DataType::Decimal(10, 6), false),
+            Field::new("c2", DataType::Float64, false),
+            Field::new("c3", DataType::Boolean, false),
+        ]));
+
+        ctx.register_csv(
+            "aggregate_simple",
+            "tests/aggregate_simple.csv",
+            CsvReadOptions::new().schema(&schema),
+        )
+        .await?;
+
+        // decimal query
+        let result = plan_and_collect(&mut ctx, "select min(c1) from aggregate_simple")
+            .await
+            .unwrap();
+        println!("{:?}", result);
+        Ok(())
+    }
+
     #[tokio::test]
     async fn create_external_table_with_timestamps() {
         let mut ctx = ExecutionContext::new();

datafusion/src/physical_plan/aggregates.rs

@@ -28,15 +28,16 @@
 
 use super::{
     functions::{Signature, Volatility},
-    type_coercion::{coerce, data_types},
     Accumulator, AggregateExpr, PhysicalExpr,
 };
 use crate::error::{DataFusionError, Result};
+use crate::physical_plan::coercion_rule::aggregate_rule::{coerce_exprs, coerce_types};
 use crate::physical_plan::distinct_expressions;
 use crate::physical_plan::expressions;
 use arrow::datatypes::{DataType, Field, Schema, TimeUnit};
 use expressions::{avg_return_type, sum_return_type};
 use std::{fmt, str::FromStr, sync::Arc};
+
 /// the implementation of an aggregate function
 pub type AccumulatorFunctionImplementation =
     Arc<dyn Fn() -> Result<Box<dyn Accumulator>> + Send + Sync>;
@@ -87,96 +88,125 @@ impl FromStr for AggregateFunction {
                 return Err(DataFusionError::Plan(format!(
                     "There is no built-in function named {}",
                     name
-                )))
+                )));
             }
         })
     }
 }
 
-/// Returns the datatype of the scalar function
+/// Returns the datatype of the aggregate function.
+/// This is used to get the returned data type for aggregate logical expr.
 pub fn return_type(fun: &AggregateFunction, arg_types: &[DataType]) -> Result<DataType> {
     // Note that this function *must* return the same type that the respective physical expression returns
     // or the execution panics.
 
-    // verify that this is a valid set of data types for this function
-    data_types(arg_types, &signature(fun))?;
+    let coerced_data_types = coerce_types(fun, arg_types, &signature(fun))?;
 
     match fun {
+        // TODO liukun4515
+        // If the datafusion is compatible with PostgreSQL, the returned data type should be INT64.
         AggregateFunction::Count | AggregateFunction::ApproxDistinct => {
             Ok(DataType::UInt64)
         }
-        AggregateFunction::Max | AggregateFunction::Min => Ok(arg_types[0].clone()),
-        AggregateFunction::Sum => sum_return_type(&arg_types[0]),
-        AggregateFunction::Avg => avg_return_type(&arg_types[0]),
+        AggregateFunction::Max | AggregateFunction::Min => {
+            // For min and max agg function, the returned type is same as input type.
+            // The coerced_data_types is same with input_types.
+            Ok(coerced_data_types[0].clone())
+        }
+        AggregateFunction::Sum => sum_return_type(&coerced_data_types[0]),
+        AggregateFunction::Avg => avg_return_type(&coerced_data_types[0]),
         AggregateFunction::ArrayAgg => Ok(DataType::List(Box::new(Field::new(
             "item",
-            arg_types[0].clone(),
+            coerced_data_types[0].clone(),
             true,
         )))),
     }
 }
 
 /// Create a physical (function) expression.
-/// This function errors when `args`' can't be coerced to a valid argument type of the function.
+/// This function errors when `input_phy_exprs`' can't be coerced to a valid argument type of the function.
 pub fn create_aggregate_expr(
     fun: &AggregateFunction,
     distinct: bool,
-    args: &[Arc<dyn PhysicalExpr>],
+    input_phy_exprs: &[Arc<dyn PhysicalExpr>],
     input_schema: &Schema,
     name: impl Into<String>,
 ) -> Result<Arc<dyn AggregateExpr>> {
     let name = name.into();
-    let arg = coerce(args, input_schema, &signature(fun))?;
-    if arg.is_empty() {
+    // get the coerced phy exprs if some expr need try cast
+    let coerced_phy_exprs =
+        coerce_exprs(fun, input_phy_exprs, input_schema, &signature(fun))?;
+    if coerced_phy_exprs.is_empty() {
         return Err(DataFusionError::Plan(format!(
             "Invalid or wrong number of arguments passed to aggregate: '{}'",
             name,
         )));
     }
-    let arg = arg[0].clone();
-
-    let arg_types = args
+    let coerced_types = coerced_phy_exprs
         .iter()
         .map(|e| e.data_type(input_schema))
         .collect::<Result<Vec<_>>>()?;
+    let first_coerced_phy_expr = coerced_phy_exprs[0].clone();
 
-    let return_type = return_type(fun, &arg_types)?;
+    // get the result data type for this aggregate function
+    let input_phy_types = input_phy_exprs
+        .iter()
+        .map(|e| e.data_type(input_schema))
+        .collect::<Result<Vec<_>>>()?;
+    let return_type = return_type(fun, &input_phy_types)?;
 
     Ok(match (fun, distinct) {
-        (AggregateFunction::Count, false) => {
-            Arc::new(expressions::Count::new(arg, name, return_type))
-        }
+        (AggregateFunction::Count, false) => Arc::new(expressions::Count::new(
+            first_coerced_phy_expr,
+            name,
+            return_type,
+        )),
         (AggregateFunction::Count, true) => {
             Arc::new(distinct_expressions::DistinctCount::new(
-                arg_types,
-                args.to_vec(),
+                coerced_types,
+                coerced_phy_exprs,
                 name,
                 return_type,
             ))
         }
-        (AggregateFunction::Sum, false) => {
-            Arc::new(expressions::Sum::new(arg, name, return_type))
-        }
+        (AggregateFunction::Sum, false) => Arc::new(expressions::Sum::new(
+            first_coerced_phy_expr,
+            name,
+            return_type,
+        )),
         (AggregateFunction::Sum, true) => {
             return Err(DataFusionError::NotImplemented(
                 "SUM(DISTINCT) aggregations are not available".to_string(),
             ));
         }
-        (AggregateFunction::ApproxDistinct, _) => Arc::new(
-            expressions::ApproxDistinct::new(arg, name, arg_types[0].clone()),
-        ),
-        (AggregateFunction::ArrayAgg, _) => {
-            Arc::new(expressions::ArrayAgg::new(arg, name, arg_types[0].clone()))
-        }
-        (AggregateFunction::Min, _) => {
-            Arc::new(expressions::Min::new(arg, name, return_type))
-        }
-        (AggregateFunction::Max, _) => {
-            Arc::new(expressions::Max::new(arg, name, return_type))
-        }
-        (AggregateFunction::Avg, false) => {
-            Arc::new(expressions::Avg::new(arg, name, return_type))
+        // TODO optimize the expr and make arg more meaningful
+        (AggregateFunction::ApproxDistinct, _) => {
+            Arc::new(expressions::ApproxDistinct::new(
+                first_coerced_phy_expr,
+                name,
+                coerced_types[0].clone(),
+            ))
         }
+        (AggregateFunction::ArrayAgg, _) => Arc::new(expressions::ArrayAgg::new(
+            first_coerced_phy_expr,
+            name,
+            coerced_types[0].clone(),
+        )),
+        (AggregateFunction::Min, _) => Arc::new(expressions::Min::new(
+            first_coerced_phy_expr,
+            name,
+            return_type,
+        )),
+        (AggregateFunction::Max, _) => Arc::new(expressions::Max::new(
+            first_coerced_phy_expr,
+            name,
+            return_type,
+        )),
+        (AggregateFunction::Avg, false) => Arc::new(expressions::Avg::new(
+            first_coerced_phy_expr,
+            name,
+            return_type,
+        )),
         (AggregateFunction::Avg, true) => {
             return Err(DataFusionError::NotImplemented(
                 "AVG(DISTINCT) aggregations are not available".to_string(),
@@ -209,6 +239,7 @@ static TIMESTAMPS: &[DataType] = &[
 
 static DATES: &[DataType] = &[DataType::Date32, DataType::Date64];
 
+// TODO liukun4515
 /// the signatures supported by the function `fun`.
 pub fn signature(fun: &AggregateFunction) -> Signature {
     // note: the physical expression must accept the type returned by this function or the execution panics.
@@ -244,6 +275,29 @@ mod tests {
 
         let observed = return_type(&AggregateFunction::Max, &[DataType::Int32])?;
         assert_eq!(DataType::Int32, observed);
+
+        let observed = return_type(&AggregateFunction::Min, &[DataType::Decimal(10, 6)])?;
+        assert_eq!(DataType::Decimal(10, 6), observed);
+
+        let observed =
+            return_type(&AggregateFunction::Max, &[DataType::Decimal(28, 13)])?;
+        assert_eq!(DataType::Decimal(28, 13), observed);
+
+        // TODO liukun4515
+        // Add test to create agg expr
+        let fun = AggregateFunction::Min;
+        let input_schema =
+            Schema::new(vec![Field::new("c1", DataType::Decimal(10, 6), false)]);
+        // let input_schema =
+        //     Schema::new(vec![Field::new("c1", DataType::Int32, false)]);
+        let expr: Arc<dyn PhysicalExpr> =
+            Arc::new(expressions::Column::new_with_schema("c1", &input_schema).unwrap());
+        let args = vec![expr];
+        let result_agg =
+            create_aggregate_expr(&fun, false, &args[0..1], &input_schema, "c1")?;
+        // assert the filed type
+        assert_eq!(input_schema.field(0), &result_agg.field().unwrap());
+        assert_eq!("min(c1)", result_agg.name());
         Ok(())
     }
 
@@ -267,6 +321,10 @@ mod tests {
 
         let observed = return_type(&AggregateFunction::Count, &[DataType::Int8])?;
         assert_eq!(DataType::UInt64, observed);
+
+        let observed =
+            return_type(&AggregateFunction::Count, &[DataType::Decimal(28, 13)])?;
+        assert_eq!(DataType::UInt64, observed);
         Ok(())
     }