feat: Fix output bytes metric in hash agg

datafusion/common/src/utils/memory.rs

@@ -165,6 +165,16 @@ pub fn get_record_batch_memory_size(batch: &RecordBatch) -> usize {
 /// batch's buffers are kept alive by the batch even when only a sub-range is
 /// referenced, so counting unique buffers in full reflects the memory the
 /// batches actually retain.
+///
+/// # Every counted batch must stay alive while the counter is in use
+///
+/// Buffers are identified by allocation address. Once a counted batch is
+/// dropped, the allocator may hand its address to a new allocation, which this
+/// counter would then wrongly skip as already counted. This makes the counter
+/// suitable for accounting batches an operator retains (a join build side, a
+/// buffered input), and unsuitable for a stream of batches that are handed
+/// downstream and forgotten. For that case, count each materialized batch once
+/// before slicing it instead.
 #[derive(Debug, Default)]
 pub struct RecordBatchMemoryCounter {
     /// Start addresses of `Buffer`s that have already been counted (instead of

datafusion/physical-expr-common/src/metrics/baseline.rs

@@ -129,6 +129,11 @@ impl BaselineMetrics {
         &self.output_batches
     }
 
+    /// return the metric for the total number of output bytes produced
+    pub fn output_bytes(&self) -> &Count {
+        &self.output_bytes
+    }
+
     /// Returns a derived metric that summarizes how unevenly `output_rows`
     /// are distributed across partitions.
     ///
@@ -203,6 +208,18 @@ impl BaselineMetrics {
         self.output_rows.add(num_rows);
     }
 
+    /// Record the rows and memory of a materialized batch whose contents will
+    /// be emitted downstream as one or more zero-copy slices.
+    ///
+    /// Slices share the materialized batch's buffers, so recording each slice
+    /// with [`RecordOutput::record_output`] would count those buffers once per
+    /// slice. Instead, call this once on the materialized batch, and for each
+    /// emitted slice only bump [`Self::output_batches`].
+    pub fn record_output_bytes(&self, batch: &RecordBatch) {
+        self.output_rows.add(batch.num_rows());
+        self.output_bytes.add(get_record_batch_memory_size(batch));
+    }
+
     /// If not previously recorded `done()`, record
     pub fn try_done(&self) {
         if self.end_time.value().is_none() {

datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs

@@ -38,6 +38,7 @@ use crate::aggregates::{
     AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
     group_id_array, max_duplicate_ordinal,
 };
+use crate::metrics::BaselineMetrics;
 
 use super::AggregateTableMetrics;
 
@@ -246,10 +247,17 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
     ///
     /// This is a temporary solution until blocked state management is implemented:
     /// Issue: <https://github.com/apache/datafusion/issues/7065>
+    ///
+    /// Output is materialized once and then handed out in `batch_size` slices
+    /// that share the materialized buffers. `baseline_metrics` therefore has
+    /// rows and bytes recorded once, when the batch is materialized, and its
+    /// batch count bumped once per returned slice. Callers must not record
+    /// the returned batches again.
     pub(super) fn next_output_batch_inner(
         &mut self,
         materialize_accumulator_fn: MaterializeAccumulatorFn,
         accumulator_phase: AccumulatorPhase,
+        baseline_metrics: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
         let output_schema = Arc::clone(&self.output_schema);
         let batch_size = self.batch_size;
@@ -278,6 +286,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
 
                     let batch = RecordBatch::try_new(output_schema, columns)?;
                     debug_assert!(batch.num_rows() > 0);
+                    baseline_metrics.record_output_bytes(&batch);
                     MaterializedAggregateOutput::new(batch)
                 }
                 AggregateHashTableState::OutputtingMaterialized(output) => output,
@@ -290,6 +299,9 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
             };
 
         let batch = output.next_batch(batch_size);
+        if batch.is_some() {
+            baseline_metrics.output_batches().add(1);
+        }
         if output.is_exhausted() {
             self.state = AggregateHashTableState::Done;
         } else {

datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs

@@ -25,6 +25,7 @@ use crate::aggregates::AggregateExec;
 use crate::aggregates::group_values::AccumulatorPhase;
 
 use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator};
+use crate::metrics::BaselineMetrics;
 
 /// Implementation specific to final aggregation, where the table stores partial
 /// aggregate states and the input rows are also partial states.
@@ -58,10 +59,12 @@ impl AggregateHashTable<FinalMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        baseline_metrics: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
         self.next_output_batch_inner(
             HashAggregateAccumulator::evaluate_to_columns,
             AccumulatorPhase::Evaluate,
+            baseline_metrics,
         )
     }
 

datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs

@@ -25,6 +25,7 @@ use crate::aggregates::AggregateExec;
 use crate::aggregates::group_values::AccumulatorPhase;
 
 use super::common::{AggregateHashTable, HashAggregateAccumulator, PartialReduceMarker};
+use crate::metrics::BaselineMetrics;
 
 /// Methods specific to the aggregate hash table used in the partial-reduce stage.
 impl AggregateHashTable<PartialReduceMarker> {
@@ -52,10 +53,12 @@ impl AggregateHashTable<PartialReduceMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        baseline_metrics: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
         self.next_output_batch_inner(
             HashAggregateAccumulator::state,
             AccumulatorPhase::State,
+            baseline_metrics,
         )
     }
 

datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs

@@ -30,6 +30,7 @@ use super::common::{
     AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState,
     HashAggregateAccumulator, PartialMarker, PartialSkipMarker,
 };
+use crate::metrics::BaselineMetrics;
 
 /// Implementation specific to partial aggregation, where the table stores
 /// partial aggregate states and the input rows are raw rows.
@@ -63,10 +64,12 @@ impl AggregateHashTable<PartialMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        baseline_metrics: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
         self.next_output_batch_inner(
             HashAggregateAccumulator::state,
             AccumulatorPhase::State,
+            baseline_metrics,
         )
     }
 

datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs

@@ -23,6 +23,7 @@ use crate::aggregates::AggregateExec;
 use crate::aggregates::group_values::AccumulatorPhase;
 
 use super::common::{AggregateHashTable, HashAggregateAccumulator, SingleMarker};
+use crate::metrics::BaselineMetrics;
 
 /// Implementation specific to single aggregation, where the table stores final
 /// aggregate values and the input rows are raw rows.
@@ -57,10 +58,12 @@ impl AggregateHashTable<SingleMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        baseline_metrics: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
         self.next_output_batch_inner(
             HashAggregateAccumulator::evaluate_to_columns,
             AccumulatorPhase::Evaluate,
+            baseline_metrics,
         )
     }
 

datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs

@@ -835,9 +835,10 @@ impl Stream for GroupedHashAggregateStream {
                     // Empty record batches should not be emitted.
                     // They need to be treated as  [`Option<RecordBatch>`]es and handled separately
                     debug_assert!(output_batch.num_rows() > 0);
-                    return Poll::Ready(Some(Ok(
-                        output_batch.record_output(&self.baseline_metrics)
-                    )));
+                    // Rows and bytes were recorded for the whole materialized
+                    // batch in `emit`; only count the emitted slice here.
+                    self.baseline_metrics.output_batches().add(1);
+                    return Poll::Ready(Some(Ok(output_batch)));
                 }
 
                 ExecutionState::Done => {
@@ -1109,6 +1110,13 @@ impl GroupedHashAggregateStream {
         let batch = RecordBatch::try_new(schema, output)?;
         debug_assert!(batch.num_rows() > 0);
 
+        if !spilling {
+            // The batch is emitted downstream in `batch_size` slices that share
+            // its buffers. Record rows and bytes once here; `ProducingOutput`
+            // counts each emitted slice in `output_batches`.
+            self.baseline_metrics.record_output_bytes(&batch);
+        }
+
         Ok(Some(batch))
     }
 

datafusion/physical-plan/src/aggregates/hash_stream.rs

@@ -597,6 +597,10 @@ impl PartialHashAggregateStream {
         mut remaining_groups: RecordBatch,
         emitter: &mut TryEmitter<RecordBatch, DataFusionError>,
     ) -> Result<()> {
+        // The slices below share this batch's buffers: record rows and bytes
+        // once here, and count each emitted slice in `output_batches`.
+        self.baseline_metrics.record_output_bytes(&remaining_groups);
+
         while remaining_groups.num_rows() > self.batch_size {
             // More batch to output, continue in the current state.
             let output = remaining_groups.slice(0, self.batch_size);
@@ -609,17 +613,15 @@ impl PartialHashAggregateStream {
             self.reduction_factor.add_part(output.num_rows());
             debug_assert!(output.num_rows() > 0);
 
-            emitter
-                .emit(output.record_output(&self.baseline_metrics))
-                .await;
+            self.baseline_metrics.output_batches().add(1);
+            emitter.emit(output).await;
         }
 
         self.reduction_factor.add_part(remaining_groups.num_rows());
         debug_assert!(remaining_groups.num_rows() > 0);
 
-        emitter
-            .emit(remaining_groups.record_output(&self.baseline_metrics))
-            .await;
+        self.baseline_metrics.output_batches().add(1);
+        emitter.emit(remaining_groups).await;
 
         Ok(())
     }
@@ -636,7 +638,10 @@ impl PartialHashAggregateStream {
         let mut timer = elapsed_compute.timer();
 
         loop {
-            let Some(batch) = hash_table.next_output_batch()? else {
+            // The table records output metrics itself: rows and bytes once
+            // per materialization, batch count once per slice.
+            let Some(batch) = hash_table.next_output_batch(&self.baseline_metrics)?
+            else {
                 // Only reachable when the table held no groups at all: a
                 // non-empty table always reports its last batch together with
                 // the `Done` state, which the `try_resize` below already zeroes.
@@ -653,9 +658,7 @@ impl PartialHashAggregateStream {
             self.reduction_factor.add_part(batch.num_rows());
 
             timer.done();
-            emitter
-                .emit(batch.record_output(&self.baseline_metrics))
-                .await;
+            emitter.emit(batch).await;
             timer = elapsed_compute.timer();
         }
     }
@@ -958,7 +961,10 @@ impl FinalHashAggregateStream {
         hash_table.start_output()?;
 
         loop {
-            let Some(batch) = hash_table.next_output_batch()? else {
+            // The table records output metrics itself: rows and bytes once
+            // per materialization, batch count once per slice.
+            let Some(batch) = hash_table.next_output_batch(&self.baseline_metrics)?
+            else {
                 // Only reachable when the table held no groups at all: a
                 // non-empty table always reports its last batch together with
                 // the `Done` state, which the `try_resize` below already zeroes.
@@ -972,9 +978,7 @@ impl FinalHashAggregateStream {
             self.reservation.try_resize(hash_table.memory_size())?;
 
             timer.done();
-            emitter
-                .emit(batch.record_output(&self.baseline_metrics))
-                .await;
+            emitter.emit(batch).await;
             timer = elapsed_compute.timer();
         }
     }