Skip to content

Tracking Issue: Define the RowFn API #9129

Description

@connortsui20

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

Parent Epic: #9128

Design

RowFn describes a strict scalar function through the typed rows it reads and the output it writes. Every implementation receives the standard ScalarFnVTable automatically. A type that needs custom vtable hooks can keep them on a separate public type and delegate to a private RowFn kernel through row_fn_return_dtype and execute_rows.

pub trait RowFn: 'static + Sized + Clone + Send + Sync {
    type Options: 'static + Send + Sync + Clone + Debug + Display + PartialEq + Eq + Hash;

    /// The arguments in display order. Its length is the exact arity.
    const ARG_NAMES: &'static [&'static str];

    /// Whether every legal dispatch and encoding-aware execution is infallible.
    const INFALLIBLE: bool;

    fn id(&self) -> ScalarFnId;

    /// Defaults to a non-serializable function.
    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: &Self::Options,
        args: &[DType],
        visitor: V,
    ) -> VortexResult<V::VisitResult>;

    // Optionally bypass the row loop when an encoding already has a better answer.
    fn reduce_encoded(
        &self,
        options: &Self::Options,
        args: &[ArrayRef],
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<ArrayRef>> {
        _ = (options, args, ctx);
        Ok(None)
    }
}

There is no argument or return witness. ARG_NAMES supplies the exact arity. The output type selected by dispatch supplies the result dtype. Planning reads dense safety and decode infallibility from the concrete argument tuple.

INFALLIBLE stays function-wide because is_fallible has no input dtypes. Every fallible decode or row result selected by dispatch requires INFALLIBLE = false. The same rule covers semantic errors from reduce_encoded.

Options persistence also belongs to the function. The default is non-serializable. Registered functions can define their own wire representation through serialize and deserialize.

RowFn has strict semantics. A null in any input produces a null output. The current output forms are total over valid inputs. Their output validity is the conjunction of the input validities. This excludes functions such as list_sum and variant_get, which can return null for valid inputs.

reduce_encoded is the escape hatch for an encoding with a better bulk answer. Its result must match the planned dtype and row count. It cannot introduce nulls outside the rows that lifting masks. Its errors are immediately user-visible, so like a dense row closure it must be total over stored payloads behind null rows.

RowVisitor: select element types and an output form

RowVisitor<Options> is sealed and has three method families:

  • visit and visit_prepared return one owned output value per row.
  • visit_into and visit_prepared_into write through an OutputSink<Options>.
  • visit_deferred and visit_prepared_deferred return owned values with OR-reduced failure evidence.

The prepared forms receive each batch-constant argument as Some(value). A non-constant argument appears as None. This lets a kernel move constant-only work out of the row loop.

Owned-output visits require IndexedElementTuple, which every sealed ElementTuple implements. Sink visits accept ElementTuple directly.

InputElement and ElementTuple: validate, decode, and read input rows
pub unsafe trait InputElement: 'static {
    type Column;
    type View<'a>: ViewLen;
    type Elem<'a>;

    const DENSE_SAFE: bool;
    const DECODE_INFALLIBLE: bool;

    fn validate(dtype: &DType) -> VortexResult<()>;
    fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column>;
    fn can_decode_null_tolerant(array: &ArrayRef) -> VortexResult<bool>;
    fn decode_null_tolerant(
        array: ArrayRef,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<Self::Column>>;
    fn get(column: &Self::Column, index: usize) -> Self::Elem<'_>;
    fn view(column: &Self::Column) -> Self::View<'_>;

    fn get_from_view<'a>(column: &Self::View<'a>, index: usize) -> Self::Elem<'a>;
    unsafe fn get_from_view_unchecked<'a>(
        column: &Self::View<'a>,
        index: usize,
    ) -> Self::Elem<'a>;
}

InputElement is unsafe because the shared indexed loop performs one length check before unchecked row reads. Implementations must make every index below ViewLen::len legal for get_from_view_unchecked.

