Better fixed-width buffer filtering

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

Author: robert3005

vortex-array/Cargo.toml

@@ -248,6 +248,10 @@ harness = false
 name = "filter_bool"
 harness = false
 
+[[bench]]
+name = "filter_fixed_width"
+harness = false
+
 [[bench]]
 name = "list_length"
 harness = false

vortex-array/benches/filter_fixed_width.rs

@@ -0,0 +1,158 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Compact fixed-width filter benchmarks.
+//!
+//! The cases cover the dispatch dimensions without a large Cartesian product so CodSpeed's
+//! instruction-count simulation remains inexpensive.
+
+#![expect(
+    clippy::cast_possible_truncation,
+    clippy::cast_precision_loss,
+    clippy::unwrap_used
+)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use vortex_array::ArrayRef;
+use vortex_array::Canonical;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::DecimalArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DecimalDType;
+use vortex_array::dtype::i256;
+use vortex_buffer::BitBuffer;
+use vortex_mask::Mask;
+use vortex_session::VortexSession;
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
+
+// Keep each case small: the sweep has 24 cases and the targeted sections add only seven more.
+const LEN: usize = 16_384;
+const DENSITIES: &[f64] = &[0.01, 0.5, 0.8, 0.95];
+const CACHED_DENSITIES: &[f64] = &[0.01, 0.1];
+
+#[derive(Clone, Copy, Debug)]
+enum Pattern {
+    Random,
+    Runs,
+    Contiguous,
+}
+
+const PATTERNS: &[Pattern] = &[Pattern::Random, Pattern::Runs, Pattern::Contiguous];
+
+fn random_mask(density: f64) -> Mask {
+    let threshold = (density * u64::MAX as f64) as u64;
+    let mut state = 0x1234_5678_9abc_def0u64;
+    Mask::from_buffer(BitBuffer::from_iter((0..LEN).map(|_| {
+        state ^= state << 13;
+        state ^= state >> 7;
+        state ^= state << 17;
+        state <= threshold
+    })))
+}
+
+fn pattern_mask(pattern: Pattern) -> Mask {
+    match pattern {
+        Pattern::Random => random_mask(0.5),
+        Pattern::Runs => Mask::from_iter((0..LEN).map(|index| (index / 32).is_multiple_of(2))),
+        Pattern::Contiguous => Mask::from_slices(LEN, vec![(LEN / 4, LEN * 3 / 4)]),
+    }
+}
+
+fn bench_filter(
+    bencher: Bencher,
+    array: ArrayRef,
+    make_mask: impl Fn() -> Mask + Sync,
+    cache_indices: bool,
+) {
+    bencher
+        .with_inputs(|| {
+            let mask = make_mask();
+            if cache_indices {
+                let _ = mask.values().unwrap().indices();
+            }
+            (array.clone(), mask, SESSION.create_execution_ctx())
+        })
+        .bench_refs(|(array, mask, ctx)| {
+            divan::black_box(
+                array
+                    .clone()
+                    .filter(mask.clone())
+                    .unwrap()
+                    .execute::<Canonical>(ctx)
+                    .unwrap(),
+            );
+        });
+}
+
+fn i8_array() -> ArrayRef {
+    PrimitiveArray::from_iter((0..LEN).map(|index| index as i8)).into_array()
+}
+
+fn i16_array() -> ArrayRef {
+    PrimitiveArray::from_iter((0..LEN).map(|index| index as i16)).into_array()
+}
+
+fn i32_array() -> ArrayRef {
+    PrimitiveArray::from_iter((0..LEN).map(|index| index as i32)).into_array()
+}
+
+fn i64_array() -> ArrayRef {
+    PrimitiveArray::from_iter((0..LEN).map(|index| index as i64)).into_array()
+}
+
+fn i128_array() -> ArrayRef {
+    DecimalArray::from_iter(
+        (0..LEN).map(|index| index as i128),
+        DecimalDType::new(19, 0),
+    )
+    .into_array()
+}
+
+fn i256_array() -> ArrayRef {
+    DecimalArray::from_iter(
+        (0..LEN).map(|index| i256::from_i128(index as i128)),
+        DecimalDType::new(39, 0),
+    )
+    .into_array()
+}
+
+macro_rules! random_density_benchmark {
+    ($name:ident, $array:ident) => {
+        #[divan::bench(args = DENSITIES)]
+        fn $name(bencher: Bencher, density: f64) {
+            bench_filter(bencher, $array(), || random_mask(density), false);
+        }
+    };
+}
+
+random_density_benchmark!(random_i8, i8_array);
+random_density_benchmark!(random_i16, i16_array);
+random_density_benchmark!(random_i32, i32_array);
+random_density_benchmark!(random_i64, i64_array);
+random_density_benchmark!(random_i128, i128_array);
+random_density_benchmark!(random_i256, i256_array);
+
+#[divan::bench(args = PATTERNS)]
+fn patterns_i128(bencher: Bencher, pattern: Pattern) {
+    bench_filter(bencher, i128_array(), || pattern_mask(pattern), false);
+}
+
+#[divan::bench(args = CACHED_DENSITIES)]
+fn cached_indices_i32(bencher: Bencher, density: f64) {
+    bench_filter(bencher, i32_array(), || random_mask(density), true);
+}
+
+#[divan::bench(args = CACHED_DENSITIES)]
+fn cached_indices_i128(bencher: Bencher, density: f64) {
+    bench_filter(bencher, i128_array(), || random_mask(density), true);
+}

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());
+    }
 }