add new emit mode.

Author: Rachelint

datafusion-examples/examples/advanced_udaf.rs

@@ -307,8 +307,8 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator {
 
     /// Generate output, as specified by `emit_to` and update the intermediate state
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let counts = emit_to.take_needed(&mut self.counts);
-        let prods = emit_to.take_needed(&mut self.prods);
+        let counts = emit_to.take_needed_rows(&mut self.counts);
+        let prods = emit_to.take_needed_rows(&mut self.prods);
         let nulls = self.null_state.build(emit_to);
 
         assert_eq!(nulls.len(), prods.len());
@@ -346,10 +346,10 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator {
         let nulls = self.null_state.build(emit_to);
         let nulls = Some(nulls);
 
-        let counts = emit_to.take_needed(&mut self.counts);
+        let counts = emit_to.take_needed_rows(&mut self.counts);
         let counts = UInt32Array::new(counts.into(), nulls.clone()); // zero copy
 
-        let prods = emit_to.take_needed(&mut self.prods);
+        let prods = emit_to.take_needed_rows(&mut self.prods);
         let prods = PrimitiveArray::<Float64Type>::new(prods.into(), nulls) // zero copy
             .with_data_type(self.prod_data_type.clone());
 

datafusion/expr-common/src/groups_accumulator.rs

@@ -17,6 +17,8 @@
 
 //! Vectorized [`GroupsAccumulator`]
 
+use std::collections::VecDeque;
+
 use arrow::array::{ArrayRef, BooleanArray};
 use datafusion_common::{not_impl_err, DataFusionError, Result};
 
@@ -31,15 +33,40 @@ pub enum EmitTo {
     /// For example, if `n=10`, group_index `0, 1, ... 9` are emitted
     /// and group indexes `10, 11, 12, ...` become `0, 1, 2, ...`.
     First(usize),
+    /// Emit next block in the blocked managed groups
+    ///
+    /// The flag's meaning:
+    ///   - `true` represents new groups still will be added,
+    ///     and we need to shift the values down.
+    ///   - `false` represents no new groups will be added again,
+    ///     and we don't need to shift the values down.
+    NextBlock(bool),
 }
 
 impl EmitTo {
+    /// Remove and return `needed values` from `values`.
+    pub fn take_needed<T>(
+        &self,
+        values: &mut VecDeque<Vec<T>>,
+        is_blocked_groups: bool,
+    ) -> Vec<T> {
+        if is_blocked_groups {
+            self.take_needed_block(values)
+        } else {
+            assert_eq!(values.len(), 1);
+            self.take_needed_rows(values.back_mut().unwrap())
+        }
+    }
+
     /// Removes the number of rows from `v` required to emit the right
     /// number of rows, returning a `Vec` with elements taken, and the
     /// remaining values in `v`.
     ///
     /// This avoids copying if Self::All
-    pub fn take_needed<T>(&self, v: &mut Vec<T>) -> Vec<T> {
+    ///
+    /// NOTICE: only support emit strategies: `Self::All` and `Self::First`
+    ///
+    pub fn take_needed_rows<T>(&self, v: &mut Vec<T>) -> Vec<T> {
         match self {
             Self::All => {
                 // Take the entire vector, leave new (empty) vector
@@ -52,8 +79,23 @@ impl EmitTo {
                 std::mem::swap(v, &mut t);
                 t
             }
+            Self::NextBlock(_) => unreachable!("don't support take block in take_needed"),
         }
     }
+
+    /// Removes one block required to emit and return it
+    ///
+    /// NOTICE: only support emit strategy `Self::NextBlock`
+    ///
+    fn take_needed_block<T>(&self, blocks: &mut VecDeque<Vec<T>>) -> Vec<T> {
+        assert!(
+            matches!(self, Self::NextBlock(_)),
+            "only support take block in take_needed_block"
+        );
+        blocks
+            .pop_front()
+            .expect("should not call emit for empty blocks")
+    }
 }
 
 /// `GroupsAccumulator` implements a single aggregate (e.g. AVG) and

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

@@ -321,7 +321,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
         let vec_size_pre = self.states.allocated_size();
 
-        let states = emit_to.take_needed(&mut self.states);
+        let states = emit_to.take_needed_rows(&mut self.states);
 
         let results: Vec<ScalarValue> = states
             .into_iter()
@@ -341,7 +341,7 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter {
     // filtered_null_mask(opt_filter, &values);
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
         let vec_size_pre = self.states.allocated_size();
-        let states = emit_to.take_needed(&mut self.states);
+        let states = emit_to.take_needed_rows(&mut self.states);
 
         // each accumulator produces a potential vector of values
         // which we need to form into columns

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

@@ -228,6 +228,7 @@ impl NullState {
                 }
                 first_n_null
             }
+            EmitTo::NextBlock(_) => todo!(),
         };
         NullBuffer::new(nulls)
     }

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

@@ -117,6 +117,9 @@ where
                 }
                 first_n
             }
