From e97db31410f1636f0b68401ee5d0204dde2101e2 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 24 Aug 2026 13:52:49 +0100 Subject: [PATCH 1/4] perf: vectorize small u8 table take with NEON Use NEON TBL for u8-coded tables with at most 16 one-byte values. Signed-off-by: Joseph Isaacs --- .../src/arrays/fixed_width/take/mod.rs | 14 ++++ .../arrays/fixed_width/take/small_table.rs | 74 +++++++++++++++++++ .../src/arrays/fixed_width/take/tests.rs | 22 ++++++ 3 files changed, 110 insertions(+) create mode 100644 vortex-array/src/arrays/fixed_width/take/small_table.rs diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs index d6e2f744bdd..78684e5b385 100644 --- a/vortex-array/src/arrays/fixed_width/take/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -6,6 +6,8 @@ mod avx2; mod records; mod scalar; mod slices; +#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +mod small_table; #[cfg(test)] mod tests; @@ -36,6 +38,8 @@ use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +use crate::dtype::PType; use crate::dtype::UnsignedPType; use crate::dtype::half::f16; use crate::match_each_unsigned_integer_ptype; @@ -80,6 +84,16 @@ pub(crate) fn take_values( values: &[T], indices: &[I], ) -> Buffer { + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + if I::PTYPE == PType::U8 { + // SAFETY: the ptype dispatcher guarantees that `I::PTYPE == U8` is the concrete `u8` + // implementation, so these slices have identical layouts. + let indices = unsafe { std::slice::from_raw_parts(indices.as_ptr().cast(), indices.len()) }; + if let Some(taken) = small_table::take(values, indices) { + return taken; + } + } + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] if *HAS_AVX2 { // SAFETY: AVX2 was detected above and `FixedWidthTakeValue` guarantees an initialized byte diff --git a/vortex-array/src/arrays/fixed_width/take/small_table.rs b/vortex-array/src/arrays/fixed_width/take/small_table.rs new file mode 100644 index 00000000000..2bb858ab721 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -0,0 +1,74 @@ +// 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::BufferMut; + +use super::FixedWidthTakeValue; + +pub(crate) fn take(values: &[T], indices: &[u8]) -> Option> { + if values.is_empty() || values.len() > 16 || size_of::() != 1 || indices.len() < 64 { + return None; + } + + let mut table = [0u8; 16]; + // SAFETY: one-byte values have the same representation as bytes, and the length is <= 16. + unsafe { + std::ptr::copy_nonoverlapping( + values.as_ptr().cast::(), + table.as_mut_ptr(), + values.len(), + ); + } + + let mut output = BufferMut::::with_capacity(indices.len()); + let output_ptr = output.spare_capacity_mut().as_mut_ptr().cast::(); + // SAFETY: AArch64 always provides NEON. Both pointers advance only by complete vectors. + 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() + ); + // SAFETY: the code was checked and this reserved output position is uninitialized. + unsafe { + output_ptr + .add(offset) + .write(*values.as_ptr().add(code).cast::()) + }; + } + 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()) }; + Some(output.freeze()) +} + +unsafe fn take_vectors(table: &[u8; 16], indices: &[u8], output: *mut u8) -> (usize, u8) { + let table = unsafe { vld1q_u8(table.as_ptr()) }; + 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) }) +} diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs index ff167d27499..90f197ef7cc 100644 --- a/vortex-array/src/arrays/fixed_width/take/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -38,6 +38,28 @@ fn take_eight_byte_values() { assert_eq!(taken.as_slice(), &[20, 30, 10]); } +#[test] +fn take_small_byte_table() { + let values = [10u8, 20, 30, 40]; + let indices = (0..128).map(|index| index % 4).collect::>(); + let expected = indices + .iter() + .map(|index| values[usize::from(*index)]) + .collect::>(); + assert_eq!( + take_values(&values, &indices).as_slice(), + expected.as_slice() + ); +} + +#[test] +#[should_panic(expected = "take index")] +fn take_small_byte_table_rejects_out_of_bounds_index() { + let mut indices = vec![0u8; 128]; + indices[64] = 4; + drop(take_values(&[10u8, 20, 30, 40], &indices)); +} + #[rstest] #[case(1)] #[case(2)] From adf36d00d60c9974fbfce8899b9fda5ace71e64d Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 9 Sep 2026 16:40:23 +0000 Subject: [PATCH 2/4] test(array): gate NEON take bounds test to aarch64 `take_small_byte_table_rejects_out_of_bounds_index` asserts the NEON small-table bounds-check message, but the table path only exists on little-endian AArch64. On other targets `take_values` falls through to the AVX2 or scalar path, whose panic messages differ, so the test failed on the linux-musl runner (no AVX2: the scalar path panics with the plain slice message). Gate the test on the same cfg as the module. `take_small_byte_table` stays unconditional so every target keeps checking the results. DCO Remediation Commit for Joe Isaacs I, Joe Isaacs , hereby add my Signed-off-by to this commit: e97db31410f1636f0b68401ee5d0204dde2101e2 Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_013YND3Jq47YtL15ob3GEsVY --- vortex-array/src/arrays/fixed_width/take/tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs index 90f197ef7cc..f161af92994 100644 --- a/vortex-array/src/arrays/fixed_width/take/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -52,6 +52,9 @@ fn take_small_byte_table() { ); } +// 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"))] #[test] #[should_panic(expected = "take index")] fn take_small_byte_table_rejects_out_of_bounds_index() { From 51d6030180b412583eee37cd186e123a7c9929e0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 5 Oct 2026 15:20:14 +0100 Subject: [PATCH 3/4] refactor: dispatch small-table take through value implementations Signed-off-by: Joe Isaacs --- .../src/arrays/fixed_width/take/mod.rs | 51 ++++++++---- .../arrays/fixed_width/take/small_table.rs | 77 ++++------------- .../fixed_width/take/small_table/neon.rs | 82 +++++++++++++++++++ .../src/arrays/fixed_width/take/tests.rs | 20 +++-- 4 files changed, 145 insertions(+), 85 deletions(-) create mode 100644 vortex-array/src/arrays/fixed_width/take/small_table/neon.rs diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs index 78684e5b385..10033a928a9 100644 --- a/vortex-array/src/arrays/fixed_width/take/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -6,7 +6,6 @@ mod avx2; mod records; mod scalar; mod slices; -#[cfg(all(target_arch = "aarch64", target_endian = "little"))] mod small_table; #[cfg(test)] mod tests; @@ -38,8 +37,6 @@ use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -#[cfg(all(target_arch = "aarch64", target_endian = "little"))] -use crate::dtype::PType; use crate::dtype::UnsignedPType; use crate::dtype::half::f16; use crate::match_each_unsigned_integer_ptype; @@ -64,7 +61,12 @@ impl 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(values: &[Self], indices: &[I]) -> Buffer { + take_values_fallback(values, indices) + } +} macro_rules! impl_fixed_width_take_value { ($($ty:ty),+ $(,)?) => { @@ -75,25 +77,44 @@ macro_rules! impl_fixed_width_take_value { }; } -impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,); +impl_fixed_width_take_value!(u16, u32, u64, i16, i32, i64, f16, f32, f64,); + +// SAFETY: u8 has no padding or uninitialized bytes. +unsafe impl FixedWidthTakeValue for u8 { + fn take(values: &[Self], indices: &[I]) -> Buffer { + small_table::take(values, indices) + } +} + +// SAFETY: i8 has no padding or uninitialized bytes. +unsafe impl FixedWidthTakeValue for i8 { + fn take(values: &[Self], indices: &[I]) -> Buffer { + small_table::take(values, indices) + } +} // SAFETY: Byte arrays have no padding and every byte is initialized. -unsafe impl FixedWidthTakeValue for [u8; N] {} +unsafe impl FixedWidthTakeValue for [u8; N] { + fn take(values: &[Self], indices: &[I]) -> Buffer { + if N == 1 { + small_table::take(values, indices) + } else { + take_values_fallback(values, indices) + } + } +} pub(crate) fn take_values( values: &[T], indices: &[I], ) -> Buffer { - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - if I::PTYPE == PType::U8 { - // SAFETY: the ptype dispatcher guarantees that `I::PTYPE == U8` is the concrete `u8` - // implementation, so these slices have identical layouts. - let indices = unsafe { std::slice::from_raw_parts(indices.as_ptr().cast(), indices.len()) }; - if let Some(taken) = small_table::take(values, indices) { - return taken; - } - } + T::take(values, indices) +} +fn take_values_fallback( + values: &[T], + indices: &[I], +) -> Buffer { #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] if *HAS_AVX2 { // SAFETY: AVX2 was detected above and `FixedWidthTakeValue` guarantees an initialized byte diff --git a/vortex-array/src/arrays/fixed_width/take/small_table.rs b/vortex-array/src/arrays/fixed_width/take/small_table.rs index 2bb858ab721..1d911ed922e 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -1,74 +1,25 @@ // 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. +//! Small byte-table take with a fallback for unsupported targets and inputs. -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; +#[cfg(all(target_arch = "aarch64", target_endian = "little"))] +mod neon; use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use super::FixedWidthTakeValue; - -pub(crate) fn take(values: &[T], indices: &[u8]) -> Option> { - if values.is_empty() || values.len() > 16 || size_of::() != 1 || indices.len() < 64 { - return None; - } - - let mut table = [0u8; 16]; - // SAFETY: one-byte values have the same representation as bytes, and the length is <= 16. - unsafe { - std::ptr::copy_nonoverlapping( - values.as_ptr().cast::(), - table.as_mut_ptr(), - values.len(), - ); - } - - let mut output = BufferMut::::with_capacity(indices.len()); - let output_ptr = output.spare_capacity_mut().as_mut_ptr().cast::(); - // SAFETY: AArch64 always provides NEON. Both pointers advance only by complete vectors. - 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() - ); - // SAFETY: the code was checked and this reserved output position is uninitialized. - unsafe { - output_ptr - .add(offset) - .write(*values.as_ptr().add(code).cast::()) - }; +use super::take_values_fallback; +use crate::dtype::UnsignedPType; + +pub(crate) fn take( + values: &[T], + indices: &[I], +) -> Buffer { + #[cfg(all(target_arch = "aarch64", target_endian = "little"))] + if let Some(taken) = neon::take(values, indices) { + return taken; } - 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()) }; - Some(output.freeze()) -} -unsafe fn take_vectors(table: &[u8; 16], indices: &[u8], output: *mut u8) -> (usize, u8) { - let table = unsafe { vld1q_u8(table.as_ptr()) }; - 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) }) + take_values_fallback(values, indices) } diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs new file mode 100644 index 00000000000..53bab2167ab --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs @@ -0,0 +1,82 @@ +// 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::BufferMut; + +use super::super::FixedWidthTakeValue; +use crate::dtype::PType; +use crate::dtype::UnsignedPType; + +pub(super) fn take( + values: &[T], + indices: &[I], +) -> Option> { + if I::PTYPE != PType::U8 + || values.is_empty() + || values.len() > 16 + || size_of::() != 1 + || indices.len() < 64 + { + return None; + } + + // 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::::with_capacity(indices.len()); + let spare = output.spare_capacity_mut(); + let output_ptr = spare.as_mut_ptr().cast::(); + // 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()) }; + Some(output.freeze()) +} + +unsafe fn take_vectors( + table: &[T; 16], + indices: &[u8], + output: *mut u8, +) -> (usize, u8) { + let table = unsafe { vld1q_u8(table.as_ptr().cast::()) }; + 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) }) +} diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs index f161af92994..010b6e8af2b 100644 --- a/vortex-array/src/arrays/fixed_width/take/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -38,10 +38,14 @@ fn take_eight_byte_values() { assert_eq!(taken.as_slice(), &[20, 30, 10]); } -#[test] -fn take_small_byte_table() { +#[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 = (0..128).map(|index| index % 4).collect::>(); + let indices = (0u8..4).cycle().take(len).collect::>(); let expected = indices .iter() .map(|index| values[usize::from(*index)]) @@ -55,11 +59,13 @@ fn take_small_byte_table() { // 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"))] -#[test] +#[rstest] +#[case::vector(64)] +#[case::tail(128)] #[should_panic(expected = "take index")] -fn take_small_byte_table_rejects_out_of_bounds_index() { - let mut indices = vec![0u8; 128]; - indices[64] = 4; +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)); } From c4d00c5cd895f5430ad1cf15852ec3ffabe246db Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 5 Oct 2026 15:22:58 +0100 Subject: [PATCH 4/4] refactor: gate byte take implementations in one module Signed-off-by: Joe Isaacs --- .../src/arrays/fixed_width/take/mod.rs | 25 ---------------- .../arrays/fixed_width/take/small_table.rs | 28 ++++++++--------- .../fixed_width/take/small_table/neon.rs | 30 ++++++++++++++++--- 3 files changed, 39 insertions(+), 44 deletions(-) diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs index 10033a928a9..9d5efd21ceb 100644 --- a/vortex-array/src/arrays/fixed_width/take/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -79,31 +79,6 @@ macro_rules! impl_fixed_width_take_value { impl_fixed_width_take_value!(u16, u32, u64, i16, i32, i64, f16, f32, f64,); -// SAFETY: u8 has no padding or uninitialized bytes. -unsafe impl FixedWidthTakeValue for u8 { - fn take(values: &[Self], indices: &[I]) -> Buffer { - small_table::take(values, indices) - } -} - -// SAFETY: i8 has no padding or uninitialized bytes. -unsafe impl FixedWidthTakeValue for i8 { - fn take(values: &[Self], indices: &[I]) -> Buffer { - small_table::take(values, indices) - } -} - -// SAFETY: Byte arrays have no padding and every byte is initialized. -unsafe impl FixedWidthTakeValue for [u8; N] { - fn take(values: &[Self], indices: &[I]) -> Buffer { - if N == 1 { - small_table::take(values, indices) - } else { - take_values_fallback(values, indices) - } - } -} - pub(crate) fn take_values( values: &[T], indices: &[I], diff --git a/vortex-array/src/arrays/fixed_width/take/small_table.rs b/vortex-array/src/arrays/fixed_width/take/small_table.rs index 1d911ed922e..cea9c429ab6 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table.rs @@ -3,23 +3,21 @@ //! Small byte-table take with a fallback for unsupported targets and inputs. -#[cfg(all(target_arch = "aarch64", target_endian = "little"))] -mod neon; +cfg_if::cfg_if! { + if #[cfg(all(target_arch = "aarch64", target_endian = "little"))] { + mod neon; + } else { + mod fallback { + use super::super::FixedWidthTakeValue; -use vortex_buffer::Buffer; + // SAFETY: u8 has no padding or uninitialized bytes. + unsafe impl FixedWidthTakeValue for u8 {} -use super::FixedWidthTakeValue; -use super::take_values_fallback; -use crate::dtype::UnsignedPType; + // SAFETY: i8 has no padding or uninitialized bytes. + unsafe impl FixedWidthTakeValue for i8 {} -pub(crate) fn take( - values: &[T], - indices: &[I], -) -> Buffer { - #[cfg(all(target_arch = "aarch64", target_endian = "little"))] - if let Some(taken) = neon::take(values, indices) { - return taken; + // SAFETY: Byte arrays have no padding and every byte is initialized. + unsafe impl FixedWidthTakeValue for [u8; N] {} + } } - - take_values_fallback(values, indices) } diff --git a/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs index 53bab2167ab..9192b94c3da 100644 --- a/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs +++ b/vortex-array/src/arrays/fixed_width/take/small_table/neon.rs @@ -15,20 +15,42 @@ use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use super::super::FixedWidthTakeValue; +use super::super::take_values_fallback; use crate::dtype::PType; use crate::dtype::UnsignedPType; -pub(super) fn take( +// SAFETY: u8 has no padding or uninitialized bytes. +unsafe impl FixedWidthTakeValue for u8 { + fn take(values: &[Self], indices: &[I]) -> Buffer { + take(values, indices) + } +} + +// SAFETY: i8 has no padding or uninitialized bytes. +unsafe impl FixedWidthTakeValue for i8 { + fn take(values: &[Self], indices: &[I]) -> Buffer { + take(values, indices) + } +} + +// SAFETY: Byte arrays have no padding and every byte is initialized. +unsafe impl FixedWidthTakeValue for [u8; N] { + fn take(values: &[Self], indices: &[I]) -> Buffer { + take(values, indices) + } +} + +fn take( values: &[T], indices: &[I], -) -> Option> { +) -> Buffer { if I::PTYPE != PType::U8 || values.is_empty() || values.len() > 16 || size_of::() != 1 || indices.len() < 64 { - return None; + return take_values_fallback(values, indices); } // SAFETY: the sealed index type is u8, as checked above. @@ -61,7 +83,7 @@ pub(super) fn take( ); // SAFETY: the vector loop and scalar remainder initialized every output value. unsafe { output.set_len(indices.len()) }; - Some(output.freeze()) + output.freeze() } unsafe fn take_vectors(