From 456a39a24e2bb0a99b21d81971a8556c6e52fc19 Mon Sep 17 00:00:00 2001 From: Fangchen Li Date: Tue, 3 Nov 2026 22:27:20 -0800 Subject: [PATCH 1/2] fix: apply SingleDistinctToGroupBy to aggregates that carry an alias The rule only matched a bare Expr::AggregateFunction, so an alias directly on the aggregate, as the DataFrame API produces, turned the rewrite off and gave the same query a different plan than SQL. Look through aliases when matching and strip them in the rewrite; the outer projection restores output names from the schema. Also correct the can_strip_every_distinct doc comment, which said this rule does not rewrite aliased aggregates. Its behavior is unchanged. Closes #23401 Co-Authored-By: Claude Opus 5.5 --- datafusion/core/tests/dataframe/mod.rs | 34 +++ .../src/eliminate_aggregate_distinct.rs | 17 +- .../src/single_distinct_to_groupby.rs | 223 +++++++++++++++++- 3 files changed, 264 insertions(+), 10 deletions(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 49635ed110e7f..f0d36b8a22f8a 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -838,6 +838,40 @@ async fn aggregate_assert_no_empty_batches() -> Result<()> { Ok(()) } +// The DataFrame API puts an alias on the aggregate itself, where SQL leaves the +// aggregate bare and aliases it in a projection. Both have to reach the same plan. +#[tokio::test] +async fn aliased_single_distinct_aggregate_is_rewritten() -> Result<()> { + let ctx = SessionContext::new(); + let df = ctx + .read_batch(record_batch!( + ("a", Int32, [1, 1, 2, 2]), + ("b", Int32, [1, 1, 1, 2]) + )?)? + .aggregate(vec![col("a")], vec![count_distinct(col("b")).alias("n")])?; + + assert_snapshot!(df.clone().into_optimized_plan()?, @r" + Projection: ?table?.a, count(alias1) AS n + Aggregate: groupBy=[[?table?.a]], aggr=[[count(alias1)]] + Aggregate: groupBy=[[?table?.a, ?table?.b AS alias1]], aggr=[[]] + TableScan: ?table? projection=[a, b] + "); + + assert_batches_eq!( + &[ + "+---+---+", + "| a | n |", + "+---+---+", + "| 1 | 1 |", + "| 2 | 2 |", + "+---+---+", + ], + &df.sort_by(vec![col("a")])?.collect().await? + ); + + Ok(()) +} + #[tokio::test] async fn test_aggregate_with_pk() -> Result<()> { // create the dataframe diff --git a/datafusion/optimizer/src/eliminate_aggregate_distinct.rs b/datafusion/optimizer/src/eliminate_aggregate_distinct.rs index d2b68853a2327..9a88830d50685 100644 --- a/datafusion/optimizer/src/eliminate_aggregate_distinct.rs +++ b/datafusion/optimizer/src/eliminate_aggregate_distinct.rs @@ -87,14 +87,15 @@ impl OptimizerRule for EliminateAggregateDistinct { /// `Insensitive`. /// /// If only some of them were stripped, a `Sensitive` `DISTINCT` could stay -/// beside a stripped aggregate. A stripped aggregate carries an alias, and -/// [`crate::single_distinct_to_groupby::SingleDistinctToGroupBy`] does not -/// rewrite a node whose `aggr_expr` contains an alias. So `min(DISTINCT x), -/// count(DISTINCT x)` would lose the rewrite that `count` needs. When every -/// `DISTINCT` is `Insensitive`, none is left after stripping, so that rule has -/// nothing to rewrite. This is conservative: `min(DISTINCT x), -/// count(DISTINCT y)` keeps the `min` flag. That flag costs nothing at run -/// time, because `min` ignores `is_distinct` when it selects its accumulator. +/// beside a stripped aggregate, which +/// [`crate::single_distinct_to_groupby::SingleDistinctToGroupBy`] then has to +/// compute in its inner group by. While both keep the flag, `min(DISTINCT x), +/// count(DISTINCT x)` reads the distinct column directly and needs no inner +/// aggregate. When every `DISTINCT` is `Insensitive`, none is left after +/// stripping, so that rule has nothing to rewrite. This is conservative: +/// `min(DISTINCT x), count(DISTINCT y)` keeps the `min` flag. That flag costs +/// nothing at run time, because `min` ignores `is_distinct` when it selects its +/// accumulator. fn can_strip_every_distinct(aggr_expr: &[Expr]) -> Result { let mut found_distinct = false; let mut all_insensitive = true; diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 8b176e904e9b9..e246c1e2e7ce9 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -15,6 +15,7 @@ use datafusion_common::{ }; use datafusion_expr::builder::project; use datafusion_expr::expr::AggregateFunctionParams; +use datafusion_expr::expr_rewriter::unalias; use datafusion_expr::expr_schema::ExprSchemable; use datafusion_expr::{ AggregateUDF, Expr, cast, col, @@ -118,6 +119,13 @@ fn is_single_distinct_agg( let mut distinct_aggs = vec![]; let mut has_count_rollup = false; for expr in aggr_expr { + // Peel aliases on the aggregate itself. The DataFrame API puts them there + // (e.g. `count(col("b")).distinct().alias("n")`), and so do rewrites that + // have to preserve an aggregate's output name. + let mut expr = expr; + while let Expr::Alias(alias) = expr { + expr = alias.expr.as_ref(); + } if let Expr::AggregateFunction(AggregateFunction { func, params: @@ -284,7 +292,9 @@ impl OptimizerRule for SingleDistinctToGroupBy { // zero that `sum` reports as NULL over an empty input. let (outer_aggr_exprs, outer_proj_exprs): (Vec, Vec) = aggr_expr .into_iter() - .map(|aggr_expr| match aggr_expr { + // Strip aliases (see `is_single_distinct_agg`); the outer + // projection below restores output names from the schema. + .map(|aggr_expr| match unalias(aggr_expr) { Expr::AggregateFunction(AggregateFunction { func, params: @@ -371,7 +381,7 @@ impl OptimizerRule for SingleDistinctToGroupBy { Ok((outer, proj)) } } - _ => Ok((aggr_expr.clone(), aggr_expr)), + other => Ok((other.clone(), other)), }) .collect::>>()? .into_iter() @@ -738,6 +748,191 @@ mod tests { ) } + #[test] + fn aliased_single_distinct_and_groupby() -> Result<()> { + let table_scan = test_table_scan()?; + + // DataFrame-API shape: alias applied directly to the aggregate expr. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_distinct(col("b")).alias("n")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, count(alias1) AS n [a:UInt32, n:Int64] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1)]] [a:UInt32, count(alias1):Int64] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn nested_aliased_single_distinct_and_groupby() -> Result<()> { + let table_scan = test_table_scan()?; + + // Chained aliases nest `Expr::Alias`; all levels must be peeled. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![count_distinct(col("b")).alias("x").alias("n")], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, count(alias1) AS n [a:UInt32, n:Int64] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1)]] [a:UInt32, count(alias1):Int64] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn aliased_single_distinct_no_groupby() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![count_distinct(col("b")).alias("n")], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: count(alias1) AS n [n:Int64] + Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64] + Aggregate: groupBy=[[test.b AS alias1]], aggr=[[]] [alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn aliased_single_distinct_expr() -> Result<()> { + let table_scan = test_table_scan()?; + + // Alias on top of a distinct aggregate over a complex argument. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![count_distinct(lit(2) * col("b")).alias("n")], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: count(alias1) AS n [n:Int64] + Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64] + Aggregate: groupBy=[[Int32(2) * test.b AS alias1]], aggr=[[]] [alias1:Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn aliased_group_by_with_expr() -> Result<()> { + let table_scan = test_table_scan()?; + + // Aliased distinct aggregate combined with a complex group-by expression. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a") + lit(1)], + vec![count_distinct(col("c")).alias("n")], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: group_alias_0 AS test.a + Int32(1), count(alias1) AS n [test.a + Int32(1):Int64, n:Int64] + Aggregate: groupBy=[[group_alias_0]], aggr=[[count(alias1)]] [group_alias_0:Int64, count(alias1):Int64] + Aggregate: groupBy=[[test.a + Int32(1) AS group_alias_0, test.c AS alias1]], aggr=[[]] [group_alias_0:Int64, alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn aliased_distinct_and_common() -> Result<()> { + let table_scan = test_table_scan()?; + + // Mix of an aliased distinct aggregate and an aliased common aggregate. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("c")], + vec![ + min(col("a")).alias("mn"), + count_distinct(col("b")).alias("n"), + ], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.c, min(alias2) AS mn, count(alias1) AS n [c:UInt32, mn:UInt32;N, n:Int64] + Aggregate: groupBy=[[test.c]], aggr=[[min(alias2), count(alias1)]] [c:UInt32, min(alias2):UInt32;N, count(alias1):Int64] + Aggregate: groupBy=[[test.c, test.b AS alias1]], aggr=[[min(test.a) AS alias2]] [c:UInt32, alias1:UInt32, alias2:UInt32;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn partially_aliased_distinct_and_common() -> Result<()> { + let table_scan = test_table_scan()?; + + // Only some aggregate exprs carry an alias; the rule must still fire. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("c")], + vec![min(col("a")), count_distinct(col("b")).alias("n")], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.c, min(alias2) AS min(test.a), count(alias1) AS n [c:UInt32, min(test.a):UInt32;N, n:Int64] + Aggregate: groupBy=[[test.c]], aggr=[[min(alias2), count(alias1)]] [c:UInt32, min(alias2):UInt32;N, count(alias1):Int64] + Aggregate: groupBy=[[test.c, test.b AS alias1]], aggr=[[min(test.a) AS alias2]] [c:UInt32, alias1:UInt32, alias2:UInt32;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn aliased_two_distinct_fields_and_groupby() -> Result<()> { + let table_scan = test_table_scan()?; + + // Two distinct aggregates on different fields: not eligible even when + // aliased. The plan must be left unchanged (aliases retained). + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + count_distinct(col("b")).alias("nb"), + count_distinct(col("c")).alias("nc"), + ], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b) AS nb, count(DISTINCT test.c) AS nc]] [a:UInt32, nb:Int64, nc:Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + #[test] fn single_distinct_and_groupby() -> Result<()> { let table_scan = test_table_scan()?; @@ -1206,6 +1401,30 @@ mod tests { ) } + #[test] + fn aliased_count_star_and_distinct_without_groupby() -> Result<()> { + let table_scan = test_table_scan_utf8_b()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![count_star().alias("n"), count_distinct(col("b")).alias("d")], + )? + .build()?; + + // Should work. The aliases don't stop the `count` from being rolled up + // and both are restored by the projection. + assert_optimized_plan_equal!( + plan, + @r" + Projection: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS n, count(alias1) AS d [n:Int64, d:Int64] + Aggregate: groupBy=[[]], aggr=[[sum(alias2), count(alias1)]] [sum(alias2):Int64;N, count(alias1):Int64] + Aggregate: groupBy=[[test.b AS alias1]], aggr=[[count(Int64(1)) AS alias2]] [alias1:Utf8, alias2:Int64] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] + " + ) + } + #[test] fn count_star_min_max_sum_and_distinct_with_groupby() -> Result<()> { let table_scan = test_table_scan_utf8_b()?; From 5982f10fa9c6d42dbfef6ecb48e7b9708510d6be Mon Sep 17 00:00:00 2001 From: Fangchen Li Date: Wed, 4 Nov 2026 19:10:31 -0800 Subject: [PATCH 2/2] fix: deduplicate aliased aggregates in SingleDistinctToGroupBy Stripping aliases from the rewritten outer aggregates turned aggregates that differed only in their alias, such as `count_distinct(b)` aliased as `n1` and `n2`, into the same `count(alias1)`. `Aggregate` rejected the repeated output name, so optimization failed with a duplicate field error on queries that succeed without the rewrite. Compute each outer aggregate once and let the projection select it under every original name, which is the plan SQL already produces. The cast check in that projection read types by position, so it now maps each original aggregate to its deduplicated position. Co-Authored-By: Claude Opus 5.5 --- datafusion/core/tests/dataframe/mod.rs | 44 ++++++ .../src/single_distinct_to_groupby.rs | 126 ++++++++++++++++-- 2 files changed, 158 insertions(+), 12 deletions(-) diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index f0d36b8a22f8a..1ed52130cf427 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -872,6 +872,50 @@ async fn aliased_single_distinct_aggregate_is_rewritten() -> Result<()> { Ok(()) } +/// Stripping the aliases leaves two identical aggregates, which an `Aggregate` +/// can't output together. They must share one, as the SQL planner's do. The +/// `sum` follows them so its position in the outer aggregate differs from its +/// position in the original one; `b` is `Int64` so that `sum` needs no cast, +/// which would count as a second distinct field and skip the rewrite. +#[tokio::test] +async fn identical_aliased_single_distinct_aggregates_are_computed_once() -> Result<()> { + let ctx = SessionContext::new(); + let df = ctx + .read_batch(record_batch!( + ("a", Int32, [1, 1, 2, 2]), + ("b", Int64, [1, 1, 1, 2]) + )?)? + .aggregate( + vec![col("a")], + vec![ + count_distinct(col("b")).alias("n1"), + count_distinct(col("b")).alias("n2"), + sum_distinct(col("b")).alias("s"), + ], + )?; + + assert_snapshot!(df.clone().into_optimized_plan()?, @r" + Projection: ?table?.a, count(alias1) AS n1, count(alias1) AS n2, sum(alias1) AS s + Aggregate: groupBy=[[?table?.a]], aggr=[[count(alias1), sum(alias1)]] + Aggregate: groupBy=[[?table?.a, ?table?.b AS alias1]], aggr=[[]] + TableScan: ?table? projection=[a, b] + "); + + assert_batches_eq!( + &[ + "+---+----+----+---+", + "| a | n1 | n2 | s |", + "+---+----+----+---+", + "| 1 | 1 | 1 | 1 |", + "| 2 | 2 | 2 | 3 |", + "+---+----+----+---+", + ], + &df.sort_by(vec![col("a")])?.collect().await? + ); + + Ok(()) +} + #[tokio::test] async fn test_aggregate_with_pk() -> Result<()> { // create the dataframe diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index e246c1e2e7ce9..007b121440c4c 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -24,6 +24,7 @@ use datafusion_expr::{ logical_plan::{Aggregate, LogicalPlan}, when, }; +use indexmap::IndexSet; /// single distinct to group by optimizer rule /// ```text @@ -387,6 +388,16 @@ impl OptimizerRule for SingleDistinctToGroupBy { .into_iter() .unzip(); + // With their aliases stripped, aggregates that differed only in + // their alias are the same expression, and an `Aggregate` cannot + // output a name twice. Compute each one once: the projection + // below still selects it under every original name. + let mut unique_outer_aggr_exprs = IndexSet::new(); + let outer_aggr_positions: Vec = outer_aggr_exprs + .into_iter() + .map(|expr| unique_outer_aggr_exprs.insert_full(expr).0) + .collect(); + // construct the inner AggrPlan let inner_agg = LogicalPlan::Aggregate(Aggregate::try_new( input, @@ -402,7 +413,7 @@ impl OptimizerRule for SingleDistinctToGroupBy { let outer_aggr = LogicalPlan::Aggregate(Aggregate::try_new( Arc::new(inner_agg), outer_group_exprs, - outer_aggr_exprs, + unique_outer_aggr_exprs.into_iter().collect(), )?); // so the aggregates are displayed in the same way even after the rewrite @@ -417,17 +428,26 @@ impl OptimizerRule for SingleDistinctToGroupBy { } None => group_expr, }) - .chain(outer_proj_exprs.into_iter().enumerate().map(|(idx, expr)| { - let idx = idx + group_size; - let (qualifier, field) = schema.qualified_field(idx); - let current_type = outer_aggr.schema().field(idx).data_type(); - let expr = if current_type != field.data_type() { - cast(expr, field.data_type().clone()) - } else { - expr - }; - expr.alias_qualified(qualifier.cloned(), field.name()) - })) + .chain( + outer_proj_exprs + .into_iter() + .zip(outer_aggr_positions) + .enumerate() + .map(|(idx, (expr, position))| { + let (qualifier, field) = + schema.qualified_field(idx + group_size); + let current_type = outer_aggr + .schema() + .field(position + group_size) + .data_type(); + let expr = if current_type != field.data_type() { + cast(expr, field.data_type().clone()) + } else { + expr + }; + expr.alias_qualified(qualifier.cloned(), field.name()) + }), + ) .collect(); Ok(Transformed::yes(project(outer_aggr, alias_expr)?)) @@ -768,6 +788,88 @@ mod tests { ) } + #[test] + fn identical_distinct_aggregates_with_different_aliases() -> Result<()> { + let table_scan = test_table_scan()?; + + // Both become `count(alias1)` once unaliased, so the outer aggregate + // computes it once. The trailing `sum` sits at a different position in + // the outer aggregate than in the original one. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + count_distinct(col("b")).alias("n1"), + count_distinct(col("b")).alias("n2"), + sum_distinct(col("b")).alias("s"), + ], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, count(alias1) AS n1, count(alias1) AS n2, sum(alias1) AS s [a:UInt32, n1:Int64, n2:Int64, s:UInt64;N] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias1)]] [a:UInt32, count(alias1):Int64, sum(alias1):UInt64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn identical_distinct_aggregates_one_unaliased() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![ + count_distinct(col("b")), + count_distinct(col("b")).alias("n"), + ], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: count(alias1) AS count(DISTINCT test.b), count(alias1) AS n [count(DISTINCT test.b):Int64, n:Int64] + Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64] + Aggregate: groupBy=[[test.b AS alias1]], aggr=[[]] [alias1:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn identical_common_aggregates_with_different_aliases() -> Result<()> { + let table_scan = test_table_scan()?; + + // Non-distinct aggregates read their own `aliasN` column, so they stay + // separate outer aggregates and need no deduplication. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + count_distinct(col("b")).alias("n"), + sum(col("c")).alias("x"), + sum(col("c")).alias("y"), + ], + )? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, count(alias1) AS n, sum(alias2) AS x, sum(alias3) AS y [a:UInt32, n:Int64, x:UInt64;N, y:UInt64;N] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias2), sum(alias3)]] [a:UInt32, count(alias1):Int64, sum(alias2):UInt64;N, sum(alias3):UInt64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2, sum(test.c) AS alias3]] [a:UInt32, alias1:UInt32, alias2:UInt64;N, alias3:UInt64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + #[test] fn nested_aliased_single_distinct_and_groupby() -> Result<()> { let table_scan = test_table_scan()?;