DataTypes#

@datatype_constructor
@classmethod
def decimal128(cls, precision: int, scale: int) -> DataType:
    """Fixed-precision decimal."""
    return cls._from_pydatatype(PyDataType.decimal128(precision, scale))

@datatype_constructor
@classmethod
def duration(cls, timeunit: TimeUnit | str) -> DataType:
    """Duration DataType."""
    if isinstance(timeunit, str):
        timeunit = TimeUnit.from_str(timeunit)
    return cls._from_pydatatype(PyDataType.duration(timeunit._timeunit))

@datatype_constructor
@classmethod
def embedding(cls, dtype: DataType, size: int) -> DataType:
    """Create an Embedding DataType: embeddings are fixed size arrays, where each element in the array has a **numeric** ``dtype`` and each array has a fixed length of ``size``.

    Args:
        dtype: DataType of each element in the list (must be numeric)
        size: length of each list
    """
    if not isinstance(size, int) or size <= 0:
        raise ValueError("The size for a embedding must be a positive integer, but got: ", size)
    return cls._from_pydatatype(PyDataType.embedding(dtype._dtype, size))

@datatype_constructor
@classmethod
def extension(cls, name: str, storage_dtype: DataType, metadata: str | None = None) -> DataType:
    return cls._from_pydatatype(PyDataType.extension(name, storage_dtype._dtype, metadata))

@datatype_constructor
@classmethod
def file(cls, media_type: MediaType = MediaType.unknown()) -> DataType:
    """Create a File DataType: a type which refers to a file object."""
    return cls._from_pydatatype(PyDataType.file(media_type._media_type))

@datatype_constructor
@classmethod
def fixed_size_binary(cls, size: int) -> DataType:
    """Create a FixedSizeBinary DataType: A fixed-size string of bytes."""
    if not isinstance(size, int) or size <= 0:
        raise ValueError("The size for a fixed-size binary must be a positive integer, but got: ", size)
    return cls._from_pydatatype(PyDataType.fixed_size_binary(size))

@datatype_constructor
@classmethod
def fixed_size_list(cls, dtype: DataType, size: int) -> DataType:
    """Create a FixedSizeList DataType: Fixed-size list, where each element in the list has type ``dtype`` and each list has length ``size``.

    Args:
        dtype: DataType of each element in the list
        size: length of each list
    """
    if not isinstance(size, int) or size <= 0:
        raise ValueError("The size for a fixed-size list must be a positive integer, but got: ", size)
    return cls._from_pydatatype(PyDataType.fixed_size_list(dtype._dtype, size))

