From da087a0614087e77ea564cd7833a6a98c5545b77 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:00:04 -0500 Subject: [PATCH 1/3] perf: walk only the groups a batch has rows for in GroupsAccumulatorAdapter To build the `take` index for a batch, `GroupsAccumulatorAdapter` walked every group that exists and skipped the empty ones. That cost one iteration per group on every batch, whatever the batch touched. With 1.5M groups in a partition and 8192-row batches, that is about 180 iterations per input row before any aggregation happens, and it is why an aggregate with no native `GroupsAccumulator` slows down as a `GROUP BY` gains distinct values. The push loop already reaches each group that has rows, and a group's scratch indices are empty exactly when this batch has not reached it yet, so recording the group there costs a comparison on a cache line the loop already has. Build the `take` index from that list instead of from a scan over every group. Rows keep their relative order within a group, so each accumulator sees exactly the rows, in the order, it saw before. The list is in order of first appearance rather than ascending group order, which the code below it does not depend on. Co-Authored-By: Claude Opus 5 --- .../src/aggregate/groups_accumulator.rs | 178 ++++++++++++++++-- 1 file changed, 159 insertions(+), 19 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 6704b068acf0b..3034d51e92763 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -200,34 +200,34 @@ impl GroupsAccumulatorAdapter { assert_eq!(values[0].len(), group_indices.len()); + // groups_with_rows holds a list of group indexes that have any rows + // that need to be accumulated, stored in order of first appearance in + // this batch + let mut groups_with_rows = vec![]; + // figure out which input rows correspond to which groups. - // Note that self.state.indices starts empty for all groups - // (it is cleared out below) + // Note that self.state.indices starts empty for all groups (it is + // cleared out below), so an empty one is exactly a group that this + // batch has not reached yet for (idx, group_index) in group_indices.iter().enumerate() { - self.states[*group_index].indices.push(idx as u32); + let state = &mut self.states[*group_index]; + if state.indices.is_empty() { + groups_with_rows.push(*group_index); + } + state.indices.push(idx as u32); } - // groups_with_rows holds a list of group indexes that have - // any rows that need to be accumulated, stored in order of - // group_index - - let mut groups_with_rows = vec![]; - // batch_indices holds indices into values, each group is contiguous - let mut batch_indices = vec![]; + let mut batch_indices = Vec::with_capacity(group_indices.len()); // offsets[i] is index into batch_indices where the rows for - // group_index i starts - let mut offsets = vec![0]; + // groups_with_rows[i] start + let mut offsets = Vec::with_capacity(groups_with_rows.len() + 1); + offsets.push(0); let mut offset_so_far = 0; - for (group_index, state) in self.states.iter_mut().enumerate() { - let indices = &state.indices; - if indices.is_empty() { - continue; - } - - groups_with_rows.push(group_index); + for &group_index in groups_with_rows.iter() { + let indices = &self.states[group_index].indices; batch_indices.extend_from_slice(indices); offset_so_far += indices.len(); offsets.push(offset_so_far); @@ -539,6 +539,146 @@ mod tests { use arrow::array::{AsArray, Int64Array}; use arrow::datatypes::{DataType, Int64Type}; + /// Counts the rows it is handed. `MaxAccumulator` cannot see a row that + /// reaches it twice, so the tests need an accumulator that can. + #[derive(Debug, Default)] + struct RowCountAccumulator { + rows: i64, + } + + impl Accumulator for RowCountAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.rows += values[0].len() as i64; + Ok(()) + } + fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.update_batch(values) + } + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(Some(self.rows))) + } + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Int64(Some(self.rows))]) + } + fn size(&self) -> usize { + size_of::() + } + } + + fn row_count_adapter() -> GroupsAccumulatorAdapter { + GroupsAccumulatorAdapter::new(|| { + Ok(Box::new(RowCountAccumulator::default()) as Box) + }) + } + + fn max_adapter() -> GroupsAccumulatorAdapter { + GroupsAccumulatorAdapter::new(|| { + Ok(Box::new(MaxAccumulator::try_new(&DataType::Int64)?) + as Box) + }) + } + + /// Every row must reach its own group's accumulator, whatever order the + /// groups appear in within a batch and however many of them have no rows. + #[test] + fn adapter_routes_rows_to_their_own_group() -> Result<()> { + // only the first `GROUPS_WITH_ROWS` of them ever get a row, so the + // batches leave a tail of groups untouched + const TOTAL_NUM_GROUPS: usize = 32; + const GROUPS_WITH_ROWS: usize = 23; + + let mut accumulator = max_adapter(); + let mut expected = vec![None; TOTAL_NUM_GROUPS]; + + // an empty batch, then batches that interleave and revisit groups + for batch in 0..5 { + let mut group_indices = vec![]; + let mut values = vec![]; + for row in 0..(batch * 31) { + let group_index = (row * 5 + batch) % GROUPS_WITH_ROWS; + let value = ((row * 7 + batch * 13) % 97) as i64; + group_indices.push(group_index); + values.push(value); + expected[group_index] = expected[group_index].max(Some(value)); + } + + let values: ArrayRef = Arc::new(Int64Array::from(values)); + accumulator.update_batch( + &[values], + &group_indices, + None, + TOTAL_NUM_GROUPS, + )?; + } + + assert_eq!( + accumulator + .evaluate(EmitTo::All)? + .as_primitive::(), + &Int64Array::from(expected) + ); + Ok(()) + } + + /// Rows the filter drops must not reach any accumulator, including when it + /// drops every row a group has. + #[test] + fn adapter_routes_filtered_rows() -> Result<()> { + let mut accumulator = max_adapter(); + + let values: ArrayRef = Arc::new(Int64Array::from(vec![5, 9, 1, 7, 3, 8])); + let group_indices = [2, 0, 2, 1, 0, 2]; + let filter = BooleanArray::from(vec![true, false, true, false, true, false]); + accumulator.update_batch(&[values], &group_indices, Some(&filter), 4)?; + + // group 0 keeps 3, group 1 loses its only row, group 2 keeps 5 and 1, + // and group 3 never had a row + assert_eq!( + accumulator + .evaluate(EmitTo::All)? + .as_primitive::(), + &Int64Array::from(vec![Some(3), None, Some(5), None]) + ); + Ok(()) + } + + /// Each row must reach its own group's accumulator exactly once. A group + /// recorded twice for one batch would hand that group its rows twice. + #[test] + fn adapter_routes_each_row_exactly_once() -> Result<()> { + const TOTAL_NUM_GROUPS: usize = 16; + + let mut accumulator = row_count_adapter(); + let mut expected = vec![0i64; TOTAL_NUM_GROUPS]; + + for batch in 0..4 { + let mut group_indices = vec![]; + for row in 0..(batch * 29) { + // revisit groups within a batch, and leave a tail untouched + let group_index = (row * 3 + batch) % 11; + group_indices.push(group_index); + expected[group_index] += 1; + } + + let values: ArrayRef = + Arc::new(Int64Array::from(vec![1i64; group_indices.len()])); + accumulator.update_batch( + &[values], + &group_indices, + None, + TOTAL_NUM_GROUPS, + )?; + } + + assert_eq!( + accumulator + .evaluate(EmitTo::All)? + .as_primitive::(), + &Int64Array::from(expected) + ); + Ok(()) + } + #[test] fn adapter_preserving_evaluation_uses_accumulator_contract() -> Result<()> { let mut accumulator = GroupsAccumulatorAdapter::new(|| { From 81e9a2cdc6998e779bd0f07de2a8cf6ef7ebf107 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 9 Sep 2026 11:27:57 -0500 Subject: [PATCH 2/3] bench: clickbench_extended q15/q16 for GroupsAccumulatorAdapter `covar_samp` has no native `GroupsAccumulator`, so it runs through `GroupsAccumulatorAdapter`. The aggregate itself is cheap, which leaves the adapter's routing as the dominant cost and makes these queries a direct measure of it. q15 groups by `"UserID"` (about 17.6M groups) and q16 by `"RegionID"` (about 9,000), so the pair covers both ends of the cardinality range that the adapter has to stay fast at. `corr` over the same two columns is the native-`GroupsAccumulator` control. The outer `MAX` keeps the stored result small. Co-Authored-By: Claude Opus 5 --- benchmarks/queries/clickbench/README.md | 41 ++++++++++++++++++- .../queries/clickbench/extended/q15.sql | 8 ++++ .../queries/clickbench/extended/q16.sql | 8 ++++ 3 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 benchmarks/queries/clickbench/extended/q15.sql create mode 100644 benchmarks/queries/clickbench/extended/q16.sql diff --git a/benchmarks/queries/clickbench/README.md b/benchmarks/queries/clickbench/README.md index 2fd31c0b4c630..1d8c574604c1d 100644 --- a/benchmarks/queries/clickbench/README.md +++ b/benchmarks/queries/clickbench/README.md @@ -260,7 +260,6 @@ FROM hits WHERE "URL" < 'zzzz'; ``` - ### Q14: Grouped `COUNT(DISTINCT )` beside a non-distinct `COUNT(*)` **Question**: "For the top 10 search phrases by hit count, how many distinct @@ -301,6 +300,46 @@ LIMIT 10; +### Q15: Correlation of ad slot dimensions per user + +**Question**: "How correlated are the width and height of the screens each user +browses from?" + +**Important Query Properties**: `covar_samp` over a very high cardinality +`GROUP BY` (about 17.6M distinct `UserID`s). `covar_samp` has no native +`GroupsAccumulator`, so both `AggregateExec`s hold one `Accumulator` per group +behind `GroupsAccumulatorAdapter` and every input batch is routed through it. +That makes the query a direct measure of the adapter's per-batch cost at high +cardinality. The aggregate is cheap and the outer `MAX` keeps the result small, +so the adapter dominates the runtime. Q16 is the low cardinality twin, and +`corr` in place of `covar_samp` is the native-`GroupsAccumulator` control. + +```sql +SELECT MAX(c) FROM ( + SELECT covar_samp("ResolutionWidth", "ResolutionHeight") as c + FROM hits + GROUP BY "UserID" +); +``` + +### Q16: Correlation of ad slot dimensions per region + +**Question**: "How correlated are the width and height of the screens browsed +from in each region?" + +**Important Query Properties**: the same `covar_samp` aggregate as Q15 through +`GroupsAccumulatorAdapter`, but over about 9,000 groups instead of 17.6M. It is +the low cardinality control for Q15: work that trades per-group cost for +per-batch cost has to leave this one alone. + +```sql +SELECT MAX(c) FROM ( + SELECT covar_samp("ResolutionWidth", "ResolutionHeight") as c + FROM hits + GROUP BY "RegionID" +); +``` + ## Data Notes Here are some interesting statistics about the data used in the queries diff --git a/benchmarks/queries/clickbench/extended/q15.sql b/benchmarks/queries/clickbench/extended/q15.sql new file mode 100644 index 0000000000000..e648b8dd97992 --- /dev/null +++ b/benchmarks/queries/clickbench/extended/q15.sql @@ -0,0 +1,8 @@ +-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591 +-- set datafusion.execution.parquet.binary_as_string = true + +SELECT MAX(c) FROM ( + SELECT covar_samp("ResolutionWidth", "ResolutionHeight") as c + FROM hits + GROUP BY "UserID" +); diff --git a/benchmarks/queries/clickbench/extended/q16.sql b/benchmarks/queries/clickbench/extended/q16.sql new file mode 100644 index 0000000000000..f3db0ae2df8ec --- /dev/null +++ b/benchmarks/queries/clickbench/extended/q16.sql @@ -0,0 +1,8 @@ +-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591 +-- set datafusion.execution.parquet.binary_as_string = true + +SELECT MAX(c) FROM ( + SELECT covar_samp("ResolutionWidth", "ResolutionHeight") as c + FROM hits + GROUP BY "RegionID" +); From c41c95406817a4deb488c1d198bb5e78b6816c2d Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:26:26 -0500 Subject: [PATCH 3/3] bench: measure GroupsAccumulatorAdapter across GROUP BY cardinalities There was no benchmark that isolated what `GroupsAccumulatorAdapter` costs, so a change to its routing could only be judged through a whole query, where the parquet scan and the aggregate itself hide the effect. Sweep the group count from 64 to 1M with two accumulators. `routing` wraps an accumulator that only counts the rows it is handed, so what it measures is the adapter and nothing else: the upper bound on what a change to the routing can move. `covar_samp` wraps a real aggregate with no native `GroupsAccumulator`, so it shows how much of that upper bound a query sees. Co-Authored-By: Claude Opus 5 --- datafusion/functions-aggregate/Cargo.toml | 4 + .../benches/groups_accumulator_adapter.rs | 163 ++++++++++++++++++ 2 files changed, 167 insertions(+) create mode 100644 datafusion/functions-aggregate/benches/groups_accumulator_adapter.rs diff --git a/datafusion/functions-aggregate/Cargo.toml b/datafusion/functions-aggregate/Cargo.toml index ae6bf70088f58..ba9347cddae58 100644 --- a/datafusion/functions-aggregate/Cargo.toml +++ b/datafusion/functions-aggregate/Cargo.toml @@ -80,6 +80,10 @@ name = "min_max_bytes" name = "approx_distinct" harness = false +[[bench]] +name = "groups_accumulator_adapter" +harness = false + [[bench]] name = "first_last" harness = false diff --git a/datafusion/functions-aggregate/benches/groups_accumulator_adapter.rs b/datafusion/functions-aggregate/benches/groups_accumulator_adapter.rs new file mode 100644 index 0000000000000..e75d4cd374fd8 --- /dev/null +++ b/datafusion/functions-aggregate/benches/groups_accumulator_adapter.rs @@ -0,0 +1,163 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! What it costs to run an [`Accumulator`] through [`GroupsAccumulatorAdapter`], +//! over a sweep of `GROUP BY` cardinalities. +//! +//! Two accumulators, because the adapter's cost and the aggregate's cost move +//! in opposite directions as the group count grows: +//! +//! * `routing` wraps an accumulator that only counts the rows it is handed, so +//! what the benchmark measures is the adapter routing the batch and nothing +//! else. It is the upper bound on what a change to the routing can do. +//! * `covar_samp` wraps a real aggregate that has no native `GroupsAccumulator`, +//! so it shows how much of that upper bound a query actually sees. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array}; +use arrow::datatypes::{DataType, Field, Schema}; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{Accumulator, EmitTo, GroupsAccumulator}; +use datafusion_functions_aggregate::covariance::covar_samp_udaf; +use datafusion_physical_expr::GroupsAccumulatorAdapter; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; +use datafusion_physical_expr::expressions::col; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; + +const BATCH_SIZE: usize = 8192; +const NUM_BATCHES: usize = 128; +const CARDINALITIES: [usize; 5] = [64, 1024, 16_384, 262_144, 1_000_000]; + +/// The cheapest thing that is still an [`Accumulator`]: it only counts the rows +/// it is handed, so what a benchmark over it measures is the adapter. +#[derive(Debug, Default)] +struct CountingAccumulator { + count: i64, +} + +impl Accumulator for CountingAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.count += values[0].len() as i64; + Ok(()) + } + fn merge_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.update_batch(values) + } + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(Some(self.count))) + } + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Int64(Some(self.count))]) + } + fn size(&self) -> usize { + size_of::() + } +} + +/// A factory for the `covar_samp` [`Accumulator`], which is what the aggregate +/// operator falls back to because `covar_samp` has no native +/// `GroupsAccumulator`. +fn covar_samp_factory() -> impl Fn() -> Result> + Send + 'static { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, true), + Field::new("b", DataType::Float64, true), + ])); + let args = vec![col("a", &schema).unwrap(), col("b", &schema).unwrap()]; + let agg = Arc::new( + AggregateExprBuilder::new(covar_samp_udaf(), args) + .schema(schema) + .alias("covar_samp(a, b)") + .build() + .unwrap(), + ); + move || agg.create_accumulator() +} + +/// `num_columns` copies of the same values, so the same batches drive both a +/// one-argument and a two-argument aggregate. +fn make_batches( + num_groups: usize, + num_columns: usize, +) -> Vec<(Vec, Vec)> { + let mut rng = StdRng::seed_from_u64(7); + (0..NUM_BATCHES) + .map(|_| { + let group_indices: Vec = (0..BATCH_SIZE) + .map(|_| rng.random_range(0..num_groups)) + .collect(); + let values: ArrayRef = Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::())) + .collect::(), + ); + (vec![values; num_columns], group_indices) + }) + .collect() +} + +/// Feeds every batch into a fresh adapter and emits every group, which is what +/// one partition of an `AggregateExec` does. +fn run( + factory: impl Fn() -> Result> + Send + 'static, + data: &[(Vec, Vec)], + num_groups: usize, +) { + let mut accumulator = GroupsAccumulatorAdapter::new(factory); + for (values, group_indices) in data { + accumulator + .update_batch(values, group_indices, None, num_groups) + .unwrap(); + } + black_box(accumulator.evaluate(EmitTo::All).unwrap()); +} + +fn adapter_benchmark(c: &mut Criterion) { + for (name, num_columns) in [("routing", 1), ("covar_samp", 2)] { + let mut group = c.benchmark_group(format!("adapter_{name}")); + group.sample_size(20); + group.throughput(Throughput::Elements((BATCH_SIZE * NUM_BATCHES) as u64)); + + for num_groups in CARDINALITIES { + let data = make_batches(num_groups, num_columns); + group.bench_with_input( + BenchmarkId::from_parameter(num_groups), + &data, + |b, data| { + b.iter(|| { + if num_columns == 1 { + run( + || Ok(Box::new(CountingAccumulator::default()) as _), + data, + num_groups, + ) + } else { + run(covar_samp_factory(), data, num_groups) + } + }) + }, + ); + } + group.finish(); + } +} + +criterion_group!(benches, adapter_benchmark); +criterion_main!(benches);