Add casts

Author: alamb

datafusion/functions-aggregate/src/sum.rs

@@ -33,6 +33,7 @@ use datafusion_common::types::{
 };
 use datafusion_common::{HashMap, Result, ScalarValue, exec_err, not_impl_err};
 use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
+use datafusion_expr::expr_fn::cast;
 use datafusion_expr::function::{
     AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs,
 };
@@ -391,12 +392,22 @@ fn sum_simplifier(mut agg: AggregateFunction, info: &SimplifyContext) -> Result<
         return Ok(Expr::AggregateFunction(agg));
     }
 
+    let lit_type = match &lit {
+        Expr::Literal(value, _) => value.data_type(),
+        _ => unreachable!("SplitResult::Split guarantees literal side"),
+    };
+    if lit_type == DataType::Null {
+        return Ok(Expr::AggregateFunction(agg));
+    }
+
     // Rewrite to SUM(arg)
     agg.params.args = vec![arg.clone()];
     let sum_agg = Expr::AggregateFunction(agg);
 
+    let count_agg = cast(crate::count::count(arg), lit_type);
+
     // sum(arg) + scalar * COUNT(arg)
-    Ok(sum_agg + (lit * crate::count::count(arg)))
+    Ok(sum_agg + (lit * count_agg))
 }
 
 fn has_common_rewrite_arg(arg: &Expr, info: &SimplifyContext) -> bool {

datafusion/sqllogictest/test_files/aggregates_simplify.slt

@@ -106,9 +106,10 @@ query TT
 EXPLAIN SELECT SUM(column1 + 1), SUM(column1 + 2) FROM sum_simplify_t;
 ----
 logical_plan
-01)Projection: sum(sum_simplify_t.column1) + count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1) + Int64(2) * count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(2))
-02)--Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
-03)----TableScan: sum_simplify_t projection=[column1]
+01)Projection: sum(sum_simplify_t.column1) + __common_expr_1 AS sum(sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1) + Int64(2) * __common_expr_1 AS sum(sum_simplify_t.column1 + Int64(2))
+02)--Projection: CAST(count(sum_simplify_t.column1) AS Int64) AS __common_expr_1, sum(sum_simplify_t.column1)
+03)----Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
+04)------TableScan: sum_simplify_t projection=[column1]
 physical_plan
 01)ProjectionExec: expr=[sum(sum_simplify_t.column1)@0 + count(sum_simplify_t.column1)@1 as sum(sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1)@0 + 2 * count(sum_simplify_t.column1)@1 as sum(sum_simplify_t.column1 + Int64(2))]
 02)--AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]
@@ -124,9 +125,10 @@ query TT
 EXPLAIN SELECT SUM(1 + column1), SUM(column1 + 2) FROM sum_simplify_t;
 ----
 logical_plan
-01)Projection: sum(sum_simplify_t.column1) + count(sum_simplify_t.column1) AS sum(Int64(1) + sum_simplify_t.column1), sum(sum_simplify_t.column1) + Int64(2) * count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(2))
-02)--Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
-03)----TableScan: sum_simplify_t projection=[column1]
+01)Projection: sum(sum_simplify_t.column1) + __common_expr_1 AS sum(Int64(1) + sum_simplify_t.column1), sum(sum_simplify_t.column1) + Int64(2) * __common_expr_1 AS sum(sum_simplify_t.column1 + Int64(2))
+02)--Projection: CAST(count(sum_simplify_t.column1) AS Int64) AS __common_expr_1, sum(sum_simplify_t.column1)
+03)----Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
+04)------TableScan: sum_simplify_t projection=[column1]
 physical_plan
 01)ProjectionExec: expr=[sum(sum_simplify_t.column1)@0 + count(sum_simplify_t.column1)@1 as sum(Int64(1) + sum_simplify_t.column1), sum(sum_simplify_t.column1)@0 + 2 * count(sum_simplify_t.column1)@1 as sum(sum_simplify_t.column1 + Int64(2))]
 02)--AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]
@@ -233,20 +235,24 @@ physical_plan
 03)----DataSourceExec: partitions=1, partition_sizes=[1]
 
 # volatile aggregate arguments
-query error DataFusion error: Arrow error: Invalid argument error: Invalid arithmetic operation: Float64 \* Int64
+query B
 SELECT SUM(random() + 1) < SUM(random() + 2) FROM sum_simplify_t;
+----
+true
 
 query TT
 EXPLAIN SELECT SUM(random() + 1) < SUM(random() + 2) FROM sum_simplify_t;
 ----
 logical_plan
-01)Projection: sum(random()) + Float64(2) * count(random()) > sum(random()) + CAST(count(random()) AS Float64) AS sum(random() + Int64(1)) < sum(random() + Int64(2))
-02)--Aggregate: groupBy=[[]], aggr=[[sum(random()), count(random())]]
-03)----TableScan: sum_simplify_t projection=[]
+01)Projection: sum(random()) + Float64(2) * __common_expr_1 > sum(random()) + __common_expr_1 AS sum(random() + Int64(1)) < sum(random() + Int64(2))
+02)--Projection: CAST(count(random()) AS Float64) AS __common_expr_1, sum(random())
+03)----Aggregate: groupBy=[[]], aggr=[[sum(random()), count(random())]]
+04)------TableScan: sum_simplify_t projection=[]
 physical_plan
