Repository: dask/dask
Pull request #11881: migrate delayed unpack_collections
Author: login: fjetter
Created: 2025-04-09T12:21:48Z | Merged: 2025-04-16T08:16:11Z

Closes https://github.com/dask/dask/issues/11879

Relates to https://github.com/dask/dask/issues/11854 as well

This includes two changes

1. It removes `_blockwise_unpack_collections_task_spec` which is a version of `delayed.unpack_collections` that was using the `Task` class already. This was introduced when migrating blockwise to keep the changes small since hybrid tasks can cause problems.
2. It changes how collection finalization works when the collection is encountered as an input (e.g. to a delayed object but also to others).
- It is converting it to a HLGExpr which automatically will perform the container specific optimization step upon materialization.
- It also takes care that when multiple collections are encountered, they are optimized together.
- We are ensuring that whatever result comes out of this optimization is unique and the results are not being reused. This is a common problem we've already encountered in dask-expr which can cause major memory blow up otherwise.




Diff:
diff --git a/dask/_expr.py b/dask/_expr.py
index b99210e4e5a..60aa9a0148e 100644
--- a/dask/_expr.py
+++ b/dask/_expr.py
@@ -12,7 +12,7 @@
 import toolz
 
 import dask
-from dask._task_spec import Task
+from dask._task_spec import Task, convert_legacy_graph
 from dask.tokenize import _tokenize_deterministic
 from dask.typing import Key
 from dask.utils import ensure_dict, funcname, import_required
@@ -1268,7 +1268,7 @@ def _simplify_down(self):
 
     @property
     def _name(self):
-        return f"finalize-{self.deterministic_token}"
+        return f"finalize-{super()._name}"
 
     def __dask_graph__(self):
         # The baseclass __dask_graph__ will not just materialize this layer but
@@ -1299,3 +1299,49 @@ def hlg(self):
 
     def __dask_keys__(self):
         return [self._name]
