diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 7c076bbb7b3..2af2eacf238 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -79,6 +79,8 @@ cudarc = ["dep:cudarc"] table-display = ["dep:tabled"] _test-harness = ["dep:goldenfile", "dep:rstest", "dep:rstest_reuse"] serde = ["dep:serde", "vortex-buffer/serde", "vortex-mask/serde"] +# Exposes experimental row-function APIs without compatibility guarantees. +unstable_row_fns = [] [dev-dependencies] divan = { workspace = true } @@ -90,7 +92,11 @@ rstest = { workspace = true } serde_json = { workspace = true } serde_test = { workspace = true } test-with = { workspace = true } -vortex-array = { path = ".", features = ["_test-harness", "table-display"] } +vortex-array = { path = ".", features = [ + "_test-harness", + "table-display", + "unstable_row_fns", +] } [[bench]] name = "aggregate_max" diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 003dbc2ca48..6be34ce1f34 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -6,6 +6,12 @@ //! This module contains the [`ScalarFnVTable`] trait and all built-in scalar function //! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions //! at each node. +//! +//! Strict functions with row-at-a-time kernels can implement `unstable::row::RowFn`. It handles +//! decoding, constants, null propagation, output construction, and validity. This API requires the +//! `unstable_row_fns` feature and has no compatibility guarantees. Implement [`ScalarFnVTable`] +//! directly for columnar kernels and functions that alias an input or can produce null from valid +//! inputs. use vortex_session::registry::Id; @@ -35,6 +41,12 @@ pub use options::*; mod signature; pub use signature::*; +#[cfg(feature = "unstable_row_fns")] +pub mod unstable; +#[cfg(not(feature = "unstable_row_fns"))] +#[allow(dead_code, unused_imports)] +pub(crate) mod unstable; + pub mod fns; pub mod internal; pub mod session; diff --git a/vortex-array/src/scalar_fn/unstable/mod.rs b/vortex-array/src/scalar_fn/unstable/mod.rs new file mode 100644 index 00000000000..6849c0bc661 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/mod.rs @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Experimental scalar-function APIs without compatibility guarantees. +//! +//! These APIs can change or disappear without a deprecation period. External users must enable +//! the corresponding `unstable_*` Cargo feature before importing them. + +pub mod row; diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs new file mode 100644 index 00000000000..bcb3a008488 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -0,0 +1,36 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Experimental support for strict scalar functions computed one row at a time. +//! +//! This module is experimental and has no compatibility guarantees. External users must enable +//! the `unstable_row_fns` Cargo feature before importing it. +//! +//! 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. +//! +//! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits +//! reduce compact failure evidence in that loop and retry only valid rows when null payloads may +//! have caused the failure. + +mod row_fn; +pub use row_fn::RowFn; + +mod types; +pub use types::ElementTuple; +pub use types::IndexedElementTuple; +pub use types::InitializedElement; +pub use types::InputElement; +pub use types::OutputElement; +pub use types::OutputSink; +pub use types::SinkResult; +pub use types::UninitElementSink; + +mod visitor; +pub use visitor::RowVisitor; + +mod vtable; +pub use vtable::execute_rows; +pub use vtable::row_fn_return_dtype; diff --git a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs new file mode 100644 index 00000000000..8a982c4fb37 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -0,0 +1,74 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`RowFn`] contract for scalar functions whose natural kernel computes one row at a time. +//! +//! Implementations declare their arity and fallibility, then use [`RowFn::dispatch`] to select the +//! typed row signature for each supported dtype combination. Optional methods provide +//! serialization without putting persistence plumbing in the row kernel. + +use std::fmt::Debug; +use std::fmt::Display; +use std::hash::Hash; + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::VortexSession; + +use super::visitor::RowVisitor; +use crate::dtype::DType; +use crate::scalar_fn::ScalarFnId; + +/// A scalar function computed one row at a time. +/// +/// 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 +/// vtable hooks can delegate its row kernel through [`row_fn_return_dtype`] and [`execute_rows`]. +/// +/// [`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 +pub trait RowFn: 'static + Sized + Clone + Send + Sync { + /// Options for this function, or [`EmptyOptions`](crate::scalar_fn::EmptyOptions) for none. + type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash; + + /// The arguments in display order. Its length is the function's exact arity. + const ARG_NAMES: &'static [&'static str]; + + /// Whether any dispatch can raise a semantic error. + /// + /// See [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible) for a + /// more detailed explanation of semantic errors. + /// + /// The framework checks dispatched element and result types. A conservative `true` is allowed. + const FALLIBLE: bool; + + /// Returns the ID of the scalar function. + fn id(&self) -> ScalarFnId; + + /// Serialize this function's options, or return `None` when the function is not serializable. + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + _ = options; + Ok(None) + } + + /// Restore options written by [`serialize`](Self::serialize). + fn deserialize( + &self, + _metadata: &[u8], + _session: &VortexSession, + ) -> VortexResult { + vortex_bail!("Expression {} is not deserializable", self.id()) + } + + /// Choose element types for these input dtypes and 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. + fn dispatch>( + &self, + options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult; +} 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 new file mode 100644 index 00000000000..28105c893cf --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/bool.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::validity::Validity; + +// SAFETY: the per-row view is a bit buffer, and its reported length is the buffer length. +unsafe impl InputElement for bool { + type Column = BitBuffer; + type View<'a> = &'a BitBuffer; + type Elem<'a> = bool; + + // Every bit of the buffer is readable, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + vortex_ensure!( + matches!(dtype, DType::Bool(_)), + "expected a Bool column, got {dtype}", + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_bit_buffer()) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> bool { + column.value(index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.len() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> bool + where + Self: 'a, + { + view.value(index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> bool + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { view.value_unchecked(index) } + } +} + +impl OutputElement for bool { + fn element_dtype() -> DType { + DType::Bool(Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + // `From>` uses the bulk bit-packing path. + BoolArray::new(BitBuffer::from(values), Validity::NonNullable).into_array() + } +} 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 new file mode 100644 index 00000000000..2764840da7a --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Typed decoding and row access for one input column. +//! +//! [`InputElement`] separates invocation-wide decoding from the checked and unchecked access paths +//! used by row kernels. + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; + +/// An element type that can be read row-wise out of an input column. +/// +/// # Safety +/// +/// For every view returned by [`view`](Self::view), every index below +/// [`view_len`](Self::view_len) **must** satisfy the safety contract of +/// [`get_from_view_unchecked`](Self::get_from_view_unchecked). Shared execution relies on this +/// proof to perform unchecked reads after one pre-loop length check. +pub unsafe trait InputElement: 'static { + /// The decoded column representation supporting `O(1)` row access. + type Column; + + /// The row-loop view of a decoded column. + /// + /// This can borrow a cheaper representation than [`Column`](Self::Column). Primitive elements, + /// for example, expose a slice so its pointer and length are loop invariants rather than + /// re-reading a [`Buffer`](vortex_buffer::Buffer) descriptor for every row. + type View<'a>; + + /// The borrowed element value handed to a row closure. + type Elem<'a>; + + /// Whether every dense decode and access path tolerates rows that are null in the input. + /// + /// Arrays guarantee payloads only for valid rows. Set this to `true` only when every decode and + /// access method remains safe for null rows. Dense execution can pass unspecified values from + /// null rows to the row closure. + const DENSE_SAFE: bool; + + /// Whether [`decode`](Self::decode) can fail on _legal_ input data. + /// + /// This excludes infrastructural failures such as IO or allocation. + const DECODE_FALLIBLE: bool; + + /// Validate that `dtype` is an acceptable input column dtype for this element type. + fn validate(dtype: &DType) -> VortexResult<()>; + + /// Decode `array` into its column representation. + /// + /// Called once per row-kernel invocation, including deferred-error retries. Hoist dtype checks, + /// downcasts, and other invocation-invariant work into this method. + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array. + /// + /// The conservative default declines. An implementation whose ordinary decode is safe and + /// infallible over null payloads can return `true`. Other implementations can inspect `array` + /// and opt in only for supported representations. + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(false) + } + + /// Decode `array` _without_ assuming every row is valid, or return `Ok(None)` when this element + /// cannot decode this particular array. + /// + /// Override this for a non-dense-safe representation that can still place safe placeholders in + /// null slots. The skip-invalid executor never reads those slots. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if Self::can_decode_null_tolerant(&array)? { + Self::decode(array, ctx).map(Some) + } else { + Ok(None) + } + } + + /// Read one row without repeating batch-constant work from [`decode`](Self::decode). + fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>; + + /// Borrow the representation used when this argument varies within the batch. + /// + /// Called once before the hot loop. Constants do not use this view because the tuple adapter + /// keeps their one-row decoded representation separate. + fn view(column: &Self::Column) -> Self::View<'_>; + + /// Number of rows addressable through a [`View`](Self::View). + /// + /// Every index below this length must be valid for + /// [`get_from_view_unchecked`](Self::get_from_view_unchecked). + fn view_len(view: &Self::View<'_>) -> usize; + + /// Read one row from a [`View`](Self::View). + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a; + + /// Read one row without checking that `index` is in bounds. + /// + /// # Safety + /// + /// `index` must be less than [`view_len`](Self::view_len) for `view`. + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a, + { + Self::get_from_view(view, index) + } +} 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 new file mode 100644 index 00000000000..51d66594332 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The element types a row function can read and produce. +//! +//! [`InputElement::Elem`] can borrow from its decoded column. Owned row computations return an +//! [`OutputElement`]. Runtime-shaped outputs use an +//! [`OutputSink`](crate::scalar_fn::unstable::row::OutputSink). + +mod bool; + +mod input; +pub use input::InputElement; + +mod output; +pub use output::OutputElement; + +mod primitive; + +mod tuple; +pub use tuple::ElementTuple; +pub use tuple::IndexedElementTuple; 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 new file mode 100644 index 00000000000..d1c75b7054a --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Owned scalar values that can be collected into all-valid output columns. +//! +//! [`OutputElement`] describes fixed-dtype values returned independently by each row invocation. + +use crate::ArrayRef; +use crate::dtype::DType; + +/// An owned row value that can be built into an all-valid column. +pub trait OutputElement: 'static + Sized { + /// The dtype of columns built from this element type. **Must** be non-nullable: nullability is + /// derived from the inputs by batch execution. + /// + /// Because this method takes no arguments, the dtype must be a property of the Rust type. Use + /// an [`OutputSink`] when the output dtype depends on function options or input dtypes. + /// + /// [`OutputSink`]: crate::scalar_fn::unstable::row::OutputSink + fn element_dtype() -> DType; + + /// Build a column from one value per row. Called once per batch. + fn build(values: Vec) -> ArrayRef; +} 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 new file mode 100644 index 00000000000..9bc5db0c8e4 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::InputElement; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::validity::Validity; + +// SAFETY: the per-row view is a native slice, and its reported length is the slice length. +unsafe impl InputElement for T { + type Column = Buffer; + type View<'a> = &'a [T]; + type Elem<'a> = T; + + // Every lane of the buffer holds a `T`, valid or not. + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + let expected = T::PTYPE; + let DType::Primitive(ptype, _) = dtype else { + vortex_bail!("expected a {expected} column, got {dtype}"); + }; + vortex_ensure_eq!( + *ptype, + expected, + "expected a {expected} column, got {dtype}" + ); + Ok(()) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + Ok(array.execute::(ctx)?.into_buffer::()) + } + + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } + + fn get(column: &Self::Column, index: usize) -> T { + column[index] + } + + fn view(column: &Self::Column) -> Self::View<'_> { + column.as_slice() + } + + fn view_len(view: &Self::View<'_>) -> usize { + view.len() + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> T + where + Self: 'a, + { + view[index] + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> T + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { *view.get_unchecked(index) } + } +} + +impl OutputElement for T { + fn element_dtype() -> DType { + DType::Primitive(T::PTYPE, Nullability::NonNullable) + } + + fn build(values: Vec) -> ArrayRef { + PrimitiveArray::new(values, Validity::NonNullable).into_array() + } +} 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 new file mode 100644 index 00000000000..b0a3e709696 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -0,0 +1,396 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Decoding and row access for tuples of input element types. +//! +//! [`ElementTuple`] combines per-column [`InputElement`] implementations, preserves batch +//! constants outside the hot loop, and supports row functions with up to twelve arguments. + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Constant; +use crate::arrays::Extension; +use crate::arrays::Masked; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; + +/// One decoded input, collapsed to a single row when it is constant for the batch. +pub struct ArgColumn( + /// The decoded column, classified by whether it stores one value per row. + ArgColumnKind, +); + +enum ArgColumnKind { + PerRow(T::Column), + Constant(T::Column), +} + +impl ArgColumn { + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + // An empty input has no row 0 to slice, and its row loop runs zero times either way. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?))); + } + + Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) + } + + fn decode_null_tolerant(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult> { + // Batch execution short-circuits null constants before selecting a strategy, so a + // constant reaching this path is non-null and can use the ordinary decode. + if let Some(constant) = batch_constant(&array) + && !array.is_empty() + { + return Ok(Some(Self(ArgColumnKind::Constant(T::decode( + constant.slice(0..1)?, + ctx, + )?)))); + } + + Ok(T::decode_null_tolerant(array, ctx)? + .map(ArgColumnKind::PerRow) + .map(Self)) + } + + fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult { + // Batch execution short-circuits null constants before selecting this path, so a + // non-empty constant can always use the ordinary decode. + if batch_constant(array).is_some() && !array.is_empty() { + return Ok(true); + } + + T::can_decode_null_tolerant(array) + } + + fn get(&self, index: usize) -> T::Elem<'_> { + match &self.0 { + ArgColumnKind::PerRow(column) => T::get(column, index), + ArgColumnKind::Constant(column) => T::get(column, 0), + } + } + + fn per_row_column(&self) -> Option<&T::Column> { + match &self.0 { + ArgColumnKind::PerRow(column) => Some(column), + ArgColumnKind::Constant(_) => None, + } + } + + fn addresses_rows(&self, row_count: usize) -> bool { + // A constant is always read at index zero, so it addresses any batch length. + match &self.0 { + ArgColumnKind::PerRow(column) => T::view_len(&T::view(column)) == row_count, + ArgColumnKind::Constant(_) => true, + } + } + + fn constant(&self) -> Option> { + match &self.0 { + ArgColumnKind::PerRow(_) => None, + ArgColumnKind::Constant(column) => Some(T::get(column, 0)), + } + } +} + +/// Return the batch-constant array, looking through masked and extension wrappers. +/// +/// Batch execution owns mask validity, so a masked constant can expose its constant child here. An +/// extension over constant storage remains wrapped to preserve its extension dtype. +pub fn batch_constant(array: &ArrayRef) -> Option { + if array.is::() { + return Some(array.clone()); + } + + if let Some(masked) = array.as_opt::() { + return Some(masked.child().clone()).filter(|child| child.is::()); + } + + array + .as_opt::() + .is_some_and(|ext| ext.storage_array().is::()) + .then(|| array.clone()) +} + +/// Typed argument tuples for arities zero through twelve. +/// +/// This trait is sealed. Add a new row representation by implementing [`InputElement`] and placing +/// it in one of the supplied tuples. +pub trait ElementTuple: 'static + private::Sealed { + /// The decoded column representations. + type Columns; + + /// Borrowed views of decoded columns when every argument stores one value per row. + type Views<'a>; + + /// The borrowed row of element values. + type Elems<'a>; + + /// The batch-constant element values. + /// + /// `Some` carries the value of a batch-constant argument. `None` marks a per-row argument. A + /// [`RowVisitor`] passes these values to its prepare closure so constant work can leave the row + /// loop. + /// + /// [`RowVisitor`]: crate::scalar_fn::unstable::row::RowVisitor + type ConstElems<'a>; + + /// The number of arguments. + const ARITY: usize; + + /// Whether every argument is [`InputElement::DENSE_SAFE`]. + const DENSE_SAFE: bool; + + /// Whether _any_ argument is [`InputElement::DECODE_FALLIBLE`]. + const DECODE_FALLIBLE: bool; + + /// Validate the input dtypes and exact arity. + fn validate(dtypes: &[DType]) -> VortexResult<()>; + + /// Decode every input column once for one row-kernel invocation. + /// + /// A dense deferred-error retry starts another invocation over filtered valid rows. + fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Whether every input can be decoded without assuming that all rows are valid. + /// + /// The tuple checks this before decoding any column, so a decline does not discard work from + /// earlier arguments. + fn can_decode_null_tolerant(args: &dyn ExecutionArgs) -> VortexResult; + + /// Decode every input column once while tolerating null rows. + /// + /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid + /// strategy calls this once per batch. + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult>; + + /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; + + /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// + /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple + /// gives the optimizer ordinary contiguous column access without a per-row constant check. + fn per_row_views(columns: &Self::Columns) -> Option>; + + /// Whether every view contains exactly `row_count` rows. + /// + /// The executor calls this once before the all-per-row hot loop. A successful check gives LLVM + /// a dominating equality between the loop bound and every source length, which lets it optimize + /// the tuple access as one fixed-length traversal. + fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; + + /// Whether every per-row argument contains exactly `row_count` rows. + /// + /// This is the mixed-shape equivalent of [`view_lens_match`](Self::view_lens_match) when + /// [`per_row_views`](Self::per_row_views) declines. It runs once before the hot loop for the + /// same LLVM optimization. A batch constant is exempt because decoding collapsed it to one row. + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; + + /// Read one row from borrowed views. + fn get_from_views<'a>(views: &Self::Views<'a>, index: usize) -> Self::Elems<'a>; + + /// Read one row from borrowed views without checking bounds. + /// + /// # Safety + /// + /// `index` must be in bounds for every column. + unsafe fn get_from_views_unchecked<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a>; + + /// Read the batch-constant elements out of the decoded columns once for one row-kernel + /// invocation. + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; +} + +impl private::Sealed for () {} + +impl ElementTuple for () { + type Columns = (); + type Views<'a> = (); + type Elems<'a> = (); + type ConstElems<'a> = (); + + const ARITY: usize = 0; + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + 0, + "expected 0 argument dtypes, got {}", + dtypes.len(), + ); + Ok(()) + } + + fn decode(_args: &dyn ExecutionArgs, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(()) + } + + fn can_decode_null_tolerant(_args: &dyn ExecutionArgs) -> VortexResult { + Ok(true) + } + + fn decode_null_tolerant( + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(())) + } + + fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} + + fn per_row_views(_columns: &Self::Columns) -> Option> { + Some(()) + } + + fn view_lens_match(_views: &Self::Views<'_>, _row_count: usize) -> bool { + true + } + + fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { + true + } + + fn get_from_views<'a>(_views: &Self::Views<'a>, _index: usize) -> Self::Elems<'a> {} + + unsafe fn get_from_views_unchecked<'a>( + _views: &Self::Views<'a>, + _index: usize, + ) -> Self::Elems<'a> { + } + + fn constants(_columns: &Self::Columns) -> Self::ConstElems<'_> {} +} + +macro_rules! element_tuple { + ($arity:literal; $($t:ident : $idx:tt),+) => { + impl<$($t: InputElement),+> private::Sealed for ($($t,)+) {} + + impl<$($t: InputElement),+> ElementTuple for ($($t,)+) { + type Columns = ($(ArgColumn<$t>,)+); + type Views<'a> = ($($t::View<'a>,)+); + type Elems<'a> = ($($t::Elem<'a>,)+); + type ConstElems<'a> = ($(Option<$t::Elem<'a>>,)+); + + const ARITY: usize = $arity; + const DENSE_SAFE: bool = $($t::DENSE_SAFE &&)+ true; + const DECODE_FALLIBLE: bool = $($t::DECODE_FALLIBLE ||)+ false; + + fn validate(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure_eq!( + dtypes.len(), + $arity, + "expected {} argument dtypes, got {}", + $arity, + dtypes.len(), + ); + + $($t::validate(&dtypes[$idx])?;)+ + Ok(()) + } + + fn decode( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + Ok(($(ArgColumn::<$t>::decode(args.get($idx)?, ctx)?,)+)) + } + + fn can_decode_null_tolerant(args: &dyn ExecutionArgs) -> VortexResult { + Ok($({ + let array = args.get($idx)?; + ArgColumn::<$t>::can_decode_null_tolerant(&array)? + } &&)+ true) + } + + fn decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if !Self::can_decode_null_tolerant(args)? { + return Ok(None); + } + + Ok(Some(( + $(match ArgColumn::<$t>::decode_null_tolerant(args.get($idx)?, ctx)? { + Some(column) => column, + None => return Ok(None), + },)+ + ))) + } + + fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_> { + ($(columns.$idx.get(index),)+) + } + + fn per_row_views(columns: &Self::Columns) -> Option> { + Some(($($t::view(columns.$idx.per_row_column()?),)+)) + } + + fn view_lens_match( + views: &Self::Views<'_>, + row_count: usize, + ) -> bool { + $($t::view_len(&views.$idx) == row_count &&)+ true + } + + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { + $(columns.$idx.addresses_rows(row_count) &&)+ true + } + + fn get_from_views<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a> { + ($($t::get_from_view(&views.$idx, index),)+) + } + + unsafe fn get_from_views_unchecked<'a>( + views: &Self::Views<'a>, + index: usize, + ) -> Self::Elems<'a> { + // SAFETY: forwarded from this method's contract. + ($(unsafe { $t::get_from_view_unchecked(&views.$idx, index) },)+) + } + + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { + ($(columns.$idx.constant(),)+) + } + } + }; +} + +element_tuple!(1; A: 0); +element_tuple!(2; A: 0, B: 1); +element_tuple!(3; A: 0, B: 1, C: 2); +element_tuple!(4; A: 0, B: 1, C: 2, D: 3); +element_tuple!(5; A: 0, B: 1, C: 2, D: 3, E: 4); +element_tuple!(6; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5); +element_tuple!(7; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6); +element_tuple!(8; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7); +element_tuple!(9; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8); +element_tuple!(10; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9); +element_tuple!(11; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10); +element_tuple!(12; A: 0, B: 1, C: 2, D: 3, E: 4, F: 5, G: 6, H: 7, I: 8, J: 9, K: 10, L: 11); + +mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs new file mode 100644 index 00000000000..d612d874935 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Validated indexed access to tuples of decoded input columns. +//! +//! [`IndexedElementTuple`] adapts row arguments to the lane-kernel interface after batch execution +//! proves that every input covers the requested row range. + +use vortex_compute::lane_kernels::IndexedSource; +use vortex_compute::lane_kernels::LaneZip; + +use super::ElementTuple; +use crate::scalar_fn::unstable::row::InputElement; + +/// An argument tuple that supports a validated dense indexed traversal. +/// +/// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's +/// unchecked view access after batch execution validates every decoded column length once. +pub trait IndexedElementTuple: ElementTuple { + /// The source shared execution uses for a dense all-per-row loop. + /// + /// Its length must be the common view length. For every valid index it must preserve row order, + /// return the same value as [`ElementTuple::get_from_views`], and uphold the unchecked read + /// contract of [`IndexedSource`]. + type Source<'a>: IndexedSource>; + + /// Build a source from views already validated to cover the complete batch. + /// + /// # Safety + /// + /// Every view in `views` **must** address exactly `row_count` rows. Violating this requirement + /// can make a safe lane kernel read outside a column's allocation. + unsafe fn indexed_source<'a>(views: Self::Views<'a>, row_count: usize) -> Self::Source<'a>; +} + +/// Indexed access to one element view. +pub struct ElementSource<'a, T: InputElement> { + view: T::View<'a>, +} + +impl<'a, T: InputElement> ElementSource<'a, T> { + fn new(view: T::View<'a>) -> Self { + Self { view } + } +} + +impl<'a, T: InputElement> IndexedSource for ElementSource<'a, T> { + type Item = T::Elem<'a>; + + fn len(&self) -> usize { + T::view_len(&self.view) + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the source length is the number of rows addressable by `view`, and the caller + // guarantees that `index` is below that length. + unsafe { T::get_from_view_unchecked(&self.view, index) } + } +} + +/// An indexed element source yielding the one-tuples expected by a unary row closure. +pub struct UnaryTupleSource(Source); + +impl IndexedSource for UnaryTupleSource { + type Item = (Source::Item,); + + fn len(&self) -> usize { + self.0.len() + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: forwarded from this method's contract. + (unsafe { self.0.get_unchecked(index) },) + } +} + +/// Indexed access to the views of an element tuple. +pub struct ElementTupleSource<'a, Args: ElementTuple> { + views: Args::Views<'a>, + row_count: usize, +} + +impl<'a, Args: ElementTuple> IndexedSource for ElementTupleSource<'a, Args> { + type Item = Args::Elems<'a>; + + fn len(&self) -> usize { + self.row_count + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the caller guarantees that `index` is below `row_count`. Batch execution checks + // that every view addresses exactly `row_count` rows before constructing this source. + unsafe { Args::get_from_views_unchecked(&self.views, index) } + } +} + +impl IndexedElementTuple for () { + type Source<'a> = ElementTupleSource<'a, ()>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, row_count: usize) -> Self::Source<'a> { + ElementTupleSource { views, row_count } + } +} + +impl IndexedElementTuple for (A,) { + type Source<'a> = UnaryTupleSource>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, _row_count: usize) -> Self::Source<'a> { + UnaryTupleSource(ElementSource::new(views.0)) + } +} + +impl IndexedElementTuple for (A, B) { + type Source<'a> = LaneZip, ElementSource<'a, B>>; + + unsafe fn indexed_source<'a>(views: Self::Views<'a>, _row_count: usize) -> Self::Source<'a> { + LaneZip::new(ElementSource::new(views.0), ElementSource::new(views.1)) + } +} + +macro_rules! indexed_element_tuple { + ($($t:ident),+) => { + impl<$($t: InputElement),+> IndexedElementTuple for ($($t,)+) { + type Source<'a> = ElementTupleSource<'a, ($($t,)+)>; + + unsafe fn indexed_source<'a>( + views: Self::Views<'a>, + row_count: usize, + ) -> Self::Source<'a> { + ElementTupleSource { views, row_count } + } + } + }; +} + +indexed_element_tuple!(A, B, C); +indexed_element_tuple!(A, B, C, D); +indexed_element_tuple!(A, B, C, D, E); +indexed_element_tuple!(A, B, C, D, E, F); +indexed_element_tuple!(A, B, C, D, E, F, G); +indexed_element_tuple!(A, B, C, D, E, F, G, H); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K); +indexed_element_tuple!(A, B, C, D, E, F, G, H, I, J, K, L); 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 new file mode 100644 index 00000000000..a2c143704a0 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -0,0 +1,16 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Combines [`InputElement`](super::InputElement)s into typed row argument lists. +//! +//! [`ElementTuple`] owns decoding, constant classification, and row access for supported arities. +//! [`IndexedElementTuple`] adds the validated indexed source used by vectorizable dense loops. + +mod element_tuple; +pub use element_tuple::ElementTuple; + +mod indexed; +pub use indexed::IndexedElementTuple; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs new file mode 100644 index 00000000000..044ad15fd4b --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; + +use super::ElementTuple; +use super::element_tuple::batch_constant; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::Constant; +use crate::arrays::ConstantArray; +use crate::arrays::ExtensionArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::extension::datetime::TimeUnit; +use crate::extension::datetime::Timestamp; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; +use crate::validity::Validity; + +static DECODE_CALLS: AtomicUsize = AtomicUsize::new(0); + +macro_rules! i64_test_element { + ($element:ident, $decode_fallible:literal $(, $can_decode:item)?) => { + struct $element; + + // SAFETY: the view and unchecked access delegate to the `i64` implementation. + unsafe impl InputElement for $element { + type Column = Buffer; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = $decode_fallible; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + DECODE_CALLS.fetch_add(1, Ordering::Relaxed); + ::decode(array, ctx) + } + + $($can_decode)? + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + ::view(column) + } + + fn view_len(view: &Self::View<'_>) -> usize { + ::view_len(view) + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_from_view(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { ::get_from_view_unchecked(view, index) } + } + } + }; +} + +i64_test_element!( + DecodeProbe, + false, + fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult { + Ok(true) + } +); +i64_test_element!(DenseFallible, true); + +#[test] +fn test_null_tolerant_decline_precedes_decoding() -> VortexResult<()> { + DECODE_CALLS.store(0, Ordering::Relaxed); + let first = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let second = PrimitiveArray::from_iter([3_i64, 4]).into_array(); + let args = VecExecutionArgs::new(vec![first, second], 2); + let mut ctx = array_session().create_execution_ctx(); + + let columns = <(DecodeProbe, DenseFallible)>::decode_null_tolerant(&args, &mut ctx)?; + + assert!(columns.is_none()); + assert_eq!(DECODE_CALLS.load(Ordering::Relaxed), 0); + Ok(()) +} + +#[test] +fn test_batch_constant_unwraps_filtered_masked_constant() -> VortexResult<()> { + let child = ConstantArray::new(7_i64, 3).into_array(); + let masked = + MaskedArray::try_new(child, Validity::from_iter([true, false, true]))?.into_array(); + let filtered = masked.filter(Mask::from_iter([true, true, false]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered masked constant must remain batch-constant"); + }; + + assert!(constant.is::()); + Ok(()) +} + +#[test] +fn test_batch_constant_preserves_filtered_extension() -> VortexResult<()> { + let ext_dtype = Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(); + let extension = + ExtensionArray::new(ext_dtype, ConstantArray::new(7_i64, 3).into_array()).into_array(); + let filtered = extension.filter(Mask::from_iter([true, false, true]))?; + + let Some(constant) = batch_constant(&filtered) else { + vortex_bail!("filtered extension storage must remain batch-constant"); + }; + + assert_eq!(constant.dtype(), extension.dtype()); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs new file mode 100644 index 00000000000..ce119f32915 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -0,0 +1,22 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Input decoding and output construction for row functions. +//! +//! [`element`] defines the Rust values decoded from input columns and built into simple output +//! columns. [`sink`] handles outputs that need row handles or batch-wide state. [`result`] defines +//! the immediate and deferred outcomes returned by sink-writing row closures. + +mod element; +pub use element::ElementTuple; +pub use element::IndexedElementTuple; +pub use element::InputElement; +pub use element::OutputElement; + +mod result; +pub use result::SinkResult; + +mod sink; +pub use sink::InitializedElement; +pub use sink::OutputSink; +pub use sink::UninitElementSink; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/result.rs b/vortex-array/src/scalar_fn/unstable/row/types/result.rs new file mode 100644 index 00000000000..6163687932d --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/result.rs @@ -0,0 +1,73 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Return types for sink-writing row closures. +//! +//! [`SinkResult`] lets the executor handle initialized sinks and sinks that require an +//! [`InitializedElement`] token, with either infallible or immediate-error callbacks. + +use vortex_error::VortexResult; + +use super::InitializedElement; + +/// The result of writing one row: success or an immediate error. +/// +/// This trait is sealed. Row functions choose one of its supplied implementations. +pub trait SinkResult: 'static + private::Sealed { + /// The [`OutputSink::WriteToken`](super::OutputSink::WriteToken) carried by a success. + type WriteToken: 'static; + + /// Whether this return type can carry an error. + const FALLIBLE: bool; + + /// Convert this row's outcome into immediate success or failure. + fn into_result(self) -> VortexResult<()>; +} + +impl private::Sealed for () {} + +impl SinkResult for () { + type WriteToken = (); + const FALLIBLE: bool = false; + + fn into_result(self) -> VortexResult<()> { + Ok(()) + } +} + +impl private::Sealed for InitializedElement {} + +impl SinkResult for InitializedElement { + type WriteToken = InitializedElement; + const FALLIBLE: bool = false; + + fn into_result(self) -> VortexResult<()> { + Ok(()) + } +} + +impl private::Sealed for VortexResult<()> {} + +impl SinkResult for VortexResult<()> { + type WriteToken = (); + const FALLIBLE: bool = true; + + fn into_result(self) -> VortexResult<()> { + self + } +} + +impl private::Sealed for VortexResult {} + +impl SinkResult for VortexResult { + type WriteToken = InitializedElement; + const FALLIBLE: bool = true; + + fn into_result(self) -> VortexResult<()> { + self.map(|_| ()) + } +} + +mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs new file mode 100644 index 00000000000..6cc3ce06d30 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -0,0 +1,217 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Output builders for row kernels that cannot return independent owned values. +//! +//! [`OutputSink`] allocates batch-wide state and lends one row handle to each callback. +//! [`UninitElementSink`] is the fixed-width implementation used when avoiding output +//! initialization matters. + +use std::mem::MaybeUninit; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::OutputElement; + +/// A column allocated once per batch that a row closure writes into, one row at a time. +/// +/// A sink can use function options and input dtypes to build a runtime-shaped output or own shared +/// batch state. The executor passes each row slot into an [`Fn`] closure. +/// +/// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. +/// Skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an +/// initializer. +/// +/// # Errors +/// +/// Lifecycle methods report only incidental failures such as allocation. A semantic error that +/// depends on input values **must** come from the row callback through a fallible [`SinkResult`], +/// or [`RowFn::FALLIBLE`] cannot protect optimizations such as dictionary push-down. +/// +/// # Safety +/// +/// An implementation must uphold all of these requirements: +/// +/// - Every index in `0..row_count(rows)` **must** identify one distinct row owned by this sink. +/// - A row must either be initialized before the callback or require a +/// [`WriteToken`] that safe code cannot produce without initializing that exact row. Evidence for +/// an uninitialized row **must not** be safely forgeable, reusable, or substitutable. +/// - An initializer returned by [`skipped_rows_initializer`] **must** initialize every row. +/// - `Self` and every borrowed [`Rows`] view **must** remain safe to drop if decoding, +/// preparation, skipped-row initialization, or a row callback returns an error or unwinds. The +/// executor can abandon a sink after any prefix of rows. +/// - [`finish`] **must** be sound once every visited callback returned its required token and the +/// skipped-row initializer, when present, ran successfully. +/// +/// [`Rows`]: Self::Rows +/// [`WriteToken`]: Self::WriteToken +/// [`finish`]: Self::finish +/// [`row_count`]: Self::row_count +/// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE +/// [`SinkResult`]: crate::scalar_fn::unstable::row::SinkResult +/// [`skipped_rows_initializer`]: Self::skipped_rows_initializer +pub unsafe trait OutputSink: 'static + Sized { + /// A loop-local view of all output rows. + /// + /// Borrowed once before execution so the sink's buffer descriptor and shape become loop + /// invariants rather than being re-read through `&mut Self` for every row. + type Rows<'a> + where + Self: 'a; + + /// The place a row closure writes one row through, borrowed from the sink. + type Row<'a> + where + Self: 'a; + + /// Proof that a successful row closure left its row handle initialized. + /// + /// Use `()` for initialized row handles. A sink exposing uninitialized storage uses a token + /// returned after initialization. If a sink uses the token to justify unsafe code, safe code + /// **must not** be able to construct one without establishing the invariant. + type WriteToken: 'static; + + /// The operation that initializes every output position before skip-invalid execution. + /// + /// `Some` enables skip-invalid execution. The initializer **must** make every row safe to + /// finish. Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// + /// `None` makes the executor fall back to filtering the inputs. + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + None + } + + /// The dtype of the column this sink builds, given the function options and input dtypes. + /// + /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the + /// result, and masks the null rows. + fn output_dtype(options: &Options, args: &[DType]) -> VortexResult; + + /// Allocate a sink for `rows` rows. + fn with_capacity(rows: usize) -> VortexResult; + + /// Borrow all output rows for the hot loop. + fn rows(&mut self) -> Self::Rows<'_>; + + /// The number of rows addressable through [`row_unchecked`](Self::row_unchecked). + fn row_count(rows: &Self::Rows<'_>) -> usize; + + /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. + /// + /// # Safety + /// + /// `index` must be less than [`row_count`](Self::row_count) for `rows`. + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; + + /// Finish into the built column, whose dtype **must** be this sink's + /// [`output_dtype`](Self::output_dtype). Called once per batch. + /// + /// # Safety + /// + /// The executor must have completed every row callback successfully, and each callback must + /// have returned this sink's [`WriteToken`](Self::WriteToken). When skipped rows are allowed, + /// the initializer returned by + /// [`skipped_rows_initializer`](Self::skipped_rows_initializer) must have run before traversal. + unsafe fn finish(self) -> VortexResult; +} + +/// Proof that one uninitialized element row was initialized. +/// +/// The private field prevents safe construction without calling [`write`](Self::write): +/// +/// ```compile_fail,E0423 +/// use vortex_array::scalar_fn::unstable::row::InitializedElement; +/// +/// let _evidence = InitializedElement(()); +/// ``` +#[must_use = "return this token from the row closure to prove that it initialized the output"] +pub struct InitializedElement( + /// Private so constructing initialization evidence requires an unsafe operation. + (), +); + +impl InitializedElement { + /// Write `value` into an uninitialized row and return its proof token. + /// + /// # Safety + /// + /// `row` must be the [`UninitElementSink`] row supplied to the current callback. The caller + /// must return the token from that callback. Using another row or returning the token from + /// another callback can cause undefined behavior. + #[inline] + pub unsafe fn write(row: &mut MaybeUninit, value: T) -> Self { + row.write(value); + + Self(()) + } +} + +/// An element sink that leaves dense output uninitialized before the row loop. +/// +/// The row closure must return the [`InitializedElement`] from [`InitializedElement::write`] on +/// success. The token is zero-sized, so the proof adds no runtime row state. +/// +/// Skip-invalid execution initializes placeholders before omitting rows. Errors and unwinds are +/// safe because `values` keeps length zero until `finish`. The `T: Copy` bound means that +/// initialized spare-capacity elements require no destruction. +pub struct UninitElementSink { + /// Spare storage written in increasing row order. + values: Vec, + + /// The number of slots exposed to the row loop and initialized before finishing. + row_count: usize, +} + +// SAFETY: the row slice covers exactly the reserved spare-capacity range, so each accepted index +// names one distinct slot. Safe code cannot construct `InitializedElement`. Its unsafe constructor +// writes the supplied slot and requires the caller to return that exact evidence. The +// skipped-row initializer writes `T::default()` into every slot before masked traversal. +unsafe impl OutputSink + for UninitElementSink +{ + type Rows<'a> = &'a mut [MaybeUninit]; + type Row<'a> = &'a mut MaybeUninit; + type WriteToken = InitializedElement; + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + for row in rows.iter_mut() { + row.write(T::default()); + } + }) + } + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(T::element_dtype()) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self { + values: Vec::with_capacity(rows), + row_count: rows, + }) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.values.spare_capacity_mut()[..self.row_count] + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(mut self) -> VortexResult { + // SAFETY: the caller guarantees every slot in `0..row_count` was initialized, and + // `with_capacity` reserved every slot in that range. + unsafe { self.values.set_len(self.row_count) }; + + Ok(T::build(self.values)) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs new file mode 100644 index 00000000000..12d5bd7c7ea --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/check.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Contract checks shared by planning and execution visits. +//! +//! Const assertions reject invalid generic visits during compilation. The validators compare a +//! selected visit with the input dtypes during planning and return its output dtype. + +use std::mem::needs_drop; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::dtype::DType; +use crate::scalar_fn::unstable::row::ElementTuple; +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::SinkResult; + +/// Assert the no-drop contract that makes partially initialized output safe to abandon on unwind. +pub(crate) const fn assert_owned_output_needs_no_drop() { + assert!( + !needs_drop::(), + "owned row outputs must not require drop glue" + ); +} + +const fn assert_input_visit_contract() { + assert!( + Args::ARITY == F::ARG_NAMES.len(), + "the visited argument tuple must have the arity declared by RowFn::ARG_NAMES", + ); + // Dictionary push-down can evaluate values that no input row references. Every dispatch must + // therefore match the function-wide fallibility declaration. + assert!( + !Args::DECODE_FALLIBLE || F::FALLIBLE, + "RowFn::FALLIBLE must be true when input decoding can fail", + ); +} + +pub(super) const fn assert_owned_visit_contract() +where + Function: RowFn, + Args: IndexedElementTuple, + Out: OutputElement, +{ + assert_input_visit_contract::(); + assert_owned_output_needs_no_drop::(); +} + +pub(super) const fn assert_sink_visit_contract() +where + Function: RowFn, + Args: ElementTuple, + ApplyResult: SinkResult, +{ + assert_input_visit_contract::(); + assert!( + !ApplyResult::FALLIBLE || Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result can fail", + ); +} + +pub(super) const fn assert_deferred_visit_contract() +where + Function: RowFn, + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, +{ + assert_owned_visit_contract::(); + assert!( + Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result defers failure evidence", + ); + assert!( + size_of::() <= size_of::(), + "failure evidence must be no wider than the value, or it bounds the vector width", + ); +} + +pub(super) fn validate_owned_visit( + dtypes: &[DType], +) -> VortexResult { + Args::validate(dtypes)?; + + let dtype = Out::element_dtype(); + vortex_ensure!( + !dtype.is_nullable(), + "row output elements must declare a non-nullable dtype, got {dtype}", + ); + + Ok(dtype) +} + +pub(super) fn validate_sink_visit( + options: &Options, + dtypes: &[DType], +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, +{ + Args::validate(dtypes)?; + + let dtype = Sink::output_dtype(options, dtypes)?; + vortex_ensure!( + !dtype.is_nullable(), + "row output sinks must declare a non-nullable dtype, got {dtype}", + ); + + Ok(dtype) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs new file mode 100644 index 00000000000..57da5f4691b --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -0,0 +1,14 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visits that plan or execute the concrete row signature selected by [`RowFn::dispatch`]. +//! +//! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch + +mod check; + +mod plan; +pub(super) use plan::BatchPlanner; + +mod row_visitor; +pub use row_visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs new file mode 100644 index 00000000000..c522376d491 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -0,0 +1,173 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Plans the concrete signature selected by [`RowFn::dispatch`]. +//! +//! [`BatchPlanner`] validates input and output dtypes, then records the output dtype and +//! null-handling policy that execution must reproduce. + +use std::marker::PhantomData; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +use super::RowVisitor; +use super::check::assert_deferred_visit_contract; +use super::check::assert_owned_visit_contract; +use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; +use super::row_visitor::private; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::ElementTuple; +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::SinkResult; + +/// A planning visitor that validates dtypes and selects the nullable execution policy. +pub(crate) struct BatchPlanner<'a, F: RowFn> { + dtypes: &'a [DType], + + options: &'a F::Options, + + /// Ties the planner to the function used by its compile-time contract checks. + function: PhantomData, +} + +impl<'a, F: RowFn> BatchPlanner<'a, F> { + pub(crate) fn new(dtypes: &'a [DType], options: &'a F::Options) -> Self { + Self { + dtypes, + options, + function: PhantomData, + } + } +} + +impl private::Sealed for BatchPlanner<'_, F> {} + +impl RowVisitor for BatchPlanner<'_, F> { + type VisitResult = BatchPlan; + + fn visit_prepared( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + const { assert_owned_visit_contract::() }; + Ok(BatchPlan { + output_dtype: validate_owned_visit::(self.dtypes)?, + policy: RowPolicy::for_owned_output::(), + }) + } + + fn visit_prepared_into( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + Ok(BatchPlan { + output_dtype: validate_sink_visit::(self.options, self.dtypes)?, + policy: RowPolicy::for_sink::(), + }) + } + + fn visit_prepared_deferred( + self, + _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + _apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + _finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + const { assert_deferred_visit_contract::() }; + Ok(BatchPlan { + output_dtype: validate_owned_visit::(self.dtypes)?, + policy: RowPolicy::for_deferred_output::(), + }) + } +} + +/// The execution policy and output dtype selected by a planning visit. +pub(crate) struct BatchPlan { + /// The non-nullable dtype built by the selected output capability. + pub(crate) output_dtype: DType, + + /// How this concrete dispatch executes nullable rows. + // TODO(connor)[RowFn]: Remove this allowance when the execution backend from #9130 consumes + // this policy. + #[allow(dead_code)] + pub(crate) policy: RowPolicy, +} + +impl BatchPlan { + /// Return the output dtype widened with strict input nullability. + pub(crate) fn result_dtype(self, args: &[DType]) -> DType { + let nullability = self.output_dtype.nullability() + | Nullability::from(args.iter().any(DType::is_nullable)); + + self.output_dtype.with_nullability(nullability) + } +} + +/// The nullable execution policy derived from one concrete dispatch. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum RowPolicy { + /// Evaluate all rows and mask the result. + Dense, + + /// Evaluate all rows, retrying only valid rows if a deferred error is raised. + DenseWithRetry, + + /// Execute only valid rows, trying skip-invalid execution before filtering. + ValidOnly, +} + +impl RowPolicy { + /// The policy for an infallible owned output. + pub(crate) const fn for_owned_output() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + Self::Dense + } else { + Self::ValidOnly + } + } + + /// The policy for an owned output carrying batch-deferred failure evidence. + pub(crate) const fn for_deferred_output() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + Self::DenseWithRetry + } else { + Self::ValidOnly + } + } + + /// The policy for a sink-writing output. + pub(crate) const fn for_sink() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { + Self::Dense + } else { + Self::ValidOnly + } + } +} 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 new file mode 100644 index 00000000000..0e7abaa4322 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -0,0 +1,298 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The typed dispatch interface implemented by [`RowFn`] planning and execution. +//! +//! [`RowVisitor`] lets a function select its concrete input and output capabilities without +//! exposing framework-specific planning or execution state. +//! +//! [`RowFn`]: crate::scalar_fn::unstable::row::RowFn + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +use crate::scalar_fn::unstable::row::ElementTuple; +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::SinkResult; + +/// A planning or execution visit at concrete input and output types. +/// +/// Only the framework implements this trait. The `visit_prepared*` methods derive shared state +/// from constant arguments before visiting any rows. Every visit verifies that the argument tuple +/// matches [`RowFn::ARG_NAMES`] and that fallible decoding agrees with [`RowFn::FALLIBLE`]. +/// +/// [`RowFn::ARG_NAMES`]: crate::scalar_fn::unstable::row::RowFn::ARG_NAMES +/// [`RowFn::FALLIBLE`]: crate::scalar_fn::unstable::row::RowFn::FALLIBLE +pub trait RowVisitor: private::Sealed + Sized { + /// The framework result of visiting one concrete row signature. + /// + /// This is a batch plan or execution result, not a per-row output. + type VisitResult; + + /// Visit an infallible row computation that returns one output value per row. + /// + /// `apply` must not panic or have side effects. Dense execution can pass unspecified values + /// from null rows. + /// + /// The framework verifies that `Out` does not require drop glue. + /// + /// # Examples + /// + /// Apply infallible wrapping arithmetic. + /// + /// ```ignore + /// visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) + /// ``` + /// + /// Dispatch an equality helper over its primitive element type. + /// + /// ```ignore + /// fn visit_equal(visitor: V) -> VortexResult + /// where + /// T: NativePType, + /// V: RowVisitor, + /// { + /// visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_eq(rhs)) + /// } + /// ``` + fn visit( + self, + apply: impl Fn(Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + { + self.visit_prepared::(|_| (), move |&(), args| apply(args)) + } + + /// The prepared form of [`visit`](Self::visit), with the same prerequisites. + /// + /// # Examples + /// + /// Test whether each string occurs in its allowed-values list. The prepare closure builds one + /// lookup table for a batch-constant list. The row closure scans a varying list directly. + /// + /// ```ignore + /// visitor.visit_prepared::< + /// (StringRow, StringListRow), + /// bool, + /// Option, + /// >( + /// |(_value, allowed_values)| allowed_values.map(PreparedAllowedValues::new), + /// |prepared_allowed_values, (value, allowed_values)| { + /// match prepared_allowed_values { + /// Some(allowed_values) => allowed_values.contains(value), + /// None => allowed_values.iter().any(|allowed| allowed == value), + /// } + /// }, + /// ) + /// ``` + fn visit_prepared( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement; + + /// Visit a row computation that writes through a row handle from an output sink. + /// + /// `apply` must not panic or have side effects except for writes to the supplied row handle. + /// Dense execution can pass unspecified values from null rows. + /// + /// On success, `apply` must return the write token for the supplied row handle. A token from + /// another row, sink, or local cell can violate the safety contract of [`OutputSink::finish`]. + /// + /// A fallible `ApplyResult` requires + /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) to be `true`. + /// + /// # Examples + /// + /// Checked integer division reports errors immediately and writes successful rows into + /// uninitialized output. The cold, non-inlined helper keeps error construction out of the row + /// callback. + /// + /// ```ignore + /// #[cold] + /// #[inline(never)] + /// fn integer_division_error() -> VortexError { + /// vortex_err!(InvalidArgument: "integer division by zero or overflow") + /// } + /// + /// visitor.visit_into::<(i64, i64), UninitElementSink, _>( + /// |(lhs, rhs), output| { + /// let Some(value) = lhs.checked_div(rhs) else { + /// return Err(integer_division_error()); + /// }; + /// + /// // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + /// Ok(unsafe { InitializedElement::write(output, value) }) + /// }, + /// ) + /// ``` + fn visit_into( + self, + apply: impl Fn(Args::Elems<'_>, >::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + self.visit_prepared_into::( + |_| (), + move |&(), args, row| apply(args, row), + ) + } + + /// The prepared form of [`visit_into`](Self::visit_into), with the same prerequisites. + /// + /// # Examples + /// + /// Compute the cosine similarity of each vector pair: their dot product divided by their + /// magnitudes. The prepare closure computes each batch-constant vector's magnitude once. + /// + /// ```ignore + /// visitor.visit_prepared_into::< + /// (TensorRow, TensorRow), + /// UninitElementSink, + /// ConstantVectorMagnitudes, + /// InitializedElement, + /// >( + /// |(lhs, rhs)| ConstantVectorMagnitudes { + /// lhs: lhs.map(vector_magnitude), + /// rhs: rhs.map(vector_magnitude), + /// }, + /// |constant_magnitudes, (lhs, rhs), output| { + /// let similarity = + /// cosine_similarity_with_constant_magnitudes(constant_magnitudes, lhs, rhs); + /// + /// // SAFETY: `output` is the `UninitElementSink` row supplied for this callback. + /// unsafe { InitializedElement::write(output, similarity) } + /// }, + /// ) + /// ``` + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn( + &Prepared, + Args::Elems<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>; + + /// Visit a row computation that returns an owned output value and deferred failure evidence. + /// + /// `apply` must not panic or have side effects. Dense execution can pass unspecified values + /// from null rows. + /// + /// The executor OR-reduces `Fail` across rows and passes the result to `finish_failure`. + /// [`Default::default`] **must** mean success, including for an empty batch. The compiler + /// cannot check this requirement. + /// + /// [`RowFn::FALLIBLE`](crate::scalar_fn::unstable::row::RowFn::FALLIBLE) **must** be `true`. + /// `Out` must not require drop glue. `Fail` must be no wider than `Out`, or failure tracking + /// reduces the vector width. The framework checks these requirements. + /// + /// # Examples + /// + /// Checked addition returns a wrapping value and compact overflow flag without branching. The + /// executor reduces the flags after the loop. The cold, non-inlined helper keeps error + /// construction out of the row loop. + /// + /// ```ignore + /// #[cold] + /// #[inline(never)] + /// fn integer_addition_error() -> VortexError { + /// vortex_err!(InvalidArgument: "integer overflow in checked add") + /// } + /// + /// visitor.visit_deferred::<(i64, i64), i64, bool>( + /// // `overflowing_add` returns `(i64, bool)`. + /// |(lhs, rhs)| lhs.overflowing_add(rhs), + /// |overflowed| { + /// if overflowed { + /// return Err(integer_addition_error()); + /// } + /// + /// Ok(()) + /// }, + /// ) + /// ``` + fn visit_deferred( + self, + apply: impl Fn(Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign, + { + self.visit_prepared_deferred::( + |_| (), + move |&(), args| apply(args), + finish_failure, + ) + } + + /// The prepared form of [`visit_deferred`](Self::visit_deferred), with the same prerequisites. + /// + /// # Examples + /// + /// Rescale each unscaled decimal by multiplying it by `10^scale`. The prepare closure computes + /// the multiplier once for a batch-constant scale. Each row returns the rescaled value and an + /// overflow flag, which the executor reduces after the loop. + /// + /// ```ignore + /// #[cold] + /// #[inline(never)] + /// fn decimal_rescaling_overflow() -> VortexError { + /// vortex_err!(InvalidArgument: "decimal rescaling overflowed") + /// } + /// + /// visitor.visit_prepared_deferred::< + /// (i64, DecimalScale), + /// i64, + /// Option, + /// bool, + /// >( + /// |(_value, scale)| scale.map(PreparedDecimalScale::new), + /// |prepared_scale, (value, scale)| match prepared_scale { + /// Some(scale) => scale.apply_checked(value), + /// None => PreparedDecimalScale::new(scale).apply_checked(value), + /// }, + /// |overflowed| { + /// if overflowed { + /// return Err(decimal_rescaling_overflow()); + /// } + /// + /// Ok(()) + /// }, + /// ) + /// ``` + fn visit_prepared_deferred( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, + ) -> VortexResult + where + Args: IndexedElementTuple, + Out: OutputElement, + Fail: Copy + Default + BitOrAssign; +} + +pub(super) mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs new file mode 100644 index 00000000000..b6c14585c48 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Adapts [`RowFn`] implementations to the scalar-function interface. +//! +//! The blanket [`ScalarFnVTable`] implementation supplies common arity, validity, fallibility, and +//! execution behavior. [`row_fn_return_dtype`] and [`execute_rows`] expose the same planning and +//! execution paths to public vtables that delegate to a private row kernel. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_session::VortexSession; + +use super::row_fn::RowFn; +use super::visitor::BatchPlanner; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::expr::Expression; +use crate::expr::union_child_validities; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; + +impl ScalarFnVTable for F { + type Options = F::Options; + + fn id(&self) -> ScalarFnId { + RowFn::id(self) + } + + fn serialize(&self, options: &Self::Options) -> VortexResult>> { + RowFn::serialize(self, options) + } + + fn deserialize(&self, metadata: &[u8], session: &VortexSession) -> VortexResult { + RowFn::deserialize(self, metadata, session) + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(F::ARG_NAMES.len()) + } + + fn child_name(&self, _options: &Self::Options, child_index: usize) -> ChildName { + ChildName::from(F::ARG_NAMES[child_index]) + } + + fn return_dtype(&self, options: &Self::Options, args: &[DType]) -> VortexResult { + row_fn_return_dtype(self, options, args) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + execute_rows(self, options, args, ctx) + } + + fn validity( + &self, + _options: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + F::FALLIBLE + } +} + +/// Compute the return dtype of a [`RowFn`] kernel without invoking its blanket vtable. +pub fn row_fn_return_dtype( + function: &F, + options: &F::Options, + args: &[DType], +) -> VortexResult { + ensure_arity(function, args.len())?; + + let plan = function.dispatch(options, args, BatchPlanner::::new(args, options))?; + + Ok(plan.result_dtype(args)) +} + +/// Execute a [`RowFn`] without using its blanket [`ScalarFnVTable`] implementation. +/// +/// A type cannot implement both [`RowFn`] and [`ScalarFnVTable`] because every `RowFn` receives the +/// standard vtable automatically. Existing vtables can keep their custom hooks on one type and +/// delegate row execution to a private `RowFn` kernel through this function. +pub fn execute_rows( + function: &F, + _options: &F::Options, + args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, +) -> VortexResult { + ensure_arity(function, args.num_inputs())?; + + // TODO(connor)[RowFn]: Replace this temporary error with the execution backend in #9129. + vortex_bail!( + "Row function {} does not yet have an execution backend", + RowFn::id(function) + ) +} + +/// Validate the number of arguments before calling user-defined dispatch code. +fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { + let expected = F::ARG_NAMES.len(); + vortex_ensure_eq!( + actual, + expected, + "row function {} requires arity {expected}, got {actual}", + RowFn::id(function), + ); + + Ok(()) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexError; + use vortex_error::VortexResult; + use vortex_session::registry::CachedId; + + use super::execute_rows; + use super::row_fn_return_dtype; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::dtype::DType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnId; + use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::RowFn; + use crate::scalar_fn::unstable::row::RowVisitor; + + #[derive(Clone)] + struct IndexingRowFn; + + impl RowFn for IndexingRowFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.indexing_row_fn"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + _ = &args[0]; + + visitor.visit::<(i64,), i64>(|(value,)| value) + } + } + + #[test] + fn test_return_dtype_rejects_wrong_arity_before_dispatch() { + let error = row_fn_return_dtype(&IndexingRowFn, &EmptyOptions, &[]) + .expect_err("wrong arity must fail before dispatch"); + + assert_arity_error(error); + } + + #[test] + fn test_execute_rejects_wrong_arity_before_dispatch() { + let args = VecExecutionArgs::new(vec![], 0); + let mut ctx = array_session().create_execution_ctx(); + let error = execute_rows(&IndexingRowFn, &EmptyOptions, &args, &mut ctx) + .expect_err("wrong arity must fail before dispatch"); + + assert_arity_error(error); + } + + #[track_caller] + fn assert_arity_error(error: VortexError) { + assert!( + error.to_string().contains("requires arity 1, got 0"), + "unexpected error: {error}", + ); + } +} diff --git a/vortex/Cargo.toml b/vortex/Cargo.toml index 6a2a840a500..57392ed627d 100644 --- a/vortex/Cargo.toml +++ b/vortex/Cargo.toml @@ -90,6 +90,8 @@ tokio = [ zstd = ["dep:vortex-zstd", "vortex-file?/zstd"] pretty = ["vortex-array/table-display"] serde = ["vortex-array/serde", "vortex-buffer/serde", "vortex-mask/serde"] +# Exposes experimental row-function APIs without compatibility guarantees. +unstable_row_fns = ["vortex-array/unstable_row_fns"] # This feature enabled unstable encodings for which we don't guarantee stability. unstable_encodings = [ "dep:vortex-tensor",