Defered cloning in with_validity (#1104)

jorgecarleitao · Jun 26, 2022 · 5b66f5f · 5b66f5f
1 parent b679b06
commit 5b66f5f
Show file tree

Hide file tree

Showing 17 changed files with 176 additions and 109 deletions.
diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs
@@ -219,17 +219,23 @@ impl<O: Offset> BinaryArray<O> {
         std::sync::Arc::new(self)
     }
 
-    /// Clones this [`BinaryArray`] with a different validity.
+    /// Returns this [`BinaryArray`] with a new validity.
     /// # Panic
     /// Panics iff `validity.len() != self.len()`.
     #[must_use]
-    pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
+    }
+
+    /// Sets the validity of this [`BinaryArray`].
+    /// # Panics
+    /// This function panics iff `values.len() != self.len()`.
+    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
         if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
-            panic!("validity's length must be equal to the array's length")
+            panic!("validity must be equal to the array's length")
         }
-        let mut arr = self.clone();
-        arr.validity = validity;
-        arr
+        self.validity = validity;
     }
 
     /// Creates an empty [`BinaryArray`], i.e. whose `.len` is zero.
@@ -449,7 +455,7 @@ impl<O: Offset> Array for BinaryArray<O> {
         Box::new(self.slice_unchecked(offset, length))
     }
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
-        Box::new(self.with_validity(validity))
+        Box::new(self.clone().with_validity(validity))
     }
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())

diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
@@ -175,14 +175,13 @@ impl BooleanArray {
         }
     }
 
-    /// Clones this [`BooleanArray`], returning one with the provided validity.
+    /// Returns this [`BooleanArray`] with a new validity.
     /// # Panic
     /// This function panics iff `validity.len() != self.len()`.
     #[must_use]
-    pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
-        let mut array = self.clone();
-        array.set_validity(validity);
-        array
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
     }
 
     /// Sets the validity of this [`BooleanArray`].
@@ -412,7 +411,7 @@ impl Array for BooleanArray {
         Box::new(self.slice_unchecked(offset, length))
     }
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
-        Box::new(self.with_validity(validity))
+        Box::new(self.clone().with_validity(validity))
     }
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())

diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs
@@ -109,16 +109,20 @@ impl<K: DictionaryKey> DictionaryArray<K> {
         }
     }
 
-    /// Sets the validity bitmap on this [`Array`].
+    /// Returns this [`DictionaryArray`] with a new validity.
     /// # Panic
     /// This function panics iff `validity.len() != self.len()`.
-    pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
-        if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
-            panic!("validity should be as least as large as the array")
-        }
-        let mut arr = self.clone();
-        arr.values = arr.values.with_validity(validity);
-        arr
+    #[must_use]
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
+    }
+
+    /// Sets the validity of the keys of this [`DictionaryArray`].
+    /// # Panics
+    /// This function panics iff `validity.len() != self.len()`.
+    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
+        self.keys.set_validity(validity);
     }
 }
 
@@ -213,7 +217,7 @@ impl<K: DictionaryKey> Array for DictionaryArray<K> {
         Box::new(self.slice_unchecked(offset, length))
     }
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
-        Box::new(self.with_validity(validity))
+        Box::new(self.clone().with_validity(validity))
     }
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())

diff --git a/src/array/fixed_size_binary/mod.rs b/src/array/fixed_size_binary/mod.rs
@@ -144,17 +144,23 @@ impl FixedSizeBinaryArray {
         }
     }
 
-    /// Sets the validity bitmap on this [`FixedSizeBinaryArray`].
+    /// Returns this [`FixedSizeBinaryArray`] with a new validity.
     /// # Panic
     /// This function panics iff `validity.len() != self.len()`.
     #[must_use]
-    pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
+    }
+
+    /// Sets the validity of this [`FixedSizeBinaryArray`].
+    /// # Panics
+    /// This function panics iff `validity.len() != self.len()`.
+    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
         if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
