Skip to content

Conversation

@KKould
Copy link
Member

@KKould KKould commented Nov 4, 2025

Extend @udf to enable it to return multiple rows of data.

@udf(
    input_types=["INT"],
    result_type=[("num", "INT"), ("double_num", "INT")],
    batch_mode=True,
)
def expand_numbers(nums: List[int]):
    import pyarrow as pa

    schema = pa.schema(
        [pa.field("num", pa.int32(), nullable=False), pa.field("double_num", pa.int32(), nullable=False)]
    )
    return pa.RecordBatch.from_arrays(
        [pa.array(nums, type=pa.int32()), pa.array([n * 2 for n in nums], type=pa.int32())],
        schema=schema,
    )

@KKould KKould closed this Nov 9, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant