Avoid the sum/count rewrite when not helpful

Author: alamb

datafusion/expr/src/simplify.rs

@@ -38,6 +38,7 @@ pub struct SimplifyContext {
     schema: DFSchemaRef,
     query_execution_start_time: Option<DateTime<Utc>>,
     config_options: Arc<ConfigOptions>,
+    aggregate_exprs: Option<Arc<Vec<Expr>>>,
 }
 
 impl Default for SimplifyContext {
@@ -46,6 +47,7 @@ impl Default for SimplifyContext {
             schema: Arc::new(DFSchema::empty()),
             query_execution_start_time: None,
             config_options: Arc::new(ConfigOptions::default()),
+            aggregate_exprs: None,
         }
     }
 }
@@ -78,6 +80,12 @@ impl SimplifyContext {
         self
     }
 
+    /// Set aggregate expressions from the containing aggregate node, if any.
+    pub fn with_aggregate_exprs(mut self, aggregate_exprs: Arc<Vec<Expr>>) -> Self {
+        self.aggregate_exprs = Some(aggregate_exprs);
+        self
+    }
+
     /// Returns the schema
     pub fn schema(&self) -> &DFSchemaRef {
         &self.schema
@@ -108,6 +116,11 @@ impl SimplifyContext {
     pub fn config_options(&self) -> &Arc<ConfigOptions> {
         &self.config_options
     }
+
+    /// Returns aggregate expressions from the containing aggregate node, if any.
+    pub fn aggregate_exprs(&self) -> Option<&[Expr]> {
+        self.aggregate_exprs.as_deref().map(Vec::as_slice)
+    }
 }
 
 /// Was the expression simplified?

datafusion/functions-aggregate/src/sum.rs

@@ -362,7 +362,7 @@ impl AggregateUDFImpl for Sum {
 ///
 /// Backstory: TODO
 ///
-fn sum_simplifier(mut agg: AggregateFunction, _info: &SimplifyContext) -> Result<Expr> {
+fn sum_simplifier(mut agg: AggregateFunction, info: &SimplifyContext) -> Result<Expr> {
     // Explicitly destructure to ensure we check all relevant fields
     let AggregateFunctionParams {
         args,
@@ -381,23 +381,71 @@ fn sum_simplifier(mut agg: AggregateFunction, _info: &SimplifyContext) -> Result
         return Ok(Expr::AggregateFunction(agg));
     }
 
-    // otherwise check the arguments if they are <col> <op> scalar
-    let (arg, lit) = match SplitResult::new(agg.params.args.swap_remove(0)) {
-        SplitResult::Original(expr) => {
-            agg.params.args.push(expr); // put it back
-            return Ok(Expr::AggregateFunction(agg));
-        }
+    // otherwise check the arguments if they are <arg> + <literal>
+    let (arg, lit) = match SplitResult::new(agg.params.args[0].clone()) {
+        SplitResult::Original => return Ok(Expr::AggregateFunction(agg)),
         SplitResult::Split { arg, lit } => (arg, lit),
     };
 
+    if !has_common_rewrite_arg(&arg, info) {
+        return Ok(Expr::AggregateFunction(agg));
+    }
+
     // Rewrite to SUM(arg)
-    agg.params.args.push(arg.clone());
+    agg.params.args = vec![arg.clone()];
     let sum_agg = Expr::AggregateFunction(agg);
 
     // sum(arg) + scalar * COUNT(arg)
     Ok(sum_agg + (lit * crate::count::count(arg)))
 }
 
+fn has_common_rewrite_arg(arg: &Expr, info: &SimplifyContext) -> bool {
+    let Some(aggregate_exprs) = info.aggregate_exprs() else {
+        // Only apply this rewrite in the context of an Aggregate node where
+        // sibling aggregate expressions are known.
+        return false;
+    };
+
+    aggregate_exprs
+        .iter()
+        .filter_map(sum_rewrite_candidate_arg)
+        .filter(|candidate_arg| candidate_arg == arg)
+        .take(2)
+        .count()
+        > 1
+}
+
+fn sum_rewrite_candidate_arg(expr: &Expr) -> Option<Expr> {
+    let Expr::AggregateFunction(aggregate_fn) = expr.clone().unalias_nested().data else {
+        return None;
+    };
+    if !aggregate_fn.func.name().eq_ignore_ascii_case("sum") {
+        return None;
+    }
+
+    let AggregateFunctionParams {
+        args,
+        distinct,
+        filter,
+        order_by,
+        null_treatment,
+    } = &aggregate_fn.params;
+
+    if *distinct
+        || filter.is_some()
+        || !order_by.is_empty()
+        || null_treatment.is_some()
+        || args.len() != 1
+    {
+        return None;
+    }
+
+    match SplitResult::new(args[0].clone()) {
+        SplitResult::Split { arg, .. } => Some(arg),
+        SplitResult::Original => None,
+    }
+}
+
 /// Result of trying to split an expression into an arg and constant
 #[derive(Debug, Clone)]
 enum SplitResult {
@@ -408,16 +456,16 @@ enum SplitResult {
     /// When `op` is `+`
     Split { arg: Expr, lit: Expr },
     /// If the expression is something else
-    Original(Expr),
+    Original,
 }
 
 impl SplitResult {
     fn new(expr: Expr) -> Self {
         let Expr::BinaryExpr(BinaryExpr { left, op, right }) = expr else {
-            return Self::Original(expr);
+            return Self::Original;
         };
         if op != Operator::Plus {
-            return Self::Original(Expr::BinaryExpr(BinaryExpr { left, op, right }));
+            return Self::Original;
         }
 
         match (left.as_ref(), right.as_ref()) {
@@ -429,7 +477,7 @@ impl SplitResult {
                 arg: *left,
                 lit: *right,
             },
-            _ => Self::Original(Expr::BinaryExpr(BinaryExpr { left, op, right })),
+            _ => Self::Original,
         }
     }
 }

datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs

@@ -103,6 +103,12 @@ impl SimplifyExpressions {
             .with_config_options(config.options())
             .with_query_execution_start_time(config.query_execution_start_time());
 
+        let info = if let LogicalPlan::Aggregate(Aggregate { aggr_expr, .. }) = &plan {
+            info.with_aggregate_exprs(Arc::new(aggr_expr.clone()))
+        } else {
+            info
+        };
+
         // Inputs have already been rewritten (due to bottom-up traversal handled by Optimizer)
         // Just need to rewrite our own expressions
 
@@ -141,9 +147,7 @@ impl SimplifyExpressions {
             }
         })?;
 
-        transformed.transform_data(|plan| {
-            rewrite_aggregate_non_aggregate_aggr_expr(plan)
-        })
+        transformed.transform_data(rewrite_aggregate_non_aggregate_aggr_expr)
     }
 }
 
@@ -345,7 +349,31 @@ mod tests {
         assert_optimized_plan_equal!(
             plan,
             @r"
-        Projection: sum(test.a) + Int64(2) * count(test.a) AS sum(test.a + Int64(2))
+        Aggregate: groupBy=[[]], aggr=[[sum(test.a + Int64(2))]]
+          TableScan: test
+        "
+        )?;
+        Ok(())
+    }
+
+    #[test]
+    fn test_simplify_udaf_to_non_aggregate_expr_with_common_sum_arg() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
+        let table_scan = table_scan(Some("test"), &schema, None)?
+            .build()
+            .expect("building scan");
+
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                Vec::<Expr>::new(),
+                vec![sum(col("a") + lit(2i64)), sum(col("a") + lit(3i64))],
+            )?
+            .build()?;
+
+        assert_optimized_plan_equal!(
+            plan,
+            @r"
+        Projection: sum(test.a) + Int64(2) * count(test.a) AS sum(test.a + Int64(2)), sum(test.a) + Int64(3) * count(test.a) AS sum(test.a + Int64(3))
           Aggregate: groupBy=[[]], aggr=[[sum(test.a), count(test.a)]]
             TableScan: test
         "

datafusion/sqllogictest/test_files/aggregates_simplify.slt

@@ -32,13 +32,11 @@ query TT
 EXPLAIN SELECT SUM(column1 + 1) 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))
