Repository: dask/dask
Pull request #11886 (2025-04-14 10:32:21)

Ensure only classes without custom init are singletons
I still don't entirely understand what is happening but it seems like we're able to sneak futures/objects into our expressions without the garbage collector learning about it. This is somehow related to our singleton / object deduplication logic but it still eludes me what is happening. I believe this change makes sense regardless of what is actually happening since objects with `__init__` are somewhat unpredictable and we shouldn't try to be smart about them.

Below are [graphs of backreferences](https://mg.pov.lt/objgraph/) of one of the futures objects right before we call `del futures` in the script https://github.com/dask/distributed/issues/9041 
The first shows the `main` branch. As you can see, the `Future` object is only referenced by the `list` in the local context. Therefore, once the list is finalized, the futures are finalized and the scheduler is notified about the futures being released.

However, without the deduplication (or if we remove the custom `__init__`) the DataFrame assumes ownership of the `Future`

With singleton deduplication

![futures](https://github.com/user-attachments/assets/61e641fa-d2eb-48ff-8c86-b9dbc8e20eb6)

without singleton deduplication
![futures](https://github.com/user-attachments/assets/36868541-34ec-426d-8444-0f5d3b3d1c26)

closes https://github.com/dask/distributed/issues/9041

Note: I'll dig into this a little more, trying to understand what's wrong and add a better test case


TomAugspurger: > This is somehow related to our singleton / object deduplication logic but it still eludes me what is happening.

I'm glad I wasn't the only one who found this confusing. Thanks for looking into it and LMK if there's anything I can help with testing.



Patch:
diff --git a/dask/_expr.py b/dask/_expr.py
index 9bdb4c72e..b99210e4e 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
@@ -837,7 +839,7 @@ class SingletonExpr(Expr):
             cls._instances = weakref.WeakValueDictionary()
         inst = super().__new__(cls, *args, _determ_token=_determ_token, **kwargs)
         _name = inst._name
-        if _name in cls._instances:
+        if _name in cls._instances and cls.__init__ == object.__init__:
             return cls._instances[_name]
 
         cls._instances[_name] = inst
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)
diff --git a/dask/tests/test_expr.py b/dask/tests/test_expr.py
index c244f7fdd..c477be558 100644
--- a/dask/tests/test_expr.py
+++ b/dask/tests/test_expr.py
@@ -5,7 +5,7 @@ import pickle
 
 import pytest
 
-from dask._expr import Expr
+from dask._expr import Expr, SingletonExpr
 
 
 class MyExpr(Expr):
@@ -85,3 +85,59 @@ def test_pickle_cached_properties():
         await c.submit(f, expr)
 
     test()
+
+
+class MySingleton(SingletonExpr): ...
+
+
+class MySingletonWithCustomInit(SingletonExpr):
+    def __init__(self, *args, **kwargs): ...
+
+
+class MySingletonInheritsCustomInit(MySingletonWithCustomInit): ...
+
+
+class Mixin:
+    def __init__(self, *args, **kwargs): ...
+
+
+class MySingletonInheritsCustomInitAsMixin(MySingleton, Mixin): ...
+
+
+def test_singleton_expr():
+    assert MySingleton(1, 2) is MySingleton(1, 2)
+    # We don't want to deduplicate if there is an __init__ since that may
+    # mutatate our singleton reference and we have no way to know
+    assert MySingletonWithCustomInit(1, 2) is not MySingletonWithCustomInit(1, 2)
+    assert MySingletonInheritsCustomInit(1, 2) is not MySingletonInheritsCustomInit(
+        1, 2
+    )
+    assert MySingletonInheritsCustomInitAsMixin(
+        1, 2
+    ) is not MySingletonInheritsCustomInitAsMixin(1, 2)
+
+
+@pytest.mark.slow()
+def test_refcounting_futures():
+    pd = pytest.importorskip("pandas")
+    dd = pytest.importorskip("dask.dataframe")
+    distributed = pytest.importorskip("distributed")
+
+    # See https://github.com/dask/distributed/issues/9041
+    # Didn't reproduce with any of our fixtures
+    with distributed.Client(
+        n_workers=2, worker_class=distributed.Worker, dashboard_address=":0"
+    ) as client:
+
+        def gen(i):
+            return pd.DataFrame({"A": [i]}, index=[i])
+
+        futures = [client.submit(gen, i) for i in range(3)]
+
+        meta = gen(0)[:0]
+        df = dd.from_delayed(futures, meta)
+        df.compute()
+
+        del futures
+
+        df.compute()