Feature: Add more Filter<_> implementations (#6173)

Allows us to filter by a `MaskValues` if we want to, as well as by
indices or slices directly.

This is just a nice to have for using the code in `vortex-compute` on
top of the incoming changes in
#6152

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-buffer/src/bit/mod.rs

@@ -26,36 +26,40 @@ pub use buf_mut::*;
 pub use view::*;
 
 /// Get the bit value at `index` out of `buf`.
+///
+/// # Panics
+///
+/// Panics if `index` is not between 0 and length of `buf * 8`.
 #[inline(always)]
 pub fn get_bit(buf: &[u8], index: usize) -> bool {
     buf[index / 8] & (1 << (index % 8)) != 0
 }
 
-/// Get the bit value at `index` out of `buf` without bounds checking
+/// Get the bit value at `index` out of `buf` without bounds checking.
 ///
 /// # Safety
 ///
-/// `index` must be between 0 and length of `buf`
+/// `index` must be between 0 and length of `buf * 8`.
 #[inline(always)]
 pub unsafe fn get_bit_unchecked(buf: *const u8, index: usize) -> bool {
     (unsafe { *buf.add(index / 8) } & (1 << (index % 8))) != 0
 }
 
-/// Set the bit value at `index` in `buf` without bounds checking
+/// Set the bit value at `index` in `buf` without bounds checking.
 ///
 /// # Safety
 ///
-/// `index` must be between 0 and length of `buf`
+/// `index` must be between 0 and length of `buf * 8`.
 #[inline(always)]
 pub unsafe fn set_bit_unchecked(buf: *mut u8, index: usize) {
     unsafe { *buf.add(index / 8) |= 1 << (index % 8) };
 }
 
-/// Unset the bit value at `index` in `buf` without bounds checking
+/// Unset the bit value at `index` in `buf` without bounds checking.
 ///
 /// # Safety
 ///
-/// `index` must be between 0 and length of `buf`
+/// `index` must be between 0 and length of `buf * 8`.
 #[inline(always)]
 pub unsafe fn unset_bit_unchecked(buf: *mut u8, index: usize) {
     unsafe { *buf.add(index / 8) &= !(1 << (index % 8)) };

vortex-compute/src/filter/bitbuffer.rs

@@ -9,55 +9,177 @@ use vortex_buffer::get_bit_unchecked;
 use vortex_buffer::set_bit_unchecked;
 use vortex_buffer::unset_bit_unchecked;
 use vortex_mask::Mask;
+use vortex_mask::MaskValues;
 
 use crate::filter::Filter;
 
 impl Filter<Mask> for &BitBuffer {
     type Output = BitBuffer;
 
     fn filter(self, selection_mask: &Mask) -> BitBuffer {
+        // We delegate checking that the mask length is equal to self to the `MaskValues`
+        // filter implementation below.
+
+        match selection_mask {
+            Mask::AllTrue(_) => self.clone(),
+            Mask::AllFalse(_) => BitBuffer::empty(),
+            Mask::Values(v) => self.filter(v.as_ref()),
+        }
+    }
+}
+
+impl Filter<MaskValues> for &BitBuffer {
+    type Output = BitBuffer;
+
+    fn filter(self, mask_values: &MaskValues) -> BitBuffer {
         assert_eq!(
-            selection_mask.len(),
+            mask_values.len(),
             self.len(),
             "Selection mask length must equal the mask length"
         );
 
-        match selection_mask {
-            Mask::AllTrue(_) => self.clone(),
-            Mask::AllFalse(_) => BitBuffer::empty(),
-            Mask::Values(v) => {
-                filter_indices(self.inner().as_ref(), self.offset(), v.indices()).freeze()
+        self.filter(mask_values.indices())
+    }
+}
+
+impl Filter<[usize]> for &BitBuffer {
+    type Output = BitBuffer;
+
+    /// Filters by indices.
+    ///
+    /// The caller should ensure that the indices are strictly increasing, otherwise the resulting
+    /// buffer might have strange values.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any index is out of bounds. With the additional constraint that the indices are
+    /// strictly increasing, the length of the indices must be less than or equal to the length of
+    /// `self`.
+    fn filter(self, indices: &[usize]) -> BitBuffer {
+        let bools = self.inner().as_ref();
+        let bit_offset = self.offset();
+
+        // FIXME(ngates): this is slower than it could be!
+        BitBufferMut::collect_bool(indices.len(), |idx| {
+            let idx = *unsafe { indices.get_unchecked(idx) };
+            get_bit(bools, bit_offset + idx) // Panics if out of bounds.
+        })
+        .freeze()
+    }
+}
+
+impl Filter<[(usize, usize)]> for &BitBuffer {
+    type Output = BitBuffer;
+
+    /// Filters by ranges of indices.
+    ///
+    /// The caller should ensure that the ranges are strictly increasing, otherwise the resulting
+    /// buffer might have strange values.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any range is out of bounds. With the additional constraint that the ranges are
+    /// strictly increasing, the length of the `slices` array must be less than or equal to the
+    /// length of `self`.
+    fn filter(self, slices: &[(usize, usize)]) -> BitBuffer {
+        let bools = self.inner().as_ref();
+        let bit_offset = self.offset();
+        let output_len: usize = slices.iter().map(|(start, end)| end - start).sum();
+
+        let mut out = BitBufferMut::with_capacity(output_len);
+
+        // FIXME(ngates): this is slower than it could be!
+        for &(start, end) in slices {
+            for idx in start..end {
+                out.append(get_bit(bools, bit_offset + idx)); // Panics if out of bounds.
             }
         }
+
+        out.freeze()
     }
 }
 
 impl Filter<Mask> for &mut BitBufferMut {
     type Output = ();
 
     fn filter(self, selection_mask: &Mask) {
-        assert_eq!(
-            selection_mask.len(),
-            self.len(),
-            "Selection mask length must equal the mask length"
-        );
+        // We delegate checking that the mask length is equal to self to the `MaskValues`
+        // filter implementation below.
 
         match selection_mask {
             Mask::AllTrue(_) => {}
             Mask::AllFalse(_) => self.clear(),
-            Mask::Values(v) => {
-                *self = filter_indices(self.inner().as_slice(), self.offset(), v.indices())
-            }
+            Mask::Values(v) => self.filter(v.as_ref()),
         }
     }
 }
 
-fn filter_indices(bools: &[u8], bit_offset: usize, indices: &[usize]) -> BitBufferMut {
-    // FIXME(ngates): this is slower than it could be!
-    BitBufferMut::collect_bool(indices.len(), |idx| {
-        let idx = *unsafe { indices.get_unchecked(idx) };
-        get_bit(bools, bit_offset + idx)
-    })
+impl Filter<MaskValues> for &mut BitBufferMut {
+    type Output = ();
+
+    fn filter(self, mask_values: &MaskValues) {
+        assert_eq!(
+            mask_values.len(),
+            self.len(),
+            "Selection mask length must equal the mask length"
+        );
+
+        // BitBufferMut filtering always uses indices for simplicity.
+        self.filter(mask_values.indices())
+    }
+}
+
+impl Filter<[usize]> for &mut BitBufferMut {
+    type Output = ();
+
+    /// Filters by indices.
+    ///
+    /// The caller should ensure that the indices are strictly increasing, otherwise the resulting
+    /// buffer might have strange values.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any index is out of bounds. With the additional constraint that the indices are
+    /// strictly increasing, the length of the indices must be less than or equal to the length of
+    /// `self`.
+    fn filter(self, indices: &[usize]) {
+        let bools = self.inner().as_slice();
+        let bit_offset = self.offset();
+
+        // FIXME(ngates): this is slower than it could be!
+        *self = BitBufferMut::collect_bool(indices.len(), |idx| {
+            let idx = *unsafe { indices.get_unchecked(idx) };
+            get_bit(bools, bit_offset + idx) // Panics if out of bounds.
+        });
+    }
+}
+
+impl Filter<[(usize, usize)]> for &mut BitBufferMut {
+    type Output = ();
+
+    /// Filters by ranges of indices.
+    ///
+    /// The caller should ensure that the ranges are strictly increasing, otherwise the resulting
+    /// buffer might have strange values.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any range is out of bounds. With the additional constraint that the ranges are
+    /// strictly increasing, the length of the `slices` array must be less than or equal to the
+    /// length of `self`.
+    fn filter(self, slices: &[(usize, usize)]) {
+        let bools = self.inner().as_slice();
+        let bit_offset = self.offset();
+        let output_len: usize = slices.iter().map(|(start, end)| end - start).sum();
+
+        let mut out = BitBufferMut::with_capacity(output_len);
+        for &(start, end) in slices {
+            for idx in start..end {
+                out.append(get_bit(bools, bit_offset + idx)); // Panics if out of bounds.
+            }
+        }
+        *self = out;
+    }
 }
 
 impl<const NB: usize> Filter<BitView<'_, NB>> for &BitBuffer {
@@ -121,7 +243,7 @@ mod test {
     #[test]
     fn filter_bool_by_index_test() {
         let buf = bitbuffer![1 1 0];
-        let filtered = filter_indices(buf.inner().as_ref(), 0, &[0, 2]).freeze();
+        let filtered = (&buf).filter([0usize, 2].as_slice());
         assert_eq!(2, filtered.len());
         assert_eq!(filtered, bitbuffer![1 0])
     }