Repository: modin-project/modin

Issue:
PERF: `drain` function should be serialized only once



Discussion:

Resolving patch:
diff --git a/modin/core/dataframe/pandas/partitioning/axis_partition.py b/modin/core/dataframe/pandas/partitioning/axis_partition.py
--- a/modin/core/dataframe/pandas/partitioning/axis_partition.py
+++ b/modin/core/dataframe/pandas/partitioning/axis_partition.py
@@ -213,3 +213,22 @@ def deploy_func_between_two_axis_partitions(
         # for func, we fold args into kwargs. This is a bit of a hack, but it works.
         result = func(lt_frame, rt_frame, *kwargs.pop("args", ()), **kwargs)
         return split_result_of_axis_func_pandas(axis, num_splits, result)
+
+    @classmethod
+    def drain(cls, df: pandas.DataFrame, call_queue: list):
+        """
+        Execute all operations stored in the call queue on the pandas object (helper function).
+
+        Parameters
+        ----------
+        df : pandas.DataFrame
+        call_queue : list
+            Call queue that needs to be executed on pandas DataFrame.
+
+        Returns
+        -------
+        pandas.DataFrame
+        """
+        for func, args, kwargs in call_queue:
+            df = func(df, *args, **kwargs)
+        return df
diff --git a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/virtual_partition.py b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/virtual_partition.py
--- a/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/virtual_partition.py
+++ b/modin/core/execution/dask/implementations/pandas_on_dask/partitioning/virtual_partition.py
@@ -418,14 +418,12 @@ def drain_call_queue(self, num_splits=None):
         num_splits : int, default: None
             The number of times to split the result object.
         """
-
-        def drain(df):
-            for func, args, kwargs in self.call_queue:
-                df = func(df, *args, **kwargs)
-            return df
-
+        # TODO: Need to check if `drain_call_queue` speeds up if helper
+        # `drain` function is serialized only once.
         drained = super(PandasOnDaskDataframeVirtualPartition, self).apply(
-            drain, num_splits=num_splits
+            PandasDataframeAxisPartition.drain,
+            num_splits=num_splits,
+            call_queue=self.call_queue,
         )
         self._list_of_block_partitions = drained
         self.call_queue = []
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py
@@ -25,6 +25,7 @@
 
 
 _DEPLOY_AXIS_FUNC = ray.put(PandasDataframeAxisPartition.deploy_axis_func)
+_DRAIN = ray.put(PandasDataframeAxisPartition.drain)
 
 
 class PandasOnRayDataframeVirtualPartition(PandasDataframeAxisPartition):
@@ -427,14 +428,8 @@ def drain_call_queue(self, num_splits=None):
         num_splits : int, default: None
             The number of times to split the result object.
         """
-
-        def drain(df):
-            for func, args, kwargs in self.call_queue:
-                df = func(df, *args, **kwargs)
-            return df
-
         drained = super(PandasOnRayDataframeVirtualPartition, self).apply(
-            drain, num_splits=num_splits
+            _DRAIN, num_splits=num_splits, call_queue=self.call_queue
         )
         self._list_of_block_partitions = drained
         self.call_queue = []


Test patch: