Repository: dask/dask

Issue:
Make it easier to enable linear fusion for delayed objects
Delayed can be used to express embarrassingly parallel chains of operations. Certain configurations benefit from processing one chain completely before starting a new one (high memory usage for example). This is currently not possible by default because other tasks will almost certainly run before follow up tasks are scheduled. Linear fusion is the obvious solution for this which ensures that there is only a single task in the dask sense to express the whole chain. Enabling this is currently a real PITA:

```

import dask
from dask.optimization import cull, fuse
from dask.utils import ensure_dict
from dask import config, core

def fuse_delayed(dsk, keys, *args, **kwargs):
    dependencies = dsk.get_all_dependencies()
    dsk = ensure_dict(dsk)

    fuse_subgraphs = config.get("optimization.fuse.subgraphs")
    if fuse_subgraphs is None:
        fuse_subgraphs = True
    dsk, _ = fuse(
        dsk,
        keys,
        dependencies=dependencies,
        fuse_subgraphs=fuse_subgraphs,
    )
    dsk, _ = cull(dsk, keys)
    return dsk
    
dask.config.set({"delayed_optimize": fuse_delayed})
```

This is pretty ugly. Even if we ignore this, you still have to put a function into the dask config which isn't great either.

Ideally, we add a configuration like

```
dask.config.set({"enable-delayed-fusion": True})
```

Even better, we turn this on by default, but one step at a time

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 22m 36s :stopwatch: - 2m 57s
 13 236 tests + 1   12 173 :white_check_mark: ±0   1 062 :zzz: ±0   1 :x: + 1 
138 084 runs  +13  118 977 :white_check_mark: ±0  19 094 :zzz: ±0  13 :x: +13 

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

Results for commit 6351c4b6. ± Comparison against base commit 8eac2fa9.

[test-results]:data:application/gzip;base64,H4sIAGhMEWcC/03MyQ6DIBSF4VcxrLtgkqEv0wBCQqrSMKxM370QFLs83735D+D8ahN4Tog8JpCKz/daSlTZh71tjGZWqV5zv2My9isVY/oTJze+/achZHiYU35tNsDGGGIVWCWWvacFFPTcI42E5HzgmZZQjscrTS74T5uwbT7XARiZkaGaOSWsIgwJA1EVqCi3VumFYm6spgx8f0mwQrwbAQAA


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/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/dask-schema.yaml b/dask/dask-schema.yaml
index 4be821187..267bbeccf 100644
--- a/dask/dask-schema.yaml
+++ b/dask/dask-schema.yaml
@@ -244,6 +244,13 @@ properties:
               False, then the top-most key will be used. For advanced usage, a function
               to create the new name is also accepted.
 
+          delayed:
+            type: boolean
+            description: |
+              Whether to fuse adjacent delayed calls together. This can be useful to ensure 
+              that an independent chain of delayed calls is executed in a single task to
+              reduce the memory footprint on a single worker.
+
   admin:
     type: object
     properties:
diff --git a/dask/dask.yaml b/dask/dask.yaml
index 5ba27ce89..ce6b93cfe 100644
--- a/dask/dask.yaml
+++ b/dask/dask.yaml
@@ -42,6 +42,7 @@ optimization:
     max-depth-new-edges: null  # ave_width * 1.5
     subgraphs: null  # true for dask.dataframe, false for everything else
     rename-keys: true
+    delayed: false
 
 admin:
   traceback:
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/dask/delayed.py b/dask/delayed.py
index 1ff6a2dba..e43fab1da 100644
--- a/dask/delayed.py
+++ b/dask/delayed.py
@@ -22,10 +22,12 @@ from dask.base import tokenize as _tokenize
 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,
     apply,
+    ensure_dict,
     funcname,
     is_namedtuple_instance,
     methodcaller,
@@ -522,9 +524,29 @@ def right(method):
     return partial(_swap, method)
 
 
+def _fuse_delayed(dsk, keys, **kwargs):
+    dependencies = dsk.get_all_dependencies()
+    dsk = ensure_dict(dsk)
+
+    fuse_subgraphs = config.get("optimization.fuse.subgraphs")
+    if fuse_subgraphs is None:
+        fuse_subgraphs = True
+    dsk, _ = fuse(
+        dsk,
+        keys,
+        dependencies=dependencies,
+        fuse_subgraphs=fuse_subgraphs,
+    )
+    return dsk
+
+
 def optimize(dsk, keys, **kwargs):
     if not isinstance(keys, (list, set)):
         keys = [keys]
+
+    if config.get("optimization.fuse.delayed"):
+        dsk = _fuse_delayed(dsk, keys, **kwargs)
+
     if not isinstance(dsk, HighLevelGraph):
         dsk = HighLevelGraph.from_collections(id(dsk), dsk, dependencies=())
     dsk = dsk.cull(set(flatten(keys)))
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_delayed.py b/dask/tests/test_delayed.py
index b9ae588cf..4ce1dbc8e 100644
--- a/dask/tests/test_delayed.py
+++ b/dask/tests/test_delayed.py
@@ -16,6 +16,7 @@ from tlz import merge
 import dask
 import dask.bag as db
 from dask import compute
+from dask.base import collections_to_dsk
 from dask.delayed import Delayed, delayed, to_task_dask
 from dask.highlevelgraph import HighLevelGraph
 from dask.utils_test import inc
@@ -855,3 +856,28 @@ def test_delayed_function_attributes_forwarded():
     assert add.__name__ == "add"
     assert add.__doc__ == "This is a docstring"
     assert add.__wrapped__(1, 2) == 3
+
+
+def test_delayed_fusion():
+    @delayed
+    def test(i):
+        return i + 1
+
+    @delayed
+    def test2(i):
+        return i + 2
+
+    @delayed
+    def test3(i):
+        return i + 3
+
+    obj = test3(test2(test(10)))
+    dsk = dict(collections_to_dsk([obj]))
+    assert len(dsk) == 3
+
+    obj2 = test3(test2(test(10)))
+    with dask.config.set({"optimization.fuse.delayed": True}):
+        dsk2 = dict(collections_to_dsk([obj]))
+        result = dask.compute(obj2)
+    assert len(dsk2) == 2
+    assert dask.compute(obj) == result
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 f9c993dec..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
 
@@ -764,6 +764,24 @@ def test_deterministic_tokenization_respected():
         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")