Refactor RayPoolExecutor to return futures and use ray.get in S3 Select

Author: jaidisido

awswrangler/_threading.py

@@ -17,8 +17,7 @@
 
 
 def _get_executor(use_threads: Union[bool, int]) -> Union["_ThreadPoolExecutor", "_RayPoolExecutor"]:
-    executor = _RayPoolExecutor if config.distributed else _ThreadPoolExecutor
-    return executor(use_threads)  # type: ignore
+    return _RayPoolExecutor() if config.distributed else _ThreadPoolExecutor(use_threads)  # type: ignore
 
 
 class _ThreadPoolExecutor:

awswrangler/_utils.py

@@ -22,7 +22,6 @@
 
 if TYPE_CHECKING or config.distributed:
     import ray  # pylint: disable=unused-import
-
     from awswrangler.distributed._utils import _arrow_refs_to_df  # pylint: disable=ungrouped-imports
 
 _logger: logging.Logger = logging.getLogger(__name__)

awswrangler/distributed/__init__.py

@@ -1,8 +1,9 @@
 """Distributed Module."""
 
-from awswrangler.distributed._distributed import initialize_ray, ray_remote  # noqa
+from awswrangler.distributed._distributed import initialize_ray, ray_get, ray_remote  # noqa
 
 __all__ = [
-    "ray_remote",
     "initialize_ray",
+    "ray_get",
+    "ray_remote",
 ]

awswrangler/distributed/_distributed.py

@@ -4,7 +4,7 @@
 import os
 import sys
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, List, Optional
 
 import psutil
 
@@ -14,6 +14,24 @@
     import ray  # pylint: disable=import-error
 
 
+def ray_get(futures: List[Any]) -> List[Any]:
+    """
+    Run ray.get on futures if distributed.
+
+    Parameters
+    ----------
+    futures : List[Any]
+        List of Ray futures
+
+    Returns
+    -------
+    List[Any]
+    """
+    if config.distributed:
+        return ray.get(futures)
+    return futures
+
+
 def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]:
     """
     Decorate callable to wrap within ray.remote.

awswrangler/distributed/_pool.py

@@ -2,23 +2,19 @@
 
 import itertools
 import logging
-from typing import Any, Callable, List, Optional, Union
+from typing import Any, Callable, List
 
 import boto3
-from ray.util.multiprocessing import Pool
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
 class _RayPoolExecutor:
-    def __init__(self, processes: Optional[Union[bool, int]] = None):
-        self._exec: Pool = Pool(processes=None if isinstance(processes, bool) else processes)
+    def __init__(self) -> None:
+        pass
 
     def map(self, func: Callable[..., List[str]], _: boto3.Session, *args: Any) -> List[Any]:
-        """Map function and its args to Ray pool."""
-        futures = []
+        """Map func and return ray futures."""
         _logger.debug("Ray map: %s", func)
-        # Discard boto3.Session object & call the fn asynchronously
-        for arg in zip(itertools.repeat(None), *args):
-            futures.append(self._exec.apply_async(func, arg))
-        return [f.get() for f in futures]
+        # Discard boto3.Session object & return futures
+        return list(func(*arg) for arg in zip(itertools.repeat(None), *args))

awswrangler/s3/_select.py

@@ -13,7 +13,7 @@
 
 from awswrangler import _data_types, _utils, exceptions
 from awswrangler._threading import _get_executor
-from awswrangler.distributed import ray_remote
+from awswrangler.distributed import ray_get, ray_remote
 from awswrangler.s3._describe import size_objects
 from awswrangler.s3._list import _path2list
 from awswrangler.s3._read import _get_path_ignore_suffix
@@ -71,6 +71,7 @@ def _select_object_content(
     return _utils.list_to_arrow_table(mapping=payload_records)
 
 
+@ray_remote
 def _select_query(
     path: str,
     executor: Any,
@@ -272,5 +273,5 @@ def select_query(
 
     arrow_kwargs = _data_types.pyarrow2pandas_defaults(use_threads=use_threads, kwargs=arrow_additional_kwargs)
     executor = _get_executor(use_threads=use_threads)
-    tables = _flatten_list([_select_query(path=path, executor=executor, **select_kwargs) for path in paths])
+    tables = _flatten_list(ray_get([_select_query(path=path, executor=executor, **select_kwargs) for path in paths]))
     return _utils.table_refs_to_df(tables=tables, kwargs=arrow_kwargs)