Skip to content

Commit fe1ca24

Browse files
committed
feat(datafusion): add grouped update metric and update_batch_grouped for accumulators
This change introduces two new methods to the Accumulator trait: 1. `grouped_update_batch_metric`: Returns an optional metric that can be used to time grouped updates once per batch instead of per group. 2. `update_batch_grouped`: Updates state when called by a grouped accumulator adapter, with support for using the grouped update metric. The GroupsAccumulatorAdapter now uses these new methods to avoid timing every per-group call, recording one interval for the full batch instead. This reduces the overhead of metric collection when there are many groups. DistinctArrayAggAccumulator is updated to take advantage of the new grouped update functionality, skipping per-group timing for deduplication operations. This refactor improves performance for grouped aggregations by making metric collection proportional to batch count rather than group cardinality.
1 parent 8fc8b2d commit fe1ca24

5 files changed

Lines changed: 104 additions & 56 deletions

File tree

‎datafusion/expr-common/src/accumulator.rs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,24 @@ pub trait Accumulator: Send + Sync + Debug + std::any::Any {
8484
/// running sum.
8585
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()>;
8686

87+
/// Returns an optional metric timed once per grouped adapter input batch.
88+
///
89+
/// A grouped accumulator adapter uses this for aggregate-owned work it
90+
/// dispatches to one accumulator per group. The default preserves the
91+
/// usual per-accumulator update path.
92+
fn grouped_update_batch_metric(&self) -> Option<Arc<dyn AggregateMetric>> {
93+
None
94+
}
95+
96+
/// Updates state when called by a grouped accumulator adapter.
97+
///
98+
/// The default delegates to [`Self::update_batch`]. Implementations that
99+
/// return a [`Self::grouped_update_batch_metric`] can avoid timing every
100+
/// per-group call; the adapter records one interval for the full batch.
101+
fn update_batch_grouped(&mut self, values: &[ArrayRef]) -> Result<()> {
102+
self.update_batch(values)
103+
}
104+
87105
/// Returns the final aggregate value.
88106
///
89107
/// For example, the `SUM` accumulator maintains a running sum,

‎datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs‎

Lines changed: 42 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ pub mod nulls;
2424
pub mod prim_op;
2525

2626
use std::mem::{size_of, size_of_val};
27+
use std::sync::Arc;
28+
use std::time::Instant;
2729

2830
use arrow::array::new_empty_array;
2931
use arrow::{
@@ -33,7 +35,7 @@ use arrow::{
3335
datatypes::UInt32Type,
3436
};
3537
use datafusion_common::{Result, ScalarValue, arrow_datafusion_err};
36-
use datafusion_expr_common::accumulator::Accumulator;
38+
use datafusion_expr_common::accumulator::{Accumulator, AggregateMetric};
3739
use datafusion_expr_common::groups_accumulator::{
3840
EmitTo, GroupSelection, GroupsAccumulator,
3941
};
@@ -102,6 +104,9 @@ pub struct GroupsAccumulatorAdapter {
102104
/// bottleneck in earlier implementations when there were many
103105
/// distinct groups.
104106
allocation_bytes: usize,
107+
108+
/// Optional aggregate-owned metric timed once for a grouped update batch.
109+
grouped_update_metric: Option<Arc<dyn AggregateMetric>>,
105110
}
106111

107112
struct AccumulatorState {
@@ -139,6 +144,7 @@ impl GroupsAccumulatorAdapter {
139144
factory: Box::new(factory),
140145
states: vec![],
141146
allocation_bytes: 0,
147+
grouped_update_metric: None,
142148
}
143149
}
144150

@@ -152,6 +158,9 @@ impl GroupsAccumulatorAdapter {
152158
let new_accumulators = total_num_groups - self.states.len();
153159
for _ in 0..new_accumulators {
154160
let accumulator = (self.factory)()?;
161+
if self.grouped_update_metric.is_none() {
162+
self.grouped_update_metric = accumulator.grouped_update_batch_metric();
163+
}
155164
let state = AccumulatorState::new(accumulator);
156165
self.add_allocation(state.size());
157166
self.states.push(state);
@@ -191,6 +200,7 @@ impl GroupsAccumulatorAdapter {
191200
group_indices: &[usize],
192201
opt_filter: Option<&BooleanArray>,
193202
total_num_groups: usize,
203+
time_grouped_update: bool,
194204
f: F,
195205
) -> Result<()>
196206
where
@@ -245,25 +255,37 @@ impl GroupsAccumulatorAdapter {
245255
// RecordBatch(es)
246256
let iter = groups_with_rows.iter().zip(offsets.windows(2));
247257

258+
let grouped_update_metric = time_grouped_update
259+
.then(|| self.grouped_update_metric.as_ref().cloned())
260+
.flatten();
261+
let start = grouped_update_metric.as_ref().map(|_| Instant::now());
262+
248263
let mut sizes_pre = 0;
249264
let mut sizes_post = 0;
250-
for (&group_idx, offsets) in iter {
251-
let state = &mut self.states[group_idx];
252-
sizes_pre += state.size();
253-
254-
let values_to_accumulate = slice_and_maybe_filter(
255-
&values,
256-
opt_filter.as_ref().map(|f| f.as_boolean()),
257-
offsets,
258-
)?;
259-
f(state.accumulator.as_mut(), &values_to_accumulate)?;
260-
261-
// clear out the state so they are empty for next
262-
// iteration
263-
state.indices.clear();
264-
sizes_post += state.size();
265-
}
265+
let result: Result<()> = (|| {
266+
for (&group_idx, offsets) in iter {
267+
let state = &mut self.states[group_idx];
268+
sizes_pre += state.size();
269+
270+
let values_to_accumulate = slice_and_maybe_filter(
271+
&values,
272+
opt_filter.as_ref().map(|f| f.as_boolean()),
273+
offsets,
274+
)?;
275+
f(state.accumulator.as_mut(), &values_to_accumulate)?;
276+
277+
// clear out the state so they are empty for next
278+
// iteration
279+
state.indices.clear();
280+
sizes_post += state.size();
281+
}
282+
Ok(())
283+
})();
266284

285+
if let (Some(metric), Some(start)) = (grouped_update_metric, start) {
286+
metric.add_duration(start.elapsed());
287+
}
288+
result?;
267289
self.adjust_allocation(sizes_pre, sizes_post);
268290
Ok(())
269291
}
@@ -310,8 +332,9 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
310332
group_indices,
311333
opt_filter,
312334
total_num_groups,
335+
true,
313336
|accumulator, values_to_accumulate| {
314-
accumulator.update_batch(values_to_accumulate)
337+
accumulator.update_batch_grouped(values_to_accumulate)
315338
},
316339
)?;
317340
Ok(())
@@ -412,6 +435,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
412435
group_indices,
413436
None,
414437
total_num_groups,
438+
false,
415439
|accumulator, values_to_accumulate| {
416440
accumulator.merge_batch(values_to_accumulate)?;
417441
Ok(())

‎datafusion/functions-aggregate/src/array_agg.rs‎

Lines changed: 33 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -857,12 +857,6 @@ pub struct DistinctArrayAggAccumulator {
857857
distinct_metric: Option<Arc<dyn AggregateMetric>>,
858858
}
859859

860-
/// Avoid per-group timer overhead when the adapter calls us with a tiny batch.
861-
///
862-
/// Batches below this threshold are not large enough for a useful timing
863-
/// sample, while timing each one adds two clock reads and an atomic update.
864-
const DISTINCT_METRIC_MIN_BATCH_SIZE: usize = 16;
865-
866860
/// Returns `true` if `dt` is, or recursively contains, a `Dictionary` type.
867861
///
868862
/// `RowConverter` always decodes to the physical (non-dictionary) type, so a
@@ -922,16 +916,12 @@ impl DistinctArrayAggAccumulator {
922916
}
923917
}
924918

925-
impl Accumulator for DistinctArrayAggAccumulator {
926-
fn set_metrics(&mut self, metrics: Arc<dyn AggregateMetrics>) {
927-
self.distinct_metric = Some(metrics.metric("distinct"));
928-
}
929-
930-
fn state(&mut self) -> Result<Vec<ScalarValue>> {
931-
Ok(vec![self.evaluate()?])
932-
}
933-
934-
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
919+
impl DistinctArrayAggAccumulator {
920+
fn update_batch_impl(
921+
&mut self,
922+
values: &[ArrayRef],
923+
record_metric: bool,
924+
) -> Result<()> {
935925
if values.is_empty() {
936926
return Ok(());
937927
}
@@ -961,7 +951,7 @@ impl Accumulator for DistinctArrayAggAccumulator {
961951
return Ok(());
962952
}
963953

964-
let distinct_metric = (col.len() >= DISTINCT_METRIC_MIN_BATCH_SIZE)
954+
let distinct_metric = record_metric
965955
.then(|| self.distinct_metric.as_ref().cloned())
966956
.flatten();
967957
let distinct_start = distinct_metric.as_ref().map(|_| Instant::now());
@@ -1016,6 +1006,28 @@ impl Accumulator for DistinctArrayAggAccumulator {
10161006
}
10171007
Ok(())
10181008
}
1009+
}
1010+
1011+
impl Accumulator for DistinctArrayAggAccumulator {
1012+
fn set_metrics(&mut self, metrics: Arc<dyn AggregateMetrics>) {
1013+
self.distinct_metric = Some(metrics.metric("distinct"));
1014+
}
1015+
1016+
fn grouped_update_batch_metric(&self) -> Option<Arc<dyn AggregateMetric>> {
1017+
self.distinct_metric.as_ref().cloned()
1018+
}
1019+
1020+
fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
1021+
self.update_batch_impl(values, true)
1022+
}
1023+
1024+
fn update_batch_grouped(&mut self, values: &[ArrayRef]) -> Result<()> {
1025+
self.update_batch_impl(values, false)
1026+
}
1027+
1028+
fn state(&mut self) -> Result<Vec<ScalarValue>> {
1029+
Ok(vec![self.evaluate()?])
1030+
}
10191031

10201032
fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> {
10211033
if states.is_empty() {
@@ -1514,19 +1526,15 @@ mod tests {
15141526
}
15151527

15161528
#[test]
1517-
fn distinct_accumulator_skips_metric_for_small_batches() -> Result<()> {
1529+
fn distinct_accumulator_records_metric_for_small_batches() -> Result<()> {
15181530
let metric_updates = Arc::new(AtomicUsize::new(0));
15191531
let mut accumulator =
15201532
DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?;
15211533
accumulator.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates))));
15221534

1523-
accumulator.update_batch(&[Arc::new(Int32Array::from(vec![
1524-
1;
1525-
DISTINCT_METRIC_MIN_BATCH_SIZE
1526-
- 1
1527-
]))])?;
1535+
accumulator.update_batch(&[Arc::new(Int32Array::from(vec![1]))])?;
15281536

1529-
assert_eq!(metric_updates.load(Ordering::Relaxed), 0);
1537+
assert_eq!(metric_updates.load(Ordering::Relaxed), 1);
15301538
Ok(())
15311539
}
15321540

@@ -1537,10 +1545,7 @@ mod tests {
15371545
DistinctArrayAggAccumulator::try_new(&DataType::Int32, None, false)?;
15381546
accumulator.set_metrics(Arc::new(CountingMetrics(Arc::clone(&metric_updates))));
15391547

1540-
accumulator.update_batch(&[Arc::new(Int32Array::from(vec![
1541-
1;
1542-
DISTINCT_METRIC_MIN_BATCH_SIZE
1543-
]))])?;
1548+
accumulator.update_batch(&[Arc::new(Int32Array::from(vec![1; 16]))])?;
15441549

15451550
assert_eq!(metric_updates.load(Ordering::Relaxed), 1);
15461551
Ok(())

‎datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1558,13 +1558,12 @@ mod tests {
15581558
batch?;
15591559
}
15601560

1561-
assert!(
1562-
aggregate_exec
1563-
.metrics()
1564-
.unwrap()
1565-
.sum_by_name("agg_expr_0_internal_distinct_time")
1566-
.is_some_and(|metric| metric.as_usize() > 0)
1567-
);
1561+
let metrics = aggregate_exec.metrics().unwrap();
1562+
let distinct_time = metrics
1563+
.iter()
1564+
.find(|metric| metric.value().name() == "agg_expr_0_internal_distinct_time")
1565+
.expect("internal distinct time metric");
1566+
assert!(distinct_time.value().as_usize() > 0);
15681567

15691568
Ok(())
15701569
}

‎docs/source/user-guide/metrics.md‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -165,9 +165,11 @@ accumulators in that partition share the same time, and normal metric display
165165
combines that time across partitions. An aggregate may request its submetric
166166
during accumulator construction, so it can appear even when its input is empty.
167167
For example, `array_agg(DISTINCT ...)` records the time spent deduplicating
168-
input values as `agg_expr_{index}_internal_distinct_time`. These submetrics
169-
complement the `update`, `merge`, `state`, and `evaluate` timers rather than
170-
subdividing or replacing them.
168+
input values as `agg_expr_{index}_internal_distinct_time`. Grouped accumulation
169+
records this once per input batch, rather than once per group, to avoid making
170+
metric collection proportional to group cardinality. These submetrics complement
171+
the `update`, `merge`, `state`, and `evaluate` timers rather than subdividing or
172+
replacing them.
171173

172174
Except for the `Summary` metric `reduction_factor`, these operator-level and
173175
per-aggregate metrics are `Dev` metrics. They appear in `EXPLAIN ANALYZE` when

0 commit comments

Comments
 (0)