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/mod.rs b/vortex-array/src/arrays/interleave/execute/mod.rs index 05dcd161f62..558db9e3309 100644 --- a/vortex-array/src/arrays/interleave/execute/mod.rs +++ b/vortex-array/src/arrays/interleave/execute/mod.rs @@ -4,19 +4,18 @@ //! 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 -//! concern handled within each kernel. Only boolean values are implemented today (see the [`bool`] module). +//! 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 -//! [`bool`]: module@crate::arrays::interleave::execute::bool mod bool; +mod primitive; use vortex_error::VortexResult; use vortex_error::vortex_panic; use super::Interleave; -use super::InterleaveArrayExt; use crate::array::Array; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; @@ -26,14 +25,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.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", - value_dtype + "interleave execution is not implemented for value dtype {}", + array.dtype() ) } } 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..e9af70a6e92 --- /dev/null +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -0,0 +1,138 @@ +// 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_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 crate::AnyColumnar; +use crate::array::Array; +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 { + array = require_child!(array, array.value(i), i + 2 => AnyColumnar); + } + + let validity = array.as_ref().validity()?; + let output = match_each_native_ptype!(array.dtype().as_ptype(), |T| { + let values = gather_values::(&array)?; + VortexResult::Ok(PrimitiveArray::new(values, validity)) + })?; + + Ok(ExecutionResult::done(output)) +} + +/// 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::() { + // 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 { + PrimitiveSource { + data: value.as_::().to_buffer::(), + len, + row_mask: usize::MAX, + } + } + }) + .collect::>(); + let branches = array.array_indices().as_::(); + let rows = array.row_indices().as_::(); + + match_each_unsigned_integer_ptype!(branches.ptype(), |A| { + match_each_unsigned_integer_ptype!(rows.ptype(), |R| { + gather(&values, branches.as_slice::(), rows.as_slice::()) + }) + }) +} + +fn gather( + values: &[PrimitiveSource], + branches: &[A], + rows: &[R], +) -> VortexResult> +where + T: NativePType, + A: AsPrimitive, + R: AsPrimitive, +{ + // `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() + ); + + 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()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_out_of_bounds_selectors() { + let values = [PrimitiveSource { + data: buffer![1u32], + len: 1, + row_mask: 0, + }]; + + assert!(gather(&values, &[1u8], &[0u8]).is_err()); + assert!(gather(&values, &[0u8], &[1u8]).is_err()); + } +} 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(()) } }