chore: implement Filter for new MaskIndices type for all vectors

Signed-off-by: Andrew Duffy <[email protected]>

Author: a10y

vortex-compute/src/filter/bitbuffer.rs

@@ -4,13 +4,13 @@
 use vortex_buffer::{BitBuffer, BitBufferMut, get_bit};
 use vortex_mask::{Mask, MaskIter};
 
-use crate::filter::Filter;
+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 for &BitBuffer {
+impl Filter<Mask> for &BitBuffer {
     type Output = BitBuffer;
 
     fn filter(self, selection_mask: &Mask) -> BitBuffer {
@@ -33,6 +33,14 @@ impl Filter for &BitBuffer {
     }
 }
 
+impl Filter<MaskIndices<'_>> for &BitBuffer {
+    type Output = BitBuffer;
+
+    fn filter(self, indices: &MaskIndices) -> BitBuffer {
+        filter_indices(self, indices)
+    }
+}
+
 fn filter_indices(bools: &BitBuffer, indices: &[usize]) -> BitBuffer {
     let buffer = bools.inner().as_ref();
     BitBuffer::collect_bool(indices.len(), |idx| {

vortex-compute/src/filter/buffer.rs

@@ -4,12 +4,12 @@
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_mask::{Mask, MaskIter};
 
-use crate::filter::Filter;
+use crate::filter::{Filter, MaskIndices};
 
 // 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 &Buffer<T> {
+impl<T: Copy> Filter<Mask> for &Buffer<T> {
     type Output = Buffer<T>;
 
     fn filter(self, selection_mask: &Mask) -> Buffer<T> {
@@ -32,7 +32,15 @@ impl<T: Copy> Filter for &Buffer<T> {
     }
 }
 
-impl<T: Copy> Filter for &mut BufferMut<T> {
+impl<T: Copy> Filter<MaskIndices<'_>> for &Buffer<T> {
+    type Output = Buffer<T>;
+
+    fn filter(self, indices: &MaskIndices) -> Buffer<T> {
+        filter_indices(self, indices)
+    }
+}
+
+impl<T: Copy> Filter<Mask> for &mut BufferMut<T> {
     type Output = ();
 
     fn filter(self, selection_mask: &Mask) {
@@ -69,7 +77,19 @@ impl<T: Copy> Filter for &mut BufferMut<T> {
     }
 }
 
-impl<T: Copy> Filter for Buffer<T> {
+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];
+        }
+
+        self.truncate(indices.len());
+    }
+}
+
+impl<T: Copy> Filter<Mask> for Buffer<T> {
     type Output = Self;
 
     fn filter(self, selection_mask: &Mask) -> Self {
@@ -91,6 +111,20 @@ impl<T: Copy> Filter for Buffer<T> {
     }
 }
 
+impl<T: Copy> Filter<MaskIndices<'_>> for Buffer<T> {
+    type Output = Self;
+
+    fn filter(self, indices: &MaskIndices<'_>) -> Self {
+        match self.try_into_mut() {
+            Ok(mut buffer_mut) => {
+                (&mut buffer_mut).filter(indices);
+                buffer_mut.freeze()
+            }
+            Err(buffer) => (&buffer).filter(indices),
+        }
+    }
+}
+
 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/mask.rs

@@ -1,11 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_mask::Mask;
+use vortex_mask::{Mask, MaskMut};
 
-use crate::filter::Filter;
+use crate::filter::{Filter, MaskIndices};
 
-impl Filter for &Mask {
+impl Filter<Mask> for &Mask {
     type Output = Mask;
 
     fn filter(self, selection_mask: &Mask) -> Mask {
@@ -27,3 +27,41 @@ impl Filter for &Mask {
         }
     }
 }
+
+impl Filter<MaskIndices<'_>> for &Mask {
+    type Output = Mask;
+
+    fn filter(self, indices: &MaskIndices<'_>) -> Mask {
+        match self {
+            Mask::AllTrue(_) => Mask::AllTrue(indices.len()),
+            Mask::AllFalse(_) => Mask::AllFalse(indices.len()),
+            Mask::Values(mask_values) => Mask::from(mask_values.bit_buffer().filter(indices)),
+        }
+    }
+}
+
+impl Filter<Mask> for &mut MaskMut {
+    type Output = ();
+
+    fn filter(self, selection_mask: &Mask) {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the mask length"
+        );
+
+        // TODO(connor): There is definitely a better way to do this (in place).
+        let filtered = self.clone().freeze().filter(selection_mask).into_mut();
+        *self = filtered;
+    }
+}
+
+impl Filter<MaskIndices<'_>> for &mut MaskMut {
+    type Output = ();
+
+    fn filter(self, indices: &MaskIndices<'_>) -> Self::Output {
+        // TODO(aduffy): Filter in-place
+        let filtered = self.clone().freeze().filter(indices).into_mut();
+        *self = filtered;
+    }
+}

vortex-compute/src/filter/mod.rs

@@ -3,15 +3,15 @@
 
 //! Filter function.
 
+use std::ops::Deref;
+
 mod bitbuffer;
 mod buffer;
 mod mask;
 mod vector;
 
-use vortex_mask::Mask;
-
 /// Function for filtering based on a selection mask.
-pub trait Filter {
+pub trait Filter<By: ?Sized> {
     /// The result type after performing the operation.
     type Output;
 
@@ -22,5 +22,37 @@ pub trait Filter {
     /// # Panics
     ///
     /// If the length of the mask does not equal the length of the value being filtered.
-    fn filter(self, selection_mask: &Mask) -> Self::Output;
+    fn filter(self, selection: &By) -> Self::Output;
+}
+
+/// A view over a set of strictly sorted indices from a bit mask.
+///
+/// Unlike other indices, `MaskIndices` are always strict-sorted, meaning they are
+/// always unique and monotonic.
+///
+/// You can treat a `MaskIndices` just like a `&[usize]` by iterating or indexing
+/// into it just like you would a slice.
+pub struct MaskIndices<'a>(&'a [usize]);
+
+impl<'a> MaskIndices<'a> {
+    /// Create new indices from a slice of strict-sorted index values.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that the indices are strict-sorted, i.e. that they
+    /// are montonic and unique.
+    ///
+    /// Users of the `Indices` type assume this and failure to uphold this guarantee
+    /// can result in UB downstream.
+    pub unsafe fn new_unchecked(indices: &'a [usize]) -> Self {
+        Self(indices)
+    }
+}
+
+impl Deref for MaskIndices<'_> {
+    type Target = [usize];
+
+    fn deref(&self) -> &Self::Target {
+        self.0
+    }
 }

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

@@ -0,0 +1,39 @@
+use vortex_mask::Mask;
+use vortex_vector::VectorOps;
+use vortex_vector::binaryview::{BinaryViewType, BinaryViewVector, BinaryViewVectorMut};
+
+use crate::filter::{Filter, MaskIndices};
+
+macro_rules! delegate_filter_impl {
+    ($mask_ty:ty) => {
+        impl<T: BinaryViewType> Filter<$mask_ty> for &BinaryViewVector<T> {
+            type Output = BinaryViewVector<T>;
+
+            fn filter(self, selection: &$mask_ty) -> Self::Output {
+                let views = self.views().filter(selection);
+                let validity = self.validity().filter(selection);
+
+                // SAFETY: we filter the views and validity using the same mask
+                unsafe {
+                    BinaryViewVector::<T>::new_unchecked(views, self.buffers().clone(), validity)
+                }
+            }
+        }
+
+        impl<T: BinaryViewType> Filter<$mask_ty> for &mut BinaryViewVectorMut<T> {
+            type Output = ();
+
+            fn filter(self, selection: &$mask_ty) -> Self::Output {
+                // SAFETY: views and validity filtered by the same mask will have
+                //  same resultant length.
+                unsafe {
+                    self.views_mut().filter(selection);
+                    self.validity_mut().filter(selection);
+                }
+            }
+        }
+    };
+}
+
+delegate_filter_impl!(Mask);
+delegate_filter_impl!(MaskIndices<'_>);

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

@@ -3,20 +3,40 @@
 
 use vortex_mask::Mask;
 use vortex_vector::VectorOps;
-use vortex_vector::bool::BoolVector;
+use vortex_vector::bool::{BoolVector, BoolVectorMut};
 
-use crate::filter::Filter;
+use crate::filter::{Filter, MaskIndices};
 
-impl Filter for &BoolVector {
-    type Output = BoolVector;
+macro_rules! delegate_filter_impl {
+    ($mask_ty:ty) => {
+        impl Filter<$mask_ty> for &BoolVector {
+            type Output = BoolVector;
 
-    fn filter(self, mask: &Mask) -> BoolVector {
-        let filtered_bits = self.bits().filter(mask);
-        let filtered_validity = self.validity().filter(mask);
+            fn filter(self, selection: &$mask_ty) -> Self::Output {
+                let filtered_bits = self.bits().filter(selection);
+                let filtered_validity = self.validity().filter(selection);
 
-        // SAFETY: We filter the bits and validity with the same mask, and since they came from an
-        // existing and valid `BoolVector`, we know that the filtered output must have the same
-        // length.
-        unsafe { BoolVector::new_unchecked(filtered_bits, filtered_validity) }
-    }
+                // SAFETY: We filter the bits and validity with the same mask, and since they came from an
+                // existing and valid `BoolVector`, we know that the filtered output must have the same
+                // length.
+                unsafe { BoolVector::new_unchecked(filtered_bits, filtered_validity) }
+            }
+        }
+
+        impl Filter<$mask_ty> for &mut BoolVectorMut {
+            type Output = ();
+
+            fn filter(self, selection: &$mask_ty) -> 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);
+                }
+            }
+        }
+    };
 }
