fixups + support partial edge chunks / resize

Author: jhamman

docs/user-guide/arrays.md

@@ -676,14 +676,6 @@ Variable chunking has some important limitations:
    property will raise a `NotImplementedError`. Use `.metadata.chunk_grid.chunk_shapes`
    instead.
 
-```python exec="true" session="arrays" source="above" result="ansi"
-# This will raise an error
-try:
-    _ = z.chunks
-except NotImplementedError as e:
-    print(f"Error: {e}")
-```
-
 ## Missing features in 3.0
 
 The following features have not been ported to 3.0 yet.

src/zarr/api/synchronous.py

@@ -13,7 +13,7 @@
 from zarr.errors import ZarrDeprecationWarning
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Sequence
+    from collections.abc import Iterable
 
     import numpy as np
     import numpy.typing as npt
@@ -29,7 +29,7 @@
     )
     from zarr.core.array_spec import ArrayConfigLike
     from zarr.core.buffer import NDArrayLike, NDArrayLikeOrScalar
-    from zarr.core.chunk_grids import ChunkGrid
+    from zarr.core.chunk_grids import ChunksLike
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
     from zarr.core.common import (
         JSON,
@@ -822,7 +822,7 @@ def create_array(
     shape: ShapeLike | None = None,
     dtype: ZDTypeLike | None = None,
     data: np.ndarray[Any, np.dtype[Any]] | None = None,
-    chunks: tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"] = "auto",
+    chunks: ChunksLike = "auto",
     shards: ShardsLike | None = None,
     filters: FiltersLike = "auto",
     compressors: CompressorsLike = "auto",
@@ -858,7 +858,7 @@ def create_array(
     data : np.ndarray, optional
         Array-like data to use for initializing the array. If this parameter is provided, the
         ``shape`` and ``dtype`` parameters must be ``None``.
-    chunks : tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"], default="auto"
+    chunks : ChunksLike, default="auto"
         Chunk shape of the array. Several formats are supported:
 
         - tuple of ints: Creates a RegularChunkGrid with uniform chunks, e.g., ``(10, 10)``

src/zarr/core/array.py

@@ -42,6 +42,7 @@
 from zarr.core.buffer.cpu import buffer_prototype as cpu_buffer_prototype
 from zarr.core.chunk_grids import (
     ChunkGrid,
+    ChunksLike,
     RegularChunkGrid,
     _auto_partition,
     _normalize_chunks,
@@ -1856,7 +1857,7 @@ async def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True)
         if delete_outside_chunks:
             # Remove all chunks outside of the new shape
             old_chunk_coords = set(self.metadata.chunk_grid.all_chunk_coords(self.metadata.shape))
-            new_chunk_coords = set(self.metadata.chunk_grid.all_chunk_coords(new_shape))
+            new_chunk_coords = set(new_metadata.chunk_grid.all_chunk_coords(new_shape))
 
             async def _delete_key(key: str) -> None:
                 await (self.store_path / key).delete()
@@ -4885,7 +4886,7 @@ async def create_array(
     shape: ShapeLike | None = None,
     dtype: ZDTypeLike | None = None,
     data: np.ndarray[Any, np.dtype[Any]] | None = None,
-    chunks: tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"] = "auto",
+    chunks: ChunksLike = "auto",
     shards: ShardsLike | None = None,
     filters: FiltersLike = "auto",
     compressors: CompressorsLike = "auto",
@@ -4919,7 +4920,7 @@ async def create_array(
     data : np.ndarray, optional
         Array-like data to use for initializing the array. If this parameter is provided, the
         ``shape`` and ``dtype`` parameters must be ``None``.
-    chunks : tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"], default="auto"
+    chunks : ChunksLike, default="auto"
         Chunk shape of the array. Several formats are supported:
 
         - tuple of ints: Creates a RegularChunkGrid with uniform chunks, e.g., ``(10, 10)``

src/zarr/core/chunk_grids.py

@@ -8,9 +8,9 @@
 import warnings
 from abc import abstractmethod
 from collections.abc import Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from functools import cached_property, reduce
-from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict, Union
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict, Union, cast
 
 import numpy as np
 
@@ -32,15 +32,16 @@
 
 
 # Type alias for chunk edge length specification
-# Can be either an integer or a run-length encoded tuple [value, count]
-ChunkEdgeLength = int | tuple[int, int]
+# Can be either an integer or a run-length encoded pair [value, count]
+# The pair can be a tuple or list (common in JSON/test code)
+ChunkEdgeLength = int | tuple[int, int] | list[int]
 
 # User-facing chunk specification types
 # Note: ChunkGrid is defined later in this file but can be used via string literal
 ChunksLike: TypeAlias = Union[
     tuple[int, ...],  # Regular chunks: (10, 10) → RegularChunkGrid
     int,  # Uniform chunks: 10 → RegularChunkGrid
-    Sequence[Sequence[int]],  # Variable chunks: [[10,20],[5,5]] → RectilinearChunkGrid
+    Sequence[Sequence[ChunkEdgeLength]],  # Variable chunks with optional RLE → RectilinearChunkGrid
     "ChunkGrid",  # Explicit ChunkGrid instance (forward reference)
     Literal["auto"],  # Auto-chunking → RegularChunkGrid
 ]
@@ -295,6 +296,7 @@ def _normalize_chunks(chunks: Any, shape: tuple[int, ...], typesize: int) -> tup
         chunks = tuple(int(chunks) for _ in shape)
 
     # handle dask-style chunks (iterable of iterables)
+    # TODO
     if all(isinstance(c, (tuple | list)) for c in chunks):
         # Check for irregular chunks and warn user
         for dim_idx, c in enumerate(chunks):
@@ -338,13 +340,20 @@ def from_dict(cls, data: dict[str, JSON] | ChunkGrid) -> ChunkGrid:
         if isinstance(data, ChunkGrid):
             return data
 
-        name_parsed, _ = parse_named_configuration(data)
+        # After isinstance check, data must be dict[str, JSON]
+        # Cast needed for older mypy versions that don't narrow types properly
+        data_dict = cast(dict[str, JSON], data)  # type: ignore[redundant-cast]
+        name_parsed, _ = parse_named_configuration(data_dict)
         if name_parsed == "regular":
-            return RegularChunkGrid._from_dict(data)
+            return RegularChunkGrid._from_dict(data_dict)
         elif name_parsed == "rectilinear":
-            return RectilinearChunkGrid._from_dict(data)
+            return RectilinearChunkGrid._from_dict(data_dict)
         raise ValueError(f"Unknown chunk grid. Got {name_parsed}.")
 
+    @abstractmethod
+    def update_shape(self, new_shape: tuple[int, ...]) -> Self:
+        pass
+
     @abstractmethod
     def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
         pass
@@ -466,6 +475,9 @@ def _from_dict(cls, data: dict[str, JSON]) -> Self:
     def to_dict(self) -> dict[str, JSON]:
         return {"name": "regular", "configuration": {"chunk_shape": tuple(self.chunk_shape)}}
 
+    def update_shape(self, new_shape: tuple[int, ...]) -> Self:
+        return self
+
     def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
         return itertools.product(
             *(range(ceildiv(s, c)) for s, c in zip(array_shape, self.chunk_shape, strict=False))
@@ -645,6 +657,46 @@ def to_dict(self) -> dict[str, JSON]:
             },
         }
 
+    def update_shape(self, new_shape: tuple[int, ...]) -> Self:
+        """TODO - write docstring"""
+
+        if len(new_shape) != len(self.chunk_shapes):
+            raise ValueError(
+                f"new_shape has {len(new_shape)} dimensions but "
+                f"chunk_shapes has {len(self.chunk_shapes)} dimensions"
+            )
+
+        new_chunk_shapes: list[tuple[int, ...]] = []
+        for dim in range(len(new_shape)):
+            old_dim_length = sum(self.chunk_shapes[dim])
+            new_dim_chunks: tuple[int, ...]
+            if new_shape[dim] == old_dim_length:
+                new_dim_chunks = self.chunk_shapes[dim]  # no changes
+
+            elif new_shape[dim] > old_dim_length:
+                # we have a decision to make on chunk size...
+                # options:
+                # - repeat the last chunk size
+                # - use the size of the new data
+                # - use some other heuristic
+                # for now, we'll use the difference in size between the old and new dim length
+                new_dim_chunks = (*self.chunk_shapes[dim], new_shape[dim] - old_dim_length)
+            else:
+                # drop chunk sizes that are not inside the shape anymore
+                total = 0
+                i = 0
+                for c in self.chunk_shapes[dim]:
+                    i += 1
+                    total += c
+                    if total >= new_shape[dim]:
+                        break
+                # keep the last chunk (it may be too long)
+                new_dim_chunks = self.chunk_shapes[dim][:i]
+
+            new_chunk_shapes.append(new_dim_chunks)
+
+        return replace(self, chunk_shapes=tuple(new_chunk_shapes))
+
     def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
         """
         Generate all chunk coordinates for the given array shape.
@@ -664,22 +716,9 @@ def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int,
         ValueError
             If array_shape doesn't match chunk_shapes
         """
-        if len(array_shape) != len(self.chunk_shapes):
-            raise ValueError(
-                f"array_shape has {len(array_shape)} dimensions but "
-                f"chunk_shapes has {len(self.chunk_shapes)} dimensions"
-            )
 
-        # Validate that chunk sizes sum to array shape
-        for axis, (arr_size, axis_chunks) in enumerate(
-            zip(array_shape, self.chunk_shapes, strict=False)
-        ):
-            chunk_sum = sum(axis_chunks)
-            if chunk_sum != arr_size:
-                raise ValueError(
-                    f"Sum of chunk sizes along axis {axis} is {chunk_sum} "
-                    f"but array shape is {arr_size}"
-                )
+        # check array_shape is compatible with chunk grid
+        self._validate_array_shape(array_shape)
 
         # Generate coordinates
         # For each axis, we have len(axis_chunks) chunks
@@ -705,22 +744,8 @@ def get_nchunks(self, array_shape: tuple[int, ...]) -> int:
         ValueError
             If array_shape doesn't match chunk_shapes
         """
-        if len(array_shape) != len(self.chunk_shapes):
-            raise ValueError(
-                f"array_shape has {len(array_shape)} dimensions but "
-                f"chunk_shapes has {len(self.chunk_shapes)} dimensions"
-            )
-
-        # Validate that chunk sizes sum to array shape
-        for axis, (arr_size, axis_chunks) in enumerate(
-            zip(array_shape, self.chunk_shapes, strict=False)
-        ):
-            chunk_sum = sum(axis_chunks)
-            if chunk_sum != arr_size:
-                raise ValueError(
-                    f"Sum of chunk sizes along axis {axis} is {chunk_sum} "
-                    f"but array shape is {arr_size}"
-                )
+        # check array_shape is compatible with chunk grid
+        self._validate_array_shape(array_shape)
 
         # Total chunks is the product of number of chunks per axis
         return reduce(operator.mul, (len(axis_chunks) for axis_chunks in self.chunk_shapes), 1)
@@ -749,10 +774,11 @@ def _validate_array_shape(self, array_shape: tuple[int, ...]) -> None:
             zip(array_shape, self.chunk_shapes, strict=False)
         ):
             chunk_sum = sum(axis_chunks)
-            if chunk_sum != arr_size:
+            if chunk_sum < arr_size:
                 raise ValueError(
                     f"Sum of chunk sizes along axis {axis} is {chunk_sum} "
-                    f"but array shape is {arr_size}"
+                    f"but array shape is {arr_size}. This is invalid for the "
+                    "RectilinearChunkGrid."
                 )
 
     @cached_property
@@ -816,12 +842,6 @@ def get_chunk_start(
         """
         self._validate_array_shape(array_shape)
 
-        if len(chunk_coord) != len(self.chunk_shapes):
-            raise IndexError(
-                f"chunk_coord has {len(chunk_coord)} dimensions but "
-                f"chunk_shapes has {len(self.chunk_shapes)} dimensions"
-            )
-
         # Validate chunk coordinates are in bounds
         for axis, (coord, axis_chunks) in enumerate(
             zip(chunk_coord, self.chunk_shapes, strict=False)
@@ -869,12 +889,6 @@ def get_chunk_shape(
         """
         self._validate_array_shape(array_shape)
 
-        if len(chunk_coord) != len(self.chunk_shapes):
-            raise IndexError(
-                f"chunk_coord has {len(chunk_coord)} dimensions but "
-                f"chunk_shapes has {len(self.chunk_shapes)} dimensions"
-            )
-
         # Validate chunk coordinates are in bounds
         for axis, (coord, axis_chunks) in enumerate(
             zip(chunk_coord, self.chunk_shapes, strict=False)
@@ -994,12 +1008,6 @@ def array_index_to_chunk_coord(
         """
         self._validate_array_shape(array_shape)
 
-        if len(array_index) != len(array_shape):
-            raise IndexError(
-                f"array_index has {len(array_index)} dimensions but "
-                f"array_shape has {len(array_shape)} dimensions"
-            )
-
         # Validate array index is in bounds
         for axis, (idx, size) in enumerate(zip(array_index, array_shape, strict=False)):
             if not (0 <= idx < size):
@@ -1047,12 +1055,6 @@ def chunks_in_selection(
         """
         self._validate_array_shape(array_shape)
 
-        if len(selection) != len(array_shape):
-            raise ValueError(
-                f"selection has {len(selection)} dimensions but "
-                f"array_shape has {len(array_shape)} dimensions"
-            )
-
         # Normalize slices and find chunk ranges for each axis
         chunk_ranges = []
         for axis, (sel, size) in enumerate(zip(selection, array_shape, strict=False)):
@@ -1248,17 +1250,17 @@ def _normalize_rectilinear_chunks(
     # Validate that chunks sum to shape for each dimension
     for i, (dim_chunks, dim_size) in enumerate(zip(chunk_shapes, shape, strict=False)):
         chunk_sum = sum(dim_chunks)
-        if chunk_sum != dim_size:
+        if chunk_sum < dim_size:
             raise ValueError(
                 f"Variable chunks along dimension {i} sum to {chunk_sum} "
-                f"but array shape is {dim_size}. Chunks must sum exactly to shape."
+                f"but array shape is {dim_size}. Chunks must sum to be greater than or equal to the shape."
             )
 
     return chunk_shapes
 
 
 def parse_chunk_grid(
-    chunks: tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"] | int,
+    chunks: ChunksLike,
     *,
     shape: ShapeLike,
     item_size: int = 1,
@@ -1277,7 +1279,7 @@ def parse_chunk_grid(
 
     Parameters
     ----------
-    chunks : tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"] | int
+    chunks : ChunksLike
         The chunks parameter to parse. Can be:
         - A ChunkGrid instance
         - A nested sequence for variable-sized chunks (supports RLE format)
@@ -1523,7 +1525,7 @@ def _validate_data_compatibility(
 
 def resolve_chunk_spec(
     *,
-    chunks: tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"] | int,
+    chunks: ChunksLike,
     shards: ShardsLike | None,
     shape: tuple[int, ...],
     dtype_itemsize: int,
@@ -1542,7 +1544,7 @@ def resolve_chunk_spec(
 
     Parameters
     ----------
-    chunks : tuple[int, ...] | Sequence[Sequence[int]] | ChunkGrid | Literal["auto"] | int
+    chunks : ChunksLike
         The chunks specification from the user. Can be:
         - A ChunkGrid instance (Zarr v3 only)
         - A nested sequence for variable-sized chunks (Zarr v3 only)