SubhamSinghal commented on code in PR #25968:
URL: https://github.com/apache/datafusion/pull/25968#discussion_r4182845485
##########
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();
+ evicted_key.extend_from_slice(
+ state.heap.worst_key().expect("classified, so
full"),
+ );
+ let (evicted, grown) = state.heap.add(k, key, at);
+ let evicted = evicted.expect("a full heap evicts");
+ 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();
+ state.ties.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]
