Repository: dask/dask
Pull request #11262 (2024-07-31 09:59:12)

Implement task-based array shuffle
- [ ] xref #11257
- [ ] Tests added / passed
- [ ] Passes `pre-commit run --all-files`

This is a first rough implementation of an array-based shuffle that uses the task implementation. Not polished and probably has a few edge cases that aren't covered yet.

TLDR: We are preserving chunk sizes (roughly) compared to take but we still have a very large graph, which strains the scheduler pretty heavily. Subsequent results are easier to process though in case of a random access pattern because the chunks stay reasonable. An improvement over random access take generally speaking, but still not very good


Let's mess around with a toy example:

```
import dask.array as da

n = 1500
arr = da.random.random((n, n, n), chunks=
(50, 50, 50))
arr
```

<img width="665" alt="Screenshot 2024-07-31 at 11 57 00" src="https://github.com/user-attachments/assets/826d3645-374e-49e9-bb2b-1863e01da591">

Our indexer has a random access pattern (and we access 5 elements at a time, but we could do random stuff here).

```
import numpy as np
from dask.core import flatten

idx = np.arange(0, n)
np.random.shuffle(idx)

indexer = []
for i in range(0, n, 5):
    indexer.append(idx[i:i+5].tolist())
take_indexer = list(flatten(indexer))

indexer
```

```
result = arr.shuffle(indexer, axis=1)
result
```

<img width="734" alt="Screenshot 2024-07-31 at 11 57 21" src="https://github.com/user-attachments/assets/f228e8a4-5b7b-4ac9-9f53-fd767bd9a05e">

Shuffling keeps the number of chunks consistent. Doing the equivalent take blows up the number of chunks

```
expected = arr[slice(None), take_indexer, slice(None)]
expected
```

<img width="662" alt="Screenshot 2024-07-31 at 11 57 35" src="https://github.com/user-attachments/assets/9531dcaf-34f4-4cdb-b9cc-b0faf03adfd8">

The number of tasks are in the same ballpark though, the growth is still exponential in a case of an all to all pattern.

 
<img width="277" alt="Screenshot 2024-07-31 at 11 57 54" src="https://github.com/user-attachments/assets/6c7d399c-ce65-4480-a1b7-db192ee07a0b">


Generating close to a million tasks for a 27k chunks array is not great.

fjetter: I assume this is pretty much equivalent to `arr.take(...).rechunk(arr.chunks)` isn't it?
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._

     15 files  ±0       15 suites  ±0   3h 32m 10s :stopwatch: + 5m 5s
 13 129 tests ±0   12 061 :white_check_mark: ±0   1 068 :zzz: ±0  0 :x: ±0 
162 795 runs  ±0  140 006 :white_check_mark: +1  22 789 :zzz:  - 1  0 :x: ±0 

Results for commit 6c848e95. ± Comparison against base commit f95007b1.

[test-results]:data:application/gzip;base64,H4sIAFoSqmYC/02Myw6DIBBFf8Ww7mIYHg79GWMpJqQ+GoRV038vUkWX59yb82GDH93K7g1Xt4atyceTnin00S/zxtgKyCqvseyCozm4W5O15QSan/Ll35sETdUNvR+zq6XOhbCE3YQ0l7TG1qida1oCgK7yn0ZsyRzuki58LdtlmnzMwLQlSc4o7ZzqjSSUAszAFSkriCwCKK7wIdj3B8lDHJAaAQAA

phofl: > I assume this is pretty much equivalent to arr.take(...).rechunk(arr.chunks) isn't it?

roughly, we can be smarter with respect to how many tasks we create, i.e. if you have a 6 x 6 array with 3 x 3 chunks and an indexer on axis 1 of

``[[0, 3, 1], [4, 2, 5]]``

take would create a task per element, because they aren't in order. We can sort and then re-order to only slice once per chunk. Shuffle also ensures that a group of elements ends up in the same chunk.

That said, worst case task graph size is roughly the same yes.

