You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/08/13 02:06:04 UTC
(datafusion) branch gh-readonly-queue/main/pr-24024-bbb552f71c93cdfadf91b374a3c4ff5f1a931ae7 created (now 21e7a437a1)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24024-bbb552f71c93cdfadf91b374a3c4ff5f1a931ae7
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 21e7a437a1 Add per-aggregate argument evaluation metrics for grouped hash aggregation (#24024)
This branch includes the following new commits:
new 21e7a437a1 Add per-aggregate argument evaluation metrics for grouped hash aggregation (#24024)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Add per-aggregate argument evaluation metrics for grouped hash aggregation (#24024)
Posted by gi...@apache.org.
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-bbb552f71c93cdfadf91b374a3c4ff5f1a931ae7
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 21e7a437a1c7d9009711494e276a5f5be9f591fc
Author: kosiew <ko...@gmail.com>
AuthorDate: Thu Aug 13 02:05:42 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 | 25 ++-
.../aggregate_hash_table/common_ordered.rs | 49 +++++-
.../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/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 ++++++-
datafusion/physical-plan/src/aggregates/mod.rs | 4 +
.../src/aggregates/ordered_final_stream.rs | 21 +--
11 files changed, 310 insertions(+), 59 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..eec454bd70 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);
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 2293e7b1b8..959c36d345 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,16 +30,40 @@ 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, EvaluatedAggregateBatch};
+#[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,
+ }
+ }
+}
+
/// Aggregate table shared by the ordered partial and final paths.
///
/// # Ordering optimization
@@ -91,6 +115,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,
@@ -140,7 +167,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,
@@ -175,7 +202,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,
@@ -204,7 +232,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);
@@ -250,8 +282,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 2c7ec01654..7c51a9399a 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs
@@ -28,4 +28,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 fd064ebffe..698f502103 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,10 +27,9 @@ 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_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.
@@ -48,7 +47,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,
@@ -59,7 +58,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 a04e4dda8f..66621603e1 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,10 +37,9 @@ use datafusion_common::Result;
use crate::aggregates::{
AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker,
- group_values::GroupByMetrics,
};
-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.
@@ -60,7 +59,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,
@@ -70,7 +69,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/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/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index a39c6f3486..b1d2ff56dd 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -1973,6 +1973,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),
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org