Add rewrite SUM(expr+C) --> SUM(expr) + COUNT(expr)*C

Author: alamb

datafusion/expr/src/expr.rs

@@ -598,7 +598,7 @@ impl Alias {
     }
 }
 
-/// Binary expression
+/// Binary expression for [`Expr::BinaryExpr`]
 #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)]
 pub struct BinaryExpr {
     /// Left-hand side of the expression

datafusion/functions-aggregate/src/sum.rs

@@ -32,12 +32,16 @@ use datafusion_common::types::{
     logical_int64, logical_uint8, logical_uint16, logical_uint32, logical_uint64,
 };
 use datafusion_common::{HashMap, Result, ScalarValue, exec_err, not_impl_err};
-use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs};
+use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams};
+use datafusion_expr::function::{
+    AccumulatorArgs, AggregateFunctionSimplification, StateFieldsArgs,
+};
+use datafusion_expr::simplify::SimplifyContext;
 use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name};
 use datafusion_expr::{
-    Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, GroupsAccumulator,
-    ReversedUDAF, SetMonotonicity, Signature, TypeSignature, TypeSignatureClass,
-    Volatility,
+    Accumulator, AggregateUDFImpl, BinaryExpr, Coercion, Documentation, Expr,
+    GroupsAccumulator, Operator, ReversedUDAF, SetMonotonicity, Signature, TypeSignature,
+    TypeSignatureClass, Volatility,
 };
 use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator;
 use datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator;
@@ -54,7 +58,7 @@ make_udaf_expr_and_func!(
 );
 
 pub fn sum_distinct(expr: Expr) -> Expr {
-    Expr::AggregateFunction(datafusion_expr::expr::AggregateFunction::new_udf(
+    Expr::AggregateFunction(AggregateFunction::new_udf(
         sum_udaf(),
         vec![expr],
         true,
@@ -346,6 +350,88 @@ impl AggregateUDFImpl for Sum {
             _ => SetMonotonicity::NotMonotonic,
         }
     }
+
+    /// Simplification Rules
+    fn simplify(&self) -> Option<AggregateFunctionSimplification> {
+        Some(Box::new(sum_simplifier))
+    }
+}
+
+/// Implement ClickBench Q29 specific optimization:
+/// `SUM(arg + constant)` --> `SUM(arg) + constant * COUNT(arg)`
+///
+/// Backstory: TODO
+///
+fn sum_simplifier(mut agg: AggregateFunction, _info: &SimplifyContext) -> Result<Expr> {
+    // Explicitly destructure to ensure we check all relevant fields
+    let AggregateFunctionParams {
+        args,
+        distinct,
+        filter,
+        order_by,
+        null_treatment,
+    } = &agg.params;
+
+    if *distinct
+        || filter.is_some()
+        || !order_by.is_empty()
+        || null_treatment.is_some()
+        || args.len() != 1
+    {
+        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));
+        }
+        SplitResult::Split { arg, lit } => (arg, lit),
+    };
+
+    // Rewrite to SUM(arg)
+    agg.params.args.push(arg.clone());
+    let sum_agg = Expr::AggregateFunction(agg);
+
+    // sum(arg) + scalar * COUNT(arg)
+    Ok(sum_agg + (lit * crate::count::count(arg)))
+}
+
+/// Result of trying to split an expression into an arg and constant
+#[derive(Debug, Clone)]
+enum SplitResult {
+    /// if the expression is either of
+    /// * `<arg> <op> <lit>`
+    /// * `<lit> <op> <arg>`
+    ///
+    /// When `op` is `+`
+    Split { arg: Expr, lit: Expr },
+    /// If the expression is something else
+    Original(Expr),
+}
+
+impl SplitResult {
+    fn new(expr: Expr) -> Self {
+        let Expr::BinaryExpr(BinaryExpr { left, op, right }) = expr else {
+            return Self::Original(expr);
+        };
+        if op != Operator::Plus {
+            return Self::Original(Expr::BinaryExpr(BinaryExpr { left, op, right }));
+        }
+
+        match (left.as_ref(), right.as_ref()) {
+            (Expr::Literal(..), _) => Self::Split {
+                arg: *right,
+                lit: *left,
+            },
+            (_, Expr::Literal(..)) => Self::Split {
+                arg: *left,
+                lit: *right,
+            },
+            _ => Self::Original(Expr::BinaryExpr(BinaryExpr { left, op, right })),
+        }
+    }
 }
 
 /// This accumulator computes SUM incrementally