Repository: dask/dask
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

Problem statement:
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

Hints from issue discussion:

Commit message: Impleent automatic rechunking for shuffle

Patch to review:
diff --git a/dask/array/_shuffle.py b/dask/array/_shuffle.py
index 6e3dd2e4ab2..6e0dd7e7b39 100644
--- a/dask/array/_shuffle.py
+++ b/dask/array/_shuffle.py
@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import copy
+import math
+from functools import reduce
 from itertools import count, product
+from operator import mul
 
 import numpy as np
 
@@ -13,7 +16,7 @@
 from dask.highlevelgraph import HighLevelGraph
 
 
-def shuffle(x, indexer: list[list[int]], axis):
+def shuffle(x, indexer: list[list[int]], axis, chunks="auto"):
     """
     Reorders one dimensions of a Dask Array based on an indexer.
 
@@ -34,6 +37,8 @@ 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.
 
     Examples
     --------
@@ -66,6 +71,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 = _resize_other_dimensions(x, max(map(len, indexer)), axis, chunks)
 
     token = tokenize(x, indexer, axis)
     out_name = f"shuffle-{token}"
@@ -79,7 +87,44 @@ 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 _resize_other_dimensions(x: Array, longest_group, axis, chunks):
+    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
+
+    max_group = reduce(mul, map(max, x.chunks))
+    new_max_group = max_group / max(x.chunks[axis]) * longest_group
+
+    new_chunks = []
+    for i in range(len(x.chunks)):
+        if i == axis:
+            new_chunks.append(x.chunks[i])
+            continue
+
+        new_chunk = []
+        # calculate what the max chunk size in this dimension is and split every
+        # chunk that is larger than that
+        dimension_max_chunk = max(x.chunks[i]) * (
+            (max_group / new_max_group) ** (1 / (len(x.chunks) - 1))
+        )
+        for c in x.chunks[i]:
+            if c > chunksize_tolerance * dimension_max_chunk:
+                factor = c / dimension_max_chunk
+                nc = [c // math.ceil(factor)] * math.ceil(factor)
+                for ii in range(c - sum(nc)):
+                    nc[ii] += 1
+                new_chunk.extend(nc)
+            else:
+                new_chunk.append(c)
+        new_chunks.append(new_chunk)
+
+    return x.rechunk(tuple(map(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 +138,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/tests/test_shuffle.py b/dask/array/tests/test_shuffle.py
index 2882cbc8752..b0b93d3ce19 100644
--- a/dask/array/tests/test_shuffle.py
+++ b/dask/array/tests/test_shuffle.py
@@ -6,6 +6,7 @@
 import dask
 import dask.array as da
 from dask.array import assert_eq, shuffle
+from dask.array._shuffle import _resize_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,19 @@ 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 = _resize_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 = _resize_other_dimensions(arr, 20, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 67), (10,) * 5, (5,) * 4)
+    assert_eq(arr, result)
+
+    result = _resize_other_dimensions(arr, 40, 1, "auto")
+    assert result.chunks == ((45, 50, 50, 38, 34, 33), (10,) * 5, (5,) * 4)
+    assert_eq(arr, result)


Code review comments:
diff_hunk: @@ -79,7 +87,44 @@ 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 _resize_other_dimensions(x: Array, longest_group, axis, chunks):
+    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
+
+    max_group = reduce(mul, map(max, x.chunks))
+    new_max_group = max_group / max(x.chunks[axis]) * longest_group
path: dask/array/_shuffle.py
text: @user1:
```suggestion
    max_axis_chunks = max(x.chunks[axis])
```

@author:
I need a bit more, but we can definitely simplify these formulas a lot

diff_hunk: @@ -79,7 +87,44 @@ 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 _resize_other_dimensions(x: Array, longest_group, axis, chunks):
+    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
+
+    max_group = reduce(mul, map(max, x.chunks))
+    new_max_group = max_group / max(x.chunks[axis]) * longest_group
+
+    new_chunks = []
+    for i in range(len(x.chunks)):
+        if i == axis:
+            new_chunks.append(x.chunks[i])
+            continue
+
+        new_chunk = []
+        # calculate what the max chunk size in this dimension is and split every
+        # chunk that is larger than that
+        dimension_max_chunk = max(x.chunks[i]) * (
+            (max_group / new_max_group) ** (1 / (len(x.chunks) - 1))
+        )
+        for c in x.chunks[i]:
+            if c > chunksize_tolerance * dimension_max_chunk:
+                factor = c / dimension_max_chunk
+                nc = [c // math.ceil(factor)] * math.ceil(factor)
path: dask/array/_shuffle.py
text: @user1:
Can you use `divmod` here?

@author:
Oh I always keep forgetting that, yep that's a good point

diff_hunk: @@ -79,7 +87,44 @@ 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 _resize_other_dimensions(x: Array, longest_group, axis, chunks):
path: dask/array/_shuffle.py
text: @user1:
```suggestion
def _resize_other_dimensions(x: Array, longest_group: int, axis: int, chunks: tuple[tuple[int, ...], ...]):
```

@author:
chunks is currently only string, but I'll add type hints.

diff_hunk: @@ -79,7 +87,44 @@ 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 _resize_other_dimensions(x: Array, longest_group, axis, chunks):
+    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
+
+    max_group = reduce(mul, map(max, x.chunks))
+    new_max_group = max_group / max(x.chunks[axis]) * longest_group
+
+    new_chunks = []
+    for i in range(len(x.chunks)):
+        if i == axis:
+            new_chunks.append(x.chunks[i])
+            continue
+
+        new_chunk = []
+        # calculate what the max chunk size in this dimension is and split every
+        # chunk that is larger than that
+        dimension_max_chunk = max(x.chunks[i]) * (
+            (max_group / new_max_group) ** (1 / (len(x.chunks) - 1))
path: dask/array/_shuffle.py
text: @user1:
```suggestion
            max_axis_chunks ** (1 / (len(x.chunks) - 1))
```

why the exponential?

@author:
Assuming 3 dimensions and we shuffle on axis=0, I want to allocate the split factor equally to the other dimensions, i.e. if the factor is 4, I want to split axis=1 and axis=2 by half each.

@user1:
ah I see . a comment would help. I think the rewrite to `max_axis_chunks` is also useful, otherwise you multiply by `new_max_group` and then divide by `new_max_group`

@author:
Yeah I rewrote this, that was actually way to complicate.

I'll add a comment