diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index 49635ed110e7f..1ed52130cf427 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -838,6 +838,84 @@ 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(()) +} + +/// 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/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..007b121440c4c 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, @@ -23,6 +24,7 @@ use datafusion_expr::{ logical_plan::{Aggregate, LogicalPlan}, when, }; +use indexmap::IndexSet; /// single distinct to group by optimizer rule /// ```text @@ -118,6 +120,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 +293,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,12 +382,22 @@ impl OptimizerRule for SingleDistinctToGroupBy { Ok((outer, proj)) } } - _ => Ok((aggr_expr.clone(), aggr_expr)), + other => Ok((other.clone(), other)), }) .collect::>>()? .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, @@ -392,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 @@ -407,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)?)) @@ -738,6 +768,273 @@ 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 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()?; + + // 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 +1503,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()?;