From 2b710bb6fd02cbf3d4936e6e3a835576eeb18d46 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 01:35:54 -0500 Subject: [PATCH 1/4] Allow a non-distinct count alongside the single distinct aggregate `SingleDistinctToGroupBy` rewrites `AGG(DISTINCT x)` into a two phase group by, which is what keeps a high cardinality distinct off the one-accumulator-per-group path in `GroupsAccumulatorAdapter`. The rule tolerated a non-distinct `sum`, `min` or `max` next to the distinct aggregate but bailed out on `count`, so the very common `count(*), count(DISTINCT x) ... GROUP BY` shape kept the unrewritten plan and its memory profile. Allow a non-distinct `count` as well. `count` is the one supported function whose outer phase is a different function: the inner group by counts the rows of each `(group, distinct value)` partition and the outer phase adds those partial counts up with `sum`, since count over a group is the sum of the counts of any partition of that group. Two details follow from that substitution: - `count` and `sum` are resolved from the session function registry and the rewrite only fires when the aggregate is that exact `count`, so a session without a registry or with its own `count` is left alone. - `count` returns a non-null 0 over an empty input while `sum` of no rows is NULL, which is reachable for an aggregate with no group by. The projection selects `CASE WHEN sum(alias) IS NOT NULL THEN sum(alias) ELSE 0 END`, which restores the 0 and keeps the column's type and nullability as `count` had them. The new sqllogictest file asserts every result twice, once with the optimizer disabled and once with it enabled, over data with NULL and all-NULL distinct values, an empty input, and `count(*)` versus `count(col)` versus `count(1)`. Co-Authored-By: Claude Opus 5 --- .../src/single_distinct_to_groupby.rs | 358 ++++++++++++++++-- .../sqllogictest/test_files/clickbench.slt | 28 +- .../test_files/single_distinct_to_groupby.slt | 304 +++++++++++++++ 3 files changed, 639 insertions(+), 51 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 00c8fab228117..42e54f2da99d7 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -28,9 +28,11 @@ use datafusion_common::{ use datafusion_expr::builder::project; use datafusion_expr::expr::AggregateFunctionParams; use datafusion_expr::{ - Expr, col, + AggregateUDF, Expr, col, expr::AggregateFunction, + lit, logical_plan::{Aggregate, LogicalPlan}, + when, }; /// single distinct to group by optimizer rule @@ -49,6 +51,33 @@ use datafusion_expr::{ /// ) /// GROUP BY a /// ``` +/// +/// A non-distinct `count` is also allowed alongside the distinct aggregate. It +/// is the one supported function whose outer phase is a *different* function: +/// the inner group-by counts rows per `(group, distinct value)` pair and the +/// outer phase adds those partial counts up with `sum`. +/// +/// ```text +/// Before: +/// SELECT a, count(*), count(DISTINCT b) +/// FROM t +/// GROUP BY a +/// +/// After: +/// SELECT a, +/// CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE 0 END, +/// count(alias1) +/// FROM ( +/// SELECT a, b as alias1, count(*) as alias2 +/// FROM t +/// GROUP BY a, b +/// ) +/// GROUP BY a +/// ``` +/// +/// The `CASE` covers the one input on which the two phases disagree: over an +/// empty input the inner group by produces no rows at all, and a `sum` of no +/// rows is NULL where `count` is 0. #[derive(Default, Debug)] pub struct SingleDistinctToGroupBy {} @@ -61,8 +90,38 @@ impl SingleDistinctToGroupBy { } } +/// The pair of functions used to compute a non-distinct `count` in two phases: +/// `count` identifies the aggregates that need the treatment, `sum` combines +/// the partial counts the inner group-by produces. +/// +/// Both are resolved from the session's function registry, so a plan built +/// without one keeps the previous behaviour of bailing out on `count`. +struct CountRollup { + count: Arc, + sum: Arc, +} + +impl CountRollup { + fn try_new(config: &dyn OptimizerConfig) -> Option { + let registry = config.function_registry()?; + Some(Self { + count: registry.udaf("count").ok()?, + sum: registry.udaf("sum").ok()?, + }) + } + + /// Whether `func` is the registry's `count`, and so decomposes into + /// `sum` over per-partition counts. + fn is_count(&self, func: &AggregateUDF) -> bool { + self.count.as_ref() == func + } +} + /// Check whether all aggregate exprs are distinct on a single field. -fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result { +fn is_single_distinct_agg( + aggr_expr: &[Expr], + count_rollup: Option<&CountRollup>, +) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; for expr in aggr_expr { @@ -89,6 +148,7 @@ fn is_single_distinct_agg(aggr_expr: &[Expr]) -> Result { } else if func.name() != "sum" && func.name().to_lowercase() != "min" && func.name().to_lowercase() != "max" + && !count_rollup.is_some_and(|rollup| rollup.is_count(func)) { return Ok(false); } @@ -120,8 +180,12 @@ impl OptimizerRule for SingleDistinctToGroupBy { fn rewrite( &self, plan: LogicalPlan, - _config: &dyn OptimizerConfig, + config: &dyn OptimizerConfig, ) -> Result, DataFusionError> { + if !matches!(plan, LogicalPlan::Aggregate(_)) { + return Ok(Transformed::no(plan)); + } + let count_rollup = CountRollup::try_new(config); match plan { LogicalPlan::Aggregate(Aggregate { input, @@ -129,7 +193,7 @@ impl OptimizerRule for SingleDistinctToGroupBy { schema, group_expr, .. - }) if is_single_distinct_agg(&aggr_expr)? + }) if is_single_distinct_agg(&aggr_expr, count_rollup.as_ref())? && !contains_grouping_set(&group_expr) => { let group_size = group_expr.len(); @@ -177,7 +241,11 @@ impl OptimizerRule for SingleDistinctToGroupBy { let mut index = 1; let mut group_fields_set = HashSet::new(); let mut inner_aggr_exprs = vec![]; - let outer_aggr_exprs = aggr_expr + // Each aggregate yields the expression the outer `Aggregate` + // computes and the expression the projection above it selects. + // They differ only for `count`, whose projection restores the + // 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 { Expr::AggregateFunction(AggregateFunction { @@ -204,18 +272,32 @@ impl OptimizerRule for SingleDistinctToGroupBy { inner_group_exprs .push(arg.alias(SINGLE_DISTINCT_ALIAS)); } - Ok(Expr::AggregateFunction(AggregateFunction::new_udf( - func, - vec![col(SINGLE_DISTINCT_ALIAS)], - false, // intentional to remove distinct here - filter, - order_by, - null_treatment, - ))) + let outer = + Expr::AggregateFunction(AggregateFunction::new_udf( + func, + vec![col(SINGLE_DISTINCT_ALIAS)], + false, // intentional to remove distinct here + filter, + order_by, + null_treatment, + )); + Ok((outer.clone(), outer)) // if the aggregate function is not distinct, we need to rewrite it like two phase aggregation } else { index += 1; let alias_str = format!("alias{index}"); + // `count` is the one function whose two phases + // use different aggregates: the inner group-by + // counts the rows of each `(group, distinct + // value)` partition and the outer phase adds + // those partial counts up. + let rollup = count_rollup + .as_ref() + .filter(|rollup| rollup.is_count(&func)); + let outer_func = match rollup { + Some(rollup) => Arc::clone(&rollup.sum), + None => Arc::clone(&func), + }; inner_aggr_exprs.push( Expr::AggregateFunction(AggregateFunction::new_udf( Arc::clone(&func), @@ -227,19 +309,34 @@ impl OptimizerRule for SingleDistinctToGroupBy { )) .alias(&alias_str), ); - Ok(Expr::AggregateFunction(AggregateFunction::new_udf( - func, - vec![col(&alias_str)], - false, - None, - vec![], - None, - ))) + let outer = + Expr::AggregateFunction(AggregateFunction::new_udf( + outer_func, + vec![col(&alias_str)], + false, + None, + vec![], + None, + )); + if rollup.is_none() { + return Ok((outer.clone(), outer)); + } + // The inner group-by produces no rows at all + // for an empty input, and `sum` reports that as + // NULL where `count` reports 0. Restore the 0, + // which also keeps the column non-nullable as + // `count` had it. + let proj = + when(outer.clone().is_not_null(), outer.clone()) + .otherwise(lit(0_i64))?; + Ok((outer, proj)) } } - _ => Ok(aggr_expr), + _ => Ok((aggr_expr.clone(), aggr_expr)), }) - .collect::>>()?; + .collect::>>()? + .into_iter() + .unzip(); // construct the inner AggrPlan let inner_agg = LogicalPlan::Aggregate(Aggregate::try_new( @@ -265,13 +362,11 @@ impl OptimizerRule for SingleDistinctToGroupBy { } None => group_expr, }) - .chain(outer_aggr_exprs.iter().cloned().enumerate().map( - |(idx, expr)| { - let idx = idx + group_size; - let (qualifier, field) = schema.qualified_field(idx); - expr.alias_qualified(qualifier.cloned(), field.name()) - }, - )) + .chain(outer_proj_exprs.into_iter().enumerate().map(|(idx, expr)| { + let idx = idx + group_size; + let (qualifier, field) = schema.qualified_field(idx); + expr.alias_qualified(qualifier.cloned(), field.name()) + })) .collect(); let outer_aggr = LogicalPlan::Aggregate(Aggregate::try_new( @@ -291,8 +386,13 @@ mod tests { use super::*; use crate::assert_optimized_plan_eq_display_indent_snapshot; use crate::test::*; + use crate::{Optimizer, OptimizerContext}; + use chrono::{DateTime, Utc}; + use datafusion_common::alias::AliasGenerator; + use datafusion_common::config::ConfigOptions; use datafusion_expr::ExprFunctionExt; use datafusion_expr::expr::GroupingSet; + use datafusion_expr::registry::{FunctionRegistry, MemoryFunctionRegistry}; use datafusion_expr::{lit, logical_plan::builder::LogicalPlanBuilder}; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::expr_fn::{count, count_distinct, max, min, sum}; @@ -310,17 +410,70 @@ mod tests { )) } + fn count_star() -> Expr { + Expr::AggregateFunction(AggregateFunction::new_udf( + count_udaf(), + vec![lit(1_i64)], + false, + None, + vec![], + None, + )) + } + + /// An [`OptimizerConfig`] that resolves functions the way a session does. + /// The rule needs `count` and `sum` from the registry to rewrite a + /// non-distinct `count`. + #[derive(Debug)] + struct RegistryOptimizerContext { + inner: OptimizerContext, + registry: MemoryFunctionRegistry, + } + + impl RegistryOptimizerContext { + fn new() -> Self { + let mut registry = MemoryFunctionRegistry::new(); + registry.register_udaf(count_udaf()).unwrap(); + registry.register_udaf(sum_udaf()).unwrap(); + Self { + inner: OptimizerContext::new(), + registry, + } + } + } + + impl OptimizerConfig for RegistryOptimizerContext { + fn query_execution_start_time(&self) -> Option> { + self.inner.query_execution_start_time() + } + + fn alias_generator(&self) -> &Arc { + self.inner.alias_generator() + } + + fn options(&self) -> Arc { + self.inner.options() + } + + fn function_registry(&self) -> Option<&dyn FunctionRegistry> { + Some(&self.registry) + } + } + macro_rules! assert_optimized_plan_equal { ( $plan:expr, @ $expected:literal $(,)? ) => {{ - let rule: Arc = Arc::new(SingleDistinctToGroupBy::new()); - assert_optimized_plan_eq_display_indent_snapshot!( - rule, - $plan, - @ $expected, - ) + let rule: Arc = + Arc::new(SingleDistinctToGroupBy::new()); + let optimizer = Optimizer::with_rules(vec![rule]); + let optimized_plan = optimizer + .optimize($plan, &RegistryOptimizerContext::new(), |_, _| {}) + .expect("failed to optimize plan"); + insta::assert_snapshot!(optimized_plan.display_indent_schema(), @ $expected); + + Ok::<(), DataFusionError>(()) }}; } @@ -523,12 +676,15 @@ mod tests { )? .build()?; - // Do nothing + // Should work: the non-distinct count becomes a sum of the counts the + // inner group-by produces per (test.a, test.b) pair assert_optimized_plan_equal!( plan, @r" - Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b), count(test.c)]] [a:UInt32, count(DISTINCT test.b):Int64, count(test.c):Int64] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + Projection: test.a, count(alias1) AS count(DISTINCT test.b), CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(test.c) [a:UInt32, count(DISTINCT test.b):Int64, count(test.c):Int64] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias2)]] [a:UInt32, count(alias1):Int64, sum(alias2):Int64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] " ) } @@ -752,4 +908,128 @@ mod tests { " ) } + + #[test] + fn count_star_and_distinct_without_groupby() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + Vec::::new(), + vec![count_star(), count_distinct(col("b"))], + )? + .build()?; + + // Should work. Without a group by the outer aggregate sees no rows at + // all for an empty input, so the projection has to turn the NULL that + // `sum` reports there back into the 0 `count` reports. + assert_optimized_plan_equal!( + plan, + @r" + Projection: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(Int64(1)), count(alias1) AS count(DISTINCT test.b) [count(Int64(1)):Int64, count(DISTINCT test.b):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:UInt32, alias2:Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn count_star_min_max_sum_and_distinct_with_groupby() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![ + count_star(), + count_distinct(col("b")), + min(col("c")), + max(col("c")), + sum(col("c")), + ], + )? + .build()?; + + // Should work: this is the shape a `count(*)` alongside a + // `count(DISTINCT ..)`, a min, a max and a sum produces + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.a, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(Int64(1)), count(alias1) AS count(DISTINCT test.b), min(alias3) AS min(test.c), max(alias4) AS max(test.c), sum(alias5) AS sum(test.c) [a:UInt32, count(Int64(1)):Int64, count(DISTINCT test.b):Int64, min(test.c):UInt32;N, max(test.c):UInt32;N, sum(test.c):UInt64;N] + Aggregate: groupBy=[[test.a]], aggr=[[sum(alias2), count(alias1), min(alias3), max(alias4), sum(alias5)]] [a:UInt32, sum(alias2):Int64;N, count(alias1):Int64, min(alias3):UInt32;N, max(alias4):UInt32;N, sum(alias5):UInt64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(Int64(1)) AS alias2, min(test.c) AS alias3, max(test.c) AS alias4, sum(test.c) AS alias5]] [a:UInt32, alias1:UInt32, alias2:Int64, alias3:UInt32;N, alias4:UInt32;N, alias5:UInt64;N] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn count_with_filter_is_not_rewritten() -> Result<()> { + let table_scan = test_table_scan()?; + + // count(a) FILTER (WHERE a > 5) + let expr = count_udaf() + .call(vec![col("a")]) + .filter(col("a").gt(lit(5))) + .build()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("c")], vec![expr, count_distinct(col("b"))])? + .build()?; + + // Do nothing: the filter would have to be applied per input row, but + // the inner aggregate has already collapsed them + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.c]], aggr=[[count(test.a) FILTER (WHERE test.a > Int32(5)), count(DISTINCT test.b)]] [c:UInt32, count(test.a) FILTER (WHERE test.a > Int32(5)):Int64, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn count_with_order_by_is_not_rewritten() -> Result<()> { + let table_scan = test_table_scan()?; + + // count(a ORDER BY a) + let expr = count_udaf() + .call(vec![col("a")]) + .order_by(vec![col("a").sort(true, false)]) + .build()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("c")], vec![expr, count_distinct(col("b"))])? + .build()?; + + // Do nothing + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.c]], aggr=[[count(test.a) ORDER BY [test.a ASC NULLS LAST], count(DISTINCT test.b)]] [c:UInt32, count(test.a) ORDER BY [test.a ASC NULLS LAST]:Int64, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn count_without_function_registry_is_not_rewritten() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("a")], vec![count_star(), count_distinct(col("b"))])? + .build()?; + + // Do nothing: without a registry the rule cannot resolve the `sum` the + // outer phase needs + let rule: Arc = + Arc::new(SingleDistinctToGroupBy::new()); + assert_optimized_plan_eq_display_indent_snapshot!( + rule, + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(Int64(1)), count(DISTINCT test.b)]] [a:UInt32, count(Int64(1)):Int64, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + ", + ) + } } diff --git a/datafusion/sqllogictest/test_files/clickbench.slt b/datafusion/sqllogictest/test_files/clickbench.slt index 4a1ef833c91db..f9e9637fad496 100644 --- a/datafusion/sqllogictest/test_files/clickbench.slt +++ b/datafusion/sqllogictest/test_files/clickbench.slt @@ -616,21 +616,25 @@ EXPLAIN SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*) AS c, COUNT(DI ---- logical_plan 01)Sort: c DESC NULLS FIRST, fetch=10 -02)--Projection: hits.SearchPhrase, min(hits.URL), min(hits.Title), count(Int64(1)) AS count(*) AS c, count(DISTINCT hits.UserID) -03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)]] -04)------SubqueryAlias: hits -05)--------Filter: hits_raw.SearchPhrase != Utf8View("") AND hits_raw.Title LIKE Utf8View("%Google%") AND hits_raw.URL NOT LIKE Utf8View("%.google.%") -06)----------TableScan: hits_raw projection=[Title, UserID, URL, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View(""), hits_raw.Title LIKE Utf8View("%Google%"), hits_raw.URL NOT LIKE Utf8View("%.google.%")] +02)--Projection: hits.SearchPhrase, min(alias2) AS min(hits.URL), min(alias3) AS min(hits.Title), CASE WHEN sum(alias4) IS NOT NULL THEN sum(alias4) ELSE Int64(0) END AS c, count(alias1) AS count(DISTINCT hits.UserID) +03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[min(alias2), min(alias3), sum(alias4), count(alias1)]] +04)------Aggregate: groupBy=[[hits.SearchPhrase, hits.UserID AS alias1]], aggr=[[min(hits.URL) AS alias2, min(hits.Title) AS alias3, count(Int64(1)) AS alias4]] +05)--------SubqueryAlias: hits +06)----------Filter: hits_raw.SearchPhrase != Utf8View("") AND hits_raw.Title LIKE Utf8View("%Google%") AND hits_raw.URL NOT LIKE Utf8View("%.google.%") +07)------------TableScan: hits_raw projection=[Title, UserID, URL, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View(""), hits_raw.Title LIKE Utf8View("%Google%"), hits_raw.URL NOT LIKE Utf8View("%.google.%")] physical_plan 01)SortPreservingMergeExec: [c@3 DESC], fetch=10 -02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(hits.URL)@1 as min(hits.URL), min(hits.Title)@2 as min(hits.Title), count(Int64(1))@3 as c, count(DISTINCT hits.UserID)@4 as count(DISTINCT hits.UserID)] -03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@3 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)] +02)--SortExec: TopK(fetch=10), expr=[c@3 DESC], preserve_partitioning=[true] +03)----ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(alias2)@1 as min(hits.URL), min(alias3)@2 as min(hits.Title), CASE WHEN sum(alias4)@3 IS NOT NULL THEN sum(alias4)@3 ELSE 0 END as c, count(alias1)@4 as count(DISTINCT hits.UserID)] +04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(alias2), min(alias3), sum(alias4), count(alias1)] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@3 as SearchPhrase], aggr=[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)] -07)------------FilterExec: SearchPhrase@3 != AND Title@0 LIKE %Google% AND URL@2 NOT LIKE %.google.% -08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[Title, UserID, URL, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND Title@2 LIKE %Google% AND URL@13 NOT LIKE %.google.%, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] +06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(alias2), min(alias3), sum(alias4), count(alias1)] +07)------------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase, alias1@1 as alias1], aggr=[min(hits.URL) as alias2, min(hits.Title) as alias3, count(1) as alias4] +08)--------------RepartitionExec: partitioning=Hash([SearchPhrase@0, alias1@1], 4), input_partitions=4 +09)----------------AggregateExec: mode=Partial, gby=[SearchPhrase@3 as SearchPhrase, UserID@1 as alias1], aggr=[min(hits.URL) as alias2, min(hits.Title) as alias3, count(1) as alias4] +10)------------------FilterExec: SearchPhrase@3 != AND Title@0 LIKE %Google% AND URL@2 NOT LIKE %.google.% +11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[Title, UserID, URL, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND Title@2 LIKE %Google% AND URL@13 NOT LIKE %.google.%, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query TTTII SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*) AS c, COUNT(DISTINCT "UserID") FROM hits WHERE "Title" LIKE '%Google%' AND "URL" NOT LIKE '%.google.%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; diff --git a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt new file mode 100644 index 0000000000000..d7c47a8626760 --- /dev/null +++ b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt @@ -0,0 +1,304 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Tests for the `single_distinct_aggregation_to_group_by` rule, which rewrites a +# single `AGG(DISTINCT x)` into a two phase group by. Every result here is +# asserted twice: once with the optimizer disabled +# (`datafusion.optimizer.max_passes = 0`, so the aggregate runs as written) and +# once with it enabled. The two must agree. + +statement ok +CREATE TABLE t(g INT, x INT, v INT); + +# g = 1: four rows, two distinct x, one NULL x, one NULL v +# g = 2: two rows, x is NULL throughout +# g = 3: a single row +statement ok +INSERT INTO t VALUES + (1, 10, 100), + (1, 10, NULL), + (1, NULL, 5), + (1, 20, 7), + (2, NULL, NULL), + (2, NULL, 3), + (3, 30, 1); + +statement ok +CREATE TABLE empty_t(g INT, x INT, v INT); + +statement ok +set datafusion.explain.logical_plan_only = true; + +######## +# The rewrite fires for a non-distinct count next to the distinct aggregate +######## + +query TT +EXPLAIN SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS records, count(alias1) AS distinct_x +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]] +03)----Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2]] +04)------TableScan: t projection=[g, x] + +# The production shape: a count next to min, max, sum and the distinct count +query TT +EXPLAIN +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x, + min(v) AS min_v, max(v) AS max_v, sum(v) AS sum_v +FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS records, count(alias1) AS distinct_x, min(alias3) AS min_v, max(alias4) AS max_v, sum(alias5) AS sum_v +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1), min(alias3), max(alias4), sum(alias5)]] +03)----Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2, min(t.v) AS alias3, max(t.v) AS alias4, sum(CAST(t.v AS Int64)) AS alias5]] +04)------TableScan: t projection=[g, x, v] + +# A count with a FILTER still blocks the rewrite: the filter is per input row, +# and the inner aggregate has already collapsed those rows +query TT +EXPLAIN SELECT g, count(*) FILTER (WHERE v > 1) AS records, count(DISTINCT x) AS distinct_x FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, count(Int64(1)) FILTER (WHERE t.v > Int64(1)) AS count(*) FILTER (WHERE t.v > Int64(1)) AS records, count(DISTINCT t.x) AS distinct_x +02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)) FILTER (WHERE t.v > Int32(1)) AS count(Int64(1)) FILTER (WHERE t.v > Int64(1)), count(DISTINCT t.x)]] +03)----TableScan: t projection=[g, x, v] + +# An unsupported non-distinct aggregate still blocks the rewrite +query TT +EXPLAIN SELECT g, avg(v) AS avg_v, count(DISTINCT x) AS distinct_x FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, avg(t.v) AS avg_v, count(DISTINCT t.x) AS distinct_x +02)--Aggregate: groupBy=[[t.g]], aggr=[[avg(CAST(t.v AS Float64)), count(DISTINCT t.x)]] +03)----TableScan: t projection=[g, x, v] + +statement ok +set datafusion.explain.logical_plan_only = false; + +######## +# count(*) vs count(col) vs count(1), grouped +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIIII +SELECT g, count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t GROUP BY g ORDER BY g; +---- +1 4 4 3 3 2 +2 2 2 0 1 0 +3 1 1 1 1 1 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIIII +SELECT g, count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t GROUP BY g ORDER BY g; +---- +1 4 4 3 3 2 +2 2 2 0 1 0 +3 1 1 1 1 1 + +######## +# count(*) vs count(col) vs count(1), no GROUP BY +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t; +---- +7 7 4 5 3 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(v) AS cnt_v, + count(DISTINCT x) AS distinct_x +FROM t; +---- +7 7 4 5 3 + +######## +# The distinct column is NULL for every row of the group +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIII +SELECT count(*) AS star, count(x) AS cnt_x, count(v) AS cnt_v, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 2; +---- +2 0 1 0 3 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIII +SELECT count(*) AS star, count(x) AS cnt_x, count(v) AS cnt_v, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 2; +---- +2 0 1 0 3 + +######## +# Empty input. Without a GROUP BY the aggregate still emits one row, and the +# counts on it must be 0 rather than the NULL an unguarded `sum` would give. +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM empty_t; +---- +0 0 0 0 NULL + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 99; +---- +0 0 0 0 NULL + +query IIII +SELECT g, count(*) AS star, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x +FROM empty_t GROUP BY g ORDER BY g; +---- + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM empty_t; +---- +0 0 0 0 NULL + +query IIIII +SELECT count(*) AS star, count(1) AS one, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x, + max(v) AS max_v +FROM t WHERE g = 99; +---- +0 0 0 0 NULL + +query IIII +SELECT g, count(*) AS star, count(x) AS cnt_x, count(DISTINCT x) AS distinct_x +FROM empty_t GROUP BY g ORDER BY g; +---- + +######## +# HAVING and ORDER BY read the rewritten count by name from the node above +######## + +statement ok +set datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x +FROM t GROUP BY g HAVING count(*) > 1 AND count(DISTINCT x) > 0 +ORDER BY count(*) DESC, g; +---- +logical_plan +01)Sort: records DESC NULLS FIRST, t.g ASC NULLS LAST +02)--Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS records, count(alias1) AS distinct_x +03)----Filter: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END > Int64(1) AND count(alias1) > Int64(0) +04)------Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]] +05)--------Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2]] +06)----------TableScan: t projection=[g, x] + +statement ok +set datafusion.explain.logical_plan_only = false; + +statement ok +set datafusion.optimizer.max_passes = 0; + +query III +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x +FROM t GROUP BY g HAVING count(*) > 1 AND count(DISTINCT x) > 0 +ORDER BY count(*) DESC, g; +---- +1 4 2 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query III +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x +FROM t GROUP BY g HAVING count(*) > 1 AND count(DISTINCT x) > 0 +ORDER BY count(*) DESC, g; +---- +1 4 2 + +######## +# The whole production shape, aggregating over a join +######## + +statement ok +CREATE TABLE failed(x INT); + +statement ok +INSERT INTO failed VALUES (10), (30); + +statement ok +set datafusion.optimizer.max_passes = 0; + +query IIIIII +SELECT t.g, count(*) AS records, count(DISTINCT t.x) AS distinct_x, + min(t.v) AS min_v, max(t.v) AS max_v, sum(t.v) AS sum_v +FROM t JOIN failed f ON t.x = f.x +GROUP BY t.g ORDER BY t.g; +---- +1 2 1 100 100 100 +3 1 1 1 1 1 + +statement ok +set datafusion.optimizer.max_passes = 3; + +query IIIIII +SELECT t.g, count(*) AS records, count(DISTINCT t.x) AS distinct_x, + min(t.v) AS min_v, max(t.v) AS max_v, sum(t.v) AS sum_v +FROM t JOIN failed f ON t.x = f.x +GROUP BY t.g ORDER BY t.g; +---- +1 2 1 100 100 100 +3 1 1 1 1 1 + +statement ok +DROP TABLE failed; + +statement ok +DROP TABLE empty_t; + +statement ok +DROP TABLE t; From f47c04523f7f84ad0993001b47ae1bc2ce92b0d5 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:14:55 -0500 Subject: [PATCH 2/4] Keep the substrait roundtrip test off the rewritten plan `aggregate_distinct_with_having` round trips `SELECT a, count(distinct b) ... HAVING count(b) > 100` through substrait and asserts the plan comes back displaying identically. It passed only because the non-distinct `count` made `SingleDistinctToGroupBy` bail out, so the plan had no aliases in it. With the rule now allowing that `count`, the query is rewritten and the assertion fails. The failure is a pre-existing substrait gap rather than anything specific to this query: substrait carries no names for an aggregate's grouping and measure expressions, so the consumer derives them from the expressions themselves and the `alias1` and `alias2` names the rule introduces are lost. Any plan the rule rewrites fails the same way, including the plain `SELECT a, count(distinct b) FROM data GROUP BY a, c` that this change does not touch. Remove the rule from the session used by this one test, so it keeps covering the un-rewritten aggregate it was written for instead of depending on the rule bailing out. Co-Authored-By: Claude Opus 5 --- .../tests/cases/roundtrip_logical_plan.rs | 34 +++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index c9f874dd9b095..6c5cd182ac8a8 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -333,8 +333,13 @@ async fn simple_aggregate() -> Result<()> { #[tokio::test] async fn aggregate_distinct_with_having() -> Result<()> { - roundtrip("SELECT a, count(distinct b) FROM data GROUP BY a, c HAVING count(b) > 100") - .await + let ctx = create_context_without_single_distinct_to_group_by().await?; + roundtrip_with_ctx( + "SELECT a, count(distinct b) FROM data GROUP BY a, c HAVING count(b) > 100", + ctx, + ) + .await?; + Ok(()) } #[tokio::test] @@ -2886,6 +2891,31 @@ async fn create_context() -> Result { create_context_with_dialect(None).await } +/// [`create_context`] with `single_distinct_aggregation_to_group_by` removed from +/// the optimizer. +/// +/// That rule rewrites a single `AGG(DISTINCT x)` into a two phase group by whose +/// inner aggregate aliases its grouping and measure expressions (`alias1`, +/// `alias2`). Substrait carries no names for those expressions, so the consumer +/// derives them from the expressions themselves and a rewritten plan does not +/// round trip to an identical plan. That holds for every output of the rule, not +/// only for the query under test. +async fn create_context_without_single_distinct_to_group_by() -> Result { + let ctx = create_context().await?; + let rules = ctx + .state() + .optimizer() + .rules + .iter() + .filter(|rule| rule.name() != "single_distinct_aggregation_to_group_by") + .cloned() + .collect(); + let state = SessionStateBuilder::new_from_existing(ctx.state()) + .with_optimizer_rules(rules) + .build(); + Ok(SessionContext::new_with_state(state)) +} + async fn create_context_with_dialect(dialect: Option) -> Result { let mut session_config = SessionConfig::default(); From d510dd45f8ad9008f9312124913d7defa503910a Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:56:42 -0500 Subject: [PATCH 3/4] Gate the count on the distinct argument lacking a groups accumulator A memory limited ClickBench run contradicted the memory rationale for the case it measures. Q22 is `SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*), COUNT(DISTINCT "UserID") ... GROUP BY "SearchPhrase"`, the only plan the change touched, and its peak memory pool reservation went from 3.1 MiB to 7.3 MiB, up 132.2%, with neighbouring queries moving by a few percent either way. `UserID` is an `Int64`, and `Count::groups_accumulator_supported` returns true for every integer type and false for everything else. So the base plan already had `PrimitiveDistinctCountGroupsAccumulator` and never went near `GroupsAccumulatorAdapter`. The rewrite replaced a compact vectorized accumulator with an inner group by on `(SearchPhrase, UserID)`, which also carries `min("URL")` and `min("Title")` at that much finer grain, and bought nothing back. Measured locally on 4M rows, 500k groups and 2M distinct pairs, peak pool reservation for `SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g`: x BIGINT 135 MiB unrewritten, 205 MiB rewritten x VARCHAR 13.6 GiB unrewritten, 255 MiB rewritten Peak RSS for the VARCHAR pair was 18.0 GiB against 433 MiB, so the pool figure is real memory rather than an accounting artifact. The rewrite pays exactly when the distinct argument would otherwise land on the adapter. `datafusion-optimizer` cannot depend on `datafusion-functions-aggregate` to read that list of types, and has no `AccumulatorArgs` to ask `groups_accumulator_supported` with. Its `Cargo.toml` names the way out, so this adds `AggregateUDFImpl::groups_accumulator_supported_for_types`, defaulting to false as `groups_accumulator_supported` does. `Count` is the only implementor and its physical method now delegates to it, so the two cannot disagree. The gate covers only the `count` this branch added. A plan that already qualified through a non-distinct `sum`, `min` or `max` is rewritten exactly as before, over any distinct argument type. The ClickBench snapshot returns to its base form, so no existing snapshot in the repository changes. `single_distinct_to_groupby.slt` now carries the same values in a `VARCHAR` and an `INT` column and asserts both sides of the gate, still under both `datafusion.optimizer.max_passes = 0` and the default. --- datafusion/expr/src/udaf.rs | 45 ++++++ datafusion/functions-aggregate/src/count.rs | 22 ++- .../src/single_distinct_to_groupby.rs | 145 ++++++++++++++++-- .../sqllogictest/test_files/clickbench.slt | 28 ++-- .../test_files/single_distinct_to_groupby.slt | 128 ++++++++++++++-- 5 files changed, 324 insertions(+), 44 deletions(-) diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index 8f7e9cc6cfc2b..b04c70f6c52b5 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -253,6 +253,16 @@ impl AggregateUDF { self.inner.groups_accumulator_supported(args) } + /// See [`AggregateUDFImpl::groups_accumulator_supported_for_types`] for more details. + pub fn groups_accumulator_supported_for_types( + &self, + arg_types: &[DataType], + is_distinct: bool, + ) -> bool { + self.inner + .groups_accumulator_supported_for_types(arg_types, is_distinct) + } + /// See [`AggregateUDFImpl::create_groups_accumulator`] for more details. pub fn create_groups_accumulator( &self, @@ -617,6 +627,32 @@ pub trait AggregateUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { false } + /// The same question as [`Self::groups_accumulator_supported`], asked with + /// only the information a logical plan carries. + /// + /// Physical planning has an [`AccumulatorArgs`] to ask with; an optimizer + /// rule does not, and cannot fabricate one, so this is how a logical + /// caller learns whether a call would get a specialized + /// [`GroupsAccumulator`] or fall back to one boxed [`Accumulator`] per + /// group in `GroupsAccumulatorAdapter`. That distinction is worth a rule + /// changing its mind over: the adapter's per-group state can be orders of + /// magnitude larger. + /// + /// The default is `false`, matching the default of + /// [`Self::groups_accumulator_supported`]. An implementation that + /// overrides that one and whose answer is decided by the argument types + /// and `DISTINCT` alone should override this one too, and have the + /// physical method call it so the two cannot disagree. An implementation + /// whose answer needs more than the argument types should leave this at + /// `false`, which claims nothing. + fn groups_accumulator_supported_for_types( + &self, + _arg_types: &[DataType], + _is_distinct: bool, + ) -> bool { + false + } + /// Return a specialized [`GroupsAccumulator`] that manages state /// for all groups. /// @@ -1552,6 +1588,15 @@ impl AggregateUDFImpl for AliasedAggregateUDFImpl { self.inner.groups_accumulator_supported(args) } + fn groups_accumulator_supported_for_types( + &self, + arg_types: &[DataType], + is_distinct: bool, + ) -> bool { + self.inner + .groups_accumulator_supported_for_types(arg_types, is_distinct) + } + fn create_groups_accumulator( &self, args: AccumulatorArgs, diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index 4c95ea431809e..ad020cfa8f179 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -346,14 +346,30 @@ impl AggregateUDFImpl for Count { } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - if args.exprs.len() != 1 { + // The answer depends on nothing but the argument types and `DISTINCT`, + // so defer to the logical form and keep one list of supported types. + let arg_types = args + .expr_fields + .iter() + .map(|field| field.data_type().clone()) + .collect::>(); + self.groups_accumulator_supported_for_types(&arg_types, args.is_distinct) + } + + fn groups_accumulator_supported_for_types( + &self, + arg_types: &[DataType], + is_distinct: bool, + ) -> bool { + if arg_types.len() != 1 { return false; } - if !args.is_distinct { + if !is_distinct { return true; } + // Keep in step with `create_distinct_count_groups_accumulator`. matches!( - args.expr_fields[0].data_type(), + arg_types[0], DataType::Int8 | DataType::Int16 | DataType::Int32 diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 42e54f2da99d7..34bf1cd039147 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -23,10 +23,12 @@ use crate::optimizer::ApplyOrder; use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::{ - DataFusionError, HashSet, Result, assert_eq_or_internal_err, tree_node::Transformed, + DFSchema, DataFusionError, HashSet, Result, assert_eq_or_internal_err, + tree_node::Transformed, }; use datafusion_expr::builder::project; use datafusion_expr::expr::AggregateFunctionParams; +use datafusion_expr::expr_schema::ExprSchemable; use datafusion_expr::{ AggregateUDF, Expr, col, expr::AggregateFunction, @@ -78,6 +80,9 @@ use datafusion_expr::{ /// The `CASE` covers the one input on which the two phases disagree: over an /// empty input the inner group by produces no rows at all, and a `sum` of no /// rows is NULL where `count` is 0. +/// +/// That `count` is allowed only when the distinct argument has no specialized +/// `GroupsAccumulator` of its own. See [`rewrite_pays_for_count`]. #[derive(Default, Debug)] pub struct SingleDistinctToGroupBy {} @@ -120,10 +125,13 @@ impl CountRollup { /// Check whether all aggregate exprs are distinct on a single field. fn is_single_distinct_agg( aggr_expr: &[Expr], + input_schema: &DFSchema, count_rollup: Option<&CountRollup>, ) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; + let mut distinct_aggs = vec![]; + let mut has_count_rollup = false; for expr in aggr_expr { if let Expr::AggregateFunction(AggregateFunction { func, @@ -145,10 +153,12 @@ fn is_single_distinct_agg( for e in args { fields_set.insert(e); } + distinct_aggs.push((func, args)); + } else if count_rollup.is_some_and(|rollup| rollup.is_count(func)) { + has_count_rollup = true; } else if func.name() != "sum" && func.name().to_lowercase() != "min" && func.name().to_lowercase() != "max" - && !count_rollup.is_some_and(|rollup| rollup.is_count(func)) { return Ok(false); } @@ -156,7 +166,45 @@ fn is_single_distinct_agg( return Ok(false); } } - Ok(aggregate_count == aggr_expr.len() && fields_set.len() == 1) + if aggregate_count != aggr_expr.len() || fields_set.len() != 1 { + return Ok(false); + } + if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs, input_schema)? { + return Ok(false); + } + Ok(true) +} + +/// Whether the rewrite is worth extending to a plan that only qualifies because +/// of the non-distinct `count`. +/// +/// The rewrite is not free: every other aggregate moves down to the inner group +/// by, which has a row per `(group, distinct value)` pair rather than per group, +/// and each one keeps its state at that finer grain. What pays for it is taking +/// the distinct aggregate off `GroupsAccumulatorAdapter`, whose one boxed +/// accumulator per group is the expensive shape. A distinct aggregate that +/// already has a specialized `GroupsAccumulator` never went near the adapter, so +/// there is nothing to buy and only the inner group by to pay for: ClickBench +/// Q22, whose `count(DISTINCT "UserID")` is over an `Int64`, measured a 132% +/// increase in peak memory pool reservation when the rewrite applied to it. +/// +/// The existing tolerance of a non-distinct `sum`, `min` or `max` predates this +/// and is left alone: narrowing it would change plans that have always been +/// rewritten, which no measurement here calls for. +fn rewrite_pays_for_count( + distinct_aggs: &[(&Arc, &Vec)], + input_schema: &DFSchema, +) -> Result { + for (func, args) in distinct_aggs { + let arg_types = args + .iter() + .map(|arg| arg.get_type(input_schema)) + .collect::>>()?; + if !func.groups_accumulator_supported_for_types(&arg_types, true) { + return Ok(true); + } + } + Ok(false) } /// Check if the first expr is [Expr::GroupingSet]. @@ -193,8 +241,11 @@ impl OptimizerRule for SingleDistinctToGroupBy { schema, group_expr, .. - }) if is_single_distinct_agg(&aggr_expr, count_rollup.as_ref())? - && !contains_grouping_set(&group_expr) => + }) if is_single_distinct_agg( + &aggr_expr, + input.schema(), + count_rollup.as_ref(), + )? && !contains_grouping_set(&group_expr) => { let group_size = group_expr.len(); // alias all original group_by exprs @@ -387,13 +438,17 @@ mod tests { use crate::assert_optimized_plan_eq_display_indent_snapshot; use crate::test::*; use crate::{Optimizer, OptimizerContext}; + use arrow::datatypes::{DataType, Field, Schema}; use chrono::{DateTime, Utc}; use datafusion_common::alias::AliasGenerator; use datafusion_common::config::ConfigOptions; use datafusion_expr::ExprFunctionExt; use datafusion_expr::expr::GroupingSet; use datafusion_expr::registry::{FunctionRegistry, MemoryFunctionRegistry}; - use datafusion_expr::{lit, logical_plan::builder::LogicalPlanBuilder}; + use datafusion_expr::{ + lit, + logical_plan::builder::{LogicalPlanBuilder, table_scan}, + }; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::expr_fn::{count, count_distinct, max, min, sum}; use datafusion_functions_aggregate::min_max::max_udaf; @@ -421,6 +476,20 @@ mod tests { )) } + /// `test` with a `Utf8` `b`, the column the distinct aggregate reads. + /// + /// `count(DISTINCT b)` over a string has no specialized + /// `GroupsAccumulator`, so it is the case a non-distinct `count` may join. + /// The `UInt32` `b` of [`test_table_scan`] is the case it may not. + fn test_table_scan_utf8_b() -> Result { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Utf8, false), + Field::new("c", DataType::UInt32, false), + ]); + table_scan(Some("test"), &schema, None)?.build() + } + /// An [`OptimizerConfig`] that resolves functions the way a session does. /// The rule needs `count` and `sum` from the registry to rewrite a /// non-distinct `count`. @@ -667,7 +736,7 @@ mod tests { #[test] fn distinct_and_common() -> Result<()> { - let table_scan = test_table_scan()?; + let table_scan = test_table_scan_utf8_b()?; let plan = LogicalPlanBuilder::from(table_scan) .aggregate( @@ -683,7 +752,55 @@ mod tests { @r" Projection: test.a, count(alias1) AS count(DISTINCT test.b), CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(test.c) [a:UInt32, count(DISTINCT test.b):Int64, count(test.c):Int64] Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias2)]] [a:UInt32, count(alias1):Int64, sum(alias2):Int64;N] - Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:Int64] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(test.c) AS alias2]] [a:UInt32, alias1:Utf8, alias2:Int64] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] + " + ) + } + + #[test] + fn distinct_and_common_over_a_natively_supported_type() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![count_distinct(col("b")), count(col("c"))], + )? + .build()?; + + // Should not work: `count(DISTINCT b)` over a `UInt32` has its own + // `GroupsAccumulator`, so the rewrite would add an inner group-by + // without taking anything off `GroupsAccumulatorAdapter` + assert_optimized_plan_equal!( + plan, + @r" + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b), count(test.c)]] [a:UInt32, count(DISTINCT test.b):Int64, count(test.c):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + + #[test] + fn distinct_over_a_natively_supported_type_without_a_count() -> Result<()> { + let table_scan = test_table_scan()?; + + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate( + vec![col("a")], + vec![count_distinct(col("b")), sum(col("c"))], + )? + .build()?; + + // Should work: the gate covers only the `count` this change added, so a + // plan that already qualified through `sum`, `min` or `max` is rewritten + // exactly as before + assert_optimized_plan_equal!( + plan, + @r" + 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] + Aggregate: groupBy=[[test.a]], aggr=[[count(alias1), sum(alias2)]] [a:UInt32, count(alias1):Int64, sum(alias2):UInt64;N] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[sum(test.c) AS alias2]] [a:UInt32, alias1:UInt32, alias2:UInt64;N] TableScan: test [a:UInt32, b:UInt32, c:UInt32] " ) @@ -911,7 +1028,7 @@ mod tests { #[test] fn count_star_and_distinct_without_groupby() -> Result<()> { - let table_scan = test_table_scan()?; + let table_scan = test_table_scan_utf8_b()?; let plan = LogicalPlanBuilder::from(table_scan) .aggregate( @@ -928,15 +1045,15 @@ mod tests { @r" Projection: CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(Int64(1)), count(alias1) AS count(DISTINCT test.b) [count(Int64(1)):Int64, count(DISTINCT test.b):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:UInt32, alias2:Int64] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + 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()?; + let table_scan = test_table_scan_utf8_b()?; let plan = LogicalPlanBuilder::from(table_scan) .aggregate( @@ -958,8 +1075,8 @@ mod tests { @r" Projection: test.a, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(Int64(1)), count(alias1) AS count(DISTINCT test.b), min(alias3) AS min(test.c), max(alias4) AS max(test.c), sum(alias5) AS sum(test.c) [a:UInt32, count(Int64(1)):Int64, count(DISTINCT test.b):Int64, min(test.c):UInt32;N, max(test.c):UInt32;N, sum(test.c):UInt64;N] Aggregate: groupBy=[[test.a]], aggr=[[sum(alias2), count(alias1), min(alias3), max(alias4), sum(alias5)]] [a:UInt32, sum(alias2):Int64;N, count(alias1):Int64, min(alias3):UInt32;N, max(alias4):UInt32;N, sum(alias5):UInt64;N] - Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(Int64(1)) AS alias2, min(test.c) AS alias3, max(test.c) AS alias4, sum(test.c) AS alias5]] [a:UInt32, alias1:UInt32, alias2:Int64, alias3:UInt32;N, alias4:UInt32;N, alias5:UInt64;N] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + Aggregate: groupBy=[[test.a, test.b AS alias1]], aggr=[[count(Int64(1)) AS alias2, min(test.c) AS alias3, max(test.c) AS alias4, sum(test.c) AS alias5]] [a:UInt32, alias1:Utf8, alias2:Int64, alias3:UInt32;N, alias4:UInt32;N, alias5:UInt64;N] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] " ) } diff --git a/datafusion/sqllogictest/test_files/clickbench.slt b/datafusion/sqllogictest/test_files/clickbench.slt index f9e9637fad496..4a1ef833c91db 100644 --- a/datafusion/sqllogictest/test_files/clickbench.slt +++ b/datafusion/sqllogictest/test_files/clickbench.slt @@ -616,25 +616,21 @@ EXPLAIN SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*) AS c, COUNT(DI ---- logical_plan 01)Sort: c DESC NULLS FIRST, fetch=10 -02)--Projection: hits.SearchPhrase, min(alias2) AS min(hits.URL), min(alias3) AS min(hits.Title), CASE WHEN sum(alias4) IS NOT NULL THEN sum(alias4) ELSE Int64(0) END AS c, count(alias1) AS count(DISTINCT hits.UserID) -03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[min(alias2), min(alias3), sum(alias4), count(alias1)]] -04)------Aggregate: groupBy=[[hits.SearchPhrase, hits.UserID AS alias1]], aggr=[[min(hits.URL) AS alias2, min(hits.Title) AS alias3, count(Int64(1)) AS alias4]] -05)--------SubqueryAlias: hits -06)----------Filter: hits_raw.SearchPhrase != Utf8View("") AND hits_raw.Title LIKE Utf8View("%Google%") AND hits_raw.URL NOT LIKE Utf8View("%.google.%") -07)------------TableScan: hits_raw projection=[Title, UserID, URL, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View(""), hits_raw.Title LIKE Utf8View("%Google%"), hits_raw.URL NOT LIKE Utf8View("%.google.%")] +02)--Projection: hits.SearchPhrase, min(hits.URL), min(hits.Title), count(Int64(1)) AS count(*) AS c, count(DISTINCT hits.UserID) +03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------Filter: hits_raw.SearchPhrase != Utf8View("") AND hits_raw.Title LIKE Utf8View("%Google%") AND hits_raw.URL NOT LIKE Utf8View("%.google.%") +06)----------TableScan: hits_raw projection=[Title, UserID, URL, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View(""), hits_raw.Title LIKE Utf8View("%Google%"), hits_raw.URL NOT LIKE Utf8View("%.google.%")] physical_plan 01)SortPreservingMergeExec: [c@3 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@3 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(alias2)@1 as min(hits.URL), min(alias3)@2 as min(hits.Title), CASE WHEN sum(alias4)@3 IS NOT NULL THEN sum(alias4)@3 ELSE 0 END as c, count(alias1)@4 as count(DISTINCT hits.UserID)] -04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(alias2), min(alias3), sum(alias4), count(alias1)] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(hits.URL)@1 as min(hits.URL), min(hits.Title)@2 as min(hits.Title), count(Int64(1))@3 as c, count(DISTINCT hits.UserID)@4 as count(DISTINCT hits.UserID)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@3 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(alias2), min(alias3), sum(alias4), count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase, alias1@1 as alias1], aggr=[min(hits.URL) as alias2, min(hits.Title) as alias3, count(1) as alias4] -08)--------------RepartitionExec: partitioning=Hash([SearchPhrase@0, alias1@1], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[SearchPhrase@3 as SearchPhrase, UserID@1 as alias1], aggr=[min(hits.URL) as alias2, min(hits.Title) as alias3, count(1) as alias4] -10)------------------FilterExec: SearchPhrase@3 != AND Title@0 LIKE %Google% AND URL@2 NOT LIKE %.google.% -11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[Title, UserID, URL, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND Title@2 LIKE %Google% AND URL@13 NOT LIKE %.google.%, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] +06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@3 as SearchPhrase], aggr=[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)] +07)------------FilterExec: SearchPhrase@3 != AND Title@0 LIKE %Google% AND URL@2 NOT LIKE %.google.% +08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[Title, UserID, URL, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND Title@2 LIKE %Google% AND URL@13 NOT LIKE %.google.%, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query TTTII SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*) AS c, COUNT(DISTINCT "UserID") FROM hits WHERE "Title" LIKE '%Google%' AND "URL" NOT LIKE '%.google.%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; diff --git a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt index d7c47a8626760..ffecec7b4ba23 100644 --- a/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt +++ b/datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt @@ -20,25 +20,30 @@ # asserted twice: once with the optimizer disabled # (`datafusion.optimizer.max_passes = 0`, so the aggregate runs as written) and # once with it enabled. The two must agree. +# +# A non-distinct `count` may only join the distinct aggregate when the distinct +# argument has no specialized `GroupsAccumulator` of its own. `x` is the +# `VARCHAR` that qualifies, `xi` the `INT` that does not; both hold the same +# values so the two sides can be compared directly. statement ok -CREATE TABLE t(g INT, x INT, v INT); +CREATE TABLE t(g INT, x VARCHAR, xi INT, v INT); # g = 1: four rows, two distinct x, one NULL x, one NULL v # g = 2: two rows, x is NULL throughout # g = 3: a single row statement ok INSERT INTO t VALUES - (1, 10, 100), - (1, 10, NULL), - (1, NULL, 5), - (1, 20, 7), - (2, NULL, NULL), - (2, NULL, 3), - (3, 30, 1); + (1, '10', 10, 100), + (1, '10', 10, NULL), + (1, NULL, NULL, 5), + (1, '20', 20, 7), + (2, NULL, NULL, NULL), + (2, NULL, NULL, 3), + (3, '30', 30, 1); statement ok -CREATE TABLE empty_t(g INT, x INT, v INT); +CREATE TABLE empty_t(g INT, x VARCHAR, xi INT, v INT); statement ok set datafusion.explain.logical_plan_only = true; @@ -88,6 +93,33 @@ logical_plan 02)--Aggregate: groupBy=[[t.g]], aggr=[[avg(CAST(t.v AS Float64)), count(DISTINCT t.x)]] 03)----TableScan: t projection=[g, x, v] +######## +# The other side of the gate: a distinct argument that has its own +# `GroupsAccumulator` keeps the count out +######## + +# `count(DISTINCT xi)` over an INT is served by +# `PrimitiveDistinctCountGroupsAccumulator`, so the rewrite would add an inner +# group by without taking anything off `GroupsAccumulatorAdapter` +query TT +EXPLAIN SELECT g, count(*) AS records, count(DISTINCT xi) AS distinct_xi FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, count(Int64(1)) AS count(*) AS records, count(DISTINCT t.xi) AS distinct_xi +02)--Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), count(DISTINCT t.xi)]] +03)----TableScan: t projection=[g, xi] + +# The gate covers only the count. A plan that already qualified through sum, +# min or max is still rewritten over the same INT column +query TT +EXPLAIN SELECT g, sum(v) AS sum_v, count(DISTINCT xi) AS distinct_xi FROM t GROUP BY g; +---- +logical_plan +01)Projection: t.g, sum(alias2) AS sum_v, count(alias1) AS distinct_xi +02)--Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]] +03)----Aggregate: groupBy=[[t.g, t.xi AS alias1]], aggr=[[sum(CAST(t.v AS Int64)) AS alias2]] +04)------TableScan: t projection=[g, xi, v] + statement ok set datafusion.explain.logical_plan_only = false; @@ -265,10 +297,10 @@ ORDER BY count(*) DESC, g; ######## statement ok -CREATE TABLE failed(x INT); +CREATE TABLE failed(x VARCHAR); statement ok -INSERT INTO failed VALUES (10), (30); +INSERT INTO failed VALUES ('10'), ('30'); statement ok set datafusion.optimizer.max_passes = 0; @@ -294,6 +326,80 @@ GROUP BY t.g ORDER BY t.g; 1 2 1 100 100 100 3 1 1 1 1 1 +######## +# Both sides of the gate produce the same results as the unoptimized plan +######## + +statement ok +set datafusion.optimizer.max_passes = 0; + +# gated: the count keeps the INT distinct unrewritten +query IIIII +SELECT g, count(*) AS records, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten: the same shape over the VARCHAR distinct +query IIIII +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten through sum alone, over the INT distinct the gate excludes +query III +SELECT g, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v +FROM t GROUP BY g ORDER BY g; +---- +1 2 112 +2 0 3 +3 1 1 + +query II +SELECT count(*) AS records, count(DISTINCT xi) AS distinct_xi FROM empty_t; +---- +0 0 + +statement ok +set datafusion.optimizer.max_passes = 3; + +# gated: the count keeps the INT distinct unrewritten +query IIIII +SELECT g, count(*) AS records, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten: the same shape over the VARCHAR distinct +query IIIII +SELECT g, count(*) AS records, count(DISTINCT x) AS distinct_x, sum(v) AS sum_v, max(v) AS max_v +FROM t GROUP BY g ORDER BY g; +---- +1 4 2 112 100 +2 2 0 3 3 +3 1 1 1 1 + +# rewritten through sum alone, over the INT distinct the gate excludes +query III +SELECT g, count(DISTINCT xi) AS distinct_xi, sum(v) AS sum_v +FROM t GROUP BY g ORDER BY g; +---- +1 2 112 +2 0 3 +3 1 1 + +query II +SELECT count(*) AS records, count(DISTINCT xi) AS distinct_xi FROM empty_t; +---- +0 0 + statement ok DROP TABLE failed; From ecd4e988517ba48318081a897b2e9126d2906e88 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:23:30 -0500 Subject: [PATCH 4/4] Do not link a private item from public documentation `cargo doc` runs with `-D warnings`, and a doc link from the public `SingleDistinctToGroupBy` to the private `rewrite_pays_for_count` is an error. Say the same thing in prose instead. --- datafusion/optimizer/src/single_distinct_to_groupby.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 34bf1cd039147..155f6b5a955fb 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -82,7 +82,9 @@ use datafusion_expr::{ /// rows is NULL where `count` is 0. /// /// That `count` is allowed only when the distinct argument has no specialized -/// `GroupsAccumulator` of its own. See [`rewrite_pays_for_count`]. +/// `GroupsAccumulator` of its own, so that the rewrite is taking the distinct +/// aggregate off `GroupsAccumulatorAdapter` rather than only adding an inner +/// group by. See `rewrite_pays_for_count` for the measurements behind that. #[derive(Default, Debug)] pub struct SingleDistinctToGroupBy {}