SubhamSinghal commented on code in PR #27136:
URL: https://github.com/apache/datafusion/pull/27136#discussion_r4501529269
##########
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:
Addressed in
https://github.com/apache/datafusion/commit/5f96ec422cf781200fb32420e09c6aaa6ed4fed8
##########
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:
Addressed in
https://github.com/apache/datafusion/commit/5f96ec422cf781200fb32420e09c6aaa6ed4fed8
--
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]