-            panic!("validity should be as least as large as the array")
+            panic!("validity must be equal to the array's length")
         }
-        let mut arr = self.clone();
-        arr.validity = validity;
-        arr
+        self.validity = validity;
     }
 }
 
@@ -273,12 +279,15 @@ impl Array for FixedSizeBinaryArray {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+
     unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice_unchecked(offset, length))
     }
+
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
-        Box::new(self.with_validity(validity))
+        Box::new(self.clone().with_validity(validity))
     }
+
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())
     }

diff --git a/src/array/fixed_size_list/mod.rs b/src/array/fixed_size_list/mod.rs
@@ -158,17 +158,23 @@ impl FixedSizeListArray {
         }
     }
 
-    /// Sets the validity bitmap on this [`FixedSizeListArray`].
+    /// Returns this [`FixedSizeListArray`] with a new validity.
     /// # Panic
     /// This function panics iff `validity.len() != self.len()`.
     #[must_use]
-    pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
+    }
+
+    /// Sets the validity of this [`FixedSizeListArray`].
+    /// # Panics
+    /// This function panics iff `validity.len() != self.len()`.
+    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
         if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
-            panic!("validity should be as least as large as the array")
+            panic!("validity must be equal to the array's length")
         }
-        let mut arr = self.clone();
-        arr.validity = validity;
-        arr
+        self.validity = validity;
     }
 }
 
@@ -268,7 +274,7 @@ impl Array for FixedSizeListArray {
         Box::new(self.slice_unchecked(offset, length))
     }
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
-        Box::new(self.with_validity(validity))
+        Box::new(self.clone().with_validity(validity))
     }
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())

diff --git a/src/array/list/mod.rs b/src/array/list/mod.rs
@@ -231,16 +231,23 @@ impl<O: Offset> ListArray<O> {
         }
     }
 
-    /// Sets the validity bitmap on this [`ListArray`].
+    /// Returns this [`ListArray`] with a new validity.
     /// # Panic
     /// This function panics iff `validity.len() != self.len()`.
-    pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
+    #[must_use]
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
+    }
+
+    /// Sets the validity of this [`ListArray`].
+    /// # Panics
+    /// This function panics iff `validity.len() != self.len()`.
+    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
         if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
-            panic!("validity should be as least as large as the array")
+            panic!("validity must be equal to the array's length")
         }
-        let mut arr = self.clone();
-        arr.validity = validity;
-        arr
+        self.validity = validity;
     }
 }
 
@@ -368,12 +375,15 @@ impl<O: Offset> Array for ListArray<O> {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+
     unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice_unchecked(offset, length))
     }
+
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
-        Box::new(self.with_validity(validity))
+        Box::new(self.clone().with_validity(validity))
     }
+
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())
     }

diff --git a/src/array/map/mod.rs b/src/array/map/mod.rs
@@ -117,6 +117,25 @@ impl MapArray {
         Self::new(data_type, Buffer::from(vec![0i32]), field, None)
     }
 
+    /// Returns this [`MapArray`] with a new validity.
+    /// # Panics
+    /// This function panics iff `validity.len() != self.len()`.
+    #[must_use]
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
+    }
+
+    /// Sets the validity of this [`MapArray`].
+    /// # Panics
+    /// This function panics iff `validity.len() != self.len()`.
+    pub fn set_validity(&mut self, validity: Option<Bitmap>) {
+        if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
+            panic!("validity's length must be equal to the array's length")
+        }
+        self.validity = validity;
+    }
+
     /// Boxes self into a [`Box<dyn Array>`].
     pub fn boxed(self) -> Box<dyn Array> {
         Box::new(self)
@@ -252,9 +271,10 @@ impl Array for MapArray {
         Box::new(self.slice_unchecked(offset, length))
     }
 
-    fn with_validity(&self, _validity: Option<Bitmap>) -> Box<dyn Array> {
-        self.to_boxed()
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.clone().with_validity(validity))
     }
+
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())
     }

