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