+            EmitTo::NextBlock(_) => {
+                unreachable!("this accumulator still not support blocked groups")
+            }
         };
 
         let nulls = self.null_state.build(emit_to);

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

@@ -115,7 +115,7 @@ where
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let values = emit_to.take_needed(&mut self.values);
+        let values = emit_to.take_needed_rows(&mut self.values);
         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());

datafusion/functions-aggregate/src/average.rs

@@ -486,8 +486,8 @@ where
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let counts = emit_to.take_needed(&mut self.counts);
-        let sums = emit_to.take_needed(&mut self.sums);
+        let counts = emit_to.take_needed_rows(&mut self.counts);
+        let sums = emit_to.take_needed_rows(&mut self.sums);
         let nulls = self.null_state.build(emit_to);
 
         assert_eq!(nulls.len(), sums.len());
@@ -526,10 +526,10 @@ where
         let nulls = self.null_state.build(emit_to);
         let nulls = Some(nulls);
 
-        let counts = emit_to.take_needed(&mut self.counts);
+        let counts = emit_to.take_needed_rows(&mut self.counts);
         let counts = UInt64Array::new(counts.into(), nulls.clone()); // zero copy
 
-        let sums = emit_to.take_needed(&mut self.sums);
+        let sums = emit_to.take_needed_rows(&mut self.sums);
         let sums = PrimitiveArray::<T>::new(sums.into(), nulls) // zero copy
             .with_data_type(self.sum_data_type.clone());
 

datafusion/functions-aggregate/src/correlation.rs

@@ -448,6 +448,9 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
         let n = match emit_to {
             EmitTo::All => self.count.len(),
             EmitTo::First(n) => n,
+            EmitTo::NextBlock(_) => {
+                unreachable!("this accumulator still not support blocked groups")
+            }
         };
 
         let mut values = Vec::with_capacity(n);
@@ -501,6 +504,9 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
         let n = match emit_to {
             EmitTo::All => self.count.len(),
             EmitTo::First(n) => n,
+            EmitTo::NextBlock(_) => {
+                unreachable!("this accumulator still not support blocked groups")
+            }
         };
 
         Ok(vec![

datafusion/functions-aggregate/src/count.rs

@@ -537,7 +537,7 @@ impl GroupsAccumulator for CountGroupsAccumulator {
     }
 
     fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
-        let counts = emit_to.take_needed(&mut self.counts);
+        let counts = emit_to.take_needed_rows(&mut self.counts);
 
         // Count is always non null (null inputs just don't contribute to the overall values)
         let nulls = None;
@@ -548,7 +548,7 @@ impl GroupsAccumulator for CountGroupsAccumulator {
 
     // return arrays for counts
     fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
-        let counts = emit_to.take_needed(&mut self.counts);
+        let counts = emit_to.take_needed_rows(&mut self.counts);
         let counts: PrimitiveArray<Int64Type> = Int64Array::from(counts); // zero copy, no nulls
         Ok(vec![Arc::new(counts) as ArrayRef])
     }

datafusion/functions-aggregate/src/first_last.rs

@@ -396,14 +396,17 @@ where
     }
 
     fn take_orderings(&mut self, emit_to: EmitTo) -> Vec<Vec<ScalarValue>> {
-        let result = emit_to.take_needed(&mut self.orderings);
+        let result = emit_to.take_needed_rows(&mut self.orderings);
 
         match emit_to {
             EmitTo::All => self.size_of_orderings = 0,
             EmitTo::First(_) => {
                 self.size_of_orderings -=
                     result.iter().map(ScalarValue::size_of_vec).sum::<usize>()
             }
+            EmitTo::NextBlock(_) => {
+                unreachable!("this accumulator still not support blocked groups")
+            }
         }
 
         result
@@ -428,6 +431,9 @@ where
                 }
                 first_n
             }
+            EmitTo::NextBlock(_) => {
+                unreachable!("this group values still not support blocked groups")
+            }
         }
     }
 
@@ -481,7 +487,7 @@ where
         &mut self,
         emit_to: EmitTo,
     ) -> (ArrayRef, Vec<Vec<ScalarValue>>, BooleanBuffer) {
-        emit_to.take_needed(&mut self.min_of_each_group_buf.0);
+        emit_to.take_needed_rows(&mut self.min_of_each_group_buf.0);
         self.min_of_each_group_buf
             .1
             .truncate(self.min_of_each_group_buf.0.len());
@@ -579,7 +585,7 @@ where
     }
 
     fn take_vals_and_null_buf(&mut self, emit_to: EmitTo) -> ArrayRef {
-        let r = emit_to.take_needed(&mut self.vals);
+        let r = emit_to.take_needed_rows(&mut self.vals);
 
         let null_buf = NullBuffer::new(Self::take_need(&mut self.null_builder, emit_to));