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)
8182import threading
8283from collections .abc import Callable , Generator , Iterator , Sequence
8384from contextlib import AbstractContextManager
84- from dataclasses import dataclass , field
85+ from dataclasses import dataclass , field , replace
8586from queue import Queue
8687from typing import IO , Any , Literal , cast
8788
110111 CopyFromContext ,
111112 CopyToContext ,
112113 InitRequest ,
114+ PlanResponse ,
115+ ScanSplit ,
113116 TableBufferingCombineRequest ,
114117 TableBufferingDestructorRequest ,
115118 TableBufferingProcessRequest ,
119+ TableFunctionPlanRequest ,
116120 VgiProtocol ,
117121)
118122from 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