Function patterns

Each of the five VGI function shapes, with a complete, runnable worker for each — so you can find the shape that fits your problem. (Do the tutorial first.)

Which shape do I need?

Scalar1 row → 1 value

Tableargs → N rows

Table-in-outN rows → M rows

AggregateN rows → 1 value

Bufferingstream → [state] → stream

At a glance — pick the shape whose input/output cardinality matches your problem:

Pattern	Shape	Use it when…	SQL
Scalar	1 row → 1 value	you transform each row independently	`SELECT f(col) FROM t`
Table	args → N rows	you generate rows from arguments	`SELECT * FROM f(args)`
Table-in-out	N rows → M rows	you reshape/filter a streamed table	`SELECT * FROM f((SELECT …))`
Aggregate	grouped rows → 1 row/group	you accumulate per `GROUP BY` group	`SELECT f(col) FROM t GROUP BY k`
Buffering	stream → [state] → stream	you must see every row first (sort, top-k, full reduction)	`SELECT * FROM f((SELECT …))`

Table-in-out has one variant worth knowing about up front: blended, where the positional arguments are the input columns, so f(x, y) reads like an ordinary function call and works inside LATERAL. It’s a per-row map — no finalize — rather than a sixth shape.

Scalar

One row in, one row out. Operate on the whole column with pyarrow.compute; the Annotated types are the schema.

calc_scalar_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""Stage 1 of the tutorial: a worker with a single scalar function.

Run it from a DuckDB-compatible engine (Haybarn shown here)::

    uvx haybarn-cli
    ATTACH 'calc' (TYPE vgi, LOCATION 'uv run calc_scalar_worker.py');
    SELECT calc.double(21);   -- 42
"""

from typing import Annotated

import pyarrow as pa
import pyarrow.compute as pc

from vgi import Param, Returns, ScalarFunction, Worker
from vgi.catalog import Catalog, Schema


class Double(ScalarFunction):
    """Double each input value (one row in, one row out)."""

    @classmethod
    def compute(
        cls,
        value: Annotated[pa.Int64Array, Param(doc="Values to double")],
    ) -> Annotated[pa.Int64Array, Returns()]:
        """Multiply the whole column by 2."""
        return pc.multiply(value, 2)


class CalcWorker(Worker):
    """A worker exposing the ``calc`` catalog with one scalar function."""

    catalog = Catalog(
        name="calc",
        schemas=[Schema(name="main", functions=[Double])],
    )


if __name__ == "__main__":
    CalcWorker().run()

ATTACH 'calc' (TYPE vgi, LOCATION 'uv run calc_scalar_worker.py');
SELECT calc.double(n) FROM (VALUES (1), (2), (3)) AS t(n);

Input

n
1
2
3

Output

double(n)
2
4
6

Scalars aren’t just for numbers — any column type works. A string transform looks the same; here compute joins Hello, + name + ! across the column:

greeting_scalar_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""Stage 1 of the tutorial: a worker with a single scalar function.

Run it from a DuckDB-compatible engine (Haybarn shown here)::

    uvx haybarn-cli
    ATTACH 'greetings' (TYPE vgi, LOCATION 'uv run greeting_scalar_worker.py');
    SELECT greetings.greeting('Alice');
"""

from typing import Annotated

import pyarrow as pa
import pyarrow.compute as pc

from vgi import Param, Returns, ScalarFunction, Worker
from vgi.catalog import Catalog, Schema


class Greeting(ScalarFunction):
    """Return a friendly greeting for each name (one row in, one row out)."""

    @classmethod
    def compute(
        cls,
        name: Annotated[pa.StringArray, Param(doc="Column of names to greet")],
    ) -> Annotated[pa.StringArray, Returns()]:
        """Join ``Hello, `` + name + ``!`` element-wise across the column."""
        return pc.binary_join_element_wise("Hello, ", name, "!", "")


class GreetingWorker(Worker):
    """A worker exposing the ``greetings`` catalog with one scalar function."""

    catalog = Catalog(
        name="greetings",
        schemas=[Schema(name="main", functions=[Greeting])],
    )


if __name__ == "__main__":
    GreetingWorker().run()

SELECT greetings.greeting(name) FROM (VALUES ('Alice'), ('Bob')) AS t(name);

Input

name
Alice
Bob

Output

greeting(name)
Hello, Alice!
Hello, Bob!

Table

