This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new aff5850c0a Add per-aggregate argument evaluation metrics for grouped
hash aggregation (#24024)
aff5850c0a is described below
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]
