|
| 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