RFC to support non-Python/separate environment UDFs

[kevinzwang]

Background

Daft provides a convenient interface to run Python code within a Daft query with our user-defined functions (UDFs). But if you wanted to run UDFs that are not written in Python, or even if they are Python but require a different Python environment than the one running Daft, there is no native way to accomplish that.

Goals

Provide simple and performant methods to run user-defined code that have one or more of the following requirements:

1. The code is not written in Python and does not have native Python bindings
2. The code requires a separate environment or setup steps

In addition, these APIs should be executor agnostic. In other words, they should work both in the native and distributed runners, and not require, say, specific Ray features.

Out of scope:

• UDFs for specific languages. We'd like to build a common API across languages/executables first, and after that we may consider providing abstractions on top of that to better support specific languages. But that should come after the work here.

Proposal

To accomplish these goals I am proposing a three-layered solution:

1. shell command UDFs - for simply running a shell command where inputs are given as command args and stdout is captured as output
2. IPC UDFs - spin up a process via a shell command and communicate inputs via stdin and outputs via stdout.
3. Container UDFs - spin up a container via a specified image. Communication through the same protocol as IPC UDFs

Shell Command UDF

Daft should provide a built-in run_process command that takes in one or more arguments (literals or expressions) that will be executed as a spawned process per row.

Example:

import daft
from daft.functions import run_process

df = daft.from_pydict({
  "file_path": ["a.txt", "b.txt"]
})

# use cat to display the contents of each local file
df.select(run_process("cat", df["file_path"]))

IPC UDF

I am not yet sure what the actual Python interface will look like for this, but it will allow users to pass in a shell command to run to start up the process, and the input columns to give to the process. Something like:

run_ipc_process("./my_shell_cmd", [col("a"), col("b")])

Then, Daft will send data through stdin and listen to stdout for outputs. We'll specify a thin protocol for this, which will support either or both:

• row-wise inputs and outputs where the values are all strings (similar to shell command UDF)
• batch inputs and outputs with Arrow buffers

Container UDF

This will be the same as the IPC UDF, except instead of providing a shell command, users can provide a container image to run and communicate to.

[jaychia]

This is cool and we can maybe check out Clickhouse UDFs for inspiration: https://clickhouse.com/docs/sql-reference/functions/udf

Very Unix philosophy - just relies on stdin and stdout

[universalmind303]

A few months ago I did a POC on IPC based udfs.

see #4441

For IPC you don't even need the overhead of flight and grpc. You can just write the arrow buffer via stdin.

For like a remote ipc, arrow flight might make more sense though.

[universalmind303]

But how would the udf implementer know how to interpret that argument?

for example you can use many complex types inside your udfs.

def read_file(file: daft.File) -> str: ...
def read_bytes(b: bytes) -> str: ...
def read_image(img: Image) -> str: ...

We use Literal and conversion rules to convert Literal -> (daft.File | bytes | Image | whatever else).

For batched udf's its easy enough to just go through arrow like you suggested, but if we were to implement rowwise ipc udfs, then there has to be a more structured output to indicate the input/output datatypes.

[kevinzwang]

Hm, maybe we can just constrain it to basic types for now? I am okay with restricting things like file and image types to Daft Python

[kevinzwang]

If someone wanted to, they could also do a batch_size=1 batch UDF as well

[universalmind303]

Hm, maybe we can just constrain it to basic types for now? I am okay with restricting things like file and image types to Daft Python

I don't think that solves the problem though. There needs to be some serialization format.
By "just convert it to a string", does that mean use the Display of that datatype? Does it mean using the "sql display" Or how exactly does that value get converted to a string.

Assuming you meant the simplest conversion equivalent of like Display:

ipc_udf(
  lit("1"),
  lit(1),
  command="path/to/executable", 
  return_dtype=DataType.string()
)

the udf would just receive both of these as "1" with no type information, and would have to guess/infer what the original datatype was.

One alternative would be to make you explicitly pass in the function signature when calling the function.

ipc_udf(
  lit("1"),
  lit(1),
  command="path/to/executable", 
  input_dtypes=[DataType.string(), DataType.int64()],
  return_dtype=DataType.string()
)

but that still doesn't answer the question of how you're proposing going from Literal -> string

[kevinzwang]

We could use Display, we could also use Python repr or our JSON encoding as well. I'm aware that it does create some level of ambiguity but that's just something that shell commands have to contend with in general, so I feel like it's not out of the ordinary. Adding additional schemas or structure would require a lot of complexity, both on our side to create stable protocols, as well as on the user side to use those protocols.

I'm open to having optional input and output dtypes that you can pass in though, but I don't know if I want it to be required.

[universalmind303]

Something else I feel is worth at least calling out is UDF's implemented via FFI.

[kevinzwang]

Yep, thanks for mentioning this. I was considering this as well but I feel that there are several aspects to it that make it less appealing than the other options:

1. The FFI UDF has to run on the same process/environment as Daft
2. We would need to build a stable C ABI for Daft
3. This ABI will need bindings for each supported language for it to be easily used, which is going to be a decent amount of maintenance burden on either us or the user

The upside is that it would potentially be faster since we can eliminate data copies and IPC. But this could also be accomplished by building a Python library with an Arrow interface, backed by the language of your choice, just like we do with Rust and PyO3. Then you can just use that library in a standard Python UDF, which incurs a slight overhead of having to go through Python, but for batch cases it should not be high.

[kesavkolla]

WASI has become mature specification. IMO it truly offers cross programming language benefits with very little overhead. All popular languages are adding support for WASI.

[kevinzwang]

Hi @kesavkolla, I am not super familiar with WASI. Could you elaborate on how a webassembly UDF would operate?

[kesavkolla]

Technically it's like FFI but rather than investing own ABI wasm become that layer and runtime layer will take care of all sandboxing security etc... Wasm component model offers ability to create the interface glue. Daft will become host of wasm runtime and execute the wasm with the predefined component specification. Developers use their own language tools to build the wasm components.

[kesavkolla]

One more benefit with wasm is distribution with OCI image format. Developers can use public/private registry to publish the wasm components. It will become easy for Daft cloud also. Configure the image registry pull credentials and volla Daft can pull it and execute the components.

[kevinzwang]

Thanks for the explanation! If there is interest in the community on running code with WASI I am happy to explore this as well

[plotor]

Thanks to kevin. We've indeed encountered some customer needs in this area, such as Shell UDFs and UDF environment isolation. For Container UDFs, I'd like to know if the containers are managed by Daft? Specifically, are container start-up, shutdown, discovery, and elastic scaling controlled manually by the user, or is it controlled by Daft?

[kevinzwang]

Hi @plotor, thanks for the questions! Yes, the intention is for Daft to manage the lifecycle of the containers. There will likely be knobs that you can tune just like with UDFs for concurrency and resources, but users should not need to spin up their own containers, they just need to specify the image they want to use and how.