fix: distributed write text regression, change to singledispatch, add repartitioning utility (#1611)

* Fix distributed write text regression

* Try out singledispatch

* Minor fixes

* Refactoring

* Fix write args order

* Fix imports

* Fix import Modin df

Author: kukushking

awswrangler/distributed/__init__.py

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

awswrangler/distributed/_distributed.py

@@ -1,16 +1,22 @@
 """Distributed Module (PRIVATE)."""
 
+import logging
 import multiprocessing
 import os
 import sys
 import warnings
+from functools import wraps
 from typing import TYPE_CHECKING, Any, Callable, List, Optional
 
 from awswrangler._config import apply_configs, config
 
 if config.distributed or TYPE_CHECKING:
     import psutil
     import ray  # pylint: disable=import-error
+    from modin.distributed.dataframe.pandas import from_partitions, unwrap_partitions
+    from modin.pandas import DataFrame as ModinDataFrame
+
+_logger: logging.Logger = logging.getLogger(__name__)
 
 
 def ray_get(futures: List[Any]) -> List[Any]:
@@ -46,13 +52,42 @@ def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]:
     """
     if config.distributed:
 
+        @wraps(function)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
             return ray.remote(function).remote(*args, **kwargs)
 
         return wrapper
     return function
 
 
+def modin_repartition(function: Callable[..., Any]) -> Callable[..., Any]:
+    """
+    Decorate callable to repartition Modin data frame.
+
+    By default, repartition along row (axis=0) axis.
+    This avoids a situation where columns are split along multiple blocks.
+
+    Parameters
+    ----------
+    function : Callable[..., Any]
+        Callable as input to ray.remote
+
+    Returns
+    -------
+    Callable[..., Any]
+    """
+
+    @wraps(function)
+    def wrapper(df, *args: Any, axis=0, row_lengths=None, **kwargs: Any) -> Any:
+        if config.distributed and isinstance(df, ModinDataFrame) and axis is not None:
+            # Repartition Modin data frame along row (axis=0) axis
+            # to avoid a situation where columns are split along multiple blocks
+            df = from_partitions(unwrap_partitions(df, axis=axis), axis=axis, row_lengths=row_lengths)
+        return function(df, *args, **kwargs)
+
+    return wrapper
+
+
 @apply_configs
 def initialize_ray(
     address: Optional[str] = None,

awswrangler/distributed/_distributed.pyi

@@ -10,6 +10,7 @@ from awswrangler._config import apply_configs, config
 
 def ray_get(futures: List[Any]) -> List[Any]: ...
 def ray_remote(function: Callable[..., Any]) -> Callable[..., Any]: ...
+def modin_repartition(function: Callable[..., Any]) -> Callable[..., Any]: ...
 def initialize_ray(
     address: Optional[str] = None,
     redis_password: Optional[str] = None,

awswrangler/s3/__init__.py

@@ -1,5 +1,6 @@
 """Amazon S3 Read Module."""
 
+from awswrangler._config import config
 from awswrangler.s3._copy import copy_objects, merge_datasets  # noqa
 from awswrangler.s3._delete import delete_objects  # noqa
 from awswrangler.s3._describe import describe_objects, get_bucket_region, size_objects  # noqa
@@ -45,3 +46,18 @@
     "download",
     "upload",
 ]
+
+if config.distributed:
+    from modin.pandas import DataFrame as ModinDataFrame
+
+    from awswrangler.s3._write_dataset import (  # pylint: disable=ungrouped-imports
+        _to_buckets,
+        _to_buckets_distributed,
+        _to_partitions,
+        _to_partitions_distributed,
+    )
+    from awswrangler.s3._write_parquet import _to_parquet, _to_parquet_distributed  # pylint: disable=ungrouped-imports
+
+    _to_parquet.register(ModinDataFrame, _to_parquet_distributed)
+    _to_buckets.register(ModinDataFrame, _to_buckets_distributed)
+    _to_partitions.register(ModinDataFrame, _to_partitions_distributed)

awswrangler/s3/_write_concurrent.py

@@ -25,27 +25,33 @@ def __init__(self, use_threads: Union[bool, int]):
 
     @staticmethod
     def _caller(
-        func: Callable[..., pd.DataFrame], boto3_primitives: _utils.Boto3PrimitivesType, func_kwargs: Dict[str, Any]
+        func: Callable[..., pd.DataFrame],
+        boto3_primitives: _utils.Boto3PrimitivesType,
+        *args: Any,
+        func_kwargs: Dict[str, Any],
     ) -> pd.DataFrame:
         boto3_session: boto3.Session = _utils.boto3_from_primitives(primitives=boto3_primitives)
         func_kwargs["boto3_session"] = boto3_session
         _logger.debug("Calling: %s", func)
-        return func(**func_kwargs)
+        return func(*args, **func_kwargs)
 
-    def write(self, func: Callable[..., List[str]], boto3_session: boto3.Session, **func_kwargs: Any) -> None:
+    def write(
+        self, func: Callable[..., List[str]], boto3_session: boto3.Session, *args: Any, **func_kwargs: Any
+    ) -> None:
         """Write File."""
         if self._exec is not None:
             _utils.block_waiting_available_thread(seq=self._futures, max_workers=self._cpus)
             _logger.debug("Submitting: %s", func)
             future = self._exec.submit(
                 _WriteProxy._caller,
-                func=func,
-                boto3_primitives=_utils.boto3_to_primitives(boto3_session=boto3_session),
+                func,
+                _utils.boto3_to_primitives(boto3_session=boto3_session),
+                *args,
                 func_kwargs=func_kwargs,
             )
             self._futures.append(future)
         else:
-            self._results += func(boto3_session=boto3_session, **func_kwargs)
+            self._results += func(*args, boto3_session=boto3_session, **func_kwargs)
 
     def close(self) -> List[str]:
         """Close the proxy."""

awswrangler/s3/_write_dataset.py

@@ -1,6 +1,7 @@
 """Amazon S3 Write Dataset (PRIVATE)."""
 
 import logging
+from functools import singledispatch
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import boto3
@@ -13,19 +14,13 @@
 
 if config.distributed:
     import modin.pandas as pd
-    from modin.distributed.dataframe.pandas import from_partitions, unwrap_partitions
     from modin.pandas import DataFrame as ModinDataFrame
 else:
     import pandas as pd
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
-def _get_subgroup_prefix(keys: Tuple[str, None], partition_cols: List[str], path_root: str) -> str:
-    subdir = "/".join([f"{name}={val}" for name, val in zip(partition_cols, keys)])
-    return f"{path_root}{subdir}/"
-
-
 def _get_bucketing_series(df: pd.DataFrame, bucketing_info: Tuple[List[str], int]) -> pd.Series:
     bucket_number_series = df.astype("O").apply(
         lambda row: _get_bucket_number(bucketing_info[1], [row[col_name] for col_name in bucketing_info[0]]),
@@ -75,6 +70,11 @@ def _get_value_hash(value: Union[str, int, bool]) -> int:
     )
 
 
+def _get_subgroup_prefix(keys: Tuple[str, None], partition_cols: List[str], path_root: str) -> str:
+    subdir = "/".join([f"{name}={val}" for name, val in zip(partition_cols, keys)])
+    return f"{path_root}{subdir}/"
+
+
 def _delete_objects(
     keys: Tuple[str, None],
     path_root: str,
@@ -168,7 +168,7 @@ def _write_partitions_distributed(
         )
     else:
         paths = write_func(  # type: ignore
-            df=df_group.drop(partition_cols, axis="columns"),
+            df_group.drop(partition_cols, axis="columns"),
             path_root=prefix,
             filename_prefix=filename_prefix,
             boto3_session=boto3_session,
@@ -178,10 +178,11 @@ def _write_partitions_distributed(
     return prefix, df_group.name, paths
 
 
+@singledispatch
 def _to_partitions(
+    df: pd.DataFrame,
     func: Callable[..., List[str]],
     concurrent_partitioning: bool,
-    df: pd.DataFrame,
     path_root: str,
     use_threads: Union[bool, int],
     mode: str,
@@ -221,8 +222,8 @@ def _to_partitions(
         )
         if bucketing_info:
             _to_buckets(
+                subgroup,
                 func=func,
-                df=subgroup,
                 path_root=prefix,
                 bucketing_info=bucketing_info,
                 boto3_session=boto3_session,
@@ -233,11 +234,11 @@ def _to_partitions(
             )
         else:
             proxy.write(
-                func=func,
-                df=subgroup,
+                func,
+                boto3_session,
+                subgroup,
                 path_root=prefix,
                 filename_prefix=filename_prefix,
-                boto3_session=boto3_session,
                 use_threads=use_threads,
                 **func_kwargs,
             )
@@ -247,9 +248,9 @@ def _to_partitions(
 
 
 def _to_partitions_distributed(  # pylint: disable=unused-argument
+    df: pd.DataFrame,
     func: Callable[..., List[str]],
     concurrent_partitioning: bool,
-    df: pd.DataFrame,
     path_root: str,
     use_threads: Union[bool, int],
     mode: str,
@@ -283,7 +284,7 @@ def _to_partitions_distributed(  # pylint: disable=unused-argument
         boto3_session=None,
         **func_kwargs,
     )
-    paths: List[str] = [path for metadata in df_write_metadata.values for _, _, path in metadata]
+    paths: List[str] = [path for metadata in df_write_metadata.values for _, _, paths in metadata for path in paths]
     partitions_values: Dict[str, List[str]] = {
         prefix: list(str(p) for p in partitions) if isinstance(partitions, tuple) else [str(partitions)]
         for metadata in df_write_metadata.values
@@ -292,9 +293,10 @@ def _to_partitions_distributed(  # pylint: disable=unused-argument
     return paths, partitions_values
 
 
+@singledispatch
 def _to_buckets(
-    func: Callable[..., List[str]],
     df: pd.DataFrame,
+    func: Callable[..., List[str]],
     path_root: str,
     bucketing_info: Tuple[List[str], int],
     filename_prefix: str,
@@ -307,11 +309,11 @@ def _to_buckets(
     df_groups = df.groupby(by=_get_bucketing_series(df=df, bucketing_info=bucketing_info))
     for bucket_number, subgroup in df_groups:
         _proxy.write(
-            func=func,
-            df=subgroup,
+            func,
+            boto3_session,
+            subgroup,
             path_root=path_root,
             filename_prefix=f"{filename_prefix}_bucket-{bucket_number:05d}",
-            boto3_session=boto3_session,
             use_threads=use_threads,
             **func_kwargs,
         )
@@ -322,8 +324,8 @@ def _to_buckets(
 
 
 def _to_buckets_distributed(  # pylint: disable=unused-argument
-    func: Callable[..., List[str]],
     df: pd.DataFrame,
+    func: Callable[..., List[str]],
     path_root: str,
     bucketing_info: Tuple[List[str], int],
     filename_prefix: str,
@@ -335,7 +337,7 @@ def _to_buckets_distributed(  # pylint: disable=unused-argument
     df_groups = df.groupby(by=_get_bucketing_series(df=df, bucketing_info=bucketing_info))
     paths: List[str] = []
     df_paths = df_groups.apply(
-        func,
+        func.dispatch(ModinDataFrame),  # type: ignore
         path_root=path_root,
         filename_prefix=filename_prefix,
         boto3_session=None,
@@ -398,24 +400,14 @@ def _to_dataset(
         else:
             delete_objects(path=path_root, use_threads=use_threads, boto3_session=boto3_session)
 
-    _to_partitions_fn: Callable[..., Tuple[List[str], Dict[str, List[str]]]] = _to_partitions
-    _to_buckets_fn: Callable[..., List[str]] = _to_buckets
-    if config.distributed and isinstance(df, ModinDataFrame):
-        # Ensure Modin dataframe is partitioned along row axis
-        # It avoids a situation where columns are split along multiple blocks
-        df = from_partitions(unwrap_partitions(df, axis=0), axis=0)
-
-        _to_partitions_fn = _to_partitions_distributed
-        _to_buckets_fn = _to_buckets_distributed
-
     # Writing
     partitions_values: Dict[str, List[str]] = {}
     paths: List[str]
     if partition_cols:
-        paths, partitions_values = _to_partitions_fn(
+        paths, partitions_values = _to_partitions(
+            df,
             func=func,
             concurrent_partitioning=concurrent_partitioning,
-            df=df,
             path_root=path_root,
             use_threads=use_threads,
             mode=mode,
@@ -433,9 +425,9 @@ def _to_dataset(
             **func_kwargs,
         )
     elif bucketing_info:
-        paths = _to_buckets_fn(
+        paths = _to_buckets(
+            df,
             func=func,
-            df=df,
             path_root=path_root,
             use_threads=use_threads,
             bucketing_info=bucketing_info,
@@ -446,7 +438,7 @@ def _to_dataset(
         )
     else:
         paths = func(
-            df=df,
+            df,
             path_root=path_root,
             filename_prefix=filename_prefix,
             use_threads=use_threads,