wehub-resource-sync
117 lines
3.9 KiB
Python
117 lines
3.9 KiB
Python
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
|