PR feedback and distributed module test

Author: jaidisido

awswrangler/_distributed.py

@@ -4,18 +4,20 @@
 
 import importlib.util
 from collections import defaultdict
-from enum import Enum
+from enum import Enum, unique
 from functools import wraps
 from typing import Any, Callable, Dict, Optional
 
 
+@unique
 class EngineEnum(Enum):
     """Execution engine enum."""
 
     RAY = "ray"
     PYTHON = "python"
 
 
+@unique
 class MemoryFormatEnum(Enum):
     """Memory format enum."""
 
@@ -26,26 +28,26 @@ class MemoryFormatEnum(Enum):
 class Engine:
     """Execution engine configuration class."""
 
-    _name: Optional[str] = None
+    _enum: Optional[Enum] = None
     _registry: Dict[str, Dict[str, Callable[..., Any]]] = defaultdict(dict)
 
     @classmethod
-    def get_installed(cls) -> str:
+    def get_installed(cls) -> Enum:
         """Get the installed distribution engine.
 
         This is the engine that can be imported.
 
         Returns
         -------
-        str
+        EngineEnum
             The distribution engine installed.
         """
         if importlib.util.find_spec("ray"):
-            return EngineEnum.RAY.value
-        return EngineEnum.PYTHON.value
+            return EngineEnum.RAY
+        return EngineEnum.PYTHON
 
     @classmethod
-    def get(cls) -> str:
+    def get(cls) -> Enum:
         """Get the configured distribution engine.
 
         This is the engine currently configured. If None, the installed engine is returned.
@@ -55,25 +57,25 @@ def get(cls) -> str:
         str
             The distribution engine configured.
         """
-        return cls._name if cls._name else cls.get_installed()
+        return cls._enum if cls._enum else cls.get_installed()
 
     @classmethod
     def set(cls, name: str) -> None:
         """Set the distribution engine."""
-        cls._name = name
+        cls._enum = EngineEnum._member_map_[name.upper()]  # pylint: disable=protected-access,no-member
 
     @classmethod
     def dispatch_func(cls, source_func: Callable[..., Any], value: Optional[Any] = None) -> Callable[..., Any]:
         """Dispatch a func based on value or the distribution engine and the source function."""
         try:
-            return cls._registry[value or cls.get()][source_func.__name__]
+            return cls._registry[value or cls.get().value][source_func.__name__]
         except KeyError:
-            return source_func
+            return getattr(source_func, "_source_func", source_func)
 
     @classmethod
     def register_func(cls, source_func: Callable[..., Any], destination_func: Callable[..., Any]) -> Callable[..., Any]:
         """Register a func based on the distribution engine and source function."""
-        cls._registry[cls.get()][source_func.__name__] = destination_func
+        cls._registry[cls.get().value][source_func.__name__] = destination_func
         return destination_func
 
     @classmethod
@@ -85,13 +87,13 @@ def wrapper(*args: Any, **kw: Dict[str, Any]) -> Any:
             return cls.dispatch_func(func)(*args, **kw)
 
         # Save the original function
-        wrapper._source_func = func  # type: ignore  # pylint: pylint: disable=protected-access
+        wrapper._source_func = func  # type: ignore  # pylint: disable=protected-access
         return wrapper
 
     @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
+        engine_name = cls.get_installed().value if not name else name
         cls.set(engine_name)
         cls._registry.clear()
 
@@ -103,7 +105,7 @@ def register(cls, name: Optional[str] = None) -> None:
     @classmethod
     def initialize(cls, name: Optional[str] = None) -> None:
         """Initialize the distribution engine."""
-        engine_name = cls.get_installed() if not name else name
+        engine_name = cls.get_installed().value if not name else name
         if engine_name == EngineEnum.RAY.value:
             from awswrangler.distributed.ray import initialize_ray
 
@@ -114,40 +116,40 @@ def initialize(cls, name: Optional[str] = None) -> None:
 class MemoryFormat:
     """Memory format configuration class."""
 
