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 78dcb554a07..ab98c9480e0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -12,9 +13,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::PrimitiveArray; use crate::assert_arrays_eq; @@ -25,10 +28,17 @@ use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::ArgView; +use crate::scalar_fn::unstable::row::DenseRows; +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::PackedBoolOutput; use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowKernel; +use crate::scalar_fn::unstable::row::RowKernelOutput; use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::VecOutput; use crate::scalar_fn::unstable::row::execute_rows; use crate::validity::Validity; @@ -50,6 +60,205 @@ struct ValidOnlyIdentity; #[derive(Clone)] struct InvalidKernelOutput; +#[derive(Clone)] +struct PackedPositive; + +struct PackedPositiveKernel; + +#[derive(Clone)] +struct WrongLengthKernelOutput; + +struct WrongLengthKernel; + +struct WrongLengthOutput { + row_count: usize, +} + +#[derive(Clone)] +struct WrongDTypeKernelOutput; + +struct WrongDTypeKernel; + +struct WrongDTypeOutput { + row_count: usize, +} + +#[derive(Clone)] +struct RetainedViewIdentity; + +struct RetainedViewIdentityKernel; + +struct ChangingViewI64; + +struct ChangingViewColumn { + values: Buffer, + view_count: AtomicUsize, +} + +#[derive(Clone)] +struct ValidOnlyAssociatedOutput; + +struct ValidOnlyAssociatedOutputKernel; + +struct CountingBoolOutput(Vec); + +static ASSOCIATED_OUTPUT_CONVERSIONS: AtomicUsize = AtomicUsize::new(0); + +// SAFETY: each returned slice is stable for its lifetime, and unchecked access is valid below that +// slice's length. Returning a shorter slice on later calls exercises the executor's obligation to +// retain the exact view whose length it validated. +unsafe impl InputElement for ChangingViewI64 { + type Column = ChangingViewColumn; + 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 { + Ok(ChangingViewColumn { + values: ::decode(array, ctx)?, + view_count: AtomicUsize::new(0), + }) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_> { + column.values[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + let values = column.values.as_slice(); + if column.view_count.fetch_add(1, Ordering::Relaxed) == 0 { + values + } else { + &values[..0] + } + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> { + view[index] + } +} + +impl RowKernel<(ChangingViewI64,)> for RetainedViewIdentityKernel { + type Element = i64; + type Output = VecOutput; + + fn eval(&self, (value,): (i64,)) -> Self::Element { + value + } +} + +impl RowKernel<(ChangingViewI64,)> for ValidOnlyAssociatedOutputKernel { + type Element = bool; + type Output = CountingBoolOutput; + + fn eval(&self, (value,): (i64,)) -> Self::Element { + value > 0 + } +} + +impl RowKernelOutput for CountingBoolOutput { + type Element = bool; + + fn from_values(values: Vec) -> VortexResult { + ASSOCIATED_OUTPUT_CONVERSIONS.fetch_add(1, Ordering::Relaxed); + Ok(Self(values)) + } + + fn finish(self) -> VortexResult { + Ok(BoolArray::from_iter(self.0).into_array()) + } +} + +impl RowKernel<(i64,)> for PackedPositiveKernel { + type Element = bool; + type Output = PackedBoolOutput; + + fn eval(&self, (value,): (i64,)) -> Self::Element { + value > 0 + } + + fn collect_dense(&self, rows: DenseRows<'_, (i64,)>) -> VortexResult { + let mut output = PackedBoolOutput::zeroed(rows.len()); + + match rows.inputs().0.view() { + ArgView::Column(values) => { + for (word_index, values) in values.chunks(64).enumerate() { + output.words_mut()[word_index] = + values.iter().enumerate().fold(0, |word, (bit, value)| { + word | (u64::from(*value > 0) << bit) + }); + } + } + ArgView::Constant(value) => { + if value[0] > 0 { + output.words_mut().fill(u64::MAX); + } + } + } + + Ok(output) + } +} + +impl RowKernel<(i64,)> for WrongLengthKernel { + type Element = bool; + type Output = WrongLengthOutput; + + fn eval(&self, (_value,): (i64,)) -> Self::Element { + true + } +} + +impl RowKernelOutput for WrongLengthOutput { + type Element = bool; + + fn from_values(values: Vec) -> VortexResult { + Ok(Self { + row_count: values.len(), + }) + } + + fn finish(self) -> VortexResult { + Ok( + BoolArray::from_iter(std::iter::repeat_n(true, self.row_count.saturating_sub(1))) + .into_array(), + ) + } +} + +impl RowKernel<(i64,)> for WrongDTypeKernel { + type Element = bool; + type Output = WrongDTypeOutput; + + fn eval(&self, (_value,): (i64,)) -> Self::Element { + true + } +} + +impl RowKernelOutput for WrongDTypeOutput { + type Element = bool; + + fn from_values(values: Vec) -> VortexResult { + Ok(Self { + row_count: values.len(), + }) + } + + fn finish(self) -> VortexResult { + Ok(PrimitiveArray::from_iter(std::iter::repeat_n(1_i64, self.row_count)).into_array()) + } +} + /// Produces a null row to exercise output validation at the row-function boundary. #[derive(Default)] struct NullProducingI64(i64); @@ -178,6 +387,111 @@ impl RowFn for InvalidKernelOutput { } } +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_kernel::<(i64,), _>(PackedPositiveKernel) + } +} + +impl RowFn for WrongLengthKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.wrong_length_kernel_output"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_kernel::<(i64,), _>(WrongLengthKernel) + } +} + +impl RowFn for WrongDTypeKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.wrong_dtype_kernel_output"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_kernel::<(i64,), _>(WrongDTypeKernel) + } +} + +impl RowFn for RetainedViewIdentity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.retained_view_identity"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_kernel::<(ChangingViewI64,), _>(RetainedViewIdentityKernel) + } +} + +impl RowFn for ValidOnlyAssociatedOutput { + 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_associated_output"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_kernel::<(ChangingViewI64,), _>(ValidOnlyAssociatedOutputKernel) + } +} + #[test] fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<()> { static ID: CachedId = CachedId::new("test.finalize_kernel_output"); @@ -224,6 +538,111 @@ fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { Ok(()) } +#[test] +fn test_dense_kernel_writes_packed_output_words() -> 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(()) +} + +#[test] +fn test_dense_kernel_handles_batch_constant_input() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 65).into_array(); + let args = VecExecutionArgs::new(vec![input], 65); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&PackedPositive, &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_dense_kernel_retains_the_validated_views() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&RetainedViewIdentity, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_kernel_uses_its_associated_output() -> VortexResult<()> { + ASSOCIATED_OUTPUT_CONVERSIONS.store(0, Ordering::Relaxed); + + let input = PrimitiveArray::new(vec![1_i64, -1, 2], 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(&ValidOnlyAssociatedOutput, &EmptyOptions, &args, &mut ctx)?; + let expected = BoolArray::from_iter([Some(true), None, Some(true)]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(ASSOCIATED_OUTPUT_CONVERSIONS.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn test_dense_kernel_rejects_wrong_output_length() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&WrongLengthKernelOutput, &EmptyOptions, &args, &mut ctx) { + Err(error) => error.to_string(), + Ok(_) => vortex_bail!("a RowKernelOutput with the wrong length passed validation"), + }; + + assert!( + error.contains("test.wrong_length_kernel_output"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("must contain 2 rows, got 1"), + "the boundary error must identify the wrong output length, got {error}", + ); + Ok(()) +} + +#[test] +fn test_dense_kernel_rejects_wrong_output_dtype() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&WrongDTypeKernelOutput, &EmptyOptions, &args, &mut ctx) { + Err(error) => error.to_string(), + Ok(_) => vortex_bail!("a RowKernelOutput with the wrong dtype passed validation"), + }; + + assert!( + error.contains("test.wrong_dtype_kernel_output"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("output dtype must match bool") && error.contains("got i64"), + "the boundary error must identify the wrong output dtype, got {error}", + ); + 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/kernel.rs b/vortex-array/src/scalar_fn/unstable/row/execute/kernel.rs new file mode 100644 index 00000000000..2de459d3fc6 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/kernel.rs @@ -0,0 +1,97 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution for typed infallible row kernels. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::MaskValuesRef; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::DenseRows; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::RowKernel; +use crate::scalar_fn::unstable::row::RowKernelOutput; +use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop; + +/// Decode and validate the inputs, then delegate only dense collection to the kernel. +pub(in crate::scalar_fn::unstable::row) fn execute_kernel( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + kernel: Kernel, +) -> VortexResult +where + Args: IndexedElementTuple, + Kernel: RowKernel, +{ + let inputs = Args::decode(args, ctx)?; + let rows = DenseRows::::new(&inputs, args.row_count())?; + let output = kernel.collect_dense(rows)?; + + output.finish() +} + +/// Decode nullable inputs, then store one kernel output for each valid row. +pub(in crate::scalar_fn::unstable::row) fn execute_kernel_valid_rows( + args: &dyn ExecutionArgs, + valid: &MaskValuesRef, + ctx: &mut ExecutionCtx, + kernel: Kernel, +) -> VortexResult> +where + Args: IndexedElementTuple, + Kernel: RowKernel, +{ + const { assert_owned_output_needs_no_drop::() }; + + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + + let row_count = args.row_count(); + let valid_rows = valid.bit_buffer(); + vortex_ensure_eq!( + valid_rows.len(), + row_count, + "the validity mask must address exactly {row_count} rows, got {}", + valid_rows.len(), + ); + + let mut values: Vec = std::iter::repeat_with(Kernel::Element::default) + .take(row_count) + .collect(); + + 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", + ); + + valid_rows.for_each_set_index(|index| { + // SAFETY: the tuple-wide length check proved every view has `row_count` rows, and mask + // indices are below `row_count`. Nullary tuples do not access an input view. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + + // SAFETY: the mask length check proved that every set index is below `row_count`. + unsafe { *values.get_unchecked_mut(index) = kernel.eval(elements) }; + }); + } else { + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + valid_rows.for_each_set_index(|index| { + let value = kernel.eval(Args::get(&columns, index)); + + // SAFETY: the mask length check proved that every set index is below `row_count`. + unsafe { *values.get_unchecked_mut(index) = value }; + }); + } + + let output = Kernel::Output::from_values(values)?; + output.finish().map(Some) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs index 3ed0b4801f4..c552507a74d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -6,6 +6,10 @@ //! [`owned`] stores one independent value per row and reduces compact failure evidence. [`sink`] //! drives output builders whose row handles may share batch state. +mod kernel; +pub(super) use kernel::execute_kernel; +pub(super) use kernel::execute_kernel_valid_rows; + mod owned; pub(super) use owned::execute_owned; pub(super) use owned::execute_owned_infallible; diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index bd73916c85c..d08cac2997b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -9,7 +9,10 @@ //! A [`RowFn`] describes the typed operation while the framework owns columnar concerns such as //! decoding, constant handling, null propagation, allocation, and validity. Its //! [`RowFn::dispatch`] implementation uses a [`RowVisitor`] to select an [`ElementTuple`] and -//! either an [`OutputElement`] or [`OutputSink`] for each supported dtype combination. +//! output contract for each supported dtype combination. An [`OutputElement`] returns one owned +//! value per row, an [`OutputSink`] writes through a shared builder, and a [`RowKernel`] can attach +//! an associated dense output representation while retaining scalar semantics for other execution +//! policies. //! //! Unlike a general strict function, a [`RowFn`] cannot produce null from valid inputs. //! @@ -29,6 +32,9 @@ mod row_fn; pub use row_fn::RowFn; mod types; +pub use types::ArgColumn; +pub use types::ArgView; +pub use types::DenseRows; pub use types::ElementTuple; pub use types::FailureEvidence; pub use types::IndexedElementTuple; @@ -36,8 +42,12 @@ pub use types::InitializedElement; pub use types::InputElement; pub use types::OutputElement; pub use types::OutputSink; +pub use types::PackedBoolOutput; +pub use types::RowKernel; +pub use types::RowKernelOutput; pub use types::SinkResult; pub use types::UninitElementSink; +pub use types::VecOutput; pub use types::ViewLen; mod visitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index e456886e618..3bf9de41066 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -26,10 +26,12 @@ use crate::scalar_fn::ScalarFnId; /// propagation but permits valid inputs to produce null. The framework derives output validity /// only from input validity. /// -/// A dispatched [`OutputElement`] or [`OutputSink`] describes the non-nullable values produced for -/// valid rows. The framework widens that dtype when an input dtype is nullable, attaches the -/// input-derived validity, and casts the finished array to the widened dtype. Implementations do -/// not construct nullable placeholders for invalid rows. +/// A dispatched [`OutputElement`], [`OutputSink`], or [`RowKernel`] describes the non-nullable +/// values produced for valid rows. A `RowKernel` also selects an associated output representation +/// and can override dense collection without changing its scalar semantics. The framework widens +/// the output dtype when an input dtype is nullable, attaches the input-derived validity, and casts +/// the finished array to the widened dtype. Implementations do not construct nullable placeholders +/// for invalid rows. /// /// Declare argument names and use [`dispatch`](Self::dispatch) to select element and output types. /// Every implementation receives the standard [`ScalarFnVTable`]. A public type that needs custom @@ -37,6 +39,7 @@ use crate::scalar_fn::ScalarFnId; /// /// [`OutputElement`]: crate::scalar_fn::unstable::row::OutputElement /// [`OutputSink`]: crate::scalar_fn::unstable::row::OutputSink +/// [`RowKernel`]: crate::scalar_fn::unstable::row::RowKernel /// [`ScalarFnVTable`]: crate::scalar_fn::ScalarFnVTable /// [`execute_rows`]: crate::scalar_fn::unstable::row::execute_rows /// [`row_fn_return_dtype`]: crate::scalar_fn::unstable::row::row_fn_return_dtype @@ -73,7 +76,8 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync { vortex_bail!("Expression {} is not deserializable", self.id()) } - /// Choose element types for these input dtypes and visit the framework with them. + /// Choose element types and an output contract for these input dtypes, then visit the + /// framework with them. /// /// Planning and execution both call this method, so its result **must** depend only on /// `options` and `args`. Cross-argument dtype validation belongs here. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 148f784bcba..38925b4dde0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -59,6 +59,18 @@ pub unsafe trait InputElement: 'static { /// and other invocation-invariant work into this method. fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + /// Decode one representative row from an input that is constant for the batch. + /// + /// The default preserves the ordinary decode contract by slicing the input to one row first. + /// Implementations can override this when their constant representation supports cheaper + /// scalar extraction. The returned column **must** contain exactly one row. + fn decode_batch_constant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Self::decode(array.slice(0..1)?, ctx) + } + /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array. /// /// The conservative default declines. An implementation whose ordinary decode is safe and diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 73a0c1d47d1..9c3d229c63e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -18,6 +18,8 @@ pub use output::OutputElement; mod primitive; mod tuple; +pub use tuple::ArgColumn; +pub use tuple::ArgView; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; pub use tuple::batch_const; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs index ed5b7b0888e..99c59684329 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs @@ -9,6 +9,7 @@ use vortex_error::vortex_ensure_eq; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; +use crate::arrays::Constant; use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::dtype::NativePType; @@ -44,6 +45,19 @@ unsafe impl InputElement for T { Ok(array.execute::(ctx)?.into_buffer::()) } + fn decode_batch_constant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if let Some(constant) = array.as_opt::() + && let Some(value) = constant.scalar().as_primitive().try_typed_value::()? + { + return Ok(Buffer::full(value, 1)); + } + + Self::decode(array.slice(0..1)?, ctx) + } + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { Ok(true) } diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index b251ede1575..236e9895daa 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -35,6 +35,17 @@ enum ArgColumnKind { Const(T::Column), } +/// A typed decoded input classified by how the dense collector addresses it. +/// +/// This classification does not prove that the view covers a particular batch length. +pub enum ArgView<'a, T: InputElement> { + /// A view with one value for every row in the batch. + Column(T::View<'a>), + + /// A one-row view whose value is broadcast across the batch. + Constant(T::View<'a>), +} + impl ArgColumn { fn try_from_const(column: T::Column) -> VortexResult { let decoded_len = T::view(&column).len(); @@ -52,7 +63,7 @@ impl ArgColumn { if let Some(const_array) = batch_const(&array) && !array.is_empty() { - return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?); + return Self::try_from_const(T::decode_batch_constant(const_array, ctx)?); } Ok(Self(ArgColumnKind::Column(T::decode(array, ctx)?))) @@ -64,7 +75,7 @@ impl ArgColumn { if let Some(const_array) = batch_const(&array) && !array.is_empty() { - return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?).map(Some); + return Self::try_from_const(T::decode_batch_constant(const_array, ctx)?).map(Some); } Ok(T::decode_null_tolerant(array, ctx)? @@ -96,6 +107,14 @@ impl ArgColumn { } } + /// Borrow the decoded input with its batch-constant classification. + pub fn view(&self) -> ArgView<'_, T> { + match &self.0 { + ArgColumnKind::Column(column) => ArgView::Column(T::view(column)), + ArgColumnKind::Const(column) => ArgView::Constant(T::view(column)), + } + } + fn addresses_rows(&self, row_count: usize) -> bool { // A constant is validated when constructed and is always read at index zero. match &self.0 { diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index cc312a3fb14..09ab9fb0683 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -7,6 +7,8 @@ //! [`IndexedElementTuple`] adds the validated indexed source used by vectorizable dense loops. mod element_tuple; +pub use element_tuple::ArgColumn; +pub use element_tuple::ArgView; pub use element_tuple::ElementTuple; pub use element_tuple::batch_const; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/kernel.rs b/vortex-array/src/scalar_fn/unstable/row/types/kernel.rs new file mode 100644 index 00000000000..58067b9337d --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/kernel.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Typed infallible row kernels with an optional dense collector. +//! +//! [`RowKernel::eval`] defines the operation's scalar semantics. The executor uses that method for +//! validity-aware traversal and converts the values through the associated [`RowKernelOutput`]. +//! Dense execution can override [`RowKernel::collect_dense`] with a representation-specific bulk +//! kernel. + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::RowKernelOutput; + +/// A validated dense batch of typed row inputs. +/// +/// The executor decodes the columns and retains the exact views whose lengths it validates before +/// constructing this value. Specialized collectors can inspect the typed decoded inputs through +/// [`inputs`](Self::inputs). This type does not expose the untyped execution arguments or output +/// arrays. +pub struct DenseRows<'a, Args: IndexedElementTuple> { + inputs: &'a Args::Columns, + views: Option>, + row_count: usize, +} + +impl<'a, Args: IndexedElementTuple> DenseRows<'a, Args> { + pub(crate) fn new(inputs: &'a Args::Columns, row_count: usize) -> VortexResult { + let views = Args::views_if_no_consts(inputs); + + if let Some(views) = &views { + vortex_ensure!( + Args::view_lens_match(views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + } else { + vortex_ensure!( + Args::decoded_lens_match(inputs, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + } + + Ok(Self { + inputs, + views, + row_count, + }) + } + + /// Return the typed decoded inputs. + /// + /// A specialized collector that borrows a new view from these inputs must validate that view + /// before using it for unchecked traversal. [`collect`](Self::collect) instead uses the exact + /// views retained during construction. + pub fn inputs(&self) -> &'a Args::Columns { + self.inputs + } + + /// Return the number of input rows. + pub fn len(&self) -> usize { + self.row_count + } + + /// Return whether the batch contains no rows. + pub fn is_empty(&self) -> bool { + self.row_count == 0 + } + + /// Collect one initialized value per row with the framework's dense traversal. + pub fn collect(self, apply: impl Fn(Args::Elems<'_>) -> Out) -> Vec { + if let Some(views) = self.views { + let mut values = Vec::::with_capacity(self.row_count); + let output = &mut values.spare_capacity_mut()[..self.row_count]; + + // SAFETY: `new` checked that these exact retained views address `row_count` rows. The + // `InputElement` contract keeps their lengths stable while they exist. + let source = unsafe { Args::indexed_source(views, self.row_count) }; + source.map_into(output, apply); + + // SAFETY: normal completion of `map_into` initializes every output slot exactly once. + unsafe { values.set_len(self.row_count) }; + + return values; + } + + let mut values = Vec::with_capacity(self.row_count); + for index in 0..self.row_count { + values.push(apply(Args::get(self.inputs, index))); + } + + values + } +} + +/// One semantic row operation with a selectable owned output representation. +pub trait RowKernel: Sized +where + Args: IndexedElementTuple, +{ + /// The logical value produced for one row. + type Element: OutputElement; + + /// The complete initialized output batch used by every execution policy. + type Output: RowKernelOutput; + + /// Evaluate one row using the kernel's portable reference semantics. + fn eval(&self, args: Args::Elems<'_>) -> Self::Element; + + /// Collect a validated dense batch. + /// + /// The default retains the framework's vectorizable scalar traversal, then converts the values + /// into the associated output. Override this method to write a packed or otherwise specialized + /// representation directly. The output must preserve row order and contain the same observable + /// values as calling [`eval`](Self::eval) once for each row. Dense execution can pass + /// unspecified payloads from null input rows; their output values can also be arbitrary because + /// batch execution masks them before returning the array. + fn collect_dense(&self, rows: DenseRows<'_, Args>) -> VortexResult { + Self::Output::from_values(rows.collect(|args| self.eval(args))) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/kernel_output.rs b/vortex-array/src/scalar_fn/unstable/row/types/kernel_output.rs new file mode 100644 index 00000000000..00da2a0b8e0 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/kernel_output.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Safe owned output values for infallible row kernels. +//! +//! A [`RowKernelOutput`] owns a completely initialized batch. The executor chooses which rows run +//! and validates the resulting array. A dense kernel can use [`PackedBoolOutput`] to write native +//! mask words directly without exposing uninitialized storage. + +use vortex_buffer::BitBuffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::validity::Validity; + +/// A completely initialized output batch produced by an infallible [`RowKernel`]. +/// +/// [`from_values`](Self::from_values) is the portable fallback used by the default dense +/// collector. Specialized kernels can construct their associated output directly. +/// +/// [`RowKernel`]: crate::scalar_fn::unstable::row::RowKernel +pub trait RowKernelOutput: Sized { + /// The logical value produced for one row. + type Element: OutputElement; + + /// Construct the output from one initialized value per row. + /// + /// Valid-only execution supplies [`Default::default`] placeholders at invalid rows. Batch + /// execution masks those rows before returning the array. + /// The output must preserve the input length and row order. + fn from_values(values: Vec) -> VortexResult; + + /// Construct the all-valid output array. + /// + /// The array must preserve the output's row count and match + /// [`OutputElement::element_dtype`] except for outer nullability. Per-row evaluation is + /// infallible, so errors from this method must not report value-dependent semantic failures. + fn finish(self) -> VortexResult; +} + +/// A kernel output backed by one native Rust value per row. +pub struct VecOutput { + values: Vec, +} + +impl RowKernelOutput for VecOutput { + type Element = T; + + fn from_values(values: Vec) -> VortexResult { + Ok(Self { values }) + } + + fn finish(self) -> VortexResult { + Ok(T::build(self.values)) + } +} + +/// A boolean kernel output backed by initialized native mask words. +/// +/// Dense SIMD kernels can write AVX-512 mask registers directly into [`words_mut`](Self::words_mut). +/// Bit `i % 64` of word `i / 64` stores row `i`, with the least-significant bit storing the first +/// row in each word. +/// Unused tail bits can contain any value; [`finish`](RowKernelOutput::finish) clears them before +/// constructing the boolean array. +pub struct PackedBoolOutput { + words: BufferMut, + row_count: usize, +} + +impl PackedBoolOutput { + /// Allocate an all-false output with `row_count` initialized bits. + pub fn zeroed(row_count: usize) -> Self { + Self { + words: BufferMut::zeroed(row_count.div_ceil(64)), + row_count, + } + } + + /// Return the logical number of output rows. + pub fn len(&self) -> usize { + self.row_count + } + + /// Return whether the output contains no rows. + pub fn is_empty(&self) -> bool { + self.row_count == 0 + } + + /// Return the initialized native words that store the output bits in row order. + pub fn words_mut(&mut self) -> &mut [u64] { + self.words.as_mut_slice() + } +} + +impl RowKernelOutput for PackedBoolOutput { + type Element = bool; + + fn from_values(values: Vec) -> VortexResult { + let mut output = Self::zeroed(values.len()); + + for (index, value) in values.into_iter().enumerate() { + if value { + output.words[index / 64] |= 1_u64 << (index % 64); + } + } + + Ok(output) + } + + fn finish(mut self) -> VortexResult { + if let Some(last_word) = self.words.last_mut() + && !self.row_count.is_multiple_of(64) + { + *last_word &= (1_u64 << (self.row_count % 64)) - 1; + } + + for word in self.words.iter_mut() { + *word = word.to_le(); + } + + let mut bytes = self.words.into_byte_buffer(); + bytes.truncate(self.row_count.div_ceil(8)); + let values = BitBuffer::new(bytes.freeze(), self.row_count); + + Ok(BoolArray::new(values, Validity::NonNullable).into_array()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index 4fc0c323ad4..9ea71d71a75 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -9,12 +9,23 @@ //! outcomes. mod element; +pub use element::ArgColumn; +pub use element::ArgView; pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; pub(super) use element::batch_const; +mod kernel; +pub use kernel::DenseRows; +pub use kernel::RowKernel; + +mod kernel_output; +pub use kernel_output::PackedBoolOutput; +pub use kernel_output::RowKernelOutput; +pub use kernel_output::VecOutput; + mod result; pub use result::FailureEvidence; pub use result::SinkResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index 8fc7209bc19..a83bbbef1e8 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -29,7 +29,10 @@ use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowKernel; use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::execute::execute_kernel; +use crate::scalar_fn::unstable::row::execute::execute_kernel_valid_rows; use crate::scalar_fn::unstable::row::execute::execute_owned; use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; use crate::scalar_fn::unstable::row::execute::execute_owned_infallible_valid_rows; @@ -83,6 +86,22 @@ impl private::Sealed for ExecuteRows<'_, '_, F> {} impl RowVisitor for ExecuteRows<'_, '_, F> { type VisitResult = ArrayRef; + fn visit_kernel(self, kernel: Kernel) -> VortexResult + where + Args: IndexedElementTuple, + Kernel: RowKernel, + { + const { assert_owned_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + execute_kernel::(self.args, self.ctx, kernel) + } + fn visit_prepared( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, @@ -213,6 +232,22 @@ impl private::Sealed for ExecuteValidRows<'_, '_, F> {} impl RowVisitor for ExecuteValidRows<'_, '_, F> { type VisitResult = Option; + fn visit_kernel(self, kernel: Kernel) -> VortexResult + where + Args: IndexedElementTuple, + Kernel: RowKernel, + { + const { assert_owned_visit_contract::() }; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; + + execute_kernel_valid_rows::(self.args, &self.valid, self.ctx, kernel) + } + fn visit_prepared( self, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 23a0d602d79..8d0b6cf683c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -15,6 +15,7 @@ use crate::scalar_fn::unstable::row::FailureEvidence; use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowKernel; use crate::scalar_fn::unstable::row::SinkResult; /// A planning or execution visit at concrete input and output types. @@ -68,6 +69,19 @@ pub trait RowVisitor: private::Sealed + Sized { self.visit_prepared::(|_| (), move |&(), args| apply(args)) } + /// Visit an infallible kernel with an associated output representation. + /// + /// Planning uses [`RowKernel::Element`] to derive the dtype. Validity-aware execution evaluates + /// rows through [`RowKernel::eval`] before constructing [`RowKernel::Output`]. Dense execution + /// can use [`RowKernel::collect_dense`] instead. + fn visit_kernel(self, kernel: Kernel) -> VortexResult + where + Args: IndexedElementTuple, + Kernel: RowKernel, + { + self.visit::(move |args| kernel.eval(args)) + } + /// The prepared form of [`visit`](Self::visit), with the same prerequisites. /// /// # Examples