cleanups for scalar functions

Author: rustyconover

CLAUDE.md

@@ -40,6 +40,9 @@ VGI (Vector Gateway Interface) provides an Apache Arrow-based protocol for conne
 │                              ▼                                      │
 │  ┌───────────────────────────────────────────────────────────────┐  │
 │  │                      Worker Process                           │  │
+│  │  SCALAR FUNCTION (ScalarFunction)                             │  │
+│  │  - compute(batch): Transform each row to single output column │  │
+│  │                           OR                                  │  │
 │  │  TABLE FUNCTION (TableFunctionGenerator)                      │  │
 │  │  - process(): Generator yielding output batches (no input)    │  │
 │  │                           OR                                  │  │
@@ -52,13 +55,15 @@ VGI (Vector Gateway Interface) provides an Apache Arrow-based protocol for conne
 
 | Type | Base Class | Input | Use Case |
 |------|------------|-------|----------|
+| **Scalar Function** | `ScalarFunction` | Batches | Per-row transforms (1:1 row mapping, single output column) |
 | **Table Function** | `TableFunctionGenerator` | None | Generate data (sequences, ranges) |
 | **Table-In-Out Function** | `TableInOutFunction` | Batches | Transform, filter, aggregate |
 
 ### Key Components
 
 - **Worker** (`vgi/worker.py`): Subprocess that hosts functions
 - **Client** (`vgi/client/client.py`): Spawns workers, streams data
+- **ScalarFunction** (`vgi/scalar_function.py`): Base for scalar functions
 - **TableFunctionGenerator** (`vgi/table_function.py`): Base for table functions
 - **TableInOutFunction** (`vgi/table_in_out_function.py`): Base for table-in-out functions
 
@@ -67,7 +72,8 @@ VGI (Vector Gateway Interface) provides an Apache Arrow-based protocol for conne
 ```
 vgi/
   __init__.py              # Package exports
-  function.py              # Invocation, OutputSpec, Arguments, GlobalInitResult
+  function.py              # Invocation, OutputSpec, Arguments, FunctionType
+  scalar_function.py       # ScalarFunction, ScalarFunctionGenerator
   table_function.py        # TableFunctionGenerator, CardinalityInfo, Output
   table_in_out_function.py # TableInOutFunction, TableInOutGeneratorFunction
   metadata.py              # Function metadata for introspection
@@ -76,6 +82,7 @@ vgi/
   client/
     client.py              # Client class
   examples/
+    scalar.py              # Example scalar functions
     table.py               # Example table functions
     table_in_out.py        # Example table-in-out functions
     worker.py              # ExampleWorker with registry
@@ -89,6 +96,32 @@ vgi-client --input data.parquet --function echo --server vgi-example-worker
 vgi-client --input data.parquet --function sum_all_columns --server vgi-example-worker
 ```
 
+## Creating a Scalar Function (Per-Row Transform)
+
+```python
+import pyarrow as pa
+import pyarrow.compute as pc
+from vgi import ScalarFunction, Arg
+
+class DoubleColumn(ScalarFunction):
+    """Double the value in a specified column."""
+
+    column = Arg[str](0, doc="Column to double")
+
+    @property
+    def output_type(self) -> pa.DataType:
+        # Output type matches input column type
+        return self.input_schema.field(self.column).type
+
+    def compute(self, batch: pa.RecordBatch) -> pa.Array:
+        return pc.multiply(batch.column(self.column), 2)
+```
+
+### Key Constraints for Scalar Functions:
+- **1:1 row mapping**: Output must have exactly the same number of rows as input
+- **Single column output**: Output schema has exactly one column named "result"
+- **No finalize phase**: All processing happens in compute()
+
 ## Creating a Table-In-Out Function (Recommended)
 
 ```python
@@ -182,6 +215,9 @@ if __name__ == "__main__":
 ### Imports
 
 ```python
+# Scalar Functions (per-row transform)
+from vgi import ScalarFunction, Arg, Worker
+
 # Table Functions (no input)
 from vgi import TableFunctionGenerator, Output, Arg, Worker
 
