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..781d15711be --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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; + +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(crate) 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(crate) 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 original input dtypes used to select the row implementation. + pub(crate) fn dtypes(&self) -> &'a [DType] { + self.dtypes + } + + /// Return the non-nullable dtype built by the selected output capability. + pub(crate) fn output_dtype(&self) -> &'a DType { + self.output_dtype + } + + /// Return the nullable execution policy selected during planning. + pub(crate) 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/execute/constant.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs new file mode 100644 index 00000000000..8cb89a14bf3 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::scalar_fn::unstable::row::execute::RowExecution; + +impl Batch { + /// Evaluate one row of constant inputs and broadcast the validated result. + pub(super) 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()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs new file mode 100644 index 00000000000..d3e0ef103c2 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +impl Batch { + /// Run every stored payload, then attach the input validity without materializing its mask. + pub(super) 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 in `Batch::execute`, before the kernel ran. + Validity::AllInvalid => Ok(self.all_null()), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs new file mode 100644 index 00000000000..871c7fc1219 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/filter_scatter.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure_eq; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::arrays::MaskedArray; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::Nullability; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +impl Batch { + /// Filter to valid rows, run the kernel, then scatter into a null-padded output. + pub(super) 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()) + } + + /// 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 output must contain {} filtered rows, got {}", + self.id, + valid.true_count(), + values.len(), + ); + + let AllOr::Some(slices) = valid.slices() else { + // The caller handles the all-true and all-false masks. + vortex_bail!( + "scatter_valid requires valid and invalid rows, got an all-valid or all-invalid 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()) + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs new file mode 100644 index 00000000000..f29fbcaad27 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -0,0 +1,76 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Selects a batch execution strategy. +//! +//! [`Batch::execute`] handles universal fast paths, then delegates to dense or valid-only +//! execution. + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::Batch; +use super::RowPolicy; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::arrays::Constant; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::scalar_fn::unstable::row::types::batch_constant; +use crate::validity::Validity; + +mod constant; +mod dense; +mod filter_scatter; +mod valid_only; + +mod output; +pub(crate) use output::finalize_kernel_output; + +impl Batch { + /// Apply constant folding and null handling around `kernel`. + /// + /// When the mask contains valid and invalid rows, `try_unfiltered` may avoid filtering. + /// `Ok(None)` filters the valid rows and scatters the output back. Every result is checked + /// against the planned shape and dtype. + pub(crate) 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), + } + } +} diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs new file mode 100644 index 00000000000..47fd8cdaf46 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; + +use super::super::Batch; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnId; + +impl Batch { + pub(super) fn all_null(&self) -> ArrayRef { + ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() + } + + pub(super) 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. + pub(super) 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) + } +} + +/// 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(crate) 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 must produce only valid rows, got at least one null row", + ); + + 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 output must contain {expected_len} rows, got {}", + values.len(), + ); + vortex_ensure!( + values.dtype().eq_ignore_nullability(result_dtype), + "the {id} kernel output dtype must match {result_dtype} ignoring nullability, got {}", + 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/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs new file mode 100644 index 00000000000..d14fed87370 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::super::Batch; +use super::super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::BoolArray; +use crate::builtins::ArrayBuiltins; +use crate::scalar_fn::unstable::row::execute::RowExecution; +use crate::validity::Validity; + +/// The result of resolving batch validity. +enum ResolvedValidity { + /// The output for an all-valid or all-null batch. + Output(ArrayRef), + + /// A mask with both valid and invalid rows. + PartiallyValid(Mask), +} + +impl Batch { + /// Resolve validity, try unfiltered execution, then fall back to filtering. + pub(super) 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)? { + ResolvedValidity::Output(output) => return Ok(output), + ResolvedValidity::PartiallyValid(valid) => valid, + }; + + if let Some(result) = self.try_execute_unfiltered(try_unfiltered, &valid, ctx)? { + return Ok(result); + } + + self.filter_and_scatter(kernel, &valid, ctx) + } + + /// Materialize validity and handle all-valid or all-null batches. + 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(ResolvedValidity::Output(values)); + } + + if valid.all_false() { + return Ok(ResolvedValidity::Output(self.all_null())); + } + + Ok(ResolvedValidity::PartiallyValid(valid)) + } + + /// 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) + } +} 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..a246ba3e30b --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -0,0 +1,63 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution around a strict 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, 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. + +use smallvec::SmallVec; + +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + +mod args; +pub(super) use args::BorrowedExecutionArgs; + +mod execute; +pub(super) use execute::finalize_kernel_output; + +mod planning; + +pub(super) use super::visitor::BatchPlan; +pub(super) use super::visitor::RowPolicy; + +/// One batch of inputs and the metadata needed before its row kernel runs. +pub(crate) 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, +} + +#[cfg(test)] +mod tests; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs new file mode 100644 index 00000000000..24bb31b2709 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure_eq; + +use super::Batch; +use super::BatchPlan; +use super::args::BorrowedExecutionArgs; +use crate::ArrayRef; +use crate::dtype::DType; +use crate::scalar_fn::ExecutionArgs; +use crate::scalar_fn::ScalarFnId; +use crate::validity::Validity; + +impl Batch { + /// Collect the inputs and derive their dtypes, 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 would vacuously pass. + pub(crate) 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, + }) + } + + /// Pair an input view with this batch's planning metadata. + pub(super) 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, + ) + } +} 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..7ad4b69c4e7 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -0,0 +1,793 @@ +// 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::Buffer; +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::ExecutionCtx; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::MaskedArray; +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::ExecutionArgs; +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; +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 AddShort(ShortVisit); + +/// An element whose decode drops the last row, standing in for an invalid element implementation. +struct ShortDecodeI64; + +#[derive(Clone)] +struct Identity; + +#[derive(Clone)] +struct ValidOnlyIdentity; + +#[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, +} + +#[derive(Clone, Copy)] +enum ShortVisit { + Owned, + Sink, +} + +// 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 = false; + 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 { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn output_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) -> VortexResult { + vortex_bail!("the planning-only test sink must not be allocated") + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 + } + + 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 output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(BufferMut::zeroed(rows))) + } + + fn rows(&mut self) -> Self::Rows<'_> { + self.0.as_mut_slice() + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a> { + // SAFETY: required by this method's contract. + unsafe { rows.get_unchecked_mut(index) } + } + + unsafe fn finish(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] = &[]; + const FALLIBLE: bool = false; + + 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"]; + const FALLIBLE: bool = false; + + 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 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 { + match self.0 { + ShortVisit::Owned => { + visitor.visit::<(ShortDecodeI64, i64), i64>(|(lhs, rhs)| lhs + rhs) + } + ShortVisit::Sink => { + visitor.visit_into::<(ShortDecodeI64, i64), I64Sink, _>(|(lhs, rhs), output| { + *output = lhs + rhs; + }) + } + } + } +} + +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"]; + const FALLIBLE: bool = false; + + 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 ValidOnlyIdentity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const FALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.valid_only_identity"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_into::<(i64,), I64Sink, VortexResult<()>>(|(value,), output| { + *output = value; + Ok(()) + }) + } +} + +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"); + *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"]; + const FALLIBLE: bool = false; + + 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![1_i64, 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_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(ShortVisit::Sink), &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(()) +} + +#[rstest] +#[case::owned(ShortVisit::Owned)] +#[case::sink(ShortVisit::Sink)] +fn test_short_constant_decode_is_rejected(#[case] visit: ShortVisit) -> VortexResult<()> { + let lhs = ConstantArray::new(10_i64, 64).into_array(); + let rhs = PrimitiveArray::from_iter(0..64_i64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 64); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(visit), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short constant decode passed the pre-loop length check"), + }; + + assert!( + error + .to_string() + .contains("batch-constant input must contain exactly 1 row, got 0"), + "unexpected error: {error}", + ); + Ok(()) +} + +#[test] +fn test_short_constant_null_tolerant_decode_is_rejected() -> VortexResult<()> { + let lhs = MaskedArray::try_new( + ConstantArray::new(10_i64, 4).into_array(), + Validity::from_iter([true, false, true, false]), + )? + .into_array(); + let rhs = PrimitiveArray::from_iter(0..4_i64).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 4); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&AddShort(ShortVisit::Sink), &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("a short null-tolerant constant decode passed validation"), + }; + + assert!( + error + .to_string() + .contains("decoded batch-constant input must contain exactly 1 row, got 0"), + "unexpected error: {error}", + ); + 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 error = match execute_rows(&RetryConstantAdd, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_bail!("valid-row overflow must remain observable"), + }; + + assert!( + error.to_string().contains("checked add overflowed"), + "unexpected error: {error}", + ); + 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.get(0)?)), + |_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.get(0)?)), + |_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_output_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(()) +} + +#[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); + 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 must produce only 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([7_i64, 7, 7]).into_array(); + + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) +} + +#[test] +fn test_valid_only_empty_batch_preserves_nonnullable_dtype() -> VortexResult<()> { + let input = PrimitiveArray::from_iter(std::iter::empty::()).into_array(); + let args = VecExecutionArgs::new(vec![input], 0); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&ValidOnlyIdentity, &EmptyOptions, &args, &mut ctx)?; + + assert_eq!(actual.len(), 0); + assert_eq!(actual.dtype(), &DType::from(i64::PTYPE)); + 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..d3809f939be --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -0,0 +1,121 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! 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; + +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 failure accumulator for infallible owned visits. +#[derive(Clone, Copy, Default)] +struct NoFailure; + +impl BitOrAssign for NoFailure { + fn bitor_assign(&mut self, _rhs: Self) {} +} + +/// Decode every input column, then store one output per row from an infallible kernel. +pub(crate) 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, then store outputs and combine per-row failure evidence. +pub(crate) 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, +{ + // The output vector stays at length zero until every slot is initialized so that an unwind + // abandons partially initialized spare capacity. This no-drop assertion proves that no + // initialized value requires a destructor to run. + const { assert_owned_output_needs_no_drop::() }; + + let columns = Args::decode(args, ctx)?; + let prepared = prepare(Args::constants(&columns)); + + let row_count = args.row_count(); + let mut values = Vec::::with_capacity(row_count); + let output = &mut values.spare_capacity_mut()[..row_count]; + + let failure = if let Some(views) = Args::views_no_constants(&columns) { + // Keep this validation beside the views so LLVM sees their common length here. + 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. + let source = unsafe { Args::indexed_source(views, row_count) }; + + source.map_checked_into(output, |elements| apply(&prepared, elements)) + } else { + // Keep this proof branch-local. Shared validation prevents LLVM from specializing this + // loop for each batch-constant arrangement, leaving it scalar under multiple CGUs without + // LTO. The exact pass interaction is unknown. + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + let mut accumulated = Fail::default(); + + // Iterate over `output` directly. A `0..row_count` range reuses the address-taken value + // from the validation error formatter and retains an output bounds check. + for (index, slot) in output.iter_mut().enumerate() { + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the loop. + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + + slot.write(value); + accumulated |= row_failure; + } + + accumulated + }; + + // SAFETY: normal completion of either execution path initializes `0..row_count` exactly + // once, and `values` was allocated with at least `row_count` capacity. + unsafe { values.set_len(row_count) }; + + // Defer failures so batch execution can retry with 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..4fd2b42476b --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/execute/sink.rs @@ -0,0 +1,441 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Executes row kernels that write through an [`OutputSink`]. +//! +//! Dense execution visits every row. Skip-invalid execution initializes skipped output rows and +//! visits only rows that are valid in every input. Skip-invalid execution declines when either the +//! input representation or sink cannot support that path. + +use vortex_buffer::BitBuffer; +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::RowExecution; +use crate::ExecutionCtx; +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; + +/// Verify that every decoded input addresses exactly `row_count` rows. +/// +/// Unlike the owned executor, the paths with and without batch constants can share this check +/// without losing sink-loop vectorization under multiple CGUs without LTO. The exact pass +/// interaction is unknown. +fn verify_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 inputs once, then write one sink row for each input row. +/// +/// The executor owns the sink and passes each output row to `apply`. This keeps `apply` as [`Fn`]. +/// Capturing the sink would require [`FnMut`] and put its buffer metadata behind loop-carried +/// mutable closure state, which can prevent LLVM from treating that metadata as loop-invariant. +pub(crate) fn execute_sink( + args: &dyn ExecutionArgs, + 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 columns = Args::decode(args, ctx)?; + let views = Args::views_no_constants(&columns); + + let row_count = args.row_count(); + verify_lengths::(&columns, views.as_ref(), row_count)?; + + let constants = Args::constants(&columns); + let prepared = prepare(constants); + + let mut sink = >::with_capacity(row_count)?; + + // Keep `rows` scoped so its borrow ends before `finish`, which consumes the sink. + { + let mut rows = >::rows(&mut sink); + + // This equality proves to LLVM that `0..row_count` is in bounds for `rows`. + let sink_row_count = >::row_count(&rows); + vortex_ensure_eq!( + sink_row_count, + row_count, + "the output sink must address exactly {row_count} rows, got {sink_row_count}", + ); + + if let Some(views) = views { + for index in 0..row_count { + // SAFETY: `verify_lengths` proved every view has `row_count` rows before the loop. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + apply(&prepared, elements, output).into_result()?; + } + } else { + for index in 0..row_count { + // SAFETY: the sink row-count check above proved every loop index is in bounds. + let output = + unsafe { >::row_unchecked(&mut rows, index) }; + + // LLVM unswitches the batch-constant checks in `Args::get` before vectorizing the + // loop. + apply(&prepared, Args::get(&columns, index), output).into_result()?; + } + } + } + + // SAFETY: every row callback completed successfully, so each returned the required write token. + unsafe { >::finish(sink) }.map(RowExecution::Output) +} + +/// Write only the rows set in `valid`, or decline when the inputs or sink cannot support +/// skip-invalid execution. +/// +/// `Ok(None)` signals batch execution to filter every input to the valid rows, run the dense +/// kernel, and scatter the results back into a null-padded array. +pub(crate) fn execute_sink_valid_rows( + args: &dyn ExecutionArgs, + 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>, +{ + let Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + mut sink, + }) = setup_sink_valid_rows::(args, valid, ctx)? + else { + return Ok(None); + }; + + let views = Args::views_no_constants(&columns); + verify_lengths::(&columns, views.as_ref(), row_count)?; + + let constants = Args::constants(&columns); + let prepared = prepare(constants); + + // Keep `rows` scoped so its borrow ends before `finish`. With multiple CGUs and no LTO, using + // `drop(rows)` duplicates `Args::get` in every sparse callback. + { + // Initialize every slot before visiting only valid rows. + let mut rows = >::rows(&mut sink); + initialize_skipped_rows(&mut rows); + + // The initializer can change addressability. Recheck it so LLVM can prove every mask + // index is in bounds. + let initialized_row_count = >::row_count(&rows); + vortex_ensure_eq!( + initialized_row_count, + row_count, + "the initialized output sink must address exactly {row_count} rows, got {initialized_row_count}", + ); + + if let Some(views) = views { + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check 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) }; + + // SAFETY: `verify_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).into_result() + })?; + } else { + valid_rows.try_for_each_set_index(|index| { + // SAFETY: the post-initialization row-count check 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) }; + + apply(&prepared, Args::get(&columns, index), output).into_result() + })?; + } + } + + // SAFETY: the initializer completed before traversal, and every visited callback completed + // successfully and returned the required write token. + unsafe { >::finish(sink) } + .map(RowExecution::Output) + .map(Some) +} + +/// State resolved before preparing the skip-invalid row loop. +struct ValidRowsSetup<'valid, Args, Sink, Options> +where + Args: ElementTuple, + Sink: OutputSink, +{ + initialize_skipped_rows: for<'rows> fn(&mut >::Rows<'rows>), + columns: Args::Columns, + valid_rows: &'valid BitBuffer, + row_count: usize, + sink: Sink, +} + +/// Resolve the capabilities, inputs, sink, and validity mask for skip-invalid execution. +fn setup_sink_valid_rows<'valid, Args, Sink, Options>( + args: &dyn ExecutionArgs, + valid: &'valid Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult>> +where + Args: ElementTuple, + Sink: OutputSink, +{ + // The initializer both declares support for skipping rows and initializes those rows. + let Some(initialize_skipped_rows) = >::skipped_rows_initializer() + else { + return Ok(None); + }; + + // Null-tolerant decoding exposes values behind nulls without filtering. Decline when any input + // cannot provide those values safely. + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + + let row_count = args.row_count(); + + // Keep allocation before the validity and length checks. With multiple CGUs and no LTO, + // moving it later inlines `Args::get` into every sparse callback, duplicating its bounds + // checks. + let sink = >::with_capacity(row_count)?; + + // Batch execution resolves all-valid and all-null inputs before selecting this path. + let AllOr::Some(valid_rows) = valid.bit_buffer() else { + vortex_bail!( + "execute_sink_valid_rows requires valid and invalid rows, got an all-valid or all-invalid mask" + ); + }; + vortex_ensure_eq!( + valid_rows.len(), + row_count, + "the validity mask must address exactly {row_count} rows, got {}", + valid_rows.len(), + ); + + Ok(Some(ValidRowsSetup { + initialize_skipped_rows, + columns, + valid_rows, + row_count, + sink, + })) +} + +#[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 { + type Rows<'a> = (); + type Row<'a> = (); + type WriteToken = (); + + fn output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(_rows: usize) -> VortexResult { + Err(vortex_err!( + "a non-skipping sink must decline before allocation" + )) + } + + fn rows(&mut self) -> Self::Rows<'_> {} + + fn row_count(_rows: &Self::Rows<'_>) -> usize { + 0 + } + + 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")) + } + } + + // 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 output_dtype(_options: &Options, _args: &[DType]) -> VortexResult { + Ok(DType::from(i64::PTYPE)) + } + + fn with_capacity(rows: usize) -> VortexResult { + Ok(Self(vec![0; rows])) + } + + fn rows(&mut self) -> Self::Rows<'_> { + &mut self.0 + } + + fn row_count(rows: &Self::Rows<'_>) -> usize { + rows.len() + } + + 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(); + 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, + &valid, + &mut ctx, + |_| (), + |_, _, _| (), + )?; + + 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, + &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, + &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("initialized output sink must address exactly 2 rows, got 1"), + "unexpected error: {error}", + ); + + 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..91f225ca071 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -11,10 +11,21 @@ //! [`RowFn::dispatch`] implementation uses a [`RowVisitor`] to select an [`ElementTuple`] and //! either an [`OutputElement`] or [`OutputSink`] for each supported dtype combination. //! +//! Unlike a general strict function, a [`RowFn`] cannot produce null from valid inputs. +//! +//! A _partially valid_ batch contains both valid and invalid rows. _Skip-invalid_ runs the kernel +//! only for valid rows without changing row positions. _Filter-and-scatter_ compacts valid rows, +//! runs the kernel, and restores their positions. +//! //! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits //! reduce compact failure evidence in that loop and retry only valid rows when null payloads may //! have caused the failure. +mod 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/row_fn.rs b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs index 8a982c4fb37..831cc25f251 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -19,7 +19,12 @@ use super::visitor::RowVisitor; use crate::dtype::DType; use crate::scalar_fn::ScalarFnId; -/// A scalar function computed one row at a time. +/// A strict scalar function whose row kernel cannot produce null from valid inputs. +/// +/// This is stronger than +/// [`ScalarFnVTable::is_strict`](crate::scalar_fn::ScalarFnVTable::is_strict), which requires null +/// propagation but permits valid inputs to produce null. The framework derives output validity +/// only from input validity. /// /// Declare argument names and use [`dispatch`](Self::dispatch) to select element and output types. /// Every implementation receives the standard [`ScalarFnVTable`]. A public type that needs custom diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs index 2764840da7a..c9140c00b8c 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/input.rs @@ -68,7 +68,7 @@ pub unsafe trait InputElement: 'static { /// cannot decode this particular array. /// /// Override this for a non-dense-safe representation that can still place safe placeholders in - /// null slots. The skip-invalid executor never reads those slots. + /// null slots. Valid-row execution never reads those slots. fn decode_null_tolerant( array: ArrayRef, ctx: &mut ExecutionCtx, 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/element_tuple.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs index b0a3e709696..7c1baed168a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs @@ -27,20 +27,31 @@ pub struct ArgColumn( ); enum ArgColumnKind { + /// One decoded value per batch row; executors validate the exact length before traversal. PerRow(T::Column), + + /// Exactly one decoded row, established by [`ArgColumn::try_from_constant`]. Constant(T::Column), } impl ArgColumn { + fn try_from_constant(column: T::Column) -> VortexResult { + let decoded_len = T::view_len(&T::view(&column)); + vortex_ensure_eq!( + decoded_len, + 1, + "a decoded batch-constant input must contain exactly 1 row, got {decoded_len}", + ); + + Ok(Self(ArgColumnKind::Constant(column))) + } + fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { // An empty input has no row 0 to slice, and its row loop runs zero times either way. if let Some(constant) = batch_constant(&array) && !array.is_empty() { - return Ok(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?); } Ok(Self(ArgColumnKind::PerRow(T::decode(array, ctx)?))) @@ -52,10 +63,7 @@ impl ArgColumn { if let Some(constant) = batch_constant(&array) && !array.is_empty() { - return Ok(Some(Self(ArgColumnKind::Constant(T::decode( - constant.slice(0..1)?, - ctx, - )?)))); + return Self::try_from_constant(T::decode(constant.slice(0..1)?, ctx)?).map(Some); } Ok(T::decode_null_tolerant(array, ctx)? @@ -88,14 +96,14 @@ impl ArgColumn { } fn addresses_rows(&self, row_count: usize) -> bool { - // A constant is always read at index zero, so it addresses any batch length. + // A constant is validated when constructed and is always read at index zero. match &self.0 { ArgColumnKind::PerRow(column) => T::view_len(&T::view(column)) == row_count, ArgColumnKind::Constant(_) => true, } } - fn constant(&self) -> Option> { + fn constant_value(&self) -> Option> { match &self.0 { ArgColumnKind::PerRow(_) => None, ArgColumnKind::Constant(column) => Some(T::get(column, 0)), @@ -130,7 +138,7 @@ pub trait ElementTuple: 'static + private::Sealed { /// The decoded column representations. type Columns; - /// Borrowed views of decoded columns when every argument stores one value per row. + /// Borrowed views of decoded columns with no batch constants. type Views<'a>; /// The borrowed row of element values. @@ -138,9 +146,9 @@ pub trait ElementTuple: 'static + private::Sealed { /// The batch-constant element values. /// - /// `Some` carries the value of a batch-constant argument. `None` marks a per-row argument. A - /// [`RowVisitor`] passes these values to its prepare closure so constant work can leave the row - /// loop. + /// `Some` carries the value of a batch-constant argument. `None` marks a non-constant argument. + /// A [`RowVisitor`] passes these values to its prepare closure so constant work can leave the + /// row loop. /// /// [`RowVisitor`]: crate::scalar_fn::unstable::row::RowVisitor type ConstElems<'a>; @@ -170,34 +178,38 @@ pub trait ElementTuple: 'static + private::Sealed { /// Decode every input column once while tolerating null rows. /// - /// Return `Ok(None)` when an argument has no null-tolerant representation. The skip-invalid - /// strategy calls this once per batch. + /// Return `Ok(None)` when an argument has no null-tolerant representation. Valid-row execution + /// calls this once per batch. fn decode_null_tolerant( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult>; /// Read the row of elements at `index`. Must be `O(1)`: it is called in the row loop. + /// + /// Each argument selects either its batch-constant value or row `index`. Keep that selection + /// visible in the loop so LLVM can unswitch it before vectorizing. fn get(columns: &Self::Columns, index: usize) -> Self::Elems<'_>; - /// Borrow every decoded column directly, or `None` when any argument is batch-constant. + /// Borrow the decoded columns when none is batch-constant. /// - /// This is selected once outside the hot loop. Keeping `ArgColumn` out of the resulting tuple - /// gives the optimizer ordinary contiguous column access without a per-row constant check. - fn per_row_views(columns: &Self::Columns) -> Option>; + /// Returns `None` if any column is batch-constant. Otherwise, omitting [`ArgColumn`] from the + /// returned tuple removes constant checks from the row loop. + fn views_no_constants(columns: &Self::Columns) -> Option>; /// Whether every view contains exactly `row_count` rows. /// - /// The executor calls this once before the all-per-row hot loop. A successful check gives LLVM - /// a dominating equality between the loop bound and every source length, which lets it optimize - /// the tuple access as one fixed-length traversal. + /// The executor calls this once before the loop used when no input is batch-constant. A + /// successful check gives LLVM a dominating equality between the loop bound and every source + /// length, which lets it optimize the tuple access as one fixed-length traversal. fn view_lens_match(views: &Self::Views<'_>, row_count: usize) -> bool; - /// Whether every per-row argument contains exactly `row_count` rows. + /// Whether every non-constant argument contains exactly `row_count` rows. /// - /// This is the mixed-shape equivalent of [`view_lens_match`](Self::view_lens_match) when - /// [`per_row_views`](Self::per_row_views) declines. It runs once before the hot loop for the - /// same LLVM optimization. A batch constant is exempt because decoding collapsed it to one row. + /// This is the equivalent of [`view_lens_match`](Self::view_lens_match) when the columns include + /// batch constants. It runs once before the hot loop for the same LLVM optimization. A batch + /// constant is exempt because its [`ArgColumn`] constructor already validated the one-row + /// representation produced by decoding. fn decoded_lens_match(columns: &Self::Columns, row_count: usize) -> bool; /// Read one row from borrowed views. @@ -257,7 +269,7 @@ impl ElementTuple for () { fn get(_columns: &Self::Columns, _index: usize) -> Self::Elems<'_> {} - fn per_row_views(_columns: &Self::Columns) -> Option> { + fn views_no_constants(_columns: &Self::Columns) -> Option> { Some(()) } @@ -341,7 +353,7 @@ macro_rules! element_tuple { ($(columns.$idx.get(index),)+) } - fn per_row_views(columns: &Self::Columns) -> Option> { + fn views_no_constants(columns: &Self::Columns) -> Option> { Some(($($t::view(columns.$idx.per_row_column()?),)+)) } @@ -372,7 +384,7 @@ macro_rules! element_tuple { } fn constants(columns: &Self::Columns) -> Self::ConstElems<'_> { - ($(columns.$idx.constant(),)+) + ($(columns.$idx.constant_value(),)+) } } }; diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs index d612d874935..832dbf9ee19 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/tuple/indexed.rs @@ -17,7 +17,7 @@ use crate::scalar_fn::unstable::row::InputElement; /// Every [`ElementTuple`] implements this trait. Its source delegates each lane read to the tuple's /// unchecked view access after batch execution validates every decoded column length once. pub trait IndexedElementTuple: ElementTuple { - /// The source shared execution uses for a dense all-per-row loop. + /// The source used when no input is batch-constant. /// /// Its length must be the common view length. For every valid index it must preserve row order, /// return the same value as [`ElementTuple::get_from_views`], and uphold the unchecked read 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/types/sink.rs b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs index 6cc3ce06d30..22c90b8da02 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -21,8 +21,7 @@ use crate::scalar_fn::unstable::row::OutputElement; /// batch state. The executor passes each row slot into an [`Fn`] closure. /// /// Rows arrive in increasing index order. Ordinary execution visits `0..row_count` exactly once. -/// Skip-invalid execution can omit invalid rows when [`skipped_rows_initializer`] returns an -/// initializer. +/// Execution can omit invalid rows when [`skipped_rows_initializer`] returns an initializer. /// /// # Errors /// @@ -73,10 +72,11 @@ pub unsafe trait OutputSink: 'static + Sized { /// **must not** be able to construct one without establishing the invariant. type WriteToken: 'static; - /// The operation that initializes every output position before skip-invalid execution. + /// The operation that initializes every output position before + /// [skip-invalid execution](crate::scalar_fn::unstable::row). /// - /// `Some` enables skip-invalid execution. The initializer **must** make every row safe to - /// finish. Callbacks overwrite valid rows, and batch execution masks skipped rows. + /// `Some` enables this strategy. The initializer **must** make every row safe to finish. + /// Callbacks overwrite valid rows, and batch execution masks skipped rows. /// /// `None` makes the executor fall back to filtering the inputs. fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { @@ -153,7 +153,7 @@ impl InitializedElement { /// The row closure must return the [`InitializedElement`] from [`InitializedElement::write`] on /// success. The token is zero-sized, so the proof adds no runtime row state. /// -/// Skip-invalid execution initializes placeholders before omitting rows. Errors and unwinds are +/// When execution omits invalid rows, it initializes placeholders first. Errors and unwinds are /// safe because `values` keeps length zero until `finish`. The `T: Copy` bound means that /// initialized spare-capacity elements require no destruction. pub struct UninitElementSink { 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..b003f70a395 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -0,0 +1,306 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Visitors that execute dense and skip-invalid row loops. +//! +//! 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, so the batch layer +//! filters the inputs and retries with [`ExecuteRows`]. + +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::check::validate_owned_visit; +use super::check::validate_sink_visit; +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 runtime visit that decodes every column once and runs the selected row loop. +pub(crate) 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, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { + pub(crate) 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, + } + } +} + +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_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + 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_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink::( + self.args, 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_plan( + self.output_dtype, + self.policy, + validate_owned_visit::(self.dtypes)?, + RowPolicy::for_deferred_output::(), + )?; + + execute_owned::( + self.args, + self.ctx, + prepare, + apply, + finish_failure, + ) + } +} + +/// The runtime visit that executes valid rows over the original input columns. +/// +/// Only output sinks have a contract for skipped output positions. Owned visits therefore decline +/// so batch execution can filter the valid inputs and scatter the output back. +pub(crate) 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, + + /// The nullable execution policy selected by the planning visit. + policy: RowPolicy, + + /// The conjoined validity, containing both valid and invalid rows. + valid: &'args Mask, + + /// The execution context used to decode the input columns. + ctx: &'ctx mut ExecutionCtx, +} + +impl<'args, 'ctx, F: RowFn> ExecuteValidRows<'args, 'ctx, F> { + pub(crate) fn new( + args: &'args dyn ExecutionArgs, + dtypes: &'args [DType], + options: &'args F::Options, + output_dtype: &'args DType, + policy: RowPolicy, + valid: &'args Mask, + ctx: &'ctx mut ExecutionCtx, + ) -> Self { + Self { + args, + dtypes, + options, + output_dtype, + policy, + valid, + ctx, + } + } +} + +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_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. + 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_plan( + self.output_dtype, + self.policy, + validate_sink_visit::(self.options, self.dtypes)?, + RowPolicy::for_sink::(), + )?; + + execute_sink_valid_rows::( + self.args, 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_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) + } +} + +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_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/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..5360ced573b 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)); @@ -139,7 +136,7 @@ pub(crate) enum RowPolicy { /// Evaluate all rows, retrying only valid rows if a deferred error is raised. DenseWithRetry, - /// Execute only valid rows, trying skip-invalid execution before filtering. + /// Execute only valid rows over the original inputs before filtering. ValidOnly, } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs index 0e7abaa4322..1a9e8628bc7 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/row_visitor.rs @@ -74,7 +74,7 @@ pub trait RowVisitor: private::Sealed + Sized { /// # Examples /// /// Test whether each string occurs in its allowed-values list. The prepare closure builds one - /// lookup table for a batch-constant list. The row closure scans a varying list directly. + /// lookup table for a batch-constant list. The row closure scans a per-row list directly. /// /// ```ignore /// visitor.visit_prepared::< diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index b6c14585c48..257ac3ca0b5 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; @@ -69,6 +76,8 @@ impl ScalarFnVTable for F { union_child_validities(expression) } + // `RowFn` is stricter than `ScalarFnVTable::is_strict`: its kernel cannot produce null from + // valid inputs, so batch execution derives output validity only from input validity. fn is_strict(&self, _options: &Self::Options) -> bool { true } @@ -98,20 +107,39 @@ 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, ) } -/// 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!( @@ -124,26 +152,101 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +fn execute_row_kernel( + function: &F, + options: &F::Options, + args: BorrowedExecutionArgs<'_>, + ctx: &mut ExecutionCtx, +) -> VortexResult { + function.dispatch( + options, + args.dtypes(), + ExecuteRows::::new( + &args, + args.dtypes(), + options, + args.output_dtype(), + args.policy(), + ctx, + ), + ) +} + +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.dtypes(), + options, + args.output_dtype(), + args.policy(), + valid, + ctx, + ), + ) +} + +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, + change: DispatchChange, + } + + #[derive(Clone, Copy)] + enum DispatchChange { + Policy, + Element, + } + impl RowFn for IndexingRowFn { type Options = EmptyOptions; @@ -168,6 +271,35 @@ 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 { + match self.change { + DispatchChange::Policy => visitor + .visit_deferred::<(i64,), i64, bool>(|(value,)| (value, false), |_| Ok(())), + DispatchChange::Element => visitor.visit::<(u64,), u64>(|(value,)| value), + } + } + } + } + #[test] fn test_return_dtype_rejects_wrong_arity_before_dispatch() { let error = row_fn_return_dtype(&IndexingRowFn, &EmptyOptions, &[]) @@ -186,6 +318,52 @@ 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)), + change: DispatchChange::Policy, + }; + 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}", + ); + } + + #[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!( diff --git a/vortex-buffer/src/bit/buf.rs b/vortex-buffer/src/bit/buf.rs index aac8ead42fe..1da0b037f2a 100644 --- a/vortex-buffer/src/bit/buf.rs +++ b/vortex-buffer/src/bit/buf.rs @@ -479,6 +479,33 @@ impl BitBuffer { } } + /// Fallible variant of [`for_each_set_index`](Self::for_each_set_index). + /// + /// Stops and returns the first error from `f`. + #[inline] + pub fn try_for_each_set_index(&self, mut f: F) -> Result<(), E> + where + F: FnMut(usize) -> Result<(), E>, + { + let mut base = 0usize; + for word in self.chunks().iter_padded() { + if word == u64::MAX { + for k in 0..64 { + f(base + k)?; + } + } else { + let mut w = word; + while w != 0 { + f(base + w.trailing_zeros() as usize)?; + w &= w - 1; + } + } + base += 64; + } + + Ok(()) + } + /// Created a new BitBuffer with offset reset to 0 pub fn sliced(&self) -> Self { if self.offset.is_multiple_of(8) { @@ -970,12 +997,21 @@ mod tests { #[case(65)] #[case(200)] #[case(1000)] - fn test_for_each_set_index_matches_set_indices(#[case] len: usize) { + fn test_set_index_visitors_match_set_indices(#[case] len: usize) { let buf = BitBuffer::collect_bool(len, |i| i % 5 == 0 || i % 7 == 0); let expected: Vec = buf.set_indices().collect(); + let mut got = Vec::new(); buf.for_each_set_index(|i| got.push(i)); assert_eq!(got, expected); + + let mut fallible_got = Vec::new(); + let result = buf.try_for_each_set_index(|i| { + fallible_got.push(i); + Ok::<(), ()>(()) + }); + assert_eq!(result, Ok(())); + assert_eq!(fallible_got, expected); } #[rstest] @@ -998,6 +1034,33 @@ mod tests { assert_eq!(got, (0..130).collect::>()); } + #[test] + fn test_try_for_each_set_index_stops_on_error() { + for (buffer, stop) in [ + (BitBuffer::new_set(130), 65), + (BitBuffer::collect_bool(130, |i| i % 3 == 0), 66), + ] { + let mut visited = Vec::new(); + let result = buffer.try_for_each_set_index(|index| { + visited.push(index); + if index == stop { + return Err(index); + } + + Ok(()) + }); + + assert_eq!(result, Err(stop)); + assert_eq!( + visited, + buffer + .set_indices() + .take_while(|&i| i <= stop) + .collect::>() + ); + } + } + #[test] fn test_map_cmp_conditional() { // map_cmp with conditional logic based on index and bit value