fix(client): clean up parallel table-function scans closed early

_table_function_parallel reads each worker's producer stream on a reader
thread. Closing the generator before the scan finished left those threads
reading and the streams open; the Client.stop() that followed drained the
same stream from the main thread. Two readers on one pipe hung about one
run in five, and a pooled worker abandoned mid-stream corrupted the next
call on it (TransportError reading `bind`).

The generator body now runs in try/finally. On an early exit (generator
closed, or a worker failed) it sets a stop event; each reader notices it
within one batch (producer streams advance one tick per batch) and cancels
its stream from its own thread, so the stream never has two readers. The
threads are joined, then streams and secondary workers released as on
normal completion. A reader still alive after the join timeout gets its
worker killed rather than a second reader on the pipe.

Found by vgi-polars (scans cut short by .head() across splits).

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

Author: rustyconover

tests/test_table_function_early_close.py

@@ -0,0 +1,81 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Closing a ``table_function`` generator before its stream ends must not hang.
+
+``Client.table_function`` reads each worker's producer stream on a background
+thread. Before the fix, closing the generator early left those threads reading,
+and the ``Client.stop()`` that followed drained the same stream from the main
+thread: two readers on one pipe, and whichever missed the end of stream blocked
+forever (roughly one run in five). Each case here repeats the early close under
+a watchdog, then checks the pooled worker still serves a complete scan.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable
+
+import pyarrow as pa
+import pytest
+
+from tests.conftest import SUBPROCESS_FIXTURE_WORKER
+from vgi.arguments import Arguments
+from vgi.client.client import Client, _default_pool
+
+_ITERATIONS = 25
+_WATCHDOG_SECONDS = 120.0
+
+
+def _run_with_watchdog(body: Callable[[], None]) -> None:
+    """Run ``body`` on a daemon thread; fail (rather than hang the suite) if it doesn't finish."""
+    errors: list[BaseException] = []
+
+    def target() -> None:
+        try:
+            body()
+        except BaseException as e:  # noqa: BLE001 - re-raised on the test thread below
+            errors.append(e)
+
+    thread = threading.Thread(target=target, daemon=True)
+    thread.start()
+    thread.join(_WATCHDOG_SECONDS)
+    assert not thread.is_alive(), "closing a table_function scan early hung"
+    if errors:
+        raise errors[0]
+
+
+def _close_after_first_batch(function_name: str, arguments: Arguments) -> None:
+    for _ in range(_ITERATIONS):
+        with Client(SUBPROCESS_FIXTURE_WORKER, pool=_default_pool) as client:
+            scan = client.table_function(function_name=function_name, schema_path=["main"], arguments=arguments)
+            first = next(scan)
+            assert first.num_rows > 0
+            scan.close()
+
+
+def _full_scan_still_works() -> None:
+    with Client(SUBPROCESS_FIXTURE_WORKER, pool=_default_pool) as client:
+        batches = list(
+            client.table_function(
+                function_name="sequence",
+                schema_path=["main"],
+                arguments=Arguments(positional=(pa.scalar(4),)),
+            )
+        )
+    assert pa.Table.from_batches(batches).column("n").to_pylist() == [0, 1, 2, 3]
+
+
+@pytest.mark.parametrize(
+    ("function_name", "arguments"),
+    [
+        # One worker, many small batches.
+        ("sequence", Arguments(positional=(pa.scalar(100_000),), named={"batch_size": pa.scalar(10)})),
+        # Several workers (one per partition), each with batches still to come.
+        ("partitioned_sequence", Arguments(positional=(pa.scalar(30_000),))),
+    ],
+    ids=["single_worker", "multi_worker"],
+)
+def test_early_close_does_not_hang(function_name: str, arguments: Arguments) -> None:
+    """A scan closed after its first batch stops cleanly, and the worker pool stays usable."""
+    _run_with_watchdog(lambda: _close_after_first_batch(function_name, arguments))
+    _run_with_watchdog(_full_scan_still_works)

vgi/client/client.py

@@ -3273,22 +3273,34 @@ def _table_function_parallel(
         # AnnotatedBatch (not just .batch) so custom_metadata survives the
         # cross-thread handoff for batch_metadata_callback below.
         output_queue: Queue[AnnotatedBatch | BaseException | None] = Queue()
+        # Set when the consumer stops before every worker finished (the
+        # generator was closed early, or a worker failed). Each reader thread is
+        # its stream's only reader, so it is also the one that ends the stream:
+        # a producer stream advances one tick per batch, so a reader notices the
+        # flag within one batch and cancels from its own thread. Ending it from
+        # the consumer's thread instead would put two readers on one pipe.
+        stop = threading.Event()
 
         def read_worker_output(worker: WorkerConnection) -> None:
             """Thread function that reads all output from a single worker."""
             try:
-                if worker.stream is None:
+                stream = worker.stream
+                if stream is None:
                     output_queue.put(None)
                     return
 
-                for output in worker.stream:
+                for output in stream:
                     _logger.debug(
                         "received_output_from_worker worker_index=%s num_rows=%s",
                         worker.worker_index,
                         output.batch.num_rows,
                     )
                     if output.batch.num_rows > 0:
                         output_queue.put(output)
+                    if stop.is_set():
+                        _logger.debug("table_function_worker_cancelled worker_index=%s", worker.worker_index)
+                        stream.cancel()
+                        break
 
                 output_queue.put(None)  # Signal completion
             except StopIteration:
@@ -3308,41 +3320,59 @@ def read_worker_output(worker: WorkerConnection) -> None:
             thread.start()
             threads.append(thread)
 
-        # Collect outputs from all workers until all are done
-        workers_finished = 0
-        while workers_finished < num_workers:
-            result = output_queue.get()
-
-            # Check for exceptions from worker threads
-            if isinstance(result, BaseException):
-                if isinstance(result, RpcError):
-                    raise ClientError.from_rpc_error(result) from result
-                raise ClientError(f"Worker thread failed: {result}") from result
-
-            # None signals a worker finished
-            if result is None:
-                workers_finished += 1
-                _logger.debug(
-                    "worker_finished workers_finished=%s total_workers=%s",
-                    workers_finished,
-                    num_workers,
-                )
-                continue
-
-            if batch_metadata_callback is not None:
-                batch_metadata_callback(result.custom_metadata)
-            yield result.batch
-
-        self._join_threads(threads)
-        _logger.debug("all_table_function_workers_complete")
+        # Collect outputs from all workers until all are done. The finally
+        # block also runs when the consumer closes the generator early
+        # (GeneratorExit at the yield) or a worker fails: the reader threads are
+        # told to stop and joined before any stream is touched here.
+        completed = False
+        try:
+            workers_finished = 0
+            while workers_finished < num_workers:
+                result = output_queue.get()
+
+                # Check for exceptions from worker threads
+                if isinstance(result, BaseException):
+                    if isinstance(result, RpcError):
+                        raise ClientError.from_rpc_error(result) from result
+                    raise ClientError(f"Worker thread failed: {result}") from result
+
+                # None signals a worker finished
+                if result is None:
+                    workers_finished += 1
+                    _logger.debug(
+                        "worker_finished workers_finished=%s total_workers=%s",
+                        workers_finished,
+                        num_workers,
+                    )
+                    continue
 
-        # Close streams and secondary workers
-        for worker in all_workers:
-            if worker.stream is not None:
-                worker.stream.close()
-                worker.stream = None
-        self._close_secondary_workers()
-        _logger.debug("parallel_table_function_complete")
+                if batch_metadata_callback is not None:
+                    batch_metadata_callback(result.custom_metadata)
+                yield result.batch
+            completed = True
+        finally:
+            if not completed:
+                stop.set()
+            self._join_threads(threads)
+            _logger.debug("all_table_function_workers_complete completed=%s", completed)
+
+            # Close streams and secondary workers. A stream a reader cancelled is
+            # already closed, so close() is a no-op for it.
+            for worker, thread in zip(all_workers, threads, strict=True):
+                if thread.is_alive():
+                    # Still reading after the join timeout: the stream's state is
+                    # unknown and its reader still owns it. Kill the worker so
+                    # neither that reader nor a later stop() blocks on the pipe.
+                    _logger.warning("table_function_worker_killed worker_index=%s", worker.worker_index)
+                    if worker.proc is not None:
+                        worker.proc.kill()
+                    worker.stream = None
+                    continue
+                if worker.stream is not None:
+                    worker.stream.close()
+                    worker.stream = None
+            self._close_secondary_workers()
+            _logger.debug("parallel_table_function_complete")
 
     # ==================================================================
     # Custom COPY formats