Skip to content

Commit f1fa4cf

Browse files
rustyconoverclaude
andcommitted
Move input schema validation to hooks and add SumColumnsFunction example
- Move _validate_input_schema_requirement() to ScalarFunctionGenerator - Move _validate_input_schema_requirement() to TableInOutGenerator - Add SumColumnsFunction demonstrating varargs with type_bound - Register SumColumnsFunction in ExampleWorker - Add test for varargs indicator in CLI schema output 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent c87ada0 commit f1fa4cf

5 files changed

Lines changed: 122 additions & 4 deletions

File tree

‎tests/client/test_cli_catalog_functions.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,53 @@ def test_table_functions_have_correct_type(self, example_worker: str) -> None:
285285
assert by_name["range"]["function_type"] == "table"
286286
assert by_name["sum_all_columns"]["function_type"] == "table"
287287

288+
def test_varargs_function_shows_varargs_in_arguments(
289+
self, example_worker: str
290+
) -> None:
291+
"""Varargs functions show varargs indicator in arguments."""
292+
runner = CliRunner()
293+
294+
# Attach and get contents
295+
attach_result = runner.invoke(
296+
cli,
297+
["catalog", "attach", "example", "--server", example_worker],
298+
)
299+
attach_id = json.loads(attach_result.output)["attach_id"]
300+
301+
contents_result = runner.invoke(
302+
cli,
303+
[
304+
"catalog",
305+
"schema",
306+
"contents",
307+
"main",
308+
"--attach-id",
309+
attach_id,
310+
"--server",
311+
example_worker,
312+
],
313+
)
314+
assert contents_result.exit_code == 0
315+
316+
# Parse and find sum_columns function
317+
lines = contents_result.output.strip().split("\n")
318+
items = [json.loads(line) for line in lines if line.strip()]
319+
by_name = {item["name"]: item for item in items}
320+
321+
# Verify sum_columns function exists
322+
assert "sum_columns" in by_name, "sum_columns function not found"
323+
324+
sum_func = by_name["sum_columns"]
325+
326+
# Verify arguments include varargs indicator
327+
arguments = sum_func.get("arguments", [])
328+
columns_arg = next(
329+
(a for a in arguments if a.get("name") == "columns"),
330+
None,
331+
)
332+
assert columns_arg is not None, "columns argument not found"
333+
assert columns_arg.get("varargs") is True, "varargs indicator missing"
334+
288335

289336
class TestCLISchemaList:
290337
"""Tests for listing schemas via CLI."""

‎vgi/examples/scalar.py‎

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,9 @@
2525
from vgi.scalar_function import ScalarFunction
2626

2727
__all__ = [
28-
"DoubleColumnFunction",
2928
"AddNumericColumnsFunction",
29+
"DoubleColumnFunction",
30+
"SumColumnsFunction",
3031
"UpperCaseFunction",
3132
]
3233

@@ -225,3 +226,65 @@ def catalog_output_type(cls) -> pa.DataType | type[AnyArrow]:
225226
def compute(self, batch: pa.RecordBatch) -> pa.Array[Any]:
226227
"""Convert the column values to uppercase."""
227228
return pc.utf8_upper(batch.column(self.column)) # type: ignore[no-matching-overload]
229+
230+
231+
class SumColumnsFunction(ScalarFunction):
232+
"""Sums values from multiple numeric columns.
233+
234+
Uses varargs with type_bound to accept any number of numeric columns
235+
and validates that all columns are addable types at bind time.
236+
237+
Example:
238+
Input: a=[1, 2], b=[10, 20], c=[100, 200]
239+
Args: columns=('a', 'b', 'c')
240+
Output: result=[111, 222]
241+
242+
"""
243+
244+
class Meta:
245+
"""Function metadata."""
246+
247+
name = "sum_columns"
248+
description = "Sum values from multiple numeric columns"
249+
examples = [
250+
FunctionExample(
251+
sql="SELECT sum_columns(price, tax, shipping) FROM orders",
252+
description="Calculate total cost from multiple columns",
253+
),
254+
]
255+
256+
# Varargs with type_bound validates all columns are numeric
257+
columns = Arg[AnyArrow](
258+
0,
259+
varargs=True,
260+
type_bound=_is_addable_type,
261+
doc="Columns to sum (must be numeric)",
262+
)
263+
264+
_output_type: pa.DataType
265+
266+
def bind(self) -> None:
267+
"""Compute output type from first column, promoted for overflow safety."""
268+
# With varargs=True, self.columns is a tuple of column names
269+
first_col = self.columns[0] # type: ignore[index]
270+
first_type = self.input_schema.field(first_col).type
271+
self._output_type = _promote_for_addition(first_type)
272+
273+
@classmethod
274+
def catalog_output_type(cls) -> pa.DataType | type[AnyArrow]:
275+
"""Output type depends on input column types."""
276+
return AnyArrow
277+
278+
@property
279+
def output_type(self) -> pa.DataType:
280+
"""Return the computed output type based on first column."""
281+
return self._output_type
282+
283+
def compute(self, batch: pa.RecordBatch) -> pa.Array[Any]:
284+
"""Sum values from all specified columns."""
285+
# With varargs=True, self.columns is a tuple of column names at runtime
286+
columns: tuple[str, ...] = self.columns # type: ignore[assignment]
287+
result = batch.column(columns[0])
288+
for col_name in columns[1:]:
289+
result = pc.add(result, batch.column(col_name))
290+
return result

‎vgi/examples/worker.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from vgi.examples.scalar import (
1717
AddNumericColumnsFunction,
1818
DoubleColumnFunction,
19+
SumColumnsFunction,
1920
UpperCaseFunction,
2021
)
2122
from vgi.examples.table import (
@@ -74,8 +75,9 @@ class ExampleWorker(Worker):
7475
ProjectedDataFunction,
7576
SettingsAwareFunction,
7677
# ScalarFunctionGenerator - transform to single-column output
77-
DoubleColumnFunction,
7878
AddNumericColumnsFunction,
79+
DoubleColumnFunction,
80+
SumColumnsFunction,
7981
UpperCaseFunction,
8082
]
8183

‎vgi/scalar_function.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -235,7 +235,10 @@ def __init__(
235235
):
236236
"""Initialize the scalar function with invocation data and logger."""
237237
super().__init__(invocation=invocation, logger=logger)
238-
if invocation.input_schema is None:
238+
239+
def _validate_input_schema_requirement(self) -> None:
240+
"""Validate that input_schema is provided for scalar functions."""
241+
if self.invocation.input_schema is None:
239242
raise ValueError(
240243
f"{type(self).__name__} requires an input schema, but none was "
241244
f"provided. ScalarFunction processes input batches and requires "

‎vgi/table_in_out_function.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -509,7 +509,10 @@ def __init__(
509509
):
510510
"""Initialize the function with invocation data and logger."""
511511
super().__init__(invocation=invocation, logger=logger)
512-
if invocation.input_schema is None:
512+
513+
def _validate_input_schema_requirement(self) -> None:
514+
"""Validate that input_schema is provided for table-in-out functions."""
515+
if self.invocation.input_schema is None:
513516
raise ValueError(
514517
f"{type(self).__name__} requires an input schema, but none was "
515518
f"provided. TableInOutGenerator processes input batches and "

0 commit comments

Comments
 (0)