Skip to content

Commit dbd29a6

Browse files
rustyconoverclaude
andcommitted
Add NamedParamsEchoFunction for named parameter test coverage
Table function that echoes all named parameter values (VARCHAR, BIGINT, DOUBLE, BOOLEAN) directly in output columns for easy assertion. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent f3e47dc commit dbd29a6

2 files changed

Lines changed: 108 additions & 0 deletions

File tree

‎vgi/examples/table.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
DoubleSequenceFunction - Generates a sequence of floats 0.0..n-1
1010
GeneratorExceptionFunction - Demonstrates exception handling
1111
LoggingGeneratorFunction - Demonstrates log message emission
12+
NamedParamsEchoFunction - Echoes named parameter values in output columns
1213
NestedSequenceFunction - Generates a sequence with nested struct/list columns
1314
PartitionedSequenceFunction - Demonstrates multi-worker parallel execution
1415
ProjectedDataFunction - Demonstrates projection pushdown
@@ -47,6 +48,7 @@
4748
"DoubleSequenceFunction",
4849
"GeneratorExceptionFunction",
4950
"LoggingGeneratorFunction",
51+
"NamedParamsEchoFunction",
5052
"NestedSequenceFunction",
5153
"PartitionedSequenceFunction",
5254
"ProjectedDataFunction",
@@ -184,6 +186,110 @@ class Meta:
184186
NUMPY_DTYPE: ClassVar[type[np.generic]] = np.int64
185187

186188

189+
@dataclass(slots=True, frozen=True)
190+
class NamedParamsEchoFunctionArgs:
191+
"""Arguments for NamedParamsEchoFunction."""
192+
193+
count: Annotated[int, Arg(0, doc="Number of rows to generate", ge=0)]
194+
greeting: Annotated[str, Arg("greeting", default="hello", doc="Greeting text echoed in output")]
195+
multiplier: Annotated[int, Arg("multiplier", default=1, doc="Multiplier for value column")]
196+
scale: Annotated[float, Arg("scale", default=1.0, doc="Scale factor for float_value column")]
197+
enabled: Annotated[bool, Arg("enabled", default=True, doc="Boolean echoed in output")]
198+
199+
200+
@init_single_worker
201+
@bind_fixed_schema
202+
@_cardinality_from_count
203+
class NamedParamsEchoFunction(TableFunctionGenerator[NamedParamsEchoFunctionArgs, CountdownState]):
204+
"""Echoes named parameter values directly in output columns.
205+
206+
USE CASE
207+
--------
208+
Testing that named parameters of various types (VARCHAR, BIGINT, DOUBLE,
209+
BOOLEAN) are correctly passed from DuckDB to the worker. Each named
210+
parameter value is echoed directly in an output column, making it easy
211+
to assert correctness.
212+
213+
SCHEMA
214+
------
215+
Output: {"id": int64, "greeting": string, "value": int64, "float_value": float64, "enabled": bool}
216+
217+
Example:
218+
-------
219+
SELECT * FROM named_params_echo(3)
220+
Returns: rows with id=0..2, greeting='hello', value=id*1, float_value=id*1.0, enabled=true
221+
222+
SELECT * FROM named_params_echo(3, greeting := 'hi', multiplier := 10)
223+
Returns: rows with id=0..2, greeting='hi', value=id*10, float_value=id*1.0, enabled=true
224+
225+
"""
226+
227+
FunctionArguments = NamedParamsEchoFunctionArgs
228+
229+
class Meta:
230+
"""Metadata for NamedParamsEchoFunction."""
231+
232+
name = "named_params_echo"
233+
description = "Echoes named parameter values in output columns"
234+
categories = ["generator", "testing"]
235+
tags = {"category": "testing", "type": "params"}
236+
examples = [
237+
FunctionExample(
238+
sql="SELECT * FROM named_params_echo(3)",
239+
description="Echo default parameter values for 3 rows",
240+
),
241+
FunctionExample(
242+
sql="SELECT * FROM named_params_echo(3, greeting := 'hi', multiplier := 10)",
243+
description="Echo custom greeting and multiplier",
244+
),
245+
]
246+
247+
FIXED_SCHEMA: ClassVar[pa.Schema] = schema(
248+
{
249+
"id": pa.int64(),
250+
"greeting": pa.string(),
251+
"value": pa.int64(),
252+
"float_value": pa.float64(),
253+
"enabled": pa.bool_(),
254+
}
255+
)
256+
257+
BATCH_SIZE: ClassVar[int] = 1000
258+
259+
@classmethod
260+
def initial_state(cls, params: ProcessParams[NamedParamsEchoFunctionArgs]) -> CountdownState:
261+
"""Create initial state with remaining count."""
262+
return CountdownState(remaining=params.args.count)
263+
264+
@classmethod
265+
def process(
266+
cls,
267+
params: ProcessParams[NamedParamsEchoFunctionArgs],
268+
state: CountdownState,
269+
out: OutputCollector,
270+
) -> None:
271+
"""Generate rows echoing named parameter values."""
272+
if state.remaining <= 0:
273+
out.finish()
274+
return
275+
276+
size = min(state.remaining, cls.BATCH_SIZE)
277+
ids = list(range(state.current_index, state.current_index + size))
278+
279+
data: dict[str, list[int] | list[str] | list[float] | list[bool]] = {
280+
"id": ids,
281+
"greeting": [params.args.greeting] * size,
282+
"value": [i * params.args.multiplier for i in ids],
283+
"float_value": [i * params.args.scale for i in ids],
284+
"enabled": [params.args.enabled] * size,
285+
}
286+
287+
out.emit(pa.RecordBatch.from_pydict(data, schema=params.output_schema))
288+
289+
state.current_index += size
290+
state.remaining -= size
291+
292+
187293
@dataclass(slots=True, frozen=True)
188294
class NestedSequenceFunctionArguments:
189295
"""Arguments for NestedSequenceFunction."""

‎vgi/examples/worker.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
DoubleSequenceFunction,
5353
GeneratorExceptionFunction,
5454
LoggingGeneratorFunction,
55+
NamedParamsEchoFunction,
5556
NestedSequenceFunction,
5657
PartitionedSequenceFunction,
5758
ProjectedDataFunction,
@@ -111,6 +112,7 @@ class Settings:
111112
DoubleSequenceFunction,
112113
GeneratorExceptionFunction,
113114
LoggingGeneratorFunction,
115+
NamedParamsEchoFunction,
114116
NestedSequenceFunction,
115117
PartitionedSequenceFunction,
116118
ProjectedDataFunction,

0 commit comments

Comments
 (0)