SubhamSinghal commented on code in PR #25468:
URL: https://github.com/apache/datafusion/pull/25468#discussion_r4052466562
##########
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:
addressed in
https://github.com/apache/datafusion/commit/e580f216bdfe07ddaa14131339223bed4f933e9f
--
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]
