From 54fb2630d9fa32702151874bfd73ea19eed3deca Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 12:12:50 +0800 Subject: [PATCH 01/17] =?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 e4199f5354114cf031cfbefc903162eadb7e44b8 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 16:51:07 +0800 Subject: [PATCH 02/17] 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 90b6f659b09d73049f92c01e4157e3a249fcce3e Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 12:33:27 +0800 Subject: [PATCH 03/17] empty commit From dc938f0d7d68d2976cd28ee1c81f330fe5dc8e85 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:15:20 +0800 Subject: [PATCH 04/17] 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 561e89c22d889fd71c552b8fc27327e05819f371 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:26:11 +0800 Subject: [PATCH 05/17] 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 2f45a329cb9225adcafeaeea85e65a55e648dca7 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:38:18 +0800 Subject: [PATCH 06/17] 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 83c49b52709b548c2a062932895808a5b678fcce Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 13:40:07 +0800 Subject: [PATCH 07/17] 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 c5e654295ae043cb6e3f67bfbeae36ca25671c99 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 14:02:04 +0800 Subject: [PATCH 08/17] 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 02d67629d37694b358f949486d9621945f9b0f3f Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Tue, 8 Sep 2026 14:12:17 +0800 Subject: [PATCH 09/17] =?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 ff26434525ae9476724145cf28b314ef071f0329 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 10:59:31 +0800 Subject: [PATCH 10/17] =?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 78e3cfce24518237c7f5c262f3c8200e585bb391 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:11:21 +0800 Subject: [PATCH 11/17] =?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 b6e945020605d385333a5a533b043726b632b0a0 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:16:27 +0800 Subject: [PATCH 12/17] =?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 e3c11e7a266ae9956f9c7c3e007d13301279522a Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:30:37 +0800 Subject: [PATCH 13/17] =?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 8078bd7cb7ac50b663d691b898736cdfb2e42406 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 22:05:58 +0800 Subject: [PATCH 14/17] 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 e660086ce1b59a8c681b99758c71dca0079f9246 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 11:36:49 +0800 Subject: [PATCH 15/17] 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 8fc8b2d8687579d0a13b27efe099f00a8618fb66 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 22:05:58 +0800 Subject: [PATCH 16/17] 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. From fe1ca2432935210f71da16357e74787697e372b4 Mon Sep 17 00:00:00 2001 From: Siew Kam Onn Date: Wed, 9 Sep 2026 22:44:23 +0800 Subject: [PATCH 17/17] 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