refactor: switch BooleanBufferBuilder to NullBufferBuilder in unit te…

…sts of multi_group_by (#14325) Co-authored-by: Cheng-Yuan-Lai <a186235@g,ail.com>
apache · Jan 27, 2025 · eabbbaf · eabbbaf
1 parent c1ca005
commit eabbbaf
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 37 deletions.
diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs
@@ -405,7 +405,7 @@ mod tests {
 
     use crate::aggregates::group_values::multi_group_by::bytes::ByteGroupValueBuilder;
     use arrow_array::{ArrayRef, StringArray};
-    use arrow_buffer::{BooleanBufferBuilder, NullBuffer};
+    use arrow_buffer::NullBufferBuilder;
     use datafusion_physical_expr::binary_map::OutputType;
 
     use super::GroupColumn;
@@ -602,16 +602,15 @@ mod tests {
         .into_parts();
 
         // explicitly build a boolean buffer where one of the null values also happens to match
-        let mut boolean_buffer_builder = BooleanBufferBuilder::new(6);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(false); // this sets Some("bar") to null above
-        boolean_buffer_builder.append(false);
-        boolean_buffer_builder.append(false);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        let nulls = NullBuffer::new(boolean_buffer_builder.finish());
+        let mut nulls = NullBufferBuilder::new(6);
+        nulls.append_non_null();
+        nulls.append_null(); // this sets Some("bar") to null above
+        nulls.append_null();
+        nulls.append_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
         let input_array =
-            Arc::new(StringArray::new(offsets, buffer, Some(nulls))) as ArrayRef;
+            Arc::new(StringArray::new(offsets, buffer, nulls.finish())) as ArrayRef;
 
         // Check
         let mut equal_to_results = vec![true; builder.len()];

diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs
@@ -548,7 +548,7 @@ mod tests {
     use arrow::array::AsArray;
     use arrow::datatypes::StringViewType;
     use arrow_array::{ArrayRef, StringViewArray};
-    use arrow_buffer::{BooleanBufferBuilder, NullBuffer};
+    use arrow_buffer::NullBufferBuilder;
 
     use super::GroupColumn;
 
@@ -751,22 +751,21 @@ mod tests {
         .into_parts();
 
         // explicitly build a boolean buffer where one of the null values also happens to match
-        let mut boolean_buffer_builder = BooleanBufferBuilder::new(9);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(false); // this sets Some("bar") to null above
-        boolean_buffer_builder.append(false);
-        boolean_buffer_builder.append(false);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        let nulls = NullBuffer::new(boolean_buffer_builder.finish());
+        let mut nulls = NullBufferBuilder::new(9);
+        nulls.append_non_null();
+        nulls.append_null(); // this sets Some("bar") to null above
+        nulls.append_null();
+        nulls.append_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
         let input_array =
-            Arc::new(StringViewArray::new(views, buffer, Some(nulls))) as ArrayRef;
+            Arc::new(StringViewArray::new(views, buffer, nulls.finish())) as ArrayRef;
 
         // Check
         let mut equal_to_results = vec![true; input_array.len()];

diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs
@@ -214,7 +214,7 @@ mod tests {
     use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder;
     use arrow::datatypes::Int64Type;
     use arrow_array::{ArrayRef, Int64Array};
-    use arrow_buffer::{BooleanBufferBuilder, NullBuffer};
+    use arrow_buffer::NullBufferBuilder;
     use arrow_schema::DataType;
 
     use super::GroupColumn;
@@ -304,16 +304,15 @@ mod tests {
             Int64Array::from(vec![Some(1), Some(2), None, None, Some(1), Some(3)])
                 .into_parts();
 
-        // explicitly build a boolean buffer where one of the null values also happens to match
-        let mut boolean_buffer_builder = BooleanBufferBuilder::new(6);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(false); // this sets Some(2) to null above
-        boolean_buffer_builder.append(false);
-        boolean_buffer_builder.append(false);
-        boolean_buffer_builder.append(true);
-        boolean_buffer_builder.append(true);
-        let nulls = NullBuffer::new(boolean_buffer_builder.finish());
-        let input_array = Arc::new(Int64Array::new(values, Some(nulls))) as ArrayRef;
+        // explicitly build a null buffer where one of the null values also happens to match
+        let mut nulls = NullBufferBuilder::new(6);
+        nulls.append_non_null();
+        nulls.append_null(); // this sets Some(2) to null above
+        nulls.append_null();
+        nulls.append_null();
+        nulls.append_non_null();
+        nulls.append_non_null();
+        let input_array = Arc::new(Int64Array::new(values, nulls.finish())) as ArrayRef;
 
         // Check
         let mut equal_to_results = vec![true; builder.len()];