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..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; diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index 7e9a4ae525ea3..aa8be52a18fc3 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,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> { + 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, 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..5d2825e51ae09 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, OnceLock}; +use std::time::Duration; use arrow::array::new_empty_array; use arrow::{ @@ -32,8 +34,8 @@ use arrow::{ compute::take_arrays, datatypes::UInt32Type, }; -use datafusion_common::{Result, ScalarValue, arrow_datafusion_err}; -use datafusion_expr_common::accumulator::Accumulator; +use datafusion_common::{Result, ScalarValue, arrow_datafusion_err, instant::Instant}; +use datafusion_expr_common::accumulator::{Accumulator, AggregateMetric}; use datafusion_expr_common::groups_accumulator::{ EmitTo, GroupSelection, GroupsAccumulator, }; @@ -102,8 +104,19 @@ 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. + /// + /// The cache is shared with factories that install metrics on only the + /// first accumulator. This avoids resolving the same expression metric for + /// every group while allowing `convert_to_state(&self)` to use it. + grouped_update_metric: Arc>>>, } +/// Maximum number of prepared group inputs retained while timing an +/// aggregate-owned grouped subphase. +const GROUPED_METRIC_PREPARATION_CHUNK_SIZE: usize = 64; + struct AccumulatorState { /// [`Accumulator`] that stores the per-group state accumulator: Box, @@ -132,6 +145,21 @@ impl GroupsAccumulatorAdapter { /// Create a new adapter that will create a new [`Accumulator`] /// for each group, using the specified factory function pub fn new(factory: F) -> Self + where + F: Fn() -> Result> + Send + 'static, + { + Self::new_with_grouped_update_metric_cache(factory, Arc::new(OnceLock::new())) + } + + /// Creates an adapter with a metric cache shared by its accumulator factory. + /// + /// The factory must initialize the cache before returning its first + /// metric-enabled accumulator. Later accumulators may be created without + /// metrics because this adapter records the grouped batch timing. + pub fn new_with_grouped_update_metric_cache( + factory: F, + grouped_update_metric: Arc>>>, + ) -> Self where F: Fn() -> Result> + Send + 'static, { @@ -139,6 +167,7 @@ impl GroupsAccumulatorAdapter { factory: Box::new(factory), states: vec![], allocation_bytes: 0, + grouped_update_metric, } } @@ -152,6 +181,8 @@ impl GroupsAccumulatorAdapter { let new_accumulators = total_num_groups - self.states.len(); for _ in 0..new_accumulators { let accumulator = (self.factory)()?; + self.grouped_update_metric + .get_or_init(|| accumulator.grouped_update_batch_metric()); let state = AccumulatorState::new(accumulator); self.add_allocation(state.size()); self.states.push(state); @@ -191,6 +222,7 @@ impl GroupsAccumulatorAdapter { group_indices: &[usize], opt_filter: Option<&BooleanArray>, total_num_groups: usize, + time_grouped_update: bool, f: F, ) -> Result<()> where @@ -240,30 +272,63 @@ impl GroupsAccumulatorAdapter { let values = take_arrays(values, &batch_indices, None)?; let opt_filter = get_filter_at_indices(opt_filter, &batch_indices)?; - // invoke each accumulator with the appropriate rows, first - // pulling the input arguments for this group into their own - // RecordBatch(es) - let iter = groups_with_rows.iter().zip(offsets.windows(2)); + let grouped_update_metric = if time_grouped_update { + self.grouped_update_metric.get().and_then(Clone::clone) + } else { + None + }; let mut sizes_pre = 0; let mut sizes_post = 0; - 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 mut aggregate_duration = Duration::ZERO; + let result: Result<()> = (|| { + // Keep preparation bounded to avoid retaining one filtered array per + // group. Time only accumulator invocation: slicing and filtering + // are adapter work, not aggregate-owned subphase work. + for (chunk_index, groups) in groups_with_rows + .chunks(GROUPED_METRIC_PREPARATION_CHUNK_SIZE) + .enumerate() + { + let first_offset = chunk_index * GROUPED_METRIC_PREPARATION_CHUNK_SIZE; + let values_to_accumulate = groups + .iter() + .enumerate() + .map(|(index, &group_idx)| { + let values = slice_and_maybe_filter( + &values, + opt_filter.as_ref().map(|f| f.as_boolean()), + &offsets[first_offset + index..first_offset + index + 2], + )?; + Ok((group_idx, values)) + }) + .collect::>>()?; + + let start = grouped_update_metric.as_ref().map(|_| Instant::now()); + let chunk_result: Result<()> = (|| { + for (group_idx, values) in values_to_accumulate { + let state = &mut self.states[group_idx]; + sizes_pre += state.size(); + f(state.accumulator.as_mut(), &values)?; + + // clear out the state so they are empty for next + // iteration + state.indices.clear(); + sizes_post += state.size(); + } + Ok(()) + })(); + if let Some(start) = start { + aggregate_duration += start.elapsed(); + } + chunk_result?; + } + Ok(()) + })(); + if let Some(metric) = grouped_update_metric { + metric.add_duration(aggregate_duration); + } + result?; self.adjust_allocation(sizes_pre, sizes_post); Ok(()) } @@ -310,8 +375,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,9 +478,9 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { group_indices, None, total_num_groups, + true, |accumulator, values_to_accumulate| { - accumulator.merge_batch(values_to_accumulate)?; - Ok(()) + accumulator.merge_batch_grouped(values_to_accumulate) }, )?; Ok(()) @@ -445,25 +511,57 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { // Each row has its respective group let mut results = vec![]; - for row_idx in 0..num_rows { - // Create the empty accumulator for converting - let mut converted_accumulator = (self.factory)()?; + let mut grouped_update_metric = None; + let mut aggregate_duration = Duration::ZERO; + for chunk_start in (0..num_rows).step_by(GROUPED_METRIC_PREPARATION_CHUNK_SIZE) { + let chunk_end = + (chunk_start + GROUPED_METRIC_PREPARATION_CHUNK_SIZE).min(num_rows); + let mut prepared = Vec::with_capacity(chunk_end - chunk_start); + + for row_idx in chunk_start..chunk_end { + // Create the empty accumulator and prepare adapter-owned input + // outside the aggregate submetric. + let accumulator = (self.factory)()?; + if row_idx == 0 { + let metric = self + .grouped_update_metric + .get_or_init(|| accumulator.grouped_update_batch_metric()); + grouped_update_metric.clone_from(metric); + } + let values_to_accumulate = + slice_and_maybe_filter(values, opt_filter, &[row_idx, row_idx + 1])?; + prepared.push((accumulator, values_to_accumulate)); + } - // Convert row to states - let values_to_accumulate = - slice_and_maybe_filter(values, opt_filter, &[row_idx, row_idx + 1])?; - converted_accumulator.update_batch(&values_to_accumulate)?; - let states = converted_accumulator.state()?; + // Time only aggregate-owned deduplication, once per bounded chunk. + let start = grouped_update_metric.as_ref().map(|_| Instant::now()); + let update_result: Result<()> = prepared.iter_mut().try_for_each( + |(accumulator, values_to_accumulate)| { + accumulator.update_batch_grouped(values_to_accumulate) + }, + ); + if let Some(start) = start { + aggregate_duration += start.elapsed(); + } + update_result?; - // Resize results to have enough columns according to the converted states - results.resize_with(states.len(), || Vec::with_capacity(num_rows)); + for (mut accumulator, _) in prepared { + let states = accumulator.state()?; - // Add the states to results - for (idx, state_val) in states.into_iter().enumerate() { - results[idx].push(state_val); + // Resize results to have enough columns according to the converted states + results.resize_with(states.len(), || Vec::with_capacity(num_rows)); + + // Add the states to results + for (idx, state_val) in states.into_iter().enumerate() { + results[idx].push(state_val); + } } } + if let Some(metric) = grouped_update_metric { + metric.add_duration(aggregate_duration); + } + let arrays = results .into_iter() .map(ScalarValue::iter_to_array) @@ -533,12 +631,165 @@ pub(crate) fn slice_and_maybe_filter( #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; + use std::time::Duration; use super::*; use crate::min_max::MaxAccumulator; - use arrow::array::{AsArray, Int64Array}; + use arrow::array::{AsArray, BooleanArray, Int64Array}; use arrow::datatypes::{DataType, Int64Type}; + #[derive(Debug)] + struct CountingMetric(Arc); + + impl AggregateMetric for CountingMetric { + fn add_duration(&self, _duration: Duration) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[derive(Debug)] + struct DurationMetric(Arc); + + impl AggregateMetric for DurationMetric { + fn add_duration(&self, duration: Duration) { + self.0 + .fetch_add(duration.as_nanos() as u64, Ordering::Relaxed); + } + } + + #[derive(Debug)] + struct TimedAccumulator { + metric: Arc, + } + + impl Accumulator for TimedAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + self.metric.add_duration(Duration::ZERO); + Ok(()) + } + + fn grouped_update_batch_metric(&self) -> Option> { + Some(Arc::clone(&self.metric)) + } + + fn update_batch_grouped(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(None)) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Int64(None)]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + } + + #[test] + fn adapter_grouped_update_records_one_metric_after_filtering() -> Result<()> { + let metric_updates = Arc::new(AtomicUsize::new(0)); + let mut accumulator = GroupsAccumulatorAdapter::new({ + let metric_updates = Arc::clone(&metric_updates); + move || { + Ok(Box::new(TimedAccumulator { + metric: Arc::new(CountingMetric(Arc::clone(&metric_updates))), + }) as Box) + } + }); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 4])); + let filter = BooleanArray::from(vec![true, false, true, false]); + accumulator.update_batch(&[values], &[0, 0, 1, 1], Some(&filter), 2)?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn adapter_convert_to_state_records_metric_once() -> Result<()> { + let metric_updates = Arc::new(AtomicUsize::new(0)); + let accumulator = GroupsAccumulatorAdapter::new({ + let metric_updates = Arc::clone(&metric_updates); + move || { + Ok(Box::new(TimedAccumulator { + metric: Arc::new(CountingMetric(Arc::clone(&metric_updates))), + }) as Box) + } + }); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + accumulator.convert_to_state(&[values], None)?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[derive(Debug)] + struct SlowStateAccumulator { + metric: Arc, + } + + impl Accumulator for SlowStateAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn grouped_update_batch_metric(&self) -> Option> { + Some(Arc::clone(&self.metric)) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(None)) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result> { + std::thread::sleep(Duration::from_millis(50)); + Ok(vec![ScalarValue::Int64(None)]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + } + + #[test] + fn adapter_convert_to_state_excludes_state_materialization_from_metric() -> Result<()> + { + let recorded_nanos = Arc::new(AtomicU64::new(0)); + let accumulator = GroupsAccumulatorAdapter::new({ + let recorded_nanos = Arc::clone(&recorded_nanos); + move || { + Ok(Box::new(SlowStateAccumulator { + metric: Arc::new(DurationMetric(Arc::clone(&recorded_nanos))), + }) as Box) + } + }); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let filter = BooleanArray::from(vec![true, false]); + accumulator.convert_to_state(&[values], Some(&filter))?; + + assert!( + recorded_nanos.load(Ordering::Relaxed) + < Duration::from_millis(25).as_nanos() as u64, + "internal metric must exclude state materialization" + ); + Ok(()) + } + #[test] fn adapter_preserving_evaluation_uses_accumulator_contract() -> Result<()> { let mut accumulator = GroupsAccumulatorAdapter::new(|| { diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 97a2520eb2b9a..95a8d6c5274f3 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.clone()) + .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,22 +1001,73 @@ impl Accumulator for DistinctArrayAggAccumulator { } } } + if let (Some(metric), Some(start)) = (distinct_metric, distinct_start) { + metric.add_duration(start.elapsed()); + } Ok(()) } +} - fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { +impl DistinctArrayAggAccumulator { + fn merge_batch_impl( + &mut self, + states: &[ArrayRef], + record_metric: bool, + ) -> Result<()> { if states.is_empty() { return Ok(()); } assert_eq_or_internal_err!(states.len(), 1, "expects single state"); - // The DISTINCT state is `List`. - states[0] + let distinct_metric = record_metric + .then(|| self.distinct_metric.clone()) + .flatten(); + let distinct_start = distinct_metric.as_ref().map(|_| Instant::now()); + + // The DISTINCT state is `List`. This calls the update + // implementation once per state row, so record the submetric once for + // the entire merge rather than once per row. + let result = states[0] .as_list::() .iter() .flatten() - .try_for_each(|val| self.update_batch(&[val])) + .try_for_each(|val| self.update_batch_impl(&[val], false)); + + if let (Some(metric), Some(start)) = (distinct_metric, distinct_start) { + metric.add_duration(start.elapsed()); + } + result + } +} + +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.clone() + } + + 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<()> { + self.merge_batch_impl(states, true) + } + + fn merge_batch_grouped(&mut self, states: &[ArrayRef]) -> Result<()> { + self.merge_batch_impl(states, false) } fn evaluate(&mut self) -> Result { @@ -1467,12 +1525,92 @@ impl Accumulator for OrderSensitiveArrayAggAccumulator { #[cfg(test)] mod tests { use super::*; - use arrow::array::{ListBuilder, StringBuilder}; + use arrow::array::{Int32Builder, ListBuilder, StringBuilder}; use arrow::datatypes::Schema; use datafusion_common::cast::as_generic_string_array; use datafusion_common::internal_err; 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_records_merge_metric_once() -> Result<()> { + let mut builder = ListBuilder::new(Int32Builder::new()); + for value in [1, 2, 3] { + builder.append_value([Some(value)]); + } + let state: ArrayRef = Arc::new(builder.finish()); + + let metric_updates = Arc::new(AtomicUsize::new(0)); + let mut target = + DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?; + target.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates)))); + target.merge_batch(&[state])?; + + assert_eq!(metric_updates.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[test] + fn distinct_accumulator_preserves_unwind_auto_traits() { + fn assert_unwind_traits() { + } + + 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 +1898,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..75de50091b86e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -24,7 +24,7 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_expr::{AggregateMetrics, EmitTo, GroupsAccumulator}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use crate::PhysicalExpr; @@ -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..41fd121ef1b67 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -26,7 +26,7 @@ use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_common::assert_or_internal_err; use datafusion_execution::memory_pool::proxy::VecAllocExt; -use datafusion_expr::EmitTo; +use datafusion_expr::{AggregateMetrics, EmitTo}; use crate::InputOrderMode; use crate::PhysicalExpr; @@ -52,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..a0e0e67de7694 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -18,6 +18,109 @@ //! 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()), + } + } + + fn new_metric(&self, subphase: &'static str) -> Arc { + let time = MetricBuilder::new(&self.metrics) + .with_new_label("aggregate", self.aggregate_label.clone()) + .subset_time( + format!("agg_expr_{}_internal_{}_time", self.index, subphase), + self.partition, + ); + Arc::new(AggregateSubMetric { time }) + } +} + +#[derive(Debug)] +struct AggregateSubMetric { + time: Time, +} + +impl AggregateMetric for AggregateSubMetric { + fn add_duration(&self, duration: Duration) { + self.time.add_duration_exact(duration); + } +} + +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 +324,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 +336,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 +344,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..20769559a2006 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -17,7 +17,7 @@ //! Hash aggregation -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::task::{Context, Poll}; use std::vec; @@ -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,18 +651,35 @@ 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!( "Creating GroupsAccumulatorAdapter for {}: {agg_expr:?}", agg_expr.name() ); + let grouped_update_metric = Arc::new(OnceLock::new()); + let factory_metric = Arc::clone(&grouped_update_metric); let agg_expr_captured = Arc::clone(agg_expr); - let factory = move || agg_expr_captured.create_accumulator(); - Ok(Box::new(GroupsAccumulatorAdapter::new(factory))) + let factory = move || { + let mut accumulator = agg_expr_captured.create_accumulator()?; + if factory_metric.get().is_none() { + accumulator.set_metrics(Arc::clone(&metrics)); + // This factory is called serially by one adapter, so the first + // accumulator is the only one that resolves the metric handle. + let _ = factory_metric.set(accumulator.grouped_update_batch_metric()); + } + Ok(accumulator) + }; + Ok(Box::new( + GroupsAccumulatorAdapter::new_with_grouped_update_metric_cache( + factory, + grouped_update_metric, + ), + )) } } @@ -1489,12 +1512,169 @@ 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_expr::AggregateMetric; + use datafusion_functions_aggregate::{array_agg::array_agg_udaf, count::count_udaf}; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + #[derive(Debug)] + struct CountingMetric(Arc); + + impl AggregateMetric for CountingMetric { + fn add_duration(&self, _duration: Duration) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[derive(Debug)] + struct CountingMetrics { + resolutions: Arc, + durations: Arc, + } + + impl AggregateMetrics for CountingMetrics { + fn metric(&self, _subphase: &'static str) -> Arc { + self.resolutions.fetch_add(1, Ordering::Relaxed); + Arc::new(CountingMetric(Arc::clone(&self.durations))) + } + } + + #[test] + fn legacy_grouped_distinct_resolves_metric_once_for_conversion_and_update() + -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::UInt32, + false, + )])); + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .distinct() + .alias("distinct_values") + .build()?, + ); + let metric_resolutions = Arc::new(AtomicUsize::new(0)); + let metric_durations = Arc::new(AtomicUsize::new(0)); + let mut accumulator = create_group_accumulator( + &aggregate_expr, + Arc::new(CountingMetrics { + resolutions: Arc::clone(&metric_resolutions), + durations: Arc::clone(&metric_durations), + }), + )?; + let values: ArrayRef = Arc::new(UInt32Array::from(vec![1, 2, 3])); + + accumulator.convert_to_state(&[Arc::clone(&values)], None)?; + accumulator.update_batch(&[values], &[0, 1, 2], None, 3)?; + + assert_eq!(metric_resolutions.load(Ordering::Relaxed), 1); + assert_eq!(metric_durations.load(Ordering::Relaxed), 2); + Ok(()) + } + + #[test] + fn legacy_grouped_distinct_merge_records_metric_once() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::UInt32, + false, + )])); + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(array_agg_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .distinct() + .alias("distinct_values") + .build()?, + ); + let metric_resolutions = Arc::new(AtomicUsize::new(0)); + let metric_durations = Arc::new(AtomicUsize::new(0)); + let mut accumulator = create_group_accumulator( + &aggregate_expr, + Arc::new(CountingMetrics { + resolutions: Arc::clone(&metric_resolutions), + durations: Arc::clone(&metric_durations), + }), + )?; + + let mut state = ListBuilder::new(UInt32Builder::new()); + for value in [1, 2, 3] { + state.append_value([Some(value)]); + } + accumulator.merge_batch(&[Arc::new(state.finish())], &[1, 2, 3], 4)?; + + assert_eq!(metric_resolutions.load(Ordering::Relaxed), 1); + assert_eq!(metric_durations.load(Ordering::Relaxed), 1); + Ok(()) + } + + #[tokio::test] + async fn grouped_hash_stream_reports_distinct_array_agg_submetric() -> Result<()> { + 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. 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