perf: Optimize multi_group_by when there are a lot of unique groups

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

@@ -22,7 +22,7 @@ use arrow::array::{
     GenericBinaryArray, GenericByteArray, GenericStringArray, OffsetSizeTrait,
 };
 use arrow::buffer::{OffsetBuffer, ScalarBuffer};
-use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType};
+use arrow::datatypes::{ArrowNativeType, ByteArrayType, DataType, GenericBinaryType};
 use datafusion_common::utils::proxy::VecAllocExt;
 use datafusion_common::{DataFusionError, Result};
 use datafusion_physical_expr_common::binary_map::{OutputType, INITIAL_BUFFER_CAPACITY};
@@ -171,6 +171,74 @@ where
         Ok(())
     }
 
+    fn append_array_slice_inner<B>(
+        &mut self,
+        array: &ArrayRef,
+        start: usize,
+        length: usize,
+    ) -> Result<()>
+    where
+        B: ByteArrayType,
+    {
+        let array = array.as_bytes::<B>();
+
+        let offsets = array.offsets();
+        let bytes = array.value_data();
+
+        // Is the slice all nulls
+        let all_nulls: bool;
+
+        // 1. Append the nulls
+        if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
+            let nulls_slice = nulls.slice(start, length);
+            all_nulls = nulls_slice.null_count() == length;
+            self.nulls.append_buffer(&nulls_slice);
+        } else {
+            all_nulls = false;
+            self.nulls.append_n(length, false);
+        }
+
+        let values_start = offsets[start].as_usize();
+        let values_end = offsets[start + length].as_usize();
+
+        // 2. Append the offsets
+
+        // Must be before adding the bytes to the buffer
+        let mut last_offset = self.buffer.len();
+
+        if all_nulls {
+            // If all nulls, we can just repeat the last offset
+            for _ in 0..length {
+                self.offsets.push(O::usize_as(last_offset));
+            }
+        } else {
+            for start_and_end_values in offsets[start..=start + length].windows(2) {
+                let length = start_and_end_values[1] - start_and_end_values[0];
+                last_offset += length.as_usize();
+                self.offsets.push(O::usize_as(last_offset));
+            }
+        }
+
+        // 3. Append the bytes
+
+        // Only if not all nulls append the actual bytes
+        if !all_nulls {
+            // Note: if the array have nulls we might copy some bytes that are not used.
+
+            // Add all the bytes for the values directly to the byte buffer
+            self.buffer.append_slice(&bytes[values_start..values_end]);
+
+            if self.buffer.len() > self.max_buffer_size {
+                return Err(DataFusionError::Execution(format!(
+                    "offset overflow, buffer size > {}",
+                    self.max_buffer_size
+                )));
+            }
+        }
+
+        Ok(())
+    }
+
     fn do_equal_to_inner<B>(
         &self,
         lhs_row: usize,
@@ -327,6 +395,37 @@ where
         Ok(())
     }
 
+    fn append_array_slice(
+        &mut self,
+        column: &ArrayRef,
+        start: usize,
+        length: usize,
+    ) -> Result<()> {
+        match self.output_type {
+            OutputType::Binary => {
+                debug_assert!(matches!(
+                    column.data_type(),
+                    DataType::Binary | DataType::LargeBinary
+                ));
+                self.append_array_slice_inner::<GenericBinaryType<O>>(
+                    column, start, length,
+                )?
+            }
+            OutputType::Utf8 => {
+                debug_assert!(matches!(
+                    column.data_type(),
+                    DataType::Utf8 | DataType::LargeUtf8
+                ));
+                self.append_array_slice_inner::<GenericStringType<O>>(
+                    column, start, length,
+                )?
+            }
+            _ => unreachable!("View types should use `ArrowBytesViewMap`"),
+        };
+
+        Ok(())
+    }
+
     fn len(&self) -> usize {
         self.offsets.len() - 1
     }

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

