Feature: add in-place filter for BufferMut + benchmarks (#5238)

Author: connortsui20

Cargo.lock

@@ -34,3 +34,10 @@ num-traits = { workspace = true }
 [features]
 default = ["arrow"]
 arrow = ["dep:arrow-array", "dep:arrow-buffer", "dep:arrow-schema"]
+
+[dev-dependencies]
+divan = { workspace = true }
+
+[[bench]]
+name = "filter_buffer_mut"
+harness = false

vortex-compute/Cargo.toml

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! In-place filter benchmarks for `BufferMut`.
+
+use divan::Bencher;
+use vortex_buffer::BufferMut;
+use vortex_compute::filter::Filter;
+use vortex_mask::Mask;
+
+fn main() {
+    divan::main();
+}
+
+const BUFFER_SIZE: usize = 1024;
+
+const SELECTIVITIES: &[f64] = &[
+    0.01, 0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90, 0.99,
+];
+
+fn create_test_buffer<T>(size: usize) -> BufferMut<T>
+where
+    T: Copy + Default + From<u8>,
+{
+    let mut buffer = BufferMut::with_capacity(size);
+    for i in 0..size {
+        #[expect(clippy::cast_possible_truncation)]
+        buffer.push(T::from((i % 256) as u8));
+    }
+    buffer
+}
+
+fn generate_mask(len: usize, selectivity: f64) -> Mask {
+    #[expect(clippy::cast_possible_truncation)]
+    #[expect(clippy::cast_sign_loss)]
+    let num_selected = ((len as f64) * selectivity).round() as usize;
+
+    let mut selection = vec![false; len];
+    let mut indices: Vec<usize> = (0..len).collect();
+
+    // Simple deterministic shuffle.
+    for i in (1..len).rev() {
+        let j = (i * 7 + 13) % (i + 1);
+        indices.swap(i, j);
+    }
+
+    for i in 0..num_selected.min(len) {
+        selection[indices[i]] = true;
+    }
+
+    Mask::from_iter(selection)
+}
+
+#[derive(Copy, Clone, Default)]
+#[allow(dead_code)]
+struct LargeElement([u8; 32]);
+
+impl From<u8> for LargeElement {
+    fn from(value: u8) -> Self {
+        LargeElement([value; 32])
+    }
+}
+
+#[divan::bench(types = [u8, u32, u64, LargeElement], args = SELECTIVITIES, sample_count = 1000)]
+fn filter_selectivity<T: Copy + Default + From<u8>>(bencher: Bencher, selectivity: f64) {
+    let mask = generate_mask(BUFFER_SIZE, selectivity);
+    bencher
+        .with_inputs(|| {
+            let buffer = create_test_buffer::<T>(BUFFER_SIZE);
+            (buffer, mask.clone())
+        })
+        .bench_values(|(mut buffer, mask)| {
+            buffer.filter(&mask);
+            divan::black_box(buffer);
+        });
+}

vortex-compute/benches/filter_buffer_mut.rs

@@ -12,21 +12,48 @@ const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8;
 impl<T: Copy> Filter for &Buffer<T> {
     type Output = Buffer<T>;
 
-    fn filter(self, mask: &Mask) -> Buffer<T> {
-        assert_eq!(mask.len(), self.len());
-        match mask {
+    fn filter(self, selection_mask: &Mask) -> Buffer<T> {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the buffer length"
+        );
+
+        match selection_mask {
             Mask::AllTrue(_) => self.clone(),
             Mask::AllFalse(_) => Buffer::empty(),
             Mask::Values(v) => match v.threshold_iter(FILTER_SLICES_SELECTIVITY_THRESHOLD) {
                 MaskIter::Indices(indices) => filter_indices(self.as_slice(), indices),
                 MaskIter::Slices(slices) => {
-                    filter_slices(self.as_slice(), mask.true_count(), slices)
+                    filter_slices(self.as_slice(), selection_mask.true_count(), slices)
                 }
             },
         }
     }
 }
 
+impl<T: Copy> Filter for Buffer<T> {
+    type Output = Self;
+
+    fn filter(self, selection_mask: &Mask) -> Self {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the buffer length"
+        );
+
+        // If we have exclusive access, we can perform the filter in place.
+        match self.try_into_mut() {
+            Ok(mut buffer_mut) => {
+                (&mut buffer_mut).filter(selection_mask);
+                buffer_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&Buffer` impl).
+            Err(buffer) => (&buffer).filter(selection_mask),
+        }
+    }
+}
+
 fn filter_indices<T: Copy>(values: &[T], indices: &[usize]) -> Buffer<T> {
     Buffer::<T>::from_trusted_len_iter(indices.iter().map(|&idx| values[idx]))
 }

vortex-compute/src/filter/buffer.rs

@@ -0,0 +1,193 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use core::ptr;
+
+use vortex_buffer::BufferMut;
+use vortex_mask::Mask;
+
+use crate::filter::Filter;
+
+impl<T: Copy> Filter for &mut BufferMut<T> {
+    type Output = ();
+
+    fn filter(self, selection_mask: &Mask) {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the buffer length"
+        );
+
+        match selection_mask {
+            Mask::AllTrue(_) => {}
+            Mask::AllFalse(_) => self.clear(),
+            Mask::Values(values) => {
+                // We choose to _always_ use slices here because iterating over indices will have
+                // strictly more loop iterations than slices, and the overhead over batched
+                // `ptr::copy(len)` is not worth it.
+                let slices = values.slices();
+
+                // SAFETY: We checked above that the selection mask has the same length as the
+                // buffer.
+                let new_len = unsafe { filter_slices_in_place(self.as_mut_slice(), slices) };
+
+                debug_assert!(
+                    new_len <= self.len(),
+                    "The new length was somehow larger after filter"
+                );
+
+                // Truncate the buffer to the new length.
+                // SAFETY: The new length cannot be larger than the old length, so all values must
+                // be initialized.
+                unsafe { self.set_len(new_len) };
+            }
+        }
+    }
+}
+
+/// Filters a buffer in-place using slice ranges to determine which values to keep.
+///
+/// Returns the new length of the buffer.
+///
+/// # Safety
+///
+/// The slice ranges must be in the range of the `buffer`.
+#[must_use = "The caller should set the new length of the buffer"]
+unsafe fn filter_slices_in_place<T: Copy>(buffer: &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 {
+        // Note that we could add an if statement here that checks `if read_idx != write_idx`, but
+        // it's probably better to just avoid the branch misprediction.
+
+        let len = end - start;
+
+        // SAFETY: The safety contract enforces that all ranges are within bounds.
+        unsafe {
+            ptr::copy(
+                buffer.as_ptr().add(start),
+                buffer.as_mut_ptr().add(write_pos),
+                len,
+            )
+        };
+
+        write_pos += len;
+    }
+
+    write_pos
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::{BufferMut, buffer_mut};
+    use vortex_mask::Mask;
+
+    use super::*;
+
+    #[test]
+    fn test_filter_all_true() {
+        let mut buf = buffer_mut![1u32, 2, 3, 4, 5];
+        let mask = Mask::new_true(5);
+
+        buf.filter(&mask);
+        assert_eq!(buf.as_slice(), &[1, 2, 3, 4, 5]);
+    }
+
+    #[test]
+    fn test_filter_all_false() {
+        let mut buf = buffer_mut![1u32, 2, 3, 4, 5];
+        let mask = Mask::new_false(5);
+
+        buf.filter(&mask);
+        assert!(buf.is_empty());
+    }
+
+    #[test]
+    fn test_filter_sparse() {
+        let mut buf = buffer_mut![10u32, 20, 30, 40, 50];
+        // Select indices 0, 2, 4 (sparse selection).
+        let mask = Mask::from_iter([true, false, true, false, true]);
+
+        buf.filter(&mask);
+        assert_eq!(buf.as_slice(), &[10, 30, 50]);
+    }
+
+    #[test]
+    fn test_filter_dense() {
+        let mut buf = buffer_mut![1u32, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        // Dense selection (80% selected).
+        let mask = Mask::from_iter([true, true, true, true, false, true, true, true, false, true]);
+
+        buf.filter(&mask);
+        assert_eq!(buf.as_slice(), &[1, 2, 3, 4, 6, 7, 8, 10]);
+    }
+
+    #[test]
+    fn test_filter_single_element_kept() {
+        let mut buf = buffer_mut![1u32, 2, 3, 4, 5];
+        let mask = Mask::from_iter([false, false, true, false, false]);
+
+        buf.filter(&mask);
+        assert_eq!(buf.as_slice(), &[3]);
+    }
+
+    #[test]
+    fn test_filter_first_last() {
+        let mut buf = buffer_mut![1u32, 2, 3, 4, 5];
+        let mask = Mask::from_iter([true, false, false, false, true]);
+
+        buf.filter(&mask);
+        assert_eq!(buf.as_slice(), &[1, 5]);
+    }
+
+    #[test]
+    fn test_filter_alternating() {
+        let mut buf = buffer_mut![1u32, 2, 3, 4, 5, 6];
+        let mask = Mask::from_iter([true, false, true, false, true, false]);
+
+        buf.filter(&mask);
+        assert_eq!(buf.as_slice(), &[1, 3, 5]);
+    }
+
+    #[test]
+    fn test_filter_empty_buffer() {
+        let mut buf: BufferMut<u32> = BufferMut::with_capacity(0);
+        let mask = Mask::new_false(0);
+
+        buf.filter(&mask);
+        assert!(buf.is_empty());
+    }
+
+    #[test]
+    fn test_filter_contiguous_regions() {
+        let mut buf = buffer_mut![1u32, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        // Two contiguous regions: [0..3] and [7..10].
+        let mask = Mask::from_iter([
+            true, true, true, false, false, false, false, true, true, true,
+        ]);
+
+        buf.filter(&mask);
+        assert_eq!(buf.as_slice(), &[1, 2, 3, 8, 9, 10]);
+    }
+
+    #[test]
+    fn test_filter_large_buffer() {
+        let mut buf: BufferMut<u32> = BufferMut::from_iter(0..1000);
+        // Keep every third element.
+        let mask = Mask::from_iter((0..1000).map(|i| i % 3 == 0));
+
+        buf.filter(&mask);
+        let expected: Vec<u32> = (0..1000).filter(|i| i % 3 == 0).collect();
+        assert_eq!(buf.as_slice(), &expected[..]);
+    }
+
+    #[test]
+    #[should_panic(expected = "Selection mask length must equal the buffer length")]
+    fn test_filter_length_mismatch() {
+        let mut buf = buffer_mut![1u32, 2, 3];
+        let mask = Mask::new_true(5); // Wrong length.
+
+        buf.filter(&mask);
+    }
+}

vortex-compute/src/filter/buffer_mut.rs

@@ -6,6 +6,7 @@
 mod bitbuffer;
 mod bool;
 mod buffer;
+mod buffer_mut;
 mod mask;
 
 use vortex_mask::Mask;
@@ -22,5 +23,5 @@ pub trait Filter {
     /// # Panics
     ///
     /// If the length of the mask does not equal the length of the value being filtered.
-    fn filter(self, mask: &Mask) -> Self::Output;
+    fn filter(self, selection_mask: &Mask) -> Self::Output;
 }