Copilot commented on code in PR #25534:
URL: https://github.com/apache/datafusion/pull/25534#discussion_r4056555589
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -1384,9 +1368,139 @@ impl PartitionedTopK {
self.row_converter
.append(&mut self.scratch_rows, &ob_arrays)?;
- // 4. Per-partition: take the sub-batch, walk indices, dispatch
- // qualifying rows into the partition's heap.
- let k = self.k;
+ Ok(groups)
+ }
+
+ /// Take back the drained scratch map from [`Self::encode_and_group`]
+ /// (capacity retained) and record the batch's heap replacements.
+ fn finish_batch(&mut self, groups: HashMap<Vec<u8>, Vec<u32>>,
replacements: usize) {
+ self.partition_groups = groups;
+ if replacements > 0 {
+ self.metrics.row_replacements.add(replacements);
+ }
+ }
+
+ /// Emit every partition in partition-key order as a stream of coalesced
+ /// `RecordBatch`es. `emit_partition` pushes one partition's rows, already
+ /// in ORDER BY order, into the coalescer.
+ fn emit<S>(
+ self,
+ states: HashMap<Vec<u8>, S>,
+ mut emit_partition: impl FnMut(S, &mut BatchCoalescer,
&BaselineMetrics) -> Result<()>,
+ ) -> Result<SendableRecordBatchStream> {
+ let baseline = &self.metrics.baseline;
+ let _timer = baseline.elapsed_compute().timer();
Review Comment:
This helper now keeps `self.reservation` alive while it sorts states and
materializes all output batches. The previous per-operator `emit`
implementations explicitly dropped `reservation` before doing this work, so the
retained-state bytes were released from the memory pool during emission;
keeping them charged here can unnecessarily exhaust a tight pool and block
otherwise valid concurrent reservations. Drop the reservation before starting
the emission work.
--
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]
