perf: interleave retained rows at emit instead of a take per group entry

[SubhamSinghal]

Which issue does this PR close?

Part of the enable_window_topn work (#24060).

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

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.97222% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.38%. Comparing base (0e292dc) to head (2c76116).
⚠️ Report is 14 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/topk/mod.rs	90.97%	3 Missing and 10 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25468      +/-   ##
==========================================
+ Coverage   82.35%   82.38%   +0.03%     
==========================================
  Files        1137     1138       +1     
  Lines      432746   433834    +1088     
  Branches   432746   433834    +1088     
==========================================
+ Hits       356375   357410    +1035     
- Misses      54843    54857      +14     
- Partials    21528    21567      +39     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[jayzhan211]

Thanks @SubhamSinghal , some suggestions

[jayzhan211]

Chunks are already exactly batch_size rows (last one ≤), so BatchCoalescer only re-copies full batches. Optional simplification; I measured no perf difference either way, and it drops the coalescer's view-buffer GC, so feel free to ignore.

-        let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size);
+        let mut out: Vec<Result<RecordBatch>> = Vec::new();
@@
                         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();
                         }
@@
         if !indices.is_empty() {
-            coalescer.push_batch(interleave_record_batch(&batch_refs, &indices)?)?;
-        }
-        coalescer.finish_buffered_batch()?;
-
-        let mut out: Vec<Result<RecordBatch>> = 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));
         }

[SubhamSinghal]

addressed in e580f21

[jayzhan211]

Suggested change

	/// adjust it; the `dense_rank_contents_bytes_tracks_recompute` test
	/// adjust it; `test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute`

[SubhamSinghal]

addressed in e580f21

[jayzhan211]

Suggested change

	// Chunk at `batch_size` so the operator emits the same batch sizes
	// as before and never materializes all retained rows at once.
	// Chunk at `batch_size` so the operator emits the same batch sizes
	// as before and no single `interleave` output exceeds `batch_size`.

[SubhamSinghal]

addressed in e580f21

[jayzhan211]

@SubhamSinghal

Emit is O(chunks × store batches × batch_size) for dictionary columns

batch_refs = every batch in the store, passed to every per-chunk interleave_record_batch. interleave_dictionaries does per-input-array work per call (key mask per dictionary + indices.iter().filter(|(a, _)| *a == a_idx); the fallback concats all B value arrays into each output) → emit is O(chunks × B × batch_size) for Dictionary columns; base was linear.

Measured (release, 8192-row batches, Dictionary<Int32, Utf8> with a per-batch dictionary, 8 partitions × 4 ob values, all rows retained), emit only:

batches	main	PR	PR + fix
128	0.010s	0.052s	0.012s
512	0.064s	0.735s	0.081s
2048	0.263s	16.1s	0.335s

Fix: build the refs per chunk from only the batches it touches (neutral on the tiny-entry shape: 1.29s vs 1.29s dict, 0.58s vs 0.58s utf8).

-        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);
-        }
+        // Rebuilt per chunk: `interleave` does per-input-array work for
+        // some types (dictionaries), so only pass batches the chunk uses.
+        let mut batch_refs: Vec<&RecordBatch> = Vec::new();
+        let mut batch_id_pos: HashMap<u32, usize> = HashMap::new();

-                    let array_pos = batch_id_pos[&entry.batch_id];
+                    let batch = &store
+                        .get(entry.batch_id)
+                        .expect("retained batch_id present in store")
+                        .batch;
+                    let mut array_pos = None;
                     for row in entry.row_indices {
-                        indices.push((array_pos, row as usize));
+                        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;
                         }

Please also add an emit test with a Dictionary column spanning ≥2 chunks.

Tiny entries (1 row per GroupEntry, 1024 partitions), 512 batches:

┌────────────────────────────────────────┬────────────────────┬──────────────────┐
│                 d type                 │ base insert / emit │ PR insert / emit │
├────────────────────────────────────────┼────────────────────┼──────────────────┤
│ Utf8                                   │ 8.06s / 2.32s      │ 1.01s / 0.63s    │
├────────────────────────────────────────┼────────────────────┼──────────────────┤
│ Dict<Int32,Utf8>, per-batch dictionary │ 7.05s / 2.83s      │ 0.71s / 1.23s    │
└────────────────────────────────────────┴────────────────────┴──────────────────┘

Both PR claims hold on this shape: insert is ~8x faster (the size() walk is gone) and emit is 2.3–3.7x faster.

Large entries (8 partitions × 4 ob values, 256 rows per entry), emit only:

┌─────────┬─────────────────┬─────────────────┬──────────────────────┐
│ batches │ Utf8 base / PR  │ Dict base / PR  │ Dict, PR + fix below │
├─────────┼─────────────────┼─────────────────┼──────────────────────┤
│ 128     │ 0.018s / 0.017s │ 0.010s / 0.052s │ 0.012s               │
├─────────┼─────────────────┼─────────────────┼──────────────────────┤
│ 512     │ 0.106s / 0.109s │ 0.064s / 0.735s │ 0.081s               │
├─────────┼─────────────────┼─────────────────┼──────────────────────┤
│ 2048    │ 0.417s / 0.481s │ 0.263s / 16.1s  │ 0.335s               │
└─────────┴─────────────────┴─────────────────┴──────────────────────┘

Fix vs PR on the tiny-entry shape (interleaved, 2 runs each), 512 batches:

┌────────┬───────────────┬───────────────┐
│ d type │    PR emit    │ PR + fix emit │
├────────┼───────────────┼───────────────┤
│ Utf8   │ 0.59s / 0.58s │ 0.62s / 0.58s │
├────────┼───────────────┼───────────────┤
│ Dict   │ 1.29s / 1.29s │ 1.32s / 1.26s │
└────────┴───────────────┴───────────────┘

[SubhamSinghal]

@jayzhan211 Thanks for review, addressed latest feedback in 2c76116