From d56d815b9e1896611ef7d869be6df9394608ab6a Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Tue, 18 Aug 2026 16:40:15 -0400 Subject: [PATCH] Execute owned RowFn outputs over valid rows Signed-off-by: Connor Tsui --- .../src/scalar_fn/unstable/row/batch/tests.rs | 50 +++++++++ .../src/scalar_fn/unstable/row/execute/mod.rs | 2 + .../scalar_fn/unstable/row/execute/owned.rs | 101 ++++++++++++++++++ .../unstable/row/types/element/output.rs | 5 +- .../scalar_fn/unstable/row/visitor/execute.rs | 29 +++-- 5 files changed, 177 insertions(+), 10 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs index 101b6932737..dfa83cc7017 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/tests.rs @@ -15,6 +15,7 @@ use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::ConstantArray; use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::NativePType; use crate::dtype::Nullability; @@ -29,6 +30,9 @@ use crate::scalar_fn::unstable::row::RowVisitor; use crate::scalar_fn::unstable::row::execute_rows; use crate::validity::Validity; +#[derive(Clone)] +struct DeferredAdd; + #[derive(Clone)] struct ValidOnlyIdentity; @@ -36,6 +40,7 @@ struct ValidOnlyIdentity; struct InvalidKernelOutput; /// Produces a null row to exercise output validation at the row-function boundary. +#[derive(Default)] struct NullProducingI64(i64); impl OutputElement for NullProducingI64 { @@ -82,6 +87,36 @@ unsafe impl OutputSink for I64Sink { } } +impl RowFn for DeferredAdd { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"]; + const INFALLIBLE: bool = false; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.deferred_add"); + *ID + } + + fn dispatch>( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_deferred::<(i64, i64), i64, bool>( + |(lhs, rhs)| lhs.overflowing_add(rhs), + |overflowed| { + if overflowed { + vortex_bail!(InvalidArgument: "deferred addition overflowed"); + } + + Ok(()) + }, + ) + } +} + impl RowFn for ValidOnlyIdentity { type Options = EmptyOptions; @@ -173,6 +208,21 @@ fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> { Ok(()) } +#[test] +fn test_deferred_owned_execution_skips_invalid_rows() -> VortexResult<()> { + let validity = Validity::from_iter([true, false]); + let lhs = PrimitiveArray::new(vec![1_i64, i64::MAX], validity.clone()).into_array(); + let rhs = ConstantArray::new(1_i64, 2).into_array(); + let args = VecExecutionArgs::new(vec![lhs, rhs], 2); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&DeferredAdd, &EmptyOptions, &args, &mut ctx)?; + let expected = PrimitiveArray::new(vec![2_i64, 0], validity).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(); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs index b2f74a515aa..8adecd4206d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/mod.rs @@ -9,6 +9,8 @@ mod owned; pub(super) use owned::execute_owned; pub(super) use owned::execute_owned_infallible; +pub(super) use owned::execute_owned_infallible_valid_rows; +pub(super) use owned::execute_owned_valid_rows; mod sink; pub(super) use sink::execute_sink; diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs index cf459f14148..ece0a5c371a 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/owned.rs @@ -11,7 +11,11 @@ use std::ops::BitOrAssign; use vortex_compute::lane_kernels::IndexedSourceExt; 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 crate::ArrayRef; use crate::ExecutionCtx; @@ -49,6 +53,103 @@ where ) } +/// Decode nullable inputs, then store one output for each valid row from an infallible kernel. +pub(crate) fn execute_owned_infallible_valid_rows( + args: &dyn ExecutionArgs, + valid: &Mask, + ctx: &mut ExecutionCtx, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, +) -> VortexResult> +where + Args: IndexedElementTuple, + Out: OutputElement, +{ + execute_owned_valid_rows::( + args, + valid, + ctx, + prepare, + move |prepared, args| (apply(prepared, args), NoFailure), + |_| Ok(()), + ) +} + +/// Decode nullable inputs, then store outputs and combine failure evidence for valid rows. +pub(crate) fn execute_owned_valid_rows( + args: &dyn ExecutionArgs, + valid: &Mask, + 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: FailureEvidence, +{ + const { assert_owned_output_needs_no_drop::() }; + + let Some(columns) = Args::decode_null_tolerant(args, ctx)? else { + return Ok(None); + }; + + let row_count = args.row_count(); + let AllOr::Some(valid_rows) = valid.bit_buffer() else { + vortex_bail!( + "execute_owned_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(), + ); + + let prepared = prepare(Args::const_values(&columns)); + let mut values: Vec = std::iter::repeat_with(Out::default) + .take(row_count) + .collect(); + let mut failure = Fail::default(); + + if let Some(views) = Args::views_if_no_consts(&columns) { + vortex_ensure!( + Args::view_lens_match(&views, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + valid_rows.for_each_set_index(|index| { + // SAFETY: the tuple-wide length check proved every view has `row_count` rows, and mask + // indices are below `row_count`. Nullary tuples do not access an input view. + let elements = unsafe { Args::get_from_views_unchecked(&views, index) }; + let (value, row_failure) = apply(&prepared, elements); + + // SAFETY: the mask length check proved that every set index is below `row_count`. + unsafe { *values.get_unchecked_mut(index) = value }; + failure |= row_failure; + }); + } else { + vortex_ensure!( + Args::decoded_lens_match(&columns, row_count), + "a decoded row input does not address exactly {row_count} rows", + ); + + valid_rows.for_each_set_index(|index| { + let (value, row_failure) = apply(&prepared, Args::get(&columns, index)); + + // SAFETY: the mask length check proved that every set index is below `row_count`. + unsafe { *values.get_unchecked_mut(index) = value }; + failure |= row_failure; + }); + } + + finish_failure(failure)?; + + Ok(Some(Out::build(values))) +} + /// Decode every input column, then store outputs and combine per-row failure evidence. pub(crate) fn execute_owned( args: &dyn ExecutionArgs, 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 f7104f47650..1034bd46498 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 @@ -9,7 +9,10 @@ use crate::ArrayRef; use crate::dtype::DType; /// An owned row value that can be built into an all-valid column. -pub trait OutputElement: 'static + Sized { +/// +/// Skip-invalid execution uses [`Default`] only as a placeholder for invalid rows. Batch execution +/// masks those rows before returning the output. +pub trait OutputElement: 'static + Sized + Default { /// The dtype of columns built from this element type. **Must** be non-nullable: nullability is /// derived from the inputs by batch execution. /// 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 345d26a2532..7016aedaf23 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -32,6 +32,8 @@ use crate::scalar_fn::unstable::row::RowFn; use crate::scalar_fn::unstable::row::SinkResult; 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_owned_infallible_valid_rows; +use crate::scalar_fn::unstable::row::execute::execute_owned_valid_rows; use crate::scalar_fn::unstable::row::execute::execute_sink; use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; @@ -159,8 +161,8 @@ 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 rejects that unsupported signature. +/// Owned outputs initialize skipped positions with [`Default::default`]. Output sinks use +/// their own skipped-row initializer. pub(crate) struct ExecuteValidRows<'args, 'ctx, F: RowFn> { /// The original inputs for this kernel invocation. args: &'args dyn ExecutionArgs, @@ -213,8 +215,8 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { fn visit_prepared( self, - _prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, - _apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out, ) -> VortexResult where Args: IndexedElementTuple, @@ -228,7 +230,9 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { RowPolicy::for_owned_output::(), )?; - Ok(None) + execute_owned_infallible_valid_rows::( + self.args, self.valid, self.ctx, prepare, apply, + ) } fn visit_prepared_into( @@ -260,9 +264,9 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { 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<()>, + prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, + apply: impl Fn(&Prepared, Args::Elems<'_>) -> (Out, Fail), + finish_failure: impl FnOnce(Fail) -> VortexResult<()>, ) -> VortexResult where Args: IndexedElementTuple, @@ -277,7 +281,14 @@ impl RowVisitor for ExecuteValidRows<'_, '_, F> { RowPolicy::for_deferred_output::(), )?; - Ok(None) + execute_owned_valid_rows::( + self.args, + self.valid, + self.ctx, + prepare, + apply, + finish_failure, + ) } }