Repository: dask/dask
Pull request #11888 (2025-04-14 13:17:55)

Remove custom initializers for delayed expressions
See https://github.com/dask/dask/pull/11886 

My best working thesis right now is that these two classes are not just implementing an initializer (which I think should be fine, generally. After all, the entire "please do not mutate the expression since otherwise our singleton cache will be dirty" is true for the entire code base) but they are setting `operands` explicitly.

Let's look at `DelayedsExpr` which is just setting operands to the input arguments. That's **almost** what `__new__` does with the exception of new unpacking collections, see https://github.com/dask/dask/blob/0fa5e18d511c49f1a9cd5f98c675a9f6cd2fc02f/dask/_expr.py#L67

Additionally, the `unpack_collections` is not implemented thoroughly and only cover the most trivial cases. In particular, the case of `Delayed` is explicitly not handled since `hasattr` is always True for a Delayed. Therefore, the `__new__` implementation will set `operands` to a list of `[op.expr for op in operands]` instead of just `operands`. Therefore, the class is created in a corrupt state since `expr` attribute on a Delayed object does not refer to the expr of the collection (yet) but rather generates a new `DelayedAttr`. That's likely the reason why the init was introduced in the first place.

Subsequently, we're calculating the `_name` of the instance to cache it right away which triggers a tokenization of the faulty operands and stores the deterministic token in the cache. Therefore, after the `__init__` operates on the instance, we have an instance that does no longer match its token.

At this point I'm lost again. Since this inconsistency should consistently applied to all instances, it should still work as expected. Regardless, cleaning this up fixes the problem and removes some anomaly in the code.


Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 9bdb4c72e..05aee4019 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -34,10 +34,12 @@ OptimizerStage: TypeAlias = Literal[
 
 
 def _unpack_collections(o):
+    from dask.delayed import Delayed
+
     if isinstance(o, Expr):
         return o
 
-    if hasattr(o, "expr"):
+    if hasattr(o, "expr") and not isinstance(o, Delayed):
         return o.expr
     else:
         return o
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..fcd464d46 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
@@ -3051,15 +3053,8 @@ class PartitionsFiltered(Expr):
 
 
 class _DelayedExpr(Expr):
-    # Wraps a Delayed object to make it an Expr for now. This is hacky and we should
-    # integrate this properly...
-    # TODO
     _parameters = ["obj"]
 
-    def __init__(self, obj, _determ_token=None):
-        self.obj = obj
-        self.operands = [obj]
-
     def __str__(self):
         return f"{type(self).__name__}({str(self.obj)})"
 
@@ -3082,11 +3077,6 @@ class _DelayedExpr(Expr):
 
 
 class DelayedsExpr(Expr):
-    _parameters = []
-
-    def __init__(self, *delayed_objects, _determ_token=None):
-        self.operands = delayed_objects
-
     def __str__(self):
         return f"{type(self).__name__}({str(self.operands[0])})"
 
@@ -3804,19 +3794,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)