From a061da6fac851c378f9a3af69cb16db1912d2a7b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 14:06:47 -0400 Subject: [PATCH 1/9] Implement RowFn batch execution Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 93 +++ .../scalar_fn/unstable/row/batch/execution.rs | 481 ++++++++++++++ .../src/scalar_fn/unstable/row/batch/mod.rs | 25 + .../src/scalar_fn/unstable/row/batch/tests.rs | 587 ++++++++++++++++++ .../src/scalar_fn/unstable/row/execute/mod.rs | 19 + .../scalar_fn/unstable/row/execute/outcome.rs | 43 ++ .../scalar_fn/unstable/row/execute/owned.rs | 124 ++++ .../scalar_fn/unstable/row/execute/sink.rs | 262 ++++++++ .../src/scalar_fn/unstable/row/mod.rs | 5 + .../unstable/row/types/element/mod.rs | 1 + .../unstable/row/types/element/tuple/mod.rs | 1 + .../unstable/row/types/element/tuple/tests.rs | 2 +- .../src/scalar_fn/unstable/row/types/mod.rs | 1 + .../scalar_fn/unstable/row/visitor/execute.rs | 263 ++++++++ .../src/scalar_fn/unstable/row/visitor/mod.rs | 7 + .../scalar_fn/unstable/row/visitor/plan.rs | 5 +- .../src/scalar_fn/unstable/row/vtable.rs | 149 ++++- 17 files changed, 2054 insertions(+), 14 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/args.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execution.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/tests.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/owned.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/execute/sink.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs new file mode 100644 index 00000000000..d5ffe042613 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! A borrowed execution view passed to one row-kernel invocation. + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::visitor::RowPolicy; + +/// A borrowed [`ExecutionArgs`] view with the planning metadata selected for its row kernel. +/// +/// `arrays` may be filtered or sliced, while `dtypes` and `output_dtype` always describe the +/// original planned batch. Keeping them together prevents an execution path from pairing an input +/// view with unrelated planning metadata. +#[derive(Clone, Copy)] +pub(in crate::scalar_fn::unstable::row) struct BorrowedExecutionArgs<'a> { + /// The input arrays for this kernel invocation. + arrays: &'a [ArrayRef], + + /// The number of rows in this kernel invocation. + row_count: usize, + + /// The original input dtypes used to select the row implementation. + dtypes: &'a [DType], + + /// The non-nullable dtype built by the selected output capability. + output_dtype: &'a DType, + + /// The nullable execution policy selected during planning. + policy: RowPolicy, +} + +impl<'a> BorrowedExecutionArgs<'a> { + /// Pair one input view with the planning metadata selected for its batch. + pub(in crate::scalar_fn::unstable::row) fn new( + arrays: &'a [ArrayRef], + row_count: usize, + dtypes: &'a [DType], + output_dtype: &'a DType, + policy: RowPolicy, + ) -> Self { + Self { + arrays, + row_count, + dtypes, + output_dtype, + policy, + } + } + + /// Return the concrete arrays used by this row-kernel invocation. + pub(in crate::scalar_fn::unstable::row) fn arrays(&self) -> &'a [ArrayRef] { + self.arrays + } + + /// Return the original input dtypes used to select the row implementation. + pub(in crate::scalar_fn::unstable::row) fn dtypes(&self) -> &'a [DType] { + self.dtypes + } + + /// Return the non-nullable dtype built by the selected output capability. + pub(in crate::scalar_fn::unstable::row) fn output_dtype(&self) -> &'a DType { + self.output_dtype + } + + /// Return the nullable execution policy selected during planning. + pub(in crate::scalar_fn::unstable::row) fn policy(&self) -> RowPolicy { + self.policy + } +} + +impl ExecutionArgs for BorrowedExecutionArgs<'_> { + fn get(&self, index: usize) -> VortexResult { + self.arrays.get(index).cloned().ok_or_else(|| { + vortex_err!( + "row-function input index must be less than {}, got {index}", + self.arrays.len(), + ) + }) + } + + fn num_inputs(&self) -> usize { + self.arrays.len() + } + + fn row_count(&self) -> usize { + self.row_count + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs new file mode 100644 index 00000000000..671b8f4ecb0 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -0,0 +1,481 @@ +// 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::BatchPlan; +use super::RowPolicy; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::Constant; +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::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::types::batch_constant; +use crate::validity::Validity; + +/// How far [`Batch::resolve_validity`] got before a mixed-mask strategy became necessary. +enum ResolvedMask { + /// The all-valid or all-null batch was answered without a mixed-mask strategy. + Decided(ArrayRef), + + /// A mask with both set and unset bits, which a strategy must now execute. + Mixed(Mask), +} + +/// One batch of inputs and the metadata needed before its row kernel runs. +pub struct Batch { + /// The function being executed, named in the errors this raises. + id: ScalarFnId, + + /// The number of rows in the original execution scope. + row_count: usize, + + /// The input columns, collected once: constant folding inspects them and the filter strategy + /// filters them. + inputs: SmallVec<[ArrayRef; 4]>, + + /// The input dtypes, collected with the columns and reused by both planning and execution. + arg_dtypes: SmallVec<[DType; 4]>, + + /// The conjoined input validity, so a row of the output is valid iff it is valid in every + /// input. Conjoining is lazy, and nothing materializes it unless the null handling asks. + validity: Validity, + + /// The dtype the function declares for these inputs, which the kernel's output is reconciled + /// against. Already widened to nullable if any input is nullable. + result_dtype: DType, + + /// The non-nullable dtype the dispatched output capability builds, computed while planning. + output_dtype: DType, + + /// How the concrete dispatch executes nullable rows. + policy: RowPolicy, +} + +impl Batch { + /// Collect the inputs and derive their dtype, validity, and execution policy. + /// + /// **Not** for a nullary function: with no inputs there is no validity to propagate and no + /// per-row work to fold, and the all-constant check below would vacuously pass. + pub fn new( + id: ScalarFnId, + args: &dyn ExecutionArgs, + plan: impl FnOnce(&[DType]) -> VortexResult, + ) -> VortexResult { + let row_count = args.row_count(); + let inputs: SmallVec<[ArrayRef; 4]> = (0..args.num_inputs()) + .map(|index| args.get(index)) + .collect::>()?; + + for (index, input) in inputs.iter().enumerate() { + vortex_ensure_eq!( + input.len(), + row_count, + "the {id} input {index} must have {row_count} rows, got {}", + input.len(), + ); + } + + let arg_dtypes: SmallVec<[DType; 4]> = + inputs.iter().map(|input| input.dtype().clone()).collect(); + let plan = plan(&arg_dtypes)?; + let result_dtype = plan.result_dtype(&arg_dtypes); + + let mut validity = Validity::NonNullable; + for input in &inputs { + validity = validity.and(input.validity()?)?; + } + + Ok(Self { + id, + row_count, + inputs, + arg_dtypes, + validity, + result_dtype, + output_dtype: plan.output_dtype, + policy: plan.policy, + }) + } + + /// Add null propagation, constant folding, and strategy selection around `kernel`. + /// + /// The kernel may ignore input validity. It receives valid-only rows when required, and its + /// output **must** match the planned dtype up to nullability. `try_unfiltered` receives the + /// originals plus a mixed validity mask; `Ok(None)` selects filter-and-scatter. + pub fn execute( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + // Strictness: an all-null batch has no observable row work. Keep the literal-constant + // check explicit alongside the conjoined validity invariant. + if matches!(self.validity, Validity::AllInvalid) + || self.inputs.iter().any(|input| { + input + .as_opt::() + .is_some_and(|constant| constant.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.row_count > 0 + && self.validity.definitely_no_nulls() + && self + .inputs + .iter() + .all(|input| batch_constant(input).is_some()) + { + return self.broadcast_one_row(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, false, ctx), + RowPolicy::DenseWithRetry => self.execute_dense(kernel, true, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_unfiltered, ctx), + } + } + + /// Evaluate a single row of all-constant inputs and broadcast its value. + /// + /// Reconciling the row's dtype before reading the scalar keeps this path on the same + /// kernel/declaration agreement check as the dense and filter paths, rather than letting `cast` + /// paper over a disagreement. + fn broadcast_one_row( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let one_row: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.slice(0..1)) + .collect::>()?; + + let result = VortexResult::from(kernel(self.execution_args(&one_row, 1), ctx)?)?; + let result = self.validate_kernel_output(result, 1, ctx)?; + let result = self.finalize_output(result, 1)?; + let scalar = result.execute_scalar(0, ctx)?; + + Ok(ConstantArray::new(scalar, self.row_count).into_array()) + } + + /// Run the kernel over every row, including the rows behind nulls, then mask its result. + /// + /// The arguments reach the kernel untouched, so the inputs keep their original encoding, and + /// the conjoined validity is handed to `mask` as an array rather than materialized into a + /// [`Mask`] first. + fn execute_dense( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + retry_deferred_error: bool, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let values = match kernel(self.execution_args(&self.inputs, self.row_count), ctx)? { + RowExecution::Output(values) => values, + RowExecution::DeferredError(error) if retry_deferred_error => { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Unlike `resolve_validity`, all-true preserves the deferred error and all-false + // suppresses evidence that came entirely from null rows. An empty loop cannot + // produce deferred evidence, so the ambiguous empty mask cannot reach this arm. + if valid.all_true() { + return Err(error); + } + if valid.all_false() { + return Ok(self.all_null()); + } + + // Deferred retry receives only the dense kernel. Filter first so the retry cannot + // evaluate null rows again. + return self.filter_and_scatter(kernel, &valid, ctx); + } + RowExecution::DeferredError(error) => return Err(error), + }; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + + match self.validity.clone() { + Validity::NonNullable | Validity::AllValid => { + self.finalize_output(values, self.row_count) + } + Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // Handled by the guard above, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } + + /// Materialize validity and answer all-valid or all-null batches before selecting a mixed-mask + /// strategy. + fn resolve_validity( + &self, + kernel: &impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; + + // Check all-true before all-false: an empty mask is both, and must not be treated as + // all-null (a zero-length non-nullable execution keeps its non-nullable dtype). + if valid.all_true() { + let values = VortexResult::from(kernel( + self.execution_args(&self.inputs, self.row_count), + ctx, + )?)?; + let values = self.validate_kernel_output(values, self.row_count, ctx)?; + let values = self.finalize_output(values, self.row_count)?; + + return Ok(ResolvedMask::Decided(values)); + } + + if valid.all_false() { + return Ok(ResolvedMask::Decided(self.all_null())); + } + + Ok(ResolvedMask::Mixed(valid)) + } + + /// Resolve validity, try unfiltered execution, then fall back to filtering. + fn execute_valid_only( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let valid = match self.resolve_validity(&kernel, ctx)? { + ResolvedMask::Decided(result) => return Ok(result), + ResolvedMask::Mixed(valid) => valid, + }; + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Try execution against the original inputs, then mask a returned full-length result. + fn try_execute_unfiltered( + &self, + try_unfiltered: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(execution) = try_unfiltered( + self.execution_args(&self.inputs, self.row_count), + valid, + ctx, + )? + else { + return Ok(None); + }; + let values = VortexResult::from(execution)?; + let values = self.validate_kernel_output(values, valid.len(), ctx)?; + + 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(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let filtered: SmallVec<[ArrayRef; 4]> = self + .inputs + .iter() + .map(|input| input.filter(valid.clone())) + .collect::>()?; + + let values = VortexResult::from(kernel( + self.execution_args(&filtered, valid.true_count()), + ctx, + )?)?; + let values = self.validate_kernel_output(values, valid.true_count(), ctx)?; + + self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) + } + + /// An all-null result of the function's declared return dtype. + fn all_null(&self) -> ArrayRef { + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + } + + /// Pair an input view with this batch's planning metadata. + fn execution_args<'b>( + &'b self, + arrays: &'b [ArrayRef], + row_count: usize, + ) -> BorrowedExecutionArgs<'b> { + BorrowedExecutionArgs::new( + arrays, + row_count, + &self.arg_dtypes, + &self.output_dtype, + self.policy, + ) + } + + /// Finalize an output against this batch's expected length and declared return dtype. + fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { + reconcile_output(self.id, &self.result_dtype, expected_len, values) + } + + /// Validate the output from a row kernel before batch validity is attached. + fn validate_kernel_output( + &self, + values: ArrayRef, + expected_len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + finalize_kernel_output(self.id, &self.output_dtype, expected_len, values, ctx) + } + + /// 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 nullable gathered array cannot be wrapped because a `Masked` child must be all valid. + // The general masking pass unions its nulls with the batch validity instead. + 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 the output produced directly by a row kernel. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` when ignoring +/// nullability, and every produced row **must** be valid. Batch execution owns strict null +/// propagation and attaches input-derived validity only after this boundary. +pub fn finalize_kernel_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + vortex_ensure!( + values.all_valid(ctx)?, + "the {id} row kernel produced nulls for valid rows", + ); + + cast_output_nullability(result_dtype, values) +} + +/// Reconcile an output with the function's declared shape and nullability. +fn reconcile_output( + id: ScalarFnId, + result_dtype: &DType, + expected_len: usize, + values: ArrayRef, +) -> VortexResult { + validate_output(id, result_dtype, expected_len, &values)?; + + cast_output_nullability(result_dtype, values) +} + +/// Validate an output's shape and logical dtype without executing a nullability cast. +fn validate_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(), + ); + + Ok(()) +} + +/// Cast only the output nullability after its shape, dtype, and validity are accepted. +fn cast_output_nullability(result_dtype: &DType, values: ArrayRef) -> VortexResult { + if values.dtype() == result_dtype { + Ok(values) + } else { + values.cast(result_dtype.clone()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs new file mode 100644 index 00000000000..4bb238564ba --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -0,0 +1,25 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution around a non-null row kernel. +//! +//! A row kernel handles typed values for one row. This module adds the columnar concerns around it: +//! planning the output and null strategy, preserving batch constants and encodings, propagating +//! strict validity, selecting an execution strategy, and validating the finished output. +//! +//! [`BatchPlan`] carries the nullable execution strategy selected by a concrete dispatch. [`Batch`] +//! applies that strategy, and [`BorrowedExecutionArgs`] pairs each kernel invocation with its +//! planning metadata. + +mod args; +pub(super) use args::BorrowedExecutionArgs; + +mod execution; +pub(super) use execution::Batch; +pub(super) use execution::finalize_kernel_output; + +pub(super) use super::visitor::BatchPlan; +pub(super) use super::visitor::RowPolicy; + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs new file mode 100644 index 00000000000..6919028715e --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -0,0 +1,587 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use rstest::rstest; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; + +use super::super::execute::RowExecution; +use super::Batch; +use super::BatchPlan; +use super::RowPolicy; +use super::finalize_kernel_output; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; +use crate::scalar_fn::unstable::row::row_fn_return_dtype; +use crate::validity::Validity; + +#[derive(Clone)] +struct RetryConstantAdd; + +#[derive(Clone)] +struct NullarySeven; + +#[derive(Clone)] +struct AddThree; + +#[derive(Clone)] +struct Identity; + +#[derive(Clone)] +struct SinkOptions; + +struct OptionsCheckingSink; + +#[derive(Clone)] +struct InvalidKernelOutput; + +/// Deliberately violates [`OutputElement::build`] to test validation at the public boundary. +struct NullProducingI64(i64); + +#[derive(Clone)] +struct PreparedAdd { + visit: PreparedVisit, + prepares: Arc, +} + +#[derive(Clone, Copy)] +enum PreparedVisit { + Owned, + Sink, + Deferred, +} + +// SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or +// `finish` through the executor. The row-initialization requirements are therefore vacuous. +unsafe impl OutputSink for OptionsCheckingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn sink_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { + if !enabled { + vortex_bail!(InvalidArgument: "the test sink is disabled"); + } + + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + vortex_bail!("the planning-only test sink must not be allocated") + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { + true + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} + + unsafe fn finish(self) -> VortexResult { + vortex_bail!("the planning-only test sink must not finish") + } +} + +impl OutputElement for NullProducingI64 { + fn element_dtype() -> DType { + DType::from(i64::PTYPE) + } + + fn build(values: Vec) -> ArrayRef { + let values: Vec<_> = values.into_iter().map(|value| value.0).collect(); + let validity = Validity::from_iter((0..values.len()).map(|index| index != 0)); + + PrimitiveArray::new(values, validity).into_array() + } +} + +struct I64Sink(BufferMut); + +// SAFETY: every row is initialized by `BufferMut::zeroed`, and the sink exposes exactly that +// initialized slice. The `()` write token therefore proves no additional invariant. +unsafe impl OutputSink for I64Sink { + type Rows<'a> = &'a mut [i64]; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::new(self.0.freeze(), Validity::NonNullable).into_array()) + } +} + +impl RowFn for NullarySeven { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.nullary_seven"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), I64Sink, _>(|(), output| { + *output = 7; + }) + } +} + +impl RowFn for AddThree { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["first", "second", "third"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_three"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64, i64, i64), i64>(|(first, second, third)| first + second + third) + } +} + +impl RowFn for RetryConstantAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.retry_constant_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(u8, u8), u8, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "checked add overflowed")); + } + + Ok(()) + }, + ) + } +} + +impl RowFn for Identity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.identity"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), i64>(|(value,)| value) + } +} + +impl RowFn for SinkOptions { + type Options = bool; + + const ARG_NAMES: &'static [&'static str] = &[]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.sink_options"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(), OptionsCheckingSink, _>(|(), ()| ()) + } +} + +impl RowFn for InvalidKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.invalid_kernel_output"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit::<(i64,), NullProducingI64>(|(value,)| NullProducingI64(value)) + } +} + +impl RowFn for PreparedAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.prepared_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + let prepares = Arc::clone(&self.prepares); + let prepare = move |(_lhs, rhs): (Option, Option)| { + prepares.fetch_add(1, Ordering::Relaxed); + rhs + }; + + match self.visit { + PreparedVisit::Owned => visitor + .visit_prepared::<(i64, i64), i64, _>(prepare, |constant_rhs, (lhs, rhs)| { + lhs.wrapping_add(constant_rhs.unwrap_or(rhs)) + }), + PreparedVisit::Sink => visitor.visit_prepared_into::<(i64, i64), I64Sink, _, ()>( + prepare, + |constant_rhs, (lhs, rhs), output| { + *output = lhs.wrapping_add(constant_rhs.unwrap_or(rhs)); + }, + ), + PreparedVisit::Deferred => visitor.visit_prepared_deferred::<(i64, i64), i64, _, bool>( + prepare, + |constant_rhs, (lhs, rhs)| lhs.overflowing_add(constant_rhs.unwrap_or(rhs)), + |failed| { + if failed { + return Err(vortex_err!(InvalidArgument: "prepared add overflowed")); + } + + Ok(()) + }, + ), + } + } +} + +#[test] +fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.row_batch"); + + let input = PrimitiveArray::new(vec![1i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let result = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::Dense, + }) + }); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![u8::MAX, 1], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let result = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx); + + assert!(result.is_err()); + Ok(()) +} + +#[test] +fn test_dense_retry_suppresses_null_row_failure() -> VortexResult<()> { + let lhs = + PrimitiveArray::new(vec![1, u8::MAX], Validity::from_iter([true, false])).into_array(); + let rhs = ConstantArray::new(1_u8, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_u8, 0], Validity::from_iter([true, false])); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} + +#[test] +fn test_constant_input_broadcasts_one_row() -> VortexResult<()> { + let input = ConstantArray::new(7_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&Identity, &EmptyOptions, &args, &mut ctx)?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[rstest] +#[case::all_valid([true, true])] +#[case::all_invalid([false, false])] +fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.resolve_validity"); + + let validity = Validity::Array(BoolArray::from_iter(validity).into_array()); + let input = PrimitiveArray::new(vec![4_i64, 5], validity).into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 2); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_filters_and_scatters() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.filter_and_scatter"); + + let input = PrimitiveArray::new( + vec![10_i64, 20, 30, 40], + Validity::from_iter([true, false, true, false]), + ) + .into_array(); + let args = VecExecutionArgs::new(vec![input.clone()], 4); + let batch = Batch::new(*ID, &args, |_| { + Ok(BatchPlan { + output_dtype: DType::from(i64::PTYPE), + policy: RowPolicy::ValidOnly, + }) + })?; + let mut ctx = array_session().create_execution_ctx(); + + let actual = batch.execute( + |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |_args, _valid, _ctx| Ok(None), + &mut ctx, + )?; + + assert_arrays_eq!(&actual, &input, &mut ctx); + Ok(()) +} + +#[test] +fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.finalize_kernel_output"); + + let values = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let result_dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + let mut ctx = array_session().create_execution_ctx(); + + let actual = finalize_kernel_output(*ID, &result_dtype, 2, values.clone(), &mut ctx)?; + let expected = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + assert_eq!(actual.dtype(), &result_dtype); + assert_arrays_eq!(&actual, &expected, &mut ctx); + + assert!(finalize_kernel_output(*ID, &result_dtype, 3, values, &mut ctx).is_err()); + + let bools = BoolArray::from_iter([true, false]).into_array(); + assert!(finalize_kernel_output(*ID, &result_dtype, 2, bools, &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn test_sink_dtype_receives_function_options() -> VortexResult<()> { + assert_eq!( + row_fn_return_dtype(&SinkOptions, &true, &[])?, + DType::from(i64::PTYPE) + ); + assert!(row_fn_return_dtype(&SinkOptions, &false, &[]).is_err()); + Ok(()) +} + +#[test] +fn test_nonnullable_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + + assert_invalid_kernel_output(input) +} + +#[test] +fn test_all_valid_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); + + assert_invalid_kernel_output(input) +} + +#[track_caller] +fn assert_invalid_kernel_output(input: ArrayRef) -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + let execution = execute_rows(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); + let error = match execution { + Err(error) => error, + Ok(output) => match output.execute::(&mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("an invalid row kernel output passed boundary validation"), + }, + }; + let error = error.to_string(); + + assert!( + error.contains("test.invalid_kernel_output"), + "the boundary error must name the function, got {error}", + ); + assert!( + error.contains("row kernel produced nulls for valid rows"), + "the boundary error must identify invalid row output, got {error}", + ); + Ok(()) +} + +#[rstest] +#[case::owned_constant(PreparedVisit::Owned, true)] +#[case::owned_per_row(PreparedVisit::Owned, false)] +#[case::sink_constant(PreparedVisit::Sink, true)] +#[case::sink_per_row(PreparedVisit::Sink, false)] +#[case::deferred_constant(PreparedVisit::Deferred, true)] +#[case::deferred_per_row(PreparedVisit::Deferred, false)] +fn test_prepared_visits( + #[case] visit: PreparedVisit, + #[case] constant_rhs: bool, +) -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter([1_i64, 2]).into_array(); + let rhs = if constant_rhs { + ConstantArray::new(3_i64, 2).into_array() + } else { + PrimitiveArray::from_iter([3_i64, 4]).into_array() + }; + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let prepares = Arc::new(AtomicUsize::new(0)); + let function = PreparedAdd { + visit, + prepares: Arc::clone(&prepares), + }; + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&function, &EmptyOptions, &args, &mut ctx)?; + let expected = if constant_rhs { + PrimitiveArray::from_iter([4_i64, 5]).into_array() + } else { + PrimitiveArray::from_iter([4_i64, 6]).into_array() + }; + + assert_arrays_eq!(&actual, &expected, &mut ctx); + assert_eq!(prepares.load(Ordering::Relaxed), 1); + Ok(()) +} + +#[test] +fn test_nullary_row_function_broadcasts() -> VortexResult<()> { + let args = VecExecutionArgs::new(vec![], 3); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&NullarySeven, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([7i64, 7, 7]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_owned_execution_traverses_three_per_row_inputs() -> VortexResult<()> { + let args = VecExecutionArgs::new( + vec![ + PrimitiveArray::from_iter([1_i64, 2, 3]).into_array(), + PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(), + PrimitiveArray::from_iter([100_i64, 200, 300]).into_array(), + ], + 3, + ); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&AddThree, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::from_iter([111_i64, 222, 333]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs new file mode 100644 index 00000000000..73722549f41 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -0,0 +1,19 @@ +// 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 reduces compact failure evidence. [`sink`] +//! drives output builders whose row handles may share batch state. Both return [`RowExecution`], +//! which distinguishes a completed array from a deferred error that batch validity may suppress. + +mod owned; +pub(super) use owned::execute_owned; +pub(super) use owned::execute_owned_infallible; + +mod outcome; +pub use outcome::RowExecution; + +mod sink; +pub(super) use sink::execute_sink; +pub(super) use sink::execute_sink_valid_rows; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs new file mode 100644 index 00000000000..fc013e7a317 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/outcome.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The result of a completed row loop before batch-level null handling. +//! +//! [`RowExecution`] preserves deferred failure evidence until batch execution can determine whether +//! the failing payload belonged to a valid row. + +use vortex_error::VortexError; +use vortex_error::VortexResult; + +use crate::ArrayRef; + +/// 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 can evaluate null payloads, so its deferred error is not always observable. Batch +/// execution can retry only valid rows to discard errors caused by null payloads. A plain +/// `VortexResult` cannot distinguish these errors from failures that a retry cannot fix. +/// +/// 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/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs new file mode 100644 index 00000000000..3ac0c553891 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -0,0 +1,124 @@ +// 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::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::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 stores one value per row, the indexed source removes argument-shape + // dispatch from the hot loop and lets the lane kernel optimize the traversal as one + // operation. Keep view construction and its length proof in this branch. Hoisting them + // through the shared validation helper changed mixed-constant add, subtract, and multiply + // from 9.219, 9.229, and 18.94 us to 30.46, 31.11, and 37.73 us on a Ryzen 9 7950X with + // rustc 1.91.0 and LLVM 21.1.2. + // Restoring this placement recovered the fast code under the 16-CGU, no-LTO bench profile. + if let Some(views) = Args::per_row_views(&columns) { + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // SAFETY: `view_lens_match` proved every view addresses exactly `row_count` rows + // immediately above. + failure = unsafe { Args::indexed_source(views, row_count) } + .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 per-row inputs. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + // Keep the output-slot iterator as the loop bound. `row_count` is address-taken by the + // validation error formatting above. With rustc 1.97.1 and LLVM 22.1.6 under 16 CGUs + // without LTO, indexing `output` by a `0..row_count` range retains an early-exit bounds + // check and prevents vectorization of mixed constant and per-row arithmetic. Recheck + // the optimized IR and mixed-constant benchmarks before restoring that range loop. + let mut accumulated = Fail::default(); + for (index, slot) in output.iter_mut().enumerate() { + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + slot.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/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs new file mode 100644 index 00000000000..fce891cf613 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -0,0 +1,262 @@ +// 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::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::SinkResult; + +/// Ensure that every decoded input addresses the complete row loop. +fn ensure_decoded_lengths( + columns: &Args::Columns, + views: Option<&Args::Views<'_>>, + row_count: usize, +) -> VortexResult<()> { + let lengths_match = match views { + Some(views) => Args::view_lens_match(views, row_count), + None => Args::decoded_lens_match(columns, row_count), + }; + vortex_ensure!( + lengths_match, + "a decoded row input does not address exactly {row_count} rows", + ); + + Ok(()) +} + +/// 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<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + let row_count = args.row_count(); + let mut sink = >::with_capacity(row_count, sink_dtype)?; + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + + { + // 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 = >::rows(&mut sink); + vortex_ensure!( + >::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + // The all-per-row representation removes argument-shape dispatch from the hot loop. The + // mixed path instead reads collapsed batch constants at row zero. + if let Some(views) = views { + for index in 0..row_count { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before + // the loop. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + // SAFETY: `row_count_matches` proved the sink addresses every loop index. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + apply(&prepared, elements, output).into_result()?; + } + } else { + for index in 0..row_count { + // SAFETY: `row_count_matches` proved the sink addresses every loop index. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + apply(&prepared, Args::get(&columns, index), output).into_result()?; + } + } + } + + finish_sink::(sink) +} + +/// Run a prepared sink over only the rows set in `valid`, or decline when the sink cannot skip. +pub fn execute_sink_valid_rows( + args: &dyn ExecutionArgs, + sink_dtype: &DType, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>, >::Row<'_>) -> ApplyResult, +) -> VortexResult> +where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, +{ + // Decline before input decoding or sink allocation when this sink cannot initialize rows that + // the mask skips. The capability and the operation are the same function pointer. + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering the inputs first. An + // element representation may decline when it cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + let prepared = prepare(Args::constants(&columns)); + let row_count = args.row_count(); + let mut sink = >::with_capacity(row_count, sink_dtype)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid) = valid.bit_buffer() else { + vortex_bail!("execute_sink_valid_rows requires a mixed mask"); + }; + vortex_ensure!( + valid.len() == row_count, + "the validity mask does not address exactly {row_count} rows", + ); + + { + let mut rows = >::rows(&mut sink); + vortex_ensure!( + >::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows", + ); + + let views = Args::per_row_views(&columns); + ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; + + // The loop writes only valid indices, but the sink still finishes a full-length output. + // Initialize placeholders now; batch execution masks them before the result escapes. + initialize_skipped_rows(&mut rows); + + // Mask traversal is callback-based and cannot return a `VortexResult`. Record the first + // immediate error, turn later callbacks into no-ops, and return before finishing the sink. + let mut error = None; + valid.for_each_set_index(|index| { + if error.is_some() { + return; + } + + // SAFETY: `row_count_matches` proved that the sink addresses every mask index, which + // is below the mask's validated `row_count`. + let output = unsafe { >::row_unchecked(&mut rows, index) }; + let result = match &views { + Some(views) => { + // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows, and + // mask indices are below `row_count`. + let elements = unsafe { Args::get_from_views_unchecked(views, index) }; + apply(&prepared, elements, output) + } + None => apply(&prepared, Args::get(&columns, index), output), + }; + if let Err(err) = result.into_result() { + error = Some(err); + } + }); + + if let Some(error) = error { + return Err(error); + } + } + + finish_sink::(sink).map(Some) +} + +fn finish_sink(sink: S) -> VortexResult +where + S: OutputSink, +{ + // SAFETY: callers reach this helper only after every completed callback returned the sink's + // write token. Skipped-row traversal also ran the sink's initializer before visiting its mask. + // The sink contract defines how that evidence establishes initialization of its row storage. + unsafe { >::finish(sink) }.map(RowExecution::Output) +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_mask::Mask; + + use super::execute_sink_valid_rows; + use crate::ArrayRef; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::PrimitiveArray; + use crate::dtype::DType; + use crate::dtype::NativePType; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::OutputSink; + use crate::validity::Validity; + + struct NonSkippingSink; + + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or + // `finish` through the executor. The row-initialization requirements are therefore vacuous. + unsafe impl OutputSink for NonSkippingSink { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + Err(vortex_err!( + "a non-skipping sink must decline before allocation" + )) + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { + true + } + + unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { + } + + unsafe fn finish(self) -> VortexResult { + Err(vortex_err!("a non-skipping sink must not finish")) + } + } + + #[test] + fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([true, false]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + assert!(execution.is_none()); + Ok(()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index bcb3a008488..e62f16639bc 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -15,6 +15,11 @@ //! reduce compact failure evidence in that loop and retry only valid rows when null payloads may //! have caused the failure. +mod execute; +pub use execute::RowExecution; + +mod batch; + mod row_fn; pub use row_fn::RowFn; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs index 51d66594332..7a8a8e92e41 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs @@ -20,3 +20,4 @@ mod primitive; mod tuple; pub use tuple::ElementTuple; pub use tuple::IndexedElementTuple; +pub use tuple::batch_constant; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs index a2c143704a0..69b5cf686f6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs @@ -8,6 +8,7 @@ mod element_tuple; pub use element_tuple::ElementTuple; +pub use element_tuple::batch_constant; mod indexed; pub use indexed::IndexedElementTuple; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs index 044ad15fd4b..560e1e48f4f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/tests.rs @@ -10,7 +10,7 @@ use vortex_error::vortex_bail; use vortex_mask::Mask; use super::ElementTuple; -use super::element_tuple::batch_constant; +use super::batch_constant; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs index ce119f32915..e47f195410c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -12,6 +12,7 @@ pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; +pub(super) use element::batch_constant; mod result; pub use result::SinkResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs new file mode 100644 index 00000000000..f5a93ed65d9 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -0,0 +1,263 @@ +// 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 nullable policy derived during 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_error::vortex_ensure_eq; +use vortex_mask::Mask; + +use super::RowPolicy; +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::row_visitor::private; +use crate::ExecutionCtx; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::unstable::row::ElementTuple; +use crate::scalar_fn::unstable::row::IndexedElementTuple; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::SinkResult; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::execute::execute_owned; +use crate::scalar_fn::unstable::row::execute::execute_owned_infallible; +use crate::scalar_fn::unstable::row::execute::execute_sink; +use crate::scalar_fn::unstable::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 nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// 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, + policy: RowPolicy, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + policy, + 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::() }; + ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + + 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<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_sink::())?; + + 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::() }; + ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + + 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 nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// 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, + policy: RowPolicy, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + output_dtype, + policy, + 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::() }; + ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + + // 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<'_>, + >::Row<'_>, + ) -> ApplyResult, + ) -> VortexResult + where + Args: ElementTuple, + Sink: OutputSink, + ApplyResult: SinkResult>::WriteToken>, + { + const { assert_sink_visit_contract::() }; + ensure_policy(self.policy, RowPolicy::for_sink::())?; + + 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::() }; + ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + + // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. + Ok(None) + } +} + +/// Validate that execution selected the nullable policy used to build the batch plan. +fn ensure_policy(planned: RowPolicy, actual: RowPolicy) -> VortexResult<()> { + vortex_ensure_eq!( + actual, + planned, + "row dispatch must select the planned nullable execution policy: planned {planned:?}, got {actual:?}", + ); + + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs index 57da5f4691b..c7f9baf6a62 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -6,9 +6,16 @@ //! [`RowFn::dispatch`]: crate::scalar_fn::unstable::row::RowFn::dispatch 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::BatchPlan; pub(super) use plan::BatchPlanner; +pub(super) use plan::RowPolicy; mod row_visitor; pub use row_visitor::RowVisitor; diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs index c522376d491..214d81c29f4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -114,15 +114,12 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. - // TODO(connor)[RowFn]: Remove this allowance when the execution backend from #9130 consumes - // this policy. - #[allow(dead_code)] pub(crate) policy: RowPolicy, } impl BatchPlan { /// Return the output dtype widened with strict input nullability. - pub(crate) fn result_dtype(self, args: &[DType]) -> DType { + pub(crate) fn result_dtype(&self, args: &[DType]) -> DType { let nullability = self.output_dtype.nullability() | Nullability::from(args.iter().any(DType::is_nullable)); diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index b6c14585c48..55e2a197c10 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -4,12 +4,13 @@ //! Adapts [`RowFn`] implementations to the scalar-function interface. //! //! The blanket [`ScalarFnVTable`] implementation supplies common arity, validity, fallibility, and -//! execution behavior. [`row_fn_return_dtype`] and [`execute_rows`] expose the same planning and -//! execution paths to public vtables that delegate to a private row kernel. +//! execution behavior. The visitor layer validates and executes the concrete signature selected by +//! dispatch. [`row_fn_return_dtype`] and [`execute_rows`] expose the same paths to public vtables +//! that delegate to a private row kernel. use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; use vortex_session::VortexSession; use super::row_fn::RowFn; @@ -24,6 +25,12 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::unstable::row::batch::Batch; +use crate::scalar_fn::unstable::row::batch::BorrowedExecutionArgs; +use crate::scalar_fn::unstable::row::batch::finalize_kernel_output; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::visitor::ExecuteRows; +use crate::scalar_fn::unstable::row::visitor::ExecuteValidRows; impl ScalarFnVTable for F { type Options = F::Options; @@ -98,16 +105,36 @@ pub fn row_fn_return_dtype( /// delegate row execution to a private `RowFn` kernel through this function. pub fn execute_rows( function: &F, - _options: &F::Options, + options: &F::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { ensure_arity(function, args.num_inputs())?; - // TODO(connor)[RowFn]: Replace this temporary error with the execution backend in #9129. - vortex_bail!( - "Row function {} does not yet have an execution backend", - RowFn::id(function) + // Nullary functions have no input validity to propagate, so they skip batch execution. + if args.num_inputs() == 0 { + let plan = function.dispatch(options, &[], BatchPlanner::::new(&[], options))?; + let result_dtype = plan.result_dtype(&[]); + let nullary_args = + BorrowedExecutionArgs::new(&[], args.row_count(), &[], &plan.output_dtype, plan.policy); + + let execution = execute_row_kernel(function, options, nullary_args, ctx)?; + let values = VortexResult::from(execution)?; + + return finalize_kernel_output( + RowFn::id(function), + &result_dtype, + args.row_count(), + values, + ctx, + ); + } + + let batch = prepare_batch(function, options, args)?; + batch.execute( + |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, valid, ctx| try_execute_rows_unfiltered(function, options, args, valid, ctx), + ctx, ) } @@ -124,26 +151,82 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +/// Execute the row loop selected by dispatch. +fn execute_row_kernel( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + function.dispatch( + options, + args.dtypes(), + ExecuteRows::::new(&args, args.output_dtype(), args.policy(), 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: BorrowedExecutionArgs<'_>, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + function.dispatch( + options, + args.dtypes(), + ExecuteValidRows::::new(&args, args.output_dtype(), args.policy(), valid, ctx), + ) +} + +/// Prepare the batch inputs and execution plan for `function`. +fn prepare_batch( + function: &F, + options: &F::Options, + args: &dyn ExecutionArgs, +) -> VortexResult { + Batch::new(RowFn::id(function), args, |arg_dtypes| { + function.dispatch( + options, + arg_dtypes, + BatchPlanner::::new(arg_dtypes, options), + ) + }) +} + #[cfg(test)] mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_session::registry::CachedId; use super::execute_rows; use super::row_fn_return_dtype; + use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::PrimitiveArray; use crate::dtype::DType; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::RowVisitor; + use crate::validity::Validity; #[derive(Clone)] struct IndexingRowFn; + #[derive(Clone)] + struct ChangingDispatchRowFn { + dispatches: Arc, + } + impl RowFn for IndexingRowFn { type Options = EmptyOptions; @@ -168,6 +251,31 @@ mod tests { } } + impl RowFn for ChangingDispatchRowFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.changing_dispatch_row_fn"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { + visitor.visit::<(i64,), i64>(|(value,)| value) + } else { + visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) + } + } + } + #[test] fn test_return_dtype_rejects_wrong_arity_before_dispatch() { let error = row_fn_return_dtype(&IndexingRowFn, &EmptyOptions, &[]) @@ -186,6 +294,29 @@ mod tests { assert_arity_error(error); } + #[test] + fn test_execute_rejects_dispatch_that_changes_after_planning() { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = execute_rows(&function, &EmptyOptions, &args, &mut ctx) + .expect_err("dispatch must not change after planning"); + let message = error.to_string(); + + assert!( + message.contains("row dispatch must select the planned nullable execution policy"), + "unexpected error: {error}", + ); + assert!( + message.contains("planned Dense, got DenseWithRetry"), + "unexpected error: {error}", + ); + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From edab9a7f5730a57e1deb23e59d0136e894c4e3e8 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 17:01:48 -0400 Subject: [PATCH 2/9] Clarify RowFn kernel invocation docs Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/unstable/row/execute/owned.rs | 4 ++-- vortex-array/src/scalar_fn/unstable/row/execute/sink.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index 3ac0c553891..57a0b3bf67a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -24,7 +24,7 @@ impl BitOrAssign for NoFailure { fn bitor_assign(&mut self, _rhs: Self) {} } -/// Decode every input column once, then store one infallible owned output per row. +/// Decode every input column for one kernel invocation, then store one infallible output per row. pub fn execute_owned_infallible( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, @@ -44,7 +44,7 @@ where ) } -/// Decode every input column once, then store owned row outputs and reduce deferred failures. +/// Decode every input column for one kernel invocation, then store outputs and reduce failures. pub fn execute_owned( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index fce891cf613..606a6fa262f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -35,7 +35,7 @@ fn ensure_decoded_lengths( Ok(()) } -/// Decode every input column once, allocate the sink once, then write one row at a time. +/// Decode every input column and allocate one sink for one kernel invocation. /// /// 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. From 893a1cfd949de1d690174c51eb9fbb2ce4ac3e30 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 19:23:37 -0400 Subject: [PATCH 3/9] Harden RowFn execution edge cases Signed-off-by: Connor Tsui --- vortex-array/src/arrays/masked/tests.rs | 31 +++++ vortex-array/src/arrays/masked/vtable/mod.rs | 5 + .../src/scalar_fn/unstable/row/batch/args.rs | 5 +- .../scalar_fn/unstable/row/batch/execution.rs | 6 +- .../src/scalar_fn/unstable/row/batch/tests.rs | 23 +--- .../scalar_fn/unstable/row/execute/owned.rs | 6 +- .../scalar_fn/unstable/row/execute/sink.rs | 129 ++++++++++++++++-- .../scalar_fn/unstable/row/visitor/execute.rs | 110 +++++++++++---- .../src/scalar_fn/unstable/row/vtable.rs | 59 +++++++- 9 files changed, 311 insertions(+), 63 deletions(-) diff --git a/vortex-array/src/arrays/masked/tests.rs b/vortex-array/src/arrays/masked/tests.rs index 92ec4eb474f..55af652725b 100644 --- a/vortex-array/src/arrays/masked/tests.rs +++ b/vortex-array/src/arrays/masked/tests.rs @@ -4,17 +4,22 @@ use rstest::rstest; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use super::*; use crate::Canonical; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::Array; use crate::array_session; +use crate::arrays::ConstantArray; use crate::arrays::ListViewArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; +use crate::dtype::NativePType; use crate::dtype::Nullability; +use crate::scalar::Scalar; use crate::validity::Validity; #[rstest] @@ -55,6 +60,32 @@ fn test_canonical_dtype_matches_array_dtype() -> VortexResult<()> { Ok(()) } +#[test] +fn test_try_from_parts_rejects_null_child() -> VortexResult<()> { + let child = PrimitiveArray::from_iter([1_i64]).into_array(); + let masked = MaskedArray::try_new(child, Validity::AllValid)?; + let mut parts = match masked.try_into_parts() { + Ok(parts) => parts, + Err(_) => vortex_bail!("the uniquely owned masked array must expose its parts"), + }; + let dtype = DType::Primitive(i64::PTYPE, Nullability::Nullable); + parts.slots[MaskedSlots::CHILD] = + Some(ConstantArray::new(Scalar::null(dtype), parts.len).into_array()); + + let error = match Array::::try_from_parts(parts) { + Err(error) => error, + Ok(_) => vortex_bail!("rebuilding must reject a child containing nulls"), + }; + + assert!( + error + .to_string() + .contains("MaskedArray children must not have nulls"), + "unexpected error: {error}", + ); + Ok(()) +} + #[test] fn test_masked_child_with_validity() { // When validity has nulls, masked_child should apply inverted mask. diff --git a/vortex-array/src/arrays/masked/vtable/mod.rs b/vortex-array/src/arrays/masked/vtable/mod.rs index c7e32a7ee3c..c0608388135 100644 --- a/vortex-array/src/arrays/masked/vtable/mod.rs +++ b/vortex-array/src/arrays/masked/vtable/mod.rs @@ -73,6 +73,7 @@ impl VTable for Masked { *ID } + #[allow(clippy::disallowed_methods)] fn validate( &self, _data: &MaskedData, @@ -92,6 +93,10 @@ impl VTable for Masked { child.dtype().as_nullable() == *dtype, "MaskedArray dtype does not match child and validity" ); + vortex_ensure!( + child.all_valid(&mut legacy_session().create_execution_ctx())?, + "MaskedArray children must not have nulls", + ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index d5ffe042613..e3759b0dcf1 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! A borrowed execution view passed to one row-kernel invocation. +//! Execution arguments paired with the metadata selected during planning. +//! +//! [`BorrowedExecutionArgs`] can point at original, sliced, or filtered arrays while retaining the +//! dtypes, output dtype, and null policy of the original batch plan. use vortex_error::VortexResult; use vortex_error::vortex_err; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs index 671b8f4ecb0..0f880bf606e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Null propagation, constant folding, and strategy execution for one columnar batch. +//! Applies columnar semantics around one typed row kernel invocation. +//! +//! [`Batch`] owns strict null propagation, constant broadcasting, execution strategy selection, and +//! output validation. The row kernel therefore handles only decoded values and its selected output +//! capability. use smallvec::SmallVec; use vortex_error::VortexResult; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 6919028715e..35d948a8833 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -479,22 +479,13 @@ fn test_sink_dtype_receives_function_options() -> VortexResult<()> { Ok(()) } -#[test] -fn test_nonnullable_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { - let input = PrimitiveArray::from_iter([1_i64, 2]).into_array(); - - assert_invalid_kernel_output(input) -} - -#[test] -fn test_all_valid_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { - let input = PrimitiveArray::new(vec![1_i64, 2], Validity::AllValid).into_array(); - - assert_invalid_kernel_output(input) -} - -#[track_caller] -fn assert_invalid_kernel_output(input: ArrayRef) -> VortexResult<()> { +#[rstest] +#[case::nonnullable(Validity::NonNullable)] +#[case::all_valid(Validity::AllValid)] +fn test_kernel_output_rejects_nulls_at_function_boundary( + #[case] validity: Validity, +) -> VortexResult<()> { + let input = PrimitiveArray::new(vec![1_i64, 2], validity).into_array(); let args = VecExecutionArgs::new(vec![input], 2); let mut ctx = array_session().create_execution_ctx(); let execution = execute_rows(&InvalidKernelOutput, &EmptyOptions, &args, &mut ctx); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index 57a0b3bf67a..11fff70417a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Execution that stores one owned output value per row. +//! Executes row kernels that return one independent owned value per row. +//! +//! [`execute_owned`] decodes inputs once, prepares constant state, writes into spare vector +//! capacity, and reduces compact failure evidence without putting error construction in the hot +//! loop. [`execute_owned_infallible`] removes that failure path for infallible kernels. use std::ops::BitOrAssign; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index 606a6fa262f..c32491abedb 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -1,7 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Execution that writes through an output sink. +//! Executes row kernels that write through an [`OutputSink`]. +//! +//! Dense execution visits every row. Skip-invalid execution can instead initialize omitted output +//! positions and visit only the set bits of a mixed validity mask, falling back when either the +//! input representation or sink lacks that capability. use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -17,7 +21,6 @@ use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::SinkResult; -/// Ensure that every decoded input addresses the complete row loop. fn ensure_decoded_lengths( columns: &Args::Columns, views: Option<&Args::Views<'_>>, @@ -133,17 +136,17 @@ where { let mut rows = >::rows(&mut sink); - vortex_ensure!( - >::row_count_matches(&rows, row_count), - "the output sink does not address exactly {row_count} rows", - ); let views = Args::per_row_views(&columns); ensure_decoded_lengths::(&columns, views.as_ref(), row_count)?; - // The loop writes only valid indices, but the sink still finishes a full-length output. - // Initialize placeholders now; batch execution masks them before the result escapes. + // Initialize every slot before skipping rows. Recheck addressability afterward because the + // initializer mutably borrows the row representation. initialize_skipped_rows(&mut rows); + vortex_ensure!( + >::row_count_matches(&rows, row_count), + "the output sink does not address exactly {row_count} rows after initializing skipped 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. @@ -191,24 +194,31 @@ where #[cfg(test)] mod tests { use vortex_error::VortexResult; + use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_mask::Mask; + use super::RowExecution; use super::execute_sink_valid_rows; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::PrimitiveArray; + use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::NativePType; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::VecExecutionArgs; + use crate::scalar_fn::unstable::row::InitializedElement; use crate::scalar_fn::unstable::row::OutputSink; + use crate::scalar_fn::unstable::row::UninitElementSink; use crate::validity::Validity; struct NonSkippingSink; + struct ShrinkingSink(Vec); + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or // `finish` through the executor. The row-initialization requirements are therefore vacuous. unsafe impl OutputSink for NonSkippingSink { @@ -240,6 +250,45 @@ mod tests { } } + // SAFETY: the initializer deliberately shrinks the row collection to exercise the executor's + // post-initialization length check. If execution incorrectly continues, safe indexing in + // `row_unchecked` panics instead of accessing invalid memory. + unsafe impl OutputSink for ShrinkingSink { + type Rows<'a> = &'a mut Vec; + type Row<'a> = &'a mut i64; + type WriteToken = (); + + fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { + Some(|rows| { + rows.pop(); + }) + } + + fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + Ok(Self(vec![0; rows])) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.0 + } + + fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { + rows.len() == row_count + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + &mut rows[index] + } + + unsafe fn finish(self) -> VortexResult { + Ok(PrimitiveArray::from_iter(self.0).into_array()) + } + } + #[test] fn test_non_skipping_sink_declines_before_allocation() -> VortexResult<()> { let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); @@ -259,4 +308,68 @@ mod tests { assert!(execution.is_none()); Ok(()) } + + #[test] + fn test_skip_invalid_sink_initializes_and_writes_addressed_rows() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20, 30]).into_array(); + let args = VecExecutionArgs::new(vec![input], 3); + let valid = Mask::from_iter([true, false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let execution = execute_sink_valid_rows::< + (i64,), + (), + UninitElementSink, + InitializedElement, + EmptyOptions, + >( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + // SAFETY: `output` is the row supplied to this callback. + unsafe { InitializedElement::write(output, value * 2) } + }, + )?; + let Some(RowExecution::Output(actual)) = execution else { + vortex_bail!("the skip-invalid sink must produce an output"); + }; + let expected = PrimitiveArray::from_iter([20_i64, 0, 60]); + + assert_arrays_eq!(&actual, expected.as_ref(), &mut ctx); + Ok(()) + } + + #[test] + fn test_skip_invalid_sink_rechecks_rows_after_initialization() -> VortexResult<()> { + let input = PrimitiveArray::from_iter([10_i64, 20]).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let valid = Mask::from_iter([false, true]); + let mut ctx = array_session().create_execution_ctx(); + + let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( + &args, + &DType::from(i64::PTYPE), + &valid, + &mut ctx, + |_| (), + |_, (value,), output| { + *output = value; + }, + ); + + let error = match result { + Err(error) => error, + Ok(_) => vortex_bail!("the sink must reject rows changed by its initializer"), + }; + assert!( + error + .to_string() + .contains("after initializing skipped rows"), + "unexpected error: {error}", + ); + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index f5a93ed65d9..010dc079599 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -3,11 +3,10 @@ //! Visitors that execute dense and skip-invalid row loops. //! -//! Each method verifies that execution selected the nullable policy derived during 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. +//! Each visit revalidates its concrete signature and checks that its output dtype and null policy +//! match the plan before entering a row loop. [`ExecuteValidRows`] can decline unsupported +//! skip-invalid execution so the batch layer filters the inputs and retries with [`ExecuteRows`]. -use std::marker::PhantomData; use std::ops::BitOrAssign; use vortex_error::VortexResult; @@ -19,6 +18,8 @@ use super::RowVisitor; use super::check::assert_deferred_visit_contract; use super::check::assert_owned_visit_contract; use super::check::assert_sink_visit_contract; +use super::check::validate_owned_visit; +use super::check::validate_sink_visit; use super::row_visitor::private; use crate::ExecutionCtx; use crate::dtype::DType; @@ -36,10 +37,16 @@ use crate::scalar_fn::unstable::row::execute::execute_sink; use crate::scalar_fn::unstable::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> { +pub struct ExecuteRows<'args, 'ctx, F: RowFn> { /// The inputs for this kernel invocation. args: &'args dyn ExecutionArgs, + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + /// The output dtype computed by the planning visit. output_dtype: &'args DType, @@ -48,29 +55,29 @@ pub struct ExecuteRows<'args, 'ctx, F> { /// 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> { +impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { pub fn new( args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, output_dtype: &'args DType, policy: RowPolicy, ctx: &'ctx mut ExecutionCtx, ) -> Self { Self { args, + dtypes, + options, output_dtype, policy, ctx, - function: PhantomData, } } } -impl private::Sealed for ExecuteRows<'_, '_, F> {} +impl private::Sealed for ExecuteRows<'_, '_, F> {} impl RowVisitor for ExecuteRows<'_, '_, F> { type VisitResult = RowExecution; @@ -85,7 +92,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; execute_owned_infallible::(self.args, self.ctx, prepare, apply) } @@ -105,7 +117,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { ApplyResult: SinkResult>::WriteToken>, { const { assert_sink_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_sink::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; execute_sink::( self.args, @@ -128,7 +145,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { Fail: Copy + Default + BitOrAssign, { const { assert_deferred_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; execute_owned::( self.args, @@ -144,10 +166,16 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { /// /// 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> { +pub struct ExecuteValidRows<'args, 'ctx, F: RowFn> { /// The original inputs for this kernel invocation. args: &'args dyn ExecutionArgs, + /// The input dtypes used by the planning visit. + dtypes: &'args [DType], + + /// The function options used to derive a sink's runtime dtype. + options: &'args F::Options, + /// The output dtype computed by the planning visit. output_dtype: &'args DType, @@ -159,14 +187,13 @@ pub struct ExecuteValidRows<'args, 'ctx, F> { /// 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> { +impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> { pub fn new( args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, output_dtype: &'args DType, policy: RowPolicy, valid: &'args Mask, @@ -174,16 +201,17 @@ impl<'args, 'ctx, F> ExecuteValidRows<'args, 'ctx, F> { ) -> Self { Self { args, + dtypes, + options, output_dtype, policy, valid, ctx, - function: PhantomData, } } } -impl private::Sealed for ExecuteValidRows<'_, '_, F> {} +impl private::Sealed for ExecuteValidRows<'_, '_, F> {} impl RowVisitor for ExecuteValidRows<'_, '_, F> { type VisitResult = Option; @@ -198,7 +226,12 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { Out: OutputElement, { const { assert_owned_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_owned_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_owned_output::(), + )?; // Owned execution has no sink that can initialize skipped output positions. Decline so // batch execution filters the inputs and retries with the dense visitor. @@ -220,7 +253,12 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { ApplyResult: SinkResult>::WriteToken>, { const { assert_sink_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_sink::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; execute_sink_valid_rows::( self.args, @@ -244,19 +282,33 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { Fail: Copy + Default + BitOrAssign, { const { assert_deferred_visit_contract::() }; - ensure_policy(self.policy, RowPolicy::for_deferred_output::())?; + ensure_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; // Deferred owned execution has the same skipped-output limitation as `visit_prepared`. Ok(None) } } -/// Validate that execution selected the nullable policy used to build the batch plan. -fn ensure_policy(planned: RowPolicy, actual: RowPolicy) -> VortexResult<()> { +fn ensure_plan( + planned_output: &DType, + planned_policy: RowPolicy, + actual_output: DType, + actual_policy: RowPolicy, +) -> VortexResult<()> { + vortex_ensure_eq!( + actual_policy, + planned_policy, + "row dispatch must select the planned nullable execution policy: planned {planned_policy:?}, got {actual_policy:?}", + ); vortex_ensure_eq!( - actual, - planned, - "row dispatch must select the planned nullable execution policy: planned {planned:?}, got {actual:?}", + actual_output, + *planned_output, + "row dispatch must select the planned output dtype: planned {planned_output}, got {actual_output}", ); Ok(()) diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 55e2a197c10..9675219d547 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -138,7 +138,6 @@ pub fn execute_rows( ) } -/// Validate the number of arguments before calling user-defined dispatch code. fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { let expected = F::ARG_NAMES.len(); vortex_ensure_eq!( @@ -151,7 +150,6 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } -/// Execute the row loop selected by dispatch. fn execute_row_kernel( function: &F, options: &F::Options, @@ -161,11 +159,17 @@ fn execute_row_kernel( function.dispatch( options, args.dtypes(), - ExecuteRows::::new(&args, args.output_dtype(), args.policy(), ctx), + ExecuteRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + ctx, + ), ) } -/// Try execution against the original inputs, returning `None` when batch execution must filter. fn try_execute_rows_unfiltered( function: &F, options: &F::Options, @@ -176,11 +180,18 @@ fn try_execute_rows_unfiltered( function.dispatch( options, args.dtypes(), - ExecuteValidRows::::new(&args, args.output_dtype(), args.policy(), valid, ctx), + ExecuteValidRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + valid, + ctx, + ), ) } -/// Prepare the batch inputs and execution plan for `function`. fn prepare_batch( function: &F, options: &F::Options, @@ -225,6 +236,13 @@ mod tests { #[derive(Clone)] struct ChangingDispatchRowFn { dispatches: Arc, + change: DispatchChange, + } + + #[derive(Clone, Copy)] + enum DispatchChange { + Policy, + Element, } impl RowFn for IndexingRowFn { @@ -271,7 +289,11 @@ mod tests { if self.dispatches.fetch_add(1, Ordering::Relaxed) == 0 { visitor.visit::<(i64,), i64>(|(value,)| value) } else { - visitor.visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())) + match self.change { + DispatchChange::Policy => visitor + .visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())), + DispatchChange::Element => visitor.visit::<(u64,), u64>(|(value,)| value), + } } } } @@ -298,6 +320,7 @@ mod tests { fn test_execute_rejects_dispatch_that_changes_after_planning() { let function = ChangingDispatchRowFn { dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Policy, }; let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); let args = VecExecutionArgs::new(vec![input], 2); @@ -317,6 +340,28 @@ mod tests { ); } + #[test] + fn test_execute_revalidates_element_types_after_planning() -> VortexResult<()> { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Element, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&function, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_error::vortex_bail!("dispatch must preserve its planned element types"), + }; + + assert!( + error.to_string().contains("expected a u64 column"), + "unexpected error: {error}", + ); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From 76c5988809de645c9896ba71b946871a616575f2 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 12 Aug 2026 19:25:28 -0400 Subject: [PATCH 4/9] Tighten RowFn executor documentation Signed-off-by: Connor Tsui --- .../scalar_fn/unstable/row/batch/execution.rs | 24 +++++-------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs index 0f880bf606e..93aad88810e 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execution.rs @@ -121,11 +121,10 @@ impl Batch { }) } - /// Add null propagation, constant folding, and strategy selection around `kernel`. + /// Apply constant folding and null handling 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 - /// originals plus a mixed validity mask; `Ok(None)` selects filter-and-scatter. + /// For a mixed validity mask, `try_unfiltered` may avoid filtering; `Ok(None)` selects + /// filter-and-scatter. Every kernel result is checked against the planned shape and dtype. pub fn execute( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -167,11 +166,7 @@ impl Batch { } } - /// 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. + /// Evaluate one row of constant inputs and broadcast the validated result. fn broadcast_one_row( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -191,11 +186,7 @@ impl Batch { Ok(ConstantArray::new(scalar, self.row_count).into_array()) } - /// Run the kernel over every row, including the rows behind nulls, then mask its result. - /// - /// The arguments reach the kernel untouched, so the inputs keep their original encoding, and - /// the conjoined validity is handed to `mask` as an array rather than materialized into a - /// [`Mask`] first. + /// Run every stored payload, then attach the input validity without materializing its mask. fn execute_dense( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -314,8 +305,7 @@ impl Batch { .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. + /// Filter to valid rows, run the kernel, then scatter into a null-padded output. fn filter_and_scatter( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -337,7 +327,6 @@ impl Batch { self.finalize_output(self.scatter_valid(values, valid)?, valid.len()) } - /// An all-null result of the function's declared return dtype. fn all_null(&self) -> ArrayRef { ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } @@ -357,7 +346,6 @@ impl Batch { ) } - /// Finalize an output against this batch's expected length and declared return dtype. fn finalize_output(&self, values: ArrayRef, expected_len: usize) -> VortexResult { reconcile_output(self.id, &self.result_dtype, expected_len, values) } From 83355fff89f43bb7e57243b50a499ecf9a2bb750 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 10:18:13 -0400 Subject: [PATCH 5/9] Adapt RowFn framework to explicit contracts Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 15 ++++++++++----- .../src/scalar_fn/unstable/row/execute/sink.rs | 18 ++++++------------ .../scalar_fn/unstable/row/visitor/execute.rs | 13 ++----------- 3 files changed, 18 insertions(+), 28 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 35d948a8833..f801fa35e4c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -82,7 +82,7 @@ unsafe impl OutputSink for OptionsCheckingSink { type Row<'a> = (); type WriteToken = (); - fn sink_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { + fn output_dtype(enabled: &bool, _args: &[DType]) -> VortexResult { if !enabled { vortex_bail!(InvalidArgument: "the test sink is disabled"); } @@ -90,7 +90,7 @@ unsafe impl OutputSink for OptionsCheckingSink { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(_rows: usize) -> VortexResult { vortex_bail!("the planning-only test sink must not be allocated") } @@ -129,11 +129,11 @@ unsafe impl OutputSink for I64Sink { type Row<'a> = &'a mut i64; type WriteToken = (); - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(rows: usize) -> VortexResult { Ok(Self(BufferMut::zeroed(rows))) } @@ -159,6 +159,7 @@ impl RowFn for NullarySeven { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.nullary_seven"); @@ -181,6 +182,7 @@ impl RowFn for AddThree { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["first", "second", "third"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.add_three"); @@ -231,6 +233,7 @@ impl RowFn for Identity { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.identity"); @@ -251,6 +254,7 @@ impl RowFn for SinkOptions { type Options = bool; const ARG_NAMES: &'static [&'static str] = &[]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.sink_options"); @@ -271,6 +275,7 @@ impl RowFn for InvalidKernelOutput { type Options = EmptyOptions; const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = false; fn id(&self) -> ScalarFnId { static ID: CachedId = CachedId::new("test.invalid_kernel_output"); @@ -470,7 +475,7 @@ fn test_finalize_kernel_output_validates_shape_and_dtype() -> VortexResult<()> { } #[test] -fn test_sink_dtype_receives_function_options() -> VortexResult<()> { +fn test_output_dtype_receives_function_options() -> VortexResult<()> { assert_eq!( row_fn_return_dtype(&SinkOptions, &true, &[])?, DType::from(i64::PTYPE) diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index c32491abedb..2c6ebcebd4a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -15,7 +15,6 @@ use vortex_mask::Mask; use super::RowExecution; use crate::ExecutionCtx; -use crate::dtype::DType; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::unstable::row::ElementTuple; use crate::scalar_fn::unstable::row::OutputSink; @@ -44,7 +43,6 @@ fn ensure_decoded_lengths( /// 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<'_>, >::Row<'_>) -> ApplyResult, @@ -55,7 +53,7 @@ where ApplyResult: SinkResult>::WriteToken>, { let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count, sink_dtype)?; + let mut sink = >::with_capacity(row_count)?; let columns = Args::decode(args, ctx)?; let prepared = prepare(Args::constants(&columns)); let views = Args::per_row_views(&columns); @@ -98,7 +96,6 @@ where /// 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, @@ -123,7 +120,7 @@ where }; let prepared = prepare(Args::constants(&columns)); let row_count = args.row_count(); - let mut sink = >::with_capacity(row_count, sink_dtype)?; + let mut sink = >::with_capacity(row_count)?; // Batch execution resolves all-valid and all-null inputs before selecting this path. let AllOr::Some(valid) = valid.bit_buffer() else { @@ -226,11 +223,11 @@ mod tests { type Row<'a> = (); type WriteToken = (); - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(_rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(_rows: usize) -> VortexResult { Err(vortex_err!( "a non-skipping sink must decline before allocation" )) @@ -264,11 +261,11 @@ mod tests { }) } - fn sink_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { Ok(DType::from(i64::PTYPE)) } - fn with_capacity(rows: usize, _dtype: &DType) -> VortexResult { + fn with_capacity(rows: usize) -> VortexResult { Ok(Self(vec![0; rows])) } @@ -298,7 +295,6 @@ mod tests { let execution = execute_sink_valid_rows::<(i64,), (), NonSkippingSink, (), EmptyOptions>( &args, - &DType::from(i64::PTYPE), &valid, &mut ctx, |_| (), @@ -324,7 +320,6 @@ mod tests { EmptyOptions, >( &args, - &DType::from(i64::PTYPE), &valid, &mut ctx, |_| (), @@ -351,7 +346,6 @@ mod tests { let result = execute_sink_valid_rows::<(i64,), (), ShrinkingSink, (), EmptyOptions>( &args, - &DType::from(i64::PTYPE), &valid, &mut ctx, |_| (), diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index 010dc079599..24d2cd1cf6c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -125,11 +125,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { )?; execute_sink::( - self.args, - self.output_dtype, - self.ctx, - prepare, - apply, + self.args, self.ctx, prepare, apply, ) } @@ -261,12 +257,7 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { )?; execute_sink_valid_rows::( - self.args, - self.output_dtype, - self.valid, - self.ctx, - prepare, - apply, + self.args, self.valid, self.ctx, prepare, apply, ) } From 05fa94b3ebf5307ad699c4f383769858ef506feb Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:23:09 -0400 Subject: [PATCH 6/9] Use sink row counts directly Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 8 ++++---- .../src/scalar_fn/unstable/row/execute/sink.rs | 18 +++++++++--------- 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index f801fa35e4c..b3ce797bc79 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -96,8 +96,8 @@ unsafe impl OutputSink for OptionsCheckingSink { fn rows(&mut self) -> Self::Rows<'_> {} - fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { - true + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 } unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> {} @@ -141,8 +141,8 @@ unsafe impl OutputSink for I64Sink { self.0.as_mut_slice() } - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { - rows.len() == row_count + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs index 2c6ebcebd4a..45b65520d3a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -64,7 +64,7 @@ where // scope releases the borrow before `finish_sink` consumes the sink. let mut rows = >::rows(&mut sink); vortex_ensure!( - >::row_count_matches(&rows, row_count), + >::row_count(&rows) == row_count, "the output sink does not address exactly {row_count} rows", ); @@ -75,14 +75,14 @@ where // SAFETY: `ensure_decoded_lengths` proved every view has `row_count` rows before // the loop. let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; - // SAFETY: `row_count_matches` proved the sink addresses every loop index. + // SAFETY: `row_count` proved the sink addresses every loop index. let output = unsafe { >::row_unchecked(&mut rows, index) }; apply(&prepared, elements, output).into_result()?; } } else { for index in 0..row_count { - // SAFETY: `row_count_matches` proved the sink addresses every loop index. + // SAFETY: `row_count` proved the sink addresses every loop index. let output = unsafe { >::row_unchecked(&mut rows, index) }; apply(&prepared, Args::get(&columns, index), output).into_result()?; @@ -141,7 +141,7 @@ where // initializer mutably borrows the row representation. initialize_skipped_rows(&mut rows); vortex_ensure!( - >::row_count_matches(&rows, row_count), + >::row_count(&rows) == row_count, "the output sink does not address exactly {row_count} rows after initializing skipped rows", ); @@ -153,7 +153,7 @@ where return; } - // SAFETY: `row_count_matches` proved that the sink addresses every mask index, which + // SAFETY: `row_count` proved that the sink addresses every mask index, which // is below the mask's validated `row_count`. let output = unsafe { >::row_unchecked(&mut rows, index) }; let result = match &views { @@ -235,8 +235,8 @@ mod tests { fn rows(&mut self) -> Self::Rows<'_> {} - fn row_count_matches(_rows: &Self::Rows<'_>, _row_count: usize) -> bool { - true + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 } unsafe fn row_unchecked<'a>(_rows: &'a mut Self::Rows<'_>, _index: usize) -> Self::Row<'a> { @@ -273,8 +273,8 @@ mod tests { &mut self.0 } - fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool { - rows.len() == row_count + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() } unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { From d5499fb8725fb1174e3cc59a954ae269bdf1834b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:23:14 -0400 Subject: [PATCH 7/9] Simplify RowFn internal visibility Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index e3759b0dcf1..d1f5b67fbd0 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -20,7 +20,7 @@ use crate::scalar_fn::unstable::row::visitor::RowPolicy; /// original planned batch. Keeping them together prevents an execution path from pairing an input /// view with unrelated planning metadata. #[derive(Clone, Copy)] -pub(in crate::scalar_fn::unstable::row) struct BorrowedExecutionArgs<'a> { +pub(crate) struct BorrowedExecutionArgs<'a> { /// The input arrays for this kernel invocation. arrays: &'a [ArrayRef], @@ -39,7 +39,7 @@ pub(in crate::scalar_fn::unstable::row) struct BorrowedExecutionArgs<'a> { impl<'a> BorrowedExecutionArgs<'a> { /// Pair one input view with the planning metadata selected for its batch. - pub(in crate::scalar_fn::unstable::row) fn new( + pub(crate) fn new( arrays: &'a [ArrayRef], row_count: usize, dtypes: &'a [DType], @@ -56,22 +56,22 @@ impl<'a> BorrowedExecutionArgs<'a> { } /// Return the concrete arrays used by this row-kernel invocation. - pub(in crate::scalar_fn::unstable::row) fn arrays(&self) -> &'a [ArrayRef] { + pub(crate) fn arrays(&self) -> &'a [ArrayRef] { self.arrays } /// Return the original input dtypes used to select the row implementation. - pub(in crate::scalar_fn::unstable::row) fn dtypes(&self) -> &'a [DType] { + pub(crate) fn dtypes(&self) -> &'a [DType] { self.dtypes } /// Return the non-nullable dtype built by the selected output capability. - pub(in crate::scalar_fn::unstable::row) fn output_dtype(&self) -> &'a DType { + pub(crate) fn output_dtype(&self) -> &'a DType { self.output_dtype } /// Return the nullable execution policy selected during planning. - pub(in crate::scalar_fn::unstable::row) fn policy(&self) -> RowPolicy { + pub(crate) fn policy(&self) -> RowPolicy { self.policy } } From 11ce2c424e7dc020086c91a5efb3a3f667a5c1dc Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 11:24:19 -0400 Subject: [PATCH 8/9] Restore decoded-length regression coverage Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index b3ce797bc79..5e01b3da94b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -6,6 +6,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use rstest::rstest; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use vortex_error::vortex_bail; @@ -18,6 +19,7 @@ use super::BatchPlan; use super::RowPolicy; use super::finalize_kernel_output; use crate::ArrayRef; +use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -31,6 +33,7 @@ use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::InputElement; use crate::scalar_fn::unstable::row::OutputElement; use crate::scalar_fn::unstable::row::OutputSink; use crate::scalar_fn::unstable::row::RowFn; @@ -48,6 +51,12 @@ struct NullarySeven; #[derive(Clone)] struct AddThree; +#[derive(Clone)] +struct AddShort; + +/// An element whose decode drops the last row, standing in for an invalid element implementation. +struct ShortDecodeI64; + #[derive(Clone)] struct Identity; @@ -75,6 +84,54 @@ enum PreparedVisit { Deferred, } +// SAFETY: the view and unchecked access delegate to the `i64` implementation. The implementation +// deliberately returns a short column so the executor's pre-loop length guard can be tested. +unsafe impl InputElement for ShortDecodeI64 { + type Column = Buffer; + type View<'a> = &'a [i64]; + type Elem<'a> = i64; + + const DENSE_SAFE: bool = true; + const DECODE_FALLIBLE: bool = false; + + fn validate(dtype: &DType) -> VortexResult<()> { + ::validate(dtype) + } + + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let column = ::decode(array, ctx)?; + + Ok(column.slice(0..column.len().saturating_sub(1))) + } + + fn get(column: &Self::Column, index: usize) -> i64 { + ::get(column, index) + } + + fn view(column: &Self::Column) -> Self::View<'_> { + ::view(column) + } + + fn view_len(view: &Self::View<'_>) -> usize { + ::view_len(view) + } + + fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + ::get_from_view(view, index) + } + + unsafe fn get_from_view_unchecked<'a>(view: &Self::View<'a>, index: usize) -> i64 + where + Self: 'a, + { + // SAFETY: forwarded from this method's contract. + unsafe { ::get_from_view_unchecked(view, index) } + } +} + // SAFETY: `with_capacity` always returns an error, so no sink value can reach `rows`, `row`, or // `finish` through the executor. The row-initialization requirements are therefore vacuous. unsafe impl OutputSink for OptionsCheckingSink { @@ -199,6 +256,29 @@ impl RowFn for AddThree { } } +impl RowFn for AddShort { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const FALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.add_short"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(ShortDecodeI64, i64), I64Sink, _>(|(lhs, rhs), output| { + *output = lhs + rhs; + }) + } +} + impl RowFn for RetryConstantAdd { type Options = EmptyOptions; @@ -358,6 +438,27 @@ fn test_batch_rejects_input_length_mismatch() -> VortexResult<()> { Ok(()) } +#[test] +fn test_short_decode_beside_constant_is_rejected() -> VortexResult<()> { + let lhs = PrimitiveArray::from_iter(0..64_i64).into_array(); + let rhs = ConstantArray::new(10_i64, 64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 64); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short decoded column passed the pre-loop length check"), + }; + + assert!( + error + .to_string() + .contains("does not address exactly 64 rows"), + "unexpected error: {error}", + ); + Ok(()) +} + #[test] fn test_dense_retry_preserves_valid_row_failure() -> VortexResult<()> { let lhs = From 219e00aa359fb58c72f4e4523fb50f888684aa6c Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 13 Aug 2026 12:09:53 -0400 Subject: [PATCH 9/9] Remove unused RowFn batch array access Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/unstable/row/batch/args.rs | 5 ----- vortex-array/src/scalar_fn/unstable/row/batch/tests.rs | 5 +++-- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index d1f5b67fbd0..781d15711be 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -55,11 +55,6 @@ impl<'a> BorrowedExecutionArgs<'a> { } } - /// Return the concrete arrays used by this row-kernel invocation. - pub(crate) fn arrays(&self) -> &'a [ArrayRef] { - self.arrays - } - /// Return the original input dtypes used to select the row implementation. pub(crate) fn dtypes(&self) -> &'a [DType] { self.dtypes diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 5e01b3da94b..6c92fd52c40 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -31,6 +31,7 @@ use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::VecExecutionArgs; use crate::scalar_fn::unstable::row::InputElement; @@ -518,7 +519,7 @@ fn test_resolve_validity_array_masks(#[case] validity: [bool; 2]) -> VortexResul let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), |_args, _valid, _ctx| Ok(None), &mut ctx, )?; @@ -546,7 +547,7 @@ fn test_valid_only_filters_and_scatters() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let actual = batch.execute( - |args, _ctx| Ok(RowExecution::Output(args.arrays()[0].clone())), + |args, _ctx| Ok(RowExecution::Output(args.get(0)?)), |_args, _valid, _ctx| Ok(None), &mut ctx, )?;