alamb commented on code in PR #23657:
URL: https://github.com/apache/datafusion/pull/23657#discussion_r3626896203
##########
datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs:
##########
@@ -234,6 +250,43 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
+ self.buffer.group_indices.allocated_size()
}
+ pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics {
+ self.group_by_metrics.clone()
+ }
+
+ /// Takes every intermediate aggregate state and resets the table so it can
+ /// continue with a new ordered input segment.
+ ///
+ /// Unlike normal ordered emission, this operation is allowed to take the
+ /// active (incomplete) groups. Partial aggregation can pass those states
to
+ /// its final stage, while final aggregation sorts and spills them before
+ /// replay.
+ pub(in crate::aggregates) fn take_state_batch(
+ &mut self,
+ ) -> Result<Option<RecordBatch>> {
+ if self.buffer.group_values.is_empty() {
+ return Ok(None);
+ }
+
+ let mut output = self.buffer.group_values.emit(EmitTo::All)?;
+ for acc in &mut self.buffer.accumulators {
+ output.extend(acc.state(EmitTo::All)?);
+ }
+
+ let batch = RecordBatch::try_new(Arc::clone(&self.state_schema),
output)?;
+ debug_assert!(batch.num_rows() > 0);
+
+ // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the
Review Comment:
One thought I had was maybe it would be worth moving all this resetting of
fields of self.buffer into a method on `OrderedAggregateTableBuffer`? likewise
for instead of directly manipulating `self.buffer.group_values` we could add a
method like `OrderedAggregateTableBuffer::emit` to ecapsulate the fields more
On the other hand, maybe that just obscures the intent more 🤷 I don't feel
strongly -- just a thought
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -47,13 +69,46 @@ pub(crate) struct OrderedFinalAggregateStream {
state: Option<OrderedFinalAggregateState>,
}
+/// Spill configuration and accumulated runs for partially ordered final
+/// aggregation. Each file is one fully group-key-sorted intermediate-state
run;
Review Comment:
I read this comment as suggesting there is one file per group that is
spilled (e.g each group is spilled into its own file).
I think it is more like a file per spill event (aka the all the groups in
memory at the time of a spill, hen sorted by group key)
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -23,22 +23,44 @@ use std::task::{Context, Poll};
use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
-use datafusion_common::Result;
+use datafusion_common::{DataFusionError, Result, internal_err};
use datafusion_execution::TaskContext;
use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion_physical_expr::PhysicalSortExpr;
+use datafusion_physical_expr::expressions::Column;
+use datafusion_physical_expr_common::sort_expr::LexOrdering;
use futures::stream::{Stream, StreamExt};
use super::AggregateExec;
use super::aggregate_hash_table::{FinalMarker, OrderedAggregateTable};
+use super::group_values::GroupByMetrics;
use crate::aggregates::AggregateMode;
use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics};
+use crate::sorts::IncrementalSortIterator;
+use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
+use crate::spill::spill_manager::SpillManager;
use crate::stream::EmptyRecordBatchStream;
use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream};
/// Final aggregate stream for `InputOrderMode::Sorted` and
/// `InputOrderMode::PartiallySorted`.
///
/// See comments at [`super::ordered_partial_stream`] for details.
+///
+/// # Spilling
+///
+/// This section is only for implementation notes, for background, see
[`super::ordered_partial_stream`]
Review Comment:
I think the context is actually on `OrderedPartialAggregateStream`
```suggestion
/// This section is only for implementation notes, for background, see
[`super::ordered_partial_stream::OrderedPartialAggregateStream`]
```
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -47,13 +69,46 @@ pub(crate) struct OrderedFinalAggregateStream {
state: Option<OrderedFinalAggregateState>,
}
+/// Spill configuration and accumulated runs for partially ordered final
+/// aggregation. Each file is one fully group-key-sorted intermediate-state
run;
+/// all runs are replayed together after the original input is exhausted.
+struct OrderedFinalSpillContext {
+ /// Aggregate configuration
+ agg: AggregateExec,
+ /// Task context
+ context: Arc<TaskContext>,
+ /// Original partition index
+ partition: usize,
+ /// Target batch size from configuration
+ batch_size: usize,
+ /// Full group-key ordering, such ordering with be kept in: a) individual
spill
+ /// files, b) order after final merging and streaming aggregate
+ spill_expr: LexOrdering,
+ /// Spill I/O and metrics manager.
+ spill_manager: SpillManager,
+ /// Fully sorted spill runs waiting to be merged.
+ spills: Vec<SortedSpillFile>,
+}
+
/// See comments at `poll_next()` for details.
enum OrderedFinalAggregateState {
ReadingInput {
table: OrderedAggregateTable<FinalMarker>,
+ spill_context: Option<Box<OrderedFinalSpillContext>>,
},
- DrainingFinal {
+ Spilling {
table: OrderedAggregateTable<FinalMarker>,
+ spill_context: Box<OrderedFinalSpillContext>,
+ },
+ ProducingOutput {
+ table: OrderedAggregateTable<FinalMarker>,
+ },
+ PreparingMergeInput {
Review Comment:
it might help readability to have a sentence or two comment explaining what
the difference is between `PreparingMergeInput` and `MergingSpills`
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -64,6 +119,135 @@ type OrderedFinalAggregateStateTransition = ControlFlow<
OrderedFinalAggregateState,
>;
+impl OrderedFinalSpillContext {
+ fn new(
+ agg: &AggregateExec,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ batch_size: usize,
+ input_order_mode: &InputOrderMode,
+ spill_schema: &SchemaRef,
+ spill_metrics: SpillMetrics,
+ ) -> Result<Self> {
+ let group_schema = agg.group_by.group_schema(spill_schema)?;
+ let output_ordering = agg.cache.output_ordering();
+ let InputOrderMode::PartiallySorted(order_indices) = input_order_mode
else {
+ return internal_err!("Ordered final spill requires partially
ordered input");
+ };
+ let spill_indices = order_indices.iter().copied().chain(
+ (0..group_schema.fields().len()).filter(|idx|
!order_indices.contains(idx)),
+ );
+ let spill_sort_exprs = spill_indices.map(|idx| {
+ let field = group_schema.field(idx);
+ let output_expr = Column::new(field.name(), idx);
+ let sort_options = output_ordering
+ .and_then(|ordering| ordering.get_sort_options(&output_expr))
+ .unwrap_or_default();
+ PhysicalSortExpr::new(Arc::new(output_expr), sort_options)
+ });
+ let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else {
+ return internal_err!("Ordered final spill expression is empty");
+ };
+
+ let spill_manager = SpillManager::new(
+ context.runtime_env(),
+ spill_metrics,
+ Arc::clone(spill_schema),
+ )
+ .with_compression_type(context.session_config().spill_compression());
+
+ Ok(Self {
+ agg: agg.clone(),
+ context: Arc::clone(context),
+ partition,
+ batch_size,
+ spill_expr,
+ spill_manager,
+ spills: vec![],
+ })
+ }
+
+ fn has_spills(&self) -> bool {
+ !self.spills.is_empty()
+ }
+
+ /// Sort and spill the aggregated groups. Memory reservation should be
cleared
+ /// by the caller of this function.
+ ///
+ /// See [`OrderedFinalAggregateStream`] for spilling details.
+ ///
+ /// It must not be called without rows to spill in the aggregate table,
otherwise
+ /// it returns `Ok(false)`, and the caller propagates the error.
+ fn spill_table(
+ &mut self,
+ table: &mut OrderedAggregateTable<FinalMarker>,
+ ) -> Result<bool> {
+ let Some(batch) = table.take_state_batch()? else {
+ return Ok(false);
+ };
+
+ let sorted_iter =
+ IncrementalSortIterator::new(batch, self.spill_expr.clone(),
self.batch_size);
Review Comment:
TIL `IncrementalSortIterator` (from @EmilyMatt in #20314 it seems ❤️ )
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -180,36 +497,183 @@ impl OrderedFinalAggregateStream {
next_state,
))
}
+ // Can't do early emit, continue aggregating.
Ok(None) => {
- // Ordered variant doesn't support memory-limited
- // execution, so it errors when memory reservation
fails.
- if let Err(e) =
self.reservation.try_resize(table.memory_size()) {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput {
table },
- ));
- }
-
- // Can't do early emit, continue aggregating.
ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput {
table,
+ spill_context,
})
}
Err(e) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
)),
}
}
Poll::Ready(Some(Err(e))) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
)),
Poll::Ready(None) => {
self.close_input();
- table.input_done();
-
ControlFlow::Continue(OrderedFinalAggregateState::DrainingFinal { table })
+ match spill_context {
+ Some(spill_context) if spill_context.has_spills() => {
+ ControlFlow::Continue(
+ OrderedFinalAggregateState::PreparingMergeInput {
+ table,
+ spill_context,
+ },
+ )
+ }
+ _ => {
+ table.input_done();
+ ControlFlow::Continue(
+ OrderedFinalAggregateState::ProducingOutput {
table },
+ )
+ }
+ }
+ }
+ }
+ }
+
+ /// Sorts and spills one complete in-memory state run, then resumes input.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_spilling(
+ &mut self,
+ original_state: OrderedFinalAggregateState,
+ ) -> OrderedFinalAggregateStateTransition {
+ let OrderedFinalAggregateState::Spilling {
+ mut table,
+ mut spill_context,
+ } = original_state
+ else {
+ unreachable!("expected spilling state")
+ };
+
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let mut result = spill_context.spill_table(&mut table);
+
+ // Spilling shrinks the aggregate table and releases its accumulated
+ // memory. Update the reservation accordingly.
+ if self.reservation.try_resize(table.memory_size()).is_err() {
+ // Fold spill error and reservation error into one
+ result = internal_err!("Reservation after spilling should succeed")
+ }
+
+ timer.done();
+
+ match result {
+ // Finished spilling the aggregate table, continue aggregating
from input
+ Ok(true) =>
ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context: Some(spill_context),
+ }),
+ Ok(false) => ControlFlow::Break((
+ Poll::Ready(Some(internal_err!(
+ "Ordered final aggregation entered Spilling with an empty
table"
+ ))),
+ OrderedFinalAggregateState::Done,
+ )),
+ Err(e) => ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ )),
+ }
+ }
+
+ /// 1. Spills the last in-memory run.
+ /// 2. Constructs a globally ordered input stream by applying a
sort-preserving
+ /// merge to all spills.
+ /// 3. Constructs a replay stream: an ordered aggregate stream over the
fully
+ /// ordered input constructed from the spills.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_preparing_merge_input(
+ &mut self,
+ original_state: OrderedFinalAggregateState,
+ ) -> OrderedFinalAggregateStateTransition {
+ let OrderedFinalAggregateState::PreparingMergeInput {
+ mut table,
+ mut spill_context,
+ } = original_state
+ else {
+ unreachable!("expected preparing merge input state")
Review Comment:
it could be more defensive and return an Err here (an make the function
return Result<..>) as a wa to avoid potential future refactors introducing
panics
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -86,6 +270,35 @@ impl OrderedFinalAggregateStream {
partition: usize,
input: SendableRecordBatchStream,
input_order_mode: &InputOrderMode,
+ ) -> Result<Self> {
+ let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
+ let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition);
+ let spill_metrics = SpillMetrics::new(&agg.metrics, partition);
+ Self::new_with_input_and_metrics(
+ agg,
+ context,
+ partition,
+ input,
+ input_order_mode,
+ baseline_metrics,
+ group_by_metrics,
+ Some(spill_metrics),
+ )
+ }
+
+ #[expect(
Review Comment:
maybe eventually it would be good to make a bulder or something ? That is
probably too much code for now though
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -64,6 +119,135 @@ type OrderedFinalAggregateStateTransition = ControlFlow<
OrderedFinalAggregateState,
>;
+impl OrderedFinalSpillContext {
+ fn new(
+ agg: &AggregateExec,
+ context: &Arc<TaskContext>,
+ partition: usize,
+ batch_size: usize,
+ input_order_mode: &InputOrderMode,
+ spill_schema: &SchemaRef,
+ spill_metrics: SpillMetrics,
+ ) -> Result<Self> {
+ let group_schema = agg.group_by.group_schema(spill_schema)?;
+ let output_ordering = agg.cache.output_ordering();
+ let InputOrderMode::PartiallySorted(order_indices) = input_order_mode
else {
+ return internal_err!("Ordered final spill requires partially
ordered input");
+ };
+ let spill_indices = order_indices.iter().copied().chain(
+ (0..group_schema.fields().len()).filter(|idx|
!order_indices.contains(idx)),
+ );
+ let spill_sort_exprs = spill_indices.map(|idx| {
+ let field = group_schema.field(idx);
+ let output_expr = Column::new(field.name(), idx);
+ let sort_options = output_ordering
+ .and_then(|ordering| ordering.get_sort_options(&output_expr))
+ .unwrap_or_default();
+ PhysicalSortExpr::new(Arc::new(output_expr), sort_options)
+ });
+ let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else {
+ return internal_err!("Ordered final spill expression is empty");
+ };
+
+ let spill_manager = SpillManager::new(
+ context.runtime_env(),
+ spill_metrics,
+ Arc::clone(spill_schema),
+ )
+ .with_compression_type(context.session_config().spill_compression());
+
+ Ok(Self {
+ agg: agg.clone(),
+ context: Arc::clone(context),
+ partition,
+ batch_size,
+ spill_expr,
+ spill_manager,
+ spills: vec![],
+ })
+ }
+
+ fn has_spills(&self) -> bool {
+ !self.spills.is_empty()
+ }
+
+ /// Sort and spill the aggregated groups. Memory reservation should be
cleared
+ /// by the caller of this function.
+ ///
+ /// See [`OrderedFinalAggregateStream`] for spilling details.
+ ///
+ /// It must not be called without rows to spill in the aggregate table,
otherwise
Review Comment:
Could you also document what the return value means? Is it "table was
empty"? If it is an error, why not just return the error rather than also a
boolean?
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -180,36 +497,183 @@ impl OrderedFinalAggregateStream {
next_state,
))
}
+ // Can't do early emit, continue aggregating.
Ok(None) => {
- // Ordered variant doesn't support memory-limited
- // execution, so it errors when memory reservation
fails.
- if let Err(e) =
self.reservation.try_resize(table.memory_size()) {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput {
table },
- ));
- }
-
- // Can't do early emit, continue aggregating.
ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput {
table,
+ spill_context,
})
}
Err(e) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
)),
}
}
Poll::Ready(Some(Err(e))) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
)),
Poll::Ready(None) => {
self.close_input();
- table.input_done();
-
ControlFlow::Continue(OrderedFinalAggregateState::DrainingFinal { table })
+ match spill_context {
+ Some(spill_context) if spill_context.has_spills() => {
+ ControlFlow::Continue(
+ OrderedFinalAggregateState::PreparingMergeInput {
+ table,
+ spill_context,
+ },
+ )
+ }
+ _ => {
+ table.input_done();
+ ControlFlow::Continue(
+ OrderedFinalAggregateState::ProducingOutput {
table },
+ )
+ }
+ }
+ }
+ }
+ }
+
+ /// Sorts and spills one complete in-memory state run, then resumes input.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_spilling(
+ &mut self,
+ original_state: OrderedFinalAggregateState,
+ ) -> OrderedFinalAggregateStateTransition {
+ let OrderedFinalAggregateState::Spilling {
+ mut table,
+ mut spill_context,
+ } = original_state
+ else {
+ unreachable!("expected spilling state")
+ };
+
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let mut result = spill_context.spill_table(&mut table);
+
+ // Spilling shrinks the aggregate table and releases its accumulated
+ // memory. Update the reservation accordingly.
+ if self.reservation.try_resize(table.memory_size()).is_err() {
+ // Fold spill error and reservation error into one
+ result = internal_err!("Reservation after spilling should succeed")
Review Comment:
I think this can happen if other plans are allocating from the memory
manager at the same time (e.g. the aggregate releases the memory but then
something in another plan reserves it)
In other words, I think we should just propagte the "resources exhausted"
error here
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -180,36 +497,183 @@ impl OrderedFinalAggregateStream {
next_state,
))
}
+ // Can't do early emit, continue aggregating.
Ok(None) => {
- // Ordered variant doesn't support memory-limited
- // execution, so it errors when memory reservation
fails.
- if let Err(e) =
self.reservation.try_resize(table.memory_size()) {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput {
table },
- ));
- }
-
- // Can't do early emit, continue aggregating.
ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput {
table,
+ spill_context,
})
}
Err(e) => ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
Review Comment:
Rather than putting the state back into ReadingInput, should we move it to
the terminal state (Done) instead? That would could potentially release the
intermediate state quickly on error rather than waiting for dropping the plan
I realize this PR follows the existing pattern, but it occured to me while
reviewing
##########
datafusion/physical-plan/src/aggregates/ordered_final_stream.rs:
##########
@@ -157,21 +409,86 @@ impl OrderedFinalAggregateStream {
if let Err(e) = result {
return ControlFlow::Break((
Poll::Ready(Some(Err(e))),
- OrderedFinalAggregateState::ReadingInput { table },
+ OrderedFinalAggregateState::ReadingInput {
+ table,
+ spill_context,
+ },
));
}
+ // Check memory reservation, and potentially spill.
let timer = elapsed_compute.timer();
- let result = table.next_output_batch();
+ let resize_result =
+ self.reservation
+ .try_resize(Self::reservation_size_for_table(
+ &table,
+ spill_context.as_deref(),
+ ));
timer.done();
+ match resize_result {
+ Ok(()) => {}
+ Err(e @ DataFusionError::ResourcesExhausted(_)) => {
+ let Some(spill_context) = spill_context else {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ };
+ if table.is_empty() {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ }
+ return ControlFlow::Continue(
+ OrderedFinalAggregateState::Spilling {
+ table,
+ spill_context,
+ },
+ );
+ }
+ Err(e) => {
+ return ControlFlow::Break((
+ Poll::Ready(Some(Err(e))),
+ OrderedFinalAggregateState::Done,
+ ));
+ }
+ }
+
+ let result = if spill_context
+ .as_ref()
+ .is_some_and(|spill_context| spill_context.has_spills())
+ {
+ // Once one incomplete run is spilled, every remaining
state
+ // must participate in replay so no group is finalized
twice.
+ Ok(None)
+ } else {
+ let timer = elapsed_compute.timer();
+ let result = table.next_output_batch();
+ timer.done();
+ result
+ };
match result {
// Some finalized groups can be emitted. Yield them, then
// continue aggregating input in the current state.
Ok(Some(batch)) => {
- let next_state =
- OrderedFinalAggregateState::ReadingInput { table };
- self.resize_reservation_for_state(&next_state);
+ if let Err(e) =
Review Comment:
This pattern (of handing the "can't reserve what is needed") is repeated a
few times and I think contributes to the length of the `poll_inner` here. I
wonder if there is some way to factor it out into a new function to make it
easier to read / keep it encapsulated 🤔
--
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]
