REFACTOR-#5005: Use finalize method instead of list comprehension +…

… `drain_call_queue` (#5006)

Signed-off-by: Myachev <[email protected]>

docs/release_notes/release_notes-0.16.0.rst

@@ -104,6 +104,7 @@ Key Features and Updates
   * REFACTOR-#4949: Remove code duplication in `default2pandas/dataframe.py` and `default2pandas/any.py` (#4950)
   * REFACTOR-#4976: Rename `RayTask` to `RayWrapper` in accordance with Dask (#4977)
   * REFACTOR-#4885: De-duplicated take_2d_labels_or_positional methods (#4883)
+  * REFACTOR-#5005: Use `finalize` method instead of list comprehension + `drain_call_queue` (#5006)
   * REFACTOR-#5001: Remove `jenkins` stuff (#5002)
   * REFACTOR-#4970: Rewrite base implementations of a partition' `width/length` (#4971)  
   * REFACTOR-#4942: Remove `call` method in favor of `register` due to duplication (4943)

modin/core/dataframe/pandas/partitioning/partition_manager.py

@@ -67,7 +67,7 @@ def wait(cls, *args, **kwargs):
             # finish before kicking off the next one. Instead, we want to
             # 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()]
+            cls.finalize(partitions)
             # The partition manager invokes the relevant .wait() method under
             # the hood, which should wait in parallel for all computations to finish
             cls.wait_partitions(partitions.flatten())
@@ -1313,7 +1313,6 @@ def get_right_block(right_partitions, row_idx, col_idx):
         )
 
     @classmethod
-    @wait_computations_if_benchmark_mode
     def finalize(cls, partitions):
         """
         Perform all deferred calls on partitions.