Repository: dask/dask (python)

Issue:
Automatically rechunk in array-shuffle if groups are too large
shuffling an array can create big chunks along the shuffled dimension. Keeping the other chunk sizes constant will potentially blow up chunk sizes along the way.

The proper solution here is to rechunk on the other dimensions to be able to ensure that chunk sizes stay reasonable.

We should be smart here to ensure that the rechunking is only a splitting of existing chunks and not an all to all communication is possible, even if we don't end up with ideal chunk sizes.

Additionally, we should expose a user option to give us explicit chunks if needed (although I hope that we can make the default smart enough that this is rarely needed)

cc @dcherian 

Pull request:
Implement automatic rechunking for shuffle
- [ ] Closes #11282
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

This makes the shuffle automatically split the other dimensions by a dynamic factor so that the maximum chunksize stays the same after the shuffle. The implementation ensures that we only split chunks along the other dimensions so that we can avoid different input chunks having to communicate with each other

cc @dcherian is this roughly what you imagined? We can add the option to customise how we split as a follow up

Patch:
diff --git a/continuous_integration/environment-3.12.yaml b/continuous_integration/environment-3.12.yaml
index 425969ef5..341fa3722 100644
--- a/continuous_integration/environment-3.12.yaml
+++ b/continuous_integration/environment-3.12.yaml
@@ -23,13 +23,14 @@ dependencies:
   - moto<5
   # Optional dependencies
   - mimesis
-  - numpy
+  - numpy >=2  # only tested here
   - pandas
   - numba
   - flask
   - fastparquet
   - h5py
-  - pytables
+  # Temporarily removing to allow `numpy >=2` to be installed
+  # - pytables
   - zarr
   # `tiledb-py=0.17.5` lead to strange seg faults in CI, However 0.18 is needed for 3.11
   # https://github.com/dask/dask/pull/9569
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 6e3dd2e4a..fa1c6ccd6 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -1,7 +1,9 @@
 from __future__ import annotations
 
 import copy
+import math
 from itertools import count, product
+from typing import Literal
 
 import numpy as np
 
@@ -13,7 +15,7 @@ from dask.base import tokenize
 from dask.highlevelgraph import HighLevelGraph
 
 
