jayzhan211 commented on code in PR #25468:
URL: https://github.com/apache/datafusion/pull/25468#discussion_r4052397492
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2314,26 +2361,60 @@ impl PartitionedTopKDenseRank {
let mut coalescer = BatchCoalescer::new(Arc::clone(&schema),
batch_size);
Review Comment:
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.
```diff
- 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));
}
```
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2314,26 +2361,60 @@ 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.
Review Comment:
```suggestion
// Chunk at `batch_size` so the operator emits the same batch sizes
// as before and no single `interleave` output exceeds `batch_size`.
```
##########
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<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 `dense_rank_contents_bytes_tracks_recompute` test
Review Comment:
```suggestion
/// adjust it;
`test_partitioned_topk_dense_rank_contents_bytes_tracks_recompute`
```
--
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]
