Copilot commented on code in PR #25406:
URL: https://github.com/apache/datafusion/pull/25406#discussion_r4035916762
##########
datafusion/physical-plan/src/sorts/builder.rs:
##########
@@ -157,34 +157,7 @@ impl BatchBuilder {
// Remove consumed indices, keeping any remaining for the next call.
self.indices.drain(..rows_to_emit);
- // Only clean up fully-consumed batches when all indices are drained,
- // because remaining indices may still reference earlier batches.
- // In the overflow/partial-emit case this may retain some extra memory
- // across a few drain polls, but avoids costly index scanning on the
- // hot path. The retention is bounded and short-lived since leftover
- // rows are drained over subsequent polls.
- if self.indices.is_empty() {
- // New cursors are only created once the previous cursor for the
stream
- // is finished. This means all remaining rows from all but the
last batch
- // for each stream have been yielded to the newly created record
batch
- //
- // We can therefore drop all but the last batch for each stream
- let mut batch_idx = 0;
- let mut retained = 0;
- self.batches.retain(|(stream_idx, batch)| {
- let stream_cursor = &mut self.cursors[*stream_idx];
- let retain = stream_cursor.batch_idx == batch_idx;
- batch_idx += 1;
-
- if retain {
- stream_cursor.batch_idx = retained;
- retained += 1;
- } else {
- self.batches_mem_used -=
get_record_batch_memory_size(batch);
- }
- retain
- });
- }
+ self.retain_live_batches();
Review Comment:
This runs the mark/remap path for every emitted batch, including the normal
case where `drain` emptied `indices`. That adds three vector
allocations/initializations and a full cursor scan to the sort output hot path,
although early release is only needed for partial emits. Please preserve the
previous empty-index cleanup as a fast path and use `retain_live_batches` only
when pending indices remain.
--
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]
