From 69018aedaec6452fc4599700223a2965cf1dc902 Mon Sep 17 00:00:00 2001
From: Arty-Maly <Arty-Maly@users.noreply.github.com>
Date: Tue, 14 Feb 2023 12:59:05 -0800
Subject: [PATCH] Added apply_validity and set_validity to mutable utf8 array
 (#1406)

add apply_validity and set_validity to mutable utf8 array

Co-authored-by: artem.malyshev <artem.malyshev@bytedance.com>
---
 src/array/utf8/mutable.rs      | 21 +++++++++++++++++
 tests/it/array/utf8/mutable.rs | 41 ++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/src/array/utf8/mutable.rs b/src/array/utf8/mutable.rs
index a967615e8bc..ca2013eac74 100644
--- a/src/array/utf8/mutable.rs
+++ b/src/array/utf8/mutable.rs
@@ -212,6 +212,27 @@ impl<O: Offset> MutableUtf8Array<O> {
     pub fn values_iter(&self) -> MutableUtf8ValuesIter<O> {
         self.values.iter()
     }
+
+    /// Sets the validity.
+    /// # Panic
+    /// Panics iff the validity's len is not equal to the existing values' length.
+    pub fn set_validity(&mut self, validity: Option<MutableBitmap>) {
+        if let Some(validity) = &validity {
+            assert_eq!(self.values.len(), validity.len())
+        }
+        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`].
+    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)))
+        }
+    }
 }
 
 impl<O: Offset> MutableUtf8Array<O> {
diff --git a/tests/it/array/utf8/mutable.rs b/tests/it/array/utf8/mutable.rs
index 964337e7b2c..b33fb59966e 100644
--- a/tests/it/array/utf8/mutable.rs
+++ b/tests/it/array/utf8/mutable.rs
@@ -199,3 +199,44 @@ fn extend_from_self() {
         MutableUtf8Array::<i32>::from([Some("aa"), None, Some("aa"), None])
     );
 }
+
+#[test]
+fn test_set_validity() {
+    let mut array = MutableUtf8Array::<i32>::from([Some("Red"), Some("Green"), Some("Blue")]);
+    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 = MutableUtf8Array::<i32>::from([Some("Red"), Some("Green"), Some("Blue")]);
+    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 = MutableUtf8Array::<i32>::from([Some("Red"), Some("Green"), Some("Blue")]);
+
+    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));
+}