diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs index 7643ad58c07..901ac673160 100644 --- a/src/array/utf8/mod.rs +++ b/src/array/utf8/mod.rs @@ -515,6 +515,17 @@ impl Utf8Array { { MutableUtf8Array::::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 Bitmap>(&mut self, f: F) { + if let Some(validity) = std::mem::take(&mut self.validity) { + self.set_validity(Some(f(validity))) + } + } } impl Array for Utf8Array { @@ -579,3 +590,25 @@ impl Default for Utf8Array { Utf8Array::new(data_type, Default::default(), Default::default(), None) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_apply_validity() { + let mut array = Utf8Array::::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)); + } +} \ No newline at end of file