Key table-in-out state by substream, not by process

TableInOutFunction.process saved each connection's state under os.getpid().
Storage is scoped to the execution, which every connection of a fanned-out
scan shares, so the key has to tell those connections apart — and a pid only
does when each connection is its own process. Under the launcher, TCP or HTTP
one process serves them all: they overwrote each other's state and finish()
undercounted (substream_partial_sum: 194850 for 1999000).

The key is now InitRequest.substream_id, the client-minted id that already
exists to key per-substream state and that the DuckDB extension sends on every
substream. The Python client never sent one; it now mints one per fan-out
connection's INPUT stream (each its own substream of the one execution) and
carries the primary's on the finalize. A client that sends none is keyed by
pid, as before.

The partial-sum tests run over every client transport (the HTTP leg is one
process serving the whole fan-out), plus a check that a client sending no
substream_id still sums over subprocess.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: rustyconover

tests/table_in_out/generator/test_substream_partial_sum_function.py

@@ -2,29 +2,37 @@
 
 """Tests for SubstreamPartialSumFunction (parallel streaming finalize, A4).
 
-At the single-worker client level there is exactly one substream, so ``finish()``
-sees all input and emits one partial row equal to the whole sum. The per-substream
-fan-out (many partials re-aggregated by an outer ``SELECT sum()``) is exercised by
-the C++ integration test ``table_in_out/parallel_finalize.test`` over both
-transports; here we assert the worker-side transform/finish contract in isolation.
+The Python client fans the input across its connections, each its own substream
+of one execution, and finalizes once: ``finish()`` receives every substream's
+state and emits one partial row equal to the whole sum. The DuckDB-side fan-out
+(many partials re-aggregated by an outer ``SELECT sum()``) is exercised by the
+C++ integration test ``table_in_out/parallel_finalize.test`` over both transports.
+
+Run over every client transport because the state is keyed per substream: over
+HTTP one process serves every connection of the fan-out, which is where a
+per-process key let the connections overwrite each other's state.
 """
 
 from __future__ import annotations
 
+from typing import Any
+
 import pyarrow as pa
+import pytest
 
 from tests.conftest import filter_non_empty, make_schema
 from vgi.client import Client
+from vgi.client import client as client_module
 
 
 class TestSubstreamPartialSumFunction:
     """Worker-side contract for substream_partial_sum."""
 
-    def test_partial_sum_single_batch(self, fixture_worker: str) -> None:
+    def test_partial_sum_single_batch(self, client_transport: Any) -> None:
         """One substream (single client worker) → one partial = the whole sum."""
         schema = make_schema([pa.field("n", pa.int64())])
         batch = pa.RecordBatch.from_pydict({"n": [1, 2, 3, 4, 5]}, schema=schema)
-        with Client(fixture_worker) as client:
+        with client_transport() as client:
             output = list(
                 client.table_in_out_function(
                     function_name="substream_partial_sum",
@@ -36,14 +44,14 @@ def test_partial_sum_single_batch(self, fixture_worker: str) -> None:
         assert len(non_empty) == 1
         assert non_empty[0].to_pydict() == {"n": [15]}
 
-    def test_partial_sum_many_batches(self, fixture_worker: str) -> None:
+    def test_partial_sum_many_batches(self, client_transport: Any) -> None:
         """Accumulates across many input batches, emits one partial at finalize."""
         schema = make_schema([pa.field("n", pa.int64())])
         batches = [
             pa.RecordBatch.from_pydict({"n": list(range(i * 100, (i + 1) * 100))}, schema=schema) for i in range(20)
         ]
         expected = sum(range(0, 2000))
-        with Client(fixture_worker) as client:
+        with client_transport() as client:
             output = list(
                 client.table_in_out_function(
                     function_name="substream_partial_sum",
@@ -55,11 +63,11 @@ def test_partial_sum_many_batches(self, fixture_worker: str) -> None:
         assert len(non_empty) == 1
         assert non_empty[0].to_pydict() == {"n": [expected]}
 
-    def test_partial_sum_empty_input(self, fixture_worker: str) -> None:
+    def test_partial_sum_empty_input(self, client_transport: Any) -> None:
         """No input rows → the single substream's partial is 0."""
         schema = make_schema([pa.field("n", pa.int64())])
         empty = pa.RecordBatch.from_pydict({"n": []}, schema=schema)
-        with Client(fixture_worker) as client:
+        with client_transport() as client:
             output = list(
                 client.table_in_out_function(
                     function_name="substream_partial_sum",
@@ -71,3 +79,25 @@ def test_partial_sum_empty_input(self, fixture_worker: str) -> None:
         # A single substream with zero rows still finalizes → one 0 partial.
         assert len(non_empty) == 1
         assert non_empty[0].to_pydict() == {"n": [0]}
+
+    def test_a_client_sending_no_substream_id_still_sums_over_subprocess(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """An older client sends no ``substream_id``; the worker keys it by process, as before.
+
+        Over subprocess every connection is its own process, so that key still
+        separates them — the behavior such a client had before the key moved.
+        """
+        monkeypatch.setattr(client_module, "_substream_id_for", lambda phase: None)
+        schema = make_schema([pa.field("n", pa.int64())])
+        batches = [
+            pa.RecordBatch.from_pydict({"n": list(range(i * 100, (i + 1) * 100))}, schema=schema) for i in range(20)
+        ]
+        # A plain subprocess worker whatever transport the suite runs on.
+        with Client("vgi-fixture-worker") as client:
+            output = list(
+                client.table_in_out_function(
+                    function_name="substream_partial_sum",
+                    schema_path=["main"],
+                    input=iter(batches),
+                )
+            )
+        assert filter_non_empty(output)[0].to_pydict() == {"n": [sum(range(0, 2000))]}

vgi/_test_fixtures/table_in_out.py

@@ -138,9 +138,9 @@ class SubstreamPartialSumFunction(TableInOutFunction[SingleTableArguments, Subst
     their finalize outputs, so the caller re-aggregates with an outer
     ``SELECT sum(...)`` to get the global total — correct no matter how the rows
     were partitioned across substreams. Each substream's ``finish()`` reads only
-    its OWN worker's accumulated state (keyed by the substream's execution_id;
-    ``params.substream_id`` is the stable client-owned key available for workers
-    that manage cross-backend state themselves). This is the per-substream
+    its OWN execution's accumulated state — the framework keeps one row per
+    ``params.substream_id`` in storage scoped to that execution, so it is found
+    even when a finalize lands on a different backend. This is the per-substream
     finalize contract A4 enables — it is NOT a global cross-substream combine
     (that is a ``TableBufferingFunction``; see ``SumAllColumnsSimpleDistributed``).
 
@@ -191,8 +191,9 @@ def finish(
         params: ProcessParams[SingleTableArguments],
         states: list[SubstreamPartialSumState],
     ) -> list[pa.RecordBatch]:
-        # `states` are THIS substream's accumulated states (one per worker pid that
-        # handled this substream's batches); their sum is this substream's partial.
+        # `states` are this execution's accumulated states, one per substream that
+        # saw input (the Python client fans one execution across several); their
+        # sum is this finalize's partial.
         total = sum(st.total for st in states)
         name = params.output_schema.names[0]
         return [pa.RecordBatch.from_pydict({name: [total]}, schema=params.output_schema)]

vgi/client/client.py

@@ -189,6 +189,20 @@ def close(self) -> None:
 _LAUNCH_SCHEME = "launch:"
 
 _DEFAULT_ACCEPTED_MAX_RESPONSE_BYTES = 256 * 1024 * 1024
+
+
+def _substream_id_for(phase: TableInOutFunctionInitPhase | None) -> bytes | None:
+    """Mint the ``InitRequest.substream_id`` for a connection's table-in-out INPUT stream.
+
+    Each connection a table-in-out call fans out to is its own substream of the
+    one execution — the worker keeps one accumulated state per substream, and
+    the single finalize drains them all. Sharing an id (or sending none, which
+    the worker keys by process) lets connections served by one process
+    overwrite each other's state. ``None`` for every other kind of init.
+    """
+    return os.urandom(16) if phase == TableInOutFunctionInitPhase.INPUT else None
+
+
 _MIN_ACCEPTED_MAX_RESPONSE_BYTES = 65536
 _MAX_SAFE_HTTP_BYTES = (1 << 53) - 1
 
@@ -221,11 +235,14 @@ class WorkerConnection:
         stream: The active streaming session, if any.
         proc: The worker subprocess for direct (non-pooled) subprocess transport.
         connection: The RPC connection for direct subprocess transport.
+        substream_id: The ``InitRequest.substream_id`` this connection's
+            table-in-out stream was opened with, or None.
     """
 
     proxy: VgiProtocol
     worker_index: int = 0
     stream: StreamSession | None = None
+    substream_id: bytes | None = None
     # Subprocess transport, direct (non-pooled).
     proc: subprocess.Popen[bytes] | None = None
     connection: RpcConnection[VgiProtocol] | None = None
@@ -1637,6 +1654,7 @@ def _do_init(
         finalize_state_id: bytes | None = None,
         split_tokens: list[bytes] | None = None,
         join_keys: list[pa.RecordBatch] | None = None,
+        substream_id: bytes | None = None,
     ) -> StreamSession:
         """Call init on a worker proxy and return a `StreamSession`.
 
@@ -1661,6 +1679,8 @@ def _do_init(
                 looked up by column name — see
                 ``PushdownFilters.get_join_keys_column``), or `None` when not
                 applicable.
+            substream_id: The table-in-out substream this init opens or
+                finalizes (see ``_substream_id_for``), or `None`.
 
         Returns:
             `StreamSession` for data exchange or production.
@@ -1681,6 +1701,7 @@ def _do_init(
             init_opaque_data=init_opaque_data,
             finalize_state_id=finalize_state_id,
             split_tokens=split_tokens,
+            substream_id=substream_id,
         )
         try:
             stream: StreamSession = proxy.init(request=init_request)  # type: ignore[assignment]
@@ -1769,6 +1790,7 @@ def _initialize_stream_common(
         )
         bind_response = self._do_bind(self._primary.proxy, bind_request, bind_result_callback)
 
+        substream_id = _substream_id_for(phase)
         stream = self._do_init(
             self._primary.proxy,
             bind_request,
@@ -1780,8 +1802,10 @@ def _initialize_stream_common(
             execution_id=split_execution_id,
             init_opaque_data=split_init_opaque_data,
             join_keys=join_keys,
+            substream_id=substream_id,
         )
         self._primary.stream = stream
+        self._primary.substream_id = substream_id
 
         init_response = stream.typed_header(GlobalInitResponse)
         max_workers = 1 if split_tokens is not None else self._determine_max_workers(init_response.max_workers)
@@ -1857,6 +1881,7 @@ def _spawn_additional_workers(
 
         def do_init(worker: WorkerConnection) -> None:
             try:
+                worker.substream_id = _substream_id_for(phase)
                 stream = self._do_init(
                     worker.proxy,
                     bind_request,
@@ -1867,6 +1892,7 @@ def do_init(worker: WorkerConnection) -> None:
                     join_keys=join_keys,
                     execution_id=global_init_response.execution_id,
                     init_opaque_data=global_init_response.opaque_data,
+                    substream_id=worker.substream_id,
                 )
                 worker.stream = stream
             except Exception as e:
@@ -2648,6 +2674,7 @@ def _finalize_primary_worker(
             phase=TableInOutFunctionInitPhase.FINALIZE,
             execution_id=init_response.execution_id,
             init_opaque_data=init_response.opaque_data,
+            substream_id=self._primary.substream_id,
         )
 
         try:

vgi/table_in_out_function.py

@@ -360,7 +360,8 @@ def finish(
 
         Args:
             params: The process parameters — function args, settings, secrets.
-            states: The accumulated per-partition states from ``transform()``.
+            states: The accumulated states from ``transform()``, one per
+                substream of this execution that saw input.
 
         Returns:
             List of pa.RecordBatch to emit as final output.
@@ -412,11 +413,17 @@ def process(
         """
         result = cls.transform(batch, params, state)
 
-        # Save state for distributed processing (upsert semantics)
+        # Save state for distributed processing (upsert semantics), one row per
+        # substream. Storage is scoped to the execution, which every connection
+        # of a fanned-out scan shares, and one process serves many of those
+        # connections under the launcher, TCP or HTTP — so a per-process key let
+        # them overwrite each other and finish() undercounted. The pid remains
+        # only for a client that sends no substream_id, which it keys as before.
         if state is not None:
-            params.storage.state_put(
-                FrameworkNS.TIO_STATE, BoundStorage.pack_int_key(os.getpid()), state.serialize_to_bytes()
-            )
+            key = params.substream_id
+            if key is None:
+                key = BoundStorage.pack_int_key(os.getpid())
+            params.storage.state_put(FrameworkNS.TIO_STATE, key, state.serialize_to_bytes())
 
         # Handle single batch or list of batches — exchange must emit exactly one
         if isinstance(result, list):