SubhamSinghal commented on code in PR #22885:
URL: https://github.com/apache/datafusion/pull/22885#discussion_r3524503672
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1414,6 +1459,358 @@ impl PartitionedTopK {
}
}
+/// A run of rows from a single source [`RecordBatch`] that tied at the
+/// boundary when inserted. Stored as `(batch, indices)` and materialized
+/// at emit time via [`take_record_batch`].
+#[derive(Debug)]
+struct TieEntry {
+ batch: RecordBatch,
+ /// Indices into `batch` of the rows tied at the (then-current)
+ /// boundary. Always non-empty by construction.
+ row_indices: Vec<u32>,
+ /// `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<TieEntry>` 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.
+struct RankPartitionState {
+ heap: TopKHeap,
+ ties: Vec<TieEntry>,
+}
+
+impl RankPartitionState {
+ fn size(&self) -> usize {
+ let ties_buffer = self.ties.capacity() * size_of::<TieEntry>();
+ let ties_contents: usize = self
+ .ties
+ .iter()
+ .map(|t| t.row_indices.capacity() * size_of::<u32>() +
t.batch_bytes)
+ .sum();
+ self.heap.size() + ties_buffer + ties_contents
+ }
+}
+
+/// Sibling to [`PartitionedTopK`] implementing `RANK()` semantics.
+///
+/// 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.
+///
+/// # 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:
+///
+/// - 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`)
+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<Arc<dyn PhysicalExpr>>,
+ /// 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.
+ states: HashMap<OwnedRow, RankPartitionState>,
+ k: usize,
+ batch_size: usize,
+}
+
+impl PartitionedTopKRank {
+ #[expect(clippy::too_many_arguments)]
+ pub(crate) fn try_new(
+ partition_id: usize,
+ schema: SchemaRef,
+ partition_exprs: Vec<Arc<dyn PhysicalExpr>>,
+ partition_sort_fields: Vec<SortField>,
+ order_expr: LexOrdering,
+ k: usize,
+ batch_size: usize,
+ runtime: &Arc<RuntimeEnv>,
+ metrics: &ExecutionPlanMetricsSet,
+ ) -> Result<Self> {
+ 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(),
+ k,
+ batch_size,
+ })
+ }
+
+ /// 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.
+ pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> {
+ let baseline = self.metrics.baseline.clone();
+ let _timer = baseline.elapsed_compute().timer();
+
+ let num_rows = batch.num_rows();
+ if num_rows == 0 {
+ return Ok(());
+ }
+
+ // Captured once so the per-tie push from this batch can reuse
+ // it (computing `get_record_batch_memory_size` is O(cols ×
+ // buffer walk) and we'd otherwise pay it per push and again
+ // per `try_resize` call).
+ let input_batch_bytes = get_record_batch_memory_size(batch);
+
+ // 1. Evaluate + encode partition columns into the reusable
+ // scratch (cleared then appended).
+ let pk_arrays: Vec<ArrayRef> = self
+ .partition_exprs
+ .iter()
+ .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows)))
+ .collect::<Result<_>>()?;
+ self.partition_scratch_rows.clear();
+ self.partition_converter
+ .append(&mut self.partition_scratch_rows, &pk_arrays)?;
+ let pk_rows = &self.partition_scratch_rows;
+
+ // 2. Demultiplex row indices by partition key (per-batch).
+ let mut groups: HashMap<OwnedRow, Vec<u32>> = HashMap::new();
+ for i in 0..num_rows {
+ groups
+ .entry(pk_rows.row(i).owned())
+ .or_default()
+ .push(i as u32);
+ }
+
+ // 3. Evaluate ORDER BY columns on the full batch and encode ONCE.
+ let ob_arrays: Vec<ArrayRef> = self
+ .expr
+ .iter()
+ .map(|e| e.expr.evaluate(batch).and_then(|v|
v.into_array(num_rows)))
+ .collect::<Result<_>>()?;
+ 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;
+ let mut replacements: usize = 0;
+
+ for (pk, indices) in groups {
+ let state = self.states.entry(pk).or_insert_with(||
RankPartitionState {
+ heap: TopKHeap::new(k),
+ ties: Vec::new(),
+ });
+
+ // Equal indices for THIS batch only. 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 entry: Option<RecordBatchEntry> = None;
+
+ for &orig_idx in &indices {
+ 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 {
+ Some(Ordering::Equal) => {
+ equal_indices.push(orig_idx);
+ continue;
+ }
+ Some(Ordering::Greater) => continue,
+ Some(Ordering::Less) | None => {
+ // Heap path: heap not yet full, or row strictly
+ // better than the current boundary.
+ let entry_ref = entry.get_or_insert_with(|| {
+ state.heap.register_batch(batch.clone())
+ });
+ if let Some(EvictedRow {
+ batch: evicted_batch,
+ index: evicted_index,
+ row_bytes: evicted_bytes,
+ }) = state.heap.add(entry_ref, row, orig_idx as usize)
+ {
+ // 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 — evicted row is tied
+ // at the (unchanged) boundary; push as a
+ // single-row entry.
+ let batch_bytes =
+
get_record_batch_memory_size(&evicted_batch);
+ state.ties.push(TieEntry {
+ batch: evicted_batch,
+ row_indices: vec![evicted_index as u32],
+ batch_bytes,
+ });
+ }
+ }
+ replacements += 1;
+ }
+ }
+ }
+
+ if let Some(e) = 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() {
+ state.ties.push(TieEntry {
Review Comment:
created issue for this: https://github.com/apache/datafusion/issues/23326
--
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]
