diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 7ebf32c3edfc..ad1a61134ade 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -168,6 +168,9 @@ pub(crate) struct PartialHashAggregateStream { /// Tracks partial aggregation row reduction, matching `GroupedHashAggregateStream`. reduction_factor: metrics::RatioMetrics, + /// Number of times accumulated states were emitted due to memory pressure. + early_emit_count: metrics::Count, + /// Tracks whether partial aggregation should switch to direct state conversion. skip_aggregation_probe: Option, @@ -394,6 +397,8 @@ impl PartialHashAggregateStream { let reduction_factor = MetricBuilder::new(&agg.metrics) .with_type(metrics::MetricType::Summary) .ratio_metrics("reduction_factor", partition); + let early_emit_count = + MetricBuilder::new(&agg.metrics).counter("early_emit_count", partition); let hash_table = AggregateHashTable::::new( agg, @@ -435,6 +440,7 @@ impl PartialHashAggregateStream { baseline_metrics, reservation, reduction_factor, + early_emit_count, skip_aggregation_probe, group_values_soft_limit: agg.limit_options().map(|config| config.limit()), hash_table: Some(hash_table), @@ -484,6 +490,7 @@ impl PartialHashAggregateStream { ) })?; + self.early_emit_count.add(1); timer.done(); self.emit_on_memory_pressure( materialized_group_states, @@ -1100,6 +1107,45 @@ mod tests { total_output_groups, num_groups, "Unexpected number of groups", ); + assert_eq!( + aggregate_exec + .metrics() + .unwrap() + .sum_by_name("early_emit_count") + .unwrap() + .as_usize(), + 0 + ); + + // Disable skip aggregation so the same input is emitted on memory pressure. + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1024, 1.0) + .build_arc()?; + let session_config = task_ctx.session_config().clone().set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(2.0)), + ); + let no_skip_task_ctx = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(session_config), + ); + let mut stream = + PartialHashAggregateStream::new(&aggregate_exec, &no_skip_task_ctx, 0)? + .into_stream(); + while let Some(result) = stream.next().await { + result?; + } + + assert_eq!( + aggregate_exec + .metrics() + .unwrap() + .sum_by_name("early_emit_count") + .unwrap() + .as_usize(), + 1 + ); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index ba08c4b00319..06d542a38999 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4698,6 +4698,16 @@ mod tests { let stream: SendableRecordBatchStream = stream.into(); let output = collect(stream).await?; + assert_eq!( + partial_reduce + .metrics() + .unwrap() + .sum_by_name("early_emit_count") + .unwrap() + .as_usize(), + num_input_batches + ); + // The table is flushed after every input batch, so each of the three // groups is emitted once per input batch instead of being merged into a // single row. Each flush is sliced into batches of 2 and 1 rows. diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d8f1447bc952..cc762146a13d 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -35,7 +35,7 @@ use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{AggregateHashTable, PartialReduceMarker}; -use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::metrics::{BaselineMetrics, Count, MetricBuilder, RecordOutput, SpillMetrics}; use crate::stream::EmptyRecordBatchStream; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; @@ -93,6 +93,9 @@ pub(crate) struct PartialReduceHashAggregateStream { /// Memory reservation for group keys and accumulators. reservation: MemoryReservation, + /// Number of times accumulated states were emitted due to memory pressure. + early_emit_count: Count, + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level /// state for emitting output batches. state: Option, @@ -193,6 +196,8 @@ impl PartialReduceHashAggregateStream { // Preserve the existing aggregate metric surface for this plan node. let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let early_emit_count = + MetricBuilder::new(&agg.metrics).counter("early_emit_count", partition); let hash_table = AggregateHashTable::::new( agg, @@ -211,6 +216,7 @@ impl PartialReduceHashAggregateStream { batch_size, baseline_metrics, reservation, + early_emit_count, state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }), }) } @@ -316,12 +322,15 @@ impl PartialReduceHashAggregateStream { let state_batch_result = original_state.hash_table_mut().take_state_batch(); match state_batch_result { - Ok(Some(remaining_groups)) => ControlFlow::Continue( - PartialReduceHashAggregateState::EmittingOnMemoryPressure { - hash_table: original_state.into_hash_table(), - remaining_groups, - }, - ), + Ok(Some(remaining_groups)) => { + self.early_emit_count.add(1); + ControlFlow::Continue( + PartialReduceHashAggregateState::EmittingOnMemoryPressure { + hash_table: original_state.into_hash_table(), + remaining_groups, + }, + ) + } // No accumulated group to emit, so early emission cannot release any // memory: report the original error. Ok(None) => Self::break_with_err(oom), diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 79c6cb7a5155..53d515bf7fd1 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -208,9 +208,9 @@ FROM ( ---- 100000 5000050000 -# Assert spill happened in the final aggregation. -# In multi-partitions configuration, 'spilled_rows' is not deterministic, so assert -# the unit to be 'K' +# Assert spill happened in the final aggregation and the partial aggregation +# reports memory-pressure emissions. Their exact counts are not deterministic, +# so assert only the spilled_rows unit and the presence of early_emit_count. query TT EXPLAIN ANALYZE SELECT count(*), sum(total) @@ -223,6 +223,8 @@ FROM ( 06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spilled_rows=K,] +08)--------------AggregateExec: mode=Partial, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[early_emit_count=] + # Restore settings to slt runner defaults statement ok