Skip to content

Commit 724b19c

Browse files
committed
feat(client): expose splits via table_function_plan + split_tokens
Splits were entirely unreachable from Client — no wrapper for on_plan() existed at all, and table_function()/_do_init had no split_tokens parameter to redeem one, even though the wire protocol (TableFunctionPlanRequest/PlanResponse/InitRequest.split_tokens) and the VgiProtocol.table_function_plan RPC already supported it. The DuckDB C++ extension drives this today; the reference Python client never grew the equivalent. Adds: - Client.table_function_plan(...): bind() + on_plan(), deserializing the wire's list[bytes] into typed ScanSplit objects (PlanResponse.splits is documented as carrying either shape across its lifetime). - Client.table_function(split_tokens=..., split_execution_id=..., split_init_opaque_data=...): redeems specific split tokens instead of an ordinary whole-scan init, forcing single-worker mode (the server's max_workers header describes whole-table fan-out, not one already-named unit of work) and echoing the plan's execution_id/init_opaque_data so a worker whose splits share cross-process BoundStorage state can find it. Additive only — no protocol/wire version change, since every field and RPC method used here already existed on the wire; only the Python client hadn't caught up. tests/client/test_split_client.py: 12 tests against the split_sequence/ split_zero fixtures, all 3 client transports (subprocess-pooled, subprocess-direct, http). Zero regressions in the existing suite.
1 parent 550d0b4 commit 724b19c

2 files changed

Lines changed: 273 additions & 3 deletions

File tree

‎tests/client/test_split_client.py‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# Copyright 2026 Query Farm LLC - https://query.farm
2+
3+
# ruff: noqa: S101, D101, D102, D103
4+
"""`Client.table_function_plan` + `Client.table_function(split_tokens=...)`.
5+
6+
Splits were previously entirely unreachable from `Client` — no wrapper for
7+
`on_plan()` existed at all, and `table_function()`/`_do_init` had no
8+
`split_tokens` parameter to redeem one, even though the wire protocol
9+
(`TableFunctionPlanRequest`/`PlanResponse`/`InitRequest.split_tokens`) and the
10+
`VgiProtocol.table_function_plan` RPC already supported it — the DuckDB C++
11+
extension drives this today, but the reference Python client never grew the
12+
equivalent. Exercised against `split_sequence` (the split-capable twin of
13+
`sequence`, `vgi/_test_fixtures/table/splits.py`) over every client transport.
14+
"""
15+
16+
from __future__ import annotations
17+
18+
from typing import Any
19+
20+
import pyarrow as pa
21+
import pytest
22+
23+
from vgi.arguments import Arguments
24+
from vgi.protocol import ScanSplit
25+
26+
MAIN = "main"
27+
28+
29+
@pytest.fixture
30+
def client(client_transport: Any) -> Any:
31+
with client_transport() as c:
32+
c.catalog_attach(name="example", data_version_spec=None, implementation_version=None)
33+
yield c
34+
35+
36+
def test_plan_returns_typed_scan_splits(client: Any) -> None:
37+
args = Arguments(named={"n": pa.scalar(20), "splits": pa.scalar(4)})
38+
plan = client.table_function_plan(function_name="split_sequence", schema_name=MAIN, arguments=args)
39+
assert len(plan.splits) == 4
40+
for split in plan.splits:
41+
assert isinstance(split, ScanSplit)
42+
assert split.token # framework-stamped, non-empty
43+
assert split.estimated_rows == 5
44+
45+
46+
def test_redeeming_each_split_in_order_reproduces_the_whole_scan(client: Any) -> None:
47+
"""split_sequence redeemed split-by-split must equal sequence(n)'s output.
48+
49+
Also threads ``split_execution_id``/``split_init_opaque_data`` from the
50+
plan on every redemption — this fixture doesn't need cross-split shared
51+
state to be correct, but this proves the wire path accepts and round-trips
52+
them without erroring, for workers that do.
53+
"""
54+
args = Arguments(named={"n": pa.scalar(37), "splits": pa.scalar(5)})
55+
plan = client.table_function_plan(function_name="split_sequence", schema_name=MAIN, arguments=args)
56+
57+
all_rows: list[int] = []
58+
for split in plan.splits:
59+
batches = list(
60+
client.table_function(
61+
function_name="split_sequence",
62+
schema_name=MAIN,
63+
arguments=args,
64+
split_tokens=[split.token],
65+
split_execution_id=plan.execution_id,
66+
split_init_opaque_data=plan.init_opaque_data,
67+
)
68+
)
69+
for batch in batches:
70+
all_rows.extend(batch.column("n").to_pylist())
71+
72+
assert sorted(all_rows) == list(range(37))
73+
74+
75+
def test_split_zero_yields_no_splits_and_no_rows(client: Any) -> None:
76+
args = Arguments(named={"n": pa.scalar(10), "splits": pa.scalar(4)})
77+
plan = client.table_function_plan(function_name="split_zero", schema_name=MAIN, arguments=args)
78+
assert plan.splits == []
79+
80+
81+
def test_redeeming_a_split_forces_single_worker(client: Any) -> None:
82+
"""`split_tokens` bypasses `_spawn_additional_workers` — no additional worker connections open."""
83+
args = Arguments(named={"n": pa.scalar(8), "splits": pa.scalar(2)})
84+
plan = client.table_function_plan(function_name="split_sequence", schema_name=MAIN, arguments=args)
85+
list(
86+
client.table_function(
87+
function_name="split_sequence",
88+
schema_name=MAIN,
89+
arguments=args,
90+
split_tokens=[plan.splits[0].token],
91+
)
92+
)
93+
assert client._additional_workers == []

