From 423cee5c8c71e16b313395495b889577a141af53 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sat, 11 Jul 2026 00:15:06 +0100 Subject: [PATCH 1/2] Better fixed-width buffer filtering Signed-off-by: Robert Kruszewski --- .../src/arrays/filter/execute/buffer.rs | 199 +++++++++++++++++- .../arrays/filter/execute/byte_compress.rs | 21 +- .../src/arrays/filter/execute/decimal.rs | 37 +++- vortex-array/src/arrays/filter/execute/mod.rs | 47 +++++ .../src/arrays/filter/execute/primitive.rs | 37 +--- .../src/arrays/filter/execute/slice.rs | 198 +++++++++++++---- .../src/arrays/filter/execute/struct_.rs | 3 + .../src/arrays/filter/execute/take/rank.rs | 5 +- vortex-array/src/arrays/filter/rules.rs | 5 + 9 files changed, 454 insertions(+), 98 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index 73ee3026654..b8244da0983 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -3,32 +3,130 @@ //! Buffer-level filter dispatch. //! -//! Provides [`filter_buffer`] which filters a [`Buffer`] by [`MaskValues`], attempting an -//! in-place filter when the buffer has exclusive ownership. +//! Reuses suitable cached mask representations and otherwise filters directly from the bitmap, +//! avoiding temporary index or range allocations for one-shot filters. + +use std::mem::size_of; use vortex_buffer::Buffer; use vortex_mask::MaskValues; +use crate::arrays::filter::execute::byte_compress; use crate::arrays::filter::execute::slice; +const CACHED_INDICES_MAX_DENSITY: f64 = 0.5; +const IN_PLACE_MIN_DENSITY: f64 = 0.5; +const MIN_SLICES_AVERAGE_RUN_LENGTH: usize = 8; + /// Filter a [`Buffer`] by [`MaskValues`], returning a new buffer. /// -/// This will attempt to filter in-place (via [`Buffer::try_into_mut`]) when the buffer has -/// exclusive ownership, avoiding an extra allocation. +/// Dense uniquely owned buffers are compacted in place. Shared and sparse buffers allocate an +/// exact-sized output and choose between cached indices, cached long ranges, bitmap iteration, +/// and byte-compress based on the mask and element width. pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { - match buffer.try_into_mut() { - Ok(mut buffer_mut) => { - let new_len = slice::filter_slice_mut_by_mask_values(buffer_mut.as_mut_slice(), mask); - buffer_mut.truncate(new_len); - buffer_mut.freeze() + assert_eq!(buffer.len(), mask.len()); + + let buffer = if mask.density() >= IN_PLACE_MIN_DENSITY { + match buffer.try_into_mut() { + Ok(mut buffer_mut) => { + let new_len = filter_slice_in_place(buffer_mut.as_mut_slice(), mask); + buffer_mut.truncate(new_len); + return buffer_mut.freeze(); + } + Err(buffer) => buffer, } - // Otherwise, allocate a new buffer and fill it in. - Err(buffer) => slice::filter_slice_by_mask_values(buffer.as_slice(), mask), + } else { + buffer + }; + + filter_slice(buffer.as_slice(), mask) +} + +fn filter_slice(values: &[T], mask: &MaskValues) -> Buffer { + if let Some(slices) = useful_cached_slices(mask) { + return slice::filter_slice_by_slices(values, slices, mask.true_count()); + } + + if mask.density() <= CACHED_INDICES_MAX_DENSITY + && let Some(indices) = mask.cached_indices() + { + return slice::filter_slice_by_indices(values, indices); + } + + if mask.density() >= byte_compress_density_threshold::() { + byte_compress::filter_buffer(values, mask) + } else { + slice::filter_slice_by_bitmap(values, mask) + } +} + +fn filter_slice_in_place(values: &mut [T], mask: &MaskValues) -> usize { + if let Some(slices) = useful_cached_slices(mask) { + return slice::filter_slice_mut_by_slices(values, slices); + } + + if mask.density() <= CACHED_INDICES_MAX_DENSITY + && let Some(indices) = mask.cached_indices() + { + return slice::filter_slice_mut_by_indices(values, indices); + } + + slice::filter_slice_mut_by_bitmap(values, mask) +} + +fn useful_cached_slices(mask: &MaskValues) -> Option<&[(usize, usize)]> { + mask.cached_slices().filter(|slices| { + slices.len() == 1 || mask.true_count() / slices.len() >= MIN_SLICES_AVERAGE_RUN_LENGTH + }) +} + +fn byte_compress_density_threshold() -> f64 { + let width = size_of::(); + + // The scalar byte-LUT has a later crossover on AArch64, where the word-at-a-time set-bit walk + // is particularly efficient. Wider values also favor the word walk until all-set bytes become + // common. Keep these cases covered by `benches/filter_fixed_width.rs`. + if cfg!(target_arch = "aarch64") { + return match width { + 8 => 0.75, + _ => 0.9, + }; + } + + match width { + 1 => 0.0, + 2 | 4 => 0.5, + 8 => 0.75, + _ => 0.875, + } +} + +/// Materialize sparse indices when enough sibling arrays will reuse the same mask. +pub(super) fn prepare_mask_for_reuse(mask: &MaskValues, consumers: usize) { + if consumers <= 1 || mask.cached_indices().is_some() || mask.cached_slices().is_some() { + return; + } + + let density_threshold = if consumers >= 3 { 0.1 } else { 0.05 }; + if mask.density() > density_threshold { + return; } + + let first = mask.bit_buffer().select(0); + let last = mask.bit_buffer().select(mask.true_count() - 1); + if first + .zip(last) + .is_some_and(|(first, last)| last - first + 1 == mask.true_count()) + { + return; + } + + let _ = mask.indices(); } #[cfg(test)] mod tests { + use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_buffer::buffer; use vortex_mask::Mask; @@ -80,4 +178,83 @@ mod tests { let result = filter_buffer(buf, mask_values(&mask)); assert_eq!(result, buffer![1u32, 2, 3, 4, 6, 7, 8, 10]); } + + #[test] + fn test_filter_shared_buffer_by_cached_indices() { + let buf = Buffer::from(BufferMut::from_iter(0u64..16)); + let shared = buf.clone(); + let mask = Mask::from_indices(16, [1, 5, 9, 15]); + + let result = filter_buffer(buf, mask_values(&mask)); + assert_eq!(result, buffer![1u64, 5, 9, 15]); + assert_eq!(shared.len(), 16); + } + + #[test] + fn test_filter_shared_buffer_by_cached_slices() { + let buf = Buffer::from(BufferMut::from_iter(0u32..32)); + let shared = buf.clone(); + let mask = Mask::from_slices(32, vec![(3, 15), (20, 30)]); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (3u32..15).chain(20..30).collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + assert_eq!(shared.len(), 32); + } + + #[test] + fn test_filter_shared_dense_bitmap() { + let buf = Buffer::from(BufferMut::from_iter(0u16..128)); + let shared = buf.clone(); + let mask = Mask::from_iter((0..128).map(|index| index != 63)); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (0u16..128).filter(|&value| value != 63).collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + assert_eq!(shared.len(), 128); + } + + #[test] + fn test_filter_unaligned_bitmap_words() { + const LEN: usize = 151; + const OFFSET: usize = 5; + + let backing = BitBuffer::from_iter( + std::iter::repeat_n(false, OFFSET).chain((0..LEN).map(|index| index % 7 == 2)), + ); + let mask = Mask::from_buffer(BitBuffer::new_with_offset( + backing.inner().clone(), + LEN, + OFFSET, + )); + let buf = Buffer::from(BufferMut::from_iter(0u64..LEN as u64)); + + let result = filter_buffer(buf, mask_values(&mask)); + let expected = (0u64..LEN as u64) + .filter(|value| value % 7 == 2) + .collect::>(); + assert_eq!(result.as_slice(), expected.as_slice()); + } + + #[test] + fn test_prepare_sparse_mask_for_sibling_reuse() { + let two_consumers = Mask::from_iter((0..100).map(|index| index % 20 == 0)); + let values = mask_values(&two_consumers); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 2); + assert!(values.cached_indices().is_some()); + + let three_consumers = Mask::from_iter((0..100).map(|index| index % 10 == 0)); + let values = mask_values(&three_consumers); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 2); + assert!(values.cached_indices().is_none()); + prepare_mask_for_reuse(values, 3); + assert!(values.cached_indices().is_some()); + + let contiguous = Mask::from_iter((0..100).map(|index| (20..25).contains(&index))); + let values = mask_values(&contiguous); + prepare_mask_for_reuse(values, 3); + assert!(values.cached_indices().is_none()); + } } diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs index 6c58ce4f33b..4ef8c520ce0 100644 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ b/vortex-array/src/arrays/filter/execute/byte_compress.rs @@ -1,20 +1,16 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Byte-level compress for primitive filtering using a `1 << 8 = 256`-entry lookup table. +//! Byte-level compress for fixed-width filtering using a `1 << 8 = 256`-entry lookup table. //! //! For each byte of the mask (8 bits -> 8 source elements), a precomputed //! permutation table compacts the selected bytes in a single indexed copy, //! avoiding the overhead of materializing indices or slices. -use std::mem::size_of; - use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_mask::MaskValues; -const BYTE_COMPRESS_DENSITY_THRESHOLD: f64 = 0.5; - /// For each mask byte (0..256), stores the element indices to keep and the count. /// /// `BYTE_COMPRESS_LUT[mask_byte]` = `([i0, i1, ..., i7], popcount)` where @@ -45,10 +41,10 @@ static BYTE_COMPRESS_LUT: &[([u8; 8], u8); 256] = &{ /// /// Processes the mask one byte at a time (8 source elements per byte), /// using a precomputed permutation to compact selected elements. -pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { - debug_assert_eq!(buffer.len(), mask.len()); +pub(super) fn filter_buffer(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer { + let src = buffer.as_ref(); + debug_assert_eq!(src.len(), mask.len()); - let src = buffer.as_slice(); let true_count = mask.true_count(); if true_count == 0 { @@ -59,14 +55,7 @@ pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Bu let mask_bytes = mask_buffer.inner().as_ref(); let mask_offset = mask_buffer.offset(); - // Fast path: byte-wide values benefit from avoiding index materialization more often. Wider - // values need enough selected values to justify scanning every mask byte directly. - if size_of::() == 1 || mask.density() >= BYTE_COMPRESS_DENSITY_THRESHOLD { - return filter_bitpacked(src, mask_bytes, mask_offset, true_count); - } - - // Slow path: lower-density wide values are better handled by the generic path. - super::slice::filter_slice_by_mask_values(src, mask) + filter_bitpacked(src, mask_bytes, mask_offset, true_count) } fn filter_bitpacked( diff --git a/vortex-array/src/arrays/filter/execute/decimal.rs b/vortex-array/src/arrays/filter/execute/decimal.rs index 5ffb4a963b2..11ff451ea0f 100644 --- a/vortex-array/src/arrays/filter/execute/decimal.rs +++ b/vortex-array/src/arrays/filter/execute/decimal.rs @@ -32,13 +32,48 @@ pub fn filter_decimal(array: &DecimalArray, mask: &Arc) -> DecimalAr } #[cfg(test)] -mod test { +mod tests { + use rstest::rstest; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::filter::execute::decimal::DecimalArray; use crate::compute::conformance::filter::test_filter_conformance; use crate::dtype::DecimalDType; + use crate::dtype::i256; + + #[rstest] + #[case(DecimalArray::from_iter( + [1i8, 2, 3, 4, 5], + DecimalDType::new(2, 0), + ))] + #[case(DecimalArray::from_iter( + [10i16, 20, 30, 40, 50], + DecimalDType::new(3, 0), + ))] + #[case(DecimalArray::from_iter( + [100i32, 200, 300, 400, 500], + DecimalDType::new(5, 0), + ))] + #[case(DecimalArray::from_iter( + [1_000i64, 2_000, 3_000, 4_000, 5_000], + DecimalDType::new(10, 0), + ))] + #[case(DecimalArray::from_iter( + [10_000i128, 20_000, 30_000, 40_000, 50_000], + DecimalDType::new(19, 0), + ))] + #[case(DecimalArray::from_iter( + [1i128, 2, 3, 4, 5].map(i256::from_i128), + DecimalDType::new(39, 0), + ))] + fn test_filter_decimal_physical_type_conformance(#[case] array: DecimalArray) { + test_filter_conformance( + &array.into_array(), + &mut array_session().create_execution_ctx(), + ); + } #[test] fn test_filter_decimal128_conformance() { diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 52b0373ac49..faba9ff9aaf 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -5,6 +5,7 @@ //! //! The main entrypoint is [`execute_filter`] which filters any [`Canonical`] array. +use std::ops::Range; use std::sync::Arc; use vortex_error::VortexExpect; @@ -48,6 +49,16 @@ fn filter_validity(validity: Validity, mask: &Arc) -> Validity { .vortex_expect("Somehow unable to wrap filter around a validity array") } +pub(super) fn contiguous_filter_range(mask: &Mask) -> Option> { + let start = mask.first()?; + let end = mask.last()?.checked_add(1)?; + (end - start == mask.true_count()).then_some(start..end) +} + +pub(super) fn prepare_mask_for_reuse(mask: &MaskValues, consumers: usize) { + buffer::prepare_mask_for_reuse(mask, consumers); +} + /// Check for some fast-path execution conditions before calling [`execute_filter`]. pub(super) fn execute_filter_fast_paths( array: ArrayView<'_, Filter>, @@ -65,6 +76,11 @@ pub(super) fn execute_filter_fast_paths( return Ok(Some(array.child().clone())); } + // Filtering by one contiguous range is exactly a slice and can remain zero-copy. + if let Some(range) = contiguous_filter_range(array.filter_mask()) { + return array.child().slice(range).map(Some); + } + // 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(); @@ -120,3 +136,34 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca } } } + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::*; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + + #[test] + fn contiguous_filter_executes_as_zero_copy_slice() -> VortexResult<()> { + let array = PrimitiveArray::from_iter(0i32..8); + let original = array.to_buffer::(); + let filtered = array + .into_array() + .filter(Mask::from_slices(8, vec![(2, 6)]))? + .execute::(&mut array_session().create_execution_ctx())?; + let filtered_values = filtered.to_buffer::(); + + assert_eq!(filtered_values.as_slice(), &[2, 3, 4, 5]); + assert_eq!(filtered_values.as_ptr(), original.as_ptr().wrapping_add(2)); + Ok(()) + } + + #[test] + fn fragmented_filter_is_not_a_contiguous_range() { + let mask = Mask::from_indices(8, [1, 2, 5, 6]); + assert_eq!(contiguous_filter_range(&mask), None); + } +} diff --git a/vortex-array/src/arrays/filter/execute/primitive.rs b/vortex-array/src/arrays/filter/execute/primitive.rs index b7abd8efa4b..08b22d2fd22 100644 --- a/vortex-array/src/arrays/filter/execute/primitive.rs +++ b/vortex-array/src/arrays/filter/execute/primitive.rs @@ -8,12 +8,8 @@ use vortex_mask::MaskValues; use crate::arrays::PrimitiveArray; use crate::arrays::filter::execute::buffer; -use crate::arrays::filter::execute::byte_compress; use crate::arrays::filter::execute::filter_validity; -use crate::dtype::NativePType; -use crate::dtype::PType; use crate::match_each_native_ptype; -use crate::validity::Validity; pub fn filter_primitive(array: &PrimitiveArray, mask: &Arc) -> PrimitiveArray { let validity = array @@ -21,34 +17,13 @@ pub fn filter_primitive(array: &PrimitiveArray, mask: &Arc) -> Primi .vortex_expect("primitive validity should be derivable"); let filtered_validity = filter_validity(validity, mask); - match array.ptype() { - // Byte-compress avoids materializing indices/slices and processes 8 elements per mask byte. - PType::U8 => filter_byte_compress::(array, filtered_validity, mask), - PType::I8 => filter_byte_compress::(array, filtered_validity, mask), - PType::U16 => filter_byte_compress::(array, filtered_validity, mask), - PType::I16 => filter_byte_compress::(array, filtered_validity, mask), - PType::U32 => filter_byte_compress::(array, filtered_validity, mask), - PType::I32 => filter_byte_compress::(array, filtered_validity, mask), - _ => match_each_native_ptype!(array.ptype(), |T| { - let filtered_buffer = buffer::filter_buffer(array.to_buffer::(), mask.as_ref()); + match_each_native_ptype!(array.ptype(), |T| { + let filtered_buffer = buffer::filter_buffer(array.to_buffer::(), mask.as_ref()); - // SAFETY: We filter both the validity and the buffer with the same mask, so they must - // have the same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } - }), - } -} - -fn filter_byte_compress( - array: &PrimitiveArray, - filtered_validity: Validity, - mask: &Arc, -) -> PrimitiveArray { - let filtered_buffer = byte_compress::filter_buffer(array.to_buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must have the - // same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } + // SAFETY: We filter both the validity and the buffer with the same mask, so they must have + // the same length. + unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } + }) } #[cfg(test)] diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index 1528d272b28..d51de41244e 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -3,47 +3,130 @@ //! Core slice-level filtering algorithms. //! -//! Provides both immutable and mutable (in-place) filtering of typed slices by various mask -//! representations: indices and ranges (slices). +//! Provides both immutable and mutable (in-place) filtering of typed slices by cached mask +//! representations or directly from the mask bitmap. use std::ptr; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_mask::MaskIter; use vortex_mask::MaskValues; -// This is modeled after the constant with the equivalent name in arrow-rs. -pub(super) const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8; +#[inline] +fn for_each_mask_word(mask: &MaskValues, mut f: impl FnMut(u64, usize, usize)) { + let bits = mask.bit_buffer(); + let unaligned = bits.unaligned_chunks(); + let lead = unaligned.lead_padding(); + let mut base = 0; + + if let Some(prefix) = unaligned.prefix() { + let len = (64 - lead).min(mask.len()); + f(prefix >> lead, base, len); + base += len; + } + + for &word in unaligned.chunks() { + f(word, base, 64); + base += 64; + } + + if let Some(suffix) = unaligned.suffix() { + let len = mask.len() - base; + f(suffix, base, len); + base += len; + } + + debug_assert_eq!(base, mask.len()); +} + +#[inline] +fn low_bits_mask(len: usize) -> u64 { + debug_assert!(len <= 64); + if len == 64 { + u64::MAX + } else { + (1u64 << len) - 1 + } +} // --------------------------------------------------------------------------- // Immutable slice filtering // --------------------------------------------------------------------------- -/// Filter a slice by [`MaskValues`], dispatching to the indices or slices path based on a -/// selectivity threshold. -pub(super) fn filter_slice_by_mask_values(slice: &[T], mask: &MaskValues) -> Buffer { +/// Filter a slice from the mask bitmap without materializing indices or ranges. +pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> Buffer { assert_eq!( mask.len(), slice.len(), "Selection mask length must equal the buffer length" ); - match mask.threshold_iter(FILTER_SLICES_SELECTIVITY_THRESHOLD) { - MaskIter::Indices(indices) => filter_slice_by_indices(slice, indices), - MaskIter::Slices(slices) => filter_slice_by_slices(slice, slices), - } + let output_len = mask.true_count(); + let mut out = BufferMut::::with_capacity(output_len); + let src_ptr = slice.as_ptr(); + let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + let mut write_pos = 0; + + for_each_mask_word(mask, |word, word_start, word_len| { + let all_selected = low_bits_mask(word_len); + debug_assert_eq!(word & !all_selected, 0); + if word == all_selected { + // SAFETY: a full mask word selects `word_len` in-bounds source values and the output + // was allocated for every selected value. + unsafe { + ptr::copy_nonoverlapping(src_ptr.add(word_start), out_ptr.add(write_pos), word_len); + } + write_pos += word_len; + } else { + let mut selected = word; + while selected != 0 { + let index = word_start + selected.trailing_zeros() as usize; + // SAFETY: set bits are limited to `word_len`, and the output was allocated for + // exactly `mask.true_count()` values. + unsafe { + out_ptr.add(write_pos).write(*src_ptr.add(index)); + } + write_pos += 1; + selected &= selected - 1; + } + } + }); + + debug_assert_eq!(write_pos, output_len); + // SAFETY: every output slot was initialized exactly once above. + unsafe { out.set_len(output_len) }; + out.freeze() } /// Filter a slice by a set of strictly increasing indices. -fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { - Buffer::::from_trusted_len_iter(indices.iter().map(|&idx| slice[idx])) +pub(super) fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { + debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); + + let mut out = BufferMut::::with_capacity(indices.len()); + let src_ptr = slice.as_ptr(); + let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); + + for (write_pos, &index) in indices.iter().enumerate() { + // SAFETY: mask indices are validated when the mask is constructed and the output has one + // slot allocated for every index. + unsafe { out_ptr.add(write_pos).write(*src_ptr.add(index)) }; + } + + // SAFETY: the loop initialized every output slot. + unsafe { out.set_len(indices.len()) }; + out.freeze() } /// Filter a slice by a set of strictly increasing `(start, end)` ranges. -fn filter_slice_by_slices(slice: &[T], slices: &[(usize, usize)]) -> Buffer { - let output_len: usize = slices.iter().map(|(start, end)| end - start).sum(); - +pub(super) fn filter_slice_by_slices( + slice: &[T], + slices: &[(usize, usize)], + output_len: usize, +) -> Buffer { + debug_assert_eq!( + output_len, + slices.iter().map(|(start, end)| end - start).sum::() + ); let mut out = BufferMut::::with_capacity(output_len); for (start, end) in slices { out.extend_from_slice(&slice[*start..*end]); @@ -56,40 +139,83 @@ fn filter_slice_by_slices(slice: &[T], slices: &[(usize, usize)]) -> Bu // Mutable (in-place) slice filtering // --------------------------------------------------------------------------- -/// Filter a mutable slice in-place by [`MaskValues`], returning the new valid length. -/// -/// We always use the slices path here because iterating over indices will have strictly more -/// loop iterations than slices (more branches), and the overhead of batched `ptr::copy(len)` is -/// not that high. -pub(super) fn filter_slice_mut_by_mask_values( - slice: &mut [T], - mask: &MaskValues, -) -> usize { +/// Filter a mutable slice in-place from the mask bitmap, returning the new valid length. +pub(super) fn filter_slice_mut_by_bitmap(slice: &mut [T], mask: &MaskValues) -> usize { assert_eq!( slice.len(), mask.len(), "Mask length must equal the slice length" ); - filter_slice_mut_by_slices(slice, mask.slices()) + let ptr = slice.as_mut_ptr(); + let mut write_pos = 0; + + for_each_mask_word(mask, |word, word_start, word_len| { + let all_selected = low_bits_mask(word_len); + debug_assert_eq!(word & !all_selected, 0); + if word == all_selected { + if write_pos != word_start { + // SAFETY: source and destination are in bounds and may overlap while compacting + // toward the start of the same allocation. + unsafe { ptr::copy(ptr.add(word_start), ptr.add(write_pos), word_len) }; + } + write_pos += word_len; + } else { + let mut selected = word; + while selected != 0 { + let index = word_start + selected.trailing_zeros() as usize; + if write_pos != index { + // SAFETY: set bits are limited to `word_len` and stable compaction guarantees + // `write_pos <= index`. + unsafe { ptr::copy(ptr.add(index), ptr.add(write_pos), 1) }; + } + write_pos += 1; + selected &= selected - 1; + } + } + }); + + debug_assert_eq!(write_pos, mask.true_count()); + write_pos +} + +/// Filter a mutable slice in-place by strictly increasing indices. +pub(super) fn filter_slice_mut_by_indices(slice: &mut [T], indices: &[usize]) -> usize { + debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); + + let ptr = slice.as_mut_ptr(); + for (write_pos, &index) in indices.iter().enumerate() { + if write_pos != index { + // SAFETY: mask indices are in bounds and stable compaction guarantees + // `write_pos <= index`. + unsafe { ptr::copy(ptr.add(index), ptr.add(write_pos), 1) }; + } + } + indices.len() } /// Filter a mutable slice in-place by a set of `(start, end)` ranges, returning the new length. -fn filter_slice_mut_by_slices(slice: &mut [T], slices: &[(usize, usize)]) -> usize { +pub(super) fn filter_slice_mut_by_slices( + slice: &mut [T], + slices: &[(usize, usize)], +) -> usize { let mut write_pos = 0; // For each range in the selection, copy all of the elements to the current write position. for &(start, end) in slices { let len = end - start; - // SAFETY: Slices should be within bounds. - unsafe { - ptr::copy( - slice.as_ptr().add(start), - slice.as_mut_ptr().add(write_pos), - len, - ) - }; + if write_pos != start { + // SAFETY: mask slices are in bounds and source and destination may overlap while + // compacting toward the start of the same allocation. + unsafe { + ptr::copy( + slice.as_ptr().add(start), + slice.as_mut_ptr().add(write_pos), + len, + ) + }; + } write_pos += len; } diff --git a/vortex-array/src/arrays/filter/execute/struct_.rs b/vortex-array/src/arrays/filter/execute/struct_.rs index e4a8157d2e2..85866fde196 100644 --- a/vortex-array/src/arrays/filter/execute/struct_.rs +++ b/vortex-array/src/arrays/filter/execute/struct_.rs @@ -10,9 +10,12 @@ use vortex_mask::MaskValues; use crate::ArrayRef; use crate::arrays::StructArray; use crate::arrays::filter::execute::filter_validity; +use crate::arrays::filter::execute::prepare_mask_for_reuse; use crate::arrays::struct_::StructArrayExt; pub fn filter_struct(array: &StructArray, mask: &Arc) -> StructArray { + prepare_mask_for_reuse(mask, array.struct_fields().nfields()); + let filtered_validity = filter_validity( array .validity() diff --git a/vortex-array/src/arrays/filter/execute/take/rank.rs b/vortex-array/src/arrays/filter/execute/take/rank.rs index efbfa768df3..d0f966609b8 100644 --- a/vortex-array/src/arrays/filter/execute/take/rank.rs +++ b/vortex-array/src/arrays/filter/execute/take/rank.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_bail; use vortex_mask::AllOr; use vortex_mask::Mask; +use super::super::contiguous_filter_range; use super::small_take_rank_lookup_len; use crate::arrays::PrimitiveArray; use crate::dtype::IntegerPType; @@ -153,7 +154,5 @@ pub(in crate::arrays::filter) fn contiguous_sequential_take_range Option { - let start = filter.first()?; - let end = filter.last()?.checked_add(1)?; - (end - start == filter.true_count()).then_some(start) + contiguous_filter_range(filter).map(|range| range.start) } diff --git a/vortex-array/src/arrays/filter/rules.rs b/vortex-array/src/arrays/filter/rules.rs index ffa5c64bd61..a0d4313ed08 100644 --- a/vortex-array/src/arrays/filter/rules.rs +++ b/vortex-array/src/arrays/filter/rules.rs @@ -14,6 +14,7 @@ use crate::arrays::StructArray; use crate::arrays::filter::FilterArrayExt; use crate::arrays::filter::FilterReduce; use crate::arrays::filter::FilterReduceAdaptor; +use crate::arrays::filter::execute::prepare_mask_for_reuse; use crate::arrays::struct_::StructDataParts; use crate::optimizer::rules::ArrayReduceRule; use crate::optimizer::rules::ParentRuleSet; @@ -66,6 +67,10 @@ impl ArrayReduceRule for FilterStructRule { .. } = struct_array.into_owned().into_data_parts(); + if let Some(values) = mask.values() { + prepare_mask_for_reuse(values, fields.len()); + } + let filtered_validity = validity.filter(mask)?; let filtered_fields = fields From 8633f2647f4a2e9e07b69f6eadc126b1dd5e2c81 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 13 Jul 2026 10:01:29 +0100 Subject: [PATCH 2/2] lessassert Signed-off-by: Robert Kruszewski --- vortex-array/src/arrays/filter/execute/slice.rs | 8 -------- 1 file changed, 8 deletions(-) diff --git a/vortex-array/src/arrays/filter/execute/slice.rs b/vortex-array/src/arrays/filter/execute/slice.rs index d51de41244e..d6ec94db062 100644 --- a/vortex-array/src/arrays/filter/execute/slice.rs +++ b/vortex-array/src/arrays/filter/execute/slice.rs @@ -100,8 +100,6 @@ pub(super) fn filter_slice_by_bitmap(slice: &[T], mask: &MaskValues) -> /// Filter a slice by a set of strictly increasing indices. pub(super) fn filter_slice_by_indices(slice: &[T], indices: &[usize]) -> Buffer { - debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); - let mut out = BufferMut::::with_capacity(indices.len()); let src_ptr = slice.as_ptr(); let out_ptr = out.spare_capacity_mut().as_mut_ptr().cast::(); @@ -123,10 +121,6 @@ pub(super) fn filter_slice_by_slices( slices: &[(usize, usize)], output_len: usize, ) -> Buffer { - debug_assert_eq!( - output_len, - slices.iter().map(|(start, end)| end - start).sum::() - ); let mut out = BufferMut::::with_capacity(output_len); for (start, end) in slices { out.extend_from_slice(&slice[*start..*end]); @@ -181,8 +175,6 @@ pub(super) fn filter_slice_mut_by_bitmap(slice: &mut [T], mask: &MaskVa /// Filter a mutable slice in-place by strictly increasing indices. pub(super) fn filter_slice_mut_by_indices(slice: &mut [T], indices: &[usize]) -> usize { - debug_assert!(indices.last().is_none_or(|&index| index < slice.len())); - let ptr = slice.as_mut_ptr(); for (write_pos, &index) in indices.iter().enumerate() { if write_pos != index {