@classmethod
def from_arrow_type(cls, arrow_type: pa.lib.DataType, python_fallback: builtins.bool = True) -> DataType:
    """Maps a PyArrow DataType to a Daft DataType."""
    if pa.types.is_int8(arrow_type):
        return cls.int8()
    elif pa.types.is_int16(arrow_type):
        return cls.int16()
    elif pa.types.is_int32(arrow_type):
        return cls.int32()
    elif pa.types.is_int64(arrow_type):
        return cls.int64()
    elif pa.types.is_uint8(arrow_type):
        return cls.uint8()
    elif pa.types.is_uint16(arrow_type):
        return cls.uint16()
    elif pa.types.is_uint32(arrow_type):
        return cls.uint32()
    elif pa.types.is_uint64(arrow_type):
        return cls.uint64()
    elif pa.types.is_float16(arrow_type):
        return cls.float16()
    elif pa.types.is_float32(arrow_type):
        return cls.float32()
    elif pa.types.is_float64(arrow_type):
        return cls.float64()
    elif pa.types.is_string(arrow_type) or pa.types.is_large_string(arrow_type):
        return cls.string()
    elif pa.types.is_binary(arrow_type) or pa.types.is_large_binary(arrow_type):
        return cls.binary()
    elif pa.types.is_fixed_size_binary(arrow_type):
        return cls.fixed_size_binary(arrow_type.byte_width)
    elif pa.types.is_boolean(arrow_type):
        return cls.bool()
    elif pa.types.is_null(arrow_type):
        return cls.null()
    elif pa.types.is_decimal128(arrow_type):
        return cls.decimal128(arrow_type.precision, arrow_type.scale)
    elif pa.types.is_date32(arrow_type):
        return cls.date()
    elif pa.types.is_date64(arrow_type):
        return cls.timestamp(TimeUnit.ms())
    elif pa.types.is_time64(arrow_type):
        timeunit = TimeUnit.from_str(pa.type_for_alias(str(arrow_type)).unit)
        return cls.time(timeunit)
    elif pa.types.is_timestamp(arrow_type):
        timeunit = TimeUnit.from_str(arrow_type.unit)
        return cls.timestamp(timeunit=timeunit, timezone=arrow_type.tz)
    elif pa.types.is_duration(arrow_type):
        timeunit = TimeUnit.from_str(arrow_type.unit)
        return cls.duration(timeunit=timeunit)
    elif pa.types.is_list(arrow_type) or pa.types.is_large_list(arrow_type):
        assert isinstance(arrow_type, (pa.ListType, pa.LargeListType))
        field = arrow_type.value_field
        return cls.list(cls.from_arrow_type(field.type, python_fallback))
    elif pa.types.is_fixed_size_list(arrow_type):
        assert isinstance(arrow_type, pa.FixedSizeListType)
        field = arrow_type.value_field
        return cls.fixed_size_list(cls.from_arrow_type(field.type, python_fallback), arrow_type.list_size)
    elif pa.types.is_struct(arrow_type):
        assert isinstance(arrow_type, pa.StructType)
        fields = [arrow_type[i] for i in range(arrow_type.num_fields)]
        return cls.struct({field.name: cls.from_arrow_type(field.type, python_fallback) for field in fields})
    elif pa.types.is_interval(arrow_type):
        return cls.interval()
    elif pa.types.is_map(arrow_type):
        assert isinstance(arrow_type, pa.MapType)
        return cls.map(
            key_type=cls.from_arrow_type(arrow_type.key_type, python_fallback),
            value_type=cls.from_arrow_type(arrow_type.item_type, python_fallback),
        )
    elif isinstance(arrow_type, pa.FixedShapeTensorType):
        scalar_dtype = cls.from_arrow_type(arrow_type.value_type, python_fallback)
        return cls.tensor(scalar_dtype, tuple(arrow_type.shape))
    elif pa.types.is_union(arrow_type):
        assert isinstance(arrow_type, pa.UnionType)
        mode = "dense" if arrow_type.mode == "dense" else "sparse"
        field_dict = {
            arrow_type.field(i).name: cls.from_arrow_type(arrow_type.field(i).type, python_fallback)
            for i in range(arrow_type.num_fields)
        }
        type_ids = list(arrow_type.type_codes)
        return cls.union(field_dict, type_ids, mode)
    # Only check for PyExtensionType if pyarrow version is < 21.0.0
    if hasattr(pa, "PyExtensionType") and isinstance(arrow_type, getattr(pa, "PyExtensionType")):
        # TODO(Clark): Add a native cross-lang extension type representation for PyExtensionTypes.
        raise ValueError(
            "pyarrow extension types that subclass pa.PyExtensionType can't be used in Daft, since they can't be "
            f"used in non-Python Arrow implementations and Daft uses the Rust Arrow implementation: {arrow_type}"
        )
    elif isinstance(arrow_type, pa.BaseExtensionType):
        name = arrow_type.extension_name

        if (get_or_create_runner().name == "ray") and (
            type(arrow_type).__reduce__ == pa.BaseExtensionType.__reduce__
        ):
            raise ValueError(
                f"You are attempting to use a Extension Type: {arrow_type} with the default pyarrow `__reduce__` which breaks pickling for Extensions"
                "To fix this, implement your own `__reduce__` on your extension type"
                "For more details see this issue: "
                "https://github.com/apache/arrow/issues/35599"
            )
        try:
            metadata = arrow_type.__arrow_ext_serialize__().decode()
        except AttributeError:
            metadata = None

        if name == "daft.super_extension":
            assert metadata is not None
            return cls._from_pydatatype(PyDataType.from_json(metadata))
        else:
            return cls.extension(
                name,
                cls.from_arrow_type(arrow_type.storage_type, python_fallback),
                metadata,
            )
    else:
        if python_fallback:
            # Fall back to a Python object type.
            # TODO(Clark): Add native support for remaining Arrow types.
            return cls.python()
        else:
            raise TypeError(f"Unsupported Arrow type: {arrow_type}")

@classmethod
def from_numpy_dtype(cls, np_type: np.dtype[Any]) -> DataType:
    """Maps a Numpy datatype to a Daft DataType."""
    arrow_type = pa.from_numpy_dtype(np_type)
    return cls.from_arrow_type(arrow_type)

@classmethod
def from_sql(cls, sql_type: str) -> DataType:
    """Construct a Daft DataType from a SQL type."""
    return cls._from_pydatatype(sql_datatype(sql_type))