@@ -88,6 +88,24 @@ pub trait GroupColumn: Send + Sync {
     /// The vectorized version `append_val`
     fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()>;
 
+    /// Append slice of values from `array`, starting at `start` for `length` rows
+    ///
+    /// This is a special case of `vectorized_append` when the rows are continuous
+    ///
+    /// You should implement this to optimize large copies of contiguous values.
+    ///
+    /// This does not get the sliced array even though it would be more user-friendly
+    /// to allow optimization that avoid the additional computation that can happen in a slice
+    fn append_array_slice(
+        &mut self,
+        array: &ArrayRef,
+        start: usize,
+        length: usize,
+    ) -> Result<()> {
+        let rows = (start..start + length).collect::<Vec<_>>();
+        self.vectorized_append(array, &rows)
+    }
+
     /// Returns the number of rows stored in this builder
     fn len(&self) -> usize;
 
@@ -233,6 +251,16 @@ struct VectorizedOperationBuffers {
     /// The `vectorized append` row indices buffer
     append_row_indices: Vec<usize>,
 
+    /// The last row index in `append_row_indices`
+    ///
+    /// Will be None if `append_row_indices` is empty
+    last_append_row_index: Option<usize>,
+
+    /// If all the values in `append_row_indices` are continuous
+    /// i.e. `append_row_indices[i] + 1 == append_row_indices[i + 1]`
+    /// this is used to optimize the `vectorized_append` operation
+    are_row_indices_continuous: bool,
+
     /// The `vectorized_equal_to` row indices buffer
     equal_to_row_indices: Vec<usize>,
 
@@ -250,12 +278,28 @@ struct VectorizedOperationBuffers {
 
 impl VectorizedOperationBuffers {
     fn clear(&mut self) {
-        self.append_row_indices.clear();
+        self.clear_append_row_indices();
         self.equal_to_row_indices.clear();
         self.equal_to_group_indices.clear();
         self.equal_to_results.clear();
         self.remaining_row_indices.clear();
     }
+
+    fn add_append_row_index(&mut self, row: usize) {
+        self.are_row_indices_continuous = self.are_row_indices_continuous
+            && self
+                .last_append_row_index
+                .is_none_or(|last| last + 1 == row);
+        self.last_append_row_index = Some(row);
+
+        self.append_row_indices.push(row);
+    }
+
+    fn clear_append_row_indices(&mut self) {
+        self.append_row_indices.clear();
+        self.are_row_indices_continuous = true;
+        self.last_append_row_index = None;
+    }
 }
 
 impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
@@ -498,7 +542,7 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
         batch_hashes: &[u64],
         groups: &mut [usize],
     ) {
-        self.vectorized_operation_buffers.append_row_indices.clear();
+        self.vectorized_operation_buffers.clear_append_row_indices();
         self.vectorized_operation_buffers
             .equal_to_row_indices
             .clear();
@@ -528,9 +572,7 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
                 );
 
                 // Add row index to `vectorized_append_row_indices`
-                self.vectorized_operation_buffers
-                    .append_row_indices
-                    .push(row);
+                self.vectorized_operation_buffers.add_append_row_index(row);
 
                 // Set group index to row in `groups`
                 groups[row] = current_group_idx;
@@ -578,11 +620,24 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
         }
 
         let iter = self.group_values.iter_mut().zip(cols.iter());
-        for (group_column, col) in iter {
-            group_column.vectorized_append(
-                col,
-                &self.vectorized_operation_buffers.append_row_indices,
-            )?;
+        if self.vectorized_operation_buffers.are_row_indices_continuous
+            && !self
+                .vectorized_operation_buffers
+                .append_row_indices
+                .is_empty()
+        {
+            let start = self.vectorized_operation_buffers.append_row_indices[0];
+            let length = self.vectorized_operation_buffers.append_row_indices.len();
+            for (group_column, col) in iter {
+                group_column.append_array_slice(col, start, length)?;
+            }
+        } else {
+            for (group_column, col) in iter {
+                group_column.vectorized_append(
+                    col,
+                    &self.vectorized_operation_buffers.append_row_indices,
+                )?;
+            }
         }
 
         Ok(())

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

@@ -175,6 +175,35 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
         Ok(())
     }
 
+    fn append_array_slice(
+        &mut self,
+        array: &ArrayRef,
+        start: usize,
+        length: usize,
+    ) -> Result<()> {
+        let array = array.as_primitive::<T>();
+
+        if NULLABLE {
+            if let Some(nulls) = array.nulls().filter(|n| n.null_count() > 0) {
+                self.nulls.append_buffer(&nulls.slice(start, length));
+            } else {
+                self.nulls.append_n(length, false);
+            }
+        } else {
+            assert_eq!(
+                array.null_count(),
+                0,
+                "unexpected nulls in non nullable input"
+            );
+            self.nulls.append_n(length, false);
+        }
+
+        self.group_values
+            .extend_from_slice(&array.values()[start..start + length]);
+
+        Ok(())
+    }
+
     fn len(&self) -> usize {
         self.group_values.len()
     }

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

@@ -61,6 +61,13 @@ impl MaybeNullBufferBuilder {
         }
     }
 
+    /// Append [`NullBuffer`] to this [`NullBufferBuilder`]
+    ///
+    /// This is useful when you want to concatenate two null buffers.
+    pub fn append_buffer(&mut self, other: &NullBuffer) {
+        self.nulls.append_buffer(other);
+    }
+
     /// return the number of heap allocated bytes used by this structure to store boolean values
     pub fn allocated_size(&self) -> usize {
         // NullBufferBuilder builder::allocated_size returns capacity in bits