ray-project--ray/python/ray/data/util/data_batch_conversion.py

import warnings
from enum import Enum
from typing import TYPE_CHECKING, Dict, List, Union

import numpy as np

from ray.air.data_batch_type import DataBatchType
from ray.data.constants import TENSOR_COLUMN_NAME
from ray.data.util.expression_utils import _get_setting_with_copy_warning
from ray.util.annotations import DeveloperAPI

if TYPE_CHECKING:
    import pandas as pd

# TODO: Consolidate data conversion edges for arrow bug workaround.
try:
    import pyarrow
except ImportError:
    pyarrow = None

# Lazy import to avoid ray init failures without pandas installed and allow
# dataset to import modules in this file.
_pandas = None


def _lazy_import_pandas():
    global _pandas
    if _pandas is None:
        import pandas

        _pandas = pandas
    return _pandas


@DeveloperAPI
class BatchFormat(str, Enum):
    PANDAS = "pandas"
    # TODO: Remove once Arrow is deprecated as user facing batch format
    ARROW = "arrow"
    NUMPY = "numpy"  # Either a single numpy array or a Dict of numpy arrays.
    CUDF = "cudf"


_CUDF_UNSET = object()
_cudf = _CUDF_UNSET


def _lazy_import_cudf():
    """Lazy import cudf, returning the module or None if not installed."""
    global _cudf
    if _cudf is _CUDF_UNSET:
        try:
            import cudf

            _cudf = cudf
        except ImportError:
            _cudf = None
    return _cudf


def _convert_batch_type_to_pandas(
    data: DataBatchType,
    cast_tensor_columns: bool = False,
) -> "pd.DataFrame":
    """Convert the provided data to a Pandas DataFrame.

    Args:
        data: Data of type DataBatchType
        cast_tensor_columns: Whether tensor columns should be cast to NumPy ndarrays.

    Returns:
        A pandas Dataframe representation of the input data.

    """
    pd = _lazy_import_pandas()

    if isinstance(data, np.ndarray):
        data = pd.DataFrame({TENSOR_COLUMN_NAME: _ndarray_to_column(data)})
    elif isinstance(data, dict):
        tensor_dict = {}
        for col_name, col in data.items():
            if not isinstance(col, np.ndarray):
                raise ValueError(
                    "All values in the provided dict must be of type "
                    f"np.ndarray. Found type {type(col)} for key {col_name} "
                    f"instead."
                )
            tensor_dict[col_name] = _ndarray_to_column(col)
        data = pd.DataFrame(tensor_dict)
    elif pyarrow is not None and isinstance(data, pyarrow.Table):
        data = data.to_pandas()
    else:
        # Handle cudf.DataFrame (lazy check to avoid import when not used)
        cudf = _lazy_import_cudf()
        if cudf is not None and isinstance(data, cudf.DataFrame):
            data = data.to_pandas()
        if not isinstance(data, pd.DataFrame):
            raise ValueError(
                f"Received data of type: {type(data)}, but expected it to be one "
                f"of {DataBatchType}"
            )
    if cast_tensor_columns:
        data = _cast_tensor_columns_to_ndarrays(data)
    return data


def _convert_pandas_to_batch_type(
    data: "pd.DataFrame",
    type: BatchFormat,
    cast_tensor_columns: bool = False,
) -> DataBatchType:
    """Convert the provided Pandas dataframe to the provided ``type``.

    Args:
        data: A Pandas DataFrame
        type: The specific ``BatchFormat`` to convert to.
        cast_tensor_columns: Whether tensor columns should be cast to our tensor
            extension type.

    Returns:
        The input data represented with the provided type.
    """
    if cast_tensor_columns:
        data = _cast_ndarray_columns_to_tensor_extension(data)

    if type == BatchFormat.PANDAS:
        return data
    elif type == BatchFormat.NUMPY:
        if len(data.columns) == 1:
            # If just a single column, return as a single numpy array.
            return data.iloc[:, 0].to_numpy()
        else:
            # Else return as a dict of numpy arrays.
            output_dict = {}
            for column in data:
                output_dict[column] = data[column].to_numpy()
            return output_dict
    elif type == BatchFormat.ARROW:
        if not pyarrow:
            raise ValueError(
                "Attempted to convert data to Pyarrow Table but Pyarrow "
                "is not installed. Please do `pip install pyarrow` to "
                "install Pyarrow."
            )
        return pyarrow.Table.from_pandas(data)
    elif type == BatchFormat.CUDF:
        cudf = _lazy_import_cudf()
        if cudf is None:
            raise ValueError(
                "Attempted to convert data to cuDF DataFrame but cuDF "
                "is not installed. Please do `pip install cudf-cu12` to "
                "install cuDF (GPU required)."
            )
        return cudf.from_pandas(data)
    else:
        raise ValueError(
            f"Received type {type}, but expected it to be one of {DataBatchType}"
        )


