extract EmitBlocksContext to common, and prevent all new updates during emitting.

Author: Rachelint

datafusion/expr-common/src/groups_accumulator.rs

@@ -64,6 +64,64 @@ impl EmitTo {
     }
 }
 
+/// Emitting context used in blocked management
+#[derive(Debug, Default)]
+pub struct EmitBlocksContext {
+    /// Mark if it is during blocks emitting, if so states can't
+    /// be updated until all blocks are emitted
+    pub emitting: bool,
+
+    /// Idx of next emitted block
+    pub next_emit_block: usize,
+
+    /// Number of blocks needed to emit
+    pub num_blocks: usize,
+}
+
+impl EmitBlocksContext {
+    #[inline]
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    #[inline]
+    pub fn start_emit(&mut self, num_blocks: usize) {
+        self.emitting = true;
+        self.num_blocks = num_blocks;
+    }
+
+    #[inline]
+    pub fn emitting(&self) -> bool {
+        self.emitting
+    }
+
+    #[inline]
+    pub fn all_emitted(&self) -> bool {
+        self.next_emit_block == self.num_blocks
+    }
+
+    #[inline]
+    pub fn cur_emit_block(&self) -> usize {
+        assert!(self.emitting, "must start emit first");
+        self.next_emit_block
+    }
+
+    #[inline]
+    pub fn advance_emit_block(&mut self) {
+        assert!(self.emitting, "must start emit first");
+        if self.next_emit_block < self.num_blocks {
+            self.next_emit_block += 1;
+        }
+    }
+
+    #[inline]
+    pub fn reset(&mut self) {
+        self.emitting = false;
+        self.next_emit_block = 0;
+        self.num_blocks = 0;
+    }
+}
+
 /// `GroupsAccumulator` implements a single aggregate (e.g. AVG) and
 /// stores the state for *all* groups internally.
 ///

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

@@ -26,6 +26,7 @@ use arrow::datatypes::{i256, DataType};
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
 use datafusion_execution::memory_pool::proxy::VecAllocExt;
+use datafusion_expr::groups_accumulator::EmitBlocksContext;
 use datafusion_expr::EmitTo;
 use datafusion_functions_aggregate_common::aggregate::groups_accumulator::group_index_operations::{
     BlockedGroupIndexOperations, FlatGroupIndexOperations, GroupIndexOperations,
@@ -101,8 +102,6 @@ pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> {
     /// The values for each group index
     values: Vec<Vec<T::Native>>,
 
-    next_emit_block_id: usize,
-
     /// The random state used to generate hashes
     random_state: RandomState,
 
@@ -116,7 +115,14 @@ pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> {
     ///
     block_size: Option<usize>,
 
+    /// Number of current storing groups
+    ///
+    /// We maintain it to avoid the expansive dynamic computation in
+    /// `blocked approach`.
     num_groups: usize,
+
+    /// Context used in emitting in `blocked approach`
+    emit_blocks_ctx: EmitBlocksContext,
 }
 
 impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> {
@@ -132,11 +138,11 @@ impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> {
             data_type,
             map: HashTable::with_capacity(128),
             values,
-            next_emit_block_id: 0,
             null_group: None,
             random_state: Default::default(),
             block_size: None,
             num_groups: 0,
+            emit_blocks_ctx: EmitBlocksContext::new(),
         }
     }
 }
@@ -184,7 +190,6 @@ where
 
     fn len(&self) -> usize {
         self.num_groups
-        // self.values.iter().map(|block| block.len()).sum::<usize>()
     }
 
     fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
@@ -263,14 +268,27 @@ where
                     .block_size
                     .expect("only support EmitTo::Next in blocked group values");
 
+                // To mark the emitting has already started, and prevent new updates
+                if !self.emit_blocks_ctx.emitting() {
+                    let num_blocks = self.num_groups.div_ceil(block_size);
+                    self.emit_blocks_ctx.start_emit(num_blocks);
+                }
+
                 // Similar as `EmitTo:All`, we will clear the old index infos both
                 // in `map` and `null_group`
                 self.map.clear();
 
-                // Get current emit block id firstly
-                let emit_block_id = self.next_emit_block_id;
+                // Get current emit block idx firstly
+                let emit_block_id = self.emit_blocks_ctx.cur_emit_block();
                 let emit_blk = std::mem::take(&mut self.values[emit_block_id]);
-                self.next_emit_block_id += 1;
+                // And then we advance the block idx
+                self.emit_blocks_ctx.advance_emit_block();
+                // Finally we check if all blocks emitted, if so, we reset the
+                // emit context to allow new updates
+                if self.emit_blocks_ctx.all_emitted() {
+                    self.emit_blocks_ctx.reset();
+                    self.values.clear();
+                }
 
                 // Check if `null` is in current block
                 let null_block_pair_opt = self.null_group.map(|group_index| {
@@ -430,8 +448,9 @@ mod tests {
 
     use crate::aggregates::group_values::single_group_by::primitive::GroupValuesPrimitive;
     use crate::aggregates::group_values::GroupValues;
-    use arrow::array::{AsArray, UInt32Array};
+    use arrow::array::{AsArray, RecordBatch, UInt32Array};
     use arrow::datatypes::{DataType, UInt32Type};
+    use arrow_schema::Schema;
     use datafusion_expr::EmitTo;
     use datafusion_functions_aggregate_common::aggregate::groups_accumulator::group_index_operations::{
         BlockedGroupIndexOperations, GroupIndexOperations,
@@ -567,6 +586,7 @@ mod tests {
         assert_eq!(actual, expected);
 
         // Insert case 1.1~1.2 + Emit case 2.2
+        group_values.clear_shrink(&RecordBatch::new_empty(Arc::new(Schema::empty())));
         group_values
             .intern(&[Arc::clone(&data2) as _], &mut group_indices)
             .unwrap();