Remove duplicate code from new rechunk implementation (#702)

* Remove duplicate code from new rechunk implementation

* Test rechunking plan chunk sizes for ERA5

Author: tomwhite

cubed/core/ops.py

@@ -1126,6 +1126,13 @@ def rechunk_new(x, chunks, *, min_mem=None):
     cubed.Array
         An array with the desired chunks.
     """
+    out = x
+    for copy_chunks, target_chunks in _rechunk_plan(x, chunks, min_mem=min_mem):
+        out = _rechunk(out, copy_chunks, target_chunks)
+    return out
+
+
+def _rechunk_plan(x, chunks, *, min_mem=None):
     if isinstance(chunks, dict):
         chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
         for i in range(x.ndim):
@@ -1165,7 +1172,6 @@ def rechunk_new(x, chunks, *, min_mem=None):
         max_mem=rechunker_max_mem,
     )
 
-    out = x
     for i, stage in enumerate(stages):
         last_stage = i == len(stages) - 1
         read_chunks, int_chunks, write_chunks = stage
@@ -1174,12 +1180,10 @@ def rechunk_new(x, chunks, *, min_mem=None):
         target_chunks_ = target_chunks if last_stage else write_chunks
 
         if read_chunks == write_chunks:
-            out = _rechunk(out, read_chunks, target_chunks_)
+            yield read_chunks, target_chunks_
         else:
-            intermediate = _rechunk(out, read_chunks, int_chunks)
-            out = _rechunk(intermediate, write_chunks, target_chunks_)
-
-    return out
+            yield read_chunks, int_chunks
+            yield write_chunks, target_chunks_
 
 
 def _rechunk(x, copy_chunks, target_chunks):
@@ -1217,62 +1221,6 @@ def selection_function(out_key):
     )
 
 
-def rechunk_plan(x, chunks, *, min_mem=None):
-    if isinstance(chunks, dict):
-        chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
-        for i in range(x.ndim):
-            if i not in chunks:
-                chunks[i] = x.chunks[i]
-            elif chunks[i] is None:
-                chunks[i] = x.chunks[i]
-    if isinstance(chunks, (tuple, list)):
-        chunks = tuple(lc if lc is not None else rc for lc, rc in zip(chunks, x.chunks))
-
-    normalized_chunks = normalize_chunks(chunks, x.shape, dtype=x.dtype)
-    if x.chunks == normalized_chunks:
-        return x
-    # normalizing takes care of dict args for chunks
-    target_chunks = to_chunksize(normalized_chunks)
-
-    # merge chunks special case
-    if all(c1 % c0 == 0 for c0, c1 in zip(x.chunksize, target_chunks)):
-        return merge_chunks(x, target_chunks)
-
-    spec = x.spec
-    source_chunks = to_chunksize(normalize_chunks(x.chunks, x.shape, dtype=x.dtype))
-
-    # rechunker doesn't take account of uncompressed and compressed copies of the
-    # input and output array chunk/selection, so adjust appropriately
-    rechunker_max_mem = (spec.allowed_mem - spec.reserved_mem) // 5
-    if min_mem is None:
-        min_mem = min(rechunker_max_mem // 20, x.nbytes)
-    stages = multistage_rechunking_plan(
-        shape=x.shape,
-        source_chunks=source_chunks,
-        target_chunks=target_chunks,
-        itemsize=x.dtype.itemsize,
-        min_mem=min_mem,
-        max_mem=rechunker_max_mem,
-    )
-
-    source_chunks = x.chunksize
-    for i, stage in enumerate(stages):
-        last_stage = i == len(stages) - 1
-        read_chunks, int_chunks, write_chunks = stage
-
-        # Use target chunks for last stage
-        target_chunks_ = target_chunks if last_stage else write_chunks
-
-        if read_chunks == write_chunks:
-            yield source_chunks, read_chunks, target_chunks_
-            source_chunks = target_chunks_
-        else:
-            yield source_chunks, read_chunks, int_chunks
-            source_chunks = int_chunks
-            yield source_chunks, write_chunks, target_chunks_
-            source_chunks = target_chunks_
-
-
 def merge_chunks(x, chunks):
     """Merge multiple chunks into one."""
     target_chunksize = chunks

cubed/tests/test_mem_utilization_rechunk.py

@@ -66,10 +66,10 @@ def test_rechunk_era5(tmp_path, spec, executor):
 
     x = cubed.random.random(shape, dtype=xp.float32, chunks=source_chunks, spec=spec)
 
-    from cubed.core.ops import rechunk_plan
+    from cubed.core.ops import _rechunk_plan
 
     i = 0
-    for source_chunks, copy_chunks, target_chunks in rechunk_plan(x, target_chunks):
+    for copy_chunks, target_chunks in _rechunk_plan(x, target_chunks):
         # Find the smallest shape that contains the three chunk sizes
         # This will be a lot less than the full ERA5 shape (350640, 721, 1440),
         # making it suitable for running in a test
@@ -85,4 +85,5 @@ def test_rechunk_era5(tmp_path, spec, executor):
 
         run_operation(tmp_path, executor, f"rechunk_era5_stage_{i}", b)
 
+        source_chunks = target_chunks
         i += 1

cubed/tests/test_rechunk.py

@@ -52,3 +52,25 @@ def test_rechunk_era5(
         d["pipeline"].config.num_output_blocks[0] for _, d in rechunks
     )
     assert max_output_blocks == expected_max_output_blocks
+
+
+def test_rechunk_era5_chunk_sizes():
+    # from https://github.com/pangeo-data/rechunker/pull/89
+    shape = (350640, 721, 1440)
+    source_chunks = (31, 721, 1440)
+    target_chunks = (350640, 10, 10)
+
+    spec = cubed.Spec(allowed_mem="2.5GB")
+
+    a = xp.empty(shape, dtype=xp.float32, chunks=source_chunks, spec=spec)
+
+    from cubed.core.ops import _rechunk_plan
+
+    assert list(_rechunk_plan(a, target_chunks)) == [
+        ((93, 721, 1440), (93, 173, 396)),
+        ((1447, 173, 396), (1447, 173, 396)),
+        ((1447, 173, 396), (1447, 41, 109)),
+        ((22528, 41, 109), (22528, 41, 109)),
+        ((22528, 41, 109), (22528, 10, 30)),
+        ((350640, 10, 30), (350640, 10, 10)),
+    ]