@@ -221,6 +257,17 @@ output_schema = schema_like(self.input_schema, rename={"old": "new"})
 
 ### Method Override Summary
 
+**ScalarFunction:**
+
+| Method | When to Override | Default |
+|--------|------------------|---------|
+| `output_type` | Define output column type | Required |
+| `compute(batch)` | Transform batch to single array | Required |
+| `setup()` | Acquire resources | No-op |
+| `teardown()` | Release resources | No-op |
+
+**TableInOutFunction:**
+
 | Method | When to Override | Default |
 |--------|------------------|---------|
 | `output_schema` | Change output columns | Returns input_schema |
@@ -232,17 +279,22 @@ output_schema = schema_like(self.input_schema, rename={"old": "new"})
 ### Pattern Decision Tree
 
 ```
-Need to implement a VGI function?
-│
-├─ Does the function receive input data?
-│  │
-│  ├─ NO → Use TableFunctionGenerator
-│  │       Override process() to yield Output batches
-│  │
-│  └─ YES → Use TableInOutFunction
-│           ├─ Transform each batch? → Override transform()
-│           ├─ Aggregate results? → Accumulate in transform(), emit in finish()
-│           └─ Need generator control? → See docs/generator-api.md
+How will your function be used in SQL?
+
+1. SELECT my_func(col1, col2) FROM table
+   → SCALAR FUNCTION: Returns one value per input row
+   → Use ScalarFunction, override output_type and compute()
+   → Example: upper(), abs(), concat()
+
+2. SELECT * FROM my_func(args)
+   → TABLE FUNCTION: Generates rows from arguments (no input table)
+   → Use TableFunctionGenerator, override process()
+   → Example: range(), read_csv(), glob()
+
+3. SELECT * FROM my_func(args, (SELECT * FROM input_table))
+   → TABLE-IN-OUT FUNCTION: Transforms input rows to output rows
+   → Use TableInOutFunction, override transform() and optionally finish()
+   → Example: filtering, enrichment, aggregation
 ```
 
 ## Additional Documentation

tests/client/test_cli.py