-def shuffle(x, indexer: list[list[int]], axis):
+def shuffle(x, indexer: list[list[int]], axis: int, chunks: Literal["auto"] = "auto"):
     """
     Reorders one dimensions of a Dask Array based on an indexer.
 
@@ -34,6 +36,12 @@ def shuffle(x, indexer: list[list[int]], axis):
         each group will end up in exactly one chunk.
     axis: int
         The axis to shuffle along.
+    chunks: "auto"
+        Hint on how to rechunk if single groups are becoming too large. The default is
+        to split chunks along the other dimensions evenly to keep the chunksize
+        consistent. The rechunking is done in a way that ensures that non all-to-all
+        network communication is necessary, chunks are only split and not combined with
+        other chunks.
 
     Examples
     --------
@@ -66,6 +74,9 @@ def shuffle(x, indexer: list[list[int]], axis):
             f"Shuffling only allowed with known chunk sizes. {unknown_chunk_message}"
         )
     assert isinstance(axis, int), "axis must be an integer"
+    _validate_indexer(x.chunks, indexer, axis)
+
+    x = _rechunk_other_dimensions(x, max(map(len, indexer)), axis, chunks)
 
     token = tokenize(x, indexer, axis)
     out_name = f"shuffle-{token}"
@@ -79,7 +90,54 @@ def shuffle(x, indexer: list[list[int]], axis):
     return Array(graph, out_name, chunks, meta=x)
 
 
-def _shuffle(chunks, indexer, axis, in_name, out_name, token):
+def _rechunk_other_dimensions(
+    x: Array, longest_group: int, axis: int, chunks: Literal["auto"]
+) -> Array:
+    assert chunks == "auto", "Only auto is supported for now"
+    chunksize_tolerance = config.get("array.shuffle.chunksize-tolerance")
+
+    if longest_group <= max(x.chunks[axis]) * chunksize_tolerance:
+        # We are staying below our threshold, so don't rechunk
+        return x
+
+    # How much the chunksizes on our shuffle axis will increase
+    chunksize_inc_factor = longest_group / max(x.chunks[axis])
+
+    new_chunks = []
+    for i in range(len(x.chunks)):
+        if i == axis:
+            new_chunks.append(x.chunks[i])
+            continue
+
+        new_chunksizes = []
+        # calculate what the max chunk size in this dimension is and split every
+        # chunk that is larger than that. We split the increase factor evenly
+        # between all dimensions that are not shuffled.
+        up_chunksize_limit_for_dim = max(x.chunks[i]) / (
+            chunksize_inc_factor ** (1 / (len(x.chunks) - 1))
+        )
+        for c in x.chunks[i]:
+            if c > chunksize_tolerance * up_chunksize_limit_for_dim:
+                factor = math.ceil(c / up_chunksize_limit_for_dim)
+
+                # Ensure that we end up at least with chunksize 1
+                factor = min(factor, c)
+
+                chunksize, remainder = divmod(c, factor)
+                nc = [chunksize] * factor
+                for ii in range(remainder):
+                    # Add remainder parts to the first few chunks
+                    nc[ii] += 1
+                new_chunksizes.extend(nc)
+
+            else:
+                new_chunksizes.append(c)
+        new_chunks.append(tuple(new_chunksizes))
+
+    return x.rechunk(tuple(new_chunks))
+
+
+def _validate_indexer(chunks, indexer, axis):
     if not isinstance(indexer, list) or not all(isinstance(i, list) for i in indexer):
         raise ValueError("indexer must be a list of lists of positional indices")
 
@@ -93,6 +151,10 @@ def _shuffle(chunks, indexer, axis, in_name, out_name, token):
             f"Indexer contains out of bounds index. Dimension only has {sum(chunks[axis])} elements."
         )
 
+
+def _shuffle(chunks, indexer, axis, in_name, out_name, token):
+    _validate_indexer(chunks, indexer, axis)
+
     if len(indexer) == len(chunks[axis]):
         # check if the array is already shuffled the way we want
         ctr = 0
diff --git a/dask/array/blockwise.py b/dask/array/blockwise.py
index 4a2a7aed7..300120e8b 100644
--- a/dask/array/blockwise.py
+++ b/dask/array/blockwise.py
@@ -122,7 +122,7 @@ def blockwise(
 
     >>> z = blockwise(sequence_dot, '', a, 'i', b, 'i', dtype='f8')
     >>> z.compute()
-    250
+    np.int64(250)
 
     Add new single-chunk dimensions with the ``new_axes=`` keyword, including
     the length of the new dimension.  New dimensions will always be in a single
diff --git a/dask/array/core.py b/dask/array/core.py
index 3009dcb11..a35c05f76 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -17,7 +17,7 @@ from itertools import product, zip_longest
 from numbers import Integral, Number
 from operator import add, mul
 from threading import Lock
-from typing import Any, TypeVar, Union, cast
+from typing import Any, Literal, TypeVar, Union, cast
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -2765,7 +2765,8 @@ class Array(DaskMethodsMixin):
     def shuffle(
         self,
         indexer: list[list[int]],
-        axis,
+        axis: int,
+        chunks: Literal["auto"] = "auto",
     ):
         """Reorders one dimensions of a Dask Array based on an indexer.
 