‎vgi/client/client.py‎

Lines changed: 180 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
client.schema_contents(...) — list tables/views/functions/macros
5353
client.scalar_function(...) — invoke a scalar
5454
client.table_function(...) — invoke a table function
55+
client.table_function_plan(...) — plan a table function into named, redeemable splits
5556
client.table_in_out_function(...) — invoke a table-in-out function
5657
client.aggregate_function(...) — run a grouped/global aggregate
5758
client.aggregate_session(...) — raw aggregate RPCs (update/combine/window)
@@ -81,7 +82,7 @@
8182
import threading
8283
from collections.abc import Callable, Generator, Iterator, Sequence
8384
from contextlib import AbstractContextManager
84-
from dataclasses import dataclass, field
85+
from dataclasses import dataclass, field, replace
8586
from queue import Queue
8687
from typing import IO, Any, Literal, cast
8788

@@ -110,9 +111,12 @@
110111
CopyFromContext,
111112
CopyToContext,
112113
InitRequest,
114+
PlanResponse,
115+
ScanSplit,
113116
TableBufferingCombineRequest,
114117
TableBufferingDestructorRequest,
115118
TableBufferingProcessRequest,
119+
TableFunctionPlanRequest,
116120
VgiProtocol,
117121
)
118122
from vgi.table_function import TableInOutFunctionInitPhase
@@ -1144,6 +1148,7 @@ def _do_init(
11441148
execution_id: bytes | None = None,
11451149
init_opaque_data: bytes | None = None,
11461150
finalize_state_id: bytes | None = None,
1151+
split_tokens: list[bytes] | None = None,
11471152
) -> StreamSession:
11481153
"""Call init on a worker proxy and return a `StreamSession`.
11491154
@@ -1160,6 +1165,9 @@ def _do_init(
11601165
the primary worker's init response.
11611166
finalize_state_id: For ``TABLE_BUFFERING_FINALIZE`` init, the
11621167
opaque finalize partition key this producer stream serves.
1168+
split_tokens: Split tokens (from a prior `table_function_plan()`
1169+
call) to redeem in this init, or `None` for an ordinary
1170+
whole-scan init.
11631171
11641172
Returns:
11651173
`StreamSession` for data exchange or production.
@@ -1178,6 +1186,7 @@ def _do_init(
11781186
execution_id=execution_id,
11791187
init_opaque_data=init_opaque_data,
11801188
finalize_state_id=finalize_state_id,
1189+
split_tokens=split_tokens,
11811190
)
11821191
try:
11831192
stream: StreamSession = proxy.init(request=init_request) # type: ignore[assignment]
@@ -1201,6 +1210,9 @@ def _initialize_stream_common(
12011210
phase: TableInOutFunctionInitPhase | None,
12021211
bind_result_callback: Callable[[BindResponse], None] | None,
12031212
copy_from: CopyFromContext | None = None,
1213+
split_tokens: list[bytes] | None = None,
1214+
split_execution_id: bytes | None = None,
1215+
split_init_opaque_data: bytes | None = None,
12041216
) -> tuple[BindRequest, BindResponse, GlobalInitResponse]:
12051217
"""Run the canonical bind → init → fan-out-workers sequence.
12061218
@@ -1219,6 +1231,24 @@ def _initialize_stream_common(
12191231
Centralizing this keeps HTTP/subprocess differences and protocol
12201232
changes (e.g. future scoped-secret re-bind, init hints) in one
12211233
place.
1234+
1235+
``split_tokens``, when given, redeems those specific named units of
1236+
work from a prior ``table_function_plan()`` call instead of an
1237+
ordinary whole-scan init — and forces single-worker mode (skips
1238+
``_spawn_additional_workers`` by clamping ``max_workers`` to 1): the
1239+
server-advertised ``max_workers`` header describes fan-out for
1240+
reading the *whole* table, which doesn't apply to redeeming one
1241+
already-named unit of work. A caller wanting split-level parallelism
1242+
drives multiple splits through multiple `Client`/thread instances
1243+
itself (mirroring how `VgiCatalog._exchange_client()` in vgi-polars
1244+
gives each thread its own `Client`), not via this method's own
1245+
worker fan-out. ``split_execution_id``/``split_init_opaque_data``
1246+
(from that same ``PlanResponse``) are echoed on this init exactly
1247+
like a secondary worker's would be — ``PlanResponse.execution_id``'s
1248+
docstring: "echoed on every split init. Scopes cross-process
1249+
BoundStorage exactly as it does elsewhere" — so a worker whose splits
1250+
share cross-process state via ``BoundStorage`` can find it, even
1251+
though this is a *first* (not secondary) init for this connection.
12221252
"""
12231253
assert self._primary is not None, "primary worker not started"
12241254

@@ -1242,11 +1272,14 @@ def _initialize_stream_common(
12421272
projection_ids=projection_ids,
12431273
pushdown_filters_batch=pushdown_filters_batch,
12441274
phase=phase,
1275+
split_tokens=split_tokens,
1276+
execution_id=split_execution_id,
1277+
init_opaque_data=split_init_opaque_data,
12451278
)
12461279
self._primary.stream = stream
12471280

12481281
init_response = stream.typed_header(GlobalInitResponse)
1249-
max_workers = self._determine_max_workers(init_response.max_workers)
1282+
max_workers = 1 if split_tokens is not None else self._determine_max_workers(init_response.max_workers)
12501283

12511284
self._spawn_additional_workers(
12521285
max_workers,
@@ -2014,6 +2047,125 @@ def _finalize_primary_worker(
20142047
finally:
20152048
finalize_stream.close()
20162049

2050+
def table_function_plan(
2051+
self,
2052+
*,
2053+
function_name: str,
2054+
schema_name: str,
2055+
arguments: Arguments | None = None,
2056+
projection_ids: list[int] | None = None,
2057+
pushdown_filters: bytes | None = None,
2058+
settings: dict[str, Any] | None = None,
2059+
transaction_opaque_data: bytes | None = None,
2060+
target_split_bytes: int | None = None,
2061+
min_splits: int | None = None,
2062+
max_splits_per_response: int | None = None,
2063+
cursor: bytes | None = None,
2064+
) -> PlanResponse:
2065+
"""Plan a table-function scan into named, independently redeemable splits.
2066+
2067+
Runs `bind()` then `on_plan()` and returns the resulting
2068+
[`PlanResponse`][], whose `splits` are each individually redeemable by
2069+
:meth:`table_function` via its `split_tokens` argument — from this
2070+
process or, since a split names work rather than describing it ("these
2071+
three files at version 47", not "rows 0-999 of whatever this returns
2072+
now"), any other. Workers that don't opt in via `supports_splits`
2073+
(`FunctionInfo.supports_splits`) inherit a framework default (commonly
2074+
one split for the whole scan) — check that flag first if the caller
2075+
cares whether real parallelism/checkpointing is available versus a
2076+
single degenerate split.
2077+
2078+
A response's `next_cursors` is normally empty or one entry; more than
2079+
one means the plan is paginated across parallel, disjoint enumeration
2080+
branches — the *caller* is responsible for that disjointness (VGI
2081+
itself does not verify it; see the "Split disjointness is a worker
2082+
contract" note in the VGI extension's own docs). For a single
2083+
sequential caller, following `next_cursors` one at a time (via
2084+
`cursor=`) and concatenating each response's `splits` is always
2085+
correct, whether the plan doled out one cursor or several.
2086+
2087+
Args:
2088+
function_name: Name of the table function to plan.
2089+
schema_name: Name of the catalog schema that declares the function.
2090+
arguments: Optional [`Arguments`][] container. Defaults to empty
2091+
`Arguments()`.
2092+
projection_ids: Optional list of column indices for projection —
2093+
threaded into the plan so split sizing can account for it.
2094+
pushdown_filters: Optional byte string of filter predicates,
2095+
same wire format as :meth:`table_function`'s.
2096+
settings: Optional dictionary of settings/pragmas.
2097+
transaction_opaque_data: Optional transaction identifier.
2098+
target_split_bytes: Requested split size — the primary sizing
2099+
lever; the client can't see per-split cost and will treat
2100+
returned splits as interchangeable units.
2101+
min_splits: Parallelism floor — ask for at least this many splits
2102+
even for a small table, so a caller with idle readers has
2103+
enough units to hand them.
2104+
max_splits_per_response: Pagination cap on this one response
2105+
(distinct from `min_splits`, which is a sizing hint, not a
2106+
pagination control).
2107+
cursor: Resume point from a previous response's `next_cursors`,
2108+
or `None` to start a fresh plan.
2109+
2110+
Returns:
2111+
[`PlanResponse`][] with `splits` (each an individually redeemable
2112+
[`ScanSplit`][], carrying the `token` to pass back into
2113+
`table_function(split_tokens=...)`) and `next_cursors` for
2114+
pagination.
2115+
2116+
Raises:
2117+
[`ClientError`][]: If the client is not started, communication with
2118+
the worker fails, or the worker returns an exception.
2119+
2120+
"""
2121+
if arguments is None:
2122+
arguments = Arguments()
2123+
2124+
if self._primary is None:
2125+
raise ClientError("Client not started. Call start() or use context manager.")
2126+
2127+
try:
2128+
pushdown_filters_batch = self._deserialize_pushdown_filters(pushdown_filters)
2129+
2130+
bind_request = self._make_bind_request(
2131+
function_name=function_name,
2132+
schema_name=schema_name,
2133+
arguments=arguments,
2134+
function_type=FunctionType.TABLE,
2135+
settings=settings,
2136+
secrets=None,
2137+
transaction_opaque_data=transaction_opaque_data,
2138+
)
2139+
bind_response = self._do_bind(self._primary.proxy, bind_request, None)
2140+
2141+
plan_request = TableFunctionPlanRequest(
2142+
bind_call=bind_request,
2143+
bind_opaque_data=bind_response.opaque_data,
2144+
projection_ids=projection_ids,
2145+
pushdown_filters=pushdown_filters_batch,
2146+
target_split_bytes=target_split_bytes,
2147+
min_splits=min_splits,
2148+
max_splits_per_response=max_splits_per_response,
2149+
cursor=cursor,
2150+
)
2151+
try:
2152+
response: PlanResponse = self._primary.proxy.table_function_plan(request=plan_request) # type: ignore[assignment]
2153+
except RpcError as e:
2154+
raise ClientError.from_rpc_error(e) from e
2155+
2156+
# The wire carries each split as a serialized blob in a
2157+
# list<binary> column (see PlanResponse.splits' docstring: "Both
2158+
# [ScanSplit objects and serialized blobs] inhabit this field
2159+
# across its lifetime") — deserialize them here so callers get
2160+
# typed ScanSplit objects (with .token ready to redeem) rather
2161+
# than raw bytes they'd all have to know to decode themselves.
2162+
deserialized_splits = [
2163+
s if isinstance(s, ScanSplit) else ScanSplit.deserialize_from_bytes(s) for s in response.splits
2164+
]
2165+
return replace(response, splits=deserialized_splits)
2166+
except ClientError as e:
2167+
raise self._client_error_with_stderr(e) from e.__cause__
2168+
20172169
def table_function(
20182170
self,
20192171
*,
@@ -2026,6 +2178,9 @@ def table_function(
20262178
settings: dict[str, Any] | None = None,
20272179
transaction_opaque_data: bytes | None = None,
20282180
copy_from: CopyFromContext | None = None,
2181+
split_tokens: list[bytes] | None = None,
2182+
split_execution_id: bytes | None = None,
2183+
split_init_opaque_data: bytes | None = None,
20292184
) -> Generator[pa.RecordBatch]:
20302185
"""Invoke a table function (source function) and stream output batches.
20312186
@@ -2035,6 +2190,8 @@ def table_function(
20352190
20362191
For parallel processing (max_workers > 1), output is read from all
20372192
workers concurrently using threads. Output order is non-deterministic.
2193+
This is unrelated to (and mutually exclusive in effect with) redeeming
2194+
splits — see ``split_tokens`` below.
20382195
20392196
Args:
20402197
function_name: Name of the function to invoke. Must exist in the
@@ -2055,10 +2212,27 @@ def table_function(
20552212
copy_from: Optional [`CopyFromContext`][] marking this scan as a
20562213
``COPY ... FROM`` read. Prefer :meth:`copy_from`, which builds
20572214
the context for you.
2215+
split_tokens: Redeem these specific split tokens (from a prior
2216+
:meth:`table_function_plan` call) instead of an ordinary
2217+
whole-scan init. Forces single-worker mode for this call — the
2218+
server's advertised ``max_workers`` describes fan-out for
2219+
reading the whole table, not for one already-named unit of
2220+
work. To read multiple splits in parallel, drive them through
2221+
multiple `Client`/thread instances yourself; to read them
2222+
sequentially (still sound and replayable, just without
2223+
concurrency), call this once per split token in a loop.
2224+
split_execution_id: When redeeming a split, the originating
2225+
``PlanResponse.execution_id`` — echoed on this init so a
2226+
worker whose splits share cross-process state via
2227+
``BoundStorage`` can find it. `None` for an ordinary
2228+
whole-scan init.
2229+
split_init_opaque_data: When redeeming a split, the originating
2230+
``PlanResponse.init_opaque_data``, echoed the same way.
20582231
20592232
Yields:
20602233
Output `RecordBatch`es from the function. In parallel mode
2061-
(max_workers > 1), output order is non-deterministic.
2234+
(max_workers > 1, only possible when ``split_tokens`` is `None`),
2235+
output order is non-deterministic.
20622236
20632237
Raises:
20642238
[`ClientError`][]: If the client is not started, communication with the
@@ -2088,6 +2262,9 @@ def table_function(
20882262
phase=None,
20892263
bind_result_callback=bind_result_callback,
20902264
copy_from=copy_from,
2265+
split_tokens=split_tokens,
2266+
split_execution_id=split_execution_id,
2267+
split_init_opaque_data=split_init_opaque_data,
20912268
)
20922269

20932270
# Read output from all workers in parallel

0 commit comments

Comments
 (0)