kumarUjjawal commented on code in PR #27136:
URL: https://github.com/apache/datafusion/pull/27136#discussion_r4500485764
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -8062,55 +7891,48 @@ mod tests {
Arc::new(Int32Array::from(vals)),
],
)?)?;
- assert_eq!(state.states.len(), PARTITIONS);
-
- let key_bytes: usize = state.states.keys().map(|pk|
pk.capacity()).sum();
- let scratch_bytes = state.partition_groups.capacity()
- * (size_of::<Vec<u8>>() + size_of::<Vec<u32>>())
- + state.ob_runs.capacity() * (size_of::<Vec<u8>>() +
size_of::<Vec<u32>>());
+ assert_eq!(state.core.partitions.len(), PARTITIONS);
+
+ // The store holds a copy of every key too, in the gathered `pk`
+ // column, and so does the partition-key scratch buffer, which still
+ // holds this batch's encoding. Both are taken out: either alone is
+ // about `key_bytes`, so leaving one in would cover the keys whether or
+ // not `size()` charges them.
+ let key_bytes: usize =
state.core.partitions.keys().map(Vec::len).sum();
assert!(key_bytes >= PARTITIONS * KEY_WIDTH, "key bytes {key_bytes}");
- assert!(scratch_bytes > 0, "scratch tables never allocated");
-
- // Reconstruct the total from its parts. Both terms above have to
- // appear for this to balance, so dropping either from `size()`
- // fails here rather than being absorbed by the slack in some
- // other term.
- let expected = size_of::<PartitionedTopKDenseRank>()
- + state.row_converter.size()
- + state.partition_converter.size()
- + state.scratch_rows.size()
- + state.partition_scratch_rows.size()
- + key_bytes
- + state.states.values().map(|s| s.size()).sum::<usize>()
- + state.states.capacity()
- * (size_of::<Vec<u8>>() + size_of::<DenseRankPartitionState>())
- + scratch_bytes
- + state.store.size();
- assert_eq!(state.size(), expected);
+ let outside_copies = state.size()
+ - state.core.store.size()
+ - state.core.partition_scratch_rows.size();
+ assert!(
Review Comment:
`test_partitioned_topk_dense_rank_size_charges_partition_keys` used to
rebuild `size()` term by term with `assert_eq!`. Now it only checks a lower
bound (`outside_copies >= key_bytes`), so it can no longer catch over-charging.
Failure scenario: If `allocated_size()` adds `index_bytes` twice, or a later
change charges `partitions.capacity() * size_of::<S>()` twice, this test still
passes. The bookkeeping test also misses it, because it uses
`core.allocated_size()` on both sides of its equality. The reservation is then
too high and can hit a false ResourcesExhausted error.
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2621,110 +2626,84 @@ 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)]
+/// Every retained row of the partition, grouped by ORDER BY value: `groups`
+/// maps each of the at most `k` distinct-smallest values seen so far to where
+/// every row holding it lives in the store. Like [`RankPartitionState`], this
+/// holds only store coordinates: the output columns of every retained row live
+/// in the operator's one shared [`RecordBatchStore`].
+#[derive(Debug, Default)]
struct DenseRankPartitionState {
- groups: HashMap<Vec<u8>, Vec<GroupEntry>>,
+ /// Keyed by the row-encoded ORDER BY bytes, capped at `k` keys. Each
+ /// group's rows are in arrival order, which `emit` keeps.
+ groups: HashMap<Vec<u8>, Vec<StoreRef>>,
Review Comment:
`DenseRankPartitionState` keeps every ob key twice: in a `HashMap` and in a
mirrored `BinaryHeap`. One `BTreeMap<Vec<u8>, Vec<StoreRef>>` would do the same
work. `last_key_value`/`pop_last` give the boundary and the eviction, iteration
is already in sorted order for emit, and there is one key copy.
Failure scenario: Each new group costs two `key.to_vec()` allocations,
`evict_max` hashes the key again, `emit` sorts every partition's groups, and
the code has to keep the 'keys mirrors groups' invariant by hand. With wide
ORDER BY keys and many partitions, the second key copy roughly doubles the
per-group key memory charged to the pool.
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2827,379 +2761,154 @@ impl PartitionedTopKDenseRank {
metrics: &ExecutionPlanMetricsSet,
ranked_schema: Option<SchemaRef>,
) -> 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 {
- emit_ranks: ranked_schema.is_some(),
- output_schema: ranked_schema.unwrap_or_else(|| Arc::clone(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,
+ core: PartitionedTopKCore::try_new(
+ "PartitionedTopKDenseRank",
+ partition_id,
+ schema,
+ partition_exprs,
+ partition_sort_fields,
+ order_expr,
+ k,
+ batch_size,
+ runtime,
+ metrics,
+ ranked_schema,
+ )?,
+ groups_bytes: 0,
})
}
- /// 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`].
+ /// Encode the partition and ORDER BY columns once each for the whole
batch,
+ /// run every row through its partition's group admission, then register
+ /// the admitted rows in the shared store with one gather.
pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> {
Review Comment:
`PartitionedTopKDenseRank::insert_batch` copies most of
`PartitionedTopKRank::insert_batch`. `push_row` is the same as
`RankPartitionState::push_tie`. `groups_bytes` repeats `ties_bytes`, because
`core.heaps_bytes` can only grow.
Failure scenario: There are now three copies of the setup: splitting the
borrows, `pending`, the `interned` flag and `index_bytes`, and folding in
`live_slots` and the running totals. A fix to one copy (for example to
`live_slots` or `index_bytes` accounting) has to be made in all three by hand.
A core helper that runs the row pass with a per-operator closure, plus a
running total in the core that can also shrink, would remove both extra
operator fields.
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2827,379 +2761,154 @@ impl PartitionedTopKDenseRank {
metrics: &ExecutionPlanMetricsSet,
ranked_schema: Option<SchemaRef>,
) -> 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 {
- emit_ranks: ranked_schema.is_some(),
- output_schema: ranked_schema.unwrap_or_else(|| Arc::clone(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,
+ core: PartitionedTopKCore::try_new(
+ "PartitionedTopKDenseRank",
+ partition_id,
+ schema,
+ partition_exprs,
+ partition_sort_fields,
+ order_expr,
+ k,
+ batch_size,
+ runtime,
+ metrics,
+ ranked_schema,
+ )?,
+ groups_bytes: 0,
})
}
- /// 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`].
+ /// Encode the partition and ORDER BY columns once each for the whole
batch,
+ /// run every row through its partition's group admission, then register
+ /// the admitted rows in the shared store with one gather.
pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> {
- let elapsed_compute = self.metrics.baseline.elapsed_compute().clone();
+ let elapsed_compute =
self.core.metrics.baseline.elapsed_compute().clone();
let _timer = 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)?;
+ // 1. Encode both column sets once for the whole batch.
+ self.core.encode(batch)?;
- let k = self.k;
+ // 2. One pass over the rows in batch order, as in
+ // `PartitionedTopK::insert_batch`: a partition's groups are only
+ // touched by its own rows, and after any prefix of them the
+ // retained set is the rows of the K distinct-smallest values, so
+ // interleaving partitions cannot change a decision. An admission
+ // records which row it kept; phase 3 gathers them all at once.
+ let core = &mut self.core;
+ let k = core.k;
let mut replacements: usize = 0;
+ core.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 = core.store.pending();
+ {
+ let pk_rows = &core.partition_scratch_rows;
+ let ob_rows = &core.scratch_rows;
+ let states = &mut core.partitions;
+ let admitted_rows = &mut core.admitted_rows;
+ let store = &mut core.store;
+ // Accumulated locally and folded in once.
+ let mut interned_bytes = 0usize;
+ let mut groups_grown = 0usize;
+ let mut groups_freed = 0usize;
+ let mut new_slots = 0usize;
+ let mut released_slots = 0usize;
- // 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();
- // Once the partition holds its full K distinct ob values, the
- // largest of them is the bar a new value must beat, and that bar
- // only ever gets stricter: with K values held there is no free
- // slot left for a new (possibly larger) one, and the only way in
- // from here is to evict that largest and put something strictly
- // smaller in its place. So a row worse than today's bar is worse
- // than every later bar too — it is not held now and can never be
- // admitted — and can be dropped before it costs a bucket. Rows
- // *equal* to the bar must still go through: that value is one of
- // the K being held, so its rows belong to a live group.
- let boundary: Option<&[u8]> = if state.groups.len() == k {
- state.keys.peek().map(Vec::as_slice)
- } else {
- None
- };
- for &orig_idx in &indices {
- let ob_row = self.scratch_rows.row(orig_idx as usize);
- if boundary.is_some_and(|b| ob_row.as_ref() > b) {
- continue;
- }
- 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;
- let before = entries.capacity() * size_of::<GroupEntry>();
- let entry = GroupEntry {
- row_indices: run_indices,
- batch_id,
- };
- state.contents_bytes +=
DenseRankPartitionState::entry_bytes(&entry);
- entries.push(entry);
- // The push may have grown the buffer.
- state.contents_bytes +=
- entries.capacity() * size_of::<GroupEntry>() - before;
- 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;
+ DenseRankPartitionState::default()
+ });
+ if interned {
+ interned_bytes += pk.as_ref().len();
}
-
- // Case B: new ob, room available.
- if state.groups.len() < k {
- batch_entry.uses += 1;
- state.insert_new_group(ob_key, run_indices, batch_id);
+ let key = ob_rows.row(row);
+ let key = key.as_ref();
+ // Once the partition holds its full K distinct ob values, the
+ // largest of them is the bar a new value must beat, and that
+ // bar only ever gets stricter: the only way in from here is to
+ // evict that largest and put something strictly smaller in its
+ // place. So a row worse than today's bar can never be
+ // admitted. Rows *equal* to it must still go through: that
+ // value is one of the K being held, so its rows belong to a
+ // live group.
+ let full = state.groups.len() == k;
+ if full && state.keys.peek().is_some_and(|max| key >
max.as_slice()) {
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. `keys` and `groups` own
- // separate copies of the key bytes, so both are
- // uncharged — with their own capacities, which is why
- // `remove_entry` is used to recover the map's copy
- // rather than assuming it matches the popped one.
- let evicted_key = state.keys.pop().expect("max key
present");
- let (map_key, evicted) = state
- .groups
- .remove_entry(&evicted_key)
- .expect("keys mirrors groups");
- state.contents_bytes -= evicted_key.capacity()
- + map_key.capacity()
- + evicted.capacity() * size_of::<GroupEntry>()
- + evicted
- .iter()
- .map(DenseRankPartitionState::entry_bytes)
- .sum::<usize>();
- 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);
+ // 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,
+ };
+ if let Some(rows) = state.groups.get_mut(key) {
+ groups_grown += DenseRankPartitionState::push_row(rows,
at);
+ } else {
+ if full {
+ // Every row of the boundary group, from this batch or
+ // earlier ones, now ranks below K.
+ let (evicted, freed) = state.evict_max();
+ for at in &evicted {
Review Comment:
Evicting a group now calls `store.release`, and so a store hash lookup in
`unuse`, once per row. The old code called `unuse` once per GroupEntry, which
covered one batch's run of rows.
Failure scenario: Heavy-tie DENSE_RANK: a boundary group holds 100K tied
rows from 10 earlier batches. One eviction does 100K `batches.get_mut` lookups
where 10 calls that drop a count would do. Rows from the same batch sit next to
each other in arrival order, so the loop could collapse runs with the same
`batch_id` and release each run's count in one call (for example a
`release_n(pending, id, n)`).
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2827,379 +2761,154 @@ impl PartitionedTopKDenseRank {
metrics: &ExecutionPlanMetricsSet,
ranked_schema: Option<SchemaRef>,
) -> 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 {
- emit_ranks: ranked_schema.is_some(),
- output_schema: ranked_schema.unwrap_or_else(|| Arc::clone(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,
+ core: PartitionedTopKCore::try_new(
+ "PartitionedTopKDenseRank",
+ partition_id,
+ schema,
+ partition_exprs,
+ partition_sort_fields,
+ order_expr,
+ k,
+ batch_size,
+ runtime,
+ metrics,
+ ranked_schema,
+ )?,
+ groups_bytes: 0,
})
}
- /// 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`].
+ /// Encode the partition and ORDER BY columns once each for the whole
batch,
+ /// run every row through its partition's group admission, then register
+ /// the admitted rows in the shared store with one gather.
pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> {
- let elapsed_compute = self.metrics.baseline.elapsed_compute().clone();
+ let elapsed_compute =
self.core.metrics.baseline.elapsed_compute().clone();
let _timer = 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)?;
+ // 1. Encode both column sets once for the whole batch.
+ self.core.encode(batch)?;
- let k = self.k;
+ // 2. One pass over the rows in batch order, as in
+ // `PartitionedTopK::insert_batch`: a partition's groups are only
+ // touched by its own rows, and after any prefix of them the
+ // retained set is the rows of the K distinct-smallest values, so
+ // interleaving partitions cannot change a decision. An admission
+ // records which row it kept; phase 3 gathers them all at once.
+ let core = &mut self.core;
+ let k = core.k;
let mut replacements: usize = 0;
+ core.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 = core.store.pending();
+ {
+ let pk_rows = &core.partition_scratch_rows;
+ let ob_rows = &core.scratch_rows;
+ let states = &mut core.partitions;
+ let admitted_rows = &mut core.admitted_rows;
+ let store = &mut core.store;
+ // Accumulated locally and folded in once.
+ let mut interned_bytes = 0usize;
+ let mut groups_grown = 0usize;
+ let mut groups_freed = 0usize;
+ let mut new_slots = 0usize;
+ let mut released_slots = 0usize;
- // 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();
- // Once the partition holds its full K distinct ob values, the
- // largest of them is the bar a new value must beat, and that bar
- // only ever gets stricter: with K values held there is no free
- // slot left for a new (possibly larger) one, and the only way in
- // from here is to evict that largest and put something strictly
- // smaller in its place. So a row worse than today's bar is worse
- // than every later bar too — it is not held now and can never be
- // admitted — and can be dropped before it costs a bucket. Rows
- // *equal* to the bar must still go through: that value is one of
- // the K being held, so its rows belong to a live group.
- let boundary: Option<&[u8]> = if state.groups.len() == k {
- state.keys.peek().map(Vec::as_slice)
- } else {
- None
- };
- for &orig_idx in &indices {
- let ob_row = self.scratch_rows.row(orig_idx as usize);
- if boundary.is_some_and(|b| ob_row.as_ref() > b) {
- continue;
- }
- 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;
- let before = entries.capacity() * size_of::<GroupEntry>();
- let entry = GroupEntry {
- row_indices: run_indices,
- batch_id,
- };
- state.contents_bytes +=
DenseRankPartitionState::entry_bytes(&entry);
- entries.push(entry);
- // The push may have grown the buffer.
- state.contents_bytes +=
- entries.capacity() * size_of::<GroupEntry>() - before;
- 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;
+ DenseRankPartitionState::default()
+ });
+ if interned {
+ interned_bytes += pk.as_ref().len();
}
-
- // Case B: new ob, room available.
- if state.groups.len() < k {
- batch_entry.uses += 1;
- state.insert_new_group(ob_key, run_indices, batch_id);
+ let key = ob_rows.row(row);
+ let key = key.as_ref();
+ // Once the partition holds its full K distinct ob values, the
+ // largest of them is the bar a new value must beat, and that
+ // bar only ever gets stricter: the only way in from here is to
+ // evict that largest and put something strictly smaller in its
+ // place. So a row worse than today's bar can never be
+ // admitted. Rows *equal* to it must still go through: that
+ // value is one of the K being held, so its rows belong to a
+ // live group.
+ let full = state.groups.len() == k;
+ if full && state.keys.peek().is_some_and(|max| key >
max.as_slice()) {
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. `keys` and `groups` own
- // separate copies of the key bytes, so both are
- // uncharged — with their own capacities, which is why
- // `remove_entry` is used to recover the map's copy
- // rather than assuming it matches the popped one.
- let evicted_key = state.keys.pop().expect("max key
present");
- let (map_key, evicted) = state
- .groups
- .remove_entry(&evicted_key)
- .expect("keys mirrors groups");
- state.contents_bytes -= evicted_key.capacity()
- + map_key.capacity()
- + evicted.capacity() * size_of::<GroupEntry>()
- + evicted
- .iter()
- .map(DenseRankPartitionState::entry_bytes)
- .sum::<usize>();
- 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);
+ // 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,
+ };
+ if let Some(rows) = state.groups.get_mut(key) {
+ groups_grown += DenseRankPartitionState::push_row(rows,
at);
+ } else {
+ if full {
+ // Every row of the boundary group, from this batch or
+ // earlier ones, now ranks below K.
+ let (evicted, freed) = state.evict_max();
+ for at in &evicted {
+ store.release(&mut pending, at.batch_id);
}
+ released_slots += evicted.len();
+ replacements += evicted.len();
+ groups_freed += freed;
}
- batch_entry.uses += 1;
- state.insert_new_group(ob_key, run_indices, batch_id);
+ groups_grown += state.insert_group(key, at);
}
- // else: ob >= max — drop the whole run.
+ new_slots += 1;
+ pending.uses += 1;
+ admitted_rows.push(row as u32);
}
- // Return the drained scratch map (capacity retained) for the
- // next partition to reuse.
- self.ob_runs = runs;
+ core.index_bytes += interned_bytes;
+ self.groups_bytes = self.groups_bytes + groups_grown -
groups_freed;
+ core.live_slots = core.live_slots + new_slots - released_slots;
}
- // 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())?;
+ // 3. Gather this batch's admissions into the store, then compact it.
+ core.finish_batch(batch, pending, replacements)?;
+ self.core.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.
+ /// Drain every partition in partition-key order and return the rows as a
+ /// stream of `RecordBatch`es ordered by `(partition_keys, order_keys)`.
+ /// Within a partition the distinct ob keys are sorted (byte-comparable
+ /// encoding == sort order), and a group's rows keep their arrival order.
pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> {
- let Self {
- output_schema,
- emit_ranks,
- 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 out: Vec<Result<RecordBatch>> = Vec::new();
-
- // Gather every retained row with a single `interleave_record_batch`
- // per output batch rather than one `take_record_batch` per
- // `GroupEntry`. A group entry holds only the rows one source batch
- // contributed at one ob value, so entries are numerous and tiny —
- // with P partitions, K distinct ob values and B contributing
- // batches there are up to P × K × B of them, and gathering each
- // one separately builds and tears down that many `RecordBatch`es.
- // `interleave` takes `(batch_pos, row)` pairs across *different*
- // source batches in one call, which is exactly the shape here.
- //
- // The pairs are pushed in emit order — partitions in sorted key
- // order, ob values ascending within a partition, entries in
- // insertion order within an ob value — so the interleaved output
- // is already ordered and needs no post-sort.
- //
- // `batch_refs` holds only the batches the current chunk references
- // and is rebuilt per chunk: for Dictionary columns `interleave`
- // does work for every input array it is handed, referenced or not,
- // so passing the whole store to each chunk made emit
- // O(chunks × store batches). A batch's slot is resolved once per
- // entry, not per row.
- let mut batch_refs: Vec<&RecordBatch> = Vec::new();
- let mut batch_id_pos: HashMap<u32, usize> = HashMap::new();
- let mut indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size);
- let mut sorted_obs: Vec<(Vec<u8>, Vec<GroupEntry>)> = Vec::new();
- // Parallel to `indices`: each row's `DENSE_RANK()`, when emitted.
- let mut ranks: Option<Vec<u64>> =
- emit_ranks.then(|| Vec::with_capacity(batch_size));
-
- // Chunk at `batch_size` so the operator emits the same batch sizes
- // as before and no single `interleave` output exceeds `batch_size`.
- for pk in sorted_pks {
- let DenseRankPartitionState {
- groups,
- keys: _,
- contents_bytes: _,
- } = states.remove(&pk).expect("key from states.keys()");
- // Sort the <= K distinct ob keys so rows emit ascending
- // (byte-comparable encoding == sort order).
- sorted_obs.extend(groups);
- sorted_obs.sort_by(|a, b| a.0.cmp(&b.0));
+ self.core.emit_with(|state| {
+ let mut groups: Vec<(Vec<u8>, Vec<StoreRef>)> =
+ state.groups.into_iter().collect();
Review Comment:
Keep dense-rank emission bookkeeping bounded by chunk size. For large tie
groups, routing emission through `emit_with` introduces substantial eager
allocation: `EmitState::stream` materializes every retained row's coordinates
and optional rank before returning the stream. With N identical `(Int32,
Int32)` rows and ranking enabled, accumulation needs approximately 16N bytes,
but emission reserves approximately 32N bytes (payload plus coordinates and
ranks). A bounded pool can therefore accept every input batch and then fail
during `emit()`. The previous dense-rank emitter kept coordinate/rank scratch
bounded by `batch_size`. Preserve a group iterator and generate this metadata
per output chunk, and cover a large all-ties workload with a bounded-pool
regression.
--
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]
