This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24996-4d87bafcdf1e0a038ffe0fa329ec5bc880e0ea07 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0da2151ca4e4af984d89223a7c9696b186d88c72 Author: Yash Bhisikar <[email protected]> AuthorDate: Fri Sep 11 08:15:52 2026 +0000 feat: add early emit count metric to partial aggregates (#24996) ## Which issue does this PR close? - Closes #24995. ## Why make this change? When DataFusion is working on a partial `GROUP BY` and memory fills up, it sends out the work completed so far and starts collecting again. Today, `EXPLAIN ANALYZE` does not show that this happened. This makes it hard to tell why the partial `GROUP BY` produced more rows than expected. ## What changed? This PR adds an `early_emit_count` number to `EXPLAIN ANALYZE`. It goes up once each time we send out partial results because memory is full. If one set of partial results is sent in several smaller pieces, the number still goes up only once. This change does not affect query results and does not add new logs. ## How was it tested? Tests cover each way DataFusion runs a partial `GROUP BY`. They check that: - the number goes up once for each memory-related send - splitting one send into smaller pieces does not increase it more than once - other ways of sending partial results do not increase it ## Are there any user-facing changes? `EXPLAIN ANALYZE` can now show `early_emit_count` for partial `GROUP BY` steps. --- .../physical-plan/src/aggregates/hash_stream.rs | 46 ++++++++++++++++++++++ datafusion/physical-plan/src/aggregates/mod.rs | 10 +++++ .../src/aggregates/partial_reduce_stream.rs | 23 +++++++---- .../test_files/aggregate_memory_spill.slt | 8 ++-- 4 files changed, 77 insertions(+), 10 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 2bcff82c80..152082d88d 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -163,6 +163,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<SkipAggregationProbe>, @@ -389,6 +392,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::<PartialMarker>::new( agg, @@ -433,6 +438,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), @@ -476,6 +482,7 @@ impl PartialHashAggregateStream { ) })?; + self.early_emit_count.add(1); timer.done(); self.emit_on_memory_pressure( materialized_group_states, @@ -1129,6 +1136,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 b628931fdd..3ed93e09ce 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -4831,6 +4831,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 286a324e96..d8ea2f711e 100644 --- a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -30,7 +30,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}; @@ -88,6 +88,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<PartialReduceHashAggregateState>, @@ -188,6 +191,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::<PartialReduceMarker>::new( agg, @@ -210,6 +215,7 @@ impl PartialReduceHashAggregateStream { batch_size, baseline_metrics, reservation, + early_emit_count, state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }), }) } @@ -315,12 +321,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 79c6cb7a51..53d515bf7f 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 ( <slt:ignore> 06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>spilled_rows=<slt:ignore>K,<slt:ignore>] <slt:ignore> +08)--------------AggregateExec: mode=Partial, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>early_emit_count=<slt:ignore>] +<slt:ignore> # Restore settings to slt runner defaults statement ok --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