-01)ProjectionExec: expr=[sum(random())@0 + 2 * count(random())@1 > sum(random())@0 + CAST(count(random())@1 AS Float64) as sum(random() + Int64(1)) < sum(random() + Int64(2))]
-02)--AggregateExec: mode=Single, gby=[], aggr=[sum(random()), count(random())]
-03)----DataSourceExec: partitions=1, partition_sizes=[1]
+01)ProjectionExec: expr=[sum(random())@1 + 2 * __common_expr_1@0 > sum(random())@1 + __common_expr_1@0 as sum(random() + Int64(1)) < sum(random() + Int64(2))]
+02)--ProjectionExec: expr=[CAST(count(random())@1 AS Float64) as __common_expr_1, sum(random())@0 as sum(random())]
+03)----AggregateExec: mode=Single, gby=[], aggr=[sum(random()), count(random())]
+04)------DataSourceExec: partitions=1, partition_sizes=[1]
 
 # Checks grouped aggregates with explicit ORDER BY return deterministic row order.
 query III
@@ -261,9 +267,10 @@ EXPLAIN SELECT column2, SUM(column1 + 1), SUM(column1 + 2) FROM sum_simplify_t G
 ----
 logical_plan
 01)Sort: sum_simplify_t.column2 DESC NULLS LAST
-02)--Projection: sum_simplify_t.column2, sum(sum_simplify_t.column1) + count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1) + Int64(2) * count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(2))
-03)----Aggregate: groupBy=[[sum_simplify_t.column2]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
-04)------TableScan: sum_simplify_t projection=[column1, column2]
+02)--Projection: sum_simplify_t.column2, sum(sum_simplify_t.column1) + __common_expr_1 AS sum(sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1) + Int64(2) * __common_expr_1 AS sum(sum_simplify_t.column1 + Int64(2))
+03)----Projection: CAST(count(sum_simplify_t.column1) AS Int64) AS __common_expr_1, sum_simplify_t.column2, sum(sum_simplify_t.column1)
+04)------Aggregate: groupBy=[[sum_simplify_t.column2]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
+05)--------TableScan: sum_simplify_t projection=[column1, column2]
 physical_plan
 01)SortPreservingMergeExec: [column2@0 DESC NULLS LAST]
 02)--SortExec: expr=[column2@0 DESC NULLS LAST], preserve_partitioning=[true]
@@ -284,7 +291,7 @@ EXPLAIN SELECT SUM(1 + column1), SUM(column1 + 1) FROM sum_simplify_t;
 ----
 logical_plan
 01)Projection: __common_expr_1 AS sum(Int64(1) + sum_simplify_t.column1), __common_expr_1 AS sum(sum_simplify_t.column1 + Int64(1))
-02)--Projection: sum(sum_simplify_t.column1) + count(sum_simplify_t.column1) AS __common_expr_1
+02)--Projection: sum(sum_simplify_t.column1) + CAST(count(sum_simplify_t.column1) AS Int64) AS __common_expr_1
 03)----Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
 04)------TableScan: sum_simplify_t projection=[column1]
 physical_plan
@@ -314,15 +321,16 @@ EXPLAIN SELECT arrow_typeof(SUM(val + 1)), SUM(val + 1), SUM(val + 2) FROM tbl;
 ----
 logical_plan
 01)Projection: arrow_typeof(sum(tbl.val + Int64(1))), sum(tbl.val + Int64(1)), sum(tbl.val + Int64(2))
-02)--Projection: sum(tbl.val) + count(tbl.val) AS sum(tbl.val + Int64(1)), sum(tbl.val) + Int64(2) * count(tbl.val) AS sum(tbl.val + Int64(2))
-03)----Aggregate: groupBy=[[]], aggr=[[sum(__common_expr_1 AS tbl.val), count(__common_expr_1 AS tbl.val)]]
-04)------Projection: CAST(tbl.val AS Int64) AS __common_expr_1
-05)--------TableScan: tbl projection=[val]
+02)--Projection: sum(tbl.val) + __common_expr_1 AS sum(tbl.val + Int64(1)), sum(tbl.val) + Int64(2) * __common_expr_1 AS sum(tbl.val + Int64(2))
+03)----Projection: CAST(count(tbl.val) AS Int64) AS __common_expr_1, sum(tbl.val)
+04)------Aggregate: groupBy=[[]], aggr=[[sum(__common_expr_2 AS tbl.val), count(__common_expr_2 AS tbl.val)]]
+05)--------Projection: CAST(tbl.val AS Int64) AS __common_expr_2
+06)----------TableScan: tbl projection=[val]
 physical_plan
 01)ProjectionExec: expr=[arrow_typeof(sum(tbl.val + Int64(1))@0) as arrow_typeof(sum(tbl.val + Int64(1))), sum(tbl.val + Int64(1))@0 as sum(tbl.val + Int64(1)), sum(tbl.val + Int64(2))@1 as sum(tbl.val + Int64(2))]
 02)--ProjectionExec: expr=[sum(tbl.val)@0 + count(tbl.val)@1 as sum(tbl.val + Int64(1)), sum(tbl.val)@0 + 2 * count(tbl.val)@1 as sum(tbl.val + Int64(2))]
 03)----AggregateExec: mode=Single, gby=[], aggr=[sum(tbl.val), count(tbl.val)]
-04)------ProjectionExec: expr=[CAST(val@0 AS Int64) as __common_expr_1]
+04)------ProjectionExec: expr=[CAST(val@0 AS Int64) as __common_expr_2]
 05)--------DataSourceExec: partitions=1, partition_sizes=[2]
 
 # Checks equivalent rewritten form (SUM + COUNT terms) matches transformed SUM semantics.