SubhamSinghal opened a new pull request, #25468:
URL: https://github.com/apache/datafusion/pull/25468
## Which issue does this PR close?
Part of the `enable_window_topn` work (#13586).
## Rationale for this change
Profiling `PartitionedTopKDenseRank` (`/usr/bin/sample`, 10M-row fixture,
100K partitions) showed two costs that dominate the operator, neither of
which is inherent to DENSE_RANK semantics:
1. **Emit was 61% of busy samples**, 87% of it in one tiny
`take_record_batch` per `GroupEntry` plus the coalescer copy that
follows. A `GroupEntry` holds only the rows one source batch
contributed at one ob value, so with P partitions, K distinct ob
values and B contributing batches there are up to P x K x B of them.
`TopKHeap::emit_with_state` already solves this exact shape with a
single `interleave_record_batch`; DENSE_RANK just wasn't using it.
2. **`size()` was 8.3% self-time.** `DenseRankPartitionState::size()`
walked every group and every entry in every partition, and `size()`
is called once per `insert_batch` to resize the reservation — so the
walk costs O(entries) per batch. `TopKHeap` already avoids this with
a running `owned_bytes` total.
## What changes are included in this PR?
Two commits, both internal to `PartitionedTopKDenseRank`:
- **emit**: gather retained rows with `interleave_record_batch`, chunked
at `batch_size` so the operator emits the same batch sizes as before
and never materializes all retained rows at once. Pairs are pushed in
emit order (partitions sorted, ob ascending, entries in insertion
order), so the output needs no post-sort.
- **size**: `DenseRankPartitionState` gains a `contents_bytes` running
total, maintained at the four mutation sites (append to an existing ob
group, insert with room, evict-then-insert, and the buffer growth each
can trigger).
No signature, plan-shape, config or semantic change. `PartitionedTopK`
and `PartitionedTopKRank` are untouched.
## Are these changes tested?
Yes.
- 91 `topk` + 42 `partitioned_topk` unit tests, `window_topn.slt`,
`dev/rust_lint.sh` (fmt + clippy) all pass.
- New `test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute`
asserts the incremental total against a full recompute **after every
batch** across 64 randomized shapes, reusing the `DiffShape` harness
from the correctness differential test. It asserts the workload
actually evicted, so the eviction case can't silently go unchecked.
Mutation-checked: zeroing one term in the eviction charge fails this
test and no other.
- Emit order is **byte-identical** to before at `target_partitions = 1`
(this matters — `PartitionedTopKExec` advertises its output ordering,
and breaking it would either produce wrong ranks or force a SortExec
back into the plan).
- `count(*)` / `sum(ob)` / `sum(rn)` / `sum(pk)` digests unchanged.
## Are there any user-facing changes?
Faster `DENSE_RANK() OVER (PARTITION BY ...)` top-N when
`datafusion.optimizer.enable_window_topn` is enabled. No API change.
### Benchmarks
`benchmarks/queries/h2o/window.sql` q24–q29 (the DENSE_RANK sweep), 10M-row
`J1_1e7_1e7_NA` parquet, `dfbench h2o`, 3 reps round-robin, medians in ms.
Both binaries `--release` from the same base in one session;
`enable_window_topn` via `DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN`.
| query | shape | flag off | main, flag on | branch, flag on | branch vs
main |
|---|---|--:|--:|--:|--:|
| q24–q28 | 100 – 10K partitions | 191–244 | 45–65 | 44–66 | within noise |
| q29 | 100K partitions | 193.8 | 213.1 | 149.6 | **1.42x** |
--
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]
