Repository: modin-project/modin

Issue:
Add trigger_execution to dataframe and series
Sometimes I need to "force" execution (trigger execution of all delayed computations).

It is useful for benchmarking. Currently benchmarks have to do something like [this](https://github.com/intel-ai/omniscripts/blob/master/utils/trigger_execution.py) to trigger computations. As you can see it is dependent on backend. It feels like this should be part of modin API to avoid clients to go so deep inside of the library.


Discussion:
@Egor-Krivov yes, this is a feature we mentioned [here](https://github.com/modin-project/modin/pull/5020#pullrequestreview-1116775616) while writing the [benchmarking guide](https://modin.readthedocs.io/en/latest/usage_guide/benchmarking.html). That guide now says

> Sometimes, if you don’t have a natural end-point to your workflow, you can just call `repr` on the workflow’s final Modin objects.

but also points out that `repr` can be misleading, too.
Thank you for the link. From the link I get that __repr__ is not completely reliable, since it will only block on first and last rows. So the feature is still relevant for benchmarking.

> The problem is that the repr blocks only on getting the first few and the last few rows, but the slow execution is for row 5001, which Modin is computing asynchronously in the background even after repr finishes.

Patch:
diff --git a/modin/conftest.py b/modin/conftest.py
--- a/modin/conftest.py
+++ b/modin/conftest.py
@@ -175,6 +175,10 @@ def __init__(self, modin_frame):
     def finalize(self):
         self._modin_frame.finalize()
 
+    def execute(self):
+        self.finalize()
+        self._modin_frame.wait_computations()
+
     @classmethod
     def from_pandas(cls, df, data_cls):
         return cls(data_cls.from_pandas(df))
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -4106,6 +4106,10 @@ def finalize(self):
         """
         self._partition_mgr_cls.finalize(self._partitions)
 
+    def wait_computations(self):
+        """Wait for all computations to complete without materializing data."""
+        self._partition_mgr_cls.wait_partitions(self._partitions.flatten())
+
     def __dataframe__(self, nan_as_null: bool = False, allow_copy: bool = True):
         """
         Get a Modin DataFrame that implements the dataframe exchange protocol.
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -317,6 +317,11 @@ def finalize(self):
         """Finalize constructing the dataframe calling all deferred functions which were used to build it."""
         pass
 
+    @abc.abstractmethod
+    def execute(self):
+        """Wait for all computations to complete without materializing data."""
+        pass
+
     # END Data Management Methods
 
     # To/From Pandas
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -278,6 +278,10 @@ def lazy_execution(self):
     def finalize(self):
         self._modin_frame.finalize()
 
+    def execute(self):
+        self.finalize()
+        self._modin_frame.wait_computations()
+
     def to_pandas(self):
         return self._modin_frame.to_pandas()
 
diff --git a/modin/experimental/core/storage_formats/hdk/query_compiler.py b/modin/experimental/core/storage_formats/hdk/query_compiler.py
--- a/modin/experimental/core/storage_formats/hdk/query_compiler.py
+++ b/modin/experimental/core/storage_formats/hdk/query_compiler.py
@@ -183,6 +183,14 @@ def finalize(self):
         # TODO: implement this for HDK storage format
         raise NotImplementedError()
 
+    def execute(self):
+        self._modin_frame._execute()
+
+    def force_import(self):
+        """Force table import."""
+        # HDK-specific method
+        self._modin_frame.force_import()
+
     def to_pandas(self):
         return self._modin_frame.to_pandas()
 
diff --git a/modin/utils.py b/modin/utils.py
--- a/modin/utils.py
+++ b/modin/utils.py
@@ -27,6 +27,7 @@
 from typing import (
     Any,
     Callable,
+    Iterable,
     List,
     Mapping,
     Optional,
@@ -606,6 +607,27 @@ def try_cast_to_pandas(obj: Any, squeeze: bool = False) -> Any:
     return obj
 
 
+def execute(*objs: Iterable[Any], trigger_hdk_import: bool = False) -> None:
+    """
+    Trigger the lazy computations for each obj in `objs`, if any, and wait for them to complete.
+
+    Parameters
+    ----------
+    *objs : Iterable[Any]
+        A collection of objects to trigger lazy computations.
+    trigger_hdk_import : bool, default: False
+        Trigger import execution. Makes sense only for HDK storage format.
+        Safe to use with other storage formats.
+    """
+    for obj in objs:
+        if not hasattr(obj, "_query_compiler"):
+            continue
+        query_compiler = obj._query_compiler
+        query_compiler.execute()
+        if trigger_hdk_import and hasattr(query_compiler, "force_import"):
+            query_compiler.force_import()
+
+
 def wrap_into_list(*args: Any, skipna: bool = True) -> List[Any]:
     """
     Wrap a sequence of passed values in a flattened list.


Test patch:
diff --git a/modin/test/test_executions_api.py b/modin/test/test_executions_api.py
--- a/modin/test/test_executions_api.py
+++ b/modin/test/test_executions_api.py
@@ -25,6 +25,7 @@ def test_base_abstract_methods():
         "__init__",
         "free",
         "finalize",
+        "execute",
         "to_pandas",
         "from_pandas",
         "from_arrow",
diff --git a/modin/test/test_utils.py b/modin/test/test_utils.py
--- a/modin/test/test_utils.py
+++ b/modin/test/test_utils.py
@@ -13,7 +13,9 @@
 
 import json
 from textwrap import dedent, indent
+from unittest.mock import Mock, patch
 
+import numpy as np
 import pandas
 import pytest
 
@@ -336,3 +338,33 @@ def test_assert_dtypes_equal():
     df1 = df1.astype({"a": "str"})
     with pytest.raises(AssertionError):
         assert_dtypes_equal(df1, df2)
+
+
+def test_execute():
+    data = np.random.rand(100, 64)
+    modin_df, pandas_df = create_test_dfs(data)
+    partitions = modin_df._query_compiler._modin_frame._partitions.flatten()
+    mgr_cls = modin_df._query_compiler._modin_frame._partition_mgr_cls
+
+    # check modin case
+    with patch.object(mgr_cls, "wait_partitions", new=Mock()):
+        modin.utils.execute(modin_df)
+        mgr_cls.wait_partitions.assert_called_once()
+        assert (mgr_cls.wait_partitions.call_args[0] == partitions).all()
+
+    # check pandas case without error
+    with patch.object(mgr_cls, "wait_partitions", new=Mock()):
+        modin.utils.execute(pandas_df)
+        mgr_cls.wait_partitions.assert_not_called()
+
+    # muke sure `trigger_hdk_import=True` doesn't broke anything
+    # when using other storage formats
+    with patch.object(mgr_cls, "wait_partitions", new=Mock()):
+        modin.utils.execute(modin_df, trigger_hdk_import=True)
+        mgr_cls.wait_partitions.assert_called_once()
+
+    # check several modin dataframes
+    with patch.object(mgr_cls, "wait_partitions", new=Mock()):
+        modin.utils.execute(modin_df, modin_df[modin_df.columns[:4]])
+        mgr_cls.wait_partitions.assert_called
+        assert mgr_cls.wait_partitions.call_count == 2