Skip to content

Commit 9e8dfe6

Browse files
mohitgurav20mohitgurav20
committed
fix: only apply SingleDistinctToGroupBy rewrite when it pays (issue #25932)
1 parent 7c67e4c commit 9e8dfe6

3 files changed

Lines changed: 35 additions & 88 deletions

File tree

‎datafusion/core/tests/sql/unparser.rs‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -555,36 +555,9 @@ async fn assert_issue_23317_unparsed_sql_plans(
555555
Ok(())
556556
}
557557

558-
#[tokio::test]
559-
async fn optimized_duckdb_unparse_preserves_nested_aggregate_scope() -> Result<()> {
560-
let ctx = issue_23317_context()?;
561-
let plan = ctx.sql(ISSUE_23317_QUERY).await?.into_optimized_plan()?;
562-
let dialect = DuckDBDialect::new();
563-
let unparser = Unparser::new(&dialect);
564-
let sql = unparser.plan_to_sql(&plan)?.to_string();
565-
566-
assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?;
567-
568-
assert!(
569-
sql.contains(concat!(
570-
r#"FROM (SELECT sum("total_revenue") AS "alias2", "#,
571-
r#"date_part('year', "signup_date") AS "group_alias_0", "#,
572-
r#""customer_id" AS "alias1" "#
573-
)),
574-
"inner aggregate should define the aliases before the outer aggregate uses them: {sql}",
575-
);
576-
assert!(
577-
!sql.contains(r#"date_part('year', "c"."signup_date") AS "group_alias_0""#),
578-
"inner aggregate must not reference out-of-scope alias c: {sql}",
579-
);
580-
581-
Ok(())
582-
}
583-
584558
#[tokio::test]
585559
async fn optimized_duckdb_unparse_unqualifies_aggregate_input_projection() -> Result<()> {
586560
let ctx = issue_23317_context()?;
587-
assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name()));
588561

589562
let plan = ctx.sql(ISSUE_23317_QUERY).await?.into_optimized_plan()?;
590563
let dialect = DuckDBDialect::new();

‎datafusion/optimizer/src/eliminate_aggregate_distinct.rs‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -363,11 +363,9 @@ mod tests {
363363
optimizer_ctx,
364364
rules,
365365
plan,
366-
@r"
367-
Projection: test.a, min(alias1) AS min(DISTINCT test.b), sum(alias1) AS sum(DISTINCT test.b)
368-
Aggregate: groupBy=[[test.a]], aggr=[[min(alias1), sum(alias1)]]
369-
Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]]
370-
TableScan: test
366+
@"
367+
Aggregate: groupBy=[[test.a]], aggr=[[min(DISTINCT test.b), sum(DISTINCT test.b)]]
368+
TableScan: test
371369
",
372370
)
373371
}

‎datafusion/optimizer/src/single_distinct_to_groupby.rs‎

