From 2af67d85ed9493535ae0d3c395f71e0420553664 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 12:12:50 +0800 Subject: [PATCH 01/24] =?UTF-8?q?feat:=20add=20data=E2=80=91free=20SQL=20h?= =?UTF-8?q?arness=20for=20array=5Fagg=5Fdistinct=20benchmark?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added harness implementation: - `benchmarks/sql_benchmarks/array_agg_distinct/array_agg_distinct.suite` – defines the benchmark suite, test parameters, and execution configuration for the data‑free SQL harness. - `benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark` – contains the specific query benchmark (`q01`) that exercises the `array_agg(DISTINCT …)` workload without requiring any input data. - Workload characteristics: - Simulates **2 M range rows** → **1 M groups**. - Each group contains **2 rows** with **2 distinct values**, providing a realistic yet data‑free test scenario for aggregation performance. - add bench.sh wrapper for array_agg_distinct --- benchmarks/bench.sh | 12 ++++++++++++ .../array_agg_distinct/array_agg_distinct.suite | 11 +++++++++++ .../array_agg_distinct/benchmarks/q01.benchmark | 11 +++++++++++ 3 files changed, 34 insertions(+) create mode 100644 benchmarks/sql_benchmarks/array_agg_distinct/array_agg_distinct.suite create mode 100644 benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 419fd5be3ad2b..3fb9dcb37616f 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -164,6 +164,7 @@ nlj: Benchmark for simple nested loop joins, testing various hj: Benchmark for simple hash joins, testing various join scenarios smj: Benchmark for simple sort merge joins, testing various join scenarios dict: Benchmark for dictionary-encoded group-by scenarios +array_agg_distinct: 1M-group, two-row-per-group array_agg(DISTINCT) benchmark compile_profile: Compile and execute TPC-H across selected Cargo profiles, reporting timing and binary size @@ -651,6 +652,9 @@ main() { dict) run_dict ;; + array_agg_distinct) + run_array_agg_distinct + ;; compile_profile) run_compile_profile "${PROFILE_ARGS[@]}" ;; @@ -1661,6 +1665,14 @@ run_dict() { debug_run $CARGO_COMMAND --bin dfbench -- dict --iterations 5 -o "${RESULTS_FILE}" ${QUERY_ARG} ${LATENCY_ARG} } +# Runs the data-free high-cardinality array_agg(DISTINCT) SQL benchmark. +run_array_agg_distinct() { + echo "Running array_agg_distinct benchmark..." + debug_run env BENCH_NAME=array_agg_distinct \ + ${QUERY:+BENCH_QUERY="${QUERY}"} \ + bash -c "$SQL_CARGO_COMMAND" +} + compare_benchmarks() { BASE_RESULTS_DIR="${SCRIPT_DIR}/results" diff --git a/benchmarks/sql_benchmarks/array_agg_distinct/array_agg_distinct.suite b/benchmarks/sql_benchmarks/array_agg_distinct/array_agg_distinct.suite new file mode 100644 index 0000000000000..b114854f14f54 --- /dev/null +++ b/benchmarks/sql_benchmarks/array_agg_distinct/array_agg_distinct.suite @@ -0,0 +1,11 @@ +description = "High-cardinality array_agg(DISTINCT) SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- array_agg_distinct" +description = "Run the high-cardinality array_agg(DISTINCT) benchmark." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- array_agg_distinct --query 1 --iterations 5 --output /tmp/array_agg_distinct.json" +description = "Run five iterations and write comparable JSON results." diff --git a/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..b1987a14656ee --- /dev/null +++ b/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark @@ -0,0 +1,11 @@ +name Q01 +group array_agg_distinct + +expect_plan AggregateExec + +run +-- 1M groups, 2 rows/group, and 2 distinct values/group. `range` is end-exclusive. +-- This is data-free so comparisons isolate grouped array_agg(DISTINCT) execution. +SELECT value / 2 AS k, array_agg(DISTINCT value % 2) AS distinct_values +FROM range(2000000) +GROUP BY value / 2; From 8a0786eb0702d4ee70308e85d4caad09e5783920 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 16:51:07 +0800 Subject: [PATCH 02/24] chore: adjust benchmark scope from 1M to 100K groups Reduce the data-free grouped array_agg(DISTINCT) workload while preserving its two-rows-per-group and two-distinct-values-per-group shape. --- benchmarks/bench.sh | 2 +- .../array_agg_distinct/benchmarks/q01.benchmark | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 3fb9dcb37616f..bc2ac0781a422 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -164,7 +164,7 @@ nlj: Benchmark for simple nested loop joins, testing various hj: Benchmark for simple hash joins, testing various join scenarios smj: Benchmark for simple sort merge joins, testing various join scenarios dict: Benchmark for dictionary-encoded group-by scenarios -array_agg_distinct: 1M-group, two-row-per-group array_agg(DISTINCT) benchmark +array_agg_distinct: 100K-group, two-row-per-group array_agg(DISTINCT) benchmark compile_profile: Compile and execute TPC-H across selected Cargo profiles, reporting timing and binary size diff --git a/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark index b1987a14656ee..aaef5414112cc 100644 --- a/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark +++ b/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark @@ -4,8 +4,8 @@ group array_agg_distinct expect_plan AggregateExec run --- 1M groups, 2 rows/group, and 2 distinct values/group. `range` is end-exclusive. +-- 100K groups, 2 rows/group, and 2 distinct values/group. `range` is end-exclusive. -- This is data-free so comparisons isolate grouped array_agg(DISTINCT) execution. SELECT value / 2 AS k, array_agg(DISTINCT value % 2) AS distinct_values -FROM range(2000000) +FROM range(200000) GROUP BY value / 2; From d91528de76d8a42be578c298784bb2c9d945a776 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 11:46:47 +0800 Subject: [PATCH 03/24] bench: scale array_agg distinct groups --- .../array_agg_distinct/benchmarks/q01.benchmark | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark index aaef5414112cc..b1987a14656ee 100644 --- a/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark +++ b/benchmarks/sql_benchmarks/array_agg_distinct/benchmarks/q01.benchmark @@ -4,8 +4,8 @@ group array_agg_distinct expect_plan AggregateExec run --- 100K groups, 2 rows/group, and 2 distinct values/group. `range` is end-exclusive. +-- 1M groups, 2 rows/group, and 2 distinct values/group. `range` is end-exclusive. -- This is data-free so comparisons isolate grouped array_agg(DISTINCT) execution. SELECT value / 2 AS k, array_agg(DISTINCT value % 2) AS distinct_values -FROM range(200000) +FROM range(2000000) GROUP BY value / 2; From 7a7bd212ef829da16e43c5b19bf1e25a292a5335 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 21:39:07 +0800 Subject: [PATCH 04/24] empty commit From b741b38a55544900eee0c999bc1155903104a5c8 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:15:20 +0800 Subject: [PATCH 05/24] feat: add optional AggregateMetric API with lazy internal timers and comprehensive tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Introduce optional `AggregateMetric(s)` API with default no‑op setters for backward compatibility. - Add lazy‑stable internal metrics: `agg_expr_{i}_internal_{subphase}_time` for fine‑grained phase tracking. - Wire the metrics across all execution paths: stream, grouped, hash, ordered, and replay. - Implement `array_agg(DISTINCT)` distinct‑timer to measure distinct‑aggregation latency. - Extend test coverage: - Partition merge scenarios. - Repeated DISTINCT expression handling. - Update documentation to reflect new API, metric naming, and wiring details. --- datafusion/expr-common/src/accumulator.rs | 26 ++++ .../expr-common/src/groups_accumulator.rs | 9 ++ datafusion/expr/src/lib.rs | 4 +- .../functions-aggregate/src/array_agg.rs | 22 +++- datafusion/physical-expr/src/aggregate.rs | 22 +++- .../aggregates/aggregate_hash_table/common.rs | 23 +++- .../aggregate_hash_table/common_ordered.rs | 15 ++- .../aggregates/aggregate_hash_table/mod.rs | 9 +- .../aggregate_hash_table/partial_table.rs | 1 + .../src/aggregates/aggregate_stream.rs | 88 +++++++++++++- .../src/aggregates/group_values/metrics.rs | 111 +++++++++++++++++- .../src/aggregates/group_values/mod.rs | 2 +- .../src/aggregates/grouped_hash_stream.rs | 19 ++- .../physical-plan/src/aggregates/mod.rs | 14 ++- docs/source/user-guide/metrics.md | 11 ++ 15 files changed, 350 insertions(+), 26 deletions(-) diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index 7e9a4ae525ea3..47678137612fe 100644 --- a/datafusion/expr-common/src/accumulator.rs +++ b/datafusion/expr-common/src/accumulator.rs @@ -20,6 +20,26 @@ use arrow::array::ArrayRef; use datafusion_common::{Result, ScalarValue, internal_err}; use std::fmt::Debug; +use std::sync::Arc; +use std::time::Duration; + +/// A metric owned by one aggregate implementation. +/// +/// Aggregate implementations use this interface for optional internal +/// subphases. The execution engine owns metric registration and aggregation. +pub trait AggregateMetric: Debug + Send + Sync { + /// Adds elapsed time to this metric. + fn add_duration(&self, duration: Duration); +} + +/// Factory for optional metrics owned by one aggregate expression. +/// +/// `subphase` must be a stable static identifier. An implementation may request +/// no metrics. The execution engine assigns the aggregate expression identity. +pub trait AggregateMetrics: Debug + Send + Sync { + /// Returns the metric for an aggregate-owned internal subphase. + fn metric(&self, subphase: &'static str) -> Arc; +} /// Tracks an aggregate function's state. /// @@ -49,6 +69,12 @@ use std::fmt::Debug; /// [`merge_batch`]: Self::merge_batch /// [window function]: https://en.wikipedia.org/wiki/Window_function_(SQL) pub trait Accumulator: Send + Sync + Debug + std::any::Any { + /// Supplies optional metrics owned by this aggregate expression. + /// + /// The default preserves compatibility for accumulators without internal + /// submetrics. + fn set_metrics(&mut self, _metrics: Arc) {} + /// Updates the accumulator's state from its input. /// /// `values` contains the arguments to this aggregate function. diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index c682d7caf1b68..38cf7de245623 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -19,6 +19,9 @@ use arrow::array::{ArrayRef, BooleanArray}; use datafusion_common::{Result, exec_err, not_impl_err, utils::split_vec_min_alloc}; +use std::sync::Arc; + +use crate::accumulator::AggregateMetrics; /// Describes how many rows should be emitted during grouping. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -188,6 +191,12 @@ impl<'a> GroupSelection<'a> { /// [`Accumulator`]: crate::accumulator::Accumulator /// [Aggregating Millions of Groups Fast blog]: https://arrow.apache.org/blog/2023/08/05/datafusion_fast_grouping/ pub trait GroupsAccumulator: Send + std::any::Any { + /// Supplies optional metrics owned by this aggregate expression. + /// + /// The default preserves compatibility for accumulators without internal + /// submetrics. + fn set_metrics(&mut self, _metrics: Arc) {} + /// Updates the accumulator's state from its arguments, encoded as /// a vector of [`ArrayRef`]s. /// diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index a904422989942..07a4faa1eec92 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -101,7 +101,9 @@ pub use datafusion_doc::{ DocSection, Documentation, DocumentationBuilder, aggregate_doc_sections, scalar_doc_sections, window_doc_sections, }; -pub use datafusion_expr_common::accumulator::Accumulator; +pub use datafusion_expr_common::accumulator::{ + Accumulator, AggregateMetric, AggregateMetrics, +}; pub use datafusion_expr_common::columnar_value::ColumnarValue; pub use datafusion_expr_common::groups_accumulator::{ EmitTo, GroupSelection, GroupsAccumulator, diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 97a2520eb2b9a..aee21799ea009 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -38,13 +38,14 @@ use datafusion_common::utils::{ SingleRowListArrayBuilder, compare_rows, get_row_at_idx, take_function_args, }; use datafusion_common::{ - Result, ScalarValue, assert_eq_or_internal_err, exec_err, internal_err, + Result, ScalarValue, assert_eq_or_internal_err, exec_err, instant::Instant, + internal_err, }; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, Signature, - Volatility, + Accumulator, AggregateMetric, AggregateMetrics, AggregateUDFImpl, Documentation, + EmitTo, GroupsAccumulator, Signature, Volatility, }; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filter_to_nulls; use datafusion_functions_aggregate_common::merge_arrays::merge_ordered_arrays; @@ -853,6 +854,7 @@ pub struct DistinctArrayAggAccumulator { datatype: DataType, sort_options: Option, ignore_nulls: bool, + distinct_metric: Option>, } /// Returns `true` if `dt` is, or recursively contains, a `Dictionary` type. @@ -888,6 +890,7 @@ impl DistinctArrayAggAccumulator { datatype: datatype.clone(), sort_options, ignore_nulls, + distinct_metric: None, }) } @@ -914,6 +917,10 @@ impl DistinctArrayAggAccumulator { } impl Accumulator for DistinctArrayAggAccumulator { + fn set_metrics(&mut self, metrics: Arc) { + self.distinct_metric = Some(metrics.metric("distinct")); + } + fn state(&mut self) -> Result> { Ok(vec![self.evaluate()?]) } @@ -948,6 +955,8 @@ impl Accumulator for DistinctArrayAggAccumulator { return Ok(()); } + let distinct_metric = self.distinct_metric.clone(); + let distinct_start = Instant::now(); self.ensure_state(col.data_type())?; // Encode the entire incoming batch into rows_buffer in one pass. @@ -994,6 +1003,9 @@ impl Accumulator for DistinctArrayAggAccumulator { } } } + if let Some(metric) = distinct_metric { + metric.add_duration(distinct_start.elapsed()); + } Ok(()) } @@ -1167,7 +1179,9 @@ impl Accumulator for DistinctArrayAggAccumulator { } fn size(&self) -> usize { - size_of_val(self) + // `distinct_metric` is execution-owned observability state, not + // accumulator state retained by this aggregate. + (size_of_val(self) - size_of_val(&self.distinct_metric)) + self .state .as_ref() diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index 5f045ca8c7277..ec9b2c46dd8e6 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -53,7 +53,7 @@ use datafusion_expr::expr::{ }; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{AggregateUDF, Expr, ReversedUDAF, SetMonotonicity}; -use datafusion_expr_common::accumulator::Accumulator; +use datafusion_expr_common::accumulator::{Accumulator, AggregateMetrics}; use datafusion_expr_common::groups_accumulator::GroupsAccumulator; use datafusion_expr_common::type_coercion::aggregates::check_arg_count; use datafusion_functions_aggregate_common::accumulator::{ @@ -747,6 +747,16 @@ impl AggregateFunctionExpr { self.fun.accumulator(acc_args) } + /// Creates an accumulator and supplies optional aggregate-owned metrics. + pub fn create_accumulator_with_metrics( + &self, + metrics: Arc, + ) -> Result> { + let mut accumulator = self.create_accumulator()?; + accumulator.set_metrics(metrics); + Ok(accumulator) + } + /// the field of the final result of this aggregation. pub fn state_fields(&self) -> Result> { let args = StateFieldsArgs { @@ -928,6 +938,16 @@ impl AggregateFunctionExpr { self.fun.create_groups_accumulator(args) } + /// Creates a groups accumulator and supplies optional aggregate-owned metrics. + pub fn create_groups_accumulator_with_metrics( + &self, + metrics: Arc, + ) -> Result> { + let mut accumulator = self.create_groups_accumulator()?; + accumulator.set_metrics(metrics); + Ok(accumulator) + } + /// Construct an expression that calculates the aggregate in reverse. /// Typically the "reverse" expression is itself (e.g. SUM, COUNT). /// For aggregates that do not support calculation in reverse, 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 75468198f51d7..d30b63eca9557 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -88,6 +88,9 @@ pub(in crate::aggregates) struct AggregateHashTable { /// Per-aggregate timing metrics for accumulator operations. pub(super) aggregate_accumulator_metrics: Arc, + /// Optional internal metrics owned by each aggregate expression. + pub(super) aggregate_submetrics: Vec>, + /// Raw input schema, used to evaluate expressions and synthesize empty /// grouping-set rows. pub(super) input_schema: SchemaRef, @@ -123,6 +126,7 @@ impl AggregateHashTable { } let input_schema = agg.input().schema(); + let metrics = AggregateTableMetrics::new(agg, partition); let aggregate_arguments = aggregate_expressions( &agg.aggr_expr, &agg.mode, @@ -133,13 +137,16 @@ impl AggregateHashTable { .iter() .zip(aggregate_arguments) .zip(filters) - .map(|((agg_expr, arguments), filter)| { - let accumulator = create_group_accumulator(agg_expr)?; + .zip(metrics.submetrics.iter()) + .map(|(((agg_expr, arguments), filter), submetrics)| { + let accumulator = + create_group_accumulator(agg_expr, Arc::clone(submetrics))?; Ok(HashAggregateAccumulator::new( Arc::clone(agg_expr), arguments, filter, accumulator, + Arc::clone(submetrics), )) }) .collect::>()?; @@ -147,12 +154,11 @@ impl AggregateHashTable { let group_schema = agg.group_by.group_schema(&input_schema)?; let group_values = new_group_values(group_schema, &GroupOrdering::None)?; - let metrics = AggregateTableMetrics::new(agg, partition); - Ok(Self { group_by_metrics: metrics.group_by, aggregate_argument_metrics: metrics.aggregate_arguments, aggregate_accumulator_metrics: metrics.accumulator, + aggregate_submetrics: metrics.submetrics, input_schema, output_schema, state_schema, @@ -502,6 +508,9 @@ pub(super) struct HashAggregateAccumulator { /// Accumulator state for all groups for one aggregate expression. accumulator: Box, + + /// Optional internal metrics owned by this aggregate expression. + submetrics: Arc, } pub(super) type AggregateAccumulator = HashAggregateAccumulator; @@ -639,24 +648,28 @@ impl HashAggregateAccumulator { arguments: Vec>, filter: Option>, accumulator: Box, + submetrics: Arc, ) -> Self { Self { aggregate_expr, arguments, filter, accumulator, + submetrics, } } /// Construct a new accumulator with the same definition, but with empty internal /// state buffers (empty [`GroupsAccumulator`]). pub(super) fn empty_like(&self) -> Result { - let accumulator = create_group_accumulator(&self.aggregate_expr)?; + let accumulator = + create_group_accumulator(&self.aggregate_expr, Arc::clone(&self.submetrics))?; Ok(Self::new( Arc::clone(&self.aggregate_expr), self.arguments.clone(), self.filter.clone(), accumulator, + Arc::clone(&self.submetrics), )) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index 418f3f376b492..ab4f27e0289b4 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -52,6 +52,7 @@ pub(in crate::aggregates) struct OrderedAggregateTableMetrics { pub(super) group_by: GroupByMetrics, pub(super) aggregate_arguments: AggregateArgumentMetrics, pub(super) accumulator: Arc, + pub(super) submetrics: Vec>, } impl OrderedAggregateTableMetrics { @@ -61,6 +62,7 @@ impl OrderedAggregateTableMetrics { group_by: metrics.group_by, aggregate_arguments: metrics.aggregate_arguments, accumulator: metrics.accumulator, + submetrics: metrics.submetrics, } } @@ -71,6 +73,7 @@ impl OrderedAggregateTableMetrics { group_by: table.group_by_metrics.clone(), aggregate_arguments: table.aggregate_argument_metrics.clone(), accumulator: Arc::clone(&table.aggregate_accumulator_metrics), + submetrics: table.aggregate_submetrics.clone(), } } } @@ -138,6 +141,9 @@ pub(in crate::aggregates) struct OrderedAggregateTable { /// Per-aggregate timing metrics for accumulator operations. pub(super) aggregate_accumulator_metrics: Arc, + /// Optional internal metrics owned by each aggregate expression. + pub(super) aggregate_submetrics: Vec>, + /// Group keys, ordering state, and accumulator states. pub(super) buffer: OrderedAggregateTableBuffer, @@ -207,13 +213,16 @@ impl OrderedAggregateTable { .iter() .zip(aggregate_arguments) .zip(filters) - .map(|((agg_expr, arguments), filter)| { - let accumulator = create_group_accumulator(agg_expr)?; + .zip(metrics.submetrics.iter()) + .map(|(((agg_expr, arguments), filter), submetrics)| { + let accumulator = + create_group_accumulator(agg_expr, Arc::clone(submetrics))?; Ok(AggregateAccumulator::new( Arc::clone(agg_expr), arguments, filter, accumulator, + Arc::clone(submetrics), )) }) .collect::>()?; @@ -225,6 +234,7 @@ impl OrderedAggregateTable { group_by_metrics: metrics.group_by, aggregate_argument_metrics: metrics.aggregate_arguments, aggregate_accumulator_metrics: metrics.accumulator, + aggregate_submetrics: metrics.submetrics, buffer: OrderedAggregateTableBuffer { group_by: Arc::clone(&agg.group_by), group_ordering, @@ -308,6 +318,7 @@ impl OrderedAggregateTable { group_by: self.group_by_metrics.clone(), aggregate_arguments: self.aggregate_argument_metrics.clone(), accumulator: Arc::clone(&self.aggregate_accumulator_metrics), + submetrics: self.aggregate_submetrics.clone(), } } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs index 902a859bac96d..29ef4a662f090 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -29,9 +29,10 @@ use std::sync::Arc; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, + GroupByMetrics, aggregate_sub_metrics, }; use crate::aggregates::{AggregateExec, AggregateMode, aggregate_metric_label}; +use datafusion_expr::AggregateMetrics; pub(super) fn accumulator_phases(mode: &AggregateMode) -> &'static [AccumulatorPhase] { match mode { @@ -63,6 +64,7 @@ pub(super) struct AggregateTableMetrics { pub(super) group_by: GroupByMetrics, pub(super) aggregate_arguments: AggregateArgumentMetrics, pub(super) accumulator: Arc, + pub(super) submetrics: Vec>, } impl AggregateTableMetrics { @@ -75,6 +77,11 @@ impl AggregateTableMetrics { Self { group_by: GroupByMetrics::new(&agg.metrics, partition), + submetrics: aggregate_sub_metrics( + &agg.metrics, + partition, + aggregate_labels.iter().cloned(), + ), aggregate_arguments: AggregateArgumentMetrics::new( &agg.metrics, partition, 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..1870051e81ce8 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 @@ -91,6 +91,7 @@ impl AggregateHashTable { aggregate_accumulator_metrics: Arc::clone( &self.aggregate_accumulator_metrics, ), + aggregate_submetrics: self.aggregate_submetrics.clone(), input_schema: Arc::clone(&self.input_schema), output_schema: Arc::clone(&self.output_schema), state_schema: Arc::clone(&self.state_schema), diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index 23f74e6352a15..d689428247963 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -19,11 +19,12 @@ use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, + aggregate_sub_metrics, }; use crate::aggregates::{ AccumulatorItem, AggrDynFilter, AggregateInputMode, AggregateMode, AggregateOutputMode, DynamicFilterAggregateType, aggregate_expressions, - aggregate_metric_label, create_accumulators, + aggregate_metric_label, create_accumulators_with_metrics, }; use crate::metrics::{BaselineMetrics, RecordOutput}; use crate::stream::EmptyRecordBatchStream; @@ -304,12 +305,19 @@ impl AggregateStream { AggregateInputMode::Raw => agg_filter_expr, AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(), }; - let accumulators = create_accumulators(&agg.aggr_expr)?; let aggregate_labels = agg .aggr_expr .iter() .map(|agg_expr| aggregate_metric_label(agg_expr)) .collect::>(); + let accumulators = create_accumulators_with_metrics( + &agg.aggr_expr, + &aggregate_sub_metrics( + &agg.metrics, + partition, + aggregate_labels.iter().cloned(), + ), + )?; let aggregate_argument_metrics = AggregateArgumentMetrics::new( &agg.metrics, partition, @@ -573,9 +581,9 @@ mod tests { use crate::metrics::{MetricValue, MetricsSet}; use crate::test::TestMemoryExec; use crate::{ExecutionPlan, collect}; - use arrow::array::Float64Array; + use arrow::array::{Float64Array, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_functions_aggregate::sum::sum_udaf; + use datafusion_functions_aggregate::{array_agg::array_agg_udaf, sum::sum_udaf}; use datafusion_physical_expr::aggregate::{ AggregateExprBuilder, AggregateFunctionExpr, }; @@ -594,6 +602,20 @@ mod tests { )) } + fn distinct_array_aggregate( + schema: &SchemaRef, + column: &str, + alias: &str, + ) -> Result> { + Ok(Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col(column, schema)?]) + .schema(Arc::clone(schema)) + .distinct() + .alias(alias) + .build()?, + )) + } + fn aggregate_metrics(metrics: &MetricsSet, phase: &str) -> Vec<(String, String)> { let mut result = metrics .iter() @@ -706,6 +728,64 @@ mod tests { Ok(()) } + #[tokio::test] + async fn aggregate_stream_reports_distinct_array_agg_submetrics() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 1, 2])), + Arc::new(UInt32Array::from(vec![3, 3, 4])), + ], + )?; + let input = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::default(), + vec![ + distinct_array_aggregate(&schema, "a", "first")?, + distinct_array_aggregate(&schema, "b", "second")?, + ], + vec![None, None], + input, + schema, + )?); + + let _ = collect( + Arc::clone(&aggregate) as Arc, + Arc::new(TaskContext::default()), + ) + .await?; + + assert_eq!( + aggregate_metrics(&aggregate.metrics().unwrap(), "internal_distinct"), + vec![ + ( + "agg_expr_0_internal_distinct_time".to_string(), + "first".to_string(), + ), + ( + "agg_expr_1_internal_distinct_time".to_string(), + "second".to_string(), + ), + ] + ); + assert_eq!( + aggregate_metrics(&aggregate.metrics().unwrap(), "update").len(), + 2 + ); + assert_eq!( + aggregate_metrics(&aggregate.metrics().unwrap(), "evaluate").len(), + 2 + ); + + Ok(()) + } + #[tokio::test] async fn aggregate_stream_reports_partial_and_final_phases() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index d6405bb403adc..8ed0fac6dd750 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -18,6 +18,89 @@ //! Metrics for the various group-by implementations. use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; +use datafusion_expr::{AggregateMetric, AggregateMetrics}; +use parking_lot::Mutex; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +/// Lazily registers optional internal metrics for one aggregate expression. +/// +/// The physical aggregate operator owns the expression index and label. Aggregate +/// implementations can only supply stable subphase identifiers through +/// [`AggregateMetrics`]. +#[derive(Debug)] +pub(crate) struct AggregateSubMetrics { + metrics: ExecutionPlanMetricsSet, + partition: usize, + index: usize, + aggregate_label: String, + subphase_times: Mutex>, +} + +impl AggregateSubMetrics { + pub(crate) fn new( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + index: usize, + aggregate_label: impl Into, + ) -> Self { + Self { + metrics: metrics.clone(), + partition, + index, + aggregate_label: aggregate_label.into(), + subphase_times: Mutex::new(HashMap::new()), + } + } +} + +#[derive(Debug)] +struct AggregateSubMetric { + time: Time, +} + +impl AggregateMetric for AggregateSubMetric { + fn add_duration(&self, duration: Duration) { + self.time.add_duration(duration); + } +} + +impl AggregateMetrics for AggregateSubMetrics { + fn metric(&self, subphase: &'static str) -> Arc { + let mut subphase_times = self.subphase_times.lock(); + let time = subphase_times + .entry(subphase) + .or_insert_with(|| { + MetricBuilder::new(&self.metrics) + .with_new_label("aggregate", self.aggregate_label.clone()) + .subset_time( + format!("agg_expr_{}_internal_{}_time", self.index, subphase), + self.partition, + ) + }) + .clone(); + Arc::new(AggregateSubMetric { time }) + } +} + +pub(crate) fn aggregate_sub_metrics( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + aggregate_labels: impl IntoIterator, +) -> Vec> +where + T: Into, +{ + aggregate_labels + .into_iter() + .enumerate() + .map(|(index, label)| { + Arc::new(AggregateSubMetrics::new(metrics, partition, index, label)) + as Arc + }) + .collect() +} #[derive(Clone)] pub(crate) struct AggregateArgumentMetrics { @@ -221,7 +304,7 @@ impl GroupByMetrics { #[cfg(test)] mod tests { - use super::GroupByMetrics; + use super::{GroupByMetrics, aggregate_sub_metrics}; use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use crate::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet}; use crate::test::TestMemoryExec; @@ -240,6 +323,32 @@ mod tests { }; use datafusion_physical_expr::expressions::col; use std::sync::Arc; + use std::time::Duration; + + #[test] + fn aggregate_submetrics_merge_across_partitions() { + let metrics = ExecutionPlanMetricsSet::new(); + let partition_0 = aggregate_sub_metrics(&metrics, 0, ["array_agg(DISTINCT a)"]); + let partition_1 = aggregate_sub_metrics(&metrics, 1, ["array_agg(DISTINCT a)"]); + + partition_0[0] + .metric("distinct") + .add_duration(Duration::from_nanos(1)); + partition_1[0] + .metric("distinct") + .add_duration(Duration::from_nanos(2)); + + let metrics = metrics.clone_inner(); + let metric_name = "agg_expr_0_internal_distinct_time"; + assert_eq!(metrics.sum_by_name(metric_name).unwrap().as_usize(), 3); + assert!(metrics.iter().all(|metric| { + metric.value().name() == metric_name + && metric.labels().iter().any(|label| { + label.name() == "aggregate" + && label.value() == "array_agg(DISTINCT a)" + }) + })); + } /// Helper function to verify all three GroupBy metrics exist and have non-zero values fn assert_groupby_metrics(metrics: &MetricsSet) { diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index b764fd7792b39..80e6b335e3188 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -51,7 +51,7 @@ mod null_builder; pub(crate) use metrics::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, + GroupByMetrics, aggregate_sub_metrics, }; /// Stores the group values during hash aggregation. diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 24bb4c16d887c..b8fcd59a4aaf3 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -27,7 +27,7 @@ use super::skip_partial::SkipAggregationProbe; use super::{AggregateExec, format_human_display}; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, - GroupByMetrics, GroupValues, new_group_values, + GroupByMetrics, GroupValues, aggregate_sub_metrics, new_group_values, }; use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ @@ -51,7 +51,7 @@ use datafusion_common::{ use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; -use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr::{AggregateMetrics, EmitTo, GroupsAccumulator}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{GroupsAccumulatorAdapter, PhysicalSortExpr}; @@ -412,6 +412,11 @@ impl GroupedHashAggregateStream { partition, aggregate_labels.iter().cloned(), ); + let aggregate_submetrics = aggregate_sub_metrics( + &agg.metrics, + partition, + aggregate_labels.iter().cloned(), + ); let aggregate_accumulator_metrics = AggregateAccumulatorMetrics::new( &agg.metrics, partition, @@ -445,7 +450,8 @@ impl GroupedHashAggregateStream { // Instantiate the accumulators let accumulators: Vec<_> = aggregate_exprs .iter() - .map(create_group_accumulator) + .zip(aggregate_submetrics) + .map(|(agg_expr, metrics)| create_group_accumulator(agg_expr, metrics)) .collect::>()?; let group_schema = agg_group_by.group_schema(&agg.input().schema())?; @@ -645,9 +651,10 @@ impl GroupedHashAggregateStream { /// [`GroupsAccumulatorAdapter`] if not. pub(crate) fn create_group_accumulator( agg_expr: &Arc, + metrics: Arc, ) -> Result> { if agg_expr.groups_accumulator_supported() { - agg_expr.create_groups_accumulator() + agg_expr.create_groups_accumulator_with_metrics(metrics) } else { // Note in the log when the slow path is used debug!( @@ -655,7 +662,9 @@ pub(crate) fn create_group_accumulator( agg_expr.name() ); let agg_expr_captured = Arc::clone(agg_expr); - let factory = move || agg_expr_captured.create_accumulator(); + let factory = move || { + agg_expr_captured.create_accumulator_with_metrics(Arc::clone(&metrics)) + }; Ok(Box::new(GroupsAccumulatorAdapter::new(factory))) } } diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 774c08535c8e5..3457a0460b0fa 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -187,7 +187,7 @@ use datafusion_common::{ assert_eq_or_internal_err, internal_err, not_impl_err, }; use datafusion_execution::TaskContext; -use datafusion_expr::{Accumulator, Aggregate}; +use datafusion_expr::{Accumulator, Aggregate, AggregateMetrics}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr, lit}; @@ -3024,6 +3024,18 @@ pub fn create_accumulators( .collect() } +pub(crate) fn create_accumulators_with_metrics( + aggr_expr: &[Arc], + aggregate_metrics: &[Arc], +) -> Result> { + debug_assert_eq!(aggr_expr.len(), aggregate_metrics.len()); + aggr_expr + .iter() + .zip(aggregate_metrics) + .map(|(expr, metrics)| expr.create_accumulator_with_metrics(Arc::clone(metrics))) + .collect() +} + /// returns a vector of ArrayRefs, where each entry corresponds to either the /// final value (mode = Final, FinalPartitioned and Single) or states (mode = Partial) pub fn finalize_aggregation( diff --git a/docs/source/user-guide/metrics.md b/docs/source/user-guide/metrics.md index 5bb4895a3da1e..5898b6d8b40a1 100644 --- a/docs/source/user-guide/metrics.md +++ b/docs/source/user-guide/metrics.md @@ -154,6 +154,17 @@ when it is evaluated per aggregate. The legacy grouped hash path evaluates filters collectively, so its per-aggregate `arguments` timers cover argument expressions only; their sum need not equal `aggregate_arguments_time`. +Aggregate implementations can also expose optional internal submetrics. These +use the `agg_expr_{index}_internal_{subphase}_time` naming and `aggregate` +label. The `internal` segment keeps them separate from call-boundary timers; +`subphase` is a stable identifier owned and documented by the aggregate +implementation. They are registered lazily only when an aggregate requests +them, so aggregates without internal submetrics add no metrics. For example, +`array_agg(DISTINCT ...)` records the time spent deduplicating input values as +`agg_expr_{index}_internal_distinct_time`. These submetrics complement the +`update`, `merge`, `state`, and `evaluate` timers rather than subdividing or +replacing them. + Except for the `Summary` metric `reduction_factor`, these operator-level and per-aggregate metrics are `Dev` metrics. They appear in `EXPLAIN ANALYZE` when `datafusion.explain.analyze_level` includes `Dev` (the default), but are omitted From fe701404036be914d8d8510b61f42d9f1d7451b4 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:26:11 +0800 Subject: [PATCH 06/24] feat(metrics): safe refactors to avoid redundant allocations and improve performance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cache one adapter per subphase; no repeated wrapper allocation - Introduce a single reusable adapter instance per subphase, eliminating the need to allocate multiple wrapper objects. This reduces memory churn and improves performance during metric collection. Skip `Arc` clone/clock read when no metric - Detect when there is no active metric to record and skip the unnecessary `Arc` clone and system clock reads. This lowers CPU overhead for subphases that don't emit metrics. Make submetric implementation details private - Move internal helpers and type-specific logic for submetrics behind `pub(super)` or module‑level privacy boundaries. This hides implementation details from external users, enhancing encapsulation and reducing the risk of misuse. --- .../functions-aggregate/src/array_agg.rs | 7 ++-- .../src/aggregates/group_values/metrics.rs | 38 +++++++++---------- 2 files changed, 22 insertions(+), 23 deletions(-) diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index aee21799ea009..01f939fb9aeac 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -955,8 +955,7 @@ impl Accumulator for DistinctArrayAggAccumulator { return Ok(()); } - let distinct_metric = self.distinct_metric.clone(); - let distinct_start = Instant::now(); + let distinct_start = self.distinct_metric.is_some().then(Instant::now); self.ensure_state(col.data_type())?; // Encode the entire incoming batch into rows_buffer in one pass. @@ -1003,8 +1002,8 @@ impl Accumulator for DistinctArrayAggAccumulator { } } } - if let Some(metric) = distinct_metric { - metric.add_duration(distinct_start.elapsed()); + if let (Some(metric), Some(start)) = (&self.distinct_metric, distinct_start) { + metric.add_duration(start.elapsed()); } Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index 8ed0fac6dd750..a0c18f8be4bfc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -30,16 +30,16 @@ use std::time::Duration; /// implementations can only supply stable subphase identifiers through /// [`AggregateMetrics`]. #[derive(Debug)] -pub(crate) struct AggregateSubMetrics { +struct AggregateSubMetrics { metrics: ExecutionPlanMetricsSet, partition: usize, index: usize, aggregate_label: String, - subphase_times: Mutex>, + subphase_metrics: Mutex>>, } impl AggregateSubMetrics { - pub(crate) fn new( + fn new( metrics: &ExecutionPlanMetricsSet, partition: usize, index: usize, @@ -50,7 +50,7 @@ impl AggregateSubMetrics { partition, index, aggregate_label: aggregate_label.into(), - subphase_times: Mutex::new(HashMap::new()), + subphase_metrics: Mutex::new(HashMap::new()), } } } @@ -68,19 +68,16 @@ impl AggregateMetric for AggregateSubMetric { impl AggregateMetrics for AggregateSubMetrics { fn metric(&self, subphase: &'static str) -> Arc { - let mut subphase_times = self.subphase_times.lock(); - let time = subphase_times - .entry(subphase) - .or_insert_with(|| { - MetricBuilder::new(&self.metrics) - .with_new_label("aggregate", self.aggregate_label.clone()) - .subset_time( - format!("agg_expr_{}_internal_{}_time", self.index, subphase), - self.partition, - ) - }) - .clone(); - Arc::new(AggregateSubMetric { time }) + let mut subphase_metrics = self.subphase_metrics.lock(); + Arc::clone(subphase_metrics.entry(subphase).or_insert_with(|| { + let time = MetricBuilder::new(&self.metrics) + .with_new_label("aggregate", self.aggregate_label.clone()) + .subset_time( + format!("agg_expr_{}_internal_{}_time", self.index, subphase), + self.partition, + ); + Arc::new(AggregateSubMetric { time }) + })) } } @@ -334,13 +331,16 @@ mod tests { partition_0[0] .metric("distinct") .add_duration(Duration::from_nanos(1)); - partition_1[0] + partition_0[0] .metric("distinct") .add_duration(Duration::from_nanos(2)); + partition_1[0] + .metric("distinct") + .add_duration(Duration::from_nanos(3)); let metrics = metrics.clone_inner(); let metric_name = "agg_expr_0_internal_distinct_time"; - assert_eq!(metrics.sum_by_name(metric_name).unwrap().as_usize(), 3); + assert_eq!(metrics.sum_by_name(metric_name).unwrap().as_usize(), 6); assert!(metrics.iter().all(|metric| { metric.value().name() == metric_name && metric.labels().iter().any(|label| { From 1e4acdebdabff6dbb89dbe3a43bdabe790faa280 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:38:18 +0800 Subject: [PATCH 07/24] fix: assert both internal_distinct timers are >0 in test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Updated the test to verify that both `internal_distinct` timers are positive (`>0`) - This resolves the blocker where timers could be zero, causing test failures - Ensures correct initialization and behavior of the timer logic - Improves the reliability and confidence of timer‑related functionality --- .../physical-plan/src/aggregates/aggregate_stream.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index d689428247963..3af318030d472 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -761,8 +761,9 @@ mod tests { ) .await?; + let metrics = aggregate.metrics().unwrap(); assert_eq!( - aggregate_metrics(&aggregate.metrics().unwrap(), "internal_distinct"), + aggregate_metrics(&metrics, "internal_distinct"), vec![ ( "agg_expr_0_internal_distinct_time".to_string(), @@ -774,6 +775,15 @@ mod tests { ), ] ); + for index in 0..2 { + assert!( + metrics + .sum_by_name(&format!("agg_expr_{index}_internal_distinct_time")) + .expect("internal distinct time metric") + .as_usize() + > 0 + ); + } assert_eq!( aggregate_metrics(&aggregate.metrics().unwrap(), "update").len(), 2 From 3105148b5344100dde8969184ca8587942690b0f Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:40:07 +0800 Subject: [PATCH 08/24] docs: update metrics.md with identity/cardinality, accumulator timer, partition display, and empty input behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Clarify identity/cardinality format as **(expr index, subphase, partition)** - Note that replacement accumulators share a single timer - Explain how partitions are combined in the normal display - Document that construction‑time requests can cause metrics to appear on empty input --- docs/source/user-guide/metrics.md | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/source/user-guide/metrics.md b/docs/source/user-guide/metrics.md index 5898b6d8b40a1..46993d64ab9eb 100644 --- a/docs/source/user-guide/metrics.md +++ b/docs/source/user-guide/metrics.md @@ -159,11 +159,15 @@ use the `agg_expr_{index}_internal_{subphase}_time` naming and `aggregate` label. The `internal` segment keeps them separate from call-boundary timers; `subphase` is a stable identifier owned and documented by the aggregate implementation. They are registered lazily only when an aggregate requests -them, so aggregates without internal submetrics add no metrics. For example, -`array_agg(DISTINCT ...)` records the time spent deduplicating input values as -`agg_expr_{index}_internal_distinct_time`. These submetrics complement the -`update`, `merge`, `state`, and `evaluate` timers rather than subdividing or -replacing them. +them, so aggregates without internal submetrics add no metrics. Registration +is per `(aggregate expression index, subphase, partition)`: replacement +accumulators in that partition share the same time, and normal metric display +combines that time across partitions. An aggregate may request its submetric +during accumulator construction, so it can appear even when its input is empty. +For example, `array_agg(DISTINCT ...)` records the time spent deduplicating +input values as `agg_expr_{index}_internal_distinct_time`. These submetrics +complement the `update`, `merge`, `state`, and `evaluate` timers rather than +subdividing or replacing them. Except for the `Summary` metric `reduction_factor`, these operator-level and per-aggregate metrics are `Dev` metrics. They appear in `EXPLAIN ANALYZE` when From b26e09d6b92120cdb57286e7c5fd79fba272318b Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 14:02:04 +0800 Subject: [PATCH 09/24] fix(DistinctArrayAggAccumulator): correct `size()` calculation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Counts the retained metric‑handle field in the accumulator’s size. - Adds a regression test to verify the size calculation under various inputs. - Updates the exact distinct‑size expectation to match the corrected behavior. --- datafusion/functions-aggregate/src/array_agg.rs | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 01f939fb9aeac..e43b47bc6a0b7 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -1178,9 +1178,7 @@ impl Accumulator for DistinctArrayAggAccumulator { } fn size(&self) -> usize { - // `distinct_metric` is execution-owned observability state, not - // accumulator state retained by this aggregate. - (size_of_val(self) - size_of_val(&self.distinct_metric)) + size_of_val(self) + self .state .as_ref() @@ -1487,6 +1485,15 @@ mod tests { use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::Column; + #[test] + fn distinct_accumulator_size_includes_metric_handle() -> Result<()> { + let accumulator = + DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; + + assert_eq!(accumulator.size(), size_of_val(&accumulator)); + Ok(()) + } + #[test] fn no_duplicates_no_distinct() -> Result<()> { let (mut acc1, mut acc2) = ArrayAggAccumulatorBuilder::string().build_two()?; @@ -1773,7 +1780,7 @@ mod tests { acc2.update_batch(&[string_list_data([vec!["e", "f", "g"]])])?; acc1 = merge(acc1, acc2)?; - assert_eq!(acc1.size(), 2274); + assert_eq!(acc1.size(), 2290); Ok(()) } From 5cfb71892077923ef5937ad6b12c82f1074bf377 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 14:12:17 +0800 Subject: [PATCH 10/24] =?UTF-8?q?feat(aggregates):=20add=202=E2=80=91parti?= =?UTF-8?q?tion=20execution=20test=20and=20fix=20array=5Fagg=20distinct=20?= =?UTF-8?q?identities=20and=20timer=20handling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add new 2‑partition execution test for `aggregate_stream.rs`. - Fix repeated `array_agg(DISTINCT ...)` identities and labels to be consistent. - Ensure exact partition identity handling: 0/1 per repeated expression. - Merge timer total calculation into a single metric. - Assert that normal `sum` operations do not produce any internal metric. --- .../src/aggregates/aggregate_stream.rs | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs index 3af318030d472..38ec175f25a21 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs @@ -796,6 +796,103 @@ mod tests { Ok(()) } + #[tokio::test] + async fn aggregate_stream_merges_submetrics_across_partitions() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::Float64, false), + ])); + let batches = [ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 1, 2])), + Arc::new(UInt32Array::from(vec![3, 3, 4])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![5, 5, 6])), + Arc::new(UInt32Array::from(vec![7, 7, 8])), + Arc::new(Float64Array::from(vec![4.0, 5.0, 6.0])), + ], + )?, + ]; + let input = TestMemoryExec::try_new_exec( + &[vec![batches[0].clone()], vec![batches[1].clone()]], + Arc::clone(&schema), + None, + )?; + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::default(), + vec![ + distinct_array_aggregate(&schema, "a", "first")?, + distinct_array_aggregate(&schema, "b", "second")?, + sum_aggregate(&schema, "c", "plain")?, + ], + vec![None, None, None], + input, + schema, + )?); + + let context = Arc::new(TaskContext::default()); + for partition in 0..2 { + let _ = crate::common::collect( + aggregate.execute(partition, Arc::clone(&context))?, + ) + .await?; + } + + let metrics = aggregate.metrics().unwrap(); + assert_eq!( + aggregate_metrics(&metrics, "internal_distinct"), + vec![ + ( + "agg_expr_0_internal_distinct_time".to_string(), + "first".to_string(), + ), + ( + "agg_expr_0_internal_distinct_time".to_string(), + "first".to_string(), + ), + ( + "agg_expr_1_internal_distinct_time".to_string(), + "second".to_string(), + ), + ( + "agg_expr_1_internal_distinct_time".to_string(), + "second".to_string(), + ), + ] + ); + let mut internal_partitions = metrics + .iter() + .filter(|metric| { + matches!(metric.value(), MetricValue::Time { name, .. } if name.ends_with("_internal_distinct_time")) + }) + .map(|metric| metric.partition()) + .collect::>(); + internal_partitions.sort_unstable(); + assert_eq!( + internal_partitions, + vec![Some(0), Some(0), Some(1), Some(1)] + ); + assert!( + metrics + .sum_by_name("agg_expr_0_internal_distinct_time") + .is_some_and(|metric| metric.as_usize() > 0) + ); + assert!(!metrics.iter().any(|metric| { + matches!(metric.value(), MetricValue::Time { name, .. } if name.starts_with("agg_expr_2_internal_")) + })); + + Ok(()) + } + #[tokio::test] async fn aggregate_stream_reports_partial_and_final_phases() -> Result<()> { let schema = Arc::new(Schema::new(vec![ From 62c6f88868974a916e67d48257730d0eaa231c19 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 10:59:31 +0800 Subject: [PATCH 11/24] =?UTF-8?q?feat(metrics):=20add=20lock=E2=80=91free?= =?UTF-8?q?=20OnceLock=20fast=20path=20for=20AggregateSubMetrics?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Implement a lock‑free `OnceLock` fast path for the first/common subphase in `AggregateSubMetrics`, eliminating per‑group `Mutex` and `HashMap` lookups for distinct subphases. - Keep the locked map only for handling extra / uncommon subphases, reducing contention and improving performance. - Preserve zero‑duration timers internally, preventing the 1 ns‑per‑call inflation that previously affected fast paths. - Ensure time‑merging logic respects zero‑duration metrics, avoiding spurious non‑zero contributions. - Add comprehensive tests covering caching behavior, zero‑duration scenarios, and multi‑subphase usage. --- .../physical-expr-common/src/metrics/value.rs | 16 ++- .../src/aggregates/group_values/metrics.rs | 121 +++++++++++++++--- 2 files changed, 120 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-expr-common/src/metrics/value.rs b/datafusion/physical-expr-common/src/metrics/value.rs index 37ab5194b2cc4..6a4ec87b0e29e 100644 --- a/datafusion/physical-expr-common/src/metrics/value.rs +++ b/datafusion/physical-expr-common/src/metrics/value.rs @@ -203,9 +203,21 @@ impl Time { self.nanos.fetch_add(more_nanos.max(1), Ordering::Relaxed); } - /// Add the number of nanoseconds of other `Time` to self + /// Add a duration without rounding it up to one nanosecond. + /// + /// Use for metrics that can record many tiny independent operations, where + /// a minimum per recording would materially inflate the total. + pub fn add_duration_exact(&self, duration: Duration) { + self.nanos + .fetch_add(duration.as_nanos() as usize, Ordering::Relaxed); + } + + /// Add the number of nanoseconds of other `Time` to self. + /// + /// Unlike recording a new interval, merging must preserve an existing + /// zero-duration measurement. pub fn add(&self, other: &Time) { - self.add_duration(Duration::from_nanos(other.value() as u64)) + self.add_duration_exact(Duration::from_nanos(other.value() as u64)) } /// return a scoped guard that adds the amount of time elapsed diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index a0c18f8be4bfc..5bba9c19faaeb 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -21,7 +21,7 @@ use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; use datafusion_expr::{AggregateMetric, AggregateMetrics}; use parking_lot::Mutex; use std::collections::HashMap; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::time::Duration; /// Lazily registers optional internal metrics for one aggregate expression. @@ -35,7 +35,10 @@ struct AggregateSubMetrics { partition: usize, index: usize, aggregate_label: String, - subphase_metrics: Mutex>>, + /// The first subphase is the common case. Keep it lock-free because an + /// accumulator adapter creates one accumulator per group. + first_subphase_metric: OnceLock<(&'static str, Arc)>, + additional_subphase_metrics: Mutex>>, } impl AggregateSubMetrics { @@ -50,7 +53,8 @@ impl AggregateSubMetrics { partition, index, aggregate_label: aggregate_label.into(), - subphase_metrics: Mutex::new(HashMap::new()), + first_subphase_metric: OnceLock::new(), + additional_subphase_metrics: Mutex::new(HashMap::new()), } } } @@ -62,22 +66,43 @@ struct AggregateSubMetric { impl AggregateMetric for AggregateSubMetric { fn add_duration(&self, duration: Duration) { - self.time.add_duration(duration); + self.time.add_duration_exact(duration); + } +} + +impl AggregateSubMetrics { + fn new_metric(&self, subphase: &'static str) -> Arc { + let time = MetricBuilder::new(&self.metrics) + .with_new_label("aggregate", self.aggregate_label.clone()) + .subset_time( + format!("agg_expr_{}_internal_{}_time", self.index, subphase), + self.partition, + ); + Arc::new(AggregateSubMetric { time }) } } impl AggregateMetrics for AggregateSubMetrics { fn metric(&self, subphase: &'static str) -> Arc { - let mut subphase_metrics = self.subphase_metrics.lock(); - Arc::clone(subphase_metrics.entry(subphase).or_insert_with(|| { - let time = MetricBuilder::new(&self.metrics) - .with_new_label("aggregate", self.aggregate_label.clone()) - .subset_time( - format!("agg_expr_{}_internal_{}_time", self.index, subphase), - self.partition, - ); - Arc::new(AggregateSubMetric { time }) - })) + if let Some((registered_subphase, metric)) = self.first_subphase_metric.get() + && *registered_subphase == subphase + { + return Arc::clone(metric); + } + + let (registered_subphase, metric) = self + .first_subphase_metric + .get_or_init(|| (subphase, self.new_metric(subphase))); + if *registered_subphase == subphase { + return Arc::clone(metric); + } + + let mut additional_subphase_metrics = self.additional_subphase_metrics.lock(); + Arc::clone( + additional_subphase_metrics + .entry(subphase) + .or_insert_with(|| self.new_metric(subphase)), + ) } } @@ -301,7 +326,7 @@ impl GroupByMetrics { #[cfg(test)] mod tests { - use super::{GroupByMetrics, aggregate_sub_metrics}; + use super::{AggregateSubMetrics, GroupByMetrics, aggregate_sub_metrics}; use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use crate::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet}; use crate::test::TestMemoryExec; @@ -313,6 +338,7 @@ mod tests { use datafusion_execution::TaskContext; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::AggregateMetrics; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::aggregate::{ @@ -322,6 +348,71 @@ mod tests { use std::sync::Arc; use std::time::Duration; + #[test] + fn aggregate_submetrics_cache_first_subphase() { + let metric_set = ExecutionPlanMetricsSet::new(); + let metrics = + AggregateSubMetrics::new(&metric_set, 0, 0, "array_agg(DISTINCT a)"); + + metrics.metric("distinct"); + + assert_eq!( + metrics + .first_subphase_metric + .get() + .map(|(subphase, _)| *subphase), + Some("distinct") + ); + } + + #[test] + fn aggregate_submetrics_preserve_zero_duration() { + let metrics = ExecutionPlanMetricsSet::new(); + let submetrics = aggregate_sub_metrics(&metrics, 0, ["array_agg(DISTINCT a)"]); + + submetrics[0] + .metric("distinct") + .add_duration(Duration::ZERO); + + assert_eq!( + metrics + .clone_inner() + .sum_by_name("agg_expr_0_internal_distinct_time") + .unwrap() + .as_usize(), + 0 + ); + } + + #[test] + fn aggregate_submetrics_support_multiple_subphases() { + let metrics = ExecutionPlanMetricsSet::new(); + let submetrics = aggregate_sub_metrics(&metrics, 0, ["array_agg(DISTINCT a)"]); + + submetrics[0] + .metric("distinct") + .add_duration(Duration::from_nanos(1)); + submetrics[0] + .metric("sort") + .add_duration(Duration::from_nanos(2)); + + let metrics = metrics.clone_inner(); + assert_eq!( + metrics + .sum_by_name("agg_expr_0_internal_distinct_time") + .unwrap() + .as_usize(), + 1 + ); + assert_eq!( + metrics + .sum_by_name("agg_expr_0_internal_sort_time") + .unwrap() + .as_usize(), + 2 + ); + } + #[test] fn aggregate_submetrics_merge_across_partitions() { let metrics = ExecutionPlanMetricsSet::new(); From ba29605f48cc3d8d4d52192efadbee088887eed4 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:11:21 +0800 Subject: [PATCH 12/24] =?UTF-8?q?fix:=20restore=20Time::add=20min=E2=80=91?= =?UTF-8?q?1ns=20behavior=20and=20remove=20exact=E2=80=91duration=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Restored shared Time::add min‑1ns behavior. - Removed exact‑duration API/use. - Added regression test merge‑zero → recorded 1ns. - Updated submetric zero test. --- .../physical-expr-common/src/metrics/value.rs | 24 ++++++++----------- .../src/aggregates/group_values/metrics.rs | 6 ++--- 2 files changed, 13 insertions(+), 17 deletions(-) diff --git a/datafusion/physical-expr-common/src/metrics/value.rs b/datafusion/physical-expr-common/src/metrics/value.rs index 6a4ec87b0e29e..acc79012468b7 100644 --- a/datafusion/physical-expr-common/src/metrics/value.rs +++ b/datafusion/physical-expr-common/src/metrics/value.rs @@ -203,21 +203,9 @@ impl Time { self.nanos.fetch_add(more_nanos.max(1), Ordering::Relaxed); } - /// Add a duration without rounding it up to one nanosecond. - /// - /// Use for metrics that can record many tiny independent operations, where - /// a minimum per recording would materially inflate the total. - pub fn add_duration_exact(&self, duration: Duration) { - self.nanos - .fetch_add(duration.as_nanos() as usize, Ordering::Relaxed); - } - - /// Add the number of nanoseconds of other `Time` to self. - /// - /// Unlike recording a new interval, merging must preserve an existing - /// zero-duration measurement. + /// Add the number of nanoseconds of other `Time` to self pub fn add(&self, other: &Time) { - self.add_duration_exact(Duration::from_nanos(other.value() as u64)) + self.add_duration(Duration::from_nanos(other.value() as u64)) } /// return a scoped guard that adds the amount of time elapsed @@ -1290,6 +1278,14 @@ mod tests { } } + #[test] + fn test_time_merge_marks_a_zero_duration_measurement_as_recorded() { + let merged = Time::new(); + merged.add(&Time::new()); + + assert_eq!(merged.value(), 1); + } + #[test] fn test_display_ratio() { let ratio_metrics = RatioMetrics::new(); diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index 5bba9c19faaeb..fdee9899d3bdf 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -66,7 +66,7 @@ struct AggregateSubMetric { impl AggregateMetric for AggregateSubMetric { fn add_duration(&self, duration: Duration) { - self.time.add_duration_exact(duration); + self.time.add_duration(duration); } } @@ -366,7 +366,7 @@ mod tests { } #[test] - fn aggregate_submetrics_preserve_zero_duration() { + fn aggregate_submetrics_record_zero_duration() { let metrics = ExecutionPlanMetricsSet::new(); let submetrics = aggregate_sub_metrics(&metrics, 0, ["array_agg(DISTINCT a)"]); @@ -380,7 +380,7 @@ mod tests { .sum_by_name("agg_expr_0_internal_distinct_time") .unwrap() .as_usize(), - 0 + 1 ); } From c684057c6b59168fcb0a243c0c797195ef4b70d9 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:16:27 +0800 Subject: [PATCH 13/24] =?UTF-8?q?fix(time):=20corrected=20split=20to=20pre?= =?UTF-8?q?serve=20exact=20duration=20adds=20and=20avoid=20per=E2=80=91bat?= =?UTF-8?q?ch=201=E2=80=AFns=20inflation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Time::add: legacy min‑1 ns merge unchanged. - Time::add_duration_exact: restored, scoped API. - Aggregate submetrics use exact adds → no per‑batch 1 ns inflation. - Tests cover both contracts. --- datafusion/physical-expr-common/src/metrics/value.rs | 9 +++++++++ .../src/aggregates/group_values/metrics.rs | 12 ++++++++---- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/datafusion/physical-expr-common/src/metrics/value.rs b/datafusion/physical-expr-common/src/metrics/value.rs index acc79012468b7..95426abc176c5 100644 --- a/datafusion/physical-expr-common/src/metrics/value.rs +++ b/datafusion/physical-expr-common/src/metrics/value.rs @@ -203,6 +203,15 @@ impl Time { self.nanos.fetch_add(more_nanos.max(1), Ordering::Relaxed); } + /// Adds a duration without rounding it up to one nanosecond. + /// + /// Use only for metrics that record many independent operations, where a + /// minimum per recording would materially inflate the total. + pub fn add_duration_exact(&self, duration: Duration) { + self.nanos + .fetch_add(duration.as_nanos() as usize, Ordering::Relaxed); + } + /// Add the number of nanoseconds of other `Time` to self pub fn add(&self, other: &Time) { self.add_duration(Duration::from_nanos(other.value() as u64)) diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index fdee9899d3bdf..518835c30473d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -66,7 +66,7 @@ struct AggregateSubMetric { impl AggregateMetric for AggregateSubMetric { fn add_duration(&self, duration: Duration) { - self.time.add_duration(duration); + self.time.add_duration_exact(duration); } } @@ -366,7 +366,7 @@ mod tests { } #[test] - fn aggregate_submetrics_record_zero_duration() { + fn aggregate_submetrics_preserve_zero_duration_per_recording() { let metrics = ExecutionPlanMetricsSet::new(); let submetrics = aggregate_sub_metrics(&metrics, 0, ["array_agg(DISTINCT a)"]); @@ -377,10 +377,14 @@ mod tests { assert_eq!( metrics .clone_inner() - .sum_by_name("agg_expr_0_internal_distinct_time") + .iter() + .find(|metric| { + metric.value().name() == "agg_expr_0_internal_distinct_time" + }) .unwrap() + .value() .as_usize(), - 1 + 0 ); } From 132a785fa47ae2bdfe1fa992c7920b60f7e18e58 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:30:37 +0800 Subject: [PATCH 14/24] =?UTF-8?q?feat(metrics):=20require=20RefUnwindSafe?= =?UTF-8?q?=20for=20AggregateMetric=20and=20add=20compile=E2=80=91time=20t?= =?UTF-8?q?est=20for=20DistinctArrayAggAccumulator?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Updated `AggregateMetric` to implement the `RefUnwindSafe` trait, ensuring safe reference semantics during unwind operations. - Added a compile‑time test (`distinct_array_agg_accumulator_unwindsafe`) that verifies `DistinctArrayAggAccumulator` maintains both `UnwindSafe` and `RefUnwindSafe` guarantees, preventing panics in error‑recovery scenarios. --- datafusion/expr-common/src/accumulator.rs | 2 +- datafusion/functions-aggregate/src/array_agg.rs | 8 ++++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index 47678137612fe..597873e303b0f 100644 --- a/datafusion/expr-common/src/accumulator.rs +++ b/datafusion/expr-common/src/accumulator.rs @@ -27,7 +27,7 @@ use std::time::Duration; /// /// Aggregate implementations use this interface for optional internal /// subphases. The execution engine owns metric registration and aggregation. -pub trait AggregateMetric: Debug + Send + Sync { +pub trait AggregateMetric: Debug + Send + Sync + std::panic::RefUnwindSafe { /// Adds elapsed time to this metric. fn add_duration(&self, duration: Duration); } diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index e43b47bc6a0b7..f0d3bddf8cf18 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -1485,6 +1485,14 @@ mod tests { use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::Column; + #[test] + fn distinct_accumulator_preserves_unwind_auto_traits() { + fn assert_unwind_traits() { + } + + assert_unwind_traits::(); + } + #[test] fn distinct_accumulator_size_includes_metric_handle() -> Result<()> { let accumulator = From dfc2f914b35a2b62113f3116dc055118971ef774 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:36:49 +0800 Subject: [PATCH 15/24] test: add legacy grouped `array_agg(DISTINCT)` metric test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Direct `GroupedHashAggregateStream` usage. - Multiple groups → adapter per‑group accumulators. - Asserts positive `agg_expr_0_internal_distinct_time`. --- .../src/aggregates/grouped_hash_stream.rs | 68 ++++++++++++++++++- 1 file changed, 66 insertions(+), 2 deletions(-) diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index b8fcd59a4aaf3..a49f549debc9e 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -1498,13 +1498,77 @@ mod tests { use crate::ExecutionPlan; use crate::InputOrderMode; use crate::test::TestMemoryExec; - use arrow::array::{Int32Array, Int64Array}; + use arrow::array::{Int32Array, Int64Array, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; - use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::{array_agg::array_agg_udaf, count::count_udaf}; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; + #[tokio::test] + async fn grouped_hash_stream_reports_distinct_array_agg_submetric() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("group", DataType::Int32, false), + Field::new("value", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 2])), + Arc::new(UInt32Array::from(vec![3, 3, 4, 4])), + ], + )?; + let input = Arc::new(TestMemoryExec::try_new( + &[vec![batch]], + Arc::clone(&schema), + None, + )?); + let group_by = PhysicalGroupBy::new_single(vec![( + col("group", &schema)?, + "group".to_string(), + )]); + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .distinct() + .alias("distinct_values") + .build()?, + ); + let output_schema = Arc::new(create_schema( + &schema, + &group_by, + std::slice::from_ref(&aggregate_expr), + AggregateMode::Single, + )?); + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![aggregate_expr], + vec![None], + input, + output_schema, + )?; + + let mut stream = GroupedHashAggregateStream::new( + &aggregate_exec, + &Arc::new(TaskContext::default()), + 0, + )?; + while let Some(batch) = stream.next().await { + batch?; + } + + assert!( + aggregate_exec + .metrics() + .unwrap() + .sum_by_name("agg_expr_0_internal_distinct_time") + .is_some_and(|metric| metric.as_usize() > 0) + ); + + Ok(()) + } + // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] From d86b9bf557fbc5d9f27869dedf80f488a7da6260 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 22:05:58 +0800 Subject: [PATCH 16/24] fix(array_agg): skip internal DISTINCT timing for small batches and improve per-group metrics - Skip internal DISTINCT timing for batches <16 to reduce overhead. - Avoid per-group `Instant::now()`, metric `Arc` clone, and atomic updates for small batches. - Metrics are still recorded for batches >=16 to retain visibility where needed. - Added threshold tests to verify the new behavior. --- .../functions-aggregate/src/array_agg.rs | 65 ++++++++++++++++++- 1 file changed, 63 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index f0d3bddf8cf18..9504bd34b8ebe 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -857,6 +857,12 @@ pub struct DistinctArrayAggAccumulator { distinct_metric: Option>, } +/// Avoid per-group timer overhead when the adapter calls us with a tiny batch. +/// +/// Batches below this threshold are not large enough for a useful timing +/// sample, while timing each one adds two clock reads and an atomic update. +const DISTINCT_METRIC_MIN_BATCH_SIZE: usize = 16; + /// Returns `true` if `dt` is, or recursively contains, a `Dictionary` type. /// /// `RowConverter` always decodes to the physical (non-dictionary) type, so a @@ -955,7 +961,10 @@ impl Accumulator for DistinctArrayAggAccumulator { return Ok(()); } - let distinct_start = self.distinct_metric.is_some().then(Instant::now); + let distinct_metric = (col.len() >= DISTINCT_METRIC_MIN_BATCH_SIZE) + .then(|| self.distinct_metric.as_ref().cloned()) + .flatten(); + let distinct_start = distinct_metric.as_ref().map(|_| Instant::now()); self.ensure_state(col.data_type())?; // Encode the entire incoming batch into rows_buffer in one pass. @@ -1002,7 +1011,7 @@ impl Accumulator for DistinctArrayAggAccumulator { } } } - if let (Some(metric), Some(start)) = (&self.distinct_metric, distinct_start) { + if let (Some(metric), Some(start)) = (distinct_metric, distinct_start) { metric.add_duration(start.elapsed()); } Ok(()) @@ -1484,6 +1493,58 @@ mod tests { use datafusion_common::internal_err; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::Column; + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[derive(Debug)] + struct CountingMetric(Arc); + + impl AggregateMetric for CountingMetric { + fn add_duration(&self, _duration: std::time::Duration) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[derive(Debug)] + struct CountingMetrics(Arc); + + impl AggregateMetrics for CountingMetrics { + fn metric(&self, _subphase: &'static str) -> Arc { + Arc::new(CountingMetric(Arc::clone(&self.0))) + } + } + + #[test] + fn distinct_accumulator_skips_metric_for_small_batches() -> Result<()> { + let metric_updates = Arc::new(AtomicUsize::new(0)); + let mut accumulator = + DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; + accumulator.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates)))); + + accumulator.update_batch(&[Arc::new(Int32Array::from(vec![ + 1; + DISTINCT_METRIC_MIN_BATCH_SIZE + - 1 + ]))])?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[test] + fn distinct_accumulator_records_metric_for_large_batches() -> Result<()> { + let metric_updates = Arc::new(AtomicUsize::new(0)); + let mut accumulator = + DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; + accumulator.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates)))); + + accumulator.update_batch(&[Arc::new(Int32Array::from(vec![ + 1; + DISTINCT_METRIC_MIN_BATCH_SIZE + ]))])?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } #[test] fn distinct_accumulator_preserves_unwind_auto_traits() { From 33b0f173915c0024c6c52c914abc1ac27138d2d2 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 22:44:23 +0800 Subject: [PATCH 17/24] feat(datafusion): add grouped update metric and update_batch_grouped for accumulators This change introduces two new methods to the Accumulator trait: 1. `grouped_update_batch_metric`: Returns an optional metric that can be used to time grouped updates once per batch instead of per group. 2. `update_batch_grouped`: Updates state when called by a grouped accumulator adapter, with support for using the grouped update metric. The GroupsAccumulatorAdapter now uses these new methods to avoid timing every per-group call, recording one interval for the full batch instead. This reduces the overhead of metric collection when there are many groups. DistinctArrayAggAccumulator is updated to take advantage of the new grouped update functionality, skipping per-group timing for deduplication operations. This refactor improves performance for grouped aggregations by making metric collection proportional to batch count rather than group cardinality. --- datafusion/expr-common/src/accumulator.rs | 18 ++++++ .../src/aggregate/groups_accumulator.rs | 60 ++++++++++++------ .../functions-aggregate/src/array_agg.rs | 61 ++++++++++--------- .../src/aggregates/grouped_hash_stream.rs | 13 ++-- docs/source/user-guide/metrics.md | 8 ++- 5 files changed, 104 insertions(+), 56 deletions(-) diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index 597873e303b0f..d05ce9ae37465 100644 --- a/datafusion/expr-common/src/accumulator.rs +++ b/datafusion/expr-common/src/accumulator.rs @@ -84,6 +84,24 @@ pub trait Accumulator: Send + Sync + Debug + std::any::Any { /// running sum. fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()>; + /// Returns an optional metric timed once per grouped adapter input batch. + /// + /// A grouped accumulator adapter uses this for aggregate-owned work it + /// dispatches to one accumulator per group. The default preserves the + /// usual per-accumulator update path. + fn grouped_update_batch_metric(&self) -> Option> { + None + } + + /// Updates state when called by a grouped accumulator adapter. + /// + /// The default delegates to [`Self::update_batch`]. Implementations that + /// return a [`Self::grouped_update_batch_metric`] can avoid timing every + /// per-group call; the adapter records one interval for the full batch. + fn update_batch_grouped(&mut self, values: &[ArrayRef]) -> Result<()> { + self.update_batch(values) + } + /// Returns the final aggregate value. /// /// For example, the `SUM` accumulator maintains a running sum, diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 6704b068acf0b..f251ace6ddcae 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -24,6 +24,8 @@ pub mod nulls; pub mod prim_op; use std::mem::{size_of, size_of_val}; +use std::sync::Arc; +use std::time::Instant; use arrow::array::new_empty_array; use arrow::{ @@ -33,7 +35,7 @@ use arrow::{ datatypes::UInt32Type, }; use datafusion_common::{Result, ScalarValue, arrow_datafusion_err}; -use datafusion_expr_common::accumulator::Accumulator; +use datafusion_expr_common::accumulator::{Accumulator, AggregateMetric}; use datafusion_expr_common::groups_accumulator::{ EmitTo, GroupSelection, GroupsAccumulator, }; @@ -102,6 +104,9 @@ pub struct GroupsAccumulatorAdapter { /// bottleneck in earlier implementations when there were many /// distinct groups. allocation_bytes: usize, + + /// Optional aggregate-owned metric timed once for a grouped update batch. + grouped_update_metric: Option>, } struct AccumulatorState { @@ -139,6 +144,7 @@ impl GroupsAccumulatorAdapter { factory: Box::new(factory), states: vec![], allocation_bytes: 0, + grouped_update_metric: None, } } @@ -152,6 +158,9 @@ impl GroupsAccumulatorAdapter { let new_accumulators = total_num_groups - self.states.len(); for _ in 0..new_accumulators { let accumulator = (self.factory)()?; + if self.grouped_update_metric.is_none() { + self.grouped_update_metric = accumulator.grouped_update_batch_metric(); + } let state = AccumulatorState::new(accumulator); self.add_allocation(state.size()); self.states.push(state); @@ -191,6 +200,7 @@ impl GroupsAccumulatorAdapter { group_indices: &[usize], opt_filter: Option<&BooleanArray>, total_num_groups: usize, + time_grouped_update: bool, f: F, ) -> Result<()> where @@ -245,25 +255,37 @@ impl GroupsAccumulatorAdapter { // RecordBatch(es) let iter = groups_with_rows.iter().zip(offsets.windows(2)); + let grouped_update_metric = time_grouped_update + .then(|| self.grouped_update_metric.as_ref().cloned()) + .flatten(); + let start = grouped_update_metric.as_ref().map(|_| Instant::now()); + let mut sizes_pre = 0; let mut sizes_post = 0; - for (&group_idx, offsets) in iter { - let state = &mut self.states[group_idx]; - sizes_pre += state.size(); - - let values_to_accumulate = slice_and_maybe_filter( - &values, - opt_filter.as_ref().map(|f| f.as_boolean()), - offsets, - )?; - f(state.accumulator.as_mut(), &values_to_accumulate)?; - - // clear out the state so they are empty for next - // iteration - state.indices.clear(); - sizes_post += state.size(); - } + let result: Result<()> = (|| { + for (&group_idx, offsets) in iter { + let state = &mut self.states[group_idx]; + sizes_pre += state.size(); + + let values_to_accumulate = slice_and_maybe_filter( + &values, + opt_filter.as_ref().map(|f| f.as_boolean()), + offsets, + )?; + f(state.accumulator.as_mut(), &values_to_accumulate)?; + + // clear out the state so they are empty for next + // iteration + state.indices.clear(); + sizes_post += state.size(); + } + Ok(()) + })(); + if let (Some(metric), Some(start)) = (grouped_update_metric, start) { + metric.add_duration(start.elapsed()); + } + result?; self.adjust_allocation(sizes_pre, sizes_post); Ok(()) } @@ -310,8 +332,9 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { group_indices, opt_filter, total_num_groups, + true, |accumulator, values_to_accumulate| { - accumulator.update_batch(values_to_accumulate) + accumulator.update_batch_grouped(values_to_accumulate) }, )?; Ok(()) @@ -412,6 +435,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { group_indices, None, total_num_groups, + false, |accumulator, values_to_accumulate| { accumulator.merge_batch(values_to_accumulate)?; Ok(()) diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 9504bd34b8ebe..c87faede7acc6 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -857,12 +857,6 @@ pub struct DistinctArrayAggAccumulator { distinct_metric: Option>, } -/// Avoid per-group timer overhead when the adapter calls us with a tiny batch. -/// -/// Batches below this threshold are not large enough for a useful timing -/// sample, while timing each one adds two clock reads and an atomic update. -const DISTINCT_METRIC_MIN_BATCH_SIZE: usize = 16; - /// Returns `true` if `dt` is, or recursively contains, a `Dictionary` type. /// /// `RowConverter` always decodes to the physical (non-dictionary) type, so a @@ -922,16 +916,12 @@ impl DistinctArrayAggAccumulator { } } -impl Accumulator for DistinctArrayAggAccumulator { - fn set_metrics(&mut self, metrics: Arc) { - self.distinct_metric = Some(metrics.metric("distinct")); - } - - fn state(&mut self) -> Result> { - Ok(vec![self.evaluate()?]) - } - - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { +impl DistinctArrayAggAccumulator { + fn update_batch_impl( + &mut self, + values: &[ArrayRef], + record_metric: bool, + ) -> Result<()> { if values.is_empty() { return Ok(()); } @@ -961,7 +951,7 @@ impl Accumulator for DistinctArrayAggAccumulator { return Ok(()); } - let distinct_metric = (col.len() >= DISTINCT_METRIC_MIN_BATCH_SIZE) + let distinct_metric = record_metric .then(|| self.distinct_metric.as_ref().cloned()) .flatten(); let distinct_start = distinct_metric.as_ref().map(|_| Instant::now()); @@ -1016,6 +1006,28 @@ impl Accumulator for DistinctArrayAggAccumulator { } Ok(()) } +} + +impl Accumulator for DistinctArrayAggAccumulator { + fn set_metrics(&mut self, metrics: Arc) { + self.distinct_metric = Some(metrics.metric("distinct")); + } + + fn grouped_update_batch_metric(&self) -> Option> { + self.distinct_metric.as_ref().cloned() + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.update_batch_impl(values, true) + } + + fn update_batch_grouped(&mut self, values: &[ArrayRef]) -> Result<()> { + self.update_batch_impl(values, false) + } + + fn state(&mut self) -> Result> { + Ok(vec![self.evaluate()?]) + } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { if states.is_empty() { @@ -1514,19 +1526,15 @@ mod tests { } #[test] - fn distinct_accumulator_skips_metric_for_small_batches() -> Result<()> { + fn distinct_accumulator_records_metric_for_small_batches() -> Result<()> { let metric_updates = Arc::new(AtomicUsize::new(0)); let mut accumulator = DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; accumulator.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates)))); - accumulator.update_batch(&[Arc::new(Int32Array::from(vec![ - 1; - DISTINCT_METRIC_MIN_BATCH_SIZE - - 1 - ]))])?; + accumulator.update_batch(&[Arc::new(Int32Array::from(vec![1]))])?; - assert_eq!(metric_updates.load(Ordering::Relaxed), 0); + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); Ok(()) } @@ -1537,10 +1545,7 @@ mod tests { DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; accumulator.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates)))); - accumulator.update_batch(&[Arc::new(Int32Array::from(vec![ - 1; - DISTINCT_METRIC_MIN_BATCH_SIZE - ]))])?; + accumulator.update_batch(&[Arc::new(Int32Array::from(vec![1; 16]))])?; assert_eq!(metric_updates.load(Ordering::Relaxed), 1); Ok(()) diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index a49f549debc9e..d3f1e7b090744 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -1558,13 +1558,12 @@ mod tests { batch?; } - assert!( - aggregate_exec - .metrics() - .unwrap() - .sum_by_name("agg_expr_0_internal_distinct_time") - .is_some_and(|metric| metric.as_usize() > 0) - ); + let metrics = aggregate_exec.metrics().unwrap(); + let distinct_time = metrics + .iter() + .find(|metric| metric.value().name() == "agg_expr_0_internal_distinct_time") + .expect("internal distinct time metric"); + assert!(distinct_time.value().as_usize() > 0); Ok(()) } diff --git a/docs/source/user-guide/metrics.md b/docs/source/user-guide/metrics.md index 46993d64ab9eb..7af82c45108ee 100644 --- a/docs/source/user-guide/metrics.md +++ b/docs/source/user-guide/metrics.md @@ -165,9 +165,11 @@ accumulators in that partition share the same time, and normal metric display combines that time across partitions. An aggregate may request its submetric during accumulator construction, so it can appear even when its input is empty. For example, `array_agg(DISTINCT ...)` records the time spent deduplicating -input values as `agg_expr_{index}_internal_distinct_time`. These submetrics -complement the `update`, `merge`, `state`, and `evaluate` timers rather than -subdividing or replacing them. +input values as `agg_expr_{index}_internal_distinct_time`. Grouped accumulation +records this once per input batch, rather than once per group, to avoid making +metric collection proportional to group cardinality. These submetrics complement +the `update`, `merge`, `state`, and `evaluate` timers rather than subdividing or +replacing them. Except for the `Summary` metric `reduction_factor`, these operator-level and per-aggregate metrics are `Dev` metrics. They appear in `EXPLAIN ANALYZE` when From 68f98fca736389f8b032c2d594c849e000a6250d Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Thu, 10 Sep 2026 12:30:47 +0800 Subject: [PATCH 18/24] fix(groups_accumulator): improve Instant usage and fix unnecessary Option clone - Replace direct `Option` clone with proper handling to avoid unnecessary allocations. - Switch to `datafusion_common::instant::Instant` for more accurate timing measurements. - Update related logic to ensure correctness and performance improvements. --- .../src/aggregate/groups_accumulator.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index f251ace6ddcae..107534908ffb3 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -25,7 +25,6 @@ pub mod prim_op; use std::mem::{size_of, size_of_val}; use std::sync::Arc; -use std::time::Instant; use arrow::array::new_empty_array; use arrow::{ @@ -34,7 +33,7 @@ use arrow::{ compute::take_arrays, datatypes::UInt32Type, }; -use datafusion_common::{Result, ScalarValue, arrow_datafusion_err}; +use datafusion_common::{Result, ScalarValue, arrow_datafusion_err, instant::Instant}; use datafusion_expr_common::accumulator::{Accumulator, AggregateMetric}; use datafusion_expr_common::groups_accumulator::{ EmitTo, GroupSelection, GroupsAccumulator, @@ -256,7 +255,7 @@ impl GroupsAccumulatorAdapter { let iter = groups_with_rows.iter().zip(offsets.windows(2)); let grouped_update_metric = time_grouped_update - .then(|| self.grouped_update_metric.as_ref().cloned()) + .then(|| self.grouped_update_metric.clone()) .flatten(); let start = grouped_update_metric.as_ref().map(|_| Instant::now()); From d087f87ab0d0764343426e38faa7e668c95781ce Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Thu, 10 Sep 2026 13:49:09 +0800 Subject: [PATCH 19/24] refactor: simplify distinct_metric cloning in DistinctArrayAggAccumulator --- datafusion/functions-aggregate/src/array_agg.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index c87faede7acc6..88e5b6a9daea2 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -952,7 +952,7 @@ impl DistinctArrayAggAccumulator { } let distinct_metric = record_metric - .then(|| self.distinct_metric.as_ref().cloned()) + .then(|| self.distinct_metric.clone()) .flatten(); let distinct_start = distinct_metric.as_ref().map(|_| Instant::now()); self.ensure_state(col.data_type())?; @@ -1014,7 +1014,7 @@ impl Accumulator for DistinctArrayAggAccumulator { } fn grouped_update_batch_metric(&self) -> Option> { - self.distinct_metric.as_ref().cloned() + self.distinct_metric.clone() } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { From 8a3f79c06e4e5ac25ae9d57ac5d9d336d4488554 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 11:07:38 +0800 Subject: [PATCH 20/24] feat(metrics): add array_agg(DISTINCT) merge, convert_to_state, and metric-count regressions - array_agg(DISTINCT) merge: one metric timer/state batch. - convert_to_state: untimed row updates + one conversion timer. - Added metric-count regressions. --- .../src/aggregate/groups_accumulator.rs | 82 ++++++++++++++++++- .../functions-aggregate/src/array_agg.rs | 36 +++++++- 2 files changed, 112 insertions(+), 6 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 107534908ffb3..ba70e659e5e51 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -468,14 +468,22 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { // Each row has its respective group let mut results = vec![]; + let mut grouped_update_metric = None; + let mut start = None; for row_idx in 0..num_rows { // Create the empty accumulator for converting let mut converted_accumulator = (self.factory)()?; + if row_idx == 0 { + grouped_update_metric = + converted_accumulator.grouped_update_batch_metric(); + start = grouped_update_metric.as_ref().map(|_| Instant::now()); + } - // Convert row to states + // Convert row to states. Use the grouped variant because the + // conversion-wide timer above accounts for the complete loop. let values_to_accumulate = slice_and_maybe_filter(values, opt_filter, &[row_idx, row_idx + 1])?; - converted_accumulator.update_batch(&values_to_accumulate)?; + converted_accumulator.update_batch_grouped(&values_to_accumulate)?; let states = converted_accumulator.state()?; // Resize results to have enough columns according to the converted states @@ -487,6 +495,10 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { } } + if let (Some(metric), Some(start)) = (grouped_update_metric, start) { + metric.add_duration(start.elapsed()); + } + let arrays = results .into_iter() .map(ScalarValue::iter_to_array) @@ -556,12 +568,78 @@ pub(crate) fn slice_and_maybe_filter( #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; use super::*; use crate::min_max::MaxAccumulator; use arrow::array::{AsArray, Int64Array}; use arrow::datatypes::{DataType, Int64Type}; + #[derive(Debug)] + struct CountingMetric(Arc); + + impl AggregateMetric for CountingMetric { + fn add_duration(&self, _duration: Duration) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[derive(Debug)] + struct TimedAccumulator { + metric: Arc, + } + + impl Accumulator for TimedAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + self.metric.add_duration(Duration::ZERO); + Ok(()) + } + + fn grouped_update_batch_metric(&self) -> Option> { + Some(Arc::clone(&self.metric)) + } + + fn update_batch_grouped(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(None)) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Int64(None)]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + } + + #[test] + fn adapter_convert_to_state_records_metric_once() -> Result<()> { + let metric_updates = Arc::new(AtomicUsize::new(0)); + let accumulator = GroupsAccumulatorAdapter::new({ + let metric_updates = Arc::clone(&metric_updates); + move || { + Ok(Box::new(TimedAccumulator { + metric: Arc::new(CountingMetric(Arc::clone(&metric_updates))), + }) as Box) + } + }); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + accumulator.convert_to_state(&[values], None)?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } + #[test] fn adapter_preserving_evaluation_uses_accumulator_contract() -> Result<()> { let mut accumulator = GroupsAccumulatorAdapter::new(|| { diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 88e5b6a9daea2..0b0a646183cba 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -1036,12 +1036,22 @@ impl Accumulator for DistinctArrayAggAccumulator { assert_eq_or_internal_err!(states.len(), 1, "expects single state"); - // The DISTINCT state is `List`. - states[0] + let distinct_metric = self.distinct_metric.clone(); + let distinct_start = distinct_metric.as_ref().map(|_| Instant::now()); + + // The DISTINCT state is `List`. This calls the update + // implementation once per state row, so record the submetric once for + // the entire merge rather than once per row. + let result = states[0] .as_list::() .iter() .flatten() - .try_for_each(|val| self.update_batch(&[val])) + .try_for_each(|val| self.update_batch_impl(&[val], false)); + + if let (Some(metric), Some(start)) = (distinct_metric, distinct_start) { + metric.add_duration(start.elapsed()); + } + result } fn evaluate(&mut self) -> Result { @@ -1499,7 +1509,7 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ListBuilder, StringBuilder}; + use arrow::array::{Int32Builder, ListBuilder, StringBuilder}; use arrow::datatypes::Schema; use datafusion_common::cast::as_generic_string_array; use datafusion_common::internal_err; @@ -1551,6 +1561,24 @@ mod tests { Ok(()) } + #[test] + fn distinct_accumulator_records_merge_metric_once() -> Result<()> { + let mut builder = ListBuilder::new(Int32Builder::new()); + for value in [1, 2, 3] { + builder.append_value([Some(value)]); + } + let state: ArrayRef = Arc::new(builder.finish()); + + let metric_updates = Arc::new(AtomicUsize::new(0)); + let mut target = + DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; + target.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates)))); + target.merge_batch(&[state])?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } + #[test] fn distinct_accumulator_preserves_unwind_auto_traits() { fn assert_unwind_traits() { From 98533fcc79d4445b53089e536ffb8aca03ecc69e Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 11:46:47 +0800 Subject: [PATCH 21/24] feat: cache aggregate metric handles for grouped accumulators --- datafusion/expr-common/src/accumulator.rs | 9 ++ .../src/aggregate/groups_accumulator.rs | 43 +++++-- .../functions-aggregate/src/array_agg.rs | 58 +++++---- .../src/aggregates/grouped_hash_stream.rs | 114 +++++++++++++++++- 4 files changed, 188 insertions(+), 36 deletions(-) diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index d05ce9ae37465..aa8be52a18fc3 100644 --- a/datafusion/expr-common/src/accumulator.rs +++ b/datafusion/expr-common/src/accumulator.rs @@ -102,6 +102,15 @@ pub trait Accumulator: Send + Sync + Debug + std::any::Any { self.update_batch(values) } + /// Merges state when called by a grouped accumulator adapter. + /// + /// The default delegates to [`Self::merge_batch`]. Implementations that + /// return a [`Self::grouped_update_batch_metric`] can avoid timing every + /// per-group merge; the adapter records one interval for the full batch. + fn merge_batch_grouped(&mut self, states: &[ArrayRef]) -> Result<()> { + self.merge_batch(states) + } + /// Returns the final aggregate value. /// /// For example, the `SUM` accumulator maintains a running sum, diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index ba70e659e5e51..0a3c95008d6d0 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -24,7 +24,7 @@ pub mod nulls; pub mod prim_op; use std::mem::{size_of, size_of_val}; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use arrow::array::new_empty_array; use arrow::{ @@ -105,7 +105,11 @@ pub struct GroupsAccumulatorAdapter { allocation_bytes: usize, /// Optional aggregate-owned metric timed once for a grouped update batch. - grouped_update_metric: Option>, + /// + /// The cache is shared with factories that install metrics on only the + /// first accumulator. This avoids resolving the same expression metric for + /// every group while allowing `convert_to_state(&self)` to use it. + grouped_update_metric: Arc>>>, } struct AccumulatorState { @@ -136,6 +140,21 @@ impl GroupsAccumulatorAdapter { /// Create a new adapter that will create a new [`Accumulator`] /// for each group, using the specified factory function pub fn new(factory: F) -> Self + where + F: Fn() -> Result> + Send + 'static, + { + Self::new_with_grouped_update_metric_cache(factory, Arc::new(OnceLock::new())) + } + + /// Creates an adapter with a metric cache shared by its accumulator factory. + /// + /// The factory must initialize the cache before returning its first + /// metric-enabled accumulator. Later accumulators may be created without + /// metrics because this adapter records the grouped batch timing. + pub fn new_with_grouped_update_metric_cache( + factory: F, + grouped_update_metric: Arc>>>, + ) -> Self where F: Fn() -> Result> + Send + 'static, { @@ -143,7 +162,7 @@ impl GroupsAccumulatorAdapter { factory: Box::new(factory), states: vec![], allocation_bytes: 0, - grouped_update_metric: None, + grouped_update_metric, } } @@ -157,9 +176,8 @@ impl GroupsAccumulatorAdapter { let new_accumulators = total_num_groups - self.states.len(); for _ in 0..new_accumulators { let accumulator = (self.factory)()?; - if self.grouped_update_metric.is_none() { - self.grouped_update_metric = accumulator.grouped_update_batch_metric(); - } + self.grouped_update_metric + .get_or_init(|| accumulator.grouped_update_batch_metric()); let state = AccumulatorState::new(accumulator); self.add_allocation(state.size()); self.states.push(state); @@ -255,7 +273,7 @@ impl GroupsAccumulatorAdapter { let iter = groups_with_rows.iter().zip(offsets.windows(2)); let grouped_update_metric = time_grouped_update - .then(|| self.grouped_update_metric.clone()) + .then(|| self.grouped_update_metric.get().and_then(Clone::clone)) .flatten(); let start = grouped_update_metric.as_ref().map(|_| Instant::now()); @@ -434,10 +452,9 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { group_indices, None, total_num_groups, - false, + true, |accumulator, values_to_accumulate| { - accumulator.merge_batch(values_to_accumulate)?; - Ok(()) + accumulator.merge_batch_grouped(values_to_accumulate) }, )?; Ok(()) @@ -474,8 +491,10 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { // Create the empty accumulator for converting let mut converted_accumulator = (self.factory)()?; if row_idx == 0 { - grouped_update_metric = - converted_accumulator.grouped_update_batch_metric(); + let metric = self + .grouped_update_metric + .get_or_init(|| converted_accumulator.grouped_update_batch_metric()); + grouped_update_metric.clone_from(metric); start = grouped_update_metric.as_ref().map(|_| Instant::now()); } diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 0b0a646183cba..95a8d6c5274f3 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -1008,6 +1008,39 @@ impl DistinctArrayAggAccumulator { } } +impl DistinctArrayAggAccumulator { + fn merge_batch_impl( + &mut self, + states: &[ArrayRef], + record_metric: bool, + ) -> Result<()> { + if states.is_empty() { + return Ok(()); + } + + assert_eq_or_internal_err!(states.len(), 1, "expects single state"); + + let distinct_metric = record_metric + .then(|| self.distinct_metric.clone()) + .flatten(); + let distinct_start = distinct_metric.as_ref().map(|_| Instant::now()); + + // The DISTINCT state is `List`. This calls the update + // implementation once per state row, so record the submetric once for + // the entire merge rather than once per row. + let result = states[0] + .as_list::() + .iter() + .flatten() + .try_for_each(|val| self.update_batch_impl(&[val], false)); + + if let (Some(metric), Some(start)) = (distinct_metric, distinct_start) { + metric.add_duration(start.elapsed()); + } + result + } +} + impl Accumulator for DistinctArrayAggAccumulator { fn set_metrics(&mut self, metrics: Arc) { self.distinct_metric = Some(metrics.metric("distinct")); @@ -1030,28 +1063,11 @@ impl Accumulator for DistinctArrayAggAccumulator { } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - if states.is_empty() { - return Ok(()); - } - - assert_eq_or_internal_err!(states.len(), 1, "expects single state"); - - let distinct_metric = self.distinct_metric.clone(); - let distinct_start = distinct_metric.as_ref().map(|_| Instant::now()); - - // The DISTINCT state is `List`. This calls the update - // implementation once per state row, so record the submetric once for - // the entire merge rather than once per row. - let result = states[0] - .as_list::() - .iter() - .flatten() - .try_for_each(|val| self.update_batch_impl(&[val], false)); + self.merge_batch_impl(states, true) + } - if let (Some(metric), Some(start)) = (distinct_metric, distinct_start) { - metric.add_duration(start.elapsed()); - } - result + fn merge_batch_grouped(&mut self, states: &[ArrayRef]) -> Result<()> { + self.merge_batch_impl(states, false) } fn evaluate(&mut self) -> Result { diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index d3f1e7b090744..20769559a2006 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -17,7 +17,7 @@ //! Hash aggregation -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::task::{Context, Poll}; use std::vec; @@ -661,11 +661,25 @@ pub(crate) fn create_group_accumulator( "Creating GroupsAccumulatorAdapter for {}: {agg_expr:?}", agg_expr.name() ); + let grouped_update_metric = Arc::new(OnceLock::new()); + let factory_metric = Arc::clone(&grouped_update_metric); let agg_expr_captured = Arc::clone(agg_expr); let factory = move || { - agg_expr_captured.create_accumulator_with_metrics(Arc::clone(&metrics)) + let mut accumulator = agg_expr_captured.create_accumulator()?; + if factory_metric.get().is_none() { + accumulator.set_metrics(Arc::clone(&metrics)); + // This factory is called serially by one adapter, so the first + // accumulator is the only one that resolves the metric handle. + let _ = factory_metric.set(accumulator.grouped_update_batch_metric()); + } + Ok(accumulator) }; - Ok(Box::new(GroupsAccumulatorAdapter::new(factory))) + Ok(Box::new( + GroupsAccumulatorAdapter::new_with_grouped_update_metric_cache( + factory, + grouped_update_metric, + ), + )) } } @@ -1501,9 +1515,103 @@ mod tests { use arrow::array::{Int32Array, Int64Array, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::AggregateMetric; use datafusion_functions_aggregate::{array_agg::array_agg_udaf, count::count_udaf}; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + #[derive(Debug)] + struct CountingMetric(Arc); + + impl AggregateMetric for CountingMetric { + fn add_duration(&self, _duration: Duration) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[derive(Debug)] + struct CountingMetrics { + resolutions: Arc, + durations: Arc, + } + + impl AggregateMetrics for CountingMetrics { + fn metric(&self, _subphase: &'static str) -> Arc { + self.resolutions.fetch_add(1, Ordering::Relaxed); + Arc::new(CountingMetric(Arc::clone(&self.durations))) + } + } + + #[test] + fn legacy_grouped_distinct_resolves_metric_once_for_conversion_and_update() + -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::UInt32, + false, + )])); + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .distinct() + .alias("distinct_values") + .build()?, + ); + let metric_resolutions = Arc::new(AtomicUsize::new(0)); + let metric_durations = Arc::new(AtomicUsize::new(0)); + let mut accumulator = create_group_accumulator( + &aggregate_expr, + Arc::new(CountingMetrics { + resolutions: Arc::clone(&metric_resolutions), + durations: Arc::clone(&metric_durations), + }), + )?; + let values: ArrayRef = Arc::new(UInt32Array::from(vec![1, 2, 3])); + + accumulator.convert_to_state(&[Arc::clone(&values)], None)?; + accumulator.update_batch(&[values], &[0, 1, 2], None, 3)?; + + assert_eq!(metric_resolutions.load(Ordering::Relaxed), 1); + assert_eq!(metric_durations.load(Ordering::Relaxed), 2); + Ok(()) + } + + #[test] + fn legacy_grouped_distinct_merge_records_metric_once() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::UInt32, + false, + )])); + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .distinct() + .alias("distinct_values") + .build()?, + ); + let metric_resolutions = Arc::new(AtomicUsize::new(0)); + let metric_durations = Arc::new(AtomicUsize::new(0)); + let mut accumulator = create_group_accumulator( + &aggregate_expr, + Arc::new(CountingMetrics { + resolutions: Arc::clone(&metric_resolutions), + durations: Arc::clone(&metric_durations), + }), + )?; + + let mut state = ListBuilder::new(UInt32Builder::new()); + for value in [1, 2, 3] { + state.append_value([Some(value)]); + } + accumulator.merge_batch(&[Arc::new(state.finish())], &[1, 2, 3], 4)?; + + assert_eq!(metric_resolutions.load(Ordering::Relaxed), 1); + assert_eq!(metric_durations.load(Ordering::Relaxed), 1); + Ok(()) + } #[tokio::test] async fn grouped_hash_stream_reports_distinct_array_agg_submetric() -> Result<()> { From 1b9cde43cb8cb015a4b95a94010eaa0046ecf8e2 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 12:16:47 +0800 Subject: [PATCH 22/24] feat(groups-accumulator-adapter): move slice/filter prep outside timing, bound prep chunks, time only accumulator calls, and emit one metric update per logical batch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Preps slice/filter outside submetric timing. - Times only accumulator calls. - Bounded 64‑group prep chunks; avoids unbounded retained arrays. - One metric update per logical batch. --- .../src/aggregate/groups_accumulator.rs | 92 ++++++++++++++----- 1 file changed, 68 insertions(+), 24 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 0a3c95008d6d0..fcfe4788c5439 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -25,6 +25,7 @@ pub mod prim_op; use std::mem::{size_of, size_of_val}; use std::sync::{Arc, OnceLock}; +use std::time::Duration; use arrow::array::new_empty_array; use arrow::{ @@ -112,6 +113,10 @@ pub struct GroupsAccumulatorAdapter { grouped_update_metric: Arc>>>, } +/// Maximum number of prepared group inputs retained while timing an +/// aggregate-owned grouped subphase. +const GROUPED_METRIC_PREPARATION_CHUNK_SIZE: usize = 64; + struct AccumulatorState { /// [`Accumulator`] that stores the per-group state accumulator: Box, @@ -267,40 +272,59 @@ impl GroupsAccumulatorAdapter { let values = take_arrays(values, &batch_indices, None)?; let opt_filter = get_filter_at_indices(opt_filter, &batch_indices)?; - // invoke each accumulator with the appropriate rows, first - // pulling the input arguments for this group into their own - // RecordBatch(es) - let iter = groups_with_rows.iter().zip(offsets.windows(2)); - let grouped_update_metric = time_grouped_update .then(|| self.grouped_update_metric.get().and_then(Clone::clone)) .flatten(); - let start = grouped_update_metric.as_ref().map(|_| Instant::now()); let mut sizes_pre = 0; let mut sizes_post = 0; + let mut aggregate_duration = Duration::ZERO; let result: Result<()> = (|| { - for (&group_idx, offsets) in iter { - let state = &mut self.states[group_idx]; - sizes_pre += state.size(); - - let values_to_accumulate = slice_and_maybe_filter( - &values, - opt_filter.as_ref().map(|f| f.as_boolean()), - offsets, - )?; - f(state.accumulator.as_mut(), &values_to_accumulate)?; - - // clear out the state so they are empty for next - // iteration - state.indices.clear(); - sizes_post += state.size(); + // Keep preparation bounded to avoid retaining one filtered array per + // group. Time only accumulator invocation: slicing and filtering + // are adapter work, not aggregate-owned subphase work. + for (chunk_index, groups) in groups_with_rows + .chunks(GROUPED_METRIC_PREPARATION_CHUNK_SIZE) + .enumerate() + { + let first_offset = chunk_index * GROUPED_METRIC_PREPARATION_CHUNK_SIZE; + let values_to_accumulate = groups + .iter() + .enumerate() + .map(|(index, &group_idx)| { + let values = slice_and_maybe_filter( + &values, + opt_filter.as_ref().map(|f| f.as_boolean()), + &offsets[first_offset + index..first_offset + index + 2], + )?; + Ok((group_idx, values)) + }) + .collect::>>()?; + + let start = grouped_update_metric.as_ref().map(|_| Instant::now()); + let chunk_result: Result<()> = (|| { + for (group_idx, values) in values_to_accumulate { + let state = &mut self.states[group_idx]; + sizes_pre += state.size(); + f(state.accumulator.as_mut(), &values)?; + + // clear out the state so they are empty for next + // iteration + state.indices.clear(); + sizes_post += state.size(); + } + Ok(()) + })(); + if let Some(start) = start { + aggregate_duration += start.elapsed(); + } + chunk_result?; } Ok(()) })(); - if let (Some(metric), Some(start)) = (grouped_update_metric, start) { - metric.add_duration(start.elapsed()); + if let Some(metric) = grouped_update_metric { + metric.add_duration(aggregate_duration); } result?; self.adjust_allocation(sizes_pre, sizes_post); @@ -592,7 +616,7 @@ mod tests { use super::*; use crate::min_max::MaxAccumulator; - use arrow::array::{AsArray, Int64Array}; + use arrow::array::{AsArray, BooleanArray, Int64Array}; use arrow::datatypes::{DataType, Int64Type}; #[derive(Debug)] @@ -640,6 +664,26 @@ mod tests { } } + #[test] + fn adapter_grouped_update_records_one_metric_after_filtering() -> Result<()> { + let metric_updates = Arc::new(AtomicUsize::new(0)); + let mut accumulator = GroupsAccumulatorAdapter::new({ + let metric_updates = Arc::clone(&metric_updates); + move || { + Ok(Box::new(TimedAccumulator { + metric: Arc::new(CountingMetric(Arc::clone(&metric_updates))), + }) as Box) + } + }); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 4])); + let filter = BooleanArray::from(vec![true, false, true, false]); + accumulator.update_batch(&[values], &[0, 0, 1, 1], Some(&filter), 2)?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } + #[test] fn adapter_convert_to_state_records_metric_once() -> Result<()> { let metric_updates = Arc::new(AtomicUsize::new(0)); From a00e324e0d4ae8c8aea0ace868c78df5d95c9f23 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 12:46:21 +0800 Subject: [PATCH 23/24] =?UTF-8?q?feat:=20convert=5Fto=5Fstate=20uses=20bou?= =?UTF-8?q?nded=2064=E2=80=91row=20prep=20chunks=20with=20timer=E2=80=91wr?= =?UTF-8?q?apped=20batch=20updates=20and=20adds=20regression=20test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Details:** - Bounded 64‑row prep chunks for `convert_to_state`. - Timer only wraps `update_batch_grouped`. - Excludes factory, slice/filter, `state/evaluate`, and result materialization steps. - Single metric update / logical conversion batch per operation. - Added RED→GREEN filtered multi‑row state‑materialization regression test. --- .../src/aggregate/groups_accumulator.rs | 133 ++++++++++++++---- 1 file changed, 109 insertions(+), 24 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index fcfe4788c5439..0803ec0b9e167 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -510,36 +510,54 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { // Each row has its respective group let mut results = vec![]; let mut grouped_update_metric = None; - let mut start = None; - for row_idx in 0..num_rows { - // Create the empty accumulator for converting - let mut converted_accumulator = (self.factory)()?; - if row_idx == 0 { - let metric = self - .grouped_update_metric - .get_or_init(|| converted_accumulator.grouped_update_batch_metric()); - grouped_update_metric.clone_from(metric); - start = grouped_update_metric.as_ref().map(|_| Instant::now()); + let mut aggregate_duration = Duration::ZERO; + for chunk_start in (0..num_rows).step_by(GROUPED_METRIC_PREPARATION_CHUNK_SIZE) { + let chunk_end = + (chunk_start + GROUPED_METRIC_PREPARATION_CHUNK_SIZE).min(num_rows); + let mut prepared = Vec::with_capacity(chunk_end - chunk_start); + + for row_idx in chunk_start..chunk_end { + // Create the empty accumulator and prepare adapter-owned input + // outside the aggregate submetric. + let accumulator = (self.factory)()?; + if row_idx == 0 { + let metric = self + .grouped_update_metric + .get_or_init(|| accumulator.grouped_update_batch_metric()); + grouped_update_metric.clone_from(metric); + } + let values_to_accumulate = + slice_and_maybe_filter(values, opt_filter, &[row_idx, row_idx + 1])?; + prepared.push((accumulator, values_to_accumulate)); } - // Convert row to states. Use the grouped variant because the - // conversion-wide timer above accounts for the complete loop. - let values_to_accumulate = - slice_and_maybe_filter(values, opt_filter, &[row_idx, row_idx + 1])?; - converted_accumulator.update_batch_grouped(&values_to_accumulate)?; - let states = converted_accumulator.state()?; + // Time only aggregate-owned deduplication, once per bounded chunk. + let start = grouped_update_metric.as_ref().map(|_| Instant::now()); + let update_result: Result<()> = prepared.iter_mut().try_for_each( + |(accumulator, values_to_accumulate)| { + accumulator.update_batch_grouped(values_to_accumulate) + }, + ); + if let Some(start) = start { + aggregate_duration += start.elapsed(); + } + update_result?; - // Resize results to have enough columns according to the converted states - results.resize_with(states.len(), || Vec::with_capacity(num_rows)); + for (mut accumulator, _) in prepared { + let states = accumulator.state()?; - // Add the states to results - for (idx, state_val) in states.into_iter().enumerate() { - results[idx].push(state_val); + // Resize results to have enough columns according to the converted states + results.resize_with(states.len(), || Vec::with_capacity(num_rows)); + + // Add the states to results + for (idx, state_val) in states.into_iter().enumerate() { + results[idx].push(state_val); + } } } - if let (Some(metric), Some(start)) = (grouped_update_metric, start) { - metric.add_duration(start.elapsed()); + if let Some(metric) = grouped_update_metric { + metric.add_duration(aggregate_duration); } let arrays = results @@ -611,7 +629,7 @@ pub(crate) fn slice_and_maybe_filter( #[cfg(test)] mod tests { use std::sync::Arc; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::time::Duration; use super::*; @@ -628,6 +646,16 @@ mod tests { } } + #[derive(Debug)] + struct DurationMetric(Arc); + + impl AggregateMetric for DurationMetric { + fn add_duration(&self, duration: Duration) { + self.0 + .fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); + } + } + #[derive(Debug)] struct TimedAccumulator { metric: Arc, @@ -703,6 +731,63 @@ mod tests { Ok(()) } + #[derive(Debug)] + struct SlowStateAccumulator { + metric: Arc, + } + + impl Accumulator for SlowStateAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn grouped_update_batch_metric(&self) -> Option> { + Some(Arc::clone(&self.metric)) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(None)) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result> { + std::thread::sleep(Duration::from_millis(50)); + Ok(vec![ScalarValue::Int64(None)]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + } + + #[test] + fn adapter_convert_to_state_excludes_state_materialization_from_metric() -> Result<()> + { + let recorded_nanos = Arc::new(AtomicU64::new(0)); + let accumulator = GroupsAccumulatorAdapter::new({ + let recorded_nanos = Arc::clone(&recorded_nanos); + move || { + Ok(Box::new(SlowStateAccumulator { + metric: Arc::new(DurationMetric(Arc::clone(&recorded_nanos))), + }) as Box) + } + }); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let filter = BooleanArray::from(vec![true, false]); + accumulator.convert_to_state(&[values], Some(&filter))?; + + assert!( + recorded_nanos.load(Ordering::Relaxed) + < Duration::from_millis(25).as_nanos() as u64, + "internal metric must exclude state materialization" + ); + Ok(()) + } + #[test] fn adapter_preserving_evaluation_uses_accumulator_contract() -> Result<()> { let mut accumulator = GroupsAccumulatorAdapter::new(|| { From 16a9573290d4a5ba4047ec65bb0b22cc77e77316 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Fri, 11 Sep 2026 13:34:18 +0800 Subject: [PATCH 24/24] refactor: simplify metric lookup and consolidate implementation blocks, import AggregateMetrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - **Clear conditional metric lookup** – Refactored the conditional logic that retrieves metrics, removing redundant checks and making the code path easier to follow. - **Merged adjacent metrics impl blocks** – Consolidated neighboring implementation blocks that handled similar metric types, reducing duplication and improving maintainability. - **Imported AggregateMetrics; removed qualified repeats** – Added the `AggregateMetrics` import and eliminated unnecessary fully‑qualified references throughout the codebase. --- .../src/aggregate/groups_accumulator.rs | 8 ++++--- .../aggregates/aggregate_hash_table/common.rs | 8 +++---- .../aggregate_hash_table/common_ordered.rs | 6 ++--- .../src/aggregates/group_values/metrics.rs | 22 +++++++++---------- 4 files changed, 22 insertions(+), 22 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 0803ec0b9e167..5d2825e51ae09 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -272,9 +272,11 @@ impl GroupsAccumulatorAdapter { let values = take_arrays(values, &batch_indices, None)?; let opt_filter = get_filter_at_indices(opt_filter, &batch_indices)?; - let grouped_update_metric = time_grouped_update - .then(|| self.grouped_update_metric.get().and_then(Clone::clone)) - .flatten(); + let grouped_update_metric = if time_grouped_update { + self.grouped_update_metric.get().and_then(Clone::clone) + } else { + None + }; let mut sizes_pre = 0; let mut sizes_post = 0; 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 d30b63eca9557..75de50091b86e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -24,7 +24,7 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr::{AggregateMetrics, EmitTo, GroupsAccumulator}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use crate::PhysicalExpr; @@ -89,7 +89,7 @@ pub(in crate::aggregates) struct AggregateHashTable { pub(super) aggregate_accumulator_metrics: Arc, /// Optional internal metrics owned by each aggregate expression. - pub(super) aggregate_submetrics: Vec>, + pub(super) aggregate_submetrics: Vec>, /// Raw input schema, used to evaluate expressions and synthesize empty /// grouping-set rows. @@ -510,7 +510,7 @@ pub(super) struct HashAggregateAccumulator { accumulator: Box, /// Optional internal metrics owned by this aggregate expression. - submetrics: Arc, + submetrics: Arc, } pub(super) type AggregateAccumulator = HashAggregateAccumulator; @@ -648,7 +648,7 @@ impl HashAggregateAccumulator { arguments: Vec>, filter: Option>, accumulator: Box, - submetrics: Arc, + submetrics: Arc, ) -> Self { Self { aggregate_expr, diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs index ab4f27e0289b4..41fd121ef1b67 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -26,7 +26,7 @@ use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_common::assert_or_internal_err; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::EmitTo; +use datafusion_expr::{AggregateMetrics, EmitTo}; use crate::InputOrderMode; use crate::PhysicalExpr; @@ -52,7 +52,7 @@ pub(in crate::aggregates) struct OrderedAggregateTableMetrics { pub(super) group_by: GroupByMetrics, pub(super) aggregate_arguments: AggregateArgumentMetrics, pub(super) accumulator: Arc, - pub(super) submetrics: Vec>, + pub(super) submetrics: Vec>, } impl OrderedAggregateTableMetrics { @@ -142,7 +142,7 @@ pub(in crate::aggregates) struct OrderedAggregateTable { pub(super) aggregate_accumulator_metrics: Arc, /// Optional internal metrics owned by each aggregate expression. - pub(super) aggregate_submetrics: Vec>, + pub(super) aggregate_submetrics: Vec>, /// Group keys, ordering state, and accumulator states. pub(super) buffer: OrderedAggregateTableBuffer, diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index 518835c30473d..a0e0e67de7694 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -57,6 +57,16 @@ impl AggregateSubMetrics { additional_subphase_metrics: Mutex::new(HashMap::new()), } } + + fn new_metric(&self, subphase: &'static str) -> Arc { + let time = MetricBuilder::new(&self.metrics) + .with_new_label("aggregate", self.aggregate_label.clone()) + .subset_time( + format!("agg_expr_{}_internal_{}_time", self.index, subphase), + self.partition, + ); + Arc::new(AggregateSubMetric { time }) + } } #[derive(Debug)] @@ -70,18 +80,6 @@ impl AggregateMetric for AggregateSubMetric { } } -impl AggregateSubMetrics { - fn new_metric(&self, subphase: &'static str) -> Arc { - let time = MetricBuilder::new(&self.metrics) - .with_new_label("aggregate", self.aggregate_label.clone()) - .subset_time( - format!("agg_expr_{}_internal_{}_time", self.index, subphase), - self.partition, - ); - Arc::new(AggregateSubMetric { time }) - } -} - impl AggregateMetrics for AggregateSubMetrics { fn metric(&self, subphase: &'static str) -> Arc { if let Some((registered_subphase, metric)) = self.first_subphase_metric.get()