Repository: dask/dask

Issue:
Deprecate legacy DataFrame implementation
We haven't done anything on the old implementation in a while and bugs are only fixed for the new one. I.e. the old version is not compatible with newer pandas releases and will break horribly with pandas 3 for example.

I think we should deprecate the mode and remove in a few releases. This would also making merging the repositories a lot easier since we don't have to content with the old implementation.

A real pain at the moment is that you still end up at the old API docs if you google things like

dask read_parquet, this is because the path to the docs didn't change (it is built based on the path of the function definition9


Thoughts?

Discussion:
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._

     13 files  ± 0       13 suites  ±0   3h 23m 36s :stopwatch: + 2m 19s
 13 235 tests + 1   12 172 :white_check_mark: ±0   1 062 :zzz: ±0  1 :x: +1 
138 068 runs  +10  118 972 :white_check_mark: +9  19 094 :zzz:  - 1  2 :x: +2 

For more details on these failures, see [this check](https://github.com/dask/dask/runs/31676649755).

Results for commit 9956e783. ± Comparison against base commit 786e8f0d.

[test-results]:data:application/gzip;base64,H4sIAGkPEWcC/03MSw7DIAxF0a1EjDvAJhjczVSIJBJqPhWBUdS9l/w7fMfWXUQX+nYWzwrUoxJzDuleTY4uhWlcNyJQoXJN+x2VPvdrzt5vT2Dwxnf4rCjpts6FfrUL2hinWEQWiXnc01aSPfaVBstbe8cjzZLr04709fNf9tMwhFSGYNbUGqtQWqqN6VBJbAyT8qBrr7QDDcROi+8PClW3nBoBAAA=


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"] = {
diff --git a/dask/core.py b/dask/core.py
index eb126cea4..afe799207 100644
--- a/dask/core.py
+++ b/dask/core.py
@@ -199,7 +199,7 @@ def keys_in_tasks(keys: Collection[Key], tasks: Iterable[Any], as_list: bool = F
             elif typ is dict:
                 work.extend(w.values())
             elif isinstance(w, GraphNode):
-                ret.extend(w.dependencies)
+                work.extend(w.dependencies)
             else:
                 try:
                     if w in keys:
diff --git a/dask/dataframe/__init__.py b/dask/dataframe/__init__.py
index ca9bf46a8..a64a2fb93 100644
--- a/dask/dataframe/__init__.py
+++ b/dask/dataframe/__init__.py
@@ -28,6 +28,13 @@ def _dask_expr_enabled() -> bool:
         return DASK_EXPR_ENABLED
 
     if use_dask_expr is False:
+        warnings.warn(
+            "The legacy Dask DataFrame implementation is deprecated and will "
+            "be removed in a future version. Set the configuration option "
+            "`dataframe.query-planning` to `True` or None to enable the new Dask "
+            "Dataframe implementation and silence this warning.",
+            FutureWarning,
+        )
         return (DASK_EXPR_ENABLED := False)
     try:
         import dask_expr  # noqa: F401
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index a1350694f..5e6c23a9d 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -1227,10 +1227,23 @@ def _apply_func_to_columns(df_like, prefix, func):
 
 
 def _finalize_mean(df, sum_column, count_column):
-    return df[sum_column] / df[count_column]
+    result = df[sum_column] / df[count_column]
+    return _adjust_for_arrow_na(result, df[count_column])
 
 
-def _finalize_var(df, count_column, sum_column, sum2_column, **kwargs):
+def _adjust_for_arrow_na(result, df, check_for_isna=False):
+    if isinstance(result.dtype, pd.ArrowDtype):
+        # Our mean computation results in np.nan here but pandas doesn't
+        if check_for_isna:
+            result[df.isna()] = pd.NA
+        else:
+            result[df == 0] = pd.NA
+    return result
+
+
+def _finalize_var(
+    df, count_column, sum_column, sum2_column, adjust_arrow=True, **kwargs
+):
     # arguments are being checked when building the finalizer. As of this moment,
     # we're only using ddof, and raising an error on other keyword args.
     ddof = kwargs.get("ddof", 1)
@@ -1243,13 +1256,20 @@ def _finalize_var(df, count_column, sum_column, sum2_column, **kwargs):
     div[div < 0] = 0
     result /= div
     result[(n - ddof) == 0] = np.nan
-
-    return result
+    if adjust_arrow:
+        return _adjust_for_arrow_na(result, div)
+    else:
+        return result
 
 
 def _finalize_std(df, count_column, sum_column, sum2_column, **kwargs):
-    result = _finalize_var(df, count_column, sum_column, sum2_column, **kwargs)
-    return np.sqrt(result)
+    result = _finalize_var(
+        df, count_column, sum_column, sum2_column, adjust_arrow=False, **kwargs
+    )
+    res = np.sqrt(result)
+    if res.dtype != result.dtype:
+        res = res.astype(result.dtype)
+    return _adjust_for_arrow_na(res, result, check_for_isna=True)
 
 
 def _cum_agg_aligned(part, cum_last, index, columns, func, initial):
diff --git a/pyproject.toml b/pyproject.toml
index 73d761cc8..8d3f0ac6d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -142,6 +142,7 @@ filterwarnings = [
     'ignore:The current Dask DataFrame implementation is deprecated.*:DeprecationWarning',
     "ignore:Minimal version of pyarrow will soon be increased to 14.0.1",
     "ignore:You are using pyarrow version .* pyarrow-hotfix",
+    "ignore:The legacy Dask DataFrame implementation is deprecated.*:FutureWarning",
 ]
 xfail_strict = true
 # pytest-timeout settings


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index ea9326b52..b39ba333b 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import contextlib
 import decimal
+import importlib
 import sys
 import warnings
 import weakref
@@ -6358,3 +6359,22 @@ def test_dataframe_into_delayed():
     result = delayed(delayed_func)(df)
     assert sum(map(len, result.dask.layers.values())) == 6
     result.compute()
+
+
+@pytest.mark.filterwarnings(
+    "error:The legacy Dask DataFrame implementation is deprecated.*:FutureWarning"
+)
+def test_import_raises_warning():
+    import dask
+
+    try:
+        with dask.config.set({"dataframe.query-planning": False}):
+            with pytest.raises(FutureWarning, match="The legacy"):
+                importlib.reload(dask.dataframe)
+    finally:
+        if dask.config.get("dataframe.query-planning") is False:
+            # Build without dask-expr and config is False
+            with pytest.raises(FutureWarning, match="The legacy"):
+                importlib.reload(dask.dataframe)
+        else:
+            importlib.reload(dask.dataframe)
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
index 11aaeecff..df5c0b45e 100644
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -3755,3 +3755,29 @@ def test_groupby_value_counts_all_na_partitions():
         ddf.groupby("A")["B"].value_counts(),
         df.groupby("A")["B"].value_counts(),
     )
+
+
+def test_agg_pyarrow_casts():
+    pytest.importorskip("pyarrow")
+    df = pd.DataFrame(
+        {
+            "x": range(15),
+            "y": pd.Series([pd.NA] * 10 + [1.0] * 5, dtype="double[pyarrow]"),
+            "group": ["a"] * 5 + ["b"] * 5 + ["c"] * 5,
+        }
+    )
+
+    ddf = dd.from_pandas(df, npartitions=2)
+    if PANDAS_GE_220:
+        # np.sqrt doesn't work before 2.2
+        additional_aggs = {"z": ("y", "std")}
+    else:
+        additional_aggs = {}
+
+    result = ddf.groupby("group").agg(
+        x=("y", "var"), y=("y", "mean"), **additional_aggs
+    )
+    expected = df.groupby("group").agg(
+        x=("y", "var"), y=("y", "mean"), **additional_aggs
+    )
+    assert_eq(result, expected)
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..46c3e6974 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -21,7 +21,7 @@ from dask._task_spec import (
     resolve_aliases,
 )
 from dask.base import tokenize
-from dask.core import reverse_dict
+from dask.core import keys_in_tasks, reverse_dict
 from dask.optimization import SubgraphCallable
 from dask.sizeof import sizeof
 
@@ -762,3 +762,36 @@ 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_keys_in_tasks():
+    b = Task("b", func, "1", "2")
+    b_legacy = Task("b", func, "1", "2")
+
+    a = Task("a", func, "1", b.ref())
+    a_legacy = (func, "1", "b")
+
+    for task in [a, a_legacy]:
+        assert not keys_in_tasks({"a"}, [task])
+
+        assert keys_in_tasks({"a", "b"}, [task]) == {"b"}
+
+        assert keys_in_tasks({"a", "b", "c"}, [task]) == {"b"}
+
+    for tasks in [[a, b], [a_legacy, b_legacy]]:
+        assert keys_in_tasks({"a", "b"}, tasks) == {"b"}
+
+
+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"}