Skip to content

perf(take): Vectorise bounds check in take_native (-8-10%)#9754

Draft
Dandandan wants to merge 1 commit intoapache:mainfrom
Dandandan:take_bounds2
Draft

perf(take): Vectorise bounds check in take_native (-8-10%)#9754
Dandandan wants to merge 1 commit intoapache:mainfrom
Dandandan:take_bounds2

Conversation

@Dandandan
Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

The non-null branch of take_native — the primitive take gather —
bounds-checks each index individually via values[index.as_usize()].
The per-lane branch dominates the loop and blocks autovectorisation.

Reducing a chunk of indices to their maximum with fold+max has no
early exit, so LLVM lowers the whole check to a SIMD max-reduction; on
aarch64 that's two ldp q + three umax.4s + one umaxv.4s — a
single bounds check per chunk, then the gather.

max_idx no longer has to stay live on the hot path for the panic
format string because the panic is moved into a #[cold] helper,
which removes a per-chunk stack spill (str x16, [sp, #16]).

Signed index types sign-extend to usize::MAX on as_usize(), so
negative indices still fail the check; panic behaviour is preserved.

What changes are included in this PR?

One-file change in arrow-select/src/take.rs:

  • CHUNK = 16 chunked max-reduction bounds check
  • #[cold] oob helper for the panic path
  • preallocated output via Vec::set_len + chunks_exact_mut so the
    gather is a straight SIMD store (no push / capacity-check
    overhead)

Output is written into a Vec<T> with uninitialised capacity and
set_len(len) up front; T: ArrowNativeType is Copy with no
Drop, and every slot is written before out is read.

Are these changes tested?

Yes — covered by existing take::tests::*, including
test_take_out_of_bounds_panic.

Measured on aarch64 (Apple Silicon) with
cargo bench --bench take_kernels -- "^take i32":

bench before after Δ
take i32 512 309 ns 279 ns −9.7%
take i32 1024 469 ns 431 ns −8.1%
take i32 null indices 542 ns 545 ns no change (unchanged branch)

Are there any user-facing changes?

No — no API change, panic behaviour on out-of-bounds indices is
preserved.

🤖 Generated with Claude Code

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) <noreply@anthropic.com>
@github-actions github-actions bot added the arrow Changes to the arrow crate label Apr 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

arrow Changes to the arrow crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant