refactor: engine register_func to handle type checking (#2309)

Author: LeonLuttenberger

awswrangler/_distributed.py

@@ -89,11 +89,10 @@ def dispatch_func(cls, source_func: FunctionType, value: Optional[EngineLiteral]
             return getattr(source_func, "_source_func", source_func)
 
     @classmethod
-    def register_func(cls, source_func: Callable[..., Any], destination_func: Callable[..., Any]) -> Callable[..., Any]:
+    def register_func(cls, source_func: FunctionType, destination_func: FunctionType) -> None:
         """Register a func based on the distribution engine and source function."""
         with cls._lock:
             cls._registry[cls.get().value][source_func.__name__] = destination_func
-        return destination_func
 
     @classmethod
     def dispatch_on_engine(cls, func: FunctionType) -> FunctionType:

awswrangler/distributed/ray/_core.py

@@ -2,7 +2,7 @@
 import logging
 import os
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, TypeVar, Union
 
 from awswrangler._config import apply_configs
 from awswrangler._distributed import EngineEnum, engine
@@ -13,6 +13,9 @@
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
+FunctionType = TypeVar("FunctionType", bound=Callable[..., Any])
+
+
 class RayLogger:
     """Create discrete Logger instance for Ray Tasks."""
 
@@ -31,10 +34,10 @@ def get_logger(self, name: Union[str, Any] = None) -> Optional[logging.Logger]:
 
 @apply_configs
 def ray_logger(
-    function: Callable[..., Any],
+    function: FunctionType,
     configure_logging: bool = True,
     logging_level: int = logging.INFO,
-) -> Callable[..., Any]:
+) -> FunctionType:
     """
     Decorate callable to add RayLogger.
 
@@ -57,7 +60,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
     return wrapper
 
 
-def ray_remote(**options: Any) -> Callable[..., Any]:
+def ray_remote(**options: Any) -> Callable[[FunctionType], FunctionType]:
     """
     Decorate with @ray.remote providing .options().
 
@@ -71,7 +74,7 @@ def ray_remote(**options: Any) -> Callable[..., Any]:
     Callable[..., Any]
     """
 
-    def remote_decorator(function: Callable[..., Any]) -> Callable[..., Any]:
+    def remote_decorator(function: FunctionType) -> FunctionType:
         """
         Decorate callable to wrap within ray.remote.
 

awswrangler/distributed/ray/_register.py

@@ -55,13 +55,11 @@ def register_ray() -> None:
     ]:
         engine.register_func(func, ray_remote()(func))
 
-    for o_f, d_f in {
-        _get_executor: _get_ray_executor,
-        _list_objects_paginate: _list_objects_s3fs,
-        _read_parquet_metadata_file: ray_remote()(_read_parquet_metadata_file_distributed),
-        ensure_worker_or_thread_count: ensure_worker_count,
-    }.items():
-        engine.register_func(o_f, d_f)  # type: ignore[arg-type]
+    # Register dispatch methods for Ray
+    engine.register_func(_get_executor, _get_ray_executor)
+    engine.register_func(_list_objects_paginate, _list_objects_s3fs)
+    engine.register_func(_read_parquet_metadata_file, ray_remote()(_read_parquet_metadata_file_distributed))
+    engine.register_func(ensure_worker_or_thread_count, ensure_worker_count)
 
     if memory_format.get() == MemoryFormatEnum.MODIN:
         from awswrangler.distributed.ray.modin._data_types import pyarrow_types_from_pandas_distributed
@@ -80,17 +78,15 @@ def register_ray() -> None:
         from awswrangler.distributed.ray.modin.s3._write_parquet import _to_parquet_distributed
         from awswrangler.distributed.ray.modin.s3._write_text import _to_text_distributed
 