@datatype_constructor
@classmethod
def image(
    cls, mode: str | ImageMode | None = None, height: int | None = None, width: int | None = None
) -> DataType:
    """Create an Image DataType: image arrays contain (height, width, channel) ndarrays of pixel values.

    Each image in the array has an :class:`~daft.ImageMode`, which describes the pixel dtype (e.g. uint8) and
    the number of image channels/bands and their logical interpretation (e.g. RGB).

    If the height, width, and mode are the same for all images in the array, specifying them when constructing
    this type is advised, since that will allow Daft to create a more optimized physical representation
    of the image array.

    If the height, width, or mode may vary across images in the array, leaving these fields unspecified when
    creating this type will cause Daft to represent this image array as a heterogeneous collection of images,
    where each image can have a different mode, height, and width. This is much more flexible, but will result
    in a less compact representation and may be make some operations less efficient.

    Args:
        mode: The mode of the image. By default, this is inferred from the underlying data.
            If height and width are specified, the mode must also be specified.
        height: The height of the image. By default, this is inferred from the underlying data.
            Must be specified if the width is specified.
        width: The width of the image. By default, this is inferred from the underlying data.
            Must be specified if the width is specified.
    """
    if isinstance(mode, str):
        mode = ImageMode.from_mode_string(mode.upper())
    if mode is not None and not isinstance(mode, ImageMode):
        raise ValueError(f"mode must be a string or ImageMode variant, but got: {mode}")
    if height is not None and width is not None:
        if not isinstance(height, int) or height <= 0:
            raise ValueError("Image height must be a positive integer, but got: ", height)
        if not isinstance(width, int) or width <= 0:
            raise ValueError("Image width must be a positive integer, but got: ", width)
    elif height is not None or width is not None:
        raise ValueError(
            f"Image height and width must either both be specified, or both not be specified, but got height={height}, width={width}"
        )
    return cls._from_pydatatype(PyDataType.image(mode, height, width))

@classmethod
def infer_from_object(cls, obj: Any) -> DataType:
    """Infer Daft DataType from a Python object."""
    from daft.series import Series

    s = Series.from_pylist([obj])
    return s.datatype()

