From e257551504aceac17eb3b2d19a7f227d7fb99963 Mon Sep 17 00:00:00 2001
From: "artem.malyshev" <artem.malyshev@bytedance.com>
Date: Thu, 16 Mar 2023 15:56:36 -0700
Subject: [PATCH] use mut_validity macro for mutable arrays to define
 set_validity, apply_validity and with_validity

---
 src/array/binary/mutable.rs      |  2 ++
 src/array/mod.rs                 | 39 +++++++++++++++++++++++++
 tests/it/array/binary/mutable.rs | 50 ++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+)

diff --git a/src/array/binary/mutable.rs b/src/array/binary/mutable.rs
index a4b6e2abddc..25bbe286c2d 100644
--- a/src/array/binary/mutable.rs
+++ b/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> {
diff --git a/src/array/mod.rs b/src/array/mod.rs
index e2eda0c5b97..91c06f02bf2 100644
--- a/src/array/mod.rs
+++ b/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 {
     () => {
diff --git a/tests/it/array/binary/mutable.rs b/tests/it/array/binary/mutable.rs
index 012e79d5e4a..0d388b24162 100644
--- a/tests/it/array/binary/mutable.rs
+++ b/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));
+}