Repository navigation
fix: skip SingleDistinctToGroupBy rewrite when single_distinct_prefer_native is set (closes #25932) - #26000
Open
mohitgurav20 wants to merge 41 commits into
Open
fix: skip SingleDistinctToGroupBy rewrite when single_distinct_prefer_native is set (closes #25932)#26000mohitgurav20 wants to merge 41 commits into
mohitgurav20 wants to merge 41 commits into
Conversation
mohitgurav20
force-pushed
the
fix/single-distinct-groupby-memory-usage
branch
from
October 4, 2026 16:36
527a544 to
9e8dfe6
Compare
mkleen
suggested changes
Oct 8, 2026
Contributor
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:
|
mohitgurav20
force-pushed
the
fix/single-distinct-groupby-memory-usage
branch
from
October 10, 2026 19:06
63e9873 to
0c76607
Compare
Contributor
Author
|
@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. |
Contributor
|
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 |
Contributor
|
@mohitgurav20 There are serious performance regressions in the clickhouse benchmark. See #26000 (comment) for more details. |
Contributor
Author
|
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.
mohitgurav20
force-pushed
the
fix/single-distinct-groupby-memory-usage
branch
from
October 17, 2026 07:46
287b932 to
ac76b0d
Compare
Contributor
Author
|
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! |
Contributor
|
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 |
Contributor
|
@mohitgurav20 Thanks! Looks good at first glance. I’ll go ahead and start the review. |
shehabgamin
pushed a commit
to lakehq/datafusion
that referenced
this pull request
Nov 8, 2026
## 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 <[email protected]>
…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
mohitgurav20
force-pushed
the
fix/single-distinct-groupby-memory-usage
branch
from
November 8, 2026 10:55
e3dd29a to
c413a0e
Compare
Contributor
Author
|
@mkleen gentle reminder sir to review whenever you are free. |
jayzhan211
requested changes
Nov 9, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
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?
Comment on lines
+128
to
+130
| | DataType::Float16 | ||
| | DataType::Float32 | ||
| | DataType::Float64 |
Contributor
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.
Suggested 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;
Comment on lines
+1943
to
+1957
| /// 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> |
Contributor
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:
Suggested change
| /// 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 |
Contributor
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".
Comment on lines
+200
to
+207
| // 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); | ||
| } |
Contributor
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
Comment on lines
+236
to
+244
| // 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 { |
Contributor
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:
Suggested change
| // 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
Contributor
Author
|
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. |
Contributor
Author
|
@mkleen @jayzhan211 i have worked on it according to the changes you said . pls feel free to review and suggest for any changes . |
Contributor
@jayzhan211 Yes, the idea was have an option to prefer memory-consumption-over-performance without setting a memory-limit. |
Contributor
…o single_distinct_prefer_native
Contributor
Author
|
@mkleen @jayzhan211 Done — consolidated everything into Grouped all-distinct queries on fixed-width, non-float types (including |
jayzhan211
requested changes
Nov 12, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
@mohitgurav20 thanks for folding this into one setting. I found two things to fix before merge: avg(DISTINCT <duration>) errors with the setting on, and the config doc says the setting lowers memory, but with many small groups it needs about 4x more. Details are inline.
Could you also update the title and description? They still name single_distinct_aggregation_prefer_memory and rewrite_pays, and "user-facing changes" is empty. Please mention that SELECT g, count(*), stddev(DISTINCT x) ... GROUP BY g now works at the default setting (it errors on main).
| | DataType::Time32(_) | ||
| | DataType::Time64(_) | ||
| | DataType::Timestamp(_, _) | ||
| | DataType::Duration(_) |
Contributor
There was a problem hiding this comment.
With the setting on, avg(DISTINCT d) over a Duration column now fails with AVG(DISTINCT) for (Duration(s) --> Duration(s)) not supported, while the rewrite returns rows. avg has no DISTINCT accumulator for durations. I ran every built-in aggregate against every type in this list with the setting on and off, and this is the only mismatch. Could you drop Duration here and add the test below? It fails without the change.
Suggested change
| | DataType::Duration(_) |
Test for single_distinct_to_groupby.slt, after DROP TABLE t_prefer_mem_float;
# `avg` has no DISTINCT accumulator for durations, so durations keep the rewrite.
query I?
SELECT g, avg(DISTINCT d) FROM (SELECT g, arrow_cast(x, 'Duration(Second)') AS d FROM t_prefer_mem) GROUP BY g ORDER BY g;
----
1 0 days 0 hours 0 mins 15 secs
2 0 days 0 hours 0 mins 10 secs
Comment on lines
+1948
to
+1953
| /// The native plan avoids materializing intermediate `(group, distinct value)` | ||
| /// pairs, lowering peak memory when no memory limit is set. However, | ||
| /// native DISTINCT must fit each group's distinct values in memory: under | ||
| /// a memory limit, the rewrite can spill those pairs, but the native | ||
| /// accumulator cannot, so queries that succeed by spilling with the rewrite | ||
| /// may fail with this option enabled. |
Contributor
There was a problem hiding this comment.
With many small groups, the native plan needs more memory, not less. At one target partition on 4M rows with 1M groups, sum(DISTINCT x) needs a 384M limit (no spilling) with this setting on vs 96M without it, and runs 3.8x slower. sum goes through GroupsAccumulatorAdapter, which keeps one accumulator per group. On the #25932 table (2,000 groups) it does halve memory (192M → 96M), so the doc should say when it helps:
Suggested change
| /// The native plan avoids materializing intermediate `(group, distinct value)` | |
| /// pairs, lowering peak memory when no memory limit is set. However, | |
| /// native DISTINCT must fit each group's distinct values in memory: under | |
| /// a memory limit, the rewrite can spill those pairs, but the native | |
| /// accumulator cannot, so queries that succeed by spilling with the rewrite | |
| /// may fail with this option enabled. | |
| /// The native plan avoids materializing intermediate `(group, distinct value)` | |
| /// pairs, which lowers peak memory when there are few groups with many | |
| /// distinct values each. With many small groups it can use more memory and | |
| /// run slower, because functions without a specialized DISTINCT | |
| /// `GroupsAccumulator` (such as `sum`) keep one accumulator per group. | |
| /// Native DISTINCT must also fit each group's distinct values in memory: | |
| /// under a memory limit, the rewrite can spill those pairs, but the native | |
| /// accumulator cannot, so queries that succeed by spilling with the rewrite | |
| /// may fail with this option enabled. |
configs.md and information_schema.slt need regenerating after this.
Comment on lines
+1484
to
+1534
| #[derive(Debug, Clone, PartialEq, Eq, Hash)] | ||
| struct UnsupportedUdaf { | ||
| signature: Signature, | ||
| } | ||
|
|
||
| impl UnsupportedUdaf { | ||
| fn new() -> Self { | ||
| Self { | ||
| signature: Signature::exact( | ||
| vec![DataType::UInt32], | ||
| Volatility::Immutable, | ||
| ), | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl AggregateUDFImpl for UnsupportedUdaf { | ||
| fn name(&self) -> &str { | ||
| "unsupported" | ||
| } | ||
|
|
||
| fn signature(&self) -> &Signature { | ||
| &self.signature | ||
| } | ||
|
|
||
| fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { | ||
| Ok(DataType::UInt32) | ||
| } | ||
|
|
||
| fn accumulator( | ||
| &self, | ||
| _acc_args: AccumulatorArgs, | ||
| ) -> Result<Box<dyn Accumulator>> { | ||
| unimplemented!("not needed for optimizer tests") | ||
| } | ||
|
|
||
| fn distinct_handling(&self) -> DistinctHandling { | ||
| DistinctHandling::Unsupported | ||
| } | ||
| } | ||
|
|
||
| fn unsupported_distinct(expr: Expr) -> Expr { | ||
| Expr::AggregateFunction(AggregateFunction::new_udf( | ||
| Arc::new(AggregateUDF::from(UnsupportedUdaf::new())), | ||
| vec![expr], | ||
| true, | ||
| None, | ||
| vec![], | ||
| None, | ||
| )) | ||
| } |
Contributor
There was a problem hiding this comment.
Optional cleanup: crate::test::udfs::DistinctHandlingTestUDAF already provides an aggregate with a chosen DistinctHandling, so this struct can go. The snapshot stays the same:
Suggested change
| #[derive(Debug, Clone, PartialEq, Eq, Hash)] | |
| struct UnsupportedUdaf { | |
| signature: Signature, | |
| } | |
| impl UnsupportedUdaf { | |
| fn new() -> Self { | |
| Self { | |
| signature: Signature::exact( | |
| vec![DataType::UInt32], | |
| Volatility::Immutable, | |
| ), | |
| } | |
| } | |
| } | |
| impl AggregateUDFImpl for UnsupportedUdaf { | |
| fn name(&self) -> &str { | |
| "unsupported" | |
| } | |
| fn signature(&self) -> &Signature { | |
| &self.signature | |
| } | |
| fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> { | |
| Ok(DataType::UInt32) | |
| } | |
| fn accumulator( | |
| &self, | |
| _acc_args: AccumulatorArgs, | |
| ) -> Result<Box<dyn Accumulator>> { | |
| unimplemented!("not needed for optimizer tests") | |
| } | |
| fn distinct_handling(&self) -> DistinctHandling { | |
| DistinctHandling::Unsupported | |
| } | |
| } | |
| fn unsupported_distinct(expr: Expr) -> Expr { | |
| Expr::AggregateFunction(AggregateFunction::new_udf( | |
| Arc::new(AggregateUDF::from(UnsupportedUdaf::new())), | |
| vec![expr], | |
| true, | |
| None, | |
| vec![], | |
| None, | |
| )) | |
| } | |
| fn unsupported_distinct(expr: Expr) -> Expr { | |
| Expr::AggregateFunction(AggregateFunction::new_udf( | |
| Arc::new(AggregateUDF::from(udfs::DistinctHandlingTestUDAF::new( | |
| "unsupported", | |
| DistinctHandling::Unsupported, | |
| ))), | |
| vec![expr], | |
| true, | |
| None, | |
| vec![], | |
| None, | |
| )) | |
| } |
Contributor
Author
|
@jayzhan211 thanks for your review I'll soon catch up with the changes sirr. |
mohitgurav20
force-pushed
the
fix/single-distinct-groupby-memory-usage
branch
from
November 12, 2026 14:39
cf8fe84 to
678fe74
Compare
Contributor
Author
|
@jayzhan211 sir i have addressed the changes just go through it once whenever u are freee . |
mohitgurav20
changed the title
jayzhan211
approved these changes
Nov 14, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
@mohitgurav20 thanks, everything from the last round is addressed. One optional doc fix inline.
Comment on lines
+1943
to
+1959
| /// 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. | ||
| /// | ||
| /// 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. | ||
| /// The native plan avoids materializing intermediate `(group, distinct value)` | ||
| /// pairs, which lowers peak memory when there are few groups with many | ||
| /// distinct values each. With many small groups it can use more memory and | ||
| /// run slower, because functions without a specialized DISTINCT | ||
| /// `GroupsAccumulator` (such as `sum`) keep one accumulator per group. | ||
| /// Native DISTINCT must also fit each group's distinct values in memory: | ||
| /// under a memory limit, the rewrite can spill those pairs, but the native | ||
| /// accumulator cannot, so queries that succeed by spilling with the rewrite | ||
| /// may fail with this option enabled. | ||
| /// | ||
| /// Queries without `GROUP BY`, and other argument types, keep the rewrite. |
Contributor
There was a problem hiding this comment.
Optional cleanup: this doc no longer says the setting only applies at target_partitions = 1 with all-DISTINCT aggregates, so with default settings turning it on does nothing. Its first paragraph also says memory always goes down, while the second says many small groups need more. Last, "the native accumulator cannot" spill was my wording and it's wrong: at one partition with 50k groups it spilled 12 times and matched the rewrite. What it can't do is split one large group. This version fixes all three (I checked that dev/update_config_docs.sh and information_schema.slt pass with it):
Suggested change
| /// 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. | |
| /// | |
| /// 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. | |
| /// The native plan avoids materializing intermediate `(group, distinct value)` | |
| /// pairs, which lowers peak memory when there are few groups with many | |
| /// distinct values each. With many small groups it can use more memory and | |
| /// run slower, because functions without a specialized DISTINCT | |
| /// `GroupsAccumulator` (such as `sum`) keep one accumulator per group. | |
| /// Native DISTINCT must also fit each group's distinct values in memory: | |
| /// under a memory limit, the rewrite can spill those pairs, but the native | |
| /// accumulator cannot, so queries that succeed by spilling with the rewrite | |
| /// may fail with this option enabled. | |
| /// | |
| /// Queries without `GROUP BY`, and other argument types, keep the rewrite. | |
| /// When `true` and `target_partitions` is 1, the `SingleDistinctToGroupBy` | |
| /// optimizer rule keeps a grouped aggregate whose aggregates are all DISTINCT | |
| /// native, instead of rewriting it into a two-phase `GROUP BY`, when every | |
| /// DISTINCT argument is an integer, boolean, decimal, date, time, timestamp | |
| /// or interval type. | |
| /// | |
| /// The native plan avoids materializing intermediate `(group, distinct value)` | |
| /// pairs, which lowers peak memory when there are few groups with many | |
| /// distinct values each. With many small groups it can use more memory and | |
| /// run slower, because functions without a specialized DISTINCT | |
| /// `GroupsAccumulator` (such as `sum`) keep one accumulator per group. | |
| /// Native DISTINCT must also fit each group's distinct values in memory: | |
| /// under a memory limit, it can spill whole groups but not split one group's | |
| /// distinct values, so a query with a few large groups that succeeds by | |
| /// spilling with the rewrite may fail with this option enabled. | |
| /// | |
| /// Queries without `GROUP BY`, and other argument types, keep the rewrite. |
Contributor
Author
There was a problem hiding this comment.
@jayzhan211 sirr i Applied the updated doc wording across config.rs, configs.md, and information_schema.slt. Verified information_schema.slt and the Prettier checks pass cleanly. Thanks for catching the spill distinction.
…group spill behavior
alamb
reviewed
Nov 14, 2026
alamb
left a comment
alamb
left a comment
Contributor
There was a problem hiding this comment.
Thank you @mohitgurav20 and @jayzhan211 . I think this code looks ok, but I think we should clean up the tests and comments if possible before merging
cc @mkleen
| /// and every aggregate explicitly supports a specialized GroupsAccumulator | ||
| /// for its DISTINCT argument types. Defaults to false, preserving the | ||
| /// two-phase GROUP BY rewrite. | ||
| /// When `true` and `target_partitions` is 1, the `SingleDistinctToGroupBy` |
Contributor
There was a problem hiding this comment.
Can we please rephrase this so it starts with the impact of this knob -- I think it is "queries with a single DISTINCT aggregate will take less memory but also run more slowly"?
i think all the background of the tradeoff is way too much detail for the config docs -- it probably belongs on the SingleDistinctToGroupBy docs, with perhaps a link from the config
|
|
||
| /// Check whether all aggregate exprs are eligible for the single-distinct-to-group-by rewrite. | ||
| /// | ||
| /// With `use_native_grouped_distinct` enabled (which requires one target partition, |
Contributor
There was a problem hiding this comment.
can you please summarieze this paragraph -- it has too much irrelevant detail (eg. restating what use_native_grouped_distinct does in English)
| Ok(()) | ||
| } | ||
|
|
||
| fn unsupported_distinct(expr: Expr) -> Expr { |
Contributor
There was a problem hiding this comment.
If these can be .slt tests, I think we should rewrite them to be slt -- the rust is far more verbose
Contributor
Author
|
thanks @alamb addressed all three points in the latest commit: Rephrased the config doc to lead with the memory vs performance impact, moving the deeper background into SingleDistinctToGroupBy docs. |
Contributor
Author
|
all checks are running green pls fell free to review this pr . |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 memory when there are few groups with many distinct values each (~538–599 MiB vs ~306 MiB in unconstrained memory runs on #25932).PR #25988 previously introduced
datafusion.optimizer.single_distinct_prefer_nativeforCOUNT(DISTINCT int). Per review feedback from @mkleen and @jayzhan211, this PR consolidates both intodatafusion.optimizer.single_distinct_prefer_nativerather than adding a second knob.At
target_partitions = 1, nativeGroupsAccumulatorexecution avoids the two-phase grouping overhead and reduces peak memory for queries with few large groups on fixed-width types (e.g.sum(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.Additionally, queries with unsupported distinct aggregates alongside non-distinct aggregates (such as
SELECT g, count(*), stddev(DISTINCT x) ... GROUP BY g) now rewrite correctly at default settings instead of erroring at runtime as on main.What changes are included in this PR?
Consolidated Configuration Knob:
datafusion.optimizer.single_distinct_prefer_native: bool(defaultfalse) to govern native single distinct execution.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 primitive types (integers, dates, decimals) supported natively.f32,f64) to preserve IEEE 754-0.0 == 0.0semantics.Durationtypes (sinceavglacks a native distinct accumulator for durations).target_partitions > 1, non-distinct companion aggregates are present, or distinct handling is unsupported (approx_median,stddev).Documentation & Schema Synchronization:
docs/source/user-guide/configs.mdandinformation_schema.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.single_distinct_to_groupby.sltcovering plans, execution, duration distinct handling, andstddev(DISTINCT x)default execution.Are there any user-facing changes?
Yes:
datafusion.optimizer.single_distinct_prefer_nativeto cover other fixed-width distinct aggregates likesum(DISTINCT int).SELECT g, count(*), stddev(DISTINCT x) ... GROUP BY gthat previously errored on main now succeed at default settings.