Repository: dask/dask
Pull request: Fix rechunking behavior for dimensions with known and unknown chunk sizes
Fixes rechunking for dimensions with known and unknown chunk sizes if the dimension is not touched; raises a useful error otherwise.

- [x] Closes #10028
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`


Problem statement:
Rechunking array with known and unknown chunk sizes fails with `KeyError`
```python
import dask.array as da
import dask.dataframe as dd
x = da.ones(shape=(4, 4), chunks=(2, 2))
y = dd.from_array(da.ones(shape=(4, 4), chunks=(2, 2))).values
z = da.concatenate([x, y])
result = z.rechunk(((2, 2, np.nan, np.nan), (4)))
result.compute()
```
fails with `KeyError: ('rechunk-merge-13faecb4faf29309353aaa9aae0f006b', 2, 0)`. Instead, I would expect rechunking with mixed known and unknown chunk sizes to
* work if we do not change the dimension with unknown chunks
* throw a descriptive error if an incompatible change is made to the dimension with unknowns
  * e.g., `((np.nan, 4), ) -> ((np.nan, 6), )`
* either work of throw a descriptive error if a compatible change is made to the dimension with unknowns
  * e.g, `((np.nan, 4), ) -> ((np.nan, 2, 2), )`

<details>
  <summary>Full traceback</summary>

```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[110], line 7
      5 z = da.concatenate([x, y])
      6 result = z.rechunk(((2, 2, np.nan, np.nan), (4)))
----> 7 result.compute()

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/base.py:314, in DaskMethodsMixin.compute(self, **kwargs)
    290 def compute(self, **kwargs):
    291     """Compute this dask collection
    292 
    293     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    312     dask.base.compute
    313     """
--> 314     (result,) = compute(self, traverse=False, **kwargs)
    315     return result

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/base.py:599, in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    596     keys.append(x.__dask_keys__())
    597     postcomputes.append(x.__dask_postcompute__())
--> 599 results = schedule(dsk, keys, **kwargs)
    600 return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/threaded.py:89, in get(dsk, keys, cache, num_workers, pool, **kwargs)
     86     elif isinstance(pool, multiprocessing.pool.Pool):
     87         pool = MultiprocessingPoolExecutor(pool)
---> 89 results = get_async(
     90     pool.submit,
     91     pool._max_workers,
     92     dsk,
     93     keys,
     94     cache=cache,
     95     get_id=_thread_get_id,
     96     pack_exception=pack_exception,
     97     **kwargs,
     98 )
    100 # Cleanup pools associated to dead threads
    101 with pools_lock:

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:525, in get_async(submit, num_workers, dsk, result, cache, get_id, rerun_exceptions_locally, pack_exception, raise_exception, callbacks, dumps, loads, chunksize, **kwargs)
    522             if finish:
    523                 finish(dsk, state, not succeeded)
--> 525 return nested_get(result, state["cache"])

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:302, in nested_get(ind, coll)
    289 """Get nested index from collection
    290 
    291 Examples
   (...)
    299 (('b', 'a'), ('a', 'b'))
    300 """
    301 if isinstance(ind, list):
--> 302     return tuple(nested_get(i, coll) for i in ind)
    303 else:
    304     return coll[ind]

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:302, in <genexpr>(.0)
    289 """Get nested index from collection
    290 
    291 Examples
   (...)
    299 (('b', 'a'), ('a', 'b'))
    300 """
    301 if isinstance(ind, list):
--> 302     return tuple(nested_get(i, coll) for i in ind)
    303 else:
    304     return coll[ind]

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:302, in nested_get(ind, coll)
    289 """Get nested index from collection
    290 
    291 Examples
   (...)
    299 (('b', 'a'), ('a', 'b'))
    300 """
    301 if isinstance(ind, list):
--> 302     return tuple(nested_get(i, coll) for i in ind)
    303 else:
    304     return coll[ind]

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:302, in <genexpr>(.0)
    289 """Get nested index from collection
    290 
    291 Examples
   (...)
    299 (('b', 'a'), ('a', 'b'))
    300 """
    301 if isinstance(ind, list):
