Add Array::with_validity (#399)

src/array/binary/mod.rs

@@ -94,6 +94,18 @@ impl<O: Offset> BinaryArray<O> {
             offset: self.offset + offset,
         }
     }
+
+    /// Sets the validity bitmap on this [`BinaryArray`].
+    /// # 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.validity = validity;
+        arr
+    }
 }
 
 // accessors
@@ -159,6 +171,9 @@ impl<O: Offset> Array for BinaryArray<O> {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl<O: Offset> std::fmt::Display for BinaryArray<O> {

src/array/boolean/mod.rs

@@ -91,6 +91,18 @@ impl BooleanArray {
     pub fn values(&self) -> &Bitmap {
         &self.values
     }
+
+    /// Sets the validity bitmap on this [`BooleanArray`].
+    /// # 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.validity = validity;
+        arr
+    }
 }
 
 impl Array for BooleanArray {
@@ -118,6 +130,9 @@ impl Array for BooleanArray {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl std::fmt::Display for BooleanArray {

src/array/dictionary/mod.rs

@@ -83,6 +83,18 @@ impl<K: DictionaryKey> DictionaryArray<K> {
         }
     }
 
+    /// Sets the validity bitmap on this [`Array`].
+    /// # 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 = Arc::from(arr.values.with_validity(validity));
+        arr
+    }
+
     /// Returns the keys of the [`DictionaryArray`]. These keys can be used to fetch values
     /// from `values`.
     #[inline]
@@ -137,6 +149,9 @@ impl<K: DictionaryKey> Array for DictionaryArray<K> {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl<K: DictionaryKey> std::fmt::Display for DictionaryArray<K>

src/array/fixed_size_binary/mod.rs

@@ -93,6 +93,18 @@ impl FixedSizeBinaryArray {
     pub fn size(&self) -> usize {
         self.size as usize
     }
+
+    /// Sets the validity bitmap on this [`FixedSizeBinaryArray`].
+    /// # 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.validity = validity;
+        arr
+    }
 }
 
 impl FixedSizeBinaryArray {
@@ -128,6 +140,9 @@ impl Array for FixedSizeBinaryArray {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl std::fmt::Display for FixedSizeBinaryArray {

src/array/fixed_size_list/mod.rs

@@ -90,6 +90,18 @@ impl FixedSizeListArray {
         self.values
             .slice(i * self.size as usize, self.size as usize)
     }
+
+    /// Sets the validity bitmap on this [`FixedSizeListArray`].
+    /// # 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.validity = validity;
+        arr
+    }
 }
 
 impl FixedSizeListArray {
@@ -131,6 +143,9 @@ impl Array for FixedSizeListArray {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl std::fmt::Display for FixedSizeListArray {

src/array/list/mod.rs

@@ -117,6 +117,18 @@ impl<O: Offset> ListArray<O> {
     pub fn values(&self) -> &Arc<dyn Array> {
         &self.values
     }
+
+    /// Sets the validity bitmap on this [`ListArray`].
+    /// # 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.validity = validity;
+        arr
+    }
 }
 
 impl<O: Offset> ListArray<O> {
@@ -174,6 +186,9 @@ impl<O: Offset> Array for ListArray<O> {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl<O: Offset> std::fmt::Display for ListArray<O> {

src/array/mod.rs

@@ -88,6 +88,11 @@ pub trait Array: std::fmt::Debug + Send + Sync {
     /// # Panic
     /// This function panics iff `offset + length >= self.len()`.
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array>;
+
+    /// Sets the validity bitmap on this [`Array`].
+    /// # Panic
+    /// This function panics iff `validity.len() < self.len()`.
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array>;
 }
 
 /// A trait describing a mutable array; i.e. an array whose values can be changed.

src/array/null.rs

@@ -63,6 +63,9 @@ impl Array for NullArray {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, _: Option<Bitmap>) -> Box<dyn Array> {
+        panic!("cannot set validity of a null array")
+    }
 }
 
 impl std::fmt::Display for NullArray {

src/array/primitive/mod.rs

@@ -94,6 +94,18 @@ impl<T: NativeType> PrimitiveArray<T> {
         }
     }
 
+    /// Sets the validity bitmap on this [`PrimitiveArray`].
+    /// # 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.validity = validity;
+        arr
+    }
+
     /// The values [`Buffer`].
     #[inline]
     pub fn values(&self) -> &Buffer<T> {
@@ -162,6 +174,9 @@ impl<T: NativeType> Array for PrimitiveArray<T> {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 /// A type definition [`PrimitiveArray`] for `i8`

src/array/struct_.rs

@@ -117,6 +117,18 @@ impl StructArray {
         }
     }
 
+    /// Sets the validity bitmap on this [`StructArray`].
+    /// # 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.validity = validity;
+        arr
+    }
+
     /// Returns the values of this [`StructArray`].
     pub fn values(&self) -> &[Arc<dyn Array>] {
         &self.values
@@ -163,6 +175,9 @@ impl Array for StructArray {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl std::fmt::Display for StructArray {

src/array/union/mod.rs

@@ -212,6 +212,9 @@ impl Array for UnionArray {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, _: Option<Bitmap>) -> Box<dyn Array> {
+        panic!("cannot set validity of a union array")
+    }
 }
 
 impl UnionArray {

src/array/utf8/mod.rs

@@ -150,6 +150,18 @@ impl<O: Offset> Utf8Array<O> {
         }
     }
 
+    /// Sets the validity bitmap on this [`Utf8Array`].
+    /// # 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.validity = validity;
+        arr
+    }
+
     /// Returns the element at index `i` as &str
     pub fn value(&self, i: usize) -> &str {
         let offsets = self.offsets.as_slice();
@@ -199,6 +211,9 @@ impl<O: Offset> Array for Utf8Array<O> {
     fn slice(&self, offset: usize, length: usize) -> Box<dyn Array> {
         Box::new(self.slice(offset, length))
     }
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
+    }
 }
 
 impl<O: Offset> std::fmt::Display for Utf8Array<O> {

tests/it/array/mod.rs

@@ -11,7 +11,9 @@ mod primitive;
 mod union;
 mod utf8;
 
-use arrow2::array::{clone, new_empty_array, new_null_array};
+use arrow2::array::{clone, new_empty_array, new_null_array, Array, PrimitiveArray};
+use arrow2::bitmap::Bitmap;
+use arrow2::datatypes::PhysicalType::Primitive;
 use arrow2::datatypes::{DataType, Field};
 
 #[test]
@@ -68,3 +70,14 @@ fn test_clone() {
         .all(|x| clone(new_null_array(x.clone(), 10).as_ref()) == new_null_array(x, 10));
     assert!(a);
 }
+
+#[test]
+fn test_with_validity() {
+    let arr = PrimitiveArray::from_slice(&[1i32, 2, 3]);
+    let validity = Bitmap::from(&[true, false, true]);
+    let arr = arr.with_validity(Some(validity));
+    let arr_ref = arr.as_any().downcast_ref::<PrimitiveArray<i32>>().unwrap();
+
+    let expected = PrimitiveArray::from(&[Some(1i32), None, Some(3)]);
+    assert_eq!(arr_ref, &expected);
+}