SubhamSinghal commented on code in PR #25968:
URL: https://github.com/apache/datafusion/pull/25968#discussion_r4182849541
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1313,6 +1314,145 @@ impl RecordBatchStore {
+ self.batches.capacity() * (size_of::<u32>() +
size_of::<RecordBatchEntry>())
+ 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<P: StoreSlots>(
+ &mut self,
+ partitions: &mut HashMap<Vec<u8>, P>,
+ live_slots: usize,
+ batch_size: usize,
+ ) -> Result<()> {
+ if self.total_rows <= live_slots * STORE_COMPACTION_RATIO {
+ return Ok(());
+ }
+
+ // Scoped so these clones, which keep the old batches alive while the
+ // compacted ones are built, drop before the store is rewritten.
+ let first_id = self.next_batch_id();
+ let (moved, chunks) = {
+ let (old, array_pos) = self.positional();
Review Comment:
Addressed in f9d7919e1fed26f901ab6b5a2730e0057b85d56d
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2179,217 +2233,192 @@ impl PartitionedTopKRank {
self.row_converter
.append(&mut self.scratch_rows, &ob_arrays)?;
- // 4. Per-partition: classify each row and dispatch.
+ // 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 k = self.k;
let mut replacements: usize = 0;
+ self.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 = self.store.pending();
+ {
+ let pk_rows = &self.partition_scratch_rows;
+ let ob_rows = &self.scratch_rows;
+ let states = &mut self.states;
+ let admitted_rows = &mut self.admitted_rows;
+ let store = &mut self.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 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<u32> = 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<RecordBatchEntry> = 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;
+ replacements += 1;
+ }
+ Some(Ordering::Less) => {
+ // Replacing the root overwrites its key in place, so
+ // keep a copy to tell whether the boundary moved.
+ evicted_key.clear();
Review Comment:
Addressed in f9d7919e1fed26f901ab6b5a2730e0057b85d56d
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
