From da02bdcd0e4ceab3fdc83f7e4b50ffa2260cb17b Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Wed, 5 Aug 2026 13:38:02 +0300 Subject: [PATCH 1/4] feat: add sqllogictest for output_bytes metric --- .../test_files/aggregate_output_bytes.slt | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 datafusion/sqllogictest/test_files/aggregate_output_bytes.slt diff --git a/datafusion/sqllogictest/test_files/aggregate_output_bytes.slt b/datafusion/sqllogictest/test_files/aggregate_output_bytes.slt new file mode 100644 index 0000000000000..f1912ab36990c --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_output_bytes.slt @@ -0,0 +1,91 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# -------------------------------------------- +# Regression test for the hash aggregation `output_bytes` metric. +# +# The hash aggregate's grouping table materializes its emitted output as a +# single RecordBatch and then slices it into `batch_size`-sized chunks across +# successive `poll_next` calls. Those slices share the same underlying +# buffers (Arrow `.slice()` is zero-copy), so each buffer must only be +# counted once in `output_bytes`, no matter how many slices reference it. +# +# Using `RecordBatch::record_output` (as the streams used to) recomputes +# memory usage from scratch on every slice, with no memory of buffers +# already counted -- so all N slices of one materialization each count the +# shared buffers again, inflating `output_bytes` by ~Nx. `RecordBatchMemoryMetrics` +# fixes this by tracking already-counted buffers across all the slices of one +# materialization, so a shared buffer is only counted on the first slice +# that references it. +# +# The two `output_bytes` values below (1264.0 B vs 1344.0 B) are close but +# not bit-for-bit equal -- `batch_size` also affects how `DataSourceExec` +# chunks the *input* to the aggregator, independently of how the aggregate's +# own output gets sliced, which shifts the exact materialized size by a few +# bytes. That's expected and not what this test checks. What matters is that +# `output_bytes` does not scale with the number of output batches: it must +# stay at `1344.0 B` when sliced into 3 batches, not jump to ~3x that value. +# -------------------------------------------- + +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.explain.analyze_level = dev; + +statement ok +CREATE TABLE agg_output_bytes_src AS +SELECT value AS group_col, 1::BIGINT AS value_col +FROM generate_series(0, 29) AS t(value); + +# batch_size = 100: the 30-row emit fits in one output batch, no slicing. +statement ok +set datafusion.execution.batch_size = 100; + +query TT +EXPLAIN ANALYZE +SELECT group_col, COUNT(value_col) FROM agg_output_bytes_src GROUP BY group_col; +---- +Plan with Metrics +01)AggregateExec: mode=Single, gby=[group_col@0 as group_col], aggr=[count(agg_output_bytes_src.value_col)], metrics=[output_rows=30, elapsed_compute=, output_bytes=1264.0 B, output_batches=1, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, aggregate_arguments_time=, aggregation_time=, emitting_time=, time_calculating_group_ids=] +02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +# batch_size = 10: the same 30-row emit is sliced into 3 output batches, all +# sharing buffers with the same underlying materialized batch. output_bytes +# must stay at 1344.0 B, not be inflated to ~3x that (see header comment). +statement ok +set datafusion.execution.batch_size = 10; + +query TT +EXPLAIN ANALYZE +SELECT group_col, COUNT(value_col) FROM agg_output_bytes_src GROUP BY group_col; +---- +Plan with Metrics +01)AggregateExec: mode=Single, gby=[group_col@0 as group_col], aggr=[count(agg_output_bytes_src.value_col)], metrics=[output_rows=30, elapsed_compute=, output_bytes=1344.0 B, output_batches=3, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, aggregate_arguments_time=, aggregation_time=, emitting_time=, time_calculating_group_ids=] +02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +statement ok +DROP TABLE agg_output_bytes_src; + +statement ok +reset datafusion.execution.batch_size; + +statement ok +reset datafusion.explain.analyze_level; + +statement ok +set datafusion.execution.target_partitions = 4; From 9dff5ce694901ef618c44c78556535e29f2fcca4 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Tue, 8 Sep 2026 14:16:52 +0300 Subject: [PATCH 2/4] feat: show the undercounting issue caused by allocation address reuse --- datafusion/common/src/utils/memory.rs | 10 ++ .../physical-plan/src/aggregates/mod.rs | 156 ++++++++++++++++++ 2 files changed, 166 insertions(+) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 119d2e73ee7c3..be6f4cb7e7439 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/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 diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index ba08c4b003195..5b8ddcef04b65 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -8467,4 +8467,160 @@ mod tests { assert!(agg.set_dynamic_filter(df).is_err()); Ok(()) } + + /// Regression test for the aggregate's `output_bytes` metric. + /// + /// The aggregate materializes each emit as one batch and hands it + /// downstream in `batch_size` slices that share the materialized buffers. + /// Two failure modes are guarded against: + /// + /// * Recording every slice with `RecordBatch::record_output` charges the + /// shared buffers once per slice, inflating `output_bytes`. + /// * Deduplicating buffers by address for the life of the stream + /// undercounts, because a downstream consumer drops each batch and the + /// allocator recycles its address for the next emit. + /// + /// The fix records rows and bytes once per materialization and only the + /// batch count per slice, so the result must not depend on whether the + /// consumer retains or drops batches, nor on how many slices an emit is + /// cut into. + /// + /// Sorted partial aggregation emits one materialization per completed key + /// range, giving many independently allocated emits. `batch_size` 8192 + /// leaves them whole; `batch_size` 100 slices each 512-row emit into 6. + #[tokio::test] + async fn output_bytes_metric_is_exact_for_sliced_and_dropped_batches() -> Result<()> { + use datafusion_common::utils::memory::RecordBatchMemoryCounter; + + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + + // 200 input batches, each with 512 distinct, strictly increasing keys. + let keys_per_batch = 512; + let input_batches = (0..200i32) + .map(|i| { + let start = i * keys_per_batch; + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from( + (start..start + keys_per_batch).collect::>(), + )), + Arc::new(Int64Array::from(vec![1i64; keys_per_batch as usize])), + ], + ) + }) + .collect::, _>>()?; + + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("key", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[input_batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input: Arc = + Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value)") + .build()?, + )]; + let new_aggregate = || -> Result> { + Ok(Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggr_expr.clone(), + vec![None], + Arc::clone(&input), + Arc::clone(&schema), + )?)) + }; + assert_eq!(new_aggregate()?.input_order_mode(), &InputOrderMode::Sorted); + + let sum_metric = |aggregate: &AggregateExec, pick: fn(&MetricValue) -> bool| { + aggregate + .metrics() + .unwrap() + .sum(|m| pick(m.value())) + .unwrap() + .as_usize() + }; + + for (batch_size, expect_sliced) in [(8192, false), (100, true)] { + let task_ctx = + Arc::new(TaskContext::default().with_session_config( + SessionConfig::new().with_batch_size(batch_size), + )); + + // Pass 1: retain every output batch. While they are all alive no + // address can be recycled, so a `RecordBatchMemoryCounter` over + // them is the exact deduplicated size of what was emitted. + let aggregate = new_aggregate()?; + let retained = collect(aggregate.execute(0, Arc::clone(&task_ctx))?).await?; + let output_batches = retained.len(); + assert!(output_batches > 1, "test needs many separate emits"); + let sliced = retained.iter().any(|b| b.num_rows() == batch_size); + assert_eq!(sliced, expect_sliced, "batch_size={batch_size}"); + let mut ground_truth = RecordBatchMemoryCounter::new(); + for batch in &retained { + ground_truth.count_batch(batch); + } + let expected_bytes = ground_truth.memory_usage(); + let expected_rows: usize = retained.iter().map(|b| b.num_rows()).sum(); + assert_eq!(expected_rows, 200 * keys_per_batch as usize); + + let reported = + sum_metric(&aggregate, |v| matches!(v, MetricValue::OutputBytes(_))); + assert_eq!( + reported, expected_bytes, + "batch_size={batch_size}, retained consumer: output_bytes" + ); + assert_eq!( + sum_metric(&aggregate, |v| matches!(v, MetricValue::OutputBatches(_))), + output_batches, + "batch_size={batch_size}, retained consumer: output_batches" + ); + assert_eq!( + sum_metric(&aggregate, |v| matches!(v, MetricValue::OutputRows(_))), + expected_rows, + "batch_size={batch_size}, retained consumer: output_rows" + ); + drop(retained); + + // Pass 2: drop each batch as it arrives, like a real downstream + // operator. The metrics must be identical to pass 1. + let aggregate = new_aggregate()?; + let mut stream = aggregate.execute(0, task_ctx)?; + let mut dropped_batches = 0usize; + while let Some(batch) = stream.next().await { + batch?; + dropped_batches += 1; + } + assert_eq!(dropped_batches, output_batches); + + let reported = + sum_metric(&aggregate, |v| matches!(v, MetricValue::OutputBytes(_))); + assert_eq!( + reported, expected_bytes, + "batch_size={batch_size}, dropping consumer: output_bytes" + ); + assert_eq!( + sum_metric(&aggregate, |v| matches!(v, MetricValue::OutputBatches(_))), + output_batches, + "batch_size={batch_size}, dropping consumer: output_batches" + ); + assert_eq!( + sum_metric(&aggregate, |v| matches!(v, MetricValue::OutputRows(_))), + expected_rows, + "batch_size={batch_size}, dropping consumer: output_rows" + ); + } + Ok(()) + } } From 9a3a593e344f508e2d2f7a28369cd836849102d3 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Tue, 8 Sep 2026 19:15:13 +0300 Subject: [PATCH 3/4] fix: add ignore fields in slt file --- datafusion/sqllogictest/test_files/aggregate_output_bytes.slt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/sqllogictest/test_files/aggregate_output_bytes.slt b/datafusion/sqllogictest/test_files/aggregate_output_bytes.slt index f1912ab36990c..ca78e63c80672 100644 --- a/datafusion/sqllogictest/test_files/aggregate_output_bytes.slt +++ b/datafusion/sqllogictest/test_files/aggregate_output_bytes.slt @@ -61,7 +61,7 @@ EXPLAIN ANALYZE SELECT group_col, COUNT(value_col) FROM agg_output_bytes_src GROUP BY group_col; ---- Plan with Metrics -01)AggregateExec: mode=Single, gby=[group_col@0 as group_col], aggr=[count(agg_output_bytes_src.value_col)], metrics=[output_rows=30, elapsed_compute=, output_bytes=1264.0 B, output_batches=1, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, aggregate_arguments_time=, aggregation_time=, emitting_time=, time_calculating_group_ids=] +01)AggregateExec: mode=Single, gby=[group_col@0 as group_col], aggr=[count(agg_output_bytes_src.value_col)], metrics=[output_rows=30, elapsed_compute=, output_bytes=1264.0 B, output_batches=1, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, agg_expr_0_arguments_time=, agg_expr_0_evaluate_time=, agg_expr_0_merge_time=, agg_expr_0_state_time=, agg_expr_0_update_time=, aggregate_arguments_time=, aggregation_time=, emitting_time=, time_calculating_group_ids=] 02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] # batch_size = 10: the same 30-row emit is sliced into 3 output batches, all @@ -75,7 +75,7 @@ EXPLAIN ANALYZE SELECT group_col, COUNT(value_col) FROM agg_output_bytes_src GROUP BY group_col; ---- Plan with Metrics -01)AggregateExec: mode=Single, gby=[group_col@0 as group_col], aggr=[count(agg_output_bytes_src.value_col)], metrics=[output_rows=30, elapsed_compute=, output_bytes=1344.0 B, output_batches=3, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, aggregate_arguments_time=, aggregation_time=, emitting_time=, time_calculating_group_ids=] +01)AggregateExec: mode=Single, gby=[group_col@0 as group_col], aggr=[count(agg_output_bytes_src.value_col)], metrics=[output_rows=30, elapsed_compute=, output_bytes=1344.0 B, output_batches=3, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, agg_expr_0_arguments_time=, agg_expr_0_evaluate_time=, agg_expr_0_merge_time=, agg_expr_0_state_time=, agg_expr_0_update_time=, aggregate_arguments_time=, aggregation_time=, emitting_time=, time_calculating_group_ids=] 02)--DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] statement ok From a1c65f0bad8812e2a7c4242762bcd811a078d439 Mon Sep 17 00:00:00 2001 From: Ariel Miculas Date: Tue, 8 Sep 2026 19:16:04 +0300 Subject: [PATCH 4/4] feat: another approach for fixing the hash agg metrics --- .../src/metrics/baseline.rs | 17 ++++++++ .../aggregates/aggregate_hash_table/common.rs | 12 ++++++ .../aggregate_hash_table/final_table.rs | 3 ++ .../partial_reduce_table.rs | 3 ++ .../aggregate_hash_table/partial_table.rs | 3 ++ .../aggregate_hash_table/single_table.rs | 3 ++ .../src/aggregates/grouped_hash_stream.rs | 14 +++++-- .../src/aggregates/hash_stream.rs | 32 ++++++++------- .../src/aggregates/partial_reduce_stream.rs | 39 +++++++++++-------- .../src/aggregates/single_stream.rs | 11 +++--- 10 files changed, 98 insertions(+), 39 deletions(-) diff --git a/datafusion/physical-expr-common/src/metrics/baseline.rs b/datafusion/physical-expr-common/src/metrics/baseline.rs index 52ad4aac9fd98..5f203e557a22c 100644 --- a/datafusion/physical-expr-common/src/metrics/baseline.rs +++ b/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() { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs index 0c481a6b3bc0b..30bf0abb9602c 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/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 AggregateHashTable { /// /// This is a temporary solution until blocked state management is implemented: /// Issue: + /// + /// 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> { let output_schema = Arc::clone(&self.output_schema); let batch_size = self.batch_size; @@ -278,6 +286,7 @@ impl AggregateHashTable { 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 AggregateHashTable { }; 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 { diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs index 307215cfd2797..7750192dc2a01 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs +++ b/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 { /// 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> { self.next_output_batch_inner( HashAggregateAccumulator::evaluate_to_columns, AccumulatorPhase::Evaluate, + baseline_metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs index 2892d059332cf..70965e4691e3e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs +++ b/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 { @@ -52,10 +53,12 @@ impl AggregateHashTable { /// 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> { self.next_output_batch_inner( HashAggregateAccumulator::state, AccumulatorPhase::State, + baseline_metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs index bbc51ae666ab7..f0e68e632cee3 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs +++ b/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 { /// 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> { self.next_output_batch_inner( HashAggregateAccumulator::state, AccumulatorPhase::State, + baseline_metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs index 2d7dc2a63d086..6f374b2538bc6 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs +++ b/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 { /// 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> { self.next_output_batch_inner( HashAggregateAccumulator::evaluate_to_columns, AccumulatorPhase::Evaluate, + baseline_metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 3f6f3f8ce815b..cabbe066347d5 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/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`]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)) } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index 7ebf32c3edfc6..1fd99b73479e0 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -597,6 +597,10 @@ impl PartialHashAggregateStream { mut remaining_groups: RecordBatch, emitter: &mut TryEmitter, ) -> 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(); } } diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs index d8f1447bc9521..9773e0e941782 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, SpillMetrics}; use crate::stream::EmptyRecordBatchStream; use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; @@ -316,12 +316,18 @@ 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)) => { + // The batch is emitted in `batch_size` slices that share its + // buffers. Record rows and bytes once here; + // `handle_emitting_on_memory_pressure` counts each slice. + self.baseline_metrics.record_output_bytes(&remaining_groups); + 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), @@ -369,10 +375,10 @@ impl PartialReduceHashAggregateStream { }; debug_assert!(output_batch.num_rows() > 0); - ControlFlow::Break(( - Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))), - next_state, - )) + // Rows and bytes were recorded for the whole materialized batch in + // `resize_or_emit_early`; only count the emitted slice here. + self.baseline_metrics.output_batches().add(1); + ControlFlow::Break((Poll::Ready(Some(Ok(output_batch))), next_state)) } /// Handle ProducingOutput state - emit merged partial aggregate state batches. @@ -392,7 +398,11 @@ impl PartialReduceHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = original_state.hash_table_mut().next_output_batch(); + // The table records output metrics itself: rows and bytes once per + // materialization, batch count once per slice. + let result = original_state + .hash_table_mut() + .next_output_batch(&self.baseline_metrics); timer.done(); match result { @@ -409,10 +419,7 @@ impl PartialReduceHashAggregateStream { original_state }; - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) + ControlFlow::Break((Poll::Ready(Some(Ok(batch))), next_state)) } Ok(None) => { let _ = self.reservation.try_resize(0); diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index 40412385efbc7..25140ad46fb39 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -42,7 +42,7 @@ use super::aggregate_hash_table::{ use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; -use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::spill_manager::SpillManager; @@ -673,7 +673,9 @@ impl SingleHashAggregateStream { let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let timer = elapsed_compute.timer(); - let result = hash_table.next_output_batch(); + // The table records output metrics itself: rows and bytes once per + // materialization, batch count once per slice. + let result = hash_table.next_output_batch(&self.baseline_metrics); timer.done(); match result { @@ -692,10 +694,7 @@ impl SingleHashAggregateStream { SingleHashAggregateState::ProducingOutput { hash_table } }; - ControlFlow::Break(( - Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), - next_state, - )) + ControlFlow::Break((Poll::Ready(Some(Ok(batch))), next_state)) } Err(e) => Self::break_with_err(e), Ok(None) => {