From 7d465b87ccddbff7696e52481de41d8d2bbe7d4b Mon Sep 17 00:00:00 2001
From: Hadrien G <knights_of_ni@gmx.com>
Date: Mon, 13 May 2024 13:14:10 +0200
Subject: [PATCH] Expose the null buffer of every builder that has one (#5754)

---
 arrow-array/src/builder/fixed_size_binary_builder.rs        | 5 +++++
 arrow-array/src/builder/fixed_size_list_builder.rs          | 5 +++++
 arrow-array/src/builder/generic_bytes_dictionary_builder.rs | 5 +++++
 arrow-array/src/builder/generic_bytes_view_builder.rs       | 5 +++++
 arrow-array/src/builder/generic_list_builder.rs             | 5 +++++
 arrow-array/src/builder/map_builder.rs                      | 5 +++++
 arrow-array/src/builder/primitive_dictionary_builder.rs     | 5 +++++
 arrow-array/src/builder/struct_builder.rs                   | 5 +++++
 8 files changed, 40 insertions(+)

diff --git a/arrow-array/src/builder/fixed_size_binary_builder.rs b/arrow-array/src/builder/fixed_size_binary_builder.rs
index 132c2e1939bb..65072a09f603 100644
--- a/arrow-array/src/builder/fixed_size_binary_builder.rs
+++ b/arrow-array/src/builder/fixed_size_binary_builder.rs
@@ -115,6 +115,11 @@ impl FixedSizeBinaryBuilder {
         let array_data = unsafe { array_data_builder.build_unchecked() };
         FixedSizeBinaryArray::from(array_data)
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.null_buffer_builder.as_slice()
+    }
 }
 
 impl ArrayBuilder for FixedSizeBinaryBuilder {
diff --git a/arrow-array/src/builder/fixed_size_list_builder.rs b/arrow-array/src/builder/fixed_size_list_builder.rs
index f65947bfff32..5dff67650687 100644
--- a/arrow-array/src/builder/fixed_size_list_builder.rs
+++ b/arrow-array/src/builder/fixed_size_list_builder.rs
@@ -208,6 +208,11 @@ where
 
         FixedSizeListArray::new(field, self.list_len, values, nulls)
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.null_buffer_builder.as_slice()
+    }
 }
 
 #[cfg(test)]
diff --git a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
index 198d4fcbeb2f..285a4f035e24 100644
--- a/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
+++ b/arrow-array/src/builder/generic_bytes_dictionary_builder.rs
@@ -299,6 +299,11 @@ where
 
         DictionaryArray::from(unsafe { builder.build_unchecked() })
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.keys_builder.validity_slice()
+    }
 }
 
 impl<K: ArrowDictionaryKeyType, T: ByteArrayType, V: AsRef<T::Native>> Extend<Option<V>>
diff --git a/arrow-array/src/builder/generic_bytes_view_builder.rs b/arrow-array/src/builder/generic_bytes_view_builder.rs
index 9accb932ae20..d043bda98cec 100644
--- a/arrow-array/src/builder/generic_bytes_view_builder.rs
+++ b/arrow-array/src/builder/generic_bytes_view_builder.rs
@@ -146,6 +146,11 @@ impl<T: ByteViewType + ?Sized> GenericByteViewBuilder<T> {
         // SAFETY: valid by construction
         unsafe { GenericByteViewArray::new_unchecked(views, completed, nulls) }
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.null_buffer_builder.as_slice()
+    }
 }
 
 impl<T: ByteViewType + ?Sized> Default for GenericByteViewBuilder<T> {
diff --git a/arrow-array/src/builder/generic_list_builder.rs b/arrow-array/src/builder/generic_list_builder.rs
index d6fff12e1ab8..6ff5f20df684 100644
--- a/arrow-array/src/builder/generic_list_builder.rs
+++ b/arrow-array/src/builder/generic_list_builder.rs
@@ -326,6 +326,11 @@ where
     pub fn offsets_slice(&self) -> &[OffsetSize] {
         self.offsets_builder.as_slice()
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.null_buffer_builder.as_slice()
+    }
 }
 
 impl<O, B, V, E> Extend<Option<V>> for GenericListBuilder<O, B>
diff --git a/arrow-array/src/builder/map_builder.rs b/arrow-array/src/builder/map_builder.rs
index b69dcac78424..1d89d427aae1 100644
--- a/arrow-array/src/builder/map_builder.rs
+++ b/arrow-array/src/builder/map_builder.rs
@@ -226,6 +226,11 @@ impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
 
         MapArray::from(array_data)
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.null_buffer_builder.as_slice()
+    }
 }
 
 impl<K: ArrayBuilder, V: ArrayBuilder> ArrayBuilder for MapBuilder<K, V> {
diff --git a/arrow-array/src/builder/primitive_dictionary_builder.rs b/arrow-array/src/builder/primitive_dictionary_builder.rs
index a64ecf0caa13..a764fa4c29c8 100644
--- a/arrow-array/src/builder/primitive_dictionary_builder.rs
+++ b/arrow-array/src/builder/primitive_dictionary_builder.rs
@@ -302,6 +302,11 @@ where
     pub fn values_slice_mut(&mut self) -> &mut [V::Native] {
         self.values_builder.values_slice_mut()
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.keys_builder.validity_slice()
+    }
 }
 
 impl<K: ArrowDictionaryKeyType, P: ArrowPrimitiveType> Extend<Option<P::Native>>
diff --git a/arrow-array/src/builder/struct_builder.rs b/arrow-array/src/builder/struct_builder.rs
index 46202bfa3950..c0e49b939f2c 100644
--- a/arrow-array/src/builder/struct_builder.rs
+++ b/arrow-array/src/builder/struct_builder.rs
@@ -381,6 +381,11 @@ impl StructBuilder {
             }
         });
     }
+
+    /// Returns the current null buffer as a slice
+    pub fn validity_slice(&self) -> Option<&[u8]> {
+        self.null_buffer_builder.as_slice()
+    }
 }
 
 #[cfg(test)]