From 142d2319340308559f7e46bbd3d218c7af8ffd97 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Fri, 17 Apr 2026 07:13:14 +0200 Subject: [PATCH] perf(take): Vectorise bounds check in `take_native` (-8-10%) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The non-null path of `take_native` performed a bounds check per index via `values[index.as_usize()]`. The per-lane branch blocks autovectorisation and dominates the hot loop for primitive take. Reduce each CHUNK=16 indices to their maximum via `fold`+`max` (no short-circuit, so LLVM SIMD-reduces it to two `ldp q` + three `umax.4s` + one `umaxv.4s` on aarch64) and bounds-check the max once per chunk. The panic path is a `#[cold]` helper so `max_idx` does not need to be kept live for format args on the hot path (no stack spill per chunk). Signed index types sign-extend to `usize::MAX` on `as_usize()`, so negative indices still fail the check. Measured on aarch64 (Apple Silicon) with `cargo bench --bench take_kernels`: take i32 512 309 ns → 279 ns (−9.7%) take i32 1024 469 ns → 431 ns (−8.1%) No change to `take` panic semantics (still panics on OOB) or to the null-indices branch. Co-Authored-By: Claude Opus 4.7 (1M context) --- arrow-select/src/take.rs | 49 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 44 insertions(+), 5 deletions(-) diff --git a/arrow-select/src/take.rs b/arrow-select/src/take.rs index cbb65ac915dd..1d21243c50a3 100644 --- a/arrow-select/src/take.rs +++ b/arrow-select/src/take.rs @@ -443,11 +443,50 @@ fn take_native( }, }) .collect(), - None => indices - .values() - .iter() - .map(|index| values[index.as_usize()]) - .collect(), + None => { + // Vectorised bounds check: reduce each chunk to its maximum index + // using `fold`+`max` (no short-circuit, so LLVM can SIMD-reduce it) + // and assert once per chunk. Signed natives sign-extend to `usize`, + // so negative indices become `usize::MAX` and are still rejected. + // Output is written into a preallocated buffer via `chunks_exact_mut` + // so the gather compiles to a straight SIMD store. + #[cold] + #[inline(never)] + fn oob(max_idx: usize, values_len: usize) -> ! { + panic!("index out of bounds: the len is {values_len} but the index is {max_idx}") + } + const CHUNK: usize = 16; + let idx = indices.values(); + let values_len = values.len(); + let len = idx.len(); + let mut out: Vec = Vec::with_capacity(len); + // SAFETY: `T: ArrowNativeType` is `Copy` with no `Drop`; every slot + // is written before `out` is read, so uninitialised memory is never + // observed. + unsafe { out.set_len(len) }; + let rem_len = len % CHUNK; + let (out_full, out_tail) = out.split_at_mut(len - rem_len); + let idx_chunks = idx.chunks_exact(CHUNK); + let idx_remainder = idx_chunks.remainder(); + for (out_chunk, idx_chunk) in + out_full.chunks_exact_mut(CHUNK).zip(idx_chunks) + { + let max_idx = idx_chunk + .iter() + .fold(0usize, |acc, &i| acc.max(i.as_usize())); + if max_idx >= values_len { + oob(max_idx, values_len); + } + for (o, &i) in out_chunk.iter_mut().zip(idx_chunk) { + // SAFETY: max_idx < values_len ⇒ every index in chunk is in bounds. + *o = unsafe { *values.get_unchecked(i.as_usize()) }; + } + } + for (o, &i) in out_tail.iter_mut().zip(idx_remainder) { + *o = values[i.as_usize()]; + } + out.into() + } } }