break[array]: mask compute to lazy mask execution

encodings/alp/public-api.lock

@@ -160,14 +160,14 @@ impl vortex_array::arrays::slice::SliceKernel for vortex_alp::ALPRDVTable
 
 pub fn vortex_alp::ALPRDVTable::slice(array: &Self::Array, range: core::ops::range::Range<usize>, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
-impl vortex_array::compute::mask::MaskKernel for vortex_alp::ALPRDVTable
-
-pub fn vortex_alp::ALPRDVTable::mask(&self, array: &vortex_alp::ALPRDArray, filter_mask: &vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
-
 impl vortex_array::expr::exprs::cast::kernel::CastReduce for vortex_alp::ALPRDVTable
 
 pub fn vortex_alp::ALPRDVTable::cast(array: &vortex_alp::ALPRDArray, dtype: &vortex_dtype::dtype::DType) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
+impl vortex_array::expr::exprs::mask::kernel::MaskReduce for vortex_alp::ALPRDVTable
+
+pub fn vortex_alp::ALPRDVTable::mask(array: &vortex_alp::ALPRDArray, mask: &vortex_array::array::ArrayRef) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
 impl vortex_array::vtable::VTable for vortex_alp::ALPRDVTable
 
 pub type vortex_alp::ALPRDVTable::Array = vortex_alp::ALPRDArray
@@ -256,10 +256,6 @@ impl vortex_array::compute::between::BetweenKernel for vortex_alp::ALPVTable
 
 pub fn vortex_alp::ALPVTable::between(&self, array: &vortex_alp::ALPArray, lower: &dyn vortex_array::array::Array, upper: &dyn vortex_array::array::Array, options: &vortex_array::compute::between::BetweenOptions) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
-impl vortex_array::compute::mask::MaskKernel for vortex_alp::ALPVTable
-
-pub fn vortex_alp::ALPVTable::mask(&self, array: &vortex_alp::ALPArray, filter_mask: &vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
-
 impl vortex_array::compute::nan_count::NaNCountKernel for vortex_alp::ALPVTable
 
 pub fn vortex_alp::ALPVTable::nan_count(&self, array: &vortex_alp::ALPArray) -> vortex_error::VortexResult<usize>
@@ -272,6 +268,14 @@ impl vortex_array::expr::exprs::cast::kernel::CastReduce for vortex_alp::ALPVTab
 
 pub fn vortex_alp::ALPVTable::cast(array: &vortex_alp::ALPArray, dtype: &vortex_dtype::dtype::DType) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
+impl vortex_array::expr::exprs::mask::kernel::MaskKernel for vortex_alp::ALPVTable
+
+pub fn vortex_alp::ALPVTable::mask(array: &vortex_alp::ALPArray, mask: &vortex_array::array::ArrayRef, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
+impl vortex_array::expr::exprs::mask::kernel::MaskReduce for vortex_alp::ALPVTable
+
+pub fn vortex_alp::ALPVTable::mask(array: &vortex_alp::ALPArray, mask: &vortex_array::array::ArrayRef) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
 impl vortex_array::vtable::VTable for vortex_alp::ALPVTable
 
 pub type vortex_alp::ALPVTable::Array = vortex_alp::ALPArray

encodings/alp/src/alp/compute/mask.rs

@@ -2,38 +2,48 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::compute::MaskKernel;
-use vortex_array::compute::MaskKernelAdapter;
-use vortex_array::compute::mask;
-use vortex_array::register_kernel;
+use vortex_array::compute::MaskReduce;
+use vortex_array::validity::Validity;
 use vortex_error::VortexResult;
-use vortex_mask::Mask;
 
 use crate::ALPArray;
 use crate::ALPVTable;
 
+impl MaskReduce for ALPVTable {
+    fn mask(array: &ALPArray, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
+        // Masking sparse patches requires reading indices, fall back to kernel.
+        if array.patches().is_some() {
+            return Ok(None);
+        }
+        let masked_encoded = array.encoded().clone().mask(mask.clone())?;
+        Ok(Some(
+            ALPArray::new(masked_encoded, array.exponents(), None).to_array(),
+        ))
+    }
+}
+
 impl MaskKernel for ALPVTable {
-    fn mask(&self, array: &ALPArray, filter_mask: &Mask) -> VortexResult<ArrayRef> {
-        let masked_encoded = mask(array.encoded(), filter_mask)?;
+    fn mask(
+        array: &ALPArray,
+        mask: &ArrayRef,
+        _ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Option<ArrayRef>> {
+        let vortex_mask = Validity::Array(mask.not()?).to_mask(array.len());
+        let masked_encoded = array.encoded().clone().mask(mask.clone())?;
         let masked_patches = array
             .patches()
-            .map(|p| p.mask(filter_mask))
+            .map(|p| p.mask(&vortex_mask))
             .transpose()?
-            .flatten()
-            .map(|patches| {
-                patches.cast_values(
-                    &array
-                        .dtype()
-                        .with_nullability(masked_encoded.dtype().nullability()),
-                )
-            })
-            .transpose()?;
-        Ok(ALPArray::new(masked_encoded, array.exponents(), masked_patches).to_array())
+            .flatten();
+        Ok(Some(
+            ALPArray::new(masked_encoded, array.exponents(), masked_patches).to_array(),
+        ))
     }
 }
 
-register_kernel!(MaskKernelAdapter(ALPVTable).lift());
-
 #[cfg(test)]
 mod test {
     use rstest::rstest;
@@ -58,4 +68,17 @@ mod test {
         let alp = alp_encode(&array.to_primitive(), None).unwrap();
         test_mask_conformance(alp.as_ref());
     }
+
+    #[test]
+    fn test_mask_alp_with_patches() {
+        use std::f64::consts::PI;
+        // PI doesn't encode cleanly with ALP, so it creates patches.
+        let values: Vec<f64> = (0..100)
+            .map(|i| if i % 4 == 3 { PI } else { 1.0 })
+            .collect();
+        let array = PrimitiveArray::from_iter(values);
+        let alp = alp_encode(&array, None).unwrap();
+        assert!(alp.patches().is_some(), "expected patches");
+        test_mask_conformance(alp.as_ref());
+    }
 }

encodings/alp/src/alp/rules.rs

@@ -5,6 +5,8 @@ use vortex_array::arrays::FilterExecuteAdaptor;
 use vortex_array::arrays::SliceExecuteAdaptor;
 use vortex_array::arrays::TakeExecuteAdaptor;
 use vortex_array::compute::CastReduceAdaptor;
+use vortex_array::compute::MaskExecuteAdaptor;
+use vortex_array::compute::MaskReduceAdaptor;
 use vortex_array::expr::CompareExecuteAdaptor;
 use vortex_array::kernel::ParentKernelSet;
 use vortex_array::optimizer::rules::ParentRuleSet;
@@ -14,9 +16,12 @@ use crate::ALPVTable;
 pub(super) const PARENT_KERNELS: ParentKernelSet<ALPVTable> = ParentKernelSet::new(&[
     ParentKernelSet::lift(&CompareExecuteAdaptor(ALPVTable)),
     ParentKernelSet::lift(&FilterExecuteAdaptor(ALPVTable)),
+    ParentKernelSet::lift(&MaskExecuteAdaptor(ALPVTable)),
     ParentKernelSet::lift(&SliceExecuteAdaptor(ALPVTable)),
     ParentKernelSet::lift(&TakeExecuteAdaptor(ALPVTable)),
 ]);
 
-pub(super) const RULES: ParentRuleSet<ALPVTable> =
-    ParentRuleSet::new(&[ParentRuleSet::lift(&CastReduceAdaptor(ALPVTable))]);
+pub(super) const RULES: ParentRuleSet<ALPVTable> = ParentRuleSet::new(&[
+    ParentRuleSet::lift(&CastReduceAdaptor(ALPVTable)),
+    ParentRuleSet::lift(&MaskReduceAdaptor(ALPVTable)),
+]);

encodings/alp/src/alp_rd/compute/mask.rs

@@ -3,32 +3,36 @@
 
 use vortex_array::ArrayRef;
 use vortex_array::IntoArray;
-use vortex_array::compute::MaskKernel;
-use vortex_array::compute::MaskKernelAdapter;
-use vortex_array::compute::mask;
-use vortex_array::register_kernel;
+use vortex_array::arrays::ScalarFnArrayExt;
+use vortex_array::compute::MaskReduce;
+use vortex_array::expr::EmptyOptions;
+use vortex_array::expr::Mask as MaskExpr;
 use vortex_error::VortexResult;
-use vortex_mask::Mask;
 
 use crate::ALPRDArray;
 use crate::ALPRDVTable;
 
-impl MaskKernel for ALPRDVTable {
-    fn mask(&self, array: &ALPRDArray, filter_mask: &Mask) -> VortexResult<ArrayRef> {
-        Ok(ALPRDArray::try_new(
-            array.dtype().as_nullable(),
-            mask(array.left_parts(), filter_mask)?,
-            array.left_parts_dictionary().clone(),
-            array.right_parts().clone(),
-            array.right_bit_width(),
-            array.left_parts_patches().cloned(),
-        )?
-        .into_array())
+impl MaskReduce for ALPRDVTable {
+    fn mask(array: &ALPRDArray, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
+        let masked_left_parts = MaskExpr.try_new_array(
+            array.left_parts().len(),
+            EmptyOptions,
+            [array.left_parts().clone(), mask.clone()],
+        )?;
+        Ok(Some(
+            ALPRDArray::try_new(
+                array.dtype().as_nullable(),
+                masked_left_parts,
+                array.left_parts_dictionary().clone(),
+                array.right_parts().clone(),
+                array.right_bit_width(),
+                array.left_parts_patches().cloned(),
+            )?
+            .into_array(),
+        ))
     }
 }
 
-register_kernel!(MaskKernelAdapter(ALPRDVTable).lift());
-
 #[cfg(test)]
 mod tests {
     use rstest::rstest;

encodings/alp/src/alp_rd/rules.rs

@@ -2,9 +2,12 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_array::compute::CastReduceAdaptor;
+use vortex_array::compute::MaskReduceAdaptor;
 use vortex_array::optimizer::rules::ParentRuleSet;
 
 use crate::alp_rd::ALPRDVTable;
 
-pub(crate) static RULES: ParentRuleSet<ALPRDVTable> =
-    ParentRuleSet::new(&[ParentRuleSet::lift(&CastReduceAdaptor(ALPRDVTable))]);
+pub(crate) static RULES: ParentRuleSet<ALPRDVTable> = ParentRuleSet::new(&[
+    ParentRuleSet::lift(&CastReduceAdaptor(ALPRDVTable)),
+    ParentRuleSet::lift(&MaskReduceAdaptor(ALPRDVTable)),
+]);

encodings/bytebool/public-api.lock

@@ -68,14 +68,14 @@ impl vortex_array::arrays::slice::SliceReduce for vortex_bytebool::ByteBoolVTabl
 
 pub fn vortex_bytebool::ByteBoolVTable::slice(array: &vortex_bytebool::ByteBoolArray, range: core::ops::range::Range<usize>) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
-impl vortex_array::compute::mask::MaskKernel for vortex_bytebool::ByteBoolVTable
-
-pub fn vortex_bytebool::ByteBoolVTable::mask(&self, array: &vortex_bytebool::ByteBoolArray, mask: &vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
-
 impl vortex_array::expr::exprs::cast::kernel::CastReduce for vortex_bytebool::ByteBoolVTable
 
 pub fn vortex_bytebool::ByteBoolVTable::cast(array: &vortex_bytebool::ByteBoolArray, dtype: &vortex_dtype::dtype::DType) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
+impl vortex_array::expr::exprs::mask::kernel::MaskReduce for vortex_bytebool::ByteBoolVTable
+
+pub fn vortex_bytebool::ByteBoolVTable::mask(array: &vortex_bytebool::ByteBoolArray, mask: &vortex_array::array::ArrayRef) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
 impl vortex_array::vtable::VTable for vortex_bytebool::ByteBoolVTable
 
 pub type vortex_bytebool::ByteBoolVTable::Array = vortex_bytebool::ByteBoolArray

encodings/bytebool/src/compute.rs

@@ -9,14 +9,12 @@ use vortex_array::IntoArray;
 use vortex_array::ToCanonical;
 use vortex_array::arrays::TakeExecute;
 use vortex_array::compute::CastReduce;
-use vortex_array::compute::MaskKernel;
-use vortex_array::compute::MaskKernelAdapter;
-use vortex_array::register_kernel;
+use vortex_array::compute::MaskReduce;
+use vortex_array::validity::Validity;
 use vortex_array::vtable::ValidityHelper;
 use vortex_dtype::DType;
 use vortex_dtype::match_each_integer_ptype;
 use vortex_error::VortexResult;
-use vortex_mask::Mask;
 
 use super::ByteBoolArray;
 use super::ByteBoolVTable;
@@ -44,14 +42,21 @@ impl CastReduce for ByteBoolVTable {
     }
 }
 
-impl MaskKernel for ByteBoolVTable {
-    fn mask(&self, array: &ByteBoolArray, mask: &Mask) -> VortexResult<ArrayRef> {
-        Ok(ByteBoolArray::new(array.buffer().clone(), array.validity().mask(mask)).into_array())
+impl MaskReduce for ByteBoolVTable {
+    fn mask(array: &ByteBoolArray, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
+        Ok(Some(
+            ByteBoolArray::new(
+                array.buffer().clone(),
+                array
+                    .validity()
+                    .clone()
+                    .and(Validity::Array(mask.clone()))?,
+            )
+            .into_array(),
+        ))
     }
 }
 
-register_kernel!(MaskKernelAdapter(ByteBoolVTable).lift());
-
 impl TakeExecute for ByteBoolVTable {
     fn take(
         array: &ByteBoolArray,

encodings/bytebool/src/rules.rs

@@ -3,11 +3,13 @@
 
 use vortex_array::arrays::SliceReduceAdaptor;
 use vortex_array::compute::CastReduceAdaptor;
+use vortex_array::compute::MaskReduceAdaptor;
 use vortex_array::optimizer::rules::ParentRuleSet;
 
 use crate::ByteBoolVTable;
 
 pub(crate) static RULES: ParentRuleSet<ByteBoolVTable> = ParentRuleSet::new(&[
-    ParentRuleSet::lift(&SliceReduceAdaptor(ByteBoolVTable)),
     ParentRuleSet::lift(&CastReduceAdaptor(ByteBoolVTable)),
+    ParentRuleSet::lift(&MaskReduceAdaptor(ByteBoolVTable)),
+    ParentRuleSet::lift(&SliceReduceAdaptor(ByteBoolVTable)),
 ]);

encodings/datetime-parts/public-api.lock

@@ -6,6 +6,8 @@ impl vortex_datetime_parts::DateTimePartsArray
 
 pub fn vortex_datetime_parts::DateTimePartsArray::days(&self) -> &vortex_array::array::ArrayRef
 
+pub fn vortex_datetime_parts::DateTimePartsArray::into_parts(self) -> vortex_datetime_parts::DateTimePartsArrayParts
+
 pub fn vortex_datetime_parts::DateTimePartsArray::seconds(&self) -> &vortex_array::array::ArrayRef
 
 pub fn vortex_datetime_parts::DateTimePartsArray::subseconds(&self) -> &vortex_array::array::ArrayRef
@@ -44,6 +46,24 @@ impl vortex_array::array::IntoArray for vortex_datetime_parts::DateTimePartsArra
 
 pub fn vortex_datetime_parts::DateTimePartsArray::into_array(self) -> vortex_array::array::ArrayRef
 
+pub struct vortex_datetime_parts::DateTimePartsArrayParts
+
+pub vortex_datetime_parts::DateTimePartsArrayParts::days: vortex_array::array::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsArrayParts::dtype: vortex_dtype::dtype::DType
+
+pub vortex_datetime_parts::DateTimePartsArrayParts::seconds: vortex_array::array::ArrayRef
+
+pub vortex_datetime_parts::DateTimePartsArrayParts::subseconds: vortex_array::array::ArrayRef
+
+impl core::clone::Clone for vortex_datetime_parts::DateTimePartsArrayParts
+
+pub fn vortex_datetime_parts::DateTimePartsArrayParts::clone(&self) -> vortex_datetime_parts::DateTimePartsArrayParts
+
+impl core::fmt::Debug for vortex_datetime_parts::DateTimePartsArrayParts
+
+pub fn vortex_datetime_parts::DateTimePartsArrayParts::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result
+
 #[repr(C)] pub struct vortex_datetime_parts::DateTimePartsMetadata
 
 pub vortex_datetime_parts::DateTimePartsMetadata::days_ptype: i32
@@ -118,10 +138,6 @@ impl vortex_array::compute::is_constant::IsConstantKernel for vortex_datetime_pa
 
 pub fn vortex_datetime_parts::DateTimePartsVTable::is_constant(&self, array: &vortex_datetime_parts::DateTimePartsArray, opts: &vortex_array::compute::is_constant::IsConstantOpts) -> vortex_error::VortexResult<core::option::Option<bool>>
 
-impl vortex_array::compute::mask::MaskKernel for vortex_datetime_parts::DateTimePartsVTable
-
-pub fn vortex_datetime_parts::DateTimePartsVTable::mask(&self, array: &vortex_datetime_parts::DateTimePartsArray, mask_array: &vortex_mask::Mask) -> vortex_error::VortexResult<vortex_array::array::ArrayRef>
-
 impl vortex_array::expr::exprs::binary::compare::CompareKernel for vortex_datetime_parts::DateTimePartsVTable
 
 pub fn vortex_datetime_parts::DateTimePartsVTable::compare(lhs: &vortex_datetime_parts::DateTimePartsArray, rhs: &dyn vortex_array::array::Array, operator: vortex_array::compute::compare::Operator, _ctx: &mut vortex_array::executor::ExecutionCtx) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
@@ -130,6 +146,10 @@ impl vortex_array::expr::exprs::cast::kernel::CastReduce for vortex_datetime_par
 
 pub fn vortex_datetime_parts::DateTimePartsVTable::cast(array: &vortex_datetime_parts::DateTimePartsArray, dtype: &vortex_dtype::dtype::DType) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
 
+impl vortex_array::expr::exprs::mask::kernel::MaskReduce for vortex_datetime_parts::DateTimePartsVTable
+
+pub fn vortex_datetime_parts::DateTimePartsVTable::mask(array: &vortex_datetime_parts::DateTimePartsArray, mask: &vortex_array::array::ArrayRef) -> vortex_error::VortexResult<core::option::Option<vortex_array::array::ArrayRef>>
+
 impl vortex_array::vtable::VTable for vortex_datetime_parts::DateTimePartsVTable
 
 pub type vortex_datetime_parts::DateTimePartsVTable::Array = vortex_datetime_parts::DateTimePartsArray

encodings/datetime-parts/src/array.rs

@@ -190,6 +190,14 @@ pub struct DateTimePartsArray {
     stats_set: ArrayStats,
 }
 
+#[derive(Clone, Debug)]
+pub struct DateTimePartsArrayParts {
+    pub dtype: DType,
+    pub days: ArrayRef,
+    pub seconds: ArrayRef,
+    pub subseconds: ArrayRef,
+}
+
 #[derive(Debug)]
 pub struct DateTimePartsVTable;
 
@@ -252,6 +260,15 @@ impl DateTimePartsArray {
         }
     }
 
+    pub fn into_parts(self) -> DateTimePartsArrayParts {
+        DateTimePartsArrayParts {
+            dtype: self.dtype,
+            days: self.days,
+            seconds: self.seconds,
+            subseconds: self.subseconds,
+        }
+    }
+
     pub fn days(&self) -> &ArrayRef {
         &self.days
     }