You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/27 07:01:56 UTC
(datafusion) branch gh-readonly-queue/main/pr-25468-431355fb0299bbe5156f88c8ce56fd91c54961b0 created (now cee15b7cca)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25468-431355fb0299bbe5156f88c8ce56fd91c54961b0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at cee15b7cca perf: interleave retained rows at emit instead of a take per group entry (#25468)
This branch includes the following new commits:
new cee15b7cca perf: interleave retained rows at emit instead of a take per group entry (#25468)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: interleave retained rows at emit instead of a take per group entry (#25468)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25468-431355fb0299bbe5156f88c8ce56fd91c54961b0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit cee15b7cca30faaf894d625d245397e4ec4320f9
Author: Subham Singhal <Su...@users.noreply.github.com>
AuthorDate: Sun Sep 27 07:01:37 2026 +0000
perf: interleave retained rows at emit instead of a take per group entry (#25468)
## Which issue does this PR close?
Part of the `enable_window_topn` work
([#24608](https://github.com/apache/datafusion/issues/24608)).
## Rationale for this change
Profiling `PartitionedTopKDenseRank` (`/usr/bin/sample`, 10M-row
fixture,
100K partitions) showed two costs that dominate the operator, neither of
which is inherent to DENSE_RANK semantics:
1. **Emit was 61% of busy samples**, 87% of it in one tiny
`take_record_batch` per `GroupEntry` plus the coalescer copy that
follows. A `GroupEntry` holds only the rows one source batch
contributed at one ob value, so with P partitions, K distinct ob
values and B contributing batches there are up to P x K x B of them.
`TopKHeap::emit_with_state` already solves this exact shape with a
single `interleave_record_batch`; DENSE_RANK just wasn't using it.
2. **`size()` was 8.3% self-time.** `DenseRankPartitionState::size()`
walked every group and every entry in every partition, and `size()`
is called once per `insert_batch` to resize the reservation — so the
walk costs O(entries) per batch. `TopKHeap` already avoids this with
a running `owned_bytes` total.
## What changes are included in this PR?
Two commits, both internal to `PartitionedTopKDenseRank`:
- **emit**: gather retained rows with `interleave_record_batch`, chunked
at `batch_size` so the operator emits the same batch sizes as before
and never materializes all retained rows at once. Pairs are pushed in
emit order (partitions sorted, ob ascending, entries in insertion
order), so the output needs no post-sort.
- **size**: `DenseRankPartitionState` gains a `contents_bytes` running
total, maintained at the four mutation sites (append to an existing ob
group, insert with room, evict-then-insert, and the buffer growth each
can trigger).
No signature, plan-shape, config or semantic change. `PartitionedTopK`
and `PartitionedTopKRank` are untouched.
## Are these changes tested?
Yes.
- 91 `topk` + 42 `partitioned_topk` unit tests, `window_topn.slt`,
`dev/rust_lint.sh` (fmt + clippy) all pass.
- New `test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute`
asserts the incremental total against a full recompute **after every
batch** across 64 randomized shapes, reusing the `DiffShape` harness
from the correctness differential test. It asserts the workload
actually evicted, so the eviction case can't silently go unchecked.
Mutation-checked: zeroing one term in the eviction charge fails this
test and no other.
- Emit order is **byte-identical** to before at `target_partitions = 1`
(this matters — `PartitionedTopKExec` advertises its output ordering,
and breaking it would either produce wrong ranks or force a SortExec
back into the plan).
- `count(*)` / `sum(ob)` / `sum(rn)` / `sum(pk)` digests unchanged.
## Are there any user-facing changes?
Faster `DENSE_RANK() OVER (PARTITION BY ...)` top-N when
`datafusion.optimizer.enable_window_topn` is enabled. No API change.
### Benchmarks
`benchmarks/queries/h2o/window.sql` q24–q29 (the DENSE_RANK sweep),
10M-row
`J1_1e7_1e7_NA` parquet, `dfbench h2o`, 3 reps round-robin, medians in
ms.
Both binaries `--release` from the same base in one session;
`enable_window_topn` via `DATAFUSION_OPTIMIZER_ENABLE_WINDOW_TOPN`.
| query | shape | flag off | main, flag on | branch, flag on | branch vs
main |
|---|---|--:|--:|--:|--:|
| q24–q28 | 100 – 10K partitions | 191–244 | 45–65 | 44–66 | within
noise |
| q29 | 100K partitions | 193.8 | 213.1 | 149.6 | **1.42x** |
---
datafusion/physical-plan/src/topk/mod.rs | 325 ++++++++++++++++++++++++++-----
1 file changed, 280 insertions(+), 45 deletions(-)
diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs
index cec3681de9..b281e49873 100644
--- a/datafusion/physical-plan/src/topk/mod.rs
+++ b/datafusion/physical-plan/src/topk/mod.rs
@@ -1951,29 +1951,75 @@ struct DenseRankPartitionState {
/// 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>>());
- let contents: usize = self
+ // 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.
+ ///
+ /// `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);
+
+ 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::<GroupEntry>();
+ 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::<GroupEntry>()
- + entries
- .iter()
- .map(|e| e.row_indices.capacity() * size_of::<u32>())
- .sum::<usize>()
+ + entries.iter().map(Self::entry_bytes).sum::<usize>()
})
.sum();
- // `keys` duplicates every key's bytes; charge for them plus the
- // heap's backing Vec (one `Vec<u8>` slot per reserved element).
- let keys_overhead: usize = self.keys.capacity() * size_of::<Vec<u8>>()
- + self.keys.iter().map(|k| k.capacity()).sum::<usize>();
- table_overhead + contents + keys_overhead
+ // `keys` duplicates every key's bytes.
+ groups + self.keys.iter().map(|k| k.capacity()).sum::<usize>()
}
}
@@ -2207,24 +2253,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::<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;
}
// 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;
}
@@ -2233,12 +2278,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::<GroupEntry>()
+ + evicted
+ .iter()
+ .map(DenseRankPartitionState::entry_bytes)
+ .sum::<usize>();
for e in &evicted {
replacements += e.row_indices.len();
if e.batch_id == batch_id {
@@ -2252,14 +2308,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.
}
@@ -2311,11 +2360,41 @@ impl PartitionedTopKDenseRank {
let mut sorted_pks: Vec<Vec<u8>> = states.keys().cloned().collect();
sorted_pks.sort();
- let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size);
+ 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);
+
+ // 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: _ } =
- 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<u8>, Vec<GroupEntry>)> =
@@ -2327,16 +2406,32 @@ impl PartitionedTopKDenseRank {
.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)?;
+ // Resolved lazily so a chunk boundary inside the entry
+ // re-registers its batch in the next chunk's slice.
+ let mut array_pos = None;
+ for row in entry.row_indices {
+ let pos = *array_pos.get_or_insert_with(|| {
+ *batch_id_pos.entry(entry.batch_id).or_insert_with(|| {
+ batch_refs.push(batch);
+ batch_refs.len() - 1
+ })
+ });
+ indices.push((pos, 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();
+ batch_refs.clear();
+ batch_id_pos.clear();
+ array_pos = None;
+ }
+ }
}
}
}
- coalescer.finish_buffered_batch()?;
-
- let mut out: Vec<Result<RecordBatch>> = Vec::new();
- while let Some(b) = coalescer.next_completed_batch() {
+ if !indices.is_empty() {
+ let b = interleave_record_batch(&batch_refs, &indices)?;
(&b).record_output(&metrics.baseline);
out.push(Ok(b));
}
@@ -2382,8 +2477,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;
@@ -3392,16 +3490,26 @@ 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<Vec<(i32, i32)>> {
let batches: Vec<RecordBatch> = stream.try_collect().await?;
let mut rows: Vec<(i32, i32)> = Vec::new();
for b in &batches {
- let pk = b.column(0).as_primitive::<arrow::datatypes::Int32Type>();
- let val = b.column(1).as_primitive::<arrow::datatypes::Int32Type>();
+ let pk = b.column(0).as_primitive::<Int32Type>();
+ let val = b.column(1).as_primitive::<Int32Type>();
for i in 0..b.num_rows() {
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)
}
@@ -3464,6 +3572,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<Schema> {
@@ -4453,6 +4597,97 @@ mod tests {
Ok(())
}
+ /// `emit` rebuilds the `interleave` batch slice per chunk, so a
+ /// Dictionary column must survive chunk boundaries that change which
+ /// source batches are in the slice — including a boundary that falls
+ /// *inside* one group entry (the four `v06` rows), whose batch must be
+ /// re-registered in the next chunk's slice. `batch_size` is 8, so the
+ /// 20 retained rows emit as 8 + 8 + 4, with the second chunk starting
+ /// mid-entry and mixing both source batches.
+ #[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<dyn PhysicalExpr> = 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<RecordBatch> {
+ let mut builder = StringDictionaryBuilder::<Int32Type>::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), all under a single
+ // partition key; `v06` is one entry of four rows at batch 1.
+ state.insert_batch(&dict_batch(&[
+ "v00", "v01", "v02", "v03", "v04", "v05", "v06", "v06", "v06", "v06",
+ ])?)?;
+ 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()).collect::<Vec<_>>(),
+ vec![8, 8, 4]
+ );
+ let mut emitted = Vec::new();
+ for b in &results {
+ let val_col = b.column(1).as_dictionary::<Int32Type>();
+ let dict_values = val_col.values().as_string::<i32>();
+ for key in val_col.keys().values() {
+ emitted.push(dict_values.value(*key as usize).to_string());
+ }
+ }
+ let expected: Vec<String> = [
+ "v00", "v01", "v02", "v03", "v04", "v05", "v06", "v06", "v06", "v06", "v10",
+ "v11", "v12", "v13", "v14", "v15", "v16", "v17", "v18", "v19",
+ ]
+ .iter()
+ .map(|s| s.to_string())
+ .collect();
+ assert_eq!(emitted, expected, "rows must 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.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org