[RFC] Arrow/Py/Daft Expression Types

[jaychia]

[RFC] Arrow/Py/Daft Expression Types

This document is a proposal for extending the Expression typing system in Daft.

RFC Summary

After the implementation of this RFC, Daft's typing system will look like the following:

1. Arrow types
   a) Arrow native types (int64, float64, list[int64] etc)
   a) Daft Arrow Extension types (image, audio, video, latlong etc)
2. Python types

Daft only supports serializing Arrow types for writing to disk and long-term storage. For storing Python types, users will first have to marshal data (e.g. DaftImage.from_pil(df["pil_image"])) into the appropriate Arrow type before leveraging Daft's tooling for saving the data.

RFC Details

Motivation

Expression Types in Daft serve the following purposes:

1. Allow for validation of the validity of Expressions at definition-time instead of failing at runtime
2. Help Daft visualize the data in each column
3. Help Daft understand how to represent data in-memory for efficient and convenient data manipulation
4. Help Daft understand how to serialize data for long-term storage
5. Daft defines common operations on these data types, such as a extracting the year from a date, or concatenating two strings (and in the future, will be able to understand how much resources it should request for such operations)

Daft currently has 2 main types:

1. "Primitive" types are a subset of the Arrow type specification represented in-memory as Arrow, and we leverage Arrow's serialization capabilities when saving data to disk
2. "Python" types are represented in-memory as Python objects, and cannot be serialized when saving data to disk

Problem

This typing system has the following shortcomings:

1. Hard to add custom complex types: since all complex types are represented as PY types, Daft currently has no support for complex types as an Arrow type
2. Cannot serialize Python types for long term storage: follows the previous point. Aside from users writing their own UDFs to serialize a PY type to a Daft BYTES type, there isn't support for native complex types at the moment.
3. Custom visualizations: Daft currently has corner-cases for displaying PIL images, but adding custom visualization code for every possible Python type is not possible

Solution

A solution needs to provide:

1. Extensibility to new complex types
   a) Should have a serializable in-memory format (e.g. Arrow, protobuf, flatbuffer)
   b) Ability to define methods/kernels on these complex types for domain-specific functionality (e.g. image resizing)
   c) Ability to marshal to/from common Python representations such as numpy, PIL etc for custom processing
   d) Visualization logic for complex types during interactive development
2. Flexibility at runtime to represent data as Python types for ease of development

User Experience

To a user, these types will be displayed on a DataFrame as follows:

1. ArrowTypes - int64, list[int64] etc
2. Complex ArrowTypes - Image, list[Image] etc
3. PythonTypes - Py[PyClass]

Here is an example of a user loading a DataFrame of videos, pulling out the first frame, and making that into a thumbnail and then saving it:

>>> df = DataFrame.read_parquet("s3://my-parquet-with-videos/**")  # loads a "video" column with type video
>>> df = df.with_column("first_frame", df["video"].video.get_first_frame())  # transforms into a "first_frame" column with type image
>>> df = df.with_column("thumbnail", df["first_frame"].image.thumbnail(128, 128))  # transforms into a "thumbnail" column with type image
>>> df.select(df["thumbnail"]).save_parquet("s3://my-parquet-with-thumbnails/")  # saves the dataframe with Image column as Parquet

Notice that with this example:

1. User does not need any UDFs, as DaftTypes come with some builtin kernels already
2. Unlike Python types, complex arrow types can be saved and loaded to/from Parquet

Not included in the above example, but crucial for DaftTypes is the ability for users to move between DaftTypes and Python types. For example, df["first_frame"].image.to_numpy() will return a PY[np.ndarray] column for users to then use Numpy to manipulate, and going the other direction should also be possible with DaftImage.from_numpy_column(df["some_numpy_column"]).

[jaychia]

Discussion ready to be ported to issue now as we decided that implementation of custom Daft types will be backed by Arrow extensions. Python types will remain as a purely in-memory abstraction for ease of data processing.