@@ -2777,7 +2778,7 @@ class Array(DaskMethodsMixin):
         """
         from dask.array._shuffle import shuffle
 
-        return shuffle(self, indexer, axis)
+        return shuffle(self, indexer, axis, chunks)
 
     @property
     def real(self):
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
index 76e8680a3..ff4e2de03 100644
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -613,11 +613,11 @@ def map_overlap(
     >>> x = da.ones(10, dtype='int')
     >>> block_args = dict(chunks=(), drop_axis=0)
     >>> da.map_blocks(func, x, **block_args).compute()
-    10
+    np.int64(10)
     >>> da.map_overlap(func, x, **block_args, boundary='reflect').compute()
-    10
+    np.int64(10)
     >>> da.map_overlap(func, x, **block_args, depth=1, boundary='reflect').compute()
-    12
+    np.int64(12)
 
     For functions that may not handle 0-d arrays, it's also possible to specify
     ``meta`` with an empty array matching the type of the expected result. In
diff --git a/dask/array/routines.py b/dask/array/routines.py
index 38fafa6da..4780f71d7 100644
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -1312,7 +1312,7 @@ def histogramdd(sample, bins, range=None, normed=None, weights=None, density=Non
     >>> w = da.random.uniform(0, 1, size=(1000,), chunks=x.chunksize[0])
     >>> h, edges = da.histogramdd(x, bins=bins, range=ranges, weights=w)
     >>> np.isclose(h.sum().compute(), w.sum().compute())
-    True
+    np.True_
 
     Using a sequence of 1D arrays as the input:
 
@@ -2270,11 +2270,11 @@ def aligned_coarsen_chunks(chunks: list[int], multiple: int) -> tuple[int, ...]:
     Examples
     --------
     >>> aligned_coarsen_chunks(chunks=(1, 2, 3), multiple=4)
-    (4, 2)
+    (np.int64(4), np.int64(2))
     >>> aligned_coarsen_chunks(chunks=(1, 20, 3, 4), multiple=4)
-    (4, 20, 4)
+    (np.int64(4), np.int64(20), np.int64(4))
     >>> aligned_coarsen_chunks(chunks=(20, 10, 15, 23, 24), multiple=10)
-    (20, 10, 20, 20, 20, 2)
+    (np.int64(20), np.int64(10), np.int64(20), np.int64(20), np.int64(20), np.int64(2))
     """
     overflow = np.array(chunks) % multiple
     excess = overflow.sum()
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
index 0a0421fd7..a1bf99e85 100644
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -531,9 +531,9 @@ def issorted(seq):
     """Is sequence sorted?
 
     >>> issorted([1, 2, 3])
-    True
+    np.True_
     >>> issorted([3, 1, 2])
-    False
+    np.False_
     """
     if len(seq) == 0:
         return True
diff --git a/dask/base.py b/dask/base.py
index 7343152c8..ecd8a98a5 100644
--- a/dask/base.py
+++ b/dask/base.py
@@ -579,9 +579,9 @@ def optimize(*args, traverse=True, **kwargs):
     >>> a2, b2 = dask.optimize(a, b)
 
     >>> a2.compute() == a.compute()
-    True
+    np.True_
     >>> b2.compute() == b.compute()
-    True
+    np.True_
     """
     collections, repack = unpack_collections(*args, traverse=traverse)
     if not collections:
@@ -634,12 +634,12 @@ def compute(
     >>> a = da.arange(10, chunks=2).sum()
     >>> b = da.arange(10, chunks=2).mean()
     >>> dask.compute(a, b)
-    (45, 4.5)
+    (np.int64(45), np.float64(4.5))
 
     By default, dask objects inside python collections will also be computed:
 
     >>> dask.compute({'a': a, 'b': b, 'c': 1})
-    ({'a': 45, 'b': 4.5, 'c': 1},)
+    ({'a': np.int64(45), 'b': np.float64(4.5), 'c': 1},)
     """
 
     collections, repack = unpack_collections(*args, traverse=traverse)
diff --git a/dask/dataframe/backends.py b/dask/dataframe/backends.py
index 1ad06959c..90ea18868 100644
--- a/dask/dataframe/backends.py
+++ b/dask/dataframe/backends.py
@@ -282,7 +282,7 @@ def make_meta_object(x, index=None):
     >>> make_meta_object(('a', 'f8'))
     Series([], Name: a, dtype: float64)
     >>> make_meta_object('i8')
