Repository: dask-contrib/dask-expr
Pull request #933 (2024-03-05 11:32:23)

Fix Copy-on-Write related bug in groupby.transform
pandas does a weird group key check here that becomes an issue if we consolidate during the shuffle. We already know that our group keys are not part of our original frame at this point, so we can simply clear the references and be done with it


Patch:
diff --git a/dask_expr/_expr.py b/dask_expr/_expr.py
index c17bef7..daccff0 100644
--- a/dask_expr/_expr.py
+++ b/dask_expr/_expr.py
@@ -1225,6 +1225,16 @@ class ColumnsSetter(RenameFrame):
         return _rename(columns, df)
 
 
+class _DeepCopy(Elemwise):
+    _parameters = ["frame"]
+    _projection_passthrough = True
+    _filter_passthrough = True
+
+    @staticmethod
+    def operation(df):
+        return df.copy(deep=True)
+
+
 class RenameSeries(Elemwise):
     _parameters = ["frame", "index", "sorted_index"]
     _defaults = {"sorted_index": False}
diff --git a/dask_expr/_groupby.py b/dask_expr/_groupby.py
index 4695ace..4b1e87c 100644
--- a/dask_expr/_groupby.py
+++ b/dask_expr/_groupby.py
@@ -62,6 +62,7 @@ from dask_expr._expr import (
     RenameFrame,
     RenameSeries,
     ToFrame,
+    _DeepCopy,
     _extract_meta,
     are_co_aligned,
     determine_column_projection,
@@ -920,11 +921,20 @@ class GroupByApply(Expr, GroupByBase):
                 if unmap_columns:
                     df = RenameFrame(df, unmap_columns)
 
+                # pandas checks if the group keys are part of the initial
+                # DataFrame through reference tracking. This is fine as long
+                # as we don't trigger a copy after the Assign above, since the
+                # blocks stay separate normally, disk shuffle triggers a copy
+                # though, and we shouldn't rely on those internals in pandas,
+                # so we can trigger a deep copy here to clear the references
+                # since we know more about the query than pandas does.
                 by = [
                     b
                     if not isinstance(b, Expr)
-                    else RenameSeries(
-                        Projection(df, f"_by_{i}"), index=self.by[i].columns[0]
+                    else _DeepCopy(
+                        RenameSeries(
+                            Projection(df, f"_by_{i}"), index=self.by[i].columns[0]
+                        )
                     )
                     for i, b in enumerate(self.by)
                 ]


Test patch:
diff --git a/dask_expr/tests/test_groupby.py b/dask_expr/tests/test_groupby.py
index d5d064d..1b96a32 100644
--- a/dask_expr/tests/test_groupby.py
+++ b/dask_expr/tests/test_groupby.py
@@ -506,6 +506,7 @@ def test_groupby_single_agg_split_out(pdf, df, api, sort, split_out):
     assert_eq(agg, expect, sort_results=not sort)
 
 
+@pytest.mark.parametrize("cow", [True, False])
 @pytest.mark.parametrize(
     "func",
     [
@@ -513,23 +514,24 @@ def test_groupby_single_agg_split_out(pdf, df, api, sort, split_out):
         lambda grouped: grouped.transform(lambda x: x.sum()),
     ],
 )
-def test_apply_or_transform_shuffle_multilevel(pdf, df, func):
-    grouper = lambda df: [df["x"] + 1, df["y"] + 1]
-
-    with pytest.warns(UserWarning):
-        # DataFrameGroupBy
-        assert_eq(func(df.groupby(grouper(df))), func(pdf.groupby(grouper(pdf))))
-
-        # SeriesGroupBy
-        assert_eq(
-            func(df.groupby(grouper(df))["z"]), func(pdf.groupby(grouper(pdf))["z"])
-        )
-
-        # DataFrameGroupBy with column slice
-        assert_eq(
-            func(df.groupby(grouper(df))[["z"]]),
-            func(pdf.groupby(grouper(pdf))[["z"]]),
-        )
+def test_apply_or_transform_shuffle_multilevel(pdf, df, func, cow):
+    with pd.option_context("mode.copy_on_write", cow):
+        grouper = lambda df: [df["x"] + 1, df["y"] + 1]
+
+        with pytest.warns(UserWarning):
+            # DataFrameGroupBy
+            assert_eq(func(df.groupby(grouper(df))), func(pdf.groupby(grouper(pdf))))
+
+            # SeriesGroupBy
+            assert_eq(
+                func(df.groupby(grouper(df))["z"]), func(pdf.groupby(grouper(pdf))["z"])
+            )
+
+            # DataFrameGroupBy with column slice
+            assert_eq(
+                func(df.groupby(grouper(df))[["z"]]),
+                func(pdf.groupby(grouper(pdf))[["z"]]),
+            )
 
 
 @pytest.mark.parametrize(