jayzhan211 commented on code in PR #27136:
URL: https://github.com/apache/datafusion/pull/27136#discussion_r4550432871
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1917,26 +1950,141 @@ impl<S: StoreSlots> PartitionedTopKCore<S> {
Ok(())
}
- /// Phase 3: gather the rows this batch contributed into a single batch and
- /// hand it to the store, then compact the store if it needs it. Only rows
- /// admitted at some point during the pass are kept, so what stays pinned
is
- /// bounded by admissions rather than by input size — and the entry is
- /// freed as soon as the last slot referencing it is released.
+ /// Run every row of `batch` through `admit` against its partition's state,
+ /// then gather the admitted rows into the store and compact it: the whole
+ /// of `insert_batch` but the admission rule, which is all that differs
+ /// between the three operators.
+ ///
+ /// Both column sets are encoded once for the whole batch, then the rows
+ /// are visited in batch order rather than grouped by partition.
+ /// Interleaving partitions cannot change a decision — a state is only read
+ /// and written by its own partition's rows, which still reach it in
+ /// ascending row order — and it costs no pass over the partitions seen so
+ /// far, which would be quadratic in the input whenever partition count
+ /// grows with it. Only keys are copied during the pass; an admission
+ /// records where its row will live, and one gather afterwards materializes
+ /// the batch's admissions once, however many partitions they went to. Only
+ /// rows admitted at some point during the pass are kept, so what stays
+ /// pinned is bounded by admissions rather than by input size, and the
+ /// entry is freed as soon as the last slot referencing it is released.
///
- /// No uses left on `pending` means every admission from this batch was
- /// released again before the pass ended, so there is nothing to keep and
- /// the id goes back to the next batch.
+ /// `admit` gets the row's partition state, its encoded ORDER BY key, the
+ /// slot the row takes if admitted, and the [`Admission`] it reports
+ /// through, and returns whether it admitted the row.
///
/// The caller resizes the reservation afterwards, since only it knows its
/// full size.
- fn finish_batch(
+ fn insert_batch_with(
&mut self,
batch: &RecordBatch,
- pending: PendingEntry,
- replacements: usize,
- ) -> Result<()> {
+ mut admit: impl FnMut(&mut S, &[u8], StoreRef, &mut Admission<'_>) ->
bool,
+ ) -> Result<()>
+ where
+ S: Default,
+ {
+ let elapsed_compute = self.metrics.baseline.elapsed_compute().clone();
+ let _timer = elapsed_compute.timer();
+
+ if batch.num_rows() == 0 {
+ return Ok(());
+ }
+ self.encode(batch)?;
+
+ self.admitted_rows.clear();
+ // The gathered batch does not exist yet, but its id does, so a slot
can
+ // point at it during the pass with no back-patching afterwards.
+ let pending_id = self.store.next_batch_id();
+ let mut admission = Admission {
+ store: &mut self.store,
+ pending_id,
+ pending_released: 0,
+ grown: 0,
+ freed: 0,
+ released_slots: 0,
+ replacements: 0,
+ };
+ // Accumulated locally and folded in once, so the running totals are
not
+ // touched per row.
+ let mut interned_bytes = 0usize;
+ for row in 0..batch.num_rows() {
Review Comment:
@SubhamSinghal thanks for applying `replace_max`! `ORDER BY ts DESC` over
rows stored by ascending `ts` is still 1.5–2x slower than `main` at 1–10
partitions (20M rows, K = 10: 652 → 1090 ms of operator compute). Walking rows
in arrival order evicts once per row here, where `main`'s hash-order buckets
evict about K·ln(n/K) times per batch. Walking a falling batch back to front
for `DENSE_RANK` brings it to 202 ms, and every other shape I measured stays
within noise of this head. Could you apply this?
```diff
trait StoreSlots {
+ /// Whether the rows a partition keeps do not depend on the order its
+ /// rows arrive in, so [`PartitionedTopKCore::insert_batch_with`] may
visit
+ /// a batch back to front. True for `DENSE_RANK`, which keeps every row
+ /// tied with a kept value; not for `ROW_NUMBER`, which keeps the tied
rows
+ /// that come first.
+ const ORDER_FREE: bool = false;
+
@@ fn insert_batch_with
let mut interned_bytes = 0usize;
- for row in 0..batch.num_rows() {
+ let num_rows = batch.num_rows();
+ // Front to back, a batch whose ORDER BY values fall, such as
+ // `ORDER BY ts DESC` over rows stored by ascending `ts`, makes
nearly
+ // every row a new best that displaces one. Back to front, the
first K
+ // per partition fill the state and the boundary check drops the
rest,
+ // which only saves work when partitions get more than K rows each.
+ let backwards = S::ORDER_FREE
+ && num_rows > self.k.saturating_mul(self.partitions.len())
+ && self.scratch_rows.row(num_rows - 1) <
self.scratch_rows.row(0);
+ for i in 0..num_rows {
+ let row = if backwards { num_rows - 1 - i } else { i };
let pk = self.partition_scratch_rows.row(row);
@@ impl StoreSlots for DenseRankPartitionState {
+ const ORDER_FREE: bool = true;
+
```
The `num_rows > k * partitions` condition skips batches where each partition
gets fewer than K rows, since those evict just as often in either direction
(without it, 1000 partitions at K = 100 ran 1.3x slower than this head).
Walking backwards reverses the order of tied rows within a group. SQL leaves
that order unspecified, but the "arrival order" docs and three tests need
updating (below).
<details>
<summary>Measurements (base, this head, with the change)</summary>
`elapsed_compute` of `PartitionedTopKExec`, median of 7 interleaved runs,
`target_partitions = 1`, 20M rows, `enable_window_topn = true`, `SELECT
count(*), sum(rn), sum(ts) FROM (SELECT pX, ts, s, dense_rank() OVER (PARTITION
BY pX ORDER BY ts DESC) AS rn FROM t) WHERE rn <= K`.
| input | partitions | K | main ms | this PR | with change |
|---|---|---|---|---|---|
| ts ascending, DESC | 1 | 10 | 652 | 1090 (1.67x) | 202 (0.31x) |
| ts ascending, DESC | 1 | 100 | 910 | 1780 (1.96x) | 218 (0.24x) |
| ts ascending, DESC | 10 | 10 | 678 | 992 (1.46x) | 208 (0.31x) |
| ts ascending, DESC | 100 | 10 | 1250 | 1090 (0.87x) | 324 (0.26x) |
| ts ascending, DESC | 100 | 100 | 2640 | 2500 (0.95x) | 2590 (0.98x) |
| ts ascending, DESC | 1000 | 10 | 2110 | 1270 (0.60x) | 1290 (0.61x) |
| ts ascending, DESC | 1000 | 100 | 4260 | 2900 (0.68x) | 2930 (0.69x) |
| ts ascending, ASC | 1 | 10 | 176 | 180 (1.03x) | 181 (1.03x) |
| ts ascending, ASC | 1000 | 10 | 339 | 228 (0.67x) | 238 (0.70x) |
| rising per batch, falling across, ASC | 1 | 10 | 618 | 183 (0.30x) | 191
(0.31x) |
| random ts, DESC | 10 | 100 | 197 | 188 (0.95x) | 193 (0.98x) |
| random ts, DESC | 1000 | 10 | 376 | 267 (0.71x) | 281 (0.75x) |
In a profile of this head at 1 partition, K = 100, the heap sift in
`PeekMut::drop` takes 38% of the operator and hashing (`get_mut` miss,
`remove_entry`, `insert`) 43%, so making each eviction cheaper can't close a
1.9x gap. ROW_NUMBER and RANK are at 0.93–1.05x of main, so the
`insert_batch_with` refactor is fine.
</details>
<details>
<summary>Tests: two new, three adjusted (all 152 topk tests, the window/topk
sqllogictests and clippy pass; removing each of the three conditions fails a
test)</summary>
```rs
/// A batch whose ORDER BY values fall is visited back to front, so it
/// evicts nothing: front to back, every row would evict the one before.
#[tokio::test]
async fn test_partitioned_topk_dense_rank_falling_batch_evicts_nothing()
-> Result<()>
{
let (schema, mut state) = build_partitioned_topk_dense_rank(2)?;
let vals: Vec<i32> = (0..1000).rev().collect();
state.insert_batch(&pk_val_batch(&schema, vec![1; 1000], vals)?)?;
assert_eq!(state.core.metrics.row_replacements.value(), 0);
assert_eq!(pk_val_rows(state.emit()?).await?, vec![(1, 0), (1, 1)]);
Ok(())
}
/// A falling batch with no more than K rows per partition seen is walked
/// front to back, since each partition would keep all of its rows either
/// way. The check goes by that average, so a batch that puts all its
rows
/// in one partition is still walked in order.
#[tokio::test]
async fn
test_partitioned_topk_dense_rank_falling_batch_over_many_partitions()
-> Result<()> {
let (schema, mut state) = build_partitioned_topk_dense_rank(1)?;
state.insert_batch(&pk_val_batch(&schema, vec![1, 2, 3, 4],
vec![100; 4])?)?;
// 4 rows <= K * 4 partitions: 4 evicts the 100, then 3, 2 and 1 each
// evict the one before.
state.insert_batch(&pk_val_batch(&schema, vec![1; 4], vec![4, 3, 2,
1])?)?;
assert_eq!(state.core.metrics.row_replacements.value(), 4);
let rows = pk_val_rows(state.emit()?).await?;
assert_eq!(rows, vec![(1, 1), (2, 100), (3, 100), (4, 100)]);
Ok(())
}
```
`evicts_same_call_group` and `evicts_group_spanning_batches` test eviction
within a batch on falling input, so each needs a trailing row that keeps the
batch walked front to back. `store_bounded_when_survivors_spread` asserts the
arrival order of tied rows, so it compares them as a set:
```diff
- let vals: Vec<i32> = (0..32).rev().collect();
- state.insert_batch(&pk_val_batch(&schema, vec![1; 32], vals)?)?;
+ // The trailing 40 ends the batch above its start, so it is visited
+ // front to back and every falling value evicts the one before it.
+ let mut vals: Vec<i32> = (0..32).rev().collect();
+ vals.push(40);
+ state.insert_batch(&pk_val_batch(&schema, vec![1; 33], vals)?)?;
@@
- state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![10,
5])?)?;
+ // The trailing 10 loses to the boundary and keeps the batch from
+ // falling, so it is visited front to back.
+ state.insert_batch(&pk_val_batch(&schema, vec![1, 1, 1], vec![10,
5, 10])?)?;
@@
+ for ties in rows.chunk_by_mut(|a, b| a.0 == b.0) {
+ ties.sort_unstable();
+ }
assert_eq!(rows, expected);
```
</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]
