add new methods for internals

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

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)
+
     def __dataframe__(self, nan_as_null: bool = False, allow_copy: bool = True):
         """
         Get a Modin DataFrame that implements the dataframe exchange protocol.

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

@@ -74,7 +74,7 @@ def wait(cls, *args, **kwargs):
             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())
+            cls.wait_partitions(partitions)
         return result
 
     return wait
@@ -930,6 +930,7 @@ def wait_partitions(cls, partitions):
         This method should be implemented in a more efficient way for engines that supports
         waiting on objects in parallel.
         """
+        partitions = partitions.flatten()
         for partition in partitions:
             partition.wait()
 

modin/core/execution/dask/implementations/pandas_on_dask/partitioning/partition_manager.py

@@ -46,6 +46,7 @@ def wait_partitions(cls, partitions):
         partitions : np.ndarray
             NumPy array with ``PandasDataframePartition``-s.
         """
+        partitions = partitions.flatten()
         DaskWrapper.wait(
             [block for partition in partitions for block in partition.list_of_blocks]
         )

modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition_manager.py

@@ -47,6 +47,7 @@ def wait_partitions(cls, partitions):
         partitions : np.ndarray
             NumPy array with ``PandasDataframePartition``-s.
         """
+        partitions = partitions.flatten()
         RayWrapper.wait(
             [block for partition in partitions for block in partition.list_of_blocks]
         )

modin/core/execution/unidist/implementations/pandas_on_unidist/partitioning/partition_manager.py

@@ -47,6 +47,7 @@ def wait_partitions(cls, partitions):
         partitions : np.ndarray
             NumPy array with ``PandasDataframePartition``-s.
         """
+        partitions = partitions.flatten()
         UnidistWrapper.wait(
             [block for partition in partitions for block in partition.list_of_blocks]
         )

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 wait_computations(self):
+        """Wait for all computations to complete without materializing data."""
+        pass
+
     # END Data Management Methods
 
     # To/From Pandas

modin/core/storage_formats/pandas/query_compiler.py

@@ -278,6 +278,9 @@ def lazy_execution(self):
     def finalize(self):
         self._modin_frame.finalize()
 
+    def wait_computations(self):
+        self._modin_frame.wait_computations()
+
     def to_pandas(self):
         return self._modin_frame.to_pandas()
 

modin/experimental/core/storage_formats/hdk/query_compiler.py

@@ -183,6 +183,9 @@ def finalize(self):
         # TODO: implement this for HDK storage format
         raise NotImplementedError()
 
+    def wait_computations(self):
+        raise NotImplementedError()
+
     def to_pandas(self):
         return self._modin_frame.to_pandas()
 

modin/test/test_executions_api.py

@@ -25,6 +25,7 @@ def test_base_abstract_methods():
         "__init__",
         "free",
         "finalize",
+        "wait_computations",
         "to_pandas",
         "from_pandas",
         "from_arrow",

modin/utils.py

@@ -606,23 +606,32 @@ def try_cast_to_pandas(obj: Any, squeeze: bool = False) -> Any:
     return obj
 
 
-def trigger_import(df: Any):
+def trigger_import(obj: Any) -> None:
     """
     Trigger import execution for DataFrames obtained by HDK engine.
 
     Parameters
     ----------
-    df : Any
-        DataFrames to trigger import.
+    obj : Any
+        An object to trigger deferred data import.
     """
-    if hasattr(df, "_query_compiler"):
-        modin_frame = df._query_compiler._modin_frame
+    if hasattr(obj, "_query_compiler"):
+        modin_frame = obj._query_compiler._modin_frame
         if hasattr(modin_frame, "force_import"):
             modin_frame.force_import()
 
 
-def wait_computations(obj: Any, *, trigger_hdk_import: bool = False):
-    """Make sure obj's computations finished."""
+def wait_computations(obj: Any, *, trigger_hdk_import: bool = False) -> None:
+    """
+    Trigger the `obj` lazy computations, if any, and wait for them to complete.
+
+    Parameters
+    ----------
+    obj : Any
+        An object to trigger lazy computations.
+    trigger_hdk_import : bool, default: False
+        Trigger import execution. Makes sense only for HDK storage format.
+    """
     if not hasattr(obj, "_query_compiler"):
         return
 
@@ -634,11 +643,7 @@ def wait_computations(obj: Any, *, trigger_hdk_import: bool = False):
         return
 
     obj._query_compiler.finalize()
-
-    partitions = obj._query_compiler._modin_frame._partitions.flatten()
-    mgr_cls = obj._query_compiler._modin_frame._partition_mgr_cls
-    if len(partitions) and hasattr(mgr_cls, "wait_partitions"):
-        mgr_cls.wait_partitions(partitions)
+    obj._query_compiler.wait_computations()
 
 
 def wrap_into_list(*args: Any, skipna: bool = True) -> List[Any]: