Allow nullable children of MaskedArray

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-array/src/arrays/bool/mod.rs

@@ -9,7 +9,6 @@ pub use array::*;
 pub mod compute;
 
 mod vtable;
-pub use vtable::BoolMaskedValidityRule;
 pub use vtable::BoolVTable;
 
 #[cfg(feature = "test-harness")]

vortex-array/src/arrays/bool/vtable/mod.rs

@@ -3,15 +3,12 @@
 
 use vortex_buffer::BufferHandle;
 use vortex_dtype::DType;
+use vortex_error::vortex_bail;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
-use vortex_error::vortex_bail;
-use vortex_vector::Vector;
 use vortex_vector::bool::BoolVector;
+use vortex_vector::Vector;
 
-use crate::DeserializeMetadata;
-use crate::ProstMetadata;
-use crate::SerializeMetadata;
 use crate::arrays::BoolArray;
 use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
@@ -21,16 +18,16 @@ use crate::vtable::ArrayVTableExt;
 use crate::vtable::NotSupported;
 use crate::vtable::VTable;
 use crate::vtable::ValidityVTableFromValidityHelper;
+use crate::DeserializeMetadata;
+use crate::ProstMetadata;
+use crate::SerializeMetadata;
 
 mod array;
 mod canonical;
 mod operations;
-pub mod operator;
 mod validity;
 mod visitor;
 
-pub use operator::BoolMaskedValidityRule;
-
 use crate::vtable::ArrayId;
 use crate::vtable::ArrayVTable;
 

vortex-array/src/arrays/bool/vtable/operator.rs

@@ -7,7 +7,6 @@ pub use array::DecimalArray;
 mod compute;
 
 mod vtable;
-pub use vtable::DecimalMaskedValidityRule;
 pub use vtable::DecimalVTable;
 
 mod utils;

vortex-array/src/arrays/decimal/mod.rs

@@ -4,20 +4,17 @@
 use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
 use vortex_buffer::BufferHandle;
+use vortex_dtype::match_each_decimal_value_type;
 use vortex_dtype::DType;
 use vortex_dtype::NativeDecimalType;
 use vortex_dtype::PrecisionScale;
-use vortex_dtype::match_each_decimal_value_type;
-use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
+use vortex_error::VortexResult;
 use vortex_scalar::DecimalType;
-use vortex_vector::Vector;
 use vortex_vector::decimal::DVector;
+use vortex_vector::Vector;
 
-use crate::DeserializeMetadata;
-use crate::ProstMetadata;
-use crate::SerializeMetadata;
 use crate::arrays::DecimalArray;
 use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
@@ -27,16 +24,16 @@ use crate::vtable::ArrayVTableExt;
 use crate::vtable::NotSupported;
 use crate::vtable::VTable;
 use crate::vtable::ValidityVTableFromValidityHelper;
+use crate::DeserializeMetadata;
+use crate::ProstMetadata;
+use crate::SerializeMetadata;
 
 mod array;
 mod canonical;
 mod operations;
-pub mod operator;
 mod validity;
 mod visitor;
 
-pub use operator::DecimalMaskedValidityRule;
-
 use crate::vtable::ArrayId;
 use crate::vtable::ArrayVTable;
 
