chore: Simplifying SortPreservingMergeStream to use generators instead of state machine#23407
Merged
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 |
4 tasks
# Conflicts: # datafusion/physical-plan/src/sorts/merge.rs
rluvaton
marked this pull request as ready for review
July 9, 2026 16:34
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
pepijnve
reviewed
Jul 9, 2026
Member
Author
|
@mbutrovich I think implementing SMJ with generators will simplify the code greatly and may even improve perf |
Member
Author
|
@Jefffrey can you please review? |
alamb
reviewed
Jul 13, 2026
alamb
left a comment
alamb
left a comment
Contributor
There was a problem hiding this comment.
Thanks for this proposal @rluvaton and @pepijnve and @stuhood
I agree with @stuhood and @pepijnve that this change while it may look simpler is subject to some opinions about if wrapping the complexity in another crate which is not widely used
I also agree with @rluvaton that the SortPreservingMerge code is about as complicated as it can get and still be undersrtandable without some refactoring -- so to make progress we will have to refactor it (basically what @Rachelint @2010YOUY01 and I are doing to the aggregate code in #22710)
One thing I think helps is separting the stream control flow (aka the poll_next implementations) from the actual logic (aka messing around with data).
A major win of this PR for simplicity I think is that it moves the logic of updating the Loser tree implementation out of the control logic. Rather than using a new crate / control flow pattern, we could instead try moving the loser tree implementation into its own struct (encapsulate the logic and make the boundaries of the different responsibilities clearer, and make the control flow easier to read)
Contributor
Tooting my own horn a bit, but could you take a quick peek at #23530 |
pull Bot
pushed a commit
to TCeason/arrow-datafusion
that referenced
this pull request
Jul 13, 2026
## Which issue does this PR close? None, related to PR apache#23407 ## Rationale for this change Manually writing `Stream` implementations can be quite tedious since it requires implementing the state machine of the stream yourself. This often results in hard to read code. The same problem exists with manual `Future` implementations and `async` was added to the Rust language to mitigate exactly this problem. Unfortunately there's no language level support yet to help with writing streams/generators. There are quite a few projects that attempt to fill this gap: - Tokio's [async_stream](https://docs.rs/async-stream/latest/async_stream/) provides a proc macro that recognises a `yield` keyword. This is a good implementation, but proc macros don't play nice with code formatting, increase compile time, and in this particular case prevent decomposition into smaller functions. - [async_fn_stream](https://docs.rs/async-fn-stream/latest/async_fn_stream/) provides an implementation of the same concept, but without the proc macro. Unfortunately this implementation chooses a heavier mechanism to communicate generated values back to the stream via a `SmallVec`. - [genawaiter](https://docs.rs/genawaiter/latest/genawaiter/sync/index.html) is a more general purpose generator library, but this can also be used to implement `Stream`s. This project does miss some of the ergonomics provided by the other two in the form of `try_` variants that help in implementing fallible streams. Since none of these variants seems like the ideal candidate, the best option might be to have a custom implementation of the concept tailored to the needs of the DataFusion project. This PR provides an initial draft of exactly that. ## What changes are included in this PR? - Adds `async_stream` and `async_try_stream` functions that create Stream implementations based on an async generator function. The initial implementation was inspired by the macro expansion produced by `async_stream`. The code was then adapted further taking inspiration from the two other libraries. Specifically, the `Emitter` terminology was taken from `async_fn_stream` and the `Arc<Mutex<Option<T>>>` value transfer mechanism was taken from `genawaiter`. `async_stream` uses a very light weight thread-local storage based mechanism to handle value transfer, but this felt too risky to use when the emitter is exposed. It makes sense in the Tokio implementation since the proc macro can manage control flow in a stricter way. I wasn't entirely sure if taking some inspiration from has code licensing implications. I applied ASLv2 for now on the added code. ## Are these changes tested? Test code was mainly adapted from the tokio implementation. ## Are there any user-facing changes? Yes, new public function `async_stream` and `async_try_stream` --------- Co-authored-by: Raz Luvaton <16746759+rluvaton@users.noreply.github.com>
Member
Author
|
@2010YOUY01 @Rachelint you might want to consider using the generator to simplify the aggregate rewrite |
Contributor
Yes this makes sense, I'll give it a try later.
I have similar opinions to @alamb on how to reduce complexity in many existing operators, I’ll try to explain them a little differently. To make an operator implementation easy to read, its entry point should look like pseudocode: There are two major gaps between this ideal and the current implementation:
Point 2 causes an issue: in order to read the code, at the same time, we have to deal with the complexity for both control flow and logic details around low-level primitives. This causes huge cognitive load. A better approach would be to encapsulate these details in a deeper module, and implement the core logic around that deep module. The goal is to make the entry-point code look as much like pseudocode as possible—e.g. the I think the second issue is the larger source of complexity, but it's often overlooked. |
Member
Author
|
@2010YOUY01 My goal for simplifying here and in other places is to make the code looks as similar as textbook psuedo code as possible and abstracting out all the implementation detail of rust and arrow. for example merge sort is pretty basic in the idea, and as such the code should look basic at first glance. @2010YOUY01 Are you opposed to having this PR? |
Contributor
|
Thanks for moving this code forward everyone |
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
7 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)