Generate rows from arguments, no input table. Declare a typed args dataclass and a FIXED_SCHEMA; process emits batches until out.finish(). The @bind_fixed_schema / @init_single_worker decorators wire up the common single-worker lifecycle. (This is the full tutorial worker — the scalar double plus the series generator.)

calc_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""The full tutorial worker: a scalar function and a table function.

Run from a DuckDB-compatible engine (Haybarn shown here)::

    uvx haybarn-cli
    ATTACH 'calc' (TYPE vgi, LOCATION 'uv run calc_worker.py');
    SELECT calc.double(21);            -- scalar -> 42
    SELECT * FROM calc.series(3);       -- table  -> 0, 1, 2
"""

from dataclasses import dataclass
from typing import Annotated, ClassVar

import pyarrow as pa
import pyarrow.compute as pc

from vgi import Arg, Param, Returns, ScalarFunction, Worker
from vgi.catalog import Catalog, Schema
from vgi.table_function import (
    OutputCollector,
    ProcessParams,
    TableFunctionGenerator,
    bind_fixed_schema,
    init_single_worker,
)


class Double(ScalarFunction):
    """Double each input value (one row in, one row out)."""

    @classmethod
    def compute(
        cls,
        value: Annotated[pa.Int64Array, Param(doc="Values to double")],
    ) -> Annotated[pa.Int64Array, Returns()]:
        """Multiply the whole column by 2."""
        return pc.multiply(value, 2)




@dataclass(slots=True, frozen=True, kw_only=True)
class SeriesArgs:
    """Arguments for :class:`Series` (one positional ``count``)."""

    count: Annotated[int, Arg(0, doc="How many numbers to generate", ge=0)]


@init_single_worker
@bind_fixed_schema
class Series(TableFunctionGenerator[SeriesArgs]):
    """Generate the integers ``0 .. count-1`` as a one-column table.

    Stateless: it emits every row in a single ``process`` call, then finishes.
    ``@bind_fixed_schema`` derives the output schema from ``FIXED_SCHEMA``;
    ``@init_single_worker`` runs the generator in a single worker.
    """

    FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema([("n", pa.int64())])

    @classmethod
    def process(cls, params: ProcessParams[SeriesArgs], state: None, out: OutputCollector) -> None:
        """Emit all rows at once, then signal completion."""
        out.emit(pa.RecordBatch.from_pydict({"n": list(range(params.args.count))}, schema=params.output_schema))
        out.finish()




class CalcWorker(Worker):
    """A worker exposing the ``calc`` catalog with both functions."""

    catalog = Catalog(
        name="calc",
        schemas=[Schema(name="main", functions=[Double, Series])],
    )


if __name__ == "__main__":
    CalcWorker().run()

ATTACH 'calc' (TYPE vgi, LOCATION 'uv run calc_worker.py');
SELECT * FROM calc.series(3);

Input

count
3

Output

n
0
1
2

Streaming with state

The tutorial’s series emits every row in a single process call. That’s fine for small results, but process is actually called repeatedly until you call out.finish() — so for large output you emit a bounded chunk per call and remember your place in a small state object. Here the state extends ArrowSerializableDataclass, which writes the encoding for you so it survives HTTP state round-trips (you can supply your own encoding instead):

series_streaming_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""A table function that streams its output in chunks, using generator state.

The tutorial's ``series`` emits everything in one call. When the output is large
you instead emit a bounded batch per ``process`` call and remember your place in
**state** — the framework calls ``process`` repeatedly until you ``out.finish()``.

    ATTACH 'calc' (TYPE vgi, LOCATION 'uv run series_streaming_worker.py');
    SELECT * FROM calc.series(1000000);   -- streamed 10k rows at a time
"""

from dataclasses import dataclass
from typing import Annotated, ClassVar

import pyarrow as pa
from vgi_rpc import ArrowSerializableDataclass

from vgi import Arg, Worker
from vgi.catalog import Catalog, Schema
from vgi.table_function import (
    OutputCollector,
    ProcessParams,
    TableFunctionGenerator,
    bind_fixed_schema,
    init_single_worker,
)

CHUNK = 10_000


@dataclass(slots=True, frozen=True, kw_only=True)
class SeriesArgs:
    """Arguments for :class:`Series` (one positional ``count``)."""

    count: Annotated[int, Arg(0, doc="How many numbers to generate", ge=0)]


