support blocked mode for PrimitiveGroupsAccumulator.

Author: Rachelint

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

@@ -20,7 +20,7 @@
 //! [`GroupsAccumulator`]: datafusion_expr_common::groups_accumulator::GroupsAccumulator
 
 use std::collections::VecDeque;
-use std::fmt::{self, Debug};
+use std::fmt::Debug;
 use std::marker::PhantomData;
 
 use arrow::array::{Array, BooleanArray, BooleanBufferBuilder, PrimitiveArray};
@@ -401,6 +401,7 @@ impl SeenValues for BlockedSeenValues {
 
 /// Adapter for supporting dynamic dispatching of [`FlatNullState`] and [`BlockedNullState`].
 /// For performance, the cost of batch-level dynamic dispatching is acceptable.
+#[derive(Debug)]
 pub enum NullStateAdapter {
     Flat(FlatNullState),
     Blocked(BlockedNullState),
@@ -479,6 +480,13 @@ impl NullStateAdapter {
         }
     }
 
+    pub fn size(&self) -> usize {
+        match self {
+            NullStateAdapter::Flat(null_state) => null_state.size(),
+            NullStateAdapter::Blocked(null_state) => null_state.size(),
+        }
+    }
+
     /// Clone and build a single [`BooleanBuffer`] from `seen_values`,
     /// only used for testing.
     #[cfg(test)]

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

@@ -15,18 +15,21 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::collections::VecDeque;
+use std::iter;
 use std::mem::size_of;
 use std::sync::Arc;
 
-use arrow::array::{ArrayRef, AsArray, BooleanArray, PrimitiveArray};
+use arrow::array::{ArrayRef, ArrowNativeTypeOp, AsArray, BooleanArray, PrimitiveArray};
 use arrow::buffer::NullBuffer;
 use arrow::compute;
 use arrow::datatypes::ArrowPrimitiveType;
 use arrow::datatypes::DataType;
 use datafusion_common::{internal_datafusion_err, DataFusionError, Result};
 use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
 
-use super::accumulate::FlatNullState;
+use crate::aggregate::groups_accumulator::accumulate::NullStateAdapter;
+use crate::aggregate::groups_accumulator::{ensure_room_enough_for_blocks, Block};
 
 /// An accumulator that implements a single operation over
 /// [`ArrowPrimitiveType`] where the accumulated state is the same as
@@ -44,7 +47,7 @@ where
     F: Fn(&mut T::Native, T::Native) + Send + Sync,
 {
     /// values per group, stored as the native type
-    values: Vec<T::Native>,
+    values: VecDeque<Vec<T::Native>>,
 
     /// The output type (needed for Decimal precision and scale)
     data_type: DataType,
@@ -53,10 +56,12 @@ where
     starting_value: T::Native,
 
     /// Track nulls in the input / filters
-    null_state: FlatNullState,
+    null_state: NullStateAdapter,
 
     /// Function that computes the primitive result
     prim_fn: F,
+
+    block_size: Option<usize>,
 }
 
 impl<T, F> PrimitiveGroupsAccumulator<T, F>
@@ -66,11 +71,12 @@ where
 {
     pub fn new(data_type: &DataType, prim_fn: F) -> Self {
         Self {
-            values: vec![],
+            values: VecDeque::new(),
             data_type: data_type.clone(),
-            null_state: FlatNullState::new(),
+            null_state: NullStateAdapter::new(None),
             starting_value: T::default_value(),
             prim_fn,
+            block_size: None,
         }
     }
 
@@ -97,16 +103,34 @@ where
         let values = values[0].as_primitive::<T>();
 
         // update values
-        self.values.resize(total_num_groups, self.starting_value);
+        if let Some(blk_size) = self.block_size {
+            let new_block = |block_size: usize| Vec::with_capacity(block_size);
+            ensure_room_enough_for_blocks(
+                &mut self.values,
+                total_num_groups,
+                blk_size,
+                new_block,
+                self.starting_value,
+            );
+        } else {
+            if self.values.is_empty() {
+                self.values.push_back(Vec::new());
+            }
+
+            self.values
+                .back_mut()
+                .unwrap()
+                .resize(total_num_groups, self.starting_value);
+        }
 
         // NullState dispatches / handles tracking nulls and groups that saw no values
         self.null_state.accumulate(
             group_indices,
             values,
             opt_filter,
             total_num_groups,
-            |_, group_index, new_value| {
-                let value = &mut self.values[group_index as usize];
+            |block_id, block_offset, new_value| {
+                let value = &mut self.values[block_id as usize][block_offset as usize];
                 (self.prim_fn)(value, new_value);
             },
         );
@@ -115,7 +139,7 @@ where
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let values = emit_to.take_needed_rows(&mut self.values);
+        let values = emit_to.take_needed(&mut self.values, self.block_size.is_some());
         let nulls = self.null_state.build(emit_to);
         let values = PrimitiveArray::<T>::new(values.into(), Some(nulls)) // no copy
             .with_data_type(self.data_type.clone());
@@ -198,4 +222,28 @@ where
     fn size(&self) -> usize {
         self.values.capacity() * size_of::<T::Native>() + self.null_state.size()
     }
+
+    fn supports_blocked_groups(&self) -> bool {
+        true
+    }
+
+    fn alter_block_size(&mut self, block_size: Option<usize>) -> Result<()> {
+        self.values.clear();
+        self.null_state = NullStateAdapter::new(block_size);
+        self.block_size = block_size;
+
+        Ok(())
+    }
+}
+
+impl<N: ArrowNativeTypeOp> Block for Vec<N> {
+    type T = N;
+
+    fn len(&self) -> usize {
+        self.len()
+    }
+
+    fn fill_default_value(&mut self, fill_len: usize, default_value: Self::T) {
+        self.extend(iter::repeat(default_value.clone()).take(fill_len));
+    }
 }