perf: vectorize small u8 table take with NEON

vortex-array/src/arrays/fixed_width/take/mod.rs

@@ -6,6 +6,7 @@ mod avx2;
 mod records;
 mod scalar;
 mod slices;
+mod small_table;
 #[cfg(test)]
 mod tests;
 
@@ -61,7 +62,16 @@ impl<V: FixedWidthArray> TakeExecute for V {
 ///
 /// Implementors must have no uninitialized bytes. The shared AVX2 gather reads the complete
 /// representation through a same-width integer lane before writing those bytes back unchanged.
-pub(crate) unsafe trait FixedWidthTakeValue: Copy {}
+pub(crate) unsafe trait FixedWidthTakeValue: Copy {
+    /// Takes values using the kernel appropriate for this value type.
+    fn take<I: UnsignedPType>(
+        values: &[Self],
+        indices: &[I],
+        allocator: &BufferAllocatorRef,
+    ) -> Buffer<Self> {
+        take_values_fallback(values, indices, allocator)
+    }
+}
 
 macro_rules! impl_fixed_width_take_value {
     ($($ty:ty),+ $(,)?) => {
@@ -72,15 +82,20 @@ macro_rules! impl_fixed_width_take_value {
     };
 }
 
-impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,);
-
-// SAFETY: Byte arrays have no padding and every byte is initialized.
-unsafe impl<const N: usize> FixedWidthTakeValue for [u8; N] {}
+impl_fixed_width_take_value!(u16, u32, u64, i16, i32, i64, f16, f32, f64,);
 
 pub(crate) fn take_values<T: FixedWidthTakeValue, I: UnsignedPType>(
     values: &[T],
     indices: &[I],
     allocator: &BufferAllocatorRef,
+) -> Buffer<T> {
+    T::take(values, indices, allocator)
+}
+
+fn take_values_fallback<T: FixedWidthTakeValue, I: UnsignedPType>(
+    values: &[T],
+    indices: &[I],
+    allocator: &BufferAllocatorRef,
 ) -> Buffer<T> {
     #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
     if *HAS_AVX2 {

vortex-array/src/arrays/fixed_width/take/small_table.rs

@@ -0,0 +1,23 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Small byte-table take with a fallback for unsupported targets and inputs.
+
+cfg_if::cfg_if! {
+    if #[cfg(all(target_arch = "aarch64", target_endian = "little"))] {
+        mod neon;
+    } else {
+        mod fallback {
+            use super::super::FixedWidthTakeValue;
+
+            // SAFETY: u8 has no padding or uninitialized bytes.
+            unsafe impl FixedWidthTakeValue for u8 {}
+
+            // SAFETY: i8 has no padding or uninitialized bytes.
+            unsafe impl FixedWidthTakeValue for i8 {}
+
+            // SAFETY: Byte arrays have no padding and every byte is initialized.
+            unsafe impl<const N: usize> FixedWidthTakeValue for [u8; N] {}
+        }
+    }
+}

vortex-array/src/arrays/fixed_width/take/small_table/neon.rs

@@ -0,0 +1,118 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! NEON byte-table take for `u8` codes and at most 16 one-byte values.
+
+use std::arch::aarch64::uint8x16_t;
+use std::arch::aarch64::vdupq_n_u8;
+use std::arch::aarch64::vld1q_u8;
+use std::arch::aarch64::vmaxq_u8;
+use std::arch::aarch64::vmaxvq_u8;
+use std::arch::aarch64::vqtbl1q_u8;
+use std::arch::aarch64::vst1q_u8;
+
+use vortex_buffer::Buffer;
+use vortex_buffer::BufferAllocatorRef;
+use vortex_buffer::BufferMut;
+
+use super::super::FixedWidthTakeValue;
+use super::super::take_values_fallback;
+use crate::dtype::PType;
+use crate::dtype::UnsignedPType;
+
+// SAFETY: u8 has no padding or uninitialized bytes.
+unsafe impl FixedWidthTakeValue for u8 {
+    fn take<I: UnsignedPType>(
+        values: &[Self],
+        indices: &[I],
+        allocator: &BufferAllocatorRef,
+    ) -> Buffer<Self> {
+        take(values, indices, allocator)
+    }
+}
+
+// SAFETY: i8 has no padding or uninitialized bytes.
+unsafe impl FixedWidthTakeValue for i8 {
+    fn take<I: UnsignedPType>(
+        values: &[Self],
+        indices: &[I],
+        allocator: &BufferAllocatorRef,
+    ) -> Buffer<Self> {
+        take(values, indices, allocator)
+    }
+}
+
+// SAFETY: Byte arrays have no padding and every byte is initialized.
+unsafe impl<const N: usize> FixedWidthTakeValue for [u8; N] {
+    fn take<I: UnsignedPType>(
+        values: &[Self],
+        indices: &[I],
+        allocator: &BufferAllocatorRef,
+    ) -> Buffer<Self> {
+        take(values, indices, allocator)
+    }
+}
+
+fn take<T: FixedWidthTakeValue, I: UnsignedPType>(
+    values: &[T],
+    indices: &[I],
+    allocator: &BufferAllocatorRef,
+) -> Buffer<T> {
+    if I::PTYPE != PType::U8
+        || values.is_empty()
+        || values.len() > 16
+        || size_of::<T>() != 1
+        || indices.len() < 64
+    {
+        return take_values_fallback(values, indices, allocator);
+    }
+
+    // SAFETY: the sealed index type is u8, as checked above.
+    let indices: &[u8] =
+        unsafe { std::slice::from_raw_parts(indices.as_ptr().cast(), indices.len()) };
+
+    let mut table = [values[0]; 16];
+    table[..values.len()].copy_from_slice(values);
+
+    let mut output = BufferMut::<T>::with_capacity_in(indices.len(), allocator.clone());
+    let spare = output.spare_capacity_mut();
+    let output_ptr = spare.as_mut_ptr().cast::<u8>();
+    // SAFETY: AArch64 always provides NEON. T is one byte with no uninitialized bytes, the table
+    // contains 16 values, and the output has capacity for every index.
+    let (offset, max_code) = unsafe { take_vectors(&table, indices, output_ptr) };
+
+    for offset in offset..indices.len() {
+        let code = usize::from(indices[offset]);
+        assert!(
+            code < values.len(),
+            "take index {code} out of bounds for length {}",
+            values.len()
+        );
+        spare[offset].write(values[code]);
+    }
+    assert!(
+        usize::from(max_code) < values.len(),
+        "take index {max_code} out of bounds for length {}",
+        values.len()
+    );
+    // SAFETY: the vector loop and scalar remainder initialized every output value.
+    unsafe { output.set_len(indices.len()) };
+    output.freeze()
+}
+
+unsafe fn take_vectors<T: FixedWidthTakeValue>(
+    table: &[T; 16],
+    indices: &[u8],
+    output: *mut u8,
+) -> (usize, u8) {
+    let table = unsafe { vld1q_u8(table.as_ptr().cast::<u8>()) };
+    let mut max_codes: uint8x16_t = unsafe { vdupq_n_u8(0) };
+    let mut offset = 0;
+    while offset + 16 <= indices.len() {
+        let codes = unsafe { vld1q_u8(indices.as_ptr().add(offset)) };
+        max_codes = unsafe { vmaxq_u8(max_codes, codes) };
+        unsafe { vst1q_u8(output.add(offset), vqtbl1q_u8(table, codes)) };
+        offset += 16;
+    }
+    (offset, unsafe { vmaxvq_u8(max_codes) })
+}

vortex-array/src/arrays/fixed_width/take/tests.rs

@@ -43,6 +43,37 @@ fn take_eight_byte_values() {
     assert_eq!(taken.as_slice(), &[20, 30, 10]);
 }
 
+#[rstest]
+#[case(64)]
+#[case(65)]
+#[case(79)]
+#[case(128)]
+fn take_small_byte_table(#[case] len: usize) {
+    let values = [10u8, 20, 30, 40];
+    let indices = (0u8..4).cycle().take(len).collect::<Vec<_>>();
+    let expected = indices
+        .iter()
+        .map(|index| values[usize::from(*index)])
+        .collect::<Vec<_>>();
+    assert_eq!(
+        take_values(&values, &indices, &allocator()).as_slice(),
+        expected.as_slice()
+    );
+}
+
+// The bounds-check message is specific to the NEON table path; other targets reach a different
+// fallback with its own message.
+#[cfg(all(target_arch = "aarch64", target_endian = "little"))]
+#[rstest]
+#[case::vector(64)]
+#[case::tail(128)]
+#[should_panic(expected = "take index")]
+fn take_small_byte_table_rejects_out_of_bounds_index(#[case] offset: usize) {
+    let mut indices = vec![0u8; 129];
+    indices[offset] = 4;
+    drop(take_values(&[10u8, 20, 30, 40], &indices, &allocator()));
+}
+
 #[rstest]
 #[case(1)]
 #[case(2)]