vgi-python 0.46.0__py3-none-any.whl → 0.47.0__py3-none-any.whl

vgi/_test_fixtures/table/cache.py

@@ -30,6 +30,7 @@ from dataclasses import dataclass
 from typing import Annotated, Any, ClassVar, cast
 
 import pyarrow as pa
+import pyarrow.compute as pc
 from vgi_rpc import ArrowSerializableDataclass
 from vgi_rpc.rpc import OutputCollector
 
@@ -945,7 +946,7 @@ class CacheBenchFunction(TableFunctionGenerator[CacheBenchArgs, _CacheCountdownS
 # sequence [0..rows) so COUNT=rows and SUM=rows*(rows-1)/2 regardless of how the
 # chunks were distributed across workers. Positional Arg(0) so the direct path
 # `vgi_table_function(w, 'cache_parallel', [rows])` honors the size; a named
-# `batch_size` (default 24000) controls batch width for the 2 GB streaming guard.
+# `batch_size` (default 24000) controls batch width for the disk-streaming guard (parallel_2gb.test_slow).
 _PC_ITEM_FMT = ">QQ"  # (start, end) as two uint64
 _PC_ITEM_SIZE = struct.calcsize(_PC_ITEM_FMT)
 
@@ -1001,6 +1002,16 @@ class CacheParallelFunction(TableFunctionGenerator[CacheParallelArgs, _CachePara
 
     FunctionArguments = CacheParallelArgs
 
+    _OFFSETS: ClassVar[dict[int, pa.Int64Array]] = {}
+
+    @classmethod
+    def _offsets(cls, n: int) -> pa.Int64Array:
+        """0 .. n-1 as int64, built once per batch size per process."""
+        arr = cls._OFFSETS.get(n)
+        if arr is None:
+            arr = cls._OFFSETS[n] = pa.array(range(n), type=pa.int64())
+        return arr
+
     @classmethod
     def on_init(cls, params: InitParams[CacheParallelArgs]) -> GlobalInitResponse:
         """Primary worker enqueues (start, end) chunks covering [0, rows)."""
@@ -1032,8 +1043,11 @@ class CacheParallelFunction(TableFunctionGenerator[CacheParallelArgs, _CachePara
             state.current_idx = state.current_start
 
         batch_end = min(state.current_idx + params.args.batch_size, state.current_end or 0)
-        values = list(range(state.current_idx, batch_end))
-        batch = pa.RecordBatch.from_pydict({"v": values}, schema=params.output_schema)
+        # v = current_idx .. batch_end-1, as one vectorised add to a cached 0..n-1
+        # (a Python list costs ~1 µs a row, which the 640 MB streaming guard pays 80M times).
+        base = cls._offsets(params.args.batch_size).slice(0, batch_end - state.current_idx)
+        values = pc.add(base, pa.scalar(state.current_idx, type=pa.int64()))
+        batch = pa.RecordBatch.from_arrays([values], schema=params.output_schema)
         cache_control = CacheControl(ttl=_DEFAULT_TTL_SECONDS) if not state.advertised else None
         cast(VgiOutputCollector, out).emit(batch, cache_control=cache_control)
         state.advertised = True

vgi/_test_fixtures/table_in_out.py

@@ -1883,6 +1883,345 @@ class BlendedExplodeFunction(RowTransformFunction[_ExplodeArgs]):
         cast("VgiOutputCollector", out).emit(out_batch, parent_rows=parent_rows)
 
 
+@dataclass(slots=True, frozen=True, kw_only=True)
+class _PagedArgs:
+    """One positional fan-out count, and how many output rows each reply carries."""
+
+    n: Annotated[int, Arg(0, doc="Fan-out count: emit rows 0..n-1 for this input row")]
+    page_size: Annotated[int, Arg("page_size", doc="Output rows per reply", default=2)] = 2
+    cache_ttl: Annotated[
+        int, Arg("cache_ttl", doc="When > 0, advertise vgi.cache.ttl on the first reply", default=0)
+    ] = 0
+    fail_on_page: Annotated[
+        int, Arg("fail_on_page", doc="When > 0, fail instead of sending that (1-based) reply", default=0)
+    ] = 0
+    revalidate: Annotated[
+        bool, Arg("revalidate", doc="Answer conditional requests with changed, paged output", default=False)
+    ] = False
+    not_modified_more: Annotated[
+        bool,
+        Arg(
+            "not_modified_more",
+            doc="With revalidate, answer a conditional request with an empty not_modified reply that says more",
+            default=False,
+        ),
+    ] = False
+
+
+@dataclass(kw_only=True)
+class _PagedState(ArrowSerializableDataclass):
+    """Where the next reply resumes: the input's fan-out counts and a position in them."""
+
+    counts: list[int] | None = None
+    row: int = 0
+    offset: int = 0
+    call_rows: int = 0  # blended_sliced: the input rows of the current call
+    page: int = 0  # replies sent for the current input
+    empty_first: bool = False  # not_modified_more: the first reply is empty and says more
+    bump: int = 0  # added to every output value (revalidate: the data changed)
+
+
+class BlendedPagedFunction(RowTransformFunction[_PagedArgs, _PagedState]):
+    """Blended 1->N fan-out returned across several replies, ``page_size`` rows each.
+
+    Same output as ``blended_explode`` (rows ``0..n-1`` per input row, with
+    ``parent_rows`` provenance), but one input batch's output is split into pages:
+    ``process()`` emits the first with ``more=True`` and ``resume()`` the rest, the
+    way a function backed by a paged API would. The state holds only the counts
+    and a position, never output rows, so it stays small in an HTTP state token.
+    ``cache_ttl > 0`` makes the first reply cacheable, for the result-cache tests.
+    ``fail_on_page`` and ``revalidate`` are described at :func:`_paged_begin`.
+    """
+
+    class Meta:
+        name = "blended_paged"
+        description = "Blended 1->N fan-out (emit 0..n-1 per input row), page_size rows per reply"
+        categories = ["blended", "test"]
+
+    @classmethod
+    def on_bind(cls, params: BindParams[_PagedArgs]) -> BindResponse:
+        if params.args.page_size < 1:
+            raise ValueError(f"page_size must be at least 1, got {params.args.page_size}")
+        if params.args.cache_ttl < 0:
+            raise ValueError(f"cache_ttl must not be negative, got {params.args.cache_ttl}")
+        return BindResponse(output_schema=schema({"i": pa.int64()}))
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_PagedArgs]) -> _PagedState:
+        return _PagedState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_PagedArgs],
+        state: _PagedState,
+        batch: pa.RecordBatch,
+        out: OutputCollector,
+    ) -> None:
+        ttl = params.args.cache_ttl
+        cc = _paged_begin(params, state, batch.column("n").to_pylist(), CacheControl(ttl=ttl) if ttl > 0 else None)
+        cls._emit_page(params, state, out, cc)
+
+    @classmethod
+    def resume(cls, params: ProcessParams[_PagedArgs], state: _PagedState, out: OutputCollector) -> None:
+        cls._emit_page(params, state, out, None)
+
+    @classmethod
+    def _emit_page(
+        cls,
+        params: ProcessParams[_PagedArgs],
+        state: _PagedState,
+        out: OutputCollector,
+        cache_control: CacheControl | None,
+    ) -> None:
+        out_vals, parent_rows, more = _paged_next(params, state)
+        cast("VgiOutputCollector", out).emit(
+            pa.record_batch({"i": pa.array(out_vals, type=pa.int64())}),
+            parent_rows=parent_rows,
+            cache_control=cache_control,
+            more=more,
+        )
+
+
+def _paged_begin(
+    params: ProcessParams[Any], state: _PagedState, counts: list[Any], cache_control: CacheControl | None
+) -> CacheControl | None:
+    """Start a new input for ``blended_paged`` / ``paged_explode``; return its first reply's cache control.
+
+    ``revalidate`` makes the output depend on the request alone. A plain request
+    gets the whole output in one reply, cacheable but immediately stale
+    (``ttl=0``, etag ``v1``, revalidatable), so the client stores it and revalidates
+    every repeat. A conditional request (the client sends ``if_none_match``) is
+    told the data changed: every value is 1000 larger, the etag is ``v2``, and the
+    output is paged by ``page_size``. A client must then return the new pages and
+    store nothing, since that output continued.
+
+    ``not_modified_more`` (with ``revalidate``) breaks the protocol on purpose: a
+    conditional request is answered with an empty ``not_modified`` reply (etag
+    ``v1``) that also says more, then the changed pages. A 304 stands for an input's
+    whole stored output, so a client must refuse it rather than serve its entry.
+    """
+    state.counts = [0 if n is None or n < 0 else int(n) for n in counts]
+    state.row = 0
+    state.offset = 0
+    state.page = 0
+    state.bump = 0
+    state.empty_first = False
+    if not params.args.revalidate:
+        return cache_control
+    if params.if_none_match is None:
+        state.bump = -1  # one reply: the whole output
+        return CacheControl(ttl=0, etag="v1", revalidatable=True)
+    state.bump = 1000
+    if params.args.not_modified_more:
+        state.empty_first = True
+        return CacheControl(not_modified=True, ttl=0, etag="v1", revalidatable=True)
+    return CacheControl(ttl=0, etag="v2", revalidatable=True)
+
+
+def _paged_next(params: ProcessParams[Any], state: _PagedState) -> tuple[list[int], list[int], bool]:
+    """The next reply of ``blended_paged`` / ``paged_explode``, or the failure ``fail_on_page`` asks for."""
+    state.page += 1
+    if state.page == params.args.fail_on_page:
+        raise RuntimeError(f"failing on page {state.page} as asked")
+    if state.page == 1 and state.empty_first:
+        return [], [], True
+    if state.bump < 0:
+        return _next_page(state, sys.maxsize)
+    out_vals, parent_rows, more = _next_page(state, params.args.page_size)
+    return [v + state.bump for v in out_vals], parent_rows, more
+
+
+def _next_page(state: _PagedState, page_size: int) -> tuple[list[int], list[int], bool]:
+    """The next ``page_size`` rows of ``blended_paged``'s output, and whether more follow.
+
+    Rows are ``0..n-1`` for each input row in order, so a page may span input rows.
+    Input rows with no output after the last non-empty page are skipped, so the
+    last page is never followed by an empty one.
+    """
+    counts = state.counts or []
+    out_vals: list[int] = []
+    parent_rows: list[int] = []
+    while state.row < len(counts) and len(out_vals) < page_size:
+        take = min(counts[state.row] - state.offset, page_size - len(out_vals))
+        out_vals.extend(range(state.offset, state.offset + take))
+        parent_rows.extend([state.row] * take)
+        state.offset += take
+        if state.offset == counts[state.row]:
+            state.row += 1
+            state.offset = 0
+    while state.row < len(counts) and counts[state.row] == 0:
+        state.row += 1
+    return out_vals, parent_rows, state.row < len(counts)
+
+
+@dataclass(slots=True, frozen=True, kw_only=True)
+class _SlicedArgs:
+    """One positional fan-out count, and how many output rows each reply carries."""
+
+    n: Annotated[int, Arg(0, doc="Fan-out count: emit rows 0..n-1 for this input row")]
+    page_size: Annotated[int, Arg("page_size", doc="Output rows per reply", default=1_000_000)] = 1_000_000
+    cache_ttl: Annotated[int, Arg("cache_ttl", doc="If > 0, per-value cacheable for this many seconds", default=0)] = 0
+
+
+class BlendedSlicedFunction(RowTransformFunction[_SlicedArgs, _PagedState]):
+    """``blended_paged`` that takes at most 3 input rows per call, and says how many it got.
+
+    Declares ``Meta.input_rows_per_call = 3``, so the client must split every input
+    chunk into calls of at most 3 rows. Each output row carries ``call_rows``, the
+    row count of the call that produced it, so a test can see the split. Output is
+    otherwise ``blended_paged``'s (rows ``0..n-1`` per input row, ``page_size``
+    rows per reply), which covers a split chunk whose calls also continue.
+    ``cache_ttl > 0`` advertises per-value caching on each call's first reply, so the
+    client stores each call's values for reuse.
+    """
+
+    class Meta:
+        name = "blended_sliced"
+        description = "Blended 1->N fan-out taking at most 3 input rows per call, reporting each call's size"
+        categories = ["blended", "test"]
+        input_rows_per_call = 3
+
+    @classmethod
+    def on_bind(cls, params: BindParams[_SlicedArgs]) -> BindResponse:
+        if params.args.page_size < 1:
+            raise ValueError(f"page_size must be at least 1, got {params.args.page_size}")
+        if params.args.cache_ttl < 0:
+            raise ValueError(f"cache_ttl must not be negative, got {params.args.cache_ttl}")
+        return BindResponse(output_schema=schema({"i": pa.int64(), "call_rows": pa.int64()}))
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_SlicedArgs]) -> _PagedState:
+        return _PagedState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_SlicedArgs],
+        state: _PagedState,
+        batch: pa.RecordBatch,
+        out: OutputCollector,
+    ) -> None:
+        if batch.num_rows > 3:
+            raise RuntimeError(f"blended_sliced declares input_rows_per_call = 3 but got {batch.num_rows} rows")
+        state.counts = [0 if n is None or n < 0 else int(n) for n in batch.column("n").to_pylist()]
+        state.row = 0
+        state.offset = 0
+        state.call_rows = batch.num_rows
+        ttl = params.args.cache_ttl
+        cls._emit_page(params, state, out, CacheControl(ttl=ttl, per_value=True) if ttl > 0 else None)
+
+    @classmethod
+    def resume(cls, params: ProcessParams[_SlicedArgs], state: _PagedState, out: OutputCollector) -> None:
+        cls._emit_page(params, state, out, None)
+
+    @classmethod
+    def _emit_page(
+        cls,
+        params: ProcessParams[_SlicedArgs],
+        state: _PagedState,
+        out: OutputCollector,
+        cache_control: CacheControl | None,
+    ) -> None:
+        out_vals, parent_rows, more = _next_page(state, params.args.page_size)
+        cast("VgiOutputCollector", out).emit(
+            pa.record_batch(
+                {
+                    "i": pa.array(out_vals, type=pa.int64()),
+                    "call_rows": pa.array([state.call_rows] * len(out_vals), type=pa.int64()),
+                }
+            ),
+            parent_rows=parent_rows,
+            cache_control=cache_control,
+            more=more,
+        )
+
+
+@dataclass(slots=True, frozen=True, kw_only=True)
+class _PagedTableArgs:
+    """An input table whose first column is the fan-out count, and the page size."""
+
+    data: Annotated[TableInput, Arg(0, doc="Input table; its first column is the fan-out count")]
+    page_size: Annotated[int, Arg("page_size", doc="Output rows per reply", default=2)] = 2
+    cache_ttl: Annotated[
+        int, Arg("cache_ttl", doc="When > 0, advertise vgi.cache.ttl on the first reply", default=0)
+    ] = 0
+    fail_on_page: Annotated[
+        int, Arg("fail_on_page", doc="When > 0, fail instead of sending that (1-based) reply", default=0)
+    ] = 0
+    revalidate: Annotated[
+        bool, Arg("revalidate", doc="Answer conditional requests with changed, paged output", default=False)
+    ] = False
+    not_modified_more: Annotated[
+        bool,
+        Arg(
+            "not_modified_more",
+            doc="With revalidate, answer a conditional request with an empty not_modified reply that says more",
+            default=False,
+        ),
+    ] = False
+
+
+class PagedExplodeFunction(TableInOutGenerator[_PagedTableArgs, _PagedState]):
+    """Classic (TABLE-input) table-in-out continuing its output across replies.
+
+    For each input row with count ``n`` (the table's first column), emits rows
+    ``i = 0..n-1``, ``page_size`` per reply, with ``more=True`` on every reply but
+    an input batch's last. The classic counterpart of ``blended_paged``: it runs
+    through the extension's stock in-out path, on a per-substream connection.
+    ``cache_ttl > 0`` makes each input's first reply cacheable. ``fail_on_page`` and
+    ``revalidate`` are described at :func:`_paged_begin`.
+    """
+
+    class Meta:
+        name = "paged_explode"
+        description = "Classic table-in-out 1->N fan-out, page_size rows per reply"
+        categories = ["table_in_out", "test"]
+
+    @classmethod
+    def on_bind(cls, params: BindParams[_PagedTableArgs]) -> BindResponse:
+        if params.args.page_size < 1:
+            raise ValueError(f"page_size must be at least 1, got {params.args.page_size}")
+        if params.args.cache_ttl < 0:
+            raise ValueError(f"cache_ttl must not be negative, got {params.args.cache_ttl}")
+        return BindResponse(output_schema=schema({"i": pa.int64()}))
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_PagedTableArgs]) -> _PagedState:
+        return _PagedState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_PagedTableArgs],
+        state: _PagedState,
+        batch: pa.RecordBatch,
+        out: OutputCollector,
+    ) -> None:
+        counts = batch.column(0).to_pylist() if batch.num_columns else []
+        ttl = params.args.cache_ttl
+        cls._emit_page(
+            params, state, out, _paged_begin(params, state, counts, CacheControl(ttl=ttl) if ttl > 0 else None)
+        )
+
+    @classmethod
+    def resume(cls, params: ProcessParams[_PagedTableArgs], state: _PagedState, out: OutputCollector) -> None:
+        cls._emit_page(params, state, out, None)
+
+    @classmethod
+    def _emit_page(
+        cls,
+        params: ProcessParams[_PagedTableArgs],
+        state: _PagedState,
+        out: OutputCollector,
+        cache_control: CacheControl | None,
+    ) -> None:
+        out_vals, _parent_rows, more = _paged_next(params, state)
+        cast("VgiOutputCollector", out).emit(
+            pa.record_batch({"i": pa.array(out_vals, type=pa.int64())}), cache_control=cache_control, more=more
+        )
+
+
 @dataclass(slots=True, frozen=True, kw_only=True)
 class _AnyArgs:
     """One positional ANY-typed input column."""

