UDF improvements

[Lundez]

Hi!

I think there's two things that could be done to improve the utility of UDF:s in the case of Deep Learning.

1. Infer the output type (even if suboptimal, i.e. int64 rather than int8).
   • It becomes a lot of code everywhere with the return_dtype of tensors
2. Enable UDFs to work on multiple columns that are supplied dict-style and can return dict-style
   • This is how some other libraries does it, i.e. HuggingFace Datasets, and it's really helpful when working with Deep Learning.
   • Applying augmentations to Object Detection Tasks now is really hard (need BBox, Image & Label). We really want to augment them together. Is the only "easy" way to do this to create a struct, really?

I hope we can have a discussion on this!

[jaychia]

For (1), I'd love to be able to have a mapping of Python types (e.g. str, int) to Daft datatypes, as well as use our SQL type string inference in more places. This really shouldn't be that hard, something like:

class DataType:
    @classmethod
    def infer_datatype(cls, datatype_specifier: str | type | DataType):
        if isinstance(datatype_specifier, DataType):
            return datatype_specifier
        else if isinstance(datatype_specifier, str):
            return from_sql_string(...)
        else:
            # custom Python mapping
            if datatype_specifier is str:
                ...

Please do make an issue for us here! I'd like for folks in the team to take a look at this.

---

For (2), could you elaborate on this ask? If you provide some examples of an API you'd like here in an issue, I think the team would love to take a look. Perhaps this + the first suggestion could look like:

@udf(return_dtype={"foo": str, "bar": int})  # Python dict can be inferred as a struct type Struct[foo: String, bar: Int64]
def myfunc(bboxes, images, labels):
    return {
        "foo": [...],
        "bar": [...],
    }

Great suggestions, would love to know your thoughts :)

[Lundez]

Hi @jaychia I'll get an issue on type inference up.

(1) sounds great! But I ment one step further, infer by the output type. I know polars does this somehow.
They define map_elements with a parameter called return_dtype that's optional, if not set they infer by first non-null value. Which I think is great DX.

(2) Albumentations is a library that has multiple transforms and augmentations for images. Here's a user-guide on object detection.
You need to augment BBox and Image in the same go as there's a probability p to augment the image and they need to have the same augmentation applied. This is also true for PyTorch Transforms (torchvision.transforms) where you would want to rescale the BBox the same way you rescale your Image.
If you'd like I can create an issue for this too.

[Lundez]

@jaychia thanks for the type part.
Now to the multi column in <> multi column out.

How would we do that best? Struct to me seems suboptimal, but it sure is a possible solution.
Based on documentation I read a UDF can take multiple input columns but only return one. Is it possible that we could add multiple as output if you return a dict?
Would it be possible to get input as a dict (this is how Ray does it I believe)?

I'll create an issue tomorrow with more structured thoughts. With use case additionally.

[jaychia]

• Multi column in should already be possible -- let me now if you're thinking of something different?
• For multi-column out, what would be the expected API here? I think my argument is that whatever we do here would end up looking very much like a struct column, and will probably be implemented under the hood with struct columns anyways

Would love to see maybe an example API that you think could be good here. I agree that the UX isn't the best it can be yet, curious to hear your thoughts!