kumarUjjawal commented on code in PR #23869:
URL: https://github.com/apache/datafusion/pull/23869#discussion_r3689341455
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1811,6 +1811,433 @@ impl PartitionedTopKRank {
}
}
+/// A run of rows from a single source [`RecordBatch`] sharing one
+/// distinct ORDER BY value. Materialized at emit time via
+/// [`take_record_batch`].
+///
+/// The batch is referenced by `batch_id` rather than held directly, so the
+/// operator-scoped [`RecordBatchStore`] can charge each distinct source
+/// batch once however many entries reference it. Holding a batch per entry
+/// and charging its bytes per entry would inflate the reservation by a
+/// factor of (partitions × K), since a single batch contributes an entry to
+/// every partition and ob group it touches.
+#[derive(Debug)]
+struct GroupEntry {
+ /// Indices into the batch identified by `batch_id`. Always non-empty
+ /// by construction.
+ row_indices: Vec<u32>,
+ /// Key into `PartitionedTopKDenseRank::store`.
+ batch_id: u32,
+}
+
+/// Per-partition state for `DENSE_RANK()` semantics.
+///
+/// A `HashMap<Vec<u8>, Vec<GroupEntry>>` keyed by the row-encoded ORDER
+/// BY bytes, capped at `k` distinct keys. Each key's `Vec<GroupEntry>`
+/// holds every row seen at that ob value, one entry per contributing
+/// source `RecordBatch`.
+#[derive(Default)]
+struct DenseRankPartitionState {
+ groups: HashMap<Vec<u8>, Vec<GroupEntry>>,
+ /// The same keys as `groups`, in a max-heap: the admission boundary
+ /// (the largest tracked ob value) is an O(1) `peek()` check, and
+ /// admission / removal are O(log K).
+ ///
+ /// INVARIANT: `keys` and `groups.keys()` hold the same set. Every
+ /// insertion into / removal from `groups` must mirror into `keys`.
+ keys: BinaryHeap<Vec<u8>>,
+}
+
+impl DenseRankPartitionState {
+ fn size(&self) -> usize {
+ let table_overhead = self.groups.capacity()
+ * (size_of::<Vec<u8>>() + size_of::<Vec<GroupEntry>>());
+ let contents: usize = self
+ .groups
+ .iter()
+ .map(|(key, entries)| {
+ key.capacity()
+ + entries.capacity() * size_of::<GroupEntry>()
+ + entries
+ .iter()
+ .map(|e| e.row_indices.capacity() * size_of::<u32>())
+ .sum::<usize>()
+ })
+ .sum();
+ // `keys` duplicates every key's bytes; charge for them plus the
+ // heap's backing Vec (one `Vec<u8>` slot per reserved element).
+ let keys_overhead: usize = self.keys.capacity() * size_of::<Vec<u8>>()
+ + self.keys.iter().map(|k| k.capacity()).sum::<usize>();
+ table_overhead + contents + keys_overhead
+ }
+}
+
+/// Sibling to [`PartitionedTopK`] / [`PartitionedTopKRank`] implementing
+/// `DENSE_RANK()` semantics.
+///
+/// Per partition, retains every row whose ORDER BY value is among the K
+/// distinct-smallest ob values seen for that partition. The total row
+/// count kept per partition is unbounded in `rows_per_distinct_value`
+/// (unlike `RANK`, which is bounded above by K + boundary ties).
+///
+/// Like [`PartitionedTopK`], the [`RowConverter`], [`MemoryReservation`],
+/// scratch [`Rows`] buffer, and [`TopKMetrics`] are shared across all
+/// partitions for this operator instance. So is the
+/// [`RecordBatchStore`]: retained rows are held as `(batch_id, indices)`
+/// so each source batch is charged once for the whole operator, however
+/// many partitions and ob groups reference it.
+///
+/// # Algorithm (per batch)
+///
+/// Evaluate + encode partition-by and order-by columns once, then group
+/// the batch's row indices by partition key. For each partition, bucket
+/// that partition's rows by distinct ob value (a within-call
+/// accumulation), then merge each bucket into the partition state. Every
+/// bucket is built from the current batch's rows, so each `GroupEntry` is
+/// pinned to the batch its `row_indices` point into.
+///
+/// For each partition, for each distinct `ob_key` run in this batch:
+/// - `ob_key` already in `state.groups` → push this batch's run as a
+/// new `GroupEntry` (one entry per contributing batch).
+/// - `ob_key` new, `state.groups.len() < k` → insert the run as a new
+/// group.
+/// - `ob_key` new, `state.groups.len() == k` → the largest tracked ob
+/// value is the admission boundary, read from the `state.keys` max-heap
+/// in O(1):
+/// - `ob_key < max` → remove the max key (evict the entire max-key
+/// group — up to many rows) and insert the run. The evicted group's
+/// row count is added to the `row_replacements` metric.
+/// - `ob_key >= max` → drop the whole run; no map mutation.
+pub(crate) struct PartitionedTopKDenseRank {
+ 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).
+ partition_scratch_rows: Rows,
+ /// One state per distinct partition key seen so far. Keyed by the
+ /// row-encoded PARTITION BY bytes (byte-comparable encoding, so the
+ /// `Vec<u8>` 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<Vec<u8>, DenseRankPartitionState>,
+ /// 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<u8>, Vec<u32>>,
+ /// Scratch map reused across partitions within a batch to bucket a
+ /// partition's rows by distinct ORDER BY value. Drained (not
+ /// reallocated) per partition so its backing table is allocated once,
+ /// not once per distinct partition key.
+ ob_runs: HashMap<Vec<u8>, Vec<u32>>,
+ /// Source batches referenced by the retained `GroupEntry`s, held once
+ /// for the whole operator with a use count per batch. This is what
+ /// keeps the reservation proportional to the batches actually pinned
+ /// rather than to the number of entries pointing at them.
+ store: RecordBatchStore,
+ k: usize,
+ batch_size: usize,
+}
+
+impl PartitionedTopKDenseRank {
+ #[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, "PartitionedTopKDenseRank requires k > 0");
+ let reservation =
+
MemoryConsumer::new(format!("PartitionedTopKDenseRank[{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(),
+ ob_runs: HashMap::new(),
+ store: RecordBatchStore::new(),
+ k,
+ batch_size,
+ })
+ }
+
+ /// Encode PARTITION BY and ORDER BY columns once, demultiplex the
+ /// batch's rows by partition key, then per partition bucket the rows
+ /// by distinct ob value and merge each bucket into the partition
+ /// state as one [`GroupEntry`].
+ 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(());
+ }
+
+ // Register this batch with the store up front. `uses` is bumped
+ // once per `GroupEntry` created below but lives on this local
+ // entry, so the store sees a single insert per batch rather than
+ // one per group — and `insert` drops batches that retained
+ // nothing without ever charging for them.
+ let mut batch_entry = self.store.register(batch.clone());
+ let batch_id = batch_entry.id;
+
+ // 1. Encode partition columns.
+ 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)?;
+
+ // 2. Group this batch's row indices by partition key.
+ // `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<u8>` 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 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)?;
+
+ let k = self.k;
+ let mut replacements: usize = 0;
+
+ // 4. Per-partition: bucket this batch's rows by distinct ob value
+ // (within-call accumulation), then merge each bucket into the
+ // partition state as a single `GroupEntry`.
+ for (pk, indices) in groups.drain() {
+ let state = self
+ .states
+ .entry(pk)
+ .or_insert_with(DenseRankPartitionState::default);
+
+ // Bucket by ob key. `ob_runs` is a reused scratch map (taken
+ // out and drained below) so its backing table is allocated
+ // once, not once per distinct partition key. `entry_ref` owns
+ // the key only on Vacant, so repeated rows of the same ob
+ // value don't re-allocate.
+ let mut runs = std::mem::take(&mut self.ob_runs);
+ runs.clear();
+ for &orig_idx in &indices {
+ let ob_row = self.scratch_rows.row(orig_idx as usize);
+ runs.entry_ref(ob_row.as_ref()).or_default().push(orig_idx);
+ }
+
+ for (ob_key, run_indices) in runs.drain() {
+ // Case A: ob already tracked — push this batch's run as a
+ // new `GroupEntry` (one entry per contributing batch).
+ if let Some(entries) = state.groups.get_mut(&ob_key) {
+ batch_entry.uses += 1;
+ entries.push(GroupEntry {
+ row_indices: run_indices,
+ batch_id,
+ });
+ continue;
+ }
+
+ // Case B: new ob, room available.
+ if state.groups.len() < k {
+ batch_entry.uses += 1;
+ state.keys.push(ob_key.clone());
+ state.groups.insert(
+ ob_key,
+ vec![GroupEntry {
+ row_indices: run_indices,
+ batch_id,
+ }],
+ );
+ continue;
+ }
+
+ // Case C: new ob, at K distinct keys. The largest tracked
+ // ob value is the admission boundary.
+ let max_key = state.keys.peek().expect("state.groups has k >=
1 keys");
+ if ob_key.as_slice() < max_key.as_slice() {
+ // Evict the entire max-key group, from both the map
+ // and its ordered mirror.
+ let evicted_key = state.keys.pop().expect("max key
present");
+ let evicted = state
+ .groups
+ .remove(&evicted_key)
+ .expect("keys mirrors groups");
+ for e in &evicted {
+ replacements += e.row_indices.len();
+ if e.batch_id == batch_id {
+ // Admitted earlier in this same call, so the
+ // store has not seen `batch_entry` yet — drop
+ // the pending use rather than calling `unuse`,
+ // which panics on an unregistered id.
+ batch_entry.uses -= 1;
+ } else {
+ self.store.unuse(e.batch_id);
+ }
+ }
+ batch_entry.uses += 1;
+ state.keys.push(ob_key.clone());
+ state.groups.insert(
+ ob_key,
+ vec![GroupEntry {
+ row_indices: run_indices,
+ batch_id,
+ }],
+ );
+ }
+ // else: ob >= max — drop the whole run.
+ }
+
+ // Return the drained scratch map (capacity retained) for the
+ // next partition to reuse.
+ self.ob_runs = runs;
+ }
+
+ // Return the drained scratch map (capacity retained) for the next
+ // batch to reuse.
+ self.partition_groups = groups;
+
+ // Charges `batch` once if any group retained rows from it.
+ self.store.insert(batch_entry);
+
+ if replacements > 0 {
+ self.metrics.row_replacements.add(replacements);
+ }
+ self.reservation.try_resize(self.size())?;
+ Ok(())
+ }
+
+ /// Drain all per-partition maps in partition-key order and return
+ /// the rows as a stream of coalesced [`RecordBatch`]es ordered by
+ /// `(partition_keys, order_keys)`. Within a partition the distinct
+ /// ob keys are sorted (byte-comparable encoding == sort order) so
+ /// emitted rows are in ob-sorted order.
+ pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> {
+ let Self {
+ schema,
+ metrics,
+ reservation: _,
+ expr: _,
+ row_converter: _,
+ scratch_rows: _,
+ partition_exprs: _,
+ partition_converter: _,
+ partition_scratch_rows: _,
+ mut states,
+ partition_groups: _,
+ ob_runs: _,
+ store,
+ k: _,
+ batch_size,
+ } = self;
+ let _timer = metrics.baseline.elapsed_compute().timer();
+
+ let mut sorted_pks: Vec<Vec<u8>> = states.keys().cloned().collect();
+ sorted_pks.sort();
+
+ let mut coalescer = BatchCoalescer::new(Arc::clone(&schema),
batch_size);
+
+ for pk in sorted_pks {
+ let DenseRankPartitionState { groups, keys: _ } =
+ states.remove(&pk).expect("key from states.keys()");
+ // Sort the <= K distinct ob keys so rows emit ascending
+ // (byte-comparable encoding == sort order).
+ let mut sorted_obs: Vec<(Vec<u8>, Vec<GroupEntry>)> =
+ groups.into_iter().collect();
+ sorted_obs.sort_by(|a, b| a.0.cmp(&b.0));
+ for (_ob, entries) in sorted_obs {
+ for entry in entries {
+ let batch = &store
+ .get(entry.batch_id)
+ .expect("retained batch_id present in store")
+ .batch;
+ let indices = UInt32Array::from(entry.row_indices);
+ let sub = take_record_batch(batch, &indices)?;
+ (&sub).record_output(&metrics.baseline);
+ coalescer.push_batch(sub)?;
+ }
+ }
+ }
+ coalescer.finish_buffered_batch()?;
+
+ let mut out: Vec<Result<RecordBatch>> = Vec::new();
+ while let Some(b) = coalescer.next_completed_batch() {
+ 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 {
+ size_of::<Self>()
+ + self.row_converter.size()
+ + self.partition_converter.size()
+ + self.scratch_rows.size()
+ + self.partition_scratch_rows.size()
+ + self.states.values().map(|s| s.size()).sum::<usize>()
Review Comment:
size() still omits the encoded key buffers owned by self.states and the
backing tables retained by partition_groups and ob_runs after draining. With
many or large partition keys, these long-lived allocations can be substantial
but remain outside the memory reservation. Could we include these buffers and
capacities in size(), or release the scratch capacities?
--
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]
