Feature: add impl owned filter impls (#5493)

Since the operator API passes around owned vectors, we probably want to
use these filter impls.

It is a good amount of boilerplate, but honestly I think it is fine
since we are not likely not adding many more vector types.

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

Author: connortsui20

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

@@ -3,10 +3,10 @@
 
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::binaryview::{
     BinaryView, BinaryViewType, BinaryViewVector, BinaryViewVectorMut,
 };
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 
@@ -42,3 +42,24 @@ where
         }
     }
 }
+
+impl<M, T: BinaryViewType> Filter<M> for BinaryViewVector<T>
+where
+    for<'a> &'a BinaryViewVector<T>: Filter<M, Output = BinaryViewVector<T>>,
+    for<'a> &'a mut BinaryViewVectorMut<T>: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection: &M) -> Self {
+        match self.try_into_mut() {
+            // If we have exclusive access, we can perform the filter in place.
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&BinaryViewVector`
+            // impl).
+            Err(vector) => (&vector).filter(selection),
+        }
+    }
+}

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

@@ -3,8 +3,8 @@
 
 use vortex_buffer::{BitBuffer, BitBufferMut};
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::bool::{BoolVector, BoolVectorMut};
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 
@@ -38,3 +38,23 @@ where
         unsafe { self.validity_mut().filter(selection) };
     }
 }
+
+impl<M> Filter<M> for BoolVector
+where
+    for<'a> &'a BoolVector: Filter<M, Output = BoolVector>,
+    for<'a> &'a mut BoolVectorMut: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection: &M) -> Self {
+        match self.try_into_mut() {
+            // If we have exclusive access, we can perform the filter in place.
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&BoolVector` impl).
+            Err(vector) => (&vector).filter(selection),
+        }
+    }
+}

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

@@ -3,7 +3,7 @@
 
 use vortex_dtype::i256;
 use vortex_vector::decimal::{DVector, DVectorMut, DecimalVector, DecimalVectorMut};
-use vortex_vector::{match_each_dvector, match_each_dvector_mut};
+use vortex_vector::{VectorMutOps, VectorOps, match_each_dvector, match_each_dvector_mut};
 
 use crate::filter::Filter;
 
@@ -38,3 +38,24 @@ where
         match_each_dvector_mut!(self, |d| { d.filter(selection) });
     }
 }
+
+impl<M> Filter<M> for DecimalVector
+where
+    for<'a> &'a DecimalVector: Filter<M, Output = DecimalVector>,
+    for<'a> &'a mut DecimalVectorMut: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection: &M) -> Self {
+        match self.try_into_mut() {
+            // If we have exclusive access, we can perform the filter in place.
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&DecimalVector`
+            // impl).
+            Err(vector) => (&vector).filter(selection),
+        }
+    }
+}

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

@@ -4,8 +4,8 @@
 use vortex_buffer::{Buffer, BufferMut};
 use vortex_dtype::NativeDecimalType;
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::decimal::{DVector, DVectorMut};
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 
@@ -39,3 +39,23 @@ where
         }
     }
 }
+
+impl<M, D: NativeDecimalType> Filter<M> for DVector<D>
+where
+    for<'a> &'a DVector<D>: Filter<M, Output = DVector<D>>,
+    for<'a> &'a mut DVectorMut<D>: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection: &M) -> Self {
+        match self.try_into_mut() {
+            // If we have exclusive access, we can perform the filter in place.
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&DVector` impl).
+            Err(vector) => (&vector).filter(selection),
+        }
+    }
+}

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

@@ -4,9 +4,9 @@
 use std::sync::Arc;
 
 use vortex_mask::{Mask, MaskMut};
-use vortex_vector::VectorOps;
 use vortex_vector::listview::{ListViewVector, ListViewVectorMut};
 use vortex_vector::primitive::{PrimitiveVector, PrimitiveVectorMut};
+use vortex_vector::{VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 
@@ -45,3 +45,24 @@ where
         }
     }
 }
