(feat) validate partitions along row axis, add warning (#1700)

* Check df is partitioned along row axis, add warning

* Fix - pass axis

* PR feedback

Author: kukushking

awswrangler/distributed/ray/modin/_core.py

@@ -1,12 +1,34 @@
 """Modin on Ray Core module (PRIVATE)."""
 # pylint: disable=import-outside-toplevel
+import logging
 from functools import wraps
 from typing import Any, Callable, Optional
 
+import numpy as np
 import pandas as pd
 from modin.distributed.dataframe.pandas import from_partitions, unwrap_partitions
 from modin.pandas import DataFrame as ModinDataFrame
 
+_logger: logging.Logger = logging.getLogger(__name__)
+
+
+def _validate_partition_shape(df: pd.DataFrame) -> bool:
+    """
+    Validate if partitions of the data frame are partitioned along row axis.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        Modin data frame
+
+    Returns
+    -------
+    bool
+    """
+    # Unwrap partitions as they are currently stored (axis=None)
+    partitions_shape = np.array(unwrap_partitions(df)).shape
+    return partitions_shape[1] == 1
+
 
 def modin_repartition(function: Callable[..., Any]) -> Callable[..., Any]:
     """
@@ -31,14 +53,23 @@ def modin_repartition(function: Callable[..., Any]) -> Callable[..., Any]:
     def wrapper(
         df: pd.DataFrame,
         *args: Any,
-        axis: int = 0,
+        axis: Optional[int] = None,
         row_lengths: Optional[int] = None,
+        validate_partitions: bool = True,
         **kwargs: Any,
     ) -> Any:
-        # Repartition Modin data frame along row (axis=0) axis
+        # Validate partitions and repartition Modin data frame along row (axis=0) axis
         # to avoid a situation where columns are split along multiple blocks
-        if isinstance(df, ModinDataFrame) and axis is not None:
-            df = from_partitions(unwrap_partitions(df, axis=axis), axis=axis, row_lengths=row_lengths)
+        if isinstance(df, ModinDataFrame):
+            if validate_partitions and not _validate_partition_shape(df):
+                _logger.warning(
+                    "Partitions of this data frame are detected to be split along column axis. "
+                    "The dataframe will be automatically repartitioned along row axis to ensure "
+                    "each partition can be processed independently."
+                )
+                axis = 0
+            if axis is not None:
+                df = from_partitions(unwrap_partitions(df, axis=axis), axis=axis, row_lengths=row_lengths)
         return function(df, *args, **kwargs)
 
     return wrapper