From bd439411cc0f67c3bdd2aa7fe2a7a822aadb03e4 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Sun, 12 Sep 2021 14:49:58 +0200
Subject: [PATCH 1/3] Add Array::with_valdity

---
 src/array/binary/mod.rs            | 17 ++++++++++++++++-
 src/array/boolean/mod.rs           | 11 +++++++++++
 src/array/dictionary/mod.rs        |  4 ++++
 src/array/fixed_size_binary/mod.rs | 17 ++++++++++++++++-
 src/array/fixed_size_list/mod.rs   | 11 +++++++++++
 src/array/list/mod.rs              | 11 +++++++++++
 src/array/mod.rs                   |  3 +++
 src/array/null.rs                  | 11 ++++++++++-
 src/array/primitive/mod.rs         | 12 +++++++++++-
 src/array/struct_.rs               | 12 +++++++++++-
 src/array/union/mod.rs             |  6 ++++++
 src/array/utf8/mod.rs              | 17 ++++++++++++++++-
 tests/it/array/mod.rs              | 15 ++++++++++++++-
 13 files changed, 140 insertions(+), 7 deletions(-)

diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs
index a268d010957..efb9406a4e5 100644
--- a/src/array/binary/mod.rs
+++ b/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> {
diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
index 8c1207e3902..25e8b764f98 100644
--- a/src/array/boolean/mod.rs
+++ b/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 {
diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs
index f9d5668d969..8cea1f4b0a0 100644
--- a/src/array/dictionary/mod.rs
+++ b/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>
diff --git a/src/array/fixed_size_binary/mod.rs b/src/array/fixed_size_binary/mod.rs
index 401b24e8ec3..6d3e8a1c60f 100644
--- a/src/array/fixed_size_binary/mod.rs
+++ b/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 {
diff --git a/src/array/fixed_size_list/mod.rs b/src/array/fixed_size_list/mod.rs
index 6c8b2f34834..ad80c4b3f41 100644
--- a/src/array/fixed_size_list/mod.rs
+++ b/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 {
diff --git a/src/array/list/mod.rs b/src/array/list/mod.rs
index 670822b2aa6..867ebde18c2 100644
--- a/src/array/list/mod.rs
+++ b/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> {
diff --git a/src/array/mod.rs b/src/array/mod.rs
index 7403af5b381..3708b3e45c4 100644
--- a/src/array/mod.rs
+++ b/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.
diff --git a/src/array/null.rs b/src/array/null.rs
index dcddf92e512..bbaea76b3f3 100644
--- a/src/array/null.rs
+++ b/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 {
diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
index 90cdc27c11c..49297b7d991 100644
--- a/src/array/primitive/mod.rs
+++ b/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`
diff --git a/src/array/struct_.rs b/src/array/struct_.rs
index d773b51e1d9..9d2c2853337 100644
--- a/src/array/struct_.rs
+++ b/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 {
diff --git a/src/array/union/mod.rs b/src/array/union/mod.rs
index f00dcdff945..e5af1374f0f 100644
--- a/src/array/union/mod.rs
+++ b/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 {
diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs
index a5cb62b57a2..559dce1792e 100644
--- a/src/array/utf8/mod.rs
+++ b/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> {
diff --git a/tests/it/array/mod.rs b/tests/it/array/mod.rs
index 63e63d01b25..325e7968548 100644
--- a/tests/it/array/mod.rs
+++ b/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);
+}

From 99a80b777401d9ab80b2330e38fd725dce51ffb6 Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Mon, 13 Sep 2021 07:50:02 +0200
Subject: [PATCH 2/3] process comments

---
 src/array/binary/mod.rs            | 30 +++++++++++++++---------------
 src/array/boolean/mod.rs           | 24 ++++++++++++++----------
 src/array/dictionary/mod.rs        | 17 ++++++++++++++---
 src/array/fixed_size_binary/mod.rs | 30 +++++++++++++++---------------
 src/array/fixed_size_list/mod.rs   | 24 ++++++++++++++----------
 src/array/list/mod.rs              | 24 ++++++++++++++----------
 src/array/mod.rs                   |  4 +++-
 src/array/null.rs                  | 12 +++---------
 src/array/primitive/mod.rs         | 25 +++++++++++++++----------
 src/array/struct_.rs               | 25 +++++++++++++++----------
 src/array/union/mod.rs             |  7 ++-----
 src/array/utf8/mod.rs              | 30 +++++++++++++++---------------
 tests/it/array/mod.rs              |  2 +-
 13 files changed, 140 insertions(+), 114 deletions(-)

diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs
index efb9406a4e5..a977c68e323 100644
--- a/src/array/binary/mod.rs
+++ b/src/array/binary/mod.rs
@@ -1,9 +1,4 @@
-use crate::{
-    bitmap::Bitmap,
-    buffer::Buffer,
-    datatypes::DataType,
-    error::{ArrowError, Result},
-};
+use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType};
 
 use super::{
     display_fmt, display_helper, specification::check_offsets, specification::Offset, Array,
@@ -99,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
@@ -164,15 +171,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
index 25e8b764f98..d7a868a4654 100644
--- a/src/array/boolean/mod.rs
+++ b/src/array/boolean/mod.rs
@@ -1,7 +1,6 @@
 use crate::{
     bitmap::Bitmap,
     datatypes::{DataType, PhysicalType},
-    error::{ArrowError, Result},
 };
 
 use super::{display_fmt, Array};
@@ -92,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 {
@@ -119,15 +130,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs
index 8cea1f4b0a0..925455b73bd 100644
--- a/src/array/dictionary/mod.rs
+++ b/src/array/dictionary/mod.rs
@@ -3,7 +3,6 @@ use std::sync::Arc;
 use crate::{
     bitmap::Bitmap,
     datatypes::DataType,
-    error::Result,
     scalar::{new_scalar, Scalar},
     types::{NativeType, NaturalDataType},
 };
@@ -84,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]
@@ -138,8 +149,8 @@ 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)
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/fixed_size_binary/mod.rs b/src/array/fixed_size_binary/mod.rs
index 6d3e8a1c60f..cdb385fbfd4 100644
--- a/src/array/fixed_size_binary/mod.rs
+++ b/src/array/fixed_size_binary/mod.rs
@@ -1,9 +1,4 @@
-use crate::{
-    bitmap::Bitmap,
-    buffer::Buffer,
-    datatypes::DataType,
-    error::{ArrowError, Result},
-};
+use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType, error::Result};
 
 use super::{display_fmt, display_helper, ffi::ToFfi, Array};
 
@@ -98,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 {
@@ -133,15 +140,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/fixed_size_list/mod.rs b/src/array/fixed_size_list/mod.rs
index ad80c4b3f41..7c31589d98d 100644
--- a/src/array/fixed_size_list/mod.rs
+++ b/src/array/fixed_size_list/mod.rs
@@ -3,7 +3,6 @@ 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};
@@ -91,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 {
@@ -132,15 +143,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/list/mod.rs b/src/array/list/mod.rs
index 867ebde18c2..576362a610b 100644
--- a/src/array/list/mod.rs
+++ b/src/array/list/mod.rs
@@ -4,7 +4,6 @@ use crate::{
     bitmap::Bitmap,
     buffer::Buffer,
     datatypes::{DataType, Field},
-    error::{ArrowError, Result},
 };
 
 use super::{
@@ -118,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> {
@@ -175,15 +186,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/mod.rs b/src/array/mod.rs
index 3708b3e45c4..fa5ea4d0794 100644
--- a/src/array/mod.rs
+++ b/src/array/mod.rs
@@ -90,7 +90,9 @@ pub trait Array: std::fmt::Debug + Send + Sync {
     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>>;
+    /// # 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.
diff --git a/src/array/null.rs b/src/array/null.rs
index bbaea76b3f3..d7c339f29e9 100644
--- a/src/array/null.rs
+++ b/src/array/null.rs
@@ -1,8 +1,4 @@
-use crate::{
-    bitmap::Bitmap,
-    datatypes::DataType,
-    error::{ArrowError, Result},
-};
+use crate::{bitmap::Bitmap, datatypes::DataType};
 
 use super::{ffi::ToFfi, Array};
 
@@ -67,10 +63,8 @@ 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(),
-        ))
+    fn with_validity(&self, _: Option<Bitmap>) -> Box<dyn Array> {
+        panic!("cannot set validity of a null array")
     }
 }
 
diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
index 49297b7d991..8e14a41e8e7 100644
--- a/src/array/primitive/mod.rs
+++ b/src/array/primitive/mod.rs
@@ -2,7 +2,7 @@ use crate::{
     bitmap::Bitmap,
     buffer::Buffer,
     datatypes::*,
-    error::{ArrowError, Result},
+    error::ArrowError,
     types::{days_ms, months_days_ns, NativeType},
 };
 
@@ -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,15 +174,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/struct_.rs b/src/array/struct_.rs
index 9d2c2853337..c2cbb8a2094 100644
--- a/src/array/struct_.rs
+++ b/src/array/struct_.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use crate::{
     bitmap::Bitmap,
     datatypes::{DataType, Field},
-    error::{ArrowError, Result},
+    error::Result,
     ffi,
 };
 
@@ -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,15 +175,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/src/array/union/mod.rs b/src/array/union/mod.rs
index e5af1374f0f..68b93002222 100644
--- a/src/array/union/mod.rs
+++ b/src/array/union/mod.rs
@@ -5,7 +5,6 @@ use crate::{
     bitmap::Bitmap,
     buffer::Buffer,
     datatypes::{DataType, Field},
-    error::{ArrowError, Result},
     scalar::{new_scalar, Scalar},
 };
 
@@ -213,10 +212,8 @@ 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(),
-        ))
+    fn with_validity(&self, _: Option<Bitmap>) -> Box<dyn Array> {
+        panic!("cannot set validity of a union array")
     }
 }
 
diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs
index 559dce1792e..cdd93f8d671 100644
--- a/src/array/utf8/mod.rs
+++ b/src/array/utf8/mod.rs
@@ -1,9 +1,4 @@
-use crate::{
-    bitmap::Bitmap,
-    buffer::Buffer,
-    datatypes::DataType,
-    error::{ArrowError, Result},
-};
+use crate::{bitmap::Bitmap, buffer::Buffer, datatypes::DataType};
 
 use super::{
     display_fmt,
@@ -155,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();
@@ -204,15 +211,8 @@ 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))
+    fn with_validity(&self, validity: Option<Bitmap>) -> Box<dyn Array> {
+        Box::new(self.with_validity(validity))
     }
 }
 
diff --git a/tests/it/array/mod.rs b/tests/it/array/mod.rs
index 325e7968548..9b367ce7c94 100644
--- a/tests/it/array/mod.rs
+++ b/tests/it/array/mod.rs
@@ -75,7 +75,7 @@ fn test_clone() {
 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 = 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)]);

From 493ddd47530a39add630a93f0cd596419f44163b Mon Sep 17 00:00:00 2001
From: Ritchie Vink <ritchie46@gmail.com>
Date: Mon, 13 Sep 2021 10:35:25 +0200
Subject: [PATCH 3/3] tighten invariants

---
 src/array/binary/mod.rs            | 4 ++--
 src/array/boolean/mod.rs           | 4 ++--
 src/array/dictionary/mod.rs        | 4 ++--
 src/array/fixed_size_binary/mod.rs | 4 ++--
 src/array/fixed_size_list/mod.rs   | 4 ++--
 src/array/list/mod.rs              | 4 ++--
 src/array/primitive/mod.rs         | 4 ++--
 src/array/struct_.rs               | 4 ++--
 src/array/utf8/mod.rs              | 4 ++--
 9 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs
index a977c68e323..607e7bd3eae 100644
--- a/src/array/binary/mod.rs
+++ b/src/array/binary/mod.rs
@@ -97,9 +97,9 @@ impl<O: Offset> BinaryArray<O> {
 
     /// Sets the validity bitmap on this [`BinaryArray`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
index d7a868a4654..dfc8134f771 100644
--- a/src/array/boolean/mod.rs
+++ b/src/array/boolean/mod.rs
@@ -94,9 +94,9 @@ impl BooleanArray {
 
     /// Sets the validity bitmap on this [`BooleanArray`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs
index 925455b73bd..e938b1e224b 100644
--- a/src/array/dictionary/mod.rs
+++ b/src/array/dictionary/mod.rs
@@ -85,9 +85,9 @@ impl<K: DictionaryKey> DictionaryArray<K> {
 
     /// Sets the validity bitmap on this [`Array`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/fixed_size_binary/mod.rs b/src/array/fixed_size_binary/mod.rs
index cdb385fbfd4..1bd31e6fef6 100644
--- a/src/array/fixed_size_binary/mod.rs
+++ b/src/array/fixed_size_binary/mod.rs
@@ -96,9 +96,9 @@ impl FixedSizeBinaryArray {
 
     /// Sets the validity bitmap on this [`FixedSizeBinaryArray`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/fixed_size_list/mod.rs b/src/array/fixed_size_list/mod.rs
index 7c31589d98d..bc7ae65e025 100644
--- a/src/array/fixed_size_list/mod.rs
+++ b/src/array/fixed_size_list/mod.rs
@@ -93,9 +93,9 @@ impl FixedSizeListArray {
 
     /// Sets the validity bitmap on this [`FixedSizeListArray`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/list/mod.rs b/src/array/list/mod.rs
index 576362a610b..1e84981df86 100644
--- a/src/array/list/mod.rs
+++ b/src/array/list/mod.rs
@@ -120,9 +120,9 @@ impl<O: Offset> ListArray<O> {
 
     /// Sets the validity bitmap on this [`ListArray`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
index 8e14a41e8e7..074047ec35d 100644
--- a/src/array/primitive/mod.rs
+++ b/src/array/primitive/mod.rs
@@ -96,9 +96,9 @@ impl<T: NativeType> PrimitiveArray<T> {
 
     /// Sets the validity bitmap on this [`PrimitiveArray`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/struct_.rs b/src/array/struct_.rs
index c2cbb8a2094..ddb18724135 100644
--- a/src/array/struct_.rs
+++ b/src/array/struct_.rs
@@ -119,9 +119,9 @@ impl StructArray {
 
     /// Sets the validity bitmap on this [`StructArray`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();
diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs
index cdd93f8d671..530858c7e25 100644
--- a/src/array/utf8/mod.rs
+++ b/src/array/utf8/mod.rs
@@ -152,9 +152,9 @@ impl<O: Offset> Utf8Array<O> {
 
     /// Sets the validity bitmap on this [`Utf8Array`].
     /// # Panic
-    /// This function panics iff `validity.len() < self.len()`.
+    /// 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()) {
+        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();