Tracking Issue: Define the `RowFn` API

[connortsui20]

This issue tracks the author-facing RowFn API in vortex-array.

Parent Epic: #9128

Design

RowFn describes a strict scalar function as a typed operation on one row. A blanket ScalarFnVTable implementation adds columnar execution. Authors do not implement another vtable.

// Default method bodies omitted.
pub trait RowFn: 'static + Sized + Clone + Send + Sync {
    type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash;

    const ARG_NAMES: &'static [&'static str];
    const FALLIBLE: bool = false;

    fn id(&self) -> ScalarFnId;
    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>>;
    fn deserialize(
        &self,
        metadata: &[u8],
        session: &VortexSession,
    ) -> VortexResult<Self::Options>;

    fn dispatch<V: RowVisitor>(
        &self,
        options: &Self::Options,
        args: &[DType],
        visitor: V,
    ) -> VortexResult<V::VisitResult>;

    fn reduce_encoded(
        &self,
        options: &Self::Options,
        args: &[ArrayRef],
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<ArrayRef>>;
}

ARG_NAMES defines the exact arity. FALLIBLE is function-wide because ScalarFnVTable::is_fallible has no input dtypes. A function can set it when only some dtype choices fail. The visitor checks each concrete dispatch against it at compile time.

dispatch validates cross-argument dtype rules and selects concrete row types. Planning and execution both call it. Its decision must depend only on options and args.

reduce_encoded is an optional bulk shortcut. It runs before row decoding when that execution path probes the given inputs. It can preserve an encoding or return a lazy result. Its output must:

• have the same row count as the invocation.
• match the planned dtype, ignoring nullability.
• not add a null where all inputs are valid.

The framework skips this hook for nullary functions.

RowVisitor: select an output form and optional prepared state

// Default method bodies omitted.
pub trait RowVisitor: Sealed + Sized {
    type VisitResult;

