Support vectorized append and compare for multi group by

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

@@ -73,19 +73,26 @@ pub struct GroupValuesColumn {
 
     /// Random state for creating hashes
     random_state: RandomState,
+
+    column_nullables_buffer: Vec<bool>,
+
+    append_rows_buffer: Vec<usize>,
 }
 
 impl GroupValuesColumn {
     /// Create a new instance of GroupValuesColumn if supported for the specified schema
     pub fn try_new(schema: SchemaRef) -> Result<Self> {
         let map = RawTable::with_capacity(0);
+        let num_cols = schema.fields.len();
         Ok(Self {
             schema,
             map,
             map_size: 0,
             group_values: vec![],
             hashes_buffer: Default::default(),
             random_state: Default::default(),
+            column_nullables_buffer: vec![false; num_cols],
+            append_rows_buffer: Vec::new(),
         })
     }
 
@@ -146,6 +153,13 @@ macro_rules! instantiate_primitive {
     };
 }
 
+fn append_col_value<C>(mut core: C, array: &ArrayRef, row: usize)
+where
+    C: FnMut(&ArrayRef, usize),
+{
+    core(array, row);
+}
+
 impl GroupValues for GroupValuesColumn {
     fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
         let n_rows = cols[0].len();
@@ -213,6 +227,14 @@ impl GroupValues for GroupValuesColumn {
         batch_hashes.resize(n_rows, 0);
         create_hashes(cols, &self.random_state, batch_hashes)?;
 
+        // 1.2 Check if columns nullable
+        for (col_idx, col) in cols.iter().enumerate() {
+            self.column_nullables_buffer[col_idx] = (col.null_count() != 0);
+        }
+
+        // 1.3 Check and record which rows of the input should be appended
+        self.append_rows_buffer.clear();
+        let mut current_group_idx = self.group_values[0].len();
         for (row, &target_hash) in batch_hashes.iter().enumerate() {
             let entry = self.map.get_mut(target_hash, |(exist_hash, group_idx)| {
                 // Somewhat surprisingly, this closure can be called even if the
@@ -249,31 +271,38 @@ impl GroupValues for GroupValuesColumn {
                     // Add new entry to aggr_state and save newly created index
                     // let group_idx = group_values.num_rows();
                     // group_values.push(group_rows.row(row));
-
-                    let mut checklen = 0;
-                    let group_idx = self.group_values[0].len();
-                    for (i, group_value) in self.group_values.iter_mut().enumerate() {
-                        group_value.append_val(&cols[i], row);
-                        let len = group_value.len();
-                        if i == 0 {
-                            checklen = len;
-                        } else {
-                            debug_assert_eq!(checklen, len);
-                        }
-                    }
+                    let prev_group_idx = current_group_idx;
 
                     // for hasher function, use precomputed hash value
                     self.map.insert_accounted(
-                        (target_hash, group_idx),
+                        (target_hash, prev_group_idx),
                         |(hash, _group_index)| *hash,
                         &mut self.map_size,
                     );
-                    group_idx
+                    self.append_rows_buffer.push(row);
+                    current_group_idx += 1;
+
+                    prev_group_idx
                 }
             };
             groups.push(group_idx);
         }
 
+        // 1.4 Vectorized append values
+        for (col_idx, col) in cols.iter().enumerate() {
+            let col_nullable = self.column_nullables_buffer[col_idx];
+            let group_value = &mut self.group_values[col_idx];
+            if col_nullable {
+                for &row in self.append_rows_buffer.iter() {
+                    group_value.append_val(&cols[col_idx], row);
+                }
+            } else {
+                for &row in self.append_rows_buffer.iter() {
+                    group_value.append_non_nullable_val(&cols[col_idx], row);
+                }
+            }
+        }
+
         Ok(())
     }
 

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

@@ -58,6 +58,9 @@ pub trait GroupColumn: Send + Sync {
     fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool;
     /// Appends the row at `row` in `array` to this builder
     fn append_val(&mut self, array: &ArrayRef, row: usize);
+
+    fn append_non_nullable_val(&mut self, array: &ArrayRef, row: usize);
+
     /// Returns the number of rows stored in this builder
     fn len(&self) -> usize;
     /// Returns the number of bytes used by this [`GroupColumn`]
@@ -113,6 +116,7 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
         self.group_values[lhs_row] == array.as_primitive::<T>().value(rhs_row)
     }
 
+
     fn append_val(&mut self, array: &ArrayRef, row: usize) {
         // Perf: skip null check if input can't have nulls
         if NULLABLE {
@@ -128,6 +132,15 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
         }
     }
 
+    fn append_non_nullable_val(&mut self, array: &ArrayRef, row: usize) {
+        if NULLABLE {
+            self.nulls.append(false);
+            self.group_values.push(array.as_primitive::<T>().value(row));
+        } else {
+            self.group_values.push(array.as_primitive::<T>().value(row));
+        }
+    }
+
     fn len(&self) -> usize {
         self.group_values.len()
     }
@@ -218,6 +231,17 @@ where
         }
     }
 
