bench: add a ClickBench extended query for a grouped COUNT(DISTINCT) over a string - #25026
Conversation
…string `COUNT(DISTINCT)` has a specialized `GroupsAccumulator` for the integer types and for no other type. Every other type falls back to `GroupsAccumulatorAdapter`, which holds one boxed `Accumulator`, and therefore one hash table, for each group, so the cost of that fallback scales with the group cardinality. Extended Q2 is the only query in either suite that puts a `COUNT(DISTINCT)` on a non-integer column, and it groups by `BrowserCountry`. Nothing exercises the adapter at a high group cardinality, and nothing covers a lone `COUNT(DISTINCT <string>)` next to a non-distinct `COUNT(*)`: standard Q8, Q10, Q11 and Q13 hold a distinct aggregate that stands alone, Q9 carries an `AVG`, and Q22, the one query that does pair a lone distinct aggregate with a non-distinct count, counts distinct `UserID`, which is an `Int64`. Extended Q14 is that query. It groups by `SearchPhrase` and counts distinct `MobilePhoneModel` beside a `COUNT(*)`. Extended queries are discovered from the directory, so this needs no change to the runner. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #25026 +/- ##
==========================================
+ Coverage 81.69% 81.85% +0.15%
==========================================
Files 1127 1127
Lines 415471 418811 +3340
Branches 415471 418811 +3340
==========================================
+ Hits 339424 342815 +3391
+ Misses 56108 56027 -81
- Partials 19939 19969 +30 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kumarUjjawal
left a comment
There was a problem hiding this comment.
Thank you @adriangb
I was thinking since we are also using the clickbench_extended should we add this to there as well?
The file only mirrored q0-q6; q7-q13 were added to benchmarks/queries/clickbench/extended/ without a corresponding sqllogictest entry, and this PR's new q14 would have widened that gap. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Good call. I noticed multiple queries were missing, so I added them all. Since this increased the scope I'll let you re-approve before merging. |
Thanks, Ping me when you add them. |
|
Sorry forgot to push! The change is now in 23a2181 |
The extended ClickBench queries exist in two runners. `benchmarks/queries/ clickbench/extended/` feeds `dfbench clickbench --queries-path`, and `benchmarks/sql_benchmarks/clickbench_extended/benchmarks/` feeds `benchmark_runner clickbench_extended`. A query added to only one of them is invisible to the other. Add the `.benchmark` file, structurally identical to q13 apart from the query itself. Suite files are discovered from the directory, so nothing else changes: `benchmark_runner clickbench_extended --list` now reports 15 queries. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… SingleDistinctToGroupBy (apache#24859) ## Which issue does this PR close? No existing issue. We found this while investigating an out-of-memory. Happy to file one if you want a changelog entry. ## Rationale for this change ### The query ```sql SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g; ``` A grouped `count(DISTINCT <string>)` next to a plain `count(*)`. It is a very common shape, and on `main` it uses several times more memory than it needs to. To reproduce, in `datafusion-cli`. This writes 4,000,000 rows in 500,000 groups, with 2,000,000 distinct `(g, x)` pairs: ```sql COPY ( SELECT value % 500000 AS g, 'id-' || CAST(value % 2000000 AS VARCHAR) AS x FROM generate_series(1, 4000000) ) TO 'repro.parquet' STORED AS PARQUET; CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'repro.parquet'; EXPLAIN FORMAT INDENT SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g; ``` ### The plan today ```text Projection: t.g, count(Int64(1)) AS count(*), count(DISTINCT t.x) Aggregate: groupBy=[[t.g]], aggr=[[count(Int64(1)), count(DISTINCT t.x)]] TableScan: t projection=[g, x] ``` A single `Aggregate` that computes the distinct count directly. `count(DISTINCT x)` has a specialized `GroupsAccumulator` for the integer types and for no others, so over a string it falls back to `GroupsAccumulatorAdapter`, which holds one boxed `Accumulator` — and therefore one hash table — **for every group**. At 500,000 groups that is 500,000 hash tables. There is a second effect once there is more than one partition. This plan's partial aggregate sees every group in every partition, so those per-group accumulators are duplicated `target_partitions` times. ### What already exists DataFusion has a rule for exactly this, `SingleDistinctToGroupBy`. It rewrites `AGG(DISTINCT x)` into a two phase group by: an inner aggregate that groups by `(group keys, x)`, so the distinct-ing is done by the hash table that group by already builds, and an outer aggregate over its output. One hash table instead of one per group, and it hash-partitions, so distinct values are split across partitions rather than replicated. The rule accepts a non-distinct `sum`, `min` or `max` alongside the distinct aggregate. **It rejects a non-distinct `count`.** So a single `count(*)` is enough to keep the slow plan. ### The change Allow that `count`. `count` is the one supported companion whose outer phase must be a *different* function. The inner group by counts the rows of each `(group, distinct value)` partition; the outer phase adds those partial counts with `sum`, because a count over a group is the sum of the counts of any partition of that group. ### The plan after ```text Projection: t.g, CASE WHEN sum(alias2) IS NOT NULL THEN sum(alias2) ELSE Int64(0) END AS count(*), count(alias1) AS count(DISTINCT t.x) Aggregate: groupBy=[[t.g]], aggr=[[sum(alias2), count(alias1)]] Aggregate: groupBy=[[t.g, t.x AS alias1]], aggr=[[count(Int64(1)) AS alias2]] TableScan: t projection=[g, x] ``` Results are identical. The `CASE` is the one place the two phases disagree: over an empty input the inner group by emits no rows, and `sum` of no rows is NULL where `count` is `0`. Restoring the `0` also preserves `count`'s non-nullability. ### When this is allowed The rewrite is not free. Every other aggregate moves down into the inner group by, which holds a row per `(group, distinct value)` pair rather than per group, and keeps its state at that finer grain. What pays for that is taking the distinct aggregate off the adapter — so if the distinct aggregate was never on the adapter, there is nothing to buy and only the inner group by to pay for. So the new `count` is gated on the distinct aggregate reporting that it has **no** specialized `GroupsAccumulator` for its argument types. ClickBench q22 is the real case: its `count(DISTINCT "UserID")` is over an `Int64`, which has one, so q22 keeps its current plan. ## What changes are included in this PR? Five files. The rule, the gate it needs, and tests. ### The rule `datafusion/optimizer/src/single_distinct_to_groupby.rs` accepts a non-distinct `count` as a companion, and gives it `sum` as its outer phase. `count` and `sum` are resolved from the session function registry, as `replace_distinct_aggregate` already does for `first_value`. The rewrite fires only for that exact `count`, compared by identity rather than by name, so a session with its own `count`, or with no registry, keeps the previous behaviour. `FILTER` and `ORDER BY` still block the rewrite. ### The gate An optimizer rule has no `AccumulatorArgs` to call `AggregateUDFImpl::groups_accumulator_supported` with, and `datafusion-optimizer` cannot depend on `datafusion-functions-aggregate` to read `count`'s type list. `datafusion/optimizer/Cargo.toml` names the intended way out: > If you want to add special handling for a specific function, use the methods on the `ScalarUDFImpl` or `AggregateUDFImpl` traits (or add a new method to those traits). So this adds one trait method: ```rust fn groups_accumulator_supported_for_types( &self, arg_types: &[DataType], is_distinct: bool, ) -> Option<bool> { None } ``` `Count` is the only implementor, and `Count::groups_accumulator_supported` now delegates to it, so there is one list of supported types rather than two that can drift. The default `None` means *the implementation does not answer this question*. It is not a third answer: a caller must not read it as either `Some(true)` or `Some(false)`. This rule rewrites only on `Some(false)`, the one answer that positively reports a call on the adapter. Every aggregate except `count` answers `None` today and so keeps its current behaviour. The gate covers only the `count` this PR adds. A plan that already qualifies through a non-distinct `sum`, `min` or `max` is rewritten as before, over any distinct argument type — see [Pre-existing regressions](#pre-existing-regressions-not-introduced-here). ### Tests - `datafusion/sqllogictest/test_files/single_distinct_to_groupby.slt`, new. - `datafusion/substrait/tests/cases/roundtrip_logical_plan.rs`: `aggregate_distinct_with_having` builds its session without this rule, so it keeps round tripping the plan shape it was written for; a companion test covers the rewritten plan's schema and results through Substrait. (The rule's aliases have no Substrait representation, so a rewritten plan does not round trip to an *identical* plan.) **No existing snapshot in the repository changes.** ## 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. A null-handling or type error therefore shows up as a result mismatch, not only as a plan difference. The table carries the same values in a `VARCHAR` column and an `INT` column, which are the two sides of the gate, and the file asserts both: the `VARCHAR` distinct rewrites with a `count` beside it, the `INT` distinct does not, and the `INT` distinct still rewrites when it qualifies through `sum`. It also asserts that an aggregate answering `None` stays unrewritten (`sum(DISTINCT v)` and `min(DISTINCT v)` beside a `count(*)` keep the plan they have on `main`). Remaining coverage: `count(*)` vs `count(1)` vs `count(col)`, grouped and ungrouped; a group whose distinct column is entirely NULL; NULLs in both the distinct and the summed column; empty input in three shapes; `HAVING` with `ORDER BY` on the rewritten count; and the same aggregates over a join. The rule's unit tests cover both sides of the gate directly, and the `None` answer twice — once with `sum(DISTINCT b)`, once with a test aggregate that leaves the new method at its default. Run locally on the rebased head, all passing: `datafusion-optimizer`, `datafusion-expr` and `datafusion-functions-aggregate` lib and integration tests, the substrait roundtrip suite, and the whole sqllogictest suite at 511 of 511 files. ## Benchmarks ### Macro: ClickBench `clickbench_extended` q14 is the one query in any suite with this shape — a lone `COUNT(DISTINCT <string>)` next to a non-distinct `COUNT(*)`, grouped by a high cardinality string. It landed on `main` in apache#25026. Three runs of this branch against merge base `16ace4f`, `DATAFUSION_RUNTIME_MEMORY_LIMIT: 16G` ([trigger](apache#24859 (comment))): | run | q14 fastest | q14 median of 5 | q14 peak pool | | --- | --- | --- | --- | | 1 | 2471 ms → 546 ms, **4.53x** | 4.05x | 4.4 GiB → 1.3 GiB, **-70.5%** | | 2 | 2534 ms → 575 ms, **4.40x** | 4.15x | 4.4 GiB → 1.3 GiB, **-69.7%** | | 3 | 3909 ms → 1165 ms, **3.36x** | 3.39x | 4.4 GiB → 1.3 GiB, **-70.3%** | 100,000,000 rows of real ClickBench data. The memory figure reproduces to a tenth of a percent across the three runs. The limit is 16G because the base side peaks at 4.4 GiB; at 4G the baseline OOMs and there is nothing to compare. Run 3 was on a contended machine — its suite total is 42.1 s against 30.8 s and 31.0 s — and it reports q0 at 1.30x slower, q7 at 1.14x and q8 at 1.13x. None of those three can change plan here (q0 has three distinct arguments; q7 and q8 have no distinct aggregate at all), so they are an in-experiment control that says what that run's noise floor was. #### The standard suite, and q22 `q22` is the only query in the standard suite that reaches the new gate, and the gate excludes it. A `clickbench_partitioned` run ([trigger](apache#24859 (comment)), `DATAFUSION_RUNTIME_MEMORY_LIMIT: 4G`): | | base `16ace4f` | this branch | | | --- | --- | --- | --- | | q22 wall clock | 975.75 ms | 976.75 ms | no change | | q22 peak pool | 2.5 MiB | 3.0 MiB | +22.0% | | suite total | 25985 ms | 26073 ms | +0.3% | | queries changed | — | — | **0 of 43** | The `+22.0%` is measurement noise on a 2.5 MiB reservation, not an effect of this PR. Three things say so: 1. **The plan is unchanged.** `EXPLAIN` on q22's exact shape against this branch produces a single `Aggregate` — the rule does not fire, because `COUNT(DISTINCT "UserID")` is over an `Int64`, which has a specialized `GroupsAccumulator`: ```text Sort: c DESC NULLS FIRST, fetch=10 Projection: hits.SearchPhrase, min(hits.URL), min(hits.Title), count(Int64(1)) AS count(*) AS c, count(DISTINCT hits.UserID) Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)]] Filter: ... TableScan: hits projection=[SearchPhrase, URL, Title, UserID] ``` 2. **The wall clock is flat** at 0.1%. 3. **The run's memory noise floor is wider than the reading.** 9 of the 43 queries move by more than 5% in both directions, including q23 at **-26.1%** and q7 at **-19.7%** — neither of which can change plan under this rule. **No other query changes plan, and that set is empty by construction, not by measurement.** The rule needs exactly one distinct argument and accepts only `sum`, `min`, `max` and now `count` as companions: | query | why it cannot move | | --- | --- | | `extended` q0, q1, q2 | three, three and four distinct arguments, so `fields_set.len() != 1` | | q4, q5, q8, q10, q11, q13 | a lone distinct aggregate with no non-distinct companion, so `main` already rewrites them | | q9 | carries an `AVG`, which the rule has never accepted | | q22 | the only standard query that reaches the new gate, and its `COUNT(DISTINCT "UserID")` is over an `Int64`, which has a specialized `GroupsAccumulator` | ### Micro: the reproduction above, swept Base is `e1ca94fb11`, 12 commits behind the `16ace4f` merge base used for the ClickBench runs above; both are from the same day, and none of the 12 changes aggregate runtime behaviour (the only one touching `count_distinct` is apache#25020, which is test-gating only). Both sides are release builds of the same harness: a `GreedyMemoryPool` of 64 GiB (never reached) wrapped in the in-tree `PeakRecordingPool`; one parquet file; the result streamed rather than collected; the pool high-water mark and `getrusage` peak RSS reported. Every figure is the median of three runs. At `target_partitions = 1` the three runs agreed to the byte in every cell; at 8, to within 1%, except two branch cells at 7% and 19%. Every shape is 4,000,000 rows, query `SELECT g, count(*), count(DISTINCT x) FROM t GROUP BY g`. Each file carries the distinct argument three times — as `Utf8`, `Utf8View` and `BIGINT` — written with an explicit `arrow_cast` so the type is the file's and not a reader setting. **`target_partitions = 8`** | groups | distinct values | argument | peak pool | peak RSS | wall clock | | --- | --- | --- | --- | --- | --- | | 500,000 | 2,000,000 | `Utf8` | 2.16x better | 2.56x better | 999 ms → 48 ms | | 500,000 | 2,000,000 | `Utf8View` | 2.68x better | 2.76x better | 1036 ms → 40 ms | | 500,000 | 1,000,000 | `Utf8` | 2.83x better | 3.40x better | 1020 ms → 41 ms | | 500,000 | 1,000,000 | `Utf8View` | 3.15x better | 3.43x better | 966 ms → 34 ms | | 10 | 2,000,000 | `Utf8` | 1.34x better | 1.85x better | 80 ms → 41 ms | | 10 | 2,000,000 | `Utf8View` | 1.87x better | 2.24x better | 73 ms → 33 ms | | 1 | 2,000,000 | `Utf8` | 1.55x better | 2.14x better | 202 ms → 45 ms | | 1 | 2,000,000 | `Utf8View` | 1.88x better | 2.51x better | 158 ms → 34 ms | | 2,000 | 4,000,000 | `Utf8` | **1.30x worse** | 2.03x better | 180 ms → 43 ms | | 2,000 | 4,000,000 | `Utf8View` | **1.10x worse** | 2.55x better | 185 ms → 38 ms | | 500,000 | 2,000,000 | `BIGINT` | 1.00x, plan unchanged | 1.01x | 43 ms → 43 ms | **`target_partitions = 1`** | groups | distinct values | argument | peak pool | peak RSS | wall clock | | --- | --- | --- | --- | --- | --- | | 500,000 | 2,000,000 | `Utf8` | 2.18x better | 1.77x better | 712 ms → 193 ms | | 500,000 | 2,000,000 | `Utf8View` | 2.70x better | 2.05x better | 684 ms → 178 ms | | 500,000 | 1,000,000 | `Utf8` | 3.31x better | 1.93x better | 626 ms → 160 ms | | 500,000 | 1,000,000 | `Utf8View` | 4.42x better | 2.34x better | 599 ms → 138 ms | | 500,000 | 4,000,000 | `Utf8` | 1.61x better | 1.70x better | 825 ms → 242 ms | | 500,000 | 4,000,000 | `Utf8View` | 1.96x better | 2.17x better | 755 ms → 232 ms | | 2,000 | 2,000,000 | `Utf8` | 1.01x better | 1.41x better | 402 ms → 161 ms | | 2,000 | 2,000,000 | `Utf8View` | 1.26x better | 1.49x better | 358 ms → 139 ms | | 2,000 | 4,000,000 | `Utf8` | 1.01x better | 1.17x better | 429 ms → 203 ms | | 2,000 | 4,000,000 | `Utf8View` | 1.26x better | 1.36x better | 407 ms → 176 ms | Three things worth calling out: - **Wall clock was not the point of this PR and is the largest effect.** At 8 partitions and 500,000 groups the query goes from about a second to about 40 ms. That is the accumulator-per-group construction cost, which grows with partitions on the unrewritten side while the rewritten side parallelizes. - **Peak RSS improves in every measured shape**, including the two where the pool peak regresses. The unrewritten plan's RSS sits far above what it reserves; millions of small independent hash tables fragment in a way one large table does not. - **The `BIGINT` row is the control for the gate.** The plan is identical on both sides, and so is every measurement. This is the q22 case, measured directly. ### Where the rewrite stops paying The crossover is in the *density* of distinct values, not the group count. It arrives only when nearly every row holds a distinct value, and the loss is bounded: 1.30x more peak pool for `Utf8` and 1.10x for `Utf8View`, at 2,000 groups over 4,000,000 fully distinct values at 8 partitions, and a wash at one partition. Those same two cells use about half the RSS and run four times faster. At 500,000 groups over 4,000,000 fully distinct values — the same density — the rewrite is 1.61x to 1.96x better again, because the unrewritten side now also pays for 500,000 accumulators. ### The gate is a proxy The gate asks which accumulator the distinct aggregate gets. That is not the true discriminator. What decides the outcome is the cost per distinct value on each side. The rewrite materializes one hash table row per distinct `(group keys, x)` pair, plus one accumulator slot per companion aggregate at that grain. It wins when the unrewritten accumulator costs more than that per value, and loses when it costs less. Proxy and discriminator agree for `count(DISTINCT x)`, which is the only case this PR opens. They disagree elsewhere. #### Pre-existing regressions, not introduced here Re-measured on the same base, over 4,000,000 rows in 2,000 groups with 2,000,000 distinct values, comparing the rewritten plan against the same query with a `count(*)` added to hold the rule off: - `sum(DISTINCT int_col)` and `avg(DISTINCT int_col)` have no distinct groups accumulator, and `main` rewrites both today. Both regress: **2.67x** more peak pool at one partition, 1.79x at eight. - `min(DISTINCT x)` is far worse. It is the same value as `min(x)`, and `min_max` correctly ignores `is_distinct`, so the unrewritten plan holds one scalar per group while the rewrite builds a hash table over every distinct pair: **543x** more peak pool at one partition, 126x at eight. That regression predates this PR and this PR does not extend it — the gate keeps every one of those functions out of the path added here. Reported upstream separately. A cost model is out of scope. ## Are there any user-facing changes? No public API break and no change to query results. `AggregateUDFImpl` gains one method with a default, which is not breaking for implementors. Plans for `SELECT ..., count(...), count(DISTINCT x) ... GROUP BY ...` change shape when `x` has no specialized `GroupsAccumulator`. `EXPLAIN` output for that shape therefore differs, and such queries should use less memory and run faster. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Which issue does this PR close?
No existing issue. This is benchmark coverage carved out of #24859 so that it can land first: a benchmark query added in the same PR that needs it cannot appear in an A/B run, because the bot compares against the merge base and the merge base does not have the query.
Rationale for this change
COUNT(DISTINCT)has a specializedGroupsAccumulatorfor the integer types and for no other type. Every other type falls back toGroupsAccumulatorAdapter, which holds one boxedAccumulator, and therefore one hash table, for each group. The cost of that fallback is per group, so the group cardinality is what decides how much it costs.Nothing in either suite measures that:
COUNT(DISTINCT)on a non-integer column at all. It groups byBrowserCountry, so it exercises the adapter at a low group cardinality, which is where the adapter is cheapest.SingleDistinctToGroupByalready rewrites them and they never reach the adapter.AVG, which that rule has never accepted.UserID, anInt64, which has a specialized accumulator.So a lone
COUNT(DISTINCT <string>)grouped by a high cardinality key, next to a non-distinctCOUNT(*), is uncovered. That is an ordinary analytics shape, it is the shape #24857 changed the memory profile of, and it is the shape #24859 proposes to rewrite. Neither of those could show its effect on any benchmark in this repository.What changes are included in this PR?
The query:
The extended ClickBench queries live in three places, and a query added to only one of them is invisible to the others. This PR adds it to all three:
benchmarks/queries/clickbench/extended/q14.sql, which feedsdfbench clickbench --queries-path. Queries are discovered from the directory, so the runner needs no change.benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q14.benchmark, which feedsbenchmark_runner clickbench_extended. Structurally identical toq13.benchmarkapart from the query. Suite files are also discovered from the directory;benchmark_runner clickbench_extended --listnow reports 15 queries.datafusion/sqllogictest/test_files/clickbench_extended.slt, which runs the extended queries against the committed ten rowclickbench_hits_10.parquetfixture.That last file had only ever mirrored q0 through q6. q7 through q13 were added to the queries directory without a matching sqllogictest entry, and q14 would have widened that gap, so this backfills q7 through q14 together. That is why the diff is larger than one query.
Having to add one query in three places is itself the problem, and the copies have already drifted: extended q6 carries a cast in its
sql_benchmarkscopy that the other two do not have. I filed #25031 for that; it is out of scope here.There is one pre-existing staleness this PR does not fix.
datafusion/core/benches/sql_planner.rsbuilds its ClickBench planning set from a hardcoded(0..=7)over the extended directory, so extended Q8 through Q13 were already outside it before this PR and Q14 joins them. That is a separate cleanup.What is the testing strategy for this PR?
The sqllogictest entries are the test.
clickbench_extended.sltruns every extended query against the committed ten row fixture and asserts its output, so q7 through q14 are now executed on every CI run rather than only by whoever runs the benchmark suite by hand.cargo test -p datafusion-sqllogictest --test sqllogictests -- clickbenchpasses, 2 of 2 files.cargo test -p datafusion-benchmarkspasses, 192 tests, andbenchmark_runner clickbench_extended --listreports the suite at 15 queries, confirming the new.benchmarkfile parses and is discovered.I do not have a
hits.parquetto hand, so I have not run the query against the full dataset. The plan shape, which is the whole value of the query, was checked separately: on an equivalent local table with the same filter, grouping,COUNT(*)and stringCOUNT(DISTINCT), the aggregate keepsaggr=[[count(Int64(1)), count(DISTINCT ...)]], which is theGroupsAccumulatorAdapterpath this query exists to measure, and it is not rewritten away.What I have not established is the effect size on the real dataset, since that depends on how many distinct
SearchPhrasevalues survive the filter. If a reviewer with the data runs it, that number is worth having on this PR.ci/scripts/doc_prettier_check.shpasses on the README change.Are there any user-facing changes?
No. This adds a benchmark query and documentation only.