diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 4d5044857b462..4cbab03d66a74 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -163,6 +163,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 @@ -280,6 +281,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 ;; @@ -681,6 +685,9 @@ main() { array_agg_distinct) run_array_agg_distinct ;; + grouped_count_distinct) + run_grouped_count_distinct + ;; compile_profile) run_compile_profile "${PROFILE_ARGS[@]}" ;; @@ -1813,6 +1820,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/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/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/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/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 new file mode 100644 index 0000000000000..16b4f0857920e --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/grouped_count_distinct.suite @@ -0,0 +1,15 @@ +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" +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." + +[[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/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..0f7c5b03622de --- /dev/null +++ b/benchmarks/sql_benchmarks/grouped_count_distinct/init/load.sql @@ -0,0 +1,13 @@ +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, + 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, + 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 6f7604bcf4012..63a7797e396ef 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1939,6 +1939,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 8b176e904e9b9..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 {} @@ -112,6 +115,7 @@ fn is_single_distinct_agg( aggr_expr: &[Expr], input_schema: &DFSchema, count_rollup: Option<&CountRollup>, + use_native_grouped_distinct: bool, ) -> Result { let mut fields_set = HashSet::new(); let mut aggregate_count = 0; @@ -154,12 +158,47 @@ fn is_single_distinct_agg( if aggregate_count != aggr_expr.len() || fields_set.len() != 1 { return Ok(false); } + // 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. + // 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 + // An unknown answer must preserve the existing rewrite. + && distinct_groups_support(&distinct_aggs, input_schema)? + .iter() + .all(|support| *support == Some(true)) + { + return Ok(false); + } if has_count_rollup && !rewrite_pays_for_count(&distinct_aggs, input_schema)? { return Ok(false); } Ok(true) } +/// Each distinct aggregate's specialized GroupsAccumulator support for its types. +fn distinct_groups_support( + distinct_aggs: &[(&Arc, &[Expr])], + input_schema: &DFSchema, +) -> 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 /// of the non-distinct `count`. /// @@ -185,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]. @@ -231,6 +261,9 @@ impl OptimizerRule for SingleDistinctToGroupBy { &aggr_expr, input.schema(), count_rollup.as_ref(), + !group_expr.is_empty() + && config.options().execution.target_partitions == 1 + && config.options().optimizer.single_distinct_prefer_native, )? && !contains_grouping_set(&group_expr) => { let group_size = group_expr.len(); @@ -572,8 +605,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, } } @@ -746,7 +782,7 @@ mod tests { .aggregate(vec![col("a")], vec![count_distinct(col("b"))])? .build()?; - // Should work + // Multiple target partitions keep deduplication distributed by both keys. assert_optimized_plan_equal!( plan, @r" diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 8178a8388a117..915e69642e260 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -105,6 +105,48 @@ FROM ( 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 new file mode 100644 index 0000000000000..1490442ab383a --- /dev/null +++ b/datafusion/sqllogictest/test_files/grouped_count_distinct.slt @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# 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 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 +SET datafusion.execution.target_partitions = 1; + +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; + +# 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; +---- +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] + +# 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; +---- +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; +---- + +# 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; + +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; +---- + +# 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 92d15197ad56b..c1aabc2946ea6 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -336,6 +336,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 @@ -504,6 +505,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 73bf1d1a35a08..2c3ca887d4a38 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -178,6 +178,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 |