--> 302     return tuple(nested_get(i, coll) for i in ind)
    303 else:
    304     return coll[ind]

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:302, in nested_get(ind, coll)
    289 """Get nested index from collection
    290 
    291 Examples
   (...)
    299 (('b', 'a'), ('a', 'b'))
    300 """
    301 if isinstance(ind, list):
--> 302     return tuple(nested_get(i, coll) for i in ind)
    303 else:
    304     return coll[ind]

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:302, in <genexpr>(.0)
    289 """Get nested index from collection
    290 
    291 Examples
   (...)
    299 (('b', 'a'), ('a', 'b'))
    300 """
    301 if isinstance(ind, list):
--> 302     return tuple(nested_get(i, coll) for i in ind)
    303 else:
    304     return coll[ind]

File /opt/homebrew/Caskroom/mambaforge/base/envs/rechunk/lib/python3.9/site-packages/dask/local.py:304, in nested_get(ind, coll)
    302     return tuple(nested_get(i, coll) for i in ind)
    303 else:
--> 304     return coll[ind]

KeyError: ('rechunk-merge-13faecb4faf29309353aaa9aae0f006b', 2, 0)
```
</details>

cc @jrbourbeau

Hints from issue discussion:

Commit message: Allow rechunking if not touching dimension with missing values

Patch to review:
diff --git a/dask/array/rechunk.py b/dask/array/rechunk.py
index 7f06df06c7f..399946826a3 100644
--- a/dask/array/rechunk.py
+++ b/dask/array/rechunk.py
@@ -10,7 +10,7 @@
 import heapq
 import math
 from functools import reduce
-from itertools import chain, count, product
+from itertools import chain, compress, count, product
 from operator import add, itemgetter, mul
 from warnings import warn
 
@@ -157,7 +157,7 @@ def _intersect_1d(breaks):
 def _old_to_new(old_chunks, new_chunks):
     """Helper to build old_chunks to new_chunks.
 
-    Handles missing values, as long as the missing dimension
+    Handles missing values, as long as the dimension with the missing chunk values
     is unchanged.
 
     Examples
@@ -169,31 +169,44 @@ def _old_to_new(old_chunks, new_chunks):
       [(2, slice(5, 10, None))],
       [(3, slice(0, 10, None)), (4, slice(0, 10, None))]]]
     """
-    old_known = [x for x in old_chunks if not any(math.isnan(y) for y in x)]
-    new_known = [x for x in new_chunks if not any(math.isnan(y) for y in x)]
 
-    n_missing = [sum(math.isnan(y) for y in x) for x in old_chunks]
-    n_missing2 = [sum(math.isnan(y) for y in x) for x in new_chunks]
+    def is_unknown(dim):
+        return any(math.isnan(chunk) for chunk in dim)
 
-    cmo = cumdims_label(old_known, "o")
-    cmn = cumdims_label(new_known, "n")
-
-    sums = [sum(o) for o in old_known]
-    sums2 = [sum(n) for n in new_known]
+    old_is_unknown = [is_unknown(dim) for dim in old_chunks]
+    new_is_unknown = [is_unknown(dim) for dim in new_chunks]
 
-    if not sums == sums2:
-        raise ValueError(f"Cannot change dimensions from {sums!r} to {sums2!r}")
-    if not n_missing == n_missing2:
+    if old_is_unknown != new_is_unknown or any(
+        new != old for new, old in compress(zip(old_chunks, new_chunks), old_is_unknown)
+    ):
         raise ValueError(
-            "Chunks must be unchanging along unknown dimensions.\n\n"
+            "Chunks must be unchanging along dimensions with missing values.\n\n"
             "A possible solution:\n  x.compute_chunk_sizes()"
         )
 
+    old_known = [dim for dim, unknown in zip(old_chunks, old_is_unknown) if not unknown]
+    new_known = [dim for dim, unknown in zip(new_chunks, new_is_unknown) if not unknown]
+
+    old_sizes = [sum(o) for o in old_known]
+    new_sizes = [sum(n) for n in new_known]
+
+    if not old_sizes == new_sizes:
+        raise ValueError(
+            f"Cannot change dimensions from {old_sizes!r} to {new_sizes!r}"
+        )
+
+    cmo = cumdims_label(old_known, "o")
+    cmn = cumdims_label(new_known, "n")
+
     old_to_new = [_intersect_1d(_breakpoints(cm[0], cm[1])) for cm in zip(cmo, cmn)]
-    for idx, missing in enumerate(n_missing):
-        if missing:
-            # Missing dimensions are always unchanged, so old -> new is everything
-            extra = [[(i, slice(0, None))] for i in range(missing)]
+    for idx, unknown in enumerate(old_is_unknown):
+        if unknown:
+            dim = old_chunks[idx]
+            # Unknown dimensions are always unchanged, so old -> new is everything
+            extra = [
+                [(i, slice(0, size if not math.isnan(size) else None))]
+                for i, size in enumerate(dim)
+            ]
             old_to_new.insert(idx, extra)
     return old_to_new
 
diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
index e236dccb1a4..8715b64ff20 100644
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -282,7 +282,7 @@ def test_rechunk_same():
     assert x is y
 
 
-def test_rechunk_same_unknown():
+def test_rechunk_same_fully_unknown():
     dd = pytest.importorskip("dask.dataframe")
     x = da.ones(shape=(10, 10), chunks=(5, 10))
     y = dd.from_array(x).values
@@ -292,8 +292,8 @@ def test_rechunk_same_unknown():
     assert y is result
 
 
-def test_rechunk_same_unknown_floats():
-    """Similar to test_rechunk_same_unknown but testing the behavior if
+def test_rechunk_same_fully_unknown_floats():
+    """Similar to test_rechunk_same_fully_unknown but testing the behavior if
     ``float("nan")`` is used instead of the recommended ``np.nan``
     """
     dd = pytest.importorskip("dask.dataframe")
@@ -304,6 +304,17 @@ def test_rechunk_same_unknown_floats():
     assert y is result
 
 
+def test_rechunk_same_partially_unknown():
+    dd = pytest.importorskip("dask.dataframe")
+    x = da.ones(shape=(10, 10), chunks=(5, 10))
+    y = dd.from_array(x).values
+    z = da.concatenate([x, y])
+    new_chunks = ((5, 5, np.nan, np.nan), (10,))
+    assert z.chunks == new_chunks
+    result = z.rechunk(new_chunks)
+    assert z is result
+
+
 def test_rechunk_with_zero_placeholders():
     x = da.ones((24, 24), chunks=((12, 12), (24, 0)))
     y = da.ones((24, 24), chunks=((12, 12), (12, 12)))
@@ -620,7 +631,7 @@ def test_rechunk_unknown_from_array():
         (da.ones(shape=(10, 10), chunks=(10, 2)), (None, (5, 5))),
     ],
 )
-def test_rechunk_unknown(x, chunks):
+def test_rechunk_with_fully_unknown_dimension(x, chunks):
     dd = pytest.importorskip("dask.dataframe")
     y = dd.from_array(x).values
     result = y.rechunk(chunks)
@@ -629,7 +640,35 @@ def test_rechunk_unknown(x, chunks):
     assert_eq(result, expected)
 
 
-def test_rechunk_unknown_explicit():
+@pytest.mark.parametrize(
+    "x, chunks",
+    [
+        (da.ones(shape=(50, 10), chunks=(25, 10)), (None, 5)),
+        (da.ones(shape=(50, 10), chunks=(25, 10)), {1: 5}),
+        (da.ones(shape=(50, 10), chunks=(25, 10)), (None, (5, 5))),
+        (da.ones(shape=(1000, 10), chunks=(5, 10)), (None, 5)),
+        (da.ones(shape=(1000, 10), chunks=(5, 10)), {1: 5}),
+        (da.ones(shape=(1000, 10), chunks=(5, 10)), (None, (5, 5))),
+        (da.ones(shape=(10, 10), chunks=(10, 10)), (None, 5)),
+        (da.ones(shape=(10, 10), chunks=(10, 10)), {1: 5}),
+        (da.ones(shape=(10, 10), chunks=(10, 10)), (None, (5, 5))),
+        (da.ones(shape=(10, 10), chunks=(10, 2)), (None, 5)),
+        (da.ones(shape=(10, 10), chunks=(10, 2)), {1: 5}),
+        (da.ones(shape=(10, 10), chunks=(10, 2)), (None, (5, 5))),
+    ],
+)
+def test_rechunk_with_partially_unknown_dimension(x, chunks):
+    dd = pytest.importorskip("dask.dataframe")
+    y = dd.from_array(x).values
+    z = da.concatenate([x, y])
+    xx = da.concatenate([x, x])
+    result = z.rechunk(chunks)
+    expected = xx.rechunk(chunks)
+    assert_chunks_match(result.chunks, expected.chunks)
+    assert_eq(result, expected)
+
+
+def test_rechunk_with_fully_unknown_dimension_explicit():
     dd = pytest.importorskip("dask.dataframe")
     x = da.ones(shape=(10, 10), chunks=(5, 2))
     y = dd.from_array(x).values
@@ -639,20 +678,37 @@ def test_rechunk_unknown_explicit():
     assert_eq(result, expected)
 
 
+def test_rechunk_with_partially_unknown_dimension_explicit():
+    dd = pytest.importorskip("dask.dataframe")
+    x = da.ones(shape=(10, 10), chunks=(5, 2))
+    y = dd.from_array(x).values
+    z = da.concatenate([x, y])
+    xx = da.concatenate([x, x])
+    result = z.rechunk(((5, 5, np.nan, np.nan), (5, 5)))
+    expected = xx.rechunk((None, (5, 5)))
+    assert_chunks_match(result.chunks, expected.chunks)
+    assert_eq(result, expected)
+
+
 def assert_chunks_match(left, right):
-    for x, y in zip(left, right):
-        if np.isnan(x).any():
-            assert np.isnan(x).all()
-        else:
-            assert x == y
+    for ldim, rdim in zip(left, right):
+        assert all(np.isnan(l) or l == r for l, r in zip(ldim, rdim))
 
 
 def test_rechunk_unknown_raises():
     dd = pytest.importorskip("dask.dataframe")
 
-    x = dd.from_array(da.ones(shape=(10, 10), chunks=(5, 5))).values
-    with pytest.raises(ValueError):
-        x.rechunk((None, (5, 5, 5)))
+    x = da.ones(shape=(10, 10), chunks=(5, 5))
+    y = dd.from_array(x).values
+    with pytest.raises(ValueError, match="Chunks do not add"):
+        y.rechunk((None, (5, 5, 5)))
+
+    with pytest.raises(ValueError, match="Chunks must be unchanging"):
+        y.rechunk(((5, 5), (5, 5)))
+
+    with pytest.raises(ValueError, match="Chunks must be unchanging"):
+        z = da.concatenate([x, y])
+        z.rechunk(((5, 3, 2, np.nan, np.nan), (5, 5)))
 
 
 def test_old_to_new_single():


Code review comments:
diff_hunk: @@ -157,7 +157,7 @@ def _intersect_1d(breaks):
 def _old_to_new(old_chunks, new_chunks):
path: dask/array/rechunk.py
text: @user1:
nit: if this is used elsewhere, I suggest to remove the leading underscore.

diff_hunk: @@ -169,31 +169,44 @@ def _old_to_new(old_chunks, new_chunks):
       [(2, slice(5, 10, None))],
       [(3, slice(0, 10, None)), (4, slice(0, 10, None))]]]
     """
-    old_known = [x for x in old_chunks if not any(math.isnan(y) for y in x)]
-    new_known = [x for x in new_chunks if not any(math.isnan(y) for y in x)]
 
-    n_missing = [sum(math.isnan(y) for y in x) for x in old_chunks]
-    n_missing2 = [sum(math.isnan(y) for y in x) for x in new_chunks]
+    def is_unknown(dim):
+        return any(math.isnan(chunk) for chunk in dim)
 
-    cmo = cumdims_label(old_known, "o")
-    cmn = cumdims_label(new_known, "n")
-
-    sums = [sum(o) for o in old_known]
-    sums2 = [sum(n) for n in new_known]
+    old_is_unknown = [is_unknown(dim) for dim in old_chunks]
+    new_is_unknown = [is_unknown(dim) for dim in new_chunks]
 
-    if not sums == sums2:
-        raise ValueError(f"Cannot change dimensions from {sums!r} to {sums2!r}")
-    if not n_missing == n_missing2:
+    if old_is_unknown != new_is_unknown or any(
+        new != old for new, old in compress(zip(old_chunks, new_chunks), old_is_unknown)
+    ):
         raise ValueError(
-            "Chunks must be unchanging along unknown dimensions.\n\n"
+            "Chunks must be unchanging along dimensions with missing values.\n\n"
             "A possible solution:\n  x.compute_chunk_sizes()"
         )
 
+    old_known = [dim for dim, unknown in zip(old_chunks, old_is_unknown) if not unknown]
+    new_known = [dim for dim, unknown in zip(new_chunks, new_is_unknown) if not unknown]
+
+    old_sizes = [sum(o) for o in old_known]
+    new_sizes = [sum(n) for n in new_known]
+
+    if not old_sizes == new_sizes:
+        raise ValueError(
+            f"Cannot change dimensions from {old_sizes!r} to {new_sizes!r}"
+        )
+
+    cmo = cumdims_label(old_known, "o")
+    cmn = cumdims_label(new_known, "n")
+
     old_to_new = [_intersect_1d(_breakpoints(cm[0], cm[1])) for cm in zip(cmo, cmn)]
-    for idx, missing in enumerate(n_missing):
-        if missing:
-            # Missing dimensions are always unchanged, so old -> new is everything
-            extra = [[(i, slice(0, None))] for i in range(missing)]
+    for idx, unknown in enumerate(old_is_unknown):
+        if unknown:
+            dim = old_chunks[idx]
+            # Unknown dimensions are always unchanged, so old -> new is everything
+            extra = [
+                [(i, slice(0, size if not math.isnan(size) else None))]
+                for i, size in enumerate(dim)
+            ]
             old_to_new.insert(idx, extra)
path: dask/array/rechunk.py
text: @user1:
This insertion is requiring on some quite delicate interplay between the way `old_to_new` is constructed and the `old_is_unknown` list. Since they are different lengths, insertion into `old_to_new` using indices from `old_is_unknown` doesn't _seem_ to make sense. But it so happens that everything is set up so the insertion exactly maintains the invariant that as we walk through `old_is_unknown`, insertion is always at the "right" place.

I wonder if it is better to make `old_to_new` of the full length before hand (with placeholders) and then replace the placeholders with a subsequent path.

@author:
It's a somewhat intricate multi-pass approach given how the code has been written, but we now populate a prepared list.

diff_hunk: @@ -169,31 +169,44 @@ def _old_to_new(old_chunks, new_chunks):
       [(2, slice(5, 10, None))],
       [(3, slice(0, 10, None)), (4, slice(0, 10, None))]]]
     """
-    old_known = [x for x in old_chunks if not any(math.isnan(y) for y in x)]
-    new_known = [x for x in new_chunks if not any(math.isnan(y) for y in x)]
 
-    n_missing = [sum(math.isnan(y) for y in x) for x in old_chunks]
-    n_missing2 = [sum(math.isnan(y) for y in x) for x in new_chunks]
+    def is_unknown(dim):
+        return any(math.isnan(chunk) for chunk in dim)
 
-    cmo = cumdims_label(old_known, "o")
-    cmn = cumdims_label(new_known, "n")
-
-    sums = [sum(o) for o in old_known]
-    sums2 = [sum(n) for n in new_known]
+    old_is_unknown = [is_unknown(dim) for dim in old_chunks]
+    new_is_unknown = [is_unknown(dim) for dim in new_chunks]
 
-    if not sums == sums2:
-        raise ValueError(f"Cannot change dimensions from {sums!r} to {sums2!r}")
-    if not n_missing == n_missing2:
+    if old_is_unknown != new_is_unknown or any(
+        new != old for new, old in compress(zip(old_chunks, new_chunks), old_is_unknown)
+    ):
         raise ValueError(
-            "Chunks must be unchanging along unknown dimensions.\n\n"
+            "Chunks must be unchanging along dimensions with missing values.\n\n"
             "A possible solution:\n  x.compute_chunk_sizes()"
         )
 
+    old_known = [dim for dim, unknown in zip(old_chunks, old_is_unknown) if not unknown]
+    new_known = [dim for dim, unknown in zip(new_chunks, new_is_unknown) if not unknown]
+
+    old_sizes = [sum(o) for o in old_known]
+    new_sizes = [sum(n) for n in new_known]
+
+    if not old_sizes == new_sizes:
path: dask/array/rechunk.py
text: @user1:
```suggestion
    if old_sizes != new_sizes:
```