diff --git a/src/array/binary/mod.rs b/src/array/binary/mod.rs index a268d010957..607e7bd3eae 100644 --- a/src/array/binary/mod.rs +++ b/src/array/binary/mod.rs @@ -94,6 +94,18 @@ impl BinaryArray { 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) -> 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 Array for BinaryArray { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } impl std::fmt::Display for BinaryArray { diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs index 8c1207e3902..dfc8134f771 100644 --- a/src/array/boolean/mod.rs +++ b/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) -> 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 { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } impl std::fmt::Display for BooleanArray { diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs index f9d5668d969..e938b1e224b 100644 --- a/src/array/dictionary/mod.rs +++ b/src/array/dictionary/mod.rs @@ -83,6 +83,18 @@ impl DictionaryArray { } } + /// Sets the validity bitmap on this [`Array`]. + /// # Panic + /// This function panics iff `validity.len() != self.len()`. + pub fn with_validity(&self, validity: Option) -> 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 Array for DictionaryArray { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } impl std::fmt::Display for DictionaryArray diff --git a/src/array/fixed_size_binary/mod.rs b/src/array/fixed_size_binary/mod.rs index 401b24e8ec3..1bd31e6fef6 100644 --- a/src/array/fixed_size_binary/mod.rs +++ b/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) -> 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 { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } 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..bc7ae65e025 100644 --- a/src/array/fixed_size_list/mod.rs +++ b/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) -> 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 { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } impl std::fmt::Display for FixedSizeListArray { diff --git a/src/array/list/mod.rs b/src/array/list/mod.rs index 670822b2aa6..1e84981df86 100644 --- a/src/array/list/mod.rs +++ b/src/array/list/mod.rs @@ -117,6 +117,18 @@ impl ListArray { pub fn values(&self) -> &Arc { &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) -> 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 ListArray { @@ -174,6 +186,9 @@ impl Array for ListArray { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } impl std::fmt::Display for ListArray { diff --git a/src/array/mod.rs b/src/array/mod.rs index 7403af5b381..fa5ea4d0794 100644 --- a/src/array/mod.rs +++ b/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; + + /// Sets the validity bitmap on this [`Array`]. + /// # Panic + /// This function panics iff `validity.len() < self.len()`. + fn with_validity(&self, validity: Option) -> Box; } /// 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..d7c339f29e9 100644 --- a/src/array/null.rs +++ b/src/array/null.rs @@ -63,6 +63,9 @@ impl Array for NullArray { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, _: Option) -> Box { + panic!("cannot set validity of a null array") + } } impl std::fmt::Display for NullArray { diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs index 90cdc27c11c..074047ec35d 100644 --- a/src/array/primitive/mod.rs +++ b/src/array/primitive/mod.rs @@ -94,6 +94,18 @@ impl PrimitiveArray { } } + /// Sets the validity bitmap on this [`PrimitiveArray`]. + /// # Panic + /// This function panics iff `validity.len() != self.len()`. + pub fn with_validity(&self, validity: Option) -> 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 { @@ -162,6 +174,9 @@ impl Array for PrimitiveArray { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } /// A type definition [`PrimitiveArray`] for `i8` diff --git a/src/array/struct_.rs b/src/array/struct_.rs index d773b51e1d9..ddb18724135 100644 --- a/src/array/struct_.rs +++ b/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) -> 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] { &self.values @@ -163,6 +175,9 @@ impl Array for StructArray { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } impl std::fmt::Display for StructArray { diff --git a/src/array/union/mod.rs b/src/array/union/mod.rs index f00dcdff945..68b93002222 100644 --- a/src/array/union/mod.rs +++ b/src/array/union/mod.rs @@ -212,6 +212,9 @@ impl Array for UnionArray { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, _: Option) -> Box { + panic!("cannot set validity of a union array") + } } impl UnionArray { diff --git a/src/array/utf8/mod.rs b/src/array/utf8/mod.rs index a5cb62b57a2..530858c7e25 100644 --- a/src/array/utf8/mod.rs +++ b/src/array/utf8/mod.rs @@ -150,6 +150,18 @@ impl Utf8Array { } } + /// Sets the validity bitmap on this [`Utf8Array`]. + /// # Panic + /// This function panics iff `validity.len() != self.len()`. + pub fn with_validity(&self, validity: Option) -> 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 Array for Utf8Array { fn slice(&self, offset: usize, length: usize) -> Box { Box::new(self.slice(offset, length)) } + fn with_validity(&self, validity: Option) -> Box { + Box::new(self.with_validity(validity)) + } } impl std::fmt::Display for Utf8Array { diff --git a/tests/it/array/mod.rs b/tests/it/array/mod.rs index 63e63d01b25..9b367ce7c94 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)); + let arr_ref = arr.as_any().downcast_ref::>().unwrap(); + + let expected = PrimitiveArray::from(&[Some(1i32), None, Some(3)]); + assert_eq!(arr_ref, &expected); +}