Skip to content

Tracking Issue: Define the RowFn API #9129

Description

@connortsui20

This is a tracking issue for defining and stabilizing the author-facing RowFn API in vortex-array.

Parent Epic: #9128

Design

RowFn describes a strict scalar function in terms of the typed rows it reads and the output slot it writes. A blanket implementation exposes it as a ScalarFnVTable, so there is no intermediate public vtable.

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 can fail while decoding or computing a row.
    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,
        args: &[DType],
        visitor: V,
    ) -> VortexResult<V::Out>;

    // 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, and the output sink chosen by dispatch supplies the result dtype. Planning reads dense safety, decode fallibility, and filtered-decode cost from the concrete argument tuple selected by dispatch.

FALLIBLE remains function-wide because is_fallible is queried without input dtypes. Every fallible decode or row result selected by dispatch requires FALLIBLE = true. Options persistence is also function-owned. The default is non-serializable, and registered functions opt into their own wire representation through serialize and deserialize.

RowFn has strict semantics: a null in any input produces a null output. The row output forms are currently also total over valid inputs, so their output validity is exactly the conjunction of the input validities. This excludes strict functions such as list_sum and variant_get, which can produce null from valid inputs.

reduce_encoded is the escape hatch for an encoding with a better bulk answer. Its result must match the declared dtype and row count, and it cannot introduce nulls outside the rows the lifting will mask. The row count it sees is the original one under dense and branch-and-skip execution, and the filtered one under filter-and-scatter.

RowVisitor: choose element types, prepare batch state, write into a sink
pub trait RowVisitor: Sealed {
    type Out;

    fn visit_prepared_into<A: ElementTuple, S: OutputSink, P, R: SinkResult>(
        self,
        prepare: impl FnOnce(A::ConstElems<'_>) -> P,
        apply: impl Fn(&P, A::Elems<'_>, S::Row<'_>) -> R,
    ) -> VortexResult<Self::Out>;
}

One method, rather than the separate ordinary, prepared, and sink-backed visits considered earlier: those were three spellings of the same executor. prepare receives the element value of every batch-constant argument and None for each one that varies by row, so a kernel can hoist work depending only on a constant operand. Passing |_| () is the unprepared case.

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

    const DENSE_SAFE: bool = false;
    const DECODE_FALLIBLE: bool = true;
    const FILTERED_DECODE_COST: usize = 0;

    fn validate(dtype: &DType) -> VortexResult<()>;

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

    /// `Ok(None)` when this element cannot read this array without assuming every row is valid.
    fn decode_null_tolerant(
        array: ArrayRef,
        ctx: &mut ExecutionCtx,
    ) -> VortexResult<Option<Self::Column>> {
        if Self::DENSE_SAFE {
            Self::decode(array, ctx).map(Some)
        } else {
            Ok(None)
        }
    }

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

ElementTuple is a sealed adapter over tuples of one through twelve InputElements. It ANDs DENSE_SAFE, ORs DECODE_FALLIBLE, adds FILTERED_DECODE_COST, and collapses a batch-constant operand to one decoded row read at stride zero. New decode primitives implement InputElement, not ElementTuple.

OutputSink and SinkResult: build the result column and report failure
pub trait OutputSink: 'static + Sized {
    /// Whether the row closure reports failure through `DeferredError` rather than `VortexResult`.
    const ERRORS_ARE_DEFERRED: bool = false;

    /// Whether this sink can finish a full-length output when some rows were never visited.
    const SUPPORTS_SKIPPED_ROWS: bool = false;

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

    /// Must be non-nullable. The lifting derives nullability from the inputs.
    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 row<'a>(rows: &'a mut Self::Rows<'_>, index: usize) -> Self::Row<'a>;

    fn finish(self, error: DeferredError) -> VortexResult<ArrayRef>;
}

sink_dtype sees the input dtypes, which lets an output whose width is runtime data fit here. ElementSink<T> adapts an OutputElement and passes the closure an &mut T. A custom sink can represent a structured output in one sink type with as many fields as it needs. There is no separate multi-sink adapter.

SinkResult is a sealed executor adapter. The row closure chooses one of the supplied forms: () for an infallible write, VortexResult<()> to stop at the first error, or an unsigned word for failure evidence OR-reduced across the batch. bool is the ordinary word. Custom output representation belongs in OutputSink.

pub trait SinkResult: 'static + Sealed {
    /// The word the executor OR-reduces in a loop-local.
    type Accumulated: 'static + Copy + Default;

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

The word is named rather than fixed because evidence sometimes has to be wider than a bit. Unsigned multiplication reports the discarded high half of its widened product, since deriving a bool from it costs a comparison that LLVM folds into llvm.umul.with.overflow, which has no vector lowering. The word must not exceed the element width, and the reduction must stay a loop-local: see #9130 for both constraints and what violating each one costs. The sink never sees the word. It is handed one DeferredError after the loop.

Example: Hypot

InputElement and OutputElement are implemented on the primitive types, so Hypot defines neither. Its complete RowFn implementation for two f64 columns 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,
        _args: &[DType],
        visitor: V,
    ) -> VortexResult<V::Out> {
        visitor.visit_prepared_into::<(f64, f64), ElementSink<f64>, _, _>(
            |_| (),
            |&(), (x, y), output| *output = x.hypot(y),
        )
    }
}

The framework derives type validation, batch decoding, constant handling, output allocation, null handling, and validity from this. A function needs a new element implementation only when it introduces a row representation the framework cannot already read or build.

Example: CosineSimilarity

CosineSimilarity is still a row computation, but a constant operand has one norm for the entire batch, which the prepare step computes once.

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

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

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 (https://github.com/vortex-data/vortex/blob/develop/vortex-tensor/src/scalar_fns/cosine_similarity.rs). At the very least, it is much harder to write incorrect code like this.

Steps

  • Stabilize RowFn, ARG_NAMES, and visit_prepared_into with representative users.
  • Stabilize the input, output, and sink contracts with conformance tests for extension types.
  • Keep function-wide fallibility consistent with every concrete dispatch.
  • Decide whether nullable row outputs are part of the initial API.
  • Preserve both the serialized metadata and the source compatibility of migrated scalar functions.
  • Document when to use RowFn, when to implement ScalarFnVTable directly, 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. InputElement defaults to DENSE_SAFE = false, DECODE_FALLIBLE = true, and FILTERED_DECODE_COST = 0. ElementTuple and SinkResult remain sealed.
  • Does OutputSink::sink_dtype need access to function options, or can option-dependent output dtypes remain unsupported initially?
  • ERRORS_ARE_DEFERRED and SUPPORTS_SKIPPED_ROWS are part of the open OutputSink contract.

Implementation history

None yet.

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