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.

[stayrascal]

Hi @kevinzwang @rchowell, one small question about We don't support calling a scalar UDF with a Series because it creates confusion and opens up problems because we are mixing domains.

how do we handle the data of logical datatype in Scalar UDF, e.g. Image, Tensor. will we create Image, Tensor python object for each logical just like daft.File, and then user can get the data within scalar udf? e.g.

@daft.func
def custom_encode_image(img: daft.Image) -> bytes
    data = img.enccode()
    # some other logic

[kevinzwang]

Hi @stayrascal, we have documented the type conversions from Daft types to Python types here: https://docs.daft.ai/en/stable/api/datatypes/#daft-to-python. Let me know if that answers your question!

[stayrascal]

Ok, got it. I'm considering how to apply scalar UDF for video data type(maybe not need) during initial the design, since video data is too big, we should use a reference as the video data instead of holding whole data in memory, it means the file/video data types are different with image/map/tensor.embedding data types their data are stored in memory with arrow format, and can be shared between rust and python with zero-copy cost.

So if we add a new reference data type, we might need to create a corresponding python object? e.g. daft.File to file datatype

[Jay-ju]

@kevinzwang ,
The current number of udfactors is set through concurrency, and they are started approximately evenly across different workers. However, due to the number of partitions, not every worker may receive data.

So I'm thinking: can udfactors apply for scheduling across nodes on Ray instead of being limited to local nodes? Of course, local node invocation should be prioritized. When it's found that some udfactors have no data, the nodes for scheduling udfactors with the concurrency count can be dynamically adjusted. I'm not sure if I've clearly expressed the problem here.
cc @colin-ho

[kevinzwang]

This sounds like a more general skew issue to me, which is applicable but not isolated to UDFs. Colin can probably provide more details on what we can build within Daft to automatically handle these issues, but currently what a user can do is a manual repartition via .into_partitions, .repartition, or our new .into_batches.

[Jay-ju]

@kevinzwang Another issue is that I don't seem to find the design of UDAF in the current design plan of UDF. Could you please describe how to handle the N->1 scenario here?

[kevinzwang]

We have not built that out yet but it is part of our roadmap to do so!

[kevinzwang]

Would be happy to hear any proposals for the design in the meantime.

[huleilei]

@kevinzwang Hi, When can the Async UDF feature be implemented? Especially for calling service APIs, this feature is highly needed. Thanks.

[rchowell]

Support for async with @daft.func was merged and released in 0.5.18.

• feat: async @daft.func #4908
• https://github.com/Eventual-Inc/Daft/releases/tag/v0.5.18

[Jay-ju]

@kevinzwang hi，

• daft.func /daft.cls can't control resource request, this is in expection ?
• in addtion, daft.cls can use gpus as resource request, but it's int, maybe use float?

[kevinzwang]

HI @Jay-ju, currently CPU and memory resources are not properly used within Daft so I chose not to expose it yet. But I do want to revisit this eventually. As for GPUs, I'm happy to convert it into a float. Feel free to file an issue/submit a PR for that!