SubhamSinghal commented on code in PR #23869:
URL: https://github.com/apache/datafusion/pull/23869#discussion_r3667302959
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1811,6 +1811,376 @@ impl PartitionedTopKRank {
}
}
+/// Per-partition state for `DENSE_RANK()` semantics.
+///
+/// A `HashMap<Vec<u8>, Vec<TieEntry>>` keyed by the row-encoded ORDER BY
+/// bytes, capped at `k` distinct keys. Each key's `Vec<TieEntry>` holds
+/// every row seen at that ob value, one entry per contributing source
+/// `RecordBatch`. `TieEntry` is reused verbatim from [`RankPartitionState`].
+struct DenseRankPartitionState {
+ groups: HashMap<Vec<u8>, Vec<TieEntry>>,
+}
+
+impl DenseRankPartitionState {
+ fn size(&self) -> usize {
+ let table_overhead =
+ self.groups.capacity() * (size_of::<Vec<u8>>() +
size_of::<Vec<TieEntry>>());
+ let contents: usize = self
+ .groups
+ .iter()
+ .map(|(key, entries)| {
+ key.capacity()
+ + entries.capacity() * size_of::<TieEntry>()
+ + entries
+ .iter()
+ .map(|e| {
+ e.row_indices.capacity() * size_of::<u32>() +
e.batch_bytes
+ })
+ .sum::<usize>()
+ })
+ .sum();
+ table_overhead + contents
+ }
+}
+
+/// 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.
+///
+/// # 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 `TieEntry` 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 `TieEntry` (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` → find the current max via
+/// an O(K) scan of `state.groups.keys()`:
+/// - `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>>,
+ 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(),
+ 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 [`TieEntry`].
+ 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 every `TieEntry` push from this batch can
+ // reuse it (avoids `get_record_batch_memory_size` per push).
+ let input_batch_bytes = get_record_batch_memory_size(batch);
+
+ // 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 `TieEntry`.
+ for (pk, indices) in groups.drain() {
+ let state =
+ self.states
+ .entry(pk)
+ .or_insert_with(|| DenseRankPartitionState {
+ groups: HashMap::new(),
+ });
+
+ // 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 `TieEntry` (one entry per contributing batch, exactly
+ // like RANK pushing one tie entry per batch).
+ if let Some(entries) = state.groups.get_mut(&ob_key) {
+ entries.push(TieEntry {
+ batch: batch.clone(),
+ row_indices: run_indices,
+ batch_bytes: input_batch_bytes,
Review Comment:
This is the DENSE_RANK instance of #23326 (same TieEntry/size()
shared-buffer over-count as RANK). Keeping storage consistent so one fix in
#23326 covers both operators.
--
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]
