Repository: ray-project/ray
File: python/ray/air/util/data_batch_conversion.py (Python)

Code before:
    """
    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.
    """
    pd = _lazy_import_pandas()
    # Pandas has moved/renamed this warning across versions, and pandas 3.x may not
    # expose it at all. If we can't find it, just don't attempt to filter it.
    SettingWithCopyWarning = None
    try:
        SettingWithCopyWarning = pd.core.common.SettingWithCopyWarning
    except Exception:
        try:
            SettingWithCopyWarning = pd.errors.SettingWithCopyWarning
        except Exception:
            SettingWithCopyWarning = None

    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.
    for col_name, col in df.items():
        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)
                    df[col_name] = TensorArray(col)

Review comment (refactor):
Let's abstract this into a util we can share to avoid duping it.

Also, let's use `has_attr` to avoid try-catches

Code after:
    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.
    for col_name, col in df.items():
        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)
                    df[col_name] = 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