Class Schema

python/ray/data/dataset.py:7662–7786

Dataset schema. Attributes: base_schema: The underlying Arrow or Pandas schema.

@PublicAPI(stability="beta")
class Schema:
    """Dataset schema.

    Attributes:
        base_schema: The underlying Arrow or Pandas schema.
    """

    def __init__(
        self,
        base_schema: Union["pyarrow.lib.Schema", "PandasBlockSchema"],
        *,
        data_context: Optional[DataContext] = None,
    ):
        """
        Initialize a :class:`Schema` wrapper around an Arrow or Pandas schema.

        Args:
            base_schema: The underlying Arrow or Pandas schema.
            data_context: The data context to use for this schema.
        """
        self.base_schema = base_schema

        # Snapshot the current context, so that the config of Datasets is always
        # determined by the config at the time it was created.
        self._context = data_context or copy.deepcopy(DataContext.get_current())

    @property
    def names(self) -> List[str]:
        """Lists the columns of this Dataset."""
        from ray.data._internal.arrow_block import _is_user_visible_column

        # ``__bsp_stub`` is a physical placeholder the read path injects
        # into zero-column blocks so ``pa.concat_tables`` doesn't collapse
        # the row count. It's not part of the user-visible schema.
        return [
            name for name in self.base_schema.names if _is_user_visible_column(name)
        ]

    @property
    def types(self) -> List[Union[type[object], "pyarrow.lib.DataType"]]:
        """Lists the types of this Dataset in Arrow format

        For non-Arrow compatible types, we return "object".
        """
        import pandas as pd
        import pyarrow as pa
        from pandas.core.dtypes.dtypes import BaseMaskedDtype

        from ray.data._internal.arrow_block import _is_user_visible_column
        from ray.data._internal.tensor_extensions.arrow import (
            create_arrow_fixed_shape_tensor_type,
        )
        from ray.data.extensions import TensorDtype

        def _convert_to_pa_type(
            dtype: Union[np.dtype, pd.ArrowDtype, BaseMaskedDtype],
        ) -> pa.DataType:
            if isinstance(dtype, pd.ArrowDtype):