new poc1.

Author: Rachelint

datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs

@@ -69,6 +69,10 @@ pub struct NullState<O: GroupIndexOperations> {
     /// pass the filter yet for group `i`
     seen_values: Blocks<BooleanBufferBuilderWrapper>,
 
+    block_indices_buffer: Vec<(u32, u32)>,
+
+    row_selection: Vec<bool>,
+
     /// phantom data for required type `<O>`
     _phantom: PhantomData<O>,
 }
@@ -120,6 +124,53 @@ impl<O: GroupIndexOperations> NullState<O> {
         });
     }
 
+    pub fn accumulate_blocks<T, F>(
+        &mut self,
+        group_indices: &[usize],
+        values: &PrimitiveArray<T>,
+        opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+        mut value_fn: F,
+    ) where
+        T: ArrowPrimitiveType + Send,
+        F: FnMut(u32, u64, T::Native) + Send,
+    {
+        // ensure the seen_values is big enough (start everything at
+        // "not seen" valid)
+        self.seen_values.resize(total_num_groups, false);
+        let seen_values = &mut self.seen_values;
+
+        // Parse indices
+        self.block_indices_buffer.clear();
+        for packed_index in group_indices.iter() {
+            let packed_index = *packed_index as u64;
+            let block_id = (packed_index >> 32) as u32;
+            let block_offset = packed_index as u32;
+            self.block_indices_buffer.push((block_id, block_offset));
+        }
+
+        // Get row selection
+        accumulate_row_selection(values, opt_filter, &mut self.row_selection);
+
+        // Accumulate values
+        let data = values.values().iter();
+        if self.row_selection.is_empty() {
+            let iter = self.block_indices_buffer.iter().zip(data);
+            for (&(block_id, block_offset), &value) in iter {
+                seen_values.set_bit(block_id, block_offset as u64, true);
+                value_fn(block_id, block_offset as u64, value);
+            }
+        } else {
+            let iter = self.block_indices_buffer.iter().zip(data);
+            for (row_idx, (&(block_id, block_offset), &value)) in iter.enumerate() {
+                if self.row_selection[row_idx] {
+                    seen_values.set_bit(block_id, block_offset as u64, true);
+                    value_fn(block_id, block_offset as u64, value);
+                }
+            }
+        }
+    }
+
     /// Invokes `value_fn(group_index, value)` for each non null, non
     /// filtered value in `values`, while tracking which groups have
     /// seen null inputs and which groups have seen any inputs, for
@@ -266,7 +317,7 @@ impl NullStateAdapter {
                 total_num_groups,
                 value_fn,
             ),
-            NullStateAdapter::Blocked(null_state) => null_state.accumulate(
+            NullStateAdapter::Blocked(null_state) => null_state.accumulate_blocks(
                 group_indices,
                 values,
                 opt_filter,
@@ -395,6 +446,8 @@ impl Default for FlatNullState {
     fn default() -> Self {
         Self {
             seen_values: Blocks::new(None),
+            block_indices_buffer: Vec::new(),
+            row_selection: Vec::new(),
             _phantom: PhantomData {},
         }
     }
@@ -427,6 +480,8 @@ impl BlockedNullState {
     pub fn new(block_size: usize) -> Self {
         Self {
             seen_values: Blocks::new(Some(block_size)),
+            block_indices_buffer: Vec::new(),
+            row_selection: Vec::new(),
             _phantom: PhantomData {},
         }
     }
@@ -648,7 +703,7 @@ pub fn accumulate<T, F>(
     }
 }
 
-pub fn accumulate_row_selection<T, F>(
+pub fn accumulate_row_selection<T>(
     values: &PrimitiveArray<T>,
     opt_filter: Option<&BooleanArray>,
     results: &mut Vec<bool>,