Repository: modin-project/modin

Issue:
FEAT: Add 'auto' mode to the lazy execution.
The `MODIN_LAZY_EXECUTION` env var should have 3 options:
`Auto` - the execution mode is chosen by the engine for each operation (default value).
`On`   - the lazy execution is performed wherever it's possible.
`Off`  - the lazy execution is disabled.


Discussion:

Resolving patch:
diff --git a/modin/config/envvars.py b/modin/config/envvars.py
--- a/modin/config/envvars.py
+++ b/modin/config/envvars.py
@@ -816,11 +816,19 @@ class ReadSqlEngine(EnvironmentVariable, type=str):
     choices = ("Pandas", "Connectorx")
 
 
-class LazyExecution(EnvironmentVariable, type=bool):
-    """Prefer the lazy execution, when it's possible."""
+class LazyExecution(EnvironmentVariable, type=str):
+    """
+    Lazy execution mode.
+
+    Supported values:
+        `Auto` - the execution mode is chosen by the engine for each operation (default value).
+        `On`   - the lazy execution is performed wherever it's possible.
+        `Off`  - the lazy execution is disabled.
+    """
 
     varname = "MODIN_LAZY_EXECUTION"
-    default = False
+    choices = ("Auto", "On", "Off")
+    default = "Auto"
 
 
 class DaskThreadsPerWorker(EnvironmentVariable, type=int):
diff --git a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
--- a/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
+++ b/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/partition.py
@@ -125,11 +125,9 @@ def apply(self, func: Callable, *args, **kwargs):
         If ``LazyExecution`` is enabled, the function is not applied immediately,
         but is added to the execution tree.
         """
-        de = DeferredExecution(self._data_ref, func, args, kwargs)
-        if LazyExecution.get():
-            return self.__constructor__(de)
         log = get_logger()
         self._is_debug(log) and log.debug(f"ENTER::Partition.apply::{self._identity}")
+        de = DeferredExecution(self._data_ref, func, args, kwargs)
         data, meta, meta_offset = de.exec()
         self._is_debug(log) and log.debug(f"EXIT::Partition.apply::{self._identity}")
         return self.__constructor__(data, meta=meta, meta_offset=meta_offset)
@@ -377,3 +375,45 @@ def _get_index_and_columns(df):  # pragma: no cover
         The number of columns.
     """
     return len(df.index), len(df.columns)
+
+
+PandasOnRayDataframePartition._eager_exec_func = PandasOnRayDataframePartition.apply
+PandasOnRayDataframePartition._lazy_exec_func = (
+    PandasOnRayDataframePartition.add_to_apply_calls
+)
+
+
+def _configure_lazy_exec(cls: LazyExecution):
+    """Configure lazy execution mode for PandasOnRayDataframePartition."""
+    mode = cls.get()
+    get_logger().debug(f"Ray lazy execution mode: {mode}")
+    if mode == "Auto":
+        PandasOnRayDataframePartition.apply = (
+            PandasOnRayDataframePartition._eager_exec_func
+        )
+        PandasOnRayDataframePartition.add_to_apply_calls = (
+            PandasOnRayDataframePartition._lazy_exec_func
+        )
+    elif mode == "On":
+
+        def lazy_exec(self, func, *args, **kwargs):
+            return self._lazy_exec_func(func, *args, length=None, width=None, **kwargs)
+
+        PandasOnRayDataframePartition.apply = lazy_exec
+        PandasOnRayDataframePartition.add_to_apply_calls = (
+            PandasOnRayDataframePartition._lazy_exec_func
+        )
+    elif mode == "Off":
+
+        def eager_exec(self, func, *args, length=None, width=None, **kwargs):
+            return self._eager_exec_func(func, *args, **kwargs)
+
+        PandasOnRayDataframePartition.apply = (
+            PandasOnRayDataframePartition._eager_exec_func
+        )
+        PandasOnRayDataframePartition.add_to_apply_calls = eager_exec
+    else:
+        raise ValueError(f"Invalid lazy execution mode: {mode}")
+
+
+LazyExecution.subscribe(_configure_lazy_exec)


Test patch:
diff --git a/modin/test/storage_formats/pandas/test_internals.py b/modin/test/storage_formats/pandas/test_internals.py
--- a/modin/test/storage_formats/pandas/test_internals.py
+++ b/modin/test/storage_formats/pandas/test_internals.py
@@ -2421,3 +2421,50 @@ def test_groupby_index_dtype(self):
             assert res_dtypes._known_dtypes["a"] == np.dtype("int64")
 
         patch.assert_not_called()
+
+
+@pytest.mark.skipif(Engine.get() != "Ray", reason="Ray specific")
+@pytest.mark.parametrize("mode", [None, "Auto", "On", "Off"])
+def test_ray_lazy_exec_mode(mode):
+    import ray
+
+    from modin.config import LazyExecution
+    from modin.core.execution.ray.common.deferred_execution import DeferredExecution
+    from modin.core.execution.ray.common.utils import ObjectIDType
+    from modin.core.execution.ray.implementations.pandas_on_ray.partitioning import (
+        PandasOnRayDataframePartition,
+    )
+
+    orig_mode = LazyExecution.get()
+    try:
+        if mode is None:
+            mode = LazyExecution.get()
+        else:
+            LazyExecution.put(mode)
+            assert mode == LazyExecution.get()
+
+        df = pandas.DataFrame({"A": [1, 2, 3]})
+        part = PandasOnRayDataframePartition(ray.put(df))
+
+        def func(df):
+            return len(df)
+
+        ray_func = ray.put(func)
+
+        if mode == "Auto":
+            assert isinstance(part.apply(ray_func)._data_ref, ObjectIDType)
+            assert isinstance(
+                part.add_to_apply_calls(ray_func)._data_ref, DeferredExecution
+            )
+        elif mode == "On":
+            assert isinstance(part.apply(ray_func)._data_ref, DeferredExecution)
+            assert isinstance(
+                part.add_to_apply_calls(ray_func)._data_ref, DeferredExecution
+            )
+        elif mode == "Off":
+            assert isinstance(part.apply(ray_func)._data_ref, ObjectIDType)
+            assert isinstance(part.add_to_apply_calls(ray_func)._data_ref, ObjectIDType)
+        else:
+            pytest.fail(f"Invalid value: {mode}")
+    finally:
+        LazyExecution.put(orig_mode)