This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22416-7a6b0626da7a733d94821b8d9d6011eb85d96593 in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0add0469eb87dbb7fd9b6f78ded41ca7c8b46b6f Author: RyanStewart <[email protected]> AuthorDate: Wed May 27 11:05:35 2026 -0700 perf: hoist split_vec_min_alloc to datafusion-common and shrink the emitted prefix (#22416) ## Which issue does this PR close? Related to #22164 and #22165. ## Rationale for this change #22165 added a `split_vec_min_alloc` helper so that `EmitTo::First(n)` allocates `min(n, len - n)` instead of always copying `n` elements. This PR finishes and corrects that work, following review: - The helper was `pub(super)` inside `datafusion-physical-plan`, so it could not be reused. It moves to `datafusion_common::utils` as the single shared copy. - The `split_off` branch returned the original allocation as the emitted prefix: short length, but original (larger) capacity. datafusion accounts memory by capacity rather than length, so that prefix did not actually release memory under pressure. The helper now calls `shrink_to_fit` on it. - Two further `EmitTo::First(n)` paths still used the `drain(..n).collect()` idiom: the min/max accumulators and `ByteViewGroupValueBuilder::take_n`. Both now route through the shared helper. ## What changes are included in this PR? - `datafusion_common::utils::split_vec_min_alloc`: the shared helper, with `shrink_to_fit` on the emitted prefix. - `datafusion-physical-plan`: the duplicate helper is removed; `bytes.rs`, `primitive.rs` and `bytes_view.rs` use the shared one. - `datafusion-functions-aggregate`: `MinMaxStructAccumulator` and `MinMaxBytesAccumulator` use it in `emit_to`. This supersedes #22205, whose `ByteViewGroupValueBuilder::take_n` change is included here. ## Are these changes tested? Yes. Unit tests for the helper (both branches and the boundaries) plus the existing `take_n` and `min_max` tests pass. ## Are there any user-facing changes? No. --------- Co-authored-by: RyanJamesStewart <[email protected]> Co-authored-by: Claude Opus 4.7 <[email protected]> --- datafusion/common/src/utils/mod.rs | 131 +++++++++++++++++++++ .../src/min_max/min_max_bytes.rs | 4 +- .../src/min_max/min_max_struct.rs | 4 +- .../group_values/multi_group_by/bytes.rs | 3 +- .../group_values/multi_group_by/bytes_view.rs | 3 +- .../aggregates/group_values/multi_group_by/mod.rs | 58 +-------- .../group_values/multi_group_by/primitive.rs | 3 +- 7 files changed, 144 insertions(+), 62 deletions(-) diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 0c667b17c3..99c5bdbc54 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -395,6 +395,137 @@ pub fn longest_consecutive_prefix<T: Borrow<usize>>( count } +/// Splits `vec` at index `n`, returning the first `n` elements and leaving the +/// remaining `vec.len() - n` elements in `vec`. +/// +/// Allocates for whichever side is smaller, so the new allocation is +/// `min(n, vec.len() - n)` rather than always `n` (as `vec.drain(0..n).collect()` +/// would). This matters when the split emits a prefix under memory pressure, +/// where `n` can be close to `vec.len()`. +pub fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize) -> Vec<T> { + if n * 2 <= vec.len() { + vec.drain(0..n).collect() + } else { + let remaining = vec.split_off(n); + std::mem::replace(vec, remaining) + } +} + +#[cfg(test)] +mod split_vec_min_alloc_tests { + use super::split_vec_min_alloc; + + #[test] + fn drain_branch() { + // n * 2 <= len -> drain+collect branch (allocates n elements) + let mut v = vec![1, 2, 3, 4, 5, 6]; + let first = split_vec_min_alloc(&mut v, 2); + assert_eq!(first, vec![1, 2]); + assert_eq!(v, vec![3, 4, 5, 6]); + } + + #[test] + fn split_off_branch() { + // remaining < n -> split_off+replace branch (allocates remaining elements) + let mut v = vec![1, 2, 3, 4, 5, 6]; + let first = split_vec_min_alloc(&mut v, 4); + assert_eq!(first, vec![1, 2, 3, 4]); + assert_eq!(v, vec![5, 6]); + } + + #[test] + fn exactly_half() { + // n * 2 == len -> drain branch (boundary) + let mut v = vec![1, 2, 3, 4]; + let first = split_vec_min_alloc(&mut v, 2); + assert_eq!(first, vec![1, 2]); + assert_eq!(v, vec![3, 4]); + } + + #[test] + fn take_all() { + let mut v = vec![1, 2, 3]; + let first = split_vec_min_alloc(&mut v, 3); + assert_eq!(first, vec![1, 2, 3]); + assert!(v.is_empty()); + } + + #[test] + fn take_none() { + let mut v = vec![1, 2, 3]; + let first = split_vec_min_alloc(&mut v, 0); + assert!(first.is_empty()); + assert_eq!(v, vec![1, 2, 3]); + } + + #[test] + fn emitted_prefix_does_not_realloc_on_push() { + // Demonstrates *why* the split-off branch must NOT call `shrink_to_fit`. + // + // Downstream callers (e.g. `multi_group_by/bytes.rs`, which does + // `first_n_offsets.push(offset_n)` right after the split) push onto the + // emitted prefix immediately. The split-off branch hands the original + // backing allocation to that prefix, so the prefix already has spare + // capacity for the very next push. + // + // If we shrank the prefix to fit, that next push would have to + // reallocate, and Vec's growth strategy would land it at a *larger* + // capacity than the original allocation we started with -- the opposite + // of the memory saving `shrink_to_fit` was meant to deliver. + + // A Vec with a known, deliberately large capacity. n*2 > len, so this + // takes the split-off branch. + let mut v: Vec<u32> = Vec::with_capacity(64); + v.extend(0..10); + let original_capacity = v.capacity(); + assert!(original_capacity >= 64); + + // Emit a prefix that is most of the Vec (n = 8, remaining = 2). + let mut prefix = split_vec_min_alloc(&mut v, 8); + assert_eq!(prefix, vec![0, 1, 2, 3, 4, 5, 6, 7]); + + // The split-off branch moved the original backing store into `prefix`, + // so it keeps the original (large) capacity -- no shrink happened. + assert_eq!( + prefix.capacity(), + original_capacity, + "split-off branch must hand the original allocation to the prefix" + ); + + // The caller's very next operation: push one element onto the prefix. + prefix.push(99); + + // Because the capacity was preserved, the push reused the existing + // allocation: post-push capacity is unchanged and still <= original. + // This is the realloc that `shrink_to_fit` would have forced. + assert_eq!( + prefix.capacity(), + original_capacity, + "push must reuse the preserved allocation (no realloc)" + ); + assert!(prefix.capacity() <= original_capacity); + + // Counter-demonstration: had we shrunk the prefix to fit (capacity 8), + // the same push would have reallocated. Vec doubles on growth, so the + // post-push capacity (16) ends up LARGER than where a length-8 prefix + // started -- and we paid a realloc for it. + let mut shrunk: Vec<u32> = prefix[..8].to_vec(); + shrunk.shrink_to_fit(); + let shrunk_capacity = shrink_then_push_capacity(&mut shrunk); + assert!( + shrunk_capacity > 8, + "shrink-to-fit then push reallocates to a larger capacity" + ); + } + + /// Helper for the counter-demonstration above: push one element and report + /// the resulting capacity. + fn shrink_then_push_capacity(v: &mut Vec<u32>) -> usize { + v.push(99); + v.capacity() + } +} + /// Creates single element [`ListArray`], [`LargeListArray`] and /// [`FixedSizeListArray`] from other arrays /// diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index e4ac7eccf5..b56c2106e3 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -27,6 +27,8 @@ use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls: use std::mem::size_of; use std::sync::Arc; +use datafusion_common::utils::split_vec_min_alloc; + /// Implements fast Min/Max [`GroupsAccumulator`] for "bytes" types ([`StringArray`], /// [`BinaryArray`], [`StringViewArray`], etc) /// @@ -493,7 +495,7 @@ impl MinMaxBytesState { ) } EmitTo::First(n) => { - let first_min_maxes: Vec<_> = self.min_max.drain(..n).collect(); + let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); let first_data_capacity: usize = first_min_maxes .iter() .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) diff --git a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs index 796fd586ca..7c94e7f573 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs @@ -30,6 +30,8 @@ use datafusion_common::{ use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls; +use datafusion_common::utils::split_vec_min_alloc; + /// Accumulator for MIN/MAX operations on Struct data types. /// /// This accumulator tracks the minimum or maximum struct value encountered @@ -282,7 +284,7 @@ impl MinMaxStructState { ) } EmitTo::First(n) => { - let first_min_maxes: Vec<_> = self.min_max.drain(..n).collect(); + let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); let first_data_capacity: usize = first_min_maxes .iter() .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index 350ec13712..c83b1da404 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -16,7 +16,7 @@ // under the License. use crate::aggregates::group_values::multi_group_by::{ - GroupColumn, Nulls, nulls_equal_to, split_vec_min_alloc, + GroupColumn, Nulls, nulls_equal_to, }; use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::{ @@ -26,6 +26,7 @@ use arrow::array::{ use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType}; use datafusion_common::utils::proxy::VecAllocExt; +use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{Result, exec_datafusion_err}; use datafusion_physical_expr_common::binary_map::{INITIAL_BUFFER_CAPACITY, OutputType}; use std::mem::size_of; diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs index 9267cf4f27..e94e4547e1 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs @@ -26,6 +26,7 @@ use arrow::array::{ use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::ByteViewType; use datafusion_common::Result; +use datafusion_common::utils::split_vec_min_alloc; use std::marker::PhantomData; use std::mem::{replace, size_of}; use std::sync::Arc; @@ -363,7 +364,7 @@ impl<B: ByteViewType> ByteViewGroupValueBuilder<B> { // // - Shift the `buffer index` of remaining non-inlined `views` // - let first_n_views = self.views.drain(0..n).collect::<Vec<_>>(); + let first_n_views = split_vec_min_alloc(&mut self.views, n); let last_non_inlined_view = first_n_views .iter() diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index cf2d4f49ae..ee2d300d9b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -107,19 +107,6 @@ pub trait GroupColumn: Send + Sync { fn take_n(&mut self, n: usize) -> ArrayRef; } -/// Splits `vec` at `n`, returning the first `n` elements and leaving the -/// remainder in `vec`. Allocates for whichever portion is smaller to minimize -/// peak memory: `drain+collect` when `n <= remaining`, `split_off+replace` -/// when `remaining < n`. -pub(super) fn split_vec_min_alloc<T>(vec: &mut Vec<T>, n: usize) -> Vec<T> { - if n * 2 <= vec.len() { - vec.drain(0..n).collect() - } else { - let remaining = vec.split_off(n); - mem::replace(vec, remaining) - } -} - /// Determines if the nullability of the existing and new input array can be used /// to short-circuit the comparison of the two values. /// @@ -1285,50 +1272,7 @@ mod tests { GroupValues, multi_group_by::GroupValuesColumn, }; - use super::{GroupIndexView, split_vec_min_alloc}; - - #[test] - fn test_split_vec_min_alloc_drain_branch() { - // n * 2 <= len → drain+collect branch (allocates n elements) - let mut v = vec![1, 2, 3, 4, 5, 6]; - let first = split_vec_min_alloc(&mut v, 2); - assert_eq!(first, vec![1, 2]); - assert_eq!(v, vec![3, 4, 5, 6]); - } - - #[test] - fn test_split_vec_min_alloc_split_off_branch() { - // remaining < n → split_off+replace branch (allocates remaining elements) - let mut v = vec![1, 2, 3, 4, 5, 6]; - let first = split_vec_min_alloc(&mut v, 4); - assert_eq!(first, vec![1, 2, 3, 4]); - assert_eq!(v, vec![5, 6]); - } - - #[test] - fn test_split_vec_min_alloc_exactly_half() { - // n * 2 == len → drain branch (boundary condition) - let mut v = vec![1, 2, 3, 4]; - let first = split_vec_min_alloc(&mut v, 2); - assert_eq!(first, vec![1, 2]); - assert_eq!(v, vec![3, 4]); - } - - #[test] - fn test_split_vec_min_alloc_take_all() { - let mut v = vec![1, 2, 3]; - let first = split_vec_min_alloc(&mut v, 3); - assert_eq!(first, vec![1, 2, 3]); - assert!(v.is_empty()); - } - - #[test] - fn test_split_vec_min_alloc_take_none() { - let mut v = vec![1, 2, 3]; - let first = split_vec_min_alloc(&mut v, 0); - assert!(first.is_empty()); - assert_eq!(v, vec![1, 2, 3]); - } + use super::GroupIndexView; #[test] fn test_intern_for_vectorized_group_values() { diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index 4aae996f68..1913371845 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -16,7 +16,7 @@ // under the License. use crate::aggregates::group_values::multi_group_by::{ - GroupColumn, Nulls, nulls_equal_to, split_vec_min_alloc, + GroupColumn, Nulls, nulls_equal_to, }; use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::ArrowNativeTypeOp; @@ -28,6 +28,7 @@ use arrow::buffer::ScalarBuffer; use arrow::datatypes::DataType; use arrow::util::bit_util::apply_bitwise_binary_op; use datafusion_common::Result; +use datafusion_common::utils::split_vec_min_alloc; use datafusion_execution::memory_pool::proxy::VecAllocExt; use std::iter; use std::sync::Arc; --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
