diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index fdd1fafba05..d92c033c2a4 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -1,16 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +//! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) boolean execute path on a +//! focused set of configurations: +//! +//! - `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 is run nullable and non-nullable. + #![expect(clippy::unwrap_used)] +use std::fmt::Display; +use std::fmt::Formatter; + use divan::Bencher; use rand::RngExt; use rand::SeedableRng; use rand::distr::Uniform; use rand::prelude::StdRng; +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; @@ -21,18 +35,74 @@ 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, +/// 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: Pattern, + 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 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 nullable in [false, true] { + for (pattern, branches) in [ + (Pattern::RoundRobin, 2), + (Pattern::Random, 2), + (Pattern::Random, 64), + ] { + out.push(Combo { + pattern, + branches, + nullable, + }); + } + } + out +} + +/// Builds the Vortex value arrays and the `u32` selector buffers for a [`Combo`]. +/// +/// 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 = (0..num_branches) + let values = (0..combo.branches) .map(|_| { - if nullable { + if combo.nullable { BoolArray::from_iter( (0..ARRAY_SIZE).map(|_| (rng.sample(bit) == 0).then_some(rng.sample(bit) == 0)), ) @@ -43,37 +113,27 @@ fn inputs( }) .collect(); - let branch = Uniform::new(0u32, u32::try_from(num_branches).unwrap()).unwrap(); + 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(|_| rng.sample(branch)).collect(); - let row_indices: Buffer = (0..ARRAY_SIZE).map(|_| rng.sample(row)).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(), - ) + let array_indices: Buffer = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + Pattern::Random => rng.sample(branch), + Pattern::RoundRobin => u32::try_from(i % combo.branches).unwrap(), }) - .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); + .collect(); + let row_indices: Buffer = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + Pattern::Random => rng.sample(row), + Pattern::RoundRobin => u32::try_from((i / combo.branches) % ARRAY_SIZE).unwrap(), + }) + .collect(); + (values, array_indices, row_indices) } -#[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(|| { ( diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index 755921abd59..fde5b161dfd 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -8,7 +8,6 @@ use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_mask::Mask; use super::super::Interleave; use super::super::InterleaveArrayExt; @@ -21,62 +20,44 @@ 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()); } + 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, @@ -85,43 +66,70 @@ 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: 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). -#[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 below are tight branchless 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" ); } - let values = - BitBufferMut::collect_bool(len, |i| value_bits[branches[i].as_()].value(rows[i].as_())); - - // A missing per-value mask means every row of that value is valid; 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_())) - }) - }); + Ok(len) +} - 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`]. +/// +/// The bounds-checked `BitBuffer::value` is slower still. +/// +/// # Safety +/// +/// `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 { + // SAFETY: `collect_bool` calls this for `i < len`, and the caller guarantees `branches[i]` and + // `rows[i]` are in bounds for `bits` / the selected buffer. + BitBufferMut::collect_bool(len, |i| unsafe { + 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 b085d9571e1..dd0b8853523 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())) } } @@ -524,6 +569,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. @@ -602,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() {