jayzhan211 commented on code in PR #26930:
URL: https://github.com/apache/datafusion/pull/26930#discussion_r4493360308
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -998,21 +986,27 @@ impl ExternalSorter {
shared_input: Option<Arc<MemoryReservation>>,
) -> Result<SendableRecordBatchStream> {
let schema = batch.schema();
- let stream =
- futures::stream::once(self.sort_run(batch, reservation,
shared_input)?)
- .try_flatten();
+ let stream = futures::stream::once(self.sort_run(
+ batch,
+ reservation,
+ shared_input,
+ false,
+ )?)
+ .try_flatten();
Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
}
/// Sorts a batch and reserves its complete output before returning a
stream.
- /// The stream releases reservations as its consumer takes each output
batch.
+ /// Reservations release as output is read; parallel preparation reads one
+ /// batch ahead to return spill workspace before the merge builds its
cursor.
Review Comment:
Optional cleanup: the old `with_lookahead` doc said a run can hold a
prefetched batch that's no longer reserved. That's worth keeping, since it's
the one place the accounting runs behind the data.
```suggestion
/// Reservations release as output is read; parallel preparation reads
one
/// batch ahead to return spill workspace before the merge builds its
cursor,
/// so each run can hold one prefetched batch whose reservation is
already
/// released.
```
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -998,21 +986,27 @@ impl ExternalSorter {
shared_input: Option<Arc<MemoryReservation>>,
) -> Result<SendableRecordBatchStream> {
let schema = batch.schema();
- let stream =
- futures::stream::once(self.sort_run(batch, reservation,
shared_input)?)
- .try_flatten();
+ let stream = futures::stream::once(self.sort_run(
+ batch,
+ reservation,
+ shared_input,
+ false,
+ )?)
+ .try_flatten();
Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream)))
}
/// Sorts a batch and reserves its complete output before returning a
stream.
- /// The stream releases reservations as its consumer takes each output
batch.
+ /// Reservations release as output is read; parallel preparation reads one
+ /// batch ahead to return spill workspace before the merge builds its
cursor.
/// If present, `shared_input` covers the runs' shared buffers and is held
/// until sorting finishes; `reservation` covers this run's sorting and
output.
fn sort_run(
&self,
batch: RecordBatch,
reservation: MemoryReservation,
shared_input: Option<Arc<MemoryReservation>>,
+ read_ahead: bool,
Review Comment:
Optional, and an alternative to my other suggestion: keep `with_lookahead`
and `buffered(2)` from main and delete only the `ReservationStream` branch. I
ran that: 200 deletions, no additions, no `read_ahead` parameter, and the same
tests and capacity-sweep results as this PR. You called the per-batch
bookkeeping premature on #26242 and the bot shows no change here, so is there a
cost `buffered(2)` still has that you want to remove? If not, I'd take the
smaller diff.
```diff
-use crate::stream::ReservationStream;
```
```diff
- 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);
- }
-
```
##########
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:
Optional cleanup: the `Option`s and `drop(batches.take())` can go. The loan
is already 0 bytes at exhaustion and the merge's reservation keeps the pool
alive, so pool bytes didn't change in any trace, and all `sorts::` tests pass
without them. The one difference is a single-run stream outliving its sorter
with workspace still retained, which only tests do: `sort()` calls
`release_unused` first and the spill path keeps the sorter alive.
```suggestion
let mut 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)?;
let mut batches = sorted_batches
.into_iter()
.inspect(move |batch| {
let size = get_record_batch_memory_size(batch);
let borrowed = size.min(workspace.size());
workspace.shrink(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 batch = batches.next()?;
if read_ahead {
batches.peek();
}
Some(Ok(batch))
});
```
--
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]
