Faster pipelines

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

Author: gatesn

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

@@ -10,7 +10,6 @@ mod toposort;
 use std::hash::{BuildHasher, Hash, Hasher};
 
 use itertools::Itertools;
-use vortex_compute::filter::MaskIndices;
 use vortex_dtype::DType;
 use vortex_error::{vortex_ensure, VortexResult};
 use vortex_mask::Mask;
@@ -310,10 +309,7 @@ impl Pipeline {
                         if selection.true_count() == 0 {
                             tail.clear();
                         } else {
-                            let mut indices = Vec::with_capacity(selection.true_count());
-                            selection.iter_ones(|i| indices.push(i));
-                            let mask_indices = unsafe { MaskIndices::new_unchecked(&indices) };
-                            tail.filter(&mask_indices);
+                            tail.filter(selection);
                             assert_eq!(tail.len(), selection.true_count());
                         }
                     }

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

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::pipeline::bit_view::BitView;
+use crate::pipeline::N;
+use vortex_compute::filter::Filter;
+
+impl<'a, T> Filter<BitView<'a>> for &'a mut [T] {
+    type Output = ();
+
+    fn filter(self, selection: &BitView<'a>) -> Self::Output {
+        // High density: use iter_zeros to compact by removing gaps
+        let mut write_idx = 0;
+        let mut read_idx = 0;
+
+        selection.iter_zeros(|zero_idx| {
+            // Copy elements from read_idx to zero_idx (exclusive) to write_idx
+            let count = zero_idx - read_idx;
+            unsafe {
+                // SAFETY: We assume that the elements are of type E and that the view is valid.
+                // Using memmove for potentially overlapping regions
+                std::ptr::copy(
+                    self.as_ptr().add(read_idx),
+                    self.as_mut_ptr().add(write_idx),
+                    count,
+                );
+                write_idx += count;
+            }
+            read_idx = zero_idx + 1;
+        });
+
+        // Copy any remaining elements after the last zero
+        unsafe {
+            std::ptr::copy(
+                self.as_ptr().add(read_idx),
+                self.as_mut_ptr().add(write_idx),
+                N - read_idx,
+            );
+        }
+    }
+}

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

@@ -0,0 +1,50 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod buffer;
+
+use crate::pipeline::bit_view::BitView;
+use vortex_compute::filter::Filter;
+use vortex_mask::MaskMut;
+use vortex_vector::primitive::{PVectorMut, PrimitiveVectorMut};
+use vortex_vector::{match_each_pvector_mut, VectorMut, VectorMutOps};
+
+impl Filter<BitView<'_>> for &mut VectorMut {
+    type Output = ();
+
+    fn filter(self, selection: &BitView<'_>) -> Self::Output {
+        match self {
+            VectorMut::Null(_) => {}
+            VectorMut::Bool(_) => {}
+            VectorMut::Decimal(_) => {}
+            VectorMut::Primitive(p) => {
+                p.filter(selection);
+            }
+            VectorMut::String(_) => {}
+            VectorMut::Binary(_) => {}
+            VectorMut::List(_) => {}
+            VectorMut::FixedSizeList(_) => {}
+            VectorMut::Struct(_) => {}
+        }
+    }
+}
+
+impl Filter<BitView<'_>> for &mut PrimitiveVectorMut {
+    type Output = ();
+
+    fn filter(self, selection: &BitView<'_>) -> Self::Output {
+        match_each_pvector_mut!(self, |v| { v.filter(selection) })
+    }
+}
+
+impl<T> Filter<BitView<'_>> for &mut PVectorMut<T> {
+    type Output = ();
+
+    fn filter(self, selection: &BitView<'_>) -> Self::Output {
+        println!("SEL(): {:?}", selection);
+        unsafe { self.elements_mut() }.as_mut().filter(selection);
+        // FIXME(ngates): filter the validity...
+        *unsafe { self.validity_mut() } = MaskMut::new_true(selection.true_count());
+        unsafe { self.set_len(selection.true_count()) };
+    }
+}

vortex-array/src/pipeline/mod.rs

@@ -3,6 +3,7 @@
 
 pub mod bit_view;
 pub mod driver;
+mod filter;
 
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_vector::{Vector, VectorMut};