@classmethod
def infer_from_type(cls, t: type | GenericAlias | UnionType) -> DataType:
    """Infer Daft DataType from a Python type."""
    # NOTE: Make sure this matches the logic in `Literal::from_pyobj` in Rust
    # NOTE: The base type for Union is hidden, so it requires special handling
    # TODO: TypeForm would cover everything: https://peps.python.org/pep-0747/

    assert isinstance(t, (type, GenericAlias, UnionType)) or typing.get_origin(t) is typing.Union, (
        f"Input to DataType.infer_from_type must be a type, found {t} (type {type(t)})"
    )

    import datetime
    import decimal
    import importlib
    from typing import is_typeddict

    import daft.file
    import daft.series

    origin_or_none = typing.get_origin(t)
    origin: type = origin_or_none if origin_or_none is not None else t  # type: ignore
    args = typing.get_args(t)

    def check_type(type_or_path: type | str) -> bool:
        """Check if `origin` is a subclass of `type_or_path`.

        Pass in a string value for `type_or_path` for types from optional dependencies.
        """
        if isinstance(type_or_path, type):
            type_obj = type_or_path
        elif isinstance(type_or_path, str):
            module_name, type_name = type_or_path.rsplit(".", 1)
            try:
                module = importlib.import_module(module_name)
                type_obj = getattr(module, type_name)
            except (ImportError, AttributeError):
                return False
        else:
            raise ValueError("`type_or_path` must be type or string")

        return issubclass(origin, type_obj)

    # NOTE: This has to be first to handle the special case of typing.Union
    if origin is typing.Union or check_type(UnionType):  # type: ignore[comparison-overlap]
        inner_types = set(DataType.infer_from_type(arg) for arg in args)
        if len(inner_types) == 1:
            return inner_types.pop()
        elif len(inner_types) == 2 and cls.null() in inner_types:
            return inner_types.difference([cls.null()]).pop()
        else:
            return cls.python()
    elif check_type(type(None)):
        return cls.null()
    elif check_type(bool):
        return cls.bool()
    elif check_type(str):
        return cls.string()
    elif check_type(bytes):
        return cls.binary()
    elif check_type(int):
        return cls.int64()
    elif check_type(float):
        return cls.float64()
    elif check_type(datetime.datetime):
        # cannot derive timezone from type
        return cls.timestamp(TimeUnit.us(), timezone=None)
    elif check_type(datetime.date):
        return cls.date()
    elif check_type(datetime.time):
        return cls.time(TimeUnit.us())
    elif check_type(datetime.timedelta):
        return cls.duration(TimeUnit.us())
    elif check_type(daft.file.VideoFile):
        return cls.file(MediaType.video())
    elif check_type(daft.file.AudioFile):
        return cls.file(MediaType.audio())
    elif check_type(daft.file.ImageFile):
        return cls.file(MediaType.image())
    elif check_type(daft.file.File):
        return cls.file(MediaType.unknown())
    elif check_type(list):
        if len(args) == 0:
            inner_dtype = cls.python()
        elif len(args) == 1:
            inner_dtype = cls.infer_from_type(args[0])
        else:
            raise TypeError(f"Python list type cannot have more than one type argument, found: {t}")

        return cls.list(inner_dtype)
    elif is_typeddict(origin):
        field_types = typing.get_type_hints(origin)
        if any(not isinstance(t, str) for t in field_types):
            warnings.warn(
                f"Expected all TypedDict keys to be strings, found: {field_types}. Defaulting to Map[Python, Python] type."
            )
            return cls.map(cls.python(), cls.python())

        field_dtypes = {k: cls.infer_from_type(v) for k, v in field_types.items()}
        return cls.struct(field_dtypes)
    elif check_type(dict):
        if len(args) == 0:
            key_dtype = cls.python()
            value_dtype = cls.python()
        elif len(args) == 2:
            key_dtype = cls.infer_from_type(args[0])
            value_dtype = cls.infer_from_type(args[1])
        else:
            raise TypeError(f"Python dict type must have exactly two type arguments, found: {t}")

        # dict type can also be turned into struct, but we are unable to derive the struct keys from the type alone
        return cls.map(key_dtype, value_dtype)
    elif check_type(tuple):
        if len(args) == 0:
            # tuple -> List[Python]
            return cls.list(cls.python())
        if len(args) == 2 and args[1] is Ellipsis:
            # tuple[inner_type, ...] -> List[inner_type]
            inner_dtype = cls.infer_from_type(args[0])
            return cls.list(inner_dtype)
        else:
            # tuple[type0, type1, ...] -> Struct[_0: type0, _1: type1, ...]
            field_dtypes = {f"_{i}": cls.infer_from_type(arg) for i, arg in enumerate(args)}
            return cls.struct(field_dtypes)
    elif check_type(decimal.Decimal):
        warnings.warn(
            "Cannot derive precision and scale from decimal.Decimal type, defaulting to DataType.python()"
        )
        return cls.python()
    elif check_type(daft.series.Series):
        warnings.warn(
            "Cannot derive inner type from daft.Series type, defaulting to DataType.python() for Series inner type"
        )
        return cls.list(cls.python())
    elif check_type("pydantic.BaseModel"):
        import pydantic

        if not (parse("2.0.0") <= parse(pydantic.__version__) < parse("3.0.0")):
            raise ValueError(
                f"Daft only supports DataType inference for Pydantic V2, found Pydantic V{pydantic.__version__}"
            )

        model: pydantic.BaseModel = origin

        serialize_by_alias = model.model_config.get("serialize_by_alias", False)

        field_dtypes = {}
        for attr_name, field_info in model.model_fields.items():
            if serialize_by_alias:
                if field_info.serialization_alias is not None:
                    serialized_name = field_info.serialization_alias
                elif field_info.alias is not None:
                    serialized_name = field_info.alias
                else:
                    serialized_name = attr_name
            else:
                serialized_name = attr_name
            field_dtypes[serialized_name] = cls.infer_from_type(field_info.annotation)

        for attr_name, field_info in model.model_computed_fields.items():
            if serialize_by_alias:
                if field_info.alias is not None:
                    serialized_name = field_info.alias
                else:
                    serialized_name = attr_name
            else:
                serialized_name = attr_name
            field_dtypes[serialized_name] = cls.infer_from_type(field_info.return_type)

        return cls.struct(field_dtypes)
    elif check_type("PIL.Image.Image"):
        return cls.image()
    elif check_type("jaxtyping.AbstractArray"):
        return cls._infer_from_jaxtyping(origin)
    elif check_type("numpy.ndarray"):
        inner_dtype = None
        if len(args) == 2:
            # https://numpy.org/doc/2.3/reference/typing.html#numpy.typing.NDArray
            numpy_dtype_args = typing.get_args(args[1])
            if len(numpy_dtype_args) == 1:
                inner_dtype = cls.infer_from_type(numpy_dtype_args[0])

        if inner_dtype is None:
            warnings.warn(
                f"Cannot derive inner type from {t}, defaulting to DataType.python() for ndarray inner type"
            )
            inner_dtype = cls.python()

        return cls.tensor(inner_dtype)

    elif check_type("torch.FloatTensor"):
        return cls.tensor(cls.float32())
    elif check_type("torch.DoubleTensor"):
        return cls.tensor(cls.float64())
    elif check_type("torch.ByteTensor"):
        return cls.tensor(cls.uint8())
    elif check_type("torch.CharTensor"):
        return cls.tensor(cls.int8())
    elif check_type("torch.ShortTensor"):
        return cls.tensor(cls.int16())
    elif check_type("torch.IntTensor"):
        return cls.tensor(cls.int32())
    elif check_type("torch.LongTensor"):
        return cls.tensor(cls.int64())
    elif check_type("torch.BoolTensor"):
        return cls.tensor(cls.bool())
    elif (
        check_type("torch.Tensor")
        or check_type("tensorflow.Tensor")
        or check_type("jax.Array")
        or check_type("cupy.ndarray")
    ):
        return cls.tensor(cls.python())
    elif check_type("numpy.generic"):
        return cls._infer_from_numpy_scalar_dtype(origin)
    elif check_type("pandas.Series"):
        warnings.warn(
            "Cannot derive inner type from pandas.Series type, defaulting to DataType.python() for Series inner type"
        )
        return cls.list(cls.python())
    else:
        return cls.python()

