[data] Move core arrow util dependencies to arrow_utils.py (#51306)

- Moves the `parse_version` logic inside of the `get_pyarrow_version`
function to avoid redundancy.
- Moves the two depended-on utils to `_private/arrow_utils.py` to reduce
dependencies from data onto the giant `_private/utils.py` file.

Once "Ray Storage" is fully deprecated, we can move these utils into
data fully.

---------

Signed-off-by: Edward Oakes <[email protected]>

Author: edoakes

python/ray/_private/arrow_utils.py

@@ -0,0 +1,92 @@
+import json
+from typing import Dict, Optional
+from urllib.parse import urlencode, unquote, urlparse, parse_qsl, urlunparse
+
+from packaging.version import Version
+from packaging.version import parse as parse_version
+
+_PYARROW_INSTALLED: Optional[bool] = None
+_PYARROW_VERSION: Optional[Version] = None
+
+
+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

python/ray/_private/storage.py

@@ -4,9 +4,10 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional
 
-from ray._private.client_mode_hook import client_mode_hook
-from ray._private.utils import _add_creatable_buckets_param_if_s3_uri, load_class
 from ray._private.auto_init_hook import wrap_auto_init
+from ray._private.client_mode_hook import client_mode_hook
+from ray._private.arrow_utils import add_creatable_buckets_param_if_s3_uri
+from ray._private.utils import load_class
 
 if TYPE_CHECKING:
     import pyarrow.fs
@@ -452,7 +453,7 @@ def _init_filesystem(create_valid_file: bool = False, check_valid_file: bool = T
     else:
         # Arrow's S3FileSystem doesn't allow creating buckets by default, so we add a
         # query arg enabling bucket creation if an S3 URI is provided.
-        _storage_uri = _add_creatable_buckets_param_if_s3_uri(_storage_uri)
+        _storage_uri = add_creatable_buckets_param_if_s3_uri(_storage_uri)
         _filesystem, _storage_prefix = pyarrow.fs.FileSystem.from_uri(_storage_uri)
 
     if os.name == "nt":

python/ray/_private/utils.py

@@ -20,7 +20,6 @@
 import tempfile
 import threading
 import time
-from urllib.parse import urlencode, unquote, urlparse, parse_qsl, urlunparse
 import warnings
 from inspect import signature
 from pathlib import Path
@@ -71,7 +70,6 @@
 win32_AssignProcessToJobObject = None
 
 ENV_DISABLE_DOCKER_CPU_WARNING = "RAY_DISABLE_DOCKER_CPU_WARNING" in os.environ
-_PYARROW_VERSION = None
 
 # This global variable is used for testing only
 _CALLED_FREQ = defaultdict(lambda: 0)
@@ -1735,90 +1733,6 @@ def get_entrypoint_name():
         return "unknown"
 
 
-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.
-    """
-    from packaging.version import parse as parse_version
-
-    pyarrow_version = _get_pyarrow_version()
-    if pyarrow_version is not None:
-        pyarrow_version = parse_version(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
-
-
-def _get_pyarrow_version() -> Optional[str]:
-    """Get the version of the installed pyarrow package, returned as a tuple of ints.
-    Returns None if the package is not found.
-    """
-    global _PYARROW_VERSION
-    if _PYARROW_VERSION is None:
-        try:
-            import pyarrow
-        except ModuleNotFoundError:
-            # pyarrow not installed, short-circuit.
-            pass
-        else:
-            if hasattr(pyarrow, "__version__"):
-                _PYARROW_VERSION = pyarrow.__version__
-    return _PYARROW_VERSION
-
-
 class DeferSigint(contextlib.AbstractContextManager):
     """Context manager that defers SIGINT signals until the context is left."""
 

python/ray/air/tests/test_arrow.py

@@ -6,7 +6,7 @@
 import pytest
 from packaging.version import parse as parse_version
 
-from ray._private.utils import _get_pyarrow_version
+from ray._private.arrow_utils import get_pyarrow_version
 from ray.air.util.tensor_extensions.arrow import (
     ArrowConversionError,
     _convert_to_pyarrow_native_array,
@@ -51,9 +51,9 @@ def test_arrow_native_list_conversion(input, disable_fallback_to_object_extensio
     upon serialization into Arrow format (and are NOT converted to numpy
     tensor using extension)"""
 
-    if isinstance(input[0], pa.Scalar) and parse_version(
-        _get_pyarrow_version()
-    ) <= parse_version("13.0.0"):
+    if isinstance(input[0], pa.Scalar) and get_pyarrow_version() <= parse_version(
+        "13.0.0"
+    ):
         pytest.skip(
             "Pyarrow < 13.0 not able to properly infer native types from its own Scalars"
         )

python/ray/air/tests/test_tensor_extension.py

@@ -6,7 +6,7 @@
 import pytest
 from packaging.version import parse as parse_version
 
-from ray._private.utils import _get_pyarrow_version
+from ray._private.arrow_utils import get_pyarrow_version
 from ray.air.util.tensor_extensions.arrow import (
     ArrowConversionError,
     ArrowTensorArray,
@@ -517,7 +517,7 @@ def test_arrow_tensor_array_getitem(chunked, restore_data_context, tensor_format
     if chunked:
         t_arr = pa.chunked_array(t_arr)
 
-    pyarrow_version = parse_version(_get_pyarrow_version())
+    pyarrow_version = get_pyarrow_version()
     if (
         chunked
         and pyarrow_version >= parse_version("8.0.0")
@@ -589,7 +589,7 @@ def test_arrow_variable_shaped_tensor_array_getitem(
     if chunked:
         t_arr = pa.chunked_array(t_arr)
 
-    pyarrow_version = parse_version(_get_pyarrow_version())
+    pyarrow_version = get_pyarrow_version()
     if (
         chunked
         and pyarrow_version >= parse_version("8.0.0")

python/ray/air/util/object_extensions/arrow.py

@@ -7,13 +7,12 @@
 
 import ray.air.util.object_extensions.pandas
 from ray._private.serialization import pickle_dumps
-from ray._private.utils import _get_pyarrow_version
+from ray._private.arrow_utils import get_pyarrow_version
 from ray.util.annotations import PublicAPI
 
 MIN_PYARROW_VERSION_SCALAR_SUBCLASS = parse_version("9.0.0")
 
-_VER = _get_pyarrow_version()
-PYARROW_VERSION = None if _VER is None else parse_version(_VER)
+PYARROW_VERSION = get_pyarrow_version()
 
 
 def _object_extension_type_allowed() -> bool:

python/ray/air/util/tensor_extensions/arrow.py

@@ -11,7 +11,7 @@
 import pyarrow as pa
 from packaging.version import parse as parse_version
 
-from ray._private.utils import _get_pyarrow_version
+from ray._private.arrow_utils import get_pyarrow_version
 from ray.air.util.tensor_extensions.utils import (
     _is_ndarray_variable_shaped_tensor,
     create_ragged_ndarray,
@@ -26,11 +26,7 @@
 from ray.util import log_once
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
-
-PYARROW_VERSION = _get_pyarrow_version()
-if PYARROW_VERSION is not None:
-    PYARROW_VERSION = parse_version(PYARROW_VERSION)
-
+PYARROW_VERSION = get_pyarrow_version()
 # Minimum version of Arrow that supports ExtensionScalars.
 # TODO(Clark): Remove conditional definition once we only support Arrow 8.0.0+.
 MIN_PYARROW_VERSION_SCALAR = parse_version("8.0.0")

python/ray/data/__init__.py

@@ -3,7 +3,7 @@
 import pandas  # noqa
 from packaging.version import parse as parse_version
 
-from ray._private.utils import _get_pyarrow_version
+from ray._private.arrow_utils import get_pyarrow_version
 from ray.data._internal.compute import ActorPoolStrategy
 from ray.data._internal.datasource.tfrecords_datasource import TFXReadOptions
 from ray.data._internal.execution.interfaces import (
@@ -80,8 +80,8 @@
     # disabled it's deserialization by default. To ensure that users can load data
     # written with earlier version of Ray Data, we enable auto-loading of serialized
     # tensor extensions.
-    pyarrow_version = _get_pyarrow_version()
-    if not isinstance(pyarrow_version, str):
+    pyarrow_version = get_pyarrow_version()
+    if pyarrow_version is None:
         # PyArrow is mocked in documentation builds. In this case, we don't need to do
         # anything.
         pass
@@ -93,7 +93,7 @@
         )
 
         if (
-            parse_version(pyarrow_version) >= parse_version("14.0.1")
+            pyarrow_version >= parse_version("14.0.1")
             and RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE
         ):
             pa.PyExtensionType.set_auto_load(True)

python/ray/data/_internal/arrow_block.py

@@ -15,7 +15,7 @@
 
 import numpy as np
 
-from ray._private.utils import _get_pyarrow_version
+from ray._private.arrow_utils import get_pyarrow_version
 from ray.air.constants import TENSOR_COLUMN_NAME
 from ray.air.util.tensor_extensions.arrow import (
     convert_to_pyarrow_array,
@@ -189,9 +189,7 @@ def _build_tensor_row(
         element = row[col_name][0]
         # TODO(Clark): Reduce this to np.asarray(element) once we only support Arrow
         # 9.0.0+.
-        pyarrow_version = _get_pyarrow_version()
-        if pyarrow_version is not None:
-            pyarrow_version = parse_version(pyarrow_version)
+        pyarrow_version = get_pyarrow_version()
         if pyarrow_version is None or pyarrow_version >= parse_version("8.0.0"):
             assert isinstance(element, pyarrow.ExtensionScalar)
             if pyarrow_version is None or pyarrow_version >= parse_version("9.0.0"):

python/ray/data/_internal/arrow_ops/transform_pyarrow.py

@@ -5,7 +5,7 @@
 from packaging.version import parse as parse_version
 
 from ray._private.ray_constants import env_integer
-from ray._private.utils import _get_pyarrow_version
+from ray._private.arrow_utils import get_pyarrow_version
 from ray.air.util.tensor_extensions.arrow import (
     INT32_OVERFLOW_THRESHOLD,
     MIN_PYARROW_VERSION_CHUNKED_ARRAY_TO_NUMPY_ZERO_COPY_ONLY,
@@ -206,7 +206,7 @@ def unify_schemas(
         schemas_to_unify = schemas
 
     try:
-        if parse_version(_get_pyarrow_version()) < MIN_PYARROW_VERSION_TYPE_PROMOTION:
+        if get_pyarrow_version() < MIN_PYARROW_VERSION_TYPE_PROMOTION:
             return pyarrow.unify_schemas(schemas_to_unify)
 
         # NOTE: By default type promotion (from "smaller" to "larger" types) is disabled,
@@ -554,7 +554,7 @@ def concat(
         # to vary b/w blocks
         #
         # NOTE: Type promotions aren't available in Arrow < 14.0
-        if parse_version(_get_pyarrow_version()) < parse_version("14.0.0"):
+        if get_pyarrow_version() < parse_version("14.0.0"):
             table = pyarrow.concat_tables(blocks, promote=True)
         else:
             arrow_promote_types_mode = "permissive" if promote_types else "default"