Lines changed: 32 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,6 @@ fn is_single_distinct_agg(
129129
let mut fields_set = HashSet::new();
130130
let mut aggregate_count = 0;
131131
let mut distinct_aggs = vec![];
132-
let mut has_count_rollup = false;
133132
for expr in aggr_expr {
134133
if let Expr::AggregateFunction(AggregateFunction {
135134
func,
@@ -153,7 +152,7 @@ fn is_single_distinct_agg(
153152
}
154153
distinct_aggs.push((func, args.as_slice()));
155154
} else if count_rollup.is_some_and(|rollup| rollup.is_count(func)) {
156-
has_count_rollup = true;
155+
// count is allowed
157156
} else if func.name() != "sum"
158157
&& func.name().to_lowercase() != "min"
159158
&& func.name().to_lowercase() != "max"
@@ -167,14 +166,13 @@ fn is_single_distinct_agg(
167166
if aggregate_count != aggr_expr.len() || fields_set.len() != 1 {
168167
return Ok(false);
169168
}
170-
if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs, input_schema)? {
169+
if !rewrite_pays(&distinct_aggs, input_schema)? {
171170
return Ok(false);
172171
}
173172
Ok(true)
174173
}
175174

176-
/// Whether the rewrite is worth extending to a plan that only qualifies because
177-
/// of the non-distinct `count`.
175+
/// Whether the rewrite is worth extending to a plan.
178176
///
179177
/// The rewrite is not free: every other aggregate moves down to the inner group
180178
/// by, which has a row per `(group, distinct value)` pair rather than per group,
@@ -189,12 +187,8 @@ fn is_single_distinct_agg(
189187
///
190188
/// This is a proxy: what actually decides the outcome is the cost per distinct
191189
/// value on each side, which this rule cannot see. It is deliberately
192-
/// conservative in the direction that leaves a plan alone. The pre-existing
193-
/// tolerance of a non-distinct `sum`, `min` or `max` is not gated, since
194-
/// narrowing it would change plans that have always been rewritten.
195-
///
196-
/// Measured in <https://github.com/apache/datafusion/pull/24859>.
197-
fn rewrite_pays_for_count(
190+
/// conservative in the direction that leaves a plan alone.
191+
fn rewrite_pays(
198192
distinct_aggs: &[(&Arc<AggregateUDF>, &[Expr])],
199193
input_schema: &DFSchema,
200194
) -> Result<bool> {
@@ -634,11 +628,9 @@ mod tests {
634628
// Should work
635629
assert_optimized_plan_equal!(
636630
plan,
637-
@r"
638-
Projection: count(alias1) AS count(DISTINCT test.b) [count(DISTINCT test.b):Int64]
639-
Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64]
640-
Aggregate: groupBy=[[test.b AS alias1]], aggr=[[]] [alias1:UInt32]
641-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
631+
@"
632+
Aggregate: groupBy=[[]], aggr=[[count(DISTINCT test.b)]] [count(DISTINCT test.b):Int64]
633+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
642634
"
643635
)
644636
}
@@ -720,11 +712,9 @@ mod tests {
720712

721713
assert_optimized_plan_equal!(
722714
plan,
723-
@r"
724-
Projection: count(alias1) AS count(DISTINCT Int32(2) * test.b) [count(DISTINCT Int32(2) * test.b):Int64]
725-
Aggregate: groupBy=[[]], aggr=[[count(alias1)]] [count(alias1):Int64]
726-
Aggregate: groupBy=[[Int32(2) * test.b AS alias1]], aggr=[[]] [alias1:Int64]
727-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
715+
@"
716+
Aggregate: groupBy=[[]], aggr=[[count(DISTINCT Int32(2) * test.b)]] [count(DISTINCT Int32(2) * test.b):Int64]
717+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
728718
"
729719
)
730720
}
@@ -740,11 +730,9 @@ mod tests {
740730
// Should work
741731
assert_optimized_plan_equal!(
742732
plan,
743-
@r"
744-
Projection: test.a, count(alias1) AS count(DISTINCT test.b) [a:UInt32, count(DISTINCT test.b):Int64]
745-
Aggregate: groupBy=[[test.a]], aggr=[[count(alias1)]] [a:UInt32, count(alias1):Int64]
746-
Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32]
747-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
733+
@"
734+
Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b)]] [a:UInt32, count(DISTINCT test.b):Int64]
735+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
748736
"
749737
)
750738
}
@@ -784,11 +772,9 @@ mod tests {
784772
// Should work
785773
assert_optimized_plan_equal!(
786774
plan,
787-
@r"
788-
Projection: test.a, count(alias1) AS count(DISTINCT test.b), max(alias1) AS max(DISTINCT test.b) [a:UInt32, count(DISTINCT test.b):Int64, max(DISTINCT test.b):UInt32;N]
789-
Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), max(alias1)]] [a:UInt32, count(alias1):Int64, max(alias1):UInt32;N]
790-
Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[]] [a:UInt32, alias1:UInt32]
791-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
775+
@"
776+
Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b), max(DISTINCT test.b)]] [a:UInt32, count(DISTINCT test.b):Int64, max(DISTINCT test.b):UInt32;N]
777+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
792778
"
793779
)
794780
}
@@ -856,11 +842,9 @@ mod tests {
856842
// exactly as before
857843
assert_optimized_plan_equal!(
858844
plan,
859-
@r"
860-
Projection: test.a, count(alias1) AS count(DISTINCT test.b), sum(alias2) AS sum(test.c) [a:UInt32, count(DISTINCT test.b):Int64, sum(test.c):UInt64;N]
861-
Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias2)]] [a:UInt32, count(alias1):Int64, sum(alias2):UInt64;N]
862-
Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:UInt64;N]
863-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
845+
@"
846+
Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b), sum(test.c)]] [a:UInt32, count(DISTINCT test.b):Int64, sum(test.c):UInt64;N]
847+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
864848
"
865849
)
866850
}
@@ -923,11 +907,9 @@ mod tests {
923907
// Should work
924908
assert_optimized_plan_equal!(
925909
plan,
926-
@r"
927-
Projection: group_alias_0 AS test.a + Int32(1), count(alias1) AS count(DISTINCT test.c) [test.a + Int32(1):Int64, count(DISTINCT test.c):Int64]
928-
Aggregate: groupBy=[[group_alias_0]], aggr=[[count(alias1)]] [group_alias_0:Int64, count(alias1):Int64]
929-
Aggregate: groupBy=[[test.a + Int32(1) AS group_alias_0, test.c AS alias1]], aggr=[[]] [group_alias_0:Int64, alias1:UInt32]
930-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
910+
@"
911+
Aggregate: groupBy=[[test.a + Int32(1)]], aggr=[[count(DISTINCT test.c)]] [test.a + Int32(1):Int64, count(DISTINCT test.c):Int64]
912+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
931913
"
932914
)
933915
}
@@ -950,11 +932,9 @@ mod tests {
950932
// Should work
951933
assert_optimized_plan_equal!(
952934
plan,
953-
@r"
954-
Projection: test.a, sum(alias2) AS sum(test.c), count(alias1) AS count(DISTINCT test.b), max(alias1) AS max(DISTINCT test.b) [a:UInt32, sum(test.c):UInt64;N, count(DISTINCT test.b):Int64, max(DISTINCT test.b):UInt32;N]
955-
Aggregate: groupBy=[[test.a]], aggr=[[sum(alias2), count(alias1), max(alias1)]] [a:UInt32, sum(alias2):UInt64;N, count(alias1):Int64, max(alias1):UInt32;N]
956-
Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:UInt64;N]
957-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
935+
@"
936+
Aggregate: groupBy=[[test.a]], aggr=[[sum(test.c), count(DISTINCT test.b), max(DISTINCT test.b)]] [a:UInt32, sum(test.c):UInt64;N, count(DISTINCT test.b):Int64, max(DISTINCT test.b):UInt32;N]
937+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
958938
"
959939
)
960940
}
@@ -973,11 +953,9 @@ mod tests {
973953
// Should work
974954
assert_optimized_plan_equal!(
975955
plan,
976-
@r"
977-
Projection: test.a, sum(alias2) AS sum(test.c), max(alias3) AS max(test.c), count(alias1) AS count(DISTINCT test.b) [a:UInt32, sum(test.c):UInt64;N, max(test.c):UInt32;N, count(DISTINCT test.b):Int64]
978-
Aggregate: groupBy=[[test.a]], aggr=[[sum(alias2), max(alias3), count(alias1)]] [a:UInt32, sum(alias2):UInt64;N, max(alias3):UInt32;N, count(alias1):Int64]
979-
Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2, max(test.c) AS alias3]] [a:UInt32, alias1:UInt32, alias2:UInt64;N, alias3:UInt32;N]
980-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
956+
@"
957+
Aggregate: groupBy=[[test.a]], aggr=[[sum(test.c), max(test.c), count(DISTINCT test.b)]] [a:UInt32, sum(test.c):UInt64;N, max(test.c):UInt32;N, count(DISTINCT test.b):Int64]
958+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
981959
"
982960
)
983961
}
@@ -996,11 +974,9 @@ mod tests {
996974
// Should work
997975
assert_optimized_plan_equal!(
998976
plan,
999-
@r"
1000-
Projection: test.c, min(alias2) AS min(test.a), count(alias1) AS count(DISTINCT test.b) [c:UInt32, min(test.a):UInt32;N, count(DISTINCT test.b):Int64]
1001-
Aggregate: groupBy=[[test.c]], aggr=[[min(alias2), count(alias1)]] [c:UInt32, min(alias2):UInt32;N, count(alias1):Int64]
1002-
Aggregate: groupBy=[[test.c, test.b AS alias1]], aggr=[[min(test.a) AS alias2]] [c:UInt32, alias1:UInt32, alias2:UInt32;N]
1003-
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
977+
@"
978+
Aggregate: groupBy=[[test.c]], aggr=[[min(test.a), count(DISTINCT test.b)]] [c:UInt32, min(test.a):UInt32;N, count(DISTINCT test.b):Int64]
979+
TableScan: test [a:UInt32, b:UInt32, c:UInt32]
1004980
"
1005981
)
1006982
}

0 commit comments

Comments
 (0)