Repository: dask/dask
Pull request #11367 (2024-09-03 10:39:41)

Ensure process_runnables is not too eager in the presence of multiple splits
Well, this isn't fixing anything visible yet but it fixes an internal branch of the algorithm that's been acting too eagerly. I'm still working on a test.

This relates to https://github.com/dask/dask/issues/11363

CPath before

![image](https://github.com/user-attachments/assets/d0efa935-3453-41b2-9ae1-58cd8adf045c)


After

![image](https://github.com/user-attachments/assets/bd3846c1-cf13-4e60-94f2-e0d4d46ba6da)

So, the initial version would start executing along the critical path until it hits the first wall. After this, it would try to process runnable tasks as good as possible. Since the root task is shared by all branches, this allows all computation branches to be considered as candidates (this is the essential dilemma of the widely shared dependencies).
Following these candidates effectively causes a depth first search. However, if we do not process them at all, this typically leaves many intermediates in memory we could otherwise release. Therefore, to allow some progress, these runnable paths are allowed to execute if they are reducing to a task, i.e. we're speculatively walking down those branches and if those speculatively executed branches reduce into a task, they are allowed to execute for real. 

What happens here is that those executable branches encounter multiple splits (we're following splits as well since they often allow us to reduce stuff). After the second split we're actually finding an intermediate reducer... This is where the algorithm on main decided to go with it. However, this leaves us still with one more task in memory than if we didn't execute that branch which is a suboptimal decision. Only if all those splits reduce, we should execute the branch.

In the case of this specific example this doesn't change anything but I believe this makes the algorithm more predictable.


<details>
<summary>Simplified graph</summary>

```python
dsk = {
        ("transpose-9c80c6b26d7ef5cfec7274ee9e6b091e", 0, 0, 0): (
            f,
            ("vindex-merge", 0, 0, 0),
        ),
        ("transpose-9c80c6b26d7ef5cfec7274ee9e6b091e", 0, 1, 0): (
            f,
            ("vindex-merge", 0, 1, 0),
        ),
        ("vindex-merge", 0, 0, 0): (
            f,
            ("vindex-slice", 0, 0, 0),
            ("vindex-slice", 1, 0, 0),
            ("vindex-slice", 2, 0, 0),
            ("vindex-slice", 3, 0, 0),
        ),
        ("vindex-slice", 2, 0, 0): (f, ("getitem-vindex", 2, 0, 0)),
        ("vindex-merge", 0, 1, 0): (
            f,
            ("vindex-slice", 0, 0, 1),
            ("vindex-slice", 1, 0, 1),
            ("vindex-slice", 2, 0, 1),
            ("vindex-slice", 3, 0, 1),
        ),
        ("vindex-slice", 2, 0, 1): (f, ("getitem-vindex", 2, 0, 1)),
        ("vindex-slice", 1, 0, 0): (f, ("getitem-vindex", 1, 0, 0)),
        ("getitem-vindex", 2, 0, 1): (
            f,
            ("shuffle-split", 321),
            ("shuffle-split", 322),
        ),
        ("vindex-slice", 3, 0, 1): (f, ("getitem-vindex", 3, 0, 1)),
        ("vindex-slice", 0, 0, 1): (f, ("getitem-vindex", 0, 0, 1)),
        ("getitem-vindex", 3, 0, 1): (
            f,
            ("shuffle-split", 299),
            ("shuffle-split", 300),
        ),
        ("getitem-vindex", 2, 0, 0): (f, ("concatenate", 0, 1, 9, 1)),
        ("vindex-slice", 3, 0, 0): (f, ("getitem-vindex", 3, 0, 0)),
        ("shuffle-split", 322): (f, ("concatenate", 0, 1, 9, 1)),
        ("vindex-slice", 0, 0, 0): (f, ("getitem-vindex", 0, 0, 0)),
        ("getitem-vindex", 0, 0, 1): (
            f,
            ("shuffle-split", 341),
            ("shuffle-split", 342),
        ),
        ("vindex-slice", 1, 0, 1): (f, ("getitem-vindex", 1, 0, 1)),
        ("shuffle-split", 321): (f, ("concatenate-getitem", 321)),
        ("shuffle-split", 299): (f, ("concatenate-getitem", 299)),
        ("getitem-vindex", 3, 0, 0): (f, ("concatenate", 0, 1, 8, 1)),
        ("getitem-vindex", 0, 0, 0): (f, ("concatenate", 0, 1, 10, 0)),
        ("getitem-vindex", 1, 0, 0): (f, ("concatenate", 0, 1, 9, 0)),
        ("concatenate-getitem", 299): (f, ("fetch-raster", 0, 0, 8, 1)),
        ("concatenate", 0, 1, 9, 1): (f, ("getitem-sub", 0, 1, 9, 1)),
        ("concatenate-getitem", 321): (f, ("fetch-raster", 0, 0, 9, 1)),
        ("fetch-raster", 0, 0, 8, 1): (f, ("asset_table", 0, 0)),
        ("concatenate", 0, 1, 8, 1): (f, ("getitem-sub", 0, 1, 8, 1)),
        ("fetch-raster", 0, 0, 9, 1): (f, ("asset_table", 0, 0)),
        ("concatenate", 0, 1, 9, 0): (f, ("getitem-sub", 0, 1, 9, 0)),
        ("shuffle-split", 300): (f, ("concatenate", 0, 1, 8, 1)),
        ("concatenate", 0, 1, 10, 0): (f, ("getitem-sub", 0, 1, 10, 0)),
        ("asset_table", 0, 0): (f, ("asset-table-data", 0, 0)),
        ("shuffle-split", 342): (f, ("concatenate", 0, 1, 10, 0)),
        ("getitem-vindex", 1, 0, 1): (
            f,
            ("shuffle-split", 319),
            ("shuffle-split", 320),
        ),
        ("getitem-sub", 0, 1, 10, 0): (f, ("fetch-raster", 0, 0, 10, 0)),
        ("getitem-sub", 0, 1, 8, 1): (f, ("fetch-raster", 0, 0, 8, 1)),
        ("shuffle-split", 341): (f, ("concatenate-getitem", 341)),
        ("getitem-sub", 0, 1, 9, 1): (f, ("fetch-raster", 0, 0, 9, 1)),
        ("concatenate-getitem", 341): (f, ("fetch-raster", 0, 0, 10, 0)),
        ("asset-table-data", 0, 0): (f,),
        ("getitem-sub", 0, 1, 9, 0): (f, ("fetch-raster", 0, 0, 9, 0)),
        ("fetch-raster", 0, 0, 10, 0): (f, ("asset_table", 0, 0)),
        ("shuffle-split", 320): (f, ("concatenate", 0, 1, 9, 0)),
        ("shuffle-split", 319): (f, ("concatenate-getitem", 319)),
        ("fetch-raster", 0, 0, 9, 0): (f, ("asset_table", 0, 0)),
        ("concatenate-getitem", 319): (f, ("fetch-raster", 0, 0, 9, 0)),
    }
```
</details>

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 4m 23s :stopwatch: + 1m 54s
 13 171 tests ±0   12 115 :white_check_mark: ±0   1 056 :zzz: ±0  0 :x: ±0 
137 272 runs  ±0  118 245 :white_check_mark:  - 1  19 027 :zzz: +1  0 :x: ±0 

Results for commit b1a7d67d. ± Comparison against base commit f5b9811d.

[test-results]:data:application/gzip;base64,H4sIAOzu1mYC/03MSw6DIBSF4a0Yxh1wrwJtN9MAFxJSHw3CyHTvlaro8P9OchbmQ+9m9mygvTVsziGdRTnqFKaxNHBZaF3TtoOCo19ztrYgAogT3+FTkAtZzevQr8YruBinuEvM43atUOHe9Rru2ImK+/WDozrscv3v67OdhiGkNZgBrUgqEq1THaEmxy1piQ7Ba03SGDS+A8G+PwpPYnsaAQAA


Patch:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index fd0a51616..b1536240c 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -207,9 +207,12 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
 
     intermediates = dict()
     merges = dict()
+    dtype = np.min_scalar_type(max(chunks[axis]))
     split_name = f"shuffle-split-{token}"
     slices = [slice(None)] * len(chunks)
     split_name_suffixes = count()
+    sorter_name = "shuffle-sorter-"
+    taker_name = "shuffle-taker-"
 
     old_blocks = np.empty([len(c) for c in chunks], dtype="O")
     for old_index in np.ndindex(old_blocks.shape):
@@ -217,7 +220,11 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
 
     for new_chunk_idx, new_chunk_taker in enumerate(new_chunks):
         new_chunk_taker = np.array(new_chunk_taker)
-        sorter = np.argsort(new_chunk_taker)
+        sorter = np.argsort(new_chunk_taker).astype(dtype)
+        sorter_key = sorter_name + tokenize(sorter)
+        # low level fusion can't deal with arrays on first position
+        merges[sorter_key] = (1, sorter)
+
         sorted_array = new_chunk_taker[sorter]
         source_chunk_nr, taker_boundary = np.unique(
             np.searchsorted(chunk_boundaries, sorted_array, side="right"),
@@ -241,16 +248,21 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
                 # Cache the takers to allow de-duplication when serializing
                 # Ugly!
                 if c in taker_cache:
-                    this_slice[axis] = taker_cache[c]
+                    taker_key = taker_cache[c]
                 else:
-                    this_slice[axis] = sorted_array[b_start:b_end] - (
-                        chunk_boundaries[c - 1] if c > 0 else 0
-                    )
+                    this_slice[axis] = (
+                        sorted_array[b_start:b_end]
+                        - (chunk_boundaries[c - 1] if c > 0 else 0)
+                    ).astype(dtype)
                     if len(source_chunk_nr) == 1:
                         this_slice[axis] = this_slice[axis][np.argsort(sorter)]
-                    taker_cache[c] = this_slice[axis]
 
-                intermediates[name] = getitem, old_blocks[chunk_key], tuple(this_slice)
+                    taker_key = taker_name + tokenize(this_slice)
+                    # low level fusion can't deal with arrays on first position
+                    intermediates[taker_key] = (1, tuple(this_slice))
+                    taker_cache[c] = taker_key
+
+                intermediates[name] = _getitem, old_blocks[chunk_key], taker_key
                 merge_keys.append(name)
 
             merge_suffix = convert_key(chunk_tuple, new_chunk_idx, axis)
@@ -258,7 +270,7 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
                 merges[(out_name,) + merge_suffix] = (
                     concatenate_arrays,
                     merge_keys,
-                    sorter,
+                    sorter_key,
                     axis,
                 )
             elif len(merge_keys) == 1:
@@ -277,9 +289,13 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
     return tuple(output_chunks), layer
 
 
+def _getitem(obj, index):
+    return getitem(obj, index[1])
+
+
 def concatenate_arrays(arrs, sorter, axis):
     concatenate = concatenate_lookup.dispatch(type(arrs[0]))
-    return np.take(concatenate(arrs, axis=axis), np.argsort(sorter), axis=axis)
+    return np.take(concatenate(arrs, axis=axis), np.argsort(sorter[1]), axis=axis)
 
 
 def convert_key(key, chunk, axis):
diff --git a/dask/array/utils.py b/dask/array/utils.py
index b515f8b84..1c7adc325 100644
--- a/dask/array/utils.py
+++ b/dask/array/utils.py
@@ -213,13 +213,18 @@ def _check_dsk(dsk):
     non_one = {k: v for k, v in freqs.items() if v != 1}
     key_collisions = set()
     # Allow redundant keys if the values are equivalent
+    collisions = set()
     for k in non_one.keys():
         for layer in dsk.layers.values():
             try:
                 key_collisions.add(tokenize(layer[k]))
             except KeyError:
                 pass
-    assert len(key_collisions) < 2, non_one
+        if len(key_collisions) >= 2:
+            collisions.add(k)
+        key_collisions.clear()
+
+    assert len(collisions) == 0, {k: non_one[k] for k in collisions}
 
 
 def assert_eq_shape(a, b, check_ndim=True, check_nan=True):
diff --git a/dask/order.py b/dask/order.py
index 19e85a056..a261f4121 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -316,9 +316,10 @@ def order(
                     add_to_result(key)
                     continue
                 path = [key]
-                branches = deque([path])
+                branches = deque([(0, path)])
+
                 while branches:
-                    path = branches.popleft()
+                    nsplits, path = branches.popleft()
                     while True:
                         # Loop invariant. Too expensive to compute at runtime
                         # assert not set(known_runnable_paths).intersection(runnable_hull)
@@ -337,6 +338,7 @@ def order(
                                 runnable_hull.discard(current)
                         elif len(path) == 1 or len(deps_upstream) == 1:
                             if len(deps_downstream) > 1:
+                                nsplits += 1
                                 for d in sorted(deps_downstream, key=sort_key):
                                     # This ensures we're only considering splitters
                                     # that are genuinely splitting and not
@@ -344,7 +346,7 @@ def order(
                                     if len(dependencies[d]) == 1:
                                         branch = path.copy()
                                         branch.append(d)
-                                        branches.append(branch)
+                                        branches.append((nsplits, branch))
                                 break
                             path.extend(deps_downstream)
                             continue
@@ -364,6 +366,13 @@ def order(
                                         pruned_branches
                                     )
                                 else:
+                                    if nsplits > 1:
+                                        path = []
+                                        for pruned in pruned_branches:
+                                            path.extend(pruned)
+                                        branches.append((nsplits - 1, path))
+                                        break
+
                                     while pruned_branches:
                                         path = pruned_branches.popleft()
                                         for k in path:


Test patch:
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 3458a443d..0b4aebbc4 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -140,3 +140,15 @@ def test_resize_other_dimensions():
     result = _rechunk_other_dimensions(arr, 4, 1, "auto")
     assert result.chunks == ((1, 1), (1,), (1, 1))
     assert_eq(arr, result)
+
+
+def test_dtype_taker(arr, darr):
+    indexer = [[1, 5, 6], [0, 3], [4, 2, 7]]
+    result = darr.shuffle(indexer, axis=1)
+    expected = arr[:, [1, 5, 6, 0, 3, 4, 2, 7]]
+    assert_eq(result, expected)
+    assert all(
+        v[1][1].dtype == np.uint8
+        for k, v in dict(result.dask).items()
+        if "shuffle-taker" in k
+    )
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
index f503c5fee..a5ea27370 100644
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -5,8 +5,6 @@ import warnings
 
 import pytest
 
-from dask.array._shuffle import concatenate_arrays
-
 np = pytest.importorskip("numpy")
 
 import dask
@@ -316,109 +314,21 @@ def test_slicing_with_newaxis():
 
 def test_take():
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [5, 1, 47, 3], axis=0)
-    expected = {
-        ("y-y", 0): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 0),
-                ("shuffle-split-y", 1),
-            ],
-            np.array([1, 3, 0, 2]),
-            0,
-        ),
-        ("shuffle-split-y", 0): (
-            getitem,
-            ("x", 0),
-            (np.array([1, 3, 5]),),
-        ),
-        ("shuffle-split-y", 1): (
-            getitem,
-            ("x", 2),
-            (np.array([7]),),
-        ),
-    }
-    np.testing.assert_equal(sorted(dsk.items()), sorted(expected.items()))
+    assert len(dsk) == 6
     assert chunks == ((4,),)
 
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20), (20, 20)], [5, 1, 47, 3], axis=0)
-    expected = {
-        ("y-y", 0, 0): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 0),
-                ("shuffle-split-y", 1),
-            ],
-            np.array([1, 3, 0, 2]),
-            0,
-        ),
-        ("y-y", 0, 1): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 2),
-                ("shuffle-split-y", 3),
-            ],
-            np.array([1, 3, 0, 2]),
-            0,
-        ),
-        ("shuffle-split-y", 0): (
-            getitem,
-            ("x", 0, 0),
-            (np.array([1, 3, 5]), slice(None)),
-        ),
-        ("shuffle-split-y", 1): (
-            getitem,
-            ("x", 2, 0),
-            (np.array([7]), slice(None)),
-        ),
-        ("shuffle-split-y", 2): (
-            getitem,
-            ("x", 0, 1),
-            (np.array([1, 3, 5]), slice(None)),
-        ),
-        ("shuffle-split-y", 3): (
-            getitem,
-            ("x", 2, 1),
-            (np.array([7]), slice(None)),
-        ),
-    }
-    np.testing.assert_equal(sorted(dsk.items()), sorted(expected.items()))
+    assert len(dsk) == 9
     assert chunks == ((4,), (20, 20))
 
 
 def test_take_sorted():
     chunks, dsk = take("y-y", "x", [(20, 20, 20, 20)], [1, 3, 5, 47], axis=0)