-    1
+    np.int64(1)
     """
 
     if is_arraylike(x) and x.shape:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 072f7af0b..39fab0858 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1408,7 +1408,7 @@ Dask Name: {name}, {layers}"""
         >>> res = ddf.x.reduction(count_greater, aggregate=lambda x: x.sum(),
         ...                       chunk_kwargs={'value': 25})
         >>> res.compute()
-        25
+        np.int64(25)
 
         Aggregate both the sum and count of a Series at the same time:
 
diff --git a/dask/dataframe/partitionquantiles.py b/dask/dataframe/partitionquantiles.py
index 9ed891e57..ca82a0d42 100644
--- a/dask/dataframe/partitionquantiles.py
+++ b/dask/dataframe/partitionquantiles.py
@@ -252,13 +252,13 @@ def percentiles_to_weights(qs, vals, length):
     between the first and second percentiles, and then scaled by length:
 
     >>> 0.5 * length * (percentiles[1] - percentiles[0])
-    125.0
+    np.float64(125.0)
 
     The second weight uses the difference of percentiles on both sides, so
     it will be twice the first weight if the percentiles are equally spaced:
 
     >>> 0.5 * length * (percentiles[2] - percentiles[0])
-    250.0
+    np.float64(250.0)
     """
     if length == 0:
         return ()
diff --git a/dask/order.py b/dask/order.py
index 4de05c2b4..62fb678b2 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -208,6 +208,9 @@ def order(
 
     _sort_keys_cache: dict[Key, tuple[int, int, int, int, str]] = {}
 
+    leafs_connected_to_loaded_roots: set[Key] = set()
+    processed_roots = set()
+
     def sort_key(x: Key) -> tuple[int, int, int, int, str]:
         try:
             return _sort_keys_cache[x]
@@ -246,6 +249,10 @@ def order(
                 result[item] = Order(i, crit_path_counter - _crit_path_counter_offset)
             else:
                 result[item] = i
+
+            if item in root_nodes:
+                processed_roots.add(item)
+
             i += 1
             # Note: This is a `set` and therefore this introduces a certain
             # randomness. However, this randomness should not have any impact on
@@ -386,6 +393,43 @@ def order(
     # writing, the most expensive part of ordering is the prep work (mostly
     # connected roots + sort_key) which can be reused for multiple orderings.
 
+    def get_target(longest_path: bool = False) -> Key:
+        # Some topologies benefit if the node with the most dependencies
+        # is used as first choice, others benefit from the opposite.
+        candidates = leaf_nodes
+        skey: Callable = sort_key
+        if reachable_hull:
+            skey = lambda k: (num_needed[k], sort_key(k))
+
+        all_leafs_accessible = len(leafs_connected_to_loaded_roots) >= len(candidates)
+        if reachable_hull and not all_leafs_accessible:
+            # Avoid this branch if we can already access all leafs. Just pick
+            # one of them without the expensive selection process in this case.
+
+            candidates = reachable_hull & candidates
+            if not candidates:
+                # We can't reach a leaf node directly, but we still have nodes
+                # with results in memory, these notes can inform our path towards
+                # a new preferred leaf node.
+                for r in processed_roots:
+                    leafs_connected_to_loaded_roots.update(leafs_connected[r])
+                processed_roots.clear()
+
+                leafs_connected_to_loaded_roots.intersection_update(leaf_nodes)
+                candidates = leafs_connected_to_loaded_roots
+            else:
+                leafs_connected_to_loaded_roots.update(candidates)
+
+        elif not reachable_hull:
+            # We reach a new and independent branch, so throw away previous branch
+            leafs_connected_to_loaded_roots.clear()
+
+        if longest_path and (not reachable_hull or all_leafs_accessible):
+            return leaf_nodes_sorted.pop()
+        else:
+            # FIXME: This can be very expensive
+            return min(candidates, key=skey)
+
     def use_longest_path() -> bool:
         size = 0
         # Heavy reducer / splitter topologies often benefit from a very
@@ -404,65 +448,9 @@ def order(
         return True
 
     longest_path = use_longest_path()
-
-    def get_target() -> Key:
-        raise NotImplementedError()
-
-    if not longest_path:
-
-        def _build_get_target() -> Callable[[], Key]:
-            occurrences: defaultdict[Key, int] = defaultdict(int)
-            for t in leaf_nodes:
-                for r in roots_connected[t]:
-                    occurrences[r] += 1
-            occurences_grouped = defaultdict(set)
-            for root, occ in occurrences.items():
-                occurences_grouped[occ].add(root)
-            occurences_grouped_sorted = {}
-            for k, v in occurences_grouped.items():
-                occurences_grouped_sorted[k] = sorted(v, key=sort_key, reverse=True)
-            del occurences_grouped, occurrences
-
-            def pick_seed() -> Key | None:
-                while occurences_grouped_sorted:
-                    key = max(occurences_grouped_sorted)
-                    picked_root = occurences_grouped_sorted[key][-1]
-                    if picked_root in result:
-                        occurences_grouped_sorted[key].pop()
-                        if not occurences_grouped_sorted[key]:
-                            del occurences_grouped_sorted[key]
-                        continue
-                    return picked_root
-                return None
-
-            def get_target() -> Key:
-                candidates = leaf_nodes
-                skey: Callable = sort_key
-
-                # We're not considering runnable_hull because if there was a
-                # runnable leaf node, process_runnables should've taken care of
-                # it already, i.e. the intersection of runnable_hull and
-                # candidates is always empty
-                if reachable_hull:
-                    skey = lambda k: (num_needed[k], sort_key(k))
-                    candidates = reachable_hull & candidates
-
-                if not candidates:
-                    if seed := pick_seed():
-                        candidates = leafs_connected[seed]
-                    else:
-                        # FIXME: This seems to be dead code (at least untested)
-                        candidates = runnable_hull or reachable_hull
-                # FIXME: This can be very expensive
-                return min(candidates, key=skey)
-
-            return get_target
-
-        get_target = _build_get_target()
-    else:
+    leaf_nodes_sorted = []
+    if longest_path:
         leaf_nodes_sorted = sorted(leaf_nodes, key=sort_key, reverse=False)
-        get_target = leaf_nodes_sorted.pop
-        del leaf_nodes_sorted
 
     # *************************************************************************
     # CORE ALGORITHM STARTS HERE
@@ -497,7 +485,7 @@ def order(
     #   can define any route through the graph that should be considered as top
     #   priority.
     #
-    #   1. Determine the target node by calling `get_target`` and append the
+    #   1. Determine the target node by calling ``get_target`` and append the
     #      target to the critical path stack
     #   2. Take the _most valuable_ (max given a `sort_key`) of its dependents
     #      and append it to the critical path stack. This key is the new target.
@@ -562,7 +550,7 @@ def order(
         assert not scrit_path
 
         # A. Build the critical path
-        target = get_target()
+        target = get_target(longest_path=longest_path)
         next_deps = dependencies[target]
         path_append(target)
 
diff --git a/docs/source/develop.rst b/docs/source/develop.rst
index bf48d4d41..57880f998 100644
--- a/docs/source/develop.rst
+++ b/docs/source/develop.rst
@@ -254,7 +254,7 @@ after the line.
 
 .. _numpydoc: https://numpydoc.readthedocs.io/en/latest/format.html#docstring-standard
 
-Docstrings are tested under Python 3.11 on GitHub Actions. You can test
+Docstrings are tested under Python 3.12 on GitHub Actions. You can test
 docstrings with pytest as follows::
 
    py.test dask --doctest-modules


Test patch:
diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
index dc1b646b8..02eea953f 100644
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -1172,3 +1172,22 @@ def test_old_to_new_with_zero():
         [[(0, slice(0, 4))], [(2, slice(0, 0))], [(2, slice(0, 2))], [(2, slice(2, 4))]]
     ]
     assert result == expected
+
+
+def test_rechunk_non_perfect_slicing_of_dimensions():
+    # GH#7859
+    # this matters -- 1060 and 1058 work
+    shape = (200, 100, 1059)
+    final_chunks = (64, 64, 64)
+
+    arr = da.coarsen(
+        da.mean,
+        da.zeros(shape, chunks=(1, -1, -1)),
+        {0: 2, 1: 2, 2: 2},
+        trim_excess=True,
+    )
+    result = arr.rechunk(*final_chunks)
+    assert_eq(arr, result)
+    result_b = arr.rechunk(final_chunks)
+    assert_eq(arr, result)
+    assert result.chunks == result_b.chunks
diff --git a/dask/array/tests/test_reshape.py b/dask/array/tests/test_reshape.py
index ebd2eae39..bc9dfa461 100644
--- a/dask/array/tests/test_reshape.py
+++ b/dask/array/tests/test_reshape.py
@@ -323,3 +323,10 @@ def test_reshape_split_out_chunks():
         (1, 1, 1, 1),
         (80, 80, 80, 80, 80, 80, 80, 80, 80, 80, 80, 80, 80, 80, 80),
     )
+
+
+def test_argwhere_reshaping_not_concating_chunks():
+    # GH#10080
+    arr = da.random.random((500, 500, 500), chunks=(100, 100, 100)) < 0
+    result = da.argwhere(arr)
+    assert result.chunks == ((np.nan,) * 125, (1, 1, 1))
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 2882cbc87..86cadcfc2 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -6,6 +6,7 @@ import pytest
 import dask
 import dask.array as da
 from dask.array import assert_eq, shuffle
+from dask.array._shuffle import _rechunk_other_dimensions
 from dask.core import flatten
 
 
@@ -20,19 +21,19 @@ def darr(arr):
 
 
 @pytest.mark.parametrize(
-    "indexer, chunks",
+    "indexer, chunks, other_chunks",
     [
-        ([[1, 5, 6], [0, 2, 3, 4, 7]], (3, 5)),
-        ([[1, 5, 6], [0, 3], [4, 2, 7]], (5, 3)),
-        ([[1], [0, 6, 5, 3, 2, 4], [7]], (1, 6, 1)),
-        ([[1, 5, 1, 5, 1, 5], [1, 6, 4, 2, 7]], (6, 5)),
+        ([[1, 5, 6], [0, 2, 3, 4, 7]], (3, 5), (2, 1)),
+        ([[1, 5, 6], [0, 3], [4, 2, 7]], (5, 3), (2, 1)),
+        ([[1], [0, 6, 5, 3, 2, 4], [7]], (1, 6, 1), (1, 1, 1)),
+        ([[1, 5, 1, 5, 1, 5], [1, 6, 4, 2, 7]], (6, 5), (1, 1, 1)),
     ],
 )
-def test_shuffle(arr, darr, indexer, chunks):
+def test_shuffle(arr, darr, indexer, chunks, other_chunks):
     result = darr.shuffle(indexer, axis=1)
     expected = arr[:, list(flatten(indexer))]
     assert_eq(result, expected)
-    assert result.chunks[0] == darr.chunks[0]
+    assert result.chunks[0] == other_chunks
     assert result.chunks[1] == chunks
 
 
@@ -94,3 +95,24 @@ def test_shuffle_no_op_with_correct_indexer():
     result = arr.shuffle(indexer, axis=0)
     assert result.dask == arr.dask
     assert_eq(arr, result)
+
+
+def test_resize_other_dimensions():
+    arr = da.random.random((250, 50), chunks=((45, 100, 38, 67), 10))
+    result = _rechunk_other_dimensions(arr, 20, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 34, 33), (10,) * 5)
+    assert_eq(arr, result)
+
+    arr = da.random.random((250, 50, 20), chunks=((45, 100, 38, 67), 10, 10))
+    result = _rechunk_other_dimensions(arr, 20, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 67), (10,) * 5, (5,) * 4)
+    assert_eq(arr, result)
+
+    result = _rechunk_other_dimensions(arr, 40, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 34, 33), (10,) * 5, (5,) * 4)
+    assert_eq(arr, result)
+
+    arr = da.random.random((250, 50, 5), chunks=((45, 100, 38, 67), 10, 1))
+    result = _rechunk_other_dimensions(arr, 40, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 34, 33), (10,) * 5, (1,) * 5)
+    assert_eq(arr, result)
diff --git a/dask/tests/test_order.py b/dask/tests/test_order.py
index 0a456cbcc..bea8caf73 100644
--- a/dask/tests/test_order.py
+++ b/dask/tests/test_order.py
@@ -256,7 +256,7 @@ def test_prefer_deep(abcde):
     |   |
     d   a
 
-    Prefer longer chains first so we should start with c
+    Prefer longer chains first so we should start with d
     """
     a, b, c, d, e = abcde
     dsk = {a: 1, b: (f, a), c: (f, b), d: 1, e: (f, d)}