def _convert_batch_type_to_numpy(
    data: DataBatchType,
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
    """Convert the provided data to a NumPy ndarray or dict of ndarrays.

    Args:
        data: Data of type DataBatchType

    Returns:
        A numpy representation of the input data.
    """
    pd = _lazy_import_pandas()

    if isinstance(data, np.ndarray):
        return data
    elif isinstance(data, dict):
        for col_name, col in data.items():
            if not isinstance(col, np.ndarray):
                raise ValueError(
                    "All values in the provided dict must be of type "
                    f"np.ndarray. Found type {type(col)} for key {col_name} "
                    f"instead."
                )
        return data
    elif pyarrow is not None and isinstance(data, pyarrow.Table):
        from ray.data._internal.arrow_ops import transform_pyarrow
        from ray.data._internal.tensor_extensions.arrow import (
            get_arrow_extension_fixed_shape_tensor_types,
        )

        column_values_ndarrays = []

        for col in data.columns:
            # Combine columnar values arrays to make these contiguous
            # (making them compatible with numpy format)
            combined_array = transform_pyarrow.combine_chunked_array(col)

            column_values_ndarrays.append(
                transform_pyarrow.to_numpy(combined_array, zero_copy_only=False)
            )

        arrow_fixed_shape_tensor_types = get_arrow_extension_fixed_shape_tensor_types()

        # NOTE: This branch is here for backwards-compatibility
        if data.column_names == [TENSOR_COLUMN_NAME] and (
            isinstance(data.schema.types[0], arrow_fixed_shape_tensor_types)
        ):
            return column_values_ndarrays[0]

        return dict(zip(data.column_names, column_values_ndarrays))
    elif isinstance(data, pd.DataFrame):
        return _convert_pandas_to_batch_type(data, BatchFormat.NUMPY)
    else:
        # Handle cudf.DataFrame via pandas path
        cudf = _lazy_import_cudf()
        if cudf is not None and isinstance(data, cudf.DataFrame):
            return _convert_pandas_to_batch_type(data.to_pandas(), BatchFormat.NUMPY)
        raise ValueError(
            f"Received data of type: {type(data)}, but expected it to be one "
            f"of {DataBatchType}"
        )


def _ndarray_to_column(arr: np.ndarray) -> Union["pd.Series", List[np.ndarray]]:
    """Convert a NumPy ndarray into an appropriate column format for insertion into a
    pandas DataFrame.

    If conversion to a pandas Series fails (e.g. if the ndarray is multi-dimensional),
    fall back to a list of NumPy ndarrays.
    """
    pd = _lazy_import_pandas()
    try:
        # Try to convert to Series, falling back to a list conversion if this fails
        # (e.g. if the ndarray is multi-dimensional).
        return pd.Series(arr)
    except ValueError:
        return list(arr)


def _unwrap_ndarray_object_type_if_needed(arr: np.ndarray) -> np.ndarray:
    """Unwrap an object-dtyped NumPy ndarray containing ndarray pointers into a single
    contiguous ndarray, if needed/possible.
    """
    if arr.dtype.type is np.object_:
        try:
            # Try to convert the NumPy ndarray to a non-object dtype.
            arr = np.array([np.asarray(v) for v in arr])
        except Exception:
            # This may fail if the subndarrays are of heterogeneous shape
            pass
    return arr


def _cast_ndarray_columns_to_tensor_extension(df: "pd.DataFrame") -> "pd.DataFrame":
    """
    Cast all NumPy ndarray columns in df to our tensor extension type, TensorArray.
    """
    # Get the SettingWithCopyWarning class if available
    SettingWithCopyWarning = _get_setting_with_copy_warning()

    from ray.data._internal.tensor_extensions.pandas import (
        TensorArray,
        column_needs_tensor_extension,
    )

    # Try to convert any ndarray columns to TensorArray columns.
    # TODO(Clark): Once Pandas supports registering extension types for type
    # inference on construction, implement as much for NumPy ndarrays and remove
    # this. See https://github.com/pandas-dev/pandas/issues/41848
    # TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of
    # column names containing tensor columns, to make this an O(# of tensor columns)
    # check rather than the current O(# of columns) check.

    # Scan dtypes rather than df.items(), which would
    # materialize a Series for every column just to read its dtype.
    # The below approach avoids the cost of a Series build for non-tensor columns.
    #
    # When column names are unique we select and assign by label.
    # With duplicate names, ``df[col_name]`` returns a DataFrame
    # rather than a Series, so we select and assign by position instead.
    columns_unique = df.columns.is_unique
    for i, (col_name, dtype) in enumerate(df.dtypes.items()):
        if (
            dtype.type is not np.object_
        ):  # Short circuit if non-object type before materializing the column
            continue
        col = df[col_name] if columns_unique else df.iloc[:, i]
        if column_needs_tensor_extension(col):
            try:
                # Suppress Pandas warnings:
                # https://github.com/ray-project/ray/issues/29270
                # We actually want in-place operations so we surpress this warning.
                # https://stackoverflow.com/a/74193599
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore", category=FutureWarning)
                    if SettingWithCopyWarning is not None:
                        warnings.simplefilter("ignore", category=SettingWithCopyWarning)
                    if columns_unique:
                        df[col_name] = TensorArray(col)
                    else:
                        df.isetitem(i, TensorArray(col))
            except Exception as e:
                raise ValueError(
                    f"Tried to cast column {col_name} to the TensorArray tensor "
                    "extension type but the conversion failed. To disable "
                    "automatic casting to this tensor extension, set "
                    "ctx = DataContext.get_current(); "
                    "ctx.enable_tensor_extension_casting = False."
                ) from e
    return df


def _cast_tensor_columns_to_ndarrays(
    df: "pd.DataFrame",
    arrow_schema: "pyarrow.Schema" = None,
) -> "pd.DataFrame":
    """Cast all tensor extension columns in df to NumPy ndarrays.

    Args:
        df: The DataFrame whose tensor columns should be converted.
        arrow_schema: If provided, used to reshape columns that were native
            ``FixedShapeTensorType`` in Arrow.  PyArrow's ``to_pandas()``
            flattens these to 1-D ndarrays; passing the original schema
            lets us restore the correct shape.

    Returns:
        The DataFrame with tensor columns converted to NumPy ndarrays.
    """
    # Get the SettingWithCopyWarning class if available
    SettingWithCopyWarning = _get_setting_with_copy_warning()
    from ray.data._internal.tensor_extensions.pandas import TensorDtype

    # Try to convert any tensor extension columns to ndarray columns.
    # TODO(Clark): Optimize this with propagated DataFrame metadata containing a list of
    # column names containing tensor columns, to make this an O(# of tensor columns)
    # check rather than the current O(# of columns) check.

    # Reshape native FixedShapeTensorType columns that were flattened by
    # to_pandas().

    if arrow_schema is not None:
        from ray.data._internal.utils.transform_pyarrow import (
            _is_native_tensor_type,
        )

        for field in arrow_schema:
            if _is_native_tensor_type(field.type) and field.name in df.columns:
                shape = tuple(field.type.shape)
                df[field.name] = [
                    arr.reshape(shape) if arr is not None else None
                    for arr in df[field.name]
                ]

    # Scan dtypes rather than df.items(), which would
    # materialize a Series for every column just to read its dtype.
    # The below approach avoids the cost of a Series build for non-tensor columns.
    #
    # When column names are unique we select and assign by label (the fast,
    # cached path). With duplicate names, ``df[col_name]`` returns a DataFrame
    # rather than a Series, so we select and assign by position instead.
    columns_unique = df.columns.is_unique
    for i, (col_name, dtype) in enumerate(df.dtypes.items()):
        if isinstance(dtype, TensorDtype):
            # Suppress Pandas warnings:
            # https://github.com/ray-project/ray/issues/29270
            # We actually want in-place operations so we surpress this warning.
            # https://stackoverflow.com/a/74193599
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", category=FutureWarning)
                if SettingWithCopyWarning is not None:
                    warnings.simplefilter("ignore", category=SettingWithCopyWarning)
                if columns_unique:
                    df[col_name] = list(df[col_name].to_numpy())
                else:
                    df.isetitem(i, list(df.iloc[:, i].to_numpy()))
    return df