wehub-resource-sync
60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
try:
|
|
import pyarrow
|
|
except ImportError:
|
|
pyarrow = None
|
|
|
|
|
|
def _is_pa_extension_type(pa_type: "pyarrow.lib.DataType") -> bool:
|
|
"""Whether the provided Arrow Table column is an extension array, using an Arrow
|
|
extension type.
|
|
"""
|
|
# NOTE: Native Tensors are also BaseExtensionType
|
|
return isinstance(pa_type, pyarrow.BaseExtensionType)
|
|
|
|
|
|
def _is_native_tensor_type(t: "pyarrow.BaseExtentionType") -> bool:
|
|
"""Whether the provided Arrow Table column is an native Tensor array"""
|
|
from ray.data.extensions import FixedShapeTensorType
|
|
|
|
return FixedShapeTensorType is not None and isinstance(t, FixedShapeTensorType)
|
|
|
|
|
|
def _concatenate_extension_column(
|
|
ca: "pyarrow.ChunkedArray", ensure_copy: bool = False
|
|
) -> "pyarrow.Array":
|
|
"""Concatenate chunks of an extension column into a contiguous array.
|
|
|
|
This concatenation is required for creating copies and for .take() to work on
|
|
extension arrays.
|
|
See https://issues.apache.org/jira/browse/ARROW-16503.
|
|
|
|
Args:
|
|
ca: The chunked array representing the extension column to be concatenated.
|
|
ensure_copy: Skip copying when ensure_copy is False and there is exactly 1 chunk.
|
|
|
|
Returns:
|
|
Array: the concatenate extension column.
|
|
"""
|
|
from ray.data._internal.tensor_extensions.arrow import (
|
|
concat_tensor_arrays,
|
|
get_arrow_extension_tensor_types,
|
|
)
|
|
|
|
if not _is_pa_extension_type(ca.type):
|
|
raise ValueError(f"Chunked array isn't an extension array: {ca.type}")
|
|
|
|
tensor_extension_types = get_arrow_extension_tensor_types()
|
|
|
|
if ca.num_chunks == 0:
|
|
# Create empty storage array.
|
|
storage = pyarrow.array([], type=ca.type.storage_type)
|
|
elif not ensure_copy and len(ca.chunks) == 1:
|
|
# Skip copying
|
|
return ca.chunks[0]
|
|
elif isinstance(ca.type, tensor_extension_types):
|
|
return concat_tensor_arrays(ca.chunks, ensure_copy)
|
|
else:
|
|
storage = pyarrow.concat_arrays([c.storage for c in ca.chunks])
|
|
|
|
return ca.type.wrap_array(storage)
|