feat(sharding): manage inner chunks as a semi-regular grid (sharding_indexed v1.1)

changes/251.feature.md

@@ -0,0 +1,11 @@
+The sharding codec no longer requires the inner `chunk_shape` to evenly divide
+the shard shape. Inner chunks are laid out on a semi-regular grid: they are
+spaced at `chunk_shape` intervals and clipped by the shard boundary, so the
+number of inner chunks per dimension is the ceiling division of the shard shape
+by `chunk_shape`, and clipped chunks are stored at their clipped shape. This
+implements version 1.1 of the `sharding_indexed` specification (see
+[zarr-specs #370](https://github.com/zarr-developers/zarr-specs/pull/370)) and
+works with regular and rectilinear chunk grids, including nested sharding.
+Note that arrays written with a non-divisible `chunk_shape` are not readable by
+implementations that only support version 1.0 of the sharding specification.
+Chunk counts retain constant-space arithmetic for uniform shard dimensions, including very large sparse arrays.

docs/user-guide/arrays.md

@@ -593,6 +593,14 @@ In this example a shard shape of (1000, 1000) and a chunk shape of (100, 100) is
 This means that `10*10` chunks are stored in each shard, and there are `10*10` shards in total.
 Without the `shards` argument, there would be 10,000 chunks stored as individual files.
 
+The chunk shape does not need to evenly divide the shard shape. Chunks are laid
+out on a regular grid within each shard and clipped by the shard boundary, so a
+shard shape of `(1000,)` with a chunk shape of `(300,)` stores four chunks per
+shard with sizes `300, 300, 300, 100`. Note that arrays that rely on this
+clipping (version 1.1 of the `sharding_indexed` specification) cannot be read by
+implementations that only support version 1.0, which requires the chunk shape
+to evenly divide the shard shape.
+
 ## Rectilinear (variable) chunk grids
 
 !!! warning "Experimental"
@@ -731,8 +739,9 @@ print("Roundtrip OK")
 
 Rectilinear chunk grids can also be used for shard boundaries when combined
 with sharding. In this case, the outer grid (shards) is rectilinear while the
-inner chunks remain regular. Each shard dimension must be divisible by the
-corresponding inner chunk size:
+inner chunks remain regular. The inner chunk size does not need to divide the
+shard sizes evenly: inner chunks are clipped by each shard's boundary, so every
+shard holds a semi-regular grid of inner chunks:
 
 ```python exec="true" session="arrays" source="above" result="ansi"
 z = zarr.create_array(

packages/zarr-metadata/src/zarr_metadata/v3/codec/sharding_indexed.py

@@ -27,8 +27,10 @@ class ShardingIndexedCodecConfiguration(TypedDict):
     """
     Configuration for the Zarr v3 `sharding_indexed` codec.
 
-    `chunk_shape` is the shape of inner chunks along each dimension;
-    it must evenly divide the shard shape.
+    `chunk_shape` is the shape of inner chunks along each dimension. It does
+    not need to evenly divide the shard shape: inner chunks are clipped by the
+    shard shape (`sharding_indexed` spec version 1.1; version 1.0 required
+    `chunk_shape` to evenly divide the shard shape).
 
     `codecs` is the codec pipeline applied to each inner chunk; exactly
     one array-to-bytes codec is required.

src/zarr/core/array.py

@@ -6,7 +6,6 @@
 from asyncio import gather
 from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass, field, replace
-from itertools import starmap
 from logging import getLogger
 from typing import (
     TYPE_CHECKING,
@@ -43,6 +42,7 @@
 from zarr.core.chunk_grids import (
     SHARDED_INNER_CHUNK_MAX_BYTES,
     ChunkGrid,
+    FixedDimension,
     _is_auto,
     _is_keep,
     _is_rectilinear_chunks,
@@ -200,6 +200,25 @@ def _chunk_sizes_from_shape(
     return tuple(result)
 
 
+def _inner_chunk_sizes_from_outer(
+    outer_chunk_sizes: tuple[tuple[int, ...], ...], inner_chunk_shape: tuple[int, ...]
+) -> tuple[tuple[int, ...], ...]:
+    """Compute dask-style inner (sub-shard) chunk sizes from outer chunk sizes.
+
+    The inner chunk grid restarts at every shard boundary and inner chunks are
+    clipped by the shard shape, so along each dimension the sizes are the
+    ceiling partition of each outer chunk's data size by the inner chunk size,
+    concatenated across the outer chunks.
+    """
+    result: list[tuple[int, ...]] = []
+    for outer_sizes, c in zip(outer_chunk_sizes, inner_chunk_shape, strict=True):
+        sizes: list[int] = []
+        for outer_size in outer_sizes:
+            sizes.extend(min(c, outer_size - i * c) for i in range(ceildiv(outer_size, c)))
+        result.append(tuple(sizes))
+    return tuple(result)
+
+
 def parse_array_metadata(data: Any) -> ArrayMetadata:
     if isinstance(data, ArrayMetadata):
         return data
@@ -900,7 +919,9 @@ def read_chunk_sizes(self) -> tuple[tuple[int, ...], ...]:
         ((20, 20),)
         """
         if (sharding_codec := _sharding_codec(self.metadata)) is not None:
-            return _chunk_sizes_from_shape(self.shape, sharding_codec.chunk_shape)
+            return _inner_chunk_sizes_from_outer(
+                self._chunk_grid.chunk_sizes, sharding_codec.chunk_shape
+            )
         return self._chunk_grid.chunk_sizes
 
     @property
@@ -1143,9 +1164,20 @@ def _chunk_grid_shape(self) -> tuple[int, ...]:
             The number of chunks along each dimension.
         """
         if (sharding_codec := _sharding_codec(self.metadata)) is not None:
-            # When sharding, count inner chunks across the whole array
+            # The inner grid restarts at every shard boundary.
             chunk_shape = sharding_codec.chunk_shape
-            return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True)))
+            counts: list[int] = []
+            for dimension, inner in zip(self._chunk_grid.dimensions, chunk_shape, strict=True):
+                if isinstance(dimension, FixedDimension):
+                    full, remainder = divmod(dimension.extent, dimension.size)
+                    counts.append(full * ceildiv(dimension.size, inner) + ceildiv(remainder, inner))
+                else:
+                    counts.append(
+                        sum(
+                            ceildiv(dimension.data_size(i), inner) for i in range(dimension.nchunks)
+                        )
+                    )
+            return tuple(counts)
         return self._chunk_grid.grid_shape
 
     @property
@@ -4117,7 +4149,7 @@ async def _shards_initialized(
 
 
 class ShardsConfigParam(TypedDict):
-    shape: tuple[int, ...]
+    shape: Sequence[int | Sequence[int]]
     index_location: IndexLocation | None
 
 
@@ -5467,7 +5499,10 @@ async def _nchunks_initialized(
         # Uniform shard shape: the exact chunk count per shard is a single
         # multiply, O(1) after the storage listing.
         chunks_per_shard = product(
-            tuple(a // b for a, b in zip(meta.chunk_grid.chunk_shape, inner_chunks, strict=True))
+            tuple(
+                ceildiv(a, b)
+                for a, b in zip(meta.chunk_grid.chunk_shape, inner_chunks, strict=True)
+            )
         )
         return (await _nshards_initialized(array)) * chunks_per_shard
     # Rectilinear shard grid: the chunk count varies per shard, so decode the
@@ -5478,7 +5513,7 @@ async def _nchunks_initialized(
         spec = grid[meta.chunk_key_encoding.decode_chunk_key(key)]
         if spec is not None:
             total += product(
-                tuple(s // c for s, c in zip(spec.codec_shape, inner_chunks, strict=True))
+                tuple(ceildiv(s, c) for s, c in zip(spec.codec_shape, inner_chunks, strict=True))
             )
     return total
 

src/zarr/core/chunk_grids.py

@@ -960,6 +960,7 @@ def resolve_outer_and_inner_chunks(
     # Extract the flat chunk shape (uniform size per dimension) for arithmetic.
     chunk_shape_flat = chunks.chunk_shape
 
+    shard_flat: tuple[int | Sequence[int], ...]
     if _is_auto(shard_shape):
         warnings.warn(
             "Automatic shard shape inference is experimental and may change without notice.",