Reduce lazy masks into arrays with all-valid metadata

encodings/bytebool/src/compute.rs

@@ -73,6 +73,8 @@ impl CastKernel for ByteBool {
 }
 
 impl MaskReduce for ByteBool {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         Ok(Some(
             ByteBool::new(

vortex-array/src/arrays/bool/compute/mask.rs

@@ -12,6 +12,8 @@ use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for Bool {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, Bool>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         Ok(Some(
             BoolArray::try_new_from_handle(

vortex-array/src/arrays/decimal/compute/mask.rs

@@ -8,25 +8,22 @@ use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::Decimal;
 use crate::arrays::DecimalArray;
-use crate::match_each_decimal_value_type;
 use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for Decimal {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, Decimal>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-        Ok(Some(match_each_decimal_value_type!(
-            array.values_type(),
-            |D| {
-                // SAFETY: masking the validity does not affect the invariants
-                unsafe {
-                    DecimalArray::new_unchecked(
-                        array.buffer::<D>(),
-                        array.decimal_dtype(),
-                        array.validity()?.and(Validity::Array(mask.clone()))?,
-                    )
-                }
-                .into_array()
-            }
-        )))
+        // SAFETY: the values and decimal type are unchanged, and masking only removes valid rows.
+        Ok(Some(unsafe {
+            DecimalArray::new_unchecked_handle(
+                array.buffer_handle().clone(),
+                array.values_type(),
+                array.decimal_dtype(),
+                array.validity()?.and(Validity::Array(mask.clone()))?,
+            )
+            .into_array()
+        }))
     }
 }

vortex-array/src/arrays/fixed_size_list/compute/mask.rs

@@ -14,6 +14,8 @@ use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for FixedSizeList {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(
         array: ArrayView<'_, FixedSizeList>,
         mask: &ArrayRef,

vortex-array/src/arrays/list/compute/mask.rs

@@ -13,12 +13,17 @@ use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for List {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, List>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-        ListArray::try_new(
-            array.elements().clone(),
-            array.offsets().clone(),
-            array.validity()?.and(Validity::Array(mask.clone()))?,
-        )
-        .map(|a| Some(a.into_array()))
+        // SAFETY: elements and offsets are unchanged, and masking only removes valid rows.
+        Ok(Some(unsafe {
+            ListArray::new_unchecked(
+                array.elements().clone(),
+                array.offsets().clone(),
+                array.validity()?.and(Validity::Array(mask.clone()))?,
+            )
+            .into_array()
+        }))
     }
 }

vortex-array/src/arrays/listview/compute/mask.rs

@@ -5,27 +5,39 @@ use vortex_error::VortexResult;
 
 use crate::ArrayRef;
 use crate::IntoArray;
+use crate::array::ArrayParts;
 use crate::array::ArrayView;
 use crate::arrays::ListView;
 use crate::arrays::ListViewArray;
 use crate::arrays::listview::ListViewArraySlotsExt;
+use crate::arrays::listview::ListViewData;
 use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for ListView {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, ListView>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-        // SAFETY: masking the validity does not affect the invariants
+        let validity = array.validity()?.and(Validity::Array(mask.clone()))?;
+        let slots = ListViewData::make_slots(
+            array.elements(),
+            array.offsets(),
+            array.sizes(),
+            &validity,
+            array.len(),
+        );
+        let parts = ArrayParts::new(
+            ListView,
+            array.dtype().as_nullable(),
+            array.len(),
+            array.data().clone(),
+        )
+        .with_slots(slots);
+
+        // SAFETY: elements, offsets, sizes, and their metadata are unchanged. Masking only removes
+        // valid rows, so the existing zero-copy-to-list guarantee still holds.
         Ok(Some(
-            unsafe {
-                ListViewArray::new_unchecked(
-                    array.elements().clone(),
-                    array.offsets().clone(),
-                    array.sizes().clone(),
-                    array.validity()?.and(Validity::Array(mask.clone()))?,
-                )
-                .with_zero_copy_to_list(array.is_zero_copy_to_list())
-            }
-            .into_array(),
+            unsafe { ListViewArray::from_parts_unchecked(parts) }.into_array(),
         ))
     }
 }

vortex-array/src/arrays/map/compute/mask.rs

@@ -15,6 +15,8 @@ use crate::scalar_fn::fns::mask::MaskKernel;
 use crate::scalar_fn::fns::mask::MaskReduce;
 
 impl MaskReduce for Map {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         let Some(entries) =
             <ListView as MaskReduce>::mask(array.entries().as_::<ListView>(), mask)?

vortex-array/src/arrays/primitive/compute/mask.rs

@@ -12,6 +12,8 @@ use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for Primitive {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, Primitive>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         // SAFETY: validity and data buffer still have same length
         Ok(Some(unsafe {

vortex-array/src/arrays/struct_/compute/mask.rs

@@ -13,6 +13,8 @@ use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for Struct {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, Struct>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         StructArray::try_new_with_dtype(
             array.iter_unmasked_fields().cloned(),

vortex-array/src/arrays/varbin/compute/mask.rs

@@ -13,16 +13,19 @@ use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for VarBin {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, VarBin>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-        Ok(Some(
-            VarBinArray::try_new(
+        // SAFETY: offsets and bytes are unchanged, and masking only removes valid rows.
+        Ok(Some(unsafe {
+            VarBinArray::new_unchecked_from_handle(
                 array.offsets().clone(),
-                array.bytes().clone(),
+                array.bytes_handle().clone(),
                 array.dtype().as_nullable(),
                 array.validity()?.and(Validity::Array(mask.clone()))?,
-            )?
-            .into_array(),
-        ))
+            )
+            .into_array()
+        }))
     }
 }
 

vortex-array/src/arrays/varbinview/compute/mask.rs

@@ -14,6 +14,8 @@ use crate::scalar_fn::fns::mask::MaskReduce;
 use crate::validity::Validity;
 
 impl MaskReduce for VarBinView {
+    const VALIDITY_IS_METADATA_ONLY: bool = true;
+
     fn mask(array: ArrayView<'_, VarBinView>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
         // SAFETY: masking the validity does not affect the invariants
         unsafe {

vortex-array/src/scalar_fn/fns/mask/kernel.rs

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Encoding-specific mask reduction and execution adapters.
+
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
 
@@ -26,9 +28,14 @@ use crate::scalar_fn::fns::mask::Mask as MaskExpr;
 ///
 /// # Preconditions
 ///
-/// The mask is guaranteed to have the same length as the array. Trivial cases
-/// (`AllValid`, `AllInvalid`, `NonNullable`) are handled by the caller before dispatch.
+/// The mask has non-nullable Boolean dtype and the same length as the array. It may be lazy.
 pub trait MaskReduce: VTable {
+    /// Whether reading this encoding's validity requires only metadata access.
+    ///
+    /// Enables lazy masks on all-valid inputs. Only opt in when `array.validity()` cannot read
+    /// buffers or execute children, including for nullable inputs.
+    const VALIDITY_IS_METADATA_ONLY: bool = false;
+
     fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>>;
 }
 
@@ -71,17 +78,21 @@ where
         if child_idx != 0 {
             return Ok(None);
         }
-        // Reduce only when the mask (child 1) is readable from metadata: a concrete `Bool` or a
-        // `Constant`. `Mask::return_dtype` guarantees the mask is `Bool(NonNullable)`, so a
-        // `Constant` here is a non-nullable Boolean. Other encodings may need execution, so leave
-        // them to the kernel.
+
         let parent_ref: ArrayRef = (*parent).clone();
         let mask_child = parent_ref
             .nth_child(1)
             .ok_or_else(|| vortex_err!("Mask expression must have 2 children"))?;
+
         if mask_child.as_opt::<Bool>().is_none() && mask_child.as_opt::<Constant>().is_none() {
-            return Ok(None);
+            let can_attach_mask =
+                V::VALIDITY_IS_METADATA_ONLY && array.validity()?.definitely_no_nulls();
+
+            if !can_attach_mask {
+                return Ok(None);
+            }
         }
+
         <V as MaskReduce>::mask(array, &mask_child)
     }
 }
@@ -120,17 +131,72 @@ mod tests {
     use vortex_buffer::buffer;
     use vortex_error::VortexResult;
 
+    use crate::ArrayEq;
+    use crate::ArrayRef;
+    use crate::EqMode;
     use crate::IntoArray;
+    use crate::arrays::BoolArray;
     use crate::arrays::ConstantArray;
     use crate::arrays::Primitive;
     use crate::arrays::PrimitiveArray;
     use crate::arrays::ScalarFn;
+    use crate::arrays::scalar_fn::ExactScalarFn;
     use crate::assert_arrays_eq;
+    use crate::builtins::ArrayBuiltins;
     use crate::dtype::Nullability;
     use crate::executor::VortexSessionExecute;
     use crate::optimizer::ArrayOptimizer;
     use crate::scalar::Scalar;
+    use crate::scalar_fn::fns::binary::Binary;
     use crate::scalar_fn::fns::mask::Mask as MaskExpr;
+    use crate::scalar_fn::fns::operators::Operator;
+    use crate::validity::Validity;
+
+    fn lazy_mask() -> VortexResult<ArrayRef> {
+        Binary::try_new(
+            BoolArray::from_iter([true, true, false]).into_array(),
+            BoolArray::from_iter([true, false, true]).into_array(),
+            Operator::And,
+        )
+        .map(IntoArray::into_array)
+    }
+
+    #[rstest]
+    #[case::non_nullable(Validity::NonNullable)]
+    #[case::all_valid(Validity::AllValid)]
+    fn all_valid_input_attaches_lazy_mask(#[case] validity: Validity) -> VortexResult<()> {
+        let input = PrimitiveArray::new(buffer![1i32, 2, 3], validity);
+        let mask = lazy_mask()?;
+        let output = input.clone().into_array().mask(mask.clone())?;
+        let primitive = output.as_::<Primitive>();
+
+        assert!(
+            primitive
+                .buffer_handle()
+                .array_eq(input.buffer_handle(), EqMode::Ptr)
+        );
+        assert_eq!(output.dtype(), &input.dtype().as_nullable());
+        assert_eq!(output.len(), input.len());
+        let Validity::Array(validity) = output.validity()? else {
+            panic!("mask must remain array-backed");
+        };
+        assert!(ArrayRef::ptr_eq(&validity, &mask));
+
+        Ok(())
+    }
+
+    #[rstest]
+    #[case::partially_valid(false)]
+    #[case::all_true_bitmap(true)]
+    fn array_backed_validity_keeps_lazy_mask(#[case] all_true: bool) -> VortexResult<()> {
+        let validity = BoolArray::from_iter([true, all_true, true]).into_array();
+        let input =
+            PrimitiveArray::new(buffer![1i32, 2, 3], Validity::Array(validity)).into_array();
+        let output = input.mask(lazy_mask()?)?;
+        assert!(output.is::<ExactScalarFn<MaskExpr>>());
+
+        Ok(())
+    }
 
     /// A constant Boolean mask child must take the metadata-only reduction path (pushing into the
     /// input encoding) rather than surviving as a `ScalarFn` wrapper that falls through to