From 0b5c2acb61e760848b8210ad11091ddc0b4c5024 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 7 Aug 2026 16:03:50 -0400 Subject: [PATCH 1/5] feat(vortex-array): execute primitive interleave arrays Signed-off-by: Nemo Yu --- .../src/arrays/interleave/execute/bool.rs | 42 +----- .../src/arrays/interleave/execute/mod.rs | 45 ++++++- .../arrays/interleave/execute/primitive.rs | 127 ++++++++++++++++++ vortex-array/src/arrays/interleave/mod.rs | 61 +++++++-- 4 files changed, 225 insertions(+), 50 deletions(-) create mode 100644 vortex-array/src/arrays/interleave/execute/primitive.rs diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index fde5b161dfd..a051f55ec5d 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -7,10 +7,10 @@ use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use super::super::Interleave; use super::super::InterleaveArrayExt; +use super::validate_selectors; use crate::array::Array; use crate::arrays::Bool; use crate::arrays::BoolArray; @@ -71,46 +71,18 @@ fn gather, R: AsPrimitive>( branches: &[A], rows: &[R], ) -> VortexResult { - let len = validate_selectors(value_bits, branches, rows)?; + let len = validate_selectors( + value_bits.len(), + |branch| value_bits[branch].len(), + 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 < value_bits.len(), - "interleave array index out of bounds" - ); - vortex_ensure!( - rows[i].as_() < value_bits[branch].len(), - "interleave row index out of bounds" - ); - } - - Ok(len) -} - /// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per /// word with [`BitBufferMut::collect_bool`]. /// diff --git a/vortex-array/src/arrays/interleave/execute/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index 05dcd161f62..e4ac1f5ef43 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -5,14 +5,16 @@ //! //! All values share a type (validated in [`Interleave::check`]), so the //! physical gather kernel is chosen from the first value. The selector types are an orthogonal -//! concern handled within each kernel. Only boolean values are implemented today (see the [`bool`] module). +//! concern handled within each kernel. //! //! [`Interleave::check`]: super::Interleave::check -//! [`bool`]: module@crate::arrays::interleave::execute::bool mod bool; +mod primitive; +use num_traits::AsPrimitive; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use super::Interleave; @@ -28,12 +30,47 @@ pub(super) fn execute( ) -> VortexResult { if array.value(0).dtype().is_boolean() { bool::execute(array, ctx) + } else if array.value(0).dtype().is_primitive() { + primitive::execute(array, ctx) } else { let value_dtype = array.value(0).dtype().clone(); vortex_panic!( - "interleave execution is only implemented for boolean values; value dtype {} is not \ - yet supported", + "interleave execution is not implemented for value dtype {}", value_dtype ) } } + +/// Validate selector lengths and bounds, returning the common output length. +/// +/// On success, `branches.len() == rows.len() == len`; for every `i < len`, +/// `branches[i] < num_values` and `rows[i] < value_len(branches[i])`. +fn validate_selectors( + num_values: usize, + value_len: F, + branches: &[A], + rows: &[R], +) -> VortexResult +where + A: AsPrimitive, + R: AsPrimitive, + F: Fn(usize) -> usize, +{ + 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!( + rows[i].as_() < value_len(branch), + "interleave row index out of bounds" + ); + } + + Ok(len) +} diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs new file mode 100644 index 00000000000..1eea5541d10 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution for primitive [`Interleave`] values. + +use num_traits::AsPrimitive; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use super::super::Interleave; +use super::super::InterleaveArrayExt; +use super::validate_selectors; +use crate::array::Array; +use crate::array::ArrayView; +use crate::arrays::Constant; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::primitive::PrimitiveArrayExt; +use crate::dtype::NativePType; +use crate::executor::ExecutionCtx; +use crate::executor::ExecutionResult; +use crate::match_each_native_ptype; +use crate::match_each_unsigned_integer_ptype; +use crate::require_child; + +pub(super) fn execute( + mut array: Array, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + let num_values = array.num_values(); + array = require_child!(array, array.array_indices(), 0 => Primitive); + array = require_child!(array, array.row_indices(), 1 => Primitive); + for i in 0..num_values { + if array.value(i).as_opt::().is_none() { + array = require_child!(array, array.value(i), i + 2 => Primitive); + } + } + + let validity = array.as_ref().validity()?; + let output = match_each_native_ptype!(array.value(0).dtype().as_ptype(), |T| { + let values = gather_values::(&array)?; + VortexResult::Ok(PrimitiveArray::new(values, validity)) + })?; + + Ok(ExecutionResult::done(output)) +} + +enum PrimitiveSource { + Buffer(Buffer), + Constant { value: T, len: usize }, +} + +impl PrimitiveSource { + fn len(&self) -> usize { + match self { + Self::Buffer(values) => values.len(), + Self::Constant { len, .. } => *len, + } + } + + fn value(&self, index: usize) -> T { + match self { + Self::Buffer(values) => values[index], + Self::Constant { value, .. } => *value, + } + } +} + +fn gather_values(array: &Array) -> VortexResult> { + let values = (0..array.num_values()) + .map(|i| { + let value = array.value(i); + if let Some(constant) = value.as_opt::() { + PrimitiveSource::Constant { + value: constant + .scalar() + .as_primitive() + .typed_value::() + // Validity carries nullness; a null constant's payload is never observed. + .unwrap_or_default(), + len: value.len(), + } + } else { + PrimitiveSource::Buffer(value.as_::().to_buffer::()) + } + }) + .collect::>(); + let branches = array.array_indices().as_::(); + let rows = array.row_indices().as_::(); + + match_each_unsigned_integer_ptype!(branches.ptype(), |A| { + gather_rows::(&values, branches.as_slice::(), rows) + }) +} + +fn gather_rows( + values: &[PrimitiveSource], + branches: &[A], + rows: ArrayView<'_, Primitive>, +) -> VortexResult> +where + T: NativePType, + A: AsPrimitive, +{ + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather(values, branches, rows.as_slice::()) + }) +} + +fn gather( + values: &[PrimitiveSource], + branches: &[A], + rows: &[R], +) -> VortexResult> +where + T: NativePType, + A: AsPrimitive, + R: AsPrimitive, +{ + let len = validate_selectors(values.len(), |branch| values[branch].len(), branches, rows)?; + let mut output = BufferMut::with_capacity(len); + for i in 0..len { + output.push(values[branches[i].as_()].value(rows[i].as_())); + } + Ok(output.freeze()) +} diff --git a/vortex-array/src/arrays/interleave/mod.rs b/vortex-array/src/arrays/interleave/mod.rs index bff03ab055f..d29981249d4 100644 --- a/vortex-array/src/arrays/interleave/mod.rs +++ b/vortex-array/src/arrays/interleave/mod.rs @@ -463,6 +463,7 @@ mod tests { use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; + use crate::dtype::PType; /// Reference (oracle) implementation of the interleave spec, used only to validate the optimized /// [execute](super::execute) path. It is intentionally simple and slow: it pulls each output @@ -719,17 +720,55 @@ mod tests { } #[test] - #[should_panic(expected = "only implemented for boolean values")] - fn non_boolean_value_execution_panics() { - // Execution dispatches on the value type: primitive values have no kernel yet. - let v0 = PrimitiveArray::from_iter([1u32]).into_array(); - let v1 = PrimitiveArray::from_iter([2u32]).into_array(); - let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array(); - let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array(); - let interleaved = InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices) - .vortex_expect("primitive values should construct") - .into_array(); + fn executes_primitive_values() -> VortexResult<()> { + let v0 = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array(); + let v1 = PrimitiveArray::from_option_iter([Some(10.0f64), None]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)?.into_array(); + let expected = + PrimitiveArray::from_option_iter([Some(1.0f64), Some(10.0), Some(2.0), None]) + .into_array(); + let mut ctx = array_session().create_execution_ctx(); + assert_arrays_eq!(interleaved, expected, &mut ctx); + Ok(()) + } + + #[test] + fn executes_primitive_constant_values() -> VortexResult<()> { + let constant = ConstantArray::new(1.0f64, 2).into_array(); + let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)? + .into_array(); + let expected = PrimitiveArray::from_iter([1.0f64, 10.0, 1.0, 20.0]).into_array(); let mut ctx = array_session().create_execution_ctx(); - interleaved.execute::(&mut ctx).ok(); + + assert_arrays_eq!(interleaved, expected, &mut ctx); + Ok(()) + } + + #[test] + fn executes_null_primitive_constant_values() -> VortexResult<()> { + let constant = ConstantArray::new( + Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)), + 2, + ) + .into_array(); + let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array(); + let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(); + let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array(); + let interleaved = + InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)? + .into_array(); + let expected = + PrimitiveArray::from_option_iter([None, Some(10.0f64), None, Some(20.0)]).into_array(); + let mut ctx = array_session().create_execution_ctx(); + + assert_arrays_eq!(interleaved, expected, &mut ctx); + Ok(()) } } From 6ad96f8864cfef46440224fdb4fac789ddff2d00 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 7 Aug 2026 16:20:28 -0400 Subject: [PATCH 2/5] refactor(vortex-array): clarify interleave primitive values Signed-off-by: Nemo Yu --- .../src/arrays/interleave/execute/primitive.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index 1eea5541d10..a42039e1f74 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -46,12 +46,13 @@ pub(super) fn execute( Ok(ExecutionResult::done(output)) } -enum PrimitiveSource { +/// Physical primitive values; nullness remains in the source array's validity. +enum PrimitiveValues { Buffer(Buffer), Constant { value: T, len: usize }, } -impl PrimitiveSource { +impl PrimitiveValues { fn len(&self) -> usize { match self { Self::Buffer(values) => values.len(), @@ -72,7 +73,7 @@ fn gather_values(array: &Array) -> VortexResult() { - PrimitiveSource::Constant { + PrimitiveValues::Constant { value: constant .scalar() .as_primitive() @@ -82,7 +83,7 @@ fn gather_values(array: &Array) -> VortexResult().to_buffer::()) + PrimitiveValues::Buffer(value.as_::().to_buffer::()) } }) .collect::>(); @@ -95,7 +96,7 @@ fn gather_values(array: &Array) -> VortexResult( - values: &[PrimitiveSource], + values: &[PrimitiveValues], branches: &[A], rows: ArrayView<'_, Primitive>, ) -> VortexResult> @@ -109,7 +110,7 @@ where } fn gather( - values: &[PrimitiveSource], + values: &[PrimitiveValues], branches: &[A], rows: &[R], ) -> VortexResult> From ca656000783b87c29eb8bdbddbb272a1c33ed61b Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Mon, 10 Aug 2026 10:04:50 -0400 Subject: [PATCH 3/5] perf(vortex-array): avoid redundant interleave bounds checks Signed-off-by: Nemo Yu --- .../arrays/interleave/execute/primitive.rs | 43 +++++++++++++++++-- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index a42039e1f74..67fb67a9797 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -60,9 +60,15 @@ impl PrimitiveValues { } } - fn value(&self, index: usize) -> T { + /// Returns the physical value at `index` without bounds checking. + /// + /// # Safety + /// + /// `index` must be less than [`Self::len`]. + unsafe fn value_unchecked(&self, index: usize) -> T { match self { - Self::Buffer(values) => values[index], + // SAFETY: the caller guarantees that `index` is in bounds. + Self::Buffer(values) => *unsafe { values.get_unchecked(index) }, Self::Constant { value, .. } => *value, } } @@ -120,9 +126,38 @@ where R: AsPrimitive, { let len = validate_selectors(values.len(), |branch| values[branch].len(), branches, rows)?; + + // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every + // `i < len` that `branches[i] < values.len()` and `rows[i] < values[branches[i]].len()`. + Ok(unsafe { gather_unchecked(len, values, branches, rows) }) +} + +/// Gathers one primitive value per output from `values[branches[i]]` at position `rows[i]`. +/// +/// # Safety +/// +/// `branches` and `rows` must both contain at least `len` elements. For every `i < len`, +/// `branches[i] < values.len()` and `rows[i] < values[branches[i]].len()`. +unsafe fn gather_unchecked( + len: usize, + values: &[PrimitiveValues], + branches: &[A], + rows: &[R], +) -> Buffer +where + T: NativePType, + A: AsPrimitive, + R: AsPrimitive, +{ let mut output = BufferMut::with_capacity(len); for i in 0..len { - output.push(values[branches[i].as_()].value(rows[i].as_())); + // SAFETY: the caller guarantees `i` is in bounds for both selectors, and that the selected + // branch and row are in bounds for `values` and the selected physical value buffer. + output.push(unsafe { + values + .get_unchecked(branches.get_unchecked(i).as_()) + .value_unchecked(rows.get_unchecked(i).as_()) + }); } - Ok(output.freeze()) + output.freeze() } From eadb2df50288e3981238181ed19be884ad02e022 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Mon, 10 Aug 2026 10:24:48 -0400 Subject: [PATCH 4/5] refactor(vortex-array): streamline primitive interleave gather Signed-off-by: Nemo Yu --- .../src/arrays/interleave/execute/mod.rs | 8 +-- .../arrays/interleave/execute/primitive.rs | 66 ++++++++++--------- 2 files changed, 37 insertions(+), 37 deletions(-) diff --git a/vortex-array/src/arrays/interleave/execute/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index e4ac1f5ef43..c43aabe914f 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -18,7 +18,6 @@ use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use super::Interleave; -use super::InterleaveArrayExt; use crate::array::Array; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; @@ -28,15 +27,14 @@ pub(super) fn execute( array: Array, ctx: &mut ExecutionCtx, ) -> VortexResult { - if array.value(0).dtype().is_boolean() { + if array.dtype().is_boolean() { bool::execute(array, ctx) - } else if array.value(0).dtype().is_primitive() { + } else if array.dtype().is_primitive() { primitive::execute(array, ctx) } else { - let value_dtype = array.value(0).dtype().clone(); vortex_panic!( "interleave execution is not implemented for value dtype {}", - value_dtype + array.dtype() ) } } diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index 67fb67a9797..f483bc015be 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -11,6 +11,7 @@ use vortex_error::VortexResult; use super::super::Interleave; use super::super::InterleaveArrayExt; use super::validate_selectors; +use crate::AnyColumnar; use crate::array::Array; use crate::array::ArrayView; use crate::arrays::Constant; @@ -32,13 +33,11 @@ pub(super) fn execute( array = require_child!(array, array.array_indices(), 0 => Primitive); array = require_child!(array, array.row_indices(), 1 => Primitive); for i in 0..num_values { - if array.value(i).as_opt::().is_none() { - array = require_child!(array, array.value(i), i + 2 => Primitive); - } + array = require_child!(array, array.value(i), i + 2 => AnyColumnar); } let validity = array.as_ref().validity()?; - let output = match_each_native_ptype!(array.value(0).dtype().as_ptype(), |T| { + let output = match_each_native_ptype!(array.dtype().as_ptype(), |T| { let values = gather_values::(&array)?; VortexResult::Ok(PrimitiveArray::new(values, validity)) })?; @@ -49,27 +48,21 @@ pub(super) fn execute( /// Physical primitive values; nullness remains in the source array's validity. enum PrimitiveValues { Buffer(Buffer), - Constant { value: T, len: usize }, + Constant(T), } impl PrimitiveValues { - fn len(&self) -> usize { - match self { - Self::Buffer(values) => values.len(), - Self::Constant { len, .. } => *len, - } - } - /// Returns the physical value at `index` without bounds checking. /// /// # Safety /// - /// `index` must be less than [`Self::len`]. + /// For [`Self::Buffer`], `index` must be less than the buffer length. [`Self::Constant`] + /// accepts any index because it has a single physical value. unsafe fn value_unchecked(&self, index: usize) -> T { match self { // SAFETY: the caller guarantees that `index` is in bounds. Self::Buffer(values) => *unsafe { values.get_unchecked(index) }, - Self::Constant { value, .. } => *value, + Self::Constant(value) => *value, } } } @@ -79,15 +72,14 @@ fn gather_values(array: &Array) -> VortexResult() { - PrimitiveValues::Constant { - value: constant + PrimitiveValues::Constant( + constant .scalar() .as_primitive() .typed_value::() // Validity carries nullness; a null constant's payload is never observed. .unwrap_or_default(), - len: value.len(), - } + ) } else { PrimitiveValues::Buffer(value.as_::().to_buffer::()) } @@ -97,11 +89,12 @@ fn gather_values(array: &Array) -> VortexResult(); match_each_unsigned_integer_ptype!(branches.ptype(), |A| { - gather_rows::(&values, branches.as_slice::(), rows) + gather_rows::(array, &values, branches.as_slice::(), rows) }) } fn gather_rows( + array: &Array, values: &[PrimitiveValues], branches: &[A], rows: ArrayView<'_, Primitive>, @@ -111,11 +104,12 @@ where A: AsPrimitive, { match_each_unsigned_integer_ptype!(rows.ptype(), |R| { - gather(values, branches, rows.as_slice::()) + gather(array, values, branches, rows.as_slice::()) }) } fn gather( + array: &Array, values: &[PrimitiveValues], branches: &[A], rows: &[R], @@ -125,10 +119,15 @@ where A: AsPrimitive, R: AsPrimitive, { - let len = validate_selectors(values.len(), |branch| values[branch].len(), branches, rows)?; + let len = validate_selectors( + values.len(), + |branch| array.value(branch).len(), + branches, + rows, + )?; - // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every - // `i < len` that `branches[i] < values.len()` and `rows[i] < values[branches[i]].len()`. + // SAFETY: `validate_selectors` proved both selector lengths and every logical source bound. + // Each `Buffer` has the same length as its source array, while `Constant` ignores the row. Ok(unsafe { gather_unchecked(len, values, branches, rows) }) } @@ -137,7 +136,8 @@ where /// # Safety /// /// `branches` and `rows` must both contain at least `len` elements. For every `i < len`, -/// `branches[i] < values.len()` and `rows[i] < values[branches[i]].len()`. +/// `branches[i] < values.len()` and, when the selected value is a [`PrimitiveValues::Buffer`], +/// `rows[i]` must be less than that buffer's length. unsafe fn gather_unchecked( len: usize, values: &[PrimitiveValues], @@ -150,14 +150,16 @@ where R: AsPrimitive, { let mut output = BufferMut::with_capacity(len); - for i in 0..len { - // SAFETY: the caller guarantees `i` is in bounds for both selectors, and that the selected - // branch and row are in bounds for `values` and the selected physical value buffer. - output.push(unsafe { - values - .get_unchecked(branches.get_unchecked(i).as_()) - .value_unchecked(rows.get_unchecked(i).as_()) - }); + for ((branch, row), slot) in branches.iter().zip(rows).zip(output.spare_capacity_mut()) { + let branch = (*branch).as_(); + let row = (*row).as_(); + // SAFETY: the caller guarantees that the selected branch and row are in bounds for + // `values` and the selected physical value buffer. + slot.write(unsafe { values.get_unchecked(branch).value_unchecked(row) }); } + + // SAFETY: the caller guarantees both selector slices have at least `len` elements, so the loop + // initialized exactly `len` output slots. + unsafe { output.set_len(len) }; output.freeze() } From a0ad6c97de7fc78bc8cfc0461c8255a2d5ad55ad Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Mon, 10 Aug 2026 12:31:47 -0400 Subject: [PATCH 5/5] refactor(vortex-array): simplify interleave selector validation Signed-off-by: Nemo Yu --- vortex-array/benches/interleave.rs | 74 ++++++++- .../src/arrays/interleave/execute/bool.rs | 42 ++++- .../src/arrays/interleave/execute/mod.rs | 38 +---- .../arrays/interleave/execute/primitive.rs | 149 +++++++----------- 4 files changed, 169 insertions(+), 134 deletions(-) diff --git a/vortex-array/benches/interleave.rs b/vortex-array/benches/interleave.rs index d92c033c2a4..a04267c0d92 100644 --- a/vortex-array/benches/interleave.rs +++ b/vortex-array/benches/interleave.rs @@ -1,8 +1,8 @@ // 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: +//! Benchmarks the Vortex [`Interleave`](vortex_array::arrays::Interleave) Boolean and primitive +//! execute paths 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. @@ -26,7 +26,9 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::array_session; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::InterleaveArray; +use vortex_array::arrays::PrimitiveArray; use vortex_buffer::Buffer; fn main() { @@ -130,6 +132,45 @@ fn vortex_inputs(combo: Combo) -> (Vec, Buffer, Buffer) { (values, array_indices, row_indices) } +fn primitive_inputs( + combo: Combo, + mixed_constants: bool, +) -> (Vec, Buffer, Buffer) { + let mut rng = StdRng::seed_from_u64(0); + let values = (0..combo.branches) + .map(|branch| { + if mixed_constants && branch % 2 == 0 { + ConstantArray::new(branch as u64, ARRAY_SIZE).into_array() + } else if combo.nullable { + PrimitiveArray::from_option_iter( + (0..ARRAY_SIZE) + .map(|row| ((row + branch) % 8 != 0).then_some((row ^ branch) as u64)), + ) + .into_array() + } else { + PrimitiveArray::from_iter((0..ARRAY_SIZE).map(|row| (row ^ branch) as u64)) + .into_array() + } + }) + .collect(); + + 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 = (0..ARRAY_SIZE) + .map(|i| match combo.pattern { + Pattern::Random => rng.sample(branch), + Pattern::RoundRobin => u32::try_from(i % combo.branches).unwrap(), + }) + .collect(); + let row_indices = (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 = combos())] fn vortex(bencher: Bencher, combo: Combo) { let (values, array_indices, row_indices) = vortex_inputs(combo); @@ -149,3 +190,32 @@ fn vortex(bencher: Bencher, combo: Combo) { }) .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); } + +fn bench_primitive(bencher: Bencher, combo: Combo, mixed_constants: bool) { + let (values, array_indices, row_indices) = primitive_inputs(combo, mixed_constants); + let session = 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(), + ) + }) + .bench_refs(|(array, ctx)| array.clone().execute::(ctx)); +} + +#[divan::bench(args = combos())] +fn primitive(bencher: Bencher, combo: Combo) { + bench_primitive(bencher, combo, false); +} + +#[divan::bench(args = combos())] +fn primitive_mixed_constants(bencher: Bencher, combo: Combo) { + bench_primitive(bencher, combo, true); +} diff --git a/vortex-array/src/arrays/interleave/execute/bool.rs b/vortex-array/src/arrays/interleave/execute/bool.rs index a051f55ec5d..fde5b161dfd 100644 --- a/vortex-array/src/arrays/interleave/execute/bool.rs +++ b/vortex-array/src/arrays/interleave/execute/bool.rs @@ -7,10 +7,10 @@ use num_traits::AsPrimitive; use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use super::super::Interleave; use super::super::InterleaveArrayExt; -use super::validate_selectors; use crate::array::Array; use crate::arrays::Bool; use crate::arrays::BoolArray; @@ -71,18 +71,46 @@ fn gather, R: AsPrimitive>( branches: &[A], rows: &[R], ) -> VortexResult { - let len = validate_selectors( - value_bits.len(), - |branch| value_bits[branch].len(), - branches, - rows, - )?; + 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 < value_bits.len(), + "interleave array index out of bounds" + ); + vortex_ensure!( + rows[i].as_() < value_bits[branch].len(), + "interleave row index out of bounds" + ); + } + + Ok(len) +} + /// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per /// word with [`BitBufferMut::collect_bool`]. /// diff --git a/vortex-array/src/arrays/interleave/execute/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index c43aabe914f..558db9e3309 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -4,7 +4,7 @@ //! Execution logic for [`Interleave`], dispatched on the value type. //! //! All values share a type (validated in [`Interleave::check`]), so the -//! physical gather kernel is chosen from the first value. The selector types are an orthogonal +//! physical gather kernel is chosen from the array's dtype. The selector types are an orthogonal //! concern handled within each kernel. //! //! [`Interleave::check`]: super::Interleave::check @@ -12,9 +12,7 @@ mod bool; mod primitive; -use num_traits::AsPrimitive; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; use vortex_error::vortex_panic; use super::Interleave; @@ -38,37 +36,3 @@ pub(super) fn execute( ) } } - -/// Validate selector lengths and bounds, returning the common output length. -/// -/// On success, `branches.len() == rows.len() == len`; for every `i < len`, -/// `branches[i] < num_values` and `rows[i] < value_len(branches[i])`. -fn validate_selectors( - num_values: usize, - value_len: F, - branches: &[A], - rows: &[R], -) -> VortexResult -where - A: AsPrimitive, - R: AsPrimitive, - F: Fn(usize) -> usize, -{ - 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!( - rows[i].as_() < value_len(branch), - "interleave row index out of bounds" - ); - } - - Ok(len) -} diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index f483bc015be..e9af70a6e92 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -6,14 +6,15 @@ use num_traits::AsPrimitive; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_buffer::buffer; use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use super::super::Interleave; use super::super::InterleaveArrayExt; -use super::validate_selectors; use crate::AnyColumnar; use crate::array::Array; -use crate::array::ArrayView; use crate::arrays::Constant; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; @@ -45,43 +46,37 @@ pub(super) fn execute( Ok(ExecutionResult::done(output)) } -/// Physical primitive values; nullness remains in the source array's validity. -enum PrimitiveValues { - Buffer(Buffer), - Constant(T), -} - -impl PrimitiveValues { - /// Returns the physical value at `index` without bounds checking. - /// - /// # Safety - /// - /// For [`Self::Buffer`], `index` must be less than the buffer length. [`Self::Constant`] - /// accepts any index because it has a single physical value. - unsafe fn value_unchecked(&self, index: usize) -> T { - match self { - // SAFETY: the caller guarantees that `index` is in bounds. - Self::Buffer(values) => *unsafe { values.get_unchecked(index) }, - Self::Constant(value) => *value, - } - } +/// Physical storage for one primitive source array. +struct PrimitiveSource { + data: Buffer, + len: usize, + // Maps logical rows to `data`: zero for constants and identity for decoded buffers. + row_mask: usize, } fn gather_values(array: &Array) -> VortexResult> { let values = (0..array.num_values()) .map(|i| { let value = array.value(i); + let len = value.len(); if let Some(constant) = value.as_opt::() { - PrimitiveValues::Constant( - constant - .scalar() - .as_primitive() - .typed_value::() - // Validity carries nullness; a null constant's payload is never observed. - .unwrap_or_default(), - ) + // Validity carries nullness; a null constant's payload is never observed. + let payload = constant + .scalar() + .as_primitive() + .typed_value::() + .unwrap_or_default(); + PrimitiveSource { + data: buffer![payload], + len, + row_mask: 0, + } } else { - PrimitiveValues::Buffer(value.as_::().to_buffer::()) + PrimitiveSource { + data: value.as_::().to_buffer::(), + len, + row_mask: usize::MAX, + } } }) .collect::>(); @@ -89,28 +84,14 @@ fn gather_values(array: &Array) -> VortexResult(); match_each_unsigned_integer_ptype!(branches.ptype(), |A| { - gather_rows::(array, &values, branches.as_slice::(), rows) - }) -} - -fn gather_rows( - array: &Array, - values: &[PrimitiveValues], - branches: &[A], - rows: ArrayView<'_, Primitive>, -) -> VortexResult> -where - T: NativePType, - A: AsPrimitive, -{ - match_each_unsigned_integer_ptype!(rows.ptype(), |R| { - gather(array, values, branches, rows.as_slice::()) + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather(&values, branches.as_slice::(), rows.as_slice::()) + }) }) } fn gather( - array: &Array, - values: &[PrimitiveValues], + values: &[PrimitiveSource], branches: &[A], rows: &[R], ) -> VortexResult> @@ -119,47 +100,39 @@ where A: AsPrimitive, R: AsPrimitive, { - let len = validate_selectors( - values.len(), - |branch| array.value(branch).len(), - branches, - rows, - )?; + // `zip` truncates to the shorter input. + vortex_ensure!( + rows.len() == branches.len(), + "interleave selectors differ in length: array_indices {}, row_indices {}", + branches.len(), + rows.len() + ); - // SAFETY: `validate_selectors` proved both selector lengths and every logical source bound. - // Each `Buffer` has the same length as its source array, while `Constant` ignores the row. - Ok(unsafe { gather_unchecked(len, values, branches, rows) }) + let output = + BufferMut::try_from_trusted_len_iter(branches.iter().zip(rows).map(|(branch, row)| { + let Some(source) = values.get((*branch).as_()) else { + vortex_bail!("interleave array index out of bounds"); + }; + let row = (*row).as_(); + vortex_ensure!(row < source.len, "interleave row index out of bounds"); + Ok(source.data[row & source.row_mask]) + }))?; + Ok(output.freeze()) } -/// Gathers one primitive value per output from `values[branches[i]]` at position `rows[i]`. -/// -/// # Safety -/// -/// `branches` and `rows` must both contain at least `len` elements. For every `i < len`, -/// `branches[i] < values.len()` and, when the selected value is a [`PrimitiveValues::Buffer`], -/// `rows[i]` must be less than that buffer's length. -unsafe fn gather_unchecked( - len: usize, - values: &[PrimitiveValues], - branches: &[A], - rows: &[R], -) -> Buffer -where - T: NativePType, - A: AsPrimitive, - R: AsPrimitive, -{ - let mut output = BufferMut::with_capacity(len); - for ((branch, row), slot) in branches.iter().zip(rows).zip(output.spare_capacity_mut()) { - let branch = (*branch).as_(); - let row = (*row).as_(); - // SAFETY: the caller guarantees that the selected branch and row are in bounds for - // `values` and the selected physical value buffer. - slot.write(unsafe { values.get_unchecked(branch).value_unchecked(row) }); - } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_out_of_bounds_selectors() { + let values = [PrimitiveSource { + data: buffer![1u32], + len: 1, + row_mask: 0, + }]; - // SAFETY: the caller guarantees both selector slices have at least `len` elements, so the loop - // initialized exactly `len` output slots. - unsafe { output.set_len(len) }; - output.freeze() + assert!(gather(&values, &[1u8], &[0u8]).is_err()); + assert!(gather(&values, &[0u8], &[1u8]).is_err()); + } }