Add associated outputs to RowFn kernels

vortex-array/src/scalar_fn/unstable/row/execute/kernel.rs

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Execution for typed infallible row kernels.
+
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+use vortex_error::vortex_ensure_eq;
+use vortex_mask::MaskValuesRef;
+
+use crate::ArrayRef;
+use crate::ExecutionCtx;
+use crate::scalar_fn::ExecutionArgs;
+use crate::scalar_fn::unstable::row::DenseRows;
+use crate::scalar_fn::unstable::row::IndexedElementTuple;
+use crate::scalar_fn::unstable::row::RowKernel;
+use crate::scalar_fn::unstable::row::RowKernelOutput;
+use crate::scalar_fn::unstable::row::visitor::assert_owned_output_needs_no_drop;
+
+/// Decode and validate the inputs, then delegate only dense collection to the kernel.
+pub(in crate::scalar_fn::unstable::row) fn execute_kernel<Args, Kernel>(
+    args: &dyn ExecutionArgs,
+    ctx: &mut ExecutionCtx,
+    kernel: Kernel,
+) -> VortexResult<ArrayRef>
+where
+    Args: IndexedElementTuple,
+    Kernel: RowKernel<Args>,
+{
+    let inputs = Args::decode(args, ctx)?;
+    let rows = DenseRows::<Args>::new(&inputs, args.row_count())?;
+    let output = kernel.collect_dense(rows)?;
+
+    output.finish()
+}
+
+/// Decode nullable inputs, then store one kernel output for each valid row.
+pub(in crate::scalar_fn::unstable::row) fn execute_kernel_valid_rows<Args, Kernel>(
+    args: &dyn ExecutionArgs,
+    valid: &MaskValuesRef,
+    ctx: &mut ExecutionCtx,
+    kernel: Kernel,
+) -> VortexResult<Option<ArrayRef>>
+where
+    Args: IndexedElementTuple,
+    Kernel: RowKernel<Args>,
+{
+    const { assert_owned_output_needs_no_drop::<Kernel::Element>() };
+
+    let Some(columns) = Args::decode_null_tolerant(args, ctx)? else {
+        return Ok(None);
+    };
+
+    let row_count = args.row_count();
+    let valid_rows = valid.bit_buffer();
+    vortex_ensure_eq!(
+        valid_rows.len(),
+        row_count,
+        "the validity mask must address exactly {row_count} rows, got {}",
+        valid_rows.len(),
+    );
+
+    let mut values: Vec<Kernel::Element> = std::iter::repeat_with(Kernel::Element::default)
+        .take(row_count)
+        .collect();
+
+    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) };
+
+            // SAFETY: the mask length check proved that every set index is below `row_count`.
+            unsafe { *values.get_unchecked_mut(index) = kernel.eval(elements) };
+        });
+    } 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 = kernel.eval(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 };
+        });
+    }
+
+    let output = Kernel::Output::from_values(values)?;
+    output.finish().map(Some)
+}

vortex-array/src/scalar_fn/unstable/row/execute/mod.rs

@@ -6,6 +6,10 @@
 //! [`owned`] stores one independent value per row and reduces compact failure evidence. [`sink`]
 //! drives output builders whose row handles may share batch state.
 
+mod kernel;
+pub(super) use kernel::execute_kernel;
+pub(super) use kernel::execute_kernel_valid_rows;
+
 mod owned;
 pub(super) use owned::execute_owned;
 pub(super) use owned::execute_owned_infallible;

vortex-array/src/scalar_fn/unstable/row/mod.rs

@@ -9,7 +9,10 @@
 //! A [`RowFn`] describes the typed operation while the framework owns columnar concerns such as
 //! decoding, constant handling, null propagation, allocation, and validity. Its
 //! [`RowFn::dispatch`] implementation uses a [`RowVisitor`] to select an [`ElementTuple`] and
-//! either an [`OutputElement`] or [`OutputSink`] for each supported dtype combination.
+//! output contract for each supported dtype combination. An [`OutputElement`] returns one owned
+//! value per row, an [`OutputSink`] writes through a shared builder, and a [`RowKernel`] can attach
+//! an associated dense output representation while retaining scalar semantics for other execution
+//! policies.
 //!
 //! Unlike a general strict function, a [`RowFn`] cannot produce null from valid inputs.
 //!
@@ -29,15 +32,22 @@ mod row_fn;
 pub use row_fn::RowFn;
 
 mod types;
+pub use types::ArgColumn;
+pub use types::ArgView;
+pub use types::DenseRows;
 pub use types::ElementTuple;
 pub use types::FailureEvidence;
 pub use types::IndexedElementTuple;
 pub use types::InitializedElement;
 pub use types::InputElement;
 pub use types::OutputElement;
 pub use types::OutputSink;
+pub use types::PackedBoolOutput;
+pub use types::RowKernel;
+pub use types::RowKernelOutput;
 pub use types::SinkResult;
 pub use types::UninitElementSink;
+pub use types::VecOutput;
 pub use types::ViewLen;
 
 mod visitor;

vortex-array/src/scalar_fn/unstable/row/row_fn.rs

@@ -26,17 +26,20 @@ 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.
+/// A dispatched [`OutputElement`], [`OutputSink`], or [`RowKernel`] describes the non-nullable
+/// values produced for valid rows. A `RowKernel` also selects an associated output representation
+/// and can override dense collection without changing its scalar semantics. The framework widens
+/// the output 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
+/// [`RowKernel`]: crate::scalar_fn::unstable::row::RowKernel
 /// [`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
