diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 00c8fab228117..6b69eb0c66bc0 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -27,6 +27,7 @@ use datafusion_common::{ }; use datafusion_expr::builder::project; use datafusion_expr::expr::AggregateFunctionParams; +use datafusion_expr::expr_rewriter::unalias; use datafusion_expr::{ Expr, col, expr::AggregateFunction, @@ -66,6 +67,12 @@ fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; for expr in aggr_expr { + // Peel aliases the DataFrame API applies to the aggregate itself (e.g. + // `count(col("b")).distinct().alias("n")`); SQL keeps `aggr_expr` bare. + let mut expr = expr; + while let Expr::Alias(alias) = expr { + expr = alias.expr.as_ref(); + } if let Expr::AggregateFunction(AggregateFunction { func, params: @@ -84,6 +91,13 @@ fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result { aggregate_count += 1; if *distinct { for e in args { + // A constant DISTINCT argument would become a constant GROUP BY + // key, which `eliminate_group_by_constant` collapses into a + // global aggregate that emits a row even on empty input, giving + // wrong results. The rewrite has no benefit here, so skip it. + if e.column_refs().is_empty() { + return Ok(false); + } fields_set.insert(e); } } else if func.name() != "sum" @@ -179,7 +193,9 @@ impl OptimizerRule for SingleDistinctToGroupBy { let mut inner_aggr_exprs = vec![]; let outer_aggr_exprs = 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: @@ -237,7 +253,7 @@ impl OptimizerRule for SingleDistinctToGroupBy { ))) } } - _ => Ok(aggr_expr), + other => Ok(other), }) .collect::>>()?; @@ -448,6 +464,216 @@ mod tests { ) } + #[test] + fn constant_distinct_arg_not_optimized() -> Result<()> { + let table_scan = test_table_scan()?; + + // A constant DISTINCT argument must not be rewritten: the constant GROUP BY + // key would be collapsed by eliminate_group_by_constant and yield a row on + // empty input (wrong result). Combined with a plain aggregate to exercise + // the non-empty inner-aggregate path. + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![count_distinct(lit(1)), max(col("a"))], + )? + .build()?; + + // Should not be optimized + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[]], aggr=[[count(DISTINCT Int32(1)), max(test.a)]] [count(DISTINCT Int32(1)):Int64, max(test.a):UInt32;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[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()?; diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index c5970bde9c954..75b892f2e1961 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -9405,3 +9405,40 @@ SET datafusion.execution.target_partitions = 4; statement ok DROP TABLE hits_raw; + +# Regression: single_distinct_to_groupby must not rewrite a DISTINCT aggregate over +# a constant. The constant GROUP BY key would be collapsed by +# eliminate_group_by_constant into a global aggregate that emits a row on empty +# input, giving wrong results. https://github.com/apache/datafusion/issues/23401 +statement ok +CREATE TABLE t_const_distinct(x INT) AS VALUES (1), (2), (3); + +# Empty input (WHERE false): DISTINCT over a constant alongside a plain aggregate. +query II +SELECT MIN(DISTINCT 37), MAX(x) FROM t_const_distinct WHERE false; +---- +NULL NULL + +query I +SELECT COUNT(DISTINCT 5) FROM t_const_distinct WHERE false; +---- +0 + +query R +SELECT AVG(DISTINCT -74) FROM t_const_distinct WHERE false; +---- +NULL + +# Non-empty input still returns the single constant value / count. +query II +SELECT MIN(DISTINCT 37), MAX(x) FROM t_const_distinct; +---- +37 3 + +query I +SELECT COUNT(DISTINCT 5) FROM t_const_distinct; +---- +1 + +statement ok +DROP TABLE t_const_distinct;