2010YOUY01 commented on code in PR #23965:
URL: https://github.com/apache/datafusion/pull/23965#discussion_r3672580948
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -5766,9 +5762,9 @@ mod tests {
#[tokio::test]
async fn test_aggregate_with_spill_if_necessary() -> Result<()> {
// test with spill
- run_test_with_spill_pool_if_necessary(2_000, true).await?;
+ run_test_with_spill_pool_if_necessary(20_000, true).await?;
Review Comment:
Here is the same as the previous note on test change.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -5664,9 +5658,11 @@ mod tests {
Field::new("b", DataType::Float64, false),
]));
+ let group_keys = [2, 3, 4, 4].repeat(1_000);
Review Comment:
The old implementation set a very tight memory budget, and the refactored
implementation might have minor implementation difference, making this test
fail.
Here it just try to scale up the test input size and memory limit, the test
target is the same.
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -6744,11 +6740,6 @@ mod tests {
matches!(root, DataFusionError::ResourcesExhausted(_)),
"Expected ResourcesExhausted, got: {root}",
);
- let msg = root.to_string();
Review Comment:
Above check on error type is enough I think, here it also assert the exact
error message
##########
datafusion/physical-plan/src/aggregates/single_stream.rs:
##########
@@ -139,162 +312,386 @@ impl SingleHashAggregateStream {
) -> Result<Self> {
debug_assert!(matches!(
agg.mode,
- super::AggregateMode::Single |
super::AggregateMode::SinglePartitioned
+ AggregateMode::Single | AggregateMode::SinglePartitioned
));
debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear);
let schema = Arc::clone(&agg.schema);
let input = agg.input.execute(partition, Arc::clone(context))?;
+ let input_schema = input.schema();
let batch_size = context.session_config().batch_size();
let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
+ let spill_metrics = SpillMetrics::new(&agg.metrics, partition);
+ let state_schema = Arc::new(create_schema(
+ input_schema.as_ref(),
+ &agg.group_by,
+ &agg.aggr_expr,
+ AggregateMode::Partial,
+ )?);
let hash_table = AggregateHashTable::<SingleMarker>::new(
agg,
partition,
Arc::clone(&schema),
+ Arc::clone(&state_schema),
batch_size,
)?;
+ let can_spill = context.runtime_env().disk_manager.tmp_files_enabled();
+ let spill_context = if can_spill {
+ Some(Box::new(SingleSpillContext::new(
+ agg,
+ context,
+ partition,
+ batch_size,
+ &state_schema,
+ spill_metrics,
+ )?))
+ } else {
+ None
+ };
+
let reservation =
MemoryConsumer::new(format!("SingleHashAggregateStream[{partition}]"))
+ .with_can_spill(can_spill)
.register(context.memory_pool());
Ok(Self {
schema,
input,
baseline_metrics,
reservation,
- state: Some(SingleHashAggregateState::ReadingInput { hash_table }),
+ state: Some(SingleHashAggregateState::ReadingInput {
+ hash_table,
+ spill_context,
+ }),
})
}
- /// 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<()> {
+ fn close_input(&mut self) {
let input_schema = self.input.schema();
self.input = Box::pin(EmptyRecordBatchStream::new(input_schema));
- hash_table.start_output()
}
- /// Handle ReadingInput state - aggregate input batches into the hash
table.
+ fn break_with_err(error: DataFusionError) ->
SingleHashAggregateStateTransition {
+ ControlFlow::Break((
+ Poll::Ready(Some(Err(error))),
+ SingleHashAggregateState::Error,
+ ))
+ }
+
+ fn break_with_internal_err(message: &str) ->
SingleHashAggregateStateTransition {
+ Self::break_with_err(internal_datafusion_err!("{message}"))
+ }
+
+ /// Reserve memory for the current aggregate table.
+ fn reservation_size_for_table(
+ hash_table: &AggregateHashTable<SingleMarker>,
+ spill_context: Option<&SingleSpillContext>,
+ ) -> usize {
+ let table_size = hash_table.memory_size();
+ if spill_context.is_some() {
+ // See `SingleHashAggregateStream` comments for how this is
estimated.
+ table_size.saturating_add(
+ hash_table
+ .building_group_count()
+ .saturating_mul(size_of::<u32>()),
+ )
+ } else {
+ table_size
+ }
+ }
+
+ /// Consumes one raw input batch and updates the single-stage hash table.
///
/// See comments at `poll_next()` for details.
///
/// Returns the next operator state with control flow decision.
fn handle_reading_input(
&mut self,
cx: &mut Context<'_>,
- mut original_state: SingleHashAggregateState,
+ original_state: SingleHashAggregateState,
) -> SingleHashAggregateStateTransition {
- debug_assert!(matches!(
- &original_state,
- SingleHashAggregateState::ReadingInput { .. }
- ));
- debug_assert!(original_state.hash_table().is_building());
+ let SingleHashAggregateState::ReadingInput {
+ mut hash_table,
+ spill_context,
+ } = original_state
+ else {
+ return Self::break_with_internal_err(
+ "Single hash aggregate stream expected ReadingInput state",
+ );
+ };
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::Pending => ControlFlow::Break((
+ Poll::Pending,
+ SingleHashAggregateState::ReadingInput {
+ hash_table,
+ spill_context,
+ },
+ )),
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);
+ let result = hash_table.aggregate_batch(&batch);
timer.done();
if let Err(e) = result {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- original_state,
- ));
+ return Self::break_with_err(e);
}
- if let Err(e) = self
- .reservation
- .try_resize(original_state.hash_table().memory_size())
- {
- return ControlFlow::Break((
- Poll::Ready(Some(Err(e))),
- original_state,
- ));
+ // Check memory reservation, and potentially spill.
+ let timer = elapsed_compute.timer();
+ let resize_result =
+ self.reservation
+ .try_resize(Self::reservation_size_for_table(
+ &hash_table,
+ spill_context.as_deref(),
+ ));
+ timer.done();
+ match resize_result {
+ Ok(()) => {}
+ Err(e @ DataFusionError::ResourcesExhausted(_)) => {
+ let Some(spill_context) = spill_context else {
+ return Self::break_with_err(e.context(
+ "Single hash aggregate cannot spill because
temporary files are not enabled in the DiskManager",
+ ));
+ };
+ if hash_table.building_group_count() == 0 {
+ return Self::break_with_internal_err(
+ "Single hash aggregate ran out of memory with
no aggregated groups",
+ );
+ }
+ return ControlFlow::Continue(
+ SingleHashAggregateState::Spilling {
+ hash_table,
+ spill_context,
+ },
+ );
+ }
+ Err(e) => {
+ return Self::break_with_err(e);
+ }
}
- ControlFlow::Continue(original_state)
+ ControlFlow::Continue(SingleHashAggregateState::ReadingInput {
+ hash_table,
+ spill_context,
+ })
}
- Poll::Ready(Some(Err(e))) => {
- ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state))
- }
- // Input ends, move to output state
+ 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(original_state.hash_table_mut());
- timer.done();
-
- match result {
- Ok(()) => {
-
ControlFlow::Continue(original_state.into_producing_output())
+ self.close_input();
+ match spill_context {
+ Some(spill_context) if spill_context.has_spills() => {
+ ControlFlow::Continue(
+ SingleHashAggregateState::PreparingMergeInput {
+ hash_table,
+ spill_context,
+ },
+ )
}
- Err(e) => {
- ControlFlow::Break((Poll::Ready(Some(Err(e))),
original_state))
+ _ => {
+ let elapsed_compute =
+ self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let result = hash_table.start_output();
+ timer.done();
+
+ match result {
+ Ok(()) => ControlFlow::Continue(
+ SingleHashAggregateState::ProducingOutput {
hash_table },
+ ),
+ Err(e) => Self::break_with_err(e),
+ }
}
}
}
}
}
- /// Handle ProducingOutput state - emit final aggregate value batches.
+ /// 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: SingleHashAggregateState,
+ ) -> SingleHashAggregateStateTransition {
+ let SingleHashAggregateState::Spilling {
+ mut hash_table,
+ mut spill_context,
+ } = original_state
+ else {
+ return Self::break_with_internal_err(
+ "Single hash aggregate stream expected Spilling state",
+ );
+ };
+
+ // Sanity check: it is impossible to OOM when the table is empty.
+ if hash_table.building_group_count() == 0 {
+ return Self::break_with_internal_err(
+ "Single hash aggregation entered Spilling with an empty table",
+ );
+ }
+
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let mut result = spill_context.spill_table(&mut hash_table);
+
+ // Spilling shrinks the aggregate table and releases its accumulated
+ // memory. Update the reservation accordingly.
+ if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) {
+ result =
+ Err(e.context("Decreasing allocation after spilling should
succeed"));
+ }
+
+ timer.done();
+
+ match result {
+ // Finished spilling the aggregate table, continue aggregating
from input.
+ Ok(()) =>
ControlFlow::Continue(SingleHashAggregateState::ReadingInput {
+ hash_table,
+ spill_context: Some(spill_context),
+ }),
+ Err(e) => Self::break_with_err(e),
+ }
+ }
+
+ /// 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 final 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: SingleHashAggregateState,
+ ) -> SingleHashAggregateStateTransition {
+ let SingleHashAggregateState::PreparingMergeInput {
+ mut hash_table,
+ mut spill_context,
+ } = original_state
+ else {
+ return Self::break_with_internal_err(
+ "Single hash aggregate stream expected PreparingMergeInput
state",
+ );
+ };
+
+ let elapsed_compute = self.baseline_metrics.elapsed_compute().clone();
+ let timer = elapsed_compute.timer();
+ let replay = match spill_context.spill_table(&mut hash_table) {
+ Ok(()) => {
+ let group_by_metrics = hash_table.group_by_metrics();
+ drop(hash_table);
+ match self.reservation.try_resize(0) {
+ Ok(()) => (*spill_context).into_replay_stream(
+ &self.baseline_metrics,
+ group_by_metrics,
+ self.reservation.new_empty(),
+ ),
+ Err(e) => Err(e),
+ }
+ }
+ Err(e) => Err(e),
+ };
+ timer.done();
+
+ match replay {
+ Ok(stream) => {
+ ControlFlow::Continue(SingleHashAggregateState::MergingSpills
{ stream })
+ }
+ Err(e) => Self::break_with_err(e),
+ }
+ }
+
+ /// Forwards output from the fully ordered stream that consumes the merged
+ /// spill runs.
+ ///
+ /// See comments at `poll_next()` for details.
+ ///
+ /// Returns the next operator state with control flow decision.
+ fn handle_merging_spills(
+ &mut self,
+ cx: &mut Context<'_>,
+ original_state: SingleHashAggregateState,
+ ) -> SingleHashAggregateStateTransition {
+ let SingleHashAggregateState::MergingSpills { mut stream } =
original_state
+ else {
+ return Self::break_with_internal_err(
+ "Single hash aggregate stream expected MergingSpills state",
+ );
+ };
+
+ match stream.poll_next_unpin(cx) {
+ Poll::Pending => ControlFlow::Break((
+ Poll::Pending,
+ SingleHashAggregateState::MergingSpills { stream },
+ )),
+ Poll::Ready(Some(Ok(batch))) => ControlFlow::Break((
+ Poll::Ready(Some(Ok(batch))),
+ SingleHashAggregateState::MergingSpills { stream },
+ )),
+ Poll::Ready(Some(Err(e))) => Self::break_with_err(e),
+ Poll::Ready(None) =>
ControlFlow::Continue(SingleHashAggregateState::Done),
+ }
+ }
+
+ /// Emits one batch after input is exhausted.
///
/// See comments at `poll_next()` for details.
///
/// Returns the next operator state with control flow decision.
fn handle_producing_output(
Review Comment:
This function has no functional changes, the diffs are
1. use internal_err instead of assert
2. simplify some control flow code with `break_with_err`
##########
datafusion/sqllogictest/test_files/aggregate_memory_spill.slt:
##########
@@ -33,6 +33,9 @@
statement ok
SET datafusion.execution.target_partitions = 1
+statement ok
+SET datafusion.execution.batch_size = 128
Review Comment:
similar to the above notes on the test change
--
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]