+
+
+class ProhibitReuse(Expr):
+    """
+    An expression that guarantees that all keys are suffixes with a unique id.
+    This can be used to break a common subexpression apart.
+    """
+
+    _parameters = ["expr"]
+
+    def __dask_keys__(self):
+        return self._modify_keys(self.expr.__dask_keys__())
+
+    @staticmethod
+    def _identity(obj):
+        return obj
+
+    @functools.cached_property
+    def _suffix(self):
+        return uuid.uuid4().hex
+
+    def _modify_keys(self, k):
+        if isinstance(k, list):
+            return [self._modify_keys(kk) for kk in k]
+        elif isinstance(k, tuple):
+            return (self._modify_keys(k[0]),) + k[1:]
+        elif isinstance(k, (int, float)):
+            k = str(k)
+        return f"{k}-{self._suffix}"
+
+    def __dask_graph__(self):
+        dsk = convert_legacy_graph(self.expr.__dask_graph__())
+
+        subs = {old_key: self._modify_keys(old_key) for old_key in dsk}
+        dsk2 = {
+            new_key: Task(
+                new_key,
+                ProhibitReuse._identity,
+                dsk.pop(old_key).substitute(subs),
+            )
+            for old_key, new_key in subs.items()
+        }
+        dsk2.update(dsk)
+        return dsk2
+
+    _layer = __dask_graph__
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 45e6c466742..f83011862b5 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -673,6 +673,9 @@ def __init__(
     def data_producer(self) -> bool:
         return self._data_producer
 
+    def has_subgraph(self) -> bool:
+        return self.func == _execute_subgraph
+
     def copy(self):
         return type(self)(
             self.key,
@@ -1119,6 +1122,10 @@ def fuse_linear_task_spec(dsk, keys):
 def cull(
     dsk: dict[KeyType, GraphNode], keys: Iterable[KeyType]
 ) -> dict[KeyType, GraphNode]:
+    if not isinstance(keys, (list, set, tuple)):
+        raise TypeError(
+            f"Expected list, set or tuple for keys, got {type(keys).__name__}"
+        )
     work = set(keys)
     seen: set[KeyType] = set()
     dsk2 = {}
diff --git a/dask/array/_array_expr/_blockwise.py b/dask/array/_array_expr/_blockwise.py
index b123494cd62..dd60900030c 100644
--- a/dask/array/_array_expr/_blockwise.py
+++ b/dask/array/_array_expr/_blockwise.py
@@ -17,8 +17,8 @@
     is_scalar_for_elemwise,
     normalize_arg,
 )
-from dask.blockwise import _blockwise_unpack_collections_task_spec
 from dask.blockwise import blockwise as core_blockwise
+from dask.delayed import unpack_collections
 from dask.layers import ArrayBlockwiseDep
 from dask.tokenize import _tokenize_deterministic
 from dask.utils import cached_property, funcname
@@ -147,7 +147,7 @@ def _layer(self):
         for arg, ind in arginds:
             if ind is None:
                 arg = normalize_arg(arg)
-                arg, collections = _blockwise_unpack_collections_task_spec(arg)
+                arg, collections = unpack_collections(arg)
                 dependencies.extend(collections)
             else:
                 if (
@@ -171,7 +171,7 @@ def _layer(self):
         kwargs2 = {}
         for k, v in self.kwargs.items():
             v = normalize_arg(v)
-            v, collections = _blockwise_unpack_collections_task_spec(v)
+            v, collections = unpack_collections(v)
             dependencies.extend(collections)
             kwargs2[k] = v
 
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index 1d2a063b709..4c0958c40dd 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -8,8 +8,8 @@
 import tlz as toolz
 
 from dask import base, utils
-from dask.blockwise import _blockwise_unpack_collections_task_spec
 from dask.blockwise import blockwise as core_blockwise
+from dask.delayed import unpack_collections
 from dask.highlevelgraph import HighLevelGraph
 from dask.layers import ArrayBlockwiseDep
 
@@ -215,7 +215,7 @@ def blockwise(
     for arg, ind in arginds:
         if ind is None:
             arg = normalize_arg(arg)
-            arg, collections = _blockwise_unpack_collections_task_spec(arg)
+            arg, collections = unpack_collections(arg)
 
             dependencies.extend(collections)
         else:
@@ -238,7 +238,7 @@ def blockwise(
     kwargs2 = {}
     for k, v in kwargs.items():
         v = normalize_arg(v)
-        v, collections = _blockwise_unpack_collections_task_spec(v)
+        v, collections = unpack_collections(v)
         dependencies.extend(collections)
         kwargs2[k] = v
 
diff --git a/dask/array/core.py b/dask/array/core.py
index d87f3164a88..b748ac14c0d 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -33,7 +33,7 @@
 xr = import_optional_dependency("xarray", errors="ignore")
 
 from dask import compute, config, core
-from dask._task_spec import List, Task, TaskRef
+from dask._task_spec import Alias, List, Task, TaskRef
 from dask.array import chunk
 from dask.array.chunk import getitem
 from dask.array.chunk_types import is_valid_array_chunk, is_valid_chunk_type
@@ -3995,15 +3995,24 @@ def from_delayed(value, shape, dtype=None, meta=None, name=None):
     """
     from dask.delayed import Delayed, delayed
 
-    if not isinstance(value, Delayed) and hasattr(value, "key"):
+    is_future = False
+    name = name or "from-value-" + tokenize(value, shape, dtype, meta)
+    if isinstance(value, TaskRef):
+        is_future = True
+    elif not isinstance(value, Delayed) and hasattr(value, "key"):
         value = delayed(value)
+    task = Alias(
+        key=(name,) + (0,) * len(shape),
+        target=value.key,
+    )
 
-    name = name or "from-value-" + tokenize(value, shape, dtype, meta)
-    dsk = {(name,) + (0,) * len(shape): value.key}
+    dsk = {task.key: task}
     chunks = tuple((d,) for d in shape)
     # TODO: value._key may not be the name of the layer in value.dask
     # This should be fixed after we build full expression graphs
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[value])
+    graph = HighLevelGraph.from_collections(
+        name, dsk, dependencies=[value] if not is_future else []
+    )
     return Array(graph, name, chunks, dtype=dtype, meta=meta)
 
 
diff --git a/dask/array/routines.py b/dask/array/routines.py
index dc30e91296e..b0f60e69091 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -9,6 +9,7 @@
 import numpy as np
 from tlz import concat, interleave, sliding_window
 
+from dask._task_spec import Task, TaskRef
 from dask.array import chunk
 from dask.array.core import (
     Array,
@@ -849,23 +850,13 @@ def searchsorted(a, v, side="left", sorter=None):
     return out
 
 
-# TODO: dask linspace doesn't support delayed values
-def _linspace_from_delayed(start, stop, num=50):
-    linspace_name = "linspace-" + tokenize(start, stop, num)
-    (start_ref, stop_ref, num_ref), deps = unpack_collections([start, stop, num])
-    if len(deps) == 0:
-        return np.linspace(start, stop, num=num)
-
-    linspace_dsk = {(linspace_name, 0): (np.linspace, start_ref, stop_ref, num_ref)}
-    linspace_graph = HighLevelGraph.from_collections(
-        linspace_name, linspace_dsk, dependencies=deps
-    )
-
-    chunks = ((np.nan,),) if is_dask_collection(num) else ((num,),)
-    return Array(linspace_graph, linspace_name, chunks, dtype=float)
+def _linspace(bins_range):
+    bins, (start, stop) = bins_range
+    return np.linspace(start, stop, num=bins + 1)
 
 
-def _block_hist(x, bins, range=None, weights=None):
+def _block_hist(x, bins_range, weights=None):
+    bins, range = bins_range
     return np.histogram(x, bins, range=range, weights=weights)[0][np.newaxis]
 
 
@@ -1006,10 +997,27 @@ def histogram(a, bins=None, range=None, normed=False, weights=None, density=None
     token = tokenize(a, bins, range, weights, density)
     name = "histogram-sum-" + token
 
+    bins_range, deps = unpack_collections([bins, range])
+
     if scalar_bins:
-        bins = _linspace_from_delayed(range[0], range[1], bins + 1)
-        # ^ NOTE `range[1]` is safe because of the above check, and the initial check
-        # that range must not be None if `scalar_bins`
+        assert range is not None
+        assert bins is not None
+        if len(deps) == 0:
+            bins = np.linspace(range[0], range[1], num=bins + 1)
+        else:
+            linspace_name = "linspace-" + tokenize(bins_range)
+
+            linspace_dsk = {
+                (linspace_name, 0): Task((linspace_name, 0), _linspace, bins_range)
+            }
+            linspace_graph = HighLevelGraph.from_collections(
+                linspace_name, linspace_dsk, dependencies=deps
+            )
+            if is_dask_collection(bins):
+                chunks = ((np.nan,),)
+            else:
+                chunks = ((bins + 1,),)
+            bins = Array(linspace_graph, linspace_name, chunks, dtype=float)
     else:
         if not isinstance(bins, (Array, np.ndarray)):
             bins = asarray(bins)
@@ -1018,12 +1026,10 @@ def histogram(a, bins=None, range=None, normed=False, weights=None, density=None
                 f"bins must be a 1-dimensional array or sequence, got shape {bins.shape}"
             )
 
-    (bins_ref, range_ref), deps = unpack_collections([bins, range])
-
     # Map the histogram to all bins, forming a 2D array of histograms, stacked for each chunk
     if weights is None:
         dsk = {
-            (name, i, 0): (_block_hist, k, bins_ref, range_ref)
+            (name, i, 0): Task((name, i, 0), _block_hist, TaskRef(k), bins_range)
             for i, k in enumerate(flatten(a.__dask_keys__()))
         }
         dtype = np.histogram([])[0].dtype
@@ -1031,7 +1037,9 @@ def histogram(a, bins=None, range=None, normed=False, weights=None, density=None
         a_keys = flatten(a.__dask_keys__())
         w_keys = flatten(weights.__dask_keys__())
         dsk = {
-            (name, i, 0): (_block_hist, k, bins_ref, range_ref, w)
+            (name, i, 0): Task(
+                (name, i, 0), _block_hist, TaskRef(k), bins_range, TaskRef(w)
+            )
             for i, (k, w) in enumerate(zip(a_keys, w_keys))
         }
         dtype = weights.dtype
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index 8374f731583..84e29572264 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -5454,12 +5454,12 @@ def test_map_blocks_large_inputs_delayed():
     b = np.ones(1000000)
 
     c = a.map_blocks(add, b)
-    assert any(b is v for v in c.dask.values())
+    assert any(b is v() for v in c.dask.values() if isinstance(v, DataNode))
     assert repr(dict(c.dask)).count(repr(b)[:10]) == 1  # only one occurrence
 
     d = a.map_blocks(lambda x, y: x + y.sum(), y=b)
     assert_eq(d, d)
-    assert any(b is v for v in d.dask.values())
+    assert any(b is v() for v in d.dask.values() if isinstance(v, DataNode))
     assert repr(dict(c.dask)).count(repr(b)[:10]) == 1  # only one occurrence
 
 
@@ -5471,12 +5471,12 @@ def func(a, b):
     b = np.ones(1000000)
 
     c = da.blockwise(func, "i", a, "i", b, None, dtype=a.dtype)
-    assert any(b is v for v in c.dask.values())
+    assert any(b is v() for v in c.dask.values() if isinstance(v, DataNode))
     assert repr(dict(c.dask)).count(repr(b)[:10]) == 1  # only one occurrence
     assert_eq(c, c)
 
     d = da.blockwise(lambda x, y: x, "i", a, "i", y=b, dtype=a.dtype)
-    assert any(b is v for v in d.dask.values())
+    assert any(b is v() for v in d.dask.values() if isinstance(v, DataNode))
     assert repr(dict(c.dask)).count(repr(b)[:10]) == 1  # only one occurrence
     assert_eq(d, d)
 
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
index d290f0478d2..5a1130cb994 100644
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -861,12 +861,12 @@ def test_histogram_bin_range_raises(bins, hist_range):
 
 @pytest.mark.parametrize("density", [True, False])
 @pytest.mark.parametrize("weighted", [True, False])
-@pytest.mark.parametrize("non_delayed_i", [None, 0, 1])
+@pytest.mark.parametrize("non_delayed_i", [None, 0])
 @pytest.mark.parametrize("delay_n_bins", [False, True])
 def test_histogram_delayed_range(density, weighted, non_delayed_i, delay_n_bins):
     n = 100
     v = np.random.default_rng().random(n)
-    vd = da.from_array(v, chunks=10)
+    vd = da.from_array(v, chunks=2)
 
     if weighted:
         weights = np.random.default_rng().random(n)
@@ -926,8 +926,11 @@ def test_histogram_delayed_bins(density, weighted):
     )
 
     assert bins_d is bins_d2
-    assert_eq(hist_d, hist)
-    assert_eq(bins_d2, bins)
+    # The HLG that is assembled from the bins and the range triggers a sanity
+    # check because they contain duplicate keys and the HLG dependencies are not
+    # reflecting this properly. Graph is perfectly fine but the check fails.
+    assert_eq(hist_d, hist, check_graph=False)
+    assert_eq(bins_d2, bins, check_graph=False)
 
 
 def test_histogram_delayed_n_bins_raises_with_density():
diff --git a/dask/array/tests/test_xarray.py b/dask/array/tests/test_xarray.py
index 9740996d514..007b4d8ada9 100644
--- a/dask/array/tests/test_xarray.py
+++ b/dask/array/tests/test_xarray.py
@@ -1,10 +1,14 @@
 from __future__ import annotations
 
+from functools import partial
+
 import numpy as np
 import pytest
 
+import dask
 import dask.array as da
 from dask.array.utils import assert_eq
+from dask.utils import tmpdir
 
 xr = pytest.importorskip("xarray")
 
@@ -82,3 +86,42 @@ def test_positional_indexer_multiple_variables():
     graph = result.__dask_graph__()
     assert len({k for k in graph if "shuffle-taker" in k}) == 4
     assert len({k for k in graph if "shuffle-sorter" in k}) == 2
+
+
+@pytest.mark.parametrize("compute", [True, False])
+def test_xarray_blockwise_fusion_store(compute):
+    def custom_scheduler_get(dsk, keys, expected, **kwargs):
+        dsk = dsk.__dask_graph__()
+        assert (
+            len(dsk) == expected
+        ), f"False number of tasks got {len(dsk)} but expected {expected}"
+        return [42 for _ in keys]
+
+    # First test that this mocking stuff works as expecged
+    with pytest.raises(AssertionError, match="False number of tasks"):
+        scheduler = partial(custom_scheduler_get, expected=42)
+        dask.compute(da.ones(10), scheduler=scheduler)
+
+    coord = da.arange(8, chunks=-1)
+    data = da.random.random((8, 8), chunks=-1) + 1
+    x = xr.DataArray(data, coords={"x": coord, "y": coord}, dims=["x", "y"])
+
+    y = ((x + 1) * 2) / 2 - 1
+
+    # Everything fused into one compute task
+    # one finalize Alias
+    expected = 2
+    scheduler = partial(custom_scheduler_get, expected=expected)
+    dask.compute(y, scheduler=scheduler)
+
+    with tmpdir() as dirname:
+        if compute:
+            with dask.config.set(scheduler=scheduler):
+                y.to_zarr(dirname, compute=True)
+        else:
+            # There's a delayed finalize store smashed on top which won't be fused by
+            # default
+            expected += 1
+            scheduler = partial(custom_scheduler_get, expected=expected)
+            stored = y.to_zarr(dirname, compute=False)
+            dask.compute(stored, scheduler=scheduler)
diff --git a/dask/array/utils.py b/dask/array/utils.py
index 9aeff17dff9..b1147f1b638 100644
--- a/dask/array/utils.py
+++ b/dask/array/utils.py
@@ -10,6 +10,7 @@
 import numpy as np
 from tlz import concat, frequencies
 
+from dask._task_spec import convert_legacy_graph
 from dask.array.numpy_compat import AxisError
 from dask.base import is_dask_collection, tokenize
 from dask.highlevelgraph import HighLevelGraph
@@ -213,8 +214,10 @@ def _check_dsk(dsk):
     key_collisions = set()
     # Allow redundant keys if the values are equivalent
     collisions = set()
+    all_keys = set(dsk)
+    layers = [convert_legacy_graph(layer, all_keys) for layer in dsk.layers.values()]
     for k in non_one.keys():
-        for layer in dsk.layers.values():
+        for layer in layers:
             try:
                 key_collisions.add(tokenize(layer[k]))
             except KeyError:
diff --git a/dask/bag/core.py b/dask/bag/core.py
index 4e9278a002f..04c5138ae69 100644
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -57,10 +57,10 @@
     replace_name_in_key,
     tokenize,
 )
-from dask.blockwise import _blockwise_unpack_collections_task_spec, blockwise
+from dask.blockwise import blockwise
 from dask.context import globalmethod
 from dask.core import flatten, istask, quote
-from dask.delayed import Delayed
+from dask.delayed import Delayed, unpack_collections
 from dask.highlevelgraph import HighLevelGraph
 from dask.sizeof import sizeof
 from dask.typing import Graph, NestedKeys, no_default
@@ -1906,17 +1906,25 @@ def from_delayed(values):
 
     if isinstance(values, Delayed):
         values = [values]
-    values = [
-        delayed(v) if not isinstance(v, Delayed) and hasattr(v, "key") else v
-        for v in values
-    ]
+    futures = [v for v in values if isinstance(v, TaskRef)]
+    if all_futures := (len(futures) == len(values)):
+        # All futures. Fast path
+        values = futures
+    else:
+        # Every Delayed generates a Layer, i.e. this path is much more expensive
+        # if there are many input values.
+        values = [
+            delayed(v) if not isinstance(v, (Delayed,)) and hasattr(v, "key") else v
+            for v in values
+        ]
 
     name = "bag-from-delayed-" + tokenize(*values)
-    names = [(name, i) for i in range(len(values))]
-    values2 = [(reify, v.key) for v in values]
-    dsk = dict(zip(names, values2))
+    tasks = [Task((name, i), reify, TaskRef(v.key)) for i, v in enumerate(values)]
+    dsk = {t.key: t for t in tasks}
 
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=values)
+    graph = HighLevelGraph.from_collections(
+        name, dsk, dependencies=values if not all_futures else ()
+    )
     return Bag(graph, name, len(values))
 
 
@@ -2095,7 +2103,7 @@ def unpack_scalar_dask_kwargs(kwargs):
     kwargs2 = {}
     dependencies = []
     for k, v in kwargs.items():
-        vv, collections = _blockwise_unpack_collections_task_spec(v)
+        vv, collections = unpack_collections(v)
         if not collections:
             kwargs2[k] = v
         else:
@@ -2268,7 +2276,7 @@ def map_partitions(func, *args, **kwargs):
             bags.append(a)
             args2.append(a)
         else:
-            a, collections = _blockwise_unpack_collections_task_spec(a)
+            a, collections = unpack_collections(a)
             args2.append(a)
             dependencies.extend(collections)
 
diff --git a/dask/blockwise.py b/dask/blockwise.py
index 43ca52d1cdb..0910f463d7c 100644
--- a/dask/blockwise.py
+++ b/dask/blockwise.py
@@ -3,7 +3,6 @@
 import itertools
 import os
 from collections.abc import Hashable, Iterable, Mapping, Sequence
-from dataclasses import fields, is_dataclass, replace
 from itertools import product
 from math import prod
 from typing import Any
@@ -11,11 +10,9 @@
 import tlz as toolz
 
 import dask
-from dask import base, utils
-from dask._task_spec import Dict, GraphNode, List, Task, TaskRef, parse_input
+from dask._task_spec import GraphNode, List, Task, TaskRef, parse_input
 from dask.base import clone_key, get_name_from_key, tokenize
 from dask.core import flatten, ishashable, keys_in_tasks, reverse_dict
-from dask.delayed import Delayed, finalize
 from dask.highlevelgraph import HighLevelGraph, Layer
 from dask.optimization import fuse
 from dask.typing import Key
@@ -1574,97 +1571,3 @@ def fuse_roots(graph: HighLevelGraph, keys: list):
             dependencies[name] = set()
 
     return HighLevelGraph(layers, dependencies)
-
-
-def _blockwise_unpack_collections_task_spec(expr):
-    # FIXME This is a copy of the delayed.unpack_collections function with the
-    # addition of the TaskSpec class. Eventually this should all be consolidated
-    # but to reduce the number of changes we'll vendor this here
-    # FIXME: There is also a dask.base version of unpack_collections that looks
-    # similar but is different. At the very least the names should be fixed
-    if isinstance(expr, Delayed):
-        return TaskRef(expr._key), (expr,)
-
-    if base.is_dask_collection(expr):
-        if hasattr(expr, "optimize"):
-            # Optimize dask-expr collections
-            expr = expr.optimize()
-
-        finalized = finalize(expr)
-        return TaskRef(finalized._key), (finalized,)
-
-    if type(expr) is type(iter(list())):
-        expr = list(expr)
-    elif type(expr) is type(iter(tuple())):
-        expr = tuple(expr)
-    elif type(expr) is type(iter(set())):
-        expr = set(expr)
-
-    typ = type(expr)
-
-    if typ in (list, tuple, set):
-        args, collections = utils.unzip(
-            (_blockwise_unpack_collections_task_spec(e) for e in expr), 2
-        )
-        collections = tuple(toolz.unique(toolz.concat(collections), key=id))
-        if not collections:
-            return expr, ()
-        args = List(*args)
-        # Ensure output type matches input type
-        if typ is not list:
-            args = Task(None, typ, args)
-        return args, collections
-
-    if typ is dict:
-        args, collections = _blockwise_unpack_collections_task_spec(
-            [[k, v] for k, v in expr.items()]
-        )
-        if not collections:
-            return expr, ()
-        return Dict(args), collections
-
-    if typ is slice:
-        args, collections = _blockwise_unpack_collections_task_spec(
-            [expr.start, expr.stop, expr.step]
-        )
-        if not collections:
-            return expr, ()
-        return Task(None, slice, *args), collections
-
-    if is_dataclass(expr):
-        args, collections = _blockwise_unpack_collections_task_spec(
-            [
-                [f.name, getattr(expr, f.name)]
-                for f in fields(expr)
-                if hasattr(expr, f.name)  # if init=False, field might not exist
-            ]
-        )
-        if not collections:
-            return expr, ()
-        try:
-            _fields = {
-                f.name: getattr(expr, f.name)
-                for f in fields(expr)
-                if hasattr(expr, f.name)
-            }
-            replace(expr, **_fields)
-        except (TypeError, ValueError) as e:
-            if isinstance(e, ValueError) or "is declared with init=False" in str(e):
-                raise ValueError(
-                    f"Failed to unpack {typ} instance. "
-                    "Note that using fields with `init=False` are not supported."
-                ) from e
-            else:
-                raise TypeError(
-                    f"Failed to unpack {typ} instance. "
-                    "Note that using a custom __init__ is not supported."
-                ) from e
-        return Task(None, typ, **dict(args)), collections
-
-    if utils.is_namedtuple_instance(expr):
-        args, collections = _blockwise_unpack_collections_task_spec([v for v in expr])
-        if not collections:
-            return expr, ()
-        return Task(None, typ, *args), collections
-
-    return expr, ()
diff --git a/dask/core.py b/dask/core.py
index dd61cc980d5..5c42bfb8c4f 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -6,7 +6,12 @@
 
 import toolz
 
-from dask._task_spec import DependenciesMapping, convert_legacy_graph, execute_graph
+from dask._task_spec import (
+    DependenciesMapping,
+    TaskRef,
+    convert_legacy_graph,
+    execute_graph,
+)
 from dask.typing import Graph, Key, NoDefault, no_default
 
 
@@ -136,6 +141,8 @@ def keys_in_tasks(keys: Collection[Key], tasks: Iterable[Any], as_list: bool = F
                 work.extend(w.values())
             elif isinstance(w, GraphNode):
                 work.extend(w.dependencies)
+            elif isinstance(w, TaskRef):
+                work.append(w.key)
             else:
                 try:
                     if w in keys:
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index fcd464d460d..0474830b089 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -20,6 +20,7 @@
 from dask._expr import FinalizeCompute
 from dask._task_spec import Alias, DataNode, Task, TaskRef, execute_graph
 from dask.array import Array
+from dask.base import collections_to_expr
 from dask.core import flatten
 from dask.dataframe import methods
 from dask.dataframe._pyarrow import to_pyarrow_string
@@ -3085,13 +3086,27 @@ def _name(self):
         return "delayed-container-" + self.deterministic_token
 
     def _layer(self) -> dict:
-        dask = {}
-        for i, obj in enumerate(self.operands):
-            dc = obj.__dask_optimize__(obj.dask, obj.key).to_dict().copy()
-            dc[(self._name, i)] = dc[obj.key]
-            dc.pop(obj.key)
-            dask.update(dc)
-        return dask
+        from dask.delayed import Delayed
+
+        if isinstance(self.operands[0], TaskRef):
+            tasks = [
+                Alias((self._name, ix), fut.key) for ix, fut in enumerate(self.operands)
+            ]
+            dsk = {t.key: t for t in tasks}
+        elif isinstance(self.operands[0], Delayed):
+            expr = collections_to_expr(self.operands).optimize()
+            keys = expr.__dask_keys__()
+            dsk = expr.__dask_graph__()
+            # Many APIs in dask-expr are not honoring __dask_keys__ but are instead
+            # assuming they can just construc the keys themselves by walking the
+            # partitions. Therefore we'll have to remap the key names and can't just
+            # expose __dask_keys__()
+            for ix, actual_key in enumerate(keys):
+                dsk[(self._name, ix)] = Alias((self._name, ix), actual_key[0])
+        else:
+            raise TypeError("Expected a Delayed or Future object")
+
+        return dsk
 
     def _divisions(self):
         return (None,) * (len(self.operands) + 1)
diff --git a/dask/dataframe/dask_expr/io/_delayed.py b/dask/dataframe/dask_expr/io/_delayed.py
index e948a4a2001..23edef9f6ab 100644
--- a/dask/dataframe/dask_expr/io/_delayed.py
+++ b/dask/dataframe/dask_expr/io/_delayed.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import functools
-from collections.abc import Iterable
+from collections.abc import Collection
 from typing import TYPE_CHECKING
 
 import pandas as pd
@@ -75,7 +75,7 @@ def identity(x):
 
 
 def from_delayed(
-    dfs: Delayed | distributed.Future | Iterable[Delayed | distributed.Future],
+    dfs: Delayed | distributed.Future | Collection[Delayed | distributed.Future],
     meta=None,
     divisions: tuple | None = None,
     prefix: str | None = None,
@@ -119,7 +119,7 @@ def read_xml(path):
     if isinstance(dfs, Delayed) or hasattr(dfs, "key"):
         dfs = [dfs]
 
-    if len(dfs) == 0:  # type: ignore
+    if len(dfs) == 0:
         raise TypeError("Must supply at least one delayed object")
 
     if meta is None:
@@ -132,16 +132,23 @@ def read_xml(path):
         )
     elif divisions is not None:
         divs = list(divisions)
-        if len(divs) != len(dfs) + 1:  # type: ignore
+        if len(divs) != len(dfs) + 1:
             raise ValueError("divisions should be a tuple of len(dfs) + 1")
 
-    dfs = [
-        delayed(df) if not isinstance(df, Delayed) and hasattr(df, "key") else df
-        for df in dfs
-    ]
+    futures = [v for v in dfs if isinstance(v, TaskRef)]
+    if len(futures) == len(dfs):
+        # All futures. Fast path
+        dfs = futures
+    else:
+        # Every Delayed generates a Layer, i.e. this path is much more expensive
+        # if there are many input values.
+        dfs = [
+            delayed(v) if not isinstance(v, (Delayed,)) and hasattr(v, "key") else v
+            for v in dfs
+        ]
 
     for item in dfs:
-        if not isinstance(item, Delayed):
+        if not (isinstance(item, Delayed) or isinstance(item, TaskRef)):
             raise TypeError("Expected Delayed object, got %s" % type(item).__name__)
 
     from dask.dataframe.dask_expr._collection import new_collection
diff --git a/dask/dataframe/dask_expr/io/tests/test_delayed.py b/dask/dataframe/dask_expr/io/tests/test_delayed.py
index 42508dcd3e0..3d1806fb56e 100644
--- a/dask/dataframe/dask_expr/io/tests/test_delayed.py
+++ b/dask/dataframe/dask_expr/io/tests/test_delayed.py
@@ -14,7 +14,8 @@
 def test_from_delayed_optimizing():
     parts = from_dict({"a": np.arange(300)}, npartitions=30).to_delayed()
     result = from_delayed(parts[0], meta=pd.DataFrame({"a": pd.Series(dtype=np.int64)}))
-    assert len(result.optimize().dask) == 2
+    # +1 due to an Alias for correct naming
+    assert len(result.optimize().dask) == 3
     assert_eq(result, pd.DataFrame({"a": pd.Series(np.arange(10))}))
 
 
@@ -28,7 +29,8 @@ def test_from_delayed(prefix):
 
     df = from_delayed(dfs, meta=pdf.head(0), divisions=None, prefix=prefix)
     assert_eq(df, pdf)
-    assert len({k[0] for k in df.optimize().dask}) == 2
+    # +1 due to an Alias for correct naming
+    assert len({k[0] for k in df.optimize().dask}) == 3
     if prefix:
         assert df._name.startswith(prefix)
 
diff --git a/dask/delayed.py b/dask/delayed.py
index 24f8894c233..b3cb09f723b 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -8,10 +8,20 @@
 from dataclasses import fields, is_dataclass, replace
 from functools import partial
 
-from tlz import concat, curry, merge, unique
-
-from dask import config
-from dask._expr import FinalizeCompute
+import toolz
+from tlz import concat, curry, merge
+
+from dask import base, config, utils
+from dask._expr import FinalizeCompute, ProhibitReuse, _ExprSequence
+from dask._task_spec import (
+    DataNode,
+    GraphNode,
+    List,
+    Task,
+    TaskRef,
+    convert_legacy_graph,
+    fuse_linear_task_spec,
+)
 from dask.base import (
     DaskMethodsMixin,
     collections_to_expr,
@@ -23,7 +33,6 @@
 from dask.context import globalmethod
 from dask.core import flatten, quote
 from dask.highlevelgraph import HighLevelGraph
-from dask.optimization import fuse
 from dask.typing import Graph, NestedKeys
 from dask.utils import (
     OperatorMethodMixin,
@@ -45,19 +54,52 @@ def finalize(collection):
     assert is_dask_collection(collection)
 
     name = "finalize-" + tokenize(collection)
-    keys = collection.__dask_keys__()
-    finalize, args = collection.__dask_postcompute__()
-    layer = {name: (finalize, keys) + args}
-    graph = HighLevelGraph.from_collections(name, layer, dependencies=[collection])
-    return Delayed(name, graph)
+    expr = collections_to_expr(collection).finalize_compute()
+    return Delayed(name, expr)
+
+
+def _convert_dask_keys(keys: NestedKeys) -> List:
+    assert isinstance(keys, list)
+    new_keys: list[List | TaskRef] = []
+    for key in keys:
+        if isinstance(key, list):
+            new_keys.append(_convert_dask_keys(key))
+        else:
+            new_keys.append(TaskRef(key))
+    return List(*new_keys)
+
 
+def _finalize_args_collections(args, collections):
+    old_keys = [c.__dask_keys__()[0] for c in collections]
+    from dask._task_spec import cull
 
-def unpack_collections(expr):
+    collections = _ExprSequence(*collections).optimize()
+    new_keys = collections.__dask_keys__()
+    dsk = convert_legacy_graph(collections.__dask_graph__())
+    collections = tuple(Delayed(k[0], cull(dsk, [k[0]])) for k in new_keys)
+    subs = {old: new[0] for old, new in zip(old_keys, new_keys) if old != new}
+    args = args.substitute(subs)
+    return args, collections
+
+
+def unpack_collections(expr, _return_collections=True):
     """Normalize a python object and merge all sub-graphs.
 
     - Replace ``Delayed`` with their keys
     - Convert literals to things the schedulers can handle
-    - Extract dask graphs from all enclosed values
+    - Extract dask graphs from all enclosed values.
+
+    Note, that the returned _task_ is not necessarily runnable and the caller is
+    responsible to deal with the output types accordingly.
+
+    The task is one of
+    - `TaskRef` as a pointer to the collection returned in collections. This is
+      not callable and should not be a top-level member of a dask task graph.
+    - A runnable task (i.e. subclass `GraphNode`) which can be embedded
+      directly into a task graph. This indicates that a dask collection was
+      encountered on a deeper nesting level and this runnable task restores the
+      input nesting with the computed dask collection replaced.
+    - The unaltered object as provided if no dask collections are found.
 
     Parameters
     ----------
@@ -65,9 +107,13 @@ def unpack_collections(expr):
         The object to be normalized. This function knows how to handle
         dask collections, as well as most builtin python types.
 
+    _optimize_collections: bool, optional
+        Internal use only!
+
+
     Returns
     -------
-    task : normalized task to be run
+    task : object
     collections : a tuple of collections
 
     Examples
@@ -77,27 +123,46 @@ def unpack_collections(expr):
     >>> b = delayed(2, 'b')
     >>> task, collections = unpack_collections([a, b, 3])
     >>> task
-    ['a', 'b', 3]
+    List((Alias('a'), Alias('b'), 3))
     >>> collections
     (Delayed('a'), Delayed('b'))
 
     >>> task, collections = unpack_collections({a: 1, b: 2})
     >>> task
-    (<class 'dict'>, [['a', 1], ['b', 2]])
+    <Task None dict(List((List((Alias('a'), 1)), List((Alias('b'), 2)))))>
     >>> collections
     (Delayed('a'), Delayed('b'))
     """
-    # FIXME: See blockwise._unpack_collections for a TaskSpec version
     if isinstance(expr, Delayed):
-        return expr._key, (expr,)
-
-    if is_dask_collection(expr):
-        if hasattr(expr, "optimize"):
-            # Optimize dask-expr collections
-            expr = expr.optimize()
-
-        finalized = finalize(expr)
-        return finalized._key, (finalized,)
+        if _return_collections:
+            return TaskRef(expr._key), (expr,)
+        else:
+            expr = collections_to_expr(expr).finalize_compute()
+            (name,) = expr.__dask_keys__()
+            return TaskRef(name), (expr,)
+
+    # FIXME: Make this not trigger materialization
+    # Currently this is checking with hasattr for __dask_graph__ which triggers
+    # a materialization
+    if base.is_dask_collection(expr):
+        if _return_collections:
+            expr2 = ProhibitReuse(collections_to_expr(expr).finalize_compute())
+            finalized = expr2.optimize()
+            # FIXME: Make this also go away
+            dsk = finalized.__dask_graph__()
+            keys = list(flatten(finalized.__dask_keys__()))
+            if len(keys) > 1:
+                # `finalize_compute` _should_ guarantee that we only have one key
+                raise RuntimeError(
+                    "Cannot unpack dask collections which don't finalize to a "
+                    f"single key. Got {type(expr)} with {keys=}",
+                )
+
+            return unpack_collections(Delayed(keys[0], dsk))
+        else:
+            expr = collections_to_expr(expr).finalize_compute()
+            (name,) = expr.__dask_keys__()
+            return TaskRef(name), (expr,)
 
     if type(expr) is type(iter(list())):
         expr = list(expr)
@@ -109,21 +174,41 @@ def unpack_collections(expr):
     typ = type(expr)
 
     if typ in (list, tuple, set):
-        args, collections = unzip((unpack_collections(e) for e in expr), 2)
-        args = list(args)
-        collections = tuple(unique(concat(collections), key=id))
+        args, collections = utils.unzip(
+            (unpack_collections(e, _return_collections=False) for e in expr), 2
+        )
+
+        collections = tuple(toolz.unique(toolz.concat(collections), key=id))
+        if not collections:
+            return expr, ()
+        args = List(*args)
+        if _return_collections:
+            args, collections = _finalize_args_collections(args, collections)
         # Ensure output type matches input type
         if typ is not list:
-            args = (typ, args)
+            args = Task(None, typ, args)
         return args, collections
 
     if typ is dict:
-        args, collections = unpack_collections([[k, v] for k, v in expr.items()])
-        return (dict, args), collections
+        args, collections = unpack_collections(
+            [[k, v] for k, v in expr.items()], _return_collections=False
+        )
+        if not collections:
+            return expr, ()
+        if _return_collections:
+            args, collections = _finalize_args_collections(args, collections)
+        return Task(None, dict, args), collections
 
     if typ is slice:
-        args, collections = unpack_collections([expr.start, expr.stop, expr.step])
-        return (slice, *args), collections
+        args, collections = unpack_collections(
+            [expr.start, expr.stop, expr.step], _return_collections=False
+        )
+        if not collections:
+            return expr, ()
+
+        if _return_collections:
+            args, collections = _finalize_args_collections(args, collections)
+        return Task(None, apply, slice, args), collections
 
     if is_dataclass(expr):
         args, collections = unpack_collections(
@@ -131,10 +216,14 @@ def unpack_collections(expr):
                 [f.name, getattr(expr, f.name)]
                 for f in fields(expr)
                 if hasattr(expr, f.name)  # if init=False, field might not exist
-            ]
+            ],
+            _return_collections=False,
         )
         if not collections:
             return expr, ()
+
+        if _return_collections:
+            args, collections = _finalize_args_collections(args, collections)
         try:
             _fields = {
                 f.name: getattr(expr, f.name)
@@ -153,15 +242,25 @@ def unpack_collections(expr):
                     f"Failed to unpack {typ} instance. "
                     "Note that using a custom __init__ is not supported."
                 ) from e
-        return (apply, typ, (), (dict, args)), collections
+        return Task(None, apply, typ, (), Task(None, dict, args)), collections
 
-    if is_namedtuple_instance(expr):
-        args, collections = unpack_collections([v for v in expr])
-        return (typ, *args), collections
+    if utils.is_namedtuple_instance(expr):
+        args, collections = unpack_collections(
+            tuple(v for v in expr), _return_collections=False
+        )
+        if not collections:
+            return expr, ()
+        if _return_collections:
+            args, collections = _finalize_args_collections(args, collections)
+        return Task(None, _reconstruct_namedtuple, typ, args), collections
 
     return expr, ()
 
 
+def _reconstruct_namedtuple(typ, fields):
+    return typ(*fields)
+
+
 def to_task_dask(expr):
     """Normalize a python object and merge all sub-graphs.
 
@@ -504,6 +603,8 @@ def delayed(obj, name=None, pure=None, nout=None, traverse=True):
         if not name:
             name = f"{type(obj).__name__}-{tokenize(task, pure=pure)}"
         layer = {name: task}
+        if isinstance(task, GraphNode):
+            task.key = name
         graph = HighLevelGraph.from_collections(name, layer, dependencies=collections)
         return Delayed(name, graph, nout)
 
@@ -517,16 +618,8 @@ def right(method):
     return partial(_swap, method)
 
 
-def _fuse_delayed(dsk, keys, **kwargs):
-    dependencies = dsk.get_all_dependencies()
-    dsk = ensure_dict(dsk)
-
-    dsk, _ = fuse(
-        dsk,
-        keys,
-        dependencies=dependencies,
-    )
-    return dsk
+def _apply2(func, *args, dask_kwargs):
+    return func(*args, **dask_kwargs)
 
 
 def optimize(dsk, keys, **kwargs):
@@ -534,7 +627,8 @@ def optimize(dsk, keys, **kwargs):
         keys = [keys]
 
     if config.get("optimization.fuse.delayed"):
-        dsk = _fuse_delayed(dsk, keys, **kwargs)
+        dsk = ensure_dict(dsk)
+        dsk = fuse_linear_task_spec(dsk, keys, **kwargs)
 
     if not isinstance(dsk, HighLevelGraph):
         dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
@@ -690,12 +784,12 @@ def call_function(func, func_token, args, kwargs, pure=None, nout=None):
     args2, collections = unzip(map(unpack_collections, args), 2)
     collections = list(concat(collections))
 
-    if kwargs:
-        dask_kwargs, collections2 = unpack_collections(kwargs)
-        collections.extend(collections2)
-        task = (apply, func, list(args2), dask_kwargs)
+    dask_kwargs, collections2 = unpack_collections(kwargs)
+    collections.extend(collections2)
+    if dask_kwargs:
+        task = Task(name, _apply2, func, *args2, dask_kwargs=dask_kwargs)
     else:
-        task = (func,) + args2
+        task = Task(name, func, *args2, **kwargs)
 
     graph = HighLevelGraph.from_collections(
         name, {name: task}, dependencies=collections
@@ -708,7 +802,7 @@ class DelayedLeaf(Delayed):
     __slots__ = ("_obj", "_pure", "_nout")
 
     def __init__(self, obj, key, pure=None, nout=None):
-        super().__init__(key, None)
+        super().__init__(key, None, length=nout)
         self._obj = obj
         self._pure = pure
         self._nout = nout
@@ -716,7 +810,7 @@ def __init__(self, obj, key, pure=None, nout=None):
     @property
     def dask(self):
         return HighLevelGraph.from_collections(
-            self._key, {self._key: self._obj}, dependencies=()
+            self._key, {self._key: DataNode(self._key, self._obj)}, dependencies=()
         )
 
     def __call__(self, *args, **kwargs):
diff --git a/dask/dot.py b/dask/dot.py
index 3b8c37ffcf7..2d5298d81ce 100644
--- a/dask/dot.py
+++ b/dask/dot.py
@@ -4,6 +4,7 @@
 import re
 from functools import partial
 
+from dask._task_spec import GraphNode, Task
 from dask.core import get_dependencies, ishashable, istask
 from dask.utils import apply, funcname, import_required, key_split
 
@@ -19,6 +20,10 @@ def task_label(task):
     >>> task_label((add, (add, 1, 2), 3))
     'add(...)'
     """
+    if isinstance(task, GraphNode):
+        if isinstance(task, Task) and task.has_subgraph():
+            return f"{task.key}(...)"
+        return task.key
     func = task[0]
     if func is apply:
         func = task[1]
diff --git a/dask/tests/test_delayed.py b/dask/tests/test_delayed.py
index 41a32817c65..0972b7eaa3e 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -16,6 +16,7 @@
 import dask
 import dask.bag as db
 from dask import compute
+from dask._task_spec import Task
 from dask.base import collections_to_expr
 from dask.delayed import Delayed, delayed, to_task_dask
 from dask.highlevelgraph import HighLevelGraph
@@ -94,7 +95,6 @@ def test_delayed():
 
     a = delayed(1)
     assert a.compute() == 1
-    assert 1 in a.dask.values()
     b = add2(add2(a, 2), 3)
     assert a.key in b.dask
 
@@ -105,6 +105,7 @@ class ANamedTuple(NamedTuple):
 
     literal = dask.delayed(3)
     with_class = dask.delayed({"a": ANamedTuple(a=literal)})
+    assert with_class.compute() == {"a": ANamedTuple(a=3)}
 
     def return_nested(obj):
         return obj["a"].a
@@ -274,7 +275,9 @@ def test_method_getattr_call_same_task():
     a = delayed([1, 2, 3])
     o = a.index(1)
     # Don't getattr the method, then call in separate task
-    assert getattr not in {v[0] for v in o.__dask_graph__().values()}
+    tasks = {v.func for v in o.__dask_graph__().values() if isinstance(v, Task)}
+    assert tasks
+    assert getattr not in tasks
 
 
 def test_np_dtype_of_delayed():
@@ -511,7 +514,8 @@ def test_nout():
 def test_nout_with_tasks(x):
     length = len(x)
     d = delayed(x, nout=length)
-    assert len(d) == len(list(d)) == length
+    assert len(d) == length
+    assert len(list(d)) == length
     assert d.compute() == x
 
 
@@ -563,6 +567,64 @@ def test_array_delayed():
     assert (delayed_arr.compute() == arr).all()
 
 
+def test_array_delayed_complex_optimization():
+    # Ensure that when multiple collections are passed to a Delayed function,
+    # they are optimized together
+    np = pytest.importorskip("numpy")
+    pytest.importorskip("dask.array")
+    from dask.array.core import from_func
+    from dask.array.utils import assert_eq
+
+    called = False
+
+    def only_once():
+        nonlocal called
+        if called:
+            raise RuntimeError("Already executed")
+
+        called = True
+        return np.arange(100).reshape((10, 10))
+
+    darr = from_func(only_once, shape=(10, 10), dtype=int)
+    a = darr + 1
+    b = darr + 2
+    val = delayed(sum)([a, b, 1])
+    assert isinstance(val, Delayed)
+    np_arr = np.arange(100).reshape((10, 10))
+    assert_eq(val.compute(), (np_arr + 1) + (np_arr + 2) + 1)
+
+
+def test_array_delayed_complex_optimization_kwargs():
+    # Ensure that collections that if multiple collections are passed to a
+    # Delayed function that they are optimized together
+    np = pytest.importorskip("numpy")
+    pytest.importorskip("dask.array")
+    from dask.array.core import from_func
+    from dask.array.utils import assert_eq
+
+    called = False
+
+    def only_once():
+        nonlocal called
+        if called:
+            raise RuntimeError("Already executed")
+
+        called = True
+        return np.arange(100).reshape((10, 10))
+
+    darr = from_func(only_once, shape=(10, 10), dtype=int)
+    a = darr + 1
+    b = darr + 2
+
+    def sum_kwargs_only(*, a, b, c):
+        return sum([a, b, c])
+
+    val = delayed(sum_kwargs_only)(a=a, b=b, c=1)
+    assert isinstance(val, Delayed)
+    np_arr = np.arange(100).reshape((10, 10))
+    assert_eq(val.compute(), (np_arr + 1) + (np_arr + 2) + 1)
+
+
 def test_array_bag_delayed():
     np = pytest.importorskip("numpy")
     da = pytest.importorskip("dask.array")
@@ -612,9 +674,7 @@ def test_delayed_method_descriptor():
 def test_delayed_callable():
     f = delayed(add, pure=True)
     v = f(1, 2)
-    assert v.dask == {v.key: (add, 1, 2)}
-
-    assert f.dask == {f.key: add}
+    assert v.compute() == 3
     assert f.compute() == add
 
 
@@ -865,7 +925,7 @@ def test3(i):
 
     obj2 = test3(test2(test(10)))
     with dask.config.set({"optimization.fuse.delayed": True}):
-        dsk2 = collections_to_expr([obj]).__dask_graph__()
+        dsk2 = collections_to_expr([obj]).optimize().__dask_graph__()
         result = dask.compute(obj2)
     assert len(dsk2) == 2
     assert dask.compute(obj) == result
diff --git a/dask/tests/test_expr.py b/dask/tests/test_expr.py
index c477be558e1..af90c712e43 100644
--- a/dask/tests/test_expr.py
+++ b/dask/tests/test_expr.py
@@ -5,7 +5,8 @@
 
 import pytest
 
-from dask._expr import Expr, SingletonExpr
+from dask._expr import Expr, ProhibitReuse, SingletonExpr, _ExprSequence
+from dask._task_spec import DataNode
 
 
 class MyExpr(Expr):
@@ -141,3 +142,22 @@ def gen(i):
         del futures
 
         df.compute()
+
+
+class FooExpr(Expr):
+    def _layer(self) -> dict:
+        return {"foo": DataNode("foo", 42)}
+
+
+def test_prohibit_reuse():
+    once = FooExpr()
+    dsk = _ExprSequence(once, ProhibitReuse(once)).optimize().__dask_graph__()
+
+    assert len(dsk) == 2
+    first = dsk.pop("foo")()
+    key, val = dsk.popitem()
+    assert key.startswith("foo") and key != "foo"
+    # We don't want to chain anything but actually _hide_ the task
+    assert not val.dependencies
+    # Task is wrapped
+    assert val() is first
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index 58534b3ac5e..8fa8fd94ca8 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -16,6 +16,16 @@ Highlights
   scheduler of a distributed cluster (if available)
 - New kwarg `force` for ``DataFrame.shuffle`` which signals the optimizer to not
   drop the shuffle during optimization.
+- (Breaking; kind of) Collections that are passed to dask methods as arguments
+  are now properly optimized. If multiple collections are passed as arguments
+  they will be optimized together. Collections passed this way are prohibited
+  from being being reused, i.e. if the collection is used again in another
+  function call it will be computed again. This pattern is used to avoid
+  pipeline breakers which typically drive memory usage. Avoiding those should
+  reduce memory pressure on the cluster but can cause runtime regressions.
+- (Special case of above point) Collections passed to Delayed objects are now
+  optimized automatically.
+
 
 Breaking changes
 ^^^^^^^^^^^^^^^^


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: 

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.


<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**dask/delayed.py:109**
* The call to finalize_compute() may inadvertently trigger materialization of the collection, potentially impacting performance; consider deferring materialization until it is absolutely necessary.
```
finalized = expr.finalize_compute().optimize()
```
</details>


state: COMMENTED
author_association: NONE
submitted_at: 2025-04-09T12:26:29Z
user: login: copilot-pull-request-reviewer[bot]
body: 

Copilot reviewed 19 out of 19 changed files in this pull request and generated 1 comment.


<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**dask/delayed.py:640**
* [nitpick] The _expr slot is added in the __slots__ list of Delayed but is never initialized in __init__. If this slot is intended for future use, consider adding a comment to explain its purpose or initialize it accordingly.
```
__slots__ = ("_key", "_dask", "_length", "_layer", "_expr")
```
</details>


state: COMMENTED
author_association: NONE
submitted_at: 2025-04-11T15:33:40Z
user: login: copilot-pull-request-reviewer[bot]
body: 

Copilot reviewed 20 out of 20 changed files in this pull request and generated 1 comment.


<details>
<summary>Comments suppressed due to low confidence (2)</summary>

**dask/delayed.py:85**
* Consider adding documentation in the function docstring to explain the internal '_return_collections' parameter and its effect on the behavior of this API.
```
def unpack_collections(expr, _return_collections=True):
```
**dask/_task_spec.py:676**
* Ensure that _execute_subgraph is imported or defined within this module; otherwise, the comparison in has_subgraph may raise a NameError at runtime.
```
def has_subgraph(self) -> bool:
```
</details>


state: COMMENTED
author_association: NONE
submitted_at: 2025-04-14T12:00:39Z
user: login: fjetter
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-15T11:27:53Z

Review comments:
diff_hunk: @@ -111,7 +111,8 @@ def return_nested(obj):
 
     final = delayed(return_nested)(with_class)
 
-    assert final.compute() == 3
+    # FIXME: this is wrong
+    assert final.compute() == [3]
path: dask/tests/test_delayed.py
user: login: Copilot
body: The assertion expecting [3] instead of 3 appears to be incorrect as indicated by the FIXME comment; please verify and correct the expected behavior.
```suggestion
    assert final.compute() == 3
```
created_at: 2025-04-09T12:26:28Z
author_association: NONE
diff_hunk: @@ -111,7 +111,8 @@ def return_nested(obj):
 
     final = delayed(return_nested)(with_class)
 
-    assert final.compute() == 3
+    # FIXME: this is wrong
+    assert final.compute() == [3]
path: dask/tests/test_delayed.py
user: login: Copilot
body: The FIXME comment indicates that the expected output for final.compute() is currently incorrect. Please clarify whether the appropriate behavior is to return 3 or [3] and update the test accordingly.
```suggestion
    # The expected output should be 3, not [3]
    assert final.compute() == 3
```
created_at: 2025-04-11T15:33:39Z
author_association: NONE
diff_hunk: @@ -563,6 +567,64 @@ def test_array_delayed():
     assert (delayed_arr.compute() == arr).all()
 
 
+def test_array_delayed_complex_optimization():
+    # Ensure that collections that if multiple collections are passed to a
+    # Delayed function that they are optimized together
path: dask/tests/test_delayed.py
user: login: Copilot
body: [nitpick] Consider rewording this comment for clarity. For example, 'Ensure that when multiple collections are passed to a Delayed function, they are optimized together.'
```suggestion
    # Ensure that when multiple collections are passed to a Delayed function,
    # they are optimized together
```
created_at: 2025-04-14T12:00:39Z
author_association: NONE
diff_hunk: @@ -808,7 +808,9 @@ def test_svd():
     y = x * 2
     u, s, v = da.linalg.svd(y)
     z = y + u
-    assert_eq(z, z)
+    # FIXME: There is a key collision
+    # Deterministic https://github.com/dask/dask/issues/9888
+    assert_eq(z, z, check_graph=False)
path: dask/array/tests/test_atop.py
user: login: fjetter
body: - [ ] Don't forget about this
created_at: 2025-04-15T11:27:53Z
author_association: MEMBER

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-09T13:05:11Z
body: ## 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 18m 20s ⏱️ + 3m 26s
 18 005 tests  -  4   16 791 ✅  -  4   1 214 💤 ±0  0 ❌ ±0 
161 119 runs   - 36  149 014 ✅  - 35  12 105 💤  - 1  0 ❌ ±0 

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

<details>
  <summary>This pull request <b>removes</b> 8 and <b>adds</b> 4 tests. <i>Note that renamed tests count towards both.</i></summary>

```
dask.array.tests.test_routines ‑ test_histogram_delayed_range[False-1-False-False]
dask.array.tests.test_routines ‑ test_histogram_delayed_range[False-1-False-True]
dask.array.tests.test_routines ‑ test_histogram_delayed_range[False-1-True-False]
dask.array.tests.test_routines ‑ test_histogram_delayed_range[False-1-True-True]
dask.array.tests.test_routines ‑ test_histogram_delayed_range[True-1-False-False]
dask.array.tests.test_routines ‑ test_histogram_delayed_range[True-1-False-True]
dask.array.tests.test_routines ‑ test_histogram_delayed_range[True-1-True-False]
dask.array.tests.test_routines ‑ test_histogram_delayed_range[True-1-True-True]
```

```
dask.array.tests.test_xarray ‑ test_xarray_blockwise_fusion_store[False]
dask.array.tests.test_xarray ‑ test_xarray_blockwise_fusion_store[True]
dask.tests.test_delayed ‑ test_array_delayed_complex_optimization
dask.tests.test_delayed ‑ test_array_delayed_complex_optimization_kwargs
```
</details>

[test-results]:data:application/gzip;base64,H4sIAPYW/WcC/02O2w6DIAxAf8XwvAeKgnY/Y6RoQuZlQXha9u+rzDHfek6T077E5OdxF/cKb5XYk48FXApD9NvKCIBSsuJlPNbQSal/3O+J6JCmRfjLh38eUkFT3DT4mV0p9WMIWzhNSGtOG+BrJ5d0gzJ3vrKk8xPZXdKZr2XalsVHBkHKdIjaKKvQIGE71o1y9SB5ICDrbOec1Vq8P/s6kP0YAQAA

:recycle: This comment has been updated with latest results.
author_association: CONTRIBUTOR
user: login: fjetter
created_at: 2025-04-11T16:17:36Z
body: Looks like this is finally converging. I pulled a thread and everything unravelled...

One thing that mildly concerns me is that I can now trigger a key collision of this type https://github.com/dask/dask/issues/9888 consistently with the `test_svd`. I do want to investigate this before we release since that could cause corrupt data, deadlocks, etc. Most notably, I don't understand where this is coming from and why my `ProhibitReuse` isn't fighting that off.
author_association: MEMBER