From 9f471782c6cb94cc98654e268d3829087c4259b2 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 08:04:31 +0000 Subject: [PATCH 01/12] bench(vortex-array): compare Interleave bool execute against arrow interleave Extends the interleave benchmark to measure the Vortex boolean execute path against the arrow-rs `interleave` kernel on identical data, across three access patterns (random, round_robin, single_branch), for N=2 and N=4 value arrays, nullable and non-nullable. Signed-off-by: Joe Isaacs --- vortex-array/benches/interleave.rs | 175 ++++++++++++++++++++++------- 1 file changed, 137 insertions(+), 38 deletions(-) diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index fdd1fafba05..d0bdc6d816a 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -1,13 +1,33 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) boolean execute path +//! against the equivalent arrow-rs [`interleave`] kernel, across several access patterns: +//! +//! - `random`: fully random `(array_index, row_index)` per output row. +//! - `round_robin`: a merge — `array_index = i % N`, `row_index = i / N`, each branch consumed +//! front-to-back; the structured pattern most amenable to both engines. +//! - `single_branch`: every row routed to branch 0 with a random row (a degenerate gather). +//! +//! Each pattern is run for `N = 2` and `N = 4`, nullable and non-nullable, against identical data +//! so the two engines are directly comparable. + #![expect(clippy::unwrap_used)] +use std::fmt::Display; +use std::fmt::Formatter; +use std::sync::Arc; + +use arrow_array::Array as ArrowArray; +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::BooleanArray; +use arrow_select::interleave::interleave; use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +use vortex_array::{array_session, ArrayRef}; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -21,59 +41,129 @@ fn main() { const ARRAY_SIZE: usize = 8_192; -/// Builds `num_branches` boolean value arrays plus random `(array_indices, row_indices)` selectors -/// describing a full random-access gather of `ARRAY_SIZE` output rows. -fn inputs( - num_branches: usize, +/// A single benchmark configuration: access pattern, branch count, and nullability. +#[derive(Clone, Copy)] +struct Combo { + pattern: &'static str, + branches: usize, nullable: bool, -) -> (Vec, Buffer, Buffer) { +} + +impl Display for Combo { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}/n{}/{}", + self.pattern, + self.branches, + if self.nullable { "null" } else { "nonnull" } + ) + } +} + +/// The full cross product of patterns × branch counts × nullability that every engine is measured +/// against, so the Vortex and arrow rows line up one-to-one in the report. +fn combos() -> Vec { + let mut out = Vec::new(); + for pattern in ["random", "round_robin", "single_branch"] { + for branches in [2, 4] { + for nullable in [false, true] { + out.push(Combo { + pattern, + branches, + nullable, + }); + } + } + } + out +} + +/// Per-branch boolean values (as `Option`, all `Some` when non-nullable) paired with the +/// `(array_index, row_index)` selectors for one [`Combo`]. +type Inputs = (Vec>>, Vec<(usize, usize)>); + +/// Generates the shared [`Inputs`] for a [`Combo`]. +/// +/// Seeded only by the combo so both engines see byte-identical data for a fair comparison. +fn gen_inputs(combo: Combo) -> Inputs { let mut rng = StdRng::seed_from_u64(0); let bit = Uniform::new(0u8, 2).unwrap(); - let values = (0..num_branches) + let values: Vec>> = (0..combo.branches) .map(|_| { - if nullable { - BoolArray::from_iter( - (0..ARRAY_SIZE).map(|_| (rng.sample(bit) == 0).then_some(rng.sample(bit) == 0)), - ) - .into_array() + (0..ARRAY_SIZE) + .map(|_| { + if combo.nullable && rng.sample(bit) == 0 { + None + } else { + Some(rng.sample(bit) == 0) + } + }) + .collect() + }) + .collect(); + + let branch = Uniform::new(0usize, combo.branches).unwrap(); + let row = Uniform::new(0usize, ARRAY_SIZE).unwrap(); + let selectors: Vec<(usize, usize)> = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + "random" => (rng.sample(branch), rng.sample(row)), + "round_robin" => (i % combo.branches, (i / combo.branches) % ARRAY_SIZE), + "single_branch" => (0, rng.sample(row)), + other => unreachable!("unknown pattern {other}"), + }) + .collect(); + + (values, selectors) +} + +/// Builds the Vortex value arrays and the `u32` selector buffers for a [`Combo`]. +fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { + let (values, selectors) = gen_inputs(combo); + let values = values + .into_iter() + .map(|vals| { + if combo.nullable { + BoolArray::from_iter(vals).into_array() } else { - BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(bit) == 0)).into_array() + BoolArray::from_iter(vals.into_iter().map(Option::unwrap)).into_array() } }) .collect(); - - let branch = Uniform::new(0u32, u32::try_from(num_branches).unwrap()).unwrap(); - let row = Uniform::new(0u32, u32::try_from(ARRAY_SIZE).unwrap()).unwrap(); - let array_indices: Buffer = (0..ARRAY_SIZE).map(|_| rng.sample(branch)).collect(); - let row_indices: Buffer = (0..ARRAY_SIZE).map(|_| rng.sample(row)).collect(); + let array_indices = selectors + .iter() + .map(|&(a, _)| u32::try_from(a).unwrap()) + .collect(); + let row_indices = selectors + .iter() + .map(|&(_, r)| u32::try_from(r).unwrap()) + .collect(); (values, array_indices, row_indices) } -#[divan::bench(args = [2, 4])] -fn interleave_bool(bencher: Bencher, num_branches: usize) { - let (values, array_indices, row_indices) = inputs(num_branches, false); - let session = vortex_array::array_session(); - bencher - .with_inputs(|| { - ( - InterleaveArray::try_new( - values.clone(), - array_indices.clone().into_array(), - row_indices.clone().into_array(), - ) - .unwrap() - .into_array(), - session.create_execution_ctx(), - ) +/// Builds the arrow value arrays and the `(usize, usize)` selectors for a [`Combo`]. +fn arrow_inputs(combo: Combo) -> (Vec, Vec<(usize, usize)>) { + let (values, selectors) = gen_inputs(combo); + let values = values + .into_iter() + .map(|vals| -> ArrowArrayRef { + if combo.nullable { + Arc::new(BooleanArray::from(vals)) + } else { + Arc::new(BooleanArray::from( + vals.into_iter().map(Option::unwrap).collect::>(), + )) + } }) - .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); + .collect(); + (values, selectors) } -#[divan::bench(args = [2, 4])] -fn interleave_bool_nullable(bencher: Bencher, num_branches: usize) { - let (values, array_indices, row_indices) = inputs(num_branches, true); - let session = vortex_array::array_session(); +#[divan::bench(args = combos())] +fn vortex(bencher: Bencher, combo: Combo) { + let (values, array_indices, row_indices) = vortex_inputs(combo); + let session = array_session(); bencher .with_inputs(|| { ( @@ -89,3 +179,12 @@ fn interleave_bool_nullable(bencher: Bencher, num_branches: usize) { }) .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); } + +#[divan::bench(args = combos())] +fn arrow(bencher: Bencher, combo: Combo) { + let (values, selectors) = arrow_inputs(combo); + bencher.bench(|| { + let refs: Vec<&dyn ArrowArray> = values.iter().map(AsRef::as_ref).collect(); + interleave(&refs, &selectors).unwrap() + }); +} From 5be0dd370dd7bb605f522f13051eb1833b419d71 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 08:04:41 +0000 Subject: [PATCH 02/12] perf(vortex-array): speed up Interleave bool gather with hoisted unchecked reads Profiling and disassembly of the gather inner loop showed the per-bit cost was dominated not by the 64-bit packing (optimal for a random gather) but by the overhead around each read: a wide `&[BitBuffer]` struct index, the redundant bounds `assert` inside `BitBuffer::value`, and on the nullable path a per-row `Option`/`Mask`-variant dispatch. Pre-validate the per-row bounds once, then hoist a raw `(ptr, bit_offset)` per value buffer and read each selected bit with `get_bit_unchecked`. Materialize each branch's validity into one full-length `BitBuffer` so the validity gather is the same uniform unchecked read. Measured ~1.2x faster non-nullable and ~1.4x nullable across random / round-robin / single-branch patterns. Also adds two two-value word-boundary regression tests exercising the gather across more than one 64-bit packing word, in the non-nullable and nullable cases, validated against the existing scalar reference oracle. Signed-off-by: Joe Isaacs --- .../src/arrays/interleave/execute/bool.rs | 52 +++++++++++++++---- vortex-array/src/arrays/interleave/mod.rs | 23 ++++++++ 2 files changed, 64 insertions(+), 11 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index 755921abd59..7c1d1a3d576 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -6,6 +6,7 @@ use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; +use vortex_buffer::get_bit_unchecked; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_mask::Mask; @@ -87,8 +88,13 @@ pub(super) fn execute( /// pair is read straight from its packed buffer. /// /// Output bits (and validity) are produced with [`BitBufferMut::collect_bool`], which packs 64 -/// results per word: every output bit is written branchlessly, avoiding a per-row `set`/`unset` -/// (each of which would bounds-check and branch on the random bit value). +/// results per word. For a random-access gather there is no word-level shortcut on the read side — +/// consecutive outputs read unrelated source words — so the work is one bit read per output. The +/// per-read overhead is what we trim: a raw `(ptr, bit_offset)` is hoisted per value buffer and the +/// bit is read with [`get_bit_unchecked`], avoiding the wide `&[BitBuffer]` struct index and the +/// redundant bounds assert in `BitBuffer::value`. Validity is materialized into one full-length +/// [`BitBuffer`] per branch so its gather is the same uniform unchecked read rather than a per-row +/// `Option`/`Mask`-variant dispatch. #[allow(clippy::too_many_arguments)] fn gather, R: AsPrimitive>( len: usize, @@ -100,7 +106,7 @@ fn gather, R: AsPrimitive>( nullable: bool, ) -> VortexResult<(BitBufferMut, Option)> { // Validate the per-row bounds once up front (returning an error rather than panicking), so the - // word-packing passes below are tight branchless loops. + // word-packing passes below are tight unchecked loops. for i in 0..len { let branch = branches[i].as_(); vortex_ensure!(branch < num_values, "interleave array index out of bounds"); @@ -110,16 +116,40 @@ fn gather, R: AsPrimitive>( ); } - let values = - BitBufferMut::collect_bool(len, |i| value_bits[branches[i].as_()].value(rows[i].as_())); + // Raw (byte pointer, bit offset) per value buffer; the offset folds the buffer's own bit offset + // into the index so the read is a single `get_bit_unchecked`. + let val_ptrs: Vec<(*const u8, usize)> = value_bits + .iter() + .map(|b| (b.inner().as_ptr(), b.offset())) + .collect(); - // A missing per-value mask means every row of that value is valid; only materialized when the - // output can be null. + // SAFETY (both passes): `i < len`, `branches`/`rows` have length `len`, and the loop above + // proved `branches[i] < num_values` and `rows[i] < value_bits[branches[i]].len()`, which equals + // the validity length for that branch. So every `get_unchecked` / `get_bit_unchecked` is in + // bounds. + let values = BitBufferMut::collect_bool(len, |i| unsafe { + let (ptr, off) = *val_ptrs.get_unchecked(branches.get_unchecked(i).as_()); + get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) + }); + + // A missing per-value mask means every row of that value is valid; validity is only materialized + // when the output can be null. let validity = nullable.then(|| { - BitBufferMut::collect_bool(len, |i| { - value_validity[branches[i].as_()] - .as_ref() - .is_none_or(|mask| mask.value(rows[i].as_())) + let validity_bits: Vec = value_validity + .iter() + .enumerate() + .map(|(j, mask)| match mask { + Some(mask) => mask.to_bit_buffer(), + None => BitBuffer::new_set(value_bits[j].len()), + }) + .collect(); + let vld_ptrs: Vec<(*const u8, usize)> = validity_bits + .iter() + .map(|b| (b.inner().as_ptr(), b.offset())) + .collect(); + BitBufferMut::collect_bool(len, |i| unsafe { + let (ptr, off) = *vld_ptrs.get_unchecked(branches.get_unchecked(i).as_()); + get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) }) }); diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index b085d9571e1..8f4db672c1e 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -524,6 +524,29 @@ mod tests { ) } + #[test] + fn interleave_binary_spans_word_boundary() -> VortexResult<()> { + // Exercises a two-value gather across more than one 64-bit packing word, with out-of-order + // routing into both branches so neither buffer is consumed contiguously. + let branch0: Vec> = (0..100).map(|i| Some(i % 3 == 0)).collect(); + let branch1: Vec> = (0..100).map(|i| Some(i % 5 == 0)).collect(); + let indices: Vec<(usize, usize)> = (0..200).map(|i| (i % 2, (i * 7) % 100)).collect(); + check(&[&branch0, &branch1], &indices) + } + + #[test] + fn interleave_binary_nullable_spans_word_boundary() -> VortexResult<()> { + // Same two-value gather, now with nulls so the validity packing is exercised too. + let branch0: Vec> = (0..70) + .map(|i| (i % 4 != 0).then_some(i % 2 == 0)) + .collect(); + let branch1: Vec> = (0..70) + .map(|i| (i % 3 != 0).then_some(i % 2 == 1)) + .collect(); + let indices: Vec<(usize, usize)> = (0..150).map(|i| (i % 2, (i * 11) % 70)).collect(); + check(&[&branch0, &branch1], &indices) + } + #[test] fn interleave_three_values() -> VortexResult<()> { // An unsigned `array_indices` routes among three values with full random access. From ea84125902e91345d44af26899b81b45e5d85249 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 09:08:52 +0000 Subject: [PATCH 03/12] bench(vortex-array): drop arrow-rs comparison from interleave benchmark Remove the arrow-rs `interleave` comparison arm and its input builder, keeping the Vortex-only benchmark across the random / round_robin / single_branch access patterns (N=2, N=4, nullable and non-nullable). No dependency changes: every arrow crate in vortex-array is used by the crate's own source. Signed-off-by: Joe Isaacs --- vortex-array/benches/interleave.rs | 111 +++++++---------------------- 1 file changed, 27 insertions(+), 84 deletions(-) diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index d0bdc6d816a..d74cc334587 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -2,26 +2,20 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) boolean execute path -//! against the equivalent arrow-rs [`interleave`] kernel, across several access patterns: +//! across several access patterns: //! //! - `random`: fully random `(array_index, row_index)` per output row. //! - `round_robin`: a merge — `array_index = i % N`, `row_index = i / N`, each branch consumed -//! front-to-back; the structured pattern most amenable to both engines. +//! front-to-back. //! - `single_branch`: every row routed to branch 0 with a random row (a degenerate gather). //! -//! Each pattern is run for `N = 2` and `N = 4`, nullable and non-nullable, against identical data -//! so the two engines are directly comparable. +//! Each pattern is run for `N = 2` and `N = 4`, nullable and non-nullable. #![expect(clippy::unwrap_used)] use std::fmt::Display; use std::fmt::Formatter; -use std::sync::Arc; -use arrow_array::Array as ArrowArray; -use arrow_array::ArrayRef as ArrowArrayRef; -use arrow_array::BooleanArray; -use arrow_select::interleave::interleave; use divan::Bencher; use rand::RngExt; use rand::SeedableRng; @@ -61,8 +55,7 @@ impl Display for Combo { } } -/// The full cross product of patterns × branch counts × nullability that every engine is measured -/// against, so the Vortex and arrow rows line up one-to-one in the report. +/// The full cross product of patterns × branch counts × nullability that the benchmark covers. fn combos() -> Vec { let mut out = Vec::new(); for pattern in ["random", "round_robin", "single_branch"] { @@ -79,87 +72,46 @@ fn combos() -> Vec { out } -/// Per-branch boolean values (as `Option`, all `Some` when non-nullable) paired with the -/// `(array_index, row_index)` selectors for one [`Combo`]. -type Inputs = (Vec>>, Vec<(usize, usize)>); - -/// Generates the shared [`Inputs`] for a [`Combo`]. +/// Builds the Vortex value arrays and the `u32` selector buffers for a [`Combo`]. /// -/// Seeded only by the combo so both engines see byte-identical data for a fair comparison. -fn gen_inputs(combo: Combo) -> Inputs { +/// Seeded only by the combo so a run is deterministic and comparable across revisions. +fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { let mut rng = StdRng::seed_from_u64(0); let bit = Uniform::new(0u8, 2).unwrap(); - let values: Vec>> = (0..combo.branches) + let values = (0..combo.branches) .map(|_| { - (0..ARRAY_SIZE) - .map(|_| { - if combo.nullable && rng.sample(bit) == 0 { - None - } else { - Some(rng.sample(bit) == 0) - } - }) - .collect() + if combo.nullable { + BoolArray::from_iter( + (0..ARRAY_SIZE).map(|_| (rng.sample(bit) == 0).then_some(rng.sample(bit) == 0)), + ) + .into_array() + } else { + BoolArray::from_iter((0..ARRAY_SIZE).map(|_| rng.sample(bit) == 0)).into_array() + } }) .collect(); - let branch = Uniform::new(0usize, combo.branches).unwrap(); - let row = Uniform::new(0usize, ARRAY_SIZE).unwrap(); - let selectors: Vec<(usize, usize)> = (0..ARRAY_SIZE) + let branch = Uniform::new(0u32, u32::try_from(combo.branches).unwrap()).unwrap(); + let row = Uniform::new(0u32, u32::try_from(ARRAY_SIZE).unwrap()).unwrap(); + let array_indices: Buffer = (0..ARRAY_SIZE) .map(|i| match combo.pattern { - "random" => (rng.sample(branch), rng.sample(row)), - "round_robin" => (i % combo.branches, (i / combo.branches) % ARRAY_SIZE), - "single_branch" => (0, rng.sample(row)), + "random" => rng.sample(branch), + "round_robin" => u32::try_from(i % combo.branches).unwrap(), + "single_branch" => 0, other => unreachable!("unknown pattern {other}"), }) .collect(); - - (values, selectors) -} - -/// Builds the Vortex value arrays and the `u32` selector buffers for a [`Combo`]. -fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { - let (values, selectors) = gen_inputs(combo); - let values = values - .into_iter() - .map(|vals| { - if combo.nullable { - BoolArray::from_iter(vals).into_array() - } else { - BoolArray::from_iter(vals.into_iter().map(Option::unwrap)).into_array() - } + let row_indices: Buffer = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + "random" | "single_branch" => rng.sample(row), + "round_robin" => u32::try_from((i / combo.branches) % ARRAY_SIZE).unwrap(), + other => unreachable!("unknown pattern {other}"), }) .collect(); - let array_indices = selectors - .iter() - .map(|&(a, _)| u32::try_from(a).unwrap()) - .collect(); - let row_indices = selectors - .iter() - .map(|&(_, r)| u32::try_from(r).unwrap()) - .collect(); (values, array_indices, row_indices) } -/// Builds the arrow value arrays and the `(usize, usize)` selectors for a [`Combo`]. -fn arrow_inputs(combo: Combo) -> (Vec, Vec<(usize, usize)>) { - let (values, selectors) = gen_inputs(combo); - let values = values - .into_iter() - .map(|vals| -> ArrowArrayRef { - if combo.nullable { - Arc::new(BooleanArray::from(vals)) - } else { - Arc::new(BooleanArray::from( - vals.into_iter().map(Option::unwrap).collect::>(), - )) - } - }) - .collect(); - (values, selectors) -} - #[divan::bench(args = combos())] fn vortex(bencher: Bencher, combo: Combo) { let (values, array_indices, row_indices) = vortex_inputs(combo); @@ -179,12 +131,3 @@ fn vortex(bencher: Bencher, combo: Combo) { }) .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); } - -#[divan::bench(args = combos())] -fn arrow(bencher: Bencher, combo: Combo) { - let (values, selectors) = arrow_inputs(combo); - bencher.bench(|| { - let refs: Vec<&dyn ArrowArray> = values.iter().map(AsRef::as_ref).collect(); - interleave(&refs, &selectors).unwrap() - }); -} From 7258d027e4b3936aab3f510852d8c57191743725 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 09:23:25 +0000 Subject: [PATCH 04/12] refactor(vortex-array): dispatch interleave gather over concrete selector types Drop the `AsPrimitive` bound on the bool gather: dispatch both selectors over their concrete unsigned width via `match_each_unsigned_integer_ptype!` and convert to an index with a plain `as usize` on the native type, passed in as `branch_at` / `row_at` accessors. Keeps the selector loads at their native width. Signed-off-by: Joe Isaacs --- .../src/arrays/interleave/execute/bool.rs | 50 +++++++++++-------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index 7c1d1a3d576..82959793745 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -3,7 +3,6 @@ //! Optimized [`Interleave`] implementation for boolean values. -use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::get_bit_unchecked; @@ -58,19 +57,30 @@ pub(super) fn execute( } // Scatter directly from the typed selector buffers — no intermediate `usize` materialization. + // Both selectors are dispatched over their concrete unsigned width, and each value is converted + // to an index with a plain `as usize` cast on that concrete type (no `AsPrimitive`). let array_indices = array.array_indices().as_::(); let row_indices = array.row_indices().as_::(); let (values, validity) = match_each_unsigned_integer_ptype!(array_indices.ptype(), |A| { match_each_unsigned_integer_ptype!(row_indices.ptype(), |R| { - gather( + let branches = array_indices.as_slice::(); + let rows = row_indices.as_slice::(); + // SAFETY: both accessors are only ever called with `i < len`, and `branches`/`rows` + // each have length `len` (the selectors are equal-length and equal to the output len). + // + // The `as usize` widens a concrete unsigned selector (`u8`..`u64`) to an index. Selector + // values index in-memory arrays, so they always fit in `usize`; the cast is lossless. + #[allow(clippy::cast_possible_truncation)] + let result = gather( len, num_values, &value_bits, &value_validity, - array_indices.as_slice::(), - row_indices.as_slice::(), + |i| unsafe { *branches.get_unchecked(i) as usize }, + |i| unsafe { *rows.get_unchecked(i) as usize }, nullable, - )? + )?; + result }) }); @@ -84,8 +94,9 @@ pub(super) fn execute( )?)) } -/// The scatter, monomorphized on the selector integer widths so each `(array_index, row_index)` -/// pair is read straight from its packed buffer. +/// The scatter, monomorphized on the selector integer widths via the `branch_at` / `row_at` +/// accessors so each `(array_index, row_index)` pair is read straight from its packed buffer at its +/// native width and cast to an index with a concrete `as usize`. /// /// Output bits (and validity) are produced with [`BitBufferMut::collect_bool`], which packs 64 /// results per word. For a random-access gather there is no word-level shortcut on the read side — @@ -96,22 +107,22 @@ pub(super) fn execute( /// [`BitBuffer`] per branch so its gather is the same uniform unchecked read rather than a per-row /// `Option`/`Mask`-variant dispatch. #[allow(clippy::too_many_arguments)] -fn gather, R: AsPrimitive>( +fn gather( len: usize, num_values: usize, value_bits: &[BitBuffer], value_validity: &[Option], - branches: &[A], - rows: &[R], + branch_at: impl Fn(usize) -> usize, + row_at: impl Fn(usize) -> usize, nullable: bool, ) -> VortexResult<(BitBufferMut, Option)> { // Validate the per-row bounds once up front (returning an error rather than panicking), so the // word-packing passes below are tight unchecked loops. for i in 0..len { - let branch = branches[i].as_(); + let branch = branch_at(i); vortex_ensure!(branch < num_values, "interleave array index out of bounds"); vortex_ensure!( - rows[i].as_() < value_bits[branch].len(), + row_at(i) < value_bits[branch].len(), "interleave row index out of bounds" ); } @@ -123,13 +134,12 @@ fn gather, R: AsPrimitive>( .map(|b| (b.inner().as_ptr(), b.offset())) .collect(); - // SAFETY (both passes): `i < len`, `branches`/`rows` have length `len`, and the loop above - // proved `branches[i] < num_values` and `rows[i] < value_bits[branches[i]].len()`, which equals - // the validity length for that branch. So every `get_unchecked` / `get_bit_unchecked` is in - // bounds. + // SAFETY (both passes): `i < len`, and the loop above proved `branch_at(i) < num_values` and + // `row_at(i) < value_bits[branch_at(i)].len()`, which equals the validity length for that + // branch. So every `get_unchecked` / `get_bit_unchecked` is in bounds. let values = BitBufferMut::collect_bool(len, |i| unsafe { - let (ptr, off) = *val_ptrs.get_unchecked(branches.get_unchecked(i).as_()); - get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) + let (ptr, off) = *val_ptrs.get_unchecked(branch_at(i)); + get_bit_unchecked(ptr, row_at(i) + off) }); // A missing per-value mask means every row of that value is valid; validity is only materialized @@ -148,8 +158,8 @@ fn gather, R: AsPrimitive>( .map(|b| (b.inner().as_ptr(), b.offset())) .collect(); BitBufferMut::collect_bool(len, |i| unsafe { - let (ptr, off) = *vld_ptrs.get_unchecked(branches.get_unchecked(i).as_()); - get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) + let (ptr, off) = *vld_ptrs.get_unchecked(branch_at(i)); + get_bit_unchecked(ptr, row_at(i) + off) }) }); From ed96c2697221aaf55a1ffaf3478d70a38088208f Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 09:26:33 +0000 Subject: [PATCH 05/12] bench(vortex-array): trim interleave benchmark to three focused configs Keep only 2-child round-robin, 2-child random, and 64-child random (each nullable and non-nullable); drop the single_branch pattern and the N=4 variants. Signed-off-by: Joe Isaacs --- vortex-array/benches/interleave.rs | 35 ++++++++++++++---------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index d74cc334587..4393d0be5e4 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -1,15 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) boolean execute path -//! across several access patterns: +//! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) boolean execute path on a +//! focused set of configurations: //! -//! - `random`: fully random `(array_index, row_index)` per output row. -//! - `round_robin`: a merge — `array_index = i % N`, `row_index = i / N`, each branch consumed -//! front-to-back. -//! - `single_branch`: every row routed to branch 0 with a random row (a degenerate gather). +//! - `round_robin`, 2 children: a merge — `array_index = i % N`, `row_index = i / N`. +//! - `random`, 2 children: fully random `(array_index, row_index)` per output row. +//! - `random`, 64 children: the same random gather spread over many value arrays. //! -//! Each pattern is run for `N = 2` and `N = 4`, nullable and non-nullable. +//! Each is run nullable and non-nullable. #![expect(clippy::unwrap_used)] @@ -55,18 +54,17 @@ impl Display for Combo { } } -/// The full cross product of patterns × branch counts × nullability that the benchmark covers. +/// The configurations the benchmark covers: 2-child round-robin, 2-child random, and 64-child +/// random — each nullable and non-nullable. fn combos() -> Vec { let mut out = Vec::new(); - for pattern in ["random", "round_robin", "single_branch"] { - for branches in [2, 4] { - for nullable in [false, true] { - out.push(Combo { - pattern, - branches, - nullable, - }); - } + for nullable in [false, true] { + for (pattern, branches) in [("round_robin", 2), ("random", 2), ("random", 64)] { + out.push(Combo { + pattern, + branches, + nullable, + }); } } out @@ -98,13 +96,12 @@ fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { .map(|i| match combo.pattern { "random" => rng.sample(branch), "round_robin" => u32::try_from(i % combo.branches).unwrap(), - "single_branch" => 0, other => unreachable!("unknown pattern {other}"), }) .collect(); let row_indices: Buffer = (0..ARRAY_SIZE) .map(|i| match combo.pattern { - "random" | "single_branch" => rng.sample(row), + "random" => rng.sample(row), "round_robin" => u32::try_from((i / combo.branches) % ARRAY_SIZE).unwrap(), other => unreachable!("unknown pattern {other}"), }) From 9557d7c7c06259490b43fa5edbe76f644b0874fc Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 09:31:34 +0000 Subject: [PATCH 06/12] refactor(vortex-array): revert interleave gather to AsPrimitive selectors Go back to the generic `gather, R: AsPrimitive>` with `.as_()` conversions, dropping the per-type `as usize` accessors. The trait conversion avoids the `cast_possible_truncation` lint without a scoped allow and matches the idiom used by the other bool kernels (e.g. take). Signed-off-by: Joe Isaacs --- .../src/arrays/interleave/execute/bool.rs | 50 ++++++++----------- 1 file changed, 20 insertions(+), 30 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index 82959793745..7c1d1a3d576 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -3,6 +3,7 @@ //! Optimized [`Interleave`] implementation for boolean values. +use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::get_bit_unchecked; @@ -57,30 +58,19 @@ pub(super) fn execute( } // Scatter directly from the typed selector buffers — no intermediate `usize` materialization. - // Both selectors are dispatched over their concrete unsigned width, and each value is converted - // to an index with a plain `as usize` cast on that concrete type (no `AsPrimitive`). let array_indices = array.array_indices().as_::(); let row_indices = array.row_indices().as_::(); let (values, validity) = match_each_unsigned_integer_ptype!(array_indices.ptype(), |A| { match_each_unsigned_integer_ptype!(row_indices.ptype(), |R| { - let branches = array_indices.as_slice::(); - let rows = row_indices.as_slice::(); - // SAFETY: both accessors are only ever called with `i < len`, and `branches`/`rows` - // each have length `len` (the selectors are equal-length and equal to the output len). - // - // The `as usize` widens a concrete unsigned selector (`u8`..`u64`) to an index. Selector - // values index in-memory arrays, so they always fit in `usize`; the cast is lossless. - #[allow(clippy::cast_possible_truncation)] - let result = gather( + gather( len, num_values, &value_bits, &value_validity, - |i| unsafe { *branches.get_unchecked(i) as usize }, - |i| unsafe { *rows.get_unchecked(i) as usize }, + array_indices.as_slice::(), + row_indices.as_slice::(), nullable, - )?; - result + )? }) }); @@ -94,9 +84,8 @@ pub(super) fn execute( )?)) } -/// The scatter, monomorphized on the selector integer widths via the `branch_at` / `row_at` -/// accessors so each `(array_index, row_index)` pair is read straight from its packed buffer at its -/// native width and cast to an index with a concrete `as usize`. +/// The scatter, monomorphized on the selector integer widths so each `(array_index, row_index)` +/// pair is read straight from its packed buffer. /// /// Output bits (and validity) are produced with [`BitBufferMut::collect_bool`], which packs 64 /// results per word. For a random-access gather there is no word-level shortcut on the read side — @@ -107,22 +96,22 @@ pub(super) fn execute( /// [`BitBuffer`] per branch so its gather is the same uniform unchecked read rather than a per-row /// `Option`/`Mask`-variant dispatch. #[allow(clippy::too_many_arguments)] -fn gather( +fn gather, R: AsPrimitive>( len: usize, num_values: usize, value_bits: &[BitBuffer], value_validity: &[Option], - branch_at: impl Fn(usize) -> usize, - row_at: impl Fn(usize) -> usize, + branches: &[A], + rows: &[R], nullable: bool, ) -> VortexResult<(BitBufferMut, Option)> { // Validate the per-row bounds once up front (returning an error rather than panicking), so the // word-packing passes below are tight unchecked loops. for i in 0..len { - let branch = branch_at(i); + let branch = branches[i].as_(); vortex_ensure!(branch < num_values, "interleave array index out of bounds"); vortex_ensure!( - row_at(i) < value_bits[branch].len(), + rows[i].as_() < value_bits[branch].len(), "interleave row index out of bounds" ); } @@ -134,12 +123,13 @@ fn gather( .map(|b| (b.inner().as_ptr(), b.offset())) .collect(); - // SAFETY (both passes): `i < len`, and the loop above proved `branch_at(i) < num_values` and - // `row_at(i) < value_bits[branch_at(i)].len()`, which equals the validity length for that - // branch. So every `get_unchecked` / `get_bit_unchecked` is in bounds. + // SAFETY (both passes): `i < len`, `branches`/`rows` have length `len`, and the loop above + // proved `branches[i] < num_values` and `rows[i] < value_bits[branches[i]].len()`, which equals + // the validity length for that branch. So every `get_unchecked` / `get_bit_unchecked` is in + // bounds. let values = BitBufferMut::collect_bool(len, |i| unsafe { - let (ptr, off) = *val_ptrs.get_unchecked(branch_at(i)); - get_bit_unchecked(ptr, row_at(i) + off) + let (ptr, off) = *val_ptrs.get_unchecked(branches.get_unchecked(i).as_()); + get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) }); // A missing per-value mask means every row of that value is valid; validity is only materialized @@ -158,8 +148,8 @@ fn gather( .map(|b| (b.inner().as_ptr(), b.offset())) .collect(); BitBufferMut::collect_bool(len, |i| unsafe { - let (ptr, off) = *vld_ptrs.get_unchecked(branch_at(i)); - get_bit_unchecked(ptr, row_at(i) + off) + let (ptr, off) = *vld_ptrs.get_unchecked(branches.get_unchecked(i).as_()); + get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) }) }); From 745702b9da70d01ff58e2248a129138914cd7f54 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 09:42:48 +0000 Subject: [PATCH 07/12] bench(vortex-array): use a Pattern enum instead of stringly-typed pattern Replace `Combo.pattern: &'static str` and the `match`/`unreachable!` on string literals with a `Pattern` enum (RoundRobin, Random), making the match exhaustive and dropping the unreachable arms. Signed-off-by: Joe Isaacs --- vortex-array/benches/interleave.rs | 36 +++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index 4393d0be5e4..0353e8fb926 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -34,10 +34,28 @@ fn main() { const ARRAY_SIZE: usize = 8_192; +/// The access pattern used to generate the `(array_index, row_index)` selectors. +#[derive(Clone, Copy)] +enum Pattern { + /// A merge: `array_index = i % N`, `row_index = i / N`. + RoundRobin, + /// Fully random `(array_index, row_index)` per output row. + Random, +} + +impl Display for Pattern { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(match self { + Pattern::RoundRobin => "round_robin", + Pattern::Random => "random", + }) + } +} + /// A single benchmark configuration: access pattern, branch count, and nullability. #[derive(Clone, Copy)] struct Combo { - pattern: &'static str, + pattern: Pattern, branches: usize, nullable: bool, } @@ -59,7 +77,11 @@ impl Display for Combo { fn combos() -> Vec { let mut out = Vec::new(); for nullable in [false, true] { - for (pattern, branches) in [("round_robin", 2), ("random", 2), ("random", 64)] { + for (pattern, branches) in [ + (Pattern::RoundRobin, 2), + (Pattern::Random, 2), + (Pattern::Random, 64), + ] { out.push(Combo { pattern, branches, @@ -94,16 +116,14 @@ fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { let row = Uniform::new(0u32, u32::try_from(ARRAY_SIZE).unwrap()).unwrap(); let array_indices: Buffer = (0..ARRAY_SIZE) .map(|i| match combo.pattern { - "random" => rng.sample(branch), - "round_robin" => u32::try_from(i % combo.branches).unwrap(), - other => unreachable!("unknown pattern {other}"), + Pattern::Random => rng.sample(branch), + Pattern::RoundRobin => u32::try_from(i % combo.branches).unwrap(), }) .collect(); let row_indices: Buffer = (0..ARRAY_SIZE) .map(|i| match combo.pattern { - "random" => rng.sample(row), - "round_robin" => u32::try_from((i / combo.branches) % ARRAY_SIZE).unwrap(), - other => unreachable!("unknown pattern {other}"), + Pattern::Random => rng.sample(row), + Pattern::RoundRobin => u32::try_from((i / combo.branches) % ARRAY_SIZE).unwrap(), }) .collect(); (values, array_indices, row_indices) From 3c3f31b69857c8a4625a7ce0ac5588ee4bca07a5 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 11 Jun 2026 10:05:40 +0000 Subject: [PATCH 08/12] refactor(vortex-array): share one bit-gather kernel for values and validity Validity is just another set of boolean buffers routed by the same selectors, so factor the hoisted-pointer unchecked gather into `gather_bits` and call it for both the values and the materialized per-branch validity, instead of duplicating the pointer-table build and `collect_bool` loop. Signed-off-by: Joe Isaacs --- .../src/arrays/interleave/execute/bool.rs | 74 +++++++++++-------- 1 file changed, 42 insertions(+), 32 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index 7c1d1a3d576..b754a55531f 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -87,14 +87,10 @@ pub(super) fn execute( /// The scatter, monomorphized on the selector integer widths so each `(array_index, row_index)` /// pair is read straight from its packed buffer. /// -/// Output bits (and validity) are produced with [`BitBufferMut::collect_bool`], which packs 64 -/// results per word. For a random-access gather there is no word-level shortcut on the read side — -/// consecutive outputs read unrelated source words — so the work is one bit read per output. The -/// per-read overhead is what we trim: a raw `(ptr, bit_offset)` is hoisted per value buffer and the -/// bit is read with [`get_bit_unchecked`], avoiding the wide `&[BitBuffer]` struct index and the -/// redundant bounds assert in `BitBuffer::value`. Validity is materialized into one full-length -/// [`BitBuffer`] per branch so its gather is the same uniform unchecked read rather than a per-row -/// `Option`/`Mask`-variant dispatch. +/// Values and validity are both gathered with [`gather_bits`]: validity is just another set of +/// boolean buffers (one per branch) routed by the same selectors, so once each branch's validity is +/// materialized into a full-length [`BitBuffer`] it runs through the identical kernel rather than a +/// per-row `Option`/`Mask`-variant dispatch. #[allow(clippy::too_many_arguments)] fn gather, R: AsPrimitive>( len: usize, @@ -106,7 +102,7 @@ fn gather, R: AsPrimitive>( nullable: bool, ) -> VortexResult<(BitBufferMut, Option)> { // Validate the per-row bounds once up front (returning an error rather than panicking), so the - // word-packing passes below are tight unchecked loops. + // word-packing passes in `gather_bits` are tight unchecked loops. for i in 0..len { let branch = branches[i].as_(); vortex_ensure!(branch < num_values, "interleave array index out of bounds"); @@ -116,21 +112,9 @@ fn gather, R: AsPrimitive>( ); } - // Raw (byte pointer, bit offset) per value buffer; the offset folds the buffer's own bit offset - // into the index so the read is a single `get_bit_unchecked`. - let val_ptrs: Vec<(*const u8, usize)> = value_bits - .iter() - .map(|b| (b.inner().as_ptr(), b.offset())) - .collect(); - - // SAFETY (both passes): `i < len`, `branches`/`rows` have length `len`, and the loop above - // proved `branches[i] < num_values` and `rows[i] < value_bits[branches[i]].len()`, which equals - // the validity length for that branch. So every `get_unchecked` / `get_bit_unchecked` is in - // bounds. - let values = BitBufferMut::collect_bool(len, |i| unsafe { - let (ptr, off) = *val_ptrs.get_unchecked(branches.get_unchecked(i).as_()); - get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) - }); + // SAFETY: the loop above proved `branches[i] < num_values == value_bits.len()` and + // `rows[i] < value_bits[branches[i]].len()` for every `i < len`. + let values = unsafe { gather_bits(len, value_bits, branches, rows) }; // A missing per-value mask means every row of that value is valid; validity is only materialized // when the output can be null. @@ -143,15 +127,41 @@ fn gather, R: AsPrimitive>( None => BitBuffer::new_set(value_bits[j].len()), }) .collect(); - let vld_ptrs: Vec<(*const u8, usize)> = validity_bits - .iter() - .map(|b| (b.inner().as_ptr(), b.offset())) - .collect(); - BitBufferMut::collect_bool(len, |i| unsafe { - let (ptr, off) = *vld_ptrs.get_unchecked(branches.get_unchecked(i).as_()); - get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) - }) + // SAFETY: each validity buffer has its value's length, so the bounds validated above hold. + unsafe { gather_bits(len, &validity_bits, branches, rows) } }); Ok((values, validity)) } + +/// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per +/// word with [`BitBufferMut::collect_bool`]. +/// +/// For a random-access gather there is no word-level shortcut on the read side — consecutive outputs +/// read unrelated source words — so the work is one bit read per output. The per-read overhead is +/// what we trim: a raw `(ptr, bit_offset)` is hoisted per buffer and the bit is read with +/// [`get_bit_unchecked`], avoiding the wide `&[BitBuffer]` struct index and the redundant bounds +/// assert in `BitBuffer::value`. +/// +/// # Safety +/// For every `i < len`, `branches[i] < bits.len()` and `rows[i] < bits[branches[i]].len()`. +unsafe fn gather_bits, R: AsPrimitive>( + len: usize, + bits: &[BitBuffer], + branches: &[A], + rows: &[R], +) -> BitBufferMut { + // Raw (byte pointer, bit offset) per buffer; the offset folds the buffer's own bit offset into + // the index so the read is a single `get_bit_unchecked`. + let ptrs: Vec<(*const u8, usize)> = bits + .iter() + .map(|b| (b.inner().as_ptr(), b.offset())) + .collect(); + + // SAFETY: `collect_bool` calls this for `i < len`, and the caller guarantees `branches[i]` and + // `rows[i]` are in bounds for `ptrs` / the selected buffer. + BitBufferMut::collect_bool(len, |i| unsafe { + let (ptr, off) = *ptrs.get_unchecked(branches.get_unchecked(i).as_()); + get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) + }) +} From 01bbd778174a4e2fbc5baac4c4174b1b031b83ba Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 17 Jun 2026 12:37:51 +0100 Subject: [PATCH 09/12] fix Signed-off-by: Joe Isaacs --- .../src/arrays/interleave/execute/bool.rs | 137 +++++++-------- vortex-array/src/arrays/interleave/mod.rs | 158 ++++++++++++++---- 2 files changed, 178 insertions(+), 117 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index b754a55531f..b4d9cff4d73 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -6,10 +6,8 @@ use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; -use vortex_buffer::get_bit_unchecked; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use super::super::Interleave; use super::super::InterleaveArrayExt; @@ -22,62 +20,47 @@ use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::match_each_unsigned_integer_ptype; use crate::require_child; -use crate::validity::Validity; /// Gathers `N` boolean values under unsigned `array_indices` / `row_indices` selectors, scattering -/// each selected bit (and its validity) into the output position it routes to. +/// each selected bit into the output position it routes to. pub(super) fn execute( array: Array, - ctx: &mut ExecutionCtx, + _ctx: &mut ExecutionCtx, ) -> VortexResult { let num_values = array.num_values(); - // Drive every value and both selectors to canonical encodings so we can operate on raw bits. + // Drive both selectors and every value to canonical encodings so we can operate on raw bits. let mut array = array; + array = require_child!(array, array.array_indices(), 0 => Primitive); + array = require_child!(array, array.row_indices(), 1 => Primitive); for i in 0..num_values { - array = require_child!(array, array.value(i), i => Bool); + array = require_child!(array, array.value(i), i + 2 => Bool); } - array = require_child!(array, array.array_indices(), num_values => Primitive); - array = require_child!(array, array.row_indices(), num_values + 1 => Primitive); - let dtype = array.as_ref().dtype().clone(); - let len = array.as_ref().len(); - let nullable = dtype.is_nullable(); - - // Materialize each value's bits, and its validity mask only when the output can be null. + // Materialize each value's bits; the selectors gather one bit per output below. let mut value_bits = Vec::with_capacity(num_values); - let mut value_validity = Vec::with_capacity(num_values); for i in 0..num_values { - let value = array.value(i).as_::(); - let bits = value.to_bit_buffer(); - let validity = nullable - .then(|| value.validity()?.execute_mask(bits.len(), ctx)) - .transpose()?; - value_bits.push(bits); - value_validity.push(validity); + value_bits.push(array.value(i).as_::().to_bit_buffer()); } + // Hold the validity as a pushed-down interleave rather than applying it: the routing pair for + // each output selects the value *and* its validity bit, so the output validity is itself an + // interleave (by these selectors) of the values' validities. This bottoms out lazily. + let validity = array.as_ref().validity()?; + // Scatter directly from the typed selector buffers — no intermediate `usize` materialization. let array_indices = array.array_indices().as_::(); let row_indices = array.row_indices().as_::(); - let (values, validity) = match_each_unsigned_integer_ptype!(array_indices.ptype(), |A| { + let values = match_each_unsigned_integer_ptype!(array_indices.ptype(), |A| { match_each_unsigned_integer_ptype!(row_indices.ptype(), |R| { gather( - len, - num_values, &value_bits, - &value_validity, array_indices.as_slice::(), row_indices.as_slice::(), - nullable, )? }) }); - let validity = match validity { - Some(bits) => Validity::from(bits.freeze()), - None => Validity::NonNullable, - }; Ok(ExecutionResult::done(BoolArray::try_new( values.freeze(), validity, @@ -86,82 +69,76 @@ pub(super) fn execute( /// The scatter, monomorphized on the selector integer widths so each `(array_index, row_index)` /// pair is read straight from its packed buffer. -/// -/// Values and validity are both gathered with [`gather_bits`]: validity is just another set of -/// boolean buffers (one per branch) routed by the same selectors, so once each branch's validity is -/// materialized into a full-length [`BitBuffer`] it runs through the identical kernel rather than a -/// per-row `Option`/`Mask`-variant dispatch. -#[allow(clippy::too_many_arguments)] fn gather, R: AsPrimitive>( - len: usize, - num_values: usize, value_bits: &[BitBuffer], - value_validity: &[Option], branches: &[A], rows: &[R], - nullable: bool, -) -> VortexResult<(BitBufferMut, Option)> { - // Validate the per-row bounds once up front (returning an error rather than panicking), so the - // word-packing passes in `gather_bits` are tight unchecked loops. +) -> VortexResult { + let len = validate_selectors(value_bits, branches, rows)?; + + // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every + // `i < len` that `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()`. + Ok(unsafe { gather_bits(len, value_bits, branches, rows) }) +} + +/// Validates the per-row selector bounds, returning the output length (`branches.len()`). +/// +/// On success, `rows.len() == branches.len() == len` and, for every `i < len`, +/// `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()` — exactly the +/// preconditions of [`gather_bits`]. Errors (rather than panics) on any out-of-bounds selector. +fn validate_selectors, R: AsPrimitive>( + value_bits: &[BitBuffer], + branches: &[A], + rows: &[R], +) -> VortexResult { + // The two selectors are validated to equal length at construction, which is the output length. + let len = branches.len(); + vortex_ensure!( + rows.len() == len, + "interleave selectors differ in length: array_indices {len}, row_indices {}", + rows.len() + ); + for i in 0..len { let branch = branches[i].as_(); - vortex_ensure!(branch < num_values, "interleave array index out of bounds"); + vortex_ensure!( + branch < value_bits.len(), + "interleave array index out of bounds" + ); vortex_ensure!( rows[i].as_() < value_bits[branch].len(), "interleave row index out of bounds" ); } - // SAFETY: the loop above proved `branches[i] < num_values == value_bits.len()` and - // `rows[i] < value_bits[branches[i]].len()` for every `i < len`. - let values = unsafe { gather_bits(len, value_bits, branches, rows) }; - - // A missing per-value mask means every row of that value is valid; validity is only materialized - // when the output can be null. - let validity = nullable.then(|| { - let validity_bits: Vec = value_validity - .iter() - .enumerate() - .map(|(j, mask)| match mask { - Some(mask) => mask.to_bit_buffer(), - None => BitBuffer::new_set(value_bits[j].len()), - }) - .collect(); - // SAFETY: each validity buffer has its value's length, so the bounds validated above hold. - unsafe { gather_bits(len, &validity_bits, branches, rows) } - }); - - Ok((values, validity)) + Ok(len) } /// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per /// word with [`BitBufferMut::collect_bool`]. /// /// For a random-access gather there is no word-level shortcut on the read side — consecutive outputs -/// read unrelated source words — so the work is one bit read per output. The per-read overhead is -/// what we trim: a raw `(ptr, bit_offset)` is hoisted per buffer and the bit is read with -/// [`get_bit_unchecked`], avoiding the wide `&[BitBuffer]` struct index and the redundant bounds -/// assert in `BitBuffer::value`. +/// read unrelated source words — so the work is one bit read per output. Each read indexes the +/// `&[BitBuffer]` slice and uses [`BitBuffer::value_unchecked`]. Benchmarking (`gather_values` in +/// `benches/interleave.rs`) showed this beats pre-hoisting a `(ptr, bit_offset)` table per buffer: +/// the table's extra indirection and per-call allocation cost more than reloading the selected +/// buffer's pointer/offset, which stay hot in its `BitBuffer` struct. The bounds-checked +/// `BitBuffer::value` is slower still. /// /// # Safety -/// For every `i < len`, `branches[i] < bits.len()` and `rows[i] < bits[branches[i]].len()`. +/// +/// `branches` and `rows` must both contain at least `len` elements. For every `i < len`, +/// `branches[i] < bits.len()` and `rows[i] < bits[branches[i]].len()`. unsafe fn gather_bits, R: AsPrimitive>( len: usize, bits: &[BitBuffer], branches: &[A], rows: &[R], ) -> BitBufferMut { - // Raw (byte pointer, bit offset) per buffer; the offset folds the buffer's own bit offset into - // the index so the read is a single `get_bit_unchecked`. - let ptrs: Vec<(*const u8, usize)> = bits - .iter() - .map(|b| (b.inner().as_ptr(), b.offset())) - .collect(); - // SAFETY: `collect_bool` calls this for `i < len`, and the caller guarantees `branches[i]` and - // `rows[i]` are in bounds for `ptrs` / the selected buffer. + // `rows[i]` are in bounds for `bits` / the selected buffer. BitBufferMut::collect_bool(len, |i| unsafe { - let (ptr, off) = *ptrs.get_unchecked(branches.get_unchecked(i).as_()); - get_bit_unchecked(ptr, rows.get_unchecked(i).as_() + off) + bits.get_unchecked(branches.get_unchecked(i).as_()) + .value_unchecked(rows.get_unchecked(i).as_()) }) } diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index 8f4db672c1e..cc32a7d59c7 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -6,8 +6,8 @@ //! //! # Specification //! -//! An [`Interleave`] array has `N + 2` children: `N` *values* followed by an `array_indices` -//! selector and a `row_indices` selector. The output has `array_indices.len()` rows, and output +//! An [`Interleave`] array has `N + 2` children: an `array_indices` selector and a `row_indices` +//! selector followed by `N` *values*. The output has `array_indices.len()` rows, and output //! row `i` comes from `values[array_indices[i]][row_indices[i]]`. //! //! Unlike a `Merge`, which consumes each branch in order under a cursor, an [`Interleave`] is @@ -86,8 +86,8 @@ pub struct Interleave; /// Per-array metadata for an [`InterleaveArray`]. /// -/// The values and selectors live in the array's slots; only the value count is stored here so the -/// selector slots can be located (`slots[num_values]` and `slots[num_values + 1]`). +/// The selectors and values live in the array's slots; the selectors occupy `slots[0]` and +/// `slots[1]`, and the `num_values` values follow at `slots[2..2 + num_values]`. #[derive(Clone, Debug)] pub struct InterleaveData { pub(crate) num_values: usize, @@ -120,21 +120,21 @@ pub trait InterleaveArrayExt: TypedArrayRef { /// The `idx`-th value array (holding the rows that `array_indices` routes to it). fn value(&self, idx: usize) -> &ArrayRef { - self.as_ref().slots()[idx] + self.as_ref().slots()[idx + 2] .as_ref() .vortex_expect("validated interleave value slot") } /// The selector routing each output row to a value array. fn array_indices(&self) -> &ArrayRef { - self.as_ref().slots()[self.num_values] + self.as_ref().slots()[0] .as_ref() .vortex_expect("validated interleave array_indices slot") } /// The selector naming each output row's position within its value array. fn row_indices(&self) -> &ArrayRef { - self.as_ref().slots()[self.num_values + 1] + self.as_ref().slots()[1] .as_ref() .vortex_expect("validated interleave row_indices slot") } @@ -215,19 +215,68 @@ impl Array { row_indices: ArrayRef, ) -> VortexResult { let dtype = Interleave::check(&values, &array_indices, &row_indices)?; - let len = array_indices.len(); - let num_values = values.len(); + // SAFETY: `check` just validated every invariant and computed the matching `dtype`. + Ok(unsafe { Self::new_unchecked(values, array_indices, row_indices, dtype) }) + } - let mut slots: ArraySlots = values.into_iter().map(Some).collect(); + /// Constructs an [`InterleaveArray`] without re-validating the spec invariants. + /// + /// This is the assembly half of [`try_new`](Self::try_new): it lays the selectors into + /// `slots[0]`/`slots[1]` and the values into `slots[2..]` and stamps `dtype`, skipping the + /// [`Interleave::check`] pass. It exists for hot internal paths — notably building the + /// pushed-down validity interleave in [`ValidityVTable::validity`] — where the inputs are known + /// to satisfy the invariants by construction and re-checking them is pure overhead. + /// + /// # Safety + /// + /// The caller must uphold every [module invariant](self): at least two `values` sharing a dtype + /// up to nullability, both selectors non-nullable unsigned integers of equal length, and `dtype` + /// equal to the value returned by [`Interleave::check`] for these arguments (the shared value + /// type with the union of the values' nullabilities). + pub unsafe fn new_unchecked( + values: Vec, + array_indices: ArrayRef, + row_indices: ArrayRef, + dtype: DType, + ) -> Self { + let mut slots: ArraySlots = ArraySlots::with_capacity(values.len() + 2); slots.push(Some(array_indices)); slots.push(Some(row_indices)); + slots.extend(values.into_iter().map(Some)); + + // SAFETY: the caller of `new_unchecked` upholds every invariant; here we only assemble the + // canonical slot layout (`array_indices`, `row_indices`, then values) that follows. + unsafe { Self::new_unchecked_slots(slots, dtype) } + } - Ok(unsafe { + /// Constructs an [`InterleaveArray`] from a pre-assembled `slots` buffer, skipping both the + /// spec re-check and the slot copy that [`new_unchecked`](Self::new_unchecked) performs. + /// + /// This is the lowest-level assembly path: it stamps `dtype` onto `slots` as-is. Callers that + /// already own a correctly laid-out [`ArraySlots`] — for example reusing a validated parent's + /// selectors while swapping its values — avoid materializing an intermediate `Vec` + /// and re-pushing it into a fresh `ArraySlots`. + /// + /// # Safety + /// + /// The caller must uphold every [module invariant](self) *and* the slot layout: `slots[0]` is + /// the `array_indices` selector, `slots[1]` is the `row_indices` selector, and `slots[2..]` are + /// the value arrays. Every slot must be `Some`, there must be at least two values + /// (`slots.len() >= 4`), and `dtype` must equal the value returned by [`Interleave::check`] for + /// these arguments. + pub unsafe fn new_unchecked_slots(slots: ArraySlots, dtype: DType) -> Self { + let num_values = slots.len() - 2; + let len = slots[0] + .as_ref() + .vortex_expect("interleave array_indices slot present") + .len(); + + unsafe { Array::from_parts_unchecked( ArrayParts::new(Interleave, dtype, len, InterleaveData { num_values }) .with_slots(slots), ) - }) + } } } @@ -259,16 +308,14 @@ impl VTable for Interleave { "InterleaveArray slots must all be present" ); - let values: Vec = slots[..data.num_values] + let array_indices = slots[0] + .clone() + .vortex_expect("validated array_indices slot"); + let row_indices = slots[1].clone().vortex_expect("validated row_indices slot"); + let values: Vec = slots[2..] .iter() .map(|s| s.clone().vortex_expect("validated value slot")) .collect(); - let array_indices = slots[data.num_values] - .clone() - .vortex_expect("validated array_indices slot"); - let row_indices = slots[data.num_values + 1] - .clone() - .vortex_expect("validated row_indices slot"); // All semantic invariants live in `check`; here we only confirm the array's cached `dtype` // and `len` agree with what the children imply. @@ -300,13 +347,11 @@ impl VTable for Interleave { None } - fn slot_name(array: ArrayView<'_, Self>, idx: usize) -> String { - if idx == array.num_values() { - "array_indices".to_string() - } else if idx == array.num_values() + 1 { - "row_indices".to_string() - } else { - format!("value_{idx}") + fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { + match idx { + 0 => "array_indices".to_string(), + 1 => "row_indices".to_string(), + _ => format!("value_{}", idx - 2), } } @@ -370,18 +415,18 @@ impl ValidityVTable for Interleave { if !array.as_ref().dtype().is_nullable() { return Ok(Validity::NonNullable); } - // The output validity is itself an interleave — by the same selectors — of the values' - // validities, expressed as non-nullable boolean arrays. This bottoms out immediately - // because the inner interleave is non-nullable. - let mut value_validities: Vec = Vec::with_capacity(array.num_values()); - for i in 0..array.num_values() { - value_validities.push(value_validity_array(array.value(i))?); + let num_values = array.num_values(); + let mut slots: ArraySlots = ArraySlots::with_capacity(num_values + 2); + slots.push(Some(array.array_indices().clone())); + slots.push(Some(array.row_indices().clone())); + for i in 0..num_values { + slots.push(Some(value_validity_array(array.value(i))?)); } - let interleaved = InterleaveArray::try_new( - value_validities, - array.array_indices().clone(), - array.row_indices().clone(), - )?; + // SAFETY: `value_validity_array` yields a non-nullable boolean array per value, + // the selectors already validated. + let interleaved = unsafe { + InterleaveArray::new_unchecked_slots(slots, DType::Bool(Nullability::NonNullable)) + }; Ok(Validity::Array(interleaved.into_array())) } } @@ -625,6 +670,45 @@ mod tests { assert!(err.to_string().contains("equal length"), "{err}"); } + #[test] + fn execute_rejects_out_of_bounds_array_index() -> VortexResult<()> { + let value = BoolArray::from_iter([true]).into_array(); + let array_indices = PrimitiveArray::from_iter([2u32]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)? + .into_array(); + + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let err = interleaved + .execute::(&mut ctx) + .err() + .vortex_expect("expected execution to reject out-of-bounds array index"); + assert!( + err.to_string().contains("array index out of bounds"), + "{err}" + ); + Ok(()) + } + + #[test] + fn execute_rejects_out_of_bounds_row_index() -> VortexResult<()> { + let value = BoolArray::from_iter([true]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u32]).into_array(); + let row_indices = PrimitiveArray::from_iter([1u32]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![value.clone(), value], array_indices, row_indices)? + .into_array(); + + let mut ctx = LEGACY_SESSION.create_execution_ctx(); + let err = interleaved + .execute::(&mut ctx) + .err() + .vortex_expect("expected execution to reject out-of-bounds row index"); + assert!(err.to_string().contains("row index out of bounds"), "{err}"); + Ok(()) + } + #[test] #[should_panic(expected = "only implemented for boolean values")] fn non_boolean_value_execution_panics() { From ccda647f5c0edb962cd4683687bdec7287e9441f Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 17 Jun 2026 13:46:08 +0100 Subject: [PATCH 10/12] fix Signed-off-by: Joe Isaacs --- vortex-array/src/arrays/interleave/execute/bool.rs | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index b4d9cff4d73..fde5b161dfd 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -43,9 +43,6 @@ pub(super) fn execute( value_bits.push(array.value(i).as_::().to_bit_buffer()); } - // Hold the validity as a pushed-down interleave rather than applying it: the routing pair for - // each output selects the value *and* its validity bit, so the output validity is itself an - // interleave (by these selectors) of the values' validities. This bottoms out lazily. let validity = array.as_ref().validity()?; // Scatter directly from the typed selector buffers — no intermediate `usize` materialization. @@ -117,13 +114,7 @@ fn validate_selectors, R: AsPrimitive>( /// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per /// word with [`BitBufferMut::collect_bool`]. /// -/// For a random-access gather there is no word-level shortcut on the read side — consecutive outputs -/// read unrelated source words — so the work is one bit read per output. Each read indexes the -/// `&[BitBuffer]` slice and uses [`BitBuffer::value_unchecked`]. Benchmarking (`gather_values` in -/// `benches/interleave.rs`) showed this beats pre-hoisting a `(ptr, bit_offset)` table per buffer: -/// the table's extra indirection and per-call allocation cost more than reloading the selected -/// buffer's pointer/offset, which stay hot in its `BitBuffer` struct. The bounds-checked -/// `BitBuffer::value` is slower still. +/// The bounds-checked `BitBuffer::value` is slower still. /// /// # Safety /// From aae3eaee4640ec5132beb0654e000cafe3254773 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 17 Jun 2026 11:49:30 +0000 Subject: [PATCH 11/12] docs(vortex-array): unlink private Interleave::check from public constructor docs The `new_unchecked` / `new_unchecked_slots` doc comments intra-doc-linked `[`Interleave::check`]`, a private item, failing `rustdoc::private_intra_doc_links` (`-D warnings`). Drop the link brackets, keeping the inline-code reference. Signed-off-by: Joe Isaacs --- vortex-array/src/arrays/interleave/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index cc32a7d59c7..dd0b8853523 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -223,7 +223,7 @@ impl Array { /// /// This is the assembly half of [`try_new`](Self::try_new): it lays the selectors into /// `slots[0]`/`slots[1]` and the values into `slots[2..]` and stamps `dtype`, skipping the - /// [`Interleave::check`] pass. It exists for hot internal paths — notably building the + /// `Interleave::check` pass. It exists for hot internal paths — notably building the /// pushed-down validity interleave in [`ValidityVTable::validity`] — where the inputs are known /// to satisfy the invariants by construction and re-checking them is pure overhead. /// @@ -231,7 +231,7 @@ impl Array { /// /// The caller must uphold every [module invariant](self): at least two `values` sharing a dtype /// up to nullability, both selectors non-nullable unsigned integers of equal length, and `dtype` - /// equal to the value returned by [`Interleave::check`] for these arguments (the shared value + /// equal to the value returned by `Interleave::check` for these arguments (the shared value /// type with the union of the values' nullabilities). pub unsafe fn new_unchecked( values: Vec, @@ -262,7 +262,7 @@ impl Array { /// The caller must uphold every [module invariant](self) *and* the slot layout: `slots[0]` is /// the `array_indices` selector, `slots[1]` is the `row_indices` selector, and `slots[2..]` are /// the value arrays. Every slot must be `Some`, there must be at least two values - /// (`slots.len() >= 4`), and `dtype` must equal the value returned by [`Interleave::check`] for + /// (`slots.len() >= 4`), and `dtype` must equal the value returned by `Interleave::check` for /// these arguments. pub unsafe fn new_unchecked_slots(slots: ArraySlots, dtype: DType) -> Self { let num_values = slots.len() - 2; From 12a1071be5f8ded156d5d1d30b505b1b4f736095 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 22 Jun 2026 23:50:45 +0000 Subject: [PATCH 12/12] style(vortex-array): split braced import in interleave benchmark `cargo +nightly fmt --all --check` (the Rust lint job's format step) requires per-item imports; split `use vortex_array::{array_session, ArrayRef}`. Signed-off-by: Joe Isaacs --- vortex-array/benches/interleave.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index 0353e8fb926..d92c033c2a4 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -20,10 +20,11 @@ use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; -use vortex_array::{array_session, ArrayRef}; +use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::InterleaveArray; use vortex_buffer::Buffer;