chore: import upstream snapshot with attribution

python/ray/data/_internal/object_extensions/arrow.py

@@ -0,0 +1,116 @@
+import pickle
+import typing
+
+import numpy as np
+import pyarrow as pa
+
+import ray.data._internal.object_extensions.pandas
+from ray._common.serialization import pickle_dumps
+from ray.data._internal.utils.arrow_utils import _check_pyarrow_version
+from ray.util.annotations import PublicAPI
+
+# First, assert Arrow version is w/in expected bounds
+_check_pyarrow_version()
+
+
+# Please see https://arrow.apache.org/docs/python/extending_types.html for more info
+@PublicAPI(stability="alpha")
+class ArrowPythonObjectType(pa.ExtensionType):
+    """Defines a new Arrow extension type for Python objects.
+    We do not require a parametrized type, so the constructor does not
+    take any arguments
+    """
+
+    def __init__(self) -> None:
+        # Defines the underlying storage type as the PyArrow LargeBinary type
+        super().__init__(pa.large_binary(), "ray.data.arrow_pickled_object")
+
+    def __arrow_ext_serialize__(self) -> bytes:
+        # Since there are no type parameters, we are free to return empty
+        return b""
+
+    @classmethod
+    def __arrow_ext_deserialize__(
+        cls, storage_type: pa.DataType, serialized: bytes
+    ) -> "ArrowPythonObjectType":
+        return ArrowPythonObjectType()
+
+    def __arrow_ext_scalar_class__(self) -> type:
+        """Returns the scalar class of the extension type. Indexing out of the
+        PyArrow extension array will return instances of this type.
+        """
+        return ArrowPythonObjectScalar
+
+    def __arrow_ext_class__(self) -> type:
+        """Returns the array type of the extension type. Selecting one array
+        out of the ChunkedArray that makes up a column in a Table with
+        this custom type will return an instance of this type.
+        """
+        return ArrowPythonObjectArray
+
+    def to_pandas_dtype(self):
+        """Pandas interoperability type. This describes the Pandas counterpart
+        to the Arrow type. See https://pandas.pydata.org/docs/development/extending.html
+        for more information.
+        """
+        return ray.data._internal.object_extensions.pandas.PythonObjectDtype()
+
+    def __reduce__(self):
+        # Earlier PyArrow versions require custom pickling behavior.
+        return self.__arrow_ext_deserialize__, (
+            self.storage_type,
+            self.__arrow_ext_serialize__(),
+        )
+
+    def __hash__(self) -> int:
+        return hash((type(self), self.storage_type.id, self.extension_name))
+
+
+@PublicAPI(stability="alpha")
+class ArrowPythonObjectScalar(pa.ExtensionScalar):
+    """Scalar class for ArrowPythonObjectType"""
+
+    def as_py(self, **kwargs) -> typing.Any:
+        # Handle None/null values
+        if self.value is None:
+            return None
+
+        if not isinstance(self.value, pa.LargeBinaryScalar):
+            raise RuntimeError(
+                f"{type(self.value)} is not the expected LargeBinaryScalar"
+            )
+        return pickle.load(pa.BufferReader(self.value.as_buffer()))
+
+
+@PublicAPI(stability="alpha")
+class ArrowPythonObjectArray(pa.ExtensionArray):
+    """Array class for ArrowPythonObjectType"""
+
+    def from_objects(
+        objects: typing.Union[np.ndarray, typing.Iterable[typing.Any]]
+    ) -> "ArrowPythonObjectArray":
+        if isinstance(objects, np.ndarray):
+            objects = objects.tolist()
+        type_ = ArrowPythonObjectType()
+        all_dumped_bytes = []
+        for obj in objects:
+            dumped_bytes = pickle_dumps(
+                obj, "Error pickling object to convert to Arrow"
+            )
+            all_dumped_bytes.append(dumped_bytes)
+        arr = pa.array(all_dumped_bytes, type=type_.storage_type)
+        return type_.wrap_array(arr)
+
+    def to_numpy(
+        self, zero_copy_only: bool = False, writable: bool = False
+    ) -> np.ndarray:
+        arr = np.empty(len(self), dtype=object)
+        arr[:] = self.to_pylist()
+        return arr
+
+
+try:
+    pa.register_extension_type(ArrowPythonObjectType())
+except pa.ArrowKeyError:
+    # Already registered
+    pass

python/ray/data/_internal/object_extensions/pandas.py

