Better fixed-width buffer filtering

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

vortex-array/src/arrays/filter/execute/buffer.rs

@@ -3,32 +3,130 @@
 
 //! Buffer-level filter dispatch.
 //!
-//! Provides [`filter_buffer`] which filters a [`Buffer<T>`] 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<T>`] 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<T: Copy>(buffer: Buffer<T>, mask: &MaskValues) -> Buffer<T> {
-    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<T: Copy>(values: &[T], mask: &MaskValues) -> Buffer<T> {
+    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::<T>() {
+        byte_compress::filter_buffer(values, mask)
+    } else {
+        slice::filter_slice_by_bitmap(values, mask)
+    }
+}
+
+fn filter_slice_in_place<T: Copy>(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<T>() -> f64 {
+    let width = size_of::<T>();
+
+    // 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::<Vec<_>>();
+        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::<Vec<_>>();
+        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::<Vec<_>>();
+        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());
+    }
 }

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<T: Copy>(buffer: Buffer<T>, mask: &MaskValues) -> Buffer<T> {
-    debug_assert_eq!(buffer.len(), mask.len());
+pub(super) fn filter_buffer<T: Copy>(buffer: impl AsRef<[T]>, mask: &MaskValues) -> Buffer<T> {
+    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<T: Copy>(buffer: Buffer<T>, 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::<T>() == 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<T: Copy>(

vortex-array/src/arrays/filter/execute/decimal.rs

@@ -32,13 +32,48 @@ pub fn filter_decimal(array: &DecimalArray, mask: &Arc<MaskValues>) -> 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() {

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<MaskValues>) -> Validity {
         .vortex_expect("Somehow unable to wrap filter around a validity array")
 }
 
+pub(super) fn contiguous_filter_range(mask: &Mask) -> Option<Range<usize>> {
+    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<MaskValues>) -> 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::<i32>();
+        let filtered = array
+            .into_array()
+            .filter(Mask::from_slices(8, vec![(2, 6)]))?
+            .execute::<PrimitiveArray>(&mut array_session().create_execution_ctx())?;
+        let filtered_values = filtered.to_buffer::<i32>();
+
+        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);
+    }
+}