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 any legal dispatch or encoding-aware execution can raise a semantic error.
    const FALLIBLE: bool = false;

    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<RowExecution>> {
        _ = (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 fallibility from the concrete argument tuple.

FALLIBLE stays function-wide because is_fallible has no input dtypes. Every fallible decode or row result selected by dispatch requires FALLIBLE = true. The same rule covers semantic errors and deferred failures 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. It can also return deferred failure evidence when null payloads can cause a false error.

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 per-row 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>;
    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 view(column: &Self::Column) -> Self::View<'_>;
    fn view_len(column: &Self::View<'_>) -> usize;
    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 view_len legal for get_from_view_unchecked.

ElementTuple is a sealed adapter over tuples of one through twelve InputElements. It combines dense safety and decode fallibility. It also collapses a batch-constant operand to one decoded row. New decode primitives implement InputElement, not ElementTuple.

OutputSink and SinkResult: build the result column
pub unsafe trait OutputSink<Options>: 'static + Sized {
    type Rows<'a> 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 sink_dtype(options: &Options, 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 row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>;
    unsafe fn finish(self) -> VortexResult<ArrayRef>;
}

sink_dtype receives function options and input dtypes. 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::sink_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