perf: make GroupsAccumulatorAdapter cost proportional to the batch, not the group count

benchmarks/queries/clickbench/README.md

@@ -260,7 +260,6 @@ FROM hits
 WHERE "URL" < 'zzzz';
 ```
 
-
 ### Q14: Grouped `COUNT(DISTINCT <string>)` 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

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"
+);

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"
+);

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

@@ -48,8 +48,8 @@ use datafusion_expr_common::groups_accumulator::{
 ///
 /// Internally, this adapter creates a new [`Accumulator`] for each group which
 /// stores the state for that group. This both requires an allocation for each
-/// Accumulator, internal indices, as well as whatever internal allocations the
-/// Accumulator itself requires.
+/// Accumulator as well as whatever internal allocations the Accumulator itself
+/// requires.
 ///
 /// For example, a `MinAccumulator` that computes the minimum string value with
 /// a [`ScalarValue::Utf8`]. That will require at least two allocations per group
@@ -89,6 +89,13 @@ use datafusion_expr_common::groups_accumulator::{
 /// The adapter minimizes the number of calls to [`Accumulator::update_batch`]
 /// by first collecting the input rows for each group into a contiguous array
 /// using [`compute::take`]
+///
+/// Routing a batch that way costs one pass over the rows of the batch and two
+/// over the groups it touches, so the per-batch cost is proportional to the
+/// batch and not to the number of groups that exist. The scratch space it needs
+/// is one `take` index per batch plus [`Self::groups_with_rows`] and
+/// [`Self::offsets`], all of which are the size of a batch: no scratch is held
+/// per group between batches.
 pub struct GroupsAccumulatorAdapter {
     factory: Box<dyn Fn() -> Result<Box<dyn Accumulator>> + Send>,
 
@@ -102,29 +109,42 @@ pub struct GroupsAccumulatorAdapter {
     /// bottleneck in earlier implementations when there were many
     /// distinct groups.
     allocation_bytes: usize,
+
+    /// scratch space: the groups that have rows in the batch being routed, in
+    /// the order they first appear in it. Empty between batches, but keeps its
+    /// capacity so that routing a batch does not allocate.
+    groups_with_rows: Vec<usize>,
+
+    /// scratch space: `offsets[i]` is where the rows of `groups_with_rows[i]`
+    /// start in the `take` index built for the batch being routed, so
+    /// `offsets` has one more element than `groups_with_rows`. Empty between
+    /// batches, but keeps its capacity.
+    offsets: Vec<usize>,
 }
 
 struct AccumulatorState {
     /// [`Accumulator`] that stores the per-group state
     accumulator: Box<dyn Accumulator>,
 
-    /// scratch space: indexes in the input array that will be fed to
-    /// this accumulator. Stores indexes as `u32` to match the arrow
-    /// `take` kernel input.
-    indices: Vec<u32>,
+    /// scratch space: while a batch is being routed, first the number of its
+    /// rows that belong to this group and then the position in the `take`
+    /// index where the next of them goes. Always 0 between batches.
+    ///
+    /// Indexes are `u32` to match the arrow `take` kernel input.
+    cursor: u32,
 }
 
 impl AccumulatorState {
     fn new(accumulator: Box<dyn Accumulator>) -> Self {
         Self {
             accumulator,
-            indices: vec![],
+            cursor: 0,
         }
     }
 
     /// Returns the amount of memory taken by this structure and its accumulator
     fn size(&self) -> usize {
-        self.accumulator.size() + size_of_val(self) + self.indices.allocated_size()
+        self.accumulator.size() + size_of_val(self)
     }
 }
 
@@ -139,6 +159,8 @@ impl GroupsAccumulatorAdapter {
             factory: Box::new(factory),
             states: vec![],
             allocation_bytes: 0,
+            groups_with_rows: vec![],
+            offsets: vec![],
         }
     }
 
@@ -200,38 +222,54 @@ impl GroupsAccumulatorAdapter {
 
         assert_eq!(values[0].len(), group_indices.len());
 
-        // figure out which input rows correspond to which groups.
-        // Note that self.state.indices starts empty for all groups
-        // (it is cleared out below)
-        for (idx, group_index) in group_indices.iter().enumerate() {
-            self.states[*group_index].indices.push(idx as u32);
+        let Self {
+            states,
+            groups_with_rows,
+            offsets,
+            ..
+        } = self;
+
+        // Bucket the rows of this batch by group, so that each accumulator can
+        // be invoked once with a contiguous run of its own rows. Every pass
+        // below is over the rows of the batch or over the groups the batch
+        // touches, never over all the groups that exist.
+        groups_with_rows.clear();
+        offsets.clear();
+
+        // count the rows of each group, recording a group the first time it is
+        // seen so that the passes that follow only visit the groups with rows
+        for &group_index in group_indices {
+            let state = &mut states[group_index];
+            if state.cursor == 0 {
+                groups_with_rows.push(group_index);
+            }
+            state.cursor += 1;
         }
 
-        // 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![];
-
-        // offsets[i] is index into batch_indices where the rows for
-        // group_index i starts
-        let mut offsets = vec![0];
-
+        // turn those counts into the offsets at which each group's rows start
+        // in batch_indices, leaving every cursor at the start of its own range
+        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);
-            batch_indices.extend_from_slice(indices);
-            offset_so_far += indices.len();
+        for &group_index in groups_with_rows.iter() {
+            let cursor = &mut states[group_index].cursor;
+            let num_rows = *cursor as usize;
+            *cursor = offset_so_far as u32;
+            offset_so_far += num_rows;
             offsets.push(offset_so_far);
         }
+
+        // batch_indices holds indices into values, each group is contiguous:
+        // scatter every row into its group's range, in input order. The cursors
+        // walk to the end of their ranges, and are reset for the next batch.
+        let mut batch_indices = vec![0u32; group_indices.len()];
+        for (idx, &group_index) in group_indices.iter().enumerate() {
+            let cursor = &mut states[group_index].cursor;
+            batch_indices[*cursor as usize] = idx as u32;
+            *cursor += 1;
+        }
+        for &group_index in groups_with_rows.iter() {
+            states[group_index].cursor = 0;
+        }
         let batch_indices = batch_indices.into();
 
         // reorder the values and opt_filter by batch_indices so that
@@ -248,7 +286,7 @@ impl GroupsAccumulatorAdapter {
         let mut sizes_pre = 0;
         let mut sizes_post = 0;
         for (&group_idx, offsets) in iter {
-            let state = &mut self.states[group_idx];
+            let state = &mut states[group_idx];
             sizes_pre += state.size();
 
             let values_to_accumulate = slice_and_maybe_filter(
@@ -258,9 +296,6 @@ impl GroupsAccumulatorAdapter {
             )?;
             f(state.accumulator.as_mut(), &values_to_accumulate)?;
 
-            // clear out the state so they are empty for next
-            // iteration
-            state.indices.clear();
             sizes_post += state.size();
         }
 
@@ -422,6 +457,8 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
 
     fn size(&self) -> usize {
         self.allocation_bytes
+            + self.groups_with_rows.allocated_size()
+            + self.offsets.allocated_size()
     }
 
     fn convert_to_state(
@@ -539,6 +576,77 @@ mod tests {
     use arrow::array::{AsArray, Int64Array};
     use arrow::datatypes::{DataType, Int64Type};
 
+    fn max_adapter() -> GroupsAccumulatorAdapter {
+        GroupsAccumulatorAdapter::new(|| {
+            Ok(Box::new(MaxAccumulator::try_new(&DataType::Int64)?)
+                as Box<dyn Accumulator>)
+        })
+    }
+
+    /// 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::<Int64Type>(),
+            &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::<Int64Type>(),
+            &Int64Array::from(vec![Some(3), None, Some(5), None])
+        );
+        Ok(())
+    }
+
     #[test]
     fn adapter_preserving_evaluation_uses_accumulator_contract() -> Result<()> {
         let mut accumulator = GroupsAccumulatorAdapter::new(|| {

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