-    expected = {
-        ("y-y", 0): (
-            concatenate_arrays,
-            [
-                ("shuffle-split-y", 0),
-                ("shuffle-split-y", 1),
-            ],
-            np.array([0, 1, 2, 3]),
-            0,
-        ),
-        ("shuffle-split-y", 0): (
-            getitem,
-            ("x", 0),
-            (np.array([1, 3, 5]),),
-        ),
-        ("shuffle-split-y", 1): (
-            getitem,
-            ("x", 2),
-            (np.array([7]),),
-        ),
-    }
-    np.testing.assert_equal(dsk, expected)
+    assert len(dsk) == 6
     assert chunks == ((4,),)
 
     chunks, dsk = take("y", "x", [(20, 20, 20, 20)], np.arange(0, 80), axis=0)
-    expected = {
-        ("y", 0): ("x", 0),
-        ("y", 1): ("x", 1),
-        ("y", 2): ("x", 2),
-        ("y", 3): ("x", 3),
-    }
-    np.testing.assert_equal(dsk, expected)
+    assert len(dsk) == 4
     assert chunks == ((20, 20, 20, 20),)
 
 
@@ -964,23 +874,23 @@ def test_take_avoids_large_chunks():
         index = np.array([0, 1] + [2] * 101 + [3])
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
         index = np.array([0] * 101 + [1, 2, 3])
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
         index = np.array([0, 1, 2] + [3] * 101)
         chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
         chunks = ((500,), (1, 1, 1, 1), (500,))
         index = np.array([0, 1, 2] + [3] * 101)
         chunks2, dsk = take("a", "b", chunks, index, axis=1)
         assert chunks2 == ((500,), (1,) * 104, (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
 
 def test_take_uses_config():
@@ -990,7 +900,7 @@ def test_take_uses_config():
         with config.set({"array.chunk-size": "10GB"}):
             chunks2, dsk = take("a", "b", chunks, index)
         assert chunks2 == ((1,) * 104, (500,), (500,))
-        assert len(dsk) == 104
+        assert len(dsk) == 106
 
 
 def test_pathological_unsorted_slicing():
@@ -1048,6 +958,13 @@ def test_shuffle_slice(size, chunks):
     b = shuffle_slice(x, index)
     assert_eq(a, b)
 
+    index = np.arange(1, len(x)).tolist()
+    index.append(0)
+    index = np.array(index)
+    a = x[index]
+    b = shuffle_slice(x, index)
+    assert_eq(a, b)
+
 
 def test_unknown_chunks_length_one():
     a = np.arange(256, dtype=int)
diff --git a/dask/array/tests/test_xarray.py b/dask/array/tests/test_xarray.py
index 4bb0739d9..9740996d5 100644
--- a/dask/array/tests/test_xarray.py
+++ b/dask/array/tests/test_xarray.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import numpy as np
 import pytest
 
 import dask.array as da
@@ -54,3 +55,30 @@ def test_polyfit_reshaping():
     arr = xr.DataArray(da.ones((10, 20, 30), chunks=(1, 5, 30)), dims=["z", "y", "x"])
     result = arr.polyfit("x", 1)
     assert result.polyfit_coefficients.chunks == ((2,), (1,) * 10, (5,) * 4)
+
+
+def test_positional_indexer_multiple_variables():
+    n = 200
+    ds = xr.Dataset(
+        data_vars=dict(
+            a=(
+                ["x", "y", "time"],
+                da.random.randint(1, 100, (10, 10, n), chunks=(-1, -1, n // 2)),
+            ),
+            b=(
+                ["x", "y", "time"],
+                da.random.randint(1, 100, (10, 10, n), chunks=(-1, -1, n // 2)),
+            ),
+        ),
+        coords=dict(
+            x=list(range(10)),
+            y=list(range(10)),
+            time=np.arange(n),
+        ),
+    )
+    indexer = np.arange(n)
+    np.random.shuffle(indexer)
+    result = ds.isel(time=indexer)
+    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
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
index e05252c62..ea9326b52 100644
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -5523,7 +5523,6 @@ def test_broadcast():
     assert_eq(ddf - (ddf.sum() + 1), df - (df.sum() + 1))
 
 
-@pytest.mark.xfail(DASK_EXPR_ENABLED, reason="array computation doesn't work yet")
 def test_scalar_with_array():
     df = pd.DataFrame({"x": [1, 2, 3, 4, 5]})
     ddf = dd.from_pandas(df, npartitions=2)
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 008795b2f..4dda35041 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -1254,13 +1254,15 @@ def test_anom_mean():
     # `mean_chunk` which is the primary reducer in this graph. Therefore we want
     # to run those as quickly as possible.
     # This is difficult to assert on but the pressure is an ok-ish proxy
-    assert max(pressure) <= 177
+    assert max(pressure) <= 178
     from collections import defaultdict
 
     count_dependents = defaultdict(set)
     for k in dict(graph).keys():
         count_dependents[len(dependents[k])].add(k)
-    n_splits = max(count_dependents)
+
+    # array-taker has the most dependents, but it's not what we want to look at
+    n_splits = sorted(count_dependents)[-2]
     # There is a transpose/stack group that is splitting into many tasks
     # see https://github.com/dask/dask/pull/10660#discussion_r1420571664
     # the name depends on the version of xarray
@@ -1285,7 +1287,7 @@ def test_anom_mean():
     max_age_transpose = max(ages_tranpose.values())
     assert max_age_transpose < 150
     assert avg_age_transpose < 100
-    assert sum(pressure) / len(pressure) < 100
+    assert sum(pressure) / len(pressure) < 101
 
 
 def test_anom_mean_raw(abcde):
@@ -2546,3 +2548,101 @@ def test_do_not_mutate_input():
     assert_topological_sort(dsk, o)
     assert dsk == dsk_copy
     assert dependencies == dependencies_copy
+
+
+def test_stackstac():
+    # see https://github.com/dask/dask/issues/11363
+    # and https://github.com/dask/dask/pull/11367
+    final_keys = [
+        ("transpose-9c80c6b26d7ef5cfec7274ee9e6b091e", 0, 0, 0),
+        ("transpose-9c80c6b26d7ef5cfec7274ee9e6b091e", 0, 1, 0),
+    ]
+    dsk = {
+        ("transpose-9c80c6b26d7ef5cfec7274ee9e6b091e", 0, 0, 0): (
+            f,
+            ("vindex-merge", 0, 0, 0),
+        ),
+        ("transpose-9c80c6b26d7ef5cfec7274ee9e6b091e", 0, 1, 0): (
+            f,
+            ("vindex-merge", 0, 1, 0),
+        ),
+        ("vindex-merge", 0, 0, 0): (
+            f,
+            ("vindex-slice", 0, 0, 0),
+            ("vindex-slice", 1, 0, 0),
+            ("vindex-slice", 2, 0, 0),
+            ("vindex-slice", 3, 0, 0),
+        ),
+        ("vindex-slice", 2, 0, 0): (f, ("getitem-vindex", 2, 0, 0)),
+        ("vindex-merge", 0, 1, 0): (
+            f,
+            ("vindex-slice", 0, 0, 1),
+            ("vindex-slice", 1, 0, 1),
+            ("vindex-slice", 2, 0, 1),
+            ("vindex-slice", 3, 0, 1),
+        ),
+        ("vindex-slice", 2, 0, 1): (f, ("getitem-vindex", 2, 0, 1)),
+        ("vindex-slice", 1, 0, 0): (f, ("getitem-vindex", 1, 0, 0)),
+        ("getitem-vindex", 2, 0, 1): (
+            f,
+            ("shuffle-split", 321),
+            ("shuffle-split", 322),
+        ),
+        ("vindex-slice", 3, 0, 1): (f, ("getitem-vindex", 3, 0, 1)),
+        ("vindex-slice", 0, 0, 1): (f, ("getitem-vindex", 0, 0, 1)),
+        ("getitem-vindex", 3, 0, 1): (
+            f,
+            ("shuffle-split", 299),
+            ("shuffle-split", 300),
+        ),
+        ("getitem-vindex", 2, 0, 0): (f, ("concatenate", 0, 1, 9, 1)),
+        ("vindex-slice", 3, 0, 0): (f, ("getitem-vindex", 3, 0, 0)),
+        ("shuffle-split", 322): (f, ("concatenate", 0, 1, 9, 1)),
+        ("vindex-slice", 0, 0, 0): (f, ("getitem-vindex", 0, 0, 0)),
+        ("getitem-vindex", 0, 0, 1): (
+            f,
+            ("shuffle-split", 341),
+            ("shuffle-split", 342),
+        ),
+        ("vindex-slice", 1, 0, 1): (f, ("getitem-vindex", 1, 0, 1)),
+        ("shuffle-split", 321): (f, ("concatenate-getitem", 321)),
+        ("shuffle-split", 299): (f, ("concatenate-getitem", 299)),
+        ("getitem-vindex", 3, 0, 0): (f, ("concatenate", 0, 1, 8, 1)),
+        ("getitem-vindex", 0, 0, 0): (f, ("concatenate", 0, 1, 10, 0)),
+        ("getitem-vindex", 1, 0, 0): (f, ("concatenate", 0, 1, 9, 0)),
+        ("concatenate-getitem", 299): (f, ("fetch-raster", 0, 0, 8, 1)),
+        ("concatenate", 0, 1, 9, 1): (f, ("getitem-sub", 0, 1, 9, 1)),
+        ("concatenate-getitem", 321): (f, ("fetch-raster", 0, 0, 9, 1)),
+        ("fetch-raster", 0, 0, 8, 1): (f, ("asset_table", 0, 0)),
+        ("concatenate", 0, 1, 8, 1): (f, ("getitem-sub", 0, 1, 8, 1)),
+        ("fetch-raster", 0, 0, 9, 1): (f, ("asset_table", 0, 0)),
+        ("concatenate", 0, 1, 9, 0): (f, ("getitem-sub", 0, 1, 9, 0)),
+        ("shuffle-split", 300): (f, ("concatenate", 0, 1, 8, 1)),
+        ("concatenate", 0, 1, 10, 0): (f, ("getitem-sub", 0, 1, 10, 0)),
+        ("asset_table", 0, 0): (f, ("asset-table-data", 0, 0)),
+        ("shuffle-split", 342): (f, ("concatenate", 0, 1, 10, 0)),
+        ("getitem-vindex", 1, 0, 1): (
+            f,
+            ("shuffle-split", 319),
+            ("shuffle-split", 320),
+        ),
+        ("getitem-sub", 0, 1, 10, 0): (f, ("fetch-raster", 0, 0, 10, 0)),
+        ("getitem-sub", 0, 1, 8, 1): (f, ("fetch-raster", 0, 0, 8, 1)),
+        ("shuffle-split", 341): (f, ("concatenate-getitem", 341)),
+        ("getitem-sub", 0, 1, 9, 1): (f, ("fetch-raster", 0, 0, 9, 1)),
+        ("concatenate-getitem", 341): (f, ("fetch-raster", 0, 0, 10, 0)),
+        ("asset-table-data", 0, 0): (f,),
+        ("getitem-sub", 0, 1, 9, 0): (f, ("fetch-raster", 0, 0, 9, 0)),
+        ("fetch-raster", 0, 0, 10, 0): (f, ("asset_table", 0, 0)),
+        ("shuffle-split", 320): (f, ("concatenate", 0, 1, 9, 0)),
+        ("shuffle-split", 319): (f, ("concatenate-getitem", 319)),
+        ("fetch-raster", 0, 0, 9, 0): (f, ("asset_table", 0, 0)),
+        ("concatenate-getitem", 319): (f, ("fetch-raster", 0, 0, 9, 0)),
+    }
+    o = order(dsk, return_stats=True)
+    _, pressure = diagnostics(dsk)
+    assert max(pressure) <= 9, pressure
+    for k in final_keys:
+        # Ensure that we're not processing the entire graph using
+        # process_runnables (fractional values) but are using the critical path
+        assert o[k].critical_path in {1, 2}