-02)--Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
-03)----TableScan: sum_simplify_t projection=[column1]
+01)Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1 + Int64(1))]]
+02)--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))]
-02)--AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]
-03)----DataSourceExec: partitions=1, partition_sizes=[1]
+01)AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1 + Int64(1))]
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
 
 
 # Mixed aggregate expressions with type validation
@@ -51,12 +49,12 @@ query TT
 EXPLAIN SELECT arrow_typeof(SUM(column1)), SUM(column1), SUM(column1 + 1) FROM sum_simplify_t;
 ----
 logical_plan
-01)Projection: arrow_typeof(sum(sum_simplify_t.column1)), sum(sum_simplify_t.column1), sum(sum_simplify_t.column1) + count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(1))
-02)--Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
+01)Projection: arrow_typeof(sum(sum_simplify_t.column1)), sum(sum_simplify_t.column1), sum(sum_simplify_t.column1 + Int64(1))
+02)--Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), sum(sum_simplify_t.column1 + Int64(1))]]
 03)----TableScan: sum_simplify_t projection=[column1]
 physical_plan
-01)ProjectionExec: expr=[arrow_typeof(sum(sum_simplify_t.column1)@0) as arrow_typeof(sum(sum_simplify_t.column1)), sum(sum_simplify_t.column1)@0 as sum(sum_simplify_t.column1), sum(sum_simplify_t.column1)@0 + count(sum_simplify_t.column1)@1 as sum(sum_simplify_t.column1 + Int64(1))]
-02)--AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]
+01)ProjectionExec: expr=[arrow_typeof(sum(sum_simplify_t.column1)@0) as arrow_typeof(sum(sum_simplify_t.column1)), sum(sum_simplify_t.column1)@0 as sum(sum_simplify_t.column1), sum(sum_simplify_t.column1 + Int64(1))@1 as sum(sum_simplify_t.column1 + Int64(1))]
+02)--AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1), sum(sum_simplify_t.column1 + Int64(1))]
 03)----DataSourceExec: partitions=1, partition_sizes=[1]
 
 # Duplicate aggregate expressions
