Repository: modin-project/modin (python)

Issue:
Using groupby while AutoSwitchBackend is on tries to pin the backend of DataFrameGroupy, raising NotImplementedError

Pull request:
FIX-#7553: Fix groupby when AutoSwitchBackend is disabled.
Prior to this commit, calling groupby() while `AutoSwitchBackend` was off would raise `NotImplementedError`. Previously, we would try to pin the result of a QueryCompilerCaster operation whenever `AutoSwitchBackend` was off, whereas we only need to pin when one of the inputs is pinned.

Resolves #7553


Patch:
diff --git a/modin/core/storage_formats/pandas/query_compiler_caster.py b/modin/core/storage_formats/pandas/query_compiler_caster.py
index a7bf3a19..3027d238 100644
--- a/modin/core/storage_formats/pandas/query_compiler_caster.py
+++ b/modin/core/storage_formats/pandas/query_compiler_caster.py
@@ -807,9 +807,10 @@ def wrap_function_in_argument_caster(
             input_qc_for_pre_op_switch = input_query_compilers[0]
             input_backend = input_qc_for_pre_op_switch.get_backend()
 
-        if not AutoSwitchBackend.get() or (
+        inputs_pinned = (
             len(input_query_compilers) < 2 and pin_target_backend is not None
-        ):
+        )
+        if not AutoSwitchBackend.get() or inputs_pinned:
             f_to_apply = _get_extension_for_method(
                 name=name,
                 extensions=extensions,
@@ -822,7 +823,7 @@ def wrap_function_in_argument_caster(
                 wrapping_function_type=wrapping_function_type,
             )
             result = f_to_apply(*args, **kwargs)
-            if isinstance(result, QueryCompilerCaster):
+            if isinstance(result, QueryCompilerCaster) and inputs_pinned:
                 result._set_backend_pinned(True, inplace=True)
             return result
 


Test patch:
diff --git a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
index 1f6ada74..688a5685 100644
--- a/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
+++ b/modin/tests/pandas/native_df_interoperability/test_compiler_caster.py
@@ -793,6 +793,63 @@ class TestSwitchBackendPostOpDependingOnDataSize:
             assert modin_result.get_backend() == expected_backend
             assert modin_df.get_backend() == "Big_Data_Cloud"
 
+    @pytest.mark.parametrize(
+        "groupby_class,operation",
+        [
+            param(
+                "DataFrameGroupBy",
+                lambda groupby: groupby.sum(),
+                id="DataFrameGroupBy",
+            ),
+            param(
+                "SeriesGroupBy",
+                lambda groupby: groupby["col1"].sum(),
+                id="SeriesGroupBy",
+            ),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "auto_switch_backend",
+        [True, False],
+        ids=lambda param: f"auto_switch_backend_{param}",
+    )
+    def test_auto_switch_config_can_disable_groupby_agg_auto_switch(
+        self,
+        operation,
+        groupby_class,
+        auto_switch_backend,
+    ):
+        num_groups = BIG_DATA_CLOUD_MIN_NUM_ROWS - 1
+        with backend_test_context(
+            test_backend="Big_Data_Cloud",
+            choices=("Big_Data_Cloud", "Small_Data_Local"),
+        ), config_context(AutoSwitchBackend=auto_switch_backend):
+            modin_groupby, pandas_groupby = (
+                df.groupby("col0")
+                for df in create_test_dfs(
+                    {
+                        "col0": list(range(num_groups)),
+                        "col1": list(range(1, num_groups + 1)),
+                    }
+                )
+            )
+
+            assert modin_groupby.get_backend() == "Big_Data_Cloud"
+            assert operation(modin_groupby).get_backend() == "Big_Data_Cloud"
+
+            register_function_for_post_op_switch(
+                class_name=groupby_class, backend="Big_Data_Cloud", method="sum"
+            )
+
+            assert modin_groupby.get_backend() == "Big_Data_Cloud"
+            modin_result = operation(modin_groupby)
+            pandas_result = operation(pandas_groupby)
+            df_equals(modin_result, pandas_result)
+            assert modin_result.get_backend() == (
+                "Small_Data_Local" if auto_switch_backend else "Big_Data_Cloud"
+            )
+            assert modin_groupby.get_backend() == "Big_Data_Cloud"
+
 
 class TestSwitchBackendPreOp:
     @pytest.mark.parametrize(
@@ -1200,19 +1257,3 @@ def test_concat_with_pin(pin_backends, expected_backend):
             df_equals(
                 result, pandas.concat([pandas.DataFrame([1] * 10)] * len(pin_backends))
             )
-
-
-def test_groupby_pinned():
-    groupby = pd.DataFrame([1, 2]).groupby(0)
-    assert not groupby.is_backend_pinned()
-
-
-@pytest.mark.xfail(strict=True, raises=NotImplementedError)
-def test_groupby_pin_backend():
-    groupby = pd.DataFrame([1, 2]).groupby(0)
-    groupby.pin_backend()
-
-
-@pytest.mark.xfail(strict=True, raises=NotImplementedError)
-def test_groupby_set_backend_pinned():
-    pd.DataFrame([1, 2]).groupby(0)._set_backend_pinned(inplace=False, pinned=True)