From c69ce874bc26829944c3375c231f46ec5bee4595 Mon Sep 17 00:00:00 2001
From: Arty-Maly <Arty-Maly@users.noreply.github.com>
Date: Tue, 17 Jan 2023 21:42:55 -0800
Subject: [PATCH] Added `Utf8Array::apply_validity` (#1367)

---
 src/array/utf8/mod.rs      | 11 +++++++++++
 tests/it/array/utf8/mod.rs | 17 +++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs
index 7643ad58c07..4ed67f88430 100644
--- a/src/array/utf8/mod.rs
+++ b/src/array/utf8/mod.rs
@@ -515,6 +515,17 @@ impl<O: Offset> Utf8Array<O> {
     {
         MutableUtf8Array::<O>::try_from_trusted_len_iter(iter).map(|x| x.into())
     }
+
+    /// 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(Bitmap) -> Bitmap>(&mut self, f: F) {
+        if let Some(validity) = std::mem::take(&mut self.validity) {
+            self.set_validity(Some(f(validity)))
+        }
+    }
 }
 
 impl<O: Offset> Array for Utf8Array<O> {
diff --git a/tests/it/array/utf8/mod.rs b/tests/it/array/utf8/mod.rs
index e73c6880536..7f48aebcfe9 100644
--- a/tests/it/array/utf8/mod.rs
+++ b/tests/it/array/utf8/mod.rs
@@ -216,3 +216,20 @@ fn iter_nth() {
     assert_eq!(array.iter().nth(1), Some(Some(" ")));
     assert_eq!(array.iter().nth(10), None);
 }
+
+#[test]
+fn test_apply_validity() {
+    let mut array = Utf8Array::<i32>::from([Some("Red"), Some("Green"), Some("Blue")]);
+    array.set_validity(Some([true, true, true].into()));
+
+    array.apply_validity(|bitmap| {
+        let mut mut_bitmap = bitmap.into_mut().right().unwrap();
+        mut_bitmap.set(1, false);
+        mut_bitmap.set(2, false);
+        mut_bitmap.into()
+    });
+
+    assert!(array.is_valid(0));
+    assert!(!array.is_valid(1));
+    assert!(!array.is_valid(2));
+}