Add source func

Author: jaidisido

awswrangler/_distributed.py

@@ -79,9 +79,11 @@ def register_func(cls, source_func: Callable[..., Any], destination_func: Callab
     @classmethod
     def dispatch_on_engine(cls, func: Callable[..., Any]) -> Callable[..., Any]:
         """Dispatch on engine function decorator."""
+
         @wraps(func)
         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
         return wrapper

awswrangler/distributed/ray/_core.py

@@ -60,11 +60,12 @@ def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]:
     Callable[..., Any]
     """
     # Access the source function if it exists
-    function = getattr(function, '_source_func', function)
+    function = getattr(function, "_source_func", function)
 
     @wraps(function)
     def wrapper(*args: Any, **kwargs: Any) -> Any:
         return ray.remote(function).remote(*args, **kwargs)  # type: ignore
+
     return wrapper
 
 

awswrangler/distributed/ray/_register.py

@@ -8,7 +8,7 @@
 from awswrangler.s3._delete import _delete_objects
 from awswrangler.s3._read_parquet import _read_parquet, _read_parquet_metadata_file
 from awswrangler.s3._read_text import _read_text
-from awswrangler.s3._select import _select_object_content
+from awswrangler.s3._select import _select_object_content, _select_query
 from awswrangler.s3._wait import _batch_paths, _wait_object_batch
 from awswrangler.s3._write_dataset import _to_buckets, _to_partitions
 from awswrangler.s3._write_parquet import _to_parquet
@@ -23,6 +23,7 @@ def register_ray() -> None:
     # S3
     engine.register_func(_delete_objects, ray_remote(_delete_objects))
     engine.register_func(_read_parquet_metadata_file, ray_remote(_read_parquet_metadata_file))
+    engine.register_func(_select_query, ray_remote(_select_query))
     engine.register_func(_select_object_content, ray_remote(_select_object_content))
     engine.register_func(_wait_object_batch, ray_remote(_wait_object_batch))
     engine.register_func(_batch_paths, _batch_paths_distributed)

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

@@ -3,9 +3,8 @@
 
 import modin.pandas as pd
 import ray
-from modin.pandas import DataFrame as ModinDataFrame
-from pandas import DataFrame as PandasDataFrame
 
+from awswrangler._distributed import engine
 from awswrangler.distributed.ray import ray_get, ray_remote
 from awswrangler.s3._write_concurrent import _WriteProxy
 from awswrangler.s3._write_dataset import _delete_objects, _get_bucketing_series, _retrieve_paths, _to_partitions
@@ -29,7 +28,7 @@ def _to_buckets_distributed(  # pylint: disable=unused-argument
     paths: List[str] = []
 
     df_paths = df_groups.apply(
-        func.dispatch(ModinDataFrame),  # type: ignore
+        engine.dispatch_func(func),  # type: ignore
         path_root=path_root,
         filename_prefix=filename_prefix,
         boto3_session=None,
@@ -126,7 +125,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 = _to_partitions.dispatch(PandasDataFrame)
+        _to_partitions_func = _to_partitions._source_func  # type: ignore
+        func = func._source_func  # type: ignore
 
         @ray_remote
         def write_partitions(df: pd.DataFrame) -> Tuple[List[str], Dict[str, List[str]]]:

awswrangler/s3/_read_parquet.py

@@ -19,7 +19,7 @@
 from awswrangler._threading import _get_executor
 from awswrangler.catalog._get import _get_partitions
 from awswrangler.catalog._utils import _catalog_id
-from awswrangler.distributed.ray import RayLogger, ray_get, ray_remote
+from awswrangler.distributed.ray import RayLogger, ray_get
 from awswrangler.s3._fs import open_s3_object
 from awswrangler.s3._list import _path2list
 from awswrangler.s3._read import (
@@ -60,7 +60,7 @@ def _pyarrow_parquet_file_wrapper(
         raise
 
 
-@ray_remote
+@engine.dispatch_on_engine
 def _read_parquet_metadata_file(
     boto3_session: boto3.Session,
     path: str,