@@ -73,7 +76,8 @@ pub trait RowFn: 'static + Sized + Clone + Send + Sync {
         vortex_bail!("Expression {} is not deserializable", self.id())
     }
 
-    /// Choose element types for these input dtypes and visit the framework with them.
+    /// Choose element types and an output contract for these input dtypes, then visit the
+    /// framework with them.
     ///
     /// Planning and execution both call this method, so its result **must** depend only on
     /// `options` and `args`. Cross-argument dtype validation belongs here.

vortex-array/src/scalar_fn/unstable/row/types/element/input.rs

@@ -59,6 +59,18 @@ pub unsafe trait InputElement: 'static {
     /// and other invocation-invariant work into this method.
     fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column>;
 
+    /// Decode one representative row from an input that is constant for the batch.
+    ///
+    /// The default preserves the ordinary decode contract by slicing the input to one row first.
+    /// Implementations can override this when their constant representation supports cheaper
+    /// scalar extraction. The returned column **must** contain exactly one row.
+    fn decode_batch_constant(
+        array: ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Self::Column> {
+        Self::decode(array.slice(0..1)?, ctx)
+    }
+
     /// Whether [`decode_null_tolerant`](Self::decode_null_tolerant) can decode this array.
     ///
     /// The conservative default declines. An implementation whose ordinary decode is safe and

vortex-array/src/scalar_fn/unstable/row/types/element/mod.rs

@@ -18,6 +18,8 @@ pub use output::OutputElement;
 mod primitive;
 
 mod tuple;
+pub use tuple::ArgColumn;
+pub use tuple::ArgView;
 pub use tuple::ElementTuple;
 pub use tuple::IndexedElementTuple;
 pub use tuple::batch_const;

vortex-array/src/scalar_fn/unstable/row/types/element/primitive.rs

@@ -9,6 +9,7 @@ use vortex_error::vortex_ensure_eq;
 use crate::ArrayRef;
 use crate::ExecutionCtx;
 use crate::IntoArray;
+use crate::arrays::Constant;
 use crate::arrays::PrimitiveArray;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
@@ -44,6 +45,19 @@ unsafe impl<T: NativePType> InputElement for T {
         Ok(array.execute::<PrimitiveArray>(ctx)?.into_buffer::<T>())
     }
 
+    fn decode_batch_constant(
+        array: ArrayRef,
+        ctx: &mut ExecutionCtx,
+    ) -> VortexResult<Self::Column> {
+        if let Some(constant) = array.as_opt::<Constant>()
+            && let Some(value) = constant.scalar().as_primitive().try_typed_value::<T>()?
+        {
+            return Ok(Buffer::full(value, 1));
+        }
+
+        Self::decode(array.slice(0..1)?, ctx)
+    }
+
     fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
         Ok(true)
     }

vortex-array/src/scalar_fn/unstable/row/types/element/tuple/element_tuple.rs

@@ -35,6 +35,17 @@ enum ArgColumnKind<T: InputElement> {
     Const(T::Column),
 }
 
+/// A typed decoded input classified by how the dense collector addresses it.
+///
+/// This classification does not prove that the view covers a particular batch length.
+pub enum ArgView<'a, T: InputElement> {
+    /// A view with one value for every row in the batch.
+    Column(T::View<'a>),
+
+    /// A one-row view whose value is broadcast across the batch.
+    Constant(T::View<'a>),
+}
+
 impl<T: InputElement> ArgColumn<T> {
     fn try_from_const(column: T::Column) -> VortexResult<Self> {
         let decoded_len = T::view(&column).len();
@@ -52,7 +63,7 @@ impl<T: InputElement> ArgColumn<T> {
         if let Some(const_array) = batch_const(&array)
             && !array.is_empty()
         {
-            return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?);
+            return Self::try_from_const(T::decode_batch_constant(const_array, ctx)?);
         }
 
         Ok(Self(ArgColumnKind::Column(T::decode(array, ctx)?)))
@@ -64,7 +75,7 @@ impl<T: InputElement> ArgColumn<T> {
         if let Some(const_array) = batch_const(&array)
             && !array.is_empty()
         {
-            return Self::try_from_const(T::decode(const_array.slice(0..1)?, ctx)?).map(Some);
+            return Self::try_from_const(T::decode_batch_constant(const_array, ctx)?).map(Some);
         }
 
         Ok(T::decode_null_tolerant(array, ctx)?
@@ -96,6 +107,14 @@ impl<T: InputElement> ArgColumn<T> {
         }
     }
 
+    /// Borrow the decoded input with its batch-constant classification.
+    pub fn view(&self) -> ArgView<'_, T> {
+        match &self.0 {
+            ArgColumnKind::Column(column) => ArgView::Column(T::view(column)),
+            ArgColumnKind::Const(column) => ArgView::Constant(T::view(column)),
+        }
+    }
+
     fn addresses_rows(&self, row_count: usize) -> bool {
         // A constant is validated when constructed and is always read at index zero.
         match &self.0 {

vortex-array/src/scalar_fn/unstable/row/types/element/tuple/mod.rs

@@ -7,6 +7,8 @@
 //! [`IndexedElementTuple`] adds the validated indexed source used by vectorizable dense loops.
 
 mod element_tuple;
+pub use element_tuple::ArgColumn;
+pub use element_tuple::ArgView;
 pub use element_tuple::ElementTuple;
 pub use element_tuple::batch_const;