New UDFs Roadmap

[kevinzwang]

Background

As Daft becomes used for more and more multimodal/AI workflows, we see some increasing patterns around the usage of UDFs, and we'd like to redesign our UDFs to better work with these patterns. This issue tracks the progress on this redesign.

The major differences between the designs of the existing (legacy) and new UDFs:

1. The default mode is to operate on a single row of data at a time, instead of a batch.
2. Instead of being separate concepts from expressions, the new UDFs will instead mirror our built-in functions as much as possible. UDFs are Daft functions, just ones that are defined by the user.
3. The legacy UDFs could be stateful using the concurrency parameter. The new UDFs will not be stateful. Instead, to do stateful things, we will introduce the concept of "resources". Resources are not covered in this roadmap but will be a separate issue.
   In addition, the scope of this work also includes some new ways to use UDFs, such as multi-column outputs, generator UDFs, async UDFs, and ergonomics around conversions between Python and Daft types.

Examples

Simple scalar UDF

@daft.func
def my_udf(input: int) -> str:  # return dtype is inferred from type hint
    return f"{input}"

# can be used to construct expressions
expr: daft.Expression = my_udf(col("a"))
df.select(expr)

# can also still be used with scalar values as a regular Python function
val = my_udf(1)
assert val == "1"

Generator UDF

@daft.func(return_dtype=DataType.list(DataType.int())) # generators return list-type expressions
def my_gen_udf(input: int):  
    for i in range(input):
        yield i

Async UDF

# should work automatically
@daft.func
def my_async_udf(input: int) -> str:
    return await api_call(input)

Batch UDF

@daft.func.batch
def my_batch_udf(input: Series) -> Series:
    return input * input

Type checking

@daft.func
def my_func(x: int):
	...
	
df = daft.from_pydict({"x": ["a", "b", "c"]})
df.select(my_func(df["x"])) # should error because column "x" is not compatible with int

Roadmap

These tasks do not necessarily need to be done in order

• MVP for scalar UDF
  • feat: rowwise (scalar) udfs #4723
  • chore: new code path for scalar UDF #4814
• Unify and document Daft <-> Python type conversions
• Generator UDF
• Async UDF
• Async generator UDF
• Type checking/casting on UDF args using argument type hints
• MVP for batch UDF
• Aggregate UDF (design TBD)
• Custom resources for stateful operations
• Deprecate and remove old UDFs

[universalmind303]

@kevinzwang I wonder if this would be better as a discussion? Colin did something similar for the flotilla roadmap, but made it a discussion instead.

[kevinzwang]

@kevinzwang I wonder if this would be better as a discussion? Colin did something similar for the flotilla roadmap, but made it a discussion instead.

What is the advantage of making it a discussion? I forgot who said this but I recall that we wanted to use issues for asks that would need changes in Daft, and discussions otherwise.

[kevinzwang]

Changed to a discussion after chatting with Jay

[rchowell]

Some Q/A on scalar function behavior:

Q: What happens if the function closes on outside state?

• Those arguments are kept as references and will work as expected.

Q: How are Series handled?

• We don't support calling a scalar UDF with a Series because it creates confusion and opens up problems because we are mixing domains.

Q: Do we support optional/defaults/variadics/kwargs?

• We do by capturing the arguments exactly as passed in, then spreading via call(*args, **kwargs) and pyo3/python handles the binding therefore runtime erroring on signature mismatch.

Q: What if I don't give a type hint e.g. def my_plus(x, y) today?

• We do not error or do any kind of signature validation, we evaluate the argument expressions in order, then convert to python before sending out with call(*args).
• If there is no return type hint, we error unless a return_dtype is given in the decorator.

Q: Can I call it like col(a).my_plus()?

• No, we do not support chaining with UDFs and are moving away from using chaining with functions.

Q: Can I create a scalar class-based UDF?

• We may have a solution for other stateful UDFs, but for now this is not necessary. You may always use closures to capture simple state objects, or by using a memoizing factory.

Q: How is a daft.DataType.Struct represented as a python value, and what about other complex types?

• We are working on this design and are introducing consistent conversions, the specifics are not currently defined.