SubhamSinghal commented on code in PR #27136:
URL: https://github.com/apache/datafusion/pull/27136#discussion_r4505137673
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2584,147 +2626,115 @@ impl PartitionedTopKRank {
/// Within a partition, heap rows come first (sorted by ob), then tie rows
/// (all sharing the boundary ob).
pub(crate) fn emit(self) -> Result<SendableRecordBatchStream> {
- self.core.emit_with(|mut state| {
- let heap = state.heap.drain_sorted();
- // `RANK` is `1 + (rows with a strictly smaller ORDER BY value)`.
- // Counting positions is only valid because the retained set is a
- // *complete* order-prefix of the partition: the heap keeps the K
- // smallest keys and the ties only add rows equal to the boundary,
so
- // no row with a smaller key than a retained row was ever dropped.
- //
- // The ties all sit at the boundary, the heap's largest key, so
they
- // share the rank of the first heap row holding it. Ties exist only
- // while the heap is full, so a heap with no rows has none.
- let boundary_rank = heap.last().map_or(0, |worst| {
- heap.partition_point(|slot| slot.key < worst.key) as u64 + 1
- });
- // Over the sorted heap, a row either ties with its predecessor and
- // repeats its rank or starts a new one at its own 1-based
position.
- let mut prev: Option<Vec<u8>> = None;
- let mut rank = 0;
- let heap = heap.into_iter().zip(1u64..).map(move |(slot,
position)| {
- if prev.as_ref() != Some(&slot.key) {
- rank = position;
- }
- prev = Some(slot.key);
- (slot.at, rank)
- });
- heap.chain(state.ties.into_iter().map(move |at| (at,
boundary_rank)))
- })
+ self.core.emit_with(
+ |mut state| {
+ let heap = state.heap.drain_sorted();
+ // `RANK` is `1 + (rows with a strictly smaller ORDER BY
value)`.
+ // Counting positions is only valid because the retained set
is a
+ // *complete* order-prefix of the partition: the heap keeps
the K
+ // smallest keys and the ties only add rows equal to the
boundary, so
+ // no row with a smaller key than a retained row was ever
dropped.
+ //
+ // The ties all sit at the boundary, the heap's largest key,
so they
+ // share the rank of the first heap row holding it. Ties exist
only
+ // while the heap is full, so a heap with no rows has none.
+ let boundary_rank = heap.last().map_or(0, |worst| {
+ heap.partition_point(|slot| slot.key < worst.key) as u64 +
1
+ });
+ // Over the sorted heap, a row either ties with its
predecessor and
+ // repeats its rank or starts a new one at its own 1-based
position.
+ let mut prev: Option<Vec<u8>> = None;
+ let mut rank = 0;
+ let heap = heap.into_iter().zip(1u64..).map(move |(slot,
position)| {
+ if prev.as_ref() != Some(&slot.key) {
+ rank = position;
+ }
+ prev = Some(slot.key);
+ (slot.at, rank)
+ });
+ heap.chain(state.ties.into_iter().map(move |at| (at,
boundary_rank)))
+ },
+ EmitOrder::Collect,
+ )
}
/// Total memory currently held by this operator, including every
/// partition's heap and ties and every batch the store still pins. O(1),
/// see [`PartitionedTopKCore::allocated_size`].
fn size(&self) -> usize {
- size_of::<Self>() + self.core.allocated_size() + self.ties_bytes
+ size_of::<Self>() + self.core.allocated_size()
}
}
-/// 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>>,
/// 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>>,
- /// Running total of the heap allocations owned by the *contents* of
- /// `groups` and `keys`: the key bytes, the per-key `Vec<GroupEntry>`
- /// buffers, and each entry's `row_indices`. Excludes the two
- /// containers' own tables, which `capacity()` reports in O(1).
- ///
- /// INVARIANT: equals `recompute_contents_bytes` (test-only, so not
- /// linkable from rustdoc). Every mutation of `groups` or `keys` must
- /// adjust it; the
- /// `test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute`
- /// test checks this against a full recompute after a randomized
- /// workload.
- contents_bytes: usize,
}
impl DenseRankPartitionState {
- fn size(&self) -> usize {
- let table_overhead = self.groups.capacity()
- * (size_of::<Vec<u8>>() + size_of::<Vec<GroupEntry>>());
- // The heap's backing Vec: one `Vec<u8>` slot per reserved element.
- let keys_overhead = self.keys.capacity() * size_of::<Vec<u8>>();
- table_overhead + self.contents_bytes + keys_overhead
- }
-
- /// Bytes owned by one entry's `row_indices` buffer.
- fn entry_bytes(entry: &GroupEntry) -> usize {
- entry.row_indices.capacity() * size_of::<u32>()
- }
-
- /// Track a previously unseen ob value, holding `run_indices` as its
- /// first (and so far only) entry, and charge everything it allocates.
+ /// Track the previously unseen ob value `key`, holding the row at `at` as
+ /// its only row, and return the bytes this newly allocated.
///
- /// `keys` and `groups` each own a copy of the key bytes; the clone's
- /// capacity is read after the fact rather than assumed equal to the
- /// original's, so the charge matches what was really allocated.
- fn insert_new_group(
- &mut self,
- ob_key: Vec<u8>,
- run_indices: Vec<u32>,
- batch_id: u32,
- ) {
- let key_copy = ob_key.clone();
- self.contents_bytes += key_copy.capacity() + ob_key.capacity();
- self.keys.push(key_copy);
+ /// `keys` and `groups` each own a copy of the key bytes. The table is
+ /// charged through `allocation_size` rather than `capacity`, which a
+ /// removal can lower without freeing anything, so the running total stays
+ /// exact across evictions.
+ fn insert_group(&mut self, key: &[u8], at: StoreRef) -> usize {
+ let table_before = self.groups.allocation_size();
+ let heap_before = self.keys.capacity();
+ let heap_key = key.to_vec();
+ let map_key = key.to_vec();
+ let rows = vec![at];
+ let grown = heap_key.capacity()
+ + map_key.capacity()
+ + rows.capacity() * size_of::<StoreRef>();
+ self.keys.push(heap_key);
+ self.groups.insert(map_key, rows);
+ grown
+ + (self.groups.allocation_size() - table_before)
+ + (self.keys.capacity() - heap_before) * size_of::<Vec<u8>>()
+ }
+
+ /// Stop tracking the largest ob value, returning its rows, which the
+ /// caller owes the store a release each for, and the bytes this freed.
+ fn evict_max(&mut self) -> (Vec<StoreRef>, usize) {
Review Comment:
Addressed in e0a4d3e893da4047f1b7ab58f2043ed22a08d08f
--
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]