@@ -339,9 +339,6 @@ def test_favor_longest_critical_path(abcde):
     assert o[e] > o[b]
 
 
-@pytest.mark.xfail(
-    reason="Second target pick does not include already computed nodes properly"
-)
 def test_run_smaller_sections(abcde):
     r"""
             aa
@@ -578,7 +575,7 @@ def test_map_overlap(abcde):
        |/  | \ | / | \|
       d1  d2  d3  d4  d5
        |       |      |
-      e1      e2      e5
+      e1      e3      e5
 
     Want to finish b1 before we start on e5
     """
@@ -1198,8 +1195,6 @@ def test_xarray_like_reduction():
 def test_array_vs_dataframe(optimize):
     xr = pytest.importorskip("xarray")
     dd = pytest.importorskip("dask.dataframe")
-    if dd._dask_expr_enabled():
-        pytest.xfail("doesn't work yet")
 
     import dask.array as da
 
@@ -1224,7 +1219,8 @@ def test_array_vs_dataframe(optimize):
     diag_df = diagnostics(
         collections_to_dsk([mean.to_dask_dataframe()], optimize_graph=optimize)
     )
-    assert max(diag_df[1]) == max(diag_array[1])
+    assert max(diag_df[1]) == (15 if dd._dask_expr_enabled() else 38)
+    assert max(diag_array[1]) == 38
     assert max(diag_array[1]) < 50
 
 