+
+delegate_filter_impl!(Mask);
+delegate_filter_impl!(MaskIndices<'_>);

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

@@ -0,0 +1,37 @@
+use vortex_mask::Mask;
+use vortex_vector::decimal::{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 {
+    type Output = DecimalVector;
+
+    fn filter(self, selection: &MaskIndices) -> 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 {
+    type Output = ();
+
+    fn filter(self, selection: &MaskIndices) -> Self::Output {
+        match_each_dvector_mut!(self, |d| { d.filter(selection) });
+    }
+}

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

@@ -0,0 +1,36 @@
+use vortex_dtype::NativeDecimalType;
+use vortex_mask::Mask;
+use vortex_vector::VectorOps;
+use vortex_vector::decimal::{DVector, DVectorMut};
+
+use crate::filter::{Filter, MaskIndices};
+
+macro_rules! delegate_filter_impl {
+    ($mask_ty:ty) => {
+        impl<D: NativeDecimalType> Filter<$mask_ty> for &DVector<D> {
+            type Output = DVector<D>;
+
+            fn filter(self, selection: &$mask_ty) -> Self::Output {
+                let elements = self.elements().filter(selection);
+                let validity = self.validity().filter(selection);
+                // SAFETY: we're filtering the elements and validity with the same mask
+                unsafe { DVector::<D>::new_unchecked(self.precision_scale(), elements, validity) }
+            }
+        }
+
+        impl<D: NativeDecimalType> Filter<$mask_ty> for &mut DVectorMut<D> {
+            type Output = ();
+
+            fn filter(self, selection: &$mask_ty) -> Self::Output {
+                // SAFETY: we filter elements and validity using the same mask
+                unsafe {
+                    self.elements_mut().filter(selection);
+                    self.validity_mut().filter(selection);
+                }
+            }
+        }
+    };
+}
+
+delegate_filter_impl!(Mask);
+delegate_filter_impl!(MaskIndices<'_>);