def is_binary(self) -> builtins.bool:
    """Check if this is a binary type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.binary  # or daft.DataType.binary()
        >>> assert dtype.is_binary()
    """
    return self._dtype.is_binary()

def is_boolean(self) -> builtins.bool:
    """Check if this is a boolean type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.bool  # or daft.DataType.bool()
        >>> assert dtype.is_boolean()
    """
    return self._dtype.is_boolean()

def is_date(self) -> builtins.bool:
    """Check if this is a date type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.date  # or daft.DataType.date()
        >>> assert dtype.is_date()
    """
    return self._dtype.is_date()

def is_decimal128(self) -> builtins.bool:
    """Check if this is a decimal128 type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.decimal128(precision=10, scale=2)
        >>> assert dtype.is_decimal128()
    """
    return self._dtype.is_decimal128()

def is_duration(self) -> builtins.bool:
    """Check if this is a duration type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.duration(timeunit="ns")
        >>> assert dtype.is_duration()
    """
    return self._dtype.is_duration()

def is_embedding(self) -> builtins.bool:
    """Check if this is an embedding type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.embedding(daft.DataType.float32(), 512)
        >>> assert dtype.is_embedding()
    """
    return self._dtype.is_embedding()

def is_extension(self) -> builtins.bool:
    """Check if this is an extension type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.extension("custom", daft.DataType.int64())
        >>> assert dtype.is_extension()
    """
    return self._dtype.is_extension()

def is_file(self) -> builtins.bool:
    """Check if this is a file type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.file()
        >>> assert dtype.is_file()
    """
    return self._dtype.is_file()

def is_fixed_shape_image(self) -> builtins.bool:
    """Check if this is a fixed shape image type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.image(mode="RGB", height=224, width=224)
        >>> assert dtype.is_fixed_shape_image()
    """
    return self._dtype.is_fixed_shape_image()

def is_fixed_shape_sparse_tensor(self) -> builtins.bool:
    """Check if this is a fixed shape sparse tensor type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.sparse_tensor(daft.DataType.float32(), shape=(2, 3))
        >>> assert dtype.is_fixed_shape_sparse_tensor()
    """
    return self._dtype.is_fixed_shape_sparse_tensor()

def is_fixed_shape_tensor(self) -> builtins.bool:
    """Check if this is a fixed shape tensor type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.tensor(daft.DataType.float32(), shape=(2, 3))
        >>> assert dtype.is_fixed_shape_tensor()
    """
    return self._dtype.is_fixed_shape_tensor()

def is_fixed_size_binary(self) -> builtins.bool:
    """Check if this is a fixed size binary type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.fixed_size_binary(size=10)
        >>> assert dtype.is_fixed_size_binary()
    """
    return self._dtype.is_fixed_size_binary()

def is_fixed_size_list(self) -> builtins.bool:
    """Check if this is a fixed size list type.

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.fixed_size_list(daft.DataType.int64(), size=10)
        >>> assert dtype.is_fixed_size_list()
    """
    return self._dtype.is_fixed_size_list()

def is_float16(self) -> builtins.bool:
    """Check if this is a 16-bit float type."""
    return self._dtype.is_float16()