@@ -0,0 +1,146 @@
+import collections.abc
+import typing
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+from pandas._libs import lib
+from pandas._typing import ArrayLike, Dtype, PositionalIndexer, TakeIndexer, npt
+
+import ray.data._internal.object_extensions.arrow
+from ray.util.annotations import PublicAPI
+
+
+# See https://pandas.pydata.org/docs/development/extending.html for more information.
+@PublicAPI(stability="alpha")
+class PythonObjectArray(pd.api.extensions.ExtensionArray):
+    """Implements the Pandas extension array interface for the Arrow object array"""
+
+    def __init__(self, values: collections.abc.Iterable[typing.Any]):
+        vals = list(values)
+        self.values = np.empty(len(vals), dtype=object)
+        self.values[:] = vals
+
+    @classmethod
+    def _from_sequence(
+        cls,
+        scalars: collections.abc.Sequence[typing.Any],
+        *,
+        dtype: typing.Union[Dtype, None] = None,
+        copy: bool = False,
+    ) -> "PythonObjectArray":
+        return PythonObjectArray(scalars)
+
+    @classmethod
+    def _from_factorized(
+        cls, values: collections.abc.Sequence[typing.Any], original: "PythonObjectArray"
+    ) -> "PythonObjectArray":
+        return PythonObjectArray(values)
+
+    def __getitem__(self, item: PositionalIndexer) -> typing.Any:
+        return self.values[item]
+
+    def __setitem__(self, key, value) -> None:
+        self.values[key] = value
+
+    def __len__(self) -> int:
+        return len(self.values)
+
+    def __eq__(self, other: object) -> ArrayLike:
+        if isinstance(other, PythonObjectArray):
+            return self.values == other.values
+        elif isinstance(other, np.ndarray):
+            return self.values == other
+        else:
+            return NotImplemented
+
+    def to_numpy(
+        self,
+        dtype: typing.Union["npt.DTypeLike", None] = None,
+        copy: bool = False,
+        na_value: object = lib.no_default,
+    ) -> np.ndarray:
+        result = self.values
+        if copy or na_value is not lib.no_default:
+            result = result.copy()
+        if na_value is not lib.no_default:
+            result[self.isna()] = na_value
+        return result
+
+    @property
+    def dtype(self) -> pd.api.extensions.ExtensionDtype:
+        return PythonObjectDtype()
+
+    @property
+    def nbytes(self) -> int:
+        return self.values.nbytes
+
+    def __arrow_array__(self, type=None):
+        return ray.data._internal.object_extensions.arrow.ArrowPythonObjectArray.from_objects(
+            self.values
+        )
+
+    def isna(self) -> np.ndarray:
+        return pd.isnull(self.values)
+
+    def take(
+        self,
+        indices: TakeIndexer,
+        *,
+        allow_fill: bool = False,
+        fill_value: typing.Any = None,
+    ) -> "PythonObjectArray":
+        if allow_fill and fill_value is None:
+            fill_value = self.dtype.na_value
+
+        result = pd.core.algorithms.take(
+            self.values, indices, allow_fill=allow_fill, fill_value=fill_value
+        )
+        return self._from_sequence(result, dtype=self.dtype)
+
+    def copy(self) -> "PythonObjectArray":
+        return PythonObjectArray(self.values)
+
+    @classmethod
+    def _concat_same_type(
+        cls, to_concat: collections.abc.Sequence["PythonObjectArray"]
+    ) -> "PythonObjectArray":
+        values_to_concat = [element.values for element in to_concat]
+        return cls(np.concatenate(values_to_concat))
+
+
+@PublicAPI(stability="alpha")
+@pd.api.extensions.register_extension_dtype
+class PythonObjectDtype(pd.api.extensions.ExtensionDtype):
+    @classmethod
+    def construct_from_string(cls, string: str):
+        if string != "python_object()":
+            raise TypeError(f"Cannot construct a '{cls.__name__}' from '{string}'")
+        return cls()
+
+    @property
+    def type(self):
+        """
+        The scalar type for the array, e.g. ``int``
+        It's expected ``ExtensionArray[item]`` returns an instance
+        of ``ExtensionDtype.type`` for scalar ``item``, assuming
+        that value is valid (not NA). NA values do not need to be
+        instances of `type`.
+        """
+        return object
+
+    @property
+    def name(self) -> str:
+        return "python_object()"
+
+    @classmethod
+    def construct_array_type(cls: type) -> type:
+        """
+        Return the array type associated with this dtype.
+        """
+        return PythonObjectArray
+
+    def __from_arrow__(
+        self, array: typing.Union[pa.Array, pa.ChunkedArray]
+    ) -> PythonObjectArray:
+        return PythonObjectArray(array.to_pylist())