Hey Daft Team! I'm looking to adopt Daft DataTypes within my own library, so I've been reviewing the datatype.py implementation but I'm a bit unsure on how to properly inherit it. Obviously I can just

class DataType(daft.DataType):
    pass

But the LanceDB team implemented a pa.DataType extension for Bfloat16, and I'm wondering if I just implement it by adding the Bfloat16 arrow extension as a new classmethod through an arrow conversion.

Hi Everett, what are you interested in doing with an extended Daft data type?

The bfloat16 is a helpful datatype for improving performance on embeddings, so I'd like to be able to support it within my own lib that adopts daft.DataType as its own DataType standard

I don't believe Daft allows for custom extension types, but what you maybe could do is use a DataType.int16() or DataType.fixed_size_binary() as your underlying data type for bfloat16 and then implement custom operators for them with UDFs

Isn't the Image & Embedding Type an extension? I see some of the logic to make sure it gets registered to PyArrow towards the end of datatype.py

_EXT_TYPE_REGISTERED = False
_STATIC_DAFT_EXTENSION = None


def _ensure_registered_super_ext_type():
    global _EXT_TYPE_REGISTERED
    global _STATIC_DAFT_EXTENSION
    if not _EXT_TYPE_REGISTERED:

        class DaftExtension(pa.ExtensionType):
            def __init__(self, dtype, metadata=b""):
                # attributes need to be set first before calling
                # super init (as that calls serialize)
                self._metadata = metadata
                super().__init__(dtype, "daft.super_extension")

            def __reduce__(self):
                return type(self).__arrow_ext_deserialize__, (self.storage_type, self.__arrow_ext_serialize__())

            def __arrow_ext_serialize__(self):
                return self._metadata

            @classmethod
            def __arrow_ext_deserialize__(cls, storage_type, serialized):
                return cls(storage_type, serialized)

        _STATIC_DAFT_EXTENSION = DaftExtension
        pa.register_extension_type(DaftExtension(pa.null()))
        import atexit

        atexit.register(lambda: pa.unregister_extension_type("daft.super_extension"))
        _EXT_TYPE_REGISTERED = True


def get_super_ext_type():
    _ensure_registered_super_ext_type()

https://github.com/Eventual-Inc/Daft/blob/4567601085ce9ef84203ce882e8a3d49d7ea5576/daft/datatype.py#L576

@jay could you chime in here? I'm not super familiar with our extension types

We don’t allow for user-defined Daft types, there isn’t really a point unless we also expose functionality for user-defined kernels, encodings etc. However we do allow you to define extension types, which you can think of as a thin wrapper on top of existing Daft types. E.g.

daft.DataType.extension("my_custom_type", daft.DataType.binary(), metadata="somemeta")

For your case specifically, if you do

daft.DataType.from_arrow_type(lancedb_custom_type)

it should correctly create the intended DataType object. I believe in this case it will basically end up calling the extension type for you. Something like this:

daft.DataType.extension("bfloat16", daft.DataType.binary(), metadata="whatever_metadata_exists_for_bfloat16")

Hmm. I'll have to think some more on this.