Allow a non-distinct count alongside the single distinct aggregate in SingleDistinctToGroupBy

[adriangb]

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

SingleDistinctToGroupBy rewrites AGG(DISTINCT x) into a two phase group by, which keeps a high cardinality distinct off the one-boxed-accumulator-per-group path in GroupsAccumulatorAdapter. The rule already tolerates a non-distinct sum, min or max next to the distinct aggregate, but bails out on count, so the very common

SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g

shape 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 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 the substitution:

• count and sum come from the session function registry (as replace_distinct_aggregate already does for first_value), and the rewrite only fires when the aggregate is that exact count, compared by identity rather than by name. 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 inner aggregate emits no rows and the outer still emits one, so SELECT count(*), count(DISTINCT x) FROM empty would return NULL, 0 instead of 0, 0. The projection selects CASE WHEN sum(alias) IS NOT NULL THEN sum(alias) ELSE 0 END, restoring the 0 and keeping the column's type and nullability as count had them.

FILTER and ORDER BY still 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_having now 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.slt asserts every result twice, once under datafusion.optimizer.max_passes = 0 and 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 covers count(*) vs count(1) vs count(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, HAVING plus ORDER BY on 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 physical SortExec: TopK(fetch=10) moves from below the projection to above it, because the sort key is now a CASE output rather than a raw aggregate column. That is order-equivalent, since the CASE is the identity on every non-NULL input and the sum is 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-targets is clean.

Benchmarks

Q22 is the only ClickBench query whose plan changes. Q9 (RegionID, SUM, COUNT(*), AVG, COUNT(DISTINCT UserID)) still bails out, because AVG disqualifies 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.py reported 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):

-2.03%   95% CI [-5.48%, +1.39%]

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, so EXPLAIN output for that shape differs, and such queries should use substantially less memory. The ClickBench Q22 plan change above is the visible example.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 84.69945% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.60%. Comparing base (da89c7c) to head (f47c045).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...fusion/optimizer/src/single_distinct_to_groupby.rs	84.69%	9 Missing and 19 partials ⚠️

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[adriangb]

run benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: 4G

[adriangbot]

🤖 Benchmark running (GKE) | trigger
Instance: c4a-highmem-16 (12 vCPU / 65 GiB) | Linux bench-c5500225568-2072-x9l4r 6.12.94+ #1 SMP Fri Jul 17 09:42:57 UTC 2026 aarch64 GNU/Linux

CPU Details (lscpu)

Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Comparing claude/single-distinct-to-groupby-allow-count (f47c045) to da89c7c (merge-base) diff

Run configuration

run benchmark clickbench_partitioned
env:
  DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"

Results will be posted here when complete

---

File an issue against this benchmark runner