Repository: dask-contrib/dask-expr
Pull request #935 (2024-03-05 11:54:30)

Fix to_dask_dataframe test after switching to dask-expr by default


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/io/tests/test_io.py b/dask_expr/io/tests/test_io.py
index d3905b6..ccf0cd0 100644
--- a/dask_expr/io/tests/test_io.py
+++ b/dask_expr/io/tests/test_io.py
@@ -225,7 +225,7 @@ def test_to_dask_dataframe(optimize):
     pdf = pd.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
     df = from_pandas(pdf, npartitions=2)
     ddf = df.to_dask_dataframe(optimize=optimize)
-    assert isinstance(ddf, dd.DataFrame)
+    assert isinstance(ddf, dd.core.DataFrame)
     assert_eq(df, ddf)
 
 
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(