Skip to content

Commit b3049b3

Browse files
rustyconoverclaude
andcommitted
Merge lateral continuation into main
Continued table-in-out output (HAVE_MORE_OUTPUT / CONTINUE ticks, resume()), Meta.input_rows_per_call, and the standard continuation fixtures every SDK implements. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2 parents f45612a + 74e989e commit b3049b3

13 files changed

Lines changed: 1147 additions & 28 deletions

File tree

‎docs/generator-api.md‎

Lines changed: 24 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -275,15 +275,33 @@ class SumFunction(TableInOutFunction):
275275
)]
276276
```
277277

278-
**Multiple outputs per input (HAVE_MORE_OUTPUT status):**
279-
```python
280-
from vgi_rpc.rpc import OutputCollector
278+
**Multiple outputs per input (`more=True` and `resume()`):**
281279

282-
def process(cls, params, state, batch, out: OutputCollector) -> None:
283-
# Emit multiple batches for one input by setting vgi.status metadata
284-
out.emit(batch, custom_metadata={b"vgi.status": b"HAVE_MORE_OUTPUT"})
280+
Each call still emits exactly one batch. To return several for one input
281+
batch, emit the first with `more=True`; the framework then calls `resume()`,
282+
with no input, for each next batch until one is emitted without `more=True`.
283+
`resume()` continues from `state` alone, so the function must declare a
284+
serializable `TState`; over HTTP it travels in the stream's state token, so
285+
store positions to resume from, not output rows. A function with `resume()`
286+
cannot set `max_workers = 1`: the client runs such a function on one connection
287+
shared by every pipeline, which cannot leave output pending.
288+
289+
```python test="skip"
290+
@classmethod
291+
def process(cls, params, state, batch, out) -> None:
292+
state.pages = plan_pages(batch) # where each page starts: small
293+
cls.resume(params, state, out)
294+
295+
@classmethod
296+
def resume(cls, params, state, out) -> None:
297+
page = fetch(state.pages.pop(0))
298+
out.emit(page.batch, parent_rows=page.parent_rows, more=bool(state.pages))
285299
```
286300