+    fn append_non_nullable_val_inner<B>(&mut self, array: &ArrayRef, row: usize)
+    where
+        B: ByteArrayType,
+    {
+        let arr = array.as_bytes::<B>();
+        self.nulls.append(false);
+        let value: &[u8] = arr.value(row).as_ref();
+        self.buffer.append_slice(value);
+        self.offsets.push(O::usize_as(self.buffer.len()));
+    }
+
     fn equal_to_inner<B>(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool
     where
         B: ByteArrayType,
@@ -287,6 +311,27 @@ where
         };
     }
 
+    fn append_non_nullable_val(&mut self, column: &ArrayRef, row: usize) {
+        // Sanity array type
+        match self.output_type {
+            OutputType::Binary => {
+                debug_assert!(matches!(
+                    column.data_type(),
+                    DataType::Binary | DataType::LargeBinary
+                ));
+                self.append_non_nullable_val_inner::<GenericBinaryType<O>>(column, row)
+            }
+            OutputType::Utf8 => {
+                debug_assert!(matches!(
+                    column.data_type(),
+                    DataType::Utf8 | DataType::LargeUtf8
+                ));
+                self.append_non_nullable_val_inner::<GenericStringType<O>>(column, row)
+            }
+            _ => unreachable!("View types should use `ArrowBytesViewMap`"),
+        };
+    } 
+
     fn len(&self) -> usize {
         self.offsets.len() - 1
     }
@@ -382,7 +427,7 @@ where
             }
             _ => unreachable!("View types should use `ArrowBytesViewMap`"),
         }
-    }
+    }   
 }
 
 /// An implementation of [`GroupColumn`] for binary view and utf8 view types.
@@ -482,6 +527,35 @@ impl<B: ByteViewType> ByteViewGroupValueBuilder<B> {
         self.views.push(view);
     }
 
+    fn append_val_non_nullable_inner(&mut self, array: &ArrayRef, row: usize)
+    where
+        B: ByteViewType,
+    {
+        let arr = array.as_byte_view::<B>();
+
+        // Not null row case
+        self.nulls.append(false);
+        let value: &[u8] = arr.value(row).as_ref();
+
+        let value_len = value.len();
+        let view = if value_len <= 12 {
+            make_view(value, 0, 0)
+        } else {
+            // Ensure big enough block to hold the value firstly
+            self.ensure_in_progress_big_enough(value_len);
+
+            // Append value
+            let buffer_index = self.completed.len();
+            let offset = self.in_progress.len();
+            self.in_progress.extend_from_slice(value);
+
+            make_view(value, buffer_index as u32, offset as u32)
+        };
+
+        // Append view
+        self.views.push(view);
+    }
+
     fn ensure_in_progress_big_enough(&mut self, value_len: usize) {
         debug_assert!(value_len > 12);
         let require_cap = self.in_progress.len() + value_len;
@@ -776,6 +850,10 @@ impl<B: ByteViewType> GroupColumn for ByteViewGroupValueBuilder<B> {
     fn append_val(&mut self, array: &ArrayRef, row: usize) {
         self.append_val_inner(array, row)
     }
+
+    fn append_non_nullable_val(&mut self, array: &ArrayRef, row: usize) {
+        self.append_val_non_nullable_inner(array, row);
+    }
 
     fn len(&self) -> usize {
         self.views.len()