feat: move split_vec_min_alloc to datafusion-common

Author: ariel-miculas

datafusion/common/src/utils/mod.rs

@@ -41,6 +41,7 @@ use std::num::NonZero;
 use std::ops::Range;
 use std::sync::{Arc, LazyLock};
 use std::thread::available_parallelism;
+use std::mem;
 
 /// Applies an optional projection to a [`SchemaRef`], returning the
 /// projected schema
@@ -569,6 +570,19 @@ pub fn base_type(data_type: &DataType) -> DataType {
     }
 }
 
+/// 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 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)
+    }
+}
+
 /// Information about how to coerce lists.
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
 pub enum ListCoercion {
@@ -1261,6 +1275,49 @@ mod tests {
         assert_eq!(set_difference([3, 4, 0], [4, 1, 2]), vec![3, 0]);
     }
 
+    #[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]);
+    }
+
     #[test]
     fn test_find_indices() -> Result<()> {
         assert_eq!(find_indices(&[0, 3, 4], [0, 3, 4])?, vec![0, 1, 2]);

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,7 +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::{Result, exec_datafusion_err};
+use datafusion_common::{Result, exec_datafusion_err, utils::split_vec_min_alloc};
 use datafusion_physical_expr_common::binary_map::{INITIAL_BUFFER_CAPACITY, OutputType};
 use itertools::izip;
 use std::mem::size_of;

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.
 ///
@@ -1275,50 +1262,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() {

datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs

@@ -16,14 +16,14 @@
 // 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;
 use arrow::array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray, cast::AsArray};
 use arrow::buffer::ScalarBuffer;
 use arrow::datatypes::DataType;
-use datafusion_common::Result;
+use datafusion_common::{Result, utils::split_vec_min_alloc};
 use datafusion_execution::memory_pool::proxy::VecAllocExt;
 use itertools::izip;
 use std::iter;