From 4ce31ec6f93d30bf9c0f917c65598444ddfc0fac Mon Sep 17 00:00:00 2001 From: hushen <190065939+918154429@users.noreply.github.com> Date: Sat, 3 Oct 2026 14:15:03 +0800 Subject: [PATCH 1/5] fix: keep grouped integer DISTINCT counts on the native accumulator path --- benchmarks/bench.sh | 16 ++ .../benchmarks/q01.benchmark | 25 +++ .../benchmarks/q02.benchmark | 25 +++ .../benchmarks/q03.benchmark | 25 +++ .../benchmarks/q04.benchmark | 25 +++ .../benchmarks/q05.benchmark | 25 +++ .../benchmarks/q06.benchmark | 25 +++ .../grouped_count_distinct.suite | 11 ++ .../grouped_count_distinct/init/cleanup.sql | 1 + .../grouped_count_distinct/init/load.sql | 9 + .../src/single_distinct_to_groupby.rs | 47 ++++- .../test_files/aggregate_memory_spill.slt | 2 +- .../test_files/aggregates_simplify.slt | 22 +-- .../sqllogictest/test_files/clickbench.slt | 108 +++++------ .../test_files/grouped_count_distinct.slt | 176 ++++++++++++++++++ datafusion/sqllogictest/test_files/joins.slt | 10 +- 16 files changed, 459 insertions(+), 93 deletions(-) create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q01.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q02.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q03.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q04.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q05.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q06.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/init/cleanup.sql create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql create mode 100644 datafusion/sqllogictest/test_files/grouped_count_distinct.slt diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 66b6dcd1b12b3..e3cdd041d2cc4 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -177,6 +177,7 @@ hj: Benchmark for simple hash joins, testing various join sc smj: Benchmark for simple sort merge joins, testing various join scenarios dict: Benchmark for dictionary-encoded group-by scenarios array_agg_distinct: 1000K-group, two-row-per-group array_agg(DISTINCT) benchmark +grouped_count_distinct: Grouped count(DISTINCT) by key type, group count and distinct cardinality compile_profile: Compile and execute TPC-H across selected Cargo profiles, reporting timing and binary size @@ -295,6 +296,9 @@ main() { # Data is generated inline by the suite's load SQL. echo "projection_subquery: no external data to generate" ;; + grouped_count_distinct) + echo "grouped_count_distinct: no external data to generate" + ;; asof_join) data_asof_join ;; @@ -697,6 +701,9 @@ main() { array_agg_distinct) run_array_agg_distinct ;; + grouped_count_distinct) + run_grouped_count_distinct + ;; compile_profile) run_compile_profile "${PROFILE_ARGS[@]}" ;; @@ -1831,6 +1838,15 @@ run_dict() { debug_run $CARGO_COMMAND --bin dfbench -- dict --iterations 5 -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG} } +# Runs grouped DISTINCT counts over inline data with varying key types and cardinalities. +run_grouped_count_distinct() { + debug_run env BENCH_NAME=grouped_count_distinct \ + BENCH_RESULTS_FILE="$(sql_results_file grouped_count_distinct)" \ + GROUPED_DISTINCT_ROWS="${GROUPED_DISTINCT_ROWS:-4000000}" \ + ${QUERY:+BENCH_QUERY="${QUERY}"} \ + bash -c "$SQL_CARGO_COMMAND" +} + # Runs the data-free high-cardinality array_agg(DISTINCT) SQL benchmark. run_array_agg_distinct() { echo "Running array_agg_distinct benchmark..." diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..ccfe7c584369f --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q01.benchmark @@ -0,0 +1,25 @@ +name Q01 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT g AS g, x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Int64, 2,000 groups, high distinct cardinality (issue #25936). +SELECT g, count(DISTINCT x) FROM grouped_distinct GROUP BY g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..fc4182415dea5 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q02.benchmark @@ -0,0 +1,25 @@ +name Q02 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT g AS g, count(DISTINCT low_x) AS n FROM grouped_distinct GROUP BY g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT g AS g, low_x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Int64, 2,000 groups, at most eight distinct values per group. +SELECT g, count(DISTINCT low_x) FROM grouped_distinct GROUP BY g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..583b3f81cd211 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q03.benchmark @@ -0,0 +1,25 @@ +name Q03 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT few_g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY few_g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT few_g AS g, x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Int64, eight groups, high distinct cardinality. +SELECT few_g, count(DISTINCT x) FROM grouped_distinct GROUP BY few_g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..1eb0af211804d --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q04.benchmark @@ -0,0 +1,25 @@ +name Q04 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT g AS g, count(DISTINCT unsigned_x) AS n FROM grouped_distinct GROUP BY g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT g AS g, unsigned_x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- UInt64, 2,000 groups, high distinct cardinality. +SELECT g, count(DISTINCT unsigned_x) FROM grouped_distinct GROUP BY g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..249d58f260b58 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q05.benchmark @@ -0,0 +1,25 @@ +name Q05 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT g AS g, count(DISTINCT string_x) AS n FROM grouped_distinct GROUP BY g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT g AS g, string_x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- VARCHAR control: keep the existing two-phase rewrite. +SELECT g, count(DISTINCT string_x) FROM grouped_distinct GROUP BY g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..c6850e58bd65a --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q06.benchmark @@ -0,0 +1,25 @@ +name Q06 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT g AS g, count(DISTINCT float_x) AS n FROM grouped_distinct GROUP BY g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT g AS g, float_x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Float64 control: no specialized DISTINCT GroupsAccumulator. +SELECT g, count(DISTINCT float_x) FROM grouped_distinct GROUP BY g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite b/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite new file mode 100644 index 0000000000000..33851a93840b5 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite @@ -0,0 +1,11 @@ +description = "Grouped COUNT(DISTINCT) across integer, string and floating-point keys, group counts and distinct cardinalities" + +[[options]] +name = "rows" +env = "GROUPED_DISTINCT_ROWS" +default = "4000000" +help = "Input rows generated outside the timed query." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 1" +description = "Compare grouped DISTINCT counts at the issue's four-million-row scale." diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/init/cleanup.sql b/benchmarks/sql_benchmarks/grouped_count_distinct/init/cleanup.sql new file mode 100644 index 0000000000000..234d5383711d4 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/init/cleanup.sql @@ -0,0 +1 @@ +DROP TABLE grouped_distinct; diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql b/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql new file mode 100644 index 0000000000000..64ea85a0931e4 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql @@ -0,0 +1,9 @@ +CREATE TABLE grouped_distinct AS +SELECT value % 2000 AS g, + value % 8 AS few_g, + (value * 48271) % 999983 AS x, + (value / 2000) % 8 AS low_x, + arrow_cast((value * 48271) % 999983, 'UInt64') AS unsigned_x, + CAST((value * 48271) % 999983 AS VARCHAR) AS string_x, + CAST((value * 48271) % 999983 AS DOUBLE) AS float_x +FROM range(${GROUPED_DISTINCT_ROWS:-4000000}); diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index f663003a9e4a4..c122aa3c88541 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -125,6 +125,7 @@ fn is_single_distinct_agg( aggr_expr: &[Expr], input_schema: &DFSchema, count_rollup: Option<&CountRollup>, + has_group_by: bool, ) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; @@ -167,12 +168,39 @@ fn is_single_distinct_agg( if aggregate_count != aggr_expr.len() || fields_set.len() != 1 { return Ok(false); } + // Grouped DISTINCT aggregates with native GroupsAccumulators do not need + // the adapter this rewrite avoids. Keep them on the direct path instead of + // building an extra aggregate with a row per (group, distinct value) pair. + // Leave mixed aggregates and global DISTINCT aggregation unchanged. + if has_group_by + && distinct_aggs.len() == aggregate_count + && all_have_groups_accumulators(&distinct_aggs, input_schema)? + { + return Ok(false); + } if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs, input_schema)? { return Ok(false); } Ok(true) } +fn all_have_groups_accumulators( + distinct_aggs: &[(&Arc, &[Expr])], + input_schema: &DFSchema, +) -> Result { + for (func, args) in distinct_aggs { + let arg_types = args + .iter() + .map(|arg| arg.get_type(input_schema)) + .collect::>>()?; + // An unknown answer must preserve the existing rewrite. + if func.groups_accumulator_supported_for_types(&arg_types, true) != Some(true) { + return Ok(false); + } + } + Ok(true) +} + /// Whether the rewrite is worth extending to a plan that only qualifies because /// of the non-distinct `count`. /// @@ -248,6 +276,7 @@ impl OptimizerRule for SingleDistinctToGroupBy { &aggr_expr, input.schema(), count_rollup.as_ref(), + !group_expr.is_empty(), )? && !contains_grouping_set(&group_expr) => { let group_size = group_expr.len(); @@ -737,14 +766,12 @@ mod tests { .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? .build()?; - // Should work + // The integer DISTINCT count already has a GroupsAccumulator. assert_optimized_plan_equal!( plan, @r" - Projection: test.a, count(alias1) AS count(DISTINCT test.b) [a:UInt32, count(DISTINCT test.b):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] + Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b)]] [a:UInt32, count(DISTINCT test.b):Int64] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] " ) } @@ -914,20 +941,20 @@ mod tests { #[test] fn group_by_with_expr() -> Result<()> { - let table_scan = test_table_scan().unwrap(); + let table_scan = test_table_scan_utf8_b()?; let plan = LogicalPlanBuilder::from(table_scan) - .aggregate(vec![col("a") + lit(1)], vec![count_distinct(col("c"))])? + .aggregate(vec![col("a") + lit(1)], vec![count_distinct(col("b"))])? .build()?; // Should work assert_optimized_plan_equal!( plan, @r" - 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] + Projection: group_alias_0 AS test.a + Int32(1), count(alias1) AS count(DISTINCT test.b) [test.a + Int32(1):Int64, count(DISTINCT test.b):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] + Aggregate: groupBy=[[test.a + Int32(1) AS group_alias_0, test.b AS alias1]], aggr=[[]] [group_alias_0:Int64, alias1:Utf8] + TableScan: test [a:UInt32, b:Utf8, c:UInt32] " ) } diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 48622127bce99..c03b9fff5a836 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -113,7 +113,7 @@ FROM ( ) ---- -04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[spill_count=7,] +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[count(DISTINCT t.v)], metrics=[spill_count=7,] # --- Case D: multiple aggregates (sum/min/max) under memory limit --- diff --git a/datafusion/sqllogictest/test_files/aggregates_simplify.slt b/datafusion/sqllogictest/test_files/aggregates_simplify.slt index 4169707913fb2..00cfc3956a232 100644 --- a/datafusion/sqllogictest/test_files/aggregates_simplify.slt +++ b/datafusion/sqllogictest/test_files/aggregates_simplify.slt @@ -543,23 +543,17 @@ query TT EXPLAIN SELECT g, count(DISTINCT v) FROM distinct_simplify_t GROUP BY g; ---- logical_plan -01)Projection: distinct_simplify_t.g, count(alias1) AS count(DISTINCT distinct_simplify_t.v) -02)--Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[count(alias1)]] -03)----Aggregate: groupBy=[[distinct_simplify_t.g, distinct_simplify_t.v AS alias1]], aggr=[[]] -04)------TableScan: distinct_simplify_t projection=[g, v] +01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[count(DISTINCT distinct_simplify_t.v)]] +02)--TableScan: distinct_simplify_t projection=[g, v] physical_plan -01)ProjectionExec: expr=[g@0 as g, count(alias1)@1 as count(DISTINCT distinct_simplify_t.v)] -02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(alias1)] -03)----RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=4 -04)------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(alias1)] -05)--------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, alias1@1 as alias1], aggr=[] -06)----------RepartitionExec: partitioning=Hash([g@0, alias1@1], 4), input_partitions=1 -07)------------AggregateExec: mode=Partial, gby=[g@0 as g, v@1 as alias1], aggr=[] -08)--------------DataSourceExec: partitions=1, partition_sizes=[1] +01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(DISTINCT distinct_simplify_t.v)] +02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(DISTINCT distinct_simplify_t.v)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] # Mixed: a node that still needs one DISTINCT keeps all of them, so this rule -# cannot change which plans SingleDistinctToGroupBy rewrites. `count` needs the -# rewrite, and the plan below shows that it still happens. +# leaves min(DISTINCT) in place. The mixed min/count aggregate still gets the +# existing SingleDistinctToGroupBy rewrite. query III SELECT g, min(DISTINCT v), count(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; ---- diff --git a/datafusion/sqllogictest/test_files/clickbench.slt b/datafusion/sqllogictest/test_files/clickbench.slt index 7cb5547383c38..175dc01ba6d8d 100644 --- a/datafusion/sqllogictest/test_files/clickbench.slt +++ b/datafusion/sqllogictest/test_files/clickbench.slt @@ -232,22 +232,18 @@ EXPLAIN SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "Reg ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.RegionID, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.RegionID]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.RegionID, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------TableScan: hits_raw projection=[RegionID, UserID] +02)--Projection: hits.RegionID, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.RegionID]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------TableScan: hits_raw projection=[RegionID, UserID] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[RegionID@0 as RegionID, count(alias1)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] -05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID, alias1@1 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([RegionID@0, alias1@1], 4), input_partitions=1 -09)----------------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID, UserID@1 as alias1], aggr=[] -10)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[RegionID, UserID], file_type=parquet +02)--ProjectionExec: expr=[RegionID@0 as RegionID, count(DISTINCT hits.UserID)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[count(DISTINCT hits.UserID)] +05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=1 +06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[count(DISTINCT hits.UserID)] +07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[RegionID, UserID], file_type=parquet query II rowsort SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "RegionID" ORDER BY u DESC LIMIT 10; @@ -290,25 +286,21 @@ EXPLAIN SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.MobilePhoneModel, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.MobilePhoneModel]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.MobilePhoneModel, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------Filter: hits_raw.MobilePhoneModel != Utf8View("") -07)------------TableScan: hits_raw projection=[UserID, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] +02)--Projection: hits.MobilePhoneModel, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.MobilePhoneModel]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------Filter: hits_raw.MobilePhoneModel != Utf8View("") +06)----------TableScan: hits_raw projection=[UserID, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(alias1)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] +02)--ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(DISTINCT hits.UserID)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([MobilePhoneModel@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel, alias1@1 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([MobilePhoneModel@0, alias1@1], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[MobilePhoneModel@1 as MobilePhoneModel, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: MobilePhoneModel@1 != -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=[UserID, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] +06)----------AggregateExec: mode=Partial, gby=[MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] +07)------------FilterExec: MobilePhoneModel@1 != +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=[UserID, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] query TI SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhoneModel" ORDER BY u DESC LIMIT 10; @@ -320,25 +312,21 @@ EXPLAIN SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.MobilePhone, hits.MobilePhoneModel, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------Filter: hits_raw.MobilePhoneModel != Utf8View("") -07)------------TableScan: hits_raw projection=[UserID, MobilePhone, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] +02)--Projection: hits.MobilePhone, hits.MobilePhoneModel, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------Filter: hits_raw.MobilePhoneModel != Utf8View("") +06)----------TableScan: hits_raw projection=[UserID, MobilePhone, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@2 DESC], fetch=10 -02)--ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(alias1)@2 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@2 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] +02)--ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(DISTINCT hits.UserID)@2 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@2 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, alias1@2 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1, alias1@2], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[MobilePhone@1 as MobilePhone, MobilePhoneModel@2 as MobilePhoneModel, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: MobilePhoneModel@2 != -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=[UserID, MobilePhone, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] +06)----------AggregateExec: mode=Partial, gby=[MobilePhone@1 as MobilePhone, MobilePhoneModel@2 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] +07)------------FilterExec: MobilePhoneModel@2 != +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=[UserID, MobilePhone, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] query ITI SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhone", "MobilePhoneModel" ORDER BY u DESC LIMIT 10; @@ -376,25 +364,21 @@ EXPLAIN SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "Se ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.SearchPhrase, count(alias1) AS u -03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[hits.SearchPhrase, hits.UserID AS alias1]], aggr=[[]] -05)--------SubqueryAlias: hits -06)----------Filter: hits_raw.SearchPhrase != Utf8View("") -07)------------TableScan: hits_raw projection=[UserID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] +02)--Projection: hits.SearchPhrase, count(DISTINCT hits.UserID) AS u +03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[count(DISTINCT hits.UserID)]] +04)------SubqueryAlias: hits +05)--------Filter: hits_raw.SearchPhrase != Utf8View("") +06)----------TableScan: hits_raw projection=[UserID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(alias1)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(alias1)] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(DISTINCT hits.UserID)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase, alias1@1 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([SearchPhrase@0, alias1@1], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[SearchPhrase@1 as SearchPhrase, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: SearchPhrase@1 != -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=[UserID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , 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@1 as SearchPhrase], aggr=[count(DISTINCT hits.UserID)] +07)------------FilterExec: SearchPhrase@1 != +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=[UserID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query TI SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY u DESC LIMIT 10; diff --git a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt new file mode 100644 index 0000000000000..0b449b723e911 --- /dev/null +++ b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt @@ -0,0 +1,176 @@ +# 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. + +# Grouped integer DISTINCT counts should keep their specialized accumulator. +# String and floating-point keys still use the two-phase rewrite. + +statement ok +CREATE TABLE numeric_counts(g INT, x INT, s VARCHAR, f DOUBLE); + +statement ok +INSERT INTO numeric_counts VALUES + (1, 1, '1', 1.0), + (1, 1, '1', 1.0), + (1, 2, '2', 2.0), + (1, NULL, NULL, NULL), + (2, NULL, NULL, NULL), + (3, 3, '3', 3.0); + +statement ok +SET datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(DISTINCT numeric_counts.x)]] +02)--TableScan: numeric_counts projection=[g, x] + +query TT +EXPLAIN SELECT g, count(DISTINCT s) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Projection: numeric_counts.g, count(alias1) AS count(DISTINCT numeric_counts.s) +02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.s AS alias1]], aggr=[[]] +04)------TableScan: numeric_counts projection=[g, s] + +query TT +EXPLAIN SELECT g, count(DISTINCT f) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Projection: numeric_counts.g, count(alias1) AS count(DISTINCT numeric_counts.f) +02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.f AS alias1]], aggr=[[]] +04)------TableScan: numeric_counts projection=[g, f] + +# The global aggregate does not use a GroupsAccumulator. +query TT +EXPLAIN SELECT count(DISTINCT x) FROM numeric_counts; +---- +logical_plan +01)Projection: count(alias1) AS count(DISTINCT numeric_counts.x) +02)--Aggregate: groupBy=[[]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[numeric_counts.x AS alias1]], aggr=[[]] +04)------TableScan: numeric_counts projection=[x] + +statement ok +SET datafusion.explain.logical_plan_only = false; + +# Specialized accumulator for Int8; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'Int8')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +# Specialized accumulator for Int16; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'Int16')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +# Specialized accumulator for Int32; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'Int32')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +# Specialized accumulator for Int64; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'Int64')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +# Specialized accumulator for UInt8; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'UInt8')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +# Specialized accumulator for UInt16; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'UInt16')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +# Specialized accumulator for UInt32; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'UInt32')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +# Specialized accumulator for UInt64; duplicate and all-NULL groups. +query II rowsort +SELECT g, count(DISTINCT arrow_cast(x, 'UInt64')) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +query III rowsort +SELECT g, count(DISTINCT x), count(DISTINCT x) AS again FROM numeric_counts GROUP BY g; +---- +1 2 2 +2 0 0 +3 1 1 + +query II rowsort +SELECT g + 1, count(DISTINCT x) FROM numeric_counts GROUP BY g + 1; +---- +2 2 +3 0 +4 1 + +query IIII rowsort +SELECT g, count(DISTINCT x), sum(x), count(*) FROM numeric_counts GROUP BY g; +---- +1 2 4 4 +2 0 NULL 1 +3 1 3 1 + +query III rowsort +SELECT g, count(DISTINCT s), count(DISTINCT f) FROM numeric_counts GROUP BY g; +---- +1 2 2 +2 0 0 +3 1 1 + +query I +SELECT count(DISTINCT x) FROM numeric_counts WHERE g = 42; +---- +0 + +query II +SELECT g, count(DISTINCT x) FROM numeric_counts WHERE g = 42 GROUP BY g; +---- + +statement ok +DROP TABLE numeric_counts; diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index f4dbd212eb8ee..bd9ab874ea083 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -1357,12 +1357,10 @@ inner join join_t2 on join_t1.t1_id = join_t2.t2_id group by join_t1.t1_id ---- logical_plan -01)Projection: join_t1.t1_id, count(alias1) AS count(DISTINCT join_t1.t1_int) -02)--Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(alias1)]] -03)----Aggregate: groupBy=[[join_t1.t1_id, join_t1.t1_int AS alias1]], aggr=[[]] -04)------LeftSemi Join: join_t1.t1_id = join_t2.t2_id -05)--------TableScan: join_t1 projection=[t1_id, t1_int] -06)--------TableScan: join_t2 projection=[t2_id] +01)Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(DISTINCT join_t1.t1_int)]] +02)--LeftSemi Join: join_t1.t1_id = join_t2.t2_id +03)----TableScan: join_t1 projection=[t1_id, t1_int] +04)----TableScan: join_t2 projection=[t2_id] # A similar query with two DISTINCT aggregates is also rewritten: each # `count(DISTINCT ...)` removes its own duplicates, so the join's duplicates From 3ae59629094ff810a0fcd1f77d666eaab53b6d4d Mon Sep 17 00:00:00 2001 From: hushen <190065939+918154429@users.noreply.github.com> Date: Sat, 3 Oct 2026 23:06:17 +0800 Subject: [PATCH 2/5] test: update TPCH Q16 plan for native distinct count --- .../test_files/tpch/plans/q16.slt.part | 58 +++++++++---------- 1 file changed, 27 insertions(+), 31 deletions(-) diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part index df7a26bf0a20d..5d85e3932e557 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part @@ -51,37 +51,33 @@ limit 10; ---- logical_plan 01)Sort: supplier_cnt DESC NULLS FIRST, part.p_brand ASC NULLS LAST, part.p_type ASC NULLS LAST, part.p_size ASC NULLS LAST, fetch=10 -02)--Projection: part.p_brand, part.p_type, part.p_size, count(alias1) AS supplier_cnt -03)----Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size]], aggr=[[count(alias1)]] -04)------Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size, partsupp.ps_suppkey AS alias1]], aggr=[[]] -05)--------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey null_aware -06)----------Projection: partsupp.ps_suppkey, part.p_brand, part.p_type, part.p_size -07)------------Inner Join: partsupp.ps_partkey = part.p_partkey -08)--------------TableScan: partsupp projection=[ps_partkey, ps_suppkey] -09)--------------Filter: part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) AND part.p_brand != Utf8View("Brand#45") AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") -10)----------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]), part.p_brand != Utf8View("Brand#45"), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%")] -11)----------SubqueryAlias: __correlated_sq_1 -12)------------Projection: supplier.s_suppkey -13)--------------Filter: supplier.s_comment LIKE Utf8View("%Customer%Complaints%") -14)----------------TableScan: supplier projection=[s_suppkey, s_comment], partial_filters=[supplier.s_comment LIKE Utf8View("%Customer%Complaints%")] +02)--Projection: part.p_brand, part.p_type, part.p_size, count(DISTINCT partsupp.ps_suppkey) AS supplier_cnt +03)----Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size]], aggr=[[count(DISTINCT partsupp.ps_suppkey)]] +04)------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey null_aware +05)--------Projection: partsupp.ps_suppkey, part.p_brand, part.p_type, part.p_size +06)----------Inner Join: partsupp.ps_partkey = part.p_partkey +07)------------TableScan: partsupp projection=[ps_partkey, ps_suppkey] +08)------------Filter: part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) AND part.p_brand != Utf8View("Brand#45") AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") +09)--------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]), part.p_brand != Utf8View("Brand#45"), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%")] +10)--------SubqueryAlias: __correlated_sq_1 +11)----------Projection: supplier.s_suppkey +12)------------Filter: supplier.s_comment LIKE Utf8View("%Customer%Complaints%") +13)--------------TableScan: supplier projection=[s_suppkey, s_comment], partial_filters=[supplier.s_comment LIKE Utf8View("%Customer%Complaints%")] physical_plan 01)SortPreservingMergeExec: [supplier_cnt@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], fetch=10 -02)--ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(alias1)@3 as supplier_cnt] -03)----SortExec: TopK(fetch=10), expr=[count(alias1)@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] +02)--ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(DISTINCT partsupp.ps_suppkey)@3 as supplier_cnt] +03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT partsupp.ps_suppkey)@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(DISTINCT partsupp.ps_suppkey)] 05)--------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] -07)------------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, alias1@3 as alias1], aggr=[] -08)--------------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2, alias1@3], 4), input_partitions=4 -09)----------------AggregateExec: mode=Partial, gby=[p_brand@1 as p_brand, p_type@2 as p_type, p_size@3 as p_size, ps_suppkey@0 as alias1], aggr=[] -10)------------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)], null_aware -11)--------------------CoalescePartitionsExec -12)----------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, p_partkey@0)], projection=[ps_suppkey@1, p_brand@3, p_type@4, p_size@5] -13)------------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 -14)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false -15)------------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 -16)--------------------------FilterExec: p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) AND p_brand@1 != Brand#45 AND p_type@2 NOT LIKE MEDIUM POLISHED% -17)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false -18)--------------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] -19)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -20)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +06)----------AggregateExec: mode=Partial, gby=[p_brand@1 as p_brand, p_type@2 as p_type, p_size@3 as p_size], aggr=[count(DISTINCT partsupp.ps_suppkey)] +07)------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)], null_aware +08)--------------CoalescePartitionsExec +09)----------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, p_partkey@0)], projection=[ps_suppkey@1, p_brand@3, p_type@4, p_size@5] +10)------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 +11)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false +12)------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 +13)--------------------FilterExec: p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) AND p_brand@1 != Brand#45 AND p_type@2 NOT LIKE MEDIUM POLISHED% +14)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +15)--------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] +16)----------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +17)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false From 2d74893df2df096301be506e410f0edece0a24c5 Mon Sep 17 00:00:00 2001 From: hushen <190065939+918154429@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:48:54 +0800 Subject: [PATCH 3/5] perf: limit native grouped DISTINCT to one target partition --- .../benchmarks/q07.benchmark | 25 ++++ .../benchmarks/q08.benchmark | 25 ++++ .../benchmarks/q09.benchmark | 27 +++++ .../grouped_count_distinct.suite | 6 +- .../grouped_count_distinct/init/load.sql | 4 +- .../src/single_distinct_to_groupby.rs | 40 ++++--- .../test_files/aggregates_simplify.slt | 22 ++-- .../sqllogictest/test_files/clickbench.slt | 108 ++++++++++-------- .../test_files/grouped_count_distinct.slt | 66 ++++++++++- datafusion/sqllogictest/test_files/joins.slt | 10 +- .../test_files/tpch/plans/q16.slt.part | 58 +++++----- 11 files changed, 286 insertions(+), 105 deletions(-) create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q07.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q08.benchmark create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q09.benchmark diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..da4691485e161 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q07.benchmark @@ -0,0 +1,25 @@ +name Q07 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT skew_g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY skew_g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT skew_g AS g, x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Int64 argument, eight integer groups; at least 90% of rows in one hot group. +SELECT skew_g, count(DISTINCT x) FROM grouped_distinct GROUP BY skew_g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..079758b046746 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q08.benchmark @@ -0,0 +1,25 @@ +name Q08 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT string_g AS g, count(DISTINCT x) AS n FROM grouped_distinct GROUP BY string_g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT string_g AS g, x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Int64 argument, eight string groups with a hot group (ClickBench Q10 shape). +SELECT string_g, count(DISTINCT x) FROM grouped_distinct GROUP BY string_g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..b4dee9b031e8f --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q09.benchmark @@ -0,0 +1,27 @@ +name Q09 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT skew_g AS g, string_g AS s, count(DISTINCT x) AS n + FROM grouped_distinct GROUP BY skew_g, string_g +), expected AS ( + SELECT g, s, count(x) AS n + FROM (SELECT DISTINCT skew_g AS g, string_g AS s, x AS x FROM grouped_distinct) + GROUP BY g, s +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Int64 argument, hot composite integer/string grouping (ClickBench Q11 shape). +SELECT skew_g, string_g, count(DISTINCT x) +FROM grouped_distinct GROUP BY skew_g, string_g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite b/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite index 33851a93840b5..725cdb058f4a6 100644 --- a/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite @@ -1,4 +1,4 @@ -description = "Grouped COUNT(DISTINCT) across integer, string and floating-point keys, group counts and distinct cardinalities" +description = "Grouped COUNT(DISTINCT) across argument types, integer/string/composite grouping keys, distinct cardinalities and skew; compare --partitions 1, 4 and 12" [[options]] name = "rows" @@ -9,3 +9,7 @@ help = "Input rows generated outside the timed query." [[examples]] command = "cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 1" description = "Compare grouped DISTINCT counts at the issue's four-million-row scale." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 12" +description = "Exercise parallel state merging, including hot integer, string and composite groups (Q07-Q09)." diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql b/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql index 64ea85a0931e4..efdbbf1a9e31b 100644 --- a/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql @@ -5,5 +5,7 @@ SELECT value % 2000 AS g, (value / 2000) % 8 AS low_x, arrow_cast((value * 48271) % 999983, 'UInt64') AS unsigned_x, CAST((value * 48271) % 999983 AS VARCHAR) AS string_x, - CAST((value * 48271) % 999983 AS DOUBLE) AS float_x + CAST((value * 48271) % 999983 AS DOUBLE) AS float_x, + CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS skew_g, + CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR) AS string_g FROM range(${GROUPED_DISTINCT_ROWS:-4000000}); diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index c122aa3c88541..0a38d74c14f80 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -125,7 +125,7 @@ fn is_single_distinct_agg( aggr_expr: &[Expr], input_schema: &DFSchema, count_rollup: Option<&CountRollup>, - has_group_by: bool, + use_native_grouped_distinct: bool, ) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; @@ -168,11 +168,13 @@ fn is_single_distinct_agg( if aggregate_count != aggr_expr.len() || fields_set.len() != 1 { return Ok(false); } - // Grouped DISTINCT aggregates with native GroupsAccumulators do not need - // the adapter this rewrite avoids. Keep them on the direct path instead of - // building an extra aggregate with a row per (group, distinct value) pair. - // Leave mixed aggregates and global DISTINCT aggregation unchanged. - if has_group_by + // With one target partition, a native GroupsAccumulator avoids building an + // extra aggregate with a row per (group, distinct value) pair. With multiple + // partitions, the rewrite also distributes deduplication by both keys; + // keeping DISTINCT instead repartitions lists of values by group and merges + // them in the final accumulator. That can concentrate work on a few groups. + // Native accumulator support alone does not establish that this is cheaper. + if use_native_grouped_distinct && distinct_aggs.len() == aggregate_count && all_have_groups_accumulators(&distinct_aggs, input_schema)? { @@ -276,7 +278,8 @@ impl OptimizerRule for SingleDistinctToGroupBy { &aggr_expr, input.schema(), count_rollup.as_ref(), - !group_expr.is_empty(), + !group_expr.is_empty() + && config.options().execution.target_partitions == 1, )? && !contains_grouping_set(&group_expr) => { let group_size = group_expr.len(); @@ -592,8 +595,11 @@ mod tests { let mut registry = MemoryFunctionRegistry::new(); registry.register_udaf(count_udaf()).unwrap(); registry.register_udaf(sum_udaf()).unwrap(); + // Keep logical-plan snapshots independent of the available CPUs. + let mut options = OptimizerContext::new().options(); + Arc::make_mut(&mut options).execution.target_partitions = 4; Self { - inner: OptimizerContext::new(), + inner: OptimizerContext::new_with_config_options(options), registry, } } @@ -766,12 +772,14 @@ mod tests { .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? .build()?; - // The integer DISTINCT count already has a GroupsAccumulator. + // Multiple target partitions keep deduplication distributed by both keys. assert_optimized_plan_equal!( plan, @r" - Aggregate: groupBy=[[test.a]], aggr=[[count(DISTINCT test.b)]] [a:UInt32, count(DISTINCT test.b):Int64] - TableScan: test [a:UInt32, b:UInt32, c:UInt32] + Projection: test.a, count(alias1) AS count(DISTINCT test.b) [a:UInt32, count(DISTINCT test.b):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] " ) } @@ -941,20 +949,20 @@ mod tests { #[test] fn group_by_with_expr() -> Result<()> { - let table_scan = test_table_scan_utf8_b()?; + let table_scan = test_table_scan().unwrap(); let plan = LogicalPlanBuilder::from(table_scan) - .aggregate(vec![col("a") + lit(1)], vec![count_distinct(col("b"))])? + .aggregate(vec![col("a") + lit(1)], vec![count_distinct(col("c"))])? .build()?; // Should work assert_optimized_plan_equal!( plan, @r" - Projection: group_alias_0 AS test.a + Int32(1), count(alias1) AS count(DISTINCT test.b) [test.a + Int32(1):Int64, count(DISTINCT test.b):Int64] + 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] 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.b AS alias1]], aggr=[[]] [group_alias_0:Int64, alias1:Utf8] - TableScan: test [a:UInt32, b:Utf8, c:UInt32] + 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] " ) } diff --git a/datafusion/sqllogictest/test_files/aggregates_simplify.slt b/datafusion/sqllogictest/test_files/aggregates_simplify.slt index 00cfc3956a232..4169707913fb2 100644 --- a/datafusion/sqllogictest/test_files/aggregates_simplify.slt +++ b/datafusion/sqllogictest/test_files/aggregates_simplify.slt @@ -543,17 +543,23 @@ query TT EXPLAIN SELECT g, count(DISTINCT v) FROM distinct_simplify_t GROUP BY g; ---- logical_plan -01)Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[count(DISTINCT distinct_simplify_t.v)]] -02)--TableScan: distinct_simplify_t projection=[g, v] +01)Projection: distinct_simplify_t.g, count(alias1) AS count(DISTINCT distinct_simplify_t.v) +02)--Aggregate: groupBy=[[distinct_simplify_t.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[distinct_simplify_t.g, distinct_simplify_t.v AS alias1]], aggr=[[]] +04)------TableScan: distinct_simplify_t projection=[g, v] physical_plan -01)AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(DISTINCT distinct_simplify_t.v)] -02)--RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 -03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(DISTINCT distinct_simplify_t.v)] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[g@0 as g, count(alias1)@1 as count(DISTINCT distinct_simplify_t.v)] +02)--AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[count(alias1)] +03)----RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=4 +04)------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(alias1)] +05)--------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g, alias1@1 as alias1], aggr=[] +06)----------RepartitionExec: partitioning=Hash([g@0, alias1@1], 4), input_partitions=1 +07)------------AggregateExec: mode=Partial, gby=[g@0 as g, v@1 as alias1], aggr=[] +08)--------------DataSourceExec: partitions=1, partition_sizes=[1] # Mixed: a node that still needs one DISTINCT keeps all of them, so this rule -# leaves min(DISTINCT) in place. The mixed min/count aggregate still gets the -# existing SingleDistinctToGroupBy rewrite. +# cannot change which plans SingleDistinctToGroupBy rewrites. `count` needs the +# rewrite, and the plan below shows that it still happens. query III SELECT g, min(DISTINCT v), count(DISTINCT v) FROM distinct_simplify_t GROUP BY g ORDER BY g; ---- diff --git a/datafusion/sqllogictest/test_files/clickbench.slt b/datafusion/sqllogictest/test_files/clickbench.slt index 175dc01ba6d8d..7cb5547383c38 100644 --- a/datafusion/sqllogictest/test_files/clickbench.slt +++ b/datafusion/sqllogictest/test_files/clickbench.slt @@ -232,18 +232,22 @@ EXPLAIN SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "Reg ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.RegionID, count(DISTINCT hits.UserID) AS u -03)----Aggregate: groupBy=[[hits.RegionID]], aggr=[[count(DISTINCT hits.UserID)]] -04)------SubqueryAlias: hits -05)--------TableScan: hits_raw projection=[RegionID, UserID] +02)--Projection: hits.RegionID, count(alias1) AS u +03)----Aggregate: groupBy=[[hits.RegionID]], aggr=[[count(alias1)]] +04)------Aggregate: groupBy=[[hits.RegionID, hits.UserID AS alias1]], aggr=[[]] +05)--------SubqueryAlias: hits +06)----------TableScan: hits_raw projection=[RegionID, UserID] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[RegionID@0 as RegionID, count(DISTINCT hits.UserID)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[count(DISTINCT hits.UserID)] -05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=1 -06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[count(DISTINCT hits.UserID)] -07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[RegionID, UserID], file_type=parquet +02)--ProjectionExec: expr=[RegionID@0 as RegionID, count(alias1)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] +05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=4 +06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] +07)------------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID, alias1@1 as alias1], aggr=[] +08)--------------RepartitionExec: partitioning=Hash([RegionID@0, alias1@1], 4), input_partitions=1 +09)----------------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID, UserID@1 as alias1], aggr=[] +10)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[RegionID, UserID], file_type=parquet query II rowsort SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "RegionID" ORDER BY u DESC LIMIT 10; @@ -286,21 +290,25 @@ EXPLAIN SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.MobilePhoneModel, count(DISTINCT hits.UserID) AS u -03)----Aggregate: groupBy=[[hits.MobilePhoneModel]], aggr=[[count(DISTINCT hits.UserID)]] -04)------SubqueryAlias: hits -05)--------Filter: hits_raw.MobilePhoneModel != Utf8View("") -06)----------TableScan: hits_raw projection=[UserID, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] +02)--Projection: hits.MobilePhoneModel, count(alias1) AS u +03)----Aggregate: groupBy=[[hits.MobilePhoneModel]], aggr=[[count(alias1)]] +04)------Aggregate: groupBy=[[hits.MobilePhoneModel, hits.UserID AS alias1]], aggr=[[]] +05)--------SubqueryAlias: hits +06)----------Filter: hits_raw.MobilePhoneModel != Utf8View("") +07)------------TableScan: hits_raw projection=[UserID, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(DISTINCT hits.UserID)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] +02)--ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(alias1)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([MobilePhoneModel@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] -07)------------FilterExec: MobilePhoneModel@1 != -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=[UserID, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] +06)----------AggregateExec: mode=Partial, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] +07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel, alias1@1 as alias1], aggr=[] +08)--------------RepartitionExec: partitioning=Hash([MobilePhoneModel@0, alias1@1], 4), input_partitions=4 +09)----------------AggregateExec: mode=Partial, gby=[MobilePhoneModel@1 as MobilePhoneModel, UserID@0 as alias1], aggr=[] +10)------------------FilterExec: MobilePhoneModel@1 != +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=[UserID, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] query TI SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhoneModel" ORDER BY u DESC LIMIT 10; @@ -312,21 +320,25 @@ EXPLAIN SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.MobilePhone, hits.MobilePhoneModel, count(DISTINCT hits.UserID) AS u -03)----Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel]], aggr=[[count(DISTINCT hits.UserID)]] -04)------SubqueryAlias: hits -05)--------Filter: hits_raw.MobilePhoneModel != Utf8View("") -06)----------TableScan: hits_raw projection=[UserID, MobilePhone, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] +02)--Projection: hits.MobilePhone, hits.MobilePhoneModel, count(alias1) AS u +03)----Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel]], aggr=[[count(alias1)]] +04)------Aggregate: groupBy=[[hits.MobilePhone, hits.MobilePhoneModel, hits.UserID AS alias1]], aggr=[[]] +05)--------SubqueryAlias: hits +06)----------Filter: hits_raw.MobilePhoneModel != Utf8View("") +07)------------TableScan: hits_raw projection=[UserID, MobilePhone, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@2 DESC], fetch=10 -02)--ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(DISTINCT hits.UserID)@2 as u] -03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@2 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] +02)--ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(alias1)@2 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@2 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[MobilePhone@1 as MobilePhone, MobilePhoneModel@2 as MobilePhoneModel], aggr=[count(DISTINCT hits.UserID)] -07)------------FilterExec: MobilePhoneModel@2 != -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=[UserID, MobilePhone, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] +06)----------AggregateExec: mode=Partial, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] +07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, alias1@2 as alias1], aggr=[] +08)--------------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1, alias1@2], 4), input_partitions=4 +09)----------------AggregateExec: mode=Partial, gby=[MobilePhone@1 as MobilePhone, MobilePhoneModel@2 as MobilePhoneModel, UserID@0 as alias1], aggr=[] +10)------------------FilterExec: MobilePhoneModel@2 != +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=[UserID, MobilePhone, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] query ITI SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhone", "MobilePhoneModel" ORDER BY u DESC LIMIT 10; @@ -364,21 +376,25 @@ EXPLAIN SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "Se ---- logical_plan 01)Sort: u DESC NULLS FIRST, fetch=10 -02)--Projection: hits.SearchPhrase, count(DISTINCT hits.UserID) AS u -03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[count(DISTINCT hits.UserID)]] -04)------SubqueryAlias: hits -05)--------Filter: hits_raw.SearchPhrase != Utf8View("") -06)----------TableScan: hits_raw projection=[UserID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] +02)--Projection: hits.SearchPhrase, count(alias1) AS u +03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[count(alias1)]] +04)------Aggregate: groupBy=[[hits.SearchPhrase, hits.UserID AS alias1]], aggr=[[]] +05)--------SubqueryAlias: hits +06)----------Filter: hits_raw.SearchPhrase != Utf8View("") +07)------------TableScan: hits_raw projection=[UserID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(DISTINCT hits.UserID)@1 as u] -03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT hits.UserID)@1 DESC], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(DISTINCT hits.UserID)] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(alias1)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@1 as SearchPhrase], aggr=[count(DISTINCT hits.UserID)] -07)------------FilterExec: SearchPhrase@1 != -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=[UserID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , 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=[count(alias1)] +07)------------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase, alias1@1 as alias1], aggr=[] +08)--------------RepartitionExec: partitioning=Hash([SearchPhrase@0, alias1@1], 4), input_partitions=4 +09)----------------AggregateExec: mode=Partial, gby=[SearchPhrase@1 as SearchPhrase, UserID@0 as alias1], aggr=[] +10)------------------FilterExec: SearchPhrase@1 != +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=[UserID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query TI SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY u DESC LIMIT 10; diff --git a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt index 0b449b723e911..35d4bf8873920 100644 --- a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt +++ b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt @@ -15,8 +15,12 @@ # specific language governing permissions and limitations # under the License. -# Grouped integer DISTINCT counts should keep their specialized accumulator. -# String and floating-point keys still use the two-phase rewrite. +# With one target partition, grouped integer DISTINCT counts keep their +# specialized accumulator. Multiple partitions retain distributed deduplication. +# String and floating-point DISTINCT arguments still use the two-phase rewrite. + +statement ok +SET datafusion.execution.target_partitions = 1; statement ok CREATE TABLE numeric_counts(g INT, x INT, s VARCHAR, f DOUBLE); @@ -172,5 +176,63 @@ query II SELECT g, count(DISTINCT x) FROM numeric_counts WHERE g = 42 GROUP BY g; ---- +# The same session must switch back to the rewrite when partitions increase. +statement ok +SET datafusion.execution.target_partitions = 4; + +statement ok +SET datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Projection: numeric_counts.g, count(alias1) AS count(DISTINCT numeric_counts.x) +02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.x AS alias1]], aggr=[[]] +04)------TableScan: numeric_counts projection=[g, x] + +statement ok +SET datafusion.explain.logical_plan_only = false; + +# Duplicates span source batches and partitions; one hot string group, a NULL +# grouping key and an all-NULL DISTINCT group exercise the fallback's results. +statement ok +CREATE TABLE parallel_counts AS +SELECT CASE WHEN value % 10 < 8 THEN 'hot' + WHEN value % 10 = 8 THEN 'null_values' END AS g, + CASE WHEN value % 10 = 8 THEN NULL ELSE value % 100 END AS x +FROM range(100000); + +query TI rowsort +SELECT g, count(DISTINCT x) FROM parallel_counts GROUP BY g; +---- +NULL 10 +hot 80 +null_values 0 + +query TII rowsort +SELECT g, count(DISTINCT x), count(DISTINCT x) AS again +FROM parallel_counts GROUP BY g; +---- +NULL 10 10 +hot 80 80 +null_values 0 0 + +# Composite grouping with the integer DISTINCT argument, as in ClickBench Q11. +query ITI rowsort +SELECT length(g), g, count(DISTINCT x) FROM parallel_counts GROUP BY length(g), g; +---- +11 null_values 0 +3 hot 80 +NULL NULL 10 + +query TI +SELECT g, count(DISTINCT x) FROM parallel_counts WHERE x = 101 GROUP BY g; +---- + +statement ok +DROP TABLE parallel_counts; + statement ok DROP TABLE numeric_counts; diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index bd9ab874ea083..f4dbd212eb8ee 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -1357,10 +1357,12 @@ inner join join_t2 on join_t1.t1_id = join_t2.t2_id group by join_t1.t1_id ---- logical_plan -01)Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(DISTINCT join_t1.t1_int)]] -02)--LeftSemi Join: join_t1.t1_id = join_t2.t2_id -03)----TableScan: join_t1 projection=[t1_id, t1_int] -04)----TableScan: join_t2 projection=[t2_id] +01)Projection: join_t1.t1_id, count(alias1) AS count(DISTINCT join_t1.t1_int) +02)--Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[join_t1.t1_id, join_t1.t1_int AS alias1]], aggr=[[]] +04)------LeftSemi Join: join_t1.t1_id = join_t2.t2_id +05)--------TableScan: join_t1 projection=[t1_id, t1_int] +06)--------TableScan: join_t2 projection=[t2_id] # A similar query with two DISTINCT aggregates is also rewritten: each # `count(DISTINCT ...)` removes its own duplicates, so the join's duplicates diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part index 5d85e3932e557..df7a26bf0a20d 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part @@ -51,33 +51,37 @@ limit 10; ---- logical_plan 01)Sort: supplier_cnt DESC NULLS FIRST, part.p_brand ASC NULLS LAST, part.p_type ASC NULLS LAST, part.p_size ASC NULLS LAST, fetch=10 -02)--Projection: part.p_brand, part.p_type, part.p_size, count(DISTINCT partsupp.ps_suppkey) AS supplier_cnt -03)----Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size]], aggr=[[count(DISTINCT partsupp.ps_suppkey)]] -04)------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey null_aware -05)--------Projection: partsupp.ps_suppkey, part.p_brand, part.p_type, part.p_size -06)----------Inner Join: partsupp.ps_partkey = part.p_partkey -07)------------TableScan: partsupp projection=[ps_partkey, ps_suppkey] -08)------------Filter: part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) AND part.p_brand != Utf8View("Brand#45") AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") -09)--------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]), part.p_brand != Utf8View("Brand#45"), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%")] -10)--------SubqueryAlias: __correlated_sq_1 -11)----------Projection: supplier.s_suppkey -12)------------Filter: supplier.s_comment LIKE Utf8View("%Customer%Complaints%") -13)--------------TableScan: supplier projection=[s_suppkey, s_comment], partial_filters=[supplier.s_comment LIKE Utf8View("%Customer%Complaints%")] +02)--Projection: part.p_brand, part.p_type, part.p_size, count(alias1) AS supplier_cnt +03)----Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size]], aggr=[[count(alias1)]] +04)------Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size, partsupp.ps_suppkey AS alias1]], aggr=[[]] +05)--------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey null_aware +06)----------Projection: partsupp.ps_suppkey, part.p_brand, part.p_type, part.p_size +07)------------Inner Join: partsupp.ps_partkey = part.p_partkey +08)--------------TableScan: partsupp projection=[ps_partkey, ps_suppkey] +09)--------------Filter: part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) AND part.p_brand != Utf8View("Brand#45") AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") +10)----------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]), part.p_brand != Utf8View("Brand#45"), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%")] +11)----------SubqueryAlias: __correlated_sq_1 +12)------------Projection: supplier.s_suppkey +13)--------------Filter: supplier.s_comment LIKE Utf8View("%Customer%Complaints%") +14)----------------TableScan: supplier projection=[s_suppkey, s_comment], partial_filters=[supplier.s_comment LIKE Utf8View("%Customer%Complaints%")] physical_plan 01)SortPreservingMergeExec: [supplier_cnt@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], fetch=10 -02)--ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(DISTINCT partsupp.ps_suppkey)@3 as supplier_cnt] -03)----SortExec: TopK(fetch=10), expr=[count(DISTINCT partsupp.ps_suppkey)@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] -04)------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(DISTINCT partsupp.ps_suppkey)] +02)--ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(alias1)@3 as supplier_cnt] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[p_brand@1 as p_brand, p_type@2 as p_type, p_size@3 as p_size], aggr=[count(DISTINCT partsupp.ps_suppkey)] -07)------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)], null_aware -08)--------------CoalescePartitionsExec -09)----------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, p_partkey@0)], projection=[ps_suppkey@1, p_brand@3, p_type@4, p_size@5] -10)------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 -11)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false -12)------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 -13)--------------------FilterExec: p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) AND p_brand@1 != Brand#45 AND p_type@2 NOT LIKE MEDIUM POLISHED% -14)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false -15)--------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] -16)----------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -17)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +06)----------AggregateExec: mode=Partial, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] +07)------------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, alias1@3 as alias1], aggr=[] +08)--------------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2, alias1@3], 4), input_partitions=4 +09)----------------AggregateExec: mode=Partial, gby=[p_brand@1 as p_brand, p_type@2 as p_type, p_size@3 as p_size, ps_suppkey@0 as alias1], aggr=[] +10)------------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)], null_aware +11)--------------------CoalescePartitionsExec +12)----------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, p_partkey@0)], projection=[ps_suppkey@1, p_brand@3, p_type@4, p_size@5] +13)------------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 +14)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false +15)------------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 +16)--------------------------FilterExec: p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) AND p_brand@1 != Brand#45 AND p_type@2 NOT LIKE MEDIUM POLISHED% +17)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +18)--------------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] +19)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +20)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false From 4a0bacd4976ed34668b0f074c9576e859de55235 Mon Sep 17 00:00:00 2001 From: hushen <190065939+918154429@users.noreply.github.com> Date: Sat, 7 Nov 2026 14:15:27 +0800 Subject: [PATCH 4/5] fix: make native grouped DISTINCT aggregation opt-in Preserve the spill-capable rewrite by default and document the native plan's memory and duplicate-heavy workload tradeoffs. Cover low-memory execution and multiple input partitions, and add the million-group duplicate benchmark. --- .../grouped_count_distinct/README.md | 30 +++++ .../benchmarks/q10.benchmark | 26 +++++ .../grouped_count_distinct.suite | 2 +- .../grouped_count_distinct/init/load.sql | 4 +- datafusion/common/src/config.rs | 12 ++ .../src/single_distinct_to_groupby.rs | 53 ++++----- .../test_files/aggregate_memory_spill.slt | 44 ++++++- .../test_files/grouped_count_distinct.slt | 107 +++++++++++++++++- .../test_files/information_schema.slt | 2 + docs/source/user-guide/configs.md | 1 + 10 files changed, 249 insertions(+), 32 deletions(-) create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/README.md create mode 100644 benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q10.benchmark diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/README.md b/benchmarks/sql_benchmarks/grouped_count_distinct/README.md new file mode 100644 index 0000000000000..65fb2de2003c9 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/README.md @@ -0,0 +1,30 @@ + + +# Grouped COUNT(DISTINCT) + +Compare the default rewrite with the native accumulator opt-in at one target +partition. The suite reads DataFusion settings from the environment; it does not +enable the opt-in itself. + +```bash +DATAFUSION_OPTIMIZER_SINGLE_DISTINCT_PREFER_NATIVE=false cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 1 --mem-pool-type greedy --memory-limit 1G +DATAFUSION_OPTIMIZER_SINGLE_DISTINCT_PREFER_NATIVE=true cargo run --release --bin benchmark_runner -- grouped_count_distinct --iterations 10 --partitions 1 --mem-pool-type greedy --memory-limit 1G +``` + +Run from `benchmarks/`. Repeat at 4 and 12 target partitions to check the +distributed rewrite, which remains selected with either flag value. The same +environment variable works with `./bench.sh run grouped_count_distinct`. + +Q01-Q09 vary argument and grouping-key types, group count, distinct cardinality, +and skew. Q10 covers the unfavorable native shape: at the default four million +rows it has one million groups, each with four copies of the same value. Each +native input row probes both the group table and the per-group distinct set; +the rewrite's outer aggregate only sees the deduplicated pairs. + +Every query checks its result against an explicit distinct-pairs-then-count +query before timing. Keep the memory pool and limit the same on both sides. +The combined correctness query needs more memory than the measured aggregate. diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..be8dcb20c5b28 --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/benchmarks/q10.benchmark @@ -0,0 +1,26 @@ +name Q10 +group grouped_count_distinct + +load sql_benchmarks/grouped_count_distinct/init/load.sql + +assert I +WITH actual AS ( + SELECT many_g AS g, count(DISTINCT dup_x) AS n FROM grouped_distinct GROUP BY many_g +), expected AS ( + SELECT g, count(x) AS n + FROM (SELECT DISTINCT many_g AS g, dup_x AS x FROM grouped_distinct) + GROUP BY g +) +SELECT NOT EXISTS (SELECT * FROM actual EXCEPT SELECT * FROM expected) + AND NOT EXISTS (SELECT * FROM expected EXCEPT SELECT * FROM actual); +---- +true + +expect_plan AggregateExec + +run +-- Int64, 1M groups with four identical values per group at the default 4M rows. +-- Native DISTINCT probes both the group table and its per-group set on each row. +SELECT many_g, count(DISTINCT dup_x) FROM grouped_distinct GROUP BY many_g; + +cleanup sql_benchmarks/grouped_count_distinct/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite b/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite index 725cdb058f4a6..16b4f0857920e 100644 --- a/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite @@ -1,4 +1,4 @@ -description = "Grouped COUNT(DISTINCT) across argument types, integer/string/composite grouping keys, distinct cardinalities and skew; compare --partitions 1, 4 and 12" +description = "Grouped COUNT(DISTINCT) across types, cardinalities, skew and duplicate-heavy groups; compare --partitions 1, 4 and 12 with DATAFUSION_OPTIMIZER_SINGLE_DISTINCT_PREFER_NATIVE=false/true" [[options]] name = "rows" diff --git a/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql b/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql index efdbbf1a9e31b..0f7c5b03622de 100644 --- a/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql @@ -7,5 +7,7 @@ SELECT value % 2000 AS g, CAST((value * 48271) % 999983 AS VARCHAR) AS string_x, CAST((value * 48271) % 999983 AS DOUBLE) AS float_x, CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS skew_g, - CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR) AS string_g + CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR) AS string_g, + value % 1000000 AS many_g, + value % 2 AS dup_x FROM range(${GROUPED_DISTINCT_ROWS:-4000000}); diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index c825dfb2a835d..c89336f55fe44 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1914,6 +1914,18 @@ config_namespace! { /// branches share the same source and compatible wrapper nodes such as identical /// projections or aliases. pub enable_unions_to_filter: bool, default = false + + /// Prefer native grouped DISTINCT aggregation when target_partitions is 1 + /// and every aggregate explicitly supports a specialized GroupsAccumulator + /// for its DISTINCT argument types. Defaults to false, preserving the + /// two-phase GROUP BY rewrite. + /// + /// The native plan can avoid materializing distinct (group, value) pairs, + /// but can be slower for many groups with duplicate values. It must fit + /// each group's distinct values in memory, so queries that succeed by + /// spilling with the rewrite may fail with this option enabled. One target + /// partition can still have multiple input partitions and merge native states. + pub single_distinct_prefer_native: bool, default = false } } diff --git a/datafusion/optimizer/src/single_distinct_to_groupby.rs b/datafusion/optimizer/src/single_distinct_to_groupby.rs index 7aabbc800d756..73efe0de34475 100644 --- a/datafusion/optimizer/src/single_distinct_to_groupby.rs +++ b/datafusion/optimizer/src/single_distinct_to_groupby.rs @@ -68,6 +68,9 @@ use datafusion_expr::{ /// `sum` of no rows is NULL. /// /// See `rewrite_pays_for_count` for when the `count` is allowed. +/// With `single_distinct_prefer_native` enabled and one target partition, grouped +/// all-DISTINCT aggregates stay native if every function explicitly supports a +/// specialized `GroupsAccumulator` for its argument types. #[derive(Default, Debug)] pub struct SingleDistinctToGroupBy {} @@ -160,10 +163,16 @@ fn is_single_distinct_agg( // partitions, the rewrite also distributes deduplication by both keys; // keeping DISTINCT instead repartitions lists of values by group and merges // them in the final accumulator. That can concentrate work on a few groups. - // Native accumulator support alone does not establish that this is cheaper. + // Even one target partition can have multiple input partitions (e.g. UNION + // ALL or an existing partitioned table), requiring Partial/Final state merging. + // This is opt-in: the rule cannot see memory limits or cost the input, and + // native DISTINCT must fit each group's distinct values in memory. if use_native_grouped_distinct && distinct_aggs.len() == aggregate_count - && all_have_groups_accumulators(&distinct_aggs, input_schema)? + // An unknown answer must preserve the existing rewrite. + && distinct_groups_support(&distinct_aggs, input_schema)? + .iter() + .all(|support| *support == Some(true)) { return Ok(false); } @@ -173,21 +182,21 @@ fn is_single_distinct_agg( Ok(true) } -fn all_have_groups_accumulators( +/// Each distinct aggregate's specialized GroupsAccumulator support for its types. +fn distinct_groups_support( distinct_aggs: &[(&Arc, &[Expr])], input_schema: &DFSchema, -) -> Result { - for (func, args) in distinct_aggs { - let arg_types = args - .iter() - .map(|arg| arg.get_type(input_schema)) - .collect::>>()?; - // An unknown answer must preserve the existing rewrite. - if func.groups_accumulator_supported_for_types(&arg_types, true) != Some(true) { - return Ok(false); - } - } - Ok(true) +) -> Result>> { + distinct_aggs + .iter() + .map(|(func, args)| { + let arg_types = args + .iter() + .map(|arg| arg.get_type(input_schema)) + .collect::>>()?; + Ok(func.groups_accumulator_supported_for_types(&arg_types, true)) + }) + .collect() } /// Whether the rewrite is worth extending to a plan that only qualifies because @@ -215,16 +224,7 @@ fn rewrite_pays_for_count( distinct_aggs: &[(&Arc, &[Expr])], 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) == Some(false) { - return Ok(true); - } - } - Ok(false) + Ok(distinct_groups_support(distinct_aggs, input_schema)?.contains(&Some(false))) } /// Check if the first expr is [Expr::GroupingSet]. @@ -262,7 +262,8 @@ impl OptimizerRule for SingleDistinctToGroupBy { input.schema(), count_rollup.as_ref(), !group_expr.is_empty() - && config.options().execution.target_partitions == 1, + && config.options().execution.target_partitions == 1 + && config.options().optimizer.single_distinct_prefer_native, )? && !contains_grouping_set(&group_expr) => { let group_size = group_expr.len(); diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 87f5db19323a1..04a268af5c205 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -102,9 +102,51 @@ FROM ( ) ---- -04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[count(DISTINCT t.v)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] +04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[spill_count=, spilled_bytes=, spilled_rows=100.0 K,] +# A few large DISTINCT groups must still finish with the default rewrite. +# Native DISTINCT would need to hold 50,000 values per group under this 1M limit. +query II +SELECT count(*), sum(d) FROM ( + SELECT v % 2 AS k, count(DISTINCT v) AS d + FROM generate_series(1, 100000) AS t(v) + GROUP BY v % 2 +); +---- +2 100000 + +statement ok +SET datafusion.optimizer.single_distinct_prefer_native = true; + +statement ok +SET datafusion.explain.physical_plan_only = true; + +# EXPLAIN does not allocate the native accumulator's state. +query TT +EXPLAIN SELECT v % 2 AS k, count(DISTINCT v) AS d +FROM generate_series(1, 100000) AS t(v) +GROUP BY v % 2; +---- +physical_plan +01)ProjectionExec: expr=[t.v % Int64(2)@0 as k, count(DISTINCT t.v)@1 as d] +02)--AggregateExec: mode=Single, gby=[v@0 % 2 as t.v % Int64(2)], aggr=[count(DISTINCT t.v)] +03)----ProjectionExec: expr=[value@0 as v] +04)------LazyMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=100000, batch_size=128] + +statement error Resources exhausted +SELECT count(*), sum(d) FROM ( + SELECT v % 2 AS k, count(DISTINCT v) AS d + FROM generate_series(1, 100000) AS t(v) + GROUP BY v % 2 +); + +statement ok +SET datafusion.explain.physical_plan_only = false; + +statement ok +SET datafusion.optimizer.single_distinct_prefer_native = false; + # --- Case D: multiple aggregates (sum/min/max) under memory limit --- # Each group holds a single v, so min(v) = max(v) = v within the group. query IIII diff --git a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt index 7dd9161032811..1a4ecbea9f580 100644 --- a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt +++ b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt @@ -2,8 +2,9 @@ # SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership # SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements -# With one target partition, grouped integer DISTINCT counts keep their -# specialized accumulator. Multiple partitions retain distributed deduplication. +# With the opt-in enabled and one target partition, grouped integer DISTINCT +# counts keep their specialized accumulator. Multiple target partitions retain +# distributed deduplication. # String and floating-point DISTINCT arguments still use the two-phase rewrite. statement ok @@ -24,6 +25,26 @@ INSERT INTO numeric_counts VALUES statement ok SET datafusion.explain.logical_plan_only = true; +# The default preserves the rewrite, even with one target partition. +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Projection: numeric_counts.g, count(alias1) AS count(DISTINCT numeric_counts.x) +02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.x AS alias1]], aggr=[[]] +04)------TableScan: numeric_counts projection=[g, x] + +query II rowsort +SELECT g, count(DISTINCT x) FROM numeric_counts GROUP BY g; +---- +1 2 +2 0 +3 1 + +statement ok +SET datafusion.optimizer.single_distinct_prefer_native = true; + query TT EXPLAIN SELECT g, count(DISTINCT x) FROM numeric_counts GROUP BY g; ---- @@ -163,7 +184,55 @@ query II SELECT g, count(DISTINCT x) FROM numeric_counts WHERE g = 42 GROUP BY g; ---- -# The same session must switch back to the rewrite when partitions increase. +# One target partition can still merge native states from several input partitions. +query II rowsort +SELECT g, count(DISTINCT x) +FROM (SELECT g, x FROM numeric_counts UNION ALL SELECT g, x FROM numeric_counts) +GROUP BY g; +---- +1 2 +2 0 +3 1 + +statement ok +SET datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT g, count(DISTINCT x) +FROM (SELECT g, x FROM numeric_counts UNION ALL SELECT g, x FROM numeric_counts) +GROUP BY g; +---- +physical_plan +01)AggregateExec: mode=Final, gby=[g@0 as g], aggr=[count(DISTINCT x)] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(DISTINCT x)] +04)------UnionExec +05)--------DataSourceExec: partitions=1, partition_sizes=[1] +06)--------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +SET datafusion.explain.physical_plan_only = false; + +# Turning the opt-in off in the same session restores the rewrite. +statement ok +SET datafusion.optimizer.single_distinct_prefer_native = false; + +statement ok +SET datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Projection: numeric_counts.g, count(alias1) AS count(DISTINCT numeric_counts.x) +02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.x AS alias1]], aggr=[[]] +04)------TableScan: numeric_counts projection=[g, x] + +statement ok +SET datafusion.optimizer.single_distinct_prefer_native = true; + +# Even with the opt-in enabled, increasing target partitions restores the rewrite. statement ok SET datafusion.execution.target_partitions = 4; @@ -218,8 +287,40 @@ query TI SELECT g, count(DISTINCT x) FROM parallel_counts WHERE x = 101 GROUP BY g; ---- +# This table retains its four input partitions when the target is reduced to one. +statement ok +SET datafusion.execution.target_partitions = 1; + +query TI rowsort +SELECT g, count(DISTINCT x) FROM parallel_counts GROUP BY g; +---- +NULL 10 +hot 80 +null_values 0 + +statement ok +SET datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT g, count(DISTINCT x) FROM parallel_counts GROUP BY g; +---- +physical_plan +01)AggregateExec: mode=Final, gby=[g@0 as g], aggr=[count(DISTINCT parallel_counts.x)] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[count(DISTINCT parallel_counts.x)] +04)------DataSourceExec: partitions=4, partition_sizes=[4, 3, 3, 3] + +statement ok +SET datafusion.explain.physical_plan_only = false; + statement ok DROP TABLE parallel_counts; statement ok DROP TABLE numeric_counts; + +statement ok +SET datafusion.optimizer.single_distinct_prefer_native = false; + +statement ok +SET datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index c7939b29d8340..e682f24cfa7a8 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -334,6 +334,7 @@ datafusion.optimizer.repartition_file_scans true datafusion.optimizer.repartition_joins true datafusion.optimizer.repartition_sorts true datafusion.optimizer.repartition_windows true +datafusion.optimizer.single_distinct_prefer_native false datafusion.optimizer.skip_failed_rules false datafusion.optimizer.subset_repartition_threshold 4 datafusion.optimizer.top_down_join_key_reordering true @@ -500,6 +501,7 @@ datafusion.optimizer.repartition_file_scans true When set to `true`, datasource datafusion.optimizer.repartition_joins true Should DataFusion repartition data using the join keys to execute joins in parallel using the provided `target_partitions` level datafusion.optimizer.repartition_sorts true Should DataFusion execute sorts in a per-partition fashion and merge afterwards instead of coalescing first and sorting globally. With this flag is enabled, plans in the form below ```text "SortExec: [a@0 ASC]", " CoalescePartitionsExec", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` would turn into the plan below which performs better in multithreaded environments ```text "SortPreservingMergeExec: [a@0 ASC]", " SortExec: [a@0 ASC]", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` datafusion.optimizer.repartition_windows true Should DataFusion repartition data using the partitions keys to execute window functions in parallel using the provided `target_partitions` level +datafusion.optimizer.single_distinct_prefer_native false Prefer native grouped DISTINCT aggregation when target_partitions is 1 and every aggregate explicitly supports a specialized GroupsAccumulator for its DISTINCT argument types. Defaults to false, preserving the two-phase GROUP BY rewrite. The native plan can avoid materializing distinct (group, value) pairs, but can be slower for many groups with duplicate values. It must fit each group's distinct values in memory, so queries that succeed by spilling with the rewrite may fail with this option enabled. One target partition can still have multiple input partitions and merge native states. datafusion.optimizer.skip_failed_rules false When set to true, the logical plan optimizer will produce warning messages if any optimization rules produce errors and then proceed to the next rule. When set to false, any rules that produce errors will cause the query to fail datafusion.optimizer.subset_repartition_threshold 4 Partition count threshold for subset satisfaction optimization. When the current partition count is >= this threshold, DataFusion will skip repartitioning if the required partitioning expression is a subset of the current partition expression such as Hash(a) satisfies Hash(a, b). When the current partition count is < this threshold, DataFusion will repartition to increase parallelism even when subset satisfaction applies. Set to 0 to always repartition (disable subset satisfaction optimization). Set to a high value to always use subset satisfaction. Example (subset_repartition_threshold = 4): ```text Hash([a]) satisfies Hash([a, b]) because (Hash([a, b]) is subset of Hash([a]) If current partitions (3) < threshold (4), repartition: AggregateExec: mode=FinalPartitioned, gby=[a, b], aggr=[SUM(x)] RepartitionExec: partitioning=Hash([a, b], 8), input_partitions=3 AggregateExec: mode=Partial, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 3) If current partitions (8) >= threshold (4), use subset satisfaction: AggregateExec: mode=SinglePartitioned, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 8) ``` datafusion.optimizer.top_down_join_key_reordering true When set to true, the physical plan optimizer will run a top down process to reorder the join keys diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index fe163e4addf5d..3593a78630319 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -189,6 +189,7 @@ The following configuration settings are available: | datafusion.optimizer.enable_sort_pushdown | true | Enable sort pushdown optimization. When enabled, attempts to push sort requirements down to data sources that can natively handle them (e.g., by reversing file/row group read order). Returns **inexact ordering**: Sort operator is kept for correctness, but optimized input enables early termination for TopK queries (ORDER BY ... LIMIT N), providing significant speedup. Memory: No additional overhead (only changes read order). Future: Will add option to detect perfectly sorted data and eliminate Sort completely. Default: true | | datafusion.optimizer.enable_leaf_expression_pushdown | true | When set to true, the optimizer will extract leaf expressions (such as `get_field`) from filter/sort/join nodes into projections closer to the leaf table scans, and push those projections down towards the leaf nodes. | | datafusion.optimizer.enable_unions_to_filter | false | When set to true, the logical optimizer will rewrite `UNION DISTINCT` branches that read from the same source and differ only by filter predicates into a single branch with a combined filter. This optimization is conservative and only applies when the branches share the same source and compatible wrapper nodes such as identical projections or aliases. | +| datafusion.optimizer.single_distinct_prefer_native | false | Prefer native grouped DISTINCT aggregation when target_partitions is 1 and every aggregate explicitly supports a specialized GroupsAccumulator for its DISTINCT argument types. Defaults to false, preserving the two-phase GROUP BY rewrite. The native plan can avoid materializing distinct (group, value) pairs, but can be slower for many groups with duplicate values. It must fit each group's distinct values in memory, so queries that succeed by spilling with the rewrite may fail with this option enabled. One target partition can still have multiple input partitions and merge native states. | | datafusion.explain.logical_plan_only | false | When set to true, the explain statement will only print logical plans | | datafusion.explain.physical_plan_only | false | When set to true, the explain statement will only print physical plans | | datafusion.explain.show_statistics | false | When set to true, the explain statement will print operator statistics for physical plans | From 764869eac9ef04b8c33578c6d93b588e41b1cd6b Mon Sep 17 00:00:00 2001 From: hushen <190065939+918154429@users.noreply.github.com> Date: Sun, 8 Nov 2026 14:59:57 +0800 Subject: [PATCH 5/5] test: cover native DISTINCT rewrite fallback guards Add EXPLAIN regressions for mixed aggregates and unknown specialized DISTINCT support with native planning enabled and one target partition. Generated-by: OpenAI Codex --- .../test_files/grouped_count_distinct.slt | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt index 1a4ecbea9f580..1490442ab383a 100644 --- a/datafusion/sqllogictest/test_files/grouped_count_distinct.slt +++ b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt @@ -70,6 +70,26 @@ logical_plan 03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.f AS alias1]], aggr=[[]] 04)------TableScan: numeric_counts projection=[g, f] +# A non-DISTINCT aggregate alongside keeps the rewrite. +query TT +EXPLAIN SELECT g, count(DISTINCT x), sum(x) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Projection: numeric_counts.g, count(alias1) AS count(DISTINCT numeric_counts.x), sum(alias2) AS sum(numeric_counts.x) +02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[count(alias1), sum(alias2)]] +03)----Aggregate: groupBy=[[numeric_counts.g, numeric_counts.x AS alias1]], aggr=[[sum(CAST(numeric_counts.x AS Int64)) AS alias2]] +04)------TableScan: numeric_counts projection=[g, x] + +# `sum` does not report specialized DISTINCT support, so it keeps the rewrite. +query TT +EXPLAIN SELECT g, sum(DISTINCT x) FROM numeric_counts GROUP BY g; +---- +logical_plan +01)Projection: numeric_counts.g, sum(alias1) AS sum(DISTINCT numeric_counts.x) +02)--Aggregate: groupBy=[[numeric_counts.g]], aggr=[[sum(alias1)]] +03)----Aggregate: groupBy=[[numeric_counts.g, CAST(numeric_counts.x AS Int64) AS alias1]], aggr=[[]] +04)------TableScan: numeric_counts projection=[g, x] + # The global aggregate does not use a GroupsAccumulator. query TT EXPLAIN SELECT count(DISTINCT x) FROM numeric_counts;