implement filter for primitive vector

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

Author: connortsui20

vortex-compute/src/filter/mask.rs

@@ -1,7 +1,7 @@
 // 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;
 
@@ -27,3 +27,19 @@ impl Filter for &Mask {
         }
     }
 }
+
+impl Filter 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;
+    }
+}

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

@@ -2,15 +2,37 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use vortex_mask::Mask;
-use vortex_vector::match_each_pvector_mut;
-use vortex_vector::primitive::PrimitiveVectorMut;
+use vortex_vector::primitive::{PrimitiveVector, PrimitiveVectorMut};
+use vortex_vector::{VectorOps, match_each_pvector, match_each_pvector_mut};
 
 use crate::filter::Filter;
 
+impl Filter for &PrimitiveVector {
+    type Output = PrimitiveVector;
+
+    fn filter(self, selection_mask: &Mask) -> PrimitiveVector {
+        match_each_pvector!(self, |v| { v.filter(selection_mask).into() })
+    }
+}
+
 impl Filter for &mut PrimitiveVectorMut {
     type Output = ();
 
     fn filter(self, selection_mask: &Mask) {
         match_each_pvector_mut!(self, |v| { v.filter(selection_mask) })
     }
 }
+
+impl Filter for PrimitiveVector {
+    type Output = Self;
+
+    fn filter(self, selection_mask: &Mask) -> Self {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the buffer length"
+        );
+
+        match_each_pvector!(self, |v| { v.filter(selection_mask).into() })
+    }
+}

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

@@ -3,11 +3,30 @@
 
 use vortex_dtype::NativePType;
 use vortex_mask::Mask;
-use vortex_vector::VectorMutOps;
-use vortex_vector::primitive::PVectorMut;
+use vortex_vector::primitive::{PVector, PVectorMut};
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 
+impl<T: NativePType> Filter for &PVector<T> {
+    type Output = PVector<T>;
+
+    fn filter(self, selection_mask: &Mask) -> PVector<T> {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the vector length"
+        );
+
+        let filtered_elements = self.elements().filter(selection_mask);
+        let filtered_validity = self.validity().filter(selection_mask);
+
+        // SAFETY: We filtered both components by the same mask, so the length invariants are
+        // upheld.
+        unsafe { PVector::new_unchecked(filtered_elements, filtered_validity) }
+    }
+}
+
 impl<T: NativePType> Filter for &mut PVectorMut<T> {
     type Output = ();
 
@@ -18,13 +37,33 @@ impl<T: NativePType> Filter for &mut PVectorMut<T> {
             "Selection mask length must equal the vector length"
         );
 
-        unimplemented!()
-
         // SAFETY: We filter the two components of the vector at the same time, so the length
         // invariants remain true.
-        // unsafe {
-        //     self.elements_mut().filter(selection_mask);
-        //     self.validity_mut().filter(selection_mask);
-        // }
+        unsafe {
+            self.elements_mut().filter(selection_mask);
+            self.validity_mut().filter(selection_mask);
+        }
+    }
+}
+
+impl<T: NativePType> Filter for PVector<T> {
+    type Output = Self;
+
+    fn filter(self, selection_mask: &Mask) -> Self {
+        assert_eq!(
+            selection_mask.len(),
+            self.len(),
+            "Selection mask length must equal the buffer length"
+        );
+
+        // If we have exclusive access, we can perform the filter in place.
+        match self.try_into_mut() {
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection_mask);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&PVector` impl).
+            Err(vector) => (&vector).filter(selection_mask),
+        }
     }
 }