use mut_validity macro for mutable arrays to define set_validity, app…

…ly_validity and with_validity

src/array/binary/mutable.rs

@@ -160,6 +160,8 @@ impl<O: Offset> MutableBinaryArray<O> {
             validity.shrink_to_fit()
         }
     }
+
+    impl_mutable_array_mut_validity!();
 }
 
 impl<O: Offset> MutableBinaryArray<O> {

src/array/mod.rs

@@ -454,6 +454,45 @@ macro_rules! impl_mut_validity {
     }
 }
 
+// macro implementing `with_validity`, `set_validity` and `apply_validity` for mutable arrays
+macro_rules! impl_mutable_array_mut_validity {
+    () => {
+        /// Returns this array with a new validity.
+        /// # Panic
+        /// Panics iff `validity.len() != self.len()`.
+        #[must_use]
+        #[inline]
+        pub fn with_validity(mut self, validity: Option<MutableBitmap>) -> Self {
+            self.set_validity(validity);
+            self
+        }
+
+        /// Sets the validity of this array.
+        /// # Panics
+        /// This function panics iff `values.len() != self.len()`.
+        #[inline]
+        pub fn set_validity(&mut self, validity: Option<MutableBitmap>) {
+            if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
+                panic!("validity must be equal to the array's length")
+            }
+            self.validity = validity;
+        }
+
+        /// Applies a function `f` to the validity of this array.
+        ///
+        /// This is an API to leverage clone-on-write
+        /// # Panics
+        /// This function panics if the function `f` modifies the length of the [`Bitmap`].
+        #[inline]
+        pub fn apply_validity<F: FnOnce(MutableBitmap) -> MutableBitmap>(&mut self, f: F) {
+            if let Some(validity) = std::mem::take(&mut self.validity) {
+                self.set_validity(Some(f(validity)))
+            }
+        }
+
+    }
+}
+
 // macro implementing `boxed` and `arced`
 macro_rules! impl_into_array {
     () => {

tests/it/array/binary/mutable.rs

@@ -163,3 +163,53 @@ fn extend_from_self() {
         MutableBinaryArray::<i32>::from([Some(b"aa"), None, Some(b"aa"), None])
     );
 }
+
+#[test]
+fn test_set_validity() {
+    let mut array = MutableBinaryArray::<i32>::new();
+    array.push(Some(b"first"));
+    array.push(Some(b"second"));
+    array.push(Some(b"third"));
+    array.set_validity(Some([false, false, true].into()));
+
+    assert!(!array.is_valid(0));
+    assert!(!array.is_valid(1));
+    assert!(array.is_valid(2));
+}
+
+#[test]
+fn test_apply_validity() {
+    let mut array = MutableBinaryArray::<i32>::new();
+    array.push(Some(b"first"));
+    array.push(Some(b"second"));
+    array.push(Some(b"third"));
+    array.set_validity(Some([true, true, true].into()));
+
+    array.apply_validity(|mut mut_bitmap| {
+        mut_bitmap.set(1, false);
+        mut_bitmap.set(2, false);
+        mut_bitmap
+    });
+
+    assert!(array.is_valid(0));
+    assert!(!array.is_valid(1));
+    assert!(!array.is_valid(2));
+}
+
+#[test]
+fn test_apply_validity_with_no_validity_inited() {
+    let mut array = MutableBinaryArray::<i32>::new();
+    array.push(Some(b"first"));
+    array.push(Some(b"second"));
+    array.push(Some(b"third"));
+
+    array.apply_validity(|mut mut_bitmap| {
+        mut_bitmap.set(1, false);
+        mut_bitmap.set(2, false);
+        mut_bitmap
+    });
+
+    assert!(array.is_valid(0));
+    assert!(array.is_valid(1));
+    assert!(array.is_valid(2));
+}