-    _name: Optional[str] = None
+    _enum: Optional[Enum] = None
 
     @classmethod
-    def get_installed(cls) -> str:
+    def get_installed(cls) -> Enum:
         """Get the installed memory format.
 
         This is the format that can be imported.
 
         Returns
         -------
-        str
+        Enum
             The memory format installed.
         """
         if importlib.util.find_spec("modin"):
-            return MemoryFormatEnum.MODIN.value
-        return MemoryFormatEnum.PANDAS.value
+            return MemoryFormatEnum.MODIN
+        return MemoryFormatEnum.PANDAS
 
     @classmethod
-    def get(cls) -> str:
+    def get(cls) -> Enum:
         """Get the configured memory format.
 
         This is the memory format currently configured. If None, the installed memory format is returned.
 
         Returns
         -------
-        str
+        Enum
             The memory format configured.
         """
-        return cls._name if cls._name else cls.get_installed()
+        return cls._enum if cls._enum else cls.get_installed()
 
     @classmethod
     def set(cls, name: str) -> None:
         """Set the memory format."""
-        cls._name = name
+        cls._enum = MemoryFormatEnum._member_map_[name.upper()]  # pylint: disable=protected-access,no-member
 
 
 engine: Engine = Engine()

awswrangler/_threading.py

@@ -34,7 +34,7 @@ def map(self, func: Callable[..., Any], boto3_session: boto3.Session, *iterables
 
 
 def _get_executor(use_threads: Union[bool, int]) -> _ThreadPoolExecutor:
-    if engine.get() == EngineEnum.RAY.value:
+    if engine.get() == EngineEnum.RAY:
         from awswrangler.distributed.ray._pool import _RayPoolExecutor  # pylint: disable=import-outside-toplevel
 
         return _RayPoolExecutor()  # type: ignore

awswrangler/distributed/ray/_core.py

@@ -7,7 +7,7 @@
 from awswrangler._config import apply_configs
 from awswrangler._distributed import EngineEnum, engine
 
-if engine.get() == EngineEnum.RAY.value or TYPE_CHECKING:
+if engine.get() == EngineEnum.RAY or TYPE_CHECKING:
     import ray
 
 _logger: logging.Logger = logging.getLogger(__name__)
@@ -26,7 +26,7 @@ def __init__(
 
     def get_logger(self, name: Union[str, Any] = None) -> Union[logging.Logger, Any]:
         """Return logger object."""
-        return logging.getLogger(name) if engine.get() == EngineEnum.RAY.value else None
+        return logging.getLogger(name) if engine.get() == EngineEnum.RAY else None
 
 
 def ray_get(futures: List[Any]) -> List[Any]:
@@ -42,7 +42,7 @@ def ray_get(futures: List[Any]) -> List[Any]:
     -------
     List[Any]
     """
-    if engine.get() == EngineEnum.RAY.value:
+    if engine.get() == EngineEnum.RAY:
         return ray.get(futures)
     return futures
 

awswrangler/distributed/ray/_register.py

@@ -26,7 +26,7 @@ def register_ray() -> None:
     engine.register_func(_select_object_content, ray_remote(_select_object_content))
     engine.register_func(_wait_object_batch, ray_remote(_wait_object_batch))
 
-    if memory_format.get() == MemoryFormatEnum.MODIN.value:
+    if memory_format.get() == MemoryFormatEnum.MODIN:
         from awswrangler.distributed.ray.modin._core import modin_repartition
         from awswrangler.distributed.ray.modin._utils import _arrow_refs_to_df
         from awswrangler.distributed.ray.modin.s3._read_parquet import _read_parquet_distributed

awswrangler/distributed/ray/modin/_core.py

@@ -24,6 +24,7 @@ def modin_repartition(function: Callable[..., Any]) -> Callable[..., Any]:
     -------
     Callable[..., Any]
     """
+    # Access the source function if it exists
     function = getattr(function, "_source_func", function)
 
     @wraps(function)

awswrangler/distributed/ray/modin/s3/_write_dataset.py

@@ -4,6 +4,7 @@
 import boto3
 import modin.pandas as pd
 import ray
+from pandas import DataFrame as PandasDataFrame
 
 from awswrangler._distributed import engine
 from awswrangler.distributed.ray import ray_get, ray_remote
@@ -123,8 +124,8 @@ def _to_partitions_distributed(  # pylint: disable=unused-argument
     if not bucketing_info:
         # If only partitioning (without bucketing), avoid expensive modin groupby
         # by partitioning and writing each block as an ordinary Pandas DataFrame
-        _to_partitions_func = getattr(_to_partitions, "_source_func", _to_partitions)
-        func = getattr(func, "_source_func", func)
+        _to_partitions_func = engine.dispatch_func(_to_partitions, PandasDataFrame)
+        func = engine.dispatch_func(func, PandasDataFrame)
 
         @ray_remote
         def write_partitions(df: pd.DataFrame) -> Tuple[List[str], Dict[str, List[str]]]:

awswrangler/oracle.py

@@ -434,7 +434,7 @@ def detect_oracle_decimal_datatype(cursor: Any) -> Dict[str, pa.DataType]:
     if isinstance(cursor, oracledb.Cursor):
         # Oracle stores DECIMAL as the NUMBER type
         for row in cursor.description:
-            if row[1] == oracledb.DB_TYPE_NUMBER and row[5] > 0:
+            if row[1] == oracledb.DB_TYPE_NUMBER and row[5] > 0:  # pylint: disable=no-member
                 dtype[row[0]] = pa.decimal128(row[4], row[5])
 
     _logger.debug("decimal dtypes: %s", dtype)

awswrangler/s3/_write.py

@@ -1,6 +1,7 @@
 """Amazon CSV S3 Write Module (PRIVATE)."""
 
 import logging
+from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple
 
 import pandas as pd
@@ -56,14 +57,14 @@ def _validate_args(
     description: Optional[str],
     parameters: Optional[Dict[str, str]],
     columns_comments: Optional[Dict[str, str]],
-    execution_engine: str,
+    execution_engine: Enum,
 ) -> None:
     if df.empty is True:
         raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
     if dataset is False:
         if path is None:
             raise exceptions.InvalidArgumentValue("If dataset is False, the `path` argument must be passed.")
-        if execution_engine == EngineEnum.PYTHON.value and path.endswith("/"):
+        if execution_engine == EngineEnum.PYTHON and path.endswith("/"):
             raise exceptions.InvalidArgumentValue(
                 "If <dataset=False>, the argument <path> should be a key, not a prefix."
             )

awswrangler/s3/_write_parquet.py

@@ -166,7 +166,7 @@ def _to_parquet(
     use_threads: Union[bool, int],
     path: Optional[str] = None,
     path_root: Optional[str] = None,
-    filename_prefix: Optional[str] = uuid.uuid4().hex,
+    filename_prefix: Optional[str] = None,
     max_rows_by_file: Optional[int] = 0,
 ) -> List[str]:
     file_path = _get_file_path(

awswrangler/s3/_write_text.py

@@ -38,7 +38,7 @@ def _to_text(  # pylint: disable=unused-argument
     s3_additional_kwargs: Optional[Dict[str, str]],
     path: Optional[str] = None,
     path_root: Optional[str] = None,
-    filename_prefix: Optional[str] = uuid.uuid4().hex,
+    filename_prefix: Optional[str] = None,
     bucketing: bool = False,
     **pandas_kwargs: Any,
 ) -> List[str]:

tests/_utils.py

@@ -12,7 +12,7 @@
 from awswrangler._distributed import EngineEnum, MemoryFormatEnum
 from awswrangler._utils import try_it
 
-if wr.engine.get() == EngineEnum.RAY.value and wr.memory_format.get() == MemoryFormatEnum.MODIN.value:
+if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
     import modin.pandas as pd
 else:
     import pandas as pd

tests/unit/test_distributed.py

@@ -0,0 +1,54 @@
+import importlib.util
+import logging
+from enum import Enum
+
+import pytest
+
+import awswrangler as wr
+from awswrangler._distributed import EngineEnum, MemoryFormatEnum
+from awswrangler.s3._write_parquet import _to_parquet
+
+logging.getLogger("awswrangler").setLevel(logging.DEBUG)
+
+pytestmark = pytest.mark.distributed
+
+
+@pytest.mark.parametrize(
+    "engine_enum",
+    [
+        pytest.param(
+            EngineEnum.RAY,
+            marks=pytest.mark.skip("ray not available") if not importlib.util.find_spec("ray") else [],
+        ),
+    ],
+)
+def test_engine(engine_enum: Enum) -> None:
+    assert wr.engine.get_installed() == engine_enum
+    assert wr.engine.get() == engine_enum
+    assert wr.engine._registry
+    assert wr.engine.dispatch_func(_to_parquet).__name__.endswith("distributed")
+    assert not wr.engine.dispatch_func(_to_parquet, "python").__name__.endswith("distributed")
+
+    wr.engine.register("python")
+    assert wr.engine.get_installed() == engine_enum
+    assert wr.engine.get() == EngineEnum.PYTHON
+    assert not wr.engine._registry
+    assert not wr.engine.dispatch_func(_to_parquet).__name__.endswith("distributed")
+
+
+@pytest.mark.parametrize(
+    "memory_format_enum",
+    [
+        pytest.param(
+            MemoryFormatEnum.MODIN,
+            marks=pytest.mark.skip("modin not available") if not importlib.util.find_spec("modin") else [],
+        ),
+    ],
+)
+def test_memory_format(memory_format_enum: Enum) -> None:
+    assert wr.memory_format.get_installed() == memory_format_enum
+    assert wr.memory_format.get() == memory_format_enum
+
+    wr.memory_format.set("pandas")
+    assert wr.memory_format.get_installed() == memory_format_enum
+    assert wr.memory_format.get() == MemoryFormatEnum.PANDAS

tests/unit/test_lakeformation.py

@@ -9,7 +9,7 @@
 
 from .._utils import ensure_data_types, ensure_data_types_csv, get_df, get_df_csv
 
-if wr.engine.get() == EngineEnum.RAY.value and wr.memory_format.get() == MemoryFormatEnum.MODIN.value:
+if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
     import modin.pandas as pd
 else:
     import pandas as pd

tests/unit/test_s3_select.py

@@ -5,7 +5,7 @@
 import awswrangler as wr
 from awswrangler._distributed import EngineEnum, MemoryFormatEnum
 
-if wr.engine.get() == EngineEnum.RAY.value and wr.memory_format.get() == MemoryFormatEnum.MODIN.value:
+if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
     import modin.pandas as pd
 else:
     import pandas as pd

tests/unit/test_s3_text.py

@@ -6,7 +6,7 @@
 import awswrangler as wr
 from awswrangler._distributed import EngineEnum, MemoryFormatEnum
 
-if wr.engine.get() == EngineEnum.RAY.value and wr.memory_format.get() == MemoryFormatEnum.MODIN.value:
+if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
     import modin.pandas as pd
 else:
     import pandas as pd

tests/unit/test_s3_text_compressed.py

@@ -13,7 +13,7 @@
 
 from .._utils import get_df_csv
 
-if wr.engine.get() == EngineEnum.RAY.value and wr.memory_format.get() == MemoryFormatEnum.MODIN.value:
+if wr.engine.get() == EngineEnum.RAY and wr.memory_format.get() == MemoryFormatEnum.MODIN:
     import modin.pandas as pd
 else:
     import pandas as pd