improve tests to cover BlockedNullState.

Author: Rachelint

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

@@ -956,12 +956,13 @@ mod test {
             })
             .collect();
 
+        // Test flat style
         Fixture {
             group_indices,
             values,
             values_with_nulls,
             filter,
-            block_size: None,
+            block_size: 3,
         }
         .run()
     }
@@ -989,8 +990,8 @@ mod test {
         /// filter (defaults to None)
         filter: BooleanArray,
 
-        ///
-        block_size: Option<usize>,
+        /// block size for testing [`BlockedNullState`]
+        block_size: usize,
     }
 
     impl Fixture {
@@ -1007,6 +1008,9 @@ mod test {
 
             let values: Vec<u32> = (0..num_values).map(|_| rng.gen()).collect();
 
+            // random block size
+            let block_size = rng.gen_range(1..num_groups);
+
             // 10% chance of false
             // 10% change of null
             // 80% chance of true
@@ -1042,7 +1046,7 @@ mod test {
                 values,
                 values_with_nulls,
                 filter,
-                block_size: None,
+                block_size,
             }
         }
 
@@ -1111,17 +1115,27 @@ mod test {
             values: &UInt32Array,
             opt_filter: Option<&BooleanArray>,
             total_num_groups: usize,
-            block_size: Option<usize>,
+            block_size: usize,
         ) {
+            // Test `accumulate` of `FlatNullState`
             Self::accumulate_values_test(
                 group_indices,
                 values,
                 opt_filter,
                 total_num_groups,
-                block_size,
+                None,
             );
-            Self::accumulate_indices_test(group_indices, values.nulls(), opt_filter);
 
+            // Test `accumulate` of `BlockedNullState`
+            Self::accumulate_values_test(
+                group_indices,
+                values,
+                opt_filter,
+                total_num_groups,
+                Some(block_size),
+            );
+
+            // Test `accumulate_boolean` of `FlatNullState`
             // Convert values into a boolean array (anything above the
             // average is true, otherwise false)
             let avg: usize = values.iter().filter_map(|v| v.map(|v| v as usize)).sum();
@@ -1132,7 +1146,20 @@ mod test {
                 &boolean_values,
                 opt_filter,
                 total_num_groups,
+                None,
             );
+
+            // Test `accumulate_boolean` of `BlockedNullState`
+            Self::accumulate_boolean_test(
+                group_indices,
+                &boolean_values,
+                opt_filter,
+                total_num_groups,
+                Some(block_size),
+            );
+
+            // Test `accumulate_indices`
+            Self::accumulate_indices_test(group_indices, values.nulls(), opt_filter);
         }
 
         /// This is effectively a different implementation of
@@ -1285,17 +1312,44 @@ mod test {
             values: &BooleanArray,
             opt_filter: Option<&BooleanArray>,
             total_num_groups: usize,
+            block_size: Option<usize>,
         ) {
             let mut accumulated_values = vec![];
-            let mut null_state = FlatNullState::new();
+            let (mut null_state, block_size, acc_group_indices) = if let Some(blk_size) =
+                block_size
+            {
+                let acc_group_indices = group_indices
+                    .iter()
+                    .copied()
+                    .map(|index| {
+                        let block_id = (index / blk_size) as u32;
+                        let block_offset = (index % blk_size) as u64;
+                        BlockedGroupIndexOperations::pack_index(block_id, block_offset)
+                            as usize
+                    })
+                    .collect::<Vec<_>>();
+                (
+                    NullStateAdapter::new(Some(blk_size)),
+                    blk_size,
+                    acc_group_indices,
+                )
+            } else {
+                (
+                    NullStateAdapter::new(None),
+                    0,
+                    group_indices.iter().copied().collect(),
+                )
+            };
 
             null_state.accumulate_boolean(
-                group_indices,
+                &acc_group_indices,
                 values,
                 opt_filter,
                 total_num_groups,
-                |_, group_index, value| {
-                    accumulated_values.push((group_index as usize, value));
+                |block_id, block_offset, value| {
+                    let flatten_index =
+                        ((block_id as u64 * block_size as u64) + block_offset) as usize;
+                    accumulated_values.push((flatten_index as usize, value));
                 },
             );
 
@@ -1332,13 +1386,13 @@ mod test {
             assert_eq!(accumulated_values, expected_values,
                        "\n\naccumulated_values:{accumulated_values:#?}\n\nexpected_values:{expected_values:#?}");
 
-            let seen_values = null_state.seen_values.builder.finish_cloned();
+            let seen_values = null_state.build_cloned_seen_values();
             mock.validate_seen_values(&seen_values);
 
             // Validate the final buffer (one value per group)
             let expected_null_buffer = mock.expected_null_buffer(total_num_groups);
 
-            let null_buffer = null_state.build(EmitTo::All);
+            let null_buffer = null_state.build_all_in_once();
 
             assert_eq!(null_buffer, expected_null_buffer);
         }