@@ -142,18 +139,18 @@ pub struct DecimalVTable;
 
 #[cfg(test)]
 mod tests {
-    use vortex_buffer::ByteBufferMut;
     use vortex_buffer::buffer;
+    use vortex_buffer::ByteBufferMut;
     use vortex_dtype::DecimalDType;
 
-    use crate::ArrayContext;
-    use crate::IntoArray;
     use crate::arrays::DecimalArray;
     use crate::arrays::DecimalVTable;
     use crate::serde::ArrayParts;
     use crate::serde::SerializeOptions;
     use crate::validity::Validity;
     use crate::vtable::ArrayVTableExt;
+    use crate::ArrayContext;
+    use crate::IntoArray;
 
     #[test]
     fn test_array_serde() {

vortex-array/src/arrays/decimal/vtable/mod.rs

@@ -1,14 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::ops::{BitAnd, Not};
 use vortex_dtype::DType;
-use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_error::VortexResult;
 
-use crate::ArrayRef;
 use crate::compute::mask;
 use crate::stats::ArrayStats;
 use crate::validity::Validity;
+use crate::{Array, ArrayRef};
 
 #[derive(Clone, Debug)]
 pub struct MaskedArray {
@@ -20,14 +21,6 @@ pub struct MaskedArray {
 
 impl MaskedArray {
     pub fn try_new(child: ArrayRef, validity: Validity) -> VortexResult<Self> {
-        if matches!(validity, Validity::NonNullable) {
-            vortex_bail!("MaskedArray must have nullable validity, got {validity:?}")
-        }
-
-        if !child.all_valid() {
-            vortex_bail!("MaskedArray children must not have nulls");
-        }
-
         if let Some(validity_len) = validity.maybe_len()
             && validity_len != child.len()
         {
@@ -50,9 +43,16 @@ impl MaskedArray {
         &self.child
     }
 
+    pub fn mask(&self) -> &Validity {
+        &self.validity
+    }
+
     pub(crate) fn masked_child(&self) -> VortexResult<ArrayRef> {
-        // Invert the validity mask - we want to set values to null where validity is false.
-        let inverted_mask = !self.validity.to_mask(self.len());
-        mask(&self.child, &inverted_mask)
+        let intersected_validity = self
+            .child
+            .validity_mask()
+            .bitand(&self.validity.to_mask(self.len()));
+        // Note: the compute function takes the validity inverted!!
+        mask(&self.child, &intersected_validity.not())
     }
 }

vortex-array/src/arrays/decimal/vtable/operator.rs

@@ -3,19 +3,19 @@
 
 use vortex_error::VortexResult;
 
-use crate::Array;
-use crate::ArrayRef;
-use crate::IntoArray;
 use crate::arrays::BoolArray;
 use crate::arrays::MaskedArray;
 use crate::arrays::MaskedVTable;
 use crate::canonical::ToCanonical;
+use crate::compute::compare;
 use crate::compute::CompareKernel;
 use crate::compute::CompareKernelAdapter;
 use crate::compute::Operator;
-use crate::compute::compare;
 use crate::register_kernel;
 use crate::vtable::ValidityHelper;
+use crate::Array;
+use crate::ArrayRef;
+use crate::IntoArray;
 
 impl CompareKernel for MaskedVTable {
     fn compare(
@@ -29,7 +29,7 @@ impl CompareKernel for MaskedVTable {
 
         // Get the boolean buffer from the comparison result
         let bool_array = compare_result.to_bool();
-        let combined_validity = bool_array.validity().clone().and(lhs.validity().clone());
+        let combined_validity = bool_array.validity().clone().and(lhs.mask().clone());
 
         // Return a plain BoolArray with the combined validity
         Ok(Some(
@@ -47,14 +47,14 @@ mod tests {
     use vortex_mask::Mask;
     use vortex_scalar::Scalar;
 
-    use crate::IntoArray;
-    use crate::ToCanonical;
     use crate::arrays::ConstantArray;
     use crate::arrays::MaskedArray;
     use crate::arrays::PrimitiveArray;
-    use crate::compute::Operator;
     use crate::compute::compare;
+    use crate::compute::Operator;
     use crate::validity::Validity;
+    use crate::IntoArray;
+    use crate::ToCanonical;
 
     #[test]
     fn test_compare_value() {

vortex-array/src/arrays/masked/array.rs

@@ -4,27 +4,26 @@
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
-use crate::ArrayRef;
-use crate::IntoArray;
 use crate::arrays::MaskedArray;
 use crate::arrays::MaskedVTable;
+use crate::compute::filter;
 use crate::compute::FilterKernel;
 use crate::compute::FilterKernelAdapter;
-use crate::compute::filter;
 use crate::register_kernel;
-use crate::vtable::ValidityHelper;
+use crate::ArrayRef;
+use crate::IntoArray;
 
 impl FilterKernel for MaskedVTable {
     fn filter(&self, array: &MaskedArray, mask: &Mask) -> VortexResult<ArrayRef> {
         // Filter the validity to get the new validity
-        let filtered_validity = array.validity().filter(mask)?;
+        let filtered_mask = array.mask().filter(mask)?;
 
         // Filter the child array
         // The child is guaranteed to have no nulls, so filtering it is straightforward
         let filtered_child = filter(&array.child, mask)?;
 
         // Construct new MaskedArray
-        Ok(MaskedArray::try_new(filtered_child, filtered_validity)?.into_array())
+        Ok(MaskedArray::try_new(filtered_child, filtered_mask)?.into_array())
     }
 }
 
@@ -34,11 +33,11 @@ register_kernel!(FilterKernelAdapter(MaskedVTable).lift());
 mod tests {
     use rstest::rstest;
 
-    use crate::IntoArray;
     use crate::arrays::MaskedArray;
     use crate::arrays::PrimitiveArray;
     use crate::compute::conformance::filter::test_filter_conformance;
     use crate::validity::Validity;
+    use crate::IntoArray;
 
     #[rstest]
     #[case(