Repository navigation
refactor: remove redundant sort-run reservation stream - #26930
Merged
jayzhan211 merged 2 commits intoNov 7, 2026
Merged
Conversation
geoffreyclaude
force-pushed
the
codex/sort-iterator-lookahead
branch
2 times, most recently
from
November 2, 2026 10:33
b04cb7e to
2be04ec
Compare
Contributor
Author
|
run benchmark sort_tpch sort_tpch10 |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/sort-iterator-lookahead (2be04ec) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpch10Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/sort-iterator-lookahead (2be04ec) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/sort-iterator-lookahead (2be04ec) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpchCPU Details (lscpu)Details
Resource Usagesort_tpch — base (merge-base)
sort_tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/sort-iterator-lookahead (2be04ec) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpch10CPU Details (lscpu)Details
Resource Usagesort_tpch10 — base (merge-base)
sort_tpch10 — branch
File an issue against this benchmark runner |
geoffreyclaude
marked this pull request as ready for review
November 2, 2026 16:43
geoffreyclaude
marked this pull request as draft
November 2, 2026 16:47
geoffreyclaude
force-pushed
the
codex/sort-iterator-lookahead
branch
from
November 4, 2026 19:44
2be04ec to
dc31005
Compare
geoffreyclaude
force-pushed
the
codex/sort-iterator-lookahead
branch
from
November 4, 2026 19:59
dc31005 to
5001e40
Compare
geoffreyclaude
marked this pull request as ready for review
November 4, 2026 20:32
jayzhan211
approved these changes
Nov 6, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
@geoffreyclaude thanks for this cleanup! I compared reservation traces and tight-memory outcomes against main, and they match, so the inline comments are all optional.
Comment on lines
+1000
to
+1001
| /// Reservations release as output is read; parallel preparation reads one | ||
| /// batch ahead to return spill workspace before the merge builds its cursor. |
Contributor
There was a problem hiding this 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.
Suggested change
| /// Reservations release as output is read; parallel preparation reads one | |
| /// batch ahead to return spill workspace before the merge builds its cursor. | |
| /// 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. |
Contributor
Author
There was a problem hiding this comment.
Restored the original with_lookahead documentation, including this caveat.
| batch: RecordBatch, | ||
| reservation: MemoryReservation, | ||
| shared_input: Option<Arc<MemoryReservation>>, | ||
| read_ahead: bool, |
Contributor
There was a problem hiding this 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.
-use crate::stream::ReservationStream;- 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);
- }
-
Contributor
Author
There was a problem hiding this comment.
Agreed, the measurements don’t justify the extra code. Adopted your smaller alternative: the PR now removes 200 lines with no additions.
Comment on lines
1042
to
1075
| 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)) | ||
| }); |
Contributor
There was a problem hiding this comment.
Optional cleanup: the Options 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.
Suggested change
| 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)) | |
| }); | |
| 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)) | |
| }); |
Contributor
Author
There was a problem hiding this comment.
Your smaller alternative takes care of this too: the custom iterator, both Options and the manual cleanup are gone.
geoffreyclaude
changed the title
jayzhan211
approved these changes
Nov 7, 2026
jayzhan211
left a comment
jayzhan211
left a comment
Contributor
There was a problem hiding this comment.
Thanks @geoffreyclaude !
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Follow-up to the lookahead discussion in #26242.
Rationale for this change
Sort runs use two different paths to return memory as their output batches are consumed. The path that handles borrowed spill workspace also works when nothing is borrowed, making the separate
ReservationStreamwrapper redundant.What changes are included in this PR?
Use the existing workspace-aware path for all sort runs and remove the unused private wrapper and its dedicated tests. Each batch returns borrowed workspace first, then releases its ordinary reservation.
What is the testing strategy for this PR?
Existing tests check lookahead under tight memory and keep pending spill output reserved, including zero, partial and full workspace loans. Output and cancellation checks cover row integrity and reservation cleanup.
Are there any user-facing changes?
Memory handoff and output ordering are preserved.
Prepared with AI assistance.