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

Author: alamb

datafusion/expr/src/expr.rs

@@ -600,7 +600,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/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

@@ -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,136 @@ 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 <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 = 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 {
+    /// 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,
+}
+
+impl SplitResult {
+    fn new(expr: Expr) -> Self {
+        let Expr::BinaryExpr(BinaryExpr { left, op, right }) = expr else {
+            return Self::Original;
+        };
+        if op != Operator::Plus {
+            return Self::Original;
+        }
+
+        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,
+        }
+    }
 }
 
 /// This accumulator computes SUM incrementally