rluvaton commented on code in PR #24016:
URL: https://github.com/apache/datafusion/pull/24016#discussion_r3685873433
##########
datafusion/physical-plan/src/aggregates/single_stream.rs:
##########
@@ -164,216 +103,73 @@ impl SingleHashAggregateStream {
input,
baseline_metrics,
reservation,
- state: Some(SingleHashAggregateState::ReadingInput { hash_table }),
+ hash_table,
})
}
- /// Moves the aggregate hash table's inner state to `Outputting`.
- ///
- /// The caller guarantees that input is fully consumed, so this function
can
- /// eagerly release the input stream.
- fn start_output(
- &mut self,
- hash_table: &mut AggregateHashTable<SingleMarker>,
- ) -> Result<()> {
- let input_schema = self.input.schema();
- self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
- hash_table.start_output()
+ pub(crate) fn into_stream(self) -> SendableRecordBatchStream {
+ let schema = Arc::clone(&self.schema);
+ let baseline_metrics = self.baseline_metrics.clone();
+ let stream =
+ Box::pin(RecordBatchStreamAdapter::new(schema,
self.create_stream()));
+
+ Box::pin(ObservedStream::new(stream, baseline_metrics, None))
}
- /// Handle ReadingInput state - aggregate input batches into the hash
table.
- ///
- /// See comments at `poll_next()` for details.
+ /// State transitions are implemented using the generator pattern; see the
+ /// comments in [`async_try_stream`].
///
- /// Returns the next operator state with control flow decision.
- fn handle_reading_input(
- &mut self,
- cx: &mut Context<'_>,
- mut original_state: SingleHashAggregateState,
- ) -> SingleHashAggregateStateTransition {
- debug_assert!(matches!(
- &original_state,
- SingleHashAggregateState::ReadingInput { .. }
- ));
- debug_assert!(original_state.hash_table().is_building());
-
- match self.input.poll_next_unpin(cx) {
- Poll::Pending => ControlFlow::Break((Poll::Pending,
original_state)),
- // Get a new input batch, aggregate it in the hash table
- Poll::Ready(Some(Ok(batch))) => {
- let elapsed_compute =
self.baseline_metrics.elapsed_compute().clone();
- let timer = elapsed_compute.timer();
- let result =
original_state.hash_table_mut().aggregate_batch(&batch);
- timer.done();
-
- if let Err(e) = result {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- original_state,
- ));
- }
-
- if let Err(e) = self
- .reservation
- .try_resize(original_state.hash_table().memory_size())
+ /// Conceptually: ReadingInput -> ProducingOutput -> Done.
+ fn create_stream(self) -> impl Stream<Item = Result<RecordBatch>> {
+ async_try_stream(|mut emitter| async move {
+ let Self {
+ mut input,
+ baseline_metrics,
+ reservation,
+ mut hash_table,
+ ..
+ } = self;
+ let elapsed_compute = baseline_metrics.elapsed_compute().clone();
+
+ debug_assert!(hash_table.is_building());
+ while let Some(batch) = input.next().await.transpose()? {
{
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- original_state,
- ));
- }
-
- ControlFlow::Continue(original_state)
- }
- Poll::Ready(Some(Err(e))) => {
- ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state))
- }
- // Input ends, move to output state
- Poll::Ready(None) => {
- let elapsed_compute =
self.baseline_metrics.elapsed_compute().clone();
- let timer = elapsed_compute.timer();
- let result =
self.start_output(original_state.hash_table_mut());
- timer.done();
-
- match result {
- Ok(()) => {
-
ControlFlow::Continue(original_state.into_producing_output())
- }
- Err(e) => {
- ControlFlow::Break((Poll::Ready(Some(Err(e))),
original_state))
- }
+ let _timer = elapsed_compute.timer();
+ hash_table.aggregate_batch(&batch)?;
}
+ reservation.try_resize(hash_table.memory_size())?;
}
- }
- }
- /// Handle ProducingOutput state - emit final aggregate value batches.
- ///
- /// See comments at `poll_next()` for details.
- ///
- /// Returns the next operator state with control flow decision.
- fn handle_producing_output(
- &mut self,
- mut original_state: SingleHashAggregateState,
- ) -> SingleHashAggregateStateTransition {
- debug_assert!(matches!(
- &original_state,
- SingleHashAggregateState::ProducingOutput { .. }
- ));
- debug_assert!(!original_state.hash_table().is_building());
-
- let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
- let timer = elapsed_compute.timer();
- let result = original_state.hash_table_mut().next_output_batch();
- timer.done();
-
- match result {
- Ok(Some(batch)) => {
- if let Err(e) = self
- .reservation
- .try_resize(original_state.hash_table().memory_size())
- {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- original_state,
- ));
- }
- debug_assert!(batch.num_rows() > 0);
- let next_state = if original_state.hash_table().is_done() {
- original_state.into_done()
- } else {
- original_state
- };
+ {
+ let _timer = elapsed_compute.timer();
- ControlFlow::Break((
-
Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))),
- next_state,
- ))
+ // Input is exhausted. Release the upstream pipeline before
draining output.
+ drop(input);
+ hash_table.start_output()?;
}
- Ok(None) => {
- let _ = self.reservation.try_resize(0);
- ControlFlow::Continue(original_state.into_done())
- }
- Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))),
original_state)),
- }
- }
-}
-impl Stream for SingleHashAggregateStream {
- type Item = Result<RecordBatch>;
+ 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)
+ })
Review Comment:
You can simplify it like this while keeping the original function structure:
```rust
/// 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(())
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
