From 626acc110d95e828d17a5ae933892aa1ad8abbea Mon Sep 17 00:00:00 2001
From: kamille <caoruiqiu.crq@antgroup.com>
Date: Sun, 24 Nov 2024 11:09:38 +0800
Subject: [PATCH] impl build_nulls_with_buffer.

---
 .../group_values/multi_group_by/primitive.rs  |  1 +
 .../aggregates/group_values/null_builder.rs   | 35 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 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
index 7e9f8ecc0115..481c85087c1a 100644
--- 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
index 9605d919bb9a..af31c38b8602 100644
--- 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))
+    }
+}