This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24024-9fdf1447ef9e797181f9b93e81cde937181093e3 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit aff5850c0aa604bb03158a137d5343e4efe73ae3 Author: kosiew <[email protected]> AuthorDate: Sun Aug 16 12:21:31 2026 +0000 Add per-aggregate argument evaluation metrics for grouped hash aggregation (#24024) #### Which issue does this PR close? * Part of #23570 --- ## Rationale for this change Grouped hash aggregation currently exposes a single `aggregate_arguments_time` metric for all aggregate argument evaluation. This makes it difficult to determine which aggregate expression is responsible for argument evaluation cost when multiple aggregates are present (for example, `SUM(a)` versus `SUM(b)`). This change adds per-aggregate argument evaluation metrics while preserving the existing operator-level metric for compatibility. --- ## What changes are included in this PR? * Add an `AggregateArgumentMetrics` helper that registers one timer per aggregate expression during operator construction. * Label per-aggregate metrics using the aggregate expression display text (or alias), allowing otherwise identical aggregate functions on different inputs (for example, `SUM(a)` and `SUM(b)`) to be distinguished. * Wrap aggregate argument evaluation with the corresponding per-aggregate timer in: * grouped hash aggregation stream * migrated hash table aggregation paths * ordered aggregation table paths * Introduce `OrderedAggregateTableMetrics` to carry both group-by and aggregate argument metrics through ordered aggregation replay/spill paths. * Factor out a reusable `aggregate_metric_label` helper for consistent metric labels. * Preserve the existing `aggregate_arguments_time` metric so whole-phase timing continues to be reported. --- ## Are these changes tested? Yes. This PR adds and updates tests including: * `test_groupby_aggregate_argument_metrics_distinguish_inputs`, which verifies that separate metrics are created for `SUM(a)` and `SUM(b)` with distinct metric names and aggregate labels. * Existing group-by metrics tests updated to use shared aggregate construction helpers. * Grouped hash aggregation stream tests updated to verify both: * the existing `aggregate_arguments_time` metric remains present and non-zero, and * the new per-aggregate metric (`agg_expr_0_arguments_time`) is emitted and records time. --- ## Are there any user-facing changes? Yes. `EXPLAIN ANALYZE` and execution plan metrics for the covered grouped hash aggregation paths now include per-aggregate argument evaluation timers (for example, `agg_expr_0_arguments_time`) with aggregate labels identifying the corresponding aggregate expression, while retaining the existing `aggregate_arguments_time` metric for compatibility. --- ## LLM-generated code disclosure This PR includes LLM-generated code and comments. All LLM-generated content has been manually reviewed. --- .../src/aggregates/aggregate_hash_table/common.rs | 29 +++- .../aggregate_hash_table/common_ordered.rs | 60 ++++++- .../src/aggregates/aggregate_hash_table/mod.rs | 2 +- .../aggregate_hash_table/ordered_final_table.rs | 7 +- .../aggregate_hash_table/ordered_partial_table.rs | 7 +- .../aggregate_hash_table/ordered_single_table.rs | 7 +- .../aggregate_hash_table/partial_table.rs | 1 + .../src/aggregates/group_values/metrics.rs | 190 ++++++++++++++++++--- .../src/aggregates/group_values/mod.rs | 2 +- .../src/aggregates/grouped_hash_stream.rs | 61 ++++++- .../physical-plan/src/aggregates/hash_stream.rs | 12 +- datafusion/physical-plan/src/aggregates/mod.rs | 4 + .../src/aggregates/ordered_final_stream.rs | 21 +-- .../src/aggregates/ordered_single_stream.rs | 13 +- .../physical-plan/src/aggregates/single_stream.rs | 13 +- 15 files changed, 343 insertions(+), 86 deletions(-) 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 91e9d6555c..fde7024f7b 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -27,11 +27,14 @@ use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use crate::PhysicalExpr; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, +}; use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ - AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, + AggregateExec, PhysicalGroupBy, aggregate_expressions, aggregate_metric_label, + evaluate_group_by, }; /// Marker for raw rows -> partial state aggregation. @@ -75,6 +78,9 @@ pub(in crate::aggregates) struct AggregateHashTable<AggrMode> { /// Grouping and accumulator-specific timing metrics. pub(super) group_by_metrics: GroupByMetrics, + /// Per-aggregate timing metrics for evaluating aggregate arguments. + pub(super) aggregate_argument_metrics: AggregateArgumentMetrics, + /// Raw input schema, used to evaluate expressions and synthesize empty /// grouping-set rows. pub(super) input_schema: SchemaRef, @@ -134,8 +140,17 @@ impl<AggrMode> AggregateHashTable<AggrMode> { let group_schema = agg.group_by.group_schema(&input_schema)?; let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let aggregate_argument_metrics = AggregateArgumentMetrics::new( + &agg.metrics, + partition, + agg.aggr_expr + .iter() + .map(|agg_expr| aggregate_metric_label(agg_expr)), + ); + Ok(Self { group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), + aggregate_argument_metrics, input_schema, output_schema, state_schema, @@ -169,7 +184,11 @@ impl<AggrMode> AggregateHashTable<AggrMode> { .building() .accumulators .iter() - .map(|acc| acc.evaluate_acc_args(batch)) + .enumerate() + .map(|(idx, acc)| { + self.aggregate_argument_metrics + .time(idx, || acc.evaluate_acc_args(batch)) + }) .collect::<Result<Vec<_>>>()?; drop(timer); @@ -288,10 +307,6 @@ impl<AggrMode> AggregateHashTable<AggrMode> { } } - pub(in crate::aggregates) fn group_by_metrics(&self) -> &GroupByMetrics { - &self.group_by_metrics - } - /// Returns the number of distinct groups accumulated so far. pub(in crate::aggregates) fn building_group_count(&self) -> usize { self.state.building().group_values.len() 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 dac0d4b7c5..22fe8f5433 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 @@ -30,19 +30,52 @@ use datafusion_expr::EmitTo; use crate::InputOrderMode; use crate::PhysicalExpr; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, +}; use crate::aggregates::grouped_hash_stream::create_group_accumulator; use crate::aggregates::order::GroupOrdering; use crate::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, - evaluate_group_by, + aggregate_metric_label, evaluate_group_by, }; use super::common::{ - AggregateAccumulator, AggregateBatchFn, EvaluatedAggregateBatch, + AggregateAccumulator, AggregateBatchFn, AggregateHashTable, EvaluatedAggregateBatch, MaterializeAccumulatorFn, }; +#[derive(Clone)] +pub(in crate::aggregates) struct OrderedAggregateTableMetrics { + pub(super) group_by: GroupByMetrics, + pub(super) aggregate_arguments: AggregateArgumentMetrics, +} + +impl OrderedAggregateTableMetrics { + pub(in crate::aggregates) fn new(agg: &AggregateExec, partition: usize) -> Self { + let aggregate_arguments = AggregateArgumentMetrics::new( + &agg.metrics, + partition, + agg.aggr_expr + .iter() + .map(|agg_expr| aggregate_metric_label(agg_expr)), + ); + Self { + group_by: GroupByMetrics::new(&agg.metrics, partition), + aggregate_arguments, + } + } + + pub(in crate::aggregates) fn from_hash_table<AggrMode>( + table: &AggregateHashTable<AggrMode>, + ) -> Self { + Self { + group_by: table.group_by_metrics.clone(), + aggregate_arguments: table.aggregate_argument_metrics.clone(), + } + } +} + /// Aggregate table shared by the ordered single, partial and final paths. /// /// # Ordering optimization @@ -100,6 +133,9 @@ pub(in crate::aggregates) struct OrderedAggregateTable<OrderedAggrMode> { /// Grouping and accumulator-specific timing metrics. pub(super) group_by_metrics: GroupByMetrics, + /// Per-aggregate timing metrics for evaluating aggregate arguments. + pub(super) aggregate_argument_metrics: AggregateArgumentMetrics, + /// Group keys, ordering state, and accumulator states. pub(super) buffer: OrderedAggregateTableBuffer, @@ -149,7 +185,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { input_order_mode: &InputOrderMode, aggregate_mode: &AggregateMode, filters: Vec<Option<Arc<dyn PhysicalExpr>>>, - group_by_metrics: GroupByMetrics, + metrics: OrderedAggregateTableMetrics, ) -> Result<Self> { assert_or_internal_err!( batch_size > 0, @@ -184,7 +220,8 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { output_schema, state_schema, batch_size, - group_by_metrics, + group_by_metrics: metrics.group_by, + aggregate_argument_metrics: metrics.aggregate_arguments, buffer: OrderedAggregateTableBuffer { group_by: Arc::clone(&agg.group_by), group_ordering, @@ -213,7 +250,11 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { .buffer .accumulators .iter() - .map(|acc| acc.evaluate_acc_args(batch)) + .enumerate() + .map(|(idx, acc)| { + self.aggregate_argument_metrics + .time(idx, || acc.evaluate_acc_args(batch)) + }) .collect::<Result<Vec<_>>>()?; drop(timer); @@ -259,8 +300,11 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> { + self.buffer.group_indices.allocated_size() } - pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics { - self.group_by_metrics.clone() + pub(in crate::aggregates) fn metrics(&self) -> OrderedAggregateTableMetrics { + OrderedAggregateTableMetrics { + group_by: self.group_by_metrics.clone(), + aggregate_arguments: self.aggregate_argument_metrics.clone(), + } } /// Takes every intermediate aggregate state and resets the table so it can 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 fbf3ccc738..435289aa30 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -29,4 +29,4 @@ pub(super) use common::{ AggregateHashTable, FinalMarker, PartialMarker, PartialReduceMarker, PartialSkipMarker, SingleMarker, }; -pub(super) use common_ordered::OrderedAggregateTable; +pub(super) use common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs index 1b7d419dd5..f3e22cdd0c 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -27,11 +27,10 @@ use datafusion_common::Result; use crate::InputOrderMode; use crate::aggregates::aggregate_hash_table::FinalMarker; -use crate::aggregates::group_values::GroupByMetrics; use crate::aggregates::{AggregateExec, AggregateMode}; use super::common::HashAggregateAccumulator; -use super::common_ordered::OrderedAggregateTable; +use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; /// Implementation specific to final aggregation, where the table stores partial /// aggregate states and the input rows are also partial states. @@ -49,7 +48,7 @@ impl OrderedAggregateTable<FinalMarker> { output_schema: SchemaRef, batch_size: usize, input_order_mode: &InputOrderMode, - group_by_metrics: GroupByMetrics, + metrics: OrderedAggregateTableMetrics, ) -> Result<Self> { Self::new_for_mode( agg, @@ -60,7 +59,7 @@ impl OrderedAggregateTable<FinalMarker> { input_order_mode, &AggregateMode::Final, vec![None; agg.aggr_expr.len()], - group_by_metrics, + metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs index 8545289f99..ea38346729 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -37,11 +37,10 @@ use datafusion_common::Result; use crate::aggregates::{ AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker, - group_values::GroupByMetrics, }; use super::common::HashAggregateAccumulator; -use super::common_ordered::OrderedAggregateTable; +use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; /// Implementation specific to partial aggregation, where the table stores /// partial aggregate states and the input rows are raw rows. @@ -61,7 +60,7 @@ impl OrderedAggregateTable<PartialMarker> { ) -> Result<Self> { let input_schema = agg.input().schema(); let state_schema = Arc::clone(&output_schema); - let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + let metrics = OrderedAggregateTableMetrics::new(agg, partition); Self::new_for_mode( agg, &input_schema, @@ -71,7 +70,7 @@ impl OrderedAggregateTable<PartialMarker> { &agg.input_order_mode, &AggregateMode::Partial, agg.filter_expr.iter().cloned().collect(), - group_by_metrics, + metrics, ) } diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs index 8ba50e2a59..db53e2822e 100644 --- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs @@ -24,11 +24,10 @@ use arrow::record_batch::RecordBatch; use datafusion_common::Result; use crate::aggregates::aggregate_hash_table::SingleMarker; -use crate::aggregates::group_values::GroupByMetrics; use crate::aggregates::{AggregateExec, AggregateMode}; use super::common::HashAggregateAccumulator; -use super::common_ordered::OrderedAggregateTable; +use super::common_ordered::{OrderedAggregateTable, OrderedAggregateTableMetrics}; /// Implementation specific to single aggregation, where the table stores final /// aggregate values and the input rows are raw rows. @@ -53,7 +52,7 @@ impl OrderedAggregateTable<SingleMarker> { )); let input_schema = agg.input().schema(); - let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + let metrics = OrderedAggregateTableMetrics::new(agg, partition); Self::new_for_mode( agg, &input_schema, @@ -63,7 +62,7 @@ impl OrderedAggregateTable<SingleMarker> { &agg.input_order_mode, &agg.mode, agg.filter_expr.iter().cloned().collect(), - group_by_metrics, + metrics, ) } 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 a64fd32536..915e7f1beb 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 @@ -86,6 +86,7 @@ impl AggregateHashTable<PartialMarker> { Ok(AggregateHashTable { group_by_metrics: self.group_by_metrics.clone(), + aggregate_argument_metrics: self.aggregate_argument_metrics.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/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index 1c6285d793..0011150970 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -19,6 +19,43 @@ use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; +#[derive(Clone)] +pub(crate) struct AggregateArgumentMetrics { + argument_times: Vec<Time>, +} + +impl AggregateArgumentMetrics { + pub(crate) fn new<T>( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + aggregate_labels: impl IntoIterator<Item = T>, + ) -> Self + where + T: Into<String>, + { + let argument_times = aggregate_labels + .into_iter() + .enumerate() + .map(|(idx, label)| { + MetricBuilder::new(metrics) + .with_new_label("aggregate", label.into()) + .subset_time(format!("agg_expr_{idx}_arguments_time"), partition) + }) + .collect(); + + Self { argument_times } + } + + pub(crate) fn time<R>(&self, index: usize, f: impl FnOnce() -> R) -> R { + debug_assert!( + index < self.argument_times.len(), + "aggregate argument metric index {index} out of range" + ); + let _timer = self.argument_times.get(index).map(Time::timer); + f() + } +} + #[derive(Clone)] pub(crate) struct GroupByMetrics { /// Time spent calculating the group IDs from the evaluated grouping columns. @@ -52,7 +89,7 @@ impl GroupByMetrics { #[cfg(test)] mod tests { use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; - use crate::metrics::MetricsSet; + use crate::metrics::{MetricValue, MetricsSet}; use crate::test::TestMemoryExec; use crate::{ExecutionPlan, collect}; use arrow::array::{Float64Array, UInt32Array}; @@ -63,7 +100,9 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::sum::sum_udaf; - use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::aggregate::{ + AggregateExprBuilder, AggregateFunctionExpr, + }; use datafusion_physical_expr::expressions::col; use std::sync::Arc; @@ -82,6 +121,55 @@ mod tests { assert!(emitting_time.unwrap().as_usize() > 0); } + fn aggregate_argument_metric_names_and_labels( + metrics: &MetricsSet, + ) -> Vec<(String, String)> { + metrics + .iter() + .filter_map(|metric| match metric.value() { + MetricValue::Time { name, .. } + if name.starts_with("agg_expr_") + && name.ends_with("_arguments_time") => + { + let aggregate_label = metric + .labels() + .iter() + .find(|label| label.name() == "aggregate")? + .value() + .to_string(); + Some((name.to_string(), aggregate_label)) + } + _ => None, + }) + .collect() + } + + fn sum_aggregate( + schema: &Arc<Schema>, + column: &str, + alias: &str, + ) -> Result<Arc<AggregateFunctionExpr>> { + Ok(Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col(column, schema)?]) + .schema(Arc::clone(schema)) + .alias(alias) + .build()?, + )) + } + + fn count_aggregate( + schema: &Arc<Schema>, + column: &str, + alias: &str, + ) -> Result<Arc<AggregateFunctionExpr>> { + Ok(Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col(column, schema)?]) + .schema(Arc::clone(schema)) + .alias(alias) + .build()?, + )) + } + #[tokio::test] async fn test_groupby_metrics_partial_mode() -> Result<()> { let schema = Arc::new(Schema::new(vec![ @@ -114,18 +202,8 @@ mod tests { PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); let aggregates = vec![ - Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("SUM(b)") - .build()?, - ), - Arc::new( - AggregateExprBuilder::new(count_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("COUNT(b)") - .build()?, - ), + sum_aggregate(&schema, "b", "SUM(b)")?, + count_aggregate(&schema, "b", "COUNT(b)")?, ]; let aggregate_exec = Arc::new(AggregateExec::try_new( @@ -153,6 +231,83 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_groupby_aggregate_argument_metrics_distinguish_inputs() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::UInt32, false), + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + ])); + + let batches = (0..5) + .map(|i| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 2])), + Arc::new(Float64Array::from(vec![ + i as f64, + (i + 1) as f64, + (i + 2) as f64, + (i + 3) as f64, + ])), + Arc::new(Float64Array::from(vec![ + (i + 4) as f64, + (i + 5) as f64, + (i + 6) as f64, + (i + 7) as f64, + ])), + ], + ) + .unwrap() + }) + .collect::<Vec<_>>(); + + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".to_string())]); + let aggregates = vec![ + sum_aggregate(&schema, "a", "SUM(a)")?, + sum_aggregate(&schema, "b", "SUM(b)")?, + ]; + + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggregates, + vec![None, None], + input, + schema, + )?); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(10 * 1024 * 1024, 1.0) + .build_arc()?; + let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime)); + let _result = + collect(Arc::clone(&aggregate_exec) as _, Arc::clone(&task_ctx)).await?; + + let metrics = aggregate_exec.metrics().unwrap(); + let mut metric_names_and_labels = + aggregate_argument_metric_names_and_labels(&metrics); + metric_names_and_labels.sort(); + assert_eq!( + metric_names_and_labels, + vec![ + ( + "agg_expr_0_arguments_time".to_string(), + "SUM(a)".to_string(), + ), + ( + "agg_expr_1_arguments_time".to_string(), + "SUM(b)".to_string(), + ), + ] + ); + + Ok(()) + } + #[tokio::test] async fn test_groupby_metrics_final_mode() -> Result<()> { let schema = Arc::new(Schema::new(vec![ @@ -183,12 +338,7 @@ mod tests { let group_by = PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); - let aggregates = vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("SUM(b)") - .build()?, - )]; + let aggregates = vec![sum_aggregate(&schema, "b", "SUM(b)")?]; // Create partial aggregate let partial_aggregate = Arc::new(AggregateExec::try_new( diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 1101d53531..bd5b92747e 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -49,7 +49,7 @@ use crate::aggregates::{ mod metrics; mod null_builder; -pub(crate) use metrics::GroupByMetrics; +pub(crate) use metrics::{AggregateArgumentMetrics, GroupByMetrics}; /// 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 99c1011994..340df57e5d 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -24,12 +24,14 @@ use std::vec; use super::order::GroupOrdering; use super::skip_partial::SkipAggregationProbe; use super::{AggregateExec, format_human_display}; -use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::group_values::{ + AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, +}; use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ AggregateInputMode, AggregateMode, AggregateOutputMode, PhysicalGroupBy, - create_schema, evaluate_group_by, evaluate_many, evaluate_optional, group_id_array, - max_duplicate_ordinal, + aggregate_metric_label, create_schema, evaluate_group_by, evaluate_optional, + group_id_array, max_duplicate_ordinal, }; use crate::metrics::{BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput}; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; @@ -52,6 +54,7 @@ use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::{GroupsAccumulatorAdapter, PhysicalSortExpr}; use datafusion_physical_expr_common::sort_expr::LexOrdering; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use crate::sorts::IncrementalSortIterator; use datafusion_common::instant::Instant; @@ -372,6 +375,9 @@ pub(crate) struct GroupedHashAggregateStream { /// Aggregation-specific metrics group_by_metrics: GroupByMetrics, + /// Per-aggregate timing metrics for evaluating aggregate arguments. + aggregate_argument_metrics: AggregateArgumentMetrics, + /// Reduction factor metric, calculated as `output_rows/input_rows` (only for partial aggregation) reduction_factor: Option<metrics::RatioMetrics>, } @@ -392,6 +398,13 @@ impl GroupedHashAggregateStream { let input = agg.input.execute(partition, Arc::clone(context))?; let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + let aggregate_argument_metrics = AggregateArgumentMetrics::new( + &agg.metrics, + partition, + agg.aggr_expr + .iter() + .map(|agg_expr| aggregate_metric_label(agg_expr)), + ); let timer = baseline_metrics.elapsed_compute().timer(); @@ -598,6 +611,7 @@ impl GroupedHashAggregateStream { exec_state, baseline_metrics, group_by_metrics, + aggregate_argument_metrics, batch_size, group_ordering, input_done: false, @@ -857,11 +871,19 @@ impl GroupedHashAggregateStream { }; // Evaluate the aggregation expressions. - let input_values = if self.spill_state.is_stream_merging { - evaluate_many(&self.spill_state.merging_aggregate_arguments, batch)? + let aggregate_arguments = if self.spill_state.is_stream_merging { + &self.spill_state.merging_aggregate_arguments } else { - evaluate_many(&self.aggregate_arguments, batch)? + &self.aggregate_arguments }; + let input_values = aggregate_arguments + .iter() + .enumerate() + .map(|(idx, expr)| { + self.aggregate_argument_metrics + .time(idx, || evaluate_expressions_to_arrays(expr, batch)) + }) + .collect::<Result<Vec<_>>>()?; drop(timer); // Evaluate the filter expressions, if any, against the inputs @@ -1368,7 +1390,17 @@ impl GroupedHashAggregateStream { /// Transforms input batch to intermediate aggregate state, without grouping it fn transform_to_states(&self, batch: &RecordBatch) -> Result<RecordBatch> { let mut group_values = evaluate_group_by(&self.group_by, batch)?; - let input_values = evaluate_many(&self.aggregate_arguments, batch)?; + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + let input_values = self + .aggregate_arguments + .iter() + .enumerate() + .map(|(idx, expr)| { + self.aggregate_argument_metrics + .time(idx, || evaluate_expressions_to_arrays(expr, batch)) + }) + .collect::<Result<Vec<_>>>()?; + drop(timer); let filter_values = evaluate_optional(&self.filter_expressions, batch)?; assert_eq_or_internal_err!( @@ -1398,6 +1430,7 @@ impl GroupedHashAggregateStream { #[cfg(test)] mod tests { use super::*; + use crate::ExecutionPlan; use crate::InputOrderMode; use crate::test::TestMemoryExec; use arrow::array::{Int32Array, Int64Array}; @@ -1439,7 +1472,7 @@ mod tests { ], )?; - let input_partitions = vec![vec![batch]]; + let input_partitions = vec![vec![batch.clone(), batch]]; // Create constrained memory to trigger early emission but not completely fail let runtime = RuntimeEnvBuilder::default() @@ -1508,10 +1541,20 @@ mod tests { } assert_eq!( - total_output_groups, num_groups, + total_output_groups, + num_groups * 2, "Unexpected number of groups", ); + let metrics = aggregate_exec.metrics().unwrap(); + let agg_arguments_time = metrics.sum_by_name("aggregate_arguments_time"); + assert!(agg_arguments_time.is_some()); + assert!(agg_arguments_time.unwrap().as_usize() > 0); + + let per_aggregate_time = metrics.sum_by_name("agg_expr_0_arguments_time"); + assert!(per_aggregate_time.is_some()); + assert!(per_aggregate_time.unwrap().as_usize() > 0); + Ok(()) } diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs index f697e5a394..2df5960188 100644 --- a/datafusion/physical-plan/src/aggregates/hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -42,9 +42,9 @@ use futures::stream::{Stream, StreamExt}; use super::AggregateExec; use super::aggregate_hash_table::{ - AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, + AggregateHashTable, FinalMarker, OrderedAggregateTableMetrics, PartialMarker, + PartialSkipMarker, }; -use super::group_values::GroupByMetrics; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::skip_partial::SkipAggregationProbe; use crate::metrics::{ @@ -382,7 +382,7 @@ impl FinalSpillContext { fn into_replay_stream( self, baseline_metrics: &BaselineMetrics, - group_by_metrics: GroupByMetrics, + metrics: OrderedAggregateTableMetrics, reservation: MemoryReservation, ) -> Result<SendableRecordBatchStream> { let Self { @@ -416,7 +416,7 @@ impl FinalSpillContext { merged, &InputOrderMode::Sorted, baseline_metrics.clone(), - group_by_metrics, + metrics, None, reservation, )?; @@ -1325,12 +1325,12 @@ impl FinalHashAggregateStream { let timer = elapsed_compute.timer(); let replay = match spill_context.spill_table(&mut hash_table) { Ok(()) => { - let group_by_metrics = hash_table.group_by_metrics().clone(); + let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); drop(hash_table); match self.reservation.try_resize(0) { Ok(()) => (*spill_context).into_replay_stream( &self.baseline_metrics, - group_by_metrics, + metrics, self.reservation.new_empty(), ), Err(e) => Err(e), diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index f9dd90f6f9..1671735ec3 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1996,6 +1996,10 @@ fn format_tree_aggregate_expr(agg: &AggregateFunctionExpr) -> Cow<'_, str> { .unwrap_or_else(|| Cow::Borrowed(agg.name())) } +fn aggregate_metric_label(agg: &AggregateFunctionExpr) -> String { + format_tree_aggregate_expr(agg).into_owned() +} + fn format_human_display<'a>( human_display: Option<&'a str>, alias: Option<&'a str>, diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs index 19deedc258..2c26b74da7 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -32,8 +32,9 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::stream::{Stream, StreamExt}; use super::AggregateExec; -use super::aggregate_hash_table::{FinalMarker, OrderedAggregateTable}; -use super::group_values::GroupByMetrics; +use super::aggregate_hash_table::{ + FinalMarker, OrderedAggregateTable, OrderedAggregateTableMetrics, +}; use crate::aggregates::AggregateMode; use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; use crate::sorts::IncrementalSortIterator; @@ -216,7 +217,7 @@ impl OrderedFinalSpillContext { fn into_replay_stream( self, baseline_metrics: &BaselineMetrics, - group_by_metrics: GroupByMetrics, + metrics: OrderedAggregateTableMetrics, reservation: MemoryReservation, ) -> Result<SendableRecordBatchStream> { let Self { @@ -250,7 +251,7 @@ impl OrderedFinalSpillContext { merged, &InputOrderMode::Sorted, baseline_metrics.clone(), - group_by_metrics, + metrics, None, reservation, )?; @@ -282,7 +283,7 @@ impl OrderedFinalAggregateStream { input_order_mode: &InputOrderMode, ) -> Result<Self> { let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); - let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + let metrics = OrderedAggregateTableMetrics::new(agg, partition); let spill_metrics = SpillMetrics::new(&agg.metrics, partition); let reservation = MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) @@ -300,7 +301,7 @@ impl OrderedFinalAggregateStream { input, input_order_mode, baseline_metrics, - group_by_metrics, + metrics, Some(spill_metrics), reservation, ) @@ -320,7 +321,7 @@ impl OrderedFinalAggregateStream { input: SendableRecordBatchStream, input_order_mode: &InputOrderMode, baseline_metrics: BaselineMetrics, - group_by_metrics: GroupByMetrics, + metrics: OrderedAggregateTableMetrics, spill_metrics: Option<SpillMetrics>, reservation: MemoryReservation, ) -> Result<Self> { @@ -359,7 +360,7 @@ impl OrderedFinalAggregateStream { Arc::clone(&schema), batch_size, input_order_mode, - group_by_metrics, + metrics, )?; Ok(Self { schema, @@ -654,12 +655,12 @@ impl OrderedFinalAggregateStream { let timer = elapsed_compute.timer(); let replay = match spill_context.spill_table(&mut table) { Ok(()) => { - let group_by_metrics = table.group_by_metrics(); + let metrics = table.metrics(); drop(table); match self.reservation.try_resize(0) { Ok(()) => (*spill_context).into_replay_stream( &self.baseline_metrics, - group_by_metrics, + metrics, self.reservation.new_empty(), ), Err(e) => Err(e), diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs index 2025bdce30..da00b42e5c 100644 --- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs @@ -31,8 +31,9 @@ use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::stream::{Stream, StreamExt}; -use super::aggregate_hash_table::{OrderedAggregateTable, SingleMarker}; -use super::group_values::GroupByMetrics; +use super::aggregate_hash_table::{ + OrderedAggregateTable, OrderedAggregateTableMetrics, SingleMarker, +}; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; @@ -276,7 +277,7 @@ impl OrderedSingleSpillContext { fn into_replay_stream( self, baseline_metrics: &BaselineMetrics, - group_by_metrics: GroupByMetrics, + metrics: OrderedAggregateTableMetrics, reservation: MemoryReservation, ) -> Result<SendableRecordBatchStream> { let Self { @@ -310,7 +311,7 @@ impl OrderedSingleSpillContext { merged, &InputOrderMode::Sorted, baseline_metrics.clone(), - group_by_metrics, + metrics, None, reservation, )?; @@ -635,12 +636,12 @@ impl OrderedSingleAggregateStream { let timer = elapsed_compute.timer(); let replay = match spill_context.spill_table(&mut table) { Ok(()) => { - let group_by_metrics = table.group_by_metrics(); + let metrics = table.metrics(); drop(table); match self.reservation.try_resize(0) { Ok(()) => (*spill_context).into_replay_stream( &self.baseline_metrics, - group_by_metrics, + metrics, self.reservation.new_empty(), ), Err(e) => Err(e), diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs index c6f25dc2cf..3e306d72a7 100644 --- a/datafusion/physical-plan/src/aggregates/single_stream.rs +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -36,8 +36,9 @@ use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::stream::{Stream, StreamExt}; -use super::aggregate_hash_table::{AggregateHashTable, SingleMarker}; -use super::group_values::GroupByMetrics; +use super::aggregate_hash_table::{ + AggregateHashTable, OrderedAggregateTableMetrics, SingleMarker, +}; use super::ordered_final_stream::OrderedFinalAggregateStream; use super::{AggregateExec, create_schema}; use crate::aggregates::AggregateMode; @@ -267,7 +268,7 @@ impl SingleSpillContext { fn into_replay_stream( self, baseline_metrics: &BaselineMetrics, - group_by_metrics: GroupByMetrics, + metrics: OrderedAggregateTableMetrics, reservation: MemoryReservation, ) -> Result<SendableRecordBatchStream> { let Self { @@ -301,7 +302,7 @@ impl SingleSpillContext { merged, &InputOrderMode::Sorted, baseline_metrics.clone(), - group_by_metrics, + metrics, None, reservation, )?; @@ -591,12 +592,12 @@ impl SingleHashAggregateStream { let timer = elapsed_compute.timer(); let replay = match spill_context.spill_table(&mut hash_table) { Ok(()) => { - let group_by_metrics = hash_table.group_by_metrics().clone(); + let metrics = OrderedAggregateTableMetrics::from_hash_table(&hash_table); drop(hash_table); match self.reservation.try_resize(0) { Ok(()) => (*spill_context).into_replay_stream( &self.baseline_metrics, - group_by_metrics, + metrics, self.reservation.new_empty(), ), Err(e) => Err(e), --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
