zhuqi-lucas opened a new pull request, #23687:
URL: https://github.com/apache/datafusion/pull/23687
## Which issue does this PR close?
Draft — cites @Dandandan's proposal on #21580 (2026-04-17) and closes an
open ClickBench-Q12/Q33 optimization gap. If interest is confirmed, 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 K`
is a hot ClickBench pattern (Q12 SearchPhrase, Q33 URL) and dominates
end-to-end time on high-cardinality data (~40M distinct URLs in
`hits.parquet`). Neither DataFusion nor DuckDB optimize it today; both
fall back to full sort of all groups.
@Dandandan proposed the direction on #21580 (2026-04-17):
> Currently many queries go like Scan → Filter → Aggregate(Partial) →
> Aggregate(Final) → TopK … Needs to first go through the full
> aggregation before it reaches TopK. **We can update bounds from
> aggregate, once we detect at least K distinct rows in a Partial
> aggregate. In practice it would mean after the first batch / batches
> arrive we already have e.g. 10 rows, so can already prune / filter
> with the much tighter bound.**
The Microsoft *Zippy* paper (VLDB'24 — "Cache-Efficient Top-k Aggregation
over High Cardinality Large Datasets") demonstrates the algorithmic
frame: identify candidate top-K groups early, cheaply prune the long
tail using a running threshold, and reports **14× COUNT speedups**
(Table 2) in the offline multi-pass setting. This PR is a
single-pass streaming adaptation for DataFusion.
## What changes are included in this PR?
### Optimizer (`physical-optimizer/src/topk_aggregation.rs`)
`TopKAggregation::transform_agg` was previously MIN/MAX-only. Extended to
recognize `count(*)`/`count(col)` (non-distinct, unfiltered, single
group-by column) at `Final`/`FinalPartitioned`/`Single`/
`SinglePartitioned` mode, marking the aggregate with
`LimitOptions::new_count(K, desc)`.
### `AggregateExec` (`physical-plan/src/aggregates/mod.rs`)
- `TopKKind::{MinMaxOrDistinct, Count}` on `LimitOptions` distinguishes
the safe capacity-K `PriorityMap` path from a new streaming path that
must keep an unbounded hash table.
- New `TopKCountSharedState`: a per-partition slot vector into which
each 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::extend` implementation
hit.
- `execute_typed` routes `TopKKind::Count` to the new stream.
### `GroupedTopKCountAggregateStream` (~500 LOC, new file)
Streaming Zippy-lite:
1. **Hash + count**: `GroupValues::intern` (any group-by column type) →
`Vec<i64>` counts indexed by group_index. `Vec<bool>` `dead` bitmap
parallel.
2. **Warmup**: first ≥ `WARMUP_GROUPS` (= max(K, 4)) live groups build
the initial local top-K; before that, gate/sweep are disarmed.
3. **New-group gate** (Zippy Algorithm 3 line 26 analogue): before
accepting a previously-unseen key, check `contrib +
remaining_upper_bound < threshold`. If so, mark the slot dead
immediately.
4. **Periodic sweep every 64K rows** (Zippy Algorithm 4 lines 12-18
analogue): mark any live group whose `count + remaining_upper < threshold`
dead. Sweep also fires once at end-of-input to catch tail groups whose
final counts fall below the last-seen threshold.
5. **`remaining_upper`** comes from
`StatisticsContext::compute(agg.input)`. `Precision::Absent` disarms
pruning (stream degrades to plain hash agg + emit-K).
6. **Global threshold** via `TopKCountSharedState`: each partition's
`current_threshold()` publishes its local top-K into its slot and
reads 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.slt` grew a `count(*)/count(col)/
count(distinct)` section: baseline vs optimized EXPLAIN comparison,
DESC + ASC + tie-break + longtail scenarios.
## Observed behaviour
Observability via `EXPLAIN ANALYZE VERBOSE` (per-partition):
- `count_topk_groups_seen`
- `count_topk_groups_gated`
- `count_topk_groups_swept_dead`
- `count_topk_sweeps_performed`
Correctness 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, I'll rerun on GKE with `adriangbot` or on my own
machine after freeing disk. Expected win concentrated on high-cardinality
skewed data (Q12/Q33/Q14). Per-`Partial` gains are out of scope (see
below).
## Known limitations / follow-ups
1. **Multi-aggregate patterns** (Q9, Q21, Q22, Q27, Q28, Q30-Q32) —
`get_count_topk_field` requires exactly one aggregate, so these fall
through unchanged. Extending this is straightforward but non-trivial:
the emit path needs to reproduce the other aggregate's state, and the
partition partial-count column index generalizes.
2. **Multi-column group-by** (Q14, Q30, Q31, Q32, Q34, Q35) — same
architecture works, but the optimizer currently rejects multi-column
`GROUP BY` at the entry point.
3. **Threshold pushdown to Partial** — currently the shared threshold
is only used at Final. Pushing it to Partial would let Partial skip
emitting long-tail rows entirely (Dandandan's original framing). Would
likely reuse `DynamicFilterPhysicalExpr` infra like MIN/MAX aggregate
already does.
4. **Early-stop partitions** — if a Final partition's `count + remaining
< shared_threshold` for every group, it can return `Poll::Ready(None)`
instead of draining. Not implemented in this PR.
5. **`Raw` input mode / Single aggregate** — `AggregateInputMode::Raw`
is currently untested; `Partial` (i.e. FinalPartitioned consuming
partial state) is the exercised path.
## References
- Dandandan's comment on #21580 (April 2026):
https://github.com/apache/datafusion/pull/21580#issuecomment-2812810017
- Zippy paper (VLDB'24): https://www.vldb.org/pvldb/vol17/p644-siddiqui.pdf
- DuckDB gap (documented, no equivalent optimization):
https://duckdb.org/2024/10/25/topn
## Are there any user-facing changes?
Behavioural: `SELECT g, count(*) FROM t GROUP BY g ORDER BY count(*)
DESC/ASC LIMIT K` now takes the new stream when
`datafusion.optimizer.enable_topk_aggregation = true` (the default).
Same results as before, plus four new per-partition metrics visible in
`EXPLAIN ANALYZE` output.
Config: no new config options in this PR.
API surface: `TopKKind`, `TopKCountSharedState`, and
`is_topk_count_pattern` are `pub` in `datafusion-physical-plan::aggregates`
for the optimizer's use.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
