diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index d00b811a387..6bb251bd808 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -129,6 +129,10 @@ harness = false name = "binary_ops" harness = false +[[bench]] +name = "row_fn_executor" +harness = false + [[bench]] name = "kleene_bool" harness = false @@ -203,6 +207,10 @@ harness = false name = "validity_is_valid" harness = false +[[bench]] +name = "strict_validity" +harness = false + [[bench]] name = "dict_unreferenced_mask" harness = false diff --git a/vortex-array/benches/row_fn_executor.rs b/vortex-array/benches/row_fn_executor.rs new file mode 100644 index 00000000000..2ce9323f09e --- /dev/null +++ b/vortex-array/benches/row_fn_executor.rs @@ -0,0 +1,280 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares owned-output, sink-writing, and hand-written primitive row loops. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::NativePType; +use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::OutputSink; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const ROWS: usize = 65_536; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +#[derive(Clone)] +struct RowWrappingAdd; + +impl RowFn for RowWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_wrapping_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64), i64>(|(lhs, rhs)| lhs.wrapping_add(rhs)) + } +} + +#[derive(Clone)] +struct RowCheckedAdd; + +impl RowFn for RowCheckedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_checked_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(i64, i64), i64, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(checked_add_error()); + } + Ok(()) + }, + ) + } +} + +/// Keep error construction out of the benchmarked success path. +#[cold] +#[inline(never)] +fn checked_add_error() -> VortexError { + vortex_err!("integer overflow in row checked add") +} + +/// A benchmark sink that writes one `i64` per row. +struct I64Sink( + /// The output values written by the row loop. + BufferMut, +); + +impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +#[derive(Clone)] +struct RowSinkWrappingAdd; + +impl RowFn for RowSinkWrappingAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.row_sink_wrapping_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64, i64), I64Sink, _>(|(lhs, rhs), out| { + *out = lhs.wrapping_add(rhs); + }) + } +} + +fn inputs() -> (ArrayRef, ArrayRef) { + let lhs = (0..ROWS) + .map(|index| index as i64) + .collect::>() + .into_array(); + let rhs = (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>() + .into_array(); + (lhs, rhs) +} + +fn constant_inputs() -> (ArrayRef, ArrayRef) { + let (lhs, _) = inputs(); + let rhs = ConstantArray::new(Scalar::from(7i64), ROWS).into_array(); + (lhs, rhs) +} + +fn nullable_inputs() -> (ArrayRef, ArrayRef) { + let lhs = PrimitiveArray::new( + (0..ROWS).map(|index| index as i64).collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(5))), + ) + .into_array(); + let rhs = PrimitiveArray::new( + (0..ROWS) + .map(|index| (index % 17) as i64) + .collect::>(), + Validity::from_iter((0..ROWS).map(|index| !index.is_multiple_of(7))), + ) + .into_array(); + (lhs, rhs) +} + +fn bench_row_fn(bencher: Bencher, function: F, make_inputs: fn() -> (ArrayRef, ArrayRef)) +where + F: RowFn, +{ + bencher + .with_inputs(make_inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + function + .clone() + .try_new_array(ROWS, EmptyOptions, [lhs, rhs]) + .unwrap() + .into_array() + .execute::(&mut ctx) + .unwrap() + }); +} + +#[divan::bench] +fn row_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, inputs); +} + +#[divan::bench] +fn row_sink_wrapping_add(bencher: Bencher) { + bench_row_fn(bencher, RowSinkWrappingAdd, inputs); +} + +#[divan::bench] +fn handrolled_sink_wrapping_add(bencher: Bencher) { + bencher + .with_inputs(inputs) + .bench_local_values(|(lhs, rhs)| { + let mut ctx = SESSION.create_execution_ctx(); + let lhs = lhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let rhs = rhs + .execute::(&mut ctx) + .unwrap() + .into_buffer::(); + let mut output = BufferMut::zeroed(ROWS); + for ((out, lhs), rhs) in output + .as_mut_slice() + .iter_mut() + .zip(lhs.as_slice()) + .zip(rhs.as_slice()) + { + *out = lhs.wrapping_add(*rhs); + } + PrimitiveArray::new(output.freeze(), Validity::NonNullable).into_array() + }); +} + +#[divan::bench] +fn row_checked_add(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, inputs); +} + +#[divan::bench] +fn row_wrapping_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_constant(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, constant_inputs); +} + +#[divan::bench] +fn row_checked_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowCheckedAdd, nullable_inputs); +} + +#[divan::bench] +fn row_wrapping_add_nullable(bencher: Bencher) { + bench_row_fn(bencher, RowWrappingAdd, nullable_inputs); +} diff --git a/vortex-array/benches/strict_validity.rs b/vortex-array/benches/strict_validity.rs new file mode 100644 index 00000000000..b2fa7a3a824 --- /dev/null +++ b/vortex-array/benches/strict_validity.rs @@ -0,0 +1,215 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks how the row lifting applies validity to a dense kernel. +//! +//! Both arms run the same kernel over the same nullable input and differ only in how the output's +//! nulls are restored: +//! +//! - `lazy` is what the lifting behind [`RowFn`] does: conjoin the input validities (itself lazy) +//! and hand the resulting boolean array straight to `mask`. +//! - `eager` is the strategy it replaced: materialize the conjunction into a `Mask`, copy it into a +//! `BitBuffer`, wrap that in a `BoolArray`, and mask with it. +//! +//! `chain` composes three of the same function, which is where a per-call materialization +//! compounds. + +#![expect(clippy::unwrap_used)] +#![expect(clippy::cast_possible_truncation)] + +use std::sync::LazyLock; + +use divan::Bencher; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::array_session; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::scalar_fn::ScalarFnFactoryExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::RowExecution; +use vortex_array::scalar_fn::RowFn; +use vortex_array::scalar_fn::RowVisitor; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +const SIZES: &[usize] = &[65_536, 1 << 20]; + +static SESSION: LazyLock = LazyLock::new(array_session); + +fn main() { + LazyLock::force(&SESSION); + divan::main(); +} + +/// The shared kernel: double every lane, ignoring validity. +fn doubled(input: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let input = input.clone().execute::(ctx)?; + let values: Buffer = input + .as_slice::() + .iter() + .map(|value| value.wrapping_mul(2)) + .collect(); + Ok(PrimitiveArray::new(values, Validity::NonNullable).into_array()) +} + +/// Dense row function: the lifting decides how validity is applied. +/// +/// Its [`reduce_encoded`](RowFn::reduce_encoded) answers with [`doubled`] before the row loop is +/// reached, so this arm runs exactly the kernel [`EagerDouble`] runs and the two differ only in +/// validity. The row closure below is what makes it a [`RowFn`] at all, and never executes. +#[derive(Clone)] +struct LazyDouble; + +impl RowFn for LazyDouble { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["input"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.lazy_double"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i32,), i32>(|(value,)| value.wrapping_mul(2)) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + doubled(&args[0], ctx).map(|output| Some(RowExecution::Output(output))) + } +} + +/// The same function, applying validity the way the adapter used to: materialize a mask first. +#[derive(Clone)] +struct EagerDouble; + +impl ScalarFnVTable for EagerDouble { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("bench.eager_double"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + ChildName::from("input") + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::I32, args[0].nullability())) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let valid = input.validity()?.execute_mask(args.row_count(), ctx)?; + let values = doubled(&input, ctx)?; + + if valid.all_true() { + return values.cast(DType::Primitive(PType::I32, input.dtype().nullability())); + } + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + values.mask(mask) + } + + 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 { + false + } +} + +/// A nullable i32 column with array-backed validity (~10% nulls). +fn nullable_input(len: usize) -> ArrayRef { + PrimitiveArray::new( + (0..len as i32).collect::>(), + Validity::from_iter((0..len).map(|index| !index.is_multiple_of(10))), + ) + .into_array() +} + +fn bench_depth(bencher: Bencher, function: F, len: usize, depth: usize) +where + F: ScalarFnVTable + Clone, +{ + bencher + .with_inputs(|| nullable_input(len)) + .bench_local_values(|input| { + let mut ctx = SESSION.create_execution_ctx(); + let mut array = input; + for _ in 0..depth { + array = function + .clone() + .try_new_array(len, EmptyOptions, [array]) + .unwrap() + .into_array(); + } + array.execute::(&mut ctx).unwrap() + }); +} + +#[divan::bench(args = SIZES)] +fn lazy(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn eager(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 1); +} + +#[divan::bench(args = SIZES)] +fn lazy_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, LazyDouble, len, 3); +} + +#[divan::bench(args = SIZES)] +fn eager_chain(bencher: Bencher, len: usize) { + bench_depth(bencher, EagerDouble, len, 3); +} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 003dbc2ca48..5e73caefdfa 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -6,6 +6,11 @@ //! This module contains the [`ScalarFnVTable`] trait and all built-in scalar function //! implementations. Expressions ([`crate::expr::Expression`]) reference scalar functions //! at each node. +//! +//! Use [`RowFn`] for strict functions whose natural kernel computes one row at a time. It derives +//! decoding, constant handling, null propagation, output construction, and validity. Implement +//! [`ScalarFnVTable`] directly when the natural kernel is columnar, aliases an input, or may +//! produce null from otherwise valid inputs. use vortex_session::registry::Id; @@ -35,6 +40,9 @@ pub use options::*; mod signature; pub use signature::*; +mod row; +pub use row::*; + pub mod fns; pub mod internal; pub mod session; diff --git a/vortex-array/src/scalar_fn/row/batch/args.rs b/vortex-array/src/scalar_fn/row/batch/args.rs new file mode 100644 index 00000000000..c908b91e00e --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/args.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Input views and planning metadata passed to a row kernel. + +use crate::ArrayRef; +use crate::dtype::DType; + +/// The arguments handed to one kernel invocation. +/// +/// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the +/// original planned batch. Keeping them together prevents an execution path from pairing an input +/// view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub struct KernelArgs<'a> { + /// The input arrays for this kernel invocation. + pub arrays: &'a [ArrayRef], + + /// The number of rows in this kernel invocation. + pub row_count: usize, + + /// The original input dtypes used to select the row implementation. + pub dtypes: &'a [DType], + + /// The non-nullable dtype built by the selected output capability. + pub output_dtype: &'a DType, +} diff --git a/vortex-array/src/scalar_fn/row/batch/execution.rs b/vortex-array/src/scalar_fn/row/batch/execution.rs new file mode 100644 index 00000000000..fd4b0401fb3 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -0,0 +1,479 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Null propagation, constant folding, and strategy execution for one columnar batch. + +use smallvec::SmallVec; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::args::KernelArgs; +use super::policy::BatchPlan; +use super::policy::RowPolicy; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::types::batch_constant; +use crate::validity::Validity; + +/// How far [`Batch::resolve_validity`] got before a mixed-mask strategy became necessary. +enum ResolvedMask { + /// The all-valid or all-null batch was answered without a mixed-mask strategy. + Decided(ArrayRef), + + /// A mask with both set and unset bits, which a strategy must now execute. + Mixed(Mask), +} + +/// One batch of inputs and the metadata needed before its row kernel runs. +pub struct Batch { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The number of rows in the original execution scope. + row_count: usize, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + +impl Batch { + /// Collect the inputs and derive their dtype, validity, and execution policy. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check below would vacuously pass. + pub fn new( + id: ScalarFnId, + args: &dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let row_count = args.row_count(); + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + for (index, input) in inputs.iter().enumerate() { + vortex_ensure_eq!( + input.len(), + row_count, + "the {id} input {index} must have {row_count} rows, got {}", + input.len(), + ); + } + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let result_dtype = plan.result_dtype(&arg_dtypes); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + row_count, + inputs, + arg_dtypes, + validity, + result_dtype, + output_dtype: plan.output_dtype, + policy: plan.policy, + }) + } + + /// Add null propagation, constant folding, and strategy selection around `kernel`. + /// + /// The kernel may ignore input validity. It receives valid-only rows when required, and its + /// output **must** match the planned dtype up to nullability. `reduce` receives the original + /// inputs exactly once, before the generic all-constant broadcast, so a function-owned encoded + /// implementation takes precedence. `try_unfiltered` receives the originals plus a mixed + /// validity mask; `Ok(None)` selects filter-and-scatter. + pub fn execute( + &self, + reduce: impl FnOnce(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult>, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.validity, Validity::AllInvalid) + || self + .inputs + .iter() + .any(|input| input.as_constant().is_some_and(|scalar| scalar.is_null())) + { + return Ok(self.all_null()); + } + + // The function-owned encoded path takes precedence over the generic all-constant + // broadcast and sees the original inputs before slicing or filtering changes them. + if let Some(execution) = reduce(self.kernel_args(&self.inputs, self.row_count), ctx)? { + match execution { + RowExecution::Output(values) => return self.finalize_reduced(values), + RowExecution::DeferredError(error) => { + return self.resolve_reduced_error(error, kernel, try_unfiltered, ctx); + } + } + } + + // All inputs constant, and their conjoined validity proves every row non-null. This sees + // through extension and masked wrappers just like argument decoding does. + if self.row_count > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), + } + } + + /// Evaluate a single row of all-constant inputs and broadcast its value. + /// + /// Reconciling the row's dtype before reading the scalar keeps this path on the same + /// kernel/declaration agreement check as the dense and filter paths, rather than letting `cast` + /// paper over a disagreement. + fn broadcast_one_row( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let result = VortexResult::from(kernel(self.kernel_args(&one_row, 1), ctx)?)?; + let scalar = self.finalize_output(result, 1)?.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.row_count).into_array()) + } + + /// Run the kernel over every row, including the rows behind nulls, then mask its result. + /// + /// The arguments reach the kernel untouched, so the inputs keep their original encoding, and + /// the conjoined validity is handed to `mask` as an array rather than materialized into a + /// [`Mask`] first. + fn execute_dense( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let values = match kernel(self.kernel_args(&self.inputs, self.row_count), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Unlike `resolve_validity`, all-true preserves the deferred error and all-false + // suppresses evidence that came entirely from null rows. An empty loop cannot + // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Deferred retry receives only the dense kernel. Filter first so the retry cannot + // evaluate null rows again. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled by the guard above, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } + + /// Materialize validity and answer all-valid or all-null batches before selecting a mixed-mask + /// strategy. + fn resolve_validity( + &self, + kernel: &impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + return self + .finalize_output( + VortexResult::from(kernel( + self.kernel_args(&self.inputs, self.row_count), + ctx, + )?)?, + self.row_count, + ) + .map(ResolvedMask::Decided); + } + + if valid.all_false() { + return Ok(ResolvedMask::Decided(self.all_null())); + } + + Ok(ResolvedMask::Mixed(valid)) + } + + /// Resolve validity, try unfiltered execution, then fall back to filtering. + fn execute_valid_only( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedMask::Decided(result) => return Ok(result), + ResolvedMask::Mixed(valid) => valid, + }; + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Try execution against the original inputs, then mask a returned full-length result. + fn try_execute_unfiltered( + &self, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(execution) = + try_unfiltered(self.kernel_args(&self.inputs, self.row_count), valid, ctx)? + else { + return Ok(None); + }; + let values = VortexResult::from(execution)?; + + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + self.finalize_output(values.mask(mask)?, valid.len()) + .map(Some) + } + + /// The filter strategy for a mixed mask: filter every input down to the rows set in `valid`, + /// run the kernel over those, and scatter its results back into a null-padded output. + fn filter_and_scatter( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let values = VortexResult::from(kernel( + self.kernel_args(&filtered, valid.true_count()), + ctx, + )?)?; + + self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) + } + + /// An all-null result of the function's declared return dtype. + fn all_null(&self) -> ArrayRef { + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + } + + /// Reconcile an encoding-aware result and apply the batch's strict input validity. + fn finalize_reduced(&self, values: ArrayRef) -> VortexResult { + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled before the encoding-aware hook runs. + Validity::AllInvalid => Ok(self.all_null()), + } + } + + /// Resolve deferred evidence from the encoded path by executing only observable rows. + fn resolve_reduced_error( + &self, + error: VortexError, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Pair an input view with this batch's planning metadata. + fn kernel_args<'b>(&'b self, arrays: &'b [ArrayRef], row_count: usize) -> KernelArgs<'b> { + KernelArgs { + arrays, + row_count, + dtypes: &self.arg_dtypes, + output_dtype: &self.output_dtype, + } + } + + /// Finalize an output against this batch's expected length and declared return dtype. + fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { + finalize_kernel_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Scatter `values` (one per set bit of `valid`, in order) back to the positions of the set + /// bits, producing an array of length `valid.len()` that is null at every unset position. + fn scatter_valid(&self, values: ArrayRef, valid: &Mask) -> VortexResult { + vortex_ensure_eq!( + values.len(), + valid.true_count(), + "the {} kernel produced {} rows for {} filtered rows", + self.id, + values.len(), + valid.true_count(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!("scatter_valid requires a mixed mask"); + }; + + // Gather indices: row i of the output reads values[rank(i)]. Rows behind nulls read index + // 0, and any in-bounds index would do since they are masked out below (values is non-empty + // here). + let mut indices = vec![0u64; valid.len()]; + let mut rank = 0u64; + for &(start, end) in slices { + for index in &mut indices[start..end] { + *index = rank; + rank += 1; + } + } + let indices = PrimitiveArray::new(indices, Validity::NonNullable).into_array(); + + let scattered = values.take(indices)?; + + // A kernel that produced nulls of its own (only `reduce_encoded` may) cannot be wrapped, + // since a `Masked` child must be all valid. Those nulls have to be unioned with the + // batch validity, which is what the general masking pass does. + if scattered.dtype().is_nullable() { + let mask = BoolArray::new(valid.to_bit_buffer(), Validity::NonNullable).into_array(); + return scattered.mask(mask); + } + + // The gathered values are all valid, so attaching validity is sufficient. + Ok(MaskedArray::try_new( + scattered, + Validity::from_mask(valid.clone(), Nullability::Nullable), + )? + .into_array()) + } +} + +/// Validate a kernel output, then cast it to the row function's declared nullability. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring +/// nullability. The kernel may omit nullability because batch execution owns strict null +/// propagation, so a nullability-only difference is cast to `result_dtype`. +pub fn finalize_kernel_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + vortex_ensure_eq!( + values.len(), + expected_len, + "the {id} kernel produced {} rows for {expected_len} input rows", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel produced {} but the function declares {result_dtype}", + values.dtype(), + ); + + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} diff --git a/vortex-array/src/scalar_fn/row/batch/mod.rs b/vortex-array/src/scalar_fn/row/batch/mod.rs new file mode 100644 index 00000000000..3dd7b08603f --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/mod.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution around a non-null row kernel. +//! +//! A row kernel handles typed values for one row. This module adds the columnar concerns around it: +//! planning the output and null strategy, preserving batch constants and encodings, propagating +//! strict validity, selecting an execution strategy, and validating the finished output. +//! +//! [`policy`] derives the nullable execution strategy from a concrete dispatch. [`execution`] +//! applies that strategy, and [`args`] pairs each kernel invocation with its planning metadata. + +mod args; +pub(super) use args::KernelArgs; + +mod execution; +pub(super) use execution::Batch; +pub(super) use execution::finalize_kernel_output; + +mod policy; +pub(super) use policy::BatchPlan; +pub(super) use policy::RowPolicy; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/row/batch/policy.rs b/vortex-array/src/scalar_fn/row/batch/policy.rs new file mode 100644 index 00000000000..d7e024ce763 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Nullable execution strategies derived from a concrete row dispatch. + +use crate::dtype::DType; +use crate::dtype::Nullability; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::SinkResult; + +/// The execution policy and output dtype selected by a planning visit. +pub struct BatchPlan { + /// The non-nullable dtype built by the selected output capability. + pub output_dtype: DType, + + /// How this concrete dispatch executes nullable rows. + pub policy: RowPolicy, +} + +impl BatchPlan { + /// Return the output dtype widened with strict input nullability. + pub 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 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 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 const fn for_deferred_output() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE { + Self::DenseWithRetry + } else { + Self::ValidOnly + } + } + + /// The policy one concrete dispatch executes nullable rows under. + /// + /// Batch execution always tries [`reduce_encoded`](crate::scalar_fn::RowFn::reduce_encoded) + /// against the original arrays before it tries the sink or filters the inputs. Skipping that + /// probe can change the result of an encoding-aware function. + pub const fn for_sink() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { + Self::Dense + } else { + Self::ValidOnly + } + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::RowPolicy; + use crate::ArrayRef; + use crate::ExecutionCtx; + use crate::dtype::DType; + use crate::scalar_fn::InputElement; + + struct SparseFallibleElement; + + // SAFETY: the varying view reports length zero, so no index satisfies the unchecked-read + // precondition. + unsafe impl InputElement for SparseFallibleElement { + type Column = (); + type Varying<'a> = (); + type Elem<'a> = (); + + const DENSE_SAFE: bool = false; + const DECODE_FALLIBLE: bool = true; + + fn validate(_dtype: &DType) -> VortexResult<()> { + Ok(()) + } + + fn decode(_array: ArrayRef, _ctx: &mut ExecutionCtx) -> VortexResult { + Ok(()) + } + + fn get(_column: &Self::Column, _index: usize) -> Self::Elem<'_> {} + + fn varying(_column: &Self::Column) -> Self::Varying<'_> {} + + fn varying_len(_column: &Self::Varying<'_>) -> usize { + 0 + } + + fn get_varying<'a>(_column: &Self::Varying<'a>, _index: usize) -> Self::Elem<'a> {} + } + + #[test] + fn test_owned_output_policy() { + assert_eq!(RowPolicy::for_owned_output::<(i64,)>(), RowPolicy::Dense); + assert_eq!( + RowPolicy::for_owned_output::<(SparseFallibleElement,)>(), + RowPolicy::ValidOnly, + ); + } + + #[test] + fn test_deferred_output_policy() { + assert_eq!( + RowPolicy::for_deferred_output::<(i64,)>(), + RowPolicy::DenseWithRetry, + ); + assert_eq!( + RowPolicy::for_deferred_output::<(SparseFallibleElement,)>(), + RowPolicy::ValidOnly, + ); + } + + #[test] + fn test_sink_policy() { + assert_eq!(RowPolicy::for_sink::<(i64,), ()>(), RowPolicy::Dense); + assert_eq!( + RowPolicy::for_sink::<(i64,), VortexResult<()>>(), + RowPolicy::ValidOnly, + ); + assert_eq!( + RowPolicy::for_sink::<(SparseFallibleElement,), ()>(), + RowPolicy::ValidOnly, + ); + } +} diff --git a/vortex-array/src/scalar_fn/row/batch/tests.rs b/vortex-array/src/scalar_fn/row/batch/tests.rs new file mode 100644 index 00000000000..a2da250e957 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/tests.rs @@ -0,0 +1,517 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use rstest::rstest; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::super::execute::RowExecution; +use super::Batch; +use super::BatchPlan; +use super::RowPolicy; +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; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::RowVisitor; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::VecExecutionArgs; +use crate::validity::Validity; + +#[derive(Clone)] +struct RetryConstantAdd; + +#[derive(Clone)] +struct NullarySeven; + +#[derive(Clone)] +struct OriginalInputReducer; + +#[derive(Clone)] +struct DeferredOriginalReducer; + +#[derive(Clone)] +struct PreparedAdd { + visit: PreparedVisit, + prepares: Arc, +} + +#[derive(Clone, Copy)] +enum PreparedVisit { + Owned, + Sink, + Deferred, +} + +struct I64Sink(BufferMut); + +impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +impl RowFn for NullarySeven { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_seven"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), I64Sink, _>(|(), output| { + *output = 7; + }) + } +} + +impl RowFn for RetryConstantAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.retry_constant_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(u8, u8), u8, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "checked add overflowed")); + } + + Ok(()) + }, + ) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 1 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(0u8, args[0].len()).into_array(), + ))); + } + + Ok(None) + } +} + +impl RowFn for OriginalInputReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.original_input_reducer"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if args[0].len() == 3 { + return Ok(Some(RowExecution::Output( + ConstantArray::new(42_i64, 3).into_array(), + ))); + } + + Ok(None) + } +} + +impl RowFn for DeferredOriginalReducer { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.deferred_original_reducer"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } + + fn reduce_encoded( + &self, + _options: &Self::Options, + _args: &[ArrayRef], + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(RowExecution::DeferredError(vortex_err!( + InvalidArgument: "encoded payload failed" + )))) + } +} + +impl RowFn for PreparedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.prepared_add"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let prepares = Arc::clone(&self.prepares); + let prepare = move |(_lhs, rhs): (Option, Option)| { + prepares.fetch_add(1, Ordering::Relaxed); + rhs + }; + + match self.visit { + PreparedVisit::Owned => visitor + .visit_prepared::<(i64, i64), i64, _>(prepare, |constant_rhs, (lhs, rhs)| { + lhs.wrapping_add(constant_rhs.unwrap_or(rhs)) + }), + PreparedVisit::Sink => visitor.visit_prepared_into::<(i64, i64), I64Sink, _, ()>( + prepare, + |constant_rhs, (lhs, rhs), output| { + *output = lhs.wrapping_add(constant_rhs.unwrap_or(rhs)); + }, + ), + PreparedVisit::Deferred => visitor.visit_prepared_deferred::<(i64, i64), i64, _, bool>( + prepare, + |constant_rhs, (lhs, rhs)| lhs.overflowing_add(constant_rhs.unwrap_or(rhs)), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "prepared add overflowed")); + } + + Ok(()) + }, + ), + } + } +} + +#[test] +fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_batch"); + + let input = PrimitiveArray::new(vec![1i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let result = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::Dense, + }) + }); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_dense_retry_does_not_reduce_filtered_inputs() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let result = ScalarFnVTable::execute(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![1, u8::MAX], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1_u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_u8, 0], Validity::from_iter([true, false])); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_defers_errors_behind_nulls() -> VortexResult<()> { + let input = + PrimitiveArray::new(vec![10_i64, 20], Validity::from_iter([true, false])).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&DeferredOriginalReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_reduce_encoded_precedes_constant_broadcast() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 3).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + let expected = ConstantArray::new(42_i64, 3).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&OriginalInputReducer, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid([true, true])] +#[case::all_invalid([false, false])] +fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.resolve_validity"); + + let validity = Validity::Array(BoolArray::from_iter(validity).into_array()); + let input = PrimitiveArray::new(vec![4_i64, 5], validity).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_filters_and_scatters() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.filter_and_scatter"); + + let input = PrimitiveArray::new( + vec![10_i64, 20, 30, 40], + Validity::from_iter([true, false, true, false]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 4); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.finalize_kernel_output"); + + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let result_dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + let mut ctx = array_session().create_execution_ctx(); + + let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone())?; + let expected = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + assert_eq!(actual.dtype(), &result_dtype); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + assert!(finalize_kernel_output(*ID, &result_dtype, 3, values).is_err()); + + let bools = BoolArray::from_iter([true, false]).into_array(); + assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools).is_err()); + Ok(()) +} + +#[rstest] +#[case::owned_constant(PreparedVisit::Owned, true)] +#[case::owned_varying(PreparedVisit::Owned, false)] +#[case::sink_constant(PreparedVisit::Sink, true)] +#[case::sink_varying(PreparedVisit::Sink, false)] +#[case::deferred_constant(PreparedVisit::Deferred, true)] +#[case::deferred_varying(PreparedVisit::Deferred, false)] +fn test_prepared_visits( + #[case] visit: PreparedVisit, + #[case] constant_rhs: bool, +) -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let rhs = if constant_rhs { + ConstantArray::new(3_i64, 2).into_array() + } else { + PrimitiveArray::from_iter([3_i64, 4]).into_array() + }; + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let prepares = Arc::new(AtomicUsize::new(0)); + let function = PreparedAdd { + visit, + prepares: Arc::clone(&prepares), + }; + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = if constant_rhs { + PrimitiveArray::from_iter([4_i64, 5]).into_array() + } else { + PrimitiveArray::from_iter([4_i64, 6]).into_array() + }; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(prepares.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[rstest] +#[case::dense(RowPolicy::Dense)] +#[case::dense_with_retry(RowPolicy::DenseWithRetry)] +#[case::valid_only(RowPolicy::ValidOnly)] +fn test_strategy_matrix(#[case] policy: RowPolicy) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_strategy"); + + let input = PrimitiveArray::new(vec![1i64, 2, 3], Validity::from_iter([true, false, true])) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 3); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |_args, _ctx| Ok(None), + |args, _ctx| Ok(RowExecution::Output(args.arrays[0].clone())), + |args, _valid, _ctx| Ok(Some(RowExecution::Output(args.arrays[0].clone()))), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_nullary_row_function_broadcasts() -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = ScalarFnVTable::execute(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([7i64, 7, 7]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/execute/mod.rs b/vortex-array/src/scalar_fn/row/execute/mod.rs new file mode 100644 index 00000000000..9c07363dd7c --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/mod.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Row-loop execution for owned outputs and output sinks. +//! +//! [`owned`] stores one independent value per row and can reduce failure evidence. [`sink`] +//! drives output builders whose row handles may refer to shared batch state. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use crate::ArrayRef; +use crate::scalar_fn::ElementTuple; + +mod owned; +pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_infallible; + +mod sink; +pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_valid_rows; + +/// The outcome of a row loop before batch execution decides whether an error is observable. +/// +/// Together with the surrounding [`VortexResult`], this represents three outcomes: +/// +/// - `Err(error)` is a non-retryable execution or immediate row error. +/// - [`Output`](Self::Output) is a successful row loop. +/// - [`DeferredError`](Self::DeferredError) is failure evidence from a completed row loop. +/// +/// A dense loop may evaluate values behind nulls. Its deferred error is therefore not necessarily +/// observable: batch execution can retry over only valid rows, suppressing an error that came from +/// a null row while preserving one from a valid row. A plain `VortexResult` would lose +/// the distinction between that retryable error and an error for which retrying cannot help. +/// +/// Once execution is known to contain only valid rows, converting this outcome into a +/// `VortexResult` turns [`DeferredError`](Self::DeferredError) into an ordinary error. +pub enum RowExecution { + /// The successfully built, full-length output column. + Output(ArrayRef), + + /// An error constructed from failure evidence reduced across a completed row loop. + DeferredError(VortexError), +} + +impl From for VortexResult { + fn from(execution: RowExecution) -> Self { + match execution { + RowExecution::Output(output) => Ok(output), + RowExecution::DeferredError(error) => Err(error), + } + } +} + +/// Ensure that every decoded input addresses the complete row loop. +pub(super) fn ensure_decoded_lengths( + columns: &Args::Columns, + varying: Option<&Args::VaryingColumns<'_>>, + row_count: usize, +) -> VortexResult<()> { + let lengths_match = match varying { + Some(varying) => Args::varying_len_matches(varying, row_count), + None => Args::decoded_lens_match(columns, row_count), + }; + vortex_ensure!( + lengths_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/execute/owned.rs b/vortex-array/src/scalar_fn/row/execute/owned.rs new file mode 100644 index 00000000000..49e33e81364 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -0,0 +1,116 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution that stores one owned output value per row. + +use std::ops::BitOrAssign; + +use vortex_compute::lane_kernels::IndexedSourceExt; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; + +use super::RowExecution; +use crate::ExecutionCtx; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::row::visitor::assert_owned_output_needs_no_drop; + +/// Zero-sized evidence used to erase failure reduction from infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column once, then store one infallible owned output per row. +pub fn execute_owned_infallible( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned::( + args, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode every input column once, then store owned row outputs and reduce deferred failures. +pub fn execute_owned( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + 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_owned_output_needs_no_drop::() }; + + // Keep the vector length at zero until every row succeeds. An unwind then abandons partially + // initialized spare capacity without treating it as initialized output. The no-drop assertion + // above proves that no initialized value requires its destructor to run. + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let failure; + + { + let output = &mut values.spare_capacity_mut()[..row_count]; + + // When every input varies, the indexed source removes argument-shape dispatch from the hot + // loop and lets the lane kernel optimize the traversal as one operation. Keep the varying + // view and its length proof in this branch: hoisting them through the shared validation + // helper produces slower mixed-constant code with LLVM 21.1.2. + // Evidence: + // https://github.com/vortex-data/vortex/blob/ef3fc1/research/rowfn-regressions-2026-08-08/README.md + if let Some(varying) = Args::varying(&columns) { + vortex_ensure!( + Args::varying_len_matches(&varying, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + failure = Args::indexed_source(&varying) + .map_checked_into(output, |elements| apply(&prepared, elements)); + } else { + // A batch-constant input was collapsed to one row during decoding. This path reads that + // row repeatedly while indexing only the inputs that vary. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated = Fail::default(); + for index in 0..row_count { + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + output[index].write(value); + accumulated |= row_failure; + } + failure = accumulated; + } + } + + // SAFETY: normal completion of either loop initializes every slot in `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Failure evidence is reduced inside the loop so its richer error construction stays cold. + // Preserve that provenance so batch execution may retry over only valid rows. + match finish_failure(failure) { + Ok(()) => Ok(RowExecution::Output(Out::build(values))), + Err(error) => Ok(RowExecution::DeferredError(error)), + } +} diff --git a/vortex-array/src/scalar_fn/row/execute/sink.rs b/vortex-array/src/scalar_fn/row/execute/sink.rs new file mode 100644 index 00000000000..42d2e3991e9 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Execution that writes through an output sink. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::RowExecution; +use super::ensure_decoded_lengths; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::SinkResult; + +/// Decode every input column once, allocate the sink once, then write one row at a time. +/// +/// The sink lives here rather than in the closure, so `apply` stays [`Fn`] and mutable output state +/// does not need to be captured by the closure. +pub fn execute_sink( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, +{ + let row_count = args.row_count(); + let mut sink = Sink::with_capacity(row_count, sink_dtype)?; + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let varying = Args::varying(&columns); + ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; + let mut accumulated = ApplyResult::Accumulated::default(); + + { + // Borrow the sink once so its shape and buffer descriptor remain loop invariants. This + // scope releases the borrow before `finish_sink` consumes the sink. + let mut rows = sink.rows(); + vortex_ensure!( + Sink::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + // The all-varying representation removes argument-shape dispatch from the hot loop. The + // mixed path instead reads collapsed batch constants at row zero. + if let Some(varying) = varying { + for index in 0..row_count { + // SAFETY: `ensure_decoded_lengths` proved every varying column has `row_count` + // rows before the loop. + let elements = unsafe { Args::get_varying_unchecked(&varying, index) }; + apply(&prepared, elements, Sink::row(&mut rows, index)) + .accumulate(&mut accumulated)?; + } + } else { + for index in 0..row_count { + apply( + &prepared, + Args::get(&columns, index), + Sink::row(&mut rows, index), + ) + .accumulate(&mut accumulated)?; + } + } + } + + finish_sink(sink) +} + +/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +pub fn execute_sink_valid_rows( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, +) -> VortexResult> +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, +{ + // Decline before input decoding or sink allocation when this sink cannot initialize rows that + // the mask skips. The capability and the operation are the same function pointer. + let Some(initialize_skipped_rows) = Sink::SKIPPED_ROWS_INITIALIZER else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An + // element representation may decline when it cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + let prepared = prepare(Args::constants(&columns)); + let row_count = args.row_count(); + let mut sink = Sink::with_capacity(row_count, sink_dtype)?; + let mut accumulated = ApplyResult::Accumulated::default(); + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid) = valid.bit_buffer() else { + vortex_bail!("execute_sink_valid_rows requires a mixed mask"); + }; + vortex_ensure!( + valid.len() == row_count, + "the validity mask does not address exactly {row_count} rows", + ); + + { + let mut rows = sink.rows(); + vortex_ensure!( + Sink::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + let varying = Args::varying(&columns); + ensure_decoded_lengths::(&columns, varying.as_ref(), row_count)?; + + // The loop writes only valid indices, but the sink still finishes a full-length output. + // Initialize placeholders now; batch execution masks them before the result escapes. + initialize_skipped_rows(&mut rows); + + // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first + // immediate error, turn later callbacks into no-ops, and return before finishing the sink. + let mut error = None; + valid.for_each_set_index(|index| { + if error.is_some() { + return; + } + + let result = match &varying { + Some(varying) => apply( + &prepared, + // SAFETY: `ensure_decoded_lengths` proved every varying column has + // `row_count` rows, and mask indices are below `row_count`. + unsafe { Args::get_varying_unchecked(varying, index) }, + Sink::row(&mut rows, index), + ), + None => apply( + &prepared, + Args::get(&columns, index), + Sink::row(&mut rows, index), + ), + }; + if let Err(err) = result.accumulate(&mut accumulated) { + error = Some(err); + } + }); + + if let Some(error) = error { + return Err(error); + } + } + + finish_sink(sink).map(Some) +} + +fn finish_sink(sink: S) -> VortexResult { + // SAFETY: callers reach this helper only after every completed callback returned the sink's + // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. + // The sink contract defines how that evidence establishes initialization of its row storage. + unsafe { sink.finish() }.map(RowExecution::Output) +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/row/execute/sink/tests.rs b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs new file mode 100644 index 00000000000..8395aa69833 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/sink/tests.rs @@ -0,0 +1,68 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use super::execute_sink_valid_rows; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::VecExecutionArgs; +use crate::validity::Validity; + +struct NonSkippingSink; + +impl OutputSink for NonSkippingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + Err(vortex_err!( + "a non-skipping sink must decline before allocation" + )) + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { + true + } + + fn row<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + unsafe fn finish(self) -> VortexResult { + Err(vortex_err!("a non-skipping sink must not finish")) + } +} + +#[test] +fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([true, false]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, ()>( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + assert!(execution.is_none()); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/row/mod.rs b/vortex-array/src/scalar_fn/row/mod.rs new file mode 100644 index 00000000000..1ad4e98b15a --- /dev/null +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions computed one row at a time. +//! +//! Start with [`RowFn`] to define an operation and [`RowVisitor`] to select one of its output +//! capabilities. [`InputElement`] and [`ElementTuple`] describe how columns become row values. +//! [`OutputElement`] covers independent owned values, while [`OutputSink`] supports outputs that +//! need row handles or shared batch state. [`SinkResult`] describes how a sink-writing closure +//! reports errors. +//! +//! The internal executor owns decoding, batch constants, null propagation, allocation, and +//! validity. A visitor's prepare closure may derive shared state from constant operands once per +//! batch. + +mod execute; +pub use execute::RowExecution; + +mod batch; + +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; diff --git a/vortex-array/src/scalar_fn/row/row_fn.rs b/vortex-array/src/scalar_fn/row/row_fn.rs new file mode 100644 index 00000000000..99b544e44aa --- /dev/null +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -0,0 +1,98 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Scalar functions computed one row at a time. + +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::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::RowExecution; +use crate::scalar_fn::ScalarFnId; + +/// A scalar function computed one row at a time. +/// +/// Declare the argument names and use [`dispatch`](Self::dispatch) to choose concrete element and +/// sink types for each accepted dtype combination. Implement +/// [`ScalarFnVTable`](crate::scalar_fn::ScalarFnVTable) directly for columnar kernels. +pub trait RowFn: 'static + Sized + Clone + Send + Sync { + /// Options for this function, if any. Use [`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 legal dispatch can raise a semantic error as defined by + /// [`ScalarFnVTable::is_fallible`](crate::scalar_fn::ScalarFnVTable::is_fallible). + /// + /// The framework checks this at compile time for every fallible dispatched element or result. + /// A conservative `true` is allowed when only some dtype choices are fallible. + const FALLIBLE: bool = false; + + /// 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. + /// + /// Plan time and run time both call this method, so the choice **must** be a pure function of + /// `options` and `args`. Cross-argument dtype validation belongs here. + fn dispatch( + &self, + options: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult; + + /// Try an encoding-aware implementation before decoding the inputs into row elements. + /// + /// `None` continues to the dispatched row loop. [`Output`](RowExecution::Output) skips that + /// loop and may remain encoded or lazy. [`DeferredError`](RowExecution::DeferredError) reruns + /// only valid rows when null payloads may have caused the failure. For non-nullary functions, + /// batch execution calls this hook at most once with the original, unfiltered arrays; slices + /// and compacted retries do not reach it. + /// + /// Like a dense row closure, this hook must be total over every stored payload, including + /// payloads behind null rows. An `Err` is immediately user-visible and is never suppressed or + /// retried through the row layer. + /// + /// # Requirements + /// + /// - `output.len()` **must** equal `args[0].len()`. + /// - The output dtype **must** match the planned dtype when ignoring nullability. + /// - The output **must not** introduce a null where every input is valid. + /// + /// The framework skips this hook for nullary functions. + fn reduce_encoded( + &self, + options: &Self::Options, + args: &[ArrayRef], + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + _ = (options, args, ctx); + Ok(None) + } +} diff --git a/vortex-array/src/scalar_fn/row/types/element/bool.rs b/vortex-array/src/scalar_fn/row/types/element/bool.rs new file mode 100644 index 00000000000..3b68dfb1e80 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/bool.rs @@ -0,0 +1,77 @@ +// 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::InputElement; +use crate::scalar_fn::OutputElement; +use crate::validity::Validity; + +// SAFETY: the varying view is a bit buffer, and its reported length is the buffer length. +unsafe impl InputElement for bool { + type Column = BitBuffer; + type Varying<'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 get(column: &Self::Column, index: usize) -> bool { + column.value(index) + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> bool + where + Self: 'a, + { + column.value(index) + } + + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> bool + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { column.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/row/types/element/mod.rs b/vortex-array/src/scalar_fn/row/types/element/mod.rs new file mode 100644 index 00000000000..232040cde12 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -0,0 +1,147 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The element types a row function can read and produce. +//! +//! [`InputElement::Elem`] may borrow from its decoded column. [`OutputElement`] is returned by an +//! owned row computation; runtime-shaped output uses an +//! [`OutputSink`](crate::scalar_fn::OutputSink). + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; + +mod bool; + +mod primitive; + +mod tuple; +pub use tuple::ElementTuple; +pub use tuple::IndexedElementTuple; +pub use tuple::batch_constant; + +/// An element type that can be read row-wise out of an input column. +/// +/// # Safety +/// +/// For every view returned by [`varying`](Self::varying), every index below +/// [`varying_len`](Self::varying_len) **must** satisfy the safety contract of +/// [`get_varying_unchecked`](Self::get_varying_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 view of a varying decoded column read by the hot row loop. + /// + /// This may 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 Varying<'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 only guarantee payloads for valid rows. This is `false` when a null row's stored + /// offset or pointer may not address anything, and `true` only when [`decode`](Self::decode), + /// [`get`](Self::get), [`varying`](Self::varying), [`varying_len`](Self::varying_len), and + /// [`get_varying`](Self::get_varying) remain safe and correct for null rows. + /// + /// Dense execution requires this of every argument; otherwise the row layer executes only + /// valid rows. + /// + /// Dense execution can pass unspecified values from null rows. The closure must be total over + /// every stored value: it cannot panic or cause side effects beyond its declared output. + const DENSE_SAFE: bool = false; + + /// Whether [`decode`](Self::decode) can fail on _legal_ input data. + /// + /// This excludes infrastructural failures such as IO or allocation. Set it when legal input may + /// contain a value that the decoder rejects. + const DECODE_FALLIBLE: bool = true; + + /// 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 batch. + /// + /// Hoist dtype checks, downcasts, and other batch-invariant work into this method. + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult; + + /// Decode `array` _without_ assuming every row is valid, or `Ok(None)` when this element + /// cannot for this particular array. + /// + /// An element with [`DENSE_SAFE`](Self::DENSE_SAFE) set **should not** override this: its + /// ordinary decode already tolerates null payloads, so the default is already correct and an + /// override just restates it. Overriding is for an element that is _not_ dense-safe but can + /// still write an arbitrary placeholder into null slots; the caller guarantees + /// [`get`](Self::get) is never called for such a row. The skip-invalid strategy uses this + /// representation to avoid filtering the input. + /// + /// Return `Ok(None)` rather than an error when an array has no null-tolerant decode; the + /// batch execution falls back to the filter strategy. + fn decode_null_tolerant( + array: ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + if Self::DENSE_SAFE { + Self::decode(array, ctx).map(Some) + } else { + Ok(None) + } + } + + /// Read the element at `index`, the one function called once per row. + /// + /// This must not repeat work that is constant across the batch; do that work in + /// [`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 varying(column: &Self::Column) -> Self::Varying<'_>; + + /// Number of rows addressable through a [`Varying`](Self::Varying) view. + /// + /// Every index below this length must be valid for + /// [`get_varying_unchecked`](Self::get_varying_unchecked). + fn varying_len(column: &Self::Varying<'_>) -> usize; + + /// Read one row from a [`Varying`](Self::Varying) view. + fn get_varying<'a>(column: &Self::Varying<'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 [`varying_len`](Self::varying_len) for `column`. + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> Self::Elem<'a> + where + Self: 'a, + { + Self::get_varying(column, index) + } +} + +/// 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. + /// + /// Taking no arguments confines an element's dtype to a property of its Rust type, so an output + /// whose dtype depends on runtime data (a tensor, whose dtype carries its shape) cannot be an + /// element. Such an output uses an [`OutputSink`](crate::scalar_fn::OutputSink), whose + /// [`sink_dtype`](crate::scalar_fn::OutputSink::sink_dtype) does see the input dtypes. + 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/row/types/element/primitive.rs b/vortex-array/src/scalar_fn/row/types/element/primitive.rs new file mode 100644 index 00000000000..9e80d937ef5 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/primitive.rs @@ -0,0 +1,83 @@ +// 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::InputElement; +use crate::scalar_fn::OutputElement; +use crate::validity::Validity; + +// SAFETY: the varying view is a native slice, and its reported length is the slice length. +unsafe impl InputElement for T { + type Column = Buffer; + type Varying<'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 get(column: &Self::Column, index: usize) -> T { + column[index] + } + + fn varying(column: &Self::Column) -> Self::Varying<'_> { + column.as_slice() + } + + fn varying_len(column: &Self::Varying<'_>) -> usize { + column.len() + } + + fn get_varying<'a>(column: &Self::Varying<'a>, index: usize) -> T + where + Self: 'a, + { + column[index] + } + + unsafe fn get_varying_unchecked<'a>(column: &Self::Varying<'a>, index: usize) -> T + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { *column.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/row/types/element/tuple.rs b/vortex-array/src/scalar_fn/row/types/element/tuple.rs new file mode 100644 index 00000000000..680ae0b4e9c --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/tuple.rs @@ -0,0 +1,475 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Argument lists built from [`InputElement`]s, and the per-argument decode behind them. + +use vortex_compute::lane_kernels::IndexedSource; +use vortex_compute::lane_kernels::LaneZip; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Extension; +use crate::arrays::Masked; +use crate::arrays::extension::ExtensionArrayExt; +use crate::arrays::masked::MaskedArraySlotsExt; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::InputElement; + +mod private { + pub trait Sealed {} +} + +/// 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 varies within the batch. + ArgColumnKind, +); + +enum ArgColumnKind { + Varying(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::Varying(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::Varying) + .map(Self)) + } + + fn get(&self, index: usize) -> T::Elem<'_> { + match &self.0 { + ArgColumnKind::Varying(column) => T::get(column, index), + ArgColumnKind::Constant(column) => T::get(column, 0), + } + } + + fn varying(&self) -> Option<&T::Column> { + match &self.0 { + ArgColumnKind::Varying(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::Varying(column) => T::varying_len(&T::varying(column)) == row_count, + ArgColumnKind::Constant(_) => true, + } + } + + fn constant(&self) -> Option> { + match &self.0 { + ArgColumnKind::Varying(_) => 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 may 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.as_constant().is_some() { + return Some(array.clone()); + } + + if let Some(masked) = array.as_opt::() { + return Some(masked.child().clone()).filter(|child| child.as_constant().is_some()); + } + + array + .as_opt::() + .is_some_and(|ext| ext.storage_array().as_constant().is_some()) + .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; + + /// Direct references to decoded columns when every argument varies within the batch. + type VaryingColumns<'a>; + + /// The borrowed row of element values. + type Elems<'a>; + + /// The batch-constant element values: [`Elems`](Self::Elems) with every argument wrapped in + /// `Option`. + /// + /// `Some` marks an argument whose operand is constant for the batch and carries the element + /// every row reads; `None` marks one that varies by row. This is what + /// [`RowVisitor`](crate::scalar_fn::RowVisitor) hands to a visit's prepare closure, so a kernel + /// can hoist work that depends only on a constant argument out of the row loop. + 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, including that `dtypes` has exactly `ARITY` entries. + /// + /// The expression layer checks the count against [`Arity`](crate::scalar_fn::Arity) before it + /// builds a call, but this is also the entry point of the public + /// [`return_dtype`](crate::scalar_fn::ScalarFnVTable::return_dtype), so the count is enforced + /// here rather than assumed. + fn validate(dtypes: &[DType]) -> VortexResult<()>; + + /// Decode every input column once. Called once per batch. + fn decode(args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> 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 varying(columns: &Self::Columns) -> Option>; + + /// Whether every varying column contains exactly `row_count` rows. + fn varying_len_matches(columns: &Self::VaryingColumns<'_>, row_count: usize) -> bool; + + /// Whether every argument that varies within the batch contains exactly `row_count` rows. + /// + /// The same guarantee as [`varying_len_matches`](Self::varying_len_matches), for the mixed case + /// [`varying`](Self::varying) declines: a batch-constant argument is exempt because it was + /// collapsed to one row, while every argument beside it still has to address the whole batch. + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; + + /// Read one row from columns already known to vary within the batch. + fn get_varying<'a>(columns: &Self::VaryingColumns<'a>, index: usize) -> Self::Elems<'a>; + + /// Read one row from varying columns without checking bounds. + /// + /// # Safety + /// + /// `index` must be in bounds for every column. + unsafe fn get_varying_unchecked<'a>( + columns: &Self::VaryingColumns<'a>, + index: usize, + ) -> Self::Elems<'a>; + + /// Read the batch-constant elements out of the decoded columns. Called once per batch. + fn constants(columns: &Self::Columns) -> Self::ConstElems<'_>; +} + +/// An argument tuple that supports a validated dense indexed traversal. +/// +/// This is separate from [`ElementTuple`] because many row elements have no contiguous source, and +/// stable Rust cannot provide a blanket fallback plus a more specific primitive implementation. +/// The trait is sealed so shared execution can rely on its unchecked-read contract. A tuple only +/// implements it when the source can be validated once and every lane can then be read +/// independently. +pub trait IndexedElementTuple: ElementTuple { + /// The source shared execution uses for a dense all-varying loop. + /// + /// Its length must be the common varying-column length. For every valid index it must preserve + /// row order, return the same value as [`ElementTuple::get_varying`], and uphold the unchecked + /// read contract of [`IndexedSource`]. + type Source<'a>: IndexedSource>; + + /// Borrow a source from columns already validated to vary within the batch. + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a>; +} + +/// An indexed native slice yielding the one-tuples expected by a unary row closure. +#[derive(Clone, Copy)] +pub struct UnaryTupleSource<'a, T>( + /// The native values read by the row loop. + &'a [T], +); + +impl IndexedSource for UnaryTupleSource<'_, T> { + type Item = (T,); + + fn len(&self) -> usize { + self.0.len() + } + + unsafe fn get_unchecked(&self, index: usize) -> Self::Item { + // SAFETY: the caller guarantees that `index` is in bounds. + (unsafe { *self.0.get_unchecked(index) },) + } +} + +impl private::Sealed for () {} + +impl ElementTuple for () { + type Columns = (); + type VaryingColumns<'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 decode_null_tolerant( + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult> { + Ok(Some(())) + } + + fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} + + fn varying(_columns: &Self::Columns) -> Option> { + Some(()) + } + + fn varying_len_matches(_columns: &Self::VaryingColumns<'_>, _row_count: usize) -> bool { + true + } + + fn decoded_lens_match(_columns: &Self::Columns, _row_count: usize) -> bool { + true + } + + fn get_varying<'a>(_columns: &Self::VaryingColumns<'a>, _index: usize) -> Self::Elems<'a> {} + + unsafe fn get_varying_unchecked<'a>( + _columns: &Self::VaryingColumns<'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 VaryingColumns<'a> = ($($t::Varying<'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 decode_null_tolerant( + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + 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 varying(columns: &Self::Columns) -> Option> { + Some(($($t::varying(columns.$idx.varying()?),)+)) + } + + fn varying_len_matches( + columns: &Self::VaryingColumns<'_>, + row_count: usize, + ) -> bool { + $($t::varying_len(&columns.$idx) == row_count &&)+ true + } + + fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool { + $(columns.$idx.addresses_rows(row_count) &&)+ true + } + + fn get_varying<'a>( + columns: &Self::VaryingColumns<'a>, + index: usize, + ) -> Self::Elems<'a> { + ($($t::get_varying(&columns.$idx, index),)+) + } + + unsafe fn get_varying_unchecked<'a>( + columns: &Self::VaryingColumns<'a>, + index: usize, + ) -> Self::Elems<'a> { + // SAFETY: forwarded from this method's contract. + ($(unsafe { $t::get_varying_unchecked(&columns.$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); + +impl IndexedElementTuple for (T,) { + type Source<'a> = UnaryTupleSource<'a, T>; + + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a> { + UnaryTupleSource(columns.0) + } +} + +impl IndexedElementTuple for (Left, Right) { + type Source<'a> = LaneZip<&'a [Left], &'a [Right]>; + + fn indexed_source<'a>(columns: &Self::VaryingColumns<'a>) -> Self::Source<'a> { + LaneZip::new(columns.0, columns.1) + } +} + +#[cfg(test)] +mod tests { + use vortex_compute::lane_kernels::IndexedSource; + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_mask::Mask; + + use super::UnaryTupleSource; + use super::batch_constant; + use crate::IntoArray; + use crate::arrays::ConstantArray; + use crate::arrays::ExtensionArray; + use crate::arrays::MaskedArray; + use crate::dtype::Nullability; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::validity::Validity; + + #[test] + fn test_unary_tuple_source_reads_one_tuple_per_row() { + let source = UnaryTupleSource(&[10, 20, 30]); + assert_eq!(source.len(), 3); + + // SAFETY: index one is within the three-element source. + assert_eq!(unsafe { source.get_unchecked(1) }, (20,)); + } + + #[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.as_constant().is_some()); + 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/row/types/mod.rs b/vortex-array/src/scalar_fn/row/types/mod.rs new file mode 100644 index 00000000000..e47f195410c --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/mod.rs @@ -0,0 +1,23 @@ +// 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; +pub(super) use element::batch_constant; + +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/row/types/result.rs b/vortex-array/src/scalar_fn/row/types/result.rs new file mode 100644 index 00000000000..cda44802c40 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/result.rs @@ -0,0 +1,81 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What a sink-writing row closure may return. + +use vortex_error::VortexResult; + +use super::InitializedElement; + +mod private { + pub trait Sealed {} +} + +/// 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; + + /// Loop-local state used while accumulating row results. + type Accumulated: 'static + Copy + Default; + + /// Whether this return type can carry an error. + const FALLIBLE: bool; + + /// Merge this row's outcome into the batch-wide reduction. + fn accumulate(self, accumulated: &mut Self::Accumulated) -> VortexResult<()>; +} + +impl private::Sealed for () {} + +impl SinkResult for () { + type WriteToken = (); + type Accumulated = (); + + const FALLIBLE: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + Ok(()) + } +} + +impl private::Sealed for InitializedElement {} + +impl SinkResult for InitializedElement { + type WriteToken = InitializedElement; + type Accumulated = (); + + const FALLIBLE: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + Ok(()) + } +} + +impl private::Sealed for VortexResult<()> {} + +impl SinkResult for VortexResult<()> { + type WriteToken = (); + type Accumulated = (); + + const FALLIBLE: bool = true; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + self + } +} + +impl private::Sealed for VortexResult {} + +impl SinkResult for VortexResult { + type WriteToken = InitializedElement; + type Accumulated = (); + + const FALLIBLE: bool = true; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + self.map(|_| ()) + } +} diff --git a/vortex-array/src/scalar_fn/row/types/sink.rs b/vortex-array/src/scalar_fn/row/types/sink.rs new file mode 100644 index 00000000000..adbfdbfa878 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -0,0 +1,165 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The column builders a row function can write its output into. + +use std::mem::MaybeUninit; + +use vortex_error::VortexResult; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::OutputElement; + +/// A column allocated once per batch that a row closure writes into, one row at a time. +/// +/// A sink may use the input dtypes to build a runtime-shaped output or own shared batch state. The +/// executor passes each row slot into an [`Fn`] closure, keeping mutable state out of its capture. +/// +/// 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`](Self::SKIPPED_ROWS_INITIALIZER) is present. +pub 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 operation that initializes every output position before skip-invalid execution. + /// + /// `None` declines skip-invalid execution before input decoding or sink allocation. A present + /// initializer **must** leave a legal arbitrary value in every row. Encoding support as the + /// initializer's presence prevents a separate capability flag from disagreeing with a no-op + /// method. + const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = None; + + /// 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 distinct + /// token returned after initialization. A sink that uses this token to justify unsafe code + /// **must** prevent safe construction that does not establish the invariant. Make construction + /// unsafe when Rust cannot tie the token to the supplied row handle. + type WriteToken: 'static; + + /// The dtype of the column this sink builds, given the function's input dtypes. + /// + /// **Must** be non-nullable: batch execution derives nullability from the inputs, widens the + /// result, and masks the null rows. + fn sink_dtype(args: &[DType]) -> VortexResult; + + /// Allocate a sink for `rows` rows of `dtype`, which is this sink's own + /// [`sink_dtype`](Self::sink_dtype). Called once per batch. + fn with_capacity(rows: usize, dtype: &DType) -> VortexResult; + + /// Borrow all output rows for the hot loop. + fn rows(&mut self) -> Self::Rows<'_>; + + /// Whether every index in `0..row_count` is addressable through [`row`](Self::row). + /// + /// Called once before the hot loop. Besides validating the sink contract, this gives the + /// optimizer the output bounds it needs to remove the bounds check hidden in each row accessor. + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool; + + /// Hand out the place to write row `index`. Must be `O(1)`: it is called in the row loop. + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>; + + /// Finish into the built column, whose dtype **must** be this sink's + /// [`sink_dtype`](Self::sink_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, + /// [`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. +#[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. Immediate failures are +/// safe because [`OutputSink::finish`] is not called after one. +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, +} + +impl OutputSink for UninitElementSink { + type Rows<'a> = &'a mut [MaybeUninit]; + + const SKIPPED_ROWS_INITIALIZER: Option fn(&mut Self::Rows<'a>)> = Some(|rows| { + for row in rows.iter_mut() { + row.write(T::default()); + } + }); + + type Row<'a> = &'a mut MaybeUninit; + type WriteToken = InitializedElement; + + fn sink_dtype(_args: &[DType]) -> VortexResult { + Ok(T::element_dtype()) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> 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_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[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/row/visitor/check.rs b/vortex-array/src/scalar_fn/row/visitor/check.rs new file mode 100644 index 00000000000..189aa49b75b --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/check.rs @@ -0,0 +1,117 @@ +// 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::ElementTuple; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::SinkResult; + +/// Assert the no-drop contract that makes partially initialized output safe to abandon on unwind. +pub(in crate::scalar_fn::row) const fn assert_owned_output_needs_no_drop() { + assert!( + !needs_drop::(), + "owned row outputs must not require drop glue" + ); +} + +/// Assert that the input arity and decode fallibility match the function-wide declarations. +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 pushdown treats an infallible function as safe to evaluate over values no code + // references, so every dispatch must fit the function-wide declaration. + assert!( + !Args::DECODE_FALLIBLE || F::FALLIBLE, + "RowFn::FALLIBLE must be true when input decoding can fail", + ); +} + +/// Assert the input contract and that owned output values do not require drop glue. +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::(); +} + +/// Assert that a sink visit obeys the input, fallibility, and deferred-error contracts. +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", + ); +} + +/// Assert the owned-output contract, fallibility declaration, and failure-evidence width bound. +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", + ); +} + +/// Validate the input dtypes and return the non-nullable dtype built by `Out`. +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) +} + +/// Validate the input dtypes and return the non-nullable dtype built by `Sink`. +pub(super) fn validate_sink_visit( + dtypes: &[DType], +) -> VortexResult { + Args::validate(dtypes)?; + + let dtype = Sink::sink_dtype(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/row/visitor/execute.rs b/vortex-array/src/scalar_fn/row/visitor/execute.rs new file mode 100644 index 00000000000..73314df9172 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/execute.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visitors that execute dense and skip-invalid row loops. +//! +//! Each method verifies that execution selected the same visit shape as planning before handing +//! its typed closures to the matching loop. Valid-row execution can decline without running a loop; +//! batch execution then filters the inputs and retries the dense loop. + +use std::marker::PhantomData; +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +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::private; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::SinkResult; +use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::execute::execute_owned; +use crate::scalar_fn::row::execute::execute_owned_infallible; +use crate::scalar_fn::row::execute::execute_sink; +use crate::scalar_fn::row::execute::execute_sink_valid_rows; + +/// The run-time visit that decodes every column once and runs the selected row loop. +pub struct ExecuteRows<'args, 'ctx, F> { + /// The inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'args, 'ctx, F> ExecuteRows<'args, 'ctx, F> { + pub fn new( + args: &'args dyn ExecutionArgs, + output_dtype: &'args DType, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + ctx, + function: PhantomData, + } + } +} + +impl private::Sealed for ExecuteRows<'_, '_, F> {} + +impl RowVisitor for ExecuteRows<'_, '_, F> { + type VisitResult = RowExecution; + + 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::() }; + + execute_owned_infallible::(self.args, self.ctx, prepare, apply) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + + execute_sink::( + self.args, + self.output_dtype, + self.ctx, + prepare, + apply, + ) + } + + 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::() }; + + execute_owned::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} + +/// The run-time visit that tries skip-invalid execution over the original input columns. +/// +/// Only output sinks have a contract for skipped output positions. Owned visits therefore decline +/// so batch execution can use its filter-and-scatter fallback. +pub struct ExecuteValidRows<'args, 'ctx, F> { + /// The original inputs for this kernel invocation. + args: &'args dyn ExecutionArgs, + + /// The output dtype computed by the planning visit. + output_dtype: &'args DType, + + /// The conjoined validity, materialized by batch execution and guaranteed mixed. + valid: &'args Mask, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'args, 'ctx, F> ExecuteValidRows<'args, 'ctx, F> { + pub fn new( + args: &'args dyn ExecutionArgs, + output_dtype: &'args DType, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + valid, + ctx, + function: PhantomData, + } + } +} + +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} + +impl RowVisitor for ExecuteValidRows<'_, '_, F> { + type VisitResult = Option; + + 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::() }; + + // Owned execution has no sink that can initialize skipped output positions. Decline so + // batch execution filters the inputs and retries with the dense visitor. + Ok(None) + } + + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + + execute_sink_valid_rows::( + self.args, + self.output_dtype, + self.valid, + self.ctx, + prepare, + apply, + ) + } + + 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::() }; + + // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. + Ok(None) + } +} diff --git a/vortex-array/src/scalar_fn/row/visitor/mod.rs b/vortex-array/src/scalar_fn/row/visitor/mod.rs new file mode 100644 index 00000000000..74122066709 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/mod.rs @@ -0,0 +1,165 @@ +// 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::RowFn::dispatch + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::SinkResult; + +mod check; +pub(super) use check::assert_owned_output_needs_no_drop; + +mod execute; +pub(super) use execute::ExecuteRows; +pub(super) use execute::ExecuteValidRows; + +mod plan; +pub(super) use plan::PlanRows; + +/// 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. +pub trait RowVisitor: private::Sealed + Sized { + /// The framework result of visiting one concrete row signature. + /// + /// This is a batch plan or execution result, not the per-row `Out` returned by + /// [`RowVisitor::visit`] and [`RowVisitor::visit_deferred`]. + type VisitResult; + + /// Visit an infallible row computation that returns one independent output value. + /// + /// `apply` must be total over every stored element value: it must not panic or have side + /// effects. Dense execution can pass unspecified values from null rows. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of + /// [`RowFn::ARG_NAMES`](crate::scalar_fn::RowFn::ARG_NAMES). + /// - [`RowFn::FALLIBLE`](crate::scalar_fn::RowFn::FALLIBLE) **must** be `true` when decoding + /// `Args` can fail. + /// - `Out` **must not** require drop glue. + 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. + 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 sink-provided row handle. + /// + /// `apply` must be total over every stored input value: it must not panic or cause side effects + /// other than writing the supplied row handle. Dense execution can pass unspecified values + /// from null rows. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of + /// [`RowFn::ARG_NAMES`](crate::scalar_fn::RowFn::ARG_NAMES). + /// - [`RowFn::FALLIBLE`](crate::scalar_fn::RowFn::FALLIBLE) **must** be `true` when decoding + /// `Args` or computing the result can fail. + fn visit_into( + self, + apply: impl Fn(Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + self.visit_prepared_into::( + |_| (), + move |&(), args, row| apply(args, row), + ) + } + + /// The prepared form of [`visit_into`](Self::visit_into), with the same prerequisites. + fn visit_prepared_into( + self, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult; + + /// Visit a row computation that returns an owned output and deferred failure evidence. + /// + /// `apply` must be total over every stored element value: it must not panic or have side + /// effects. Dense execution can pass unspecified values from null rows. + /// + /// `Fail` is OR-reduced across rows and handed to `finish_failure`. The value from + /// [`Default::default`] **must** mean success, including for an empty batch. The compiler + /// cannot check this semantic requirement. + /// + /// # Prerequisites + /// + /// The framework checks these at compile time: + /// + /// - `Args::ARITY` **must** equal the length of + /// [`RowFn::ARG_NAMES`](crate::scalar_fn::RowFn::ARG_NAMES). + /// - [`RowFn::FALLIBLE`](crate::scalar_fn::RowFn::FALLIBLE) **must** be `true`. + /// - `Out` **must not** require drop glue. + /// - `Out` **must** be at least as wide as `Fail` so failure tracking does not reduce the + /// vector width. + 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. + 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; +} + +mod private { + pub trait Sealed {} +} diff --git a/vortex-array/src/scalar_fn/row/visitor/plan.rs b/vortex-array/src/scalar_fn/row/visitor/plan.rs new file mode 100644 index 00000000000..894acf33082 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/plan.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The visitor that validates a concrete dispatch and plans its nullable execution. + +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::private; +use crate::dtype::DType; +use crate::scalar_fn::ElementTuple; +use crate::scalar_fn::IndexedElementTuple; +use crate::scalar_fn::OutputElement; +use crate::scalar_fn::OutputSink; +use crate::scalar_fn::RowFn; +use crate::scalar_fn::SinkResult; +use crate::scalar_fn::row::batch::BatchPlan; +use crate::scalar_fn::row::batch::RowPolicy; + +/// The plan-time visit that validates dtypes and derives the nullable execution policy. +pub struct PlanRows<'a, F> { + /// The input dtypes for this plan. + dtypes: &'a [DType], + + /// The visited function, carried only so the dispatch check can name its contract. + function: PhantomData, +} + +impl<'a, F> PlanRows<'a, F> { + pub fn new(dtypes: &'a [DType]) -> Self { + Self { + dtypes, + function: PhantomData, + } + } +} + +impl private::Sealed for PlanRows<'_, F> {} + +impl RowVisitor for PlanRows<'_, 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<'_>, Sink::Row<'_>) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult, + { + const { assert_sink_visit_contract::() }; + + Ok(BatchPlan { + output_dtype: validate_sink_visit::(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::(), + }) + } +} diff --git a/vortex-array/src/scalar_fn/row/vtable.rs b/vortex-array/src/scalar_fn/row/vtable.rs new file mode 100644 index 00000000000..8ed9163866f --- /dev/null +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`ScalarFnVTable`] adapter shared by every [`RowFn`]. +//! +//! The [`visitor`](super::visitor) module validates and executes the concrete row signature +//! selected by dispatch. This module connects those visits to batch execution and exposes the +//! resulting scalar function behavior to the rest of the compute stack. + +use vortex_error::VortexResult; +use vortex_mask::Mask; +use vortex_session::VortexSession; + +use super::row_fn::RowFn; +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::BorrowedExecutionArgs; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::row::batch::Batch; +use crate::scalar_fn::row::batch::KernelArgs; +use crate::scalar_fn::row::batch::finalize_kernel_output; +use crate::scalar_fn::row::execute::RowExecution; +use crate::scalar_fn::row::visitor::ExecuteRows; +use crate::scalar_fn::row::visitor::ExecuteValidRows; +use crate::scalar_fn::row::visitor::PlanRows; + +/// Implement [`ScalarFnVTable`] for every [`RowFn`]. +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 { + let plan = self.dispatch(options, args, PlanRows::::new(args))?; + + Ok(plan.result_dtype(args)) + } + + fn execute( + &self, + options: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Nullary functions have no input validity to propagate, so they skip batch execution. + if args.num_inputs() == 0 { + let result_dtype = ScalarFnVTable::return_dtype(self, options, &[])?; + let nullary_args = KernelArgs { + arrays: &[], + row_count: args.row_count(), + dtypes: &[], + output_dtype: &result_dtype, + }; + + let execution = execute_rows(self, options, nullary_args, ctx)?; + let values = VortexResult::from(execution)?; + + return finalize_kernel_output( + RowFn::id(self), + &result_dtype, + args.row_count(), + values, + ); + } + + let batch = prepare_batch(self, options, args)?; + batch.execute( + |args, ctx| self.reduce_encoded(options, args.arrays, ctx), + |args, ctx| execute_rows(self, options, args, ctx), + |args, valid, ctx| try_execute_rows_unfiltered(self, options, args, valid, ctx), + 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 + } +} + +/// Run the encoding-aware rewrite when available, or execute the selected row loop. +fn execute_rows( + function: &F, + options: &F::Options, + args: KernelArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let execution = BorrowedExecutionArgs::new(args.arrays, args.row_count); + + function.dispatch( + options, + args.dtypes, + ExecuteRows::::new(&execution, args.output_dtype, ctx), + ) +} + +/// Try execution against the original inputs, returning `None` when batch execution must filter. +fn try_execute_rows_unfiltered( + function: &F, + options: &F::Options, + args: KernelArgs<'_>, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let execution = BorrowedExecutionArgs::new(args.arrays, args.row_count); + + function.dispatch( + options, + args.dtypes, + ExecuteValidRows::::new(&execution, args.output_dtype, valid, ctx), + ) +} + +/// Prepare the batch inputs and execution plan for `function`. +fn prepare_batch( + function: &F, + options: &F::Options, + args: &dyn ExecutionArgs, +) -> VortexResult { + Batch::new(RowFn::id(function), args, |arg_dtypes| { + function.dispatch(options, arg_dtypes, PlanRows::::new(arg_dtypes)) + }) +} diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 5d3561ff039..30f38439dd5 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -196,8 +196,7 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// Returns whether this scalar function is strict. /// /// A strict function returns null for a row when any argument is null for that row. This - /// matches [PostgreSQL's `STRICT` convention](https://www.postgresql.org/docs/current/sql-createfunction.html) - /// for null propagation. + /// matches [PostgreSQL's `STRICT` convention][postgres-strict] for null propagation. /// /// Return `true` only when this holds for every argument. `add` is strict, but Kleene `AND` /// is not because `false AND null` returns `false`. `is_null` is also not strict. @@ -212,6 +211,8 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// /// This property applies only to the scalar function, not its child expressions. Nullary /// functions are vacuously strict. The default is conservatively `false`. + /// + /// [postgres-strict]: https://www.postgresql.org/docs/current/sql-createfunction.html fn is_strict(&self, options: &Self::Options) -> bool { _ = options; false @@ -328,20 +329,22 @@ pub trait ExecutionArgs { fn row_count(&self) -> usize; } -/// A concrete [`ExecutionArgs`] backed by a `Vec`. -pub struct VecExecutionArgs { - inputs: Vec, +/// An [`ExecutionArgs`] view over borrowed arrays with an explicit row count. +pub(crate) struct BorrowedExecutionArgs<'a> { + /// The arrays exposed through this execution view. + inputs: &'a [ArrayRef], + + /// The row count reported for this execution view. row_count: usize, } -impl VecExecutionArgs { - /// Create a new `VecExecutionArgs`. - pub fn new(inputs: Vec, row_count: usize) -> Self { +impl<'a> BorrowedExecutionArgs<'a> { + pub(crate) fn new(inputs: &'a [ArrayRef], row_count: usize) -> Self { Self { inputs, row_count } } } -impl ExecutionArgs for VecExecutionArgs { +impl ExecutionArgs for BorrowedExecutionArgs<'_> { fn get(&self, index: usize) -> VortexResult { self.inputs.get(index).cloned().ok_or_else(|| { vortex_err!( @@ -361,6 +364,36 @@ impl ExecutionArgs for VecExecutionArgs { } } +/// A concrete [`ExecutionArgs`] backed by a `Vec`. +pub struct VecExecutionArgs { + /// The owned arrays exposed through this execution view. + inputs: Vec, + + /// The row count reported for this execution view. + row_count: usize, +} + +impl VecExecutionArgs { + /// Create a new `VecExecutionArgs`. + pub fn new(inputs: Vec, row_count: usize) -> Self { + Self { inputs, row_count } + } +} + +impl ExecutionArgs for VecExecutionArgs { + fn get(&self, index: usize) -> VortexResult { + BorrowedExecutionArgs::new(&self.inputs, self.row_count).get(index) + } + + fn num_inputs(&self) -> usize { + BorrowedExecutionArgs::new(&self.inputs, self.row_count).num_inputs() + } + + fn row_count(&self) -> usize { + BorrowedExecutionArgs::new(&self.inputs, self.row_count).row_count() + } +} + #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct EmptyOptions; impl Display for EmptyOptions {