301+
On the wire the reply carries `vgi.status: HAVE_MORE_OUTPUT` and the client
302+
asks for the next batch with a zero-row input carrying `vgi.status: CONTINUE`.
303+
While output is pending that tick is the only valid input.
304+
287305
**In-band logging via OutputCollector:**
288306
```python
289307
from vgi_rpc.rpc import OutputCollector

‎docs/metadata.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ info = SumValuesFunction.describe()
8080
| `tags` | `dict[str, str]` | `{}` | Custom key-value tags |
8181
| `examples` | `list` | `[]` | SQL examples (str or FunctionExample) |
8282
| `max_workers` | `int\|None` | `None` (unlimited) | Max parallel workers |
83+
| `input_rows_per_call` | `int\|None` | `None` (a whole chunk) | `RowTransformFunction` only: most input rows per worker call; the client splits each chunk into calls of at most this many rows |
8384
| `stability` | `FunctionStability` | `CONSISTENT` | Output determinism |
8485
| `null_handling` | `NullHandling` | `DEFAULT` | NULL input behavior |
8586
| `argument_monotonicity` | `list[ArgumentMonotonicity]\|None` | `None` | Scalar-only monotonicity claims in argument declaration order |

‎examples/paged_api_worker.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
# /// script
2+
# requires-python = ">=3.13"
3+
# dependencies = ["vgi-python"]
4+
# ///
5+
"""A lateral function backed by a paged API, returning one API page per reply.
6+
7+
``orders(customer_id)`` lists a customer's orders from an API that returns them a
8+
page at a time behind a cursor. Rather than fetching every page before replying,
9+
``process()`` fetches the first page and emits it with ``more=True``; the framework
10+
then calls ``resume()`` for each next page. The state holds only where to resume
11+
(which input row, which page cursor), never fetched rows, so it is small enough to
12+
travel in an HTTP state token.
13+
14+
ATTACH 'shop' (TYPE vgi, LOCATION 'uv run paged_api_worker.py');
15+
SELECT c.id, o.order_id, o.total
16+
FROM (VALUES (1), (2), (3)) c(id), shop.orders(c.id) o;
17+
"""
18+
19+
from dataclasses import dataclass
20+
from typing import Annotated, Any, cast
21+
22+
import pyarrow as pa
23+
from vgi_rpc import ArrowSerializableDataclass
24+
25+
from vgi import Arg, Worker
26+
from vgi.catalog import Catalog, Schema
27+
from vgi.invocation import BindResponse
28+
from vgi.table_function import BindParams, ProcessParams
29+
from vgi.table_in_out_function import OutputCollector, RowTransformFunction
30+
31+
PAGE_SIZE = 5
32+
33+
34+
def fetch_orders_page(customer_id: int, cursor: str) -> tuple[list[dict[str, Any]], str]:
35+
"""Stand-in for an HTTP API: one page of a customer's orders and the next cursor.
36+
37+
Customer ``n`` has ``7 * n`` orders. An empty next cursor means the last page.
38+
"""
39+
start = int(cursor or "0")
40+
total = 7 * customer_id
41+
orders = [
42+
{"order_id": customer_id * 1000 + k, "total": round(10.0 + k * 2.5, 2)}
43+
for k in range(start, min(start + PAGE_SIZE, total))
44+
]
45+
end = start + len(orders)
46+
return orders, str(end) if end < total else ""
47+
48+
49+
@dataclass(slots=True, frozen=True, kw_only=True)
50+
class OrdersArgs:
51+
"""One positional input column: the customer to list orders for."""
52+
53+
customer_id: Annotated[int, Arg(0, doc="Customer whose orders to list")]
54+
55+
56+
@dataclass(kw_only=True)
57+
class OrdersCursor(ArrowSerializableDataclass):
58+
"""Where the next reply resumes: the input's customer ids, a row, and that row's page cursor."""
59+
60+
customer_ids: list[int] | None = None
61+
row: int = 0
62+
cursor: str = ""
63+
64+
65+
class Orders(RowTransformFunction[OrdersArgs, OrdersCursor]):
66+
"""List each input customer's orders, one API page per reply."""
67+
68+
@classmethod
69+
def on_bind(cls, params: BindParams[OrdersArgs]) -> BindResponse:
70+
"""A fixed output schema."""
71+
return BindResponse(output_schema=pa.schema([("order_id", pa.int64()), ("total", pa.float64())]))
72+
73+
@classmethod
74+
def initial_state(cls, params: ProcessParams[OrdersArgs]) -> OrdersCursor:
75+
"""An empty cursor; ``process()`` fills it for each input batch."""
76+
return OrdersCursor()
77+
78+
@classmethod
79+
def process(
80+
cls, params: ProcessParams[OrdersArgs], state: OrdersCursor, batch: pa.RecordBatch, out: OutputCollector
81+
) -> None:
82+
"""Start a new input batch at its first row, then emit the first page."""
83+
state.customer_ids = [-1 if c is None else c for c in batch.column("customer_id").to_pylist()]
84+
state.row = 0
85+
state.cursor = ""
86+
cls.resume(params, state, out)
87+
88+
@classmethod
89+
def resume(cls, params: ProcessParams[OrdersArgs], state: OrdersCursor, out: OutputCollector) -> None:
90+
"""Emit the next non-empty page, with ``more=True`` unless it is the last."""
91+
ids = state.customer_ids or []
92+
orders: list[dict[str, Any]] = []
93+
row = state.row
94+
while state.row < len(ids) and not orders:
95+
row = state.row
96+
orders, state.cursor = fetch_orders_page(ids[row], state.cursor) if ids[row] > 0 else ([], "")
97+
if not state.cursor: # that customer is done: the next page starts the next row
98+
state.row += 1
99+
# Every row of a page came from one input row; parent_rows says which.
100+
cast(Any, out).emit(
101+
pa.RecordBatch.from_pylist(orders, schema=params.output_schema),
102+
parent_rows=[row] * len(orders),
103+
more=state.row < len(ids),
104+
)
105+
106+
107+
class ShopWorker(Worker):
108+
"""A worker exposing the ``shop`` catalog with the paged ``orders`` function."""
109+
110+
catalog = Catalog(
111+
name="shop",
112+
schemas=[Schema(path=["main"], functions=[Orders])],
113+
)
114+
115+
116+
if __name__ == "__main__":
117+
ShopWorker().run()

0 commit comments

Comments
 (0)