only use slices

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-compute/src/filter/buffer_mut.rs

@@ -4,16 +4,12 @@
 use core::ptr;
 
 use vortex_buffer::BufferMut;
-use vortex_mask::{Mask, MaskIter};
+use vortex_mask::Mask;
 
 use crate::filter::Filter;
 
 // TODO(connor): Implement `Filter` for more combinations of `Filter`
 
-// TODO(connor): Figure out if this threshold makes sense for in-place filter.
-// This is modeled after the constant with the equivalent name in arrow-rs.
-const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8;
-
 impl<T: Copy> Filter for &mut BufferMut<T> {
     type Output = ();
 
@@ -27,54 +23,32 @@ impl<T: Copy> Filter for &mut BufferMut<T> {
         match selection_mask {
             Mask::AllTrue(_) => {}
             Mask::AllFalse(_) => self.clear(),
-            // SAFETY: We checked above that the selection mask has the same length as the buffer.
-            Mask::Values(values) => unsafe {
-                match values.threshold_iter(FILTER_SLICES_SELECTIVITY_THRESHOLD) {
-                    MaskIter::Indices(indices) => filter_indices_in_place(self, indices),
-                    MaskIter::Slices(slices) => filter_slices_in_place(self, slices),
-                }
-            },
+            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, slices) };
+
+                // Truncate the buffer to the new length.
+                self.truncate(new_len);
+            }
         }
     }
 }
 
-/// Filters a buffer in-place using indices to determine which values to keep.
-///
-/// # Safety
-///
-/// The indices must be in the range of the `buffer`.
-unsafe fn filter_indices_in_place<T: Copy>(buffer: &mut BufferMut<T>, indices: &[usize]) {
-    let slice = buffer.as_mut_slice();
-    let mut write_idx = 0;
-
-    // For each index in the selection, copy the element to the current write position.
-    for &read_idx in indices {
-        // Note that we could add an if statement here that checks `if read_idx != write_idx` and
-        // use `ptr::copy_nonoverlapping`, but it's probably better to just avoid the branch
-        // misprediction.
-
-        // SAFETY: Both indices are within bounds since indices come from a valid mask.
-        unsafe {
-            ptr::copy(
-                slice.as_ptr().add(read_idx),
-                slice.as_mut_ptr().add(write_idx),
-                1,
-            )
-        };
-        write_idx += 1;
-    }
-
-    // Truncate the buffer to the new length.
-    buffer.truncate(write_idx);
-}
-
 /// 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`.
-unsafe fn filter_slices_in_place<T: Copy>(buffer: &mut BufferMut<T>, slices: &[(usize, usize)]) {
-    let slice = buffer.as_mut_slice();
+#[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.
@@ -84,21 +58,19 @@ unsafe fn filter_slices_in_place<T: Copy>(buffer: &mut BufferMut<T>, slices: &[(
 
         let len = end - start;
 
-        // SAFETY: The ranges are within bounds since they come from a valid mask for the
-        // buffer.
+        // SAFETY: The safety contract enforces that all ranges are within bounds.
         unsafe {
             ptr::copy(
-                slice.as_ptr().add(start),
-                slice.as_mut_ptr().add(write_pos),
+                buffer.as_ptr().add(start),
+                buffer.as_mut_ptr().add(write_pos),
                 len,
             )
         };
 
         write_pos += len;
     }
 
-    // Truncate the buffer to the new length.
-    buffer.truncate(write_pos);
+    write_pos
 }
 
 #[cfg(test)]
@@ -213,26 +185,4 @@ mod tests {
 
         buf.filter(&mask);
     }
-
-    #[test]
-    fn test_filter_indices_direct() {
-        let mut buf = buffer_mut![100u32, 200, 300, 400, 500];
-        unsafe { filter_indices_in_place(&mut buf, &[1, 3, 4]) };
-        assert_eq!(buf.as_slice(), &[200, 400, 500]);
-    }
-
-    #[test]
-    fn test_filter_slices_direct() {
-        let mut buf = buffer_mut![1u32, 2, 3, 4, 5, 6, 7];
-        unsafe { filter_slices_in_place(&mut buf, &[(1, 3), (5, 7)]) };
-        assert_eq!(buf.as_slice(), &[2, 3, 6, 7]);
-    }
-
-    #[test]
-    fn test_filter_overlapping_slices() {
-        // Test that overlapping regions are handled correctly.
-        let mut buf = buffer_mut![1u32, 2, 3, 4, 5, 6, 7, 8];
-        unsafe { filter_slices_in_place(&mut buf, &[(2, 6)]) };
-        assert_eq!(buf.as_slice(), &[3, 4, 5, 6]);
-    }
 }