Repository navigation
fix: port PartitionedTopKDenseRank onto the shared store - #27136
SubhamSinghal wants to merge 4 commits into
Conversation
…ger pins whole input batches
|
@kosiew @jayzhan211 @kumarUjjawal can you review this PR? |
|
run benchmark h2o_small_window_parquet env:
DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN: true |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing dense-rank-shared-store (f57e006) to b6b2e0e (merge-base) diff Run configurationrun benchmark h2o_small_window_parquet
env:
DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN: "true"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing dense-rank-shared-store (f57e006) to b6b2e0e (merge-base) diff Run configurationrun benchmark h2o_small_window_parquet
env:
DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN: "true"CPU Details (lscpu)Details
Resource Usageh2o_small_window_parquet — base (merge-base)
h2o_small_window_parquet — branch
File an issue against this benchmark runner |
kosiew
left a comment
kosiew
left a comment
There was a problem hiding this comment.
Thanks for this change. Moving DENSE_RANK TopK onto the shared core is an appropriate solution to the whole-batch memory retention problem. The implementation preserves distinct-key ranking and boundary ties while reducing duplicated machinery.
The architectural boundary is sound: the shared core manages storage, reference release, compaction, and emission, while dense-rank-specific admission and group accounting remain local. This addresses the underlying retention problem rather than masking reservation growth.
I found no blocking correctness or compatibility issues, so I am leaving an optional test coverage suggestion rather than requesting changes. The supplied benchmarks show workload-dependent DENSE_RANK improvements, particularly Q27-Q29, but do not independently verify the memory results. Some ROW_NUMBER and RANK timings warrant controlled reruns before attributing them to this change. The remaining suggestion would strengthen row-identity verification after compaction without blocking further review.
Approving.
| // Compaction rewrote every slot's coordinates, so the rows emit | ||
| // resolves are the check that it repointed them: a slot left at a | ||
| // pre-compaction coordinate would surface some other partition's row. | ||
| let expected: Vec<(i32, i32)> = (0..B).flat_map(|p| [(p, 0), (p, 0)]).collect(); |
There was a problem hiding this comment.
Optional: could this survivor-spread regression include a unique, non-ordering payload column and assert the exact retained payloads after compaction? The two (p, 0) survivors per partition are indistinguishable, so duplicating one while losing the other could still satisfy the current assertions. This would verify row identity without requiring SQL ordering among ties.
kumarUjjawal
left a comment
kumarUjjawal
left a comment
There was a problem hiding this comment.
Thank you @SubhamSinghal
I have left some comments please take a look
| pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> { | ||
| let Self { | ||
| output_schema, | ||
| emit_ranks, | ||
| metrics, | ||
| reservation: _, | ||
| expr: _, | ||
| row_converter: _, | ||
| scratch_rows: _, | ||
| partition_exprs: _, | ||
| partition_converter: _, | ||
| partition_scratch_rows: _, | ||
| mut states, | ||
| partition_groups: _, | ||
| ob_runs: _, | ||
| store, | ||
| k: _, | ||
| batch_size, | ||
| } = self; | ||
| let _timer = metrics.baseline.elapsed_compute().timer(); | ||
|
|
||
| let mut sorted_pks: Vec<Vec<u8>> = states.keys().cloned().collect(); | ||
| sorted_pks.sort(); | ||
|
|
||
| let mut out: Vec<Result<RecordBatch>> = Vec::new(); | ||
|
|
||
| // Gather every retained row with a single `interleave_record_batch` | ||
| // per output batch rather than one `take_record_batch` per | ||
| // `GroupEntry`. A group entry holds only the rows one source batch | ||
| // contributed at one ob value, so entries are numerous and tiny — | ||
| // with P partitions, K distinct ob values and B contributing | ||
| // batches there are up to P × K × B of them, and gathering each | ||
| // one separately builds and tears down that many `RecordBatch`es. | ||
| // `interleave` takes `(batch_pos, row)` pairs across *different* | ||
| // source batches in one call, which is exactly the shape here. | ||
| // | ||
| // The pairs are pushed in emit order — partitions in sorted key | ||
| // order, ob values ascending within a partition, entries in | ||
| // insertion order within an ob value — so the interleaved output | ||
| // is already ordered and needs no post-sort. | ||
| // | ||
| // `batch_refs` holds only the batches the current chunk references | ||
| // and is rebuilt per chunk: for Dictionary columns `interleave` | ||
| // does work for every input array it is handed, referenced or not, | ||
| // so passing the whole store to each chunk made emit | ||
| // O(chunks × store batches). A batch's slot is resolved once per | ||
| // entry, not per row. | ||
| let mut batch_refs: Vec<&RecordBatch> = Vec::new(); | ||
| let mut batch_id_pos: HashMap<u32, usize> = HashMap::new(); | ||
| let mut indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size); | ||
| let mut sorted_obs: Vec<(Vec<u8>, Vec<GroupEntry>)> = Vec::new(); | ||
| // Parallel to `indices`: each row's `DENSE_RANK()`, when emitted. | ||
| let mut ranks: Option<Vec<u64>> = | ||
| emit_ranks.then(|| Vec::with_capacity(batch_size)); | ||
|
|
||
| // Chunk at `batch_size` so the operator emits the same batch sizes | ||
| // as before and no single `interleave` output exceeds `batch_size`. | ||
| for pk in sorted_pks { | ||
| let DenseRankPartitionState { | ||
| groups, | ||
| keys: _, | ||
| contents_bytes: _, | ||
| } = states.remove(&pk).expect("key from states.keys()"); | ||
| // Sort the <= K distinct ob keys so rows emit ascending | ||
| // (byte-comparable encoding == sort order). | ||
| sorted_obs.extend(groups); | ||
| sorted_obs.sort_by(|a, b| a.0.cmp(&b.0)); | ||
| self.core.emit_with(|state| { | ||
| let mut groups: Vec<(Vec<u8>, Vec<StoreRef>)> = | ||
| state.groups.into_iter().collect(); |
There was a problem hiding this comment.
Keep dense-rank emission bookkeeping bounded by chunk size. For large tie groups, routing emission through emit_with introduces substantial eager allocation: EmitState::stream materializes every retained row's coordinates and optional rank before returning the stream. With N identical (Int32, Int32) rows and ranking enabled, accumulation needs approximately 16N bytes, but emission reserves approximately 32N bytes (payload plus coordinates and ranks). A bounded pool can therefore accept every input batch and then fail during emit(). The previous dense-rank emitter kept coordinate/rank scratch bounded by batch_size. Preserve a group iterator and generate this metadata per output chunk, and cover a large all-ties workload with a bounded-pool regression.
| // Every row of the boundary group, from this batch or | ||
| // earlier ones, now ranks below K. | ||
| let (evicted, freed) = state.evict_max(); | ||
| for at in &evicted { |
There was a problem hiding this comment.
Evicting a group now calls store.release, and so a store hash lookup in unuse, once per row. The old code called unuse once per GroupEntry, which covered one batch's run of rows.
Failure scenario: Heavy-tie DENSE_RANK: a boundary group holds 100K tied rows from 10 earlier batches. One eviction does 100K batches.get_mut lookups where 10 calls that drop a count would do. Rows from the same batch sit next to each other in arrival order, so the loop could collapse runs with the same batch_id and release each run's count in one call (for example a release_n(pending, id, n)).
| let outside_copies = state.size() | ||
| - state.core.store.size() | ||
| - state.core.partition_scratch_rows.size(); | ||
| assert!( |
There was a problem hiding this comment.
test_partitioned_topk_dense_rank_size_charges_partition_keys used to rebuild size() term by term with assert_eq!. Now it only checks a lower bound (outside_copies >= key_bytes), so it can no longer catch over-charging.
Failure scenario: If allocated_size() adds index_bytes twice, or a later change charges partitions.capacity() * size_of::<S>() twice, this test still passes. The bookkeeping test also misses it, because it uses core.allocated_size() on both sides of its equality. The reservation is then too high and can hit a false ResourcesExhausted error.
| groups: HashMap<Vec<u8>, Vec<GroupEntry>>, | ||
| /// Keyed by the row-encoded ORDER BY bytes, capped at `k` keys. Each | ||
| /// group's rows are in arrival order, which `emit` keeps. | ||
| groups: HashMap<Vec<u8>, Vec<StoreRef>>, |
There was a problem hiding this comment.
DenseRankPartitionState keeps every ob key twice: in a HashMap and in a mirrored BinaryHeap. One BTreeMap<Vec<u8>, Vec<StoreRef>> would do the same work. last_key_value/pop_last give the boundary and the eviction, iteration is already in sorted order for emit, and there is one key copy.
Failure scenario: Each new group costs two key.to_vec() allocations, evict_max hashes the key again, emit sorts every partition's groups, and the code has to keep the 'keys mirrors groups' invariant by hand. With wide ORDER BY keys and many partitions, the second key copy roughly doubles the per-group key memory charged to the pool.
There was a problem hiding this comment.
I tried using BTreeMap but seeing regression at high cardinality. Maybe I can spend some more time on this as followup.
| /// Encode the partition and ORDER BY columns once each for the whole batch, | ||
| /// run every row through its partition's group admission, then register | ||
| /// the admitted rows in the shared store with one gather. | ||
| pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { |
There was a problem hiding this comment.
PartitionedTopKDenseRank::insert_batch copies most of PartitionedTopKRank::insert_batch. push_row is the same as RankPartitionState::push_tie. groups_bytes repeats ties_bytes, because core.heaps_bytes can only grow.
Failure scenario: There are now three copies of the setup: splitting the borrows, pending, the interned flag and index_bytes, and folding in live_slots and the running totals. A fix to one copy (for example to live_slots or index_bytes accounting) has to be made in all three by hand. A core helper that runs the row pass with a per-operator closure, plus a running total in the core that can also shrink, would remove both extra operator fields.
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
Two optional cleanups. The PR description predates 5f96ec4: it still names groups_bytes and ties_bytes, doesn't say ROW_NUMBER and RANK now go through insert_batch_with, and its CPU table was measured before that commit, so could you refresh it, say that descending input stays slower than main at low partition counts, and finish the truncated title? In topk/mod.rs around line 920, the STORE_COMPACTION_RATIO doc lists "heap slots, plus ties for RANK" and should also name DENSE_RANK's groups.
|
|
||
| /// Stop tracking the largest ob value, returning its rows, which the | ||
| /// caller owes the store a release each for, and the bytes this freed. | ||
| fn evict_max(&mut self) -> (Vec<StoreRef>, usize) { |
There was a problem hiding this comment.
@SubhamSinghal thanks for this! On input where every row is a new minimum, this runs slower than main: ORDER BY ts DESC over 20M rows stored ascending (10 partitions, rn <= 10, target_partitions = 1) goes from 930 ms to 1524 ms. Each such row now evicts a whole group, which costs six allocations or frees, two hash operations, and a heap pop and push. Reusing the evicted group's buffers through peek_mut, as PartitionHeap::replace_worst does, brings most shapes to or below main and the worst one I measured from 2.85x to 1.53x. Can you apply this? It passes the topk tests, clippy and fmt.
SET datafusion.optimizer.enable_window_topn = true;
SET datafusion.execution.target_partitions = 1;
SELECT count(*), sum(rn), sum(ts) FROM (
SELECT p10, ts, s, dense_rank() OVER (PARTITION BY p10 ORDER BY ts DESC) AS rn FROM t
) WHERE rn <= 10;- /// Stop tracking the largest ob value, returning its rows, which the
- /// caller owes the store a release each for, and the bytes this freed.
- fn evict_max(&mut self) -> (Vec<StoreRef>, usize) {
- let heap_key = self.keys.pop().expect("a full partition holds k >= 1 keys");
- let (map_key, rows) = self
+ /// Replace the largest tracked ob value with the smaller, previously
+ /// unseen `key`, holding the row at `at` as its only row, reusing the old
+ /// key buffers and letting `PeekMut` sift down, so a steady-state eviction
+ /// allocates nothing. Every old row goes back through `admission`.
+ fn replace_max(&mut self, key: &[u8], at: StoreRef, admission: &mut Admission<'_>) {
+ let table_before = self.groups.allocation_size();
+ let mut heap_key = self
+ .keys
+ .peek_mut()
+ .expect("a full partition holds k >= 1 keys");
+ let (mut map_key, mut rows) = self
.groups
- .remove_entry(&heap_key)
+ .remove_entry(heap_key.as_slice())
.expect("keys mirrors groups");
- let freed = heap_key.capacity()
+ // A group's rows from one batch sit next to each other, in arrival
+ // order, so each run of them drops its uses with one store lookup
+ // rather than one per row.
+ for run in rows.chunk_by(|a, b| a.batch_id == b.batch_id) {
+ admission.release_n(run[0].batch_id, run.len());
+ }
+ admission.replacements += rows.len();
+ admission.freed += heap_key.capacity()
+ + map_key.capacity()
+ + rows.capacity() * size_of::<StoreRef>();
+
+ heap_key.clear();
+ heap_key.extend_from_slice(key);
+ map_key.clear();
+ map_key.extend_from_slice(key);
+ if rows.len() == 1 {
+ rows[0] = at;
+ } else {
+ // Do not let a group that was large pin its capacity.
+ rows = vec![at];
+ }
+ admission.grown += heap_key.capacity()
+ map_key.capacity()
+ rows.capacity() * size_of::<StoreRef>();
- (rows, freed)
+ drop(heap_key);
+ self.groups.insert(map_key, rows);
+ admission.grown += self.groups.allocation_size() - table_before;
} admission.grown += push_slot(rows, at);
+ } else if full {
+ // Every row of the boundary group, from this batch or
+ // earlier ones, now ranks below K.
+ state.replace_max(key, at, admission);
} else {
- if full {
- // Every row of the boundary group, from this batch or
- // earlier ones, now ranks below K.
- let (evicted, freed) = state.evict_max();
- // A group's rows from one batch sit next to each other, in
- // arrival order, so each run of them drops its uses with
- // one store lookup rather than one per row.
- for run in evicted.chunk_by(|a, b| a.batch_id == b.batch_id) {
- admission.release_n(run[0].batch_id, run.len());
- }
- admission.replacements += evicted.len();
- admission.freed += freed;
- }
admission.grown += state.insert_group(key, at);Operator microbench by shape: main, this PR, this PR with the fix (ms, median of 7 interleaved runs)
P = partitions. "Descending" means every row is a new minimum.
| shape | main | PR | PR + fix |
|---|---|---|---|
| 4M rows, descending, P=1, K=1 | 199 | 261 | 133 |
| 4M rows, descending, P=1, K=10 | 240 | 398 | 264 |
| 4M rows, descending, P=1, K=100 | 259 | 455 | 348 |
| 4M rows, descending, P=1, K=1000 | 448 | 594 | 472 |
| 4M rows, descending, P=10, K=100 | 408 | 553 | 417 |
| 4M rows, descending, P=100, K=10 | 336 | 405 | 274 |
| 4M rows, descending, P=100, K=100 | 805 | 946 | 620 |
| 4M rows, descending, P=1000, K=2 | 430 | 468 | 260 |
| 8M rows, descending, P=10, K=1 | 466 | 650 | 356 |
| 8M rows, descending, P=10, K=10 | 591 | 1682 | 904 |
| 4M rows, shuffled, P=100K, K=2 | 1178 | 575 | 548 |
| 4M rows, 5 distinct values, P=1K, K=2 | 271 | 155 | 153 |
End to end, the query above over 20M rows, median of 9: main 930 ms, PR 1524 ms, PR + fix 1388 ms.
SubhamSinghal
changed the title
jayzhan211
left a comment
jayzhan211
left a comment
There was a problem hiding this comment.
Optional cleanup: in the PR description, the Performance table repeats the Q24 and Q25 / Q26 rows (once indented, once not), and "Slower: input in reverse ORDER BY order" will need new numbers if the back-to-front change goes in.
| // Accumulated locally and folded in once, so the running totals are not | ||
| // touched per row. | ||
| let mut interned_bytes = 0usize; | ||
| for row in 0..batch.num_rows() { |
There was a problem hiding this comment.
@SubhamSinghal thanks for applying replace_max! ORDER BY ts DESC over rows stored by ascending ts is still 1.5–2x slower than main at 1–10 partitions (20M rows, K = 10: 652 → 1090 ms of operator compute). Walking rows in arrival order evicts once per row here, where main's hash-order buckets evict about K·ln(n/K) times per batch. Walking a falling batch back to front for DENSE_RANK brings it to 202 ms, and every other shape I measured stays within noise of this head. Could you apply this?
trait StoreSlots {
+ /// Whether the rows a partition keeps do not depend on the order its
+ /// rows arrive in, so [`PartitionedTopKCore::insert_batch_with`] may visit
+ /// a batch back to front. True for `DENSE_RANK`, which keeps every row
+ /// tied with a kept value; not for `ROW_NUMBER`, which keeps the tied rows
+ /// that come first.
+ const ORDER_FREE: bool = false;
+
@@ fn insert_batch_with
let mut interned_bytes = 0usize;
- for row in 0..batch.num_rows() {
+ let num_rows = batch.num_rows();
+ // Front to back, a batch whose ORDER BY values fall, such as
+ // `ORDER BY ts DESC` over rows stored by ascending `ts`, makes nearly
+ // every row a new best that displaces one. Back to front, the first K
+ // per partition fill the state and the boundary check drops the rest,
+ // which only saves work when partitions get more than K rows each.
+ let backwards = S::ORDER_FREE
+ && num_rows > self.k.saturating_mul(self.partitions.len())
+ && self.scratch_rows.row(num_rows - 1) < self.scratch_rows.row(0);
+ for i in 0..num_rows {
+ let row = if backwards { num_rows - 1 - i } else { i };
let pk = self.partition_scratch_rows.row(row);
@@ impl StoreSlots for DenseRankPartitionState {
+ const ORDER_FREE: bool = true;
+The num_rows > k * partitions condition skips batches where each partition gets fewer than K rows, since those evict just as often in either direction (without it, 1000 partitions at K = 100 ran 1.3x slower than this head). Walking backwards reverses the order of tied rows within a group. SQL leaves that order unspecified, but the "arrival order" docs and three tests need updating (below).
Measurements (base, this head, with the change)
elapsed_compute of PartitionedTopKExec, median of 7 interleaved runs, target_partitions = 1, 20M rows, enable_window_topn = true, SELECT count(*), sum(rn), sum(ts) FROM (SELECT pX, ts, s, dense_rank() OVER (PARTITION BY pX ORDER BY ts DESC) AS rn FROM t) WHERE rn <= K.
| input | partitions | K | main ms | this PR | with change |
|---|---|---|---|---|---|
| ts ascending, DESC | 1 | 10 | 652 | 1090 (1.67x) | 202 (0.31x) |
| ts ascending, DESC | 1 | 100 | 910 | 1780 (1.96x) | 218 (0.24x) |
| ts ascending, DESC | 10 | 10 | 678 | 992 (1.46x) | 208 (0.31x) |
| ts ascending, DESC | 100 | 10 | 1250 | 1090 (0.87x) | 324 (0.26x) |
| ts ascending, DESC | 100 | 100 | 2640 | 2500 (0.95x) | 2590 (0.98x) |
| ts ascending, DESC | 1000 | 10 | 2110 | 1270 (0.60x) | 1290 (0.61x) |
| ts ascending, DESC | 1000 | 100 | 4260 | 2900 (0.68x) | 2930 (0.69x) |
| ts ascending, ASC | 1 | 10 | 176 | 180 (1.03x) | 181 (1.03x) |
| ts ascending, ASC | 1000 | 10 | 339 | 228 (0.67x) | 238 (0.70x) |
| rising per batch, falling across, ASC | 1 | 10 | 618 | 183 (0.30x) | 191 (0.31x) |
| random ts, DESC | 10 | 100 | 197 | 188 (0.95x) | 193 (0.98x) |
| random ts, DESC | 1000 | 10 | 376 | 267 (0.71x) | 281 (0.75x) |
In a profile of this head at 1 partition, K = 100, the heap sift in PeekMut::drop takes 38% of the operator and hashing (get_mut miss, remove_entry, insert) 43%, so making each eviction cheaper can't close a 1.9x gap. ROW_NUMBER and RANK are at 0.93–1.05x of main, so the insert_batch_with refactor is fine.
Tests: two new, three adjusted (all 152 topk tests, the window/topk sqllogictests and clippy pass; removing each of the three conditions fails a test)
/// A batch whose ORDER BY values fall is visited back to front, so it
/// evicts nothing: front to back, every row would evict the one before.
#[tokio::test]
async fn test_partitioned_topk_dense_rank_falling_batch_evicts_nothing() -> Result<()>
{
let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
let vals: Vec<i32> = (0..1000).rev().collect();
state.insert_batch(&pk_val_batch(&schema, vec![1; 1000], vals)?)?;
assert_eq!(state.core.metrics.row_replacements.value(), 0);
assert_eq!(pk_val_rows(state.emit()?).await?, vec![(1, 0), (1, 1)]);
Ok(())
}
/// A falling batch with no more than K rows per partition seen is walked
/// front to back, since each partition would keep all of its rows either
/// way. The check goes by that average, so a batch that puts all its rows
/// in one partition is still walked in order.
#[tokio::test]
async fn test_partitioned_topk_dense_rank_falling_batch_over_many_partitions()
-> Result<()> {
let (schema, mut state) = build_partitioned_topk_dense_rank(1)?;
state.insert_batch(&pk_val_batch(&schema, vec![1, 2, 3, 4], vec![100; 4])?)?;
// 4 rows <= K * 4 partitions: 4 evicts the 100, then 3, 2 and 1 each
// evict the one before.
state.insert_batch(&pk_val_batch(&schema, vec![1; 4], vec![4, 3, 2, 1])?)?;
assert_eq!(state.core.metrics.row_replacements.value(), 4);
let rows = pk_val_rows(state.emit()?).await?;
assert_eq!(rows, vec![(1, 1), (2, 100), (3, 100), (4, 100)]);
Ok(())
}evicts_same_call_group and evicts_group_spanning_batches test eviction within a batch on falling input, so each needs a trailing row that keeps the batch walked front to back. store_bounded_when_survivors_spread asserts the arrival order of tied rows, so it compares them as a set:
- let vals: Vec<i32> = (0..32).rev().collect();
- state.insert_batch(&pk_val_batch(&schema, vec![1; 32], vals)?)?;
+ // The trailing 40 ends the batch above its start, so it is visited
+ // front to back and every falling value evicts the one before it.
+ let mut vals: Vec<i32> = (0..32).rev().collect();
+ vals.push(40);
+ state.insert_batch(&pk_val_batch(&schema, vec![1; 33], vals)?)?;
@@
- state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![10, 5])?)?;
+ // The trailing 10 loses to the boundary and keeps the batch from
+ // falling, so it is visited front to back.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1, 1], vec![10, 5, 10])?)?;
@@
+ for ties in rows.chunk_by_mut(|a, b| a.0 == b.0) {
+ ties.sort_unstable();
+ }
assert_eq!(rows, expected);
Which issue does this PR close?
Rationale for this change
With
enable_window_topn = true,PartitionedTopKDenseRankregisters every whole input batch and never compacts, so one retained row pins its whole batch and memory tracks the input, not the rows kept. ROW_NUMBER and RANK avoid this throughPartitionedTopKCore(#25730, #25968); DENSE_RANK was never ported. This blocks #24608 independently of #26381.Smallest passing
--memory-limit,target_partitions = 1, #26381 data (4.1M rows, 5,000 retained):dense_rankbeforerow_numberUtf8View's remainder is #26381's view-buffer pinning, which
row_numbershows too.What changes are included in this PR?
PartitionedTopKDenseRanknow runs onPartitionedTopKCore<DenseRankPartitionState>:PartitionedTopKCore::insert_batch_with, which ROW_NUMBER and RANK now use too; each operator supplies only its admission rule as a closure.StoreSlots::ORDER_FREE): the rows it keeps do not depend on visit order, but tied rows in a group come out reversed for such a batch.Vec<StoreRef>; the store gathers only admitted rows and is compacted byRecordBatchStore::compact. Evicting a group costs one store lookup per run of same-batch rows.replace_max(from @jayzhan211) reuses the evicted group's buffers throughpeek_mut, likePartitionHeap::replace_worst.size()is O(1): one runningstate_bytesreplaces the per-operator totals and the O(partitions) sum. The groups table is charged viaHashMap::allocation_size, sincecapacity()drops on removal without freeing.batch_sizechunk at a time; DENSE_RANK hands over its groups lazily, so emit needs only one chunk of scratch beyond what accumulation held.row_replacements(rows evicted) changes value because rows are visited in batch order, not hash order: ~33K → ~3.5K on h2o at 20% retention.Performance
Operator
elapsed_compute,main=d4c2a329a, median of 11 paired reps, plan asserted. h2oJ1_1e7_1e7_NA.parquet, DENSE_RANK top-2, 14 partitions (Q26/Q28 order by Float64v2 % 10, so no ties; Q15/Q16 are the tie-heavy ones):Peak RSS is lower on every query (Q29 407 → 300 MB, Q14 353 → 241 MB, Q16 251 → 140 MB).
Slower: input sorted by the partition key. The back-to-front check compares only a batch's first and last rows, which can belong to different partitions. On
window_sorted, a batch that crosses a partition boundary ends on the next partition's best row, so it is walked worst-first: Q15 is 0.90× ofmain, and Q11–Q14 are 0.55×–0.85× of the same PR without the walk (rows evicted go from ≤ 36 to 0.3M–3.3M). A see-saw batch (rising runs, each starting lower) is also walked against its runs and lands atmain's speed (1.00×–1.15×).What is the testing strategy for this PR?
bookkeeping_tracks_recompute: 256 randomized shapes; after every batch, every slot,state_bytes,index_bytes, the reservation andsize()against a recompute, then output against brute force.size_charges_partition_keys:size()equals a term-by-term rebuild exactly, so over- and under-charging fail.store_bounded_when_survivors_spread: asserts exactly which rows (by a uniqueid) survive compaction; tied rows are compared as a set, since a batch walked back to front stores them reversed.runs_under_bounded_memory_pool: 256 partitions × 256 batches under 80× one batch; peaks at ~39×, whilemainfails at batch 40 withResourcesExhausted.emit_fits_pool_of_all_ties: 131K tied rows emit under 1.25× of what accumulation reached.falling_batch_evicts_nothingandfalling_batch_over_many_partitions(from @jayzhan211): a falling batch evicts nothing, and one with no more than K rows per partition seen is still walked front to back.row_replacementsfor DENSE_RANK and ROW_NUMBER, and the three insert guards that turn over-release into an internal error.repoint, the metric, lazy emit, an insert guard,replace_max's row swap, or any of the back-to-front walk's three conditions each fails a test.datafusion-physical-planlib tests and thewindow*,window_topn,window_topk_pushdown,window_limitsandtopksqllogictests pass.Are there any user-facing changes?
No API change; the operators are
pub(crate).DENSE_RANKwithenable_window_topn = trueuses memory proportional to the rows it keeps rather than its input, and less at emit.