feat: add early emit count metric to partial aggregates

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<SkipAggregationProbe>,
 
@@ -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::<PartialMarker>::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(())
     }

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.

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<PartialReduceHashAggregateState>,
@@ -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::<PartialReduceMarker>::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),

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