From ed9a2b5afeb02533afa77dc65eefbbe2850fa8fc Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Fri, 18 Sep 2026 12:23:04 +0530 Subject: [PATCH 1/3] interleave retained rows at emit instead of a take per group entry --- datafusion/physical-plan/src/topk/mod.rs | 207 ++++++++++++++++++----- 1 file changed, 167 insertions(+), 40 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 0ca700cb37655..61294625584e5 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1952,29 +1952,73 @@ struct DenseRankPartitionState { /// INVARIANT: `keys` and `groups.keys()` hold the same set. Every /// insertion into / removal from `groups` must mirror into `keys`. keys: BinaryHeap>, + /// Running total of the heap allocations owned by the *contents* of + /// `groups` and `keys`: the key bytes, the per-key `Vec` + /// 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 `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::>() + size_of::>()); - let contents: usize = self + // The heap's backing Vec: one `Vec` slot per reserved element. + let keys_overhead = self.keys.capacity() * size_of::>(); + 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::() + } + + /// Track a previously unseen ob value, holding `run_indices` as its + /// first (and so far only) entry, and charge everything it allocates. + /// + /// `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, + run_indices: Vec, + batch_id: u32, + ) { + let key_copy = ob_key.clone(); + self.contents_bytes += key_copy.capacity() + ob_key.capacity(); + self.keys.push(key_copy); + + let entry = GroupEntry { + row_indices: run_indices, + batch_id, + }; + self.contents_bytes += Self::entry_bytes(&entry); + let entries = vec![entry]; + self.contents_bytes += entries.capacity() * size_of::(); + self.groups.insert(ob_key, entries); + } + + /// The value [`Self::contents_bytes`] must hold, computed the slow + /// way. Used to assert the incremental accounting in tests. + #[cfg(test)] + fn recompute_contents_bytes(&self) -> usize { + let groups: usize = self .groups .iter() .map(|(key, entries)| { key.capacity() + entries.capacity() * size_of::() - + entries - .iter() - .map(|e| e.row_indices.capacity() * size_of::()) - .sum::() + + entries.iter().map(Self::entry_bytes).sum::() }) .sum(); - // `keys` duplicates every key's bytes; charge for them plus the - // heap's backing Vec (one `Vec` slot per reserved element). - let keys_overhead: usize = self.keys.capacity() * size_of::>() - + self.keys.iter().map(|k| k.capacity()).sum::(); - table_overhead + contents + keys_overhead + // `keys` duplicates every key's bytes. + groups + self.keys.iter().map(|k| k.capacity()).sum::() } } @@ -2208,24 +2252,23 @@ impl PartitionedTopKDenseRank { // new `GroupEntry` (one entry per contributing batch). if let Some(entries) = state.groups.get_mut(&ob_key) { batch_entry.uses += 1; - entries.push(GroupEntry { + let before = entries.capacity() * size_of::(); + 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::() - before; continue; } // Case B: new ob, room available. if state.groups.len() < k { batch_entry.uses += 1; - state.keys.push(ob_key.clone()); - state.groups.insert( - ob_key, - vec![GroupEntry { - row_indices: run_indices, - batch_id, - }], - ); + state.insert_new_group(ob_key, run_indices, batch_id); continue; } @@ -2234,12 +2277,23 @@ impl PartitionedTopKDenseRank { 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. + // 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 evicted = state + let (map_key, evicted) = state .groups - .remove(&evicted_key) + .remove_entry(&evicted_key) .expect("keys mirrors groups"); + state.contents_bytes -= evicted_key.capacity() + + map_key.capacity() + + evicted.capacity() * size_of::() + + evicted + .iter() + .map(DenseRankPartitionState::entry_bytes) + .sum::(); for e in &evicted { replacements += e.row_indices.len(); if e.batch_id == batch_id { @@ -2253,14 +2307,7 @@ impl PartitionedTopKDenseRank { } } batch_entry.uses += 1; - state.keys.push(ob_key.clone()); - state.groups.insert( - ob_key, - vec![GroupEntry { - row_indices: run_indices, - batch_id, - }], - ); + state.insert_new_group(ob_key, run_indices, batch_id); } // else: ob >= max — drop the whole run. } @@ -2314,9 +2361,36 @@ impl PartitionedTopKDenseRank { let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + // 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. + let mut batch_refs = Vec::with_capacity(store.len()); + let mut batch_id_pos = HashMap::with_capacity(store.len()); + for (array_pos, (batch_id, entry)) in store.batches.iter().enumerate() { + batch_refs.push(&entry.batch); + batch_id_pos.insert(*batch_id, array_pos); + } + + // Chunk at `batch_size` so the operator emits the same batch sizes + // as before and never materializes all retained rows at once. + let mut indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size); for pk in sorted_pks { - let DenseRankPartitionState { groups, keys: _ } = - states.remove(&pk).expect("key from states.keys()"); + 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). let mut sorted_obs: Vec<(Vec, Vec)> = @@ -2324,16 +2398,23 @@ impl PartitionedTopKDenseRank { sorted_obs.sort_by(|a, b| a.0.cmp(&b.0)); for (_ob, entries) in sorted_obs { for entry in entries { - let batch = &store - .get(entry.batch_id) - .expect("retained batch_id present in store") - .batch; - let indices = UInt32Array::from(entry.row_indices); - let sub = take_record_batch(batch, &indices)?; - coalescer.push_batch(sub)?; + let array_pos = batch_id_pos[&entry.batch_id]; + for row in entry.row_indices { + indices.push((array_pos, row as usize)); + if indices.len() == batch_size { + coalescer.push_batch(interleave_record_batch( + &batch_refs, + &indices, + )?)?; + indices.clear(); + } + } } } } + if !indices.is_empty() { + coalescer.push_batch(interleave_record_batch(&batch_refs, &indices)?)?; + } coalescer.finish_buffered_batch()?; let mut out: Vec> = Vec::new(); @@ -3393,6 +3474,12 @@ mod tests { /// Drain an operator's output into sorted `(pk, val)` pairs, ready to /// compare against [`DiffShape::expected`]. + /// + /// Asserts the rows arrived already in `(pk ASC, val ASC)` order before + /// sorting them. That order is load-bearing rather than cosmetic — it is + /// what `PartitionedTopKExec::compute_properties` advertises, and so what + /// lets the window above it run `mode=Sorted` with no `SortExec` — and + /// the content comparison this feeds would otherwise sort it away. async fn sorted_pk_val(stream: SendableRecordBatchStream) -> Result> { let batches: Vec = stream.try_collect().await?; let mut rows: Vec<(i32, i32)> = Vec::new(); @@ -3403,6 +3490,10 @@ mod tests { rows.push((pk.value(i), val.value(i))); } } + assert!( + rows.windows(2).all(|w| w[0] <= w[1]), + "emitted rows are not in (pk, val) order: {rows:?}" + ); rows.sort_unstable(); Ok(rows) } @@ -3465,6 +3556,42 @@ mod tests { Ok(()) } + /// `DenseRankPartitionState::contents_bytes` is maintained + /// incrementally at four mutation sites (append to an existing ob + /// group, insert a new one with room, evict-then-insert, and the + /// key/entry buffer growth each can trigger). Drift there would + /// silently give the reservation a wrong total, so check it against a + /// recompute over the same randomized workload the correctness + /// differential test uses — its shapes are tuned to exercise + /// eviction, which is the case with the most bookkeeping. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute() + -> Result<()> { + let mut saw_eviction = false; + for seed in 0..64u64 { + let shape = DiffShape::new(seed, 8); + let (schema, mut state) = build_partitioned_topk_dense_rank(shape.k)?; + for (pks, vals) in &shape.batches { + state.insert_batch(&pk_val_batch(&schema, pks.clone(), vals.clone())?)?; + // Check after every batch, not just at the end: a + // compensating pair of errors within one batch would + // survive an end-only assertion. + for (pk, partition) in &state.states { + assert_eq!( + partition.contents_bytes, + partition.recompute_contents_bytes(), + "seed {seed}, partition {pk:?}: {shape}" + ); + } + } + saw_eviction |= state.metrics.row_replacements.value() > 0; + } + // Guards the guard: if the shapes stopped evicting, case C would + // go unchecked and this test would still pass. + assert!(saw_eviction, "workload never evicted a group"); + Ok(()) + } + /// The `(pk Int32, val Int32)` schema every `PartitionedTopK*` test /// builds against. `val_nullable` is what the null-ordering tests vary. fn pk_val_schema(val_nullable: bool) -> Arc { From e580f216bdfe07ddaa14131339223bed4f933e9f Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Sat, 19 Sep 2026 12:08:08 +0530 Subject: [PATCH 2/3] drop batch coalescer and address review comments --- datafusion/physical-plan/src/topk/mod.rs | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 61294625584e5..0d1867b662fba 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -1959,8 +1959,10 @@ struct DenseRankPartitionState { /// /// INVARIANT: equals `recompute_contents_bytes` (test-only, so not /// linkable from rustdoc). Every mutation of `groups` or `keys` must - /// adjust it; the `dense_rank_contents_bytes_tracks_recompute` test - /// checks this against a full recompute after a randomized workload. + /// 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, } @@ -2359,7 +2361,7 @@ impl PartitionedTopKDenseRank { let mut sorted_pks: Vec> = states.keys().cloned().collect(); sorted_pks.sort(); - let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + let mut out: Vec> = Vec::new(); // Gather every retained row with a single `interleave_record_batch` // per output batch rather than one `take_record_batch` per @@ -2383,7 +2385,7 @@ impl PartitionedTopKDenseRank { } // Chunk at `batch_size` so the operator emits the same batch sizes - // as before and never materializes all retained rows at once. + // as before and no single `interleave` output exceeds `batch_size`. let mut indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size); for pk in sorted_pks { let DenseRankPartitionState { @@ -2402,10 +2404,9 @@ impl PartitionedTopKDenseRank { for row in entry.row_indices { indices.push((array_pos, row as usize)); if indices.len() == batch_size { - coalescer.push_batch(interleave_record_batch( - &batch_refs, - &indices, - )?)?; + let b = interleave_record_batch(&batch_refs, &indices)?; + (&b).record_output(&metrics.baseline); + out.push(Ok(b)); indices.clear(); } } @@ -2413,12 +2414,7 @@ impl PartitionedTopKDenseRank { } } if !indices.is_empty() { - coalescer.push_batch(interleave_record_batch(&batch_refs, &indices)?)?; - } - coalescer.finish_buffered_batch()?; - - let mut out: Vec> = Vec::new(); - while let Some(b) = coalescer.next_completed_batch() { + let b = interleave_record_batch(&batch_refs, &indices)?; (&b).record_output(&metrics.baseline); out.push(Ok(b)); } From 2c761166c5c5bdce6440133c29865696b49d4f77 Mon Sep 17 00:00:00 2001 From: SubhamSinghal Date: Sat, 19 Sep 2026 17:16:56 +0530 Subject: [PATCH 3/3] PartitionedTopKDenseRank: scope emit's interleave batch slice per chunk --- datafusion/physical-plan/src/topk/mod.rs | 136 +++++++++++++++++++---- 1 file changed, 116 insertions(+), 20 deletions(-) diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 0d1867b662fba..3d1615b50af57 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -2377,16 +2377,36 @@ impl PartitionedTopKDenseRank { // 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. - let mut batch_refs = Vec::with_capacity(store.len()); - let mut batch_id_pos = HashMap::with_capacity(store.len()); - for (array_pos, (batch_id, entry)) in store.batches.iter().enumerate() { - batch_refs.push(&entry.batch); - batch_id_pos.insert(*batch_id, array_pos); - } + // + // `indices` carries the global `batch_id` rather than a position + // into a fixed `batch_refs` slice: for Dictionary columns, + // `interleave` does work proportional to the *number of input + // arrays*, not just the ones the indices reference, so passing + // every batch in the store on every chunk turns emit into + // O(chunks × total batches) instead of O(chunks × batches the + // chunk actually uses). Each chunk below rebuilds a batch slice + // scoped to just its own `batch_id`s. + let mut indices: Vec<(u32, usize)> = Vec::with_capacity(batch_size); + let mut flush = |indices: &mut Vec<(u32, usize)>| -> Result<()> { + let mut batch_refs = Vec::new(); + let mut local_pos = HashMap::new(); + let mut local_indices = Vec::with_capacity(indices.len()); + for &(batch_id, row) in indices.iter() { + let pos = *local_pos.entry(batch_id).or_insert_with(|| { + batch_refs.push(&store.batches[&batch_id].batch); + batch_refs.len() - 1 + }); + local_indices.push((pos, row)); + } + let b = interleave_record_batch(&batch_refs, &local_indices)?; + (&b).record_output(&metrics.baseline); + out.push(Ok(b)); + indices.clear(); + Ok(()) + }; // Chunk at `batch_size` so the operator emits the same batch sizes // as before and no single `interleave` output exceeds `batch_size`. - let mut indices: Vec<(usize, usize)> = Vec::with_capacity(batch_size); for pk in sorted_pks { let DenseRankPartitionState { groups, @@ -2400,23 +2420,17 @@ impl PartitionedTopKDenseRank { sorted_obs.sort_by(|a, b| a.0.cmp(&b.0)); for (_ob, entries) in sorted_obs { for entry in entries { - let array_pos = batch_id_pos[&entry.batch_id]; for row in entry.row_indices { - indices.push((array_pos, row as usize)); + indices.push((entry.batch_id, row as usize)); if indices.len() == batch_size { - let b = interleave_record_batch(&batch_refs, &indices)?; - (&b).record_output(&metrics.baseline); - out.push(Ok(b)); - indices.clear(); + flush(&mut indices)?; } } } } } if !indices.is_empty() { - let b = interleave_record_batch(&batch_refs, &indices)?; - (&b).record_output(&metrics.baseline); - out.push(Ok(b)); + flush(&mut indices)?; } Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -2460,8 +2474,11 @@ impl PartitionedTopKDenseRank { mod tests { use super::*; use crate::metrics::MetricValue; - use arrow::array::{BooleanArray, Float64Array, Int32Array, StringArray}; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::array::{ + AsArray, BooleanArray, Float64Array, Int32Array, StringArray, + StringDictionaryBuilder, + }; + use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use arrow_schema::SortOptions; use datafusion_common::{assert_batches_eq, exec_datafusion_err}; use datafusion_execution::memory_pool::GreedyMemoryPool; @@ -3480,8 +3497,8 @@ mod tests { let batches: Vec = stream.try_collect().await?; let mut rows: Vec<(i32, i32)> = Vec::new(); for b in &batches { - let pk = b.column(0).as_primitive::(); - let val = b.column(1).as_primitive::(); + let pk = b.column(0).as_primitive::(); + let val = b.column(1).as_primitive::(); for i in 0..b.num_rows() { rows.push((pk.value(i), val.value(i))); } @@ -4535,6 +4552,85 @@ mod tests { Ok(()) } + /// Regression test for the `emit` rewrite: passing the *entire* store + /// to every `interleave_record_batch` call made Dictionary columns + /// O(chunks × total batches) instead of O(chunks × batches the chunk + /// actually uses), since `interleave_dictionaries` does per-input-array + /// work for every array it's handed. `batch_size` is 8, so 20 retained + /// rows spread across 2 source batches emit in 3 chunks — with several + /// chunks mixing rows from both source batches — the exact shape that + /// exercises the per-chunk batch-slice rebuild. + #[tokio::test] + async fn test_partitioned_topk_dense_rank_emit_dictionary_spans_chunks() -> Result<()> + { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new( + "val", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + ), + ])); + + let pk_expr: Arc = col("pk", schema.as_ref())?; + let partition_sort_fields = build_sort_fields( + &[PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }], + &schema, + )?; + let order_expr = LexOrdering::from([PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: SortOptions::default(), + }]); + + let mut state = PartitionedTopKDenseRank::try_new( + 0, + Arc::clone(&schema), + vec![pk_expr], + partition_sort_fields, + order_expr, + 20, // k: large enough to retain every distinct value below + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + + let dict_batch = |vals: &[&str]| -> Result { + let mut builder = StringDictionaryBuilder::::new(); + for v in vals { + builder.append_value(v); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; vals.len()])), + Arc::new(builder.finish()), + ], + )?) + }; + + // Two source batches (two distinct `batch_id`s), 10 distinct values + // each, all under a single partition key so every output chunk + // draws from both. + state.insert_batch(&dict_batch(&[ + "v00", "v01", "v02", "v03", "v04", "v05", "v06", "v07", "v08", "v09", + ])?)?; + state.insert_batch(&dict_batch(&[ + "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19", + ])?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_eq!(results.iter().map(|b| b.num_rows()).sum::(), 20); + assert_eq!(results.len(), 3, "expected 3 chunks of batch_size=8"); + let val_col = results[0].column(1).as_dictionary::(); + let dict_values = val_col.values().as_string::(); + let first_val = dict_values.value(val_col.keys().value(0) as usize); + assert_eq!(first_val, "v00", "rows must still emit in ob-sorted order"); + Ok(()) + } + /// DENSE_RANK-specific: eviction removes the entire max group when /// a strictly-smaller distinct ob arrives. Multiple rows at the /// evicted key all disappear.