DRAFT Test aggmetrics3

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:     1000K-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[@]}"
                     ;;
@@ -1665,6 +1669,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
+-- 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;

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,33 @@ 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)
+    }
+
+    /// Merges state when called by a grouped accumulator adapter.
+    ///
+    /// The default delegates to [`Self::merge_batch`]. Implementations that
+    /// return a [`Self::grouped_update_batch_metric`] can avoid timing every
+    /// per-group merge; the adapter records one interval for the full batch.
+    fn merge_batch_grouped(&mut self, states: &[ArrayRef]) -> Result<()> {
+        self.merge_batch(states)
+    }
+
     /// Returns the final aggregate value.
     ///
     /// For example, the `SUM` accumulator maintains a running sum,

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/Cargo.toml

@@ -57,3 +57,7 @@ rand = { workspace = true }
 [[bench]]
 harness = false
 name = "accumulate"
+
+[[bench]]
+harness = false
+name = "groups_accumulator_adapter"

datafusion/functions-aggregate-common/benches/groups_accumulator_adapter.rs

@@ -0,0 +1,62 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Benchmarks the uninstrumented `GroupsAccumulatorAdapter` update path.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array};
+use arrow::datatypes::DataType;
+use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
+use datafusion_expr_common::accumulator::Accumulator;
+use datafusion_expr_common::groups_accumulator::GroupsAccumulator;
+use datafusion_functions_aggregate_common::aggregate::groups_accumulator::GroupsAccumulatorAdapter;
+use datafusion_functions_aggregate_common::min_max::MaxAccumulator;
+
+const NUM_GROUPS: usize = 8_192;
+
+fn groups_accumulator_adapter(c: &mut Criterion) {
+    let values: ArrayRef = Arc::new(Int64Array::from_iter_values(0..NUM_GROUPS as i64));
+    let group_indices: Vec<_> = (0..NUM_GROUPS).collect();
+
+    c.bench_function("groups_accumulator_adapter/update_batch/8192_groups", |b| {
+        b.iter_batched(
+            || {
+                GroupsAccumulatorAdapter::new(|| {
+                    Ok(Box::new(MaxAccumulator::try_new(&DataType::Int64)?)
+                        as Box<dyn Accumulator>)
+                })
+            },
+            |mut accumulator| {
+                accumulator
+                    .update_batch(
+                        &[Arc::clone(&values)],
+                        &group_indices,
+                        None,
+                        NUM_GROUPS,
+                    )
+                    .unwrap();
+                black_box(accumulator);
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}
+
+criterion_group!(benches, groups_accumulator_adapter);
+criterion_main!(benches);