Repository: dask/dask
Issue category: Feature Request

Improve how normalize_chunks selects chunk sizes if auto is given
Specifying auto for a dimension when new chunks are selected like in rechunking allows us to choose a chunk size that we think is a good fit and efficient. 

This is for now an issue to collect thoughts from others and feedback, I am pretty sure that I haven't found all suboptimal corner cases yet.


Assume we have the following array:

```
import dask.array as da

arr = da.random.random((5000, 300, 372), chunks=(600, 150, 172))
```

Triggering a Rechunk operation to -1 and auto for dimensions 2 and 3 results in the following chunks:

```
x = arr.rechunk((-1, "auto", "auto"))
```

<img width="747" alt="Screenshot 2024-08-20 at 15 34 19" src="https://github.com/user-attachments/assets/31e5c070-fc01-40be-af63-21fa06e457f5">

```normalize_chunks`` targets the 128MiB that is the default for auto pretty aggressively and creates a significantly more complex task graph than necessary if we would compromise a little bit on the in memory size of the chunks.

Choosing 50 and (58, 57, 57) per input chunk would allows us to split every chunk into 3 pieces along every axis and thus avoid including the neighbouring chunks alongside this axis.

I am pretty sure that there are better examples where this makes things very complex, but this should illustrate the general issue.

The proposal is to compromise a little bit on the chunk size side of things to create a simpler topology if possible.




Functions to edit:
dask/array/core.py:normalize_chunks
dask/array/core.py:_compute_multiplier
dask/array/core.py:auto_chunks
Functions added: dask/array/core.py:_convert_int_chunk_to_tuple
dask/array/core.py:_split_up_single_chunk

Patch:
diff --git a/dask/array/core.py b/dask/array/core.py
index 76ab569f4ee..d7e80038c97 100644
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -22,7 +22,7 @@
 
 import numpy as np
 from numpy.typing import ArrayLike
-from tlz import accumulate, concat, first, frequencies, groupby, partition
+from tlz import accumulate, concat, first, groupby, partition
 from tlz.curried import pluck
 
 from dask import compute, config, core
@@ -3138,17 +3138,7 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
     if chunks and shape is not None:
         # allints: did we start with chunks as a simple tuple of ints?
         allints = all(isinstance(c, int) for c in chunks)
-        chunks = sum(
-            (
-                (
-                    blockdims_from_blockshape((s,), (c,))
-                    if not isinstance(c, (tuple, list))
-                    else (c,)
-                )
-                for s, c in zip(shape, chunks)
-            ),
-            (),
-        )
+        chunks = _convert_int_chunk_to_tuple(shape, chunks)
     for c in chunks:
         if not c:
             raise ValueError(
@@ -3178,6 +3168,20 @@ def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks
     )
 
 
+def _convert_int_chunk_to_tuple(shape, chunks):
+    return sum(
+        (
+            (
+                blockdims_from_blockshape((s,), (c,))
+                if not isinstance(c, (tuple, list))
+                else (c,)
+            )
+            for s, c in zip(shape, chunks)
+        ),
+        (),
+    )
+
+
 def _compute_multiplier(limit: int, dtype, largest_block: int, result):
     """
     Utility function for auto_chunk, to fin how much larger or smaller the ideal
@@ -3187,7 +3191,7 @@ def _compute_multiplier(limit: int, dtype, largest_block: int, result):
         limit
         / dtype.itemsize
         / largest_block
-        / math.prod(r for r in result.values() if r)
+        / math.prod(max(r) if isinstance(r, tuple) else r for r in result.values() if r)
     )
 
 
@@ -3216,9 +3220,7 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
     normalize_chunks: for full docstring and parameters
     """
     if previous_chunks is not None:
-        previous_chunks = tuple(
-            c if isinstance(c, tuple) else (c,) for c in previous_chunks
-        )
+        previous_chunks = _convert_int_chunk_to_tuple(shape, previous_chunks)
     chunks = list(chunks)
 
     autos = {i for i, c in enumerate(chunks) if c == "auto"}
@@ -3252,6 +3254,7 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
             )
 
     limit = max(1, limit)
+    chunksize_tolerance = config.get("array.chunk-size-tolerance")
 
     largest_block = math.prod(
         cs if isinstance(cs, Number) else max(cs) for cs in chunks if cs != "auto"
@@ -3259,47 +3262,96 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
 
     if previous_chunks:
         # Base ideal ratio on the median chunk size of the previous chunks
-        result = {a: np.median(previous_chunks[a]) for a in autos}
-
-        ideal_shape = []
-        for i, s in enumerate(shape):
-            chunk_frequencies = frequencies(previous_chunks[i])
-            mode, count = max(chunk_frequencies.items(), key=lambda kv: kv[1])
-            if mode > 1 and count >= len(previous_chunks[i]) / 2:
-                ideal_shape.append(mode)
-            else:
-                ideal_shape.append(s)
+        median_chunks = {a: np.median(previous_chunks[a]) for a in autos}
+        result = {}
 
         # How much larger or smaller the ideal chunk size is relative to what we have now
-        multiplier = _compute_multiplier(limit, dtype, largest_block, result)
-
-        last_multiplier = 0
-        last_autos = set()
-        while (
-            multiplier != last_multiplier or autos != last_autos
-        ):  # while things change
-            last_multiplier = multiplier  # record previous values
+        multiplier = _compute_multiplier(limit, dtype, largest_block, median_chunks)
+
+        def _trivial_aggregate(a):
+            autos.remove(a)
+            del median_chunks[a]
+            return True
+
+        multiplier_remaining = True
+        while multiplier_remaining:  # while things change
             last_autos = set(autos)  # record previous values
+            multiplier_remaining = False
 
             # Expand or contract each of the dimensions appropriately
             for a in sorted(autos):
-                if ideal_shape[a] == 0:
-                    result[a] = 0
-                    continue
-                proposed = result[a] * multiplier ** (1 / len(autos))
+                this_multiplier = multiplier ** (1 / len(last_autos))
+
+                proposed = median_chunks[a] * this_multiplier
+                this_chunksize_tolerance = chunksize_tolerance ** (1 / len(last_autos))
+                max_chunk_size = proposed * this_chunksize_tolerance
+
                 if proposed > shape[a]:  # we've hit the shape boundary
-                    autos.remove(a)
-                    largest_block *= shape[a]
                     chunks[a] = shape[a]
-                    del result[a]
+                    multiplier_remaining = _trivial_aggregate(a)
+                    largest_block *= shape[a]
+                    continue
+                elif this_multiplier <= 1:
+                    if max(previous_chunks[a]) == 1:
+                        multiplier_remaining = _trivial_aggregate(a)
+                        chunks[a] = tuple(previous_chunks[a])
+                        continue
+
+                    cache, dimension_result = {}, []
+                    for c in previous_chunks[a]:
+                        if c in cache:
+                            # We potentially split the same chunk over and over again
+                            # so cache that
+                            dimension_result.extend(cache[c])
+                            continue
+                        elif c <= max(max_chunk_size, 1):
+                            new_chunks = [c]
+                        else:
+                            new_chunks = _split_up_single_chunk(
+                                c, this_chunksize_tolerance, max_chunk_size, proposed
+                            )
+
+                        cache[c] = new_chunks
+                        dimension_result.extend(new_chunks)
+
+                    if max(dimension_result) == 1:
+                        # See if we can split up other axes further
+                        multiplier_remaining = _trivial_aggregate(a)
+
                 else:
-                    result[a] = round_to(proposed, ideal_shape[a])
+                    dimension_result, new_chunk = [], 0
+                    for c in previous_chunks[a]:
+                        if c + new_chunk <= proposed:
+                            # keep increasing the chunk
+                            new_chunk += c
+                        else:
+                            # We reach the boundary so start a new chunk
+                            dimension_result.append(new_chunk)
+                            new_chunk = c
+                            if new_chunk > max_chunk_size:
+                                # our new chunk is greater than the allowed, so split it up
+                                dimension_result.extend(
+                                    _split_up_single_chunk(
+                                        c,
+                                        this_chunksize_tolerance,
+                                        max_chunk_size,
+                                        proposed,
+                                    )
+                                )
+                                new_chunk = 0
+                    if new_chunk > 0:
+                        dimension_result.append(new_chunk)
+
+                result[a] = tuple(dimension_result)
 
             # recompute how much multiplier we have left, repeat
-            multiplier = _compute_multiplier(limit, dtype, largest_block, result)
+            if multiplier_remaining:
+                multiplier = _compute_multiplier(
+                    limit, dtype, largest_block, median_chunks
+                )
 
         for k, v in result.items():
-            chunks[k] = v
+            chunks[k] = v if v else 0
         return tuple(chunks)
 
     else:
@@ -3321,6 +3373,25 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
         return tuple(chunks)
 
 
+def _split_up_single_chunk(
+    c: int, this_chunksize_tolerance: float, max_chunk_size: int, proposed: int
+) -> list[int]:
+    # Calculate by what factor we have to split this chunk
+    m = c / proposed
+    if math.ceil(m) / m > this_chunksize_tolerance:
+        # We want to smooth things potentially if rounding up would change the result
+        # by a lot
+        m = math.ceil(c / max_chunk_size)
+    else:
+        m = math.ceil(m)
+    # split the chunk
+    new_c, remainder = divmod(c, min(m, c))
+    x = [new_c] * min(m, c)
+    for i in range(remainder):
+        x[i] += 1
+    return x
+
+
 def round_to(c, s):
     """Return a chunk dimension that is close to an even multiple or factor
 

Test patch: diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
index 02eea953f08..26a7e405e4f 100644
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -825,7 +825,18 @@ def test_rechunk_avoid_needless_chunking():
         (100, 50, 10, (10,) * 10),
         (100, 100, 10, (10,) * 10),
         (20, 7, 10, (7, 7, 6)),
-        (20, (1, 1, 1, 1, 6, 2, 1, 7), 5, (5, 5, 5, 5)),
+        (
+            20,
+            (1, 1, 1, 1, 6, 2, 1, 7),
+            5,
+            (4, 6, 3, 4, 3),
+        ),  # 6 is in tolerance, 7 is not
+        (21, (1,) * 21, 10, (10, 10, 1)),  # ensure that we squash together properly
+        (20, ((2, 2, 2, 9, 1, 2, 2)), 5, (4, 2, 5, 4, 5)),
+        (20, ((10, 10)), 4, (4, 3, 3, 4, 3, 3)),
+        (21, ((10, 11)), 4, (4, 3, 3, 4, 4, 3)),
+        (20, ((1, 18, 1)), 5, (1, 5, 5, 4, 4, 1)),
+        (38, ((10, 18, 10)), 5, (5, 5, 5, 5, 4, 4, 5, 5)),
     ],
 )
 def test_rechunk_auto_1d(shape, chunks, bs, expected):
@@ -834,6 +845,26 @@ def test_rechunk_auto_1d(shape, chunks, bs, expected):
     assert y.chunks == (expected,)
 
 
+@pytest.mark.parametrize(
+    "previous_chunks,bs,expected",
+    [
+        (((1, 1, 1), (10, 10, 10, 10, 10, 10, 10, 10)), 160, ((3,), (50, 30))),
+        (((2, 2), (20,)), 5, ((1, 1, 1, 1), (5, 5, 5, 5))),
+        (((1, 1), (20,)), 5, ((1, 1), (5, 5, 5, 5))),
+    ],
+)
+def test_normalize_chunks_auto_2d(previous_chunks, bs, expected):
+    shape = tuple(map(sum, previous_chunks))
+    result = normalize_chunks(
+        {0: "auto", 1: "auto"},
+        shape,
+        limit=bs,
+        dtype=np.int8,
+        previous_chunks=previous_chunks,
+    )
+    assert result == expected
+
+
 def test_rechunk_auto_2d():
     x = da.ones((20, 20), chunks=(2, 2))
     y = x.rechunk({0: -1, 1: "auto"}, block_size_limit=20 * x.dtype.itemsize)
@@ -880,7 +911,12 @@ def test_rechunk_auto_image_stack(n):
     with dask.config.set({"array.chunk-size": "1MiB"}):
         x = da.ones((n, 1000, 1000), chunks=(1, 1000, 1000), dtype="float64")
         z = x.rechunk("auto")
-        assert z.chunks == ((1,) * n, (362, 362, 276), (362, 362, 276))
+        assert z.chunks == ((1,) * n, (334, 333, 333), (334, 333, 333))
+
+    with dask.config.set({"array.chunk-size": "1MiB"}):
+        x = da.ones((n, 2000, 2000), chunks=(1, 1000, 1000), dtype="float64")
+        z = x.rechunk("auto")
+        assert z.chunks == ((1,) * n, (334, 333, 333) * 2, (334, 333, 333) * 2)
 
 
 def test_rechunk_down():
@@ -891,14 +927,14 @@ def test_rechunk_down():
 
     with dask.config.set({"array.chunk-size": "1MiB"}):
         z = y.rechunk("auto")
-        assert z.chunks == ((4,) * 25, (511, 489), (511, 489))
+        assert z.chunks == ((5,) * 20, (500, 500), (500, 500))
 
     with dask.config.set({"array.chunk-size": "1MiB"}):
         z = y.rechunk({0: "auto"})
         assert z.chunks == ((1,) * 100, (1000,), (1000,))
 
         z = y.rechunk({1: "auto"})
-        assert z.chunks == ((10,) * 10, (104,) * 9 + (64,), (1000,))
+        assert z.chunks == ((10,) * 10, (100,) * 10, (1000,))
 
 
 def test_rechunk_zero():