|
9 | 9 | DoubleSequenceFunction - Generates a sequence of floats 0.0..n-1 |
10 | 10 | GeneratorExceptionFunction - Demonstrates exception handling |
11 | 11 | LoggingGeneratorFunction - Demonstrates log message emission |
| 12 | +NamedParamsEchoFunction - Echoes named parameter values in output columns |
12 | 13 | NestedSequenceFunction - Generates a sequence with nested struct/list columns |
13 | 14 | PartitionedSequenceFunction - Demonstrates multi-worker parallel execution |
14 | 15 | ProjectedDataFunction - Demonstrates projection pushdown |
|
47 | 48 | "DoubleSequenceFunction", |
48 | 49 | "GeneratorExceptionFunction", |
49 | 50 | "LoggingGeneratorFunction", |
| 51 | + "NamedParamsEchoFunction", |
50 | 52 | "NestedSequenceFunction", |
51 | 53 | "PartitionedSequenceFunction", |
52 | 54 | "ProjectedDataFunction", |
@@ -184,6 +186,110 @@ class Meta: |
184 | 186 | NUMPY_DTYPE: ClassVar[type[np.generic]] = np.int64 |
185 | 187 |
|
186 | 188 |
|
| 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 | + |
187 | 293 | @dataclass(slots=True, frozen=True) |
188 | 294 | class NestedSequenceFunctionArguments: |
189 | 295 | """Arguments for NestedSequenceFunction.""" |
|
0 commit comments