    fn visit<Args, Out>(
        self,
        apply: impl Fn(Args::Elems<'_>) -> Out,
    ) -> VortexResult<Self::VisitResult>
    where
        Args: IndexedElementTuple,
        Out: OutputElement;

    fn visit_prepared<Args, Out, Prepared>(
        self,
        prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
        apply: impl Fn(&Prepared, Args::Elems<'_>) -> Out,
    ) -> VortexResult<Self::VisitResult>
    where
        Args: IndexedElementTuple,
        Out: OutputElement;

    fn visit_into<Args, Sink, ApplyResult>(
        self,
        apply: impl Fn(Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult,
    ) -> VortexResult<Self::VisitResult>
    where
        Args: ElementTuple,
        Sink: OutputSink,
        ApplyResult: SinkResult<WriteToken = Sink::WriteToken>;

    fn visit_prepared_into<Args, Sink, Prepared, ApplyResult>(
        self,
        prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared,
        apply: impl Fn(&Prepared, Args::Elems<'_>, Sink::Row<'_>) -> ApplyResult,
    ) -> VortexResult<Self::VisitResult>
    where
        Args: ElementTuple,
        Sink: OutputSink,
        ApplyResult: SinkResult<WriteToken = Sink::WriteToken>;

    fn visit_deferred<Args, Out, Fail>(
        self,
        apply: impl Fn(Args::Elems<'_>) -> (Out, Fail),
        finish_failure: impl FnOnce(Fail) -> VortexResult<()>,
    ) -> VortexResult<Self::VisitResult>
    where
        Args: IndexedElementTuple,
        Out: OutputElement,
        Fail: Copy + Default + BitOrAssign;

    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<()>,
    ) -> VortexResult<Self::VisitResult>
    where
        Args: IndexedElementTuple,
        Out: OutputElement,
        Fail: Copy + Default + BitOrAssign;
}

The three unprepared methods delegate to their prepared forms with unit state. A function that has nothing to prepare does not need an empty prepare closure.

RowVisitor::VisitResult is the framework result. It is a plan during dtype planning and a batch execution result at run time. By contrast, the Out generic is one owned value returned for one row.

Prepared visits receive Args::ConstElems. Each entry is Some for a batch-constant argument and None for a varying argument. This lets a function compute constant-dependent state once per batch.

Deferred visits OR-reduce Fail across the row loop. Fail::default() must mean success. Fail must not be wider than Out, so failure tracking does not reduce the vector width. The framework checks the size rule at compile time.

InputElement and OutputElement: decode and build row values

// Default method bodies omitted.
pub trait InputElement: 'static {
    type Column;
    type Varying<'a>;
    type Elem<'a>;

    const DENSE_SAFE: bool = false;
    const DECODE_FALLIBLE: bool = true;

    fn validate(dtype: &DType) -> VortexResult<()>;
    fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column>;

    fn decode_null_tolerant(
        array: ArrayRef,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<Self::Column>>;

    fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>;
    fn varying(column: &Self::Column) -> Self::Varying<'_>;
    fn varying_len(column: &Self::Varying<'_>) -> usize;

    fn get_varying<'a>(
        column: &Self::Varying<'a>,
        index: usize,
    ) -> Self::Elem<'a>
    where
        Self: 'a;
}

pub trait OutputElement: 'static + Sized {
    fn element_dtype() -> DType;
    fn build(values: Vec<Self>) -> ArrayRef;
}

InputElement decodes one column into an O(1) row representation. It declares whether dense access is safe and whether legal input can fail to decode.

OutputElement builds a column from independent owned row values. Its dtype is fixed by its Rust type. Runtime-shaped output uses OutputSink, whose sink_dtype can inspect the input dtypes.

ElementTuple and IndexedElementTuple are sealed adapters over tuples of InputElements. Function authors select a tuple through RowVisitor but do not implement either adapter.

OutputSink, SinkResult, DeferredError, and InitializedElement: write custom outputs and report errors

// Default method bodies omitted.
pub trait OutputSink: 'static + Sized {
    const ERRORS_ARE_DEFERRED: bool = false;
    const SUPPORTS_SKIPPED_ROWS: bool = false;

    type Rows<'a>
    where
        Self: 'a;

    type Row<'a>
    where
        Self: 'a;

    type WriteToken: 'static;

    fn sink_dtype(args: &[DType]) -> VortexResult<DType>;
    fn with_capacity(rows: usize, dtype: &DType) -> VortexResult<Self>;
    fn rows(&mut self) -> Self::Rows<'_>;
    fn row_count_matches(rows: &Self::Rows<'_>, row_count: usize) -> bool;
    fn initialize_skipped_rows(rows: &mut Self::Rows<'_>);
    fn row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>;
    fn finish(self, error: DeferredError) -> VortexResult<ArrayRef>;
}

pub trait SinkResult: 'static + Sealed {
    type WriteToken: 'static;
    type Accumulated: 'static + Copy + Default;

    const FALLIBLE: bool;
    const DEFERRED: bool;

    fn accumulate(self, accumulated: &mut Self::Accumulated) -> VortexResult<()>;
    fn occurred(accumulated: Self::Accumulated) -> bool;
}

pub struct DeferredError(/* private */ bool);

pub struct InitializedElement(/* private */ ());

impl InitializedElement {
    pub fn write<T>(row: &mut MaybeUninit<T>, value: T) -> Self;
}

An OutputSink can own shared batch state or provide a custom row handle. A sink can opt into deferred errors and skipped-row execution. The framework masks any placeholder values left in skipped output rows.

OutputSink::WriteToken proves that a successful closure initialized its row handle. Initialized sinks use (). UninitElementSink uses InitializedElement, which safe code can obtain only through InitializedElement::write.

SinkResult is sealed, and its WriteToken must match the sink. A sink-writing closure returns (), VortexResult<()>, VortexResult<InitializedElement>, or a Boolean or unsigned failure word. The executor keeps deferred evidence in a loop-local and passes one DeferredError to the sink after the loop.

Example: Hypot

Primitive input and output elements are already available, so Hypot only defines its row operation:

#[derive(Clone)]
struct Hypot;

impl RowFn for Hypot {
    type Options = EmptyOptions;

    const ARG_NAMES: &'static [&'static str] = &["x", "y"];

    fn id(&self) -> ScalarFnId {
        static ID: CachedId = CachedId::new("vortex.hypot");
        *ID
    }

    fn dispatch<V: RowVisitor>(
        &self,
        _options: &Self::Options,
        _args: &[DType],
        visitor: V,
    ) -> VortexResult<V::VisitResult> {
        visitor.visit::<(f64, f64), f64>(|(x, y)| x.hypot(y))
    }
}

The framework derives type validation, batch decoding, constant handling, output allocation, null handling, and validity from this definition.

Example: CosineSimilarity

CosineSimilarity is also a row computation, but a constant operand has one norm for the whole batch. The prepared visit computes that norm once.

struct ConstNorms<T> {
    lhs: Option<T>,
    rhs: Option<T>,
}

fn dispatch<V: RowVisitor>(
    &self,
    _options: &Self::Options,
    args: &[DType],
    visitor: V,
) -> VortexResult<V::VisitResult> {
    match_each_float_ptype!(tensor_element_ptype(args)?, |T| {
        visitor.visit_prepared::<(TensorRow<T>, TensorRow<T>), T, _>(
            |(lhs, rhs)| ConstNorms {
                lhs: lhs.map(l2_norm_row),
                rhs: rhs.map(l2_norm_row),
            },
            |norms, (lhs, rhs)| {
                cosine_similarity_row_prepared(norms, lhs, rhs)
            },
        )
    })
}

TensorRow<T> implements InputElement and exposes each tensor row as &[T]. The dispatch selects TensorRow<f16>, TensorRow<f32>, or TensorRow<f64> from the input dtype.

Since I was the one who implemented cosine similarity in vortex-tensor, I can confidently say that this is SIGNIFICANTLY less complex and easier to write than the implementation on develop. At the very least, it is much harder to write incorrect code like this.

Current scope

RowFn is only for strict functions. A null input makes the corresponding output null. The current row output forms cannot produce a new null from otherwise valid inputs.

Implement ScalarFnVTable directly for columnar or zero-copy kernels, shared cross-row state, or functions with non-strict validity rules.

Steps

• Define RowFn and its blanket ScalarFnVTable implementation.
• Add owned, prepared, sink, and deferred visitor forms.
• Add compile-time dispatch contract checks.
• Define the input element, output element, and output sink extension points.
• Validate the API with primitive numeric arithmetic.
• Add conformance tests for downstream element and sink implementations.
• Decide whether nullable row outputs belong in the initial API.
• Document when to use RowFn and how to add an element or sink type.
• Stabilize the public API.

Unresolved questions

• Should Option<T> row outputs be supported before stabilization?
• Does OutputSink::sink_dtype need function options?

Implementation history

The current branch separates the framework, primitive numeric integration, and focused executor benchmarks into reviewable commits.