fix: do not rewrite SUM(x + literal) for floating point arguments

datafusion/optimizer/src/simplify_expressions/linear_aggregates.rs

@@ -17,9 +17,9 @@
 
 //! Simplification to refactor multiple aggregate functions to use the same aggregate function
 
-use datafusion_common::HashMap;
+use datafusion_common::{DFSchema, HashMap};
 use datafusion_expr::expr::AggregateFunctionParams;
-use datafusion_expr::{BinaryExpr, Expr};
+use datafusion_expr::{BinaryExpr, Expr, ExprSchemable};
 use datafusion_expr_common::operator::Operator;
 
 /// Threshold of the number of aggregates that share similar arguments before
@@ -44,6 +44,10 @@ const DUPLICATE_THRESHOLD: usize = 2;
 ///
 /// See the background [`AggregateUDFImpl::simplify_expr_op_literal`] for details.
 ///
+/// The rewrite is only applied when it is exact: floating point arithmetic is
+/// not associative, so `SUM(x + 1.0)` and `SUM(x) + 1.0 * COUNT(x)` can return
+/// different values for the same input.
+///
 /// Returns `true` if any of the arguments are rewritten (modified), `false`
 /// otherwise.
 ///
@@ -54,6 +58,7 @@ const DUPLICATE_THRESHOLD: usize = 2;
 /// [`AggregateUDFImpl::simplify_expr_op_literal`]: datafusion_expr::AggregateUDFImpl::simplify_expr_op_literal
 pub(super) fn rewrite_multiple_linear_aggregates(
     agg_expr: &mut [Expr],
+    input_schema: &DFSchema,
 ) -> datafusion_common::Result<bool> {
     // map <expr>: count of expressions that have a common argument
     let mut common_args = HashMap::new();
@@ -73,6 +78,10 @@ pub(super) fn rewrite_multiple_linear_aggregates(
             continue;
         };
 
+        if expr_literal.has_floating_point_operand(input_schema)? {
+            continue;
+        }
+
         let counter = common_args.entry(expr_literal.expr()).or_insert(0);
         *counter += 1;
     }
@@ -95,6 +104,10 @@ pub(super) fn rewrite_multiple_linear_aggregates(
             continue;
         };
 
+        if expr_literal.has_floating_point_operand(input_schema)? {
+            continue;
+        }
+
         // Not enough common expressions to make it worth rewriting
         if common_args.get(expr_literal.expr()).unwrap_or(&0) < &DUPLICATE_THRESHOLD {
             continue;
@@ -226,4 +239,14 @@ impl<'a> ExprLiteral<'a> {
     fn arg_is_left(&self) -> bool {
         matches!(self, Self::ArgOpLit { .. })
     }
+
+    /// Returns true if either operand is a floating point value, in which case
+    /// rewriting the aggregate would not be exact.
+    fn has_floating_point_operand(
+        &self,
+        schema: &DFSchema,
+    ) -> datafusion_common::Result<bool> {
+        Ok(self.expr().get_type(schema)?.is_floating()
+            || self.lit().get_type(schema)?.is_floating())
+    }
 }

datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs

@@ -183,7 +183,8 @@ fn rewrite_aggregate_non_aggregate_aggr_expr(
         return Ok(Transformed::no(plan));
     };
 
-    let rewrote_aggs = rewrite_multiple_linear_aggregates(&mut aggr_expr)?;
+    let rewrote_aggs =
+        rewrite_multiple_linear_aggregates(&mut aggr_expr, input.schema())?;
 
     // Ensure that all Aggregate arguments are AggregateExpr
     if aggr_expr.iter().all(is_top_level_aggregate_expr) {

datafusion/sqllogictest/test_files/aggregates_simplify.slt

@@ -134,6 +134,59 @@ physical_plan
 02)--AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1), count(sum_simplify_t.column1)]
 03)----DataSourceExec: partitions=1, partition_sizes=[1]
 
+# Floating point arguments are not rewritten, because the rewrite is not exact
+# for floating point arithmetic. With the values below the large values cancel
+# out and `SUM(x + 1.0)` is 4, while `SUM(x) + 1.0 * COUNT(x)` would be 6.
+statement ok
+CREATE TABLE sum_simplify_float_t (x DOUBLE) AS VALUES (1e17), (3.0), (-1e17);
+
+query R
+SELECT SUM(x + 1.0) FROM sum_simplify_float_t;
+----
+4
+
+query RR
+SELECT SUM(x + 1.0), SUM(x + 2.0) FROM sum_simplify_float_t;
+----
+4 5
+
+query TT
+EXPLAIN SELECT SUM(x + 1.0), SUM(x + 2.0) FROM sum_simplify_float_t;
+----
+logical_plan
+01)Aggregate: groupBy=[[]], aggr=[[sum(sum_simplify_float_t.x + Float64(1)), sum(sum_simplify_float_t.x + Float64(2))]]
+02)--TableScan: sum_simplify_float_t projection=[x]
+physical_plan
+01)AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_float_t.x + Float64(1)), sum(sum_simplify_float_t.x + Float64(2))]
+02)--DataSourceExec: partitions=1, partition_sizes=[1]
+
+# The same applies when only the column is floating point
+query RR
+SELECT SUM(x + 1), SUM(x + 2) FROM sum_simplify_float_t;
+----
+4 5
+
+# ... or when only the literal is floating point
+query RR
+SELECT SUM(column1 + 1.0), SUM(column1 + 2.0) FROM sum_simplify_t;
+----
+7 10
+
+query TT
+EXPLAIN SELECT SUM(column1 + 1.0), SUM(column1 + 2.0) FROM sum_simplify_t;
+----
+logical_plan
+01)Aggregate: groupBy=[[]], aggr=[[sum(__common_expr_1 AS sum_simplify_t.column1 + Float64(1)), sum(__common_expr_1 AS sum_simplify_t.column1 + Float64(2))]]
+02)--Projection: CAST(sum_simplify_t.column1 AS Float64) AS __common_expr_1
+03)----TableScan: sum_simplify_t projection=[column1]
+physical_plan
+01)AggregateExec: mode=Single, gby=[], aggr=[sum(sum_simplify_t.column1 + Float64(1)), sum(sum_simplify_t.column1 + Float64(2))]
+02)--ProjectionExec: expr=[CAST(column1@0 AS Float64) as __common_expr_1]
+03)----DataSourceExec: partitions=1, partition_sizes=[1]
+
+statement ok
+DROP TABLE sum_simplify_float_t;
+
 # DISTINCT aggregates with different arguments
 query II
 SELECT SUM(DISTINCT column1 + 1), SUM(DISTINCT column1 + 2) FROM sum_simplify_t;