@dataclass(kw_only=True)
class SeriesState(ArrowSerializableDataclass):
    """Cursor remembering how many rows we've emitted across ``process`` calls.

    Extends ``ArrowSerializableDataclass`` so the cursor survives HTTP state
    round-trips (the framework requires serializable state for generators).
    """

    emitted: int = 0


@init_single_worker
@bind_fixed_schema
class Series(TableFunctionGenerator[SeriesArgs, SeriesState]):
    """Generate ``0 .. count-1`` in chunks, keeping a cursor in state."""

    FIXED_SCHEMA: ClassVar[pa.Schema] = pa.schema([("n", pa.int64())])

    @classmethod
    def initial_state(cls, params: ProcessParams[SeriesArgs]) -> SeriesState:
        """Start a fresh cursor at zero."""
        return SeriesState()

    @classmethod
    def process(cls, params: ProcessParams[SeriesArgs], state: SeriesState, out: OutputCollector) -> None:
        """Emit one bounded chunk per call; finish when the cursor reaches count."""
        if state.emitted >= params.args.count:
            out.finish()
            return
        batch_size = min(params.args.count - state.emitted, CHUNK)
        values = list(range(state.emitted, state.emitted + batch_size))
        out.emit(pa.RecordBatch.from_pydict({"n": values}, schema=params.output_schema))
        state.emitted += batch_size


class CalcWorker(Worker):
    """A worker exposing the ``calc`` catalog with the streaming series."""

    catalog = Catalog(
        name="calc",
        schemas=[Schema(name="main", functions=[Series])],
    )


if __name__ == "__main__":
    CalcWorker().run()

SELECT * FROM calc.series(1000000);   -- streamed CHUNK rows per process() call

Table-in-out

Stream an input table through, batch by batch, emitting transformed output. on_bind declares the output schema; process receives each input batch and emits results. Here we keep only rows whose value column is positive.

filter_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""A table-in-out function: keep only rows where ``value`` is positive.

Table-in-out functions stream an input table through, batch by batch, emitting
transformed output. Run from a DuckDB-compatible engine::

    ATTACH 'filters' (TYPE vgi, LOCATION 'uv run filter_worker.py');
    SELECT * FROM filters.filter_positive((SELECT * FROM my_table));
"""

from dataclasses import dataclass
from typing import Annotated

import pyarrow as pa
import pyarrow.compute as pc

from vgi import Arg, Worker
from vgi.arguments import TableInput
from vgi.catalog import Catalog, Schema
from vgi.invocation import BindResponse
from vgi.table_function import BindParams, ProcessParams
from vgi.table_in_out_function import OutputCollector, TableInOutGenerator


@dataclass(slots=True, frozen=True, kw_only=True)
class FilterArgs:
    """Arguments: a single input table."""

    data: Annotated[TableInput, Arg(0, doc="Input table to filter")]


class FilterPositive(TableInOutGenerator[FilterArgs]):
    """Emit only the input rows whose ``value`` column is greater than zero."""

    @classmethod
    def on_bind(cls, params: BindParams[FilterArgs]) -> BindResponse:
        """Output schema equals the input schema (rows are filtered, not reshaped)."""
        assert params.bind_call.input_schema is not None
        return BindResponse(output_schema=params.bind_call.input_schema)

    @classmethod
    def process(
        cls,
        params: ProcessParams[FilterArgs],
        state: None,
        batch: pa.RecordBatch,
        out: OutputCollector,
    ) -> None:
        """Filter each input batch and emit the surviving rows."""
        mask = pc.greater(batch.column("value"), pa.scalar(0, type=batch.column("value").type))
        out.emit(batch.filter(mask))


class FilterWorker(Worker):
    """A worker exposing the ``filters`` catalog."""

    catalog = Catalog(
        name="filters",
        schemas=[Schema(name="main", functions=[FilterPositive])],
    )


if __name__ == "__main__":
    FilterWorker().run()

ATTACH 'filters' (TYPE vgi, LOCATION 'uv run filter_worker.py');
SELECT * FROM filters.filter_positive((SELECT * FROM (VALUES (-2), (5), (0), (9), (-1)) AS t(value)));

Input

value
-2
5
0
9
-1

Output

value
5
9

Blended: positional args are the input columns

The filter_positive function above takes an explicit TABLE subquery, so it can only be called one way. A RowTransformFunction removes that constraint: its positional Args declare the per-row input columns directly — no synthetic TABLE placeholder — so a single registration serves a literal call, a column call, and a correlated LATERAL call:

SELECT geo.geo_encode(52.0, 13.0);                    -- literals -> one input row
SELECT * FROM t, geo.geo_encode(t.lat, t.lon);        -- columns  -> streaming input
SELECT * FROM t, LATERAL geo.geo_encode(t.lat, t.lon);

Subclassing RowTransformFunction is the signal — there’s no Meta flag, deliberately, because a flag can be forgotten on one of several same-named overloads and inheritance cannot. Positional args are read off batch by their declared name; named (str-position) args stay bind-time scalars on params.args:

class GeoEncodeFunction(RowTransformFunction[GeoArgs]):
  class Meta:
      name = "geo_encode"

  @classmethod
  def on_bind(cls, params: BindParams[GeoArgs]) -> BindResponse:
      return BindResponse(output_schema=schema({"geohash": pa.string()}))

  @classmethod
  def process(cls, params, state, batch, out) -> None:
      precision = params.args.precision          # named arg -> params.args
      lats = batch.column("latitude").to_pylist()  # positional args -> batch columns
      lons = batch.column("longitude").to_pylist()
      codes = [f"{round(la, precision)}:{round(lo, precision)}" for la, lo in zip(lats, lons)]
      out.emit(pa.record_batch({"geohash": pa.array(codes, type=pa.string())}))

Three rules follow from the call shapes it has to serve:

Map-shaped only — there is no finalize. A finalize()/finish() override is rejected at resolve_metadata, because DuckDB forbids FinalExecute under correlated LATERAL. 1→1, 1→N and 1→0 all work; anything that must accumulate belongs in a classic TableInput table-in-out or a buffering function.
No positional const args. In the column form DuckDB sweeps a constant into the input subquery, and in the literal form it’s indistinguishable from an input column. Use a named arg for optional config.
Overloads resolve by arity, so same-named blended functions with different positional counts coexist. For a varargs function the runtime column names aren’t known at declaration time — read them positionally with RowTransformFunction.input_columns(batch).

Which table-in-out do I want?

Reach for blended when your function is a per-row map that users should be able to call like an ordinary function (f(x, y)), including inside LATERAL. Keep the classic TableInput form when the function genuinely consumes a relation — when it needs a required constant, or when it accumulates across the stream and needs a finalize.

Aggregate

Accumulate input rows into per-group state, then emit one row per group. Aggregates are driven by DuckDB’s GROUP BY and run in three phases — update (fold a batch into per-group state), combine (merge partial states across parallel workers), and finalize (state → output row).

sum_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""An aggregate function: sum a column per group.

Aggregate functions accumulate input rows into per-group state, then emit one
result row per group. They are driven by DuckDB's ``GROUP BY``::

    ATTACH 'aggregates' (TYPE vgi, LOCATION 'uv run sum_worker.py');
    SELECT category, aggregates.vgi_sum(value) FROM t GROUP BY category;

The three phases:

- ``update``   — fold a batch of values into per-group state (keyed by group id)
- ``combine``  — merge two partial states for the same group (parallel workers)
- ``finalize`` — turn each group's state into its output row
"""

from dataclasses import dataclass
from typing import Annotated

import pyarrow as pa
from vgi_rpc import ArrowSerializableDataclass, ArrowType

from vgi import Worker
from vgi.aggregate_function import AggregateFunction
from vgi.arguments import Param, Returns
from vgi.catalog import Catalog, Schema
from vgi.metadata import DistinctDependence, NullHandling, OrderDependence
from vgi.table_function import ProcessParams


@dataclass(kw_only=True)
class SumState(ArrowSerializableDataclass):
    """Running total for one group. Serializable so it survives parallel combine."""

    total: Annotated[int, ArrowType(pa.int64())] = 0


