Make allow_irregular default to True for rechunk (#933)

* Change allow_irregular default to True; update tests

* Add test for explicit allow_irregular=False

* Remove no-longer-needed Zarr Python v2 check

Author: tomwhite

cubed/core/array.py

@@ -179,7 +179,7 @@ def plan(
         )
 
     def rechunk(
-        self: T_ChunkedArray, chunks, *, min_mem=None, allow_irregular=False
+        self: T_ChunkedArray, chunks, *, min_mem=None, allow_irregular=True
     ) -> T_ChunkedArray:
         """Change the chunking of this array without changing its shape or data.
 

cubed/core/ops.py

@@ -948,7 +948,7 @@ def _map_blocks(
     )
 
 
-def rechunk(x, chunks, *, min_mem=None, allow_irregular=False):
+def rechunk(x, chunks, *, min_mem=None, allow_irregular=True):
     """Change the chunking of an array without changing its shape or data.
 
     Parameters
@@ -969,7 +969,7 @@ def rechunk(x, chunks, *, min_mem=None, allow_irregular=False):
     return out
 
 
-def _rechunk_plan(x, chunks, *, min_mem=None, allow_irregular=False):
+def _rechunk_plan(x, chunks, *, min_mem=None, allow_irregular=True):
     if isinstance(chunks, dict):
         chunks = {validate_axis(c, x.ndim): v for c, v in chunks.items()}
         for i in range(x.ndim):
@@ -1052,7 +1052,7 @@ def split_chunksizes(n: int, sc: int, tc: int) -> tuple[int]:
     return tuple(np.diff(c).tolist())
 
 
-def _rechunk(x, copy_chunks, target_chunks, allow_irregular=False):
+def _rechunk(x, copy_chunks, target_chunks, allow_irregular=True):
     # rechunk x so that its target store has target_chunks, using copy_chunks as the size of chunks for copying from source to target
 
     normalized_copy_chunks = normalize_chunks(copy_chunks, x.shape, dtype=x.dtype)

cubed/core/rechunk.py

@@ -311,7 +311,7 @@ def from_plan(cls, rechunk_plan: RechunkPlan, plan: FinalizedPlan):
         )
 
 
-def rechunk_plan(x, chunks, *, min_mem=None, allow_irregular=False):
+def rechunk_plan(x, chunks, *, min_mem=None, allow_irregular=True):
     from cubed.core.ops import _rechunk_plan
 
     copy_ops = []

cubed/tests/test_array_api.py

@@ -1,6 +1,5 @@
 import numpy as np
 import pytest
-import zarr
 
 import cubed
 import cubed.array_api as xp
@@ -14,8 +13,6 @@
     skip_if_cupy,
 )
 
-ZARR_PYTHON_V2 = zarr.__version__[0] == "2"
-
 
 @pytest.fixture
 def spec(tmp_path):
@@ -662,7 +659,6 @@ def test_flip(executor, shape, chunks, axis):
 
 
 
-@pytest.mark.skipif(ZARR_PYTHON_V2, reason="fails with Zarr v2")
 def test_flip_zero_size_dim(spec):
     x = np.ones((0, 4))
     a = xp.asarray(x, chunks=(1, 2), spec=spec)

cubed/tests/test_rechunk.py

@@ -1,5 +1,4 @@
 import pytest
-import zarr
 
 import cubed
 import cubed as xp
@@ -29,7 +28,7 @@
         "expected_max_output_blocks",
     ),
     [
-        (None, 5, 7, 9),  # multistage rechunk - more stages, lower fan in/out
+        (None, 3, 16, 15),  # multistage rechunk - more stages, lower fan in/out
         (0, 1, 3771, 3460),  # single stage rechunk - very high fan in/out
     ],
 )
@@ -97,11 +96,9 @@ def test_rechunk_era5_chunk_sizes(spec):
 
     rechunk_plan = list(_rechunk_plan(a, target_chunks))
     assert rechunk_plan == [
-        ((93, 721, 1440), (93, 240, 480)),
-        ((465, 240, 480), (465, 120, 240)),
-        ((2325, 120, 240), (2325, 40, 120)),
-        ((11625, 40, 120), (11625, 20, 60)),
-        ((58125, 20, 60), (58125, 10, 30)),
+        ((93, 721, 1440), (93, 173, 396)),
+        ((1447, 173, 396), (1447, 41, 109)),
+        ((22528, 41, 109), (22528, 10, 30)),
         ((350640, 10, 30), (350640, 10, 10)),
     ]
 
@@ -140,32 +137,32 @@ def test_rechunk_hypothesis_generated_bug():
     from cubed.core.ops import _rechunk_plan
 
     rechunk_plan = list(_rechunk_plan(a, target_chunks))
-    assert rechunk_plan == [((30, 376), (15, 376)), ((15, 1001), (5, 146))]
-    for copy_chunks, target_chunks in rechunk_plan:
-        verify_chunk_compatibility(shape, copy_chunks, target_chunks)
+    assert rechunk_plan == [((38, 376), (15, 376)), ((15, 1001), (5, 146))]
 
     b = a.rechunk((5, 146))
 
     assert_array_equal(b.compute(), nxp.ones(shape))
 
 
-@pytest.mark.skipif(
-    zarr.__version__[0] == "2",
-    reason="irregular chunking is not supported for Zarr Python v2",
-)
-def test_rechunk_hypothesis_generated_bug_allow_irregular():
-    rechunk_shapes = (tuple([1001, 1001]), (38, 376), (5, 146))
-    shape, source_chunks, target_chunks = rechunk_shapes
+def test_rechunk_allow_irregular_false():
+    # Verify the regular planner still works when allow_irregular=False is
+    # passed explicitly. Regular intermediates must be integer multiples of
+    # target chunks, so copy_chunks[0]=30 (multiple of 15) instead of 38.
+    shape = (1001, 1001)
+    source_chunks = (38, 376)
+    target_chunks = (5, 146)
 
     spec = cubed.Spec(allowed_mem=8000000 / 10)
     a = xp.ones(shape, chunks=source_chunks, spec=spec)
 
     from cubed.core.ops import _rechunk_plan
 
-    rechunk_plan = list(_rechunk_plan(a, target_chunks, allow_irregular=True))
-    assert rechunk_plan == [((38, 376), (15, 376)), ((15, 1001), (5, 146))]
+    rechunk_plan = list(_rechunk_plan(a, target_chunks, allow_irregular=False))
+    assert rechunk_plan == [((30, 376), (15, 376)), ((15, 1001), (5, 146))]
+    for copy_chunks, store_chunks in rechunk_plan:
+        verify_chunk_compatibility(shape, copy_chunks, store_chunks)
 
-    b = a.rechunk((5, 146), allow_irregular=True)
+    b = a.rechunk(target_chunks, allow_irregular=False)
 
     assert_array_equal(b.compute(), nxp.ones(shape))