@@ -70,14 +68,12 @@ EXPLAIN SELECT SUM(column1 + 1) AS sum_plus_1_a, SUM(column1 + 1) AS sum_plus_1_
 ----
 logical_plan
 01)Projection: sum(sum_simplify_t.column1 + Int64(1)) AS sum_plus_1_a, sum(sum_simplify_t.column1 + Int64(1)) AS sum_plus_1_b
-02)--Projection: sum(sum_simplify_t.column1) + count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(1))
-03)----Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
-04)------TableScan: sum_simplify_t projection=[column1]
+02)--Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_t.column1 + Int64(1))]]
+03)----TableScan: sum_simplify_t projection=[column1]
 physical_plan
 01)ProjectionExec: expr=[sum(sum_simplify_t.column1 + Int64(1))@0 as sum_plus_1_a, sum(sum_simplify_t.column1 + Int64(1))@0 as sum_plus_1_b]
-02)--ProjectionExec: expr=[sum(sum_simplify_t.column1)@0 + count(sum_simplify_t.column1)@1 as sum(sum_simplify_t.column1 + Int64(1))]
-03)----AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]
-04)------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1 + Int64(1))]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
 
 
 # constant aggregate expressions
@@ -162,13 +158,21 @@ query TT
 EXPLAIN SELECT SUM(DISTINCT column1 + 1), SUM(column1 + 1) FROM sum_simplify_t;
 ----
 logical_plan
-01)Projection: sum(DISTINCT sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1) + count(sum_simplify_t.column1) AS sum(sum_simplify_t.column1 + Int64(1))
-02)--Aggregate: groupBy=[[]], aggr=[[sum(DISTINCT sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]]
-03)----TableScan: sum_simplify_t projection=[column1]
+01)Projection: sum(alias1) AS sum(DISTINCT sum_simplify_t.column1 + Int64(1)), sum(alias2) AS sum(sum_simplify_t.column1 + Int64(1))
+02)--Aggregate: groupBy=[[]], aggr=[[sum(alias1), sum(alias2)]]
+03)----Aggregate: groupBy=[[__common_expr_1 AS alias1]], aggr=[[sum(__common_expr_1) AS alias2]]
+04)------Projection: sum_simplify_t.column1 + Int64(1) AS __common_expr_1
+05)--------TableScan: sum_simplify_t projection=[column1]
 physical_plan
-01)ProjectionExec: expr=[sum(DISTINCT sum_simplify_t.column1 + Int64(1))@0 as sum(DISTINCT sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1)@1 + count(sum_simplify_t.column1)@2 as sum(sum_simplify_t.column1 + Int64(1))]
-02)--AggregateExec: mode=Single, gby=[], aggr=[sum(DISTINCT sum_simplify_t.column1 + Int64(1)), sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]
-03)----DataSourceExec: partitions=1, partition_sizes=[1]
+01)ProjectionExec: expr=[sum(alias1)@0 as sum(DISTINCT sum_simplify_t.column1 + Int64(1)), sum(alias2)@1 as sum(sum_simplify_t.column1 + Int64(1))]
+02)--AggregateExec: mode=Final, gby=[], aggr=[sum(alias1), sum(alias2)]
+03)----CoalescePartitionsExec
+04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(alias1), sum(alias2)]
+05)--------AggregateExec: mode=FinalPartitioned, gby=[alias1@0 as alias1], aggr=[alias2]
+06)----------RepartitionExec: partitioning=Hash([alias1@0], 4), input_partitions=1
+07)------------AggregateExec: mode=Partial, gby=[__common_expr_1@0 as alias1], aggr=[alias2]
+08)--------------ProjectionExec: expr=[column1@0 + 1 as __common_expr_1]
+09)----------------DataSourceExec: partitions=1, partition_sizes=[1]
 
 # FILTER clauses with different aggregate arguments
 query II