Test aggmetrics

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"

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."

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;

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<dyn AggregateMetric>;
+}
 
 /// 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<dyn AggregateMetrics>) {}
+
     /// 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<Arc<dyn AggregateMetric>> {
+        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,

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<dyn AggregateMetrics>) {}
+
     /// Updates the accumulator's state from its arguments, encoded as
     /// a vector of [`ArrayRef`]s.
     ///

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,

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<Arc<dyn AggregateMetric>>,
 }
 
 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(())