Execute owned RowFn outputs over valid rows

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,13 +30,17 @@ 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;
 
 #[derive(Clone)]
 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<Options> OutputSink<Options> 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<V: RowVisitor<Self::Options>>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        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::<i64>()).into_array();

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;

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, Out, Prepared>(
+    args: &dyn ExecutionArgs,
+    valid: &Mask,
+    ctx: &mut ExecutionCtx,
+    prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
+    apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out,
+) -> VortexResult<Option<ArrayRef>>
+where
+    Args: IndexedElementTuple,
+    Out: OutputElement,
+{
+    execute_owned_valid_rows::<Args, Out, Prepared, NoFailure>(
+        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, Out, Prepared, Fail>(
+    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<Option<ArrayRef>>
+where
+    Args: IndexedElementTuple,
+    Out: OutputElement,
+    Fail: FailureEvidence,
+{
+    const { assert_owned_output_needs_no_drop::<Out>() };
+
+    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<Out> = 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, Out, Prepared, Fail>(
     args: &dyn ExecutionArgs,

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.
     ///

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<F: RowFn> RowVisitor<F::Options> 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<F: RowFn> RowVisitor<F::Options> for ExecuteValidRows<'_, '_, F> {
 
     fn visit_prepared<Args, Out, 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<Self::VisitResult>
     where
         Args: IndexedElementTuple,
@@ -228,7 +230,9 @@ impl<F: RowFn> RowVisitor<F::Options> for ExecuteValidRows<'_, '_, F> {
             RowPolicy::for_owned_output::<Args>(),
         )?;
 
-        Ok(None)
+        execute_owned_infallible_valid_rows::<Args, Out, Prepared>(
+            self.args, self.valid, self.ctx, prepare, apply,
+        )
     }
 
     fn visit_prepared_into<Args, Sink, Prepared, ApplyResult>(
@@ -260,9 +264,9 @@ impl<F: RowFn> RowVisitor<F::Options> for ExecuteValidRows<'_, '_, F> {
 
     fn visit_prepared_deferred<Args, Out, Prepared, Fail>(
         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<Self::VisitResult>
     where
         Args: IndexedElementTuple,
@@ -277,7 +281,14 @@ impl<F: RowFn> RowVisitor<F::Options> for ExecuteValidRows<'_, '_, F> {
             RowPolicy::for_deferred_output::<Args>(),
         )?;
 
-        Ok(None)
+        execute_owned_valid_rows::<Args, Out, Prepared, Fail>(
+            self.args,
+            self.valid,
+            self.ctx,
+            prepare,
+            apply,
+            finish_failure,
+        )
     }
 }