Repository: modin-project/modin

Issue:
REFACTOR: remove code duplication for `broadcast_apply` method 



Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/pandas/partitioning/partition_manager.py b/modin/core/dataframe/pandas/partitioning/partition_manager.py
--- a/modin/core/dataframe/pandas/partitioning/partition_manager.py
+++ b/modin/core/dataframe/pandas/partitioning/partition_manager.py
@@ -335,42 +335,41 @@ def broadcast_apply(cls, axis, apply_func, left, right, other_name="r"):
             Axis to apply and broadcast over.
         apply_func : callable
             Function to apply.
-        left : NumPy 2D array
-            Left partitions.
-        right : NumPy 2D array
-            Right partitions.
+        left : np.ndarray
+            NumPy array of left partitions.
+        right : np.ndarray
+            NumPy array of right partitions.
         other_name : str, default: "r"
             Name of key-value argument for `apply_func` that
             is used to pass `right` to `apply_func`.
 
         Returns
         -------
-        NumPy array
-            An of partition objects.
+        np.ndarray
+            NumPy array of result partition objects.
 
         Notes
         -----
         This will often be overridden by implementations. It materializes the
         entire partitions of the right and applies them to the left through `apply`.
         """
-        [obj.drain_call_queue() for row in right for obj in row]
-        new_right = np.empty(shape=right.shape[axis], dtype=object)
 
-        if axis:
-            right = right.T
-
-        for i in range(len(right)):
-            new_right[i] = pandas.concat(
-                [right[i][j].get() for j in range(len(right[i]))], axis=axis ^ 1
-            )
-        right = new_right.T if axis else new_right
+        def map_func(df, *others):
+            other = pandas.concat(others, axis=axis ^ 1)
+            return apply_func(df, **{other_name: other})
 
-        new_partitions = np.array(
+        map_func = cls.preprocess_func(map_func)
+        rt_axis_parts = cls.axis_partition(right, axis ^ 1)
+        return np.array(
             [
                 [
                     part.apply(
-                        apply_func,
-                        **{other_name: right[col_idx] if axis else right[row_idx]},
+                        map_func,
+                        *(
+                            rt_axis_parts[col_idx].list_of_blocks
+                            if axis
+                            else rt_axis_parts[row_idx].list_of_blocks
+                        ),
                     )
                     for col_idx, part in enumerate(left[row_idx])
                 ]
@@ -378,8 +377,6 @@ def broadcast_apply(cls, axis, apply_func, left, right, other_name="r"):
             ]
         )
 
-        return new_partitions
-
     @classmethod
     @wait_computations_if_benchmark_mode
     def broadcast_axis_partitions(
diff --git a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py
--- a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py
+++ b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py
@@ -13,7 +13,7 @@
 
 """Module houses class that implements ``PandasDataframePartitionManager``."""
 
-import numpy as np
+from distributed.client import default_client
 
 from modin.core.dataframe.pandas.partitioning.partition_manager import (
     PandasDataframePartitionManager,
@@ -24,9 +24,6 @@
 )
 from .partition import PandasOnDaskDataframePartition
 from modin.error_message import ErrorMessage
-import pandas
-
-from distributed.client import default_client
 
 
 class PandasOnDaskDataframePartitionManager(PandasDataframePartitionManager):
@@ -80,51 +77,3 @@ def get_indices(cls, axis, partitions, index_func):
             )
         new_idx = client.gather(new_idx)
         return new_idx[0].append(new_idx[1:]) if len(new_idx) else new_idx
-
-    @classmethod
-    def broadcast_apply(cls, axis, apply_func, left, right, other_name="r"):
-        """
-        Broadcast the `right` partitions to `left` and apply `apply_func` function.
-
-        Parameters
-        ----------
-        axis : {0, 1}
-            Axis to apply and broadcast over.
-        apply_func : callable
-            Function to apply.
-        left : np.ndarray
-            NumPy array of left partitions.
-        right : np.ndarray
-            NumPy array of right partitions.
-        other_name : str, default: "r"
-            Name of key-value argument for `apply_func` that
-            is used to pass `right` to `apply_func`.
-
-        Returns
-        -------
-        np.ndarray
-            NumPy array of result partition objects.
-        """
-
-        def map_func(df, *others):
-            other = pandas.concat(others, axis=axis ^ 1)
-            return apply_func(df, **{other_name: other})
-
-        map_func = cls.preprocess_func(map_func)
-        rt_axis_parts = cls.axis_partition(right, axis ^ 1)
-        return np.array(
-            [
-                [
-                    part.apply(
-                        map_func,
-                        *(
-                            rt_axis_parts[col_idx].list_of_blocks
-                            if axis
-                            else rt_axis_parts[row_idx].list_of_blocks
-                        ),
-                    )
-                    for col_idx, part in enumerate(left[row_idx])
-                ]
-                for row_idx in range(len(left))
-            ]
-        )
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py
@@ -14,9 +14,11 @@
 """Module houses class that implements ``GenericRayDataframePartitionManager`` using Ray."""
 
 import inspect
-import numpy as np
 import threading
 
+import numpy as np
+import ray
+
 from modin.config import ProgressBar, NPartitions
 from modin.core.execution.ray.generic.partitioning.partition_manager import (
     GenericRayDataframePartitionManager,
@@ -29,9 +31,6 @@
 from modin.core.execution.ray.generic.modin_aqp import call_progress_bar
 from modin.core.storage_formats.pandas.utils import compute_chunksize
 from modin.error_message import ErrorMessage
-import pandas
-
-import ray
 
 
 def progress_bar_wrapper(f):
@@ -286,54 +285,6 @@ def rebalance_partitions(cls, partitions):
             start = stop + 1
         return np.array(new_partitions)
 
-    @classmethod
-    def broadcast_apply(cls, axis, apply_func, left, right, other_name="r"):
-        """
-        Broadcast the `right` partitions to `left` and apply `apply_func` to selected indices.
-
-        Parameters
-        ----------
-        axis : {0, 1}
-            Axis to apply and broadcast over.
-        apply_func : callable
-            Function to apply.
-        left : np.ndarray
-            NumPy 2D array of left partitions.
-        right : np.ndarray
-            NumPy 2D array of right partitions.
-        other_name : str, default: "r"
-            Name of key-value argument for `apply_func` that
-            is used to pass `right` to `apply_func`.
-
-        Returns
-        -------
-        np.ndarray
-            An array of partition objects.
-        """
-
-        def map_func(df, *others):
-            other = pandas.concat(others, axis=axis ^ 1)
-            return apply_func(df, **{other_name: other})
-
-        map_func = cls.preprocess_func(map_func)
-        rt_axis_parts = cls.axis_partition(right, axis ^ 1)
-        return np.array(
-            [
-                [
-                    part.apply(
-                        map_func,
-                        *(
-                            rt_axis_parts[col_idx].list_of_blocks
-                            if axis
-                            else rt_axis_parts[row_idx].list_of_blocks
-                        ),
-                    )
-                    for col_idx, part in enumerate(left[row_idx])
-                ]
-                for row_idx in range(len(left))
-            ]
-        )
-
     @classmethod
     @progress_bar_wrapper
     def map_partitions(cls, partitions, map_func):


Test patch: