docs: example worker for ordered buffering input (batch_index)

A buffering sink runs in parallel, so process() sees batches in arbitrary
order — and nothing in the docs said how to get the input back in order.
Meta.requires_input_batch_index is the answer, and it now works over any
source: when the source cannot supply an index (range(), VALUES) the
extension serializes the sink and numbers the batches itself rather than
aborting the pipeline, so callers no longer have to wrap input in a temp
table.

Adds examples/batch_index_worker.py, which reports the index and row count
of every buffered batch so the guarantee is directly observable, plus the
"Getting the input in order" section that embeds it and contrasts the flag
with sink_order_dependent.

Author: rustyconover

docs/how-to/function-patterns.md

@@ -213,6 +213,31 @@ ATTACH 'buffers' (TYPE vgi, LOCATION 'uv run row_count_worker.py');
 SELECT * FROM buffers.row_count((SELECT * FROM big_table));
 ```
 
+### Getting the input in order
+
+The sink runs in parallel across DuckDB threads, so `process` sees batches in no particular order.
+When output depends on input *order* — a running total, a row pattern match — set
+`Meta.requires_input_batch_index = True` and each `process` call receives `params.batch_index`, a
+monotonic ordinal you can sort the buffered batches by in `finalize`.
+
+Not every source can supply one (a base table scan can; `range()` and `VALUES` cannot). When it
+cannot, the extension serializes the sink and numbers the batches itself, so a worker always gets a
+valid index and callers never have to wrap their input. This example reports the indices it saw, one
+row per buffered batch:
+
+```python
+--8<-- "examples/batch_index_worker.py"
+```
+
+```sql
+ATTACH 'bi' (TYPE vgi, LOCATION 'uv run batch_index_worker.py');
+SELECT * FROM bi.batch_indexes((SELECT * FROM range(5000))) ORDER BY batch_index;
+```
+
+The alternative, `Meta.sink_order_dependent = True`, delivers ordered input by forcing a
+single-threaded sink. It works with any source but gives up parallel ingest, which is usually the
+bulk of a buffering query's wall time. The two flags are mutually exclusive.
+
 ??? info "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

examples/batch_index_worker.py

@@ -0,0 +1,142 @@
+# /// 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()

tests/test_examples_workers.py

@@ -100,6 +100,26 @@ def test_row_count_worker_buffering() -> None:
     assert [v for b in out for v in b.column("count").to_pylist()] == [5]
 
 
+def test_batch_index_worker_reports_input_order() -> None:
+    """The batch-index example sees a distinct index for every sunk batch."""
+    with _spawn("batch_index_worker.py") as client:
+        batches = [
+            pa.record_batch({"x": pa.array([1, 2, 3], type=pa.int64())}),
+            pa.record_batch({"x": pa.array([4, 5], type=pa.int64())}),
+            pa.record_batch({"x": pa.array([6], type=pa.int64())}),
+        ]
+        out = list(
+            client.table_buffering_function(function_name="batch_indexes", schema_name="main", input=iter(batches))
+        )
+    indexes = [v for b in out for v in b.column("batch_index").to_pylist()]
+    rows = [v for b in out for v in b.column("rows").to_pylist()]
+    # One report row per input batch, ordered by index, and none defaulted to -1.
+    assert indexes == sorted(indexes)
+    assert len(set(indexes)) == 3
+    assert min(indexes) >= 0
+    assert sum(rows) == 6
+
+
 def test_greeting_scalar_worker_string_example() -> None:
     """The string-scalar example (used in the function-patterns guide) still serves."""
     with _spawn("greeting_scalar_worker.py") as client: