From 7b716d4b310de7dad349a9dfe2adc246ef77a490 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 22 Sep 2026 08:30:22 -0500 Subject: [PATCH 1/3] perf: reuse GroupsAccumulatorAdapter scratch buffers across batches `invoke_per_accumulator` allocated `groups_with_rows` and `offsets` again for every input batch. Keep them in a `Scratch` struct on the adapter and clear them at the start of each batch. The adapter takes the buffers out of `self` for the batch and puts them back on every return path, errors included, then charges the change in their capacity to `allocation_bytes`. It releases them when an emit leaves no group, so an adapter that has emitted everything holds no memory. `batch_indices` stays a new `Vec` for each batch: `into()` moves it into the `UInt32Array` that `take_arrays` reads, so reusing it would add a copy. Closes https://github.com/apache/datafusion/issues/25604 Co-Authored-By: Claude Opus 5 --- .../src/aggregate/groups_accumulator.rs | 81 ++++++++++++++++++- 1 file changed, 78 insertions(+), 3 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index d49cae6a161c3..ead4a60bf926a 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -111,6 +111,29 @@ pub struct GroupsAccumulatorAdapter { /// Optional aggregate-owned metric timed once for a grouped update batch. grouped_update_metric: OnceLock>>, + + /// Buffers reused by every call to `invoke_per_accumulator` + scratch: Scratch, +} + +/// Per-batch buffers of [`GroupsAccumulatorAdapter`], kept between batches so +/// that each batch does not allocate them again. Their capacity is part of +/// [`GroupsAccumulatorAdapter::allocation_bytes`]. +#[derive(Default)] +struct Scratch { + /// Group indexes that have rows in the batch, in order of first + /// appearance in the batch + groups_with_rows: Vec, + + /// `offsets[i]` is the index into the reordered rows where the rows for + /// `groups_with_rows[i]` start + offsets: Vec, +} + +impl Scratch { + fn allocated_size(&self) -> usize { + self.groups_with_rows.allocated_size() + self.offsets.allocated_size() + } } /// Maximum number of prepared group inputs retained while timing an @@ -154,6 +177,7 @@ impl GroupsAccumulatorAdapter { allocation_bytes: 0, metrics: None, grouped_update_metric: OnceLock::new(), + scratch: Scratch::default(), } } @@ -235,10 +259,43 @@ impl GroupsAccumulatorAdapter { assert_eq!(values[0].len(), group_indices.len()); + // Take the scratch buffers out of `self` for the batch, and put them + // back on every return path, including errors + let mut scratch = std::mem::take(&mut self.scratch); + let scratch_size_pre = scratch.allocated_size(); + let result = self.invoke_per_accumulator_with_scratch( + &mut scratch, + values, + group_indices, + opt_filter, + f, + ); + self.adjust_allocation(scratch_size_pre, scratch.allocated_size()); + self.scratch = scratch; + result + } + + /// Body of [`Self::invoke_per_accumulator`], with its scratch buffers + fn invoke_per_accumulator_with_scratch( + &mut self, + scratch: &mut Scratch, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + f: F, + ) -> Result<()> + where + F: Fn(&mut dyn Accumulator, &[ArrayRef]) -> Result<()>, + { + let Scratch { + groups_with_rows, + offsets, + } = scratch; + // 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![]; + groups_with_rows.clear(); // figure out which input rows correspond to which groups. // Note that self.state.indices starts empty for all groups (it is @@ -267,7 +324,7 @@ impl GroupsAccumulatorAdapter { // offsets[i] is index into batch_indices where the rows for // groups_with_rows[i] start - let mut offsets = Vec::with_capacity(groups_with_rows.len() + 1); + offsets.clear(); offsets.push(0); let mut offset_so_far = 0; @@ -398,6 +455,15 @@ impl GroupsAccumulatorAdapter { self.allocation_bytes = self.allocation_bytes.saturating_sub(size) } + /// Release the scratch buffers when no group is left, so that an adapter + /// that has emitted every group holds no memory + fn free_scratch_if_empty(&mut self) { + if self.states.is_empty() { + let scratch = std::mem::take(&mut self.scratch); + self.free_allocation(scratch.allocated_size()); + } + } + /// Release the allocation held by a state that is being emitted. fn free_state_allocation(&mut self, state: &AccumulatorState) { self.free_allocation(state.size()); @@ -456,6 +522,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { let result = ScalarValue::iter_to_array(results); self.adjust_allocation(vec_size_pre, self.states.allocated_size()); + self.free_scratch_if_empty(); result } @@ -520,6 +587,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { } } self.adjust_allocation(vec_size_pre, self.states.allocated_size()); + self.free_scratch_if_empty(); Ok(arrays) } @@ -820,7 +888,12 @@ mod tests { let retained_indices = adapter.states[0].indices.allocated_size(); assert!(retained_indices > 0); assert!(adapter.states[0].indices.is_empty()); - assert_eq!(adapter.size(), allocation_before_update + retained_indices); + let retained_scratch = adapter.scratch.allocated_size(); + assert!(retained_scratch > 0); + assert_eq!( + adapter.size(), + allocation_before_update + retained_indices + retained_scratch + ); let allocation_after_first_update = adapter.size(); adapter.update_batch(&[values], &[0, 0, 0, 0], None, 1)?; @@ -1085,6 +1158,7 @@ mod tests { .iter() .map(AccumulatorState::size) .sum::() + + accumulator.scratch.allocated_size() ); accumulator.update_batch(&[values], &[0, 1], None, 2)?; @@ -1129,6 +1203,7 @@ mod tests { .iter() .map(AccumulatorState::size) .sum::() + + accumulator.scratch.allocated_size() ); Ok(()) } From e3b5efbd8526bdb3bd07990a2a678836b3a6b69e Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 22 Sep 2026 10:33:31 -0500 Subject: [PATCH 2/3] perf: reuse GroupsAccumulatorAdapter batch_indices buffer across batches Keep batch_indices in the adapter scratch buffers. The buffer moves into the UInt32Array for the take kernels without a copy, and comes back with Buffer::into_vec after the take, because the kernels only borrow it. Co-Authored-By: Claude Opus 5 --- .../src/aggregate/groups_accumulator.rs | 34 +++++++++++++++---- 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index ead4a60bf926a..1c02778e1882c 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -125,14 +125,19 @@ struct Scratch { /// appearance in the batch groups_with_rows: Vec, - /// `offsets[i]` is the index into the reordered rows where the rows for + /// `offsets[i]` is the index into `batch_indices` where the rows for /// `groups_with_rows[i]` start offsets: Vec, + + /// Indices into the batch rows, with the rows of each group contiguous + batch_indices: Vec, } impl Scratch { fn allocated_size(&self) -> usize { - self.groups_with_rows.allocated_size() + self.offsets.allocated_size() + self.groups_with_rows.allocated_size() + + self.offsets.allocated_size() + + self.batch_indices.allocated_size() } } @@ -290,6 +295,7 @@ impl GroupsAccumulatorAdapter { let Scratch { groups_with_rows, offsets, + batch_indices, } = scratch; // groups_with_rows holds a list of group indexes that have any rows @@ -320,7 +326,14 @@ impl GroupsAccumulatorAdapter { self.add_allocation(indices_allocation_delta); // batch_indices holds indices into values, each group is contiguous - let mut batch_indices = Vec::with_capacity(group_indices.len()); + batch_indices.clear(); + batch_indices + .try_reserve(group_indices.len()) + .map_err(|e| { + arrow_datafusion_err!(arrow::error::ArrowError::MemoryError( + e.to_string() + )) + })?; // offsets[i] is index into batch_indices where the rows for // groups_with_rows[i] start @@ -334,13 +347,20 @@ impl GroupsAccumulatorAdapter { offset_so_far += indices.len(); offsets.push(offset_so_far); } - let batch_indices = batch_indices.into(); + // Move the buffer into an array without a copy, for the take kernels + let indices_array = + PrimitiveArray::::from(std::mem::take(batch_indices)); // reorder the values and opt_filter by batch_indices so that // all values for each group are contiguous, then invoke the // accumulator once per group with values - let values = take_arrays(values, &batch_indices, None)?; - let opt_filter = get_filter_at_indices(opt_filter, &batch_indices)?; + let values = take_arrays(values, &indices_array, None)?; + let opt_filter = get_filter_at_indices(opt_filter, &indices_array)?; + + // The take kernels only borrow the indices, so the array holds the + // only reference to the buffer and we can get it back without a copy + let (_, indices_buffer, _) = indices_array.into_parts(); + *batch_indices = indices_buffer.into_inner().into_vec().unwrap_or_default(); let grouped_update_metric = self.grouped_metric(); @@ -888,6 +908,8 @@ mod tests { let retained_indices = adapter.states[0].indices.allocated_size(); assert!(retained_indices > 0); assert!(adapter.states[0].indices.is_empty()); + // The indices buffer comes back from the take kernels without a copy + assert!(adapter.scratch.batch_indices.capacity() >= 4); let retained_scratch = adapter.scratch.allocated_size(); assert!(retained_scratch > 0); assert_eq!( From e61b6b278ee1f7026f13dc877fb8a2e18b2c9f7d Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 22 Sep 2026 10:47:55 -0500 Subject: [PATCH 3/3] perf: shrink oversized GroupsAccumulatorAdapter scratch buffers The scratch buffers keep the capacity of the largest batch that they have seen. One unusually large batch thus made the adapter keep that memory, and charge it to the pool, for all the batches after it. After each batch, shrink a buffer whose capacity is more than 64 Ki entries and more than 4 times what the batch used. It shrinks to the smaller of the two, so it never keeps more than either limit allows. The buffer is cleared first, so the shrink has nothing to copy. Co-Authored-By: Claude Opus 5 --- .../src/aggregate/groups_accumulator.rs | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index 1c02778e1882c..dbf2af512b9f7 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -133,12 +133,41 @@ struct Scratch { batch_indices: Vec, } +/// A scratch buffer keeps its capacity after a batch unless that capacity is +/// more than this many entries and more than [`SCRATCH_RETAIN_RATIO`] times +/// what the batch used. Then it shrinks to the smaller of the two. This stops +/// one unusually large batch from holding memory for all the batches after it. +const MAX_RETAINED_SCRATCH_ENTRIES: usize = 64 * 1024; + +/// See [`MAX_RETAINED_SCRATCH_ENTRIES`] +const SCRATCH_RETAIN_RATIO: usize = 4; + impl Scratch { fn allocated_size(&self) -> usize { self.groups_with_rows.allocated_size() + self.offsets.allocated_size() + self.batch_indices.allocated_size() } + + /// Shrink each buffer that is much larger than the batch that just used + /// it. Call this after the batch, before the buffers are cleared for the + /// next batch. + fn release_oversized(&mut self) { + fn release_if_oversized(buffer: &mut Vec) { + let keep = SCRATCH_RETAIN_RATIO * buffer.len(); + if buffer.capacity() > MAX_RETAINED_SCRATCH_ENTRIES + && buffer.capacity() > keep + { + // The batch is done with the contents, so clear them first + // and the shrink has nothing to copy + buffer.clear(); + buffer.shrink_to(MAX_RETAINED_SCRATCH_ENTRIES.min(keep)); + } + } + release_if_oversized(&mut self.groups_with_rows); + release_if_oversized(&mut self.offsets); + release_if_oversized(&mut self.batch_indices); + } } /// Maximum number of prepared group inputs retained while timing an @@ -275,6 +304,7 @@ impl GroupsAccumulatorAdapter { opt_filter, f, ); + scratch.release_oversized(); self.adjust_allocation(scratch_size_pre, scratch.allocated_size()); self.scratch = scratch; result @@ -1255,6 +1285,49 @@ mod tests { Ok(()) } + /// One large batch must not make the adapter keep its scratch capacity + /// for the smaller batches after it. + #[test] + fn adapter_releases_oversized_scratch() -> Result<()> { + const NUM_GROUPS: usize = 2 * MAX_RETAINED_SCRATCH_ENTRIES; + let mut adapter = GroupsAccumulatorAdapter::new(|| { + Ok(Box::new(MaxAccumulator::try_new(&DataType::Int64)?) + as Box) + }); + + // every row is a different group, so the scratch buffers need one + // entry for each row + let group_indices: Vec = (0..NUM_GROUPS).collect(); + let values: ArrayRef = + Arc::new(Int64Array::from_iter_values(0..NUM_GROUPS as i64)); + adapter.update_batch(&[values], &group_indices, None, NUM_GROUPS)?; + let large_scratch = adapter.scratch.allocated_size(); + assert!(large_scratch >= NUM_GROUPS * size_of::()); + let size_after_large_batch = adapter.size(); + + // a batch with one row shrinks each large buffer to + // SCRATCH_RETAIN_RATIO times what it used: one group, two offsets and + // one row index + let values: ArrayRef = Arc::new(Int64Array::from(vec![1])); + adapter.update_batch(&[values], &[0], None, NUM_GROUPS)?; + let small_scratch = adapter.scratch.allocated_size(); + assert!(small_scratch < large_scratch); + assert_eq!( + adapter.scratch.groups_with_rows.capacity(), + SCRATCH_RETAIN_RATIO + ); + assert_eq!(adapter.scratch.offsets.capacity(), 2 * SCRATCH_RETAIN_RATIO); + assert_eq!( + adapter.scratch.batch_indices.capacity(), + SCRATCH_RETAIN_RATIO + ); + assert_eq!( + adapter.size(), + size_after_large_batch - large_scratch + small_scratch + ); + Ok(()) + } + #[test] fn adapter_preserving_evaluation_uses_accumulator_contract() -> Result<()> { let mut accumulator = GroupsAccumulatorAdapter::new(|| {