Repository: modin-project/modin

Issue:
BUG: benchmark mode should wait for all partitions in parallel
We are waiting for all partitions in sequence, but it's more accurate to wait for them all at once, e.g with `ray.wait(partitions, num_returns=len(partitions))` for Ray. The performance difference is probably negligible because iterating across all the partitions should be really fast.

https://github.com/modin-project/modin/blob/c17dde71ae8114723a13279e36bf5532dcba328a/modin/core/dataframe/pandas/partitioning/partition_manager.py#L66


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
@@ -53,9 +53,9 @@ def wait_computations_if_benchmark_mode(func):
     if BenchmarkMode.get():
 
         @wraps(func)
-        def wait(*args, **kwargs):
+        def wait(cls, *args, **kwargs):
             """Wait for computation results."""
-            result = func(*args, **kwargs)
+            result = func(cls, *args, **kwargs)
             if isinstance(result, tuple):
                 partitions = result[0]
             else:
@@ -67,12 +67,11 @@ def wait(*args, **kwargs):
             # serially kick off all the deferred computations so that they can
             # all run asynchronously, then wait on all the results.
             [part.drain_call_queue() for part in partitions.flatten()]
-            # need to go through all the values of the map iterator
-            # since `wait` does not return anything, we need to explicitly add
-            # the return `True` value from the lambda
-            # TODO(https://github.com/modin-project/modin/issues/4491): Wait
-            # for all the partitions in parallel.
-            all(map(lambda partition: partition.wait() or True, partitions.flatten()))
+            # The partition manager invokes the relevant .wait() method under
+            # the hood, which should wait in parallel for all computations to finish
+            # (We can't just add a `cls` argument to this `wait` function, since doing so
+            # seems to be incompatible with the way function decorators work)
+            cls.wait_partitions(partitions.flatten())
             return result
 
         return wait
@@ -810,6 +809,26 @@ def get_objects_from_partitions(cls, partitions):
         """
         return [partition.get() for partition in partitions]
 
+    @classmethod
+    def wait_partitions(cls, partitions):
+        """
+        Wait on the objects wrapped by `partitions`, without materializing them.
+
+        This method will block until all computations in the list have completed.
+
+        Parameters
+        ----------
+        partitions : np.ndarray
+            NumPy array with ``PandasDataframePartition``-s.
+
+        Notes
+        -----
+        This method should be implemented in a more efficient way for engines that supports
+        waiting on objects in parallel.
+        """
+        for partition in partitions:
+            partition.wait()
+
     @classmethod
     def get_indices(cls, axis, partitions, index_func=None):
         """
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,6 +13,8 @@
 
 """Module houses class that implements ``PandasDataframePartitionManager``."""
 
+from dask.distributed import wait
+
 from modin.core.dataframe.pandas.partitioning.partition_manager import (
     PandasDataframePartitionManager,
 )
@@ -48,3 +50,17 @@ def get_objects_from_partitions(cls, partitions):
             The objects wrapped by `partitions`.
         """
         return DaskWrapper.materialize([partition._data for partition in partitions])
+
+    @classmethod
+    def wait_partitions(cls, partitions):
+        """
+        Wait on the objects wrapped by `partitions` in parallel, without materializing them.
+
+        This method will block until all computations in the list have completed.
+
+        Parameters
+        ----------
+        partitions : np.ndarray
+            NumPy array with ``PandasDataframePartition``-s.
+        """
+        wait([partition._data for partition in partitions], return_when="ALL_COMPLETED")
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
@@ -109,6 +109,22 @@ def get_objects_from_partitions(cls, partitions):
         """
         return ray.get([partition._data for partition in partitions])
 
+    @classmethod
+    def wait_partitions(cls, partitions):
+        """
+        Wait on the objects wrapped by `partitions` in parallel, without materializing them.
+
+        This method will block until all computations in the list have completed.
+
+        Parameters
+        ----------
+        partitions : np.ndarray
+            NumPy array with ``PandasDataframePartition``-s.
+        """
+        ray.wait(
+            [partition._data for partition in partitions], num_returns=len(partitions)
+        )
+
     @classmethod
     @progress_bar_wrapper
     def map_partitions(cls, partitions, map_func):


Test patch: