support min,max for decimal data type

liukun4515 · liukun4515 · commit 450ccbbb9817 · 2021-11-27T10:01:48.000+08:00
diff --git a/datafusion-examples/examples/csv_decimal_sql.rs b/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(())
+}
diff --git a/datafusion/Cargo.toml b/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"
diff --git a/datafusion/src/execution/context.rs b/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();
diff --git a/datafusion/src/physical_plan/aggregates.rs b/datafusion/src/physical_plan/aggregates.rs
@@ -209,6 +209,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 +245,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 +291,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(())
     }
 
diff --git a/datafusion/src/physical_plan/expressions/coercion.rs b/datafusion/src/physical_plan/expressions/coercion.rs
@@ -20,7 +20,7 @@
 use arrow::datatypes::DataType;
 
 /// Determine if a DataType is signed numeric or not
-pub fn is_signed_numeric(dt: &DataType) -> bool {
+pub(crate) fn is_signed_numeric(dt: &DataType) -> bool {
     matches!(
         dt,
         DataType::Int8
@@ -29,12 +29,13 @@ pub fn is_signed_numeric(dt: &DataType) -> bool {
             | DataType::Int64
             | DataType::Float16
             | DataType::Float32
-            | DataType::Float64
+            | DataType::Float64 // TODO liukun4515
+                                // | DataType::Decimal(_,_)
     )
 }
 
 /// Determine if a DataType is numeric or not
-pub fn is_numeric(dt: &DataType) -> bool {
+fn is_numeric(dt: &DataType) -> bool {
     is_signed_numeric(dt)
         || match dt {
             DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => {
@@ -125,6 +126,11 @@ pub fn numerical_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<Da
         return Some(lhs_type.clone());
     }
 
+    // TODO liukun4515
+    // In the decimal data type, what to do if the metadata of decimal type is diff.
+    // add decimal data type, diff operator we should have diff rule to do coercion.
+    // first step, we can just support decimal type in case which left and right datatype are the same
+
     // these are ordered from most informative to least informative so
     // that the coercion removes the least amount of information
     match (lhs_type, rhs_type) {
@@ -170,6 +176,8 @@ pub fn order_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option<DataTy
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::arrow::datatypes::DataType::Int8;
+    use arrow::datatypes::DataType::{Float32, Float64, Int16, Int32, Int64};
 
     #[test]
     fn test_dictionary_type_coersion() {
@@ -192,4 +200,75 @@ mod tests {
         let rhs_type = Dictionary(Box::new(Int8), Box::new(Utf8));
         assert_eq!(dictionary_coercion(&lhs_type, &rhs_type), Some(Utf8));
     }
+
+    #[test]
+    fn test_is_signed_numeric() {
+        assert!(is_signed_numeric(&DataType::Int8));
+        assert!(is_signed_numeric(&DataType::Int16));
+        assert!(is_signed_numeric(&DataType::Int32));
+        assert!(is_signed_numeric(&DataType::Int64));
+        assert!(is_signed_numeric(&DataType::Float16));
+        assert!(is_signed_numeric(&DataType::Float32));
+        assert!(is_signed_numeric(&DataType::Float64));
+
+        // decimal data type
+        assert!(is_signed_numeric(&DataType::Decimal(12, 2)));
+        assert!(is_signed_numeric(&DataType::Decimal(14, 10)));
+
+        // negative test
+        assert!(!is_signed_numeric(&DataType::UInt64));
+        assert!(!is_signed_numeric(&DataType::UInt16));
+    }
+
+    #[test]
+    fn test_is_numeric() {
+        assert!(is_numeric(&DataType::Int8));
+        assert!(is_numeric(&DataType::Int16));
+        assert!(is_numeric(&DataType::Int32));
+        assert!(is_numeric(&DataType::Int64));
+        assert!(is_numeric(&DataType::Float16));
+        assert!(is_numeric(&DataType::Float32));
+        assert!(is_numeric(&DataType::Float64));
+
+        // decimal data type
+        assert!(is_numeric(&DataType::Decimal(12, 2)));
+        assert!(is_numeric(&DataType::Decimal(14, 10)));
+
+        // unsigned test
+        assert!(is_numeric(&DataType::UInt8));
+        assert!(is_numeric(&DataType::UInt16));
+        assert!(is_numeric(&DataType::UInt32));
+        assert!(is_numeric(&DataType::UInt64));
+
+        // negative test
+        assert!(!is_numeric(&DataType::Boolean));
+        assert!(!is_numeric(&DataType::Date32));
+    }
+
+    #[test]
+    fn test_numerical_coercion() {
+        // negative test
+        assert_eq!(
+            None,
+            numerical_coercion(&DataType::Float64, &DataType::Binary)
+        );
+        assert_eq!(
+            None,
+            numerical_coercion(&DataType::Float64, &DataType::Utf8)
+        );
+
+        // positive test
+        let test_types = vec![Int8, Int16, Int32, Int64, Float32, Float64];
+        let mut index = test_types.len();
+        while index > 0 {
+            let this_type = &test_types[index - 1];
+            for i in 0..index {
+                assert_eq!(
+                    Some(this_type.clone()),
+                    numerical_coercion(this_type, &test_types[i])
+                );
+            }
+            index -= 1;
+        }
+    }
 }
diff --git a/datafusion/src/physical_plan/expressions/min_max.rs b/datafusion/src/physical_plan/expressions/min_max.rs
@@ -487,6 +487,36 @@ mod tests {
     use arrow::datatypes::*;
     use arrow::record_batch::RecordBatch;
 
+    #[test]
+    fn min_decimal() -> Result<()> {
+        todo!()
+    }
+
+    #[test]
+    fn max_decimal() -> Result<()> {
+        todo!()
+    }
+
+    #[test]
+    fn min_decimal_with_nulls() -> Result<()> {
+        todo!()
+    }
+
+    #[test]
+    fn max_decimal_with_nulls() -> Result<()> {
+        todo!()
+    }
+
+    #[test]
+    fn min_decimal_with_all_nulls() -> Result<()> {
+        todo!()
+    }
+
+    #[test]
+    fn max_decimal_with_all_nulls() -> Result<()> {
+        todo!()
+    }
+
     #[test]
     fn max_i32() -> Result<()> {
         let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]));
diff --git a/datafusion/src/physical_plan/functions.rs b/datafusion/src/physical_plan/functions.rs
@@ -144,6 +144,7 @@ impl Signature {
 }
 
 ///A function's volatility, which defines the functions eligibility for certain optimizations
+///Ref from postgresql https://www.postgresql.org/docs/current/xfunc-volatility.html
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
 pub enum Volatility {
     /// Immutable - An immutable function will always return the same output when given the same input. An example of this is [BuiltinScalarFunction::Cos].
diff --git a/datafusion/src/physical_plan/type_coercion.rs b/datafusion/src/physical_plan/type_coercion.rs
@@ -76,6 +76,7 @@ pub fn data_types(
     if current_types.is_empty() {
         return Ok(vec![]);
     }
+
     let valid_types = get_valid_types(&signature.type_signature, current_types)?;
 
     if valid_types
@@ -144,6 +145,8 @@ fn maybe_data_types(
     valid_types: &[DataType],
     current_types: &[DataType],
 ) -> Option<Vec<DataType>> {
+    // TODO liukun4515
+
     if valid_types.len() != current_types.len() {
         return None;
     }
@@ -155,7 +158,6 @@ fn maybe_data_types(
         if current_type == valid_type {
             new_type.push(current_type.clone())
         } else {
-            // attempt to coerce
             if can_coerce_from(valid_type, current_type) {
                 new_type.push(valid_type.clone())
             } else {
@@ -171,9 +173,11 @@ fn maybe_data_types(
 /// (losslessly converted) into a value of `type_to`
 ///
 /// See the module level documentation for more detail on coercion.
-pub fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
+fn can_coerce_from(type_into: &DataType, type_from: &DataType) -> bool {
     use self::DataType::*;
+    // TODO liukun4515
     match type_into {
+        // TODO, decimal data type, we just support the decimal
         Int8 => matches!(type_from, Int8),
         Int16 => matches!(type_from, Int8 | Int16 | UInt8),
         Int32 => matches!(type_from, Int8 | Int16 | Int32 | UInt8 | UInt16),
diff --git a/datafusion/src/sql/planner.rs b/datafusion/src/sql/planner.rs
diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs

Original file line number	Diff line number	Diff line change
`@@ -144,6 +144,7 @@ impl Signature {`
`144`	`144`	`}`
`145`	`145`
`146`	`146`	`///A function's volatility, which defines the functions eligibility for certain optimizations`
	`147`	`+///Ref from postgresql https://www.postgresql.org/docs/current/xfunc-volatility.html`
`147`	`148`	`#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]`
`148`	`149`	`pub enum Volatility {`
`149`	`150`	`/// Immutable - An immutable function will always return the same output when given the same input. An example of this is [BuiltinScalarFunction::Cos].`