PartitionedTopKRank over-accounts memory when a single batch has boundary ties across many partitions
- Dominant language
- Rust
- Stars
- 9.3k
- Forks
- 2.4k
- Avg merge
- 3d 7h
- Merged PRs (30d)
- 344
Description
### Describe the bug
PartitionedTopKRank (introduced in #22885 to extend WindowTopN to RANK) tracks per-partition boundary ties in a Vec:
```
// datafusion/physical-plan/src/topk/mod.rs
struct TieEntry {
batch: RecordBatch,
row_indices: Vec,
batch_bytes: usize, // = get_record_batch_memory_size(&batch) at push time
}
struct RankPartitionState {
heap: TopKHeap,
ties: Vec,
}
impl RankPartitionState {
fn size(&self) -> usize {
let ties_buffer = self.ties.capacity() * size_of::();
let ties_contents: usize = self.ties.iter()
.map(|t| t.row_indices.capacity() * size_of::() + t.batch_bytes)
.sum();
self.heap.size() + ties_buffer + ties_contents
}
}
```
When a single input batch produces boundary ties in N distinct partitions, insert_batch pushes N separate TieEntrys that each clone the same RecordBatch and record batch_bytes = get_record_batch_memory_size(batch). Because Arrow RecordBatch columns are Arc-shared, the actual retained memory is 1× the batch, but the operator's memory reservation reports N× the batch.
Consequence: MemoryReservation::try_resize is called with a wildly inflated size, potentially triggering spurious ResourcesExhausted or spurious spill/throttle in downstream operators sharing the pool. Worse-case fan-out is min(distinct_partition_keys_in_batch, batch.num_rows()).
The same over-count exists on a smaller factor for the heap-eviction tie path (state.ties.push(...) at ~line 1713) where the evicted batch is also retained by TopKHeap.store — bounded by K per partition, but multiplied across partitions.
### To Reproduce
```
set datafusion.optimizer.enable_window_topn = true;
-- N distinct partition keys, each with all rows tying at the same ORDER BY value
select *
from (
select pk, val, rank() over (partition by pk order by val) as rk
from ties_input -- 100k rows, 10k distinct pk, val=1 for every row
) t
where rk <= 3;
```
Instrument RankPartitionState::size() sum vs actual RSS: size() will report roughly distinct_pk × input_batch_bytes while RSS reflects one copy.
### Expected behavior
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.