fix: output_bytes metric in hash aggregation

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

@@ -20,7 +20,10 @@
 use std::{borrow::Cow, collections::BTreeMap, sync::Arc, task::Poll};
 
 use arrow::record_batch::RecordBatch;
-use datafusion_common::{Result, utils::memory::get_record_batch_memory_size};
+use datafusion_common::{
+    Result,
+    utils::memory::{RecordBatchMemoryCounter, get_record_batch_memory_size},
+};
 
 use super::{
     Count, ExecutionPlanMetricsSet, Metric, MetricBuilder, MetricsSet, Time, Timestamp,
@@ -312,6 +315,24 @@ impl SplitMetrics {
     }
 }
 
+#[derive(Debug, Default)]
+pub struct RecordBatchMemoryMetrics(RecordBatchMemoryCounter);
+
+impl RecordBatchMemoryMetrics {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Similar to RecordBatch.record_output, but deduplicating across batches to avoid
+    /// output_size inflation due to shared memory
+    pub fn record_output(&mut self, batch: &RecordBatch, bm: &BaselineMetrics) {
+        bm.record_output(batch.num_rows());
+        let n_bytes = self.0.count_batch(batch);
+        bm.output_bytes.add(n_bytes);
+        bm.output_batches.add(1);
+    }
+}
+
 /// Trait for things that produce output rows as a result of execution.
 pub trait RecordOutput {
     /// Record that some number of output rows have been produced

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

@@ -37,7 +37,9 @@ use std::{
 
 // public exports
 
-pub use baseline::{BaselineMetrics, RecordOutput, SpillMetrics, SplitMetrics};
+pub use baseline::{
+    BaselineMetrics, RecordBatchMemoryMetrics, RecordOutput, SpillMetrics, SplitMetrics,
+};
 pub use builder::MetricBuilder;
 pub use custom::CustomMetricValue;
 pub use elapsed_compute::{ElapsedComputeFuture, ElapsedComputeFutureExt};

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

@@ -33,6 +33,7 @@ use crate::aggregates::order::GroupOrdering;
 use crate::aggregates::{
     AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
 };
+use crate::metrics::{BaselineMetrics, RecordBatchMemoryMetrics};
 
 /// Marker for raw rows -> partial state aggregation.
 pub(in crate::aggregates) struct PartialMarker;
@@ -219,6 +220,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
     pub(super) fn next_output_batch_inner(
         &mut self,
         materialize_accumulator_fn: MaterializeAccumulatorFn,
+        bm: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
         let output_schema = Arc::clone(&self.output_schema);
         let batch_size = self.batch_size;
@@ -253,7 +255,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
                 }
             };
 
-        let batch = output.next_batch(batch_size);
+        let batch = output.next_batch(batch_size, bm);
         if output.is_exhausted() {
             self.state = AggregateHashTableState::Done;
         } else {
@@ -431,14 +433,26 @@ pub(super) enum AggregateHashTableState {
 pub(super) struct MaterializedAggregateOutput {
     batch: RecordBatch,
     offset: usize,
+    /// Deduplicates buffer bytes across the slices sliced from `batch`, since
+    /// they share the same underlying buffers and must only be counted once
+    /// in `output_bytes`.
+    record_batch_metrics: RecordBatchMemoryMetrics,
 }
 
 impl MaterializedAggregateOutput {
     pub(super) fn new(batch: RecordBatch) -> Self {
-        Self { batch, offset: 0 }
+        Self {
+            batch,
+            offset: 0,
+            record_batch_metrics: RecordBatchMemoryMetrics::new(),
+        }
     }
 
-    pub(super) fn next_batch(&mut self, batch_size: usize) -> Option<RecordBatch> {
+    pub(super) fn next_batch(
+        &mut self,
+        batch_size: usize,
+        bm: &BaselineMetrics,
+    ) -> Option<RecordBatch> {
         debug_assert!(batch_size > 0);
         if self.is_exhausted() {
             return None;
@@ -447,6 +461,7 @@ impl MaterializedAggregateOutput {
         let length = batch_size.min(self.batch.num_rows() - self.offset);
         let batch = self.batch.slice(self.offset, length);
         self.offset += length;
+        self.record_batch_metrics.record_output(&batch, bm);
         Some(batch)
     }
 
@@ -633,11 +648,22 @@ mod tests {
             vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
         )?;
         let mut output = MaterializedAggregateOutput::new(batch);
-
-        assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]);
-        assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]);
-        assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![5]);
-        assert!(output.next_batch(2).is_none());
+        let metrics_set = crate::metrics::ExecutionPlanMetricsSet::new();
+        let bm = BaselineMetrics::new(&metrics_set, 0);
+
+        assert_eq!(
+            int32_values(&output.next_batch(2, &bm).unwrap(), 0),
+            vec![1, 2]
+        );
+        assert_eq!(
+            int32_values(&output.next_batch(2, &bm).unwrap(), 0),
+            vec![3, 4]
+        );
+        assert_eq!(
+            int32_values(&output.next_batch(2, &bm).unwrap(), 0),
+            vec![5]
+        );
+        assert!(output.next_batch(2, &bm).is_none());
         assert!(output.is_exhausted());
 
         Ok(())

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

@@ -20,6 +20,7 @@ use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
 
 use crate::aggregates::AggregateExec;
+use crate::metrics::BaselineMetrics;
 
 use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator};
 
@@ -54,8 +55,9 @@ impl AggregateHashTable<FinalMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        bm: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
+        self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns, bm)
     }
 
     /// Final aggregation consumes partial aggregate states and merges them into

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

@@ -20,6 +20,7 @@ use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
 
 use crate::aggregates::AggregateExec;
+use crate::metrics::BaselineMetrics;
 
 use super::common::{AggregateHashTable, HashAggregateAccumulator, PartialReduceMarker};
 
@@ -48,8 +49,9 @@ impl AggregateHashTable<PartialReduceMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        bm: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_inner(HashAggregateAccumulator::state)
+        self.next_output_batch_inner(HashAggregateAccumulator::state, bm)
     }
 
     /// Partial-reduce aggregation consumes partial aggregate states and merges

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

@@ -27,6 +27,7 @@ use datafusion_common::{Result, assert_eq_or_internal_err};
 use crate::aggregates::group_values::new_group_values;
 use crate::aggregates::order::GroupOrdering;
 use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal};
+use crate::metrics::BaselineMetrics;
 
 use super::common::{
     AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState,
@@ -64,8 +65,9 @@ impl AggregateHashTable<PartialMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        bm: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_inner(HashAggregateAccumulator::state)
+        self.next_output_batch_inner(HashAggregateAccumulator::state, bm)
     }
 
     pub(in crate::aggregates) fn can_skip_aggregation(&self) -> bool {

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

@@ -20,6 +20,7 @@ use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
 
 use crate::aggregates::AggregateExec;
+use crate::metrics::BaselineMetrics;
 
 use super::common::{AggregateHashTable, HashAggregateAccumulator, SingleMarker};
 
@@ -54,8 +55,9 @@ impl AggregateHashTable<SingleMarker> {
     /// exhausted, and an internal error if polled in the `Building` state.
     pub(in crate::aggregates) fn next_output_batch(
         &mut self,
+        bm: &BaselineMetrics,
     ) -> Result<Option<RecordBatch>> {
-        self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
+        self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns, bm)
     }
 
     /// Single aggregation consumes raw input rows and updates the table's