-        for o_f, d_f in {
-            pyarrow_types_from_pandas: pyarrow_types_from_pandas_distributed,
-            _read_parquet: _read_parquet_distributed,
-            _read_text: _read_text_distributed,
-            _to_buckets: _to_buckets_distributed,
-            _to_parquet: _to_parquet_distributed,
-            _to_partitions: _to_partitions_distributed,
-            _to_text: _to_text_distributed,
-            copy_df_shallow: _copy_modin_df_shallow,
-            is_pandas_frame: _is_pandas_or_modin_frame,
-            split_pandas_frame: _split_modin_frame,
-            table_refs_to_df: _arrow_refs_to_df,
-        }.items():
-            engine.register_func(o_f, d_f)  # type: ignore[arg-type]
+        # Register dispatch methods for Modin
+        engine.register_func(pyarrow_types_from_pandas, pyarrow_types_from_pandas_distributed)
+        engine.register_func(_read_parquet, _read_parquet_distributed)
+        engine.register_func(_read_text, _read_text_distributed)
+        engine.register_func(_to_buckets, _to_buckets_distributed)
+        engine.register_func(_to_parquet, _to_parquet_distributed)
+        engine.register_func(_to_partitions, _to_partitions_distributed)
+        engine.register_func(_to_text, _to_text_distributed)
+        engine.register_func(copy_df_shallow, _copy_modin_df_shallow)
+        engine.register_func(is_pandas_frame, _is_pandas_or_modin_frame)
+        engine.register_func(split_pandas_frame, _split_modin_frame)
+        engine.register_func(table_refs_to_df, _arrow_refs_to_df)

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

@@ -1,5 +1,5 @@
 """Modin on Ray S3 read parquet module (PRIVATE)."""
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import modin.pandas as pd
 import pyarrow as pa
@@ -37,7 +37,7 @@ def _read_parquet_distributed(  # pylint: disable=unused-argument
     s3_additional_kwargs: Optional[Dict[str, Any]],
     arrow_kwargs: Dict[str, Any],
     bulk_read: bool,
-) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
+) -> pd.DataFrame:
     dataset_kwargs = {}
     if coerce_int96_timestamp_unit:
         dataset_kwargs["coerce_int96_timestamp_unit"] = coerce_int96_timestamp_unit

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

@@ -121,14 +121,14 @@ def _read_text_distributed(  # pylint: disable=unused-argument
     read_format: str,
     paths: List[str],
     path_root: Optional[str],
+    use_threads: Union[bool, int],
+    s3_client: Optional["S3Client"],
     s3_additional_kwargs: Optional[Dict[str, str]],
     dataset: bool,
     ignore_index: bool,
     parallelism: int,
     version_ids: Optional[Dict[str, str]],
     pandas_kwargs: Dict[str, Any],
-    use_threads: Union[bool, int],
-    s3_client: Optional["S3Client"],
 ) -> pd.DataFrame:
     try:
         configuration: Dict[str, Any] = _parse_configuration(  # type: ignore[assignment]

awswrangler/s3/_read_parquet.py

@@ -331,7 +331,7 @@ def _read_parquet(  # pylint: disable=W0613
     s3_additional_kwargs: Optional[Dict[str, Any]],
     arrow_kwargs: Dict[str, Any],
     bulk_read: bool,
-) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
+) -> pd.DataFrame:
     executor: _BaseExecutor = _get_executor(use_threads=use_threads)
     tables = executor.map(
         _read_parquet_file,

awswrangler/s3/_read_text.py

@@ -52,7 +52,7 @@ def _read_text(  # pylint: disable=W0613
     parallelism: int,
     version_ids: Optional[Dict[str, str]],
     pandas_kwargs: Dict[str, Any],
-) -> Union[pd.DataFrame, Iterator[pd.DataFrame]]:
+) -> pd.DataFrame:
     parser_func = _resolve_format(read_format)
     executor: _BaseExecutor = _get_executor(use_threads=use_threads)
     tables = executor.map(

awswrangler/s3/_write_parquet.py

@@ -171,7 +171,7 @@ def _to_parquet_chunked(
 
 
 @engine.dispatch_on_engine
-def _to_parquet(
+def _to_parquet(  # pylint: disable=unused-argument
     df: pd.DataFrame,
     schema: pa.Schema,
     index: bool,
@@ -187,6 +187,7 @@ def _to_parquet(
     path_root: Optional[str] = None,
     filename_prefix: Optional[str] = None,
     max_rows_by_file: Optional[int] = 0,
+    bucketing: bool = False,
 ) -> List[str]:
     s3_client = s3_client if s3_client else _utils.client(service_name="s3")
     file_path = _get_file_path(