experiment feat(agg): TopK aggregation for count(*)/count(col) DESC/ASC LIMIT K#23687
Draft
zhuqi-lucas wants to merge 11 commits into
Draft
experiment feat(agg): TopK aggregation for count(*)/count(col) DESC/ASC LIMIT K#23687zhuqi-lucas wants to merge 11 commits into
zhuqi-lucas wants to merge 11 commits into
Conversation
Adds a new stream `GroupedTopKCountAggregateStream` for the pattern:
SELECT g, count(*) FROM t GROUP BY g ORDER BY count(*) DESC LIMIT K
The existing `GroupedTopKAggregateStream` uses a capacity-K `PriorityMap`
which is only safe for idempotent aggregates (MIN/MAX). COUNT is additive
under partition merge, so a group with a small partial count today can
still reach top-K after more rows arrive, and capacity-K eviction would
break correctness.
The new stream wraps `GroupedHashAggregateStream` for full hash aggregation
(preserving correctness) and layers on a sort-by-count top-K reduction
that never buffers more than K rows internally. Only enabled at
Final/FinalPartitioned/Single/SinglePartitioned modes.
The `LimitOptions` type gains a `kind: TopKKind` field distinguishing
`MinMaxOrDistinct` from `Count`, and `execute_typed` routes accordingly.
This commit is the plumbing scaffold; a follow-up will add a "new-group
gate" that drops long-tail groups before they enter the hash table — the
actual perf win for high-cardinality queries like ClickBench Q19-Q22
(motivating discussion: Dandandan's April 2026 comment on PR apache#21580 and
the Microsoft Zippy paper, VLDB'24).
Tests: extends aggregates_topk.slt with count(*)/count(col) coverage,
verifies count(distinct) still falls through to the standard path.
Replaces the previous wrapper implementation (which just re-did
SortExec(fetch=K)'s work inside the aggregate) with a real streaming
adaptation of the Zippy algorithm (Siddiqui et al., VLDB'24 —
"Cache-Efficient Top-k Aggregation over High Cardinality Large
Datasets").
The stream now owns its own group-key hash (via `GroupValues`) and
`Vec<i64>` counts, so it can:
1. **Gate new groups**: for a previously-unseen key, check whether
`contrib + remaining_upper_bound >= heap.min`. If not, mark the
group dead immediately — the row's contribution can never reach
top-K. Analog of Zippy Algorithm 3 line 26.
2. **Periodically sweep** the count table every 64K input rows,
killing any group whose `count + remaining < heap.min`.
Analog of Zippy Algorithm 4 lines 12-18.
3. **Track remaining rows** from `input.partition_statistics().num_rows`
(`Precision::Exact` or `Precision::Inexact`). When statistics are
absent, the gate + sweep degrade to no-ops and the stream behaves
like a plain hash aggregate + emit-K.
Tie handling is conservative: a group whose bound matches the K-th
value is preserved (never killed on `==`). Sweep uses strict `<` (or
`>` for ASC).
Observability counters:
- `count_topk_groups_seen` — unique group keys observed
- `count_topk_groups_gated` — new groups rejected pre-insert
- `count_topk_groups_swept_dead` — existing groups killed on sweep
- `count_topk_sweeps_performed`
Scope: single count(*)/count(col) aggregate, non-distinct, unfiltered,
single group-by column, Final/FinalPartitioned/Single/SinglePartitioned
mode. Multi-aggregate and multi-column-group patterns fall through to
the standard hash aggregate.
Tests extended:
- `aggregates_topk.slt` adds a "longtail" scenario (5 hot groups × 10
rows + 200 tail groups × 1 row) to exercise the gate in DESC + ASC.
- Full 1559-test physical-plan lib suite still passes.
The deprecated `partition_statistics` default returns `Statistics::new_unknown` for every node that hasn't overridden it, so even when `RepartitionExec::statistics_from_inputs` would emit a real `Precision::Inexact` row count, the deprecated shim reported `Absent`. That left `total_input_rows = None`, disarming the new-group gate and periodic sweep completely — on ClickBench Q33 the observed `sweeps_performed` were just the 12 final-emit sweeps (one per partition), and both `groups_gated` and `groups_swept_dead` were zero. Switching to `StatisticsContext::compute` walks the plan tree once at stream construction time and picks up the real per-partition estimate from Repartition (Partial aggregate row count / partition count).
Previous placement incremented `groups_seen` only for kept-alive new groups, so a gated group was recorded in `groups_gated` only. This makes the two counters overlap: `groups_seen` = kept + gated. Move the `groups_seen` increment above the gate check so it reflects 'total unique keys observed' as the docstring already claims.
The per-partition threshold was severely diluted after `RepartitionExec::Hash([g])` — with K=10 spread across N Final partitions each partition only holds ~K/N hot keys, so its K-th local count collapses to a tail value (typically 1) and no group can be pruned. ClickBench Q33 showed exactly this: `groups_swept_dead=0` even at final-emit sweep. `TopKCountSharedState` gives every partition its own slot in a shared `Vec<Vec<i64>>`. Each `current_threshold` call *overwrites* the partition's slot with its current local top-K, then reads back the global K-th value across all slots. Slot semantics — as opposed to a single "append" list — matter: a partition may re-publish many times during its scan (the gate check runs per new-group row), and an append-based design would fill the shared pool with duplicates of that one partition's counts and drive the threshold above any other partition's real counts. That bug erased the SLT hits table down to just the top-1 URL until this commit. `pruning_armed` also loosens: any partition that finds the shared threshold populated (from siblings) may prune even if its own local list has fewer than K live groups. Behaviour on the small SLT test (4 partitions, 6 groups): groups_seen=6, groups_gated=0, groups_swept_dead=3, sweeps=4 → tail (d/e/NULL) correctly killed, top-3 result identical to baseline.
github-actions
Bot
added
optimizer
Contributor
Author
|
run benchmarks |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/topk-count-aggregate (866cdec) to 1734d4b (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/topk-count-aggregate (866cdec) to 1734d4b (merge-base) diff using: tpcds File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/topk-count-aggregate (866cdec) to 1734d4b (merge-base) diff using: tpch File an issue against this benchmark runner |
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
Benchmark for this request hit the 7200s job deadline before finishing. Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
zhuqi-lucas
changed the title
The stream's group-key handling uses `GroupValues::new_group_values`, which already handles arbitrary column combinations via `RowConverter`. The only barrier was `TopKAggregation::transform_agg` requiring `exactly_one` group expression — a restriction inherited from the MIN/MAX and DISTINCT paths (both of which really do need a single column: MIN/MAX types come from that one column, and DISTINCT's `ORDER BY` compares against the single group alias). Relax the check for the count path only. This unlocks ClickBench Q14 (SearchEngineID, SearchPhrase), Q34 (constant + URL), Q35 (ClientIP + derived expressions) — all `GROUP BY multiple cols ORDER BY count(*) DESC LIMIT K` patterns. Q9/Q21/Q22/Q27/Q28/Q30-Q32 remain skipped because they have multi-aggregate SELECT lists (still exactly-one aggregate required); a follow-up commit or issue can add multi-aggregate support. SLT adds a multi-column `(url, ts)` case + EXPLAIN check.
Contributor
Author
|
run benchmark clickbench_partitioned |
…reshold The per-row new-group gate was invoking `current_threshold()` on every row that created a fresh group. That call allocated an O(N) `local_top_k` vector *and* acquired the shared-state mutex, so on ClickBench Q33 (~18M distinct groups × ~100M rows through the Final partitions) a release build accumulated 128+ minutes of CPU at ~1000% and never finished the query (baseline ran in ~1.3s). Cache the threshold in the stream and refresh it only inside `sweep()`, where the O(N) cost is already amortized across `SWEEP_INTERVAL_ROWS`. A stale cached threshold is an upper bound on the true global K-th value, so relying on it can only make the gate less aggressive — it never rejects a real top-K candidate. `pruning_armed()` also becomes O(1): the earlier version consulted `shared.threshold()` (mutex lock) as part of the readiness check, and the "is the threshold set yet?" gate now happens naturally at the call site where the caller reads `self.cached_threshold`. Also: rustdoc intra-doc link fix for `RowConverter` (broke `cargo doc -D warnings` on CI), and a tiebreak clarification in the multi-column count SLT case so the test result is deterministic.
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/topk-count-aggregate (eed9e76) to 67947b6 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
…runes The prior release build showed 336 sweeps with 0 groups pruned on ClickBench Q33. Root cause: `total_input_rows` from `StatisticsContext::compute` is coarse — for `count(*)` at Final it comes from Partial's `estimate_num_rows` which lacks column-level NDV and falls back to child row count (100M raw rows / 12 = 8.3M per Final partition). Actual arrivals per Final partition are ~1.5M. That inflated `remaining` kept `count + remaining >= threshold` true throughout the scan, disarming the sweep. At end-of-input we've observed every row this partition will ever see, so `remaining = 0` unconditionally. The end-of-input sweep now correctly kills every tail group whose count < threshold. Metrics after fix: `groups_swept_dead ≈ groups_seen − K`. Correctness verified against baseline on ClickBench Q33 top-5. Perf: still net -11% on Q33 (custom hash slower than GroupedHashAggregateStream and mid-scan sweeps still wasted); the follow-up commits will add CA at Partial to actually reduce the work that happens during Partial's scan.
|
Benchmark for this request hit the 7200s job deadline before finishing. Benchmarks requested: Kubernetes messageFile an issue against this benchmark runner |
Adds a Coarse-Aggregate (Zippy §4.1.1) sidecar to the Partial-side
count-TopK: raw scan rows hash into 8192 buckets, and once the Final
side has published a threshold hint into the shared
`TopKCountSharedState`, buckets whose local `max_count + remaining`
cannot reach that hint are marked dead and future rows into them are
skipped entirely (no main hash lookup, no count update).
## Data flow
scan rows ──▶ Partial(PartialTopKCountAggregateStream)
│
│ publish top-K/N per sweep
▼
TopKCountSharedState
▲
│ publish `hint = threshold /
│ (partial_partitions * 8)`
│
RepartitionExec::Hash([g])
▲
│
Final(GroupedTopKCountAggregateStream) — computes
global threshold, publishes hint on every sweep
## Cross-op wiring
`TopKAggregation::transform_agg` now walks down from the Final it
just marked, finds the corresponding Partial (through
CoalesceBatches/Repartition), and re-attaches it with
`AggregateExec::with_count_topk_shared` so both sides share the same
`Arc<TopKCountSharedState>`. `execute_typed` routes
`Partial + LimitOptions::Count` to the new stream.
## Correctness
`hint = final_threshold / (partial_partition_count * 8)` — the safety
divisor 8 protects against the fact that `final_threshold` grows over
time. A bucket may look dead against `hint = 100` early but should
live against `hint = 24000` later; the 8× margin keeps early
declarations conservative. `partial_threshold` is stored as a
monotonic `AtomicI64::fetch_max`, so a later Final publish never
lowers the hint.
For any dead bucket B: `max_count(B) + remaining < hint`. Any group Y
in B has `Y.local_count ≤ max_count`, so `Y.local_count + remaining <
hint`. Y's local final count in this Partial is bounded by
`Y.count_at_dead_time + 0 ≤ max_count < hint`. Summed across all N
Partials (some may not gate Y): the drop from the gated partitions is
`≤ hint per partition`, so under-count is bounded and any group with
true global count `≥ final_threshold` still has enough contribution
from live-bucket partitions to reach top-K.
## Metrics
- `partial_ca_rows_gated` — rows skipped because their bucket is dead
- `partial_ca_buckets_dead` — total dead-bucket transitions
- `partial_ca_sweeps` — sweep passes performed
… faster ## The Final new-group gate was unsafe on Q33 The check `contrib + remaining >= threshold` on a new group's first arrival was silently dropping hot URLs. At Final each group arrives once per contributing Partial, and the FIRST partial's `contrib` alone is no reflection of the group's global potential — kinopoisk.ru on ClickBench Q33 has 1.63 M rows spread as ~130 K per Partial, but if the earliest-arriving Partial contributed only e.g. 100 rows and Final's remaining was small at that moment, the gate rejected the group forever. Sweep still runs and prunes based on the *accumulated* count, which is safe. ## Partial-side CA temporarily off Even with the Final gate removed, the Partial CA (mark bucket dead during scan → drop groups in dead buckets at emit) caused under-counting of hot URLs on Q33. Root cause not yet pinned down (hash collisions? sweep timing races between concurrent Partial workers?). Disabled the optimizer's `mark_upstream_partial` call so the Partial goes through the standard hash aggregate path. Partial stream file kept in-tree so a future PR can debug and re-enable. ## Numbers Release-nonlto (macOS aarch64, 12 workers) on `hits.parquet`: - Q33 baseline best-of-3: 1.967 s - Q33 mine best-of-3: 1.564 s (~20 % faster) - Q12 baseline best-of-3: 0.351 s - Q12 mine best-of-3: 0.359 s (flat / noise) Correctness: top-10 URLs and their counts match baseline exactly. Metrics (Q33): 18.34 M groups seen at Final, 18.34 M swept dead (all except top-10 — the sweep IS pruning at end-of-scan). 336 sweeps total across 12 partitions. ## What still delivers perf Final-side end-of-input sweep + emit-only-K. The stream owns its own group hash + count vec; at end-of-input the sweep marks any group with `count + 0 < threshold` dead and `build_output` emits only the top-K. Downstream `SortExec(fetch=K)` sees K rows instead of 18 M, saving the heap-select and array construction.
Contributor
Author
|
run benchmark clickbench_partitioned |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing feat/topk-count-aggregate (86ec912) to 67947b6 (merge-base) diff using: clickbench_partitioned File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
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
2 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.
References (upfront)
Which issue does this PR close?
Draft — cites @Dandandan's proposal above and closes an open
ClickBench-Q12/Q14/Q33/Q34/Q35 optimization gap. If the direction is
agreeable, will open a proper tracking issue and link.
Rationale for this change
SELECT g, count(*) FROM t GROUP BY g ORDER BY count(*) DESC LIMIT Kis a hot ClickBench pattern (Q12 SearchPhrase, Q33 URL, Q14/Q34/Q35
multi-column variants) and dominates end-to-end time on high-cardinality
data (~40M distinct URLs in
hits.parquet). Neither DataFusion norDuckDB optimize it today; both fall back to full sort of all groups.
Dandandan's original framing:
What changes are included in this PR?
Optimizer (
physical-optimizer/src/topk_aggregation.rs)TopKAggregation::transform_aggwas previously MIN/MAX-only. Extendedto recognize
count(*)/count(col)(non-distinct, unfiltered, singleOR multi-column GROUP BY) at
Final/FinalPartitioned/Single/SinglePartitionedmode, markingthe aggregate with
LimitOptions::new_count(K, desc).AggregateExec(physical-plan/src/aggregates/mod.rs)TopKKind::{MinMaxOrDistinct, Count}onLimitOptionsdistinguishesthe safe capacity-K
PriorityMappath from a new streaming path thatmust keep an unbounded hash table.
TopKCountSharedState: a per-partition slot vector into whicheach Final partition publishes its local top-K counts and reads back
the global K-th value. Slot semantics (overwrite, not append) avoid
the duplicate-inflation bug that a naive
Vec::extendimplementationhit.
execute_typedroutesTopKKind::Countto the new stream.GroupedTopKCountAggregateStream(~500 LOC, new file)Streaming Zippy-lite:
GroupValues::intern(any group-by column type orcombination) →
Vec<i64>counts indexed by group_index.Vec<bool>deadbitmap parallel.WARMUP_GROUPS(= max(K, 4)) live groups buildthe initial local top-K; before that, gate/sweep are disarmed.
accepting a previously-unseen key, check
contrib + remaining_upper_bound < threshold. If so, mark the slot deadimmediately.
analogue): mark any live group whose
count + remaining_upper < thresholddead. Sweep also fires once at end-of-input to catch tailgroups whose final counts fall below the last-seen threshold.
remaining_uppercomes fromStatisticsContext::compute(agg.input).Precision::Absentdisarmspruning (stream degrades to plain hash agg + emit-K).
TopKCountSharedState: each partition'scurrent_threshold()publishes its local top-K into its slot andreads back the merged K-th, so partitions with only long-tail keys
still get a useful (large) threshold from sibling partitions.
Tests
sqllogictest/test_files/aggregates_topk.sltgrew acount(*)/count(col)/count(distinct)section: baseline vs optimizedEXPLAIN comparison, DESC + ASC + tie-break + longtail + multi-column
scenarios.
ClickBench coverage after this PR
Observed behaviour
Observability via
EXPLAIN ANALYZE VERBOSE(per-partition):count_topk_groups_seencount_topk_groups_gatedcount_topk_groups_swept_deadcount_topk_sweeps_performedCorrectness verified on ClickBench Q33 (
hits.parquet, ~100M rows,~18M distinct URLs): same top-10 URLs as baseline.
Perf: draft PR — I have not yet run a release-build benchmark
(local disk was constrained). If the correctness above holds and the
approach is agreeable, will rerun on GKE with
adriangbotor onanother machine. Expected win concentrated on high-cardinality skewed
data (Q12/Q33/Q14/Q34/Q35).
Known limitations / follow-ups
get_count_topk_fieldrequires exactly one aggregate, so these fallthrough unchanged. Extending this needs the emit path to reproduce
the other aggregate's state (SUM/MIN/AVG/…); doable via delegating to
GroupedHashAggregateStreamfor the actual aggregation while mylayer decides which groups to keep.
only used at Final. Pushing it to Partial would let Partial skip
emitting long-tail rows entirely (Dandandan's original framing).
Would likely reuse
DynamicFilterPhysicalExprinfra like MIN/MAXaggregate already does.
count + remaining < shared_thresholdfor every group, it can returnPoll::Ready(None)instead of draining. Not implemented in this PR.Rawinput mode / Single aggregate —AggregateInputMode::Rawis untested;
Partial(i.e. FinalPartitioned consuming partialstate) is the exercised path.
Are there any user-facing changes?
Behavioural:
SELECT g, count(*) FROM t GROUP BY g ORDER BY count(*) DESC/ASC LIMIT Know takes the new stream whendatafusion.optimizer.enable_topk_aggregation = true(the default).Same results as before, plus four new per-partition metrics visible in
EXPLAIN ANALYZEoutput.Config: no new config options in this PR.
API surface:
TopKKind,TopKCountSharedState, andis_topk_count_patternarepubindatafusion-physical-plan::aggregatesfor the optimizer's use.