mrocklin: A couple questions / thoughts (but I haven't thought about this as much, so feel free to disregard)

### Multi-stage

Let's imagine a pathological case where we have two chunks and and indexer that switches back and forth between them

```
chunks = ((100000, 100000),)
indexer = [0, 100000, 1, 100001, 2, 100002]
```

Optimally we should be able to do this in around 6 tasks I think, assuming we're willing to do the take in a couple of stages, in worst case we'll do it in 200,000 tasks.  What happens with this implementation?

Also, it may not make sense to invest in the multi-stage implementation here for various reasons, I'm just making sure that we have that option on our radar.

### Lots of Rechunking 

If we did the multi-stage thing, then suddenly we're strongly incentivized to rechunk ahead of time to make the take dimension very large.  This could be useful for two reasons:

1.  The input chunk sizes are too small (I think that we should be willing to override small input chunk sizes for intermediate results)
2. The input chunking direction isn't optimal for the take operation


Just some thoughts.  It could be that we've already been thinking about these or that they're bad for some reason.  No need to anchor on them, just bringing them up in case they're useful.
mrocklin: > We do this in 6 tasks at the moment. We sort the indexer and then resort after selecting the data

Oh grand!

So I guess we're experiencing here a task blowup of shape/chunkshape for the take dimension, in this case that's 1500/50 == 30x.  So 27k * 30 == ~800k, which tracks.  This is a lot better than the number of chunk transitions in the indexer, which is what we had before.  This seems great!
dcherian: Looks like the approach in https://github.com/dask/dask/pull/3901?
phofl: Yep very similar approach. Task graph will be smaller but both leverage the same technique.

This is not the state that I would like to end up in though.


A more general question: Is this the kind of API could be used in Xarray for a group by-transform or group by-reduce if you don't want to do a tree reduction? Basically what we discussed here: https://github.com/pydata/xarray/issues/9220

Patch:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 58cac45bb..12849f6f1 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -32,6 +32,7 @@ def _array_expr_enabled() -> bool:
 
 try:
     from dask.array import backends, fft, lib, linalg, ma, overlap, random
+    from dask.array._shuffle import shuffle
     from dask.array.blockwise import atop, blockwise
     from dask.array.chunk_types import register_chunk_type
     from dask.array.core import (
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
new file mode 100644
index 000000000..37aacb4bb
--- /dev/null
+++ b/dask/array/_shuffle.py
@@ -0,0 +1,196 @@
+from __future__ import annotations
+
+import copy
+from itertools import count, product
+
+import numpy as np
+
+from dask import config
+from dask.array.chunk import getitem
+from dask.array.core import Array, unknown_chunk_message
+from dask.base import tokenize
+from dask.highlevelgraph import HighLevelGraph
+
+
+def shuffle(x, indexer: list[list[int]], axis):
+    """
+    Reorders one dimensions of a Dask Array based on an indexer.
+
+    The indexer defines a list of positional groups that will end up in the same chunk
+    together. A single group is in at most one chunk on this dimension, but a chunk
+    might contain multiple groups to avoid fragmentation of the array.
+
+    The algorithm tries to balance the chunksizes as much as possible to ideally keep the
+    number of chunks consistent or at least manageable.
+
+    Parameters
+    ----------
+    x: dask array
+        Array to be shuffled.
+    indexer:  list[list[int]]
+        The indexer that determines which elements along the dimension will end up in the
+        same chunk. Multiple groups can be in the same chunk to avoid fragmentation, but
+        each group will end up in exactly one chunk.
+    axis: int
+        The axis to shuffle along.
+
+    Examples
+    --------
+    >>> import dask.array as da
+    >>> import numpy as np
+    >>> arr = np.array([[1, 2, 3, 4, 5, 6, 7, 8], [9, 10, 11, 12, 13, 14, 15, 16]])
+    >>> x = da.from_array(arr, chunks=(2, 4))
+
+    Separate the elements in different groups.
+
+    >>> y = x.shuffle([[6, 5, 2], [4, 1], [3, 0, 7]], axis=1)
+
+    The shuffle algorihthm will combine the first 2 groups into a single chunk to keep
+    the number of chunks small.
+
+    The tolerance of increasing the chunk size is controlled by the configuration
+    "array.shuffle.chunksize-tolerance". The default value is 1.25.
+
+    >>> y.chunks
+    ((2,), (5, 3))
+
+    The array was reordered along axis 1 according to the positional indexer that was given.
+
+    >>> y.compute()
+    array([[ 7,  6,  3,  5,  2,  4,  1,  8],
+           [15, 14, 11, 13, 10, 12,  9, 16]])
+    """
+    if np.isnan(x.shape).any():
+        raise ValueError(
+            f"Shuffling only allowed with known chunk sizes. {unknown_chunk_message}"
+        )
+    assert isinstance(axis, int), "axis must be an integer"
+
+    token = tokenize(x, indexer, axis)
+    out_name = f"shuffle-{token}"
+
+    chunks, layer = _shuffle(x.chunks, indexer, axis, x.name, out_name, token)
+    graph = HighLevelGraph.from_collections(out_name, layer, dependencies=[x])
+
+    return Array(graph, out_name, chunks, meta=x)
+
+
+def _shuffle(chunks, indexer, axis, in_name, out_name, token):
+    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")
+
+    if not axis <= len(chunks):
+        raise ValueError(
+            f"Axis {axis} is out of bounds for array with {len(chunks)} axes"
+        )
+
+    if max(map(max, indexer)) >= sum(chunks[axis]):
+        raise IndexError(
+            f"Indexer contains out of bounds index. Dimension only has {sum(chunks[axis])} elements."
+        )
+    indexer = copy.deepcopy(indexer)
+
+    chunksize_tolerance = config.get("array.shuffle.chunksize-tolerance")
+    chunk_size_limit = int(sum(chunks[axis]) / len(chunks[axis]) * chunksize_tolerance)
+
+    # Figure out how many groups we can put into one chunk
+    current_chunk, new_chunks = [], []
+    for idx in indexer:
+        if len(current_chunk) + len(idx) > chunk_size_limit and len(current_chunk) > 0:
+            new_chunks.append(current_chunk)
+            current_chunk = idx.copy()
+        else:
+            current_chunk.extend(idx)
+            if len(current_chunk) > chunk_size_limit / chunksize_tolerance:
+                new_chunks.append(current_chunk)
+                current_chunk = []
+    if len(current_chunk) > 0:
+        new_chunks.append(current_chunk)
+
+    chunk_boundaries = np.cumsum(chunks[axis])
+
+    # Get existing chunk tuple locations
+    chunk_tuples = list(
+        product(*(range(len(c)) for i, c in enumerate(chunks) if i != axis))
+    )
+
+    intermediates = dict()
+    merges = dict()
+    split_name = f"shuffle-split-{token}"
+    slices = [slice(None)] * len(chunks)
+    split_name_suffixes = count()
+
+    old_blocks = np.empty([len(c) for c in chunks], dtype="O")
+    for old_index in np.ndindex(old_blocks.shape):
+        old_blocks[old_index] = (in_name,) + old_index
+
+    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)
+        sorted_array = new_chunk_taker[sorter]
+        source_chunk_nr, taker_boundary = np.unique(
+            np.searchsorted(chunk_boundaries, sorted_array, side="right"),
+            return_index=True,
+        )
+        taker_boundary = taker_boundary.tolist()
+        taker_boundary.append(len(new_chunk_taker))
+
+        taker_cache = {}
+        for chunk_tuple in chunk_tuples:
+            merge_keys = []
+
+            for c, b_start, b_end in zip(
+                source_chunk_nr, taker_boundary[:-1], taker_boundary[1:]
+            ):
+                # insert our axis chunk id into the chunk_tuple
+                chunk_key = convert_key(chunk_tuple, c, axis)
+                name = (split_name, next(split_name_suffixes))
+                this_slice = slices.copy()
+
+                # Cache the takers to allow de-duplication when serializing
+                # Ugly!
+                if c in taker_cache:
+                    this_slice[axis] = taker_cache[c]
+                else:
+                    this_slice[axis] = sorted_array[b_start:b_end] - (
+                        chunk_boundaries[c - 1] if c > 0 else 0
+                    )
+                    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)
+                merge_keys.append(name)
+
+            merge_suffix = convert_key(chunk_tuple, new_chunk_idx, axis)
+            if len(merge_keys) > 1:
+                merges[(out_name,) + merge_suffix] = (
+                    concatenate_arrays,
+                    merge_keys,
+                    sorter,
+                    axis,
+                )
+            elif len(merge_keys) == 1:
+                merges[(out_name,) + merge_suffix] = intermediates.pop(merge_keys[0])
+            else:
+                raise NotImplementedError
+
+    output_chunks = []
+    for i, c in enumerate(chunks):
+        if i == axis:
+            output_chunks.append(tuple(map(len, new_chunks)))
+        else:
+            output_chunks.append(c)
+
+    layer = {**merges, **intermediates}
+    return tuple(output_chunks), layer
+
+
+def concatenate_arrays(arrs, sorter, axis):
+    return np.take(np.concatenate(arrs, axis=axis), np.argsort(sorter), axis=axis)
+
+
+def convert_key(key, chunk, axis):
+    key = list(key)
+    key.insert(axis, chunk)
+    return tuple(key)
diff --git a/dask/array/core.py b/dask/array/core.py
index 59f10fecc..b710fe491 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -2762,6 +2762,23 @@ class Array(DaskMethodsMixin):
 
         return rechunk(self, chunks, threshold, block_size_limit, balance, method)
 
