diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 419fd5be3ad2b..bc2ac0781a422 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: 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 @@ -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..aaef5414112cc --- /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 +-- 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(200000) +GROUP BY value / 2; diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index 7e9a4ae525ea3..d05ce9ae37465 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 + std::panic::RefUnwindSafe { + /// 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. @@ -58,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/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-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 97a2520eb2b9a..c87faede7acc6 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, }) } @@ -913,12 +916,12 @@ impl DistinctArrayAggAccumulator { } } -impl Accumulator for DistinctArrayAggAccumulator { - 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(()); } @@ -948,6 +951,10 @@ impl Accumulator for DistinctArrayAggAccumulator { return Ok(()); } + let distinct_metric = record_metric + .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. @@ -994,8 +1001,33 @@ impl Accumulator for DistinctArrayAggAccumulator { } } } + if let (Some(metric), Some(start)) = (distinct_metric, distinct_start) { + metric.add_duration(start.elapsed()); + } 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() { @@ -1473,6 +1505,68 @@ 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_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]))])?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + 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; 16]))])?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[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 = + DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; + + assert_eq!(accumulator.size(), size_of_val(&accumulator)); + Ok(()) + } #[test] fn no_duplicates_no_distinct() -> Result<()> { @@ -1760,7 +1854,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(()) } diff --git a/datafusion/physical-expr-common/src/metrics/value.rs b/datafusion/physical-expr-common/src/metrics/value.rs index 37ab5194b2cc4..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)) @@ -1278,6 +1287,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-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..38ec175f25a21 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,171 @@ 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?; + + 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_1_internal_distinct_time".to_string(), + "second".to_string(), + ), + ] + ); + 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 + ); + assert_eq!( + aggregate_metrics(&aggregate.metrics().unwrap(), "evaluate").len(), + 2 + ); + + 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![ diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index d6405bb403adc..518835c30473d 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -18,6 +18,111 @@ //! 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, OnceLock}; +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)] +struct AggregateSubMetrics { + metrics: ExecutionPlanMetricsSet, + partition: usize, + index: usize, + aggregate_label: String, + /// 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 { + fn new( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + index: usize, + aggregate_label: impl Into, + ) -> Self { + Self { + metrics: metrics.clone(), + partition, + index, + aggregate_label: aggregate_label.into(), + first_subphase_metric: OnceLock::new(), + additional_subphase_metrics: Mutex::new(HashMap::new()), + } + } +} + +#[derive(Debug)] +struct AggregateSubMetric { + time: Time, +} + +impl AggregateMetric for AggregateSubMetric { + fn add_duration(&self, 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 { + 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)), + ) + } +} + +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 +326,7 @@ impl GroupByMetrics { #[cfg(test)] mod tests { - use super::GroupByMetrics; + use super::{AggregateSubMetrics, GroupByMetrics, aggregate_sub_metrics}; use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use crate::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet}; use crate::test::TestMemoryExec; @@ -233,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::{ @@ -240,6 +346,104 @@ mod tests { }; use datafusion_physical_expr::expressions::col; 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_per_recording() { + 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() + .iter() + .find(|metric| { + metric.value().name() == "agg_expr_0_internal_distinct_time" + }) + .unwrap() + .value() + .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(); + 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_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(), 6); + 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..d3f1e7b090744 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))) } } @@ -1489,13 +1498,76 @@ 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?; + } + + 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(()) + } + // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] 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..7af82c45108ee 100644 --- a/docs/source/user-guide/metrics.md +++ b/docs/source/user-guide/metrics.md @@ -154,6 +154,23 @@ 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. 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`. 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 `datafusion.explain.analyze_level` includes `Dev` (the default), but are omitted