feat: shrink_to_fit in multi-group-by

Author: ariel-miculas

datafusion/expr-common/src/groups_accumulator.rs

@@ -18,7 +18,7 @@
 //! Vectorized [`GroupsAccumulator`]
 
 use arrow::array::{ArrayRef, BooleanArray};
-use datafusion_common::{Result, not_impl_err};
+use datafusion_common::{Result, not_impl_err, utils::split_vec_min_alloc};
 
 /// Describes how many rows should be emitted during grouping.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -45,13 +45,7 @@ impl EmitTo {
                 // Take the entire vector, leave new (empty) vector
                 std::mem::take(v)
             }
-            Self::First(n) => {
-                // get end n+1,.. values into t
-                let mut t = v.split_off(*n);
-                // leave n+1,.. in v
-                std::mem::swap(v, &mut t);
-                t
-            }
+            Self::First(n) => split_vec_min_alloc(v, *n),
         }
     }
 }

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

@@ -188,8 +188,10 @@ impl<const NULLABLE: bool> GroupColumn for BooleanGroupValueBuilder<NULLABLE> {
 
         // take only first n values from the original builder
         new_builder.truncate(n);
+        let mut buffer = new_builder.finish();
+        buffer.shrink_to_fit();
 
-        Arc::new(BooleanArray::new(new_builder.finish(), first_n_nulls))
+        Arc::new(BooleanArray::new(buffer, first_n_nulls))
     }
 }
 

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

@@ -401,7 +401,8 @@ where
         // Find out a way to avoid copying buffer but split the original one into two.
         remaining_buffer.append_slice(&self.buffer.as_slice()[first_remaining_offset..]);
         self.buffer.truncate(first_remaining_offset);
-        let values = self.buffer.finish();
+        let mut values = self.buffer.finish();
+        values.shrink_to_fit();
         self.buffer = remaining_buffer;
 
         match self.output_type {

datafusion/physical-plan/src/aggregates/group_values/null_builder.rs

@@ -87,7 +87,11 @@ impl MaybeNullBufferBuilder {
 
         // take only first n values from the original builder
         new_builder.truncate(n);
-        new_builder.finish()
+        let mut buffer = new_builder.finish();
+        if let Some(b) = buffer.as_mut() {
+            b.shrink_to_fit()
+        }
+        buffer
     }
 
     /// Returns true if this builder might have any nulls