class Sum(AggregateFunction[SumState]):
    """Sum an int64 column, grouped by DuckDB's ``GROUP BY`` columns."""

    class Meta:
        """Function metadata: name, description, and aggregate semantics."""

        name = "vgi_sum"
        description = "Sum integer values per group"
        null_handling = NullHandling.DEFAULT
        order_dependent = OrderDependence.NOT_ORDER_DEPENDENT
        distinct_dependent = DistinctDependence.NOT_DISTINCT_DEPENDENT

    @classmethod
    def initial_state(cls, params: ProcessParams[None]) -> SumState:
        """One fresh accumulator per group."""
        return SumState()

    @classmethod
    def update(
        cls,
        states: dict[int, SumState],
        group_ids: pa.Int64Array,
        value: Annotated[pa.Int64Array, Param(doc="Column to sum")],
    ) -> None:
        """Fold a batch of values into each group's running total."""
        table = pa.table({"gid": group_ids, "value": value})
        grouped = table.group_by("gid").aggregate([("value", "sum")])
        for i in range(grouped.num_rows):
            gid: int = grouped.column("gid")[i].as_py()
            val = grouped.column("value_sum")[i].as_py()
            if val is not None:
                states[gid] = SumState(total=states[gid].total + val)

    @classmethod
    def combine(cls, source: SumState, target: SumState, params: ProcessParams[None]) -> SumState:
        """Merge two partial sums for the same group."""
        return SumState(total=source.total + target.total)

    @classmethod
    def finalize(
        cls,
        group_ids: pa.Int64Array,
        states: dict[int, SumState],
        params: ProcessParams[None],
    ) -> Annotated[pa.RecordBatch, Returns(pa.int64())]:
        """Emit one total per group."""
        results = [s.total if (s := states[gid.as_py()]) is not None else None for gid in group_ids]
        return pa.record_batch({"result": pa.array(results, type=pa.int64())})


class AggregateWorker(Worker):
    """A worker exposing the ``aggregates`` catalog."""

    catalog = Catalog(
        name="aggregates",
        schemas=[Schema(name="main", functions=[Sum])],
    )


if __name__ == "__main__":
    AggregateWorker().run()

ATTACH 'aggregates' (TYPE vgi, LOCATION 'uv run sum_worker.py');
SELECT category, aggregates.vgi_sum(value) AS total
FROM (VALUES (0, 10), (0, 5), (1, 1), (1, 2), (1, 3)) AS t(category, value)
GROUP BY category ORDER BY category;

Input

category	value
0	10
0	5
1	1
1	2
1	3

Output

category	total
0	15
1	6

State has to be turnable into bytes

Aggregate state is merged across parallel workers and round-trips through storage between calls, so TState must supply serialize_to_bytes() and deserialize_from_bytes(). Extending ArrowSerializableDataclass writes both for you — annotate fields with ArrowType(...) so the wire type is explicit — and that’s what sum_worker.py above does.

It is not the only option, and for a small state it’s an expensive one: an aggregate serializes once per group, per batch, and a one-row Arrow stream pays for a schema message, a batch message and an end-of-stream marker whatever the payload. A two-integer state packs into 16 bytes with struct. The same rule and the same trade-off apply to table generators and table-in-out functions — see State doesn’t have to be Arrow.

Buffering

When a function must see the whole input before it can produce any output — a global sort, top-k, or a full reduction — use a buffering function. Unlike table-in-out (which emits per input batch), it runs in three phases: process (the sink — stash each batch, return a state_id), combine (reduce all the partials once), and finalize (the source — stream the result out). Because the phases can run in different worker processes, state can’t live in memory — it goes in params.storage, a shared store keyed to this call.

row_count_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""A buffering function: count every input row, then emit one total.

A buffering function must see the *whole* input before it can produce output —
the basis for sorts, top-k, and full-stream reductions. It runs in three phases:

- **sink** (``process``): called per input batch; stash a partial in shared
  storage and return a ``state_id``.
- **combine**: called once after all input; reduce the partials into a result.
- **source** (``finalize``): called per tick to stream the result out.

State crosses process boundaries, so it lives in ``params.storage`` (scoped by
``execution_id``), not in memory.

    ATTACH 'buffers' (TYPE vgi, LOCATION 'uv run row_count_worker.py');
    SELECT * FROM buffers.row_count((SELECT * FROM big_table));