vgi/_test_fixtures/worker.py

@@ -279,6 +279,8 @@ from vgi._test_fixtures.table_in_out import (
     BlendedAnyVarargsFunction,
     BlendedDropFunction,
     BlendedExplodeFunction,
+    BlendedPagedFunction,
+    BlendedSlicedFunction,
     BufferEmitWideFunction,
     BufferInputFunction,
     CachedDoubleFunction,
@@ -304,6 +306,7 @@ from vgi._test_fixtures.table_in_out import (
     MultiBatchFinishFunction,
     OrderedBufferInputFunction,
     OrderedSourceFunction,
+    PagedExplodeFunction,
     ProjectableBlendedFunction,
     RepeatInputsFunction,
     RowSumFunction,
@@ -445,6 +448,9 @@ _EXAMPLE_CATALOG = Catalog(
                 BlendedAnyFunction,
                 BlendedAnyVarargsFunction,
                 BlendedExplodeFunction,
+                BlendedPagedFunction,
+                BlendedSlicedFunction,
+                PagedExplodeFunction,
                 CachedExplodeFunction,
                 ProjectableBlendedFunction,
                 HostileProvenanceFunction,

vgi/catalog/catalog_interface.py

@@ -797,6 +797,9 @@ class FunctionInfo(CatalogSchemaObject, ArrowSerializableDataclass):
             subclassing; the C++ extension reads it to enter the in-out
             registration branch with real-typed args and drive the literal
             single-row scan-mode.
+        input_rows_per_call: For a blended ``RowTransformFunction``, the most
+            input rows one worker call receives; the client splits each input
+            chunk into calls of at most this many rows. Null sends a whole chunk.
         required_settings: Settings required by the function.
         required_secrets: Secrets required by the function (each entry has
             secret_type, optional secret_name, optional scope).
@@ -865,6 +868,10 @@ class FunctionInfo(CatalogSchemaObject, ArrowSerializableDataclass):
     # branch with real-typed args and drive the literal single-row scan-mode.
     input_from_args: bool = False
 
+    # Blended only: the most input rows per worker call; the client splits each
+    # input chunk into calls of at most this many rows. None = a whole chunk.
+    input_rows_per_call: int | None = None
+
     required_settings: list[str] = field(default_factory=list)
 
     required_secrets: list[SecretLookupEntry] = field(default_factory=list)
@@ -3698,6 +3705,7 @@ class ReadOnlyCatalogInterface(CatalogInterface):
             sink_order_dependent=meta.sink_order_dependent,
             requires_input_batch_index=meta.requires_input_batch_index,
             input_from_args=False if is_scalar else meta.input_from_args,
+            input_rows_per_call=None if is_scalar else meta.input_rows_per_call,
             # Settings
             required_settings=meta.required_settings,
             # Secrets

vgi/client/client.py

@@ -124,6 +124,9 @@ from vgi.invocation import (
     GlobalInitResponse,
 )
 from vgi.protocol import (
+    CONTINUE,
+    HAVE_MORE_OUTPUT,
+    STATUS_KEY,
     BindRequest,
     CopyFromContext,
     CopyToContext,
@@ -2147,9 +2150,10 @@ class Client(CatalogClientMixin, AggregateClientMixin):
     ) -> tuple[list[pa.RecordBatch], list[list[int]] | None]:
         """Send a batch to a worker and collect all output batches.
 
-        Sends the input batch via stream.exchange(), then checks the vgi.status
-        metadata. If the worker returns HAVE_MORE_OUTPUT, sends the same input
-        again. Continues until NEED_MORE_INPUT or no status (scalar functions).
+        Sends the input batch via stream.exchange(), then checks the reply's
+        vgi.status. While the worker returns HAVE_MORE_OUTPUT, sends a zero-row tick
+        (``vgi.status: CONTINUE``) for its next batch. Stops at NEED_MORE_INPUT
+        or no status (scalar functions).
 
         Args:
             worker: The worker connection to use. Must have stream initialized.
@@ -2181,17 +2185,26 @@ class Client(CatalogClientMixin, AggregateClientMixin):
 
         output_batches: list[pa.RecordBatch] = []
         parent_rows_batches: list[list[int]] | None = [] if decode_parent_rows else None
+        tick = False
 
         while True:
             _logger.debug(
-                "sending_batch_to_worker worker_index=%s batch_index=%s num_rows=%s",
+                "sending_batch_to_worker worker_index=%s batch_index=%s num_rows=%s tick=%s",
                 worker.worker_index,
                 batch_index,
                 input_batch.num_rows,
+                tick,
             )
 
             try:
-                annotated = self._maybe_externalize_input_batch(input_batch)
+                if tick:
+                    # More output for the same input: ask for it with a zero-row tick.
+                    annotated = AnnotatedBatch(
+                        batch=input_batch.slice(0, 0),
+                        custom_metadata=pa.KeyValueMetadata({STATUS_KEY: CONTINUE}),
+                    )
+                else:
+                    annotated = self._maybe_externalize_input_batch(input_batch)
                 output = worker.stream.exchange(annotated)
             except RpcError as e:
                 raise ClientError.from_rpc_error(e) from e
@@ -2216,13 +2229,13 @@ class Client(CatalogClientMixin, AggregateClientMixin):
                 except RuntimeError as e:
                     raise ClientError(str(e)) from e
 
-            # Check vgi.status for table-in-out status
+            # HAVE_MORE_OUTPUT: this input has more output, fetched with a tick.
             status = None
             if output.custom_metadata:
-                status = output.custom_metadata.get(b"vgi.status")
+                status = output.custom_metadata.get(STATUS_KEY)
 
-            # status is None for scalar functions (no status metadata)
-            if status == b"HAVE_MORE_OUTPUT":
+            if status == HAVE_MORE_OUTPUT:
+                tick = True
                 continue
             elif status == b"NEED_MORE_INPUT" or status is None:
                 break

vgi/http/demo_storage.py

@@ -118,8 +118,9 @@ class BlobResource:
         resp.content_type = "application/octet-stream"
         resp.set_header("Accept-Ranges", "none")
         if content_encoding:
+            # HTTP clients decode this standard header automatically. Sending
+            # the custom VGI header too makes clients try to decode it again.
             resp.set_header("Content-Encoding", content_encoding)
-            resp.set_header("X-VGI-Content-Encoding", content_encoding)
 
     def on_head(self, req: Any, resp: Any, blob_id: str) -> None:  # noqa: D102
         # Mirror on_get headers (Content-Length/-Type/-Encoding) without a body.
@@ -137,7 +138,6 @@ class BlobResource:
         resp.set_header("Accept-Ranges", "none")
         if content_encoding:
             resp.set_header("Content-Encoding", content_encoding)
-            resp.set_header("X-VGI-Content-Encoding", content_encoding)
 
     def on_put(self, req: Any, resp: Any, blob_id: str) -> None:  # noqa: D102
         # vgi_rpc's _CompressionMiddleware drains ``req.bounded_stream`` when

vgi/metadata.py

@@ -551,6 +551,10 @@ class ResolvedMetadata:
             positional args ARE its per-row input columns, so one registration
             serves the literal / column / LATERAL call shapes. Detected from
             ``RowTransformFunction`` subclassing.
+        input_rows_per_call: Only for a ``RowTransformFunction``: the most input
+            rows one worker call receives. The client splits each input chunk into
+            calls of at most this many rows (``1`` = a row per call). None (the
+            default) sends a whole chunk per call.
 
     """
 
@@ -647,6 +651,11 @@ class ResolvedMetadata:
     # drive the literal single-row scan-mode. function_type stays TABLE.
     input_from_args: bool = False
 
+    # RowTransformFunction only: the most input rows one worker call receives; the
+    # client splits each input chunk into calls of at most this many rows. None =
+    # a whole chunk per call.
+    input_rows_per_call: int | None = None
+
     def to_dict(self) -> dict[str, Any]:
         """Convert to dictionary for JSON serialization."""
         return {
@@ -691,6 +700,7 @@ class ResolvedMetadata:
             "sink_order_dependent": self.sink_order_dependent,
             "requires_input_batch_index": self.requires_input_batch_index,
             "input_from_args": self.input_from_args,
+            "input_rows_per_call": self.input_rows_per_call,
         }
 
     @staticmethod
@@ -746,6 +756,7 @@ class ResolvedMetadata:
             sink_order_dependent=d.get("sink_order_dependent", False),
             requires_input_batch_index=d.get("requires_input_batch_index", False),
             input_from_args=d.get("input_from_args", False),
+            input_rows_per_call=d.get("input_rows_per_call"),
         )
 
 
@@ -1146,6 +1157,7 @@ _VALID_META_ATTRIBUTES: frozenset[str] = frozenset(
         "filters_exactly_applied",  # worker applies pushed filters exactly (engine may drop its own filter)
         "supports_positions",  # addressable positions in the data (incremental / streaming reads)
         "split_token_ttl_seconds",  # None = unbounded; a client refuses a TTL below its scheduling horizon
+        "input_rows_per_call",  # RowTransformFunction: most input rows per worker call (None = a whole chunk)
         "partition_kind",  # opt-in to PartitionColumns mode for Hive-style partitioning
         # Table-in-out specific: explicit override for the has_finalize auto-detection.
         # Set to True or False to force the emitted ``in_out_function_final``
@@ -1347,6 +1359,9 @@ def resolve_metadata(cls: type) -> ResolvedMetadata:
                 f"TableInput mode for a REQUIRED constant, or a named arg for optional config."
             )
 
+    _validate_resume_state(cls, attrs.get("max_workers"))
+    input_rows_per_call = _validate_input_rows_per_call(cls, attrs.get("input_rows_per_call"), input_from_args)
+
     # Merge annotation-derived setting/secret keys into required_settings/required_secrets
     meta_required_settings: list[str] = list(attrs.get("required_settings", []))
 
@@ -1452,9 +1467,24 @@ def resolve_metadata(cls: type) -> ResolvedMetadata:
         sink_order_dependent=bool(attrs.get("sink_order_dependent", False)),
         requires_input_batch_index=bool(attrs.get("requires_input_batch_index", False)),
         input_from_args=input_from_args,
+        input_rows_per_call=input_rows_per_call,
     )
 
 
+def _validate_input_rows_per_call(cls: type, value: Any, input_from_args: bool) -> int | None:
+    """Check ``Meta.input_rows_per_call``: a positive int, on a RowTransformFunction only."""
+    if value is None:
+        return None
+    if not input_from_args:
+        raise TypeError(
+            f"{cls.__name__}: Meta.input_rows_per_call applies only to a RowTransformFunction, whose "
+            f"per-row input the client can split into calls"
+        )
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValueError(f"{cls.__name__}: Meta.input_rows_per_call must be a positive int, got {value!r}")
+    return value
+
+
 def _validate_partition_kind(cls: type, kind: PartitionKind) -> PartitionKind:
     """Cross-check ``Meta.partition_kind`` against the bind schema.
 
@@ -1536,6 +1566,35 @@ def _detect_has_finalize(cls: type, function_type: CatalogFunctionType) -> bool:
     return cls.has_finalize_override()
 
 
+def _validate_resume_state(cls: type, max_workers: object) -> None:
+    """Refuse a table-in-out ``resume()`` the client cannot drive.
+
+    ``resume()`` receives no input, so it continues from ``state`` alone, and over
+    HTTP that state reaches the next turn only through the stream's state token.
+    Without a serializable state the cursor would be silently lost between turns.
+
+    ``max_workers = 1`` is refused too: the client then feeds every pipeline's
+    input through one shared connection, which cannot leave output pending while
+    another pipeline's input waits.
+    """
+    # Lazy imports to avoid a circular dependency.
+    from vgi.protocol import _resolve_state_type
+    from vgi.table_in_out_function import TableInOutGenerator
+
+    if not (isinstance(cls, type) and issubclass(cls, TableInOutGenerator)) or not cls.has_resume_override():
+        return
+    if max_workers == 1:
+        raise TypeError(
+            f"{cls.__name__}: resume() (continued output) is not supported with max_workers = 1; "
+            f"the client runs such a function on one shared connection."
+        )
+    if _resolve_state_type(cls) is None:
+        raise TypeError(
+            f"{cls.__name__}: resume() continues from its state alone, so the class must declare a "
+            f"serializable TState (an ArrowSerializableDataclass), e.g. RowTransformFunction[Args, State]."
+        )
+
+
 def _detect_input_from_args(cls: type) -> bool:
     """Whether ``cls`` is a blended ("UNNEST-style") table-in-out function.
 
@@ -1646,6 +1705,7 @@ _METADATA_SCHEMA = pa.schema(
         pa.field("sink_order_dependent", pa.bool_()),
         pa.field("requires_input_batch_index", pa.bool_()),
         pa.field("input_from_args", pa.bool_()),
+        pa.field("input_rows_per_call", pa.int64(), nullable=True),
     ]
 )