+
+impl<M> Filter<M> for ListViewVector
+where
+    for<'a> &'a ListViewVector: Filter<M, Output = ListViewVector>,
+    for<'a> &'a mut ListViewVectorMut: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection: &M) -> Self {
+        match self.try_into_mut() {
+            // If we have exclusive access, we can perform the filter in place.
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&ListViewVector`
+            // impl).
+            Err(vector) => (&vector).filter(selection),
+        }
+    }
+}

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

@@ -1,6 +1,36 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Filter implementations for vector types.
+//!
+//! This module provides [`Filter`] trait implementations for all vector types with three distinct
+//! patterns optimized for different ownership scenarios:
+//!
+//! ### 1. Reference Implementation (`&Vector`)
+//!
+//! Filters by allocating new memory and copying selected elements. Returns a new owned vector.
+//! This is the base implementation that all other patterns can fall back to.
+//!
+//! ### 2. Mutable Reference Implementation (`&mut VectorMut`)
+//!
+//! Filters in-place when exclusive mutable access is available, avoiding allocation. Returns `()`
+//! as the vector is modified directly. This is the most efficient when you already have a mutable
+//! vector (it is only less efficient if the vector is very small and the output vector is already
+//! in the L1 cache).
+//!
+//! ### 3. Owned Implementation (`Vector`)
+//!
+//! Uses [`VectorOps::try_into_mut`] to check for exclusive ownership. If successful, performs
+//! in-place filtering via the mutable implementation and calls [`VectorMutOps::freeze`] to convert
+//! back. Otherwise, delegates to the reference implementation and makes an allocation.
+//!
+//! ## Breaking Recursive Trait Bounds
+//!
+//! To allow all vector types to implement filter generically over a "mask" type `M`, we must break
+//! the recursive trait bounds (e.g. from [`StructVector`] requiring `Vector: Filter<M>` for its
+//! fields) by manually implementing [`Filter`] for [`Vector`] and [`VectorMut`] for each concrete
+//! mask type in this file.
+
 use vortex_buffer::BitView;
 use vortex_mask::Mask;
 use vortex_vector::{Vector, VectorMut, match_each_vector, match_each_vector_mut};
@@ -18,9 +48,8 @@ mod primitive;
 mod pvector;
 mod struct_;
 
-// To allow all vector types to implement filter generically over `M`, we must break the recursive
-// trait bounds (e.g. from StructVector requiring Vector: Filter<M> for its fields) by manually
-// implementing Filter for Vector and VectorMut for each concrete mask type here.
+// We manually implement Filter for Vector and VectorMut for each concrete mask type here to break
+// the recursive trait bounds.
 
 impl Filter<Mask> for &Vector {
     type Output = Vector;

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

@@ -39,6 +39,22 @@ impl<const NB: usize> Filter<BitView<'_, NB>> for &mut NullVectorMut {
     }
 }
 
+impl Filter<Mask> for NullVector {
+    type Output = Self;
+
+    fn filter(self, selection: &Mask) -> Self {
+        (&self).filter(selection)
+    }
+}
+
+impl<const NB: usize> Filter<BitView<'_, NB>> for NullVector {
+    type Output = Self;
+
+    fn filter(self, selection: &BitView<'_, NB>) -> Self {
+        (&self).filter(selection)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use vortex_mask::Mask;

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

@@ -3,7 +3,7 @@
 
 use vortex_dtype::half::f16;
 use vortex_vector::primitive::{PVector, PVectorMut, PrimitiveVector, PrimitiveVectorMut};
-use vortex_vector::{match_each_pvector, match_each_pvector_mut};
+use vortex_vector::{VectorMutOps, VectorOps, match_each_pvector, match_each_pvector_mut};
 
 use crate::filter::Filter;
 
@@ -48,3 +48,24 @@ where
         match_each_pvector_mut!(self, |v| { v.filter(selection) })
     }
 }
+
+impl<M> Filter<M> for PrimitiveVector
+where
+    for<'a> &'a PrimitiveVector: Filter<M, Output = PrimitiveVector>,
+    for<'a> &'a mut PrimitiveVectorMut: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection: &M) -> Self {
+        match self.try_into_mut() {
+            // If we have exclusive access, we can perform the filter in place.
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&PrimitiveVector`
+            // impl).
+            Err(vector) => (&vector).filter(selection),
+        }
+    }
+}

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

@@ -51,8 +51,8 @@ where
     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() {
+            // If we have exclusive access, we can perform the filter in place.
             Ok(mut vector_mut) => {
                 (&mut vector_mut).filter(selection_mask);
                 vector_mut.freeze()

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

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use vortex_mask::{Mask, MaskMut};
 use vortex_vector::struct_::{StructVector, StructVectorMut};
-use vortex_vector::{Vector, VectorMut, VectorOps};
+use vortex_vector::{Vector, VectorMut, VectorMutOps, VectorOps};
 
 use crate::filter::Filter;
 
@@ -49,3 +49,23 @@ where
         }
     }
 }
+
+impl<M> Filter<M> for StructVector
+where
+    for<'a> &'a StructVector: Filter<M, Output = StructVector>,
+    for<'a> &'a mut StructVectorMut: Filter<M, Output = ()>,
+{
+    type Output = Self;
+
+    fn filter(self, selection: &M) -> Self {
+        match self.try_into_mut() {
+            // If we have exclusive access, we can perform the filter in place.
+            Ok(mut vector_mut) => {
+                (&mut vector_mut).filter(selection);
+                vector_mut.freeze()
+            }
+            // Otherwise, allocate a new buffer and fill it in (delegate to the `&StructVector` impl).
+            Err(vector) => (&vector).filter(selection),
+        }
+    }
+}