DataTypes#

@classmethod
def binary(cls) -> DataType:
    """Create a Binary DataType: A string of bytes."""
    return cls._from_pydatatype(PyDataType.binary())

@classmethod
def bool(cls) -> DataType:
    """Create the Boolean DataType: Either ``True`` or ``False``."""
    return cls._from_pydatatype(PyDataType.bool())

@classmethod
def date(cls) -> DataType:
    """Create a Date DataType: A date with a year, month and day."""
    return cls._from_pydatatype(PyDataType.date())

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

@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))

@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))

@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))

@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))

@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 float32(cls) -> DataType:
    """Create a 32-bit float DataType."""
    return cls._from_pydatatype(PyDataType.float32())

@classmethod
def float64(cls) -> DataType:
    """Create a 64-bit float DataType."""
    return cls._from_pydatatype(PyDataType.float64())

@classmethod
def from_arrow_type(cls, arrow_type: pa.lib.DataType) -> 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_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))
    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), 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) 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),
            value_type=cls.from_arrow_type(arrow_type.item_type),
        )
    elif isinstance(arrow_type, getattr(pa, "FixedShapeTensorType", ())):
        scalar_dtype = cls.from_arrow_type(arrow_type.value_type)
        return cls.tensor(scalar_dtype, tuple(arrow_type.shape))
    elif isinstance(arrow_type, 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 Arrow2 implementation: {arrow_type}"
        )
    elif isinstance(arrow_type, pa.BaseExtensionType):
        name = arrow_type.extension_name

        if (get_context().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),
                metadata,
            )
    else:
        # Fall back to a Python object type.
        # TODO(Clark): Add native support for remaining Arrow types.
        return cls.python()

@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 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 int16(cls) -> DataType:
    """Create an 16-bit integer DataType."""
    return cls._from_pydatatype(PyDataType.int16())

@classmethod
def int32(cls) -> DataType:
    """Create an 32-bit integer DataType."""
    return cls._from_pydatatype(PyDataType.int32())

@classmethod
def int64(cls) -> DataType:
    """Create an 64-bit integer DataType."""
    return cls._from_pydatatype(PyDataType.int64())

@classmethod
def int8(cls) -> DataType:
    """Create an 8-bit integer DataType."""
    return cls._from_pydatatype(PyDataType.int8())

@classmethod
def interval(cls) -> DataType:
    """Interval DataType."""
    return cls._from_pydatatype(PyDataType.interval())

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

    Examples:
        >>> import daft
        >>> dtype = 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()
        >>> 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()
        >>> 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_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_float32(self) -> builtins.bool:
    """Check if this is a 32-bit float type.

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

def is_float64(self) -> builtins.bool:
    """Check if this is a 64-bit float type.

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

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

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

def is_int16(self) -> builtins.bool:
    """Check if this is a 16-bit integer type.

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

def is_int32(self) -> builtins.bool:
    """Check if this is a 32-bit integer type.

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

def is_int64(self) -> builtins.bool:
    """Check if this is a 64-bit integer type.

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

def is_int8(self) -> builtins.bool:
    """Check if this is an 8-bit integer type.

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

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

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

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

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

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

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.list(daft.DataType.int64())
        >>> assert dtype.is_list()
    """
    return self._dtype.is_list()

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

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.bool()
        >>> assert not dtype.is_logical()
    """
    return self._dtype.is_logical()

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

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.map(daft.DataType.string(), daft.DataType.int64())
        >>> assert dtype.is_map()
    """
    return self._dtype.is_map()

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

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.null()
        >>> dtype.is_null()
        True
    """
    return self._dtype.is_null()

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

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

def is_python(self) -> builtins.bool:
    """Check if this is a python object type.

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

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

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.sparse_tensor(daft.DataType.float32())
        >>> assert dtype.is_sparse_tensor()
    """
    return self._dtype.is_sparse_tensor()

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

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

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

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.struct({"a": daft.DataType.int64()})
        >>> assert dtype.is_struct()
    """
    return self._dtype.is_struct()

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

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

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

    Examples:
        >>> import daft
        >>> dtype = daft.DataType.tensor(daft.DataType.float32())
        >>> assert dtype.is_tensor()
    """
    return self._dtype.is_tensor()

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

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