"""

from dataclasses import dataclass
from typing import Annotated

import pyarrow as pa
from vgi_rpc import ArrowSerializableDataclass

from vgi import Arg, Worker
from vgi.arguments import TableInput
from vgi.catalog import Catalog, Schema
from vgi.invocation import BindResponse
from vgi.table_buffering_function import OutputCollector, TableBufferingFunction, TableBufferingParams
from vgi.table_function import BindParams

_RESULT = pa.schema([("count", pa.int64())])


@dataclass(slots=True, frozen=True, kw_only=True)
class RowCountArgs:
    """Arguments: a single input table to count."""

    data: Annotated[TableInput, Arg(0, doc="Input table")]


@dataclass(kw_only=True)
class DrainState(ArrowSerializableDataclass):
    """Per-finalize-stream cursor: emit the total once, then finish."""

    done: bool = False


class RowCount(TableBufferingFunction[RowCountArgs, DrainState]):
    """Count all input rows and emit a single ``count`` row."""

    class Meta:
        """Function metadata."""

        name = "row_count"

    @classmethod
    def on_bind(cls, params: BindParams[RowCountArgs]) -> BindResponse:
        """Output is one int64 column regardless of input shape."""
        return BindResponse(output_schema=_RESULT)

    @classmethod
    def process(cls, batch: pa.RecordBatch, params: TableBufferingParams[RowCountArgs]) -> bytes:
        """Sink: stash this batch's row count; one bucket per execution."""
        params.storage.state_append(b"counts", b"", batch.num_rows.to_bytes(8, "little"))
        return params.execution_id

    @classmethod
    def combine(cls, state_ids: list[bytes], params: TableBufferingParams[RowCountArgs]) -> list[bytes]:
        """Combine: sum the partial counts into a single result."""
        total = sum(int.from_bytes(v, "little") for _id, v in params.storage.state_log_scan(b"counts", b""))
        params.storage.state_append(b"result", b"", total.to_bytes(8, "little"))
        return [params.execution_id]

    @classmethod
    def initial_finalize_state(cls, finalize_state_id: bytes, params: TableBufferingParams[RowCountArgs]) -> DrainState:
        """One cursor per finalize stream."""
        return DrainState()

    @classmethod
    def finalize(
        cls,
        params: TableBufferingParams[RowCountArgs],
        finalize_state_id: bytes,
        state: DrainState,
        out: OutputCollector,
    ) -> None:
        """Source: emit the total once, then signal completion."""
        if state.done:
            out.finish()
            return
        rows = params.storage.state_log_scan(b"result", b"")
        total = int.from_bytes(rows[-1][1], "little") if rows else 0
        out.emit(pa.RecordBatch.from_pydict({"count": [total]}, schema=params.output_schema))
        state.done = True


class BufferWorker(Worker):
    """A worker exposing the ``buffers`` catalog."""

    catalog = Catalog(
        name="buffers",
        schemas=[Schema(name="main", functions=[RowCount])],
    )


if __name__ == "__main__":
    BufferWorker().run()

ATTACH 'buffers' (TYPE vgi, LOCATION 'uv run row_count_worker.py');
SELECT * FROM buffers.row_count((SELECT * FROM (VALUES (1), (2), (3), (4), (5)) AS t(x)));

Input

x
1
2
3
4
5

Output

count
5

Buffering vs. table-in-out

Both consume a relation, but a table-in-out function emits per input batch and never holds the whole input — use it for streaming transforms (filter, enrich, reshape). Reach for buffering only when output genuinely depends on every row.

Recovering input order

A buffering sink runs in parallel across DuckDB threads, so batches arrive in no particular order. Set Meta.requires_input_batch_index = True and each process() call also receives params.batch_index — a globally-unique monotonic index — which is what lets you put the input back in order. That’s the prerequisite for anything order-sensitive: row pattern matching, a running total, LAG-style windows over the buffered stream.

batch_index_worker.py

# /// script
# requires-python = ">=3.13"
# dependencies = ["vgi-python"]
# ///
"""A buffering function that asks for input order, and reports what it got.

A buffering sink normally runs in parallel across DuckDB threads, so batches
arrive in no particular order. Setting ``Meta.requires_input_batch_index`` asks
for DuckDB's per-chunk index alongside each batch, which is what lets a worker
put the input back in order — the prerequisite for anything order-sensitive,
like row pattern matching or a running total.

Not every source can supply one: a base table scan can, while ``range()`` and
``VALUES`` cannot. When it cannot, the extension serializes the sink and numbers
the batches itself, so a worker sees a valid monotonic index either way and never
has to care which route produced it.

This function emits one row per buffered batch — its index and its row count — so
the guarantee is directly observable:

    ATTACH 'bi' (TYPE vgi, LOCATION 'uv run batch_index_worker.py');
    SELECT * FROM bi.batch_indexes((SELECT * FROM range(5000))) ORDER BY batch_index;
"""

