chore: Simplifying SortPreservingMergeStream to use generators instead of state machine#23407
Open
rluvaton wants to merge 21 commits into
Open
chore: Simplifying SortPreservingMergeStream to use generators instead of state machine#23407rluvaton wants to merge 21 commits into
SortPreservingMergeStream to use generators instead of state machine#23407rluvaton wants to merge 21 commits into
Conversation
Member
Author
|
run benchmarks |
Member
Author
|
run benchmark sort tpch10 |
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
This comment was marked as outdated.
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch10 — base (merge-base)
tpch10 — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagesort — base (merge-base)
sort — branch
File an issue against this benchmark runner |
Member
Author
|
run benchmark sort tpch10 |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing use-yield (18262df) to b790763 (merge-base) diff using: sort File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpch10 — base (merge-base)
tpch10 — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagesort — base (merge-base)
sort — branch
File an issue against this benchmark runner |
Member
Author
Damn I would never have caught that |
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagesort — base (merge-base)
sort — branch
File an issue against this benchmark runner |
rluvaton
commented
Jul 9, 2026
Comment on lines
+310
to
361
| fn initialize_all_partitions( | ||
| &mut self, | ||
| uninitiated_partitions: &mut Vec<usize>, | ||
| cx: &mut Context, | ||
| ) -> Poll<Result<()>> { | ||
| // Once all partitions have set their corresponding cursors for the loser tree, | ||
| // we skip the following block. Until then, this function may be called multiple | ||
| // times and can return Poll::Pending if any partition returns Poll::Pending. | ||
|
|
||
| assert_eq!( | ||
| self.loser_tree.len(), | ||
| 0, | ||
| "loser tree must be empty when initializing" | ||
| ); | ||
|
|
||
| // Manual indexing since we're iterating over the vector and shrinking it in the loop | ||
| let mut idx = 0; | ||
| while idx < uninitiated_partitions.len() { | ||
| let partition_idx = uninitiated_partitions[idx]; | ||
| match self.maybe_poll_stream(cx, partition_idx) { | ||
| Poll::Ready(Err(e)) => { | ||
| return Poll::Ready(Err(e)); | ||
| } | ||
| Poll::Pending => { | ||
| // The polled stream is pending which means we're already set up to | ||
| // be woken when necessary | ||
| // Try the next stream | ||
| idx += 1; | ||
| } | ||
| _ => { | ||
| // The polled stream is ready | ||
| // Remove it from uninitiated_partitions | ||
| // Don't bump idx here, since a new element will have taken its | ||
| // place which we'll try in the next loop iteration | ||
| // swap_remove will change the partition poll order, but that shouldn't | ||
| // make a difference since we're waiting for all streams to be ready. | ||
| uninitiated_partitions.swap_remove(idx); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| return Poll::Ready(self.emit_in_progress_batch().transpose()); | ||
| if uninitiated_partitions.is_empty() { | ||
| Poll::Ready(Ok(())) | ||
| } else { | ||
| // There are still uninitiated partitions so return pending. | ||
| // We only get here if we've polled all uninitiated streams and at least one of them | ||
| // returned pending itself. That means we will be woken as soon as one of the | ||
| // streams would like to be polled again. | ||
| // There is no need to reschedule ourselves eagerly. | ||
| Poll::Pending | ||
| } | ||
| } |
Member
Author
There was a problem hiding this comment.
Extracted this to a separate pr so the diff will be smaller:
rluvaton
changed the title
yield like machinery SortPreservingMergeStream to use generators instead of state machine
4 tasks
# Conflicts: # datafusion/physical-plan/src/sorts/merge.rs
pepijnve
reviewed
Jul 9, 2026
| let streams = | ||
| FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); | ||
| return Ok(Box::pin(SortPreservingMergeStream::new( | ||
| return Ok(SortPreservingMergeStream::new( |
Contributor
There was a problem hiding this comment.
Since SortPreservingMergeStream no longer has a Stream implementation and you need to turn it into a stream, we might want to give it a different name.
Would it be possible to return the impl Stream directly from new to avoid the somewhat odd construct?
Member
Author
There was a problem hiding this comment.
Since
SortPreservingMergeStreamis no longer has aStreamimplementation and you need to turn it into a stream, we might want to give it a different name.Would it be possible to return the
impl Streamdirectly fromnewto avoid the somewhat odd construct?
I wanted to reduce the amount of changes in this PR as possible so it will be easier to review and see the actual difference between generators and current implementation.
this can be a followup PR
pepijnve
reviewed
Jul 9, 2026
| datafusion-spark = { path = "datafusion/spark", version = "54.0.0" } | ||
| datafusion-sql = { path = "datafusion/sql", version = "54.0.0" } | ||
| datafusion-substrait = { path = "datafusion/substrait", version = "54.0.0" } | ||
| genawaiter = { version = "0.99.1", default-features = false } |
Contributor
There was a problem hiding this comment.
I've noticed some reluctance in other PRs to pull in additional dependencies since the dependency graph is quite large already. Not sure what the policy is to accept or reject new ones.
Member
Author
There was a problem hiding this comment.
The good thing is that this does not have any dependencies at all
- no runtime dependencies
- no compile-time dependencies either, with
default-features = falseFrom
genawaiterREADME
The only dependencies that we are using with the enabled futures feature already exists
pepijnve
reviewed
Jul 9, 2026
| let cloned_metrics = self.metrics.clone(); | ||
|
|
||
| let mut this = self; | ||
| let stream = Gen::new(|co| async move { |
Contributor
There was a problem hiding this comment.
I tried out the try_stream! variant of this code locally to get an idea of what it would look like. The points made in the PR description are certainly valid. I did find the try_stream! variant a bit more intuitive to read though since the error handling is a lot more natural. The construct here seems a bit clunky where the Co essentially always receives an Ok in run and the Err case gets handled here.
Contributor
There was a problem hiding this comment.
Looking at the implementation of async-stream and comparing to sync::Gen, the async-stream implementation uses a thread local Cell while sync::Gen uses a mutex. It's only one mutex acquire per batch of course, but this does incur some cost compared to the code it's replacing. Is this a concern if the intention would be to apply this pattern more generally?
Member
Author
There was a problem hiding this comment.
async-stream have that because it is a macro you cant move the yield sender between functions.
the async-fn-stream crate that aim to have the same behavior and logic as async-stream but in fact store SmallVec for multiple emits per poll so this is not an option.
the Mutex in async-stream is just to make it safe in rust
there shouldn't be any contention as the mutex is just for step lock (it does still can add an overhead of course).
there is an added cost which is why I run the benchmark and it did not show any regressions.
Contributor
There was a problem hiding this comment.
All of the downsides of async-stream that you list in the Description are accurate, but I would still recommend it, both for the error handling differences, as well as because it does not introduce the requirement that the stream be Send.
Contributor
There was a problem hiding this comment.
A possible alternative to both could be a DIY solution for DataFusion. I tried to sketch out an alternative loosely based on the various existing implementations https://github.com/pepijnve/datafusion/blob/4b4954cd50a26101d4e40d25c13d9068b00b8d2c/datafusion/execution/src/async_stream.rs There's probably still room for improvement here, the main thing I wanted to get across is that this looks much nicer imo at the call site in SortPreservingMergeStream (see https://github.com/pepijnve/datafusion/blob/4b4954cd50a26101d4e40d25c13d9068b00b8d2c/datafusion/physical-plan/src/sorts/merge.rs#L214)
The main downside I see is that without a proc macro it's really hard to make this entirely abuse proof. I started with the thread local approach used by the async-stream macro, but that's unsound when the emitter is exposed. I went for the Arc<Mutex> then rather than the SmallVec that async-fn-stream uses.
Doesn't have to be this implementation, but perhaps something like this could be a "have your cake and eat it" solution?
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
4 participants
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?
Rationale for this change
I'm trying to make some Sort optimization but it is hard to get stuff merged where the code is already complex even if adding a simple optimization (like when only 1 stream is left just return those batches).
This makes the code simpler as you can now read it in straight line flow control, while not sacrificing performance
I did not use
RecordBatchReceiverStreamBuildersince I want pull based and not push based streamWhat changes are included in this PR?
Added
genawaiterdependency, replace the SortPersevingMergeStream main state machine loop with generator yieldAre these changes tested?
existing tests
Are there any user-facing changes?
no
Why not using
async-streamcrate from tokio?Why not using
async-streamcrate from tokio (that we already have transitive dependency on) or other crates that are macro based.Because couple of reasons:
cargo fmtdoes not format macros body, so you manually need to format and validate that the style is keptyieldis not a function calli. When you read the code, you cant go to definition of the
yieldkeyword since it is not a functionii. Cant put debugger on that keyword (from what I remember)
iii. there is some hidden code that you need to
yieldfrom the macro body (so you cant extract functions withyieldfor example)