Repository: dask/dask
Pull request #11784 (2025-02-25 16:47:35)

Ensure concat optimize project does not raise
Our way of subclassing isn't working well. The subclasses can not/should not define a different set of paramters since otherwise all logic on the base class will just break. In this case it can be provoked by adding a projection after an optimize call

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ± 0        9 suites  ±0   3h 31m 3s ⏱️ -30s
 17 803 tests + 4   16 589 ✅ + 4   1 214 💤 ±0  0 ❌ ±0 
159 301 runs  +36  147 144 ✅ +34  12 157 💤 +2  0 ❌ ±0 

Results for commit 78714485. ± Comparison against base commit 06c1bb72.

[test-results]:data:application/gzip;base64,H4sIALP9vWcC/02Myw6DIBBFf8Ww7oLH4EB/plGUhFSlQVg1/fcisVOX59zc82Y+LPPO7p29dWwvIRNMJQ05xK2ikH2vqqpjPmaBhhM/9uLcIXtt7F8+w6s9BZDzQ1iq4yTmlGI6TSpbS2uruDiZ0oACgCSlNf7cJd34WnZxXUOuwNAcHaO59uPAlR0nLxGlUm60AJJLDYBu8uzzBSimrf8YAQAA


Patch:
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index d89f00a45..7aa07d4a4 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -5519,10 +5519,10 @@ def concat(
         Concat(
             join,
             ignore_order,
-            kwargs,
             axis,
             ignore_unknown_divisions,
             interleave_partitions,
+            kwargs,
             *dfs,
         )
     )
diff --git a/dask/dataframe/dask_expr/_concat.py b/dask/dataframe/dask_expr/_concat.py
index 7e9eec52d..c96287132 100644
--- a/dask/dataframe/dask_expr/_concat.py
+++ b/dask/dataframe/dask_expr/_concat.py
@@ -28,10 +28,10 @@ class Concat(Expr):
     _parameters = [
         "join",
         "ignore_order",
-        "_kwargs",
         "axis",
         "ignore_unknown_divisions",
         "interleave_partitions",
+        "_kwargs",
     ]
     _defaults = {
         "join": "outer",
@@ -167,6 +167,8 @@ class Concat(Expr):
                     self.join,
                     self.ignore_order,
                     self.axis,
+                    self.ignore_unknown_divisions,
+                    self.interleave_partitions,
                     self._kwargs,
                     *cast_dfs,
                 )
@@ -220,6 +222,8 @@ class Concat(Expr):
                 self.join,
                 self.ignore_order,
                 self.axis,
+                self.ignore_unknown_divisions,
+                self.interleave_partitions,
                 self._kwargs,
                 *cast_dfs,
             )
@@ -228,6 +232,8 @@ class Concat(Expr):
             self.join,
             self.ignore_order,
             self.axis,
+            self.ignore_unknown_divisions,
+            self.interleave_partitions,
             self._kwargs,
             *cast_dfs,
         )
@@ -264,12 +270,13 @@ class Concat(Expr):
             result = type(self)(
                 self.join,
                 self.ignore_order,
-                self._kwargs,
                 self.axis,
                 self.ignore_unknown_divisions,
                 self.interleave_partitions,
+                self._kwargs,
                 *frames,
             )
+
             if result.columns == _convert_to_list(parent.operand("columns")):
                 if result.ndim == parent.ndim:
                     return result
@@ -280,8 +287,6 @@ class Concat(Expr):
 
 
 class StackPartition(Concat):
-    _parameters = ["join", "ignore_order", "axis", "_kwargs"]
-    _defaults = {"join": "outer", "ignore_order": False, "_kwargs": {}, "axis": 0}
 
     def _layer(self):
         dsk, i = {}, 0


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_concat.py b/dask/dataframe/dask_expr/tests/test_concat.py
index eb05a403f..f9afd8fd6 100644
--- a/dask/dataframe/dask_expr/tests/test_concat.py
+++ b/dask/dataframe/dask_expr/tests/test_concat.py
@@ -380,3 +380,20 @@ def test_concat_mixed_dtype_columns():
     result = concat([df1, df2]).drop(["_y"], axis=1)
     expected = pd.concat([df1.compute(), df2.compute()]).drop(columns="_y")
     assert_eq(result, expected)
+
+
+@pytest.mark.parametrize("axis", [0, 1])
+@pytest.mark.parametrize("interleave_partitions", [True, False])
+def test_concat_optimize_project(axis, interleave_partitions):
+    df1 = from_dict({"a": range(10), "b": range(10)}, npartitions=2)
+    df2 = from_dict({"c": [5, 2, 3] + list(range(7))}, npartitions=3)
+    df2.clear_divisions()
+    concatenated = concat(
+        [df1, df2], axis=axis, interleave_partitions=interleave_partitions
+    )
+    optimized = concatenated.optimize()
+    optimized_project = optimized[["a", "c"]]
+    assert_eq(optimized_project, concatenated[["a", "c"]])
+    assert (
+        optimized_project.optimize()._name == concatenated[["a", "c"]].optimize()._name
+    )