+    def shuffle(
+        self,
+        indexer: list[list[int]],
+        axis,
+    ):
+        """Reorders one dimensions of a Dask Array based on an indexer.
+
+        Refer to :func:`dask.array.shuffle` for full documentation.
+
+        See Also
+        --------
+        dask.array.shuffle : equivalent function
+        """
+        from dask.array._shuffle import shuffle
+
+        return shuffle(self, indexer, axis)
+
     @property
     def real(self):
         from dask.array.ufunc import real
diff --git a/dask/dask-schema.yaml b/dask/dask-schema.yaml
index 12bf2ee06..e05aac038 100644
--- a/dask/dask-schema.yaml
+++ b/dask/dask-schema.yaml
@@ -135,6 +135,17 @@ properties:
               The graph growth factor above which task-based shuffling introduces
               an intermediate step.
 
+      shuffle:
+        type: object
+        properties:
+
+          chunksize-tolerance:
+            type: number
+            description: |
+              Upper tolerance for the shuffle algorithm when creating output chunks.
+              Default is 1.25. This means that the shuffle algorithm can exceed
+              the average input chunk size along this dimension by 25%.
+
       svg:
         type: object
         properties:
diff --git a/dask/dask.yaml b/dask/dask.yaml
index 06509fe87..6e35c10e4 100644
--- a/dask/dask.yaml
+++ b/dask/dask.yaml
@@ -24,6 +24,8 @@ array:
   rechunk:
     method: "tasks"  # Rechunking method to use
     threshold: 4
