Added Array::with_validity

src/array/binary/mod.rs

@@ -1,4 +1,9 @@
-use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType};
+use crate::{
+    bitmap::Bitmap,
+    buffer::Buffer,
+    datatypes::DataType,
+    error::{ArrowError, Result},
+};
 
 use super::{
     display_fmt, display_helper, specification::check_offsets, specification::Offset, Array,
@@ -159,6 +164,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 impl<O: Offset> std::fmt::Display for BinaryArray<O> {

src/array/boolean/mod.rs

@@ -1,6 +1,7 @@
 use crate::{
     bitmap::Bitmap,
     datatypes::{DataType, PhysicalType},
+    error::{ArrowError, Result},
 };
 
 use super::{display_fmt, Array};
@@ -118,6 +119,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 impl std::fmt::Display for BooleanArray {

src/array/dictionary/mod.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use crate::{
     bitmap::Bitmap,
     datatypes::DataType,
+    error::Result,
     scalar::{new_scalar, Scalar},
     types::{NativeType, NaturalDataType},
 };
@@ -137,6 +138,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>) -> Result<Box<dyn Array>> {
+        self.values.with_validity(validity)
+    }
 }
 
 impl<K: DictionaryKey> std::fmt::Display for DictionaryArray<K>

src/array/fixed_size_binary/mod.rs

@@ -1,4 +1,9 @@
-use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType, error::Result};
+use crate::{
+    bitmap::Bitmap,
+    buffer::Buffer,
+    datatypes::DataType,
+    error::{ArrowError, Result},
+};
 
 use super::{display_fmt, display_helper, ffi::ToFfi, Array};
 
@@ -128,6 +133,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 impl std::fmt::Display for FixedSizeBinaryArray {

src/array/fixed_size_list/mod.rs

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use crate::{
     bitmap::Bitmap,
     datatypes::{DataType, Field},
+    error::{ArrowError, Result},
 };
 
 use super::{display_fmt, ffi::ToFfi, new_empty_array, new_null_array, Array};
@@ -131,6 +132,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 impl std::fmt::Display for FixedSizeListArray {

src/array/list/mod.rs

@@ -4,6 +4,7 @@ use crate::{
     bitmap::Bitmap,
     buffer::Buffer,
     datatypes::{DataType, Field},
+    error::{ArrowError, Result},
 };
 
 use super::{
@@ -174,6 +175,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 impl<O: Offset> std::fmt::Display for ListArray<O> {

src/array/mod.rs

@@ -88,6 +88,9 @@ 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`].
+    fn with_validity(&self, validity: Option<Bitmap>) -> Result<Box<dyn Array>>;
 }
 
 /// A trait describing a mutable array; i.e. an array whose values can be changed.

src/array/null.rs

@@ -1,4 +1,8 @@
-use crate::{bitmap::Bitmap, datatypes::DataType};
+use crate::{
+    bitmap::Bitmap,
+    datatypes::DataType,
+    error::{ArrowError, Result},
+};
 
 use super::{ffi::ToFfi, Array};
 
@@ -63,6 +67,11 @@ 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>) -> Result<Box<dyn Array>> {
+        Err(ArrowError::Other(
+            "cannot set validity of a null array".into(),
+        ))
+    }
 }
 
 impl std::fmt::Display for NullArray {

src/array/primitive/mod.rs

@@ -2,7 +2,7 @@ use crate::{
     bitmap::Bitmap,
     buffer::Buffer,
     datatypes::*,
-    error::ArrowError,
+    error::{ArrowError, Result},
     types::{days_ms, months_days_ns, NativeType},
 };
 
@@ -162,6 +162,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 /// A type definition [`PrimitiveArray`] for `i8`

src/array/struct_.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use crate::{
     bitmap::Bitmap,
     datatypes::{DataType, Field},
-    error::Result,
+    error::{ArrowError, Result},
     ffi,
 };
 
@@ -163,6 +163,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 impl std::fmt::Display for StructArray {

src/array/union/mod.rs

@@ -5,6 +5,7 @@ use crate::{
     bitmap::Bitmap,
     buffer::Buffer,
     datatypes::{DataType, Field},
+    error::{ArrowError, Result},
     scalar::{new_scalar, Scalar},
 };
 
@@ -212,6 +213,11 @@ 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>) -> Result<Box<dyn Array>> {
+        Err(ArrowError::Other(
+            "cannot set validity of a union array".into(),
+        ))
+    }
 }
 
 impl UnionArray {

src/array/utf8/mod.rs

@@ -1,4 +1,9 @@
-use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType};
+use crate::{
+    bitmap::Bitmap,
+    buffer::Buffer,
+    datatypes::DataType,
+    error::{ArrowError, Result},
+};
 
 use super::{
     display_fmt,
@@ -199,6 +204,16 @@ 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>) -> Result<Box<dyn Array>> {
+        if matches!(&validity, Some(bitmap) if bitmap.len() < self.len()) {
+            return Err(ArrowError::InvalidArgumentError(
+                "validity should be as least as large as the array".into(),
+            ));
+        }
+        let mut arr = self.clone();
+        arr.validity = validity;
+        Ok(Box::new(arr))
+    }
 }
 
 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)).unwrap();
+    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);
+}