From 675cc4659fbe0fd6b77de1d5d04ca7427857aff6 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Sun, 6 Sep 2026 09:15:35 -0400 Subject: [PATCH 1/2] perf(optimizer): avoid redundant aggregate planning work Reduce work in EliminateDuplicatedExpr with two independent optimizations: 1. Before deduplication, return immediately for zero or one grouping expression. These lists cannot contain duplicate entries, so avoid constructing the deduplication set as well. 2. After deduplication, reuse the Aggregate if no grouping expressions were removed. This avoids repeating validation, schema construction, and functional-dependency derivation. This improves query optimizer performance on the TPC-H and TPC-DS queries by about 5%. --- .../src/eliminate_duplicated_expr.rs | 93 +++++++++++++++++-- 1 file changed, 85 insertions(+), 8 deletions(-) diff --git a/datafusion/optimizer/src/eliminate_duplicated_expr.rs b/datafusion/optimizer/src/eliminate_duplicated_expr.rs index 97aa6e1d8480d..8f1e024328e41 100644 --- a/datafusion/optimizer/src/eliminate_duplicated_expr.rs +++ b/datafusion/optimizer/src/eliminate_duplicated_expr.rs @@ -112,8 +112,12 @@ impl OptimizerRule for EliminateDuplicatedExpr { fetch: sort.fetch, }))) } - LogicalPlan::Aggregate(agg) => { + LogicalPlan::Aggregate(mut agg) => { let len = agg.group_expr.len(); + // With zero or one group expression, there is nothing to deduplicate. + if len <= 1 { + return Ok(Transformed::no(LogicalPlan::Aggregate(agg))); + } let unique_exprs: Vec = agg .group_expr @@ -122,14 +126,15 @@ impl OptimizerRule for EliminateDuplicatedExpr { .into_iter() .collect(); - let transformed = if len != unique_exprs.len() { - Transformed::yes - } else { - Transformed::no - }; + if len == unique_exprs.len() { + // No duplicate expressions were found, so we can reuse the + // existing schema (including functional dependencies). + agg.group_expr = unique_exprs; + return Ok(Transformed::no(LogicalPlan::Aggregate(agg))); + } Aggregate::try_new(agg.input, unique_exprs, agg.aggr_expr) - .map(|f| transformed(LogicalPlan::Aggregate(f))) + .map(|f| Transformed::yes(LogicalPlan::Aggregate(f))) } _ => Ok(Transformed::no(plan)), } @@ -145,7 +150,10 @@ mod tests { use crate::OptimizerContext; use crate::assert_optimized_plan_eq_snapshot; use crate::test::*; - use datafusion_expr::{col, logical_plan::builder::LogicalPlanBuilder}; + use datafusion_expr::{ + GroupingSet, col, lit, logical_plan::builder::LogicalPlanBuilder, + }; + use datafusion_functions_aggregate::expr_fn::sum; use std::sync::Arc; macro_rules! assert_optimized_plan_equal { @@ -200,4 +208,73 @@ mod tests { TableScan: test ") } + + /// A no-op must preserve both the plan and its already computed schema. + fn assert_unchanged_aggregate(plan: LogicalPlan) -> Result<()> { + let schema = Arc::clone(plan.schema()); + let result = EliminateDuplicatedExpr::new() + .rewrite(plan.clone(), &OptimizerContext::new())?; + assert!(!result.transformed); + assert_eq!(result.data, plan); + assert!(Arc::ptr_eq(result.data.schema(), &schema)); + Ok(()) + } + + #[test] + fn unchanged_aggregates_reuse_schema() -> Result<()> { + for group_expr in [vec![], vec![col("a")], vec![col("b"), col("a") + lit(1u32)]] { + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .aggregate(group_expr, vec![sum(col("c")).alias("total")])? + .build()?; + assert_unchanged_aggregate(plan)?; + } + Ok(()) + } + + #[test] + fn unchanged_grouping_sets_reuse_schema() -> Result<()> { + // Repeated grouping sets must retain their multiplicity. + let grouping_set = GroupingSet::GroupingSets(vec![ + vec![col("a"), col("b")], + vec![col("a")], + vec![col("a")], + vec![], + ]); + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .aggregate(vec![Expr::GroupingSet(grouping_set)], vec![sum(col("c"))])? + .build()?; + assert_unchanged_aggregate(plan) + } + + #[test] + fn eliminate_duplicate_group_exprs() -> Result<()> { + let input = Arc::new(test_table_scan()?); + let a = col("test.a"); + let b = col("test.b"); + let computed = a.clone() + lit(1u32); + let aggr_expr = vec![sum(col("test.c")).alias("total")]; + let plan = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::clone(&input), + vec![ + b.clone(), + computed.clone(), + b.clone(), + a.clone(), + computed.clone(), + ], + aggr_expr.clone(), + )?); + let expected = LogicalPlan::Aggregate(Aggregate::try_new( + input, + // Keep the order of the first occurrence of each expression. + vec![b, computed, a], + aggr_expr, + )?); + + let result = + EliminateDuplicatedExpr::new().rewrite(plan, &OptimizerContext::new())?; + assert!(result.transformed); + assert_eq!(result.data, expected); + Ok(()) + } } From 8c303c47e90e247f53a5c476f31cc762829cfc5a Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Sun, 6 Sep 2026 10:04:37 -0400 Subject: [PATCH 2/2] Tweak comment --- datafusion/optimizer/src/eliminate_duplicated_expr.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/optimizer/src/eliminate_duplicated_expr.rs b/datafusion/optimizer/src/eliminate_duplicated_expr.rs index 8f1e024328e41..f2f489ab825c2 100644 --- a/datafusion/optimizer/src/eliminate_duplicated_expr.rs +++ b/datafusion/optimizer/src/eliminate_duplicated_expr.rs @@ -233,7 +233,7 @@ mod tests { #[test] fn unchanged_grouping_sets_reuse_schema() -> Result<()> { - // Repeated grouping sets must retain their multiplicity. + // Repeat "a" to check that this rule leaves the GroupingSet's contents unchanged. let grouping_set = GroupingSet::GroupingSets(vec![ vec![col("a"), col("b")], vec![col("a")],