Allow a non-distinct count alongside the single distinct aggregate in SingleDistinctToGroupBy - #24859
Draft
adriangb wants to merge 2 commits into
Draft
Allow a non-distinct count alongside the single distinct aggregate in SingleDistinctToGroupBy#24859adriangb wants to merge 2 commits into
adriangb wants to merge 2 commits into
Conversation
`SingleDistinctToGroupBy` rewrites `AGG(DISTINCT x)` into a two phase group by, which is what keeps a high cardinality distinct off the one-accumulator-per-group path in `GroupsAccumulatorAdapter`. The rule tolerated a non-distinct `sum`, `min` or `max` next to the distinct aggregate but bailed out on `count`, so the very common `count(*), count(DISTINCT x) ... GROUP BY` shape kept the unrewritten plan and its memory profile. Allow a non-distinct `count` as well. `count` is the one supported function whose outer phase is a different function: the inner group by counts the rows of each `(group, distinct value)` partition and the outer phase adds those partial counts up with `sum`, since count over a group is the sum of the counts of any partition of that group. Two details follow from that substitution: - `count` and `sum` are resolved from the session function registry and the rewrite only fires when the aggregate is that exact `count`, so a session without a registry or with its own `count` is left alone. - `count` returns a non-null 0 over an empty input while `sum` of no rows is NULL, which is reachable for an aggregate with no group by. The projection selects `CASE WHEN sum(alias) IS NOT NULL THEN sum(alias) ELSE 0 END`, which restores the 0 and keeps the column's type and nullability as `count` had them. The new sqllogictest file asserts every result twice, once with the optimizer disabled and once with it enabled, over data with NULL and all-NULL distinct values, an empty input, and `count(*)` versus `count(col)` versus `count(1)`. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`aggregate_distinct_with_having` round trips `SELECT a, count(distinct b) ... HAVING count(b) > 100` through substrait and asserts the plan comes back displaying identically. It passed only because the non-distinct `count` made `SingleDistinctToGroupBy` bail out, so the plan had no aliases in it. With the rule now allowing that `count`, the query is rewritten and the assertion fails. The failure is a pre-existing substrait gap rather than anything specific to this query: substrait carries no names for an aggregate's grouping and measure expressions, so the consumer derives them from the expressions themselves and the `alias1` and `alias2` names the rule introduces are lost. Any plan the rule rewrites fails the same way, including the plain `SELECT a, count(distinct b) FROM data GROUP BY a, c` that this change does not touch. Remove the rule from the session used by this one test, so it keeps covering the un-rewritten aggregate it was written for instead of depending on the rule bailing out. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
optimizer
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24859 +/- ##
========================================
Coverage 81.60% 81.60%
========================================
Files 1123 1123
Lines 408898 409051 +153
Branches 408898 409051 +153
========================================
+ Hits 333670 333810 +140
+ Misses 55625 55624 -1
- Partials 19603 19617 +14 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
Author
|
run benchmark clickbench_partitioned |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/single-distinct-to-groupby-allow-count (f47c045) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/single-distinct-to-groupby-allow-count (f47c045) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
adriangb
added a commit
to pydantic/datafusion
that referenced
this pull request
Sep 1, 2026
The two grouped `COUNT(DISTINCT <string>)` memory limit tests only reach the per group accumulators while `single_distinct_aggregation_to_group_by` declines to rewrite the query. They leant on `count(*)` for that, which the rule rejects only because `count` is missing from the `sum`/`min`/`max` allow list. apache#24859 proposes adding `count` to that list, which would rewrite the query, remove the accumulators, and leave both tests passing at any memory limit while still looking like they test something. Aggregate `avg(payload)` over a new `Int64` column instead. `avg` cannot be added to that list: the rule re-aggregates its own partial results over the deduplicated inner group by, and averaging per group averages of different sizes gives the wrong answer. That is why ClickBench Q9 keeps its distinct aggregate under apache#24859. Verified from the physical plan with apache#24859 cherry-picked on top of this branch: the `avg` query still plans as `aggr=[count(DISTINCT t.value), avg(t.payload)]`, while the `count(*)` query becomes `aggr=[count(alias1), sum(alias2)]` over an inner `GROUP BY group_key, value`, and drops from needing ~1.9 MB to ~0.9 MB. Re-swept both thresholds against the base commit. `Utf8` needs ~35.5 MB before and ~1.9 MB after; `Utf8View` needs ~123 MB before and ~2.5 MB after, so the 8 MB and 16 MB limits keep at least 4x margin on each side and are unchanged.
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
3 participants
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?
No existing issue. This was found while investigating a production out of memory. Happy to file one if you would like it tracked for the changelog.
Rationale for this change
SingleDistinctToGroupByrewritesAGG(DISTINCT x)into a two phase group by, which keeps a high cardinality distinct off the one-boxed-accumulator-per-group path inGroupsAccumulatorAdapter. The rule already tolerates a non-distinctsum,minormaxnext to the distinct aggregate, but bails out oncount, so the very commonshape keeps the unrewritten plan and its memory profile.
We hit this in production: a query of exactly that shape drove a process running DataFusion to 10.98 GB and death, and the single
count(*)was the only reason the rewrite did not apply.What changes are included in this PR?
This allows a non-distinct
countas well.countis the one supported function whose outer phase is a different function: the inner group by counts the rows of each(group, distinct value)partition, and the outer phase adds those partial counts up withsum, sincecountover a group is the sum of the counts of any partition of that group.Two details follow from the substitution:
countandsumcome from the session function registry (asreplace_distinct_aggregatealready does forfirst_value), and the rewrite only fires when the aggregate is that exactcount, compared by identity rather than by name. A session without a registry, or with its owncount, is left alone.countreturns a non-null0over an empty input whilesumof no rows is NULL, which is reachable for an aggregate with noGROUP BY: the inner aggregate emits no rows and the outer still emits one, soSELECT count(*), count(DISTINCT x) FROM emptywould returnNULL, 0instead of0, 0. The projection selectsCASE WHEN sum(alias) IS NOT NULL THEN sum(alias) ELSE 0 END, restoring the0and keeping the column's type and nullability ascounthad them.FILTERandORDER BYstill block the rewrite.Files touched beyond the rule itself:
datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt: the new coverage described below.datafusion/sqllogictest/test_files/clickbench.slt: the one existing snapshot in the repository that changes, discussed below.datafusion/substrait/tests/cases/roundtrip_logical_plan.rs:aggregate_distinct_with_havingnow builds its session with this rule removed, so it keeps round tripping the plan shape the test was written for.What is the testing strategy for this PR?
single_distinct_to_groupby.sltasserts every result twice, once underdatafusion.optimizer.max_passes = 0and once under the default, with identical expected blocks, so a null-handling or type error surfaces as a result mismatch rather than only a plan diff. It coverscount(*)vscount(1)vscount(col)grouped and ungrouped, a group whose distinct column is entirely NULL, a group with NULLs in both the distinct and summed columns, empty input three ways,HAVINGplusORDER BYon the rewritten count, and the production join shape.Exactly one existing snapshot in the repository changes: the ClickBench Q22
EXPLAIN, which is this shape verbatim. Its result block directly beneath, running on real ClickBench parquet, is unchanged. The physicalSortExec: TopK(fetch=10)moves from below the projection to above it, because the sort key is now aCASEoutput rather than a raw aggregate column. That is order-equivalent, since theCASEis the identity on every non-NULL input and thesumis never NULL in a grouped aggregate.Run locally: the full sqllogictest suite, plus
datafusion --test core_integration(1079),--test tpcds_planning(198) and-p datafusion --lib(444).cargo clippy -p datafusion-optimizer --all-targetsis clean.Benchmarks
Q22 is the only ClickBench query whose plan changes. Q9 (
RegionID, SUM, COUNT(*), AVG, COUNT(DISTINCT UserID)) still bails out, becauseAVGdisqualifies it.Measured on
clickbench_partitioned(100 files, ~100M rows), release builds of this branch and of the base commit it sat on at the time of the run, on a 12-core machine.A run-level A/B could not resolve a change this small here. Comparing the base binary against itself with
compare.pyreported 9 queries faster, 28 slower and 6 unchanged, with swings up to 1.58x, and two runs of the same base-vs-branch comparison gave opposite verdicts (11 faster / 21 slower, then 30 faster / 6 slower, with a 1.97x swing). Those tables measure background load, not the patch, because one arm is a full 43-query pass of about four minutes and load drifts between the arms.Instead the arms were paired per query, running base and branch back to back and alternating which goes first, over 40 repetitions. The 42 queries whose plans are unchanged then serve as an in-experiment control for residual bias.
Q22, net of the control bias (difference in differences, bootstrap CI):
The interval includes zero, so there is no measurable latency difference, and the 95% upper bound excludes a Q22 regression larger than about 1.5%. Pooled controls moved +0.36% [-0.64%, +1.05%], confirming the setup resolves effects of roughly 3% and no better.
This is latency-neutral on ClickBench, consistent with #11360, which found removing the rule entirely to be a wash. The case for the change rests on the memory behaviour of the rewritten plan, not on latency.
Not covered: memory. These runs used no
--memory-limit, so they do not exercise the spilling behaviour that motivates the rewrite. That is a separate experiment.Are there any user-facing changes?
No public API change and no change to query results. Plans for
SELECT ..., count(...), count(DISTINCT x) ... GROUP BY ...change shape, soEXPLAINoutput for that shape differs, and such queries should use substantially less memory. The ClickBench Q22 plan change above is the visible example.