@@ -29,55 +29,94 @@ VGI (Vector Gateway Interface) provides an Apache Arrow-based protocol for conne
2929
3030## Architecture
3131
32+ VGI supports two types of functions:
33+
3234```
3335┌─────────────────────────────────────────────────────────────────────┐
3436│ DuckDB / Client │
3537│ ┌───────────────────────────────────────────────────────────────┐ │
36- │ │ Client spawns worker subprocess, sends FunctionInvocation, streams │ │
37- │ │ input batches, receives output batches via Arrow IPC │ │
38+ │ │ Client spawns worker subprocess, sends Invocation, │ │
39+ │ │ streams input batches (if any) , receives output via Arrow IPC │ │
3840│ └───────────────────────────────────────────────────────────────┘ │
3941│ │ stdin/stdout │
4042│ ▼ │
4143│ ┌───────────────────────────────────────────────────────────────┐ │
4244│ │ Worker Process │ │
45+ │ │ │ │
46+ │ │ ┌─────────────────────────────────────────────────────────┐ │ │
47+ │ │ │ TABLE FUNCTION (TableFunctionGenerator) │ │ │
48+ │ │ │ - process(): Generator yielding output batches │ │ │
49+ │ │ │ - No input data, generates output (sequences, ranges) │ │ │
50+ │ │ └─────────────────────────────────────────────────────────┘ │ │
51+ │ │ OR │ │
4352│ │ ┌─────────────────────────────────────────────────────────┐ │ │
44- │ │ │ TableInOutGeneratorFunction.process() / finalize() │ │ │
45- │ │ │ - process( ): Generator receiving RecordBatch via yield │ │ │
46- │ │ │ - finalize (): Generator emitting final results │ │ │
47- │ │ │ - Yields Output with output RecordBatches │ │ │
53+ │ │ │ TABLE-IN-OUT FUNCTION (TableInOutFunction) │ │ │
54+ │ │ │ - transform(batch ): Process each input batch │ │ │
55+ │ │ │ - finish (): Emit final results after all input │ │ │
56+ │ │ │ - Transforms input data to output data │ │ │
4857│ │ └─────────────────────────────────────────────────────────┘ │ │
58+ │ │ │ │
4959│ └───────────────────────────────────────────────────────────────┘ │
5060└─────────────────────────────────────────────────────────────────────┘
5161```
5262
63+ ### Function Types
64+
65+ | Type | Base Class | Input | Use Case |
66+ | ------| ------------| -------| ----------|
67+ | ** Table Function** | ` TableFunctionGenerator ` | None | Generate data (sequences, ranges, random samples) |
68+ | ** Table-In-Out Function** | ` TableInOutFunction ` | Batches | Transform, filter, aggregate input data |
69+
5370### Key Components
5471
5572- ** Worker** (` vgi/worker.py ` ): Subprocess that hosts functions, handles protocol
56- - ** Client** (` vgi/client.py ` ): Spawns workers, streams data through functions
57- - ** TableInOutGeneratorFunction** (` vgi/table_in_out_function.py ` ): Base class for table-in-out functions
58- - ** FunctionInvocation/FunctionOutputSpec** (` vgi/function.py ` ): Protocol messages for initialization
73+ - ** Client** (` vgi/client/client.py ` ): Spawns workers, streams data through functions
74+ - ** TableFunctionGenerator** (` vgi/table_function.py ` ): Base class for table functions (no input)
75+ - ** TableInOutFunction** (` vgi/table_in_out_function.py ` ): Base class for table-in-out functions
76+ - ** Invocation/OutputSpec** (` vgi/function.py ` ): Protocol messages for initialization
5977- ** GlobalInitResult** (` vgi/function.py ` ): Shared state for parallel workers
6078
6179## Protocol Flow
6280
81+ ### Table Function (no input)
82+
83+ ```
84+ Client Worker
85+ │ │
86+ │──── Invocation (function, args) ─────▶│
87+ │ │ instantiate function
88+ │◀──── OutputSpec (output schema) ──────│
89+ │ │
90+ │──── GlobalStateInitInput ────────────▶│
91+ │◀──── GlobalInitResult ────────────────│ perform_init()
92+ │ │
93+ │◀──── Output Batch 1 ──────────────────│ process() yields
94+ │◀──── Output Batch 2 ──────────────────│
95+ │◀──── ... ─────────────────────────────│
96+ │◀──── Final Output (FINISHED) ─────────│
97+ │ │
98+ ```
99+
100+ ### Table-In-Out Function (with input)
101+
63102```
64103Client Worker
65104 │ │
66- │──── FunctionInvocation (function, args) ── ─────▶│
105+ │──── Invocation (function, args) ─────▶│
67106 │ │ instantiate function
68- │◀──── FunctionOutputSpec (output schema) ──────│
107+ │◀──── OutputSpec (output schema) ──────│
69108 │ │
70109 │──── GlobalStateInitInput ────────────▶│
71110 │◀──── GlobalInitResult ────────────────│ perform_init()
72111 │ │
73112 │──── Input Batch 1 ───────────────────▶│
74- │◀──── Output Batch 1 (NEED_MORE_INPUT)─│ process () yields
113+ │◀──── Output Batch 1 (NEED_MORE_INPUT)─│ transform () / process()
75114 │ │
76115 │──── Input Batch 2 ───────────────────▶│
77116 │◀──── Output Batch 2 (NEED_MORE_INPUT)─│
78117 │ │
79118 │──── FINALIZE (empty batch) ──────────▶│
80- │◀──── Final Output (FINISHED) ─────────│ finalize () yields
119+ │◀──── Final Output (FINISHED) ─────────│ finish () / finalize()
81120 │ │
82121```
83122
@@ -86,16 +125,19 @@ Client Worker
86125```
87126vgi/
88127 __init__.py # Package exports and module docstring
89- function.py # FunctionInvocation, FunctionOutputSpec , Arguments, GlobalInitResult
90- table_function.py # CardinalityInfo, TableFunction base class
91- table_in_out_function.py # TableInOutGeneratorFunction, Output, OutputGenerator
128+ function.py # Invocation, OutputSpec , Arguments, GlobalInitResult
129+ table_function.py # TableFunctionGenerator, CardinalityInfo, Output
130+ table_in_out_function.py # TableInOutFunction, TableInOutGeneratorFunction
92131 metadata.py # Function metadata for introspection and registration
93132 schema_utils.py # Schema builder helpers (schema, schema_like)
133+ ipc_utils.py # Arrow IPC serialization utilities
94134 worker.py # Worker base class
95- client.py # Client class and CLI
96- util.py # Serialization utilities
135+ client/
136+ client.py # Client class
137+ __main__.py # CLI entry point
97138 examples/
98- table_in_out.py # Example functions (Echo, BufferInput, SumAllColumns, etc.)
139+ table.py # Table functions (Sequence, Range, RandomSample, etc.)
140+ table_in_out.py # Table-in-out functions (Echo, BufferInput, SumAllColumns)
99141 worker.py # ExampleWorker with registry
100142```
101143
@@ -197,7 +239,92 @@ vgi-client --input data.parquet --function sum_all_columns --server vgi-example-
197239vgi-client --input data.parquet --function repeat_inputs --args ' [3]' --server vgi-example-worker
198240```
199241
200- ## Creating a Custom Function (Simple API - Recommended)
242+ ## Creating a Table Function (Data Generation)
243+
244+ Use ` TableFunctionGenerator ` when you need to generate data without input. Override ` process() ` to yield output batches:
245+
246+ ``` python
247+ import pyarrow as pa
248+ from vgi import TableFunctionGenerator, Output, Arg
249+
250+ class SequenceFunction (TableFunctionGenerator ):
251+ """ Generate a sequence of integers from 0 to n-1."""
252+
253+ class Meta :
254+ name = " sequence"
255+ max_workers = 1 # Single worker for sequential output
256+
257+ count = Arg[int ](0 , doc = " Number of integers to generate" )
258+
259+ @ property
260+ def output_schema (self ) -> pa.Schema:
261+ return pa.schema([(" n" , pa.int64())])
262+
263+ def process (self ):
264+ """ Yield output batches (no input received)."""
265+ for i in range (0 , self .count, 1000 ): # Batch by 1000
266+ batch_size = min (1000 , self .count - i)
267+ values = list (range (i, i + batch_size))
268+ yield Output(
269+ pa.RecordBatch.from_pydict({" n" : values}, schema = self .output_schema)
270+ )
271+ ```
272+
273+ ### TableFunctionGenerator Methods
274+
275+ | Method | When to Override | Default |
276+ | --------| ------------------| ---------|
277+ | ` process() ` | Always - generate output | Required |
278+ | ` output_schema ` | Define output columns | Required |
279+ | ` cardinality() ` | Provide row count estimates | Returns None |
280+ | ` setup() ` | Acquire resources | No-op |
281+ | ` teardown() ` | Release resources | No-op |
282+ | ` perform_init() ` | Distributed init (primary) | Default impl |
283+ | ` retrieve_init() ` | Distributed init (secondary) | Default impl |
284+
285+ ### Table Function Patterns
286+
287+ ** Simple sequence:**
288+ ``` python
289+ def process (self ):
290+ for i in range (self .count):
291+ yield Output(pa.RecordBatch.from_pydict(
292+ {" n" : [i]}, schema = self .output_schema
293+ ))
294+ ```
295+
296+ ** Batched output (recommended for large outputs):**
297+ ``` python
298+ BATCH_SIZE = 1000
299+
300+ def process (self ):
301+ for start in range (0 , self .count, self .BATCH_SIZE ):
302+ end = min (start + self .BATCH_SIZE , self .count)
303+ values = list (range (start, end))
304+ yield Output(pa.RecordBatch.from_pydict(
305+ {" n" : values}, schema = self .output_schema
306+ ))
307+ ```
308+
309+ ** Parallel generation with work queue:**
310+ ``` python
311+ def perform_init (self , init_input ):
312+ # Primary worker: populate work queue
313+ work_items = [chunk.serialize() for chunk in self .create_chunks()]
314+ self .enqueue_work(work_items)
315+ return GlobalInitResult(self .init_identifier)
316+
317+ def process (self ):
318+ # All workers: pull from queue until empty
319+ while True :
320+ work = self .dequeue_work()
321+ if work is None :
322+ break
323+ for batch in self .generate_chunk(work):
324+ yield Output(batch)
325+ ```
326+
327+ ## Creating a Table-In-Out Function (Simple API - Recommended)
201328
202329Use ` TableInOutFunction ` for most use cases. Override ` transform() ` for per-batch processing and ` finish() ` for final output:
203330
@@ -302,6 +429,11 @@ class DistributedSum(TableInOutFunction):
302429
303430| Use Case | Base Class |
304431| ----------| ------------|
432+ | ** Table Functions (no input)** | |
433+ | Generate sequences, ranges | ` TableFunctionGenerator ` |
434+ | Generate random/test data | ` TableFunctionGenerator ` |
435+ | Parallel data generation | ` TableFunctionGenerator ` + work queue |
436+ | ** Table-In-Out Functions (with input)** | |
305437| Transform each batch independently | ` TableInOutFunction ` |
306438| Aggregate to single result | ` TableInOutFunction ` + ` finish() ` |
307439| Buffer all input, emit on finalize | ` TableInOutFunction ` + ` finish() ` |
@@ -569,13 +701,16 @@ def finalize(self) -> OutputGenerator | None:
569701### Import Cheatsheet
570702
571703``` python
572- # Simple API (recommended for most uses)
704+ # Table Functions (no input - generate data)
705+ from vgi import TableFunctionGenerator, Output, Arg, Worker
706+
707+ # Table-In-Out Functions (transform input)
573708from vgi import TableInOutFunction, Invocation, Arg, Worker
574709
575710# Schema helpers (for output_schema definitions)
576711from vgi import schema, schema_like
577712
578- # Generator API (advanced)
713+ # Generator API (advanced table-in-out )
579714from vgi import TableInOutGeneratorFunction, Output, OutputGenerator, Invocation, Arg, Worker
580715
581716# Logging
@@ -761,7 +896,50 @@ threshold = args.get("threshold") # Required
761896count = args.get(0 , type = pa.int64()) # Raises TypeError if wrong type
762897```
763898
764- ### Function Skeleton Template (Simple API - Recommended)
899+ ### Function Skeleton Template (Table Function - No Input)
900+
901+ ``` python
902+ import pyarrow as pa
903+ from vgi import TableFunctionGenerator, Output, Arg
904+ from vgi.table_function import CardinalityInfo
905+
906+ class MyTableFunction (TableFunctionGenerator ):
907+ """ One-line description.
908+
909+ Detailed description of what this function generates.
910+ """
911+
912+ class Meta :
913+ name = " my_table_function"
914+ max_workers = 1 # Or None for parallel
915+
916+ # Declare arguments as class attributes
917+ count = Arg[int ](0 , doc = " Number of rows to generate" )
918+
919+ BATCH_SIZE = 1000 # Recommended for large outputs
920+
921+ @ property
922+ def output_schema (self ) -> pa.Schema:
923+ """ Define the output schema."""
924+ return pa.schema([(" value" , pa.int64())])
925+
926+ def cardinality (self ) -> CardinalityInfo:
927+ """ Optional: provide row count estimate."""
928+ return CardinalityInfo(estimate = self .count, max = self .count)
929+
930+ def process (self ):
931+ """ Generate output batches."""
932+ for start in range (0 , self .count, self .BATCH_SIZE ):
933+ end = min (start + self .BATCH_SIZE , self .count)
934+ values = list (range (start, end))
935+ yield Output(
936+ pa.RecordBatch.from_pydict(
937+ {" value" : values}, schema = self .output_schema
938+ )
939+ )
940+ ```
941+
942+ ### Function Skeleton Template (Table-In-Out - Recommended)
765943
766944``` python
767945import pyarrow as pa
@@ -836,28 +1014,38 @@ class MyFunction(TableInOutGeneratorFunction):
8361014```
8371015Need to implement a VGI function?
8381016│
839- ├─ No transformation needed?
840- │ └─ class Echo(TableInOutFunction): pass
841- │
842- ├─ Transform each batch independently?
843- │ └─ Override transform() → returns pa.RecordBatch
844- │
845- ├─ Produce multiple outputs per input?
846- │ └─ Override transform() → returns list[pa.RecordBatch]
847- │
848- ├─ Aggregate across all batches?
849- │ └─ Accumulate in transform(), emit in finish()
850- │ └─ Set max_processes() -> 1
851- │
852- ├─ Buffer all input, emit on finalize?
853- │ └─ Buffer in transform(), return in finish()
854- │ └─ Set max_processes() -> 1
855- │
856- ├─ Need GeneratorExit handling or distributed state?
857- │ └─ Use TableInOutGeneratorFunction (generator API)
858- │
859- └─ Need fine-grained streaming control?
860- └─ Use TableInOutGeneratorFunction (generator API)
1017+ ├─ Does the function receive input data?
1018+ │ │
1019+ │ ├─ NO (generate data from scratch)
1020+ │ │ └─ Use TableFunctionGenerator
1021+ │ │ ├─ Sequential output? Set max_workers=1
1022+ │ │ ├─ Parallel generation? Use work queue pattern
1023+ │ │ └─ Override process() to yield Output batches
1024+ │ │
1025+ │ └─ YES (transform input data)
1026+ │ │
1027+ │ ├─ No transformation needed?
1028+ │ │ └─ class Echo(TableInOutFunction): pass
1029+ │ │
1030+ │ ├─ Transform each batch independently?
1031+ │ │ └─ Override transform() → returns pa.RecordBatch
1032+ │ │
1033+ │ ├─ Produce multiple outputs per input?
1034+ │ │ └─ Override transform() → returns list[pa.RecordBatch]
1035+ │ │
1036+ │ ├─ Aggregate across all batches?
1037+ │ │ └─ Accumulate in transform(), emit in finish()
1038+ │ │ └─ Set max_processes() -> 1
1039+ │ │
1040+ │ ├─ Buffer all input, emit on finalize?
1041+ │ │ └─ Buffer in transform(), return in finish()
1042+ │ │ └─ Set max_processes() -> 1
1043+ │ │
1044+ │ ├─ Need GeneratorExit handling or distributed state?
1045+ │ │ └─ Use TableInOutGeneratorFunction (generator API)
1046+ │ │
1047+ │ └─ Need fine-grained streaming control?
1048+ │ └─ Use TableInOutGeneratorFunction (generator API)
8611049```
8621050
8631051### Status Values (in IPC metadata)
0 commit comments