@@ -795,9 +795,7 @@ def test_auto_type_without_input_uses_table(
         lines = output_file.read_text().strip().split("\n")
         assert len(lines) == 3
 
-    def test_scalar_with_add_columns(
-        self, example_worker: str, tmp_path: Path
-    ) -> None:
+    def test_scalar_with_add_columns(self, example_worker: str, tmp_path: Path) -> None:
         """Test add_columns scalar function via CLI."""
         # Create input with two columns
         batch = pa.RecordBatch.from_pydict({"a": [1, 2, 3], "b": [10, 20, 30]})
@@ -832,9 +830,7 @@ def test_scalar_with_add_columns(
         results = [json.loads(line)["result"] for line in lines]
         assert results == [11, 22, 33]
 
-    def test_scalar_with_upper_case(
-        self, example_worker: str, tmp_path: Path
-    ) -> None:
+    def test_scalar_with_upper_case(self, example_worker: str, tmp_path: Path) -> None:
         """Test upper_case scalar function via CLI."""
         batch = pa.RecordBatch.from_pydict({"name": ["alice", "bob"]})
         input_file = tmp_path / "input.parquet"

tests/scalar/test_function.py

@@ -9,23 +9,23 @@
 import structlog
 
 from vgi.arguments import Arg
-from vgi.function import Arguments, Invocation
+from vgi.function import Arguments, Invocation, InvocationType, SchemaValidationError
 from vgi.log import Level, Message
 from vgi.scalar_function import (
     Output,
-    OutputGenerator,
     ProtocolInput,
     ScalarFunction,
     ScalarFunctionGenerator,
+    ScalarOutputGenerator,
 )
-from vgi.table_function import SchemaValidationError
 
 
 def create_invocation(input_schema: pa.Schema) -> Invocation:
     """Create a test invocation with the given input schema."""
     return Invocation(
         function_name="test_function",
-        in_out_function_input_schema=input_schema,
+        input_schema=input_schema,
+        function_type=InvocationType.SCALAR,
         correlation_id="test-correlation",
         invocation_id=b"test-invocation",
         arguments=Arguments(),
@@ -43,8 +43,8 @@ class DoubleColumn(ScalarFunctionGenerator):
             def output_schema(self) -> pa.Schema:
                 return pa.schema([("result", pa.int64())])
 
-            def process(self, batch: pa.RecordBatch) -> OutputGenerator:
-                _ = yield None
+            def process(self, batch: pa.RecordBatch) -> ScalarOutputGenerator:
+                _ = yield Output(self.empty_output_batch)  # Priming yield
                 import pyarrow.compute as pc
 
                 while True:
@@ -83,12 +83,13 @@ class TestFunc(ScalarFunctionGenerator):
             def output_schema(self) -> pa.Schema:
                 return pa.schema([("result", pa.int64())])
 
-            def process(self, batch: pa.RecordBatch) -> OutputGenerator:
-                _ = yield None
+            def process(self, batch: pa.RecordBatch) -> ScalarOutputGenerator:
+                _ = yield Output(self.empty_output_batch)
 
         invocation = Invocation(
             function_name="test",
-            in_out_function_input_schema=None,
+            input_schema=None,
+            function_type=InvocationType.SCALAR,
             correlation_id="test",
             invocation_id=None,
             arguments=Arguments(),
@@ -105,8 +106,8 @@ class TwoColumnOutput(ScalarFunctionGenerator):
             def output_schema(self) -> pa.Schema:
                 return pa.schema([("a", pa.int64()), ("b", pa.int64())])
 
-            def process(self, batch: pa.RecordBatch) -> OutputGenerator:
-                _ = yield None
+            def process(self, batch: pa.RecordBatch) -> ScalarOutputGenerator:
+                _ = yield Output(self.empty_output_batch)
 
         input_schema = pa.schema([("x", pa.int64())])
         invocation = create_invocation(input_schema)
@@ -122,8 +123,8 @@ class LoggingScalar(ScalarFunctionGenerator):
             def output_schema(self) -> pa.Schema:
                 return pa.schema([("result", pa.int64())])
 
-            def process(self, batch: pa.RecordBatch) -> OutputGenerator:
-                _ = yield None
+            def process(self, batch: pa.RecordBatch) -> ScalarOutputGenerator:
+                _ = yield Output(self.empty_output_batch)  # Priming yield
                 import pyarrow.compute as pc
 
                 while True:
@@ -179,7 +180,8 @@ def compute(self, batch: pa.RecordBatch) -> pa.Array[Any]:
         input_schema = pa.schema([("x", pa.int64())])
         invocation = Invocation(
             function_name="test",
-            in_out_function_input_schema=input_schema,
+            input_schema=input_schema,
+            function_type=InvocationType.SCALAR,
             correlation_id="test",
             invocation_id=None,
             arguments=Arguments(positional=(pa.scalar("x"),)),
@@ -198,29 +200,6 @@ def compute(self, batch: pa.RecordBatch) -> pa.Array[Any]:
         assert output.batch.num_rows == 3
         assert output.batch.column("result").to_pylist() == [2, 4, 6]
 
-    def test_custom_output_name(self) -> None:
-        """Test custom output column name."""
-
-        class CustomName(ScalarFunction):
-            @property
-            def output_name(self) -> str:
-                return "doubled"
-
-            @property
-            def output_type(self) -> pa.DataType:
-                return pa.int64()
-
-            def compute(self, batch: pa.RecordBatch) -> pa.Array[Any]:
-                import pyarrow.compute as pc
-
-                return pc.multiply(batch.column("x"), 2)
-
-        input_schema = pa.schema([("x", pa.int64())])
-        invocation = create_invocation(input_schema)
-
-        func = CustomName(invocation=invocation, logger=structlog.get_logger())
-        assert func.output_schema.names == ["doubled"]
-
     def test_log_method(self) -> None:
         """Test self.log() method."""
 

tests/table/generator/test_constant_table_function.py

@@ -29,14 +29,15 @@ def test_cardinality(self) -> None:
         """Cardinality should always be 1."""
         import structlog
 
-        from vgi.function import Arguments, Invocation
+        from vgi.function import Arguments, Invocation, InvocationType
 
         invocation = Invocation(
             function_name="constant_table",
-            arguments=Arguments(positional=(pa.scalar(42),)),
-            in_out_function_input_schema=None,
+            input_schema=None,
+            function_type=InvocationType.TABLE,
             correlation_id="test",
             invocation_id=b"test",
+            arguments=Arguments(positional=(pa.scalar(42),)),
         )
         func = ConstantTableFunction(
             invocation=invocation,

tests/table/generator/test_projected_data_function.py

@@ -152,15 +152,16 @@ def test_output_schema_reflects_projection(self) -> None:
         """The output_schema property should reflect the projection."""
         import structlog
 
-        from vgi.function import Invocation
-        from vgi.table_function import GlobalStateInitInput
+        from vgi.function import Invocation, InvocationType
+        from vgi.table_function import TableFunctionInitInput
 
         invocation = Invocation(
             function_name="projected_data",
-            arguments=Arguments(positional=(pa.scalar(10),)),
-            in_out_function_input_schema=None,
+            input_schema=None,
+            function_type=InvocationType.TABLE,
             correlation_id="test",
             invocation_id=b"test",
+            arguments=Arguments(positional=(pa.scalar(10),)),
         )
         func = ProjectedDataFunction(
             invocation=invocation,
@@ -171,7 +172,7 @@ def test_output_schema_reflects_projection(self) -> None:
         assert func.output_schema == ProjectedDataFunction.FULL_SCHEMA
 
         # After setting init_data with projection, should return projected schema
-        func.init_data = GlobalStateInitInput(projection_ids=[0, 2])
+        func.init_data = TableFunctionInitInput(projection_ids=[0, 2])
         schema = func.output_schema
         assert len(schema) == 2
         assert schema.names == ["id", "value"]

tests/table/generator/test_sequence_function.py

@@ -30,14 +30,15 @@ def test_cardinality(self) -> None:
         """Cardinality should match requested count."""
         import structlog
 
-        from vgi.function import Arguments, Invocation
+        from vgi.function import Arguments, Invocation, InvocationType
 
         invocation = Invocation(
             function_name="sequence",
-            arguments=Arguments(positional=(pa.scalar(100),)),
-            in_out_function_input_schema=None,
+            input_schema=None,
+            function_type=InvocationType.TABLE,
             correlation_id="test",
             invocation_id=b"test",
+            arguments=Arguments(positional=(pa.scalar(100),)),
         )
         func = SequenceFunction(
             invocation=invocation,

tests/table/test_function.py

@@ -7,7 +7,7 @@
 import structlog
 
 from tests.utils import make_schema
-from vgi.function import Arguments, Invocation
+from vgi.function import Arguments, Invocation, InvocationType
 from vgi.table_function import (
     CardinalityInfo,
     Output,
@@ -194,10 +194,11 @@ def output_schema(self) -> pa.Schema:
 
         invocation = Invocation(
             function_name="test",
-            arguments=Arguments(),
-            in_out_function_input_schema=None,
+            input_schema=None,
+            function_type=InvocationType.TABLE,
             correlation_id="test",
             invocation_id=b"test",
+            arguments=Arguments(),
         )
         func = NoCardinalityFunction(
             invocation=invocation,
@@ -219,10 +220,11 @@ def cardinality(self) -> CardinalityInfo:
 
         invocation = Invocation(
             function_name="test",
-            arguments=Arguments(),
-            in_out_function_input_schema=None,
+            input_schema=None,
+            function_type=InvocationType.TABLE,
             correlation_id="test",
             invocation_id=b"test",
+            arguments=Arguments(),
         )
         func = CardinalityFunction(
             invocation=invocation,
@@ -311,10 +313,11 @@ def output_schema(self) -> pa.Schema:
 
         invocation = Invocation(
             function_name="test",
-            arguments=Arguments(),
-            in_out_function_input_schema=None,
+            input_schema=None,
+            function_type=InvocationType.TABLE,
             correlation_id="test",
             invocation_id=b"test",
+            arguments=Arguments(),
         )
         func = TestFunction(
             invocation=invocation,