SubhamSinghal commented on code in PR #23869:
URL: https://github.com/apache/datafusion/pull/23869#discussion_r3683430470
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1811,6 +1811,435 @@ 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`.
+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 {
Review Comment:
fixed in afc931c10aa0d15445bb8b49f4779e03a7d79ac0 using `BinaryHeap::new()`
--
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]
