diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 16f2db98554..a692037e282 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -135,6 +135,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_output" +harness = false + [[bench]] name = "kleene_bool" harness = false diff --git a/vortex-array/benches/row_fn_output.rs b/vortex-array/benches/row_fn_output.rs new file mode 100644 index 00000000000..6ba02b6df8c --- /dev/null +++ b/vortex-array/benches/row_fn_output.rs @@ -0,0 +1,129 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect( + clippy::cast_possible_truncation, + reason = "benchmark fixtures reduce values below 1024 before narrowing" +)] + +use std::marker::PhantomData; +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::VecExecutionArgs; +use vortex_array::scalar_fn::unstable::row::RowFn; +use vortex_array::scalar_fn::unstable::row::RowVisitor; +use vortex_array::scalar_fn::unstable::row::execute_rows; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +static SESSION: LazyLock = LazyLock::new(array_session); + +const ROWS: usize = 1 << 14; + +const INPUT_SHAPES: &[InputShape] = &[ + InputShape::PerRowPerRow, + InputShape::PerRowConstant, + InputShape::ConstantPerRow, +]; + +#[derive(Clone, Copy, Debug)] +enum InputShape { + PerRowPerRow, + PerRowConstant, + ConstantPerRow, +} + +trait BenchPrimitive: NativePType { + fn per_row(offset: usize) -> ArrayRef; + + fn constant() -> ArrayRef; +} + +impl BenchPrimitive for i32 { + fn per_row(offset: usize) -> ArrayRef { + PrimitiveArray::from_iter((0..ROWS).map(|index| ((index + offset) % 1024) as i32)) + .into_array() + } + + fn constant() -> ArrayRef { + ConstantArray::new(512_i32, ROWS).into_array() + } +} + +impl BenchPrimitive for i64 { + fn per_row(offset: usize) -> ArrayRef { + PrimitiveArray::from_iter((0..ROWS).map(|index| ((index + offset) % 1024) as i64)) + .into_array() + } + + fn constant() -> ArrayRef { + ConstantArray::new(512_i64, ROWS).into_array() + } +} + +#[derive(Clone)] +struct InfallibleBool(PhantomData); + +impl RowFn for InfallibleBool { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_fn_output.infallible_bool"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs)) + } +} + +#[vortex_bench_support::cpu_features] +#[divan::bench(types = [i32, i64], args = INPUT_SHAPES)] +fn infallible_bool(bencher: Bencher, &shape: &InputShape) { + let function = InfallibleBool::(PhantomData); + let args = match shape { + InputShape::PerRowPerRow => vec![T::per_row(0), T::per_row(1)], + InputShape::PerRowConstant => vec![T::per_row(0), T::constant()], + InputShape::ConstantPerRow => vec![T::constant(), T::per_row(1)], + }; + let args = VecExecutionArgs::new(args, ROWS); + + bencher + .counter(ItemsCount::new(ROWS)) + .with_inputs(|| (&args, SESSION.create_execution_ctx())) + .bench_refs(|(args, ctx)| { + execute_rows(&function, &EmptyOptions, *args, ctx) + .vortex_expect("row execution should succeed in benchmark") + }); +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 25a260846c6..99e2c074373 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -6,6 +6,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use rstest::rstest; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -15,9 +16,11 @@ use vortex_session::registry::CachedId; use super::finalize_kernel_output; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::ExtensionArray; use crate::arrays::FixedSizeListArray; @@ -35,6 +38,7 @@ use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::FixedSizeListSink; use crate::scalar_fn::unstable::row::InitializedRow; +use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; @@ -61,6 +65,51 @@ struct ValidOnlyIdentity; #[derive(Clone)] struct InvalidKernelOutput; +#[derive(Clone)] +struct PackedPositive; + +#[derive(Clone)] +struct ValidOnlyPositive; + +#[derive(Clone)] +struct NullaryTrue; + +struct ValidOnlyI64; + +// SAFETY: the view is a slice, and its reported length is the buffer length. +unsafe impl InputElement for ValidOnlyI64 { + type Column = Buffer; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = false; + const DECODE_INFALLIBLE: bool = true; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + ::decode(array, ctx) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_> { + column[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> { + view[index] + } +} + /// Produces a null row to exercise output validation at the row-function boundary. #[derive(Default)] struct NullProducingI64(i64); @@ -274,6 +323,72 @@ fn test_fixed_size_list_sink_uses_runtime_width( Ok(()) } +impl RowFn for PackedPositive { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.packed_positive"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), bool>(|(value,)| value > 0) + } +} + +impl RowFn for ValidOnlyPositive { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.valid_only_positive"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(ValidOnlyI64,), bool>(|(value,)| { + assert_ne!(value, i64::MIN, "an invalid row was evaluated"); + value > 0 + }) + } +} + +impl RowFn for NullaryTrue { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_true"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(), bool>(|()| true) + } +} + #[test] fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<()> { static ID: CachedId = CachedId::new("test.finalize_kernel_output"); @@ -320,6 +435,73 @@ fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { Ok(()) } +#[test] +fn test_bool_output_builds_packed_values() -> VortexResult<()> { + let input = PrimitiveArray::new( + vec![1_i64, -1, 2, 0, 3], + Validity::from_iter([true, true, false, true, true]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input], 5); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&PackedPositive, &EmptyOptions, &args, &mut ctx)?; + let expected = + BoolArray::from_iter([Some(true), Some(false), None, Some(false), Some(true)]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::empty(0)] +#[case::one(1)] +#[case::below_word(63)] +#[case::exact_word(64)] +#[case::above_word(65)] +#[case::multiword_remainder(130)] +fn test_bool_output_word_boundaries(#[case] len: usize) -> VortexResult<()> { + let values: Vec<_> = (0..len).map(|index| index as i64 - 32).collect(); + let input = PrimitiveArray::from_iter(values.iter().copied()).into_array(); + let args = VecExecutionArgs::new(vec![input], len); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&PackedPositive, &EmptyOptions, &args, &mut ctx)?; + let expected = BoolArray::from_iter(values.iter().map(|value| *value > 0)).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_bool_output_handles_nullary_rows() -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], 65); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&NullaryTrue, &EmptyOptions, &args, &mut ctx)?; + let expected = BoolArray::from_iter(std::iter::repeat_n(true, 65)).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_bool_output_skips_invalid_rows() -> VortexResult<()> { + let input = PrimitiveArray::new( + vec![1_i64, i64::MIN, -1], + Validity::from_iter([true, false, true]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&ValidOnlyPositive, &EmptyOptions, &args, &mut ctx)?; + let expected = BoolArray::from_iter([Some(true), None, Some(false)]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + #[test] fn test_deferred_owned_execution_retries_null_row_failure() -> VortexResult<()> { let function = DeferredAdd::default(); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index 7e73c183b0c..0a8c250ce1f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -3,9 +3,9 @@ //! Executes row kernels that return one independent owned value per row. //! -//! [`execute_owned`] decodes inputs once, prepares constant state, writes into spare vector -//! capacity, and reduces compact failure evidence without putting error construction in the hot -//! loop. [`execute_owned_infallible`] removes that failure path for infallible kernels. +//! [`execute_owned`] writes fallible row results into spare vector capacity and reduces compact +//! failure evidence outside the hot loop. [`execute_owned_infallible`] lets the output type map a +//! validated non-constant row source directly into its physical representation. use std::ops::BitOrAssign; @@ -42,13 +42,44 @@ where Args: IndexedElementTuple, Out: OutputElement, { - execute_owned::( - args, - ctx, - prepare, - move |prepared, args| (apply(prepared, args), NoFailure), - |_| Ok(()), - ) + const { assert_owned_output_needs_no_drop::() }; + + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::const_values(&columns)); + let row_count = args.row_count(); + + if let Some(views) = Args::views_if_no_consts(&columns) { + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` checked that these exact retained views address `row_count` + // rows. + let source = unsafe { Args::indexed_source(views, row_count) }; + + return Ok(Out::build_from(source, |elements| { + apply(&prepared, elements) + })); + } + + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut values = Vec::::with_capacity(row_count); + let output = &mut values.spare_capacity_mut()[..row_count]; + + for (index, slot) in output.iter_mut().enumerate() { + slot.write(apply(&prepared, Args::get(&columns, index))); + } + + // SAFETY: normal completion initializes every output slot exactly once, and `values` was + // allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + Ok(Out::build(values)) } /// Decode nullable inputs, then store one output for each valid row from an infallible kernel. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs index 323880148b5..26028335a37 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_buffer::BitBuffer; +use vortex_compute::lane_kernels::IndexedSource; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -74,4 +75,19 @@ impl OutputElement for bool { // `From>` uses the bulk bit-packing path. BoolArray::new(BitBuffer::from(values), Validity::NonNullable).into_array() } + + fn build_from(source: S, apply: F) -> ArrayRef + where + S: IndexedSource, + F: Fn(S::Item) -> Self, + { + let len = source.len(); + let values = BitBuffer::collect_bool(len, |index| { + // SAFETY: `collect_bool` only invokes this closure with `index < len`, and + // `len` is `source.len()`. + apply(unsafe { source.get_unchecked(index) }) + }); + + BoolArray::new(values, Validity::NonNullable).into_array() + } } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs index 1034bd46498..3b92107931b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs @@ -5,6 +5,9 @@ //! //! [`OutputElement`] describes fixed-dtype values returned independently by each row invocation. +use vortex_compute::lane_kernels::IndexedSource; +use vortex_compute::lane_kernels::IndexedSourceExt; + use crate::ArrayRef; use crate::dtype::DType; @@ -25,7 +28,36 @@ pub trait OutputElement: 'static + Sized + Default { /// Build an all-valid column from one value per row. /// /// The returned column must contain `values.len()` rows and match - /// [`element_dtype`](Self::element_dtype) except for outer nullability. The framework calls - /// this method once per batch. + /// [`element_dtype`](Self::element_dtype) except for outer nullability. The default + /// [`build_from`](Self::build_from) implementation and valid-row execution call this method. fn build(values: Vec) -> ArrayRef; + + /// Map a contiguous row source directly into an all-valid column. + /// + /// The default collects one value per row into a [`Vec`] before calling [`build`](Self::build). + /// An output type can override this method when its physical representation supports a more + /// efficient bulk mapping. The implementation **must** call `apply` exactly once for every + /// source row in increasing order and return the same values as the default implementation. + /// + /// An override **must not** introduce value-dependent errors or panics. Fallible operations + /// must use a fallible visitor path so that [`RowFn::INFALLIBLE`] continues to protect optimizer + /// transformations. + /// + /// [`RowFn::INFALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::INFALLIBLE + fn build_from(source: S, apply: F) -> ArrayRef + where + S: IndexedSource, + F: Fn(S::Item) -> Self, + { + let row_count = source.len(); + let mut values = Vec::::with_capacity(row_count); + let output = &mut values.spare_capacity_mut()[..row_count]; + + source.map_into(output, apply); + + // SAFETY: normal completion of `map_into` initializes every output slot exactly once. + unsafe { values.set_len(row_count) }; + + Self::build(values) + } }