Add Mojo SIMD filter kernel for primitive arrays

Extends the Mojo AOT kernel with filter-by-indices support. The
primitive filter path converts sparse masks (<80% selectivity) into
an index array, then gathers values at those positions — identical
to the take operation but with usize indices.

Four new exported symbols (vortex_filter_{1,2,4,8}byte) are added
to the Mojo kernel and wired into filter_slice_by_indices behind
cfg(vortex_mojo). Falls back to scalar when Mojo is unavailable.

All 121 existing filter tests pass with the Mojo kernel active.

Signed-off-by: Claude <noreply@anthropic.com>

https://claude.ai/code/session_01EVcJZP4ZmfvWRRg2CsgvST

Author: claude

vortex-array/kernels/take.mojo

@@ -143,3 +143,29 @@ fn take_1byte_u32idx(src: Int, idx: Int, dst: Int, n: Int):
 @export("vortex_take_1byte_u64idx")
 fn take_1byte_u64idx(src: Int, idx: Int, dst: Int, n: Int):
     _take[DType.uint8, DType.uint64, W1](src, idx, dst, n)
+
+
+# ---------------------------------------------------------------------------
+# Filter kernels (gather by usize indices from mask)
+#
+# These are used by the primitive filter path when the mask is sparse (<80%
+# selectivity). The Rust side converts the bitmap to a &[usize] index array
+# and passes it here. On x86_64 usize = u64, so these are gathers with
+# u64 element indices.
+# ---------------------------------------------------------------------------
+
+@export("vortex_filter_1byte")
+fn filter_1byte(src: Int, idx: Int, dst: Int, n: Int):
+    _take[DType.uint8, DType.uint64, W1](src, idx, dst, n)
+
+@export("vortex_filter_2byte")
+fn filter_2byte(src: Int, idx: Int, dst: Int, n: Int):
+    _take[DType.uint16, DType.uint64, W2](src, idx, dst, n)
+
+@export("vortex_filter_4byte")
+fn filter_4byte(src: Int, idx: Int, dst: Int, n: Int):
+    _take[DType.uint32, DType.uint64, W4](src, idx, dst, n)
+
+@export("vortex_filter_8byte")
+fn filter_8byte(src: Int, idx: Int, dst: Int, n: Int):
+    _take[DType.uint64, DType.uint64, W8](src, idx, dst, n)

vortex-array/src/arrays/filter/execute/slice.rs

@@ -6,6 +6,7 @@
 //! Provides both immutable and mutable (in-place) filtering of typed slices by various mask
 //! representations: indices and ranges (slices).
 
+use std::mem::size_of;
 use std::ptr;
 
 use vortex_buffer::Buffer;
@@ -37,9 +38,64 @@ pub(super) fn filter_slice_by_mask_values<T: Copy>(slice: &[T], mask: &MaskValue
 
 /// Filter a slice by a set of strictly increasing indices.
 fn filter_slice_by_indices<T: Copy>(slice: &[T], indices: &[usize]) -> Buffer<T> {
+    #[cfg(vortex_mojo)]
+    {
+        if let Some(buf) = mojo::filter_by_indices_mojo(slice, indices) {
+            return buf;
+        }
+    }
+
     Buffer::<T>::from_trusted_len_iter(indices.iter().map(|&idx| slice[idx]))
 }
 
+#[cfg(vortex_mojo)]
+mod mojo {
+    use vortex_buffer::Buffer;
+    use vortex_buffer::BufferMut;
+
+    use super::size_of;
+
+    unsafe extern "C" {
+        fn vortex_filter_1byte(src: usize, idx: usize, dst: usize, n: usize);
+        fn vortex_filter_2byte(src: usize, idx: usize, dst: usize, n: usize);
+        fn vortex_filter_4byte(src: usize, idx: usize, dst: usize, n: usize);
+        fn vortex_filter_8byte(src: usize, idx: usize, dst: usize, n: usize);
+    }
+
+    /// SIMD gather for the filter-by-indices path. Returns `None` for unsupported
+    /// element sizes so the caller falls back to scalar.
+    pub(super) fn filter_by_indices_mojo<T: Copy>(
+        slice: &[T],
+        indices: &[usize],
+    ) -> Option<Buffer<T>> {
+        let kernel: unsafe extern "C" fn(usize, usize, usize, usize) = match size_of::<T>() {
+            1 => vortex_filter_1byte,
+            2 => vortex_filter_2byte,
+            4 => vortex_filter_4byte,
+            8 => vortex_filter_8byte,
+            _ => return None,
+        };
+
+        let len = indices.len();
+        let mut buffer = BufferMut::<T>::with_capacity(len);
+        let dst = buffer.spare_capacity_mut().as_mut_ptr().cast::<T>();
+
+        // SAFETY: The Mojo kernel reads `len` indices from `indices`, gathers from
+        // `slice`, and writes `len` elements to `dst`. All pointers are valid.
+        unsafe {
+            kernel(
+                slice.as_ptr() as usize,
+                indices.as_ptr() as usize,
+                dst as usize,
+                len,
+            );
+            buffer.set_len(len);
+        }
+
+        Some(buffer.freeze())
+    }
+}
+
 /// Filter a slice by a set of strictly increasing `(start, end)` ranges.
 fn filter_slice_by_slices<T: Copy>(slice: &[T], slices: &[(usize, usize)]) -> Buffer<T> {
     let output_len: usize = slices.iter().map(|(start, end)| end - start).sum();