@@ -2100,6 +2096,204 @@ def test_xarray_map_reduce_with_slicing():
     assert max(pressure) <= 5
 
 
+@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),
+    }
+    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)]
+        dsk.update({("dummy", 0): (f, keys)})
+        dsk.update({k: (f,) for k in keys})
+
+    o = order(dsk)
+
+    assert_topological_sort(dsk, o)
+    _, pressure = diagnostics(dsk, o=o)
+    # cut the dummy tasks in the end
+    assert max(pressure[:99]) <= 7
+
+    final_nodes = sorted(
+        [("transpose", ix, jx, 0) for ix in range(2) for jx in range(2)],
+        key=o.__getitem__,
+    )
+    all_diffs = []
+    for ix in range(1, len(final_nodes)):
+        all_diffs.append(o[final_nodes[ix]] - o[final_nodes[ix - 1]])
+
+    # We process a big chunk first and then a small side-branch
+    # before we repeat this for the next independent branch
+    assert all_diffs == [10, 39, 10]
+
+
 def test_xarray_8414():
     # https://github.com/pydata/xarray/issues/8414#issuecomment-1793860552
     np = pytest.importorskip("numpy")
diff --git a/dask/tests/test_tokenize.py b/dask/tests/test_tokenize.py
index 49461e541..fc422d992 100644
--- a/dask/tests/test_tokenize.py
+++ b/dask/tests/test_tokenize.py
@@ -1271,6 +1271,7 @@ def test_tokenize_random_functions(module):
     """random.random() and other methods of the global random state do not compare as
     equal to themselves after a pickle roundtrip"""
     module = eval(module)
+    module.seed(2)
 
     a = module.random
     b = pickle.loads(pickle.dumps(a))