diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 1530c33874d..0b1b9422f86 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::mem::MaybeUninit; -use std::sync::Arc; use fastlanes::BitPacking; use vortex_array::ArrayRef; @@ -20,7 +19,7 @@ use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use super::chunked_indices; use super::take::UNPACK_CHUNK_THRESHOLD; @@ -82,7 +81,7 @@ impl FilterKernel for BitPacked { let patches = array .patches() - .map(|patches| patches.filter(&Mask::Values(Arc::clone(values)), ctx)) + .map(|patches| patches.filter(&Mask::Values(MaskValuesRef::clone(values)), ctx)) .transpose()? .flatten(); @@ -109,12 +108,12 @@ impl FilterKernel for BitPacked { /// Returns a tuple of (values buffer, validity mask). fn filter_primitive_without_patches( array: ArrayView<'_, BitPacked>, - selection: &Arc, + selection: &MaskValuesRef, ) -> VortexResult<(Buffer, Validity)> { let values = filter_with_indices(array.data(), selection.indices()); let validity = array .validity()? - .filter(&Mask::Values(Arc::clone(selection)))?; + .filter(&Mask::Values(MaskValuesRef::clone(selection)))?; Ok((values.freeze(), validity)) } @@ -176,7 +175,7 @@ fn filter_with_indices( } #[cfg(test)] -mod test { +mod tests { use std::sync::LazyLock; use vortex_array::IntoArray as _; diff --git a/vortex-array/src/arrays/filter/execute/bool.rs b/vortex-array/src/arrays/filter/execute/bool.rs index 08364e3cac5..2a9f60cb1b1 100644 --- a/vortex-array/src/arrays/filter/execute/bool.rs +++ b/vortex-array/src/arrays/filter/execute/bool.rs @@ -1,20 +1,18 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::BoolArray; use crate::arrays::bool::BoolArrayExt; use crate::arrays::filter::execute::bitbuffer; use crate::arrays::filter::execute::filter_validity; -pub fn filter_bool(array: &BoolArray, mask: &Arc) -> BoolArray { +pub fn filter_bool(array: &BoolArray, mask: &MaskValuesRef) -> BoolArray { let validity = array .validity() - .vortex_expect("bool validity should be derivable"); + .vortex_expect("validity is derivable for a valid BoolArray"); let filtered_validity = filter_validity(validity, mask); let bit_buffer = array.to_bit_buffer(); @@ -24,7 +22,7 @@ pub fn filter_bool(array: &BoolArray, mask: &Arc) -> BoolArray { } #[cfg(test)] -mod test { +mod tests { use itertools::Itertools; use rstest::rstest; use vortex_mask::Mask; diff --git a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs index 33fdf62bfd6..b73fff21a2d 100644 --- a/vortex-array/src/arrays/filter/execute/fixed_size_list.rs +++ b/vortex-array/src/arrays/filter/execute/fixed_size_list.rs @@ -1,12 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_mask::Mask; use vortex_mask::MaskIter; use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::FixedSizeListArray; use crate::arrays::filter::execute::filter_validity; @@ -24,12 +23,12 @@ const MASK_EXPANSION_DENSITY_THRESHOLD: f64 = 0.05; /// mask down to the child elements array. pub fn filter_fixed_size_list( array: &FixedSizeListArray, - selection_mask: &Arc, + selection_mask: &MaskValuesRef, ) -> FixedSizeListArray { let filtered_validity = filter_validity( array .validity() - .vortex_expect("fixed-size-list validity should be derivable"), + .vortex_expect("validity is derivable for a valid FixedSizeListArray"), selection_mask, ); @@ -38,13 +37,11 @@ pub fn filter_fixed_size_list( let list_size = array.list_size(); let new_elements = { - // We want to create a new mask specialized to the underlying `elements` of the array. if list_size != 0 { - // TODO(connor): If we can push down a "indices" or "slices" selection instead that - // would be much more performant. - let elements_mask = compute_mask_for_fsl_elements(selection_mask, list_size as usize); + // TODO(connor): Push down an indices or slices selection to avoid expanding the mask. + let elements_mask = + compute_mask_for_fsl_elements(selection_mask.as_ref(), list_size as usize); - // Allow the child array to filter itself. let new_elements = elements .filter(elements_mask) .vortex_expect("FixedSizeListArray elements are guaranteed to support filter"); @@ -59,22 +56,14 @@ pub fn filter_fixed_size_list( "degenerate FixedSizeListArray is invalid" ); - // NB: The safety comment for the `list_size == 0` case is here for clarity. - - // SAFETY: We have verified that when `list_size == 0` - // - `elements` has length 0 (since it came from a valid `FixedSizeListArray`) - // - `filtered_validity` has the correct length because we filter with the same - // `selection_mask` as the array itself elements.clone() } }; - // SAFETY: We have verified that - // - The case when `list_size == 0` is safe (see above) - // - The `new_elements` array is guaranteed to have a length that is a multiple of - // `list_size` - // - `filtered_validity` has the correct length because we filter with the same - // `selection_mask` as the array itself + // SAFETY: + // - A valid zero-width array has no elements, which the zero-width branch preserves. + // - Otherwise, the expanded selection retains `list_size` elements for each selected list. + // - Filtering validity with `selection_mask` gives it the required `new_len`. unsafe { FixedSizeListArray::new_unchecked(new_elements, list_size, filtered_validity, new_len) } @@ -120,7 +109,7 @@ fn compute_mask_for_fsl_elements(selection_mask: &MaskValues, list_size: usize) } #[cfg(test)] -mod test { +mod tests { use vortex_buffer::buffer; use vortex_mask::Mask; diff --git a/vortex-array/src/arrays/filter/execute/listview.rs b/vortex-array/src/arrays/filter/execute/listview.rs index db90f1e6680..ba15ee51998 100644 --- a/vortex-array/src/arrays/filter/execute/listview.rs +++ b/vortex-array/src/arrays/filter/execute/listview.rs @@ -1,11 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::ListViewArray; use crate::arrays::filter::execute::filter_validity; @@ -23,7 +21,7 @@ use crate::arrays::listview::ListViewArraySlotsExt; /// /// The trade-off is that we may keep unreferenced elements in memory, but this is acceptable since /// we're optimizing for read performance and the data isn't being copied. -pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) -> ListViewArray { +pub fn filter_listview(array: &ListViewArray, selection_mask: &MaskValuesRef) -> ListViewArray { let elements = array.elements(); let offsets = array.offsets(); let sizes = array.sizes(); @@ -31,7 +29,7 @@ pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) let new_validity = filter_validity( array .validity() - .vortex_expect("listview validity should be derivable"), + .vortex_expect("validity is derivable for a valid ListViewArray"), selection_mask, ); debug_assert!( @@ -40,8 +38,7 @@ pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) .is_none_or(|len| len == selection_mask.true_count()) ); - // Simply filter the offsets and sizes arrays. - let mask_for_filter = Mask::Values(Arc::clone(selection_mask)); + let mask_for_filter = Mask::Values(MaskValuesRef::clone(selection_mask)); let new_offsets = offsets .filter(mask_for_filter.clone()) .vortex_expect("ListViewArray offsets are guaranteed to support filter"); @@ -49,15 +46,15 @@ pub fn filter_listview(array: &ListViewArray, selection_mask: &Arc) .filter(mask_for_filter) .vortex_expect("ListViewArray sizes are guaranteed to support filter"); - // SAFETY: Filter operation maintains all `ListViewArray` invariants: + // SAFETY: // - Offsets and sizes are derived from existing valid child arrays. - // - Offsets and sizes have the same length (both filtered by `selection_mask`). - // - Validity matches the filtered array's nullability. + // - Filtering offsets, sizes, and validity with `selection_mask` preserves their row alignment. + // - The retained elements buffer keeps every referenced element in bounds. unsafe { ListViewArray::new_unchecked(elements.clone(), new_offsets, new_sizes, new_validity) } } #[cfg(test)] -mod test { +mod tests { use std::sync::LazyLock; use vortex_buffer::buffer; diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 037975b3f66..e0ebceefb09 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -5,12 +5,10 @@ //! //! The main entrypoint is [`execute_filter`] which filters any [`Canonical`] array. -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::ArrayRef; use crate::Canonical; @@ -25,7 +23,6 @@ use crate::arrays::MapArray; use crate::arrays::NullArray; use crate::arrays::VariantArray; use crate::arrays::extension::ExtensionArrayExt; -use crate::arrays::filter::FilterArraySlotsExt; use crate::arrays::filter::FilterReduce; use crate::arrays::fixed_width; use crate::arrays::variant::VariantArraySlotsExt; @@ -44,41 +41,28 @@ mod take; mod union; mod varbinview; -/// A helper function that lazily filters a [`Validity`] with selection mask values. -pub(crate) fn filter_validity(validity: Validity, mask: &Arc) -> Validity { +/// Lazily filters a [`Validity`] with a partially selective mask. +pub(crate) fn filter_validity(validity: Validity, mask: &MaskValuesRef) -> Validity { validity - .filter(&Mask::Values(Arc::clone(mask))) - .vortex_expect("Somehow unable to wrap filter around a validity array") + .filter(&Mask::Values(MaskValuesRef::clone(mask))) + .vortex_expect("filtering validity with a partially selective mask is valid") } -/// Check for some fast-path execution conditions before calling [`execute_filter`]. -pub(super) fn execute_filter_fast_paths( +/// Returns an all-null result when the child contains no valid values. +pub(super) fn execute_all_null_filter_fast_path( array: ArrayView<'_, Filter>, + selected_count: usize, ctx: &mut ExecutionCtx, ) -> VortexResult> { - let true_count = array.mask.true_count(); - - // If the mask selects nothing, the output is empty. - if true_count == 0 { - return Ok(Some(Canonical::empty(array.dtype()).into_array())); - } - - // If the mask selects everything, then we can just fully decompress the whole thing. - if true_count == array.mask.len() { - return Ok(Some(array.child().clone())); - } - - // Also check if the array itself is completely null, in which case we only care about the total - // number of nulls, not the values. - let child_arr = array.array(); - if child_arr + let child = array.array(); + if child .validity()? - .execute_mask(child_arr.len(), ctx)? + .execute_mask(child.len(), ctx)? .true_count() == 0 { return Ok(Some( - ConstantArray::new(Scalar::null(array.dtype().clone()), true_count).into_array(), + ConstantArray::new(Scalar::null(array.dtype().clone()), selected_count).into_array(), )); } @@ -86,7 +70,7 @@ pub(super) fn execute_filter_fast_paths( } /// Filter a canonical array by a mask, returning a new canonical array. -pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Canonical { +pub(super) fn execute_filter(canonical: Canonical, mask: &MaskValuesRef) -> Canonical { match canonical { Canonical::Null(_) => Canonical::Null(NullArray::new(mask.true_count())), Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)), @@ -103,12 +87,12 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca Canonical::Extension(a) => { let filtered_storage = a .storage_array() - .filter(Mask::Values(Arc::clone(mask))) + .filter(Mask::Values(MaskValuesRef::clone(mask))) .vortex_expect("ExtensionArray storage type somehow could not be filtered"); Canonical::Extension(ExtensionArray::new(a.ext_dtype().clone(), filtered_storage)) } Canonical::Variant(a) => { - let filter_mask = Mask::Values(Arc::clone(mask)); + let filter_mask = Mask::Values(MaskValuesRef::clone(mask)); let filtered_core_storage = a .core_storage() .filter(filter_mask.clone()) @@ -126,8 +110,8 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca } } -fn filter_map(array: &MapArray, mask: &Arc) -> MapArray { - let filter_mask = Mask::Values(Arc::clone(mask)); +fn filter_map(array: &MapArray, mask: &MaskValuesRef) -> MapArray { + let filter_mask = Mask::Values(MaskValuesRef::clone(mask)); let filtered = ::filter(array.as_view(), &filter_mask) .vortex_expect("MapArray somehow could not be filtered") .vortex_expect("Map filter reduce always produces an array"); diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index e4a8157d2e2..22bd3c9f098 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -1,26 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::ArrayRef; use crate::arrays::StructArray; use crate::arrays::filter::execute::filter_validity; use crate::arrays::struct_::StructArrayExt; -pub fn filter_struct(array: &StructArray, mask: &Arc) -> StructArray { +pub fn filter_struct(array: &StructArray, mask: &MaskValuesRef) -> StructArray { let filtered_validity = filter_validity( array .validity() - .vortex_expect("struct validity should be derivable"), + .vortex_expect("validity is derivable for a valid StructArray"), mask, ); - let mask_for_filter = Mask::Values(Arc::clone(mask)); + let mask_for_filter = Mask::Values(MaskValuesRef::clone(mask)); let fields: Vec = array .iter_unmasked_fields() .map(|field| { @@ -45,7 +43,7 @@ pub fn filter_struct(array: &StructArray, mask: &Arc) -> StructArray } #[cfg(test)] -mod test { +mod tests { use vortex_mask::Mask; use crate::IntoArray; diff --git a/vortex-array/src/arrays/filter/execute/union.rs b/vortex-array/src/arrays/filter/execute/union.rs index 0e911c8102f..81955d1371d 100644 --- a/vortex-array/src/arrays/filter/execute/union.rs +++ b/vortex-array/src/arrays/filter/execute/union.rs @@ -1,19 +1,17 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_error::VortexExpect; use vortex_mask::Mask; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::ArrayRef; use crate::arrays::UnionArray; use crate::arrays::union::UnionArrayExt; use crate::arrays::union::UnionArraySlotsExt; -pub fn filter_union(array: &UnionArray, mask: &Arc) -> UnionArray { - let filter_mask = Mask::Values(Arc::clone(mask)); +pub fn filter_union(array: &UnionArray, mask: &MaskValuesRef) -> UnionArray { + let filter_mask = Mask::Values(MaskValuesRef::clone(mask)); let type_ids = array .type_ids() diff --git a/vortex-array/src/arrays/filter/execute/varbinview.rs b/vortex-array/src/arrays/filter/execute/varbinview.rs index 58a348b41b0..ed490347599 100644 --- a/vortex-array/src/arrays/filter/execute/varbinview.rs +++ b/vortex-array/src/arrays/filter/execute/varbinview.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use vortex_buffer::Buffer; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::arrays::VarBinViewArray; use crate::arrays::filter::execute::buffer; @@ -13,7 +13,7 @@ use crate::arrays::varbinview::BinaryView; use crate::arrays::varbinview::VarBinViewArrayExt; use crate::buffer::BufferHandle; -pub fn filter_varbinview(array: &VarBinViewArray, mask: &Arc) -> VarBinViewArray { +pub fn filter_varbinview(array: &VarBinViewArray, mask: &MaskValuesRef) -> VarBinViewArray { let filtered_validity = filter_validity(array.varbinview_validity(), mask); let views = Buffer::::from_byte_buffer(array.views_handle().as_host().clone()); @@ -32,7 +32,7 @@ pub fn filter_varbinview(array: &VarBinViewArray, mask: &Arc) -> Var } #[cfg(test)] -mod test { +mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; diff --git a/vortex-array/src/arrays/filter/vtable.rs b/vortex-array/src/arrays/filter/vtable.rs index 59e27a1df09..c90cd981fc7 100644 --- a/vortex-array/src/arrays/filter/vtable.rs +++ b/vortex-array/src/arrays/filter/vtable.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::hash::Hasher; -use std::sync::Arc; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -10,6 +9,7 @@ use vortex_error::vortex_bail; use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use vortex_mask::Mask; +use vortex_mask::MaskValuesRef; use vortex_session::VortexSession; use vortex_session::registry::CachedId; @@ -31,8 +31,8 @@ use crate::array::with_empty_buffers; use crate::arrays::filter::FilterArraySlotsExt; use crate::arrays::filter::array::FilterData; use crate::arrays::filter::array::FilterSlots; +use crate::arrays::filter::execute::execute_all_null_filter_fast_path; use crate::arrays::filter::execute::execute_filter; -use crate::arrays::filter::execute::execute_filter_fast_paths; use crate::arrays::filter::rules::PARENT_RULES; use crate::arrays::filter::rules::RULES; use crate::buffer::BufferHandle; @@ -152,14 +152,24 @@ impl VTable for Filter { } fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { - if let Some(canonical) = execute_filter_fast_paths(array.as_view(), ctx)? { - return Ok(ExecutionResult::done(canonical)); - } + // Match the mask once. A zero-length mask is both all true and all false, so check the + // empty output before the unfiltered one. let mask_values = match &array.mask { - Mask::Values(v) => Arc::clone(v), - _ => unreachable!("`execute_filter_fast_paths` handles AllTrue and AllFalse"), + Mask::AllFalse(_) | Mask::AllTrue(0) => { + return Ok(ExecutionResult::done( + Canonical::empty(array.dtype()).into_array(), + )); + } + Mask::AllTrue(_) => return Ok(ExecutionResult::done(array.child().clone())), + Mask::Values(values) => MaskValuesRef::clone(values), }; + if let Some(canonical) = + execute_all_null_filter_fast_path(array.as_view(), mask_values.true_count(), ctx)? + { + return Ok(ExecutionResult::done(canonical)); + } + let array = require_child!(array, array.child(), FilterSlots::CHILD => AnyCanonical); // We rely on the optimization pass that runs prior to this execution for filter pushdown, diff --git a/vortex-array/src/arrays/fixed_width/filter.rs b/vortex-array/src/arrays/fixed_width/filter.rs index eb30229ec48..d09dcc57772 100644 --- a/vortex-array/src/arrays/fixed_width/filter.rs +++ b/vortex-array/src/arrays/fixed_width/filter.rs @@ -1,13 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use super::FixedWidthArray; use super::match_each_record_width; @@ -21,13 +20,13 @@ use crate::arrays::filter::filter_validity; #[expect(clippy::cast_possible_truncation)] mod tests; -pub(crate) fn filter(array: &Array, mask: &Arc) -> Array { +pub(crate) fn filter(array: &Array, mask: &MaskValuesRef) -> Array { let array = array.as_view(); - let values = filter_records(V::values(array), V::byte_width(array), mask); + let values = filter_records(V::values(array), V::byte_width(array), mask.as_ref()); let validity = filter_validity( array .validity() - .vortex_expect("fixed-width validity should be derivable"), + .vortex_expect("validity is derivable for a valid fixed-width array"), mask, ); with_values(array, values, mask.true_count(), validity) diff --git a/vortex-array/src/arrays/list/compute/filter.rs b/vortex-array/src/arrays/list/compute/filter.rs index af02ce23d7b..937773575ac 100644 --- a/vortex-array/src/arrays/list/compute/filter.rs +++ b/vortex-array/src/arrays/list/compute/filter.rs @@ -1,8 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::sync::Arc; - use num_traits::Zero; use vortex_buffer::BitBufferMut; use vortex_buffer::Buffer; @@ -10,7 +8,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_mask::MaskIter; -use vortex_mask::MaskValues; +use vortex_mask::MaskValuesRef; use crate::ArrayRef; use crate::Canonical; @@ -37,7 +35,7 @@ const MASK_EXPANSION_DENSITY_THRESHOLD: f64 = 0.05; /// Construct an element mask from contiguous list offsets and a selection mask. pub fn element_mask_from_offsets( offsets: &[O], - selection: &Arc, + selection: &MaskValuesRef, ) -> Mask { let first_offset = offsets.first().map_or(0, |first_offset| first_offset.as_()); let last_offset = offsets.last().map_or(0, |last_offset| last_offset.as_()); @@ -102,7 +100,7 @@ impl FilterKernel for List { ) -> VortexResult> { let selection = match mask { Mask::AllTrue(_) | Mask::AllFalse(_) => return Ok(None), - Mask::Values(v) => v, + Mask::Values(values) => values, }; let new_validity = match array.validity()? { diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index ea279cb46cf..440a7e99b5a 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -8,7 +8,6 @@ use std::fmt::Formatter; use std::sync::Arc; pub use kernel::*; -use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; @@ -202,12 +201,11 @@ pub(crate) fn zip_impl( let return_type = zip_return_dtype(if_true.dtype(), if_false.dtype())?; - if mask.all_true() { - return if_true.cast(return_type); - } - if mask.all_false() { - return if_false.cast(return_type); - } + let mask_values = match mask { + Mask::AllTrue(_) | Mask::AllFalse(0) => return if_true.cast(return_type), + Mask::AllFalse(_) => return if_false.cast(return_type), + Mask::Values(values) => values, + }; // `append_to_builder` requires exact dtype equality, so normalize branch // nullability to the output dtype before appending slices into the builder. @@ -217,8 +215,7 @@ pub(crate) fn zip_impl( zip_impl_with_builder( &if_true, &if_false, - mask.values() - .vortex_expect("zip_impl_with_builder: mask is not all-true or all-false"), + mask_values.as_ref(), builder_with_capacity(&return_type, if_true.len()), ctx, ) diff --git a/vortex-duckdb/src/exporter/validity.rs b/vortex-duckdb/src/exporter/validity.rs index 13774b33f27..736c0e9aed7 100644 --- a/vortex-duckdb/src/exporter/validity.rs +++ b/vortex-duckdb/src/exporter/validity.rs @@ -4,69 +4,69 @@ use vortex::array::ExecutionCtx; use vortex::error::VortexResult; use vortex::mask::Mask; +use vortex::mask::MaskValues; use crate::duckdb::ValidityData; use crate::duckdb::VectorBuffer; use crate::duckdb::VectorRef; use crate::exporter::ColumnExporter; -/// A [`ColumnExporter`] that wraps another exporter with a validity -/// export, allowing you to write data using something else here. +/// Exports validity before delegating values to another [`ColumnExporter`]. struct ValidityExporter { mask: Mask, - /// If the mask's bit buffer is u64-aligned with no sub-byte offset, - /// we can zero-copy it into DuckDB. We hold the ValidityData to keep - /// the underlying memory alive via DuckDB's ref-counting. + + /// Points into `mask` and keeps its bitmap alive when DuckDB can use it directly. zero_copy: Option, + exporter: Box, } -/// Returns true if the bit buffer can be zero-copied as a DuckDB validity mask. +/// Returns the zero-copy validity data for `values`, if DuckDB can read its bit buffer directly. /// -/// Requirements: -/// - No sub-byte bit offset (offset == 0) -/// - The underlying byte buffer is u64-aligned -/// - The underlying byte buffer length is a multiple of 8 (so u64 reads are in-bounds) -fn can_zero_copy_validity(mask: &Mask) -> bool { - let Mask::Values(values) = mask else { - return false; - }; - let bit_buf = values.bit_buffer(); - if bit_buf.offset() != 0 { - return false; +/// The bit buffer must satisfy these requirements: +/// +/// - Its bit offset is zero. +/// - Its byte buffer is aligned for `u64`. +/// - Its byte length is a multiple of `size_of::()`. +fn zero_copy_validity(values: &MaskValues) -> Option { + let bit_buffer = values.bit_buffer(); + if bit_buffer.offset() != 0 { + return None; + } + + let buffer = bit_buffer.inner().clone(); + let bytes = buffer.as_slice(); + let data_ptr = bytes.as_ptr(); + + // DuckDB reads `u64` words. A misaligned pointer causes undefined behavior, and a trailing + // partial word can cause an out-of-bounds read. + if !(data_ptr as usize).is_multiple_of(size_of::()) + || !bytes.len().is_multiple_of(size_of::()) + { + return None; } - let inner = bit_buf.inner(); - let slice = inner.as_slice(); - // DuckDB reads validity as u64 words, so the buffer must be u64-aligned and - // its length must be a multiple of 8 bytes to avoid out-of-bounds reads. - (slice.as_ptr() as usize).is_multiple_of(size_of::()) - && slice.len().is_multiple_of(size_of::()) + + Some(ValidityData { + shared_buffer: VectorBuffer::new(buffer), + data_ptr, + }) } pub(crate) fn new_exporter( mask: Mask, exporter: Box, ) -> Box { - if mask.all_true() { - exporter - } else { - let zero_copy = can_zero_copy_validity(&mask).then(|| { - let Mask::Values(values) = &mask else { - unreachable!() - }; - let buffer = values.bit_buffer().inner().clone(); - let data_ptr = buffer.as_slice().as_ptr(); - ValidityData { - shared_buffer: VectorBuffer::new(buffer), - data_ptr, - } - }); - Box::new(ValidityExporter { - mask, - zero_copy, - exporter, - }) - } + let zero_copy = match &mask { + Mask::AllTrue(_) | Mask::AllFalse(0) => return exporter, + Mask::AllFalse(_) => None, + Mask::Values(values) => zero_copy_validity(values.as_ref()), + }; + + Box::new(ValidityExporter { + mask, + zero_copy, + exporter, + }) } impl ColumnExporter for ValidityExporter { @@ -85,6 +85,7 @@ impl ColumnExporter for ValidityExporter { offset + len <= self.mask.len(), "cannot access outside of array" ); + if unsafe { vector.set_validity_zero_copy(&self.mask, offset, len, self.zero_copy.as_ref()) } { diff --git a/vortex-duckdb/src/exporter/vector.rs b/vortex-duckdb/src/exporter/vector.rs index 3ab0c484318..03a6cc7faab 100644 --- a/vortex-duckdb/src/exporter/vector.rs +++ b/vortex-duckdb/src/exporter/vector.rs @@ -10,15 +10,29 @@ use crate::duckdb::VectorRef; use crate::exporter::copy_from_slice; impl VectorRef { - /// Returns true if all values are null (caller can skip data export). + /// Sets validity from the selected range of `mask`. + /// + /// Returns whether all selected values are null. + /// + /// # Safety + /// + /// - `offset + len` must not exceed `mask.len()`. + /// - `len` must not exceed the vector capacity. pub unsafe fn set_validity(&mut self, mask: &Mask, offset: usize, len: usize) -> bool { unsafe { self.set_validity_zero_copy(mask, offset, len, None) } } - /// Like [`set_validity`](Self::set_validity), but attempts a zero-copy path when - /// `zero_copy` is provided and the offset is u64-aligned. + /// Sets validity from `mask`, using `zero_copy` when its bitmap and `offset` are `u64`-aligned. + /// + /// Returns whether all selected values are null. /// - /// Returns true if all values are null (caller can skip data export). + /// # Safety + /// + /// - `offset + len` must not exceed `mask.len()`. + /// - `len` must not exceed the vector capacity. + /// - A supplied `zero_copy` value must point to the start of the `Mask::Values` bitmap in + /// `mask`. The pointer must be aligned for `u64`, and the buffer must contain only complete + /// `u64` words. pub(super) unsafe fn set_validity_zero_copy( &mut self, mask: &Mask, @@ -35,27 +49,36 @@ impl VectorRef { self.set_all_false_validity(); true } - Mask::Values(arr) => { - let true_count = arr.bit_buffer().slice(offset..(offset + len)).true_count(); + Mask::Values(values) => { + let true_count = values + .bit_buffer() + .slice(offset..(offset + len)) + .true_count(); if true_count == len { self.set_all_true_validity() } else if true_count == 0 { self.set_all_false_validity() - } else if let Some(zc) = zero_copy.filter(|_| offset.is_multiple_of(64)) { + } else if let Some(validity_data) = zero_copy.filter(|_| offset.is_multiple_of(64)) + { let u64_offset = offset / 64; - // SAFETY: the underlying buffer is u64-aligned (checked in - // can_zero_copy_validity) and the VectorBuffer keeps the data alive. - // data_ptr points into the buffer at the start of the validity bitmap. - unsafe { self.set_validity_data(u64_offset, len, zc) }; + + // SAFETY: + // - `zero_copy_validity` points `data_ptr` to an aligned buffer of complete + // `u64` words. + // - `ValidityExporter::export` bounds the selected range to the mask, and this + // branch requires a `u64`-aligned offset, so the buffer contains every word. + // - `shared_buffer` keeps the bitmap alive while DuckDB reads it. + unsafe { self.set_validity_data(u64_offset, len, validity_data) }; } else { - // If zero_copy is available and offset is aligned, we should - // have taken the branch above. Assert this invariant. + // An available zero-copy buffer with an aligned offset must take the branch + // above. assert!( zero_copy.is_none() || !offset.is_multiple_of(64), "zero-copy validity available and offset {offset} is aligned \ but copy path was taken" ); - let source = arr.bit_buffer().inner().as_slice(); + + let source = values.bit_buffer().inner().as_slice(); copy_from_slice( unsafe { self.ensure_validity_slice(len) }, source, diff --git a/vortex-mask/src/intersect_by_rank.rs b/vortex-mask/src/intersect_by_rank.rs index e5a8d6a08e5..f163870d988 100644 --- a/vortex-mask/src/intersect_by_rank.rs +++ b/vortex-mask/src/intersect_by_rank.rs @@ -14,6 +14,7 @@ use vortex_error::VortexExpect; use crate::Mask; use crate::MaskValues; +use crate::MaskValuesRef; trait DepositBits { /// Whether the implementation benefits from short-circuiting on `rank_bits == 0` @@ -442,18 +443,18 @@ where intersect_mask_driven::(self_buffer, mask_indices, true_count) } -/// Check if a mask is sparse. +/// Returns whether a mask is sparse. /// -/// BitBuffer traversal uses u64, hence we conclude that one or fewer values per u64 is sparse -fn mask_is_sparse(values: &Arc) -> bool { +/// [`BitBuffer`] traversal uses `u64` words, so fewer than one selected value per word is sparse. +fn mask_is_sparse(values: &MaskValuesRef) -> bool { values.true_count().saturating_mul(64) < values.len() } -/// Check if a rank mask is sparse +/// Returns whether a rank mask is sparse. /// -/// The mask-driven path becomes worthwhile around ~3% mask density: each set -/// bit costs a select and push, but we save a per-self-chunk popcount + deposit. -fn rank_mask_is_sparse(values: &Arc) -> bool { +/// The mask-driven path becomes worthwhile at approximately 3% mask density. Each set bit costs a +/// select and push, but this avoids a popcount and deposit for each chunk of `self`. +fn rank_mask_is_sparse(values: &MaskValuesRef) -> bool { values.true_count().saturating_mul(32) < values.len() } @@ -557,7 +558,7 @@ impl Mask { } #[cfg(test)] -mod test { +mod tests { use rstest::rstest; use vortex_buffer::BitBuffer;