jayzhan211 opened a new pull request, #24591:
URL: https://github.com/apache/datafusion/pull/24591
## Rationale for this change
### The symptom
`WHERE rk <= K` over `RANK() OVER (PARTITION BY ... ORDER BY ...)` can fail
with
`Resources exhausted` on input where the same query written with `ROW_NUMBER`
succeeds — same data, same partitioning, and the same rows ultimately kept.
### Why it happened
`PartitionedTopKExec` exists to replace a full sort with per-partition
retention,
so its entire value rests on holding only what the removed filter needed. For
`RANK` that is K rows plus every row tied at the K-th ORDER BY value: a bound
expressed in **rows**.
`PartitionedTopKRank` did not store rows. It stored **references into whole
input
batches** — the heap was handed `batch.clone()`, and each boundary tie was
kept as
`(source_batch, indices)`. A reference keeps its entire source batch alive,
so
what the operator actually retained was bounded in *batches*, not in rows.
Both
failures below follow from that one mismatch.
**Every partition in a batch charged for that whole batch.** Each
partition's heap
received its own clone. Those clones share buffers, so the bytes are really
paid
for once — but the reservation sums the heaps independently and has no way
to see
the sharing, so a batch spanning P partitions was counted P times. With 500
partitions in a single 4000-row batch and no ties at all — so `RANK` retains
exactly what `ROW_NUMBER` retains:
```
row_number size() = 473,352 bytes (14.8x the input batch)
rank size() = 16,515,008 bytes (516.1x the input batch)
```
Identical retained data, reported at ~35x the cost. That gap is what pushes
the
query past a memory limit `ROW_NUMBER` stays comfortably under, and because
it is
an over-count it fails without the memory ever actually being needed.
**Ties held memory proportional to the input, not to K.** A tie entry lives
until
the boundary improves, so one tied row pinned its whole source batch for that
entire span — and ties accumulate across batches. With eight 1000-row
batches each
contributing exactly one retained row, reported size grew by the full ~8 KB
batch
every time, to keep eight rows. Unlike the over-count, this is memory
genuinely
held, and nothing in the design bounds it by K.
Worth noting: `PartitionedTopK` (`ROW_NUMBER`) already copied out each
partition's
rows instead of holding on to the input batch. `RANK` was the outlier, which
is
precisely why the two diverge so sharply on identical input.
### Why this PR resolves it
#### The fix: store the rows, not the batches they came from
Instead of keeping a reference into the input batch, the operator now copies
out
the rows it actually keeps. The heap is given a copy of just that partition's
rows, and a tie entry holds a batch containing only the tied rows.
Both failures described above follow from that one substitution, so both go
away
with it:
- **Charging the same batch once per partition** — gone, because there is
nothing
shared left to double count. Each partition's heap now owns bytes that
belong to
it alone, so adding the heaps together gives the true total instead of an
over-estimate.
- **Ties holding memory proportional to the input** — gone, because nothing
points
into the input batch any more, so it is released as soon as the operator
moves
past it. What stays in memory is the kept rows themselves: `K + ties`, the
bound
the retention rule always promised.
#### Making the copy affordable
Copying is not free, and what it replaces nearly was — `batch.clone()` only
bumps
a reference count. Worse, most of the copying would be wasted work: once a
partition's heap is full, a group of incoming rows that are all worse than
what it
already holds changes nothing at all, and with many partitions that is the
normal
case rather than the exception.
So the operator checks before it copies. It compares the group's ORDER BY
values
against the current cutoff — those values are already encoded for
comparison, so
the check allocates nothing — and moves on if every row is worse. Without
that
check, `RANK` insert takes 53.4 ms instead of 39.5 ms on the benchmark
below. It
is part of the fix rather than a separate optimization. `PartitionedTopK`
already
worked this way, for the same reason.
#### Two related changes in the same code
Neither of the following is a memory bug, but both are the same shape of
mistake:
doing work in proportion to how much data arrives rather than to how little
is
kept. The first is also what makes the copying above affordable.
**Grouping rows by partition key allocated memory for every single row.** To
sort
a batch into per-partition buckets, `ROW_NUMBER` and `RANK` copied each row's
encoded partition key onto the heap (`pk_rows.row(i).owned()`) purely to
look it
up in a map — one allocation per row — then threw the map away at the end of
the
batch. Only the *distinct* partition keys ever need storing, and there are
far
fewer of those than there are rows. `DENSE_RANK` already did it that way,
looking
keys up by reference and reusing a single map for the operator's lifetime.
This is the same handful of lines the new copying sits inside, and fixing it
is
what pays for that copying: with it, `RANK` ends up faster than before this
PR
instead of trading memory for CPU.
**`DENSE_RANK` sorted rows into buckets before deciding whether it wanted
them.**
Every incoming row was grouped by its ORDER BY value first and only then
tested
for whether it could be kept, so rows destined to be discarded still cost a
bucket.
The test can come first. A partition keeps the K smallest distinct ORDER BY
values
it has seen; once it has K of them, the largest is the bar a new value must
beat.
That bar only ever gets stricter, because with K values held there is no
free slot
for a new one — the only way in is to evict the largest and put something
smaller
in its place. So a row worse than today's bar is worse than every future bar
too,
and can be dropped on sight instead of being bucketed and then thrown away.
Rows
exactly *equal* to the bar are kept, since that value is one of the K.
## What changes are included in this PR?
All in `datafusion/physical-plan/src/topk/mod.rs`.
1. `PartitionedTopKRank`'s heap is given a per-partition `take_record_batch`
copy
instead of `batch.clone()`; `TieEntry` holds a batch of only the tied
rows (its
`row_indices` field is gone, and `emit` no longer re-reads from the
source);
boundary-evicted rows are copied out one at a time. Adds the early skip
that
keeps non-contributing partitions from paying for the copy.
2. `ROW_NUMBER` and `RANK` key their per-partition maps on the row-encoded
partition bytes and use `entry_ref` against a reused scratch map, matching
`DENSE_RANK`. The encoded key bytes are now charged to the reservation;
before,
only the `OwnedRow` struct was counted and never its buffer.
3. `DENSE_RANK` skips rows above a saturated partition's admission boundary
before
bucketing them.
## Are these changes tested?
Yes.
**Regression tests for the two failures above** — both fail on `main`, pass
here:
- `test_partitioned_topk_rank_size_is_not_per_partition_batch` — `RANK`'s
reported
size stays within 2x `ROW_NUMBER`'s for identical retained rows.
- `test_partitioned_topk_rank_ties_do_not_pin_input_batches` — the tie list
grows
by rows, not by the batches those rows arrived in.
**Randomized differential tests**, because both changed paths turn on subtle
admission rules (`RANK` must discard its whole tie list the moment the
K-th-best
value improves; `DENSE_RANK`'s new pre-filter must not swallow boundary-equal
rows):
- `test_partitioned_topk_rank_matches_bruteforce`
- `test_partitioned_topk_dense_rank_matches_bruteforce`
Each checks the operator against a brute-force `RANK`/`DENSE_RANK <= K`
reference
over 64 seeded random shapes, varying K, partition count, value cardinality
(kept
deliberately small so ties above, at, and below the boundary are frequent),
batch
count and batch size. Both were mutation-checked: removing the tie-clearing
on a
boundary shift, and making the `DENSE_RANK` pre-filter reject boundary-equal
rows,
are each caught within the first few seeds.
Existing coverage (82 `topk` unit tests, `window_topn.slt`) passes
unchanged, as
does the full extended suite:
```
cargo test --profile ci --exclude datafusion-examples --exclude
datafusion-benchmarks \
--exclude datafusion-cli --workspace --lib --tests --bins \
--features
avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption
# exit 0 — 68 test binaries, 0 failures
```
`cargo fmt --all --check` and `cargo clippy --all-targets --all-features --
-D warnings`
are clean.
### Measurements
Memory, one 4000-row batch spanning 500 partitions, no ties:
| operator | before | after |
|---|---|---|
| `rank` | 516.1x batch | 18.1x batch |
| `row_number` | 14.8x batch | 17.4x batch |
`row_number` rising is accounting, not bytes: the partition-key buffers and
the
reusable scratch are now charged, where the old calculation counted the
`OwnedRow`
struct but never its buffer. `rank` landing beside it is the point: it no
longer
costs multiples of what `row_number` costs for the same retained rows.
Insert throughput, 200 x 8192 rows over 256 partitions, K=10, fresh random
data
per batch, best of two runs:
| operator | before | after | |
|---|---|---|---|
| `row_number` | 64.5–71.9 ms | 34.6–39.5 ms | ~1.9x |
| `rank` | 61.7–64.0 ms | 34.9–35.0 ms | ~1.8x |
| `dense_rank` | 102.6–104.7 ms | 30.5–31.2 ms | ~3.4x |
## Are there any user-facing changes?
No API changes, and query results are unchanged. `WHERE rk <= K` over
`RANK()` now
completes within memory limits where it could previously fail, and all three
ranking variants get faster.
--
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]
