Mooooaarrrrr generics

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/pipeline/filter/buffer.rs

@@ -2,7 +2,6 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 pub mod driver;
-mod filter;
 
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_vector::{Vector, VectorMut};

vortex-array/src/pipeline/filter/mod.rs

@@ -52,8 +52,10 @@ impl<const NB: usize> BitView<'static, NB> {
 }
 
 impl<'a, const NB: usize> BitView<'a, NB> {
-    const N: usize = NB * 8;
-    const N_WORDS: usize = NB * 8 / (usize::BITS as usize);
+    /// The number of bits in the view.
+    pub const N: usize = NB * 8;
+    /// The number of machine words in the view.
+    pub const N_WORDS: usize = NB * 8 / (usize::BITS as usize);
 
     const _ASSERT_MULTIPLE_OF_8: () = assert!(
         NB % 8 == 0,

vortex-array/src/pipeline/mod.rs

@@ -1,7 +1,10 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_buffer::{BitBuffer, BitBufferMut, get_bit};
+use vortex_buffer::{
+    BitBuffer, BitBufferMut, BitView, get_bit, get_bit_unchecked, set_bit_unchecked,
+    unset_bit_unchecked,
+};
 use vortex_mask::Mask;
 
 use crate::filter::{Filter, MaskIndices};
@@ -70,6 +73,56 @@ fn filter_indices(bools: &[u8], bit_offset: usize, indices: &[usize]) -> BitBuff
     })
 }
 
+impl<const NB: usize> Filter<BitView<'_, NB>> for &BitBuffer {
+    type Output = BitBuffer;
+
+    fn filter(self, selection: &BitView<'_, NB>) -> BitBuffer {
+        let bits = self.inner().as_ptr();
+        let mut out = BitBufferMut::with_capacity(selection.true_count());
+        let mut out_idx = 0;
+        selection.iter_ones(|idx| {
+            let value = unsafe { get_bit_unchecked(bits, self.offset() + idx) };
+            unsafe { out.set_to_unchecked(out_idx, value) };
+            out_idx += 1;
+        });
+        out.freeze()
+    }
+}
+
+impl<const NB: usize> Filter<BitView<'_, NB>> for &mut BitBufferMut {
+    type Output = ();
+
+    fn filter(self, selection: &BitView<'_, NB>) {
+        assert_eq!(
+            self.len(),
+            BitView::<NB>::N,
+            "Selection mask length must equal the mask length"
+        );
+
+        let this = std::mem::take(self);
+
+        let offset = this.offset();
+        let mut buffer = this.into_inner();
+
+        let buffer_ptr = buffer.as_mut_ptr();
+        let mut out_idx = 0;
+        selection.iter_ones(|idx| {
+            let value = unsafe { get_bit_unchecked(buffer_ptr, offset + idx) };
+
+            // NOTE(ngates): we don't call out.set_bit_unchecked here because it's nice that we
+            //  can shift away any non-zero offset by writing directly into the bits buffer.
+            if value {
+                unsafe { set_bit_unchecked(buffer_ptr, out_idx) };
+            } else {
+                unsafe { unset_bit_unchecked(buffer_ptr, out_idx) };
+            }
+            out_idx += 1;
+        });
+
+        *self = BitBufferMut::from_buffer(buffer, 0, selection.true_count());
+    }
+}
+
 #[cfg(test)]
 mod test {
     use vortex_buffer::bitbuffer;

vortex-buffer/src/bit/view.rs

@@ -1,11 +1,31 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_buffer::{Buffer, BufferMut};
+use vortex_buffer::{BitView, Buffer, BufferMut};
 use vortex_mask::{Mask, MaskIter};
 
 use crate::filter::{Filter, MaskIndices};
 
+impl<M, T: Copy> Filter<M> for Buffer<T>
+where
+    for<'a> &'a Buffer<T>: Filter<M, Output = Buffer<T>>,
+    for<'a> &'a mut BufferMut<T>: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection_mask: &M) -> Self {
+        // 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),
+        }
+    }
+}
+
 // This is modeled after the constant with the equivalent name in arrow-rs.
 const FILTER_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.8;
 
@@ -40,72 +60,29 @@ impl<T: Copy> Filter<MaskIndices<'_>> for &Buffer<T> {
     }
 }
 
-impl<T: Copy> Filter<Mask> 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) };
-            }
-        }
-    }
-}
-
-impl<T: Copy> Filter<MaskIndices<'_>> for &mut BufferMut<T> {
-    type Output = ();
-
-    fn filter(self, indices: &MaskIndices) -> Self::Output {
-        for (write_index, &read_index) in indices.iter().enumerate() {
-            self[write_index] = self[read_index];
-        }
+impl<const NB: usize, T: Copy> Filter<BitView<'_, NB>> for &Buffer<T> {
+    type Output = Buffer<T>;
 
-        self.truncate(indices.len());
+    fn filter(self, selection: &BitView<'_, NB>) -> Self::Output {
+        // TODO(ngates): this is very very slow!
+        let elems = self.as_slice();
+        let mut out = BufferMut::<T>::with_capacity(selection.true_count());
+        selection.iter_ones(|idx| {
+            unsafe { out.push_unchecked(elems[idx]) };
+        });
+        out.freeze()
     }
 }
 
-impl<M, T: Copy> Filter<M> for Buffer<T>
+impl<M, T> Filter<M> for &mut BufferMut<T>
 where
-    for<'a> &'a Buffer<T>: Filter<M, Output = Buffer<T>>,
-    for<'a> &'a mut BufferMut<T>: Filter<M, Output = ()>,
+    for<'a> &'a mut [T]: Filter<M, Output = &'a mut [T]>,
 {
-    type Output = Self;
+    type Output = ();
 
-    fn filter(self, selection_mask: &M) -> Self {
-        // 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(self, selection_mask: &M) -> Self::Output {
+        let true_count = self.as_mut_slice().filter(selection_mask).len();
+        self.truncate(true_count);
     }
 }
 
@@ -121,42 +98,9 @@ fn filter_slices<T>(values: &[T], output_len: usize, slices: &[(usize, usize)])
     out.freeze()
 }
 
-/// 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 {
-            core::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::buffer;
+    use vortex_buffer::{BufferMut, buffer, buffer_mut};
     use vortex_mask::Mask;
 
     use super::*;
@@ -202,8 +146,6 @@ mod tests {
         assert_eq!(result, buffer![1u32, 2, 5]);
     }
 
-    use vortex_buffer::{BufferMut, buffer_mut};
-
     #[test]
     fn test_filter_all_true() {
         let mut buf = buffer_mut![1u32, 2, 3, 4, 5];