fix: reduce RANK window top-K memory from O(input) to O(K + ties) per partition

[jayzhan211]

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.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 88.42593% with 25 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.38%. Comparing base (5134a1a) to head (2d15aeb).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/topk/mod.rs	88.42%	5 Missing and 20 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #24591      +/-   ##
==========================================
- Coverage   81.38%   81.38%   -0.01%     
==========================================
  Files        1116     1116              
  Lines      397960   398118     +158     
  Branches   397960   398118     +158     
==========================================
+ Hits       323880   324007     +127     
- Misses      55120    55133      +13     
- Partials    18960    18978      +18     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.