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..3c035fc864e 100644 --- a/vortex-array/src/arrays/listview/compute/mask.rs +++ b/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> { - // 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(), )) } } 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..3009c43aebc 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; @@ -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>; } @@ -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::().is_none() && mask_child.as_opt::().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); + } } + ::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 { + 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