from dataclasses import dataclass
from typing import Annotated

import pyarrow as pa
from vgi_rpc import ArrowSerializableDataclass

from vgi import Arg, Worker
from vgi.arguments import TableInput
from vgi.catalog import Catalog, Schema
from vgi.invocation import BindResponse
from vgi.table_buffering_function import OutputCollector, TableBufferingFunction, TableBufferingParams
from vgi.table_function import BindParams

_RESULT = pa.schema([("batch_index", pa.int64()), ("rows", pa.int64())])

# One log entry per buffered batch: the index DuckDB gave us, and the batch size.
_NS = b"batches"


@dataclass(slots=True, frozen=True, kw_only=True)
class BatchIndexArgs:
    """Arguments: the input table whose batches should be reported."""

    data: Annotated[TableInput, Arg(0, doc="Input table")]


@dataclass(kw_only=True)
class DrainState(ArrowSerializableDataclass):
    """Per-finalize-stream cursor: emit the report once, then finish."""

    done: bool = False


class BatchIndexes(TableBufferingFunction[BatchIndexArgs, DrainState]):
    """Report the batch index and row count of every buffered input batch."""

    class Meta:
        """Function metadata."""

        name = "batch_indexes"
        # Ask for DuckDB's per-chunk index. Mutually exclusive with
        # sink_order_dependent, which orders the input by serializing the sink
        # instead of by numbering it.
        requires_input_batch_index = True

    @classmethod
    def on_bind(cls, params: BindParams[BatchIndexArgs]) -> BindResponse:
        """Output shape is fixed: one row per input batch."""
        return BindResponse(output_schema=_RESULT)

    @classmethod
    def process(cls, batch: pa.RecordBatch, params: TableBufferingParams[BatchIndexArgs]) -> bytes:
        """Sink: record this batch's index and size.

        ``params.batch_index`` is populated because ``Meta`` asked for it; -1
        stands in for the absent case so an older host that does not supply one
        degrades to a visible marker instead of a crash.
        """
        index = params.batch_index if params.batch_index is not None else -1
        payload = index.to_bytes(8, "little", signed=True) + batch.num_rows.to_bytes(8, "little")
        params.storage.state_append(_NS, b"", payload)
        return params.execution_id

    @classmethod
    def combine(cls, state_ids: list[bytes], params: TableBufferingParams[BatchIndexArgs]) -> list[bytes]:
        """Nothing to reduce: the log already holds one entry per batch."""
        return [params.execution_id]

    @classmethod
    def initial_finalize_state(
        cls, finalize_state_id: bytes, params: TableBufferingParams[BatchIndexArgs]
    ) -> DrainState:
        """One cursor per finalize stream."""
        return DrainState()

    @classmethod
    def finalize(
        cls,
        params: TableBufferingParams[BatchIndexArgs],
        finalize_state_id: bytes,
        state: DrainState,
        out: OutputCollector,
    ) -> None:
        """Source: emit the report, ordered by batch index."""
        if state.done:
            out.finish()
            return
        entries = [
            (
                int.from_bytes(value[:8], "little", signed=True),
                int.from_bytes(value[8:16], "little"),
            )
            for _id, value in params.storage.state_log_scan(_NS, b"")
        ]
        entries.sort()
        out.emit(
            pa.RecordBatch.from_pydict(
                {
                    "batch_index": [index for index, _rows in entries],
                    "rows": [rows for _index, rows in entries],
                },
                schema=params.output_schema,
            )
        )
        state.done = True


class BatchIndexWorker(Worker):
    """A worker exposing the ``bi`` catalog."""

    catalog = Catalog(
        name="bi",
        schemas=[Schema(name="main", functions=[BatchIndexes])],
    )


if __name__ == "__main__":
    BatchIndexWorker().run()

ATTACH 'bi' (TYPE vgi, LOCATION 'uv run batch_index_worker.py');
SELECT * FROM bi.batch_indexes((SELECT * FROM range(5000))) ORDER BY batch_index;

Not every source can supply an index — a base table scan can, while range() and VALUES cannot. When it can’t, the extension serializes the sink and numbers the batches itself, so a worker always sees a valid monotonic index and never has to care which route produced it. params.batch_index is None on every other call path.

Next steps

1. Scalar function and 2. Table function — the guided build, start to finish.
API Reference — functions — the exact contracts for each shape.