Repository: dask/dask
Pull request #11793 (2025-02-27 11:42:37)

Migrate base.unpack_collections to Task class
xref https://github.com/dask/dask/pull/11736

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 27m 45s ⏱️ -7s
 17 805 tests ±0   16 591 ✅ ±0   1 214 💤 ±0  0 ❌ ±0 
159 319 runs  ±0  147 163 ✅ ±0  12 156 💤 ±0  0 ❌ ±0 

Results for commit a12131fd. ± Comparison against base commit 5f61e423.

[test-results]:data:application/gzip;base64,H4sIAA1ZwGcC/02OSQ6DMAwAv4Jy7gFnI+5nKpqAFBVIFZIT6t9rEHU5zlgeexNjnIZV3Bu8NWKtsTCEmvsS00IIUltDioZlH0PnWubHWr3fpTUIf/mK72MTNLuxjxO5lsWQc8qnyXU50gYV4Mmc1h1YxZLTxv7cJX3wtezTPMdCIHraUTAGSe8H1O3TGauDp4tGenQanA2gcBCfL9aN2HIYAQAA


Patch:
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index c1e7704b4..1d2a063b7 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -246,7 +246,18 @@ def blockwise(
     if not out:
         out = "{}-{}".format(
             token or utils.funcname(func).strip("_"),
-            base.tokenize(func, out_ind, argindsstr, dtype, **kwargs),
+            base.tokenize(
+                func,
+                out_ind,
+                argindsstr,
+                adjust_chunks,
+                new_axes,
+                align_arrays,
+                concatenate,
+                meta,
+                dtype,
+                **kwargs,
+            ),
         )
 
     graph = core_blockwise(
diff --git a/dask/array/core.py b/dask/array/core.py
index 41c22b802..c42c1536b 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -815,7 +815,7 @@ def map_blocks(
         )
         name = token
 
-    name = f"{name or funcname(func)}-{tokenize(func, dtype, chunks, drop_axis, new_axis, *args, **kwargs)}"
+    token = f"{name or funcname(func)}"
     new_axes = {}
 
     if isinstance(drop_axis, Number):
@@ -885,7 +885,7 @@ def map_blocks(
             *concat(argpairs),
             expected_ndim=len(out_ind),
             _func=func,
-            name=name,
+            token=token,
             new_axes=new_axes,
             dtype=dtype,
             concatenate=True,
@@ -899,7 +899,7 @@ def map_blocks(
             func,
             out_ind,
             *concat(argpairs),
-            name=name,
+            token=token,
             new_axes=new_axes,
             dtype=dtype,
             concatenate=True,
diff --git a/dask/base.py b/dask/base.py
index 3992ed3fb..b7448d2df 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -18,18 +18,12 @@ from tlz import groupby, merge
 
 from dask import config, local
 from dask._compatibility import EMSCRIPTEN
+from dask._task_spec import DataNode, Dict, List, Task, TaskRef
 from dask.core import flatten
 from dask.core import get as simple_get
-from dask.core import quote
 from dask.system import CPU_COUNT
 from dask.typing import Key, SchedulerGetCallable
-from dask.utils import (
-    apply,
-    ensure_dict,
-    is_namedtuple_instance,
-    key_split,
-    shorten_traceback,
-)
+from dask.utils import ensure_dict, is_namedtuple_instance, key_split, shorten_traceback
 
 _DistributedClient = None
 _get_distributed_client = None
@@ -493,47 +487,45 @@ def unpack_collections(*args, traverse=True):
         if is_dask_collection(expr):
             tok = tokenize(expr)
             if tok not in repack_dsk:
-                repack_dsk[tok] = (getitem, collections_token, len(collections))
+                repack_dsk[tok] = Task(
+                    tok, getitem, TaskRef(collections_token), len(collections)
+                )
                 collections.append(expr)
-            return tok
+            return TaskRef(tok)
 
         tok = uuid.uuid4().hex
+        tsk: DataNode | Task
         if not traverse:
-            tsk = quote(expr)
+            tsk = DataNode(None, expr)
         else:
             # Treat iterators like lists
             typ = list if isinstance(expr, Iterator) else type(expr)
             if typ in (list, tuple, set):
-                tsk = (typ, [_unpack(i) for i in expr])
+                tsk = Task(tok, typ, List(*[_unpack(i) for i in expr]))
             elif typ in (dict, OrderedDict):
-                tsk = (typ, [[_unpack(k), _unpack(v)] for k, v in expr.items()])
+                tsk = Task(
+                    tok, typ, Dict({_unpack(k): _unpack(v) for k, v in expr.items()})
+                )
             elif dataclasses.is_dataclass(expr) and not isinstance(expr, type):
-                tsk = (
-                    apply,
+                tsk = Task(
+                    tok,
                     typ,
-                    (),
-                    (
-                        dict,
-                        [
-                            [f.name, _unpack(getattr(expr, f.name))]
-                            for f in dataclasses.fields(expr)
-                        ],
-                    ),
+                    *[_unpack(getattr(expr, f.name)) for f in dataclasses.fields(expr)],
                 )
             elif is_namedtuple_instance(expr):
-                tsk = (typ, *[_unpack(i) for i in expr])
+                tsk = Task(tok, typ, *[_unpack(i) for i in expr])
             else:
                 return expr
 
         repack_dsk[tok] = tsk
-        return tok
+        return TaskRef(tok)
 
     out = uuid.uuid4().hex
-    repack_dsk[out] = (tuple, [_unpack(i) for i in args])
+    repack_dsk[out] = Task(out, tuple, List(*[_unpack(i) for i in args]))
 
     def repack(results):
         dsk = repack_dsk.copy()
-        dsk[collections_token] = quote(results)
+        dsk[collections_token] = DataNode(collections_token, results)
         return simple_get(dsk, out)
 
     # The original `collections` is kept alive by the closure


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
index dbc79c970..c894ca695 100644
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -1857,6 +1857,16 @@ def test_map_blocks_no_array_args():
     assert_eq(x, np.arange(8, dtype=np.float32))
 
 
+def test_map_blocks_unique_name_enforce_dim():
+    def func(some_3d, block_info=None):
+        return some_3d
+
+    input_arr = da.zeros((3, 4, 5), chunks=((3,), (4,), (5,)), dtype=np.float32)
+    x = da.map_blocks(func, input_arr, enforce_ndim=True, dtype=np.float32)
+    y = da.map_blocks(func, input_arr, enforce_ndim=False, dtype=np.float32)
+    assert x._name != y._name
+
+
 def test_map_blocks_unique_name_chunks_dtype():
     def func(block_info=None):
         loc = block_info[None]["array-location"]