Mooooaarrrrr generics (#5367)

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

Author: gatesn

vortex-buffer/src/bit/buf.rs

@@ -8,7 +8,7 @@ use crate::bit::{
     BitChunks, BitIndexIterator, BitIterator, BitSliceIterator, UnalignedBitChunk,
     get_bit_unchecked,
 };
-use crate::{Alignment, BitBufferMut, Buffer, BufferMut, ByteBuffer, buffer};
+use crate::{Alignment, BitBufferMut, Buffer, ByteBuffer, buffer};
 
 /// An immutable bitset stored as a packed byte buffer.
 #[derive(Debug, Clone, Eq)]
@@ -114,37 +114,9 @@ impl BitBuffer {
         }
     }
 
-    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBuffer`
-    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, mut f: F) -> Self {
-        let mut buffer = BufferMut::with_capacity(len.div_ceil(64) * 8);
-
-        let chunks = len / 64;
-        let remainder = len % 64;
-        for chunk in 0..chunks {
-            let mut packed = 0;
-            for bit_idx in 0..64 {
-                let i = bit_idx + chunk * 64;
-                packed |= (f(i) as u64) << bit_idx;
-            }
-
-            // SAFETY: Already allocated sufficient capacity
-            unsafe { buffer.push_unchecked(packed) }
-        }
-
-        if remainder != 0 {
-            let mut packed = 0;
-            for bit_idx in 0..remainder {
-                let i = bit_idx + chunks * 64;
-                packed |= (f(i) as u64) << bit_idx;
-            }
-
-            // SAFETY: Already allocated sufficient capacity
-            unsafe { buffer.push_unchecked(packed) }
-        }
-
-        buffer.truncate(len.div_ceil(8));
-
-        Self::new(buffer.freeze().into_byte_buffer(), len)
+    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new [`BitBuffer`].
+    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
+        BitBufferMut::collect_bool(len, f).freeze()
     }
 
     /// Get the logical length of this `BoolBuffer`.

vortex-buffer/src/bit/buf_mut.rs

@@ -117,6 +117,43 @@ impl BitBufferMut {
         }
     }
 
+    /// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BitBufferMut`
+    pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, mut f: F) -> Self {
+        let mut buffer = BufferMut::with_capacity(len.div_ceil(64) * 8);
+
+        let chunks = len / 64;
+        let remainder = len % 64;
+        for chunk in 0..chunks {
+            let mut packed = 0;
+            for bit_idx in 0..64 {
+                let i = bit_idx + chunk * 64;
+                packed |= (f(i) as u64) << bit_idx;
+            }
+
+            // SAFETY: Already allocated sufficient capacity
+            unsafe { buffer.push_unchecked(packed) }
+        }
+
+        if remainder != 0 {
+            let mut packed = 0;
+            for bit_idx in 0..remainder {
+                let i = bit_idx + chunks * 64;
+                packed |= (f(i) as u64) << bit_idx;
+            }
+
+            // SAFETY: Already allocated sufficient capacity
+            unsafe { buffer.push_unchecked(packed) }
+        }
+
+        buffer.truncate(len.div_ceil(8));
+
+        Self {
+            buffer: buffer.into_byte_buffer(),
+            offset: 0,
+            len,
+        }
+    }
+
     /// Return the underlying byte buffer.
     pub fn inner(&self) -> &ByteBufferMut {
         &self.buffer

vortex-buffer/src/buffer_mut.rs

@@ -110,6 +110,16 @@ impl<T> BufferMut<T> {
         buffer
     }
 
+    /// Return the [`ByteBufferMut`] for this [`BufferMut<T>`].
+    pub fn into_byte_buffer(self) -> ByteBufferMut {
+        ByteBufferMut {
+            bytes: self.bytes,
+            length: self.length * size_of::<T>(),
+            alignment: self.alignment,
+            _marker: Default::default(),
+        }
+    }
+
     /// Get the alignment of the buffer.
     #[inline(always)]
     pub fn alignment(&self) -> Alignment {

vortex-compute/src/filter/bitbuffer.rs

@@ -2,14 +2,10 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_buffer::{BitBuffer, BitBufferMut, get_bit};
-use vortex_mask::{Mask, MaskIter};
+use vortex_mask::Mask;
 
 use crate::filter::{Filter, MaskIndices};
 
-/// If the filter density is above 80%, we use slices to filter the array instead of indices.
-// TODO(ngates): we need more experimentation to determine the best threshold here.
-const FILTER_SLICES_DENSITY_THRESHOLD: f64 = 0.8;
-
 impl Filter<Mask> for &BitBuffer {
     type Output = BitBuffer;
 
@@ -23,12 +19,29 @@ impl Filter<Mask> for &BitBuffer {
         match selection_mask {
             Mask::AllTrue(_) => self.clone(),
             Mask::AllFalse(_) => BitBuffer::empty(),
-            Mask::Values(v) => match v.threshold_iter(FILTER_SLICES_DENSITY_THRESHOLD) {
-                MaskIter::Indices(indices) => filter_indices(self, indices),
-                MaskIter::Slices(slices) => {
-                    filter_slices(self, selection_mask.true_count(), slices)
-                }
-            },
+            Mask::Values(v) => {
+                filter_indices(self.inner().as_ref(), self.offset(), v.indices()).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"
+        );
+
+        match selection_mask {
+            Mask::AllTrue(_) => {}
+            Mask::AllFalse(_) => self.clear(),
+            Mask::Values(v) => {
+                *self = filter_indices(self.inner().as_slice(), self.offset(), v.indices())
+            }
         }
     }
 }
@@ -37,25 +50,24 @@ impl Filter<MaskIndices<'_>> for &BitBuffer {
     type Output = BitBuffer;
 
     fn filter(self, indices: &MaskIndices) -> BitBuffer {
-        filter_indices(self, indices)
+        filter_indices(self.inner().as_ref(), self.offset(), indices).freeze()
     }
 }
 
-fn filter_indices(bools: &BitBuffer, indices: &[usize]) -> BitBuffer {
-    let buffer = bools.inner().as_ref();
-    BitBuffer::collect_bool(indices.len(), |idx| {
-        let idx = *unsafe { indices.get_unchecked(idx) };
-        get_bit(buffer, bools.offset() + idx)
-    })
-}
+impl Filter<MaskIndices<'_>> for &mut BitBufferMut {
+    type Output = ();
 
-fn filter_slices(buffer: &BitBuffer, output_len: usize, slices: &[(usize, usize)]) -> BitBuffer {
-    let mut builder = BitBufferMut::with_capacity(output_len);
-    for (start, end) in slices {
-        // TODO(ngates): we probably want a borrowed slice for things like this.
-        builder.append_buffer(&buffer.slice(*start..*end));
+    fn filter(self, indices: &MaskIndices) {
+        *self = filter_indices(self.inner().as_ref(), self.offset(), indices)
     }
-    builder.freeze()
+}
+
+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)
+    })
 }
 
 #[cfg(test)]
