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 cd8cef983d Add per-aggregate accumulator phase metrics (#24423)
cd8cef983d is described below
commit cd8cef983d3f4e9c5f52585de792d67d59d51f48
Author: kosiew <[email protected]>
AuthorDate: Thu Aug 20 12:34:30 2026 +0000
Add per-aggregate accumulator phase metrics (#24423)
## Which issue does this PR close?
* Part of #23570
## Rationale for this change
Aggregate argument evaluation metrics only measure the cost of producing
input arrays and do not show how much time individual aggregate
expressions spend performing accumulator work.
This PR adds per-aggregate-expression timing for accumulator phases in
migrated grouped hash aggregation, making it possible to attribute time
spent updating, merging, emitting state, converting input to state, and
evaluating accumulators to the aggregate expression performing that
work. Existing shared group-by aggregation and emission timing remains
in place.
## What changes are included in this PR?
* Adds `AggregateAccumulatorMetrics`, which creates per-aggregate timers
for accumulator phases.
* Adds phase-specific metrics for `update`, `merge`, `state`,
`convert_to_state`, and `evaluate`.
* Creates accumulator metrics once per aggregate table and shares them
with the migrated hash table implementations.
* Instruments accumulator operations at the per-aggregate/per-batch
boundary rather than per row or group.
* Adds instrumentation to both unordered and ordered migrated aggregate
tables.
* Defines the accumulator phases expected for each `AggregateMode`.
* Preserves the existing `GroupByMetrics` and aggregate argument
metrics.
* Instruments the skip-partial path's `convert_to_state` operation.
* Carries the accumulator metrics through partial-to-skip table
transitions and spill/replay paths.
## Are these changes tested?
Yes.
The patch adds or updates tests that:
* Verify the accumulator phases associated with each `AggregateMode`.
* Verify two aggregate expressions (`SUM(a)` / `SUM(b)` and `SUM(b)` /
`SUM(c)` in the respective tests) receive separate metric names and
`aggregate` labels.
* Verify partial aggregation records positive `update_time` and
`state_time` metrics.
* Verify final aggregation records positive `merge_time` and
`evaluate_time` metrics.
* Verify migrated ordered aggregation paths record the expected
accumulator phase timings.
* Verify skip-partial aggregation records `convert_to_state` timing.
* Verify a migrated single aggregate that spills records `update`,
`state`, `merge`, and `evaluate` timings.
* Explicitly enable `datafusion.execution.enable_migration_aggregate` in
the relevant tests.
The issue suggests the targeted verification command:
```bash
cargo test -p datafusion-physical-plan aggregates
```
The patch itself does not show that command being run, so this PR
description does not claim that it passed.
## Are there any user-facing changes?
Yes. Execution metrics for migrated grouped hash aggregation now expose
per-aggregate-expression accumulator timing, with metrics such as:
```text
agg_expr_0_update_time
agg_expr_0_merge_time
agg_expr_0_state_time
agg_expr_0_convert_to_state_time
agg_expr_0_evaluate_time
```
Each metric also carries the corresponding `aggregate` label.
There are no public accumulator API or aggregate semantics changes shown
in this patch.
## 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 | 52 ++++--
.../aggregate_hash_table/common_ordered.rs | 57 ++++--
.../aggregates/aggregate_hash_table/final_table.rs | 12 +-
.../src/aggregates/aggregate_hash_table/mod.rs | 128 ++++++++++++++
.../aggregate_hash_table/ordered_final_table.rs | 8 +-
.../aggregate_hash_table/ordered_partial_table.rs | 7 +-
.../aggregate_hash_table/ordered_single_table.rs | 8 +-
.../aggregate_hash_table/partial_reduce_table.rs | 12 +-
.../aggregate_hash_table/partial_table.rs | 33 +++-
.../aggregate_hash_table/single_table.rs | 12 +-
.../src/aggregates/group_values/metrics.rs | 195 ++++++++++++++++++---
.../src/aggregates/group_values/mod.rs | 5 +-
datafusion/physical-plan/src/aggregates/mod.rs | 38 +++-
13 files changed, 487 insertions(+), 80 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 fde7024f7b..f7c12b1931 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
@@ -28,15 +28,17 @@ use
datafusion_physical_expr::aggregate::AggregateFunctionExpr;
use crate::PhysicalExpr;
use crate::aggregates::group_values::{
- AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values,
+ AccumulatorPhase, AggregateAccumulatorMetrics, 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,
aggregate_metric_label,
- evaluate_group_by,
+ AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by,
};
+use super::AggregateTableMetrics;
+
/// Marker for raw rows -> partial state aggregation.
pub(in crate::aggregates) struct PartialMarker;
/// Marker for raw rows -> final value aggregation.
@@ -81,6 +83,9 @@ pub(in crate::aggregates) struct AggregateHashTable<AggrMode>
{
/// Per-aggregate timing metrics for evaluating aggregate arguments.
pub(super) aggregate_argument_metrics: AggregateArgumentMetrics,
+ /// Per-aggregate timing metrics for accumulator operations.
+ pub(super) aggregate_accumulator_metrics: Arc<AggregateAccumulatorMetrics>,
+
/// Raw input schema, used to evaluate expressions and synthesize empty
/// grouping-set rows.
pub(super) input_schema: SchemaRef,
@@ -140,17 +145,12 @@ 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)),
- );
+ let metrics = AggregateTableMetrics::new(agg, partition);
Ok(Self {
- group_by_metrics: GroupByMetrics::new(&agg.metrics, partition),
- aggregate_argument_metrics,
+ group_by_metrics: metrics.group_by,
+ aggregate_argument_metrics: metrics.aggregate_arguments,
+ aggregate_accumulator_metrics: metrics.accumulator,
input_schema,
output_schema,
state_schema,
@@ -208,8 +208,10 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
&mut self,
batch: &RecordBatch,
aggregate_fn: AggregateBatchFn,
+ accumulator_phase: AccumulatorPhase,
) -> Result<()> {
let evaluated_batch = self.evaluate_batch(batch)?;
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
let state = self.state.building_mut();
let _timer = self.group_by_metrics.aggregation_time.timer();
@@ -220,12 +222,15 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
let group_indices = &state.batch_group_indices;
let total_num_groups = state.group_values.len();
- for (acc, values) in state
+ for (idx, (acc, values)) in state
.accumulators
.iter_mut()
.zip(evaluated_batch.accumulator_args.iter())
+ .enumerate()
{
- aggregate_fn(acc, values, group_indices, total_num_groups)?;
+ accumulator_metrics.time(idx, accumulator_phase, || {
+ aggregate_fn(acc, values, group_indices, total_num_groups)
+ })?;
}
}
@@ -244,9 +249,11 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
pub(super) fn next_output_batch_inner(
&mut self,
materialize_accumulator_fn: MaterializeAccumulatorFn,
+ accumulator_phase: AccumulatorPhase,
) -> Result<Option<RecordBatch>> {
let output_schema = Arc::clone(&self.output_schema);
let batch_size = self.batch_size;
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
let mut output =
match std::mem::replace(&mut self.state,
AggregateHashTableState::Done) {
@@ -260,8 +267,12 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
let emit_to = EmitTo::All;
let timer = self.group_by_metrics.emitting_time.timer();
let mut columns = state.group_values.emit(emit_to)?;
- for acc in state.accumulators.iter_mut() {
- columns.extend(materialize_accumulator_fn(acc,
emit_to)?);
+ for (idx, acc) in
state.accumulators.iter_mut().enumerate() {
+ columns.extend(accumulator_metrics.time(
+ idx,
+ accumulator_phase,
+ || materialize_accumulator_fn(acc, emit_to),
+ )?);
}
drop(timer);
@@ -322,14 +333,19 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
&mut self,
) -> Result<Option<RecordBatch>> {
let state_schema = Arc::clone(&self.state_schema);
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
let state = self.state.building_mut();
if state.group_values.is_empty() {
return Ok(None);
}
let mut output = state.group_values.emit(EmitTo::All)?;
- for acc in &mut state.accumulators {
- output.extend(acc.state(EmitTo::All)?);
+ for (idx, acc) in state.accumulators.iter_mut().enumerate() {
+ output.extend(accumulator_metrics.time(
+ idx,
+ AccumulatorPhase::State,
+ || acc.state(EmitTo::All),
+ )?);
}
let batch = RecordBatch::try_new(state_schema, output)?;
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 22fe8f5433..4ced967a09 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
@@ -31,15 +31,17 @@ use datafusion_expr::EmitTo;
use crate::InputOrderMode;
use crate::PhysicalExpr;
use crate::aggregates::group_values::{
- AggregateArgumentMetrics, GroupByMetrics, GroupValues, new_group_values,
+ AccumulatorPhase, AggregateAccumulatorMetrics, 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,
- aggregate_metric_label, evaluate_group_by,
+ evaluate_group_by,
};
+use super::AggregateTableMetrics;
use super::common::{
AggregateAccumulator, AggregateBatchFn, AggregateHashTable,
EvaluatedAggregateBatch,
MaterializeAccumulatorFn,
@@ -49,20 +51,16 @@ use super::common::{
pub(in crate::aggregates) struct OrderedAggregateTableMetrics {
pub(super) group_by: GroupByMetrics,
pub(super) aggregate_arguments: AggregateArgumentMetrics,
+ pub(super) accumulator: Arc<AggregateAccumulatorMetrics>,
}
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)),
- );
+ let metrics = AggregateTableMetrics::new(agg, partition);
Self {
- group_by: GroupByMetrics::new(&agg.metrics, partition),
- aggregate_arguments,
+ group_by: metrics.group_by,
+ aggregate_arguments: metrics.aggregate_arguments,
+ accumulator: metrics.accumulator,
}
}
@@ -72,6 +70,7 @@ impl OrderedAggregateTableMetrics {
Self {
group_by: table.group_by_metrics.clone(),
aggregate_arguments: table.aggregate_argument_metrics.clone(),
+ accumulator: Arc::clone(&table.aggregate_accumulator_metrics),
}
}
}
@@ -136,6 +135,9 @@ pub(in crate::aggregates) struct
OrderedAggregateTable<OrderedAggrMode> {
/// Per-aggregate timing metrics for evaluating aggregate arguments.
pub(super) aggregate_argument_metrics: AggregateArgumentMetrics,
+ /// Per-aggregate timing metrics for accumulator operations.
+ pub(super) aggregate_accumulator_metrics: Arc<AggregateAccumulatorMetrics>,
+
/// Group keys, ordering state, and accumulator states.
pub(super) buffer: OrderedAggregateTableBuffer,
@@ -222,6 +224,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
batch_size,
group_by_metrics: metrics.group_by,
aggregate_argument_metrics: metrics.aggregate_arguments,
+ aggregate_accumulator_metrics: metrics.accumulator,
buffer: OrderedAggregateTableBuffer {
group_by: Arc::clone(&agg.group_by),
group_ordering,
@@ -304,6 +307,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
OrderedAggregateTableMetrics {
group_by: self.group_by_metrics.clone(),
aggregate_arguments: self.aggregate_argument_metrics.clone(),
+ accumulator: Arc::clone(&self.aggregate_accumulator_metrics),
}
}
@@ -321,9 +325,14 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
return Ok(None);
}
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
let mut output = self.buffer.group_values.emit(EmitTo::All)?;
- for acc in &mut self.buffer.accumulators {
- output.extend(acc.state(EmitTo::All)?);
+ for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() {
+ output.extend(accumulator_metrics.time(
+ idx,
+ AccumulatorPhase::State,
+ || acc.state(EmitTo::All),
+ )?);
}
let batch = RecordBatch::try_new(Arc::clone(&self.state_schema),
output)?;
@@ -369,7 +378,9 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
&mut self,
evaluated_batch: &EvaluatedAggregateBatch,
aggregate_fn: AggregateBatchFn,
+ accumulator_phase: AccumulatorPhase,
) -> Result<()> {
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
for group_values in &evaluated_batch.grouping_set_args {
let starting_num_groups = self.buffer.group_values.len();
self.buffer
@@ -385,13 +396,21 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
}
let timer = self.group_by_metrics.aggregation_time.timer();
- for (acc, values) in self
+ for (idx, (acc, values)) in self
.buffer
.accumulators
.iter_mut()
.zip(evaluated_batch.accumulator_args.iter())
+ .enumerate()
{
- aggregate_fn(acc, values, &self.buffer.group_indices,
total_num_groups)?;
+ accumulator_metrics.time(idx, accumulator_phase, || {
+ aggregate_fn(
+ acc,
+ values,
+ &self.buffer.group_indices,
+ total_num_groups,
+ )
+ })?;
}
drop(timer);
}
@@ -409,6 +428,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
pub(super) fn next_output_batch_inner(
&mut self,
materialize_accumulator_fn: MaterializeAccumulatorFn,
+ accumulator_phase: AccumulatorPhase,
) -> Result<Option<RecordBatch>> {
if self.buffer.group_values.is_empty() {
return Ok(None);
@@ -420,6 +440,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
let (emit_to, should_remove_groups) =
self.clamp_emit_to(self.buffer.group_values.len(), emit_to);
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
let timer = self.group_by_metrics.emitting_time.timer();
let mut output = self.buffer.group_values.emit(emit_to)?;
if should_remove_groups {
@@ -432,8 +453,10 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
}
}
- for acc in &mut self.buffer.accumulators {
- output.extend(materialize_accumulator_fn(acc, emit_to)?);
+ for (idx, acc) in self.buffer.accumulators.iter_mut().enumerate() {
+ output.extend(accumulator_metrics.time(idx, accumulator_phase, || {
+ materialize_accumulator_fn(acc, emit_to)
+ })?);
}
drop(timer);
diff --git
a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
index b80e15d7f8..307215cfd2 100644
---
a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
+++
b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
@@ -22,6 +22,7 @@ use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use crate::aggregates::AggregateExec;
+use crate::aggregates::group_values::AccumulatorPhase;
use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator};
@@ -58,7 +59,10 @@ impl AggregateHashTable<FinalMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
-
self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
+ self.next_output_batch_inner(
+ HashAggregateAccumulator::evaluate_to_columns,
+ AccumulatorPhase::Evaluate,
+ )
}
/// Final aggregation consumes partial aggregate states and merges them
into
@@ -67,7 +71,11 @@ impl AggregateHashTable<FinalMarker> {
&mut self,
batch: &RecordBatch,
) -> Result<()> {
- self.aggregate_batch_inner(batch,
HashAggregateAccumulator::merge_batch)
+ self.aggregate_batch_inner(
+ batch,
+ HashAggregateAccumulator::merge_batch,
+ AccumulatorPhase::Merge,
+ )
}
pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> {
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 435289aa30..902a859bac 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs
@@ -25,8 +25,136 @@ mod partial_reduce_table;
mod partial_table;
mod single_table;
+use std::sync::Arc;
+
+use crate::aggregates::group_values::{
+ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics,
+ GroupByMetrics,
+};
+use crate::aggregates::{AggregateExec, AggregateMode, aggregate_metric_label};
+
+pub(super) fn accumulator_phases(mode: &AggregateMode) -> &'static
[AccumulatorPhase] {
+ match mode {
+ AggregateMode::Partial => &[
+ AccumulatorPhase::Update,
+ AccumulatorPhase::State,
+ AccumulatorPhase::ConvertToState,
+ ],
+ AggregateMode::PartialReduce => {
+ &[AccumulatorPhase::Merge, AccumulatorPhase::State]
+ }
+ // Final and single aggregation emit intermediate states when spilling,
+ // then replay them through a final aggregate table.
+ AggregateMode::Final | AggregateMode::FinalPartitioned => &[
+ AccumulatorPhase::Merge,
+ AccumulatorPhase::State,
+ AccumulatorPhase::Evaluate,
+ ],
+ AggregateMode::Single | AggregateMode::SinglePartitioned => &[
+ AccumulatorPhase::Update,
+ AccumulatorPhase::State,
+ AccumulatorPhase::Merge,
+ AccumulatorPhase::Evaluate,
+ ],
+ }
+}
+
+pub(super) struct AggregateTableMetrics {
+ pub(super) group_by: GroupByMetrics,
+ pub(super) aggregate_arguments: AggregateArgumentMetrics,
+ pub(super) accumulator: Arc<AggregateAccumulatorMetrics>,
+}
+
+impl AggregateTableMetrics {
+ pub(super) fn new(agg: &AggregateExec, partition: usize) -> Self {
+ let aggregate_labels = agg
+ .aggr_expr
+ .iter()
+ .map(|agg_expr| aggregate_metric_label(agg_expr))
+ .collect::<Vec<_>>();
+
+ Self {
+ group_by: GroupByMetrics::new(&agg.metrics, partition),
+ aggregate_arguments: AggregateArgumentMetrics::new(
+ &agg.metrics,
+ partition,
+ aggregate_labels.clone(),
+ ),
+ accumulator: Arc::new(AggregateAccumulatorMetrics::new(
+ &agg.metrics,
+ partition,
+ aggregate_labels,
+ accumulator_phases(&agg.mode),
+ )),
+ }
+ }
+}
+
pub(super) use common::{
AggregateHashTable, FinalMarker, PartialMarker, PartialReduceMarker,
PartialSkipMarker, SingleMarker,
};
pub(super) use common_ordered::{OrderedAggregateTable,
OrderedAggregateTableMetrics};
+
+#[cfg(test)]
+mod tests {
+ use super::accumulator_phases;
+ use crate::aggregates::AggregateMode;
+ use crate::aggregates::group_values::AccumulatorPhase;
+
+ #[test]
+ fn accumulator_phases_match_aggregate_mode() {
+ let cases: [(AggregateMode, &[AccumulatorPhase]); 6] = [
+ (
+ AggregateMode::Partial,
+ &[
+ AccumulatorPhase::Update,
+ AccumulatorPhase::State,
+ AccumulatorPhase::ConvertToState,
+ ],
+ ),
+ (
+ AggregateMode::PartialReduce,
+ &[AccumulatorPhase::Merge, AccumulatorPhase::State],
+ ),
+ (
+ AggregateMode::Final,
+ &[
+ AccumulatorPhase::Merge,
+ AccumulatorPhase::State,
+ AccumulatorPhase::Evaluate,
+ ],
+ ),
+ (
+ AggregateMode::FinalPartitioned,
+ &[
+ AccumulatorPhase::Merge,
+ AccumulatorPhase::State,
+ AccumulatorPhase::Evaluate,
+ ],
+ ),
+ (
+ AggregateMode::Single,
+ &[
+ AccumulatorPhase::Update,
+ AccumulatorPhase::State,
+ AccumulatorPhase::Merge,
+ AccumulatorPhase::Evaluate,
+ ],
+ ),
+ (
+ AggregateMode::SinglePartitioned,
+ &[
+ AccumulatorPhase::Update,
+ AccumulatorPhase::State,
+ AccumulatorPhase::Merge,
+ AccumulatorPhase::Evaluate,
+ ],
+ ),
+ ];
+
+ for (mode, expected) in cases {
+ assert!(accumulator_phases(&mode) == expected, "{mode:?}");
+ }
+ }
+}
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 f3e22cdd0c..d0d0c99bb5 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,7 +27,7 @@ use datafusion_common::Result;
use crate::InputOrderMode;
use crate::aggregates::aggregate_hash_table::FinalMarker;
-use crate::aggregates::{AggregateExec, AggregateMode};
+use crate::aggregates::{AggregateExec, AggregateMode,
group_values::AccumulatorPhase};
use super::common::HashAggregateAccumulator;
use super::common_ordered::{OrderedAggregateTable,
OrderedAggregateTableMetrics};
@@ -76,6 +76,7 @@ impl OrderedAggregateTable<FinalMarker> {
self.aggregate_evaluated_batch(
&evaluated_batch,
HashAggregateAccumulator::merge_batch,
+ AccumulatorPhase::Merge,
)
}
@@ -83,6 +84,9 @@ impl OrderedAggregateTable<FinalMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
-
self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
+ self.next_output_batch_inner(
+ HashAggregateAccumulator::evaluate_to_columns,
+ AccumulatorPhase::Evaluate,
+ )
}
}
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 ea38346729..6ed93e59f3 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,6 +37,7 @@ use datafusion_common::Result;
use crate::aggregates::{
AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker,
+ group_values::AccumulatorPhase,
};
use super::common::HashAggregateAccumulator;
@@ -84,6 +85,7 @@ impl OrderedAggregateTable<PartialMarker> {
self.aggregate_evaluated_batch(
&evaluated_batch,
HashAggregateAccumulator::update_batch,
+ AccumulatorPhase::Update,
)
}
@@ -104,6 +106,9 @@ impl OrderedAggregateTable<PartialMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- self.next_output_batch_inner(HashAggregateAccumulator::state)
+ self.next_output_batch_inner(
+ HashAggregateAccumulator::state,
+ AccumulatorPhase::State,
+ )
}
}
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 db53e2822e..ce1ce647b4 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,7 +24,7 @@ use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use crate::aggregates::aggregate_hash_table::SingleMarker;
-use crate::aggregates::{AggregateExec, AggregateMode};
+use crate::aggregates::{AggregateExec, AggregateMode,
group_values::AccumulatorPhase};
use super::common::HashAggregateAccumulator;
use super::common_ordered::{OrderedAggregateTable,
OrderedAggregateTableMetrics};
@@ -76,6 +76,7 @@ impl OrderedAggregateTable<SingleMarker> {
self.aggregate_evaluated_batch(
&evaluated_batch,
HashAggregateAccumulator::update_batch,
+ AccumulatorPhase::Update,
)
}
@@ -84,6 +85,9 @@ impl OrderedAggregateTable<SingleMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
-
self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
+ self.next_output_batch_inner(
+ HashAggregateAccumulator::evaluate_to_columns,
+ AccumulatorPhase::Evaluate,
+ )
}
}
diff --git
a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
index 4dfd6a74d1..2892d05933 100644
---
a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
+++
b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
@@ -22,6 +22,7 @@ use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use crate::aggregates::AggregateExec;
+use crate::aggregates::group_values::AccumulatorPhase;
use super::common::{AggregateHashTable, HashAggregateAccumulator,
PartialReduceMarker};
@@ -52,7 +53,10 @@ impl AggregateHashTable<PartialReduceMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- self.next_output_batch_inner(HashAggregateAccumulator::state)
+ self.next_output_batch_inner(
+ HashAggregateAccumulator::state,
+ AccumulatorPhase::State,
+ )
}
/// Partial-reduce aggregation consumes partial aggregate states and merges
@@ -61,7 +65,11 @@ impl AggregateHashTable<PartialReduceMarker> {
&mut self,
batch: &RecordBatch,
) -> Result<()> {
- self.aggregate_batch_inner(batch,
HashAggregateAccumulator::merge_batch)
+ self.aggregate_batch_inner(
+ batch,
+ HashAggregateAccumulator::merge_batch,
+ AccumulatorPhase::Merge,
+ )
}
pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> {
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 915e7f1beb..54997f0537 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
@@ -24,7 +24,7 @@ use arrow::datatypes::SchemaRef;
use arrow::record_batch::RecordBatch;
use datafusion_common::{Result, assert_eq_or_internal_err};
-use crate::aggregates::group_values::new_group_values;
+use crate::aggregates::group_values::{AccumulatorPhase, new_group_values};
use crate::aggregates::order::GroupOrdering;
use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal};
@@ -66,7 +66,10 @@ impl AggregateHashTable<PartialMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
- self.next_output_batch_inner(HashAggregateAccumulator::state)
+ self.next_output_batch_inner(
+ HashAggregateAccumulator::state,
+ AccumulatorPhase::State,
+ )
}
/// In skip-partial-aggregation optimization, when a decision has been
made to skip
@@ -87,6 +90,9 @@ impl AggregateHashTable<PartialMarker> {
Ok(AggregateHashTable {
group_by_metrics: self.group_by_metrics.clone(),
aggregate_argument_metrics:
self.aggregate_argument_metrics.clone(),
+ aggregate_accumulator_metrics: Arc::clone(
+ &self.aggregate_accumulator_metrics,
+ ),
input_schema: Arc::clone(&self.input_schema),
output_schema: Arc::clone(&self.output_schema),
state_schema: Arc::clone(&self.state_schema),
@@ -107,7 +113,11 @@ impl AggregateHashTable<PartialMarker> {
&mut self,
batch: &RecordBatch,
) -> Result<()> {
- self.aggregate_batch_inner(batch,
HashAggregateAccumulator::update_batch)
+ self.aggregate_batch_inner(
+ batch,
+ HashAggregateAccumulator::update_batch,
+ AccumulatorPhase::Update,
+ )
}
pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> {
@@ -136,6 +146,7 @@ impl AggregateHashTable<PartialMarker> {
return Ok(());
}
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
let max_ordinal = max_duplicate_ordinal(state.group_by.groups());
let mut ordinals: HashMap<&[bool], usize> = HashMap::new();
let group_schema = state.group_by.group_schema(&self.input_schema)?;
@@ -171,13 +182,15 @@ impl AggregateHashTable<PartialMarker> {
if any_interned {
let total_groups = state.group_values.len();
let false_filter = BooleanArray::from(vec![false]);
- for acc in state.accumulators.iter_mut() {
+ for (idx, acc) in state.accumulators.iter_mut().enumerate() {
let null_args = acc.null_arguments(&self.input_schema)?;
let values = EvaluatedAccumulatorArgs {
arguments: null_args,
filter: Some(Arc::new(false_filter.clone())),
};
- acc.update_batch(&values, &[0], total_groups)?;
+ accumulator_metrics.time(idx, AccumulatorPhase::Update, || {
+ acc.update_batch(&values, &[0], total_groups)
+ })?;
}
}
@@ -203,13 +216,19 @@ impl AggregateHashTable<PartialSkipMarker> {
.next()
.unwrap_or_default();
+ let accumulator_metrics =
Arc::clone(&self.aggregate_accumulator_metrics);
let state = self.state.building_mut();
- for (acc, values) in state
+ for (idx, (acc, values)) in state
.accumulators
.iter_mut()
.zip(evaluated_batch.accumulator_args.iter())
+ .enumerate()
{
- output.extend(acc.convert_to_state(values)?);
+ output.extend(accumulator_metrics.time(
+ idx,
+ AccumulatorPhase::ConvertToState,
+ || acc.convert_to_state(values),
+ )?);
}
Ok(RecordBatch::try_new(
diff --git
a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
index 56d601c793..1ff05fc79d 100644
---
a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
+++
b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
@@ -20,6 +20,7 @@ use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use crate::aggregates::AggregateExec;
+use crate::aggregates::group_values::AccumulatorPhase;
use super::common::{AggregateHashTable, HashAggregateAccumulator,
SingleMarker};
@@ -57,7 +58,10 @@ impl AggregateHashTable<SingleMarker> {
pub(in crate::aggregates) fn next_output_batch(
&mut self,
) -> Result<Option<RecordBatch>> {
-
self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns)
+ self.next_output_batch_inner(
+ HashAggregateAccumulator::evaluate_to_columns,
+ AccumulatorPhase::Evaluate,
+ )
}
/// Single aggregation consumes raw input rows and updates the table's
@@ -66,7 +70,11 @@ impl AggregateHashTable<SingleMarker> {
&mut self,
batch: &RecordBatch,
) -> Result<()> {
- self.aggregate_batch_inner(batch,
HashAggregateAccumulator::update_batch)
+ self.aggregate_batch_inner(
+ batch,
+ HashAggregateAccumulator::update_batch,
+ AccumulatorPhase::Update,
+ )
}
pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> {
diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs
b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs
index 0011150970..a1ff2b7b6e 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs
@@ -56,6 +56,91 @@ impl AggregateArgumentMetrics {
}
}
+#[derive(Clone, Copy, PartialEq)]
+pub(crate) enum AccumulatorPhase {
+ Update,
+ Merge,
+ State,
+ ConvertToState,
+ Evaluate,
+}
+
+#[derive(Clone)]
+pub(crate) struct AggregateAccumulatorMetrics {
+ update_times: Option<Vec<Time>>,
+ merge_times: Option<Vec<Time>>,
+ state_times: Option<Vec<Time>>,
+ convert_to_state_times: Option<Vec<Time>>,
+ evaluate_times: Option<Vec<Time>>,
+}
+
+impl AggregateAccumulatorMetrics {
+ pub(crate) fn new<T>(
+ metrics: &ExecutionPlanMetricsSet,
+ partition: usize,
+ aggregate_labels: impl IntoIterator<Item = T>,
+ phases: &[AccumulatorPhase],
+ ) -> Self
+ where
+ T: Into<String>,
+ {
+ let aggregate_labels = aggregate_labels
+ .into_iter()
+ .map(Into::into)
+ .collect::<Vec<String>>();
+ let new_phase_metrics = |phase| {
+ aggregate_labels
+ .iter()
+ .enumerate()
+ .map(|(idx, label)| {
+ MetricBuilder::new(metrics)
+ .with_new_label("aggregate", label.clone())
+ .subset_time(format!("agg_expr_{idx}_{phase}_time"),
partition)
+ })
+ .collect()
+ };
+
+ Self {
+ update_times: phases
+ .contains(&AccumulatorPhase::Update)
+ .then(|| new_phase_metrics("update")),
+ merge_times: phases
+ .contains(&AccumulatorPhase::Merge)
+ .then(|| new_phase_metrics("merge")),
+ state_times: phases
+ .contains(&AccumulatorPhase::State)
+ .then(|| new_phase_metrics("state")),
+ convert_to_state_times: phases
+ .contains(&AccumulatorPhase::ConvertToState)
+ .then(|| new_phase_metrics("convert_to_state")),
+ evaluate_times: phases
+ .contains(&AccumulatorPhase::Evaluate)
+ .then(|| new_phase_metrics("evaluate")),
+ }
+ }
+
+ pub(crate) fn time<R>(
+ &self,
+ index: usize,
+ phase: AccumulatorPhase,
+ f: impl FnOnce() -> R,
+ ) -> R {
+ let times = match phase {
+ AccumulatorPhase::Update => self.update_times.as_ref(),
+ AccumulatorPhase::Merge => self.merge_times.as_ref(),
+ AccumulatorPhase::State => self.state_times.as_ref(),
+ AccumulatorPhase::ConvertToState =>
self.convert_to_state_times.as_ref(),
+ AccumulatorPhase::Evaluate => self.evaluate_times.as_ref(),
+ };
+ debug_assert!(
+ times.is_some_and(|times| index < times.len()),
+ "aggregate accumulator metric index {index} for uninitialized
phase"
+ );
+ let _timer = times.and_then(|times| times.get(index)).map(Time::timer);
+ f()
+ }
+}
+
#[derive(Clone)]
pub(crate) struct GroupByMetrics {
/// Time spent calculating the group IDs from the evaluated grouping
columns.
@@ -97,6 +182,7 @@ mod tests {
use arrow::record_batch::RecordBatch;
use datafusion_common::Result;
use datafusion_execution::TaskContext;
+ use datafusion_execution::config::SessionConfig;
use datafusion_execution::runtime_env::RuntimeEnvBuilder;
use datafusion_functions_aggregate::count::count_udaf;
use datafusion_functions_aggregate::sum::sum_udaf;
@@ -121,15 +207,18 @@ mod tests {
assert!(emitting_time.unwrap().as_usize() > 0);
}
- fn aggregate_argument_metric_names_and_labels(
+ fn aggregate_metric_names_and_labels(
metrics: &MetricsSet,
+ suffix: &str,
) -> 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") =>
+ if name
+ .strip_prefix("agg_expr_")
+ .and_then(|name| name.split_once('_'))
+ .is_some_and(|(_, metric)| metric == suffix) =>
{
let aggregate_label = metric
.labels()
@@ -144,6 +233,37 @@ mod tests {
.collect()
}
+ fn assert_aggregate_metric_labels(metrics: &MetricsSet, suffix: &str) {
+ let mut metric_names_and_labels =
+ aggregate_metric_names_and_labels(metrics, suffix);
+ metric_names_and_labels.sort();
+ assert_eq!(
+ metric_names_and_labels,
+ vec![
+ (format!("agg_expr_0_{suffix}"), "SUM(a)".to_string()),
+ (format!("agg_expr_1_{suffix}"), "SUM(b)".to_string()),
+ ]
+ );
+ }
+
+ fn assert_aggregate_metric_times_positive(metrics: &MetricsSet, suffix:
&str) {
+ let mut found = false;
+ for metric in metrics.iter().filter(|metric| {
+ matches!(
+ metric.value(),
+ MetricValue::Time { name, .. }
+ if name
+ .strip_prefix("agg_expr_")
+ .and_then(|name| name.split_once('_'))
+ .is_some_and(|(_, phase)| phase == suffix)
+ )
+ }) {
+ found = true;
+ assert!(metric.value().as_usize() > 0);
+ }
+ assert!(found, "expected aggregate metrics ending in {suffix}");
+ }
+
fn sum_aggregate(
schema: &Arc<Schema>,
column: &str,
@@ -283,27 +403,24 @@ mod tests {
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 task_ctx =
+ Arc::new(
+ TaskContext::default()
+ .with_runtime(runtime)
+ .with_session_config(SessionConfig::new().set_bool(
+ "datafusion.execution.enable_migration_aggregate",
+ true,
+ )),
+ );
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(),
- ),
- ]
- );
+ 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(())
}
@@ -313,6 +430,7 @@ mod tests {
let schema = Arc::new(Schema::new(vec![
Field::new("a", DataType::UInt32, false),
Field::new("b", DataType::Float64, false),
+ Field::new("c", DataType::Float64, false),
]));
let batches = (0..3)
@@ -326,6 +444,11 @@ mod tests {
(i + 1) as f64,
(i + 2) as f64,
])),
+ Arc::new(Float64Array::from(vec![
+ (i + 3) as f64,
+ (i + 4) as f64,
+ (i + 5) as f64,
+ ])),
],
)
.unwrap()
@@ -338,14 +461,17 @@ mod tests {
let group_by =
PhysicalGroupBy::new_single(vec![(col("a", &schema)?,
"a".to_string())]);
- let aggregates = vec![sum_aggregate(&schema, "b", "SUM(b)")?];
+ let aggregates = vec![
+ sum_aggregate(&schema, "b", "SUM(b)")?,
+ sum_aggregate(&schema, "c", "SUM(c)")?,
+ ];
// Create partial aggregate
let partial_aggregate = Arc::new(AggregateExec::try_new(
AggregateMode::Partial,
group_by.clone(),
aggregates.clone(),
- vec![None],
+ vec![None, None],
partial_input,
Arc::clone(&schema),
)?);
@@ -355,17 +481,38 @@ mod tests {
AggregateMode::Final,
group_by.as_final(),
aggregates,
- vec![None],
+ vec![None, None],
partial_aggregate,
schema,
)?);
- let task_ctx = Arc::new(TaskContext::default());
+ let task_ctx = Arc::new(
+ TaskContext::default().with_session_config(
+ SessionConfig::new()
+
.set_bool("datafusion.execution.enable_migration_aggregate", true),
+ ),
+ );
let _result =
collect(Arc::clone(&final_aggregate) as _,
Arc::clone(&task_ctx)).await?;
let metrics = final_aggregate.metrics().unwrap();
assert_groupby_metrics(&metrics);
+ assert_eq!(
+ aggregate_metric_names_and_labels(&metrics, "merge_time"),
+ vec![
+ ("agg_expr_0_merge_time".to_string(), "SUM(b)".to_string()),
+ ("agg_expr_1_merge_time".to_string(), "SUM(c)".to_string()),
+ ]
+ );
+ assert_eq!(
+ aggregate_metric_names_and_labels(&metrics, "evaluate_time"),
+ vec![
+ ("agg_expr_0_evaluate_time".to_string(), "SUM(b)".to_string()),
+ ("agg_expr_1_evaluate_time".to_string(), "SUM(c)".to_string()),
+ ]
+ );
+ assert_aggregate_metric_times_positive(&metrics, "merge_time");
+ assert_aggregate_metric_times_positive(&metrics, "evaluate_time");
Ok(())
}
diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
index bd5b92747e..c1b1f0c468 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -49,7 +49,10 @@ use crate::aggregates::{
mod metrics;
mod null_builder;
-pub(crate) use metrics::{AggregateArgumentMetrics, GroupByMetrics};
+pub(crate) use metrics::{
+ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics,
+ GroupByMetrics,
+};
/// Stores the group values during hash aggregation.
///
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs
b/datafusion/physical-plan/src/aggregates/mod.rs
index 97226159da..0f1b718a3c 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -3375,6 +3375,18 @@ mod tests {
Arc::new(task_ctx)
}
+ fn new_migrated_spill_ctx(batch_size: usize, max_memory: usize) ->
Arc<TaskContext> {
+ let runtime = RuntimeEnvBuilder::new()
+ .with_memory_pool(Arc::new(FairSpillPool::new(max_memory)))
+ .build_arc()
+ .unwrap();
+ Arc::new(
+ TaskContext::default()
+ .with_session_config(migrated_hash_session_config(batch_size))
+ .with_runtime(runtime),
+ )
+ }
+
fn migrated_hash_session_config(batch_size: usize) -> SessionConfig {
SessionConfig::new()
.with_batch_size(batch_size)
@@ -3388,6 +3400,16 @@ mod tests {
)
}
+ fn assert_accumulator_phase_times(aggregate: &AggregateExec, phases:
&[&str]) {
+ let metrics = aggregate.metrics().unwrap();
+ for phase in phases {
+ let time = metrics
+ .sum_by_name(&format!("agg_expr_0_{phase}_time"))
+ .unwrap_or_else(|| panic!("aggregate records {phase} time"));
+ assert!(time.as_usize() > 0);
+ }
+ }
+
fn new_finite_memory_migrated_hash_ctx(
batch_size: usize,
max_memory: usize,
@@ -4554,6 +4576,7 @@ mod tests {
let stream: SendableRecordBatchStream = stream.into();
let output = collect(stream).await?;
+ assert_accumulator_phase_times(&aggregate, &["update", "state"]);
assert_snapshot!(batches_to_sort_string(&output), @r"
+----------+-----------+-------------------------+
| sort_col | group_col | COUNT(value_col)[count] |
@@ -4632,6 +4655,7 @@ mod tests {
let stream: SendableRecordBatchStream = stream.into();
let output = collect(stream).await?;
+ assert_accumulator_phase_times(&final_aggregate, &["merge",
"evaluate"]);
assert_snapshot!(batches_to_sort_string(&output), @r"
+-----+--------------+
| key | COUNT(value) |
@@ -5661,6 +5685,7 @@ mod tests {
.map(|m| m.as_usize())
.unwrap_or(0);
assert_eq!(skipped_rows, 3);
+ assert_accumulator_phase_times(&aggregate_exec, &["convert_to_state"]);
Ok(())
}
@@ -7007,9 +7032,18 @@ mod tests {
Arc::clone(&schema),
)?);
- let task_ctx = new_spill_ctx(1, 600);
+ let task_ctx = new_migrated_spill_ctx(1, 600);
let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await?;
- assert_spill_count_metric(true, aggr);
+ assert_spill_count_metric(true, Arc::clone(&aggr));
+ let metrics = aggr.metrics().unwrap();
+ for phase in ["update", "state", "merge", "evaluate"] {
+ let time = metrics
+ .sum_by_name(&format!("agg_expr_0_{phase}_time"))
+ .unwrap_or_else(|| {
+ panic!("migrated single aggregate records {phase} time")
+ });
+ assert!(time.as_usize() > 0);
+ }
allow_duplicates! {
assert_snapshot!(batches_to_string(&result), @r"
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
