diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index edf084ad328b..691670a13db2 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -26,9 +26,7 @@ //! See issue for details: use std::mem::size_of; -use std::ops::ControlFlow; use std::sync::Arc; -use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; @@ -57,7 +55,7 @@ use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::SpillManager; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; -use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; +use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; /// Hash aggregation is implemented in two stages: partial and final. This /// stream implements the partial stage. @@ -169,47 +167,10 @@ pub(crate) struct PartialHashAggregateStream { /// be empty. See struct comments for details. group_values_soft_limit: Option, - /// Tracks the high-level stream lifecycle. The hash table owns the lower-level - /// state for emitting output batches. - state: Option, + /// The hash table owns the lower-level state for emitting output batches. + hash_table: Option>, } -/// States for partial hash aggregation processing. -enum PartialHashAggregateState { - ReadingInput { - hash_table: AggregateHashTable, - }, - /// A fully materialized partial-state batch being emitted incrementally. - EmittingOnMemoryPressure { - hash_table: AggregateHashTable, - // After each incremental emitting step, the `remaining_groups` will be updated - // with batch slicing. - remaining_groups: RecordBatch, - }, - ProducingOutput { - hash_table: AggregateHashTable, - /// If `None`, partial skip was never triggered and this state will - /// finish in `Done`. If `Some`, partial skip has triggered and the - /// stream will move to `SkippingAggregation` after these accumulated - /// groups are emitted. - skip_hash_table: Option>, - }, - SkippingAggregation { - hash_table: AggregateHashTable, - }, - Done, - /// Sentinel state to use when returning error from any other states, because: - /// - It explicitly releases state-owned resources immediately - /// - More defensive against accidentally resuming execution after error - Error, -} - -type PartialHashAggregatePoll = Poll>>; -type PartialHashAggregateStateTransition = ControlFlow< - (PartialHashAggregatePoll, PartialHashAggregateState), - PartialHashAggregateState, ->; - /// Spill configuration and accumulated runs for final hash aggregation. /// /// Each spill event drains all currently buffered groups, sorts their intermediate @@ -395,6 +356,15 @@ impl FinalSpillContext { } } +#[derive(PartialEq)] +enum HandleInputResult { + ProcessNext, + ReachedLimit, + #[expect(clippy::upper_case_acronyms)] + OOM, + SwitchToSkipAggregation, +} + impl PartialHashAggregateStream { pub fn new( agg: &AggregateExec, @@ -457,26 +427,85 @@ impl PartialHashAggregateStream { reduction_factor, skip_aggregation_probe, group_values_soft_limit: agg.limit_options().map(|config| config.limit()), - state: Some(PartialHashAggregateState::ReadingInput { hash_table }), + hash_table: Some(hash_table), }) } - fn close_input(&mut self) { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - } + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema = Arc::clone(&self.schema); - fn break_with_err(error: DataFusionError) -> PartialHashAggregateStateTransition { - ControlFlow::Break(( - Poll::Ready(Some(Err(error))), - PartialHashAggregateState::Error, - )) + Box::pin(RecordBatchStreamAdapter::new(schema, self.create_stream())) } - fn break_with_internal_err( - message: impl std::fmt::Display, - ) -> PartialHashAggregateStateTransition { - Self::break_with_err(internal_datafusion_err!("{message}")) + /// Entry point for the partial hash aggregate state machine. + /// + /// See comments in [`PartialHashAggregateStream`] for high-level ideas. + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + let mut hash_table = self + .hash_table + .take() + .expect("hash_table should not be None"); + + debug_assert!(hash_table.is_building()); + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + + let mut last_state = HandleInputResult::ProcessNext; + while let Some(batch) = self.input.next().await.transpose()? { + let timer = elapsed_compute.timer(); + last_state = self.handle_input_batch(&batch, &mut hash_table)?; + + match last_state { + HandleInputResult::ProcessNext => {} + HandleInputResult::ReachedLimit + | HandleInputResult::SwitchToSkipAggregation => { + break; + } + HandleInputResult::OOM => { + let state_batch_result = hash_table.take_state_batch(); + + // Emitting clears the aggregate table and releases its + // accumulated memory. Update the reservation accordingly. + self.reservation.try_resize(hash_table.memory_size())?; + + let materialized_group_states = state_batch_result?.ok_or_else(|| { + internal_datafusion_err!( + "Partial hash aggregate ran out of memory with no aggregated groups" + ) + })?; + + timer.done(); + self.emit_on_memory_pressure( + materialized_group_states, + &mut emitter, + ) + .await?; + } + } + } + + let timer = elapsed_compute.timer(); + + let skip_hash_table = + if last_state == HandleInputResult::SwitchToSkipAggregation { + Some(hash_table.partial_skip_table()?) + } else { + self.close_input(); + + None + }; + hash_table.start_output()?; + + timer.done(); + + self.produce_output(hash_table, &mut emitter).await?; + + if let Some(hash_table) = skip_hash_table { + self.skip_rest_of_aggregation(hash_table, emitter).await?; + } + + Ok(()) + }) } /// See comments in [`Self::group_values_soft_limit`] for details. @@ -503,477 +532,150 @@ impl PartialHashAggregateStream { .is_some_and(|probe| probe.should_skip()) } - fn start_output( - &mut self, - hash_table: &mut AggregateHashTable, - close_input: bool, - ) -> Result<()> { - if close_input { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - } - hash_table.start_output() + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); } - /// Handle ReadingInput state - aggregate input batches into the hash table. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_reading_input( + /// Aggregate input batch into the hash table + fn handle_input_batch( &mut self, - cx: &mut Context<'_>, - original_state: PartialHashAggregateState, - ) -> PartialHashAggregateStateTransition { - let PartialHashAggregateState::ReadingInput { mut hash_table } = original_state - else { - return Self::break_with_internal_err( - "Partial hash aggregate stream expected ReadingInput state", - ); - }; - debug_assert!(hash_table.is_building()); - - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - PartialHashAggregateState::ReadingInput { hash_table }, - )), - Poll::Ready(Some(Ok(batch))) => { - // ---------------------------------- - // Step 1: Aggregate the input batch - // ---------------------------------- - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let input_rows = batch.num_rows(); - self.reduction_factor.add_total(input_rows); - let result = hash_table.aggregate_batch(&batch); - timer.done(); - - if let Err(e) = result { - return Self::break_with_err(e); - } + batch: &RecordBatch, + hash_table: &mut AggregateHashTable, + ) -> Result { + // ---------------------------------- + // Step 1: Aggregate the input batch + // ---------------------------------- + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); + hash_table.aggregate_batch(batch)?; + + // -------------------------------- + // Step 2: Soft limit optimization + // -------------------------------- + if self.hit_soft_group_limit(hash_table) { + return Ok(HandleInputResult::ReachedLimit); + } - // -------------------------------- - // Step 2: Soft limit optimization - // -------------------------------- - if self.hit_soft_group_limit(&hash_table) { - let timer = elapsed_compute.timer(); - let result = self.start_output(&mut hash_table, true); - timer.done(); + // ---------------------------------------------- + // Step 3: Skip partial aggregation optimization + // ---------------------------------------------- + self.update_skip_aggregation_probe(input_rows, hash_table.building_group_count()); - if let Err(e) = result { - return Self::break_with_err(e); - } + // True branch: a decision has been made to skip partial aggregation. + if self.should_skip_aggregation() { + return Ok(HandleInputResult::SwitchToSkipAggregation); + } - return ControlFlow::Continue( - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table: None, - }, - ); - } + // ------------------------------------------------- + // Step 4: Larger-than-memory execution (early emit) + // ------------------------------------------------- + let resize_result = self.reservation.try_resize(hash_table.memory_size()); + match resize_result { + Ok(()) => Ok(HandleInputResult::ProcessNext), + Err(DataFusionError::ResourcesExhausted(_)) => Ok(HandleInputResult::OOM), + Err(e) => Err(e), + } + } - // ---------------------------------------------- - // Step 3: Skip partial aggregation optimization - // ---------------------------------------------- - self.update_skip_aggregation_probe( - input_rows, - hash_table.building_group_count(), - ); - - // True branch: a decision has been made to skip partial aggregation. - if self.should_skip_aggregation() { - let timer = elapsed_compute.timer(); - let result = match hash_table.partial_skip_table() { - Ok(skip_hash_table) => self - .start_output(&mut hash_table, false) - .map(|()| skip_hash_table), - Err(e) => Err(e), - }; - timer.done(); - - match result { - Ok(skip_hash_table) => { - // Move to `ProducingOutput` first. Its `skip_hash_table` - // field moves the stream to skip-partial aggregation after - // the accumulated batches have been output. - return ControlFlow::Continue( - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table: Some(skip_hash_table), - }, - ); - } - Err(e) => return Self::break_with_err(e), - } - } + /// emit a materialized partial-state on memory pressure + /// batch in `batch_size`(from configuration) slices + async fn emit_on_memory_pressure( + &mut self, + // After each incremental emitting step, the `remaining_groups` will be updated + // with batch slicing. + mut remaining_groups: RecordBatch, + emitter: &mut TryEmitter, + ) -> Result<()> { + while remaining_groups.num_rows() > self.batch_size { + // More batch to output, continue in the current state. + let output = remaining_groups.slice(0, self.batch_size); - // ------------------------------------------------- - // Step 4: Larger-than-memory execution (early emit) - // ------------------------------------------------- - let timer = elapsed_compute.timer(); - let resize_result = self.reservation.try_resize(hash_table.memory_size()); - timer.done(); - match resize_result { - Ok(()) => {} - Err(DataFusionError::ResourcesExhausted(_)) => { - let elapsed_compute = - self.baseline_metrics.elapsed_compute().clone(); - // Stops on drop - let _timer = elapsed_compute.timer(); - let state_batch_result = hash_table.take_state_batch(); + remaining_groups = remaining_groups.slice( + self.batch_size, + remaining_groups.num_rows() - self.batch_size, + ); - // Emitting clears the aggregate table and releases its - // accumulated memory. Update the reservation accordingly. - let resize_result = - self.reservation.try_resize(hash_table.memory_size()); - - if let Err(e) = resize_result { - return Self::break_with_err(e); - } - - let materialized_group_states = match state_batch_result { - Ok(Some(batch)) => batch, - Ok(None) => { - return Self::break_with_err(internal_datafusion_err!( - "Partial hash aggregate ran out of memory with no aggregated groups" - )); - } - Err(e) => return Self::break_with_err(e), - }; - - return ControlFlow::Continue( - PartialHashAggregateState::EmittingOnMemoryPressure { - hash_table, - remaining_groups: materialized_group_states, - }, - ); - } - Err(e) => return Self::break_with_err(e), - } + self.reduction_factor.add_part(output.num_rows()); + debug_assert!(output.num_rows() > 0); - ControlFlow::Continue(PartialHashAggregateState::ReadingInput { - hash_table, - }) - } - Poll::Ready(Some(Err(e))) => Self::break_with_err(e), - Poll::Ready(None) => { - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = self.start_output(&mut hash_table, true); - timer.done(); - - match result { - Ok(()) => ControlFlow::Continue( - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table: None, - }, - ), - Err(e) => Self::break_with_err(e), - } - } + emitter + .emit(output.record_output(&self.baseline_metrics)) + .await; } - } - /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state - /// batch in `batch_size`(from configuration) slices, then resume reading input. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_emitting_on_memory_pressure( - &mut self, - original_state: PartialHashAggregateState, - ) -> PartialHashAggregateStateTransition { - let PartialHashAggregateState::EmittingOnMemoryPressure { - hash_table, - remaining_groups: batch, - } = original_state - else { - return Self::break_with_internal_err( - "Partial hash aggregate stream expected EmittingOnMemoryPressure state", - ); - }; + self.reduction_factor.add_part(remaining_groups.num_rows()); + debug_assert!(remaining_groups.num_rows() > 0); - let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { - // Last batch to output, go back to `ReadingInput` - ( - batch, - PartialHashAggregateState::ReadingInput { hash_table }, - ) - } else { - // More batch to output, continue in the current state. - let remaining = - batch.slice(self.batch_size, batch.num_rows() - self.batch_size); - let output = batch.slice(0, self.batch_size); - ( - output, - PartialHashAggregateState::EmittingOnMemoryPressure { - hash_table, - remaining_groups: remaining, - }, - ) - }; + emitter + .emit(remaining_groups.record_output(&self.baseline_metrics)) + .await; - self.reduction_factor.add_part(output_batch.num_rows()); - debug_assert!(output_batch.num_rows() > 0); - ControlFlow::Break(( - Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))), - next_state, - )) + Ok(()) } - /// Handle ProducingOutput state - emit partial aggregate state batches. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_producing_output( + /// emit partial aggregate state batches. + async fn produce_output( &mut self, - original_state: PartialHashAggregateState, - ) -> PartialHashAggregateStateTransition { - let PartialHashAggregateState::ProducingOutput { - mut hash_table, - skip_hash_table, - } = original_state - else { - return Self::break_with_internal_err( - "Partial hash aggregate stream expected ProducingOutput state", - ); - }; + mut hash_table: AggregateHashTable, + emitter: &mut TryEmitter, + ) -> Result<()> { debug_assert!(!hash_table.is_building()); let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.next_output_batch(); - timer.done(); - - match result { - Ok(Some(batch)) => { - let _ = self.reservation.try_resize(hash_table.memory_size()); - self.reduction_factor.add_part(batch.num_rows()); - debug_assert!(batch.num_rows() > 0); - let next_state = if hash_table.is_done() { - match skip_hash_table { - Some(hash_table) => { - PartialHashAggregateState::SkippingAggregation { hash_table } - } - None => PartialHashAggregateState::Done, - } - } else { - PartialHashAggregateState::ProducingOutput { - hash_table, - skip_hash_table, - } - }; + let mut timer = elapsed_compute.timer(); - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) - } - Ok(None) => { - let _ = self.reservation.try_resize(0); - // If the previous `Aggregating` stage decided to skip partial - // aggregation, go to the `SkippingAggregation` stage; otherwise finish. - let next_state = match skip_hash_table { - Some(hash_table) => { - PartialHashAggregateState::SkippingAggregation { hash_table } - } - None => PartialHashAggregateState::Done, - }; - ControlFlow::Continue(next_state) - } - Err(e) => Self::break_with_err(e), - } - } + loop { + let Some(batch) = hash_table.next_output_batch()? else { + // Only reachable when the table held no groups at all: a + // non-empty table always reports its last batch together with + // the `Done` state, which the `try_resize` below already zeroes. + self.reservation.try_resize(0)?; + return Ok(()); + }; - /// Handle SkippingAggregation state - convert raw input directly to partial states. - /// - /// See comments at `poll_next()` for details. - /// - /// Returns the next operator state with control flow decision. - fn handle_skipping_aggregation( - &mut self, - cx: &mut Context<'_>, - original_state: PartialHashAggregateState, - ) -> PartialHashAggregateStateTransition { - let PartialHashAggregateState::SkippingAggregation { mut hash_table } = - original_state - else { - return Self::break_with_internal_err( - "Partial hash aggregate stream expected SkippingAggregation state", - ); - }; + debug_assert!(batch.num_rows() > 0); - match self.input.poll_next_unpin(cx) { - Poll::Pending => ControlFlow::Break(( - Poll::Pending, - PartialHashAggregateState::SkippingAggregation { hash_table }, - )), - Poll::Ready(Some(Ok(batch))) => { - if let Some(probe) = self.skip_aggregation_probe.as_mut() { - probe.record_skipped(&batch); - } + // The table hands over its groups as they are materialized and + // reports a size of 0 once it reaches `Done`, so this releases the + // reservation before the final batch goes downstream. + let _ = self.reservation.try_resize(hash_table.memory_size()); + self.reduction_factor.add_part(batch.num_rows()); - let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - let timer = elapsed_compute.timer(); - let result = hash_table.convert_batch_to_state(&batch); - timer.done(); - - match result { - Ok(batch) => ControlFlow::Break(( - Poll::Ready(Some( - Ok(batch.record_output(&self.baseline_metrics)), - )), - PartialHashAggregateState::SkippingAggregation { hash_table }, - )), - Err(e) => Self::break_with_err(e), - } - } - Poll::Ready(Some(Err(e))) => Self::break_with_err(e), - Poll::Ready(None) => { - let input_schema = self.input.schema(); - self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - ControlFlow::Continue(PartialHashAggregateState::Done) - } + timer.done(); + emitter + .emit(batch.record_output(&self.baseline_metrics)) + .await; + timer = elapsed_compute.timer(); } } -} -impl Stream for PartialHashAggregateStream { - type Item = Result; + /// convert raw input directly to partial states. + async fn skip_rest_of_aggregation( + &mut self, + mut hash_table: AggregateHashTable, + mut emitter: TryEmitter, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); - /// Entry point for the partial hash aggregate state machine. - /// - /// See comments in [`PartialHashAggregateStream`] for high-level ideas. - /// - /// State transition graph: - /// - /// ```text - /// (start) - /// -> ReadingInput - /// The stream starts by polling input and aggregating batches into the - /// in-memory hash table. - /// - /// ReadingInput - /// -> ReadingInput - /// Aggregate one batch, update the inner aggregate hash table, and - /// continue with the next input batch. - /// -> EmittingOnMemoryPressure - /// The table cannot reserve enough memory. Materialize all accumulated - /// partial states and begin emitting them incrementally. - /// -> ProducingOutput(skip=None) - /// Input was exhausted, or the soft group limit was reached. Move to - /// the next state to start outputting. - /// -> ProducingOutput(skip=Some) - /// Partial skip aggregation was triggered. First move to the - /// `ProducingOutput` state to drain the accumulated state, then move to - /// the `SkippingAggregation` state to convert input directly to partial - /// state without aggregation. - /// - /// EmittingOnMemoryPressure - /// -> EmittingOnMemoryPressure - /// One batch-sized slice was yielded; repeat until all materialized - /// partial states are emitted. - /// -> ReadingInput - /// The materialized states were emitted; continue with the empty table. - /// - /// ProducingOutput(skip=None) - /// -> ProducingOutput(skip=None) - /// One accumulated output batch was yielded, repeat to continue producing - /// output incrementally. - /// -> Done - /// All accumulated output was emitted. - /// - /// ProducingOutput(skip=Some) - /// -> ProducingOutput(skip=Some) - /// One accumulated output batch was yielded, repeat to continue producing - /// output incrementally. - /// -> SkippingAggregation - /// All accumulated output was emitted. Continue by converting raw - /// input batches directly to partial aggregate state. - /// - /// SkippingAggregation - /// -> SkippingAggregation - /// One `convert_to_state` batch was yielded; repeat to continue - /// processing. - /// -> Done - /// Input was exhausted. - /// - /// Any active state - /// -> Error - /// An error drops state-owned resources before it is returned. - /// - /// Error - /// -> (end) - /// - /// Done - /// -> (end) - /// ``` - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - loop { - let cur_state = self - .state - .take() - .expect("PartialHashAggregateStream state should not be None"); + while let Some(batch) = self.input.next().await.transpose()? { + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + probe.record_skipped(&batch); + } - let next_state = match cur_state { - state @ PartialHashAggregateState::ReadingInput { .. } => { - self.handle_reading_input(cx, state) - } - state @ PartialHashAggregateState::EmittingOnMemoryPressure { .. } => { - self.handle_emitting_on_memory_pressure(state) - } - state @ PartialHashAggregateState::ProducingOutput { .. } => { - self.handle_producing_output(state) - } - state @ PartialHashAggregateState::SkippingAggregation { .. } => { - self.handle_skipping_aggregation(cx, state) - } - state @ PartialHashAggregateState::Error => { - self.close_input(); - self.reservation.free(); - self.state = Some(state); - return Poll::Ready(None); - } - state @ PartialHashAggregateState::Done => { - let _ = self.reservation.try_resize(0); - self.state = Some(state); - return Poll::Ready(None); - } + let result = { + let _timer = elapsed_compute.timer(); + hash_table.convert_batch_to_state(&batch)? }; - match next_state { - ControlFlow::Continue(next_state) => { - self.state = Some(next_state); - } - ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { - debug_assert!(matches!(next_state, PartialHashAggregateState::Error)); - - // The handler has already discarded its state-owned resources. - // Release the remaining stream-owned resources before returning. - self.close_input(); - self.reservation.free(); - self.state = Some(PartialHashAggregateState::Error); - return Poll::Ready(Some(Err(e))); - } - ControlFlow::Break((poll, next_state)) => { - self.state = Some(next_state); - return poll; - } - } + emitter + .emit(result.record_output(&self.baseline_metrics)) + .await; } - } -} -impl RecordBatchStream for PartialHashAggregateStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) + self.close_input(); + + Ok(()) } } @@ -1369,7 +1071,8 @@ mod tests { // Execute and collect results let mut stream = - PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)? + .into_stream(); let mut results = Vec::new(); while let Some(result) = stream.next().await { @@ -1513,7 +1216,8 @@ mod tests { // Execute and collect results let mut stream = - PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)? + .into_stream(); let mut results = Vec::new(); while let Some(result) = stream.next().await { diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index a6b581991627..0527098e9a7f 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -713,7 +713,7 @@ impl From for SendableRecordBatchStream { fn from(stream: StreamType) -> Self { match stream { StreamType::AggregateStream(stream) => Box::pin(stream), - StreamType::PartialHash(stream) => Box::pin(stream), + StreamType::PartialHash(stream) => stream.into_stream(), StreamType::PartialReduceHash(stream) => Box::pin(stream), StreamType::FinalHash(stream) => stream.into_stream(), StreamType::SingleHash(stream) => Box::pin(stream),