apache / apache/datafusion

PartitionedTopKRank over-accounts memory when a single batch has boundary ties across many partitions

Open
#23,326 1 comment 0 reactions 1 assignee Claimed by @mkleen View on GitHub
bug
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.