make EmitTo::NextBlock simpler.

Rachelint · Rachelint · commit f60f9e2d490d · 2025-04-20T18:37:40.000+08:00
diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs
@@ -36,7 +36,7 @@ pub enum EmitTo {
     /// Emit next block in the blocked managed groups
     ///
     /// Similar as `Emit::All`, will also clear all existing group indexes
-    NextBlock(bool),
+    NextBlock,
 }
 
 impl EmitTo {
@@ -75,7 +75,7 @@ impl EmitTo {
                 std::mem::swap(v, &mut t);
                 t
             }
-            Self::NextBlock(_) => unreachable!("don't support take block in take_needed"),
+            Self::NextBlock => unreachable!("don't support take block in take_needed"),
         }
     }
 
@@ -85,7 +85,7 @@ impl EmitTo {
     ///
     fn take_needed_block<T>(&self, blocks: &mut VecDeque<Vec<T>>) -> Vec<T> {
         assert!(
-            matches!(self, Self::NextBlock(_)),
+            matches!(self, Self::NextBlock),
             "only support take block in take_needed_block"
         );
         blocks
diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs
@@ -303,7 +303,7 @@ impl SeenValues for FlatSeenValues {
                 }
                 first_n_null
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("not support block emit in flat seen values")
             }
         };
@@ -415,7 +415,7 @@ impl SeenValues for BlockedSeenValues {
     }
 
     fn emit(&mut self, emit_to: EmitTo) -> NullBuffer {
-        assert!(matches!(emit_to, EmitTo::NextBlock(_)));
+        assert!(matches!(emit_to, EmitTo::NextBlock));
 
         let mut block = self
             .blocked_builders
@@ -542,7 +542,7 @@ impl NullStateAdapter {
                 let mut return_builder = BooleanBufferBuilder::new(0);
                 let num_blocks = null_state.seen_values.blocked_builders.len();
                 for _ in 0..num_blocks {
-                    let blocked_nulls = null_state.build(EmitTo::NextBlock(true));
+                    let blocked_nulls = null_state.build(EmitTo::NextBlock);
                     for bit in blocked_nulls.inner().iter() {
                         return_builder.append(bit);
                     }
diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs
@@ -118,7 +118,7 @@ where
                 }
                 first_n
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this accumulator still not support blocked groups")
             }
         };
diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs
@@ -448,7 +448,7 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
         let n = match emit_to {
             EmitTo::All => self.count.len(),
             EmitTo::First(n) => n,
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this accumulator still not support blocked groups")
             }
         };
@@ -504,7 +504,7 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator {
         let n = match emit_to {
             EmitTo::All => self.count.len(),
             EmitTo::First(n) => n,
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this accumulator still not support blocked groups")
             }
         };
diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs
@@ -404,7 +404,7 @@ where
                 self.size_of_orderings -=
                     result.iter().map(ScalarValue::size_of_vec).sum::<usize>()
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this accumulator still not support blocked groups")
             }
         }
@@ -431,7 +431,7 @@ where
                 }
                 first_n
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this group values still not support blocked groups")
             }
         }
diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs
@@ -511,7 +511,7 @@ impl MinMaxBytesState {
                 self.total_data_bytes -= first_data_capacity;
                 (first_data_capacity, first_min_maxes)
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this accumulator still not support blocked groups")
             }
         }
diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs
@@ -1154,7 +1154,7 @@ impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> {
 
                 output
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this group values still not support blocked groups")
             }
         };
diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs
@@ -229,7 +229,7 @@ impl GroupValues for GroupValuesRows {
                 });
                 output
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this group values still not support blocked groups")
             }
         };
diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs
@@ -116,7 +116,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {
 
                 emit_group_values
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this group values still not support blocked groups")
             }
         };
diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs
@@ -117,7 +117,7 @@ impl GroupValues for GroupValuesBytesView {
 
                 emit_group_values
             }
-            EmitTo::NextBlock(_) => {
+            EmitTo::NextBlock => {
                 unreachable!("this group values still not support blocked groups")
             }
         };
diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs
@@ -115,6 +115,9 @@ pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> {
 impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> {
     pub fn new(data_type: DataType) -> Self {
         assert!(PrimitiveArray::<T>::is_compatible(&data_type));
+
+        // As a optimization, we ensure the `single block` always exist
+        // in flat mode, it can eliminate an expansive row-level empty checking
         let mut values = VecDeque::new();
         values.push_back(Vec::new());
 
@@ -245,74 +248,15 @@ where
             // ===============================================
             // Emitting in blocked mode
             // ===============================================
-            // TODO: we should consider if it is necessary to support indices modifying
-            // in `EmitTo::NextBlock`. It is only used in spilling case, maybe we can
-            // always emit all in blocked mode. So, we just need to clear the map rather
-            // than doing expansive modification for each buck in it.
-            EmitTo::NextBlock(true) => {
+            EmitTo::NextBlock => {
                 assert!(
                     self.block_size.is_some(),
                     "only support EmitTo::Next in blocked group values"
                 );
 
-                // We only emit the first block(`block_id == 0`),
-                // so erase the entries with `block_id == 0`, and decrease entries with `block_id > 0`
-                self.map.retain(|packed_idx| {
-                    let old_blk_id =
-                        BlockedGroupIndexOperations::get_block_id(*packed_idx);
-                    match old_blk_id.checked_sub(1) {
-                        // `block_id > 0`, shift `block_id` down
-                        Some(new_blk_id) => {
-                            let blk_offset =
-                                BlockedGroupIndexOperations::get_block_offset(
-                                    *packed_idx,
-                                );
-                            let new_packed_idx = BlockedGroupIndexOperations::pack_index(
-                                new_blk_id as u32,
-                                blk_offset,
-                            );
-                            *packed_idx = new_packed_idx;
-
-                            true
-                        }
-
-                        // `block_id == 0`, so remove from table
-                        None => false,
-                    }
-                });
-
-                // Similar as `non-nulls`, if `block_id > 0` we decrease, and if `block_id == 0` we erase
-                let null_block_pair_opt = self.null_group.map(|packed_idx| {
-                    (
-                        BlockedGroupIndexOperations::get_block_id(packed_idx),
-                        BlockedGroupIndexOperations::get_block_offset(packed_idx),
-                    )
-                });
-                let null_idx = match null_block_pair_opt {
-                    Some((blk_id, blk_offset)) if blk_id > 0 => {
-                        let new_blk_id = blk_id - 1;
-                        let new_packed_idx = BlockedGroupIndexOperations::pack_index(
-                            new_blk_id, blk_offset,
-                        );
-                        self.null_group = Some(new_packed_idx);
-                        None
-                    }
-                    Some((_, blk_offset)) => {
-                        self.null_group = None;
-                        Some(blk_offset as usize)
-                    }
-                    None => None,
-                };
-
-                let emit_blk = self.values.pop_front().unwrap();
-                build_primitive(emit_blk, null_idx)
-            }
-
-            EmitTo::NextBlock(false) => {
-                assert!(
-                    self.block_size.is_some(),
-                    "only support EmitTo::Next in blocked group values"
-                );
+                // Similar as `EmitTo:All`, we will clear the old index infos both
+                // in `map` and `null_group`
+                self.map.clear();
 
                 let null_block_pair_opt = self.null_group.map(|packed_idx| {
                     (
@@ -359,6 +303,25 @@ where
         self.map.clear();
         self.map.shrink_to(count, |_| 0); // hasher does not matter since the map is cleared
     }
+
+    fn supports_blocked_groups(&self) -> bool {
+        true
+    }
+
+    fn alter_block_size(&mut self, block_size: Option<usize>) -> Result<()> {
+        self.map.clear();
+        self.values.clear();
+        self.null_group = None;
+        self.block_size = block_size;
+
+        // As mentioned above, we ensure the `single block` always exist
+        // in `flat mode`
+        if block_size.is_none() {
+            self.values.push_back(Vec::new());
+        }
+
+        Ok(())
+    }
 }
 
 impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T>

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@ pub enum EmitTo {`
`36`	`36`	`/// Emit next block in the blocked managed groups`
`37`	`37`	`///`
`38`	`38`	/// Similar as `Emit::All`, will also clear all existing group indexes
`39`		`- NextBlock(bool),`
	`39`	`+ NextBlock,`
`40`	`40`	`}`
`41`	`41`
`42`	`42`	`impl EmitTo {`
`@@ -75,7 +75,7 @@ impl EmitTo {`
`75`	`75`	`std::mem::swap(v, &mut t);`
`76`	`76`	`t`
`77`	`77`	`}`
`78`		`- Self::NextBlock(_) => unreachable!("don't support take block in take_needed"),`
	`78`	`+ Self::NextBlock => unreachable!("don't support take block in take_needed"),`
`79`	`79`	`}`
`80`	`80`	`}`
`81`	`81`
`@@ -85,7 +85,7 @@ impl EmitTo {`
`85`	`85`	`///`
`86`	`86`	`fn take_needed_block<T>(&self, blocks: &mut VecDeque<Vec<T>>) -> Vec<T> {`
`87`	`87`	`assert!(`
`88`		`- matches!(self, Self::NextBlock(_)),`
	`88`	`+ matches!(self, Self::NextBlock),`
`89`	`89`	`"only support take block in take_needed_block"`
`90`	`90`	`);`
`91`	`91`	`blocks`
Original file line number	Diff line number	Diff line change
`@@ -303,7 +303,7 @@ impl SeenValues for FlatSeenValues {`
`303`	`303`	`}`
`304`	`304`	`first_n_null`
`305`	`305`	`}`
`306`		`- EmitTo::NextBlock(_) => {`
	`306`	`+ EmitTo::NextBlock => {`
`307`	`307`	`unreachable!("not support block emit in flat seen values")`
`308`	`308`	`}`
`309`	`309`	`};`
`@@ -415,7 +415,7 @@ impl SeenValues for BlockedSeenValues {`
`415`	`415`	`}`
`416`	`416`
`417`	`417`	`fn emit(&mut self, emit_to: EmitTo) -> NullBuffer {`
`418`		`- assert!(matches!(emit_to, EmitTo::NextBlock(_)));`
	`418`	`+ assert!(matches!(emit_to, EmitTo::NextBlock));`
`419`	`419`
`420`	`420`	`let mut block = self`
`421`	`421`	`.blocked_builders`
`@@ -542,7 +542,7 @@ impl NullStateAdapter {`
`542`	`542`	`let mut return_builder = BooleanBufferBuilder::new(0);`
`543`	`543`	`let num_blocks = null_state.seen_values.blocked_builders.len();`
`544`	`544`	`for _ in 0..num_blocks {`
`545`		`- let blocked_nulls = null_state.build(EmitTo::NextBlock(true));`
	`545`	`+ let blocked_nulls = null_state.build(EmitTo::NextBlock);`
`546`	`546`	`for bit in blocked_nulls.inner().iter() {`
`547`	`547`	`return_builder.append(bit);`
`548`	`548`	`}`
Original file line number	Diff line number	Diff line change
`@@ -118,7 +118,7 @@ where`
`118`	`118`	`}`
`119`	`119`	`first_n`
`120`	`120`	`}`
`121`		`- EmitTo::NextBlock(_) => {`
	`121`	`+ EmitTo::NextBlock => {`
`122`	`122`	`unreachable!("this accumulator still not support blocked groups")`
`123`	`123`	`}`
`124`	`124`	`};`
Original file line number	Diff line number	Diff line change
`@@ -404,7 +404,7 @@ where`
`404`	`404`	`self.size_of_orderings -=`
`405`	`405`	`result.iter().map(ScalarValue::size_of_vec).sum::<usize>()`
`406`	`406`	`}`
`407`		`- EmitTo::NextBlock(_) => {`
	`407`	`+ EmitTo::NextBlock => {`
`408`	`408`	`unreachable!("this accumulator still not support blocked groups")`
`409`	`409`	`}`
`410`	`410`	`}`
`@@ -431,7 +431,7 @@ where`
`431`	`431`	`}`
`432`	`432`	`first_n`
`433`	`433`	`}`
`434`		`- EmitTo::NextBlock(_) => {`
	`434`	`+ EmitTo::NextBlock => {`
`435`	`435`	`unreachable!("this group values still not support blocked groups")`
`436`	`436`	`}`
`437`	`437`	`}`
Original file line number	Diff line number	Diff line change
`@@ -511,7 +511,7 @@ impl MinMaxBytesState {`
`511`	`511`	`self.total_data_bytes -= first_data_capacity;`
`512`	`512`	`(first_data_capacity, first_min_maxes)`
`513`	`513`	`}`
`514`		`- EmitTo::NextBlock(_) => {`
	`514`	`+ EmitTo::NextBlock => {`
`515`	`515`	`unreachable!("this accumulator still not support blocked groups")`
`516`	`516`	`}`
`517`	`517`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1154,7 +1154,7 @@ impl<const STREAMING: bool> GroupValues for GroupValuesColumn<STREAMING> {`
`1154`	`1154`
`1155`	`1155`	`output`
`1156`	`1156`	`}`
`1157`		`- EmitTo::NextBlock(_) => {`
	`1157`	`+ EmitTo::NextBlock => {`
`1158`	`1158`	`unreachable!("this group values still not support blocked groups")`
`1159`	`1159`	`}`
`1160`	`1160`	`};`
Original file line number	Diff line number	Diff line change
`@@ -229,7 +229,7 @@ impl GroupValues for GroupValuesRows {`
`229`	`229`	`});`
`230`	`230`	`output`
`231`	`231`	`}`
`232`		`- EmitTo::NextBlock(_) => {`
	`232`	`+ EmitTo::NextBlock => {`
`233`	`233`	`unreachable!("this group values still not support blocked groups")`
`234`	`234`	`}`
`235`	`235`	`};`
Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,7 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesByes<O> {`
`116`	`116`
`117`	`117`	`emit_group_values`
`118`	`118`	`}`
`119`		`- EmitTo::NextBlock(_) => {`
	`119`	`+ EmitTo::NextBlock => {`
`120`	`120`	`unreachable!("this group values still not support blocked groups")`
`121`	`121`	`}`
`122`	`122`	`};`
Original file line number	Diff line number	Diff line change
`@@ -117,7 +117,7 @@ impl GroupValues for GroupValuesBytesView {`
`117`	`117`
`118`	`118`	`emit_group_values`
`119`	`119`	`}`
`120`		`- EmitTo::NextBlock(_) => {`
	`120`	`+ EmitTo::NextBlock => {`
`121`	`121`	`unreachable!("this group values still not support blocked groups")`
`122`	`122`	`}`
`123`	`123`	`};`