ElementTuple is a sealed adapter over tuples of zero through twelve InputElements. It combines dense safety and decode infallibility. Its Views<'a> associated type implements ViewLen and exists only when every argument is non-constant. The tuple length validates every borrowed view before unchecked traversal.

Columns cannot implement ViewLen because a batch constant physically contains one decoded row while logically addressing every batch row. decoded_lens_match validates the non-constant columns and exempts batch constants after their one-row representation passes validation. New decode primitives implement InputElement, not ElementTuple.

OutputSink and SinkResult: build the result column
pub unsafe trait OutputSink<Options>: 'static + Sized {
    type Rows<'a>: ViewLen where Self: 'a;
    type Row<'a> where Self: 'a;
    type WriteToken: 'static;

    fn skipped_rows_initializer() -> Option<for<'a> fn(&mut Self::Rows<'a>)> {
        None
    }

    /// Must be non-nullable. The lifting derives nullability from the inputs.
    fn return_dtype(options: &Options) -> VortexResult<DType>;
    fn with_capacity(rows: usize) -> VortexResult<Self>;
    fn rows(&mut self) -> Self::Rows<'_>;

    unsafe fn row_unchecked<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>;
    unsafe fn finish(self) -> VortexResult<ArrayRef>;
}

return_dtype receives function options. This supports option-dependent output dtypes, including decimal numeric operations.

The trait is unsafe because the executor trusts its row handles, write tokens, and finish implementation. A sink and every borrowed Rows view must also remain safe to drop after any callback prefix. This includes errors and unwinds during decoding, preparation, skipped-row initialization, and row execution.

SinkResult is a sealed adapter for (), InitializedElement, VortexResult<()>, and VortexResult<InitializedElement>. The result's WriteToken must match the sink's token. InitializedElement::write is unsafe because its zero-sized token cannot carry a lifetime that brands one row allocation.

Deferred failure evidence belongs to the owned visit_deferred methods. The evidence must not be wider than the output value. A wider reduction can reduce vector width. The accumulator must stay local to the loop because sink storage adds a loop-carried memory dependency.

Example: Hypot

Primitive types implement InputElement and OutputElement, so Hypot defines neither. Its complete row implementation is:

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: &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. A function needs a new element implementation only for a row representation that the framework cannot read or build.

Example: CosineSimilarity

CosineSimilarity is still a row computation. A constant operand has one norm for the batch, so the prepare step computes it once.

fn dispatch<V: RowVisitor<Self::Options>>(
    &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),
        )
    })
}

The dispatch selects TensorRow<f16>, TensorRow<f32>, or TensorRow<f64>. Planning derives the argument properties from the selected tuple.

Steps

  • Stabilize RowFn, ARG_NAMES, and the visitor methods with representative users.
  • Stabilize the input, output, and sink contracts with conformance tests for extension types.
  • Keep function-wide fallibility consistent with every dispatch and encoding-aware result.
  • Decide whether nullable row outputs are part of the initial API.
  • Preserve serialized metadata and source compatibility for migrated scalar functions.
  • Document when to use RowFn, when to keep a custom ScalarFnVTable, and how to add an element type.
  • Stabilize the public API.

Unresolved questions

  • Should nullable Option<T> outputs be supported before stabilization, or remain an additive follow-up?
  • InputElement, OutputElement, and OutputSink are downstream extension points. ElementTuple and SinkResult remain sealed.
  • OutputSink::return_dtype receives RowFn::Options. Decimal output dtypes can depend on the numeric operator.
  • Every RowFn receives the standard ScalarFnVTable implementation. A type that needs custom vtable hooks can delegate through row_fn_return_dtype and execute_rows from a separate public vtable type.
  • Unsafe InputElement and OutputSink contracts cover unchecked indexed reads, partial-drop safety, and sink finalization.

Implementation history

Metadata

Metadata

Assignees

Labels

tracking-issueShared implementation context for work likely to span multiple PRs.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions