chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,137 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from typing import Dict, Optional
|
||||
from urllib.parse import parse_qsl, unquote, urlencode, urlparse, urlunparse
|
||||
|
||||
from packaging.version import Version, parse as parse_version
|
||||
|
||||
_RAY_DISABLE_PYARROW_VERSION_CHECK = "RAY_DISABLE_PYARROW_VERSION_CHECK"
|
||||
|
||||
|
||||
_PYARROW_INSTALLED: Optional[bool] = None
|
||||
_PYARROW_VERSION: Optional[Version] = None
|
||||
|
||||
|
||||
# NOTE: Make sure that these lower and upper bounds stay in sync with version
|
||||
# constraints given in python/setup.py.
|
||||
# Inclusive minimum pyarrow version.
|
||||
_PYARROW_SUPPORTED_VERSION_MIN = "17.0.0"
|
||||
_PYARROW_VERSION_VALIDATED = False
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _check_pyarrow_version():
|
||||
"""Checks that Pyarrow's version is within the supported bounds."""
|
||||
global _PYARROW_VERSION_VALIDATED
|
||||
|
||||
if os.environ.get("RAY_DOC_BUILD", "0") == "1":
|
||||
return
|
||||
|
||||
if not _PYARROW_VERSION_VALIDATED:
|
||||
if os.environ.get(_RAY_DISABLE_PYARROW_VERSION_CHECK, "0") == "1":
|
||||
_PYARROW_VERSION_VALIDATED = True
|
||||
return
|
||||
|
||||
version = get_pyarrow_version()
|
||||
if version is not None:
|
||||
if version < parse_version(_PYARROW_SUPPORTED_VERSION_MIN):
|
||||
raise ImportError(
|
||||
f"Dataset requires pyarrow >= {_PYARROW_SUPPORTED_VERSION_MIN}, but "
|
||||
f"{version} is installed. Reinstall with "
|
||||
f'`pip install -U "pyarrow"`. '
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"You are using the 'pyarrow' module, but the exact version is unknown "
|
||||
"(possibly carried as an internal component by another module). Please "
|
||||
f"make sure you are using pyarrow >= {_PYARROW_SUPPORTED_VERSION_MIN} to ensure "
|
||||
"compatibility with Ray Dataset. "
|
||||
)
|
||||
|
||||
_PYARROW_VERSION_VALIDATED = True
|
||||
|
||||
|
||||
def get_pyarrow_version() -> Optional[Version]:
|
||||
"""Get the version of the pyarrow package or None if not installed."""
|
||||
global _PYARROW_INSTALLED, _PYARROW_VERSION
|
||||
if _PYARROW_INSTALLED is False:
|
||||
return None
|
||||
|
||||
if _PYARROW_INSTALLED is None:
|
||||
try:
|
||||
import pyarrow
|
||||
|
||||
_PYARROW_INSTALLED = True
|
||||
if hasattr(pyarrow, "__version__"):
|
||||
_PYARROW_VERSION = parse_version(pyarrow.__version__)
|
||||
except ModuleNotFoundError:
|
||||
_PYARROW_INSTALLED = False
|
||||
|
||||
return _PYARROW_VERSION
|
||||
|
||||
|
||||
def _add_url_query_params(url: str, params: Dict[str, str]) -> str:
|
||||
"""Add params to the provided url as query parameters.
|
||||
|
||||
If url already contains query parameters, they will be merged with params, with the
|
||||
existing query parameters overriding any in params with the same parameter name.
|
||||
|
||||
Args:
|
||||
url: The URL to add query parameters to.
|
||||
params: The query parameters to add.
|
||||
|
||||
Returns:
|
||||
URL with params added as query parameters.
|
||||
"""
|
||||
# Unquote URL first so we don't lose existing args.
|
||||
url = unquote(url)
|
||||
# Parse URL.
|
||||
parsed_url = urlparse(url)
|
||||
# Merge URL query string arguments dict with new params.
|
||||
base_params = params
|
||||
params = dict(parse_qsl(parsed_url.query))
|
||||
base_params.update(params)
|
||||
# bool and dict values should be converted to json-friendly values.
|
||||
base_params.update(
|
||||
{
|
||||
k: json.dumps(v)
|
||||
for k, v in base_params.items()
|
||||
if isinstance(v, (bool, dict))
|
||||
}
|
||||
)
|
||||
|
||||
# Convert URL arguments to proper query string.
|
||||
encoded_params = urlencode(base_params, doseq=True)
|
||||
# Replace query string in parsed URL with updated query string.
|
||||
parsed_url = parsed_url._replace(query=encoded_params)
|
||||
# Convert back to URL.
|
||||
return urlunparse(parsed_url)
|
||||
|
||||
|
||||
def add_creatable_buckets_param_if_s3_uri(uri: str) -> str:
|
||||
"""If the provided URI is an S3 URL, add allow_bucket_creation=true as a query
|
||||
parameter. For pyarrow >= 9.0.0, this is required in order to allow
|
||||
``S3FileSystem.create_dir()`` to create S3 buckets.
|
||||
|
||||
If the provided URI is not an S3 URL or if pyarrow < 9.0.0 is installed, we return
|
||||
the URI unchanged.
|
||||
|
||||
Args:
|
||||
uri: The URI that we'll add the query parameter to, if it's an S3 URL.
|
||||
|
||||
Returns:
|
||||
A URI with the added allow_bucket_creation=true query parameter, if the provided
|
||||
URI is an S3 URL; uri will be returned unchanged otherwise.
|
||||
"""
|
||||
|
||||
pyarrow_version = get_pyarrow_version()
|
||||
if pyarrow_version is not None and pyarrow_version < parse_version("9.0.0"):
|
||||
# This bucket creation query parameter is not required for pyarrow < 9.0.0.
|
||||
return uri
|
||||
parsed_uri = urlparse(uri)
|
||||
if parsed_uri.scheme == "s3":
|
||||
uri = _add_url_query_params(uri, {"allow_bucket_creation": True})
|
||||
return uri
|
||||
@@ -0,0 +1,157 @@
|
||||
"""
|
||||
Copyright 2009 Stutzbach Enterprises, LLC (daniel@stutzbachenterprises.com)
|
||||
|
||||
Redistribution and use in source and binary forms, with or without
|
||||
modification, are permitted provided that the following conditions are
|
||||
met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
2. Redistributions in binary form must reproduce the above
|
||||
copyright notice, this list of conditions and the following
|
||||
disclaimer in the documentation and/or other materials provided
|
||||
with the distribution.
|
||||
3. The name of the author may not be used to endorse or promote
|
||||
products derived from this software without specific prior written
|
||||
permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
|
||||
IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
||||
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||
DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,
|
||||
INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
|
||||
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
|
||||
HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
|
||||
STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
|
||||
IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
|
||||
POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
|
||||
Vendorized heapdict implementation.
|
||||
|
||||
This is a copy of the heapdict library to avoid external dependencies.
|
||||
Original source: https://pypi.org/project/HeapDict/
|
||||
"""
|
||||
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
try:
|
||||
from collections.abc import MutableMapping
|
||||
except ImportError:
|
||||
from collections import MutableMapping
|
||||
|
||||
KT = TypeVar("KT")
|
||||
VT = TypeVar("VT")
|
||||
|
||||
|
||||
def doc(s):
|
||||
if callable(s):
|
||||
s = s.__doc__
|
||||
|
||||
def f(g):
|
||||
g.__doc__ = s
|
||||
return g
|
||||
|
||||
return f
|
||||
|
||||
|
||||
class heapdict(MutableMapping, Generic[KT, VT]):
|
||||
__marker = object()
|
||||
|
||||
def __init__(self, *args, **kw):
|
||||
self.heap = []
|
||||
self.d = {}
|
||||
self.update(*args, **kw)
|
||||
|
||||
@doc(dict.clear)
|
||||
def clear(self):
|
||||
del self.heap[:]
|
||||
self.d.clear()
|
||||
|
||||
@doc(dict.__setitem__)
|
||||
def __setitem__(self, key, value):
|
||||
if key in self.d:
|
||||
self.pop(key)
|
||||
wrapper = [value, key, len(self)]
|
||||
self.d[key] = wrapper
|
||||
self.heap.append(wrapper)
|
||||
self._decrease_key(len(self.heap) - 1)
|
||||
|
||||
def _min_heapify(self, i):
|
||||
n = len(self.heap)
|
||||
h = self.heap
|
||||
while True:
|
||||
# calculate the offset of the left child
|
||||
l = (i << 1) + 1
|
||||
# calculate the offset of the right child
|
||||
r = (i + 1) << 1
|
||||
if l < n and h[l][0] < h[i][0]:
|
||||
low = l
|
||||
else:
|
||||
low = i
|
||||
if r < n and h[r][0] < h[low][0]:
|
||||
low = r
|
||||
|
||||
if low == i:
|
||||
break
|
||||
|
||||
self._swap(i, low)
|
||||
i = low
|
||||
|
||||
def _decrease_key(self, i):
|
||||
while i:
|
||||
# calculate the offset of the parent
|
||||
parent = (i - 1) >> 1
|
||||
if self.heap[parent][0] < self.heap[i][0]:
|
||||
break
|
||||
self._swap(i, parent)
|
||||
i = parent
|
||||
|
||||
def _swap(self, i, j):
|
||||
h = self.heap
|
||||
h[i], h[j] = h[j], h[i]
|
||||
h[i][2] = i
|
||||
h[j][2] = j
|
||||
|
||||
@doc(dict.__delitem__)
|
||||
def __delitem__(self, key):
|
||||
wrapper = self.d[key]
|
||||
while wrapper[2]:
|
||||
# calculate the offset of the parent
|
||||
parentpos = (wrapper[2] - 1) >> 1
|
||||
parent = self.heap[parentpos]
|
||||
self._swap(wrapper[2], parent[2])
|
||||
self.popitem()
|
||||
|
||||
@doc(dict.__getitem__)
|
||||
def __getitem__(self, key):
|
||||
return self.d[key][0]
|
||||
|
||||
@doc(dict.__iter__)
|
||||
def __iter__(self):
|
||||
return iter(self.d)
|
||||
|
||||
def popitem(self):
|
||||
"""D.popitem() -> (k, v), remove and return the (key, value) pair with lowest\nvalue; but raise KeyError if D is empty."""
|
||||
wrapper = self.heap[0]
|
||||
if len(self.heap) == 1:
|
||||
self.heap.pop()
|
||||
else:
|
||||
self.heap[0] = self.heap.pop()
|
||||
self.heap[0][2] = 0
|
||||
self._min_heapify(0)
|
||||
del self.d[wrapper[1]]
|
||||
return wrapper[1], wrapper[0]
|
||||
|
||||
@doc(dict.__len__)
|
||||
def __len__(self):
|
||||
return len(self.d)
|
||||
|
||||
def peekitem(self):
|
||||
"""D.peekitem() -> (k, v), return the (key, value) pair with lowest value;\n but raise KeyError if D is empty."""
|
||||
return (self.heap[0][1], self.heap[0][0])
|
||||
|
||||
|
||||
del doc
|
||||
__all__ = ["heapdict"]
|
||||
@@ -0,0 +1,139 @@
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import pyarrow
|
||||
import tensorflow as tf
|
||||
|
||||
from ray.data._internal.tensor_extensions.arrow import get_arrow_extension_tensor_types
|
||||
from ray.data.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.pandas_block import PandasBlockSchema
|
||||
|
||||
|
||||
def convert_ndarray_to_tf_tensor(
|
||||
ndarray: np.ndarray,
|
||||
dtype: Optional[tf.dtypes.DType] = None,
|
||||
type_spec: Optional[tf.TypeSpec] = None,
|
||||
) -> tf.Tensor:
|
||||
"""Convert a NumPy ndarray to a TensorFlow Tensor.
|
||||
|
||||
Args:
|
||||
ndarray: A NumPy ndarray that we wish to convert to a TensorFlow Tensor.
|
||||
dtype: A TensorFlow dtype for the created tensor; if None, the dtype will be
|
||||
inferred from the NumPy ndarray data.
|
||||
type_spec: A type spec that specifies the shape and dtype of the returned
|
||||
tensor. If you specify ``dtype``, the dtype stored in the type spec is
|
||||
ignored.
|
||||
|
||||
Returns:
|
||||
A TensorFlow Tensor.
|
||||
"""
|
||||
if dtype is None and type_spec is not None:
|
||||
dtype = type_spec.dtype
|
||||
|
||||
is_ragged = isinstance(type_spec, tf.RaggedTensorSpec)
|
||||
ndarray = _unwrap_ndarray_object_type_if_needed(ndarray)
|
||||
if is_ragged:
|
||||
return tf.ragged.constant(ndarray, dtype=dtype)
|
||||
else:
|
||||
return tf.convert_to_tensor(ndarray, dtype=dtype)
|
||||
|
||||
|
||||
def convert_ndarray_batch_to_tf_tensor_batch(
|
||||
ndarrays: Union[np.ndarray, Dict[str, np.ndarray]],
|
||||
dtypes: Optional[Union[tf.dtypes.DType, Dict[str, tf.dtypes.DType]]] = None,
|
||||
) -> Union[tf.Tensor, Dict[str, tf.Tensor]]:
|
||||
"""Convert a NumPy ndarray batch to a TensorFlow Tensor batch.
|
||||
|
||||
Args:
|
||||
ndarrays: A (dict of) NumPy ndarray(s) that we wish to convert to a TensorFlow
|
||||
Tensor.
|
||||
dtypes: A (dict of) TensorFlow dtype(s) for the created tensor; if None, the
|
||||
dtype will be inferred from the NumPy ndarray data.
|
||||
|
||||
Returns:
|
||||
A (dict of) TensorFlow Tensor(s).
|
||||
"""
|
||||
if isinstance(ndarrays, np.ndarray):
|
||||
# Single-tensor case.
|
||||
if isinstance(dtypes, dict):
|
||||
if len(dtypes) != 1:
|
||||
raise ValueError(
|
||||
"When constructing a single-tensor batch, only a single dtype "
|
||||
f"should be given, instead got: {dtypes}"
|
||||
)
|
||||
dtypes = next(iter(dtypes.values()))
|
||||
batch = convert_ndarray_to_tf_tensor(ndarrays, dtypes)
|
||||
else:
|
||||
# Multi-tensor case.
|
||||
batch = {
|
||||
col_name: convert_ndarray_to_tf_tensor(
|
||||
col_ndarray,
|
||||
dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
|
||||
)
|
||||
for col_name, col_ndarray in ndarrays.items()
|
||||
}
|
||||
|
||||
return batch
|
||||
|
||||
|
||||
def get_type_spec(
|
||||
schema: Union["pyarrow.lib.Schema", "PandasBlockSchema"],
|
||||
columns: Union[str, List[str]],
|
||||
) -> Union[tf.TypeSpec, Dict[str, tf.TypeSpec]]:
|
||||
import pyarrow as pa
|
||||
|
||||
from ray.data.extensions import TensorDtype
|
||||
|
||||
tensor_extension_types = get_arrow_extension_tensor_types()
|
||||
|
||||
assert not isinstance(schema, type)
|
||||
|
||||
dtypes: Dict[str, Union[np.dtype, pa.DataType]] = dict(
|
||||
zip(schema.names, schema.types)
|
||||
)
|
||||
|
||||
def get_dtype(dtype: Union[np.dtype, pa.DataType]) -> tf.dtypes.DType:
|
||||
if isinstance(dtype, pa.ListType):
|
||||
dtype = dtype.value_type
|
||||
if isinstance(dtype, pa.DataType):
|
||||
dtype = dtype.to_pandas_dtype()
|
||||
if isinstance(dtype, TensorDtype):
|
||||
dtype = dtype.element_dtype
|
||||
res = tf.dtypes.as_dtype(dtype)
|
||||
return res
|
||||
|
||||
def get_shape(dtype: Union[np.dtype, pa.DataType]) -> Tuple[int, ...]:
|
||||
shape = (None,)
|
||||
if isinstance(dtype, tensor_extension_types):
|
||||
dtype = dtype.to_pandas_dtype()
|
||||
if isinstance(dtype, pa.ListType):
|
||||
shape += (None,)
|
||||
elif isinstance(dtype, TensorDtype):
|
||||
shape += dtype.element_shape
|
||||
return shape
|
||||
|
||||
def get_tensor_spec(
|
||||
dtype: Union[np.dtype, pa.DataType], *, name: str
|
||||
) -> tf.TypeSpec:
|
||||
|
||||
shape, dtype = get_shape(dtype), get_dtype(dtype)
|
||||
# Batch dimension is always `None`. So, if there's more than one `None`-valued
|
||||
# dimension, then the tensor is ragged.
|
||||
is_ragged = sum(dim is None for dim in shape) > 1
|
||||
if is_ragged:
|
||||
type_spec = tf.RaggedTensorSpec(shape, dtype=dtype)
|
||||
else:
|
||||
type_spec = tf.TensorSpec(shape, dtype=dtype, name=name)
|
||||
return type_spec
|
||||
|
||||
if isinstance(columns, str):
|
||||
name, dtype = columns, dtypes[columns]
|
||||
return get_tensor_spec(dtype, name=name)
|
||||
|
||||
return {
|
||||
name: get_tensor_spec(dtype, name=name)
|
||||
for name, dtype in dtypes.items()
|
||||
if name in columns
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user