From 3004bf0c525b81ebcb6f31c0c2727608724224ce Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 19 Aug 2026 15:14:20 -0400 Subject: [PATCH 1/3] Implement RowFn batch execution Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 91 +++++++++ .../unstable/row/batch/execute/constant.rs | 34 ++++ .../unstable/row/batch/execute/dense.rs | 32 +++ .../unstable/row/batch/execute/mod.rs | 76 +++++++ .../unstable/row/batch/execute/output.rs | 100 ++++++++++ .../unstable/row/batch/execute/valid_only.rs | 109 ++++++++++ .../src/scalar_fn/unstable/row/batch/mod.rs | 65 ++++++ .../scalar_fn/unstable/row/batch/planning.rs | 77 ++++++++ .../src/scalar_fn/unstable/row/batch/tests.rs | 187 ++++++++++++++++++ .../src/scalar_fn/unstable/row/mod.rs | 7 +- .../src/scalar_fn/unstable/row/row_fn.rs | 7 + .../unstable/row/types/element/mod.rs | 1 + .../unstable/row/types/element/tuple/mod.rs | 1 + .../src/scalar_fn/unstable/row/types/mod.rs | 1 + .../src/scalar_fn/unstable/row/types/sink.rs | 2 +- .../scalar_fn/unstable/row/visitor/execute.rs | 2 +- .../src/scalar_fn/unstable/row/visitor/mod.rs | 5 +- .../scalar_fn/unstable/row/visitor/plan.rs | 2 - .../src/scalar_fn/unstable/row/vtable.rs | 185 ++++++++++++++++- 19 files changed, 966 insertions(+), 18 deletions(-) create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/args.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/mod.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/planning.rs create mode 100644 vortex-array/src/scalar_fn/unstable/row/batch/tests.rs diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs new file mode 100644 index 00000000000..9a3a0b5cd88 --- /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 or sliced 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 metadata selected for its row function. +/// +/// `arrays` can be 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 row-function invocation. + arrays: &'a [ArrayRef], + + /// The number of rows in this row-function 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..8f36ef788df --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -0,0 +1,34 @@ +// 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; + +impl Batch { + /// Execute all-constant inputs by evaluating one row and broadcasting the validated result. + pub(super) fn execute_all_constant( + &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 = + self.validate_kernel_output(kernel(self.execution_args(&one_row, 1), ctx)?, 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..3e2a23802f2 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -0,0 +1,32 @@ +// 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::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, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; + 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/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs new file mode 100644 index 00000000000..f1a167e3550 --- /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::types::batch_const; +use crate::validity::Validity; + +mod constant; +mod dense; +mod valid_only; + +mod output; +#[cfg(test)] +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_valid_rows` executes only valid rows + /// over the original inputs. Every result is checked against the planned shape and dtype. + pub(crate) fn execute( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_valid_rows: 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_const(input).is_some()) + { + return self.execute_all_constant(kernel, ctx); + } + + // A known all-valid batch does not need to materialize validity, even when its row policy + // only permits valid rows. + if self.validity.definitely_no_nulls() { + return self.execute_dense(kernel, ctx); + } + + match self.policy { + RowPolicy::Dense => self.execute_dense(kernel, ctx), + RowPolicy::ValidOnly => self.execute_valid_only(kernel, try_valid_rows, 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..f0b830c6d7b --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -0,0 +1,100 @@ +// 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() + } + + /// Validate the finished output and apply the row function's logical outer nullability. + pub(super) fn finalize_output( + &self, + values: ArrayRef, + expected_len: usize, + ) -> VortexResult { + validate_output(self.id, &self.result_dtype, expected_len, &values)?; + + cast_output_nullability(&self.result_dtype, values) + } + + /// Validate the output from a row function 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 function. +/// +/// `values` **must** contain `expected_len` rows. Its dtype must match `result_dtype` except for +/// outer 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) +} + +/// Validate an output's shape and logical dtype without executing an outer-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(), + ); + let values_with_result_nullability = + values.dtype().with_nullability(result_dtype.nullability()); + vortex_ensure!( + values_with_result_nullability == *result_dtype, + "the {id} output dtype must match {result_dtype} except for outer nullability, got {}", + values.dtype(), + ); + + Ok(()) +} + +/// Cast only the outer output nullability after validation accepts every other dtype component. +/// +/// This changes no logical dtype component other than outer nullability. An encoding that cannot +/// rewrite its nullability directly can still retain a lazy cast until execution. +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..f6fe7ab757d --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_panic; +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::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, then execute valid rows over the original inputs. + /// + /// # Panics + /// + /// Panics if the concrete row signature cannot use direct valid-row execution. Inputs must + /// support null-tolerant decoding, and output sinks must initialize skipped rows. + pub(super) fn execute_valid_only( + &self, + kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + try_valid_rows: 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_valid_rows(try_valid_rows, &valid, ctx)? { + return Ok(result); + } + + vortex_panic!( + "valid-only execution requires direct valid-row support; {} selected an unsupported signature", + self.id, + ) + } + + /// 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)?; + + // An array-backed validity can materialize to all valid even though the cheap checks in + // `Batch::execute` could not prove that. Run the full-row kernel in that case. Check + // all-true before all-false because an empty mask is both. + if valid.all_true() { + let values = 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_valid_rows( + &self, + try_valid_rows: impl FnOnce( + BorrowedExecutionArgs<'_>, + &Mask, + &mut ExecutionCtx, + ) -> VortexResult>, + valid: &Mask, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let Some(values) = try_valid_rows( + self.execution_args(&self.inputs, self.row_count), + valid, + ctx, + )? + else { + return Ok(None); + }; + 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..eb0760b6645 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -0,0 +1,65 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Batch execution for a strict row function. +//! +//! A batch is the set of same-length input columns supplied in one scalar-function call. A row +//! function handles typed values for one logical row. This module adds the columnar concerns around +//! that row function: 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; +#[cfg(test)] +pub(super) use execute::finalize_kernel_output; + +mod planning; + +pub(super) use super::visitor::BatchPlan; +pub(super) use super::visitor::RowPolicy; + +/// The same-length input columns and metadata for one row-function execution. +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 for validity, constant handling, and execution. + 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..1acf7332018 --- /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. + /// + /// This constructor **must not** be used for a nullary function. With no input columns, there + /// is no validity to propagate and the all-constant check would pass vacuously. + 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..101b6932737 --- /dev/null +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -0,0 +1,187 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; + +use super::finalize_kernel_output; +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::ConstantArray; +use crate::arrays::PrimitiveArray; +use crate::dtype::DType; +use crate::dtype::NativePType; +use crate::dtype::Nullability; +use crate::scalar::Scalar; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::VecExecutionArgs; +use crate::scalar_fn::unstable::row::OutputElement; +use crate::scalar_fn::unstable::row::OutputSink; +use crate::scalar_fn::unstable::row::RowFn; +use crate::scalar_fn::unstable::row::RowVisitor; +use crate::scalar_fn::unstable::row::execute_rows; +use crate::validity::Validity; + +#[derive(Clone)] +struct ValidOnlyIdentity; + +#[derive(Clone)] +struct InvalidKernelOutput; + +/// Produces a null row to exercise output validation at the row-function boundary. +struct NullProducingI64(i64); + +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 return_dtype(_options: &Options) -> 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() + } + + 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 ValidOnlyIdentity { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = false; + + 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 InvalidKernelOutput { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + 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)) + } +} + +#[test] +fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<()> { + static ID: CachedId = CachedId::new("test.finalize_kernel_output"); + + let element_dtype = DType::Primitive(i64::PTYPE, Nullability::NonNullable); + let values = ConstantArray::new( + Scalar::list_empty(Arc::new(element_dtype), Nullability::NonNullable), + 2, + ) + .into_array(); + let result_dtype = DType::List( + Arc::new(DType::Primitive(i64::PTYPE, Nullability::Nullable)), + Nullability::NonNullable, + ); + let mut ctx = array_session().create_execution_ctx(); + + assert!(finalize_kernel_output(*ID, &result_dtype, 2, values, &mut ctx).is_err()); + Ok(()) +} + +#[test] +fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { + 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 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(()) +} + +#[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(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/mod.rs b/vortex-array/src/scalar_fn/unstable/row/mod.rs index b863b8d24de..5551684dd4b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/mod.rs @@ -13,13 +13,16 @@ //! //! 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. +//! //! Prepared visits move work derived from constant operands outside the hot loop. Deferred visits //! reduce compact failure evidence without constructing errors in that loop. -// TODO(connor)[RowFn]: Remove this expectation when #9450 connects the batch executor. -#[expect(dead_code)] mod execute; +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 dd183d13e58..de0e25e9e46 100644 --- a/vortex-array/src/scalar_fn/unstable/row/row_fn.rs +++ b/vortex-array/src/scalar_fn/unstable/row/row_fn.rs @@ -26,10 +26,17 @@ use crate::scalar_fn::ScalarFnId; /// propagation but permits valid inputs to produce null. The framework derives output validity /// only from input validity. /// +/// A dispatched [`OutputElement`] or [`OutputSink`] describes the non-nullable values produced for +/// valid rows. The framework widens that dtype when an input dtype is nullable, attaches the +/// input-derived validity, and casts the finished array to the widened dtype. Implementations do +/// not construct nullable placeholders for invalid rows. +/// /// Declare argument names and use [`dispatch`](Self::dispatch) to select element and output types. /// Every implementation receives the standard [`ScalarFnVTable`]. A public type that needs custom /// vtable hooks can delegate its row kernel through [`row_fn_return_dtype`] and [`execute_rows`]. /// +/// [`OutputElement`]: crate::scalar_fn::unstable::row::OutputElement +/// [`OutputSink`]: crate::scalar_fn::unstable::row::OutputSink /// [`ScalarFnVTable`]: crate::scalar_fn::ScalarFnVTable /// [`execute_rows`]: crate::scalar_fn::unstable::row::execute_rows /// [`row_fn_return_dtype`]: crate::scalar_fn::unstable::row::row_fn_return_dtype 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..73a0c1d47d1 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_const; 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..cc312a3fb14 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_const; mod indexed; pub use indexed::IndexedElementTuple; 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 3b4a92e784e..4fc0c323ad4 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/mod.rs @@ -13,6 +13,7 @@ pub use element::ElementTuple; pub use element::IndexedElementTuple; pub use element::InputElement; pub use element::OutputElement; +pub(super) use element::batch_const; mod result; pub use result::FailureEvidence; 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 747d68d47f4..2bf51cf999d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/sink.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/sink.rs @@ -84,7 +84,7 @@ pub unsafe trait OutputSink: 'static + Sized { /// `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. + /// `None` makes skip-invalid execution unavailable for this sink. fn skipped_rows_initializer() -> Option fn(&mut Self::Rows<'a>)> { None } diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index 93fd640b82e..345d26a2532 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -160,7 +160,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { /// 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, -/// and batch execution decides how to handle that unsupported signature. +/// and batch execution rejects that unsupported signature. pub(crate) struct ExecuteValidRows<'args, 'ctx, F: RowFn> { /// The original inputs for this kernel invocation. args: &'args dyn ExecutionArgs, 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 1b063c788db..c7f9baf6a62 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/mod.rs @@ -8,11 +8,12 @@ mod check; pub(super) use check::assert_owned_output_needs_no_drop; -// TODO(connor)[RowFn]: Remove this expectation when #9450 constructs the execution visitors. -#[expect(dead_code)] 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; 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 a742045749d..75d5852d83f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/plan.rs @@ -114,8 +114,6 @@ pub(crate) struct BatchPlan { pub(crate) output_dtype: DType, /// How this concrete dispatch executes nullable rows. - // TODO(connor)[RowFn]: Remove this expectation when #9450 uses the planned policy. - #[expect(dead_code)] pub(crate) policy: RowPolicy, } diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 317fae1ed17..200660ab445 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -4,12 +4,14 @@ //! 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; use vortex_error::vortex_ensure_eq; +use vortex_mask::Mask; use vortex_session::VortexSession; use super::row_fn::RowFn; @@ -24,6 +26,10 @@ 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::visitor::ExecuteRows; +use crate::scalar_fn::unstable::row::visitor::ExecuteValidRows; impl ScalarFnVTable for F { type Options = F::Options; @@ -69,6 +75,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 +106,24 @@ 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())?; + vortex_ensure!( + args.num_inputs() != 0, + "row-function execution does not support nullary kernels" + ); - // 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) + let batch = prepare_batch(function, options, args)?; + batch.execute( + |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, valid, ctx| try_execute_valid_rows(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 +136,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_valid_rows( + 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 +255,35 @@ mod tests { } } + impl RowFn for ChangingDispatchRowFn { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = false; + + 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 +302,55 @@ mod tests { assert_arity_error(error); } + #[test] + fn test_execute_rejects_dispatch_that_changes_after_planning() -> VortexResult<()> { + 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 = match execute_rows(&function, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_error::vortex_bail!("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 ValidOnly"), + "unexpected error: {error}", + ); + Ok(()) + } + + #[test] + fn test_execute_revalidates_element_types_after_planning() -> VortexResult<()> { + let function = ChangingDispatchRowFn { + dispatches: Arc::new(AtomicUsize::new(0)), + change: DispatchChange::Element, + }; + let input = PrimitiveArray::new(vec![1_i64, 2], Validity::NonNullable).into_array(); + let args = VecExecutionArgs::new(vec![input], 2); + let mut ctx = array_session().create_execution_ctx(); + + let error = match execute_rows(&function, &EmptyOptions, &args, &mut ctx) { + Err(error) => error, + Ok(_) => vortex_error::vortex_bail!("dispatch must preserve its planned element types"), + }; + + assert!( + error.to_string().contains("expected a u64 column"), + "unexpected error: {error}", + ); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!( From 1125097bf0934f71cb80e29b8f6309169890075a Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 19 Aug 2026 15:15:32 -0400 Subject: [PATCH 2/3] Improve RowFn batch documentation Signed-off-by: Connor Tsui --- vortex-array/src/scalar_fn/unstable/row/batch/args.rs | 4 ++-- .../src/scalar_fn/unstable/row/batch/execute/mod.rs | 9 +++++---- vortex-array/src/scalar_fn/unstable/row/batch/mod.rs | 10 +++++----- .../src/scalar_fn/unstable/row/types/element/output.rs | 6 +++++- vortex-array/src/scalar_fn/unstable/row/vtable.rs | 8 ++++---- 5 files changed, 21 insertions(+), 16 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index 9a3a0b5cd88..29df51f7163 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -4,15 +4,15 @@ //! Execution arguments paired with the metadata selected during planning. //! //! [`BorrowedExecutionArgs`] can point at original or sliced arrays while retaining the dtypes, -//! output dtype, and null policy of the original batch plan. +//! output dtype, and execution policy of the original batch plan. use vortex_error::VortexResult; use vortex_error::vortex_err; +use super::RowPolicy; 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 metadata selected for its row function. /// 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 index f1a167e3550..091b7de8a05 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -29,8 +29,8 @@ 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_valid_rows` executes only valid rows - /// over the original inputs. Every result is checked against the planned shape and dtype. + /// For a partially valid batch, `try_valid_rows` executes only valid rows over the original + /// inputs. Every result is checked against the planned shape and dtype. pub(crate) fn execute( &self, kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, @@ -53,8 +53,9 @@ impl Batch { 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. + // All inputs are constant, and their conjoined validity proves that every row is non-null. + // The constant check sees through extension and masked wrappers, just like argument + // decoding. if self.row_count > 0 && self.validity.definitely_no_nulls() && self.inputs.iter().all(|input| batch_const(input).is_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 index eb0760b6645..22351f46cbf 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -5,7 +5,7 @@ //! //! A batch is the set of same-length input columns supplied in one scalar-function call. A row //! function handles typed values for one logical row. This module adds the columnar concerns around -//! that row function: planning the output and null strategy, preserving batch constants, +//! that row function: planning the output and null handling, preserving batch constants, //! propagating strict validity, selecting an execution strategy, and validating the finished //! output. //! @@ -46,12 +46,12 @@ pub(crate) struct Batch { /// 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. + /// The conjoined input validity. An output row is valid exactly when it is valid in every + /// input. Conjoining is lazy, and null handling materializes the mask only when required. 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. + /// The declared output dtype, widened to nullable when any input is nullable. Kernel output is + /// reconciled against this dtype. result_dtype: DType, /// The non-nullable dtype the dispatched output capability builds, computed while planning. diff --git a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs index 3ca5af5672c..f7104f47650 100644 --- a/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/types/element/output.rs @@ -19,6 +19,10 @@ pub trait OutputElement: 'static + Sized { /// [`OutputSink`]: crate::scalar_fn::unstable::row::OutputSink fn element_dtype() -> DType; - /// Build a column from one value per row. Called once per batch. + /// Build an all-valid column from one value per row. + /// + /// The returned column must contain `values.len()` rows and match + /// [`element_dtype`](Self::element_dtype) except for outer nullability. The framework calls + /// this method once per batch. fn build(values: Vec) -> ArrayRef; } diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 200660ab445..5b3813fa4f5 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -14,8 +14,12 @@ use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; use vortex_session::VortexSession; +use super::batch::Batch; +use super::batch::BorrowedExecutionArgs; use super::row_fn::RowFn; use super::visitor::BatchPlanner; +use super::visitor::ExecuteRows; +use super::visitor::ExecuteValidRows; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; @@ -26,10 +30,6 @@ 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::visitor::ExecuteRows; -use crate::scalar_fn::unstable::row::visitor::ExecuteValidRows; impl ScalarFnVTable for F { type Options = F::Options; From dae66b7b1b3345f5b047421ff3b7de674f7efeb3 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 20 Aug 2026 13:09:48 +0000 Subject: [PATCH 3/3] Rename RowFn batch execution types Rename `Batch` to `RowFnExecutionArgs` and `BorrowedExecutionArgs` to `BorrowedRowFnArgs` so both names say which layer owns them. Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/args.rs | 8 ++++---- .../unstable/row/batch/execute/constant.rs | 8 ++++---- .../unstable/row/batch/execute/dense.rs | 10 +++++----- .../unstable/row/batch/execute/mod.rs | 14 +++++++------- .../unstable/row/batch/execute/output.rs | 4 ++-- .../unstable/row/batch/execute/valid_only.rs | 18 +++++++++--------- .../src/scalar_fn/unstable/row/batch/mod.rs | 10 +++++----- .../scalar_fn/unstable/row/batch/planning.rs | 10 +++++----- .../src/scalar_fn/unstable/row/vtable.rs | 12 ++++++------ 9 files changed, 47 insertions(+), 47 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs index 29df51f7163..d45a66529a3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/args.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/args.rs @@ -3,7 +3,7 @@ //! Execution arguments paired with the metadata selected during planning. //! -//! [`BorrowedExecutionArgs`] can point at original or sliced arrays while retaining the dtypes, +//! [`BorrowedRowFnArgs`] can point at original or sliced arrays while retaining the dtypes, //! output dtype, and execution policy of the original batch plan. use vortex_error::VortexResult; @@ -20,7 +20,7 @@ use crate::scalar_fn::ExecutionArgs; /// 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> { +pub(crate) struct BorrowedRowFnArgs<'a> { /// The input arrays for this row-function invocation. arrays: &'a [ArrayRef], @@ -37,7 +37,7 @@ pub(crate) struct BorrowedExecutionArgs<'a> { policy: RowPolicy, } -impl<'a> BorrowedExecutionArgs<'a> { +impl<'a> BorrowedRowFnArgs<'a> { /// Pair one input view with the planning metadata selected for its batch. pub(crate) fn new( arrays: &'a [ArrayRef], @@ -71,7 +71,7 @@ impl<'a> BorrowedExecutionArgs<'a> { } } -impl ExecutionArgs for BorrowedExecutionArgs<'_> { +impl ExecutionArgs for BorrowedRowFnArgs<'_> { fn get(&self, index: usize) -> VortexResult { self.arrays.get(index).cloned().ok_or_else(|| { vortex_err!( 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 index 8f36ef788df..0cd60e821b6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -4,18 +4,18 @@ use smallvec::SmallVec; use vortex_error::VortexResult; -use super::super::Batch; -use super::super::args::BorrowedExecutionArgs; +use super::super::RowFnExecutionArgs; +use super::super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; -impl Batch { +impl RowFnExecutionArgs { /// Execute all-constant inputs by evaluating one row and broadcasting the validated result. pub(super) fn execute_all_constant( &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let one_row: SmallVec<[ArrayRef; 4]> = self 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 index 3e2a23802f2..9855e74b0bc 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -3,18 +3,18 @@ use vortex_error::VortexResult; -use super::super::Batch; -use super::super::args::BorrowedExecutionArgs; +use super::super::RowFnExecutionArgs; +use super::super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::ExecutionCtx; use crate::builtins::ArrayBuiltins; use crate::validity::Validity; -impl Batch { +impl RowFnExecutionArgs { /// 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, + kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; @@ -25,7 +25,7 @@ impl Batch { 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. + // Handled by the guard in `RowFnExecutionArgs::execute`, before the kernel ran. Validity::AllInvalid => Ok(self.all_null()), } } 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 index 091b7de8a05..60a23a1154d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -3,15 +3,15 @@ //! Selects a batch execution strategy. //! -//! [`Batch::execute`] handles universal fast paths, then delegates to dense or valid-only -//! execution. +//! [`RowFnExecutionArgs::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::RowFnExecutionArgs; use super::RowPolicy; -use super::args::BorrowedExecutionArgs; +use super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::ExecutionCtx; use crate::arrays::Constant; @@ -26,16 +26,16 @@ mod output; #[cfg(test)] pub(crate) use output::finalize_kernel_output; -impl Batch { +impl RowFnExecutionArgs { /// Apply constant folding and null handling around `kernel`. /// /// For a partially valid batch, `try_valid_rows` executes only valid rows over the original /// inputs. Every result is checked against the planned shape and dtype. pub(crate) fn execute( &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_valid_rows: impl FnOnce( - BorrowedExecutionArgs<'_>, + BorrowedRowFnArgs<'_>, &Mask, &mut ExecutionCtx, ) -> VortexResult>, 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 index f0b830c6d7b..0f8170630e3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -5,7 +5,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; -use super::super::Batch; +use super::super::RowFnExecutionArgs; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -15,7 +15,7 @@ use crate::dtype::DType; use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnId; -impl Batch { +impl RowFnExecutionArgs { pub(super) fn all_null(&self) -> ArrayRef { ConstantArray::new(Scalar::null(self.result_dtype.clone()), self.row_count).into_array() } 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 index f6fe7ab757d..c614e04dc3e 100644 --- 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 @@ -5,8 +5,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_mask::Mask; -use super::super::Batch; -use super::super::args::BorrowedExecutionArgs; +use super::super::RowFnExecutionArgs; +use super::super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; @@ -23,7 +23,7 @@ enum ResolvedValidity { PartiallyValid(Mask), } -impl Batch { +impl RowFnExecutionArgs { /// Resolve validity, then execute valid rows over the original inputs. /// /// # Panics @@ -32,9 +32,9 @@ impl Batch { /// support null-tolerant decoding, and output sinks must initialize skipped rows. pub(super) fn execute_valid_only( &self, - kernel: impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, try_valid_rows: impl FnOnce( - BorrowedExecutionArgs<'_>, + BorrowedRowFnArgs<'_>, &Mask, &mut ExecutionCtx, ) -> VortexResult>, @@ -58,14 +58,14 @@ impl Batch { /// Materialize validity and handle all-valid or all-null batches. fn resolve_validity( &self, - kernel: &impl Fn(BorrowedExecutionArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: &impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let valid = self.validity.clone().execute_mask(self.row_count, ctx)?; // An array-backed validity can materialize to all valid even though the cheap checks in - // `Batch::execute` could not prove that. Run the full-row kernel in that case. Check - // all-true before all-false because an empty mask is both. + // `RowFnExecutionArgs::execute` could not prove that. Run the full-row kernel in that + // case. Check all-true before all-false because an empty mask is both. if valid.all_true() { let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; let values = self.validate_kernel_output(values, self.row_count, ctx)?; @@ -85,7 +85,7 @@ impl Batch { fn try_execute_valid_rows( &self, try_valid_rows: impl FnOnce( - BorrowedExecutionArgs<'_>, + BorrowedRowFnArgs<'_>, &Mask, &mut ExecutionCtx, ) -> VortexResult>, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs index 22351f46cbf..dd6cfa9ba0f 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/mod.rs @@ -9,9 +9,9 @@ //! 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. +//! [`BatchPlan`] carries the nullable execution strategy selected by a concrete dispatch. +//! [`RowFnExecutionArgs`] applies that strategy, and [`BorrowedRowFnArgs`] pairs each kernel +//! invocation with its planning metadata. use smallvec::SmallVec; @@ -21,7 +21,7 @@ use crate::scalar_fn::ScalarFnId; use crate::validity::Validity; mod args; -pub(super) use args::BorrowedExecutionArgs; +pub(super) use args::BorrowedRowFnArgs; mod execute; #[cfg(test)] @@ -33,7 +33,7 @@ pub(super) use super::visitor::BatchPlan; pub(super) use super::visitor::RowPolicy; /// The same-length input columns and metadata for one row-function execution. -pub(crate) struct Batch { +pub(crate) struct RowFnExecutionArgs { /// The function being executed, named in the errors this raises. id: ScalarFnId, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs index 1acf7332018..fa52fd91b05 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/planning.rs @@ -5,16 +5,16 @@ use smallvec::SmallVec; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; -use super::Batch; use super::BatchPlan; -use super::args::BorrowedExecutionArgs; +use super::RowFnExecutionArgs; +use super::args::BorrowedRowFnArgs; use crate::ArrayRef; use crate::dtype::DType; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::validity::Validity; -impl Batch { +impl RowFnExecutionArgs { /// Collect the inputs and derive their dtypes, validity, and execution policy. /// /// This constructor **must not** be used for a nullary function. With no input columns, there @@ -65,8 +65,8 @@ impl Batch { &'b self, arrays: &'b [ArrayRef], row_count: usize, - ) -> BorrowedExecutionArgs<'b> { - BorrowedExecutionArgs::new( + ) -> BorrowedRowFnArgs<'b> { + BorrowedRowFnArgs::new( arrays, row_count, &self.arg_dtypes, diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 5b3813fa4f5..c17c59a96f3 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -14,8 +14,8 @@ use vortex_error::vortex_ensure_eq; use vortex_mask::Mask; use vortex_session::VortexSession; -use super::batch::Batch; -use super::batch::BorrowedExecutionArgs; +use super::batch::BorrowedRowFnArgs; +use super::batch::RowFnExecutionArgs; use super::row_fn::RowFn; use super::visitor::BatchPlanner; use super::visitor::ExecuteRows; @@ -139,7 +139,7 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { fn execute_row_kernel( function: &F, options: &F::Options, - args: BorrowedExecutionArgs<'_>, + args: BorrowedRowFnArgs<'_>, ctx: &mut ExecutionCtx, ) -> VortexResult { function.dispatch( @@ -159,7 +159,7 @@ fn execute_row_kernel( fn try_execute_valid_rows( function: &F, options: &F::Options, - args: BorrowedExecutionArgs<'_>, + args: BorrowedRowFnArgs<'_>, valid: &Mask, ctx: &mut ExecutionCtx, ) -> VortexResult> { @@ -182,8 +182,8 @@ fn prepare_batch( function: &F, options: &F::Options, args: &dyn ExecutionArgs, -) -> VortexResult { - Batch::new(RowFn::id(function), args, |arg_dtypes| { +) -> VortexResult { + RowFnExecutionArgs::new(RowFn::id(function), args, |arg_dtypes| { function.dispatch( options, arg_dtypes,