From 0a0ad0db146c1b761b3727f40f1c78ef02f62baa Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 10:42:32 -0400 Subject: [PATCH 1/3] Add the RowFn scalar function framework Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/mod.rs | 8 + vortex-array/src/scalar_fn/row/batch/args.rs | 66 +++ .../src/scalar_fn/row/batch/execution.rs | 440 ++++++++++++++++++ vortex-array/src/scalar_fn/row/batch/mod.rs | 22 + .../src/scalar_fn/row/batch/policy.rs | 102 ++++ vortex-array/src/scalar_fn/row/execute/mod.rs | 52 +++ .../src/scalar_fn/row/execute/owned.rs | 112 +++++ .../src/scalar_fn/row/execute/sink.rs | 192 ++++++++ vortex-array/src/scalar_fn/row/mod.rs | 36 ++ vortex-array/src/scalar_fn/row/row_fn.rs | 90 ++++ .../src/scalar_fn/row/types/element/bool.rs | 68 +++ .../src/scalar_fn/row/types/element/mod.rs | 129 +++++ .../scalar_fn/row/types/element/primitive.rs | 74 +++ .../src/scalar_fn/row/types/element/tuple.rs | 414 ++++++++++++++++ vortex-array/src/scalar_fn/row/types/mod.rs | 23 + .../src/scalar_fn/row/types/result.rs | 124 +++++ vortex-array/src/scalar_fn/row/types/sink.rs | 139 ++++++ .../src/scalar_fn/row/visitor/check.rs | 126 +++++ .../src/scalar_fn/row/visitor/execute.rs | 226 +++++++++ vortex-array/src/scalar_fn/row/visitor/mod.rs | 157 +++++++ .../src/scalar_fn/row/visitor/plan.rs | 104 +++++ vortex-array/src/scalar_fn/row/vtable.rs | 171 +++++++ 22 files changed, 2875 insertions(+) create mode 100644 vortex-array/src/scalar_fn/row/batch/args.rs create mode 100644 vortex-array/src/scalar_fn/row/batch/execution.rs create mode 100644 vortex-array/src/scalar_fn/row/batch/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/batch/policy.rs create mode 100644 vortex-array/src/scalar_fn/row/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/execute/owned.rs create mode 100644 vortex-array/src/scalar_fn/row/execute/sink.rs create mode 100644 vortex-array/src/scalar_fn/row/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/row_fn.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/bool.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/primitive.rs create mode 100644 vortex-array/src/scalar_fn/row/types/element/tuple.rs create mode 100644 vortex-array/src/scalar_fn/row/types/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/types/result.rs create mode 100644 vortex-array/src/scalar_fn/row/types/sink.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/check.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/execute.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/mod.rs create mode 100644 vortex-array/src/scalar_fn/row/visitor/plan.rs create mode 100644 vortex-array/src/scalar_fn/row/vtable.rs 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..262b8252928 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/args.rs @@ -0,0 +1,66 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Input views and planning metadata passed to a row kernel. + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; + +/// 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 accidentally +/// pairing an input view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub struct KernelArgs<'a> { + /// The executor-facing view, including the row count for this invocation. + pub execution: &'a dyn ExecutionArgs, + + /// The same inputs as concrete arrays for encoding-aware rewrites. + pub arrays: &'a [ArrayRef], + + /// 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, +} + +/// An [`ExecutionArgs`] view over borrowed arrays with an explicit row count. +pub(super) 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<'a> BorrowedExecutionArgs<'a> { + pub(super) fn new(inputs: &'a [ArrayRef], row_count: usize) -> Self { + Self { inputs, row_count } + } +} + +impl ExecutionArgs for BorrowedExecutionArgs<'_> { + fn get(&self, index: usize) -> VortexResult { + self.inputs.get(index).cloned().ok_or_else(|| { + vortex_err!( + "Input index {} out of bounds (num_inputs={})", + index, + self.inputs.len() + ) + }) + } + + fn num_inputs(&self) -> usize { + self.inputs.len() + } + + fn row_count(&self) -> usize { + self.row_count + } +} 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..30670f968a1 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/execution.rs @@ -0,0 +1,440 @@ +// 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::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::BorrowedExecutionArgs; +use super::args::KernelArgs; +use super::policy::BatchPlan; +use super::policy::RowPolicy; +use super::policy::skipping_beats_filtering; +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<'a> { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The arguments as the execution layer handed them over. Every path but the filter strategy + /// gives the kernel these untouched, so it sees the original encodings. + args: &'a dyn ExecutionArgs, + + /// 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<'a> Batch<'a> { + /// 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: &'a dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let nullability = plan.output_dtype.nullability() + | Nullability::from(arg_dtypes.iter().any(DType::is_nullable)); + let result_dtype = plan.output_dtype.with_nullability(nullability); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + args, + 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. `try_unfiltered` receives the + /// original inputs plus a mixed validity mask. `Ok(None)` selects filter-and-scatter. + pub fn execute( + &self, + kernel: impl Fn(KernelArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + KernelArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: any null-constant input forces an all-null result without evaluating the + // kernel. + if self + .inputs + .iter() + .any(|input| input.as_constant().is_some_and(|scalar| scalar.is_null())) + { + return Ok(self.all_null()); + } + + // 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.args.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 { + filtered_decode_cost, + } => self.execute_valid_only(kernel, try_unfiltered, filtered_decode_cost, 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 args = BorrowedExecutionArgs::new(&one_row, 1); + let result = VortexResult::from(kernel(self.kernel_args(&args, &one_row), ctx)?)?; + let scalar = self.finalize_output(result, 1)?.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.args.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 { + // Every row is null, so the kernel has nothing to contribute. + if matches!(self.validity, Validity::AllInvalid) { + return Ok(self.all_null()); + } + + let values = match kernel(self.kernel_args(self.args, &self.inputs), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self + .validity + .clone() + .execute_mask(self.args.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.args.row_count()) + } + Validity::Array(valid) => { + self.finalize_output(values.mask(valid)?, self.args.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.args.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.args, &self.inputs), ctx)?)?, + self.args.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 when worthwhile, 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>, + filtered_decode_cost: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedMask::Decided(result) => return Ok(result), + ResolvedMask::Mixed(valid) => valid, + }; + + if skipping_beats_filtering(filtered_decode_cost, &valid) + && 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.args, &self.inputs), 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 args = BorrowedExecutionArgs::new(&filtered, valid.true_count()); + let values = VortexResult::from(kernel(self.kernel_args(&args, &filtered), 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.args.row_count(), + ) + .into_array() + } + + /// Pair an input view with this batch's planning metadata. + fn kernel_args<'b>( + &'b self, + execution: &'b dyn ExecutionArgs, + arrays: &'b [ArrayRef], + ) -> KernelArgs<'b> { + KernelArgs { + execution, + arrays, + 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..1b492f1ff6f --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/mod.rs @@ -0,0 +1,22 @@ +// 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; 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..1b6f3f1f6cc --- /dev/null +++ b/vortex-array/src/scalar_fn/row/batch/policy.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Nullable execution strategies derived from a concrete row dispatch. + +use vortex_mask::Mask; + +use crate::dtype::DType; +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, +} + +/// 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, choosing between skip-invalid execution and filtering based on the + /// mask and decode cost. + ValidOnly { + /// Relative per-row decode work that filtering would avoid. + filtered_decode_cost: usize, + }, +} + +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 { + filtered_decode_cost: Args::FILTERED_DECODE_COST, + } + } + } + + /// 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 { + filtered_decode_cost: Args::FILTERED_DECODE_COST, + } + } + } + + /// The policy one concrete dispatch executes nullable rows under. + /// + /// This deliberately ignores [`OutputSink::SUPPORTS_SKIPPED_ROWS`]. Batch execution 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. + /// + /// [`OutputSink::SUPPORTS_SKIPPED_ROWS`]: crate::scalar_fn::OutputSink::SUPPORTS_SKIPPED_ROWS + pub const fn for_sink() -> Self { + if Args::DENSE_SAFE && !Args::DECODE_FALLIBLE && !ApplyResult::FALLIBLE { + if ApplyResult::DEFERRED { + Self::DenseWithRetry + } else { + Self::Dense + } + } else { + Self::ValidOnly { + filtered_decode_cost: Args::FILTERED_DECODE_COST, + } + } + } +} + +/// Minimum surviving-row fractions for skipping when filtering avoids per-row decode work. +/// The thresholds distinguish one costly decode from multiple costly decodes. +const ONE_DECODE_SKIP_MIN_SURVIVING_FRACTION: f64 = 0.50; +const MULTI_DECODE_SKIP_MIN_SURVIVING_FRACTION: f64 = 0.85; + +/// Whether skipping invalid rows should be preferred over filtering for a mixed mask. +pub(super) fn skipping_beats_filtering(filtered_decode_cost: usize, valid: &Mask) -> bool { + if filtered_decode_cost == 0 { + return true; + } + + let minimum = if filtered_decode_cost == 1 { + ONE_DECODE_SKIP_MIN_SURVIVING_FRACTION + } else { + MULTI_DECODE_SKIP_MIN_SURVIVING_FRACTION + }; + + valid.true_count() as f64 >= valid.len() as f64 * minimum +} 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..cdfad84ee7a --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/mod.rs @@ -0,0 +1,52 @@ +// 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 crate::ArrayRef; + +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), + } + } +} 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..f2246875143 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/owned.rs @@ -0,0 +1,112 @@ +// 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. + 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..abb58f95b02 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/execute/sink.rs @@ -0,0 +1,192 @@ +// 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 crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::DeferredError; +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 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) = Args::varying(&columns) { + vortex_ensure!( + Args::varying_len_matches(&varying, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + for index in 0..row_count { + apply( + &prepared, + Args::get_varying(&varying, index), + Sink::row(&mut rows, index), + ) + .accumulate(&mut accumulated)?; + } + } else { + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + for index in 0..row_count { + apply( + &prepared, + Args::get(&columns, index), + Sink::row(&mut rows, index), + ) + .accumulate(&mut accumulated)?; + } + } + } + + // Immediate failures returned above. Only reduced, deferred failure evidence reaches finish. + finish_sink(sink, DeferredError::new(ApplyResult::occurred(accumulated))) +} + +/// 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, +{ + // Batch execution needs a full-length result before applying the validity mask. Decline when + // the sink cannot leave legal placeholders in positions this loop skips. + if !Sink::SUPPORTS_SKIPPED_ROWS { + 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"); + }; + + { + 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); + let lens_match = match &varying { + Some(varying) => Args::varying_len_matches(varying, row_count), + None => Args::decoded_lens_match(&columns, row_count), + }; + vortex_ensure!( + lens_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + // 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. + Sink::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, + Args::get_varying(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); + } + } + + // Immediate failures returned above. Only reduced, deferred failure evidence reaches finish. + finish_sink(sink, DeferredError::new(ApplyResult::occurred(accumulated))).map(Some) +} + +/// Classify a sink error as retryable only when row accumulation recorded a deferred failure. +/// +/// The sink contract requires [`OutputSink::finish`] to surface recorded failure evidence. Without +/// that evidence, its error is structural and retrying over a different set of rows cannot help. +fn finish_sink( + sink: S, + deferred_error: DeferredError, +) -> VortexResult { + match sink.finish(deferred_error) { + Ok(output) => Ok(RowExecution::Output(output)), + Err(error) if deferred_error.occurred() => Ok(RowExecution::DeferredError(error)), + Err(error) => Err(error), + } +} 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..3351c24d100 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/mod.rs @@ -0,0 +1,36 @@ +// 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`] and [`DeferredError`] describe 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; + +mod batch; + +mod row_fn; +pub use row_fn::RowFn; + +mod types; +pub use types::DeferredError; +pub use types::ElementTuple; +pub use types::IndexedElementTuple; +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..68d92c192ad --- /dev/null +++ b/vortex-array/src/scalar_fn/row/row_fn.rs @@ -0,0 +1,90 @@ +// 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::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. `Some(output)` skips that loop. The output can + /// remain encoded or lazy. Filter-and-scatter execution can pass compacted inputs. + /// + /// # 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..e5c29f756b5 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/bool.rs @@ -0,0 +1,68 @@ +// 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; + +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) + } +} + +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..3aa92df9121 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/mod.rs @@ -0,0 +1,129 @@ +// 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. +pub 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. + 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; + + /// A relative unit count for per-row decode work avoided by filtering this argument first. + /// + /// Leave this at zero for bulk canonicalization. Use a positive value when filtering first + /// avoids meaningful per-row decode work. The executor adds this cost across arguments when it + /// chooses between skipping invalid rows and filtering. + const FILTERED_DECODE_COST: usize = 0; + + /// 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. + 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; +} + +/// 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..05fddbd25e4 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/primitive.rs @@ -0,0 +1,74 @@ +// 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; + +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] + } +} + +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..becde073a0b --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/element/tuple.rs @@ -0,0 +1,414 @@ +// 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; + + /// The additive cost of per-row decode work avoided by filtering the arguments first. + const FILTERED_DECODE_COST: usize; + + /// 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 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; + const FILTERED_DECODE_COST: usize = 0; + + 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> {} + + 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; + const FILTERED_DECODE_COST: usize = $($t::FILTERED_DECODE_COST +)+ 0; + + 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),)+) + } + + 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 super::UnaryTupleSource; + + #[test] + fn 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,)); + } +} 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..2998e19ccbd --- /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::DeferredError; +pub use result::SinkResult; + +mod sink; +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..efd841cc969 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/result.rs @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! What a sink-writing row closure may return. + +use std::ops::BitOrAssign; + +use vortex_error::VortexResult; + +mod private { + pub trait Sealed {} +} + +/// A value-dependent failure bit reduced across the row loop and handed to the output sink. +/// +/// Unlike [`VortexResult`], this never exits the loop. Use it when every row can write a safe +/// provisional value and report failure once at the end. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct DeferredError( + /// Whether this value records a deferred error. + bool, +); + +impl DeferredError { + /// Record whether this row encountered an error. + pub const fn new(failed: bool) -> Self { + Self(failed) + } + + /// Whether any row accumulated into this value failed. + pub const fn occurred(self) -> bool { + self.0 + } +} + +impl BitOrAssign for DeferredError { + fn bitor_assign(&mut self, rhs: Self) { + self.0 |= rhs.0; + } +} + +/// The result of writing one row: success, an immediate error, or deferred error evidence. +/// +/// The executor OR-reduces [`Accumulated`](Self::Accumulated) in a loop-local. The accumulated word +/// should be no wider than the computed element so error tracking does not constrain vector width. +/// This trait is sealed; row functions choose one of its supplied implementations. +pub trait SinkResult: 'static + private::Sealed { + /// The word this result reduces into, kept in a loop-local by the executor. + type Accumulated: 'static + Copy + Default; + + /// Whether this return type can carry an error. + const FALLIBLE: bool; + + /// Whether this result defers failure reporting until the sink finishes. + const DEFERRED: bool; + + /// Merge this row's outcome into the batch-wide reduction. + fn accumulate(self, accumulated: &mut Self::Accumulated) -> VortexResult<()>; + + /// Whether the finished reduction means some row failed. + fn occurred(accumulated: Self::Accumulated) -> bool; +} + +impl private::Sealed for () {} + +impl SinkResult for () { + type Accumulated = (); + + const FALLIBLE: bool = false; + const DEFERRED: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + Ok(()) + } + + fn occurred(_accumulated: ()) -> bool { + false + } +} + +impl private::Sealed for VortexResult<()> {} + +impl SinkResult for VortexResult<()> { + type Accumulated = (); + + const FALLIBLE: bool = true; + const DEFERRED: bool = false; + + fn accumulate(self, _accumulated: &mut ()) -> VortexResult<()> { + self + } + + fn occurred(_accumulated: ()) -> bool { + false + } +} + +/// The evidence widths a row closure may reduce. `bool` is the ordinary answer; the unsigned +/// integers exist for a kernel whose per-row comparison would cost it its vectorization. +macro_rules! impl_sink_result_word { + ($($word:ty),+ $(,)?) => { + $( + impl private::Sealed for $word {} + + impl SinkResult for $word { + type Accumulated = $word; + + const FALLIBLE: bool = false; + const DEFERRED: bool = true; + + fn accumulate(self, accumulated: &mut $word) -> VortexResult<()> { + *accumulated |= self; + Ok(()) + } + + fn occurred(accumulated: $word) -> bool { + accumulated != <$word>::default() + } + } + )+ + }; +} + +impl_sink_result_word!(bool, u8, u16, u32, u64); 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..712d209f7e9 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/types/sink.rs @@ -0,0 +1,139 @@ +// 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::DeferredError; +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 +/// [`SUPPORTS_SKIPPED_ROWS`](Self::SUPPORTS_SKIPPED_ROWS) is `true`. +pub trait OutputSink: 'static + Sized { + /// Whether this sink accepts [`DeferredError`] from its row closure instead of requiring a + /// per-row [`VortexResult`]. + /// + /// A supporting sink must return an error from [`finish`](Self::finish) when its deferred error + /// argument occurred. + const ERRORS_ARE_DEFERRED: bool = false; + + /// Whether this sink can finish a full-length output when some rows were never visited. + /// + /// A supporting sink must use [`initialize_skipped_rows`](Self::initialize_skipped_rows) to + /// leave a legal arbitrary value at every skipped row. Batch execution masks those values. + const SUPPORTS_SKIPPED_ROWS: bool = false; + + /// A loop-local view of all output rows. + /// + /// Borrowed once before execution so the sink's buffer descriptor and shape become loop + /// invariants rather than being re-read through `&mut Self` for every row. + type Rows<'a> + where + Self: 'a; + + /// The place a row closure writes one row through, borrowed from the sink. + type Row<'a> + where + Self: 'a; + + /// 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; + + /// Initialize output positions that skip-invalid execution can omit. + /// + /// Called only when [`SUPPORTS_SKIPPED_ROWS`](Self::SUPPORTS_SKIPPED_ROWS) is `true`. The + /// default is for sinks whose allocation already contains legal values. + fn initialize_skipped_rows(_rows: &mut Self::Rows<'_>) {} + + /// 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 with whether any deferred row error + /// occurred. + fn finish(self, error: DeferredError) -> VortexResult; +} + +/// An element sink that leaves dense output uninitialized before the row loop. +/// +/// 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 { + const SUPPORTS_SKIPPED_ROWS: bool = true; + + type Rows<'a> = &'a mut [MaybeUninit]; + type Row<'a> = &'a mut MaybeUninit; + + 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 initialize_skipped_rows(rows: &mut Self::Rows<'_>) { + for row in rows.iter_mut() { + row.write(T::default()); + } + } + + fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + fn finish(mut self, _error: DeferredError) -> VortexResult { + // SAFETY: dense execution writes every row, while skip-invalid execution initializes every + // row before overwriting valid ones. The executor calls `finish` only after successful + // execution, and the allocation reserved every slot in `0..row_count`. + 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..a000ce38369 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/check.rs @@ -0,0 +1,126 @@ +// 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, + Sink: OutputSink, + ApplyResult: SinkResult, +{ + assert_input_visit_contract::(); + assert!( + !ApplyResult::FALLIBLE || Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result can fail", + ); + assert!( + !ApplyResult::DEFERRED || Function::FALLIBLE, + "RowFn::FALLIBLE must be true when a row result defers failure evidence", + ); + assert!( + Sink::ERRORS_ARE_DEFERRED == ApplyResult::DEFERRED, + "OutputSink::ERRORS_ARE_DEFERRED must match SinkResult::DEFERRED", + ); +} + +/// 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..36fd03af5d2 --- /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..bf172103b5d --- /dev/null +++ b/vortex-array/src/scalar_fn/row/visitor/mod.rs @@ -0,0 +1,157 @@ +// 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 [`visit`] and + /// [`visit_deferred`](Self::visit_deferred). + type VisitResult; + + /// Visit an infallible row computation that returns one independent output value. + /// + /// # 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. + /// + /// # 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. + /// - [`OutputSink::ERRORS_ARE_DEFERRED`] **must** match [`SinkResult::DEFERRED`] for the + /// selected `Sink` and `ApplyResult`. + 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. + /// + /// `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..caa17ad651b --- /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..98c2cdc9671 --- /dev/null +++ b/vortex-array/src/scalar_fn/row/vtable.rs @@ -0,0 +1,171 @@ +// 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::dtype::Nullability; +use crate::expr::Expression; +use crate::expr::union_child_validities; +use crate::scalar_fn::Arity; +use crate::scalar_fn::ChildName; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTable; +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))?; + + // Union the output nullability with the nullability of the inputs. This is required for + // strict scalar function semantics. + let nullability = plan.output_dtype.nullability() + | Nullability::from(args.iter().any(DType::is_nullable)); + + Ok(plan.output_dtype.with_nullability(nullability)) + } + + 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 { + execution: args, + arrays: &[], + 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| 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 { + if !args.arrays.is_empty() + && let Some(reduced) = function.reduce_encoded(options, args.arrays, ctx)? + { + return Ok(RowExecution::Output(reduced)); + } + + function.dispatch( + options, + args.dtypes, + ExecuteRows::::new(args.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> { + // Try the encoding-aware path before filtering changes the inputs. The caller masks its + // full-length result with `valid` before returning it. + if let Some(reduced) = function.reduce_encoded(options, args.arrays, ctx)? { + return Ok(Some(RowExecution::Output(reduced))); + } + + function.dispatch( + options, + args.dtypes, + ExecuteValidRows::::new(args.execution, args.output_dtype, valid, ctx), + ) +} + +/// Prepare the batch inputs and execution plan for `function`. +fn prepare_batch<'args, F: RowFn>( + function: &F, + options: &F::Options, + args: &'args dyn ExecutionArgs, +) -> VortexResult> { + Batch::new(RowFn::id(function), args, |arg_dtypes| { + function.dispatch(options, arg_dtypes, PlanRows::::new(arg_dtypes)) + }) +} From 89fd28bc137a39acfa2bdc939497b21baa6e9002 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sat, 8 Aug 2026 10:42:44 -0400 Subject: [PATCH 2/3] Execute primitive numeric operators with RowFn Signed-off-by: Connor Tsui --- vortex-array/benches/binary_ops.rs | 8 + .../typed_view/primitive/numeric_operator.rs | 2 +- .../scalar_fn/fns/binary/numeric/checked.rs | 88 +---- .../src/scalar_fn/fns/binary/numeric/mod.rs | 12 +- .../scalar_fn/fns/binary/numeric/primitive.rs | 355 ++++-------------- .../src/scalar_fn/fns/binary/numeric/row.rs | 131 +++++++ .../src/scalar_fn/fns/binary/numeric/tests.rs | 9 +- 7 files changed, 236 insertions(+), 369 deletions(-) create mode 100644 vortex-array/src/scalar_fn/fns/binary/numeric/row.rs diff --git a/vortex-array/benches/binary_ops.rs b/vortex-array/benches/binary_ops.rs index 6a07d03f50b..3bd466da0b1 100644 --- a/vortex-array/benches/binary_ops.rs +++ b/vortex-array/benches/binary_ops.rs @@ -170,6 +170,14 @@ fn div_i64_nonnull(bencher: Bencher) { bench_primitive(bencher, lhs, rhs, Operator::Div); } +#[divan::bench] +fn div_i64_nullable(bencher: Bencher) { + let lhs = primitive_nullable(1_000_000, 7).into_array(); + let rhs = primitive_nullable(17, 5).into_array(); + + bench_primitive(bencher, lhs, rhs, Operator::Div); +} + #[divan::bench] fn sub_i64_constant(bencher: Bencher) { let lhs = primitive_nonnull(0).into_array(); diff --git a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs index 6b389a0554b..d7bcc8d0b4c 100644 --- a/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs +++ b/vortex-array/src/scalar/typed_view/primitive/numeric_operator.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_err; use crate::scalar_fn::fns::operators::Operator; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] /// Binary element-wise operations. pub enum NumericOperator { /// Binary element-wise addition of two arrays or of two scalars. diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs index afb709abe1c..054846b7ef7 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/checked.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Checked-lane execution for numeric kernels, driven by the shared +//! Checked-lane execution for the decimal kernels, driven by the shared //! `vortex-compute` lane kernels. - -use std::ops::BitOrAssign; +//! +//! The primitive widths do not come through here: they are computed one row at a time by +//! [`row`](super::row), which writes a value for every row and reduces failure evidence without +//! scanning the finished output. use vortex_buffer::Buffer; use vortex_buffer::BufferMut; @@ -13,33 +15,18 @@ use vortex_compute::lane_kernels::IndexedSourceExt; use vortex_mask::AllOr; use vortex_mask::Mask; -/// Evidence that a lane failed, anything other than [`Default`] meaning failure. -/// -/// `bool` is the ordinary choice; [`map_checked_into`] explains why the wider members exist and -/// asserts the width bound that membership here does **not** imply. -/// -/// [`map_checked_into`]: IndexedSourceExt::map_checked_into -pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - -impl Failure for bool {} -impl Failure for u8 {} -impl Failure for u16 {} -impl Failure for u32 {} -impl Failure for u64 {} - /// Apply the fallible `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it returns `None` on a _valid_ lane. +/// `Err(first_failing_valid_lane)` only when it returns `None` on a valid lane. /// /// `apply` also runs on invalid lanes, whose failures are masked out and whose values are /// unspecified, so it must be total: no panics or side effects on any stored lane value. /// -/// This drives the one-pass early-exit kernels, which abort at the end of the enclosing 64-lane -/// chunk. Use it when per-lane failure handling is cheap relative to the operation, as in integer -/// division. Prefer [`checked_apply_lanes`] when the failure check itself vectorizes. +/// This drives the one-pass early-exit kernels: failures abort at the end of the enclosing +/// 64-lane chunk. It suits an operation whose per-lane failure handling is cheap relative to the +/// operation itself, which is what the decimal kernels and their per-lane casts are. /// -/// `#[inline]`: this must inline into the caller that builds the closure, so that a captured -/// constant operand flattens into a register rather than living behind a pointer the loop reloads -/// on every lane, which blocks vectorization. +/// Keep this wrapper inlineable so captured constants can become loop invariants in the caller. +/// The lane kernels retain their own inlining decisions. #[inline] pub(super) fn checked_lanes( source: S, @@ -61,7 +48,6 @@ where }; let mut values = BufferMut::::with_capacity(len); - let out = &mut values.spare_capacity_mut()[..len]; match valid_bits { None => source.try_map_into(out, apply)?, @@ -73,55 +59,3 @@ where Ok(values.freeze()) } - -/// Apply the split value/failure `apply` over every lane of `source`, returning -/// `Err(first_failing_valid_lane)` only when it flags a _valid_ lane. -/// -/// The hot pass writes every value unconditionally and OR-reduces one piece of evidence, leaving -/// the loop free of per-lane selects and per-chunk exit branches. Only if a lane flagged does a -/// cold second pass re-run `apply` through the early-exit kernels, which drop null-lane failures -/// and attribute the first valid one. -/// -/// Like [`checked_lanes`], `apply` runs on invalid lanes and must be total. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -pub(super) fn checked_apply_lanes( - source: S, - valid_rows: &Mask, - mut apply: Apply, -) -> Result, usize> -where - S: IndexedSource + Copy, - T: Copy + Default, - Fail: Failure, - Apply: FnMut(S::Item) -> (T, Fail), -{ - let len = source.len(); - debug_assert_eq!(len, valid_rows.len()); - - let valid_bits = match valid_rows.bit_buffer() { - AllOr::All => None, - AllOr::None => return Ok(Buffer::zeroed(len)), - AllOr::Some(valid_bits) => Some(valid_bits), - }; - - let mut values = BufferMut::::with_capacity(len); - - let out = &mut values.spare_capacity_mut()[..len]; - if source.map_checked_into(out, &mut apply) != Fail::default() { - let mut checked = |item: S::Item| { - let (value, failure) = apply(item); - (failure == Fail::default()).then_some(value) - }; - match valid_bits { - None => source.try_map_into(out, &mut checked)?, - Some(valid_bits) => source.try_map_masked_into(valid_bits, out, &mut checked)?, - } - } - - // SAFETY: the kernels initialize every lane in `out`. - unsafe { values.set_len(len) }; - - Ok(values.freeze()) -} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs index 6dc0de0fbea..c7ae86b93c9 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/mod.rs @@ -4,16 +4,19 @@ //! Native execution of the arithmetic operators (Add/Sub/Mul/Div) of the [`Binary`] scalar //! function. There is no Arrow fallback. //! +//! The primitive widths are computed by a [`RowFn`](crate::scalar_fn::RowFn), which owns null +//! handling, constants, and validity for them; see [`row`]. Decimal keeps its own columnar +//! implementation in [`decimal`]. +//! //! [`Binary`]: super::Binary mod checked; mod decimal; mod primitive; -#[cfg(test)] -mod tests; +mod row; use decimal::execute_numeric_decimal; -use primitive::execute_numeric_primitive; +use row::execute_numeric_primitive; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -81,3 +84,6 @@ fn build_empty_result( Ok(Canonical::empty(&result_dtype).into_array()) } + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs index 8fd53d15216..42fe3fd3e03 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/primitive.rs @@ -1,73 +1,48 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_compute::lane_kernels::IndexedSource; -use vortex_compute::lane_kernels::LaneZip; -use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_mask::Mask; - -use super::checked::Failure; -use super::checked::checked_apply_lanes; -use super::checked::checked_lanes; -use crate::ArrayRef; -use crate::ExecutionCtx; -use crate::IntoArray; -use crate::arrays::ConstantArray; -use crate::arrays::PrimitiveArray; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; +//! Checked arithmetic for one primitive row. + +use std::ops::BitOrAssign; + use crate::dtype::NativePType; -use crate::dtype::PType; use crate::dtype::half::f16; -use crate::match_each_native_ptype; -use crate::scalar::NumericOperator; -use crate::scalar::Scalar; -use crate::scalar_fn::fns::binary::primitive_operand::PrimitiveOperand; -use crate::validity::Validity; -struct CheckedAdd; +/// Checked addition, failing on integer overflow. +pub(super) struct CheckedAdd; -struct CheckedSub; +/// Checked subtraction, failing on integer overflow. +pub(super) struct CheckedSub; -struct CheckedMul; +/// Checked multiplication, failing on integer overflow. +pub(super) struct CheckedMul; -struct CheckedDiv; +/// Checked division, failing on integer division by zero and on `MIN / -1`. +pub(super) struct CheckedDiv; -trait CheckedPrimitiveOp: Sized { - /// Message for the error raised when this operation fails on a valid lane. - const ERROR: &'static str; +/// OR-reducible evidence that a row failed, with [`Default`] meaning success. +pub(super) trait Failure: Copy + Default + PartialEq + BitOrAssign {} - /// Whether to check inside the value loop and exit early, rather than reducing evidence over - /// the whole batch and re-running only once some lane has flagged. - const CHECKED_VALUE_LOOP: bool = false; +impl Failure for T {} - /// How this operation reports a failing lane. See [`Failure`]. - type Failure: Failure; - - /// Compute a lane's value and its failure evidence together. - /// - /// Overflowing-style rather than `Option`, so the vectorizable kernels can write every - /// lane's value unconditionally and reduce the evidence separately. [`Self::checked`] is the - /// shape for the scalar and early-exit paths. - fn apply(lhs: T, rhs: T) -> (T, Self::Failure); +/// One arithmetic operator at one width, split into its value and failure evidence. +pub(super) trait CheckedPrimitiveOp: 'static + Sized { + /// The error reported for a batch in which some valid row failed. + const ERROR: &'static str; - /// [`Self::apply`] folded into the `Option` that the early-exit kernels take. - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - let (value, failed) = Self::apply(lhs, rhs); + /// How this operation reports a failing row. See [`Failure`]. + type Fail: Failure; - (failed == Self::Failure::default()).then_some(value) - } + /// The result of this operation, paired with evidence of whether the row failed. + fn apply(lhs: T, rhs: T) -> (T, Self::Fail); } impl CheckedPrimitiveOp for CheckedAdd { const ERROR: &'static str = "integer overflow in checked add"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.add_value(rhs), lhs.add_error(rhs)) } @@ -76,9 +51,9 @@ impl CheckedPrimitiveOp for CheckedAdd { impl CheckedPrimitiveOp for CheckedSub { const ERROR: &'static str = "integer overflow in checked sub"; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { (lhs.sub_value(rhs), lhs.sub_error(rhs)) } @@ -87,9 +62,9 @@ impl CheckedPrimitiveOp for CheckedSub { impl CheckedPrimitiveOp for CheckedMul { const ERROR: &'static str = "integer overflow in checked mul"; - type Failure = T::MulFailure; + type Fail = T::MulFailure; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, T::MulFailure) { (lhs.mul_value(rhs), lhs.mul_failure(rhs)) } @@ -97,16 +72,10 @@ impl CheckedPrimitiveOp for CheckedMul { impl CheckedPrimitiveOp for CheckedDiv { const ERROR: &'static str = "integer division by zero or overflow in checked div"; - // Integer division still lowers to scalar divides, so the split - // value/error-scan loop used to auto-vectorize add/sub/mul only adds a - // second full scan. Use the one-pass early-exit checked kernel for integer - // division, matching Arrow/Velox. Float division has no checked errors and - // stays on the split/vectorizable default path. - const CHECKED_VALUE_LOOP: bool = T::DIV_CHECKS_IN_VALUE_LOOP; - type Failure = bool; + type Fail = bool; - #[inline(always)] + #[inline] fn apply(lhs: T, rhs: T) -> (T, bool) { let failed = lhs.div_error(rhs); let value = if failed { @@ -116,151 +85,13 @@ impl CheckedPrimitiveOp for CheckedDiv { }; (value, failed) } - - #[inline(always)] - fn checked(lhs: T, rhs: T) -> Option { - lhs.div_checked(rhs) - } -} - -/// Execute a numeric operation between two primitive-typed arrays. -pub(super) fn execute_numeric_primitive( - lhs: &ArrayRef, - rhs: &ArrayRef, - op: NumericOperator, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let ptype = PType::try_from(lhs.dtype())?; - - match_each_native_ptype!(ptype, |T| { - match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), - } - }) -} - -fn execute_checked_typed( - lhs: &ArrayRef, - rhs: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: NativePType, - Op: CheckedPrimitiveOp, - Scalar: From, - Scalar: From>, -{ - let result_dtype = lhs - .dtype() - .with_nullability(lhs.dtype().nullability() | rhs.dtype().nullability()); - let lhs = PrimitiveOperand::::try_new(lhs, ctx)?; - let rhs = PrimitiveOperand::::try_new(rhs, ctx)?; - let len = lhs.len(); - debug_assert_eq!(len, rhs.len()); - - let validity = lhs.validity().and(rhs.validity())?; - let valid_rows = validity.execute_mask(len, ctx)?; - - let values = match (&lhs, &rhs) { - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_op_lanes::<_, T, Op>( - LaneZip::new(lhs.as_slice(), rhs.as_slice()), - &valid_rows, - |(lhs, rhs)| (lhs, rhs), - ), - ( - PrimitiveOperand::Array { values: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - // Capture the constant by value so it stays hoisted out of the lane loop. - let rhs = *rhs; - checked_op_lanes::<_, T, Op>(lhs.as_slice(), &valid_rows, move |lhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let lhs = *lhs; - checked_op_lanes::<_, T, Op>(rhs.as_slice(), &valid_rows, move |rhs| (lhs, rhs)) - } - ( - PrimitiveOperand::Constant { value: lhs, .. }, - PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let value = Op::checked(*lhs, *rhs) - .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - return Ok(constant_result_array(value, len, &result_dtype)); - } - (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => Ok(Buffer::zeroed(len)), - } - .map_err(|_lane| vortex_err!(InvalidArgument: "{}", Op::ERROR))?; - - primitive_result_array::(values, validity, &result_dtype) } -/// Run `Op` over the lanes of `source` in the loop shape it declares: the one-pass early-exit -/// kernel when `Op::CHECKED_VALUE_LOOP` is set, and the split value/failure kernel otherwise. -/// -/// `#[inline]`: see [`checked_lanes`]. -#[inline] -fn checked_op_lanes( - source: S, - valid_rows: &Mask, - mut to_operands: impl FnMut(S::Item) -> (T, T), -) -> Result, usize> -where - S: IndexedSource + Copy, - T: NativePType, - Op: CheckedPrimitiveOp, -{ - if Op::CHECKED_VALUE_LOOP { - checked_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::checked(lhs, rhs) - }) - } else { - checked_apply_lanes(source, valid_rows, |item| { - let (lhs, rhs) = to_operands(item); - Op::apply(lhs, rhs) - }) - } -} - -fn primitive_result_array( - values: Buffer, - validity: Validity, - dtype: &DType, -) -> VortexResult { - let array = PrimitiveArray::new(values, validity).into_array(); - if array.dtype() == dtype { - return Ok(array); - } - array.cast(dtype.clone()) -} - -fn constant_result_array(value: T, len: usize, dtype: &DType) -> ArrayRef -where - T: NativePType, - Scalar: From + From>, -{ - if dtype.is_nullable() { - ConstantArray::new(Some(value), len).into_array() - } else { - ConstantArray::new(value, len).into_array() - } -} - -trait CheckedArithmetic: NativePType { - const DIV_CHECKS_IN_VALUE_LOOP: bool; - - /// How multiplication reports a failing lane, which is width-dependent in a way the other - /// operations are not. `impl_checked_unsigned!` and `impl_checked_signed!` say why each family - /// reports what it reports. +/// Per-width checked arithmetic. Every value method **must** be total over stored lane values. +pub(super) trait CheckedArithmetic: NativePType { + /// How multiplication reports a failing row. + /// + /// This may be a word rather than `bool` when narrowing evidence would block vectorization. type MulFailure: Failure; fn add_value(self, rhs: Self) -> Self; @@ -271,16 +102,9 @@ trait CheckedArithmetic: NativePType { fn mul_failure(self, rhs: Self) -> Self::MulFailure; fn div_value(self, rhs: Self) -> Self; fn div_error(self, rhs: Self) -> bool; - fn div_checked(self, rhs: Self) -> Option; } -/// The integer arithmetic every width shares, given the two things that actually differ between -/// them: how multiplication reports a failing lane, and how add/sub/div detect one. -/// -/// Only `mul_failure` genuinely varies per width, so it is the macro's parameter and everything -/// else is written once. The `$mul_failure_ty` and body are supplied by the caller because the -/// choice between a widened high half and a `bool` is exactly the vectorization decision described -/// on [`Failure`]. +/// Generate the shared integer operations from their failure predicates. macro_rules! impl_checked_integer { ( $ty:ty, @@ -291,67 +115,57 @@ macro_rules! impl_checked_integer { = |$mf_lhs:ident, $mf_rhs:ident| $mul_failure:expr, ) => { impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = true; - type MulFailure = $mul_failure_ty; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self.wrapping_add(rhs) } - #[inline(always)] + #[inline] fn add_error(self, rhs: Self) -> bool { let ($add_lhs, $add_rhs) = (self, rhs); $add_error } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self.wrapping_sub(rhs) } - #[inline(always)] + #[inline] fn sub_error(self, rhs: Self) -> bool { let ($sub_lhs, $sub_rhs) = (self, rhs); $sub_error } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self.wrapping_mul(rhs) } - #[inline(always)] + #[inline] $(#[$mul_failure_attr])* fn mul_failure(self, rhs: Self) -> $mul_failure_ty { let ($mf_lhs, $mf_rhs) = (self, rhs); $mul_failure } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, rhs: Self) -> bool { let ($div_lhs, $div_rhs) = (self, rhs); $div_error } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - self.checked_div(rhs) - } } }; } -/// The unsigned widths. `add`, `sub` and `div` are written once here, and `widening_mul` covers -/// every width including the 64-bit one, which widens into `u128`: the discarded high half of the -/// widened product is the failure evidence, and costs none of the comparison LLVM folds into -/// `umul.with.overflow`. +/// Unsigned multiplication reports its discarded high half as failure evidence. macro_rules! impl_checked_unsigned { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_integer!( @@ -364,12 +178,7 @@ macro_rules! impl_checked_unsigned { }; } -/// The signed widths, on the same principle. `widening_mul` is the shorthand for the narrow widths, -/// whose two-sided range check over a wider product is not the shape LLVM folds into an overflow -/// intrinsic, so they vectorize while reporting a plain `bool`. The 64-bit width cannot: deriving a -/// `bool` there costs the comparison that scalarizes the loop, so `high_half_mul` hands back the -/// discarded high half as a word. Both arms derive every shift and bound from `$ty`, so -/// instantiating one at a new width cannot silently keep another width's constants. +/// Signed widths use a range check or discarded high-half evidence. macro_rules! impl_checked_signed { ($ty:ty, widening_mul: $wide:ty) => { impl_checked_signed!($ty, mul_failure: bool = |lhs, rhs| { @@ -377,9 +186,6 @@ macro_rules! impl_checked_signed { product < <$ty>::MIN as $wide || product > <$ty>::MAX as $wide }); }; - // Zero exactly when the product fits: a signed multiply overflows iff the high half of the true - // product differs from the sign extension of the half that was kept, so the two XOR to zero on - // the lanes that fit. `tests::test_multiply_overflow_boundaries` pins the boundaries. ($ty:ty, high_half_mul: $wide:ty => $failure:ty) => { impl_checked_signed!($ty, mul_failure: #[expect( clippy::cast_possible_truncation, @@ -395,7 +201,7 @@ macro_rules! impl_checked_signed { ( $ty:ty, mul_failure: $(#[$mul_failure_attr:meta])* $mul_failure_ty:ty - = |$l:ident, $r:ident| $mul_failure:expr + = |$lhs:ident, $rhs:ident| $mul_failure:expr ) => { impl_checked_integer!( $ty, @@ -408,7 +214,7 @@ macro_rules! impl_checked_signed { ((lhs ^ rhs) & (lhs ^ value)) < 0 }, div_error: |lhs, rhs| rhs == 0 || (lhs == <$ty>::MIN && rhs == -1), - mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$l, $r| $mul_failure, + mul_failure: $(#[$mul_failure_attr])* $mul_failure_ty = |$lhs, $rhs| $mul_failure, ); }; } @@ -417,54 +223,47 @@ macro_rules! impl_checked_float { ($($ty:ty),+ $(,)?) => { $( impl CheckedArithmetic for $ty { - const DIV_CHECKS_IN_VALUE_LOOP: bool = false; - type MulFailure = bool; - #[inline(always)] + #[inline] fn add_value(self, rhs: Self) -> Self { self + rhs } - #[inline(always)] + #[inline] fn add_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn sub_value(self, rhs: Self) -> Self { self - rhs } - #[inline(always)] + #[inline] fn sub_error(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn mul_value(self, rhs: Self) -> Self { self * rhs } - #[inline(always)] + #[inline] fn mul_failure(self, _rhs: Self) -> bool { false } - #[inline(always)] + #[inline] fn div_value(self, rhs: Self) -> Self { self / rhs } - #[inline(always)] + #[inline] fn div_error(self, _rhs: Self) -> bool { false } - - #[inline(always)] - fn div_checked(self, rhs: Self) -> Option { - Some(self / rhs) - } } )+ }; @@ -484,30 +283,25 @@ impl_checked_float!(f16, f32, f64); mod tests { use super::CheckedArithmetic; - /// Values whose pairwise products are worth probing: the saturating boundaries, the sign-change - /// pivots, and a spread of magnitudes that straddles the 64-bit split. const PROBES: &[i64] = &[ - 0, // - 1, // - -1, // - 2, // - -2, // - 3, // - i64::MIN, // - i64::MIN + 1, // - i64::MAX, // - i64::MAX - 1, // - 1 << 31, // - 1 << 32, // - 1 << 62, // - -(1 << 62), // - 0x7FFF_FFFF, // - -0x8000_0000, // + 0, + 1, + -1, + 2, + -2, + 3, + i64::MIN, + i64::MIN + 1, + i64::MAX, + i64::MAX - 1, + 1 << 31, + 1 << 32, + 1 << 62, + -(1 << 62), + 0x7FFF_FFFF, + -0x8000_0000, ]; - /// Every `mul_failure` impl is either a bit trick or a two-sided range check, so hold each - /// against the obvious reference: `reference` is the width's own `checked_mul`, whose `None` - /// _is_ the definition of overflow. #[track_caller] fn assert_agrees_with_checked_mul(lhs: T, rhs: T, reference: Option) { let failed = lhs.mul_failure(rhs) != ::default(); @@ -522,14 +316,11 @@ mod tests { assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); let (lhs, rhs) = (lhs as u64, rhs as u64); - assert_agrees_with_checked_mul(lhs, rhs, lhs.checked_mul(rhs)); } } } - /// The 8-bit widths are cheap enough to check exhaustively, which pins the shift of the - /// unsigned formula and the range check of the signed one against every product that exists. #[test] fn mul_failure_agrees_with_checked_mul_exhaustively_at_8_bits() { for lhs in u8::MIN..=u8::MAX { diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs new file mode 100644 index 00000000000..a8e78bdcea2 --- /dev/null +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/row.rs @@ -0,0 +1,131 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Primitive arithmetic execution through [`RowFn`]. +//! +//! `Binary` keeps its registered contract; [`NumericBinary`] is only an execution helper. Decimal +//! arithmetic remains on its existing columnar path. + +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::primitive::CheckedAdd; +use super::primitive::CheckedArithmetic; +use super::primitive::CheckedDiv; +use super::primitive::CheckedMul; +use super::primitive::CheckedPrimitiveOp; +use super::primitive::CheckedSub; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::PType; +use crate::match_each_native_ptype; +use crate::scalar::NumericOperator; +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::scalar_fn::row::UninitElementSink; + +pub(super) fn execute_numeric_primitive( + lhs: &ArrayRef, + rhs: &ArrayRef, + op: NumericOperator, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let args = VecExecutionArgs::new(vec![lhs.clone(), rhs.clone()], lhs.len()); + + ScalarFnVTable::execute(&NumericBinary, &op, &args, ctx) +} + +/// Internal row execution for the primitive arithmetic operators. +#[derive(Clone)] +struct NumericBinary; + +impl RowFn for NumericBinary { + type Options = NumericOperator; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + + // Fallibility is queried without input dtypes, so this conservatively covers integer widths. + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.numeric_binary"); + *ID + } + + fn dispatch( + &self, + op: &Self::Options, + args: &[DType], + visitor: V, + ) -> VortexResult { + let ptype = PType::try_from( + args.first() + .ok_or_else(|| vortex_err!("a numeric operator takes two operands, got none"))?, + )?; + + match_each_native_ptype!(ptype, |T| { + match op { + NumericOperator::Add => visit_checked::(visitor), + NumericOperator::Sub => visit_checked::(visitor), + NumericOperator::Mul => visit_checked::(visitor), + NumericOperator::Div => visit_div::(visitor), + } + }) + } +} + +fn visit_checked(visitor: V) -> VortexResult +where + T: NativePType, + Op: CheckedPrimitiveOp, + V: RowVisitor, +{ + visitor.visit_deferred::<(T, T), T, Op::Fail>( + |(lhs, rhs)| Op::apply(lhs, rhs), + |failure| { + if failure != ::default() { + return Err(numeric_error(Op::ERROR)); + } + + Ok(()) + }, + ) +} + +fn visit_div(visitor: V) -> VortexResult +where + T: CheckedArithmetic, + V: RowVisitor, +{ + if T::PTYPE.is_float() { + return visit_checked::(visitor); + } + + // Integer division is scalar and expensive, so deferring its cheap failure check preserves no + // vectorization. Check each divide immediately and stop at the first failure. + // Dense execution leaves output uninitialized. Nullable branches fill placeholders only when + // they need to skip invalid rows. + visitor.visit_into::<(T, T), UninitElementSink, VortexResult<()>>(|(lhs, rhs), output| { + let (value, failed) = CheckedDiv::apply(lhs, rhs); + if failed { + return Err(numeric_error(>::ERROR)); + } + + output.write(value); + Ok(()) + }) +} + +/// Keep rich error construction out of row closures so the closures remain inlineable. +#[cold] +#[inline(never)] +fn numeric_error(message: &'static str) -> VortexError { + vortex_err!(InvalidArgument: "{message}") +} diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs index 7ee98ae717e..3813c8612b3 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric/tests.rs @@ -201,8 +201,7 @@ fn test_integer_array_array_errors_on_valid_lanes() { assert!(result.is_err()); } -/// Multiply two non-nullable lanes of `lhs` by two of `rhs`, expecting `Some(product)` where the -/// product fits and `None` where the checked kernel must report overflow. +/// Assert one checked multiplication through the complete array execution path. #[track_caller] fn assert_multiply(lhs: T, rhs: T, expected: Option) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); @@ -297,13 +296,11 @@ fn test_multiply_overflow_on_null_lane_ignored( Ok(()) } -/// The hot pass OR-reduces evidence across whole 64-lane chunks before anything looks at it, so an -/// overflow in a late chunk must still be caught, and must still be suppressed when its lane is -/// null. Every other test here fits in a single chunk and cannot show either. +/// An overflow late in the batch must still be reported, unless its row is null. #[rstest] #[case::reported(true)] #[case::suppressed_by_null(false)] -fn test_multiply_overflow_survives_chunk_reduction( +fn test_multiply_overflow_survives_batch_reduction( #[case] lane_is_valid: bool, ) -> VortexResult<()> { const LEN: u32 = 1000; From 85add6465d00a7ec17f495177e000572ff027ce8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Sun, 9 Aug 2026 09:46:58 -0400 Subject: [PATCH 3/3] Run focused take_filter CodSpeed on numeric RowFn Signed-off-by: "Connor Tsui" --- .github/workflows/codspeed.yml | 158 +++------------------------------ 1 file changed, 11 insertions(+), 147 deletions(-) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 8faebf79f86..473dba30d26 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -1,20 +1,10 @@ -name: Codspeed Benchmarking +name: Codspeed RowFn Take Filter -# Concurrency control: -# - PRs: new commits on a feature branch will cancel in-progress (outdated) runs. -# - Push to develop: every commit gets its own group, so baseline runs never cancel and never -# queue behind each other. Serialising them meant a burst of merges left later commits without -# a finished baseline, so CodSpeed fell back to an older comparison base and reported changes -# unrelated to the PR being tested. -# - `workflow_dispatch`: groups by branch and queues if run on develop. concurrency: - group: ${{ github.workflow }}-${{ github.ref }}-${{ github.event_name == 'push' && github.sha || '' }} - cancel-in-progress: ${{ github.ref != 'refs/heads/develop' }} + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true on: - push: - branches: [develop] pull_request: { } - workflow_dispatch: { } permissions: contents: read @@ -22,47 +12,14 @@ permissions: env: CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 - NIGHTLY_TOOLCHAIN: nightly-2026-02-05 jobs: - changes: - name: "Detect CUDA changes" - runs-on: ubuntu-latest - timeout-minutes: 10 - permissions: - pull-requests: read - outputs: - run-cuda-benchmarks: ${{ github.event_name != 'pull_request' || steps.filter.outputs.cuda == 'true' }} - steps: - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 - - uses: dorny/paths-filter@7b450fff21473bca461d4b92ce414b9d0420d706 # v4 - id: filter - if: github.event_name == 'pull_request' - with: - filters: | - cuda: - - "vortex-cuda/**" - # Only this workflow defines the CUDA benchmark jobs. - - ".github/workflows/codspeed.yml" - bench-codspeed: - strategy: - matrix: - include: - - { shard: 1, name: "Core foundation", packages: "vortex-buffer vortex-error vortex-mask vortex-compute vortex-file" } - - { shard: 2, name: "Arrays", packages: "vortex-array", features: "--features _test-harness" } - - { shard: 3, name: "Main library", packages: "vortex" } - - { shard: 4, name: "Encodings 1", packages: "vortex-alp vortex-bytebool vortex-datetime-parts" } - - { shard: 5, name: "Encodings 2", packages: "vortex-decimal-byte-parts vortex-fastlanes vortex-fsst", features: "--features _test-harness" } - - { shard: 6, name: "Encodings 3", packages: "vortex-pco vortex-runend vortex-sequence" } - - { shard: 7, name: "Encodings 4", packages: "vortex-sparse vortex-zigzag vortex-zstd" } - - { shard: 8, name: "Storage formats & row encoding", packages: "vortex-flatbuffers vortex-proto vortex-btrblocks vortex-row" } - - { shard: 9, name: "Tensor & spatial", packages: "vortex-tensor vortex-spatial" } - name: "Benchmark with Codspeed (Shard #${{ matrix.shard }})" + name: "Benchmark take_filter with Codspeed" timeout-minutes: 30 runs-on: >- ${{ github.repository == 'vortex-data/vortex' - && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=bench-codspeed-{1}', github.run_id, matrix.shard) + && format('runs-on={0}/runner=amd64-medium/image=ubuntu24-full-x64-pre-v2/extras=s3-cache/tag=bench-codspeed-take-filter', github.run_id) || 'ubuntu-latest' }} steps: - uses: runs-on/action@v2 @@ -78,108 +35,15 @@ jobs: uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 with: tool: cargo-codspeed - - name: Build benchmarks + - name: Build take_filter env: RUSTFLAGS: "-C target-feature=+avx2" - run: cargo codspeed build ${{ matrix.features }} $(printf -- '-p %s ' ${{ matrix.packages }}) --profile bench - - name: Run benchmarks + run: >- + cargo codspeed build --features _test-harness -p vortex-array + --bench take_filter --profile bench + - name: Run take_filter uses: CodSpeedHQ/action@0ca9cbbf4623b599a6c3ed4fc8a922942705d9f1 # v5 with: - run: cargo codspeed run + run: cargo codspeed run --bench take_filter token: ${{ secrets.CODSPEED_TOKEN }} mode: "simulation" - - # Getting a GPU box is slow, in the future we can build on a box without one and only run - # on GPU machines. - bench-codspeed-cuda-build: - needs: [changes] - if: >- - always() && github.repository == 'vortex-data/vortex' && - needs.changes.outputs.run-cuda-benchmarks == 'true' - name: "Build Codspeed CUDA benchmarks" - timeout-minutes: 30 - runs-on: >- - runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/extras=s3-cache/tag=bench-codspeed-cuda-build - steps: - - uses: runs-on/action@v2 - with: - sccache: s3 - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 - - uses: ./.github/actions/setup-rust - with: - repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "true" - - name: Install Codspeed - uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 - with: - tool: cargo-codspeed - - name: Build benchmarks - run: | - cargo codspeed build -m walltime \ - --bench bitpacked_cuda \ - --bench dynamic_dispatch_cuda \ - --bench alp_cuda \ - --bench date_time_parts_cuda \ - --bench dict_cuda \ - --bench fsst_cuda \ - --bench runend_cuda \ - --profile bench - - name: Package CUB shared library - run: | - find target/release/build -path '*/out/libvortex_cub.so' \ - -exec cp {} target/codspeed/walltime/vortex-cuda/libvortex_cub.so \; - test -f target/codspeed/walltime/vortex-cuda/libvortex_cub.so - - name: Upload benchmark executables - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 - with: - name: codspeed-cuda-benchmarks - path: target/codspeed/ - retention-days: 1 - if-no-files-found: error - - bench-codspeed-cuda: - if: github.repository == 'vortex-data/vortex' - needs: [bench-codspeed-cuda-build] - strategy: - matrix: - include: - - { shard: 1, name: "Bitpacked", benches: "bitpacked_cuda" } - - { shard: 2, name: "Dynamic dispatch", benches: "dynamic_dispatch_cuda" } - - { shard: 3, name: "Standalone kernels", benches: "alp_cuda date_time_parts_cuda dict_cuda fsst_cuda runend_cuda" } - name: "Benchmark with Codspeed (CUDA Shard #${{ matrix.shard }} - ${{ matrix.name }})" - timeout-minutes: 30 - runs-on: >- - runs-on=${{ github.run_id }}/family=g5/cpu=8/image=ubuntu24-gpu-x64/extras=s3-cache/tag=bench-codspeed-cuda-${{ matrix.shard }} - steps: - - uses: runs-on/action@v2 - with: - sccache: s3 - - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7 - - uses: ./.github/actions/setup-rust - with: - repo-token: ${{ secrets.GITHUB_TOKEN }} - enable-sccache: "true" - - name: Display NVIDIA SMI details - run: | - nvidia-smi - nvidia-smi -L - nvidia-smi -q -d Memory - - name: Install Codspeed - uses: taiki-e/cache-cargo-install-action@66c9585ef5ca780ee69399975a5e911f47905995 - with: - tool: cargo-codspeed - - name: Download benchmark executables - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8 - with: - name: codspeed-cuda-benchmarks - path: target/codspeed - - name: Restore executable permissions - run: find target/codspeed -type f -exec chmod +x {} + - - name: Run benchmarks - uses: CodSpeedHQ/action@0ca9cbbf4623b599a6c3ed4fc8a922942705d9f1 # v5 - env: - CARGO_MANIFEST_DIR: ${{ github.workspace }}/vortex-cuda - with: - run: cargo codspeed run $(printf -- '--bench %s ' ${{ matrix.benches }}) - token: ${{ secrets.CODSPEED_TOKEN }} - mode: "walltime"