From 39d7bdfbdf9db13dbcd10137a3d3c1da208b2a73 Mon Sep 17 00:00:00 2001 From: RishavTiwari25 Date: Mon, 12 Oct 2026 21:13:27 +0530 Subject: [PATCH] fix: do not rewrite `SUM(x + literal)` for floating point arguments With two or more `SUM(x + )` aggregates over the same `x`, the optimizer rewrites them into `SUM(x) + * COUNT(x)`. Floating point addition is not associative, so for `Float64` this changes the result: `sum(x + 1.0)` over `1e16, 1.0, -1e16` returns `2.0` alone but `4.0` next to `sum(x + 2.0)`. The rewrite now skips aggregates whose argument or literal is a floating point type, so they are computed as written. Integer aggregates are rewritten as before. Closes #26223 Co-Authored-By: Claude Fable 5.1 --- .../simplify_expressions/linear_aggregates.rs | 27 +++++++++- .../simplify_expressions/simplify_exprs.rs | 3 +- .../test_files/aggregates_simplify.slt | 53 +++++++++++++++++++ 3 files changed, 80 insertions(+), 3 deletions(-) diff --git a/datafusion/optimizer/src/simplify_expressions/linear_aggregates.rs b/datafusion/optimizer/src/simplify_expressions/linear_aggregates.rs index 4b3b6579584bf..26e7a13c50f7a 100644 --- a/datafusion/optimizer/src/simplify_expressions/linear_aggregates.rs +++ b/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 { // map : 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 { + Ok(self.expr().get_type(schema)?.is_floating() + || self.lit().get_type(schema)?.is_floating()) + } } diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index c98ac48871ba1..d56d1ebaf49e0 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/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) { diff --git a/datafusion/sqllogictest/test_files/aggregates_simplify.slt b/datafusion/sqllogictest/test_files/aggregates_simplify.slt index 4169707913fb2..d50f6f7d4437a 100644 --- a/datafusion/sqllogictest/test_files/aggregates_simplify.slt +++ b/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;