@@ -64,20 +76,10 @@ mod test {
 
     use super::*;
 
-    #[test]
-    fn filter_bool_by_slice_test() {
-        let bits = bitbuffer![1 1 0];
-
-        let filtered = filter_slices(&bits, 2, &[(0, 1), (2, 3)]);
-        assert_eq!(2, filtered.len());
-
-        assert_eq!(filtered, bitbuffer![1 0])
-    }
-
     #[test]
     fn filter_bool_by_index_test() {
         let buf = bitbuffer![1 1 0];
-        let filtered = filter_indices(&buf, &[0, 2]);
+        let filtered = filter_indices(buf.inner().as_ref(), 0, &[0, 2]).freeze();
         assert_eq!(2, filtered.len());
         assert_eq!(filtered, bitbuffer![1 0])
     }

vortex-compute/src/filter/slice.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_buffer::BitBuffer;
+use vortex_buffer::{BitBuffer, BitBufferMut};
 use vortex_mask::{Mask, MaskMut};
 use vortex_vector::VectorOps;
 use vortex_vector::bool::{BoolVector, BoolVectorMut};
@@ -28,18 +28,15 @@ where
 
 impl<M> Filter<M> for &mut BoolVectorMut
 where
-    for<'a> &'a BitBuffer: Filter<M, Output = BitBuffer>,
+    for<'a> &'a mut BitBufferMut: Filter<M, Output = ()>,
     for<'a> &'a mut MaskMut: Filter<M, Output = ()>,
 {
     type Output = ();
 
     fn filter(self, selection: &M) -> Self::Output {
         // TODO(aduffy): how can we do this faster in-place?
-        unsafe {
-            let bits = self.bits_mut();
-            *bits = (*bits).clone().freeze().filter(selection).into_mut();
-            self.validity_mut().filter(selection);
-        }
+        unsafe { self.bits_mut().filter(selection) };
+        unsafe { self.validity_mut().filter(selection) };
     }
 }
 

vortex-compute/src/filter/vector/bool.rs

@@ -1,40 +1,40 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_mask::Mask;
-use vortex_vector::decimal::{DecimalVector, DecimalVectorMut};
+use vortex_dtype::i256;
+use vortex_vector::decimal::{DVector, DVectorMut, DecimalVector, DecimalVectorMut};
 use vortex_vector::{match_each_dvector, match_each_dvector_mut};
 
-use crate::filter::{Filter, MaskIndices};
-
-impl Filter<Mask> for &DecimalVector {
-    type Output = DecimalVector;
-
-    fn filter(self, selection: &Mask) -> Self::Output {
-        match_each_dvector!(self, |d| { d.filter(selection).into() })
-    }
-}
-
-impl Filter<MaskIndices<'_>> for &DecimalVector {
+use crate::filter::Filter;
+
+impl<M> Filter<M> for &DecimalVector
+where
+    for<'a> &'a DVector<i8>: Filter<M, Output = DVector<i8>>,
+    for<'a> &'a DVector<i16>: Filter<M, Output = DVector<i16>>,
+    for<'a> &'a DVector<i32>: Filter<M, Output = DVector<i32>>,
+    for<'a> &'a DVector<i64>: Filter<M, Output = DVector<i64>>,
+    for<'a> &'a DVector<i128>: Filter<M, Output = DVector<i128>>,
+    for<'a> &'a DVector<i256>: Filter<M, Output = DVector<i256>>,
+{
     type Output = DecimalVector;
 
-    fn filter(self, selection: &MaskIndices) -> Self::Output {
+    fn filter(self, selection: &M) -> Self::Output {
         match_each_dvector!(self, |d| { d.filter(selection).into() })
     }
 }
 
-impl Filter<Mask> for &mut DecimalVectorMut {
-    type Output = ();
-
-    fn filter(self, selection: &Mask) -> Self::Output {
-        match_each_dvector_mut!(self, |d| { d.filter(selection) });
-    }
-}
-
-impl Filter<MaskIndices<'_>> for &mut DecimalVectorMut {
+impl<M> Filter<M> for &mut DecimalVectorMut
+where
+    for<'a> &'a mut DVectorMut<i8>: Filter<M, Output = ()>,
+    for<'a> &'a mut DVectorMut<i16>: Filter<M, Output = ()>,
+    for<'a> &'a mut DVectorMut<i32>: Filter<M, Output = ()>,
+    for<'a> &'a mut DVectorMut<i64>: Filter<M, Output = ()>,
+    for<'a> &'a mut DVectorMut<i128>: Filter<M, Output = ()>,
+    for<'a> &'a mut DVectorMut<i256>: Filter<M, Output = ()>,
+{
     type Output = ();
 
-    fn filter(self, selection: &MaskIndices) -> Self::Output {
+    fn filter(self, selection: &M) -> Self::Output {
         match_each_dvector_mut!(self, |d| { d.filter(selection) });
     }
 }
@@ -48,6 +48,7 @@ mod tests {
     use vortex_vector::{VectorMutOps, VectorOps};
 
     use super::*;
+    use crate::filter::MaskIndices;
 
     #[test]
     fn test_filter_decimal_vector_with_mask() {

vortex-compute/src/filter/vector/decimal.rs

@@ -1,47 +1,26 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_mask::Mask;
 use vortex_vector::fixed_size_list::{FixedSizeListVector, FixedSizeListVectorMut};
 
-use crate::filter::{Filter, MaskIndices};
+use crate::filter::Filter;
 
 // TODO(aduffy): there really isn't a cheap way to implement these is there.
 
-impl Filter<Mask> for &FixedSizeListVector {
+impl<M> Filter<M> for &FixedSizeListVector {
     type Output = FixedSizeListVector;
 
-    fn filter(self, _selection: &Mask) -> Self::Output {
+    fn filter(self, _selection: &M) -> Self::Output {
         // We need to spread the mask out to point to offsets from
         // the inner vector type
         todo!()
     }
 }
 
-impl Filter<MaskIndices<'_>> for &FixedSizeListVector {
-    type Output = FixedSizeListVector;
-
-    fn filter(self, _selection: &MaskIndices) -> Self::Output {
-        // We need to spread the mask out to point to offsets from
-        // the inner vector type
-        todo!()
-    }
-}
-
-impl Filter<Mask> for &mut FixedSizeListVectorMut {
-    type Output = ();
-
-    fn filter(self, _selection: &Mask) -> Self::Output {
-        // We need to spread the mask out to point to offsets from
-        // the inner vector type
-        todo!()
-    }
-}
-
-impl Filter<MaskIndices<'_>> for &mut FixedSizeListVectorMut {
+impl<M> Filter<M> for &mut FixedSizeListVectorMut {
     type Output = ();
 
-    fn filter(self, _selection: &MaskIndices) -> Self::Output {
+    fn filter(self, _selection: &M) -> Self::Output {
         // We need to spread the mask out to point to offsets from
         // the inner vector type
         todo!()