+  shuffle:
+    chunksize-tolerance: 1.25  # Tolerance for the shuffle algorithm when creating output chunks.
   svg:
     size: 120  # pixels
   slicing:
diff --git a/dask/order.py b/dask/order.py
index 6e427efa0..d91e139a4 100644
--- a/dask/order.py
+++ b/dask/order.py
@@ -46,7 +46,6 @@ graph types laid out very carefully to show the kinds of situations that often
 arise, and the order we would like to be determined.
 
 """
-import copy
 from collections import defaultdict, deque, namedtuple
 from collections.abc import Iterable, Mapping, MutableMapping
 from typing import Any, Callable, Literal, NamedTuple, overload
@@ -110,14 +109,7 @@ def order(
     expected_len = len(dsk)
 
     if dependencies is None:
-        dependencies_are_copy = True
         dependencies = {k: get_dependencies(dsk, k) for k in dsk}
-    else:
-        # Below we're removing items from the sets in this dict
-        # We need a deepcopy for that but we only want to do this if necessary
-        # since this only happens for special cases.
-        dependencies_are_copy = False
-        dependencies = dict(dependencies)
 
     dependents = reverse_dict(dependencies)
 
@@ -137,7 +129,6 @@ def order(
     # way that is simpler to handle
     all_tasks = False
     n_removed_leaves = 0
-    requires_data_task = defaultdict(set)
     while not all_tasks:
         all_tasks = True
         for leaf in list(leaf_nodes):
@@ -165,23 +156,6 @@ def order(
                     if not dependents[dep]:
                         leaf_nodes.add(dep)
 
-        for root in list(root_nodes):
-            if root in leaf_nodes:
-                continue
-            if not istask(dsk[root]) and len(dependents[root]) > 1:
-                if not dependencies_are_copy:
-                    dependencies_are_copy = True
-                    dependencies = copy.deepcopy(dependencies)
-                root_nodes.remove(root)
-                for dep in dependents[root]:
-                    requires_data_task[dep].add(root)
-                    dependencies[dep].remove(root)
-                    if not dependencies[dep]:
-                        root_nodes.add(dep)
-                del dsk[root]
-                del dependencies[root]
-                del dependents[root]
-
     num_needed, total_dependencies = ndependencies(dependencies, dependents)
     if len(total_dependencies) != len(dsk):
         cycle = getcycle(dsk, None)
@@ -239,9 +213,6 @@ def order(
             if item in result:
                 continue
 
-            while requires_data_task[item]:
-                add_to_result(requires_data_task[item].pop())
-
             if return_stats:
                 result[item] = Order(i, crit_path_counter - _crit_path_counter_offset)
             else:
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index b79642613..4765378e5 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -207,6 +207,7 @@ Top level functions
    searchsorted
    select
    shape
+   shuffle
    sign
    signbit
    sin
@@ -300,6 +301,7 @@ Array
    Array.reshape
    Array.round
    Array.shape
+   Array.shuffle
    Array.size
    Array.squeeze
    Array.std
diff --git a/docs/source/support.rst b/docs/source/support.rst
index 805e333e5..bfbfa23f0 100644
--- a/docs/source/support.rst
+++ b/docs/source/support.rst
@@ -17,7 +17,7 @@ Community Meeting
     We've combined the monthly Dask Demo Day and Dask Developer Meeting into a single, monthly Dask community meeting.
 
 Join us for the monthly community meetings on the first Thursday of the month at
-10:00 US Central Time. `Join via Zoom <https://us06web.zoom.us/j/87619866741?pwd=S2RxMlRKcnVvakt4NHZoS1cwOGZoZz09>`__.
+10:00 US Central Time. `Join via Zoom <https://us06web.zoom.us/j/89383035703?pwd=WkRJSzNnRTh4T2R1ZjJuVVdJWlMxQT09>`__.
 
 Have something you'd like to share? Let us know by dropping a comment on `this GitHub issue <https://github.com/dask/community/issues/307>`__.
 


Test patch:
diff --git a/dask/array/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
new file mode 100644
index 000000000..469f22be0
--- /dev/null
+++ b/dask/array/tests/test_shuffle.py
@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import dask
+import dask.array as da
+from dask.array import assert_eq, shuffle
+from dask.core import flatten
+
+
+@pytest.fixture()
+def arr():
+    return np.arange(0, 24).reshape(8, 3).T.copy()
+
+
+@pytest.fixture()
+def darr(arr):
+    return da.from_array(arr, chunks=((2, 1), (4, 4)))
+
+
+@pytest.mark.parametrize(
+    "indexer, 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)),
+    ],
+)
+def test_shuffle(arr, darr, indexer, 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[1] == chunks
+
+
+@pytest.mark.parametrize("tol, chunks", ((1, (3, 2, 3)), (1.4, (5, 3))))
+def test_shuffle_config_tolerance(arr, darr, tol, chunks):
+    indexer = [[1, 5, 6], [0, 3], [4, 2, 7]]
+    with dask.config.set({"array.shuffle.chunksize-tolerance": tol}):
+        result = darr.shuffle(indexer, axis=1)
+    expected = arr[:, [1, 5, 6, 0, 3, 4, 2, 7]]
+    assert_eq(result, expected)
+    assert result.chunks[0] == darr.chunks[0]
+    assert result.chunks[1] == chunks
+
+
+def test_shuffle_larger_array():
+    arr = da.random.random((15, 15, 15), chunks=(5, 5, 5))
+    indexer = np.arange(0, 15)
+    np.random.shuffle(indexer)
+    indexer = [indexer[0:6], indexer[6:8], indexer[8:9], indexer[9:]]
+    indexer = list(map(list, indexer))
+    take_indexer = list(flatten(indexer))
+    assert_eq(shuffle(arr, indexer, axis=1), arr[..., take_indexer, :])
+
+
+def test_incompatible_indexer(darr):
+    with pytest.raises(ValueError, match="indexer must be a list of lists"):
+        darr.shuffle("s", axis=1)
+
+    with pytest.raises(ValueError, match="indexer must be a list of lists"):
+        darr.shuffle([1], axis=1)
+
+
+def test_unknown_chunk_sizes(darr):
+    darr._chunks = ((np.nan, 1), (4, 4))
+    with pytest.raises(
+        ValueError, match="Shuffling only allowed with known chunk sizes"
+    ):
+        darr.shuffle([[1]], axis=1)
+
+
+def test_oob_axis(darr):
+    with pytest.raises(ValueError, match="is out of bounds"):
+        darr.shuffle([[1]], axis=5)
+
+
+def test_oob_indexer(darr):
+    with pytest.raises(IndexError, match="Indexer contains out of bounds index"):
+        darr.shuffle([[16]], axis=1)