geoffreyclaude commented on code in PR #26930:
URL: https://github.com/apache/datafusion/pull/26930#discussion_r4493700645
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -1045,31 +1039,39 @@ impl ExternalSorter {
.iter()
.map(get_record_batch_memory_size)
.sum();
- let mut workspace =
+ let workspace =
merge_pool.borrow(total_sorted_size.saturating_sub(reservation.size()));
reservation
.try_resize(total_sorted_size - workspace.size())
.map_err(Self::err_with_oom_context)?;
- if workspace.size() == 0 {
- return Ok(Box::pin(ReservationStream::new(
- Arc::clone(&schema),
- Box::pin(RecordBatchStreamAdapter::new(
- Arc::clone(&schema),
-
futures::stream::iter(sorted_batches.into_iter().map(Ok)),
- )),
- reservation,
- )) as SendableRecordBatchStream);
- }
-
- // Return borrowed workspace first so the merge's cursors can reuse
- // it immediately. Both reservations also release on stream drop.
- let batches = sorted_batches.into_iter().map(move |batch| {
- let size = get_record_batch_memory_size(&batch);
- let borrowed = size.min(workspace.size());
- workspace.shrink(borrowed);
- reservation.shrink(size - borrowed);
- Ok(batch)
+ // An empty loan must not keep unused merge workspace alive.
+ let mut workspace = (workspace.size() != 0).then_some(workspace);
+ let mut batches = Some(
+ sorted_batches
+ .into_iter()
+ .inspect(move |batch| {
+ let size = get_record_batch_memory_size(batch);
+ let borrowed = workspace.as_mut().map_or(0,
|workspace| {
+ let borrowed = size.min(workspace.size());
+ workspace.shrink(borrowed);
+ borrowed
+ });
+ reservation.shrink(size - borrowed);
+ })
+ .peekable(),
+ );
+ // Output is already in memory. Peeking releases the next batch's
+ // reservation before the merge allocates its current cursor.
+ let batches = std::iter::from_fn(move || {
+ let iter = batches.as_mut()?;
+ let batch = iter.next()?;
+ if read_ahead && iter.peek().is_none() {
+ // Release the exhausted iterator's workspace loan now, so
it
+ // cannot keep unused merge memory reserved until another
poll.
+ drop(batches.take());
+ }
+ Some(Ok(batch))
});
Review Comment:
Your smaller alternative takes care of this too: the custom iterator, both
Options and the manual cleanup are gone.
--
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]
