Repository: dask/dask
Pull request #11887 (2025-04-14 11:52:44)

Fix persisting multiple DFs at the same time
Closes https://github.com/dask/dask/issues/11885


The sequence class is somewhat new and with 2025.3.0 and earlier, whenever compute or persist would encounter multiple collections they would just be optimized independently. The sequence class changes this now. However, this also causes us to mix `Expr` subclasses.
This is typically not a problem because the sequence basically seals/finalizes the expression on its head (inserting an ExprSequence somewhere in the tree of expressions is not supported). Therefore, as long as all expression look at their children, they won't notice a change.

The `simplify_up` (and `tune_up`) operation is looking at their parent, i.e. a DF expression may actually be able to see the sequence. Similarly, when we pass dependents to the expressions, this could also occur.

All `simplify_up`s are essentially switch statements that select the code to be run based on the instance type encountered. That's safe. The dependents computation (which is, so far, exlusively used in the DF land) is not robust to this. This is what this change is doing.


On top of this, there was also a bug in the persist code since it was using the keys of the unoptimized collections instead of the optimized expr. 

cc @TomAugspurger 

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 22m 20s ⏱️ + 7m 26s
 18 010 tests +1   16 796 ✅ +1   1 214 💤 ±0  0 ❌ ±0 
161 164 runs  +9  149 058 ✅ +9  12 106 💤 ±0  0 ❌ ±0 

Results for commit ffde7292. ± Comparison against base commit 0fa5e18d.

[test-results]:data:application/gzip;base64,H4sIACkA/WcC/02Myw6DIBQFf8Ww7oKLcIH+TCMCCamPBmFl+u9FQ9HlzMmZnfgwuY08O/3oyJZDamBzHFJYl4LAgNOiypiOGRSFxq8tj+MhUWq85Dt86rM5P4SpuOvpYlxjNTEvZxoBkFduaa6pUE22NMW/u6VPvpfHdZ5DKkC8t04yzQwIiwJ6YXoEYSw3zijF7SChZ4Zx8v0BcYwiShgBAAA=


Patch:
diff --git a/dask/base.py b/dask/base.py
index 020523c4c..ba2a04922 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -1011,8 +1011,8 @@ def persist(*args, traverse=True, optimize_graph=True, scheduler=None, **kwargs)
     expr = collections_to_expr(collections, optimize_graph)
     expr = expr.optimize()
     keys, postpersists = [], []
-    for a in collections:
-        a_keys = list(flatten(a.__dask_keys__()))
+    for a, akeys in zip(collections, expr.__dask_keys__(), strict=True):
+        a_keys = list(flatten(akeys))
         rebuild, state = a.__dask_postpersist__()
         keys.extend(a_keys)
         postpersists.append((rebuild, a_keys, state))
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 13dd2e467..75264acc9 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -5,8 +5,9 @@ import functools
 import numbers
 import operator
 import warnings
+import weakref
 from collections import defaultdict
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Collection, Mapping
 from typing import Any as AnyType
 
 import numpy as np
@@ -15,6 +16,7 @@ from pandas.errors import PerformanceWarning
 from tlz import merge_sorted, partition, unique
 
 from dask import _expr as core
+from dask._expr import Expr as BaseExpr
 from dask._expr import FinalizeCompute
 from dask._task_spec import Alias, DataNode, Task, TaskRef, execute_graph
 from dask.array import Array
@@ -3804,19 +3806,24 @@ class MinType:
         return True
 
 
-def determine_column_projection(expr, parent, dependents, additional_columns=None):
+def determine_column_projection(
+    expr: Expr,
+    parent: Expr,
+    dependents: dict[str, Collection[weakref.ref[BaseExpr]]],
+    additional_columns: list | None = None,
+) -> object:
     if isinstance(parent, Index):
         column_union = []
     else:
         column_union = parent.columns.copy()
-    parents = [x() for x in dependents[expr._name] if x() is not None]
+    parents: list[Expr]
+    parents = [inst for x in dependents[expr._name] if isinstance((inst := x()), Expr)]
 
     seen = set()
     for p in parents:
         if p._name in seen:
             continue
         seen.add(p._name)
-
         column_union.extend(p._projection_columns)
 
     if additional_columns is not None:


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_core.py b/dask/dataframe/dask_expr/tests/test_core.py
index b564690b9..74b9c9d57 100644
--- a/dask/dataframe/dask_expr/tests/test_core.py
+++ b/dask/dataframe/dask_expr/tests/test_core.py
@@ -3,9 +3,13 @@ from __future__ import annotations
 import pickle
 import uuid
 
+import pandas as pd
 import pytest
 
+import dask
+import dask.dataframe as dd
 from dask._expr import Expr
+from dask.dataframe.utils import assert_eq
 from dask.tokenize import tokenize
 
 
@@ -47,3 +51,15 @@ def test_expr_nondeterministic_token_pickle_roundtrip():
     assert tokenize(inst) == tok
     rt = pickle.loads(pickle.dumps(inst))
     assert tokenize(rt) == tok
+
+
+def test_combine_multiple_dataframes():
+    df = pd.DataFrame({"A": [1, 2, 3, 4], "B": [0, 0, 1, 1]})
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    X = ddf[["A", "B"]]
+    y = ddf[["B"]]
+
+    x_per, y_per = dask.persist(X, y)
+    assert_eq(x_per, X)
+    assert_eq(y_per, y)