Repository: dask/dask
Pull request #11440 (2024-10-17 15:15:29)

Convert dtype to int instead of np.uint8 for visualising large task graphs
- [ ] Closes #xxxx
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`


cmap returns np.uint8 since numpy 2, but this raises in _colorize for larger tasks graphs when visualising order stuff. I ran into this a few times with larger graphs before trimming them down.

This is the problem.

```
i = sum(v * 256 ** (len(t) - i - 1) for i, v in enumerate(t))
```

Don't really know how and if we should test this though :(

cc @fjetter 


Patch:
diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
index 55542bd03..9c9cf358d 100644
--- a/continuous_integration/environment-3.12.yaml
+++ b/continuous_integration/environment-3.12.yaml
@@ -71,6 +71,7 @@ dependencies:
   - python-xxhash
   - mmh3
   - jinja2
+  - matplotlib  # visualize test
   - pip
   - pip:
     - git+https://github.com/dask/distributed
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index 77c388735..3b563aced 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -407,7 +407,7 @@ _func_cache_reverse: MutableMapping = LRU(maxsize=1000)
 
 class GraphNode:
     key: KeyType
-    dependencies: set | frozenset
+    _dependencies: frozenset
 
     __slots__ = tuple(__annotations__)
 
@@ -417,6 +417,10 @@ class GraphNode:
     def copy(self):
         raise NotImplementedError
 
+    @property
+    def dependencies(self) -> frozenset:
+        return self._dependencies
+
     def _verify_values(self, values: tuple | dict) -> None:
         if not self.dependencies:
             return
@@ -442,7 +446,6 @@ _no_deps: frozenset = frozenset()
 
 class Alias(GraphNode):
     __weakref__: Any = None
-    _dependencies: set | None
     target: TaskRef
     __slots__ = tuple(__annotations__)
 
@@ -455,13 +458,7 @@ class Alias(GraphNode):
         if not isinstance(target, TaskRef):
             target = TaskRef(target)
         self.target = target
-        self._dependencies = None
-
-    @property
-    def dependencies(self):
-        if self._dependencies is None:
-            self._dependencies = {self.target.key}
-        return self._dependencies
+        self._dependencies = frozenset([target.key])
 
     def copy(self):
         return Alias(self.key, self.target)
@@ -509,7 +506,7 @@ class DataNode(GraphNode):
         self.key = key
         self.value = value
         self.typ = type(value)
-        self.dependencies = _no_deps
+        self._dependencies = _no_deps
 
     def inline(self, dsk) -> DataNode:
         return self
@@ -610,9 +607,9 @@ class Task(GraphNode):
         dependencies.update(_get_dependencies(self.args))
         dependencies.update(_get_dependencies(tuple(self.kwargs.values())))
         if dependencies:
-            self.dependencies = dependencies
+            self._dependencies = frozenset(dependencies)
         else:
-            self.dependencies = _no_deps
+            self._dependencies = _no_deps
         self._is_coro = None
         self._token = None
 
@@ -771,7 +768,7 @@ class Task(GraphNode):
     def __setstate__(self, state):
         self.key = state["key"]
         self.packed_func = state["packed_func"]
-        self.dependencies = state["dependencies"]
+        self._dependencies = state["dependencies"]
         self.kwargs = state["kwargs"]
         self.args = state["args"]
         self._is_coro = state["_is_coro"]
@@ -823,7 +820,9 @@ class DependenciesMapping(MutableMapping):
         else:
             deps = self.dsk[key].dependencies
         if self._removed:
-            deps -= self._removed
+            # deps is a frozenset but for good measure, let's not use -= since
+            # that _may_ perform an inplace mutation
+            deps = deps - self._removed
         return deps
 
     def __iter__(self):
diff --git a/dask/base.py b/dask/base.py
index 518c47d30..a5f3ccbfa 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -853,14 +853,16 @@ def visualize_dsk(
             values = o
         if maxval is None:
             maxval = max(1, max(values.values()))
-        colors = {k: _colorize(cmap(v / maxval, bytes=True)) for k, v in values.items()}
+        colors = {
+            k: _colorize(tuple(map(int, cmap(v / maxval, bytes=True))))
+            for k, v in values.items()
+        }
         if data_values is None:
-            data_values = values
             data_colors = colors
         else:
             data_colors = {
-                k: _colorize(cmap(v / maxval, bytes=True))
-                for k, v in data_values.items()
+                k: _colorize(tuple(map(int, cmap(v / maxval, bytes=True))))
+                for k, v in values.items()
             }
 
         kwargs["function_attributes"] = {


Test patch:
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 6f04138fb..bb01e3b0c 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -81,6 +81,186 @@ def visualize(dsk, suffix="", **kwargs):
     )
 
 
+def _rechunk_merge_graph():
+    return {
+        ("transpose", 0, 0, 0): (f, ("concat-groupby", 0, 0, 0)),
+        ("transpose", 0, 1, 0): (f, ("concat-groupby", 0, 1, 0)),
+        ("transpose", 1, 0, 0): (f, ("concat-groupby", 1, 0, 0)),
+        ("transpose", 1, 1, 0): (f, ("concat-groupby", 1, 1, 0)),
+        ("groupby-cohort", 0, 0, 0): (f, ("groupby-chunk", 0, 0, 0)),
+        ("groupby-cohort", 0, 0, 1): (f, ("groupby-chunk", 0, 0, 1)),
+        ("groupby-cohort", 1, 0, 0): (f, ("groupby-chunk", 1, 0, 0)),
+        ("groupby-cohort", 1, 0, 1): (f, ("groupby-chunk", 1, 0, 1)),
+        ("groupby-cohort-2", 0, 0, 0): (f, ("groupby-chunk-2", 0, 0, 0)),
+        ("groupby-cohort-2", 0, 0, 1): (f, ("groupby-chunk-2", 0, 0, 1)),
+        ("groupby-cohort-2", 1, 0, 0): (f, ("groupby-chunk-2", 1, 0, 0)),
+        ("groupby-cohort-2", 1, 0, 1): (f, ("groupby-chunk-2", 1, 0, 1)),
+        ("rechunk-merge", 3, 0, 0): (
+            f,
+            ("concat-shuffle", 4, 0, 0),
+            ("rechunk-split", 12),
+        ),
+        ("rechunk-merge", 0, 0, 0): (
+            f,
+            ("rechunk-split", 1),
+            ("concat-shuffle", 0, 0, 0),
+        ),
+        ("rechunk-merge", 3, 1, 0): (
+            f,
+            ("rechunk-split", 14),
+            ("concat-shuffle", 4, 1, 0),
+        ),
+        ("rechunk-merge", 2, 1, 0): (f, ("rechunk-split", 10), ("rechunk-split", 11)),
+        ("rechunk-split", 12): (f, ("concat-shuffle", 3, 0, 0)),
+        ("rechunk-merge", 0, 1, 0): (
+            f,
+            ("rechunk-split", 3),
+            ("concat-shuffle", 0, 1, 0),
+        ),
+        ("rechunk-merge", 1, 0, 0): (f, ("rechunk-split", 4), ("rechunk-split", 5)),
+        ("rechunk-merge", 1, 1, 0): (f, ("rechunk-split", 7), ("rechunk-split", 6)),
+        ("rechunk-split", 5): (f, ("concat-shuffle", 2, 0, 0)),
+        ("rechunk-split", 11): (f, ("concat-shuffle", 3, 1, 0)),
+        ("rechunk-merge", 2, 0, 0): (f, ("rechunk-split", 8), ("rechunk-split", 9)),
+        ("rechunk-split", 1): (f, ("concat-shuffle", 1, 0, 0)),
+        ("rechunk-split", 14): (f, ("concat-shuffle", 3, 1, 0)),
+        ("rechunk-split", 4): (f, ("concat-shuffle", 1, 0, 0)),
+        ("rechunk-split", 7): (f, ("concat-shuffle", 2, 1, 0)),
+        ("rechunk-split", 10): (f, ("concat-shuffle", 2, 1, 0)),
+        ("rechunk-split", 6): (f, ("concat-shuffle", 1, 1, 0)),
+        ("rechunk-split", 3): (f, ("concat-shuffle", 1, 1, 0)),
+        ("rechunk-split", 9): (f, ("concat-shuffle", 3, 0, 0)),
+        ("rechunk-split", 8): (f, ("concat-shuffle", 2, 0, 0)),
+        ("concat-shuffle", 0, 0, 0): (f, ("shuffle-split", 0), ("shuffle-split", 1)),
+        ("concat-shuffle", 0, 1, 0): (
+            f,
+            ("shuffle-split", 106),
+            ("shuffle-split", 107),
+        ),
+        ("concat-shuffle", 1, 0, 0): (
+            f,
+            ("shuffle-split", 4665),
+            ("shuffle-split", 4664),
+        ),
+        ("concat-shuffle", 1, 1, 0): (
+            f,
+            ("shuffle-split", 4770),
+            ("shuffle-split", 4771),
+        ),
+        ("concat-shuffle", 2, 0, 0): (
+            f,
+            ("shuffle-split", 9328),
+            ("shuffle-split", 9329),
+            ("shuffle-split", 9330),
+        ),
+        ("concat-shuffle", 2, 1, 0): (
+            f,
+            ("shuffle-split", 9487),
+            ("shuffle-split", 9488),
+            ("shuffle-split", 9489),
+        ),
+        ("concat-shuffle", 3, 0, 0): (
+            f,
+            ("shuffle-split", 16324),
+            ("shuffle-split", 16325),
+        ),
+        ("concat-shuffle", 3, 1, 0): (
+            f,
+            ("shuffle-split", 16430),
+            ("shuffle-split", 16431),
+        ),
+        ("concat-shuffle", 4, 0, 0): (
+            f,
+            ("shuffle-split", 20989),
+            ("shuffle-split", 20988),
+        ),
+        ("concat-shuffle", 4, 1, 0): (
+            f,
+            ("shuffle-split", 21094),
+            ("shuffle-split", 21095),
+        ),
+        ("shuffle-split", 9487): (f, ("getitem-2", 2, 1, 0)),
+        ("shuffle-split", 9489): (f, ("getitem-2", 14, 1, 0)),
+        ("shuffle-split", 106): (f, ("getitem-open", 106)),
+        ("shuffle-split", 4664): (f, ("getitem-2", 1, 0, 0)),
+        ("shuffle-split", 16431): (f, ("getitem-2", 15, 1, 0)),
+        ("shuffle-split", 16324): (f, ("getitem-2", 14, 0, 0)),
+        ("shuffle-split", 107): (f, ("getitem-2", 1, 1, 0)),
+        ("shuffle-split", 4665): (f, ("getitem-2", 2, 0, 0)),
+        ("shuffle-split", 4770): (f, ("getitem-2", 1, 1, 0)),
+        ("shuffle-split", 0): (f, ("getitem-open", 0)),
+        ("shuffle-split", 9328): (f, ("getitem-2", 2, 0, 0)),
+        ("shuffle-split", 9488): (f, ("getitem-open", 9488)),
+        ("shuffle-split", 16325): (f, ("getitem-2", 15, 0, 0)),
+        ("shuffle-split", 16430): (f, ("getitem-2", 14, 1, 0)),
+        ("shuffle-split", 20988): (f, ("getitem-2", 15, 0, 0)),
+        ("shuffle-split", 9329): (f, ("getitem-open", 9329)),
+        ("shuffle-split", 4771): (f, ("getitem-2", 2, 1, 0)),
+        ("shuffle-split", 1): (f, ("getitem-2", 1, 0, 0)),
+        ("shuffle-split", 20989): (f, ("getitem-open", 20989)),
+        ("shuffle-split", 9330): (f, ("getitem-2", 14, 0, 0)),
+        ("shuffle-split", 21094): (f, ("getitem-2", 15, 1, 0)),
+        ("shuffle-split", 21095): (f, ("getitem-open", 21095)),
+        ("getitem-2", 1, 0, 0): (f, ("open_dataset", 1, 0, 0)),
+        ("getitem-2", 14, 0, 0): (f, ("open_dataset", 14, 0, 0)),
+        ("getitem-2", 2, 1, 0): (f, ("open_dataset", 2, 1, 0)),
+        ("getitem-2", 15, 0, 0): (f, ("open_dataset", 15, 0, 0)),
+        ("getitem-2", 15, 1, 0): (f, ("open_dataset", 15, 1, 0)),
+        ("getitem-2", 2, 0, 0): (f, ("open_dataset", 2, 0, 0)),
+        ("getitem-2", 1, 1, 0): (f, ("open_dataset", 1, 1, 0)),
+        ("getitem-2", 14, 1, 0): (f, ("open_dataset", 14, 1, 0)),
+        ("groupby-chunk-2", 0, 0, 1): (f, ("rechunk-merge", 2, 0, 0)),
+        ("groupby-chunk-2", 0, 0, 0): (f, ("rechunk-merge", 0, 0, 0)),
+        ("concat-groupby", 0, 0, 0): (
+            f,
+            ("groupby-cohort-2", 0, 0, 0),
+            ("groupby-cohort-2", 0, 0, 1),
+        ),
+        ("groupby-chunk", 0, 0, 1): (f, ("rechunk-merge", 3, 0, 0)),
+        ("groupby-chunk", 0, 0, 0): (f, ("rechunk-merge", 1, 0, 0)),
+        ("concat-groupby", 1, 0, 0): (
+            f,
+            ("groupby-cohort", 0, 0, 0),
+            ("groupby-cohort", 0, 0, 1),
+        ),
+        ("groupby-chunk", 1, 0, 1): (f, ("rechunk-merge", 3, 1, 0)),
+        ("groupby-chunk", 1, 0, 0): (f, ("rechunk-merge", 1, 1, 0)),
+        ("concat-groupby", 1, 1, 0): (
+            f,
+            ("groupby-cohort", 1, 0, 0),
+            ("groupby-cohort", 1, 0, 1),
+        ),
+        ("open_dataset", 14, 1, 0): (f,),
+        ("groupby-chunk-2", 1, 0, 0): (f, ("rechunk-merge", 0, 1, 0)),
+        ("groupby-chunk-2", 1, 0, 1): (f, ("rechunk-merge", 2, 1, 0)),
+        ("concat-groupby", 0, 1, 0): (
+            f,
+            ("groupby-cohort-2", 1, 0, 1),
+            ("groupby-cohort-2", 1, 0, 0),
+        ),
+        ("getitem-open", 9329): (f,),
+        ("open_dataset", 2, 1, 0): (f,),
+        ("open_dataset", 15, 1, 0): (f),
+        ("getitem-open", 20989): (f,),
+        ("getitem-open", 0): (f,),
+        ("open_dataset", 1, 0, 0): (f,),
+        ("getitem-open", 9488): (f,),
+        ("getitem-open", 21095): (f,),
+        ("open_dataset", 2, 0, 0): (f,),
+        ("getitem-open", 106): (f,),
+        ("open_dataset", 1, 1, 0): (f,),
+        ("open_dataset", 14, 0, 0): (f),
+        ("open_dataset", 15, 0, 0): (f),
+    }
+
+
+def test_visualize_int_overflow():
+    pytest.importorskip("graphviz")
+    pytest.importorskip("matplotlib")
+    # regression test for https://github.com/dask/dask/pull/11440
+    visualize(_rechunk_merge_graph())
+
+
 def test_ordering_keeps_groups_together(abcde):
     a, b, c, d, e = abcde
     d = {(a, i): (f,) for i in range(4)}
@@ -2116,176 +2296,7 @@ def test_xarray_map_reduce_with_slicing():
 
 @pytest.mark.parametrize("use_longest_path", [True, False])
 def test_xarray_rechunk_map_reduce_cohorts(use_longest_path):
-    dsk = {
-        ("transpose", 0, 0, 0): (f, ("concat-groupby", 0, 0, 0)),
-        ("transpose", 0, 1, 0): (f, ("concat-groupby", 0, 1, 0)),
-        ("transpose", 1, 0, 0): (f, ("concat-groupby", 1, 0, 0)),
-        ("transpose", 1, 1, 0): (f, ("concat-groupby", 1, 1, 0)),
-        ("groupby-cohort", 0, 0, 0): (f, ("groupby-chunk", 0, 0, 0)),
-        ("groupby-cohort", 0, 0, 1): (f, ("groupby-chunk", 0, 0, 1)),
-        ("groupby-cohort", 1, 0, 0): (f, ("groupby-chunk", 1, 0, 0)),
-        ("groupby-cohort", 1, 0, 1): (f, ("groupby-chunk", 1, 0, 1)),
-        ("groupby-cohort-2", 0, 0, 0): (f, ("groupby-chunk-2", 0, 0, 0)),
-        ("groupby-cohort-2", 0, 0, 1): (f, ("groupby-chunk-2", 0, 0, 1)),
-        ("groupby-cohort-2", 1, 0, 0): (f, ("groupby-chunk-2", 1, 0, 0)),
-        ("groupby-cohort-2", 1, 0, 1): (f, ("groupby-chunk-2", 1, 0, 1)),
-        ("rechunk-merge", 3, 0, 0): (
-            f,
-            ("concat-shuffle", 4, 0, 0),
-            ("rechunk-split", 12),
-        ),
-        ("rechunk-merge", 0, 0, 0): (
-            f,
-            ("rechunk-split", 1),
-            ("concat-shuffle", 0, 0, 0),
-        ),
-        ("rechunk-merge", 3, 1, 0): (
-            f,
-            ("rechunk-split", 14),
-            ("concat-shuffle", 4, 1, 0),
-        ),
-        ("rechunk-merge", 2, 1, 0): (f, ("rechunk-split", 10), ("rechunk-split", 11)),
-        ("rechunk-split", 12): (f, ("concat-shuffle", 3, 0, 0)),
-        ("rechunk-merge", 0, 1, 0): (
-            f,
-            ("rechunk-split", 3),
-            ("concat-shuffle", 0, 1, 0),
-        ),
-        ("rechunk-merge", 1, 0, 0): (f, ("rechunk-split", 4), ("rechunk-split", 5)),
-        ("rechunk-merge", 1, 1, 0): (f, ("rechunk-split", 7), ("rechunk-split", 6)),
-        ("rechunk-split", 5): (f, ("concat-shuffle", 2, 0, 0)),
-        ("rechunk-split", 11): (f, ("concat-shuffle", 3, 1, 0)),
-        ("rechunk-merge", 2, 0, 0): (f, ("rechunk-split", 8), ("rechunk-split", 9)),
-        ("rechunk-split", 1): (f, ("concat-shuffle", 1, 0, 0)),
-        ("rechunk-split", 14): (f, ("concat-shuffle", 3, 1, 0)),
-        ("rechunk-split", 4): (f, ("concat-shuffle", 1, 0, 0)),
-        ("rechunk-split", 7): (f, ("concat-shuffle", 2, 1, 0)),
-        ("rechunk-split", 10): (f, ("concat-shuffle", 2, 1, 0)),
-        ("rechunk-split", 6): (f, ("concat-shuffle", 1, 1, 0)),
-        ("rechunk-split", 3): (f, ("concat-shuffle", 1, 1, 0)),
-        ("rechunk-split", 9): (f, ("concat-shuffle", 3, 0, 0)),
-        ("rechunk-split", 8): (f, ("concat-shuffle", 2, 0, 0)),
-        ("concat-shuffle", 0, 0, 0): (f, ("shuffle-split", 0), ("shuffle-split", 1)),
-        ("concat-shuffle", 0, 1, 0): (
-            f,
-            ("shuffle-split", 106),
-            ("shuffle-split", 107),
-        ),
-        ("concat-shuffle", 1, 0, 0): (
-            f,
-            ("shuffle-split", 4665),
-            ("shuffle-split", 4664),
-        ),
-        ("concat-shuffle", 1, 1, 0): (
-            f,
-            ("shuffle-split", 4770),
-            ("shuffle-split", 4771),
-        ),
-        ("concat-shuffle", 2, 0, 0): (
-            f,
-            ("shuffle-split", 9328),
-            ("shuffle-split", 9329),
-            ("shuffle-split", 9330),
-        ),
-        ("concat-shuffle", 2, 1, 0): (
-            f,
-            ("shuffle-split", 9487),
-            ("shuffle-split", 9488),
-            ("shuffle-split", 9489),
-        ),
-        ("concat-shuffle", 3, 0, 0): (
-            f,
-            ("shuffle-split", 16324),
-            ("shuffle-split", 16325),
-        ),
-        ("concat-shuffle", 3, 1, 0): (
-            f,
-            ("shuffle-split", 16430),
-            ("shuffle-split", 16431),
-        ),
-        ("concat-shuffle", 4, 0, 0): (
-            f,
-            ("shuffle-split", 20989),
-            ("shuffle-split", 20988),
-        ),
-        ("concat-shuffle", 4, 1, 0): (
-            f,
-            ("shuffle-split", 21094),
-            ("shuffle-split", 21095),
-        ),
-        ("shuffle-split", 9487): (f, ("getitem-2", 2, 1, 0)),
-        ("shuffle-split", 9489): (f, ("getitem-2", 14, 1, 0)),
-        ("shuffle-split", 106): (f, ("getitem-open", 106)),
-        ("shuffle-split", 4664): (f, ("getitem-2", 1, 0, 0)),
-        ("shuffle-split", 16431): (f, ("getitem-2", 15, 1, 0)),
-        ("shuffle-split", 16324): (f, ("getitem-2", 14, 0, 0)),
-        ("shuffle-split", 107): (f, ("getitem-2", 1, 1, 0)),
-        ("shuffle-split", 4665): (f, ("getitem-2", 2, 0, 0)),
-        ("shuffle-split", 4770): (f, ("getitem-2", 1, 1, 0)),
-        ("shuffle-split", 0): (f, ("getitem-open", 0)),
-        ("shuffle-split", 9328): (f, ("getitem-2", 2, 0, 0)),
-        ("shuffle-split", 9488): (f, ("getitem-open", 9488)),
-        ("shuffle-split", 16325): (f, ("getitem-2", 15, 0, 0)),
-        ("shuffle-split", 16430): (f, ("getitem-2", 14, 1, 0)),
-        ("shuffle-split", 20988): (f, ("getitem-2", 15, 0, 0)),
-        ("shuffle-split", 9329): (f, ("getitem-open", 9329)),
-        ("shuffle-split", 4771): (f, ("getitem-2", 2, 1, 0)),
-        ("shuffle-split", 1): (f, ("getitem-2", 1, 0, 0)),
-        ("shuffle-split", 20989): (f, ("getitem-open", 20989)),
-        ("shuffle-split", 9330): (f, ("getitem-2", 14, 0, 0)),
-        ("shuffle-split", 21094): (f, ("getitem-2", 15, 1, 0)),
-        ("shuffle-split", 21095): (f, ("getitem-open", 21095)),
-        ("getitem-2", 1, 0, 0): (f, ("open_dataset", 1, 0, 0)),
-        ("getitem-2", 14, 0, 0): (f, ("open_dataset", 14, 0, 0)),
-        ("getitem-2", 2, 1, 0): (f, ("open_dataset", 2, 1, 0)),
-        ("getitem-2", 15, 0, 0): (f, ("open_dataset", 15, 0, 0)),
-        ("getitem-2", 15, 1, 0): (f, ("open_dataset", 15, 1, 0)),
-        ("getitem-2", 2, 0, 0): (f, ("open_dataset", 2, 0, 0)),
-        ("getitem-2", 1, 1, 0): (f, ("open_dataset", 1, 1, 0)),
-        ("getitem-2", 14, 1, 0): (f, ("open_dataset", 14, 1, 0)),
-        ("groupby-chunk-2", 0, 0, 1): (f, ("rechunk-merge", 2, 0, 0)),
-        ("groupby-chunk-2", 0, 0, 0): (f, ("rechunk-merge", 0, 0, 0)),
-        ("concat-groupby", 0, 0, 0): (
-            f,
-            ("groupby-cohort-2", 0, 0, 0),
-            ("groupby-cohort-2", 0, 0, 1),
-        ),
-        ("groupby-chunk", 0, 0, 1): (f, ("rechunk-merge", 3, 0, 0)),
-        ("groupby-chunk", 0, 0, 0): (f, ("rechunk-merge", 1, 0, 0)),
-        ("concat-groupby", 1, 0, 0): (
-            f,
-            ("groupby-cohort", 0, 0, 0),
-            ("groupby-cohort", 0, 0, 1),
-        ),
-        ("groupby-chunk", 1, 0, 1): (f, ("rechunk-merge", 3, 1, 0)),
-        ("groupby-chunk", 1, 0, 0): (f, ("rechunk-merge", 1, 1, 0)),
-        ("concat-groupby", 1, 1, 0): (
-            f,
-            ("groupby-cohort", 1, 0, 0),
-            ("groupby-cohort", 1, 0, 1),
-        ),
-        ("open_dataset", 14, 1, 0): (f,),
-        ("groupby-chunk-2", 1, 0, 0): (f, ("rechunk-merge", 0, 1, 0)),
-        ("groupby-chunk-2", 1, 0, 1): (f, ("rechunk-merge", 2, 1, 0)),
-        ("concat-groupby", 0, 1, 0): (
-            f,
-            ("groupby-cohort-2", 1, 0, 1),
-            ("groupby-cohort-2", 1, 0, 0),
-        ),
-        ("getitem-open", 9329): (f,),
-        ("open_dataset", 2, 1, 0): (f,),
-        ("open_dataset", 15, 1, 0): (f),
-        ("getitem-open", 20989): (f,),
-        ("getitem-open", 0): (f,),
-        ("open_dataset", 1, 0, 0): (f,),
-        ("getitem-open", 9488): (f,),
-        ("getitem-open", 21095): (f,),
-        ("open_dataset", 2, 0, 0): (f,),
-        ("getitem-open", 106): (f,),
-        ("open_dataset", 1, 1, 0): (f,),
-        ("open_dataset", 14, 0, 0): (f),
-        ("open_dataset", 15, 0, 0): (f),
-    }
+    dsk = _rechunk_merge_graph()
     if use_longest_path:
         # ensure that we run through longes path True and False
         keys = [("open-dataset", i, 0, 0) for i in range(20, 35)]
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index da361701e..f9c993dec 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -762,3 +762,18 @@ def test_execute_tasks_in_graph():
 def test_deterministic_tokenization_respected():
     with pytest.raises(RuntimeError, match="deterministic"):
         tokenize(Task("key", func, object()), ensure_deterministic=True)
+
+
+def test_dependencies_mapping_doesnt_mutate_task():
+
+    t = Task("key", func, "a", "b")
+    t2 = Task("key2", func, "a", t.ref())
+
+    assert t2.dependencies == {"key"}
+
+    dsk = {t.key: t, t2.key: t2}
+    deps = DependenciesMapping(dsk)
+    del deps[t.key]
+    # The getitem is doing weird stuff and could mutate the task state
+    deps[t2.key]
+    assert t2.dependencies == {"key"}