[Status Update] Simplifing Streams to be more textbook-like and have less state while keeping the same perf

[rluvaton]

If you wanna look at before and after example look at:

• chore: refactor SortMergeJoin bitwise stream to generators and simplify to be textbook like as possible #23761

Status:

• SortPreservingMergeStream
  • chore: Simplifying SortPreservingMergeStream to use generators instead of state machine #23407 - Moving to generators
  • chore: simplify SortPreservingMergeStream to be as textbook-like as possible #23702 - Simplifying the code
• SortMergeJoin
  • BitwiseSortMergeJoinStream chore: refactor SortMergeJoin bitwise stream to generators and simplify to be textbook like as possible #23761 - moving to generators and simplifying the code
  • MaterializingSortMergeJoinStream - chore: refactor MaterializingSortMergeJoinStream into generators and simplify code to be textbook like as possible #23976
• Aggregate
  • OrderedPartialAggregateStream - chore(ordered-partial-aggregate): move OrderedPartialAggregateStream to generators for readability #23951 - moving to generators, the code is already pretty simple
  • OrderedFinalAggregateStream - chore(OrderedFinalAggregateStream): refactor to async generator #24008

Possible Candidates:

• Aggregate
  • SingleHashAggregateStream
  • OrderedFinalAggregateStream
  • PartialHashAggregateStream
  • PartialReduceHashAggregateStream
  • FinalHashAggregateStream
• Join
  • SortMergeJoin
    • MaterializingSortMergeJoinStream

Background

So, I start seeing that some code that on the surface should be dead simple in textbook form in reality the main entry point and the whole state handling is really complex

For example SortMergeJoin, on the surface the algorithm is simple

Algorithm

Taken from Sort-Merge Joins

But due to all the following reasons:

1. child return pending and have to keep state between calls
2. programming language as opposed to pseudo code
3. Flow is not linear as pseudo code since you can return only 1 value at a time
4. needing to handle spill
5. performance optimization

the actual implementation is really complex which means that doing any sort of PR there is hard to review or to understand

Tradeoffs

So I started with moving some stuff to async generators that solve some of the problems while adding others.
problem that it is solving:

1. Child return pending and have to keep state between calls
2. Flow is now linear since we can yield from the code and have the state and code in mind kept
3. Fewer lines - since less code needed for holding and managing the state
4. In some cases improve performance since going back to where you was in the state by making the function idompotent could be expensive

Problems that it is creating:

1. Timing is more annoying, you should pause the elapsed_compute timer between
   i. yields - since you shouldn't count the time that the parent has done work between calling you
   ii. awaits on child streams - because you don't want to count the child time
   iii. awaits on stuff that you do async like reading a spill file - because even though the on the surface this is your work, you might overcount the elapsed time and the spill reading finish earlier but you did not woke up (up for debate)
2. Reserving memory is less intuitive since the code look linear but you hand off control between awaits so the data that you hold should be reserved for
3. You need to wrap the async generator with ObservedStream to track end time and output batches/rows/etc
4. can introduce performance problems since we are adding async machinery if not used with caution

Alternative solutions

We already have RecordBatchReceiverStreamBuilder but the flow of data is different - i.e. it is push based and not pull based which means:

1. more memory is being in the buffer (have at least 1 item pending
2. You produce data even if not needed
3. In case you are spawning blocked task to do all your work you can get to something like what was fixed here: Remove waits from blocking threads reading spill files. #15654

Initial Progress

I started the first PR in

• chore: Simplifying SortPreservingMergeStream to use generators instead of state machine #23407

and @pepijnve kindly extracted a trimmed down version for the async generators from genawaiter crate that I previously used in that PR and tokio async-stream crate in:

• Add minimal genarator-like stream implementation #23530

All the discussion for why we have our own implementation and also why not having macro implementation can be found in:

• chore: Simplifying SortPreservingMergeStream to use generators instead of state machine #23407 (comment) - for why not using external crate or macro
• chore: Simplifying SortPreservingMergeStream to use generators instead of state machine #23407 - why not using macro in the pr description

the first rewrite PR allowed for having the code rewritten to match simpler form:

• https://github.com/apache/datafusion/pull/23702e

[buraksenn]

I would like to try one of them to learn more about streams while trying to simplify it. Are you open to help on this?

[rluvaton]

@buraksenn sure, try to pick something simple but complex enough that the added cost of async won't be too much

[buraksenn]

Thanks. I've looked into this a bit and will start working on SingleHashAggregateStream if you've no objections. Otherwise I can pick up what you advise too.