diff --git a/src/array/mod.rs b/src/array/mod.rs
@@ -102,9 +102,9 @@ pub trait Array: Send + Sync + dyn_clone::DynClone + 'static {
     /// The caller must ensure that `offset + length <= self.len()`
     unsafe fn slice_unchecked(&self, offset: usize, length: usize) -> Box<dyn Array>;
 
-    /// Sets the validity bitmap on this [`Array`].
+    /// Clones this [`Array`] with a new new assigned bitmap.
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// This function panics iff `validity.len() != self.len()`.
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array>;
 
     /// Clone a `&dyn Array` to an owned `Box<dyn Array>`.

diff --git a/src/array/null.rs b/src/array/null.rs
@@ -114,6 +114,7 @@ impl Array for NullArray {
     fn with_validity(&self, _: Option<Bitmap>) -> Box<dyn Array> {
         panic!("cannot set validity of a null array")
     }
+
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())
     }

diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
@@ -241,58 +241,54 @@ impl<T: NativeType> PrimitiveArray<T> {
         }
     }
 
-    /// Returns a clone of this [`PrimitiveArray`] with a new validity.
+    /// Returns this [`PrimitiveArray`] with a new validity.
     /// # Panics
     /// This function panics iff `validity.len() != self.len()`.
     #[must_use]
-    pub fn with_validity(&self, validity: Option<Bitmap>) -> Self {
-        let mut out = self.clone();
-        out.set_validity(validity);
-        out
+    pub fn with_validity(mut self, validity: Option<Bitmap>) -> Self {
+        self.set_validity(validity);
+        self
     }
 
-    /// Update the validity buffer of this [`PrimitiveArray`].
+    /// Sets the validity of this [`PrimitiveArray`].
     /// # Panics
-    /// This function panics iff `values.len() != self.len()`.
+    /// This function panics iff `validity.len() != self.len()`.
     pub fn set_validity(&mut self, validity: Option<Bitmap>) {
         if matches!(&validity, Some(bitmap) if bitmap.len() != self.len()) {
-            panic!("validity should be as least as large as the array")
+            panic!("validity's length must be equal to the array's length")
         }
         self.validity = validity;
     }
 
-    /// Returns a clone of this [`PrimitiveArray`] with new values.
+    /// Returns this [`PrimitiveArray`] with new values.
     /// # Panics
     /// This function panics iff `values.len() != self.len()`.
     #[must_use]
-    pub fn with_values(&self, values: Buffer<T>) -> Self {
-        let mut out = self.clone();
-        out.set_values(values);
-        out
+    pub fn with_values(mut self, values: Buffer<T>) -> Self {
+        self.set_values(values);
+        self
     }
 
-    /// Update the values buffer of this [`PrimitiveArray`].
+    /// Update the values of this [`PrimitiveArray`].
     /// # Panics
     /// This function panics iff `values.len() != self.len()`.
     pub fn set_values(&mut self, values: Buffer<T>) {
         assert_eq!(
             values.len(),
             self.len(),
-            "values length should be equal to this arrays length"
+            "values' length must be equal to this arrays' length"
         );
         self.values = values;
     }
 
-    /// Applies a function `f` to the validity of this array, the caller can decide to make
-    /// it mutable or not.
+    /// 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: Fn(Bitmap) -> Bitmap>(&mut self, f: F) {
         if let Some(validity) = std::mem::take(&mut self.validity) {
-            assert_eq!(validity.len(), self.values.len());
-            self.validity = Some(f(validity))
+            self.set_validity(Some(f(validity)))
         }
     }
 
@@ -462,7 +458,7 @@ impl<T: NativeType> Array for PrimitiveArray<T> {
         Box::new(self.slice_unchecked(offset, length))
     }
     fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
-        Box::new(self.with_validity(validity))
+        Box::new(self.clone().with_validity(validity))
     }
     fn to_boxed(&self) -> Box<dyn Array> {
         Box::new(self.clone())