diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index a1ff2b7b6e9f2..21e19027cebcf 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -426,7 +426,58 @@ mod tests { } #[tokio::test] - async fn test_groupby_metrics_final_mode() -> Result<()> { + async fn test_legacy_groupby_aggregate_accumulator_metrics() -> 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 batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 2])), + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])), + Arc::new(Float64Array::from(vec![5.0, 6.0, 7.0, 8.0])), + ], + )?; + let input = + TestMemoryExec::try_new_exec(&[vec![batch]], 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 task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", false), + ), + ); + let _result = + collect(Arc::clone(&aggregate_exec) as _, Arc::clone(&task_ctx)).await?; + + let metrics = aggregate_exec.metrics().unwrap(); + assert_aggregate_metric_labels(&metrics, "arguments_time"); + assert_aggregate_metric_labels(&metrics, "update_time"); + assert_aggregate_metric_labels(&metrics, "state_time"); + assert_aggregate_metric_times_positive(&metrics, "update_time"); + assert_aggregate_metric_times_positive(&metrics, "state_time"); + + Ok(()) + } + + async fn assert_groupby_metrics_final_mode( + enable_migration_aggregate: bool, + ) -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), Field::new("b", DataType::Float64, false), @@ -486,12 +537,12 @@ mod tests { schema, )?); - let task_ctx = Arc::new( - TaskContext::default().with_session_config( - SessionConfig::new() - .set_bool("datafusion.execution.enable_migration_aggregate", true), + let task_ctx = Arc::new(TaskContext::default().with_session_config( + SessionConfig::new().set_bool( + "datafusion.execution.enable_migration_aggregate", + enable_migration_aggregate, ), - ); + )); let _result = collect(Arc::clone(&final_aggregate) as _, Arc::clone(&task_ctx)).await?; @@ -516,4 +567,13 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn test_groupby_metrics_final_mode() -> Result<()> { + for enable_migration_aggregate in [true, false] { + assert_groupby_metrics_final_mode(enable_migration_aggregate).await?; + } + + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index c0253093c8a7b..cb5f8ade61c0d 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -21,11 +21,13 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::vec; +use super::aggregate_hash_table::accumulator_phases; use super::order::GroupOrdering; use super::skip_partial::SkipAggregationProbe; use super::{AggregateExec, format_human_display}; use crate::aggregates::group_values::{ - AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values, + AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, + GroupByMetrics, GroupValues, new_group_values, }; use crate::aggregates::order::GroupOrderingFull; use crate::aggregates::{ @@ -378,6 +380,9 @@ pub(crate) struct GroupedHashAggregateStream { /// Per-aggregate timing metrics for evaluating aggregate arguments. aggregate_argument_metrics: AggregateArgumentMetrics, + /// Per-aggregate timing metrics for accumulator phases. + aggregate_accumulator_metrics: AggregateAccumulatorMetrics, + /// Reduction factor metric, calculated as `output_rows/input_rows` (only for partial aggregation) reduction_factor: Option, } @@ -398,12 +403,21 @@ 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_labels = agg + .aggr_expr + .iter() + .map(|agg_expr| aggregate_metric_label(agg_expr)) + .collect::>(); let aggregate_argument_metrics = AggregateArgumentMetrics::new( &agg.metrics, partition, - agg.aggr_expr - .iter() - .map(|agg_expr| aggregate_metric_label(agg_expr)), + aggregate_labels.iter().cloned(), + ); + let aggregate_accumulator_metrics = AggregateAccumulatorMetrics::new( + &agg.metrics, + partition, + aggregate_labels, + accumulator_phases(&agg.mode), ); let timer = baseline_metrics.elapsed_compute().timer(); @@ -615,6 +629,7 @@ impl GroupedHashAggregateStream { baseline_metrics, group_by_metrics, aggregate_argument_metrics, + aggregate_accumulator_metrics, batch_size, group_ordering, input_done: false, @@ -929,7 +944,7 @@ impl GroupedHashAggregateStream { .zip(input_values.iter()) .zip(filter_values.iter()); - for ((acc, values), opt_filter) in t { + for (idx, ((acc, values), opt_filter)) in t.enumerate() { let opt_filter = opt_filter.as_ref().map(|filter| filter.as_boolean()); // Call the appropriate method on each aggregator with @@ -937,11 +952,17 @@ impl GroupedHashAggregateStream { if self.mode.input_mode() == AggregateInputMode::Raw && !self.spill_state.is_stream_merging { - acc.update_batch( - values, - group_indices, - opt_filter, - total_num_groups, + self.aggregate_accumulator_metrics.time( + idx, + AccumulatorPhase::Update, + || { + acc.update_batch( + values, + group_indices, + opt_filter, + total_num_groups, + ) + }, )?; } else { assert_or_internal_err!( @@ -951,7 +972,11 @@ impl GroupedHashAggregateStream { // if aggregation is over intermediate states, // use merge - acc.merge_batch(values, group_indices, total_num_groups)?; + self.aggregate_accumulator_metrics.time( + idx, + AccumulatorPhase::Merge, + || acc.merge_batch(values, group_indices, total_num_groups), + )?; } self.group_by_metrics .aggregation_time @@ -1058,13 +1083,21 @@ impl GroupedHashAggregateStream { } // Next output each aggregate value - for acc in self.accumulators.iter_mut() { + for (idx, acc) in self.accumulators.iter_mut().enumerate() { if self.mode.output_mode() == AggregateOutputMode::Final && !spilling { - output.push(acc.evaluate(emit_to)?) + output.push(self.aggregate_accumulator_metrics.time( + idx, + AccumulatorPhase::Evaluate, + || acc.evaluate(emit_to), + )?) } else { // Output partial state: either because we're in a non-final mode, // or because we're spilling and will merge/re-evaluate later. - output.extend(acc.state(emit_to)?) + output.extend(self.aggregate_accumulator_metrics.time( + idx, + AccumulatorPhase::State, + || acc.state(emit_to), + )?) } } drop(timer); @@ -1178,8 +1211,17 @@ impl GroupedHashAggregateStream { }) .collect::>>()?; let false_filter = BooleanArray::from(vec![false]); - for (acc, args) in self.accumulators.iter_mut().zip(null_args.iter()) { - acc.update_batch(args, &[0], Some(&false_filter), total_groups)?; + for (idx, (acc, args)) in self + .accumulators + .iter_mut() + .zip(null_args.iter()) + .enumerate() + { + self.aggregate_accumulator_metrics.time( + idx, + AccumulatorPhase::Update, + || acc.update_batch(args, &[0], Some(&false_filter), total_groups), + )?; } } @@ -1419,9 +1461,13 @@ impl GroupedHashAggregateStream { .zip(input_values.iter()) .zip(filter_values.iter()); - for ((acc, values), opt_filter) in iter { + for (idx, ((acc, values), opt_filter)) in iter.enumerate() { let opt_filter = opt_filter.as_ref().map(|filter| filter.as_boolean()); - output.extend(acc.convert_to_state(values, opt_filter)?); + output.extend(self.aggregate_accumulator_metrics.time( + idx, + AccumulatorPhase::ConvertToState, + || acc.convert_to_state(values, opt_filter), + )?); } let states_batch = RecordBatch::try_new(self.schema(), output)?; diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 193fdba4b0198..a5b274cffe2ee 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -17,13 +17,13 @@ //! A memory-conscious aggregation implementation that limits group buckets to a fixed number -use crate::aggregates::group_values::GroupByMetrics; +use crate::aggregates::group_values::{AggregateArgumentMetrics, GroupByMetrics}; use crate::aggregates::topk::priority_map::PriorityMap; #[cfg(debug_assertions)] use crate::aggregates::topk_types_supported; use crate::aggregates::{ - AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, - evaluate_many, + AggregateExec, PhysicalGroupBy, aggregate_expressions, aggregate_metric_label, + evaluate_group_by, }; use crate::metrics::BaselineMetrics; use crate::stream::EmptyRecordBatchStream; @@ -37,6 +37,7 @@ use datafusion_common::internal_datafusion_err; use datafusion_execution::TaskContext; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr_common::metrics::RecordOutput; +use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::stream::{Stream, StreamExt}; use log::{Level, trace}; use std::pin::Pin; @@ -52,6 +53,9 @@ pub struct GroupedTopKAggregateStream { input: SendableRecordBatchStream, baseline_metrics: BaselineMetrics, group_by_metrics: GroupByMetrics, + // TopK directly maintains MIN/MAX values in its priority map, so it has no + // accumulator update, merge, state, or evaluate phases to time. + aggregate_argument_metrics: AggregateArgumentMetrics, aggregate_arguments: Vec>>, group_by: Arc, priority_map: PriorityMap, @@ -71,6 +75,13 @@ impl GroupedTopKAggregateStream { let input = aggr.input.execute(partition, Arc::clone(context))?; let baseline_metrics = BaselineMetrics::new(&aggr.metrics, partition); let group_by_metrics = GroupByMetrics::new(&aggr.metrics, partition); + let aggregate_argument_metrics = AggregateArgumentMetrics::new( + &aggr.metrics, + partition, + aggr.aggr_expr + .iter() + .map(|agg_expr| aggregate_metric_label(agg_expr)), + ); let aggregate_arguments = aggregate_expressions(&aggr.aggr_expr, &aggr.mode, group_by.expr.len())?; @@ -119,6 +130,7 @@ impl GroupedTopKAggregateStream { input, baseline_metrics, group_by_metrics, + aggregate_argument_metrics, aggregate_arguments, group_by, priority_map, @@ -235,9 +247,7 @@ impl Stream for GroupedTopKAggregateStream { print_batches(std::slice::from_ref(&batch))?; } self.row_count += batch.num_rows(); - let batches = &[batch]; - let group_by_values = - evaluate_group_by(&self.group_by, batches.first().unwrap())?; + let group_by_values = evaluate_group_by(&self.group_by, &batch)?; assert_eq!( group_by_values.len(), 1, @@ -256,10 +266,16 @@ impl Stream for GroupedTopKAggregateStream { // MIN/MAX case: evaluate aggregate expressions let _timer = self.group_by_metrics.aggregate_arguments_time.timer(); - let input_values = evaluate_many( - &self.aggregate_arguments, - batches.first().unwrap(), - )?; + let input_values = self + .aggregate_arguments + .iter() + .enumerate() + .map(|(idx, expr)| { + self.aggregate_argument_metrics.time(idx, || { + evaluate_expressions_to_arrays(expr, &batch) + }) + }) + .collect::>>()?; assert_eq!(input_values.len(), 1, "Exactly 1 input required"); assert_eq!(input_values[0].len(), 1, "Exactly 1 input required"); Arc::clone(&input_values[0][0]) @@ -304,3 +320,84 @@ impl Stream for GroupedTopKAggregateStream { Poll::Pending } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::ExecutionPlan; + use crate::aggregates::{AggregateMode, LimitOptions}; + use crate::collect; + use crate::metrics::MetricValue; + use crate::test::TestMemoryExec; + use arrow::array::{Float64Array, UInt32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion_common::assert_batches_eq; + use datafusion_functions_aggregate::min_max::min_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + + #[tokio::test] + async fn test_topk_aggregate_argument_metrics() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("k", DataType::UInt32, false), + Field::new("a", DataType::Float64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 3, 4])), + Arc::new(Float64Array::from(vec![4.0, 3.0, 2.0, 1.0])), + ], + )?; + let input = + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("k", &schema)?, "k".to_string())]); + let aggregate = Arc::new( + AggregateExprBuilder::new(min_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("MIN(a)") + .build()?, + ); + let aggregate_exec = Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + group_by, + vec![aggregate], + vec![None], + input, + schema, + )? + .with_limit_options(Some(LimitOptions::new(2))), + ); + let context = Arc::new(TaskContext::default()); + let result = collect(Arc::clone(&aggregate_exec) as _, context).await?; + assert_batches_eq!( + [ + "+---+--------+", + "| k | MIN(a) |", + "+---+--------+", + "| 4 | 1.0 |", + "| 3 | 2.0 |", + "+---+--------+", + ], + &result + ); + + let metrics = aggregate_exec.metrics().unwrap(); + let argument_metric = metrics.iter().find(|metric| { + matches!( + metric.value(), + MetricValue::Time { name, .. } if name == "agg_expr_0_arguments_time" + ) && metric + .labels() + .iter() + .any(|label| label.name() == "aggregate" && label.value() == "MIN(a)") + }); + assert!(argument_metric.is_some()); + assert!(argument_metric.unwrap().value().as_usize() > 0); + + Ok(()) + } +}