diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index 67f49d58bd0a8..04072fffaf8eb 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -483,20 +483,20 @@ impl ExecutionPlan for PartitionedTopKExec { /// this operator instance. /// /// 2. **Emission** — `emit` drains all per-partition state in sorted -/// partition-key order. `ROW_NUMBER` interleaves the retained rows out -/// of its shared store in `batch_size` chunks, so its output needs no -/// coalescing; `RANK` and `DENSE_RANK` still coalesce. For -/// `RANK`, boundary-tied rows are materialized and emitted after -/// each partition's heap rows. For `DENSE_RANK`, rows are emitted -/// from a K-bounded map of distinct ob keys, sorted ascending. +/// partition-key order. `ROW_NUMBER` and `RANK` interleave the retained +/// rows out of their shared store in `batch_size` chunks, so their output +/// needs no coalescing; `DENSE_RANK` still coalesces. For `RANK`, +/// boundary-tied rows are emitted after each partition's heap rows. For +/// `DENSE_RANK`, rows are emitted from a K-bounded map of distinct ob +/// keys, sorted ascending. /// /// # Cost /// /// - Time: O(N log K) where N = total rows, K = fetch /// - Memory: O(K × P × row_size) where P = number of distinct partitions -/// plus, for RANK, the boundary ties' rows. `ROW_NUMBER` holds its rows -/// by reference into gathered batches, so its constant is the store's -/// compaction ratio, and one in-flight gather is pinned on top. +/// plus, for RANK, the boundary ties' rows. `ROW_NUMBER` and `RANK` hold +/// their rows by reference into gathered batches, so their constant is the +/// store's compaction ratio, and one in-flight gather is pinned on top. #[expect(clippy::too_many_arguments)] async fn do_partitioned_topk( partition_id: usize, diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index d4405c1f12166..d9d67243d1d8c 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -20,12 +20,12 @@ use arrow::{ array::{Array, AsArray}, compute::{ - BatchCoalescer, FilterBuilder, interleave_record_batch, prep_null_mask_filter, - take_record_batch, + FilterBuilder, interleave_record_batch, prep_null_mask_filter, take_record_batch, }, row::{RowConverter, Rows, SortField}, }; use datafusion_expr::{ColumnarValue, Operator}; +use std::borrow::Borrow; use std::mem::size_of; use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::{cmp::Ordering, collections::BinaryHeap, sync::Arc}; @@ -307,24 +307,6 @@ impl TopKDynamicFilters { // Guesstimate for memory allocation: estimated number of bytes used per row in the RowConverter const ESTIMATED_BYTES_PER_ROW: usize = 20; -/// Owned data of a row that was just evicted from a [`TopKHeap`]. -/// -/// Returned by [`TopKHeap::add`] so that callers (e.g. rank-aware -/// wrappers that retain boundary ties) can decide whether to retain -/// the evicted row externally. The underlying batch is captured -/// before the heap's internal `RecordBatchStore` decrements the -/// batch's use count, so the data remains accessible even if the -/// heap drops its internal reference to the batch. -#[derive(Debug, Clone)] -pub(crate) struct EvictedRow { - /// The record batch the evicted row came from. - pub batch: RecordBatch, - /// Row index within `batch`. - pub index: usize, - /// Encoded ORDER BY tuple for the evicted row, in [`arrow::row`] format. - pub row_bytes: Vec, -} - pub(crate) fn build_sort_fields( ordering: &[PhysicalSortExpr], schema: &SchemaRef, @@ -891,9 +873,9 @@ impl TopKMetrics { /// what it holds tracks the *input* rather than the rows retained. Both top-K /// paths bound that the same way — compact once the store holds this multiple /// of what is still referenced — and only the denominator differs: one heap's -/// length in [`TopKHeap::maybe_compact`], every partition's slots in -/// [`PartitionedTopK::compact_store`]. The multiplier avoids compacting when -/// the savings would be marginal. +/// length in [`TopKHeap::maybe_compact`], every partition's retained rows +/// (heap slots, plus ties for `RANK`) in [`RecordBatchStore::compact`]. The +/// multiplier avoids compacting when the savings would be marginal. const STORE_COMPACTION_RATIO: usize = 2; /// This structure keeps at most the *smallest* k items, using the @@ -952,16 +934,12 @@ impl TopKHeap { /// Adds `row` to this heap. If inserting this new item would /// increase the size past `k`, removes the previously smallest /// item. - /// - /// Returns `Some(EvictedRow)` if an existing row was evicted to - /// make room for `row`, or `None` if the row was inserted into a - /// non-full heap. fn add( &mut self, batch_entry: &mut RecordBatchEntry, row: impl AsRef<[u8]>, index: usize, - ) -> Option { + ) { let batch_id = batch_entry.id; batch_entry.uses += 1; @@ -972,26 +950,6 @@ impl TopKHeap { if self.inner.len() == self.k { let mut prev_min = self.inner.peek_mut().unwrap(); - // Capture evicted row data before `unuse` (which may GC the - // batch from the store) and `replace_with` (which overwrites - // `prev_min` in place). The batch comes from `self.store` for - // cross-batch evictions, or directly from `batch_entry` when - // a row evicts another row from the same in-flight batch - // (entry not yet registered in the store). - let evicted_batch = if prev_min.batch_id == batch_entry.id { - batch_entry.batch.clone() - } else { - self.store - .get(prev_min.batch_id) - .map(|entry| entry.batch.clone()) - .expect("evicted row's batch must be present in the store") - }; - let evicted = EvictedRow { - batch: evicted_batch, - index: prev_min.index, - row_bytes: prev_min.row.clone(), - }; - // Update batch use if prev_min.batch_id == batch_entry.id { batch_entry.uses -= 1; @@ -1005,14 +963,11 @@ impl TopKHeap { prev_min.replace_with(row, batch_id, index); self.owned_bytes += prev_min.owned_size(); - - Some(evicted) } else { let new_row = TopKRow::new(row, batch_id, index); self.owned_bytes += new_row.owned_size(); // put the new row into the heap self.inner.push(new_row); - None } } @@ -1033,14 +988,9 @@ impl TopKHeap { return Ok((None, topk_rows)); } - // Collect the batches into a vec and store the "batch_id -> array_pos" mapping, to then - // build the `indices` vec below. This is needed since the batch ids are not continuous. - let mut record_batches = Vec::new(); - let mut batch_id_array_pos = HashMap::new(); - for (array_pos, (batch_id, batch)) in self.store.batches.iter().enumerate() { - record_batches.push(&batch.batch); - batch_id_array_pos.insert(*batch_id, array_pos); - } + // The batch ids are not continuous, so `indices` below is built through + // the store's "batch_id -> array_pos" mapping. + let (record_batches, batch_id_array_pos) = self.store.positional(); let indices: Vec<_> = topk_rows .iter() @@ -1220,15 +1170,33 @@ impl RecordBatchStore { } /// The id the next [`Self::register`] call will assign. - /// - /// For callers that must reference a batch before they can build it — see - /// `PartitionedTopK::insert_batch`, which does not know which rows to gather - /// until it has finished deciding. Valid until the next `register`, and only - /// if nothing is left pointing at the id when that `register` is skipped. fn next_batch_id(&self) -> u32 { self.next_id } + /// Start the entry a partitioned operator's `insert_batch` fills: slots can + /// point at its id during the admission pass, before the gathered batch + /// exists. Finished by [`Self::insert_rows`]; valid until the next + /// `register`, which nothing may call in between. + fn pending(&self) -> PendingEntry { + PendingEntry { + id: self.next_batch_id(), + uses: 0, + } + } + + /// Drop one use of entry `id` for a retained row going away. A row of the + /// batch being inserted is not in the store yet, so its use comes off + /// `pending` rather than through [`Self::unuse`], which would panic on the + /// unregistered id. + fn release(&mut self, pending: &mut PendingEntry, id: u32) { + if id == pending.id { + pending.uses -= 1; + } else { + self.unuse(id); + } + } + /// Insert a record batch entry into this store, tracking its /// memory use, if it has any uses pub fn insert(&mut self, entry: RecordBatchEntry) { @@ -1240,6 +1208,35 @@ impl RecordBatchStore { } } + /// Finish `pending`: gather `rows` of `batch` into its entry and insert + /// it, or do nothing if no slot still points at it, leaving the id to the + /// next batch. + /// + /// For the partitioned operators' decide-then-gather: each slot's row is + /// its position in `rows`, which the gather keeps. Takes `pending` by + /// value so no use can be released against it once it is finished. + #[expect(clippy::needless_pass_by_value)] + fn insert_rows( + &mut self, + batch: &RecordBatch, + rows: &[u32], + pending: PendingEntry, + ) -> Result<()> { + let PendingEntry { id, uses } = pending; + if uses == 0 { + return Ok(()); + } + let gather_idx = UInt32Array::from_iter_values(rows.iter().copied()); + let mut entry = self.register(take_record_batch(batch, &gather_idx)?); + debug_assert_eq!( + entry.id, id, + "the id handed to the slots must be the id the gather got" + ); + entry.uses = uses; + self.insert(entry); + Ok(()) + } + /// Clear all values in this store, invalidating all previous batch ids fn clear(&mut self) { self.batches.clear(); @@ -1261,15 +1258,13 @@ impl RecordBatchStore { /// /// `interleave_record_batch` addresses its inputs positionally while this /// store keys them by a non-contiguous id, so every caller that interleaves - /// out of the store needs both. The batches are cloned — an `Arc` bump per - /// column — so they can outlive the store: `emit` needs that, since the - /// store is dropped as it returns. `compact_store` does not, and only - /// shares the helper. - fn positional(&self) -> (Vec, HashMap) { + /// out of the store needs both. The batches are borrowed; a caller that + /// needs them to outlive the store clones them itself. + fn positional(&self) -> (Vec<&RecordBatch>, HashMap) { let mut batches = Vec::with_capacity(self.batches.len()); let mut positions = HashMap::with_capacity(self.batches.len()); for (pos, (batch_id, entry)) in self.batches.iter().enumerate() { - batches.push(entry.batch.clone()); + batches.push(&entry.batch); positions.insert(*batch_id, pos); } (batches, positions) @@ -1313,6 +1308,206 @@ impl RecordBatchStore { + self.batches.capacity() * (size_of::() + size_of::()) + self.batches_size } + + /// Rewrite the store to hold only the rows `partitions` still reference, + /// once it holds [`STORE_COMPACTION_RATIO`]× more than the `live_slots` of + /// them, and repoint every partition at the rows' new places. + /// [`PartitionedTopK`] and [`PartitionedTopKRank`] call this at the end of + /// every `insert_batch`. + /// + /// An entry holds every row *admitted* from its input batch and is freed + /// only when the last of them is evicted, so when survivors spread thinly + /// one live row keeps a whole entry resident and residency tracks the + /// *input*, not `partitions × K`: 512 partitions of `K = 1` fed 512 batches + /// pin 131 K rows to retain 512. Nothing else bounds that. A single entry is + /// no exception — rows admitted then superseded within their own batch stay + /// in the gather unreferenced — so this does not skip a one-entry store. + /// `RANK` ties add a second way to leave rows unreferenced, since a + /// boundary move releases every tie of the partition at once. + /// + /// Amortized O(1) per admitted row: one pass over the live slots, and it + /// cannot recur until the store has taken on another `live_slots` rows. + /// + /// Rows are rewritten into `batch_size` chunks, not one batch: an entry is + /// released only when its last slot is evicted, so a single batch of every + /// live row would free nothing until every partition has churned. + /// + /// Peak residency is the old store plus the new one — chunks are built + /// before the old entries drop, and the reservation is not resized until + /// `insert_batch` returns — so a pool sized at the steady-state bound can be + /// exceeded transiently without erroring. + /// + /// All-or-nothing: plan the move and interleave, which is the only fallible + /// step, before rewriting the slots and the store. A failing interleave + /// leaves the operator as it was rather than holding slots pointing at ids + /// the store never got. + fn compact( + &mut self, + partitions: &mut HashMap, P>, + live_slots: usize, + batch_size: usize, + ) -> Result<()> { + if self.total_rows <= live_slots * STORE_COMPACTION_RATIO { + return Ok(()); + } + + // Scoped so the borrows of the old batches end before the store is + // rewritten. + let first_id = self.next_batch_id(); + // The slot's old place in walk order, so debug builds can check the + // repoint walk below visits the slots in the same order. + #[cfg(debug_assertions)] + let mut walked: Vec = Vec::with_capacity(live_slots); + let chunks = { + let (old, array_pos) = self.positional(); + + // The `n`th slot this walk visits moves to the `n`th row of the + // compacted store. The repoint walk below finds each slot's new + // place by counting, which holds because an unmodified map yields + // `values()` and `values_mut()` in the same order, and + // `StoreSlots` requires `store_rows` and `repoint` to agree. + let mut coords: Vec<(usize, usize)> = Vec::with_capacity(live_slots); + for at in partitions.values().flat_map(P::store_rows) { + let pos = *array_pos + .get(&at.batch_id) + .expect("a live slot's batch_id is present in the store"); + coords.push((pos, at.row as usize)); + #[cfg(debug_assertions)] + walked.push(at); + } + debug_assert_eq!( + coords.len(), + live_slots, + "live_slots must count exactly the slots that reference the store" + ); + #[cfg(debug_assertions)] + debug_assert_eq!( + walked + .iter() + .collect::>() + .len(), + walked.len(), + "two slots must not share a store row" + ); + + let mut chunks: Vec<(RecordBatch, usize)> = + Vec::with_capacity(coords.len().div_ceil(batch_size)); + let mut remap = vec![usize::MAX; old.len()]; + for chunk in coords.chunks(batch_size) { + chunks + .push((interleave_referenced(&old, chunk, &mut remap)?, chunk.len())); + } + chunks + }; + + let mut moved_to = 0usize; + for partition in partitions.values_mut() { + partition.repoint(|at| { + #[cfg(debug_assertions)] + assert_eq!( + walked[moved_to], at, + "store_rows and repoint must visit slots in the same order" + ); + #[cfg(not(debug_assertions))] + let _ = at; + let new_place = StoreRef { + batch_id: first_id + (moved_to / batch_size) as u32, + row: (moved_to % batch_size) as u32, + }; + moved_to += 1; + new_place + }); + } + debug_assert_eq!(moved_to, live_slots, "repoint must visit every slot"); + self.clear(); + for (chunk_idx, (batch, uses)) in chunks.into_iter().enumerate() { + let mut entry = self.register(batch); + debug_assert_eq!( + entry.id, + first_id + chunk_idx as u32, + "the ids handed to the slots must be the ids the chunks got" + ); + entry.uses = uses; + self.insert(entry); + } + Ok(()) + } +} + +/// `interleave_record_batch` over only the `batches` that `indices` +/// references. +/// +/// The kernel does per-input work for every batch it is handed — for a +/// dictionary column, a scan of all of `indices` per input — so handing it +/// the whole store would make every chunk's cost grow with the store. When a +/// chunk touches a few of many stored batches, as when input arrives +/// clustered by partition key, that is the difference between linear and +/// quadratic in the input. +/// +/// `remap` is scratch with one entry per batch, all `usize::MAX` on entry; +/// only the entries this chunk set are reset, so a chunk costs +/// `O(indices.len())` however many batches the store holds. +fn interleave_referenced>( + batches: &[B], + indices: &[(usize, usize)], + remap: &mut [usize], +) -> Result { + debug_assert_eq!(remap.len(), batches.len()); + let mut used: Vec = Vec::new(); + let local: Vec<(usize, usize)> = indices + .iter() + .map(|&(pos, row)| { + if remap[pos] == usize::MAX { + remap[pos] = used.len(); + used.push(pos); + } + (remap[pos], row) + }) + .collect(); + let refs: Vec<&RecordBatch> = used.iter().map(|&pos| batches[pos].borrow()).collect(); + let batch = interleave_record_batch(&refs, &local); + for pos in used { + remap[pos] = usize::MAX; + } + Ok(batch?) +} + +/// A partition's retained rows as [`RecordBatchStore`] coordinates: what +/// [`RecordBatchStore::compact`] reads to plan the move, and rewrites after it. +/// +/// `repoint` must visit the slots in the order `store_rows` yields them: +/// `compact` assigns new places by position in that walk rather than looking +/// each slot up. Debug builds check this. +trait StoreSlots { + /// Where every retained row lives. + fn store_rows(&self) -> impl Iterator + '_; + + /// Rewrite where every retained row lives through `moved_to`, called once + /// per slot in [`Self::store_rows`] order. + fn repoint(&mut self, moved_to: impl FnMut(StoreRef) -> StoreRef); +} + +/// The store entry a partitioned operator's `insert_batch` is filling, from +/// [`RecordBatchStore::pending`] to [`RecordBatchStore::insert_rows`]. +/// +/// Like a [`RecordBatchEntry`] whose batch does not exist yet: the rows to +/// gather are known only once the admission pass is over, but slots point at +/// `id` during it, and `uses` counts those still doing so. +#[derive(Debug)] +struct PendingEntry { + id: u32, + uses: usize, +} + +/// Store coordinates of one retained row: row `row` of the shared +/// [`RecordBatchStore`] batch `batch_id`. +/// +/// What a [`PartitionSlot`] holds besides its key, on its own. Rows tied at a +/// `RANK` boundary all share the boundary's key, so they need nothing more. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +struct StoreRef { + batch_id: u32, + row: u32, } /// One retained row: its sort key by value, its output columns by reference. @@ -1322,16 +1517,13 @@ impl RecordBatchStore { /// rows without dereferencing anything, and eviction needs no access to the /// payload at all. /// -/// `(batch_id, row)` locates the output columns in the operator's shared +/// `at` locates the output columns in the operator's shared /// [`RecordBatchStore`], which refcounts each batch by the slots pointing into /// it. A slot's existence is what keeps its batch alive. #[derive(Debug)] struct PartitionSlot { key: Vec, - /// Id in the shared store of the batch holding this row's output columns. - batch_id: u32, - /// Row of that batch. - row: u32, + at: StoreRef, } impl PartialEq for PartitionSlot { @@ -1369,36 +1561,43 @@ struct PartitionHeap { } impl PartitionHeap { - /// True if `key` belongs in the top-`k`. - /// - /// The worst retained row is the root, so this is one comparison — and - /// while the heap is not yet full every row qualifies. + /// True if `key` belongs in the top-`k`: the heap is not yet full, or + /// `key` beats the worst retained row. fn qualifies(&self, k: usize, key: &[u8]) -> bool { + matches!(self.classify(k, key), None | Some(Ordering::Less)) + } + + /// How `key` compares with the worst retained row, or `None` while the + /// heap holds fewer than `k` rows and so admits anything. + /// + /// The worst retained row is the root, so this is one comparison. `RANK` + /// needs all three outcomes, since `Equal` is a tie at the boundary; + /// `ROW_NUMBER` only needs [`Self::qualifies`]. + fn classify(&self, k: usize, key: &[u8]) -> Option { if self.inner.len() < k { - return true; - } - match self.inner.peek() { - Some(worst) => key < worst.key.as_slice(), - None => true, + return None; } + self.inner.peek().map(|worst| key.cmp(worst.key.as_slice())) + } + + /// The worst retained row's key — the admission boundary once the heap is + /// full. + fn worst_key(&self) -> Option<&[u8]> { + self.inner.peek().map(|worst| worst.key.as_slice()) } - /// Retain `key`, pointing at row `row` of store batch `batch_id`, evicting + /// Retain `key`, with its output columns at `at` in the store, evicting /// the worst retained row if the heap is already full. /// - /// Returns the evicted row's `batch_id`, which the caller owes the store an + /// Returns where the evicted row lives, which the caller owes the store an /// `unuse` for — this type does not touch the store itself, because the /// in-flight batch is not registered yet and only the caller knows that — - /// along with the bytes this call newly allocated, which the operator folds - /// into its running total. Returning the delta rather than keeping a running - /// sum per heap is what lets this type hold nothing but the heap. - fn add( - &mut self, - k: usize, - key: &[u8], - batch_id: u32, - row: u32, - ) -> (Option, usize) { + /// along with the bytes this call newly allocated, which the operator + /// folds into its running total. Returning the delta rather than keeping a + /// running sum per heap is what lets this type hold nothing but the heap. + /// `RANK` evicts through [`Self::replace_worst`] instead, since it needs + /// the evicted key back. + fn add(&mut self, k: usize, key: &[u8], at: StoreRef) -> (Option, usize) { debug_assert!(self.inner.len() <= k); if self.inner.len() == k { let mut worst = self.inner.peek_mut().expect("heap is full"); @@ -1406,16 +1605,13 @@ impl PartitionHeap { worst.key.clear(); worst.key.extend_from_slice(key); let grown = worst.key.capacity() - before; - let evicted = worst.batch_id; - worst.batch_id = batch_id; - worst.row = row; + let evicted = std::mem::replace(&mut worst.at, at); drop(worst); (Some(evicted), grown) } else { let slot = PartitionSlot { key: key.to_vec(), - batch_id, - row, + at, }; let key_bytes = slot.key.capacity(); let inner_before = self.inner.capacity(); @@ -1426,6 +1622,29 @@ impl PartitionHeap { } } + /// [`Self::add`] on a full heap that hands the evicted key back in + /// `evicted_key` rather than overwriting it: `key` is written into + /// `evicted_key`'s buffer and the two buffers swap, one copy instead of + /// two. Returns where the evicted row lives and the bytes `evicted_key` + /// grew by — the buffers only trade places, so their combined capacity + /// changes by exactly that. + fn replace_worst( + &mut self, + key: &[u8], + at: StoreRef, + evicted_key: &mut Vec, + ) -> (StoreRef, usize) { + let before = evicted_key.capacity(); + evicted_key.clear(); + evicted_key.extend_from_slice(key); + let grown = evicted_key.capacity() - before; + let mut worst = self.inner.peek_mut().expect("heap is full"); + std::mem::swap(&mut worst.key, evicted_key); + let evicted = std::mem::replace(&mut worst.at, at); + drop(worst); + (evicted, grown) + } + /// The retained slots, in unspecified order. fn slots(&self) -> impl Iterator + '_ { self.inner.iter() @@ -1435,37 +1654,37 @@ impl PartitionHeap { fn drain_sorted(&mut self) -> Vec { std::mem::take(&mut self.inner).into_sorted_vec() } +} - /// Rewrite every slot's store coordinates, leaving the keys — and so the - /// heap order — untouched. - /// - /// For [`PartitionedTopK::compact_store`], which moves every live row into - /// a fresh set of store batches and has to repoint the slots at it. - /// `into_vec` hands back the heap's own allocation and `BinaryHeap::from` - /// takes it again, so this allocates nothing. - fn repoint(&mut self, mut f: impl FnMut(&mut PartitionSlot)) { +impl StoreSlots for PartitionHeap { + fn store_rows(&self) -> impl Iterator + '_ { + self.slots().map(|slot| slot.at) + } + + /// Leaves the keys — and so the heap order — untouched. `into_vec` hands + /// back the heap's own allocation and `BinaryHeap::from` takes it again, + /// so this allocates nothing. + fn repoint(&mut self, mut moved_to: impl FnMut(StoreRef) -> StoreRef) { let mut slots = std::mem::take(&mut self.inner).into_vec(); for slot in &mut slots { - f(slot); + slot.at = moved_to(slot.at); } self.inner = BinaryHeap::from(slots); } } -/// Top-K-per-partition operator state. -/// -/// Sibling to [`TopK`]. Where `TopK` maintains a single global heap, -/// `PartitionedTopK` maintains one [`PartitionHeap`] per distinct partition -/// key while sharing its two [`RowConverter`]s — one for the partition key, -/// one for the ORDER BY key — along with the [`MemoryReservation`], the -/// scratch [`Rows`] buffers, the [`RecordBatchStore`] holding retained rows, -/// and [`TopKMetrics`] across all partitions. +/// What [`PartitionedTopK`] and [`PartitionedTopKRank`] share: the two +/// [`RowConverter`]s — one for the partition key, one for the ORDER BY key — +/// and their scratch [`Rows`], the [`MemoryReservation`], [`TopKMetrics`], one +/// per-partition state `S` per distinct partition key, and the +/// [`RecordBatchStore`] holding every retained row. /// -/// This sharing is the point of the type: with N distinct partition -/// keys, a naive `HashMap<_, TopK>` pays N × constant overhead for -/// `RowConverter::new`, `MemoryConsumer::register`, `RecordBatchStore::new`, -/// and metric counter setup. `PartitionedTopK` pays it once. -pub(crate) struct PartitionedTopK { +/// Both operators run the same three phases per batch: encode both column +/// sets once ([`Self::encode`]), one admission pass over the rows, then one +/// gather into the store ([`Self::finish_batch`]). Only the admission pass — +/// the retention rule — differs, so each operator writes its own row loop over +/// these fields between the two calls. +struct PartitionedTopKCore { schema: SchemaRef, metrics: TopKMetrics, reservation: MemoryReservation, @@ -1486,37 +1705,39 @@ pub(crate) struct PartitionedTopK { /// Scratch row buffer for partition-key encoding. Reused across /// `insert_batch` calls (cleared + appended each batch). partition_scratch_rows: Rows, - /// One heap per distinct partition key seen so far, keyed by the + /// One state per distinct partition key seen so far, keyed by the /// row-encoded PARTITION BY key. /// /// `entry_ref` owns the key only on Vacant, so a key is allocated once per /// partition for the lifetime of the operator — not once per row, and not /// once per partition per batch (which is what draining a per-batch map /// would cost). Map order is arbitrary, so `emit` sorts the keys itself. - heaps: HashMap, PartitionHeap>, + partitions: HashMap, S>, /// The batches holding every retained row's output columns, refcounted by /// the slots pointing into them. /// /// One store for the whole operator rather than one per partition, which is /// what [`TopKHeap`] would give. Each entry holds only the rows that were - /// admitted from one input batch — see `insert_batch` phase 3 — and is - /// dropped as soon as the last slot referencing it is evicted. When that - /// alone leaves it holding far more than the heaps point at, - /// `compact_store` rewrites it. + /// admitted from one input batch — see [`Self::finish_batch`] — and is + /// dropped as soon as the last slot referencing it is released. When that + /// alone leaves it holding far more than the partitions point at, + /// `RecordBatchStore::compact` rewrites it. store: RecordBatchStore, /// Rows of the batch currently being inserted that were admitted, in /// ascending row order. Reused across `insert_batch` calls. admitted_rows: Vec, - /// Rows the heaps currently hold, i.e. slots pointing into `store`. + /// Rows the partitions currently hold, i.e. slots pointing into `store`. /// /// Tracked incrementally for the same reason as `heaps_bytes`: it is the - /// denominator of `compact_store`'s ratio, read on every batch, and summing - /// `inner.len()` over the heaps would be O(partitions seen so far). Slots - /// are replaced rather than removed, so this only grows — it settles at - /// `partitions × K` once every partition's heap has filled. + /// denominator of `RecordBatchStore::compact`'s ratio, read on every + /// batch, and summing over the partitions would be O(partitions seen so + /// far). For `ROW_NUMBER` slots are replaced rather than removed, so it + /// only grows and settles at `partitions × K`; for `RANK` it also shrinks, + /// when a boundary move releases a partition's ties. live_slots: usize, /// Running sum of the bytes every [`PartitionHeap`] has allocated: each - /// slot's key and the `BinaryHeap`'s own buffer. + /// slot's key and the `BinaryHeap`'s own buffer (for `RANK`, also + /// `evicted_key`'s buffer, which swaps with heap keys). /// /// Maintained incrementally because `size()` runs on every batch: summing /// over the heaps would make it O(partitions seen so far) per batch, which @@ -1526,8 +1747,8 @@ pub(crate) struct PartitionedTopK { /// than keeping its own running sum, so the heaps stay free of a field that /// would be multiplied by the partition count. heaps_bytes: usize, - /// Running sum of the partition keys `heaps` has interned, one per distinct - /// key for the operator's lifetime. + /// Running sum of the partition keys `partitions` has interned, one per + /// distinct key for the operator's lifetime. /// /// Counted by length rather than capacity: `entry_ref` builds the owned key /// with `to_vec`, so the two are equal. Separate from [`Self::heaps_bytes`] @@ -1537,9 +1758,11 @@ pub(crate) struct PartitionedTopK { batch_size: usize, } -impl PartitionedTopK { +impl PartitionedTopKCore { + /// `name` labels the memory consumer, as `name[partition_id]`. #[expect(clippy::too_many_arguments)] - pub(crate) fn try_new( + fn try_new( + name: &str, partition_id: usize, schema: SchemaRef, partition_exprs: Vec>, @@ -1550,8 +1773,8 @@ impl PartitionedTopK { runtime: &Arc, metrics: &ExecutionPlanMetricsSet, ) -> Result { - assert!(k > 0, "PartitionedTopK requires k > 0"); - let reservation = MemoryConsumer::new(format!("PartitionedTopK[{partition_id}]")) + assert!(k > 0, "{name} requires k > 0"); + let reservation = MemoryConsumer::new(format!("{name}[{partition_id}]")) .register(&runtime.memory_pool); // Both encoders are shared by every partition, and each scratch buffer @@ -1576,7 +1799,7 @@ impl PartitionedTopK { partition_exprs, partition_converter, partition_scratch_rows, - heaps: HashMap::new(), + partitions: HashMap::new(), store: RecordBatchStore::new(), admitted_rows: Vec::new(), live_slots: 0, @@ -1587,22 +1810,12 @@ impl PartitionedTopK { }) } - /// Encode the partition and ORDER BY columns once each for the whole batch, - /// admit every qualifying row into the [`PartitionHeap`] of the partition it - /// belongs to, then register the admitted rows in the shared store. - pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { - let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); - let _timer = elapsed_compute.timer(); - + /// Phase 1: evaluate the partition and ORDER BY columns of a non-empty + /// `batch` and encode each once for the whole batch into the scratch + /// buffers. Both encodes are whole-batch kernels that do not care how the + /// rows group, which is what lets the admission pass be a single row loop. + fn encode(&mut self, batch: &RecordBatch) -> Result<()> { let num_rows = batch.num_rows(); - if num_rows == 0 { - return Ok(()); - } - - // 1. Evaluate the partition and ORDER BY columns and encode each once - // for the whole batch. Both encodes are whole-batch kernels that do - // not care how the rows group, which is what lets the admission pass - // below be a single row loop. let pk_arrays: Vec = self .partition_exprs .iter() @@ -1620,6 +1833,158 @@ impl PartitionedTopK { self.scratch_rows.clear(); self.row_converter .append(&mut self.scratch_rows, &ob_arrays)?; + Ok(()) + } + + /// Phase 3: gather the rows this batch contributed into a single batch and + /// hand it to the store, then compact the store if it needs it. Only rows + /// admitted at some point during the pass are kept, so what stays pinned is + /// bounded by admissions rather than by input size — and the entry is + /// freed as soon as the last slot referencing it is released. + /// + /// No uses left on `pending` means every admission from this batch was + /// released again before the pass ended, so there is nothing to keep and + /// the id goes back to the next batch. + /// + /// The caller resizes the reservation afterwards, since only it knows its + /// full size. + fn finish_batch( + &mut self, + batch: &RecordBatch, + pending: PendingEntry, + replacements: usize, + ) -> Result<()> { + self.store + .insert_rows(batch, &self.admitted_rows, pending)?; + if replacements > 0 { + self.metrics.row_replacements.add(replacements); + } + self.store + .compact(&mut self.partitions, self.live_slots, self.batch_size) + } + + /// Drain every partition in partition-key order, each through `drain` into + /// its retained rows in output order, and return them as a stream of + /// `RecordBatch`es ordered by `(partition_keys, order_keys)`. + /// + /// Only the order is resolved here; [`EmitState::stream`] interleaves the + /// rows out one `batch_size` chunk per poll, carrying the reservation until + /// the stream is dropped since the store's batches stay pinned until then. + fn emit_with>( + self, + drain: impl FnMut(S) -> I, + ) -> Result { + let Self { + schema, + metrics, + reservation, + partitions, + store, + live_slots, + batch_size, + .. + } = self; + + // Map order is arbitrary, so partition-key order has to be recovered + // explicitly here. + let timer = metrics.baseline.elapsed_compute().timer(); + let mut sorted: Vec<(Vec, S)> = partitions.into_iter().collect(); + sorted.sort_unstable_by(|(a, _), (b, _)| a.cmp(b)); + drop(timer); + + let slots = sorted + .into_iter() + .map(|(_key, state)| state) + .flat_map(drain); + EmitState::stream( + schema, + metrics, + reservation, + batch_size, + &store, + slots, + live_slots, + ) + } + + /// The bytes these shared fields hold on the heap, beyond the operator + /// struct itself: encoders, scratch buffers, the partition map and its + /// keys, every heap, and every batch the store still pins. + /// + /// Every term is O(1): this runs on every batch, so the per-partition + /// contributions are the running totals `heaps_bytes` and `index_bytes` + /// rather than a sum over partitions. + fn allocated_size(&self) -> usize { + self.row_converter.size() + + self.partition_converter.size() + + self.scratch_rows.size() + + self.partition_scratch_rows.size() + + self.admitted_rows.allocated_size() + + self.partitions.capacity() * (size_of::>() + size_of::()) + + self.heaps_bytes + + self.index_bytes + + self.store.size() + } +} + +/// Top-K-per-partition operator state. +/// +/// Sibling to [`TopK`]. Where `TopK` maintains a single global heap, +/// `PartitionedTopK` maintains one [`PartitionHeap`] per distinct partition +/// key while sharing everything else in [`PartitionedTopKCore`] across all +/// partitions. +/// +/// This sharing is the point of the type: with N distinct partition +/// keys, a naive `HashMap<_, TopK>` pays N × constant overhead for +/// `RowConverter::new`, `MemoryConsumer::register`, `RecordBatchStore::new`, +/// and metric counter setup. `PartitionedTopK` pays it once. +pub(crate) struct PartitionedTopK { + core: PartitionedTopKCore, +} + +impl PartitionedTopK { + #[expect(clippy::too_many_arguments)] + pub(crate) fn try_new( + partition_id: usize, + schema: SchemaRef, + partition_exprs: Vec>, + partition_sort_fields: Vec, + order_expr: LexOrdering, + k: usize, + batch_size: usize, + runtime: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result { + Ok(Self { + core: PartitionedTopKCore::try_new( + "PartitionedTopK", + partition_id, + schema, + partition_exprs, + partition_sort_fields, + order_expr, + k, + batch_size, + runtime, + metrics, + )?, + }) + } + + /// Encode the partition and ORDER BY columns once each for the whole batch, + /// admit every qualifying row into the [`PartitionHeap`] of the partition it + /// belongs to, then register the admitted rows in the shared store. + pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let elapsed_compute = self.core.metrics.baseline.elapsed_compute().clone(); + let _timer = elapsed_compute.timer(); + + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } + + // 1. Encode both column sets once for the whole batch. + self.core.encode(batch)?; // 2. One pass over the rows: find the row's partition and admit it to // that partition's heap if it qualifies. Only the key is copied @@ -1633,20 +1998,19 @@ impl PartitionedTopK { // it in ascending row order — and it costs no pass over the // partitions seen so far, which would be quadratic in the input // whenever partition count grows with it. - let k = self.k; + let core = &mut self.core; + let k = core.k; let mut replacements: usize = 0; - self.admitted_rows.clear(); + core.admitted_rows.clear(); // The gathered batch does not exist yet, but its id does, so a slot can - // point at it during the pass with no back-patching afterwards. Its - // `uses` is counted locally for the same reason. - let batch_id = self.store.next_batch_id(); - let mut uses = 0usize; + // point at it during the pass with no back-patching afterwards. + let mut pending = core.store.pending(); { - let pk_rows = &self.partition_scratch_rows; - let ob_rows = &self.scratch_rows; - let heaps = &mut self.heaps; - let admitted_rows = &mut self.admitted_rows; - let store = &mut self.store; + let pk_rows = &core.partition_scratch_rows; + let ob_rows = &core.scratch_rows; + let heaps = &mut core.partitions; + let admitted_rows = &mut core.admitted_rows; + let store = &mut core.store; // Accumulated locally and folded in once, so the running totals are // not touched per row. let mut interned_bytes = 0usize; @@ -1673,17 +2037,15 @@ impl PartitionedTopK { } // An admission's row in the gathered batch is its position in // `admitted_rows`, because the gather preserves that order. - let gather_pos = admitted_rows.len() as u32; - uses += 1; - let (evicted, grown) = heap.add(k, key.as_ref(), batch_id, gather_pos); + let at = StoreRef { + batch_id: pending.id, + row: admitted_rows.len() as u32, + }; + pending.uses += 1; + let (evicted, grown) = heap.add(k, key.as_ref(), at); admitted_bytes += grown; - // Mirrors `TopKHeap::add`: a row evicted from the batch being - // inserted is not in the store yet, so its use comes off the - // local count rather than through `unuse`, which would panic on - // an unregistered id. match evicted { - Some(evicted_id) if evicted_id == batch_id => uses -= 1, - Some(evicted_id) => store.unuse(evicted_id), + Some(evicted) => store.release(&mut pending, evicted.batch_id), // The heap was not yet full, so this slot is a new // reference into the store rather than a replaced one. None => new_slots += 1, @@ -1692,144 +2054,16 @@ impl PartitionedTopK { replacements += 1; } - self.index_bytes += interned_bytes; - self.heaps_bytes += admitted_bytes; - self.live_slots += new_slots; + core.index_bytes += interned_bytes; + core.heaps_bytes += admitted_bytes; + core.live_slots += new_slots; } - // 3. Gather the rows this batch contributed into a single batch and hand - // it to the store. Only rows admitted at some point during the pass - // are kept, so what stays pinned is bounded by admissions rather than - // by input size — and the entry is freed as soon as the last slot - // referencing it is evicted. - // - // `uses == 0` means every admission from this batch was evicted again - // before the pass ended, so there is nothing to keep and the id goes - // back to the next batch. - if uses > 0 { - let gather_idx = - UInt32Array::from_iter_values(self.admitted_rows.iter().copied()); - let mut entry = self.store.register(take_record_batch(batch, &gather_idx)?); - debug_assert_eq!( - entry.id, batch_id, - "the id handed to the slots must be the id the gather got" - ); - entry.uses = uses; - self.store.insert(entry); - } - - if replacements > 0 { - self.metrics.row_replacements.add(replacements); - } - self.compact_store()?; - self.reservation.try_resize(self.size())?; - Ok(()) - } - - /// Rewrite the store to hold only the rows a heap still points at, once it - /// holds [`STORE_COMPACTION_RATIO`]× more than that. - /// - /// An entry holds every row *admitted* from its input batch and is freed - /// only when the last of them is evicted, so when survivors spread thinly - /// one live row keeps a whole entry resident and residency tracks the - /// *input*, not `partitions × K`: 512 partitions of `K = 1` fed 512 batches - /// pin 131 K rows to retain 512. Nothing else bounds that. A single entry is - /// no exception — rows admitted then superseded within their own batch stay - /// in the gather unreferenced — so this does not skip a one-entry store. - /// - /// Amortized O(1) per admitted row: one pass over the live slots, and it - /// cannot recur until the store has taken on another `live_slots` rows. - /// - /// Rows are rewritten into `batch_size` chunks, not one batch: an entry is - /// released only when its last slot is evicted, so a single batch of every - /// live row would free nothing until every partition has churned. - /// - /// Peak residency is the old store plus the new one — chunks are built - /// before the old entries drop, and the reservation is not resized until - /// `insert_batch` returns — so a pool sized at the steady-state bound can be - /// exceeded transiently without erroring. - /// - /// All-or-nothing: plan the move, interleave, then rewrite the slots and the - /// store. A failing interleave leaves the operator as it was rather than - /// holding slots pointing at ids the store never got. - fn compact_store(&mut self) -> Result<()> { - if self.store.total_rows <= self.live_slots * STORE_COMPACTION_RATIO { - return Ok(()); - } - - // These clones keep the old batches alive while the compacted ones are - // built. - let (old, array_pos) = self.store.positional(); - - // Plan the move without touching anything: where each live row is now, - // and where it is going. Keyed by the row's current `(batch_id, row)` - // rather than by its position in this walk, so the rewrite below is free - // to visit the heaps in any order — no two slots share a store row, so - // the key identifies exactly one slot. - let first_id = self.store.next_batch_id(); - let batch_size = self.batch_size; - let mut coords: Vec<(usize, usize)> = Vec::with_capacity(self.live_slots); - let mut moved: HashMap<(u32, u32), (u32, u32)> = - HashMap::with_capacity(self.live_slots); - for heap in self.heaps.values() { - for slot in heap.slots() { - let pos = *array_pos - .get(&slot.batch_id) - .expect("a live slot's batch_id is present in the store"); - let moved_to = coords.len(); - coords.push((pos, slot.row as usize)); - moved.insert( - (slot.batch_id, slot.row), - ( - first_id + (moved_to / batch_size) as u32, - (moved_to % batch_size) as u32, - ), - ); - } - } - debug_assert_eq!( - coords.len(), - self.live_slots, - "live_slots must count exactly the slots the heaps hold" - ); - debug_assert_eq!( - moved.len(), - coords.len(), - "two slots must not share a store row" - ); - - // The only fallible step, and nothing has been mutated yet: an error - // here leaves both the heaps and the store exactly as they were. - let refs: Vec<&RecordBatch> = old.iter().collect(); - let mut compacted: Vec<(RecordBatch, usize)> = - Vec::with_capacity(coords.len().div_ceil(batch_size)); - for chunk in coords.chunks(batch_size) { - compacted.push((interleave_record_batch(&refs, chunk)?, chunk.len())); - } - drop(refs); - drop(old); - - // Infallible from here, so the store's ids and the slots agree again. - for heap in self.heaps.values_mut() { - heap.repoint(|slot| { - let (batch_id, row) = moved[&(slot.batch_id, slot.row)]; - slot.batch_id = batch_id; - slot.row = row; - }); - } - self.store.clear(); - for (chunk_idx, (batch, uses)) in compacted.into_iter().enumerate() { - let mut entry = self.store.register(batch); - debug_assert_eq!( - entry.id, - first_id + chunk_idx as u32, - "the ids handed to the slots must be the ids the chunks got" - ); - entry.uses = uses; - self.store.insert(entry); - } - Ok(()) - } + // 3. Gather this batch's admissions into the store, then compact it. + core.finish_batch(batch, pending, replacements)?; + self.core.reservation.try_resize(self.size())?; + Ok(()) + } /// Drain all heaps in partition-key order and return the rows as /// a stream of `RecordBatch`es ordered by @@ -1846,68 +2080,91 @@ impl PartitionedTopK { /// the reservation is carried into [`EmitState`] and released when the stream /// is dropped rather than when this returns. pub(crate) fn emit(self) -> Result { - let Self { - schema, - metrics, - reservation, - expr: _, - row_converter: _, - scratch_rows: _, - partition_exprs: _, - partition_converter: _, - partition_scratch_rows: _, - heaps, - store, - admitted_rows: _, - live_slots: _, - heaps_bytes: _, - index_bytes: _, - k: _, - batch_size, - } = self; - let timer = metrics.baseline.elapsed_compute().timer(); + self.core + .emit_with(|mut heap| heap.drain_sorted().into_iter().map(|slot| slot.at)) + } - // Map order is arbitrary, so partition-key order has to be recovered - // explicitly here. - let mut sorted_groups: Vec<(Vec, PartitionHeap)> = - heaps.into_iter().collect(); - sorted_groups.sort_unstable_by(|(a, _), (b, _)| a.cmp(b)); + /// Total memory currently held by this operator, including all + /// per-partition heaps and every batch the store still pins. O(1), see + /// [`PartitionedTopKCore::allocated_size`]. + fn size(&self) -> usize { + size_of::() + self.core.allocated_size() + } +} - // The batches outlive the store, which is dropped as this returns. +/// Hands out `batch_size` rows at a time, interleaved out of the store batches +/// the heaps referenced. +struct EmitState { + metrics: TopKMetrics, + /// Covers the pinned batches and `ordered` for as long as they are held. + /// + /// Carried here rather than dropped at the end of `emit` so the bytes stay + /// accounted for until the stream is, and released by this struct's drop. + /// Never resized: the store's batches are pinned until the last chunk, so + /// there is nothing to hand back as chunks are emitted. Underscored because + /// it is held only for that drop, as in `HashJoinStream` and `AsofJoinStream`. + _reservation: MemoryReservation, + batch_size: usize, + /// The store's batches, positionally indexed by `ordered`. + batches: Vec, + /// `(array_pos, row)` pairs in `(partition_keys, order_keys)` order. + ordered: Vec<(usize, usize)>, + /// Scratch for [`interleave_referenced`], one entry per batch. + remap: Vec, + pos: usize, +} + +impl EmitState { + /// Resolve `slots`, `len` of them in output order, against `store`, and + /// return the stream that emits them. Shared by [`PartitionedTopK::emit`] + /// and [`PartitionedTopKRank::emit`], which differ only in that order. + fn stream( + schema: SchemaRef, + metrics: TopKMetrics, + reservation: MemoryReservation, + batch_size: usize, + store: &RecordBatchStore, + slots: impl Iterator, + len: usize, + ) -> Result { + let timer = metrics.baseline.elapsed_compute().timer(); + // Cloned, an `Arc` bump per column, because the batches outlive the + // store, which the caller drops as `emit` returns. let (batches, batch_id_array_pos) = store.positional(); + let batches: Vec = batches.into_iter().cloned().collect(); // Flattened in output order, so the emit itself is a slice walk. - let mut ordered: Vec<(usize, usize)> = Vec::new(); - for (_key, mut heap) in sorted_groups { - for slot in heap.drain_sorted() { - let array_pos = *batch_id_array_pos - .get(&slot.batch_id) - .expect("a retained slot's batch_id is present in the store"); - ordered.push((array_pos, slot.row as usize)); - } + let mut ordered: Vec<(usize, usize)> = Vec::with_capacity(len); + for slot in slots { + let array_pos = *batch_id_array_pos + .get(&slot.batch_id) + .expect("a retained slot's batch_id is present in the store"); + ordered.push((array_pos, slot.row as usize)); } drop(timer); // What survives this function is the store's batches — pinned until the // returned stream is dropped — plus `ordered`, one 16-byte pair per - // retained row. Everything else the operator held (both scratch - // buffers, every heap and its interned key) is freed above, so the - // resize below is normally a shrink. The reservation moves into the - // stream state rather than being dropped here: releasing it while the - // store is still pinned would stop accounting for bytes that are still - // held, which is the one direction that matters. - let pinned_bytes = store.batches_size; + // retained row, and `remap`, one entry per batch. Everything else the + // operator held (its scratch buffers, every partition and its interned + // key) is freed by now, so the resize below is normally a shrink. The + // reservation moves into the stream state rather than being dropped + // here: releasing it while the store is still pinned would stop + // accounting for bytes that are still held, which is the one direction + // that matters. reservation.try_resize( size_of::() - + pinned_bytes + + store.batches_size + batches.capacity() * size_of::() - + ordered.capacity() * size_of::<(usize, usize)>(), + + ordered.capacity() * size_of::<(usize, usize)>() + + batches.len() * size_of::(), )?; let state = EmitState { metrics, _reservation: reservation, batch_size, + remap: vec![usize::MAX; batches.len()], batches, ordered, pos: 0, @@ -1921,47 +2178,6 @@ impl PartitionedTopK { ))) } - /// Total memory currently held by this operator, including all - /// per-partition heaps and every batch the store still pins. - /// - /// Every term is O(1): this runs on every batch, so the per-partition - /// contributions are the running totals `heaps_bytes` and `index_bytes` - /// rather than a sum over partitions. - fn size(&self) -> usize { - size_of::() - + self.row_converter.size() - + self.partition_converter.size() - + self.scratch_rows.size() - + self.partition_scratch_rows.size() - + self.admitted_rows.allocated_size() - + self.heaps.capacity() * (size_of::>() + size_of::()) - + self.heaps_bytes - + self.index_bytes - + self.store.size() - } -} - -/// Hands out `batch_size` rows at a time, interleaved out of the store batches -/// the heaps referenced. -struct EmitState { - metrics: TopKMetrics, - /// Covers the pinned batches and `ordered` for as long as they are held. - /// - /// Carried here rather than dropped at the end of `emit` so the bytes stay - /// accounted for until the stream is, and released by this struct's drop. - /// Never resized: the store's batches are pinned until the last chunk, so - /// there is nothing to hand back as chunks are emitted. Underscored because - /// it is held only for that drop, as in `HashJoinStream` and `AsofJoinStream`. - _reservation: MemoryReservation, - batch_size: usize, - /// The store's batches, positionally indexed by `ordered`. - batches: Vec, - /// `(array_pos, row)` pairs in `(partition_keys, order_keys)` order. - ordered: Vec<(usize, usize)>, - pos: usize, -} - -impl EmitState { /// Build the next chunk, or `None` once every retained row has been /// emitted. fn next_batch(&mut self) -> Result> { @@ -1974,55 +2190,50 @@ impl EmitState { let chunk = &self.ordered[self.pos..end]; self.pos = end; - let refs: Vec<&RecordBatch> = self.batches.iter().collect(); - let batch = interleave_record_batch(&refs, chunk)?; + let batch = interleave_referenced(&self.batches, chunk, &mut self.remap)?; (&batch).record_output(&self.metrics.baseline); Ok(Some(batch)) } } -/// Rows that tied at the boundary when inserted, materialized into a -/// batch holding *only* those rows. -/// -/// The rows are gathered eagerly rather than kept as `(source_batch, -/// indices)`: a tie entry lives until the boundary moves, so holding the -/// source batch would pin an entire input batch — and charge for it — -/// for as long as a single row of it stays tied. With ties spread across -/// many input batches that makes retained memory grow with the *input* -/// size instead of with `K + ties`. -#[derive(Debug)] -struct TieEntry { - /// The tied rows, and nothing else. Always non-empty by construction. - batch: RecordBatch, - /// `get_record_batch_memory_size(&batch)` captured at push time so - /// `RankPartitionState::size()` doesn't recurse through `batch`'s - /// columns on every `try_resize` call. - batch_bytes: usize, -} - /// Per-partition state for `RANK()` semantics. /// -/// Composes [`TopKHeap`] as the K-bounded core plus a sibling -/// `Vec` for boundary-tied rows. `RANK ≤ K` keeps the K -/// best rows by ORDER BY plus every row tied at the K-th-best -/// ORDER BY value — the boundary. So the total retained rows can -/// exceed K when ties straddle the boundary. +/// `RANK ≤ K` keeps the K best rows by ORDER BY plus every row tied at the +/// K-th-best ORDER BY value — the boundary — so a partition can retain more +/// than K rows when ties straddle it. The K best live in `heap`; the rows +/// tied with its root beyond those K live in `ties`. +/// +/// Like [`PartitionHeap`], this holds only store coordinates: the output +/// columns of every retained row, heap or tie, live in the operator's one +/// shared [`RecordBatchStore`]. +#[derive(Debug, Default)] struct RankPartitionState { - heap: TopKHeap, - ties: Vec, + heap: PartitionHeap, + /// Rows tied with `heap`'s root. They all share its key, so none is stored. + /// Only ever non-empty while `heap` is full. + 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.batch_bytes).sum(); - self.heap.size() + ties_buffer + ties_contents + /// Add a tie at `at`, returning the bytes this newly allocated for the + /// operator to fold into its running total, as [`PartitionHeap::add`] does. + fn push_tie(&mut self, at: StoreRef) -> usize { + let before = self.ties.capacity(); + self.ties.push(at); + (self.ties.capacity() - before) * size_of::() } +} - /// Push `batch`'s rows onto the tie list, charging exactly their bytes. - fn push_ties(&mut self, batch: RecordBatch) { - let batch_bytes = get_record_batch_memory_size(&batch); - self.ties.push(TieEntry { batch, batch_bytes }); +impl StoreSlots for RankPartitionState { + fn store_rows(&self) -> impl Iterator + '_ { + self.heap.store_rows().chain(self.ties.iter().copied()) + } + + fn repoint(&mut self, mut moved_to: impl FnMut(StoreRef) -> StoreRef) { + self.heap.repoint(&mut moved_to); + for tie in &mut self.ties { + *tie = moved_to(*tie); + } } } @@ -2030,57 +2241,37 @@ impl RankPartitionState { /// /// Per partition, retains the K-best rows plus every row tied at the /// K-th-best ORDER BY value (so `WHERE rk <= K` may keep more than K -/// rows when ties straddle the boundary). Like [`PartitionedTopK`], -/// the [`RowConverter`], [`MemoryReservation`], scratch [`Rows`] -/// buffer, and [`TopKMetrics`] are shared across all partitions for -/// this operator instance. +/// rows when ties straddle the boundary). Shares everything in +/// [`PartitionedTopKCore`] with [`PartitionedTopK`] and admits rows the same +/// way — decide per row, then gather everything the batch contributed once — +/// with a [`RankPartitionState`] per partition in place of a bare heap. /// /// # Algorithm (per row) /// -/// For each incoming row, compare its encoded ORDER BY bytes against -/// `heap.max()` — the K-th-best row, which is by definition the -/// admission boundary. `heap.max()` is `None` until the heap fills -/// to K rows: +/// Compare the row's encoded ORDER BY bytes against the heap's root — the +/// K-th-best row, which is by definition the admission boundary. There is +/// no boundary until the heap fills to K rows: /// -/// - heap not full (`max() == None`) → forward to the heap -/// - row's ob `==` max → push to ties (no heap call) -/// - row's ob `>` max → drop -/// - row's ob `<` max → forward to heap; on eviction, compare the -/// new `heap.max()` to the evicted row's bytes: if equal, push -/// evicted to ties (still tied at the new boundary's rank); else -/// clear ties (boundary moved up, old ties no longer satisfy -/// `rk ≤ K`) +/// - heap not full → admit to the heap +/// - row's ob `==` boundary → admit to ties +/// - row's ob `>` boundary → drop +/// - row's ob `<` boundary → admit to the heap, evicting the root. If the +/// new root still has the evicted row's key, the evicted row is still tied +/// at the boundary and moves to ties; otherwise the boundary moved up, the +/// old ties no longer satisfy `rk ≤ K`, and they are released along with +/// the evicted row. pub(crate) struct PartitionedTopKRank { - schema: SchemaRef, - metrics: TopKMetrics, - reservation: MemoryReservation, - /// ORDER BY expressions (excludes PARTITION BY). - expr: LexOrdering, - /// Encoder for ORDER BY columns. Reused across partitions. - row_converter: RowConverter, - /// Scratch row buffer reused across `insert_batch` calls. - scratch_rows: Rows, - /// PARTITION BY expressions. - partition_exprs: Vec>, - /// Encoder for the partition key. - partition_converter: RowConverter, - /// Scratch row buffer for partition-key encoding. Reused across - /// `insert_batch` calls (cleared + appended each batch) so we - /// avoid allocating a fresh `Rows` buffer every batch. - partition_scratch_rows: Rows, - /// One rank state per distinct partition key seen so far. Keyed by - /// the row-encoded PARTITION BY bytes (a byte-comparable encoding, so - /// the `Vec` hashes, compares, and sorts identically to an - /// `OwnedRow`) which lets `insert_batch` look partitions up with - /// `entry_ref` — allocating a key only on first sight of a partition - /// rather than once per row. - states: HashMap, RankPartitionState>, - /// Scratch map reused across `insert_batch` calls to group a batch's - /// row indices by partition key. Drained (not reallocated) each batch - /// so its backing table is allocated once, not per batch. - partition_groups: HashMap, Vec>, - k: usize, - batch_size: usize, + core: PartitionedTopKCore, + /// The key of the row an admission just evicted, so the new boundary can + /// be compared against it. Trades buffers with the heap root on every + /// strictly-better admission, so its capacity is counted in + /// `core.heaps_bytes`. + evicted_key: Vec, + /// Running sum of every tie list's buffer. A boundary move frees the + /// partition's buffer rather than clearing it, so a partition that once + /// held many ties does not keep their capacity for the operator's + /// lifetime; this shrinks with it. + ties_bytes: usize, } impl PartitionedTopKRank { @@ -2096,43 +2287,29 @@ impl PartitionedTopKRank { runtime: &Arc, metrics: &ExecutionPlanMetricsSet, ) -> Result { - assert!(k > 0, "PartitionedTopKRank requires k > 0"); - let reservation = - MemoryConsumer::new(format!("PartitionedTopKRank[{partition_id}]")) - .register(&runtime.memory_pool); - - let order_sort_fields = build_sort_fields(&order_expr, &schema)?; - let row_converter = RowConverter::new(order_sort_fields)?; - let scratch_rows = - row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); - - let partition_converter = RowConverter::new(partition_sort_fields)?; - let partition_scratch_rows = partition_converter - .empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); - Ok(Self { - schema, - metrics: TopKMetrics::new(metrics, partition_id), - reservation, - expr: order_expr, - row_converter, - scratch_rows, - partition_exprs, - partition_converter, - partition_scratch_rows, - states: HashMap::new(), - partition_groups: HashMap::new(), - k, - batch_size, + core: PartitionedTopKCore::try_new( + "PartitionedTopKRank", + partition_id, + schema, + partition_exprs, + partition_sort_fields, + order_expr, + k, + batch_size, + runtime, + metrics, + )?, + evicted_key: Vec::new(), + ties_bytes: 0, }) } - /// Demultiplex `batch` rows by partition key, encode the ORDER BY - /// columns once for the whole batch, and feed each partition's - /// rows through the rank classifier into its dedicated heap and - /// ties Vec. + /// Encode the partition and ORDER BY columns once each for the whole batch, + /// run every row through its partition's rank classifier, then register + /// the admitted rows in the shared store with one gather. pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { - let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); + let elapsed_compute = self.core.metrics.baseline.elapsed_compute().clone(); let _timer = elapsed_compute.timer(); let num_rows = batch.num_rows(); @@ -2140,256 +2317,134 @@ impl PartitionedTopKRank { return Ok(()); } - // 1. Evaluate + encode partition columns into the reusable - // scratch (cleared then appended). - let pk_arrays: Vec = self - .partition_exprs - .iter() - .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) - .collect::>()?; - self.partition_scratch_rows.clear(); - self.partition_converter - .append(&mut self.partition_scratch_rows, &pk_arrays)?; - - // 2. Demultiplex row indices by partition key (per-batch). - // `partition_groups` is a reused scratch map: taken out here and - // drained below, so its backing table is allocated once for the - // operator, not once per batch. `entry_ref` owns the key only on - // Vacant, so it allocates one `Vec` per distinct partition - // rather than one per row. - let mut groups = std::mem::take(&mut self.partition_groups); - groups.clear(); - { - let pk_rows = &self.partition_scratch_rows; - for i in 0..num_rows { - groups - .entry_ref(pk_rows.row(i).as_ref()) - .or_default() - .push(i as u32); - } - } - - // 3. Evaluate ORDER BY columns on the full batch and encode ONCE. - let ob_arrays: Vec = self - .expr - .iter() - .map(|e| e.expr.evaluate(batch).and_then(|v| v.into_array(num_rows))) - .collect::>()?; - self.scratch_rows.clear(); - self.row_converter - .append(&mut self.scratch_rows, &ob_arrays)?; - - // 4. Per-partition: classify each row and dispatch. - let k = self.k; + // 1. Encode both column sets once for the whole batch. + self.core.encode(batch)?; + + // 2. One pass over the rows in batch order, as in + // `PartitionedTopK::insert_batch`: a partition's heap and ties are + // only touched by its own rows, which still arrive in ascending row + // order, so interleaving partitions cannot change a decision. An + // admission, to the heap or to ties, records which row it kept; + // phase 3 gathers them all at once. + let core = &mut self.core; + let k = core.k; let mut replacements: usize = 0; + core.admitted_rows.clear(); + // The gathered batch's id is known before the batch exists, so slots + // can point at it during the pass. + let mut pending = core.store.pending(); + { + let pk_rows = &core.partition_scratch_rows; + let ob_rows = &core.scratch_rows; + let states = &mut core.partitions; + let admitted_rows = &mut core.admitted_rows; + let store = &mut core.store; + let evicted_key = &mut self.evicted_key; + // Accumulated locally and folded in once. + let mut interned_bytes = 0usize; + let mut admitted_bytes = 0usize; + let mut ties_grown = 0usize; + let mut ties_freed = 0usize; + let mut new_slots = 0usize; + let mut released_slots = 0usize; - for (pk, indices) in groups.drain() { - let state = self.states.entry(pk).or_insert_with(|| RankPartitionState { - heap: TopKHeap::new(k), - ties: Vec::new(), - }); - - // Once the heap is full, a group whose rows are *all* strictly - // worse than the boundary changes neither the heap nor the - // ties. Bail before the gather below — at high partition - // cardinality this is the common case. - if let Some(max_row) = state.heap.max() { - let boundary = max_row.row(); - if indices - .iter() - .all(|&i| self.scratch_rows.row(i as usize).as_ref() > boundary) - { - continue; + for row in 0..num_rows { + let pk = pk_rows.row(row); + let mut interned = false; + let state = states.entry_ref(pk.as_ref()).or_insert_with(|| { + interned = true; + RankPartitionState::default() + }); + if interned { + interned_bytes += pk.as_ref().len(); } - } - - // Gather this partition's rows into their own batch, as - // `PartitionedTopK` does. Registering the whole input batch - // instead would pin it — and charge for it — once per - // partition key present in the batch, so a batch spanning P - // partitions would be counted P times over. - let indices_arr = UInt32Array::from(indices); - let sub_batch = take_record_batch(batch, &indices_arr)?; - - // Indices *into `sub_batch`* of rows from this batch that tied - // at the boundary. Coalesced into a single tie entry at the end - // of the partition's loop. Discarded if the boundary moves up - // mid-loop (those rows were tied to the old boundary, which is - // now strictly worse than the new K-th). - let mut equal_indices: Vec = Vec::new(); - // Lazy-registered: only attached if at least one row reaches - // the heap from this batch in this partition. - let mut heap_entry: Option = None; - - for (sub_idx, &orig_idx) in indices_arr.values().iter().enumerate() { - let row = self.scratch_rows.row(orig_idx as usize); - - // Classify against the current K-th-best (the heap top). - // `heap.max()` returns `None` while the heap is filling, - // so unclassified rows fall through to the heap path. - let classification = state - .heap - .max() - .map(|max_row| row.as_ref().cmp(max_row.row())); - - match classification { + let key = ob_rows.row(row); + let key = key.as_ref(); + // An admission's row in the gathered batch is its position in + // `admitted_rows`, because the gather preserves that order. + let at = StoreRef { + batch_id: pending.id, + row: admitted_rows.len() as u32, + }; + match state.heap.classify(k, key) { + // Strictly worse than the boundary: drop the row. + Some(Ordering::Greater) => continue, Some(Ordering::Equal) => { - equal_indices.push(sub_idx as u32); + ties_grown += state.push_tie(at); + new_slots += 1; } - // Strictly worse than the current boundary: drop the row. - Some(Ordering::Greater) => {} - Some(Ordering::Less) | None => { - // Heap path: heap not yet full, or row strictly - // better than the current boundary. - let entry_ref = heap_entry.get_or_insert_with(|| { - state.heap.register_batch(sub_batch.clone()) - }); - if let Some(EvictedRow { - batch: evicted_batch, - index: evicted_index, - row_bytes: evicted_bytes, - }) = state.heap.add(entry_ref, row, sub_idx) - { - // Compare the new boundary (post-eviction heap - // top) against the evicted row's bytes — both - // already in encoded form, no clones needed. - let boundary_changed = state - .heap - .max() - .expect("heap was full to evict; must still be full") - .row() - != evicted_bytes.as_slice(); - if boundary_changed { - // Boundary moved up — prior ties (across - // all prior batches) and equal_indices - // accumulated earlier in THIS batch were - // tied to the old boundary, now strictly - // worse than the new K-th-best. Discard. - state.ties.clear(); - equal_indices.clear(); - } else { - // Boundary unchanged — the evicted row is - // still tied at the boundary. Gather just - // that row: holding `evicted_batch` would - // keep a whole heap batch alive for one row, - // and one such entry per input batch would - // again make memory grow with the input. - let one = UInt32Array::from(vec![evicted_index as u32]); - state.push_ties(take_record_batch(&evicted_batch, &one)?); + + None => { + let (evicted, grown) = state.heap.add(k, key, at); + debug_assert!( + evicted.is_none(), + "a heap that is not full evicts nothing" + ); + admitted_bytes += grown; + new_slots += 1; + } + Some(Ordering::Less) => { + // The old root's key comes back in `evicted_key`, to + // tell whether the boundary moved. + let (evicted, grown) = + state.heap.replace_worst(key, at, evicted_key); + admitted_bytes += grown; + if state.heap.worst_key() == Some(evicted_key.as_slice()) { + // Boundary unchanged: the evicted row is still + // tied at it. Its slot moves to the tie list, and + // its store use with it — no gather, no release. + ties_grown += state.push_tie(evicted); + new_slots += 1; + } else { + // Boundary moved up: the evicted row and every tie + // (from this batch or earlier ones) now rank + // below K. + store.release(&mut pending, evicted.batch_id); + for tie in &state.ties { + store.release(&mut pending, tie.batch_id); } + released_slots += state.ties.len(); + ties_freed += state.ties.capacity() * size_of::(); + state.ties = Vec::new(); } - replacements += 1; } } + // Every arm but `Greater` admitted the row, to the heap or to + // ties, so each counts, as every admission does in + // `PartitionedTopK`. + pending.uses += 1; + admitted_rows.push(row as u32); + replacements += 1; } - let registered_with_heap = heap_entry.is_some(); - if let Some(e) = heap_entry { - state.heap.insert_batch_entry(e); - state.heap.maybe_compact()?; - } - - // Commit this batch's ties as a single entry. - if !equal_indices.is_empty() { - // No row of this group reached the heap, so `sub_batch` is - // not registered there and reusing it here cannot - // double-charge it. Combined with every row having tied, - // `sub_batch` already *is* exactly the tie rows — the - // gather below would just copy it. - let tie_batch = if !registered_with_heap - && equal_indices.len() == sub_batch.num_rows() - { - sub_batch - } else { - take_record_batch(&sub_batch, &UInt32Array::from(equal_indices))? - }; - state.push_ties(tie_batch); - } - } - - // Return the drained scratch map (capacity retained) for the next - // batch to reuse. - self.partition_groups = groups; - - if replacements > 0 { - self.metrics.row_replacements.add(replacements); - } - self.reservation.try_resize(self.size())?; - Ok(()) - } - - /// Drain all heaps and ties in partition-key order and return the - /// rows as a stream of coalesced [`RecordBatch`]es ordered by - /// `(partition_keys, order_keys)`. Within a partition, heap rows - /// come first (sorted by ob), then tie rows (all sharing the - /// boundary ob). - pub(crate) fn emit(self) -> Result { - let Self { - schema, - metrics, - reservation: _, - expr: _, - row_converter: _, - scratch_rows: _, - partition_exprs: _, - partition_converter: _, - partition_scratch_rows: _, - mut states, - partition_groups: _, - k: _, - batch_size, - } = self; - let _timer = metrics.baseline.elapsed_compute().timer(); - - let mut sorted_pks: Vec> = states.keys().cloned().collect(); - sorted_pks.sort(); - - let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); - - for pk in sorted_pks { - let RankPartitionState { mut heap, ties } = - states.remove(&pk).expect("key from states.keys()"); - if let Some(batch) = heap.emit()? { - coalescer.push_batch(batch)?; - } - for tie in ties { - coalescer.push_batch(tie.batch)?; - } - } - coalescer.finish_buffered_batch()?; - - let mut out: Vec> = Vec::new(); - while let Some(b) = coalescer.next_completed_batch() { - (&b).record_output(&metrics.baseline); - out.push(Ok(b)); - } - - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::iter(out), - ))) - } - - /// Total memory currently held, including all per-partition states. - fn size(&self) -> usize { - // Per partition: the state plus the encoded partition key owned by - // the map. The key bytes are a heap allocation the table slot - // doesn't cover. - let states_contents: usize = self - .states - .iter() - .map(|(pk, state)| pk.capacity() + state.size()) - .sum(); - size_of::() - + self.row_converter.size() - + self.partition_converter.size() - + self.scratch_rows.size() - + self.partition_scratch_rows.size() - + states_contents - + self.states.capacity() - * (size_of::>() + size_of::()) - + self.partition_groups.capacity() - * (size_of::>() + size_of::>()) + core.index_bytes += interned_bytes; + core.heaps_bytes += admitted_bytes; + self.ties_bytes = self.ties_bytes + ties_grown - ties_freed; + core.live_slots = core.live_slots + new_slots - released_slots; + } + + // 3. Gather this batch's admissions into the store, then compact it. + core.finish_batch(batch, pending, replacements)?; + self.core.reservation.try_resize(self.size())?; + Ok(()) + } + + /// Drain every partition in partition-key order and return the rows as a + /// stream of `RecordBatch`es ordered by `(partition_keys, order_keys)`. + /// Within a partition, heap rows come first (sorted by ob), then tie rows + /// (all sharing the boundary ob). + pub(crate) fn emit(self) -> Result { + self.core.emit_with(|mut state| { + let heap = state.heap.drain_sorted().into_iter().map(|slot| slot.at); + heap.chain(state.ties) + }) + } + + /// Total memory currently held by this operator, including every + /// partition's heap and ties and every batch the store still pins. O(1), + /// see [`PartitionedTopKCore::allocated_size`]. + fn size(&self) -> usize { + size_of::() + self.core.allocated_size() + self.ties_bytes } } @@ -4376,9 +4431,6 @@ mod tests { /// batches stream through, where whole-batch retention was /// (partitions x batch) and needed ~P times as much. Only the ratio /// has to hold as the size accounting is legitimately adjusted. - /// - /// Scope: the insert path only. `emit` drops the reservation before - /// it materializes ties, so emit-time growth is out of reach here. #[tokio::test] async fn test_partitioned_topk_rank_runs_under_bounded_memory_pool() -> Result<()> { // P is the whole point: the old code pinned and charged one copy of @@ -4690,14 +4742,14 @@ mod tests { // Batch 0 fills the single partition's heap, so it is referenced and // must be pinned. state.insert_batch(&pk_val_batch(&schema, vec![0, 0], vec![50, 60])?)?; - assert_eq!(state.store.len(), 1, "the batch every slot points at"); - let first_id = *state.store.batches.keys().next().expect("one entry"); + assert_eq!(state.core.store.len(), 1, "the batch every slot points at"); + let first_id = *state.core.store.batches.keys().next().expect("one entry"); // Batch 1 is entirely worse, so nothing is admitted and nothing may be // registered — a store that registered unconditionally would grow here. state.insert_batch(&pk_val_batch(&schema, vec![0, 0], vec![70, 80])?)?; assert_eq!( - state.store.len(), + state.core.store.len(), 1, "a fully-rejected batch must not be registered" ); @@ -4706,12 +4758,12 @@ mod tests { // reference and must be dropped as batch 2 takes its place. state.insert_batch(&pk_val_batch(&schema, vec![0, 0], vec![10, 20])?)?; assert_eq!( - state.store.len(), + state.core.store.len(), 1, "superseded batch 0 must be released, not accumulated" ); assert!( - !state.store.batches.contains_key(&first_id), + !state.core.store.batches.contains_key(&first_id), "the released entry must be the superseded one" ); @@ -4719,7 +4771,7 @@ mod tests { // both the old and the new batch are legitimately referenced. state.insert_batch(&pk_val_batch(&schema, vec![0], vec![5])?)?; assert_eq!( - state.store.len(), + state.core.store.len(), 2, "a batch still holding one live row must stay pinned" ); @@ -4737,7 +4789,7 @@ mod tests { /// `val = |batch - partition|` makes batch `j` the owner of partition `j`'s /// winner, so all `B` entries keep exactly one live row and each stays /// resident holding every row it admitted. Releasing is powerless here — - /// nothing is ever fully superseded — and without `compact_store` the store + /// nothing is ever fully superseded — and without compaction the store /// pins `B(B+1)/2` rows (131,328 at `B = 512`, half the whole input) to /// retain `B`. That is the unbounded case: over-retention grows with `B`. #[tokio::test] @@ -4752,15 +4804,19 @@ mod tests { } // Every partition holds its K = 1 row, so the ratio's denominator is B. - assert_eq!(state.live_slots, B as usize); + assert_eq!(state.core.live_slots, B as usize); let pinned: usize = state + .core .store .batches .values() .map(|e| e.batch.num_rows()) .sum(); - assert_eq!(pinned, state.store.total_rows, "store row count is exact"); - // An invariant, not a measurement: `compact_store` runs at the end of + assert_eq!( + pinned, state.core.store.total_rows, + "store row count is exact" + ); + // An invariant, not a measurement: compaction runs at the end of // every `insert_batch`, so on return the store either never tripped the // guard or was just rewritten down to `live_slots`. Expressed through // the constant so tuning it cannot leave this stale. Measured 977 here @@ -4771,7 +4827,7 @@ mod tests { pinned <= bound, "{pinned} rows pinned to retain {} — above the {bound}-row bound, so \ retention tracks the input rather than the rows held", - state.live_slots + state.core.live_slots ); // Compaction rewrote every slot's coordinates, so the rows it resolves @@ -4789,7 +4845,7 @@ mod tests { /// A store holding a *single* entry must still be compacted when that entry /// is mostly dead rows. /// - /// Regression: `compact_store` used to bail on `store.len() <= 1`, on the + /// Regression: compaction used to bail on `store.len() <= 1`, on the /// reasoning that one entry cannot compact to anything smaller. It can. An /// entry holds every row *admitted* from its batch, and a row admitted and /// then superseded within that same pass stays in the gather with nothing @@ -4813,15 +4869,19 @@ mod tests { (0..ROWS).map(|i| ROWS - i).collect(), )?)?; - assert_eq!(state.live_slots, 1, "K = 1 in a single partition"); - assert_eq!(state.store.len(), 1, "one input batch, one gathered entry"); + assert_eq!(state.core.live_slots, 1, "K = 1 in a single partition"); + assert_eq!( + state.core.store.len(), + 1, + "one input batch, one gathered entry" + ); // The bound the ratio promises, which the old guard exempted this shape // from entirely: 200 rows stayed pinned to retain 1. assert!( - state.store.total_rows <= state.live_slots * STORE_COMPACTION_RATIO, + state.core.store.total_rows <= state.core.live_slots * STORE_COMPACTION_RATIO, "{} rows pinned to retain {}; a single entry escaped compaction", - state.store.total_rows, - state.live_slots + state.core.store.total_rows, + state.core.live_slots ); // Compaction repointed the surviving slot, so emit still resolves it. @@ -4834,95 +4894,103 @@ mod tests { /// /// The running totals (`live_slots`, `store.total_rows`, /// `store.batches_size`, each entry's `uses`) are maintained incrementally - /// precisely so that neither `size()` nor `compact_store`'s guard has to + /// precisely so that neither `size()` nor compaction's guard has to /// walk the partitions. Nothing else re-derives them, so a drift in any one /// is invisible: an over-counted `uses` leaks an entry forever, an /// under-counted one drops a batch rows still point at, and a wrong /// `live_slots` silently disables the pinning bound. fn assert_store_invariants(state: &PartitionedTopK, label: &str) { - // 1. `live_slots` is the number of slots the heaps actually hold. - let counted: usize = state.heaps.values().map(|h| h.inner.len()).sum(); + let slots: Vec = state + .core + .partitions + .values() + .flat_map(PartitionHeap::store_rows) + .collect(); + assert_store_matches_slots( + &state.core.store, + state.core.live_slots, + &slots, + label, + ); + + // 5. What the operator reported to the pool is what it computes now. + assert_eq!( + state.core.reservation.size(), + state.size(), + "{label}: the reservation does not match size()" + ); + } + + /// Sections 1–4 of [`assert_store_invariants`], for any operator on the + /// shared store: `slots` is every retained row's store coordinates, heap + /// or tie. + fn assert_store_matches_slots( + store: &RecordBatchStore, + live_slots: usize, + slots: &[StoreRef], + label: &str, + ) { + // 1. `live_slots` is the number of slots actually held. assert_eq!( - counted, state.live_slots, - "{label}: live_slots disagrees with the heaps" + slots.len(), + live_slots, + "{label}: live_slots disagrees with the slots held" ); - // 2. Every slot points at a live entry, and each entry's `uses` is - // exactly the number of slots pointing into it. + // 2. Every slot points at a live entry and a row inside it, and each + // entry's `uses` is exactly the number of slots pointing into it. let mut refs: HashMap = HashMap::new(); - for heap in state.heaps.values() { - for slot in heap.slots() { - assert!( - state.store.get(slot.batch_id).is_some(), + for slot in slots { + let entry = store.get(slot.batch_id).unwrap_or_else(|| { + panic!( "{label}: slot points at batch {} which the store does not hold", slot.batch_id - ); - *refs.entry(slot.batch_id).or_default() += 1; - } + ) + }); + // Out of range, emit's interleave would read out of bounds (or + // silently pick a different row). + assert!( + (slot.row as usize) < entry.batch.num_rows(), + "{label}: slot row {} is outside entry {} ({} rows)", + slot.row, + slot.batch_id, + entry.batch.num_rows() + ); + *refs.entry(slot.batch_id).or_default() += 1; } assert_eq!( refs.len(), - state.store.len(), + store.len(), "{label}: the store holds entries nothing points at" ); - for (id, entry) in &state.store.batches { + for (id, entry) in &store.batches { + let expected = refs.get(id).copied().unwrap_or(0); assert_eq!( - entry.uses, - refs.get(id).copied().unwrap_or(0), - "{label}: entry {id} has uses={} but {} slots point at it", - entry.uses, - refs.get(id).copied().unwrap_or(0) + entry.uses, expected, + "{label}: entry {id} has uses={} but {expected} slots point at it", + entry.uses ); - // A slot's row must be in range, or emit's interleave would read - // out of bounds (or silently pick a different row). - for heap in state.heaps.values() { - for slot in heap.slots().filter(|s| s.batch_id == *id) { - assert!( - (slot.row as usize) < entry.batch.num_rows(), - "{label}: slot row {} is outside entry {id} ({} rows)", - slot.row, - entry.batch.num_rows() - ); - } - } } // 3. The store's two running totals are exact, not estimates. - let rows: usize = state - .store - .batches - .values() - .map(|e| e.batch.num_rows()) - .sum(); - assert_eq!( - rows, state.store.total_rows, - "{label}: store.total_rows drifted" - ); - let bytes: usize = state - .store + let rows: usize = store.batches.values().map(|e| e.batch.num_rows()).sum(); + assert_eq!(rows, store.total_rows, "{label}: store.total_rows drifted"); + let bytes: usize = store .batches .values() .map(|e| get_record_batch_memory_size(&e.batch)) .sum(); assert_eq!( - bytes, state.store.batches_size, + bytes, store.batches_size, "{label}: store.batches_size drifted" ); - // 4. The bound `compact_store` exists to enforce. Without it, what stays + // 4. The bound compaction exists to enforce. Without it, what stays // pinned tracks the *input* rather than `partitions × K`. assert!( - state.store.total_rows <= state.live_slots * STORE_COMPACTION_RATIO, - "{label}: {} rows pinned to retain {}", - state.store.total_rows, - state.live_slots - ); - - // 5. What the operator reported to the pool is what it computes now. - assert_eq!( - state.reservation.size(), - state.size(), - "{label}: the reservation does not match size()" + store.total_rows <= live_slots * STORE_COMPACTION_RATIO, + "{label}: {} rows pinned to retain {live_slots}", + store.total_rows ); } @@ -5074,43 +5142,47 @@ mod tests { /// `register` entirely when `uses` falls back to 0. If the id were consumed /// anyway the sequence would gap harmlessly, but if `register` ran with /// `uses == 0` the store would keep an entry nothing references — and - /// `compact_store`'s `first_id` would then collide with it. + /// `RecordBatchStore::compact`'s `first_id` would then collide with it. #[tokio::test] async fn test_partitioned_topk_unused_gather_id_is_reused() -> Result<()> { let (schema, mut state) = build_partitioned_topk(1)?; // Batch 0: one partition, one row. Consumes id 0. state.insert_batch(&pk_val_batch(&schema, vec![0], vec![5])?)?; - assert_eq!(state.store.next_batch_id(), 1, "batch 0 consumed its id"); + assert_eq!( + state.core.store.next_batch_id(), + 1, + "batch 0 consumed its id" + ); // Batch 1: every value is worse than the retained 5, so nothing is // admitted at all — `uses` never rises and the id is untouched. No // compaction either: neither total_rows nor live_slots moved, and the // guard already held when the previous call returned. state.insert_batch(&pk_val_batch(&schema, vec![0, 0, 0], vec![7, 8, 9])?)?; - assert_eq!(state.store.next_batch_id(), 1, "no admission, no id"); - assert_eq!(state.store.len(), 1); + assert_eq!(state.core.store.next_batch_id(), 1, "no admission, no id"); + assert_eq!(state.core.store.len(), 1); assert_store_invariants(&state, "after a batch with no admissions"); // Batch 2: rows are admitted but each is superseded by a later row of // the *same* batch, and the first evicts the row from batch 0. The // gather therefore holds 3 rows for 1 live slot, which trips - // `compact_store` — so this consumes two ids, the gather's and the + // compaction — so this consumes two ids, the gather's and the // compacted chunk's. state.insert_batch(&pk_val_batch(&schema, vec![0, 0, 0], vec![4, 3, 2])?)?; assert_store_invariants(&state, "after a self-superseding batch"); - assert_eq!(state.store.len(), 1, "batch 0's entry was released"); + assert_eq!(state.core.store.len(), 1, "batch 0's entry was released"); assert_eq!( - state.store.total_rows, 1, + state.core.store.total_rows, 1, "compaction dropped the dead rows" ); - let after_compaction = state.store.next_batch_id(); + let after_compaction = state.core.store.next_batch_id(); // Batch 3: admitted then superseded within the batch, ending worse than // the retained 2 — so `uses` returns to 0 and the id is handed back. state.insert_batch(&pk_val_batch(&schema, vec![0, 0], vec![3, 6])?)?; assert_eq!( - state.store.next_batch_id(), + state.core.store.next_batch_id(), after_compaction, "every admission was superseded, so the id must be reusable" ); @@ -5180,7 +5252,7 @@ mod tests { )?)?; assert_eq!( - state.live_slots, + state.core.live_slots, PARTITIONS * K, "every row should have been retained" ); @@ -5189,9 +5261,9 @@ mod tests { // least that much, on top of whatever the store and the map cost. let key_bytes = PARTITIONS * K * KEY_WIDTH; assert!( - state.heaps_bytes >= key_bytes, + state.core.heaps_bytes >= key_bytes, "heaps_bytes is {} but the retained keys are at least {key_bytes} bytes", - state.heaps_bytes + state.core.heaps_bytes ); assert!( state.size() >= empty + key_bytes, @@ -5200,42 +5272,16 @@ mod tests { ); // And the interned partition keys are counted separately. assert!( - state.index_bytes >= PARTITIONS, + state.core.index_bytes >= PARTITIONS, "index_bytes is {} for {PARTITIONS} interned keys", - state.index_bytes + state.core.index_bytes ); Ok(()) } - /// Randomized differential test for `PartitionedTopKRank`. - /// - /// The retention rule is subtle — a K-bounded heap plus a boundary-tie - /// list that must be discarded the moment the K-th-best ORDER BY value - /// improves — and it turns on how partitions, ties and batch boundaries - /// interleave. 64 seeds runs in ~10 ms; deleting the tie-clear on a - /// boundary shift is caught by seed 0. - #[tokio::test] - async fn test_partitioned_topk_rank_matches_bruteforce() -> Result<()> { - for seed in 0..64u64 { - let shape = DiffShape::new(seed, 6); - let (schema, mut state) = build_partitioned_topk_rank(shape.k)?; - for (pks, vals) in &shape.batches { - state.insert_batch(&pk_val_batch(&schema, pks.clone(), vals.clone())?)?; - } - - // RANK: 1 + the number of strictly smaller rows. - let expected = shape.expected(|rows, pk, val| { - 1 + rows.iter().filter(|&&(p, v)| p == pk && v < val).count() - }); - - assert_eq!(sorted_pk_val(state.emit()?).await?, expected, "{shape}"); - } - Ok(()) - } - /// Randomized differential test for `PartitionedTopKDenseRank`. /// - /// Same harness as [`test_partitioned_topk_rank_matches_bruteforce`], + /// Same harness as [`test_partitioned_topk_rank_bookkeeping_tracks_recompute`], /// differing only in the ranking formula and a wider value domain: /// DENSE_RANK bounds *distinct* values, so a partition needs more than /// `k` of them before it evicts anything, and eviction is what the @@ -5739,7 +5785,7 @@ mod tests { /// State must accumulate across `insert_batch` calls. A row in /// batch 2 that's strictly better than the existing K-th must /// evict it; an evicted row whose bytes match the new boundary - /// becomes a `TieEntry` pinned to the prior batch. + /// moves to the tie list, still pointing into the prior batch's entry. #[tokio::test] async fn test_partitioned_topk_rank_cross_batch_eviction() -> Result<()> { let (schema, mut state) = build_partitioned_topk_rank(2)?; @@ -5778,8 +5824,8 @@ mod tests { /// `fetch = 1` is a common case (rk = 1 filter) and exercises the /// boundary-defined-immediately path: after the first admission per - /// partition, `heap.max()` is `Some`, so every subsequent row goes - /// through full Equal/Greater/Less classification. + /// partition the heap is full, so every subsequent row goes through + /// full Equal/Greater/Less classification. #[tokio::test] async fn test_partitioned_topk_rank_fetch_one() -> Result<()> { let (schema, mut state) = build_partitioned_topk_rank(1)?; @@ -5942,7 +5988,7 @@ mod tests { /// RANK-specific: heap fills with K rows tied at the same OB value, /// then more rows at that same value arrive. They take the Equal arm - /// (heap is full, `heap.max() == row`) and accumulate as ties, while + /// (heap is full, its root's key equals the row's) and accumulate as ties, while /// strictly-greater rows are dropped. All retained rows have rank 1. #[tokio::test] async fn test_partitioned_topk_rank_boundary_ties_retained() -> Result<()> { @@ -5972,9 +6018,9 @@ mod tests { Ok(()) } - /// Tie rows are emitted through the same coalescer as heap rows, so they - /// must be counted in `output_rows` once, not once as a tie batch and - /// again as part of the coalesced output batch. + /// Tie rows are emitted in the same `EmitState` chunks as heap rows, so + /// they must be counted in `output_rows` once, not once as ties and again + /// as part of the chunk that holds them. #[tokio::test] async fn test_partitioned_topk_rank_output_rows_counts_ties_once() -> Result<()> { let schema = pk_val_schema(false); @@ -6014,10 +6060,10 @@ mod tests { Ok(()) } - /// RANK-specific: heap fills with K rows tied at value V, equal_indices - /// accumulate at V, then a strictly-better row arrives whose admission - /// shifts the boundary strictly below V. The boundary-changed branch - /// must clear both `state.ties` and the in-flight `equal_indices` — + /// RANK-specific: heap fills with K rows tied at value V, ties accumulate + /// at V, then a strictly-better row arrives whose admission shifts the + /// boundary strictly below V. The boundary-changed branch must release + /// every tie, including those admitted earlier in the same batch — /// otherwise the now-rank-> K rows at value V would leak into output. #[tokio::test] async fn test_partitioned_topk_rank_boundary_shifts_clears_ties() -> Result<()> { @@ -6025,11 +6071,11 @@ mod tests { // pk=1 vals: 10, 10, 10, 5, 3 // - first two 10s fill heap (max=10) - // - third 10 → Equal → equal_indices=[2] + // - third 10 → Equal → ties=[row 2] // - 5 < 10 → admit, evict 10 → heap={5,10}, max=10 (unchanged). - // Push evicted to ties: ties=[10@curr_batch[ev_idx]]. + // The evicted 10 moves to ties: ties=[row 2, evicted]. // - 3 < 10 → admit, evict 10 → heap={3,5}, max=5 (CHANGED). - // Clear ties AND equal_indices. + // Release the evicted 10 and every tie. // Sorted RANKs: 3→1, 5→2, 10→3, 10→3, 10→3. WHERE rk ≤ 2 → [3, 5]. let batch = pk_val_batch(&schema, vec![1, 1, 1, 1, 1], vec![10, 10, 10, 5, 3])?; state.insert_batch(&batch)?; @@ -6049,11 +6095,288 @@ mod tests { Ok(()) } + /// `row_replacements` counts every admission, to the heap or to ties, as + /// it does for `PartitionedTopK`. Nothing is evicted below, so the count + /// is exactly the rows retained. + #[tokio::test] + async fn test_partitioned_topk_rank_row_replacements_counts_ties() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10, 10, 10, 20, 10, 5 + // - two 10s fill the heap (2 admissions) + // - 10 → Equal → tie (3); 20 → Greater → dropped; 10 → tie (4) + // - 5 < 10 → admit, evict 10, boundary unchanged → the evicted 10 + // moves to ties (5) + let batch = + pk_val_batch(&schema, vec![1, 1, 1, 1, 1, 1], vec![10, 10, 10, 20, 10, 5])?; + state.insert_batch(&batch)?; + assert_eq!(state.core.metrics.row_replacements.value(), 5); + + let results: Vec<_> = state.emit()?.try_collect().await?; + let emitted: usize = results.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(emitted, 5); + Ok(()) + } + + /// A boundary move frees the partition's tie list rather than clearing + /// it, so a partition that once held many ties does not keep their + /// capacity, and `ties_bytes` gives the bytes back. + #[tokio::test] + async fn test_partitioned_topk_rank_boundary_move_frees_tie_capacity() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10 ×5, 5, 3 + // - two 10s fill the heap, three more 10s are ties + // - 5 < 10 → evict 10, boundary still 10 → the evicted 10 is a tie + // - 3 < 10 → evict 10, boundary moves to 5 → every tie released + let batch = pk_val_batch(&schema, vec![1; 7], vec![10, 10, 10, 10, 10, 5, 3])?; + state.insert_batch(&batch)?; + + let partition = state + .core + .partitions + .values() + .next() + .expect("one partition"); + assert!(partition.ties.is_empty()); + assert_eq!(partition.ties.capacity(), 0, "tie capacity must be freed"); + assert_eq!(state.ties_bytes, 0); + assert_eq!( + rank_recompute(&state), + (state.core.heaps_bytes, 0, state.core.index_bytes) + ); + Ok(()) + } + + /// A `PartitionedTopKRank` over `schema`, partitioned by `pk`, ordered by + /// `val` ascending. + fn build_partitioned_topk_rank_for( + schema: &Arc, + k: usize, + batch_size: usize, + ) -> Result { + let pk_expr: Arc = col("pk", schema.as_ref())?; + let partition_sort_fields = build_sort_fields( + &[PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }], + schema, + )?; + let order_expr = LexOrdering::from([PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: SortOptions::default(), + }]); + PartitionedTopKRank::try_new( + 0, + Arc::clone(schema), + vec![pk_expr], + partition_sort_fields, + order_expr, + k, + batch_size, + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + ) + } + + /// Input in descending order, so under `ORDER BY val ASC` every row after + /// the heap fills beats the boundary and takes the eviction path that + /// swaps key buffers with `evicted_key`. Keys vary in length, so the swap + /// moves buffers of different capacities between the heap and + /// `evicted_key`, and `heaps_bytes` must still match a recompute. + #[tokio::test] + async fn test_partitioned_topk_rank_sorted_input_evicts_every_row() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new("val", DataType::Utf8, false), + ])); + let mut state = build_partitioned_topk_rank_for(&schema, 2, 8)?; + + // Two equal keys fill the heap, then strictly descending keys of + // lengths 1, 6, 1, 6, 1, 9, 1. "y" evicts one "z…" and the boundary + // stays at the other, so the evicted row moves to ties; every later + // eviction moves the boundary and releases what is tied. + let vals = [ + "zzzzzzzzz", + "zzzzzzzzz", + "y", + "xxxxxx", + "w", + "cccccc", + "c", + "bbbbbbbbb", + "a", + ]; + for (i, chunk) in vals.chunks(3).enumerate() { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; chunk.len()])), + Arc::new(StringArray::from(chunk.to_vec())), + ], + )?; + state.insert_batch(&batch)?; + if i == 0 { + let partition = state + .core + .partitions + .values() + .next() + .expect("one partition"); + assert_eq!(partition.ties.len(), 1, "\"y\" must leave a tie"); + } + let (heaps, ties, index) = rank_recompute(&state); + assert_eq!(state.core.heaps_bytes, heaps, "heaps_bytes"); + assert_eq!(state.ties_bytes, ties, "ties_bytes"); + assert_eq!(state.core.index_bytes, index, "index_bytes"); + } + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----------+", + "| pk | val |", + "+----+-----------+", + "| 1 | a |", + "| 1 | bbbbbbbbb |", + "+----+-----------+", + ], + &results + ); + Ok(()) + } + + /// `interleave_referenced` is `interleave_record_batch` restricted to the + /// batches a chunk uses: the same rows come out, and `remap` is left + /// clean for the next chunk. Every batch has its own dictionary, as when + /// each input batch was decoded separately. + #[test] + fn test_interleave_referenced_matches_interleave_over_all() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + )])); + let batches: Vec = (0..4) + .map(|b| { + let mut builder = StringDictionaryBuilder::::new(); + for row in 0..3 { + builder.append_value(format!("b{b}r{row}")); + } + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(builder.finish())], + ) + }) + .collect::>()?; + let all: Vec<&RecordBatch> = batches.iter().collect(); + let as_strings = |batch: &RecordBatch| -> Result> { + let col = arrow::compute::cast(batch.column(0), &DataType::Utf8)?; + Ok(col + .as_string::() + .iter() + .map(|v| v.unwrap().to_string()) + .collect()) + }; + + let mut remap = vec![usize::MAX; batches.len()]; + // Chunks touching one, two and all batches, out of order and revisiting + // a batch the previous chunk used at a different local position. + let chunks: [&[(usize, usize)]; 3] = [ + &[(2, 0), (2, 2)], + &[(3, 1), (2, 1), (3, 0)], + &[(1, 2), (0, 0), (3, 2), (2, 0), (1, 1)], + ]; + for chunk in chunks { + let got = interleave_referenced(&batches, chunk, &mut remap)?; + let want = interleave_record_batch(&all, chunk)?; + assert_eq!(as_strings(&got)?, as_strings(&want)?, "chunk {chunk:?}"); + assert!(remap.iter().all(|&r| r == usize::MAX), "remap left dirty"); + } + Ok(()) + } + + /// The reviewer's emit shape on #25468, for RANK: input clustered by + /// partition key, every retained row a tie so nothing is compacted, and a + /// Dictionary payload with one dictionary per input batch. The store holds + /// four batches while each 8-row chunk uses two, so every chunk exercises + /// the remap from store positions to chunk-local ones. + #[tokio::test] + async fn test_partitioned_topk_rank_emit_dictionary_chunks_use_few_batches() + -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new("val", DataType::Int32, false), + Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + ), + ])); + let mut state = build_partitioned_topk_rank_for(&schema, 2, 8)?; + + // Four batches, batch b holding partitions 2b and 2b+1, six rows each, + // every row tied at val 0: all 48 rows are retained. + for b in 0..4 { + let mut pks = Vec::new(); + let mut builder = StringDictionaryBuilder::::new(); + for pk in [2 * b, 2 * b + 1] { + for row in 0..6 { + pks.push(pk); + builder.append_value(format!("p{pk}r{row}")); + } + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(pks)), + Arc::new(Int32Array::from(vec![0; 12])), + Arc::new(builder.finish()), + ], + )?; + state.insert_batch(&batch)?; + } + assert_eq!( + state.core.store.len(), + 4, + "every input batch stays in the store" + ); + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_eq!( + results + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![8; 6] + ); + // Tied rows emit in no particular order within a partition, so compare + // per-partition sorted. + let mut emitted: Vec<(i32, String)> = Vec::new(); + for batch in &results { + let pks = batch.column(0).as_primitive::(); + let d = arrow::compute::cast(batch.column(2), &DataType::Utf8)?; + let d = d.as_string::(); + for i in 0..batch.num_rows() { + emitted.push((pks.value(i), d.value(i).to_string())); + } + } + let pks_in_order: Vec = emitted.iter().map(|(pk, _)| *pk).collect(); + assert!(pks_in_order.is_sorted(), "partitions emit in key order"); + emitted.sort(); + let expected: Vec<(i32, String)> = (0..8) + .flat_map(|pk| (0..6).map(move |row| (pk, format!("p{pk}r{row}")))) + .collect(); + assert_eq!(emitted, expected); + Ok(()) + } + /// RANK-specific: heap has multiple rows at boundary value V, then a - /// strictly-better row arrives. The heap evicts one V (popping - /// `prev_min`), but `heap.max()` is still V — boundary unchanged. - /// The evicted V row must be pushed as a `TieEntry`; without that - /// branch a `rk <= K` query would silently lose a tied row. + /// strictly-better row arrives. The heap evicts one V, but its root + /// is still V — boundary unchanged. The evicted V row must move to + /// the tie list; without that branch a `rk <= K` query would silently + /// lose a tied row. #[tokio::test] async fn test_partitioned_topk_rank_eviction_at_unchanged_boundary() -> Result<()> { let (schema, mut state) = build_partitioned_topk_rank(2)?; @@ -6082,6 +6405,183 @@ mod tests { Ok(()) } + /// The byte totals `PartitionedTopKRank::size` reads, recomputed from the + /// partitions: `(heaps_bytes, ties_bytes, index_bytes)`. `heaps_bytes` + /// includes `evicted_key`, which trades buffers with heap keys. + fn rank_recompute(state: &PartitionedTopKRank) -> (usize, usize, usize) { + let (mut heaps, mut ties, mut index) = (state.evicted_key.capacity(), 0, 0); + for (pk, partition) in &state.core.partitions { + index += pk.len(); + heaps += partition.heap.inner.capacity() * size_of::() + + partition + .heap + .slots() + .map(|slot| slot.key.capacity()) + .sum::(); + ties += partition.ties.capacity() * size_of::(); + } + (heaps, ties, index) + } + + /// Every retained row's store coordinates, heap and ties. + fn rank_slots(state: &PartitionedTopKRank) -> Vec { + state + .core + .partitions + .values() + .flat_map(RankPartitionState::store_rows) + .collect() + } + + /// Randomized differential test for `PartitionedTopKRank`, checking the + /// shared-store bookkeeping after every batch as well as the output. + /// + /// The retention rule is subtle — a K-bounded heap plus a boundary-tie + /// list that must be discarded the moment the K-th-best ORDER BY value + /// improves — and it turns on how partitions, ties and batch boundaries + /// interleave. After every batch, [`assert_store_matches_slots`] checks + /// that each retained row, heap or tie, holds exactly one use of the + /// entry it points into, and the running totals `size()` reads are + /// checked against a recompute. The use check is what catches a slot that + /// moves to the tie list without its use, a tie released against the + /// wrong entry, or an in-flight release that forgets the local count — + /// each leaves the store either panicking later or pinning a batch + /// forever, and none is visible in the emitted rows. + #[tokio::test] + async fn test_partitioned_topk_rank_bookkeeping_tracks_recompute() -> Result<()> { + let (mut saw_ties, mut saw_compaction) = (false, false); + for seed in 0..256u64 { + let shape = DiffShape::new(seed, 6); + let (schema, mut state) = build_partitioned_topk_rank(shape.k)?; + for (pks, vals) in &shape.batches { + let next_id = state.core.store.next_batch_id(); + state.insert_batch(&pk_val_batch(&schema, pks.clone(), vals.clone())?)?; + // A batch takes at most one id; any more went to compaction. + saw_compaction |= state.core.store.next_batch_id() > next_id + 1; + saw_ties |= state.core.partitions.values().any(|p| !p.ties.is_empty()); + + assert_store_matches_slots( + &state.core.store, + state.core.live_slots, + &rank_slots(&state), + &shape.to_string(), + ); + let (heaps, ties, index) = rank_recompute(&state); + assert_eq!(state.core.heaps_bytes, heaps, "heaps_bytes: {shape}"); + assert_eq!(state.ties_bytes, ties, "ties_bytes: {shape}"); + assert_eq!(state.core.index_bytes, index, "index_bytes: {shape}"); + assert_eq!( + state.core.reservation.size(), + state.size(), + "reservation vs size(): {shape}" + ); + } + + let expected = shape.expected(|rows, pk, val| { + 1 + rows.iter().filter(|&&(p, v)| p == pk && v < val).count() + }); + assert_eq!(sorted_pk_val(state.emit()?).await?, expected, "{shape}"); + } + // Guards the guard: shapes that never tie or never compact would leave + // the paths this is about unchecked. + assert!(saw_ties, "no shape retained a tie"); + assert!(saw_compaction, "no shape compacted the store"); + Ok(()) + } + + /// Ties are store slots like heap rows, so they are subject to the same + /// pinning as [`test_partitioned_topk_store_compacts_a_single_oversized_entry`] + /// guards against for `ROW_NUMBER`: a batch's entry stays resident while + /// any one of its rows is retained. + /// + /// Batch `j` gives every partition `p` a tie pair at `|j - p|`, so + /// partition `p` keeps improving until batch `p` and is then fixed at a + /// tie pair of zeros. Each batch's entry therefore keeps two live rows + /// forever while its other admissions are superseded: without + /// compaction the store would hold ~B² rows to retain 2B. + #[tokio::test] + async fn test_partitioned_topk_rank_store_bounded_when_ties_spread_thinly() + -> Result<()> { + const B: i32 = 64; + let (schema, mut state) = build_partitioned_topk_rank(1)?; + for j in 0..B { + let pks: Vec = (0..B).flat_map(|p| [p, p]).collect(); + let vals: Vec = (0..B) + .flat_map(|p| { + let v = (j - p).abs(); + [v, v] + }) + .collect(); + state.insert_batch(&pk_val_batch(&schema, pks, vals)?)?; + assert!( + state.core.store.total_rows + <= STORE_COMPACTION_RATIO * state.core.live_slots, + "batch {j}: store holds {} rows for {} retained", + state.core.store.total_rows, + state.core.live_slots + ); + } + let expected: Vec<(i32, i32)> = (0..B).flat_map(|p| [(p, 0), (p, 0)]).collect(); + assert_eq!(pk_val_rows(state.emit()?).await?, expected); + Ok(()) + } + + /// A boundary move releases every tie of the partition, which can come + /// from earlier batches and from the batch in flight, in any + /// interleaving: ties released by run must hand each use back to the + /// right place, the store for an earlier batch and the local count for + /// the in-flight one. + #[tokio::test] + async fn test_partitioned_topk_rank_boundary_move_releases_ties_across_batches() + -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // Heap {10, 10} and ties [10, 10]: four slots into one store entry. + state.insert_batch(&pk_val_batch(&schema, vec![1; 4], vec![10; 4])?)?; + assert_eq!((state.core.store.len(), state.core.live_slots), (1, 4)); + + // 10 ties, from this batch. 3 evicts a 10 but the root is still 10, so + // the evicted row moves to the ties. 1 evicts the last 10 and the root + // becomes 3, releasing the evicted row and every tie: runs of + // [first, first, this, first] batch. + state.insert_batch(&pk_val_batch(&schema, vec![1; 3], vec![10, 3, 1])?)?; + + // The first batch's entry is freed. This batch's gather holds all + // three of its admissions, two of them still retained. + assert_eq!((state.core.store.len(), state.core.live_slots), (1, 2)); + let entry = state.core.store.batches.values().next().expect("one entry"); + assert_eq!((entry.batch.num_rows(), entry.uses), (3, 2)); + + assert_eq!(pk_val_rows(state.emit()?).await?, vec![(1, 1), (1, 3)]); + Ok(()) + } + + /// Ties from every partition a batch touches share that batch's single + /// store entry, so they are charged once. + /// + /// Regression for #23326: ties used to be materialized — and charged — + /// once per partition, which with every row tied at `K = 1` across 4096 + /// partitions multiplied the reservation by the partition count. + #[tokio::test] + async fn test_partitioned_topk_rank_ties_share_one_store_entry() -> Result<()> { + const P: i32 = 4096; + let (schema, mut state) = build_partitioned_topk_rank(1)?; + let batch = + pk_val_batch(&schema, (0..4 * P).map(|i| i % P).collect(), vec![0; 16384])?; + state.insert_batch(&batch)?; + + // Every row ties: one heap row plus three ties per partition. + assert_eq!(state.core.live_slots, 4 * P as usize); + assert_eq!(state.core.store.len(), 1); + assert!( + state.core.store.batches_size <= 2 * get_record_batch_memory_size(&batch), + "store charges {} bytes for one {}-byte batch", + state.core.store.batches_size, + get_record_batch_memory_size(&batch) + ); + Ok(()) + } + // ==================================================================== // PartitionedTopKDenseRank operator tests //