jayzhan211 commented on code in PR #27136:
URL: https://github.com/apache/datafusion/pull/27136#discussion_r4501947399
##########
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:
@SubhamSinghal thanks for this! On input where every row is a new minimum,
this runs slower than `main`: `ORDER BY ts DESC` over 20M rows stored ascending
(10 partitions, `rn <= 10`, `target_partitions = 1`) goes from 930 ms to 1524
ms. Each such row now evicts a whole group, which costs six allocations or
frees, two hash operations, and a heap pop and push. Reusing the evicted
group's buffers through `peek_mut`, as `PartitionHeap::replace_worst` does,
brings most shapes to or below `main` and the worst one I measured from 2.85x
to 1.53x. Can you apply this? It passes the topk tests, clippy and fmt.
```sql
SET datafusion.optimizer.enable_window_topn = true;
SET datafusion.execution.target_partitions = 1;
SELECT count(*), sum(rn), sum(ts) FROM (
SELECT p10, ts, s, dense_rank() OVER (PARTITION BY p10 ORDER BY ts DESC)
AS rn FROM t
) WHERE rn <= 10;
```
```diff
- /// 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) {
- let heap_key = self.keys.pop().expect("a full partition holds k >=
1 keys");
- let (map_key, rows) = self
+ /// Replace the largest tracked ob value with the smaller, previously
+ /// unseen `key`, holding the row at `at` as its only row, reusing the
old
+ /// key buffers and letting `PeekMut` sift down, so a steady-state
eviction
+ /// allocates nothing. Every old row goes back through `admission`.
+ fn replace_max(&mut self, key: &[u8], at: StoreRef, admission: &mut
Admission<'_>) {
+ let table_before = self.groups.allocation_size();
+ let mut heap_key = self
+ .keys
+ .peek_mut()
+ .expect("a full partition holds k >= 1 keys");
+ let (mut map_key, mut rows) = self
.groups
- .remove_entry(&heap_key)
+ .remove_entry(heap_key.as_slice())
.expect("keys mirrors groups");
- let freed = heap_key.capacity()
+ // A group's rows from one batch sit next to each other, in arrival
+ // order, so each run of them drops its uses with one store lookup
+ // rather than one per row.
+ for run in rows.chunk_by(|a, b| a.batch_id == b.batch_id) {
+ admission.release_n(run[0].batch_id, run.len());
+ }
+ admission.replacements += rows.len();
+ admission.freed += heap_key.capacity()
+ + map_key.capacity()
+ + rows.capacity() * size_of::<StoreRef>();
+
+ heap_key.clear();
+ heap_key.extend_from_slice(key);
+ map_key.clear();
+ map_key.extend_from_slice(key);
+ if rows.len() == 1 {
+ rows[0] = at;
+ } else {
+ // Do not let a group that was large pin its capacity.
+ rows = vec![at];
+ }
+ admission.grown += heap_key.capacity()
+ map_key.capacity()
+ rows.capacity() * size_of::<StoreRef>();
- (rows, freed)
+ drop(heap_key);
+ self.groups.insert(map_key, rows);
+ admission.grown += self.groups.allocation_size() - table_before;
}
```
```diff
admission.grown += push_slot(rows, at);
+ } else if full {
+ // Every row of the boundary group, from this batch or
+ // earlier ones, now ranks below K.
+ state.replace_max(key, at, admission);
} 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();
- // A group's rows from one batch sit next to each
other, in
- // arrival order, so each run of them drops its
uses with
- // one store lookup rather than one per row.
- for run in evicted.chunk_by(|a, b| a.batch_id ==
b.batch_id) {
- admission.release_n(run[0].batch_id, run.len());
- }
- admission.replacements += evicted.len();
- admission.freed += freed;
- }
admission.grown += state.insert_group(key, at);
```
<details>
<summary>Operator microbench by shape: main, this PR, this PR with the fix
(ms, median of 7 interleaved runs)</summary>
P = partitions. "Descending" means every row is a new minimum.
| shape | main | PR | PR + fix |
|---|---|---|---|
| 4M rows, descending, P=1, K=1 | 199 | 261 | 133 |
| 4M rows, descending, P=1, K=10 | 240 | 398 | 264 |
| 4M rows, descending, P=1, K=100 | 259 | 455 | 348 |
| 4M rows, descending, P=1, K=1000 | 448 | 594 | 472 |
| 4M rows, descending, P=10, K=100 | 408 | 553 | 417 |
| 4M rows, descending, P=100, K=10 | 336 | 405 | 274 |
| 4M rows, descending, P=100, K=100 | 805 | 946 | 620 |
| 4M rows, descending, P=1000, K=2 | 430 | 468 | 260 |
| 8M rows, descending, P=10, K=1 | 466 | 650 | 356 |
| 8M rows, descending, P=10, K=10 | 591 | 1682 | 904 |
| 4M rows, shuffled, P=100K, K=2 | 1178 | 575 | 548 |
| 4M rows, 5 distinct values, P=1K, K=2 | 271 | 155 | 153 |
End to end, the query above over 20M rows, median of 9: main 930 ms, PR 1524
ms, PR + fix 1388 ms.
</details>
--
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]
