From 95a0b9bd1baaed17598988aa78f584cdb45a4a4b Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Mon, 7 Sep 2026 01:01:04 +0530 Subject: [PATCH 1/4] feat: add early emit count metric to partial aggregates --- .../src/aggregates/grouped_hash_stream.rs | 44 ++++++++++++++++++ .../src/aggregates/hash_stream.rs | 46 +++++++++++++++++++ .../physical-plan/src/aggregates/mod.rs | 20 ++++++++ .../src/aggregates/ordered_partial_stream.rs | 6 +++ .../src/aggregates/partial_reduce_stream.rs | 23 +++++++--- 5 files changed, 132 insertions(+), 7 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 3f6f3f8ce815b..41239c710bbad 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -385,6 +385,9 @@ pub(crate) struct GroupedHashAggregateStream { /// Reduction factor metric, calculated as `output_rows/input_rows` (only for partial aggregation) reduction_factor: Option, + + /// Number of times accumulated states were emitted due to memory pressure. + early_emit_count: Option, } impl GroupedHashAggregateStream { @@ -611,6 +614,9 @@ impl GroupedHashAggregateStream { } else { None }; + let early_emit_count = (agg.mode == AggregateMode::Partial).then(|| { + MetricBuilder::new(&agg.metrics).counter("early_emit_count", partition) + }); Ok(GroupedHashAggregateStream { schema: agg_schema, @@ -637,6 +643,7 @@ impl GroupedHashAggregateStream { group_values_soft_limit: agg.limit_options().map(|config| config.limit()), skip_aggregation_probe, reduction_factor, + early_emit_count, }) } } @@ -1021,6 +1028,10 @@ impl GroupedHashAggregateStream { if let Some(emit_to) = self.group_ordering.oom_emit_to(n) && let Some(batch) = self.emit(emit_to, false)? { + self.early_emit_count + .as_ref() + .expect("early emit metric exists for partial aggregation") + .add(1); return Ok(Some(ExecutionState::ProducingOutput(batch))); } Err(oom) @@ -1604,6 +1615,39 @@ mod tests { let per_aggregate_time = metrics.sum_by_name("agg_expr_0_arguments_time"); assert!(per_aggregate_time.is_some()); assert!(per_aggregate_time.unwrap().as_usize() > 0); + assert_eq!( + metrics.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 = + GroupedHashAggregateStream::new(&aggregate_exec, &no_skip_task_ctx, 0)?; + while let Some(result) = stream.next().await { + result?; + } + + assert_eq!( + aggregate_exec + .metrics() + .unwrap() + .sum_by_name("early_emit_count") + .unwrap() + .as_usize(), + 2 + ); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 7ebf32c3edfc6..ad1a61134ade0 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 ba08c4b003195..e46a5acae409c 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. @@ -5071,6 +5081,16 @@ mod tests { } } + assert_eq!( + aggregate + .metrics() + .unwrap() + .sum_by_name("early_emit_count") + .unwrap() + .as_usize(), + 1 + ); + Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 9e93a111a6466..dc1ee9227d195 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -116,6 +116,8 @@ pub(crate) struct OrderedPartialAggregateStream { reservation: MemoryReservation, baseline_metrics: BaselineMetrics, reduction_factor: metrics::RatioMetrics, + /// Number of times accumulated states were emitted due to memory pressure. + early_emit_count: metrics::Count, table: Option>, } @@ -138,6 +140,8 @@ impl OrderedPartialAggregateStream { 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 table = OrderedAggregateTable::::new( agg, @@ -159,6 +163,7 @@ impl OrderedPartialAggregateStream { reservation, baseline_metrics, reduction_factor, + early_emit_count, table: Some(table), }) } @@ -307,6 +312,7 @@ impl OrderedPartialAggregateStream { return Err(oom); }; self.reservation.try_resize(table.memory_size())?; + self.early_emit_count.add(1); Ok(Some(batch)) } diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d8f1447bc9521..cc762146a13d9 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), From e40fb0f16676540d8a87e2ad3c0acd45e8376a63 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Mon, 7 Sep 2026 01:07:57 +0530 Subject: [PATCH 2/4] style: remove redundant metric comment --- .../physical-plan/src/aggregates/ordered_partial_stream.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index dc1ee9227d195..156845da75d98 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -116,7 +116,6 @@ pub(crate) struct OrderedPartialAggregateStream { reservation: MemoryReservation, baseline_metrics: BaselineMetrics, reduction_factor: metrics::RatioMetrics, - /// Number of times accumulated states were emitted due to memory pressure. early_emit_count: metrics::Count, table: Option>, } From 4ad9204c305c6a84787e5d676b0ac05954acc6f9 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Mon, 7 Sep 2026 12:22:33 +0530 Subject: [PATCH 3/4] fix: narrow early emit metric scope --- .../src/aggregates/grouped_hash_stream.rs | 44 ------------------- .../physical-plan/src/aggregates/mod.rs | 10 ----- .../src/aggregates/ordered_partial_stream.rs | 5 --- 3 files changed, 59 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 41239c710bbad..3f6f3f8ce815b 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -385,9 +385,6 @@ pub(crate) struct GroupedHashAggregateStream { /// Reduction factor metric, calculated as `output_rows/input_rows` (only for partial aggregation) reduction_factor: Option, - - /// Number of times accumulated states were emitted due to memory pressure. - early_emit_count: Option, } impl GroupedHashAggregateStream { @@ -614,9 +611,6 @@ impl GroupedHashAggregateStream { } else { None }; - let early_emit_count = (agg.mode == AggregateMode::Partial).then(|| { - MetricBuilder::new(&agg.metrics).counter("early_emit_count", partition) - }); Ok(GroupedHashAggregateStream { schema: agg_schema, @@ -643,7 +637,6 @@ impl GroupedHashAggregateStream { group_values_soft_limit: agg.limit_options().map(|config| config.limit()), skip_aggregation_probe, reduction_factor, - early_emit_count, }) } } @@ -1028,10 +1021,6 @@ impl GroupedHashAggregateStream { if let Some(emit_to) = self.group_ordering.oom_emit_to(n) && let Some(batch) = self.emit(emit_to, false)? { - self.early_emit_count - .as_ref() - .expect("early emit metric exists for partial aggregation") - .add(1); return Ok(Some(ExecutionState::ProducingOutput(batch))); } Err(oom) @@ -1615,39 +1604,6 @@ mod tests { let per_aggregate_time = metrics.sum_by_name("agg_expr_0_arguments_time"); assert!(per_aggregate_time.is_some()); assert!(per_aggregate_time.unwrap().as_usize() > 0); - assert_eq!( - metrics.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 = - GroupedHashAggregateStream::new(&aggregate_exec, &no_skip_task_ctx, 0)?; - while let Some(result) = stream.next().await { - result?; - } - - assert_eq!( - aggregate_exec - .metrics() - .unwrap() - .sum_by_name("early_emit_count") - .unwrap() - .as_usize(), - 2 - ); Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index e46a5acae409c..06d542a389998 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -5081,16 +5081,6 @@ mod tests { } } - assert_eq!( - aggregate - .metrics() - .unwrap() - .sum_by_name("early_emit_count") - .unwrap() - .as_usize(), - 1 - ); - Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs index 156845da75d98..9e93a111a6466 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -116,7 +116,6 @@ pub(crate) struct OrderedPartialAggregateStream { reservation: MemoryReservation, baseline_metrics: BaselineMetrics, reduction_factor: metrics::RatioMetrics, - early_emit_count: metrics::Count, table: Option>, } @@ -139,8 +138,6 @@ impl OrderedPartialAggregateStream { 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 table = OrderedAggregateTable::::new( agg, @@ -162,7 +159,6 @@ impl OrderedPartialAggregateStream { reservation, baseline_metrics, reduction_factor, - early_emit_count, table: Some(table), }) } @@ -311,7 +307,6 @@ impl OrderedPartialAggregateStream { return Err(oom); }; self.reservation.try_resize(table.memory_size())?; - self.early_emit_count.add(1); Ok(Some(batch)) } From a5e32f789a5abaf1d01539b42c4be96c660d3f63 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Mon, 7 Sep 2026 12:30:00 +0530 Subject: [PATCH 4/4] test: cover early emit metric in SQL --- .../sqllogictest/test_files/aggregate_memory_spill.slt | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 79c6cb7a5155e..53d515bf7fd13 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