impl build_nulls_with_buffer.

apache · Nov 24, 2024 · 626acc1 · 626acc1
1 parent 510495d
commit 626acc1
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 1 deletion.
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
@@ -102,6 +102,7 @@ impl<T: ArrowPrimitiveType, const NULLABLE: bool> GroupColumn
 
     // fn take_from_exists(&mut self, lhs_rows: &[usize]) -> PrimitiveArray<T> {
     //     // Take value firstly
+    //     // take(values, indices, options)
     //     take(values, indices, options)
     // }
 

diff --git a/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs b/datafusion/physical-plan/src/aggregates/group_values/null_builder.rs
@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow_buffer::{BooleanBufferBuilder, MutableBuffer, NullBuffer};
+use arrow_buffer::{bit_util, BooleanBuffer, BooleanBufferBuilder, MutableBuffer, NullBuffer};
 
 /// Builder for an (optional) null mask
 ///
@@ -136,3 +136,36 @@ impl MaybeNullBufferBuilder {
         }
     }
 }
+
+pub fn build_nulls_with_buffer<I>(
+    null_iter: I,
+    nulls_len: usize,
+    mut buffer: MutableBuffer,
+) -> (Option<NullBuffer>, Option<MutableBuffer>)
+where
+    I: Iterator<Item = bool>,
+{
+    // Ensure the buffer big enough, and init to all `false`
+    buffer.clear();
+    let bytes_len = bit_util::ceil(nulls_len, 8);
+    buffer.resize(bytes_len, 0);
+
+    let null_slice = buffer.as_slice_mut();
+    let mut has_null = false;
+    null_iter.enumerate().for_each(|(idx, is_valid)| {
+        if is_valid {
+            bit_util::set_bit(null_slice, idx);
+        } else {
+            has_null = true;
+        }
+    });
+
+    if has_null {
+        let bool_buffer = BooleanBuffer::new(buffer.into(), 0, nulls_len);
+        let null_buffer = NullBuffer::new(bool_buffer);
+
+        (Some(null_buffer), None)
+    } else {
+        (None, Some(buffer))
+    }
+}