Major refactoring:

- Move distributed methods out of non-distributed modules
- Refactor dispatching
- Refactor structure of distributed modules
- Add classes for execution engine and memory format

Author: jaidisido

awswrangler/__init__.py

@@ -31,12 +31,10 @@
     timestream,
 )
 from awswrangler.__metadata__ import __description__, __license__, __title__, __version__  # noqa
-from awswrangler._config import ExecutionEngine, config  # noqa
-from awswrangler.distributed import initialize_ray
-
-if config.execution_engine == ExecutionEngine.RAY.value:
-    initialize_ray()
+from awswrangler._config import config  # noqa
+from awswrangler._distributed import engine, memory_format  # noqa
 
+engine.initialize()
 
 __all__ = [
     "athena",
@@ -60,6 +58,8 @@
     "secretsmanager",
     "sqlserver",
     "config",
+    "engine",
+    "memory_format",
     "timestream",
     "__description__",
     "__license__",

awswrangler/_config.py

@@ -1,10 +1,8 @@
 """Configuration file for AWS SDK for pandas."""
 
-import importlib.util
 import inspect
 import logging
 import os
-from enum import Enum
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, Type, Union, cast
 
 import botocore.config
@@ -18,20 +16,6 @@
 _ConfigValueType = Union[str, bool, int, botocore.config.Config, None]
 
 
-class ExecutionEngine(Enum):
-    """Execution engine enum."""
-
-    RAY = "ray"
-    PYTHON = "python"
-
-
-class MemoryFormat(Enum):
-    """Memory format enum."""
-
-    MODIN = "modin"
-    PANDAS = "pandas"
-
-
 class _ConfigArg(NamedTuple):
     dtype: Type[Union[str, bool, int, botocore.config.Config]]
     nullable: bool
@@ -69,12 +53,6 @@ class _ConfigArg(NamedTuple):
     "botocore_config": _ConfigArg(dtype=botocore.config.Config, nullable=True),
     "verify": _ConfigArg(dtype=str, nullable=True, loaded=True),
     # Distributed
-    "execution_engine": _ConfigArg(
-        dtype=str, nullable=False, loaded=True, default="ray" if importlib.util.find_spec("ray") else "python"
-    ),
-    "memory_format": _ConfigArg(
-        dtype=str, nullable=False, loaded=True, default="modin" if importlib.util.find_spec("modin") else "pandas"
-    ),
     "address": _ConfigArg(dtype=str, nullable=True),
     "redis_password": _ConfigArg(dtype=str, nullable=True),
     "ignore_reinit_error": _ConfigArg(dtype=bool, nullable=True),
@@ -440,24 +418,6 @@ def verify(self) -> Optional[str]:
     def verify(self, value: Optional[str]) -> None:
         self._set_config_value(key="verify", value=value)
 
-    @property
-    def execution_engine(self) -> str:
-        """Property execution_engine."""
-        return cast(str, self["execution_engine"])
-
-    @execution_engine.setter
-    def execution_engine(self, value: str) -> None:
-        self._set_config_value(key="execution_engine", value=value)
-
-    @property
-    def memory_format(self) -> str:
-        """Property memory_format."""
-        return cast(str, self["memory_format"])
-
-    @memory_format.setter
-    def memory_format(self, value: str) -> None:
-        self._set_config_value(key="memory_format", value=value)
-
     @property
     def address(self) -> Optional[str]:
         """Property address."""

awswrangler/_dispatch.py

@@ -0,0 +1,34 @@
+"""Dispatch engine and memory format configuration (PRIVATE)."""
+from typing import Any, Callable, Dict, Optional
+
+from awswrangler._distributed import engine
+
+
+def dispatch_on_engine(func: Callable[..., Any]) -> Callable[..., Any]:
+    """Dispatch on engine function decorator.
+
+    Transforms a function into a dispatch function,
+    which can have different behaviors based on the value of the distribution engine.
+    """
+    registry: Dict[str, Callable[..., Any]] = {}
+
+    def dispatch(value: str) -> Callable[..., Any]:
+        try:
+            return registry[value]
+        except KeyError:
+            return func
+
+    def register(value: str, func: Optional[Callable[..., Any]] = None) -> Callable[..., Any]:
+        if func is None:
+            return lambda f: register(value, f)
+        registry[value] = func
+        return func
+
+    def wrapper(*args: Any, **kw: Dict[str, Any]) -> Any:
+        return dispatch(engine.get())(*args, **kw)
+
+    wrapper.register = register  # type: ignore
+    wrapper.dispatch = dispatch  # type: ignore
+    wrapper.registry = registry  # type: ignore
+
+    return wrapper

awswrangler/_distributed.py

@@ -0,0 +1,128 @@
+"""Distributed engine and memory format configuration."""
+
+# pylint: disable=import-outside-toplevel
+
+import importlib.util
+from enum import Enum
+from typing import Optional
+
+
+class EngineEnum(Enum):
+    """Execution engine enum."""
+
+    RAY = "ray"
+    PYTHON = "python"
+
+
+class MemoryFormatEnum(Enum):
+    """Memory format enum."""
+
+    MODIN = "modin"
+    PANDAS = "pandas"
+
+
+class Engine:
+    """Execution engine configuration class."""
+
+    _name: Optional[str] = None
+
+    @classmethod
+    def get_installed(cls) -> str:
+        """Get the installed distribution engine.
+
+        This is the engine that can be imported.
+
+        Returns
+        -------
+        str
+            The distribution engine installed.
+        """
+        if importlib.util.find_spec("ray"):
+            return EngineEnum.RAY.value
+        return EngineEnum.PYTHON.value
+
+    @classmethod
+    def get(cls) -> str:
+        """Get the configured distribution engine.
+
+        This is the engine currently configured. If None, the installed engine is returned.
+
+        Returns
+        -------
+        str
+            The distribution engine configured.
+        """
+        return cls._name if cls._name else cls.get_installed()
+
+    @classmethod
+    def set(cls, name: str) -> None:
+        """Set the distribution engine."""
+        cls._name = name
+
+    @classmethod
+    def register(cls, name: Optional[str] = None) -> None:
+        """Register the distribution engine dispatch methods."""
+        engine_name = cls.get_installed() if not name else name
+
+        if engine_name == EngineEnum.RAY.value:
+            from awswrangler.distributed.ray._register import register_ray
+
+            register_ray()
+
+    @classmethod
+    def initialize(cls, name: Optional[str] = None) -> None:
+        """Initialize the distribution engine."""
+        engine_name = cls.get_installed() if not name else name
+
+        if engine_name == EngineEnum.RAY.value:
+            cls._name = EngineEnum.RAY.value
+
+            from awswrangler.distributed.ray import initialize_ray
+
+            initialize_ray()
+            cls.register(cls._name)
+        else:
+            cls._name = EngineEnum.PYTHON.value
+
+
+class MemoryFormat:
+    """Memory format configuration class."""
+
+    _name: Optional[str] = None
+
+    @classmethod
+    def get_installed(cls) -> str:
+        """Get the installed memory format.
+
+        This is the format that can be imported.
+
+        Returns
+        -------
+        str
+            The memory format installed.
+        """
+        if importlib.util.find_spec("modin"):
+            return MemoryFormatEnum.MODIN.value
+        return MemoryFormatEnum.PANDAS.value
+
+    @classmethod
+    def get(cls) -> str:
+        """Get the configured memory format.
+
+        This is the memory format currently configured. If None, the installed memory format is returned.
+
+        Returns
+        -------
+        str
+            The memory format configured.
+        """
+        return cls._name if cls._name else cls.get_installed()
+
+    @classmethod
+    def set(cls, name: str) -> None:
+        """Set the memory format."""
+        cls._name = name
+
+
+engine: Engine = Engine()
+memory_format: MemoryFormat = MemoryFormat()

awswrangler/_threading.py

@@ -3,27 +3,16 @@
 import concurrent.futures
 import itertools
 import logging
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
+from typing import Any, Callable, List, Optional, Union
 
 import boto3
 
 from awswrangler import _utils
-from awswrangler._config import ExecutionEngine, config
-
-if config.execution_engine == ExecutionEngine.RAY.value or TYPE_CHECKING:
-    from awswrangler.distributed.ray._pool import _RayPoolExecutor
+from awswrangler._distributed import EngineEnum, engine
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
-def _get_executor(use_threads: Union[bool, int]) -> Union["_ThreadPoolExecutor", "_RayPoolExecutor"]:
-    return (
-        _RayPoolExecutor()
-        if config.execution_engine == ExecutionEngine.RAY.value
-        else _ThreadPoolExecutor(use_threads)  # type: ignore
-    )
-
-
 class _ThreadPoolExecutor:
     def __init__(self, use_threads: Union[bool, int]):
         super().__init__()
@@ -42,3 +31,11 @@ def map(self, func: Callable[..., Any], boto3_session: boto3.Session, *iterables
             return list(self._exec.map(func, *args))
         # Single-threaded
         return list(map(func, *(itertools.repeat(boto3_session), *iterables)))  # type: ignore
+
+
+def _get_executor(use_threads: Union[bool, int]) -> _ThreadPoolExecutor:
+    if engine.get() == EngineEnum.RAY.value:
+        from awswrangler.distributed.ray._pool import _RayPoolExecutor  # pylint: disable=import-outside-toplevel
+
+        return _RayPoolExecutor()  # type: ignore
+    return _ThreadPoolExecutor(use_threads)

awswrangler/_utils.py

@@ -8,7 +8,7 @@
 import random
 import time
 from concurrent.futures import FIRST_COMPLETED, Future, wait
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Dict, Generator, List, Optional, Sequence, Tuple, Union, cast
 
 import boto3
 import botocore.config
@@ -19,13 +19,8 @@
 from awswrangler import _config, exceptions
 from awswrangler.__metadata__ import __version__
 from awswrangler._arrow import _table_to_df
-from awswrangler._config import ExecutionEngine, MemoryFormat, apply_configs, config
-
-if config.execution_engine == ExecutionEngine.RAY.value or TYPE_CHECKING:
-    import ray  # pylint: disable=unused-import
-
-    if config.memory_format == MemoryFormat.MODIN.value:
-        from awswrangler.distributed.ray._utils import _arrow_refs_to_df  # pylint: disable=ungrouped-imports
+from awswrangler._config import apply_configs
+from awswrangler._dispatch import dispatch_on_engine
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -413,13 +408,10 @@ def check_schema_changes(columns_types: Dict[str, str], table_input: Optional[Di
                 )
 
 
-def table_refs_to_df(
-    tables: Union[List[pa.Table], List["ray.ObjectRef"]], kwargs: Dict[str, Any]  # type: ignore
-) -> pd.DataFrame:
+@dispatch_on_engine
+def table_refs_to_df(tables: List[pa.Table], kwargs: Dict[str, Any]) -> pd.DataFrame:  # type: ignore
     """Build Pandas dataframe from list of PyArrow tables."""
-    if isinstance(tables[0], pa.Table):
-        return _table_to_df(pa.concat_tables(tables, promote=True), kwargs=kwargs)
-    return _arrow_refs_to_df(arrow_refs=tables, kwargs=kwargs)  # type: ignore
+    return _table_to_df(pa.concat_tables(tables, promote=True), kwargs=kwargs)
 
 
 def list_to_arrow_table(

awswrangler/distributed/__init__.py

@@ -1,17 +1 @@
 """Distributed Module."""
-
-from awswrangler.distributed._distributed import (  # noqa
-    RayLogger,
-    initialize_ray,
-    modin_repartition,
-    ray_get,
-    ray_remote,
-)
-
-__all__ = [
-    "RayLogger",
-    "initialize_ray",
-    "modin_repartition",
-    "ray_get",
-    "ray_remote",
-]

awswrangler/distributed/ray/__init__.py

@@ -1 +1,10 @@
 """Ray Module."""
+
+from awswrangler.distributed.ray._core import RayLogger, initialize_ray, ray_get, ray_remote  # noqa
+
+__all__ = [
+    "RayLogger",
+    "initialize_ray",
+    "ray_get",
+    "ray_remote",
+]