From 5634f7daf0585e09aa624e002855dbfd737cbe8b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 16:45:59 -0400 Subject: [PATCH 1/3] Reduce lazy masks into arrays with all-valid metadata Signed-off-by: Connor Tsui --- encodings/bytebool/src/compute.rs | 2 + vortex-array/src/arrays/bool/compute/mask.rs | 2 + .../src/arrays/decimal/compute/mask.rs | 27 +++---- .../arrays/fixed_size_list/compute/mask.rs | 2 + vortex-array/src/arrays/list/compute/mask.rs | 17 ++-- .../src/arrays/listview/compute/mask.rs | 36 ++++++--- vortex-array/src/arrays/map/compute/mask.rs | 2 + .../src/arrays/primitive/compute/mask.rs | 2 + .../src/arrays/struct_/compute/mask.rs | 2 + .../src/arrays/varbin/compute/mask.rs | 15 ++-- .../src/arrays/varbinview/compute/mask.rs | 2 + vortex-array/src/scalar_fn/fns/mask/kernel.rs | 80 +++++++++++++++++-- 12 files changed, 142 insertions(+), 47 deletions(-) diff --git a/encodings/bytebool/src/compute.rs b/encodings/bytebool/src/compute.rs index 7f9a4e803de..e2ae35dc41d 100644 --- a/encodings/bytebool/src/compute.rs +++ b/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> { Ok(Some( ByteBool::new( diff --git a/vortex-array/src/arrays/bool/compute/mask.rs b/vortex-array/src/arrays/bool/compute/mask.rs index 43890c42bc8..15179804194 100644 --- a/vortex-array/src/arrays/bool/compute/mask.rs +++ b/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> { Ok(Some( BoolArray::try_new_from_handle( diff --git a/vortex-array/src/arrays/decimal/compute/mask.rs b/vortex-array/src/arrays/decimal/compute/mask.rs index 804e6557f5c..91490a288ad 100644 --- a/vortex-array/src/arrays/decimal/compute/mask.rs +++ b/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> { - 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::(), - 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() + })) } } diff --git a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs b/vortex-array/src/arrays/fixed_size_list/compute/mask.rs index 1db564552e7..983c848c85a 100644 --- a/vortex-array/src/arrays/fixed_size_list/compute/mask.rs +++ b/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, diff --git a/vortex-array/src/arrays/list/compute/mask.rs b/vortex-array/src/arrays/list/compute/mask.rs index f2727247646..40737eb6583 100644 --- a/vortex-array/src/arrays/list/compute/mask.rs +++ b/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> { - 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() + })) } } diff --git a/vortex-array/src/arrays/listview/compute/mask.rs b/vortex-array/src/arrays/listview/compute/mask.rs index 2b5a89132fb..65fa7b96796 100644 --- a/vortex-array/src/arrays/listview/compute/mask.rs +++ b/vortex-array/src/arrays/listview/compute/mask.rs @@ -5,27 +5,37 @@ 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> { - // SAFETY: masking the validity does not affect the invariants - 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(), - )) + 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::from_parts_unchecked(parts) }.into_array())) } } diff --git a/vortex-array/src/arrays/map/compute/mask.rs b/vortex-array/src/arrays/map/compute/mask.rs index 41703a2a3ea..97218d08454 100644 --- a/vortex-array/src/arrays/map/compute/mask.rs +++ b/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> { let Some(entries) = ::mask(array.entries().as_::(), mask)? diff --git a/vortex-array/src/arrays/primitive/compute/mask.rs b/vortex-array/src/arrays/primitive/compute/mask.rs index 023310b3807..8fb7d51e1b7 100644 --- a/vortex-array/src/arrays/primitive/compute/mask.rs +++ b/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> { // SAFETY: validity and data buffer still have same length Ok(Some(unsafe { diff --git a/vortex-array/src/arrays/struct_/compute/mask.rs b/vortex-array/src/arrays/struct_/compute/mask.rs index 1b5b526e69d..68db6166455 100644 --- a/vortex-array/src/arrays/struct_/compute/mask.rs +++ b/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> { StructArray::try_new_with_dtype( array.iter_unmasked_fields().cloned(), diff --git a/vortex-array/src/arrays/varbin/compute/mask.rs b/vortex-array/src/arrays/varbin/compute/mask.rs index 0d91b6fe3c3..d47935ec7de 100644 --- a/vortex-array/src/arrays/varbin/compute/mask.rs +++ b/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> { - 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() + })) } } diff --git a/vortex-array/src/arrays/varbinview/compute/mask.rs b/vortex-array/src/arrays/varbinview/compute/mask.rs index e470dceba32..601ae402a6a 100644 --- a/vortex-array/src/arrays/varbinview/compute/mask.rs +++ b/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> { // SAFETY: masking the validity does not affect the invariants unsafe { diff --git a/vortex-array/src/scalar_fn/fns/mask/kernel.rs b/vortex-array/src/scalar_fn/fns/mask/kernel.rs index c73ca6c231d..e0dff396ee4 100644 --- a/vortex-array/src/scalar_fn/fns/mask/kernel.rs +++ b/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; @@ -15,6 +17,7 @@ use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::kernel::ExecuteParentKernel; use crate::optimizer::rules::ArrayParentReduceRule; use crate::scalar_fn::fns::mask::Mask as MaskExpr; +use crate::validity::Validity; /// Mask an array without reading buffers. /// @@ -26,9 +29,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>; } @@ -71,17 +79,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::().is_none() && mask_child.as_opt::().is_none() { - return Ok(None); + let can_attach_mask = V::VALIDITY_IS_METADATA_ONLY + && matches!(array.validity()?, Validity::AllValid | Validity::NonNullable); + + if !can_attach_mask { + return Ok(None); + } } + ::mask(array, &mask_child) } } @@ -120,17 +132,71 @@ 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 { + 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_::(); + + 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::>()); + + 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 From 6ddf55f9721d6ba9aee22164dc87528d59aeb980 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 17:14:25 -0400 Subject: [PATCH 2/3] Fix formatting in mask reductions Signed-off-by: Connor Tsui --- vortex-array/src/arrays/listview/compute/mask.rs | 4 +++- vortex-array/src/scalar_fn/fns/mask/kernel.rs | 8 ++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/arrays/listview/compute/mask.rs b/vortex-array/src/arrays/listview/compute/mask.rs index 65fa7b96796..3c035fc864e 100644 --- a/vortex-array/src/arrays/listview/compute/mask.rs +++ b/vortex-array/src/arrays/listview/compute/mask.rs @@ -36,6 +36,8 @@ impl MaskReduce for ListView { // 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::from_parts_unchecked(parts) }.into_array())) + Ok(Some( + unsafe { ListViewArray::from_parts_unchecked(parts) }.into_array(), + )) } } diff --git a/vortex-array/src/scalar_fn/fns/mask/kernel.rs b/vortex-array/src/scalar_fn/fns/mask/kernel.rs index e0dff396ee4..f8b48b07cea 100644 --- a/vortex-array/src/scalar_fn/fns/mask/kernel.rs +++ b/vortex-array/src/scalar_fn/fns/mask/kernel.rs @@ -87,7 +87,10 @@ where if mask_child.as_opt::().is_none() && mask_child.as_opt::().is_none() { let can_attach_mask = V::VALIDITY_IS_METADATA_ONLY - && matches!(array.validity()?, Validity::AllValid | Validity::NonNullable); + && matches!( + array.validity()?, + Validity::AllValid | Validity::NonNullable + ); if !can_attach_mask { return Ok(None); @@ -191,7 +194,8 @@ mod tests { #[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 input = + PrimitiveArray::new(buffer![1i32, 2, 3], Validity::Array(validity)).into_array(); let output = input.mask(lazy_mask()?)?; assert!(output.is::>()); From a8d283b439d1350e329c7bbb0793b57c38a61d81 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 23 Sep 2026 17:26:11 -0400 Subject: [PATCH 3/3] Use validity helper in mask reduction Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/fns/mask/kernel.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/mask/kernel.rs b/vortex-array/src/scalar_fn/fns/mask/kernel.rs index f8b48b07cea..3009c43aebc 100644 --- a/vortex-array/src/scalar_fn/fns/mask/kernel.rs +++ b/vortex-array/src/scalar_fn/fns/mask/kernel.rs @@ -17,7 +17,6 @@ use crate::arrays::scalar_fn::ScalarFnArrayView; use crate::kernel::ExecuteParentKernel; use crate::optimizer::rules::ArrayParentReduceRule; use crate::scalar_fn::fns::mask::Mask as MaskExpr; -use crate::validity::Validity; /// Mask an array without reading buffers. /// @@ -86,11 +85,8 @@ where .ok_or_else(|| vortex_err!("Mask expression must have 2 children"))?; if mask_child.as_opt::().is_none() && mask_child.as_opt::().is_none() { - let can_attach_mask = V::VALIDITY_IS_METADATA_ONLY - && matches!( - array.validity()?, - Validity::AllValid | Validity::NonNullable - ); + let can_attach_mask = + V::VALIDITY_IS_METADATA_ONLY && array.validity()?.definitely_no_nulls(); if !can_attach_mask { return Ok(None);