Repository: dask/dask

Issue:
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



Discussion:

Resolving patch:
diff --git a/dask/array/rechunk.py b/dask/array/rechunk.py
--- 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
 
@@ -154,48 +154,71 @@ def _intersect_1d(breaks):
     return ret
 
 
-def _old_to_new(old_chunks, new_chunks):
+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.
 
+    This function expects that the arguments have been pre-processed by
+    :func:`dask.array.core.normalize_chunks`. In particular any ``nan`` values should
+    have been replaced (and are so by :func:`dask.array.core.normalize_chunks`)
+    by the canonical ``np.nan``.
+
     Examples
     --------
     >>> old = ((10, 10, 10, 10, 10), )
     >>> new = ((25, 5, 20), )
-    >>> _old_to_new(old, new)  # doctest: +NORMALIZE_WHITESPACE
+    >>> old_to_new(old, new)  # doctest: +NORMALIZE_WHITESPACE
     [[[(0, slice(0, 10, None)), (1, slice(0, 10, None)), (2, slice(0, 5, None))],
       [(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]
 
-    cmo = cumdims_label(old_known, "o")
-    cmn = cumdims_label(new_known, "n")
+    def is_unknown(dim):
+        return any(math.isnan(chunk) for chunk in dim)
 
-    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_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)]
-            old_to_new.insert(idx, extra)
-    return old_to_new
+    old_known_indices = [i for i, unknown in enumerate(old_is_unknown) if not unknown]
+    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 old_sizes != new_sizes:
+        raise ValueError(
+            f"Cannot change dimensions from {old_sizes!r} to {new_sizes!r}"
+        )
+
+    cmos = cumdims_label(old_known, "o")
+    cmns = cumdims_label(new_known, "n")
+
+    sliced = [None] * len(old_chunks)
+    for i, cmo, cmn in zip(old_known_indices, cmos, cmns):
+        sliced[i] = _intersect_1d(_breakpoints(cmo, cmn))
+
+    for i, unknown in enumerate(old_is_unknown):
+        if unknown:
+            dim = old_chunks[i]
+            # Unknown dimensions are always unchanged, so old -> new is everything
+            extra = [
+                [(j, slice(0, size if not math.isnan(size) else None))]
+                for j, size in enumerate(dim)
+            ]
+            sliced[i] = extra
+    assert all(x is not None for x in sliced)
+    return sliced
 
 
 def intersect_chunks(old_chunks, new_chunks):
@@ -218,9 +241,7 @@ def intersect_chunks(old_chunks, new_chunks):
     new_chunks: iterable of tuples
         block sizes along each dimension (converts to new_chunks)
     """
-    old_to_new = _old_to_new(old_chunks, new_chunks)
-
-    cross1 = product(*old_to_new)
+    cross1 = product(*old_to_new(old_chunks, new_chunks))
     cross = chain(tuple(product(*cr)) for cr in cross1)
     return cross
 


Test patch:
diff --git a/dask/array/tests/test_rechunk.py b/dask/array/tests/test_rechunk.py
--- a/dask/array/tests/test_rechunk.py
+++ b/dask/array/tests/test_rechunk.py
@@ -4,18 +4,19 @@
 import pytest
 
 np = pytest.importorskip("numpy")
+import math
 
 import dask
 import dask.array as da
 from dask.array.rechunk import (
     _breakpoints,
     _intersect_1d,
-    _old_to_new,
     cumdims_label,
     divide_to_width,
     intersect_chunks,
     merge_to_number,
     normalize_chunks,
+    old_to_new,
     plan_rechunk,
     rechunk,
 )
@@ -282,7 +283,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 +293,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 +305,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 +632,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,36 +641,97 @@ 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)
+
+
+@pytest.mark.parametrize(
+    "new_chunks",
+    [
+        ((np.nan, np.nan), (5, 5)),
+        ((math.nan, math.nan), (5, 5)),
+        ((float("nan"), float("nan")), (5, 5)),
+    ],
+)
+def test_rechunk_with_fully_unknown_dimension_explicit(new_chunks):
     dd = pytest.importorskip("dask.dataframe")
     x = da.ones(shape=(10, 10), chunks=(5, 2))
     y = dd.from_array(x).values
-    result = y.rechunk(((np.nan, np.nan), (5, 5)))
+    result = y.rechunk(new_chunks)
     expected = x.rechunk((None, (5, 5)))
     assert_chunks_match(result.chunks, expected.chunks)
     assert_eq(result, expected)
 
 
+@pytest.mark.parametrize(
+    "new_chunks",
+    [
+        ((5, 5, np.nan, np.nan), (5, 5)),
+        ((5, 5, math.nan, math.nan), (5, 5)),
+        ((5, 5, float("nan"), float("nan")), (5, 5)),
+    ],
+)
+def test_rechunk_with_partially_unknown_dimension_explicit(new_chunks):
+    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(new_chunks)
+    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():
     old = ((np.nan, np.nan), (8,))
     new = ((np.nan, np.nan), (4, 4))
-    result = _old_to_new(old, new)
+    result = old_to_new(old, new)
 
     expected = [
         [[(0, slice(0, None, None))], [(1, slice(0, None, None))]],
@@ -671,7 +744,7 @@ def test_old_to_new_single():
 def test_old_to_new():
     old = ((np.nan,), (10,))
     new = ((np.nan,), (5, 5))
-    result = _old_to_new(old, new)
+    result = old_to_new(old, new)
     expected = [
         [[(0, slice(0, None, None))]],
         [[(0, slice(0, 5, None))], [(0, slice(5, 10, None))]],
@@ -684,7 +757,7 @@ def test_old_to_new_large():
     old = (tuple([np.nan] * 4), (10,))
     new = (tuple([np.nan] * 4), (5, 5))
 
-    result = _old_to_new(old, new)
+    result = old_to_new(old, new)
     expected = [
         [
             [(0, slice(0, None, None))],
@@ -699,7 +772,7 @@ def test_old_to_new_large():
 
 def test_changing_raises():
     with pytest.raises(ValueError) as record:
-        _old_to_new(((np.nan, np.nan), (4, 4)), ((np.nan, np.nan, np.nan), (4, 4)))
+        old_to_new(((np.nan, np.nan), (4, 4)), ((np.nan, np.nan, np.nan), (4, 4)))
 
     assert "unchanging" in str(record.value)
 
@@ -707,7 +780,7 @@ def test_changing_raises():
 def test_old_to_new_known():
     old = ((10, 10, 10, 10, 10),)
     new = ((25, 5, 20),)
-    result = _old_to_new(old, new)
+    result = old_to_new(old, new)
     expected = [
         [
             [(0, slice(0, 10, None)), (1, slice(0, 10, None)), (2, slice(0, 5, None))],
@@ -1071,23 +1144,23 @@ def test_intersect_chunks_with_zero():
 
 
 def test_old_to_new_with_zero():
-    from dask.array.rechunk import _old_to_new
+    from dask.array.rechunk import old_to_new
 
     old = ((4, 4),)
     new = ((4, 0, 4),)
-    result = _old_to_new(old, new)
+    result = old_to_new(old, new)
     expected = [[[(0, slice(0, 4))], [(1, slice(0, 0))], [(1, slice(0, 4))]]]
     assert result == expected
 
     old = ((4,),)
     new = ((4, 0),)
-    result = _old_to_new(old, new)
+    result = old_to_new(old, new)
     expected = [[[(0, slice(0, 4))], [(0, slice(4, 4))]]]
     assert result == expected
 
     old = ((4, 0, 4),)
     new = ((4, 0, 2, 2),)
-    result = _old_to_new(old, new)
+    result = old_to_new(old, new)
     expected = [
         [[(0, slice(0, 4))], [(2, slice(0, 0))], [(2, slice(0, 2))], [(2, slice(2, 4))]]
     ]