Repository navigation
fix: Only apply SingleDistinctToGroupBy rewrite when it pays (closes #25932) - #26000
mohitgurav20 wants to merge 38 commits into
Conversation
527a544 to
9e8dfe6
Compare
There was a problem hiding this comment.
@mohitgurav20 Thanks for working on this.
Can you make sure all sql-logic tests pass, or correct them if invalid.
You can run them locally with:
cargo test --profile=ci --test sqllogictests
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #26000 +/- ##
========================================
Coverage 82.93% 82.94%
========================================
Files 1154 1154
Lines 458002 458300 +298
Branches 458002 458300 +298
========================================
+ Hits 379856 380130 +274
+ Misses 55123 55104 -19
- Partials 23023 23066 +43 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
63e9873 to
0c76607
Compare
|
@mkleen sir i have completed the work and addressed ur suggestion pls feel free to review and would like to work if any changes needed. |
|
run benchmarks |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix/single-distinct-groupby-memory-usage (0923424) to 334f293 (merge-base) diff Run configurationrun benchmark clickbench_partitionedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix/single-distinct-groupby-memory-usage (0923424) to 334f293 (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix/single-distinct-groupby-memory-usage (0923424) to 334f293 (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix/single-distinct-groupby-memory-usage (0923424) to 334f293 (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix/single-distinct-groupby-memory-usage (0923424) to 334f293 (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix/single-distinct-groupby-memory-usage (0923424) to 334f293 (merge-base) diff Run configurationrun benchmark clickbench_partitionedCPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
@mohitgurav20 There are serious performance regressions in the clickhouse benchmark. See #26000 (comment) for more details. |
|
Thanks for running the benchmark and flagging this @mkleen. I analyzed the benchmark run in detail: Because Without the two-phase rewrite to This confirms that the two-phase rewrite is a critical parallelization optimization for lone distinct queries, even when a specialized To address this cleanly: Does that direction sound good to you? If so, I'll update the rule accordingly so we preserve ClickBench performance. |
…pache#25932) SingleDistinctToGroupBy converts single DISTINCT aggregate expressions into a two-phase GROUP BY plan. This enables parallel deduplication via RepartitionExec hash partitioning, which is critical for lone-distinct queries like COUNT(DISTINCT x). However, when companion non-distinct aggregates are present alongside the distinct aggregate (e.g. SELECT sum(y), count(DISTINCT x)), the companion aggregates move to the inner group-by at a finer grain. When the distinct aggregate already has native GroupsAccumulator support, this inner grouping adds unnecessary memory and CPU overhead without benefiting from adapter offloading. This change: 1. Gates SingleDistinctToGroupBy with rewrite_pays() only when companion non-distinct aggregates are present. 2. Preserves the two-phase rewrite for pure lone-distinct queries so ClickBench and similar workloads maintain parallel deduplication without regression. 3. Consistently handles distinct aggregates with unsupported distinct handling. 4. Updates optimizer plan snapshots and sqllogictests.
287b932 to
ac76b0d
Compare
|
Hi @mkleen, just following up on this when you get a chance. I've rebased cleanly onto latest Could you trigger a benchmark run whenever convenient? Thanks! |
|
run benchmarks |
|
Benchmark for this request failed before finishing (Kubernetes reason: Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix/single-distinct-groupby-memory-usage (bbf4ccb) to 484183d (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing fix/single-distinct-groupby-memory-usage (bbf4ccb) to 484183d (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix/single-distinct-groupby-memory-usage (bbf4ccb) to 484183d (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing fix/single-distinct-groupby-memory-usage (bbf4ccb) to 484183d (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
@mohitgurav20 Thanks! Looks good at first glance. I’ll go ahead and start the review. |
|
Great. So in which scenarios would switching of the rewrite makes sense? When does it significantly reduce memory consumption with only a minor impact on performance? Ok, this is the part i would focus on. |
|
From checking the execution plans and benchmark numbers, turning off the rewrite (
On the other hand, the benchmarks also clearly showed why the rewrite should stay enabled by default:
So the flag works best when users have memory-heavy |
|
Can we implement the flag in a way, that the memory comsumptions gets better without having any real bad performance impacts and avoid the cases where it gets worse ? |
|
Yes, we can definitely do that. From the benchmarks, the only cases that got worse were:
So under the flag, we can only skip the rewrite if:
That way, strings and queries without I can push this change if that sounds good to you! |
…trings and un-grouped distincts
…roupby-memory-usage
…roupby-memory-usage # Conflicts: # docs/source/user-guide/configs.md
|
@mohitgurav20 Thanks i will review it asap. |
|
@mkleen yeah sure sir have a look whenever u get time , and feel free to address any changes to be done by my side. |
## Which issue does this PR close? Closes apache#25936. ## Rationale for this change Grouped integer `count(DISTINCT x)` can use its native GroupsAccumulator to avoid an extra aggregate over distinct `(group, value)` pairs. This helps the original single-target-partition workload, but the choice also affects spilling and performance on other data shapes. The native plan is now **opt-in**. The default keeps the existing two-phase rewrite. Under a 1 MiB memory limit, the reviewer's query with two groups and 100,000 distinct values completes with the rewrite and returns `2, 100000`; the native plan runs out of memory. Native DISTINCT must fit each group's distinct values in memory, including when merging spilled or partial states. ## What changes are included in this PR? Add `datafusion.optimizer.single_distinct_prefer_native`, default `false`. When enabled, grouped all-DISTINCT aggregates retain the native plan only if: - `execution.target_partitions == 1`; - all aggregates share one DISTINCT argument; - every aggregate explicitly reports specialized GroupsAccumulator support for its argument types. Unknown support preserves the rewrite. Integer DISTINCT counts currently qualify; string and floating-point arguments, global aggregates, mixed aggregates and multiple target partitions keep their existing behavior. The earlier [GKE ClickBench comparison](apache#25988 (comment)) regressed with native parallel plans, so multiple target partitions continue to distribute deduplication by both the group key and distinct value. One target partition can still have multiple input partitions. `UNION ALL` and a table created with four partitions can use Partial → Coalesce → Final and rebuild the DISTINCT set in Final. This is covered by execution tests and physical-plan snapshots. Add a ten-query SQL benchmark suite covering argument types, grouping-key types, cardinality, skew and duplicates. Q10 has one million groups with four copies of the same value per group at the default four million rows. Native execution probes both the group table and the distinct set on every row; the rewrite's outer aggregate only sees deduplicated pairs. This is an unfavorable case for the native path, and the logical rule has no statistics to cost it. This overlaps with the option proposed in apache#26000. `prefer_native` describes the plan choice without promising lower memory use; the two PRs should converge on one setting. ### Local comparison, November 7 Compared main `636590b` with review-fix commit `4a0bacd4`, flag off and on, on Linux (Xeon Gold 6238T, CPU affinity 0–7), Rust 1.99.1, `release-nonlto`, four million rows, greedy / 1 GiB. All ten queries ran at 1, 4 and 12 target partitions in three rounds, reversing version order in the second round. Each round had six iterations per query, with the first discarded: 15 retained samples per version and configuration. All result assertions passed. The main executable was reused from the November 2 build; the review fix was freshly rebuilt, and all versions used the same updated ten-query suite. No builds ran during timing. Single-partition medians: | Case | Main | Flag off (default) | Flag on | On vs off elapsed time | | --- | ---: | ---: | ---: | ---: | | Q01: Int64, 2,000 groups, high NDV | 448.97 ms | 424.17 ms | 305.27 ms | −28.0% | | Q02: Int64, 2,000 groups, low NDV | 102.96 ms | 101.23 ms | 79.91 ms | −21.1% | | Q10: 1M groups, four identical values per group | 316.79 ms | 284.87 ms | 398.12 ms | +39.8% | Q10 reproduces the reviewer's regression. Its memory-pool peak also rises from 69.88 MiB with the flag off to 95.88 MiB with it on. Q01 goes the other way, from 191.45 to 168.17 MiB. Both the runtime and memory tradeoffs are reasons to keep this opt-in. The primary comparison is off/on in the same executable. Main and flag-off use the same plans but have timing differences on this shared host. At four and twelve target partitions both flag values keep the rewrite; observed off/on median differences range from −8.0% to +11.2%. Separate physical-plan checks for Q01 and Q10 passed for all three versions at all three partition settings (18 checks). ## What is the testing strategy for this PR? SQL tests cover the default and both flag values in one session, integer types and NULLs, the low-memory spill reproducer, native plan selection, `UNION ALL`, a pre-partitioned table, and the rewrite at four target partitions. The configuration documentation and information-schema snapshots include the new default. After syncing with main `4bf32cb4`, formatting, generated configuration docs, Prettier, strict all-targets/all-features Clippy and the full `dev/rust_lint.sh` passed on Linux at `0db14a06`. The extended workspace run passed 13,537 Rust tests (8 ignored, none failed) and all 547 SQL files. The CLI run passed 116 tests, and the focused SQL run passed all seven files. The committed PR range also passed the large-file check. ## Are there any user-facing changes? The new setting lets users opt into native grouped integer DISTINCT counts with one target partition. Default planning stays unchanged. --------- Co-authored-by: hushen <190065939+918154429@users.noreply.github.com>
…roupby-memory-usage # Conflicts: # datafusion/common/src/config.rs # datafusion/optimizer/src/single_distinct_to_groupby.rs # datafusion/sqllogictest/test_files/information_schema.slt # docs/source/user-guide/configs.md
e3dd29a to
c413a0e
Compare
|
@mkleen gentle reminder sir to review whenever you are free. |
jayzhan211
left a comment
There was a problem hiding this comment.
@mohitgurav20 Thanks for the updates and the benchmark numbers. Two things block this for me: with the setting on, float DISTINCT returns different results, and #25988 (merged today) already added single_distinct_prefer_native to this rule, so this would be a second setting for the same choice. Details are inline.
@mkleen one question for you: with a memory limit and spilling on, this setting never lowered the limit a query needed in my runs, and at 4 target partitions it raised it 4–12x (numbers in the config.rs comment). It only lowers peak memory when no limit is set. Is that still the case you want the setting for?
| | DataType::Float16 | ||
| | DataType::Float32 | ||
| | DataType::Float64 |
There was a problem hiding this comment.
With the setting on, float DISTINCT gives different answers: on (0.0, -0.0, 1.5), count(DISTINCT f) goes from 2 to 3 and avg(DISTINCT f) from 0.75 to 0.5. The native accumulators compare raw bits (Hashable in functions-aggregate-common/src/utils.rs), so -0.0 and 0.0 count as two values, while the rewrite's GROUP BY treats them as one. Could you keep floats on the rewrite and add the test below? It fails without this change.
| | DataType::Float16 | |
| | DataType::Float32 | |
| | DataType::Float64 |
Regression test for single_distinct_to_groupby.slt (place it while the setting is on)
# Float DISTINCT keeps the rewrite: the native accumulator counts -0.0 and 0.0 separately.
statement ok
CREATE TABLE t_prefer_mem_float AS SELECT * FROM (VALUES (1, CAST('0.0' AS DOUBLE)), (1, CAST('-0.0' AS DOUBLE)), (1, CAST('1.5' AS DOUBLE))) AS v(g, f);
query IIR
SELECT g, count(DISTINCT f), avg(DISTINCT f) FROM t_prefer_mem_float GROUP BY g;
----
1 2 0.75
statement ok
DROP TABLE t_prefer_mem_float;| /// When `true`, the `SingleDistinctToGroupBy` optimizer rule applies its cost guard | ||
| /// universally — including to lone-distinct queries that have no companion non-distinct | ||
| /// aggregates. The rewrite is skipped whenever the distinct function already has a | ||
| /// native `GroupsAccumulator`, reducing peak memory by avoiding the larger intermediate | ||
| /// `(group, distinct value)` result set that the two-phase plan produces. | ||
| /// | ||
| /// The trade-off is reduced parallelism: skipping the rewrite causes the engine to fall | ||
| /// back to serial per-group accumulation, which is typically 2–3× slower for queries | ||
| /// with high-cardinality distinct values. | ||
| /// | ||
| /// When `false` (the default), lone-distinct queries are always rewritten into the | ||
| /// two-phase `GROUP BY` so they can fan out across partitions and achieve maximum | ||
| /// throughput. | ||
| /// | ||
| /// See <https://github.com/apache/datafusion/issues/25932> |
There was a problem hiding this comment.
Under a memory limit, this setting needs more memory, not less: at 4 target partitions, sum(DISTINCT x) on the #25932 table needs 16M with the setting on and finishes within 4M without it (count: 48M vs 4M). The rewrite can spill its (group, value) pairs, but the native accumulator keeps each group's distinct values in memory, so with a few large groups it fails where the default finishes. Your numbers are right for a pool with no limit, so the doc should say that:
| /// When `true`, the `SingleDistinctToGroupBy` optimizer rule applies its cost guard | |
| /// universally — including to lone-distinct queries that have no companion non-distinct | |
| /// aggregates. The rewrite is skipped whenever the distinct function already has a | |
| /// native `GroupsAccumulator`, reducing peak memory by avoiding the larger intermediate | |
| /// `(group, distinct value)` result set that the two-phase plan produces. | |
| /// | |
| /// The trade-off is reduced parallelism: skipping the rewrite causes the engine to fall | |
| /// back to serial per-group accumulation, which is typically 2–3× slower for queries | |
| /// with high-cardinality distinct values. | |
| /// | |
| /// When `false` (the default), lone-distinct queries are always rewritten into the | |
| /// two-phase `GROUP BY` so they can fan out across partitions and achieve maximum | |
| /// throughput. | |
| /// | |
| /// See <https://github.com/apache/datafusion/issues/25932> | |
| /// When `true`, the `SingleDistinctToGroupBy` optimizer rule keeps a grouped | |
| /// DISTINCT aggregate native, instead of rewriting it into a two-phase `GROUP BY`, | |
| /// when every DISTINCT argument is a fixed-width, non-float type. This lowers peak | |
| /// memory when no memory limit is set, because the intermediate | |
| /// `(group, distinct value)` result is never built. | |
| /// | |
| /// Under a memory limit, the rewrite can spill those pairs, but the native | |
| /// accumulator keeps each group's distinct values in memory. With this option | |
| /// enabled, a query can need more memory, or fail where the rewrite finishes. | |
| /// | |
| /// Queries without `GROUP BY`, and other argument types, keep the rewrite. | |
| /// | |
| /// See <https://github.com/apache/datafusion/issues/25932> |
Smallest memory limit that finishes (release datafusion-cli, greedy pool, spilling on)
| query on the #25932 table (4M rows, 2,000 groups) | target partitions | setting off | setting on |
|---|---|---|---|
sum(DISTINCT x) |
1 | ≤4M | ≤4M (also 1M) |
sum(DISTINCT x) |
4 | ≤4M | 16M |
count(DISTINCT x) |
1 | ≤4M | ≤4M |
count(DISTINCT x) |
4 | ≤4M | 48M |
Two large groups at a 1M limit (SELECT v % 2, count(DISTINCT v) FROM generate_series(1, 100000) AS t(v) GROUP BY v % 2, batch_size = 128): the setting fails with Resources exhausted at 1 target partition for count and sum, and at 4 for count. The default finishes in every case. The same query is in aggregate_memory_spill.slt.
configs.md and information_schema.slt need regenerating afterwards.
| /// throughput. | ||
| /// | ||
| /// See <https://github.com/apache/datafusion/issues/25932> | ||
| pub single_distinct_aggregation_prefer_memory: bool, default = false |
There was a problem hiding this comment.
#25988 (merged today, not released yet) added single_distinct_prefer_native just below for the same choice, and with both settings count(DISTINCT int) at one target partition skips the rewrite under either one. Could this PR widen that setting's gate instead of adding a second setting? Because it hasn't been released, its name and doc can still change, and users would get one switch for "keep grouped DISTINCT native".
| // Aggregates with DistinctHandling::Unsupported (e.g. approx_median) cannot | ||
| // run natively with DISTINCT; they must always be rewritten. | ||
| if distinct_aggs | ||
| .iter() | ||
| .any(|(func, _)| func.distinct_handling() == DistinctHandling::Unsupported) | ||
| { | ||
| return Ok(true); | ||
| } |
There was a problem hiding this comment.
With the setting off, this return also changes main's behavior: SELECT g, count(*), stddev(DISTINCT x) FROM t GROUP BY g fails on main with STDDEV_POP(DISTINCT) aggregations are not available and returns the right rows here. That's a fix, but no test covers it at the default, and the description says the default is unchanged. Could you add this test before the SET … = true line in single_distinct_to_groupby.slt and mention the fix in the description?
# stddev cannot run DISTINCT itself, so it is rewritten even next to count(*).
query IIR rowsort
SELECT g, count(*), stddev(DISTINCT x) FROM t_prefer_mem GROUP BY g;
----
1 2 7.071067811865
2 1 NULL| // Resolve the data type of each distinct argument. | ||
| let all_fixed_width = distinct_aggs.iter().all(|(_, args)| { | ||
| args.iter().all(|arg| { | ||
| arg.get_type(input_schema) | ||
| .map(|dt| is_fixed_width_type(&dt)) | ||
| .unwrap_or(false) | ||
| }) | ||
| }); | ||
| if all_fixed_width { |
There was a problem hiding this comment.
Optional cleanup: .unwrap_or(false) silently keeps the rewrite when a type can't be resolved, while distinct_groups_support above returns the error with ?. This does the same with one pass:
| // Resolve the data type of each distinct argument. | |
| let all_fixed_width = distinct_aggs.iter().all(|(_, args)| { | |
| args.iter().all(|arg| { | |
| arg.get_type(input_schema) | |
| .map(|dt| is_fixed_width_type(&dt)) | |
| .unwrap_or(false) | |
| }) | |
| }); | |
| if all_fixed_width { | |
| let arg_types = distinct_aggs | |
| .iter() | |
| .flat_map(|(_, args)| args.iter()) | |
| .map(|arg| arg.get_type(input_schema)) | |
| .collect::<Result<Vec<_>>>()?; | |
| if arg_types.iter().all(is_fixed_width_type) { |
…roupby-memory-usage
|
Yeah that's a legit bug, thanks for catching it. The problem is is_fixed_width_type was lumping floats in with integers, which breaks down because the two codepaths disagree on -0.0/0.0 — GROUP BY treats them equal (IEEE 754), but the native accumulator hashes raw bits so they end up in separate buckets. Net result: count(DISTINCT f) changes from 2 to 3 on your example, which is wrong. Pulled floats out of is_fixed_width_type so prefer_memory never skips the rewrite for float columns. Added a plan-level unit test confirming the rewrite still fires with prefer_memory=true on a Float64 column, and the .slt test with (0.0, -0.0, 1.5) verifying the count stays at 2 under both settings. On the config overlap with #25988 — makes sense to let @mkleen weigh in first before I touch that. If the consensus is to fold both into one knob I'm happy to do that in a follow-up. |
|
@mkleen @jayzhan211 i have worked on it according to the changes you said . pls feel free to review and suggest for any changes . |
@jayzhan211 Yes, the idea was have an option to prefer memory-consumption-over-performance without setting a memory-limit. |
…o single_distinct_prefer_native
|
@mkleen @jayzhan211 Done — consolidated everything into Grouped all-distinct queries on fixed-width, non-float types (including |
Which issue does this PR close?
Closes #25932.
Rationale for this change
In queries with a single distinct aggregation,
SingleDistinctToGroupByrewrites the plan into a two-phaseGROUP BY. While this enables partition parallelism, it materializes intermediate(group, distinct_value)pairs, increasing peak intermediate memory (~538–599 MiB vs ~306 MiB in unconstrained memory runs).PR #25988 previously added
datafusion.optimizer.single_distinct_prefer_nativespecifically forCOUNT(DISTINCT int). Per review feedback from @mkleen, rather than introducing a separate knob (single_distinct_aggregation_prefer_memory), this PR consolidates both behaviors intodatafusion.optimizer.single_distinct_prefer_native.At
target_partitions = 1, nativeGroupsAccumulatorexecution avoids the two-phase grouping overhead and reduces peak memory for fixed-width non-floating types (such assum(DISTINCT int)andcount(DISTINCT int)). Whentarget_partitions > 1, queries continue using the two-phase rewrite to maintain partition fan-out and support spilling under memory limits.What changes are included in this PR?
Consolidated Configuration Knob:
datafusion.optimizer.single_distinct_prefer_native: bool(defaultfalse) to govern all native single distinct paths, dropping the separatesingle_distinct_aggregation_prefer_memorysetting.config.rsandconfigs.mdregarding memory, partition gating, and spilling trade-offs.Opt-in Native Execution Gating:
single_distinct_prefer_native = trueandtarget_partitions == 1, skips the two-phase rewrite if all distinct aggregate arguments are fixed-width, non-floating types (integers, dates, decimals) supported by nativeGroupsAccumulator.f32,f64) to preserve IEEE 754-0.0 == 0.0semantics (which differ betweenGROUP BYand raw-bit accumulator hashing).target_partitions > 1), queries with non-distinct companion aggregates, and aggregates with unsupported distinct handling (such asapprox_medianorstddev) continue with the two-phase rewrite.Documentation & Schema Synchronization:
docs/source/user-guide/configs.md.information_schema.slt,grouped_count_distinct.slt, andsingle_distinct_to_groupby.slt.Are these changes tested?
Yes:
single_distinct_to_groupby.rscovering flag states, partition gating (target_partitions == 1vs> 1), float argument exclusions, and unsupported aggregate fallbacks.EXPLAIN) and end-to-end query execution acrosssingle_distinct_to_groupby.sltandgrouped_count_distinct.slt.cargo clippy --all-targets --all-features -- -D warnings), and full optimizer test suite.Are there any user-facing changes?
Yes:
datafusion.optimizer.single_distinct_prefer_native. No breaking API changes.