Repository: modin-project/modin (python)

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.

Pull request description:
FEAT-#5221: add `execute` to trigger lazy computations and wait for them to complete
<!--
Thank you for your contribution!
Please review the contributing docs: https://modin.readthedocs.io/en/latest/development/contributing.html
if you have questions about contributing.
-->

## What do these changes do?

<!-- Please give a short brief about these changes. -->

- [x] first commit message and PR title follow format outlined [here](https://modin.readthedocs.io/en/latest/development/contributing.html#commit-message-formatting)
  > **_NOTE:_**  If you edit the PR title to match this format, you need to add another commit (even if it's empty) or amend your last commit for the CI job that checks the PR title to pick up the new PR title.
- [x] passes `flake8 modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] passes `black --check modin/ asv_bench/benchmarks scripts/doc_checker.py`
- [x] signed commit with `git commit -s` <!-- you can amend your commit with a signature via `git commit -amend -s` -->
- [x] Resolves #5221 <!-- issue must be created for each patch -->
- [x] tests added and passing
- [x] module layout described at `docs/development/architecture.rst` is up-to-date <!-- if you have added, renamed or removed files or directories please update the documentation accordingly -->


Patch:
diff --git a/docs/development/architecture.rst b/docs/development/architecture.rst
index 130813a8..22fc7f91 100644
--- a/docs/development/architecture.rst
+++ b/docs/development/architecture.rst
@@ -304,6 +304,7 @@ details. The documentation covers most modules, with more docs being added every
    ├───examples
    ├───modin
    │   ├─── :doc:`config </flow/modin/config>`
+   |   ├─── :doc:`utils </flow/modin/utils>`
    │   ├───core
    │   │   ├─── :doc:`dataframe </flow/modin/core/dataframe/index>`
    │   │   │   ├─── :doc:`algebra </flow/modin/core/dataframe/algebra>`
diff --git a/docs/flow/modin/utils.rst b/docs/flow/modin/utils.rst
new file mode 100644
index 00000000..70c622d1
--- /dev/null
+++ b/docs/flow/modin/utils.rst
@@ -0,0 +1,12 @@
+:orphan:
+
+Modin Utils
+"""""""""""
+
+Here are utilities that can be useful when working with Modin.
+
+Public API
+''''''''''
+
+.. autofunction:: modin.utils.try_cast_to_pandas
+.. autofunction:: modin.utils.execute
diff --git a/docs/usage_guide/benchmarking.rst b/docs/usage_guide/benchmarking.rst
index 7340a31d..551c9950 100644
--- a/docs/usage_guide/benchmarking.rst
+++ b/docs/usage_guide/benchmarking.rst
@@ -8,8 +8,8 @@ To benchmark a single Modin function, often turning on the
 :code:`BenchmarkMode` will suffice.
 
 There is no simple way to benchmark more complex Modin workflows, though
-benchmark mode or calling ``repr`` on Modin objects may be useful. The
-:doc:`Modin logs </usage_guide/advanced_usage/modin_logging>` may help you
+benchmark mode or calling ``modin.utils.execute`` on Modin objects may be useful.
+The :doc:`Modin logs </usage_guide/advanced_usage/modin_logging>` may help you
 identify bottlenecks in your code, and they may also help profile the execution
 of each Modin function.
 
@@ -125,9 +125,8 @@ at each Modin :doc:`layer </development/architecture>`. Log mode is more
 useful when used in conjuction with benchmark mode.
 
 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. That will typically
-block on any asynchronous computation. However, beware that ``repr`` can also
-be misleading, e.g. here:
+just call ``modin.utils.execute`` on the workflow's final Modin objects.
+That will typically block on any asynchronous computation:
 
 .. code-block:: python
 
@@ -137,6 +136,7 @@ be misleading, e.g. here:
 
     import modin.pandas as pd
     from modin.config import MinPartitionSize, NPartitions
+    import modin.utils
 
     MinPartitionSize.put(32)
     NPartitions.put(16)
@@ -149,17 +149,13 @@ be misleading, e.g. here:
     ray.init()
     df1 = pd.DataFrame(list(range(10_000)), columns=['A'])
     result = df1.map(slow_add_one)
-    %time repr(result)
-    # time.sleep(10)
+    # %time modin.utils.execute(result)
     %time result.to_parquet(BytesIO())
 .. code-block::python
 
-The ``repr`` takes only 802 milliseconds, but writing the result to a buffer
-takes 9.84 seconds. However, if you uncomment the :code:`time.sleep` before the
-:code:`to_parquet` call, the :code:`to_parquet` takes just 23.8 milliseconds!
-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.
+Writing the result to a buffer takes 9.84 seconds. However, if you uncomment
+the :code:`%time modin.utils.execute(result)` before the :code:`to_parquet`
+call, the :code:`to_parquet` takes just 23.8 milliseconds!
 
 .. note::
     If you see any Modin documentation touting Modin's speed without using
diff --git a/modin/core/dataframe/pandas/dataframe/dataframe.py b/modin/core/dataframe/pandas/dataframe/dataframe.py
index eb936991..5d2b32d9 100644
--- a/modin/core/dataframe/pandas/dataframe/dataframe.py
+++ b/modin/core/dataframe/pandas/dataframe/dataframe.py
@@ -4106,6 +4106,10 @@ class PandasDataframe(ClassLogger):
         """
         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
index ec31f98b..3bf0e06f 100644
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -317,6 +317,11 @@ class BaseQueryCompiler(ClassLogger, abc.ABC):
         """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
index 0fc992bb..eb7c2e80 100644
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -278,6 +278,10 @@ class PandasQueryCompiler(BaseQueryCompiler):
     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
index 439af624..b4a58047 100644
--- a/modin/experimental/core/storage_formats/hdk/query_compiler.py
+++ b/modin/experimental/core/storage_formats/hdk/query_compiler.py
@@ -183,6 +183,14 @@ class DFAlgQueryCompiler(BaseQueryCompiler):
         # 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/pandas/groupby.py b/modin/pandas/groupby.py
index c50e97a5..caa3086e 100644
--- a/modin/pandas/groupby.py
+++ b/modin/pandas/groupby.py
@@ -659,16 +659,20 @@ class DataFrameGroupBy(ClassLogger):
         if not isinstance(func, BuiltinFunctionType):
             func = wrap_udf_function(func)
 
-        return self._check_index(
-            self._wrap_aggregation(
-                qc_method=type(self._query_compiler).groupby_agg,
-                numeric_only=False,
-                agg_func=func,
-                agg_args=args,
-                agg_kwargs=kwargs,
-                how="group_wise",
-            )
+        apply_res = self._wrap_aggregation(
+            qc_method=type(self._query_compiler).groupby_agg,
+            numeric_only=False,
+            agg_func=func,
+            agg_args=args,
+            agg_kwargs=kwargs,
+            how="group_wise",
         )
+        reduced_index = pandas.Index([MODIN_UNNAMED_SERIES_LABEL])
+        if not isinstance(apply_res, Series) and apply_res.columns.equals(
+            reduced_index
+        ):
+            apply_res = apply_res.squeeze(axis=1)
+        return self._check_index(apply_res)
 
     @property
     def dtypes(self):
diff --git a/modin/utils.py b/modin/utils.py
index 6aba77f5..c42a0516 100644
--- a/modin/utils.py
+++ b/modin/utils.py
@@ -27,6 +27,7 @@ from textwrap import dedent, indent
 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.
diff --git a/requirements/env_hdk.yml b/requirements/env_hdk.yml
index 701f475a..551ec920 100644
--- a/requirements/env_hdk.yml
+++ b/requirements/env_hdk.yml
@@ -7,7 +7,7 @@ dependencies:
   # required dependencies
   - pandas>=2.1,<2.2
   - numpy>=1.22.4
-  - pyhdk==0.8
+  - pyhdk==0.9
   - fsspec>=2022.05.0
   - packaging>=21.0
   - psutil>=5.8.0


Test patch:
diff --git a/modin/conftest.py b/modin/conftest.py
index 87cc0083..d860807e 100644
--- a/modin/conftest.py
+++ b/modin/conftest.py
@@ -175,6 +175,10 @@ class TestQC(BaseQueryCompiler):
     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/pandas/test/test_groupby.py b/modin/pandas/test/test_groupby.py
index 30ba2c11..3bbf649e 100644
--- a/modin/pandas/test/test_groupby.py
+++ b/modin/pandas/test/test_groupby.py
@@ -2962,6 +2962,23 @@ def test_reshuffling_groupby_on_strings(modify_config):
     )
 
 
+@pytest.mark.parametrize(
+    "modify_config", [{ExperimentalGroupbyImpl: True}], indirect=True
+)
+def test_groupby_apply_series_result(modify_config):
+    # reproducer from the issue:
+    # https://github.com/modin-project/modin/issues/6632
+    df = pd.DataFrame(
+        np.random.randint(5, 10, size=5), index=[f"s{i+1}" for i in range(5)]
+    )
+    df["group"] = [1, 1, 2, 2, 3]
+
+    # res = df.groupby('group').apply(lambda x: x.name+2)
+    eval_general(
+        df, df._to_pandas(), lambda df: df.groupby("group").apply(lambda x: x.name + 2)
+    )
+
+
 ### TEST GROUPBY WARNINGS ###
 
 
diff --git a/modin/test/test_executions_api.py b/modin/test/test_executions_api.py
index b60afd1a..35013132 100644
--- 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
index 3337385d..8f1b68a6 100644
--- 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