chore(SingleHashAggregateStream): refactor to async generator implementation - #24016
Open
buraksenn wants to merge 2 commits into
Open
chore(SingleHashAggregateStream): refactor to async generator implementation#24016buraksenn wants to merge 2 commits into
SingleHashAggregateStream): refactor to async generator implementation#24016buraksenn wants to merge 2 commits into
Conversation
buraksenn
changed the title
SingleHashAggregateStream): refactor to async generator implementation
Contributor
Author
rluvaton
reviewed
Jul 30, 2026
Member
There was a problem hiding this comment.
I would keep the code in their original functions but just change to async and remove state
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #24016 +/- ##
========================================
Coverage 80.84% 80.85%
========================================
Files 1096 1099 +3
Lines 373832 374273 +441
Branches 373832 374273 +441
========================================
+ Hits 302240 302623 +383
- Misses 53563 53600 +37
- Partials 18029 18050 +21 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
rluvaton
reviewed
Jul 30, 2026
Comment on lines
+151
to
+173
| debug_assert!(!hash_table.is_building()); | ||
| loop { | ||
| let next_batch = { | ||
| let _timer = elapsed_compute.timer(); | ||
| hash_table.next_output_batch()? | ||
| }; | ||
| let Some(batch) = next_batch else { | ||
| return Ok(()); | ||
| }; | ||
|
|
||
| /// Entry point for the single hash aggregate state machine. | ||
| /// | ||
| /// See comments in [`SingleHashAggregateStream`] for high-level ideas. | ||
| /// | ||
| /// State transition graph: | ||
| /// | ||
| /// ```text | ||
| /// (start) | ||
| /// -> ReadingInput | ||
| /// The stream starts by polling raw input rows and aggregating those | ||
| /// rows into the single-stage hash table. | ||
| /// | ||
| /// ReadingInput | ||
| /// -> ReadingInput | ||
| /// Aggregate one raw input batch, update the inner aggregate hash | ||
| /// table, and continue with the next input batch. | ||
| /// | ||
| /// -> ProducingOutput | ||
| /// Input was exhausted. Move to the next state to start outputting | ||
| /// final aggregate values. | ||
| /// | ||
| /// ProducingOutput | ||
| /// -> ProducingOutput | ||
| /// One final output batch was yielded; repeat to continue producing | ||
| /// output incrementally. | ||
| /// | ||
| /// -> Done | ||
| /// All final output was emitted. | ||
| /// | ||
| /// Done | ||
| /// -> (end) | ||
| /// ``` | ||
| fn poll_next( | ||
| mut self: std::pin::Pin<&mut Self>, | ||
| cx: &mut Context<'_>, | ||
| ) -> Poll<Option<Self::Item>> { | ||
| loop { | ||
| let cur_state = self | ||
| .state | ||
| .take() | ||
| .expect("SingleHashAggregateStream state should not be None"); | ||
| reservation.try_resize(hash_table.memory_size())?; | ||
| debug_assert!(batch.num_rows() > 0); | ||
|
|
||
| let next_state = match cur_state { | ||
| state @ SingleHashAggregateState::ReadingInput { .. } => { | ||
| self.handle_reading_input(cx, state) | ||
| if hash_table.is_done() { | ||
| drop(hash_table); | ||
| reservation.free(); | ||
| emitter.emit(batch).await; | ||
| return Ok(()); | ||
| } | ||
| state @ SingleHashAggregateState::ProducingOutput { .. } => { | ||
| self.handle_producing_output(state) | ||
| } | ||
| state @ SingleHashAggregateState::Done => { | ||
| let _ = self.reservation.try_resize(0); | ||
| self.state = Some(state); | ||
| return Poll::Ready(None); | ||
| } | ||
| }; | ||
|
|
||
| match next_state { | ||
| ControlFlow::Continue(next_state) => { | ||
| self.state = Some(next_state); | ||
| continue; | ||
| } | ||
| ControlFlow::Break((poll, next_state)) => { | ||
| self.state = Some(next_state); | ||
| return poll; | ||
| } | ||
| emitter.emit(batch).await; | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl RecordBatchStream for SingleHashAggregateStream { | ||
| fn schema(&self) -> SchemaRef { | ||
| Arc::clone(&self.schema) | ||
| }) |
Member
There was a problem hiding this comment.
You can simplify it like this while keeping the original function structure:
/// Handle ProducingOutput state - emit final aggregate value batches.
///
/// See comments at [`Self::create_stream`] for details.
async fn handle_producing_output(
&mut self,
mut hash_table: AggregateHashTable<SingleMarker>,
mut emitter: TryEmitter<RecordBatch, DataFusionError>,
) -> Result<()> {
debug_assert!(!hash_table.is_building());
let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
let mut timer = elapsed_compute.timer();
while let Some(batch) = hash_table.next_output_batch()? {
self
.reservation
.try_resize(hash_table.memory_size())?;
debug_assert!(batch.num_rows() > 0);
if hash_table.is_done() {
drop(hash_table);
timer.done();
emitter.emit(batch).await;
return Ok(());
}
timer.done();
emitter.emit(batch).await;
timer = elapsed_compute.timer();
}
Ok(())
}
Contributor
Author
There was a problem hiding this comment.
I was trying to get rid of option in hash_table: Option<AggregateHashTable<SingleMarker>> but when I did that functional structure and these needed to be changed. Reintroduced Option and this functional structure again. Thanks for the review
13 tasks
Contributor
|
Could we delay those cleanup PRs after #22710 is functionally complete? Now we're in a inconsistent state, and I think it's better to fully remove the legacy implementation sooner, and next we can continue cleaning things up. I estimate it's ~3 PRs away from that, I'll update in the EPIC issue when it's ready. |
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?
Does not close but part of #23974
Rationale for this change
Simplify
SingleHashAggregateStreamby removing state required only for manual polling.What changes are included in this PR?
Convert the stream to an async generator and remove the explicit state machine.
Are these changes tested?
Yes existing tests with additional metric assert
Are there any user-facing changes?
No