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

changes/202.bugfix.md

@@ -0,0 +1 @@
+Stop emitting an `UnstableSpecificationWarning` when serializing the `struct` data type to Zarr V3 metadata. The `struct` data type now has a stable Zarr V3 specification. The legacy `structured` alias and the unspecified `null_terminated_bytes`, `raw_bytes`, and `variable_length_bytes` data types continue to warn.

changes/251.feature.md

@@ -0,0 +1,10 @@
+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.

docs/user-guide/arrays.md

@@ -594,6 +594,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"
@@ -728,8 +736,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

@@ -197,6 +197,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
@@ -883,8 +902,11 @@ def read_chunk_sizes(self) -> tuple[tuple[int, ...], ...]:
 
         codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ())
         if len(codecs) == 1 and isinstance(codecs[0], ShardingCodec):
+            # The inner chunk grid restarts at every shard boundary (inner
+            # chunks are clipped by the shard shape), so derive the sizes from
+            # the outer chunk sizes rather than the array shape alone.
             inner_chunk_shape = codecs[0].chunk_shape
-            return _chunk_sizes_from_shape(self.shape, inner_chunk_shape)
+            return _inner_chunk_sizes_from_outer(self._chunk_grid.chunk_sizes, inner_chunk_shape)
         return self._chunk_grid.chunk_sizes
 
     @property
@@ -1121,9 +1143,14 @@ def _chunk_grid_shape(self) -> tuple[int, ...]:
 
         codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ())
         if len(codecs) == 1 and isinstance(codecs[0], ShardingCodec):
-            # When sharding, count inner chunks across the whole array
+            # When sharding, count inner chunks across the whole array. The
+            # inner chunk grid restarts at every shard boundary (inner chunks
+            # are clipped by the shard shape), so count per outer chunk.
             chunk_shape = codecs[0].chunk_shape
-            return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True)))
+            return tuple(
+                sum(ceildiv(outer_size, c) for outer_size in outer_sizes)
+                for outer_sizes, c in zip(self._chunk_grid.chunk_sizes, chunk_shape, strict=True)
+            )
         return self._chunk_grid.grid_shape
 
     @property
@@ -5318,8 +5345,10 @@ async def _nchunks_initialized(
     if array.shards is None:
         chunks_per_shard = 1
     else:
+        # Inner chunks are clipped by the shard shape, so the count per shard
+        # is the ceiling division of the shard shape by the chunk shape.
         chunks_per_shard = product(
-            tuple(a // b for a, b in zip(array.shards, array.chunks, strict=True))
+            tuple(ceildiv(a, b) for a, b in zip(array.shards, array.chunks, strict=True))
         )
     return (await _nshards_initialized(array)) * chunks_per_shard
 

tests/test_codecs/test_sharding.py

@@ -32,6 +32,7 @@
 from zarr.core.buffer import NDArrayLike, default_buffer_prototype
 from zarr.core.indexing import lexicographic_order_coords
 from zarr.core.metadata.v3 import ArrayV3Metadata
+from zarr.core.sync import sync
 from zarr.storage import MemoryStore, StorePath, ZipStore
 
 from ..conftest import ArrayRequest
@@ -831,34 +832,189 @@ def test_invalid_metadata(store: Store) -> None:
             dtype=np.dtype("uint8"),
             fill_value=0,
         )
-    spath2 = StorePath(store, "invalid_inner_chunk_shape")
-    with pytest.raises(ValueError):
-        zarr.create_array(
-            spath2,
-            shape=(16, 16),
-            shards=(16, 16),
-            chunks=(8, 7),
-            dtype=np.dtype("uint8"),
-            fill_value=0,
+
+
+@pytest.mark.parametrize(
+    ("array_shape", "shard_shape", "chunk_shape"),
+    [
+        ((24,), (12,), (5,)),  # trailing inner chunk clipped to 2 elements
+        ((10,), (4,), (6,)),  # inner chunk larger than the shard: one clipped chunk per shard
+        ((23, 17), (12, 12), (5, 5)),  # array shape not divisible by the shard shape either
+        ((16, 16), (16, 16), (9, 9)),  # single shard, clipped in both dimensions
+        ((16, 16), (8, 16), (4, 4)),  # evenly divisible control case
+    ],
+)
+@pytest.mark.parametrize("index_location", ["start", "end"])
+def test_sharding_clipped_inner_chunks(
+    array_shape: tuple[int, ...],
+    shard_shape: tuple[int, ...],
+    chunk_shape: tuple[int, ...],
+    index_location: IndexLocation,
+) -> None:
+    """`chunk_shape` does not need to evenly divide the shard shape
+    (`sharding_indexed` v1.1): inner chunks are clipped by the shard boundary.
+
+    Full and partial reads and writes round-trip for reasonable combinations of
+    array, shard, and inner chunk shape.
+    """
+    data = np.arange(np.prod(array_shape), dtype="uint16").reshape(array_shape)
+    arr = zarr.create_array(
+        {},
+        shape=array_shape,
+        shards={"shape": shard_shape, "index_location": index_location},
+        chunks=chunk_shape,
+        dtype=data.dtype,
+        fill_value=0,
+    )
+    arr[...] = data
+    assert np.array_equal(arr[...], data)
+
+    # partial read crossing clipped inner chunk boundaries
+    sel = tuple(slice(1, s - 1) for s in array_shape)
+    assert np.array_equal(arr[sel], data[sel])
+
+    # partial write crossing clipped inner chunk boundaries
+    arr[sel] = 0
+    expected = data.copy()
+    expected[sel] = 0
+    assert np.array_equal(arr[...], expected)
+
+
+def test_sharding_clipped_inner_chunks_byte_layout() -> None:
+    """Spec example: a shard of shape [12] with inner chunk shape [5] holds 3
+    inner chunks stored at their clipped shapes ([5], [5], and [2]).
+
+    With an uncompressed inner codec chain the shard blob holds exactly 12 data
+    bytes (not 15) plus a 3-entry shard index.
+    """
+    store = MemoryStore()
+    arr = zarr.create_array(
+        store,
+        shape=(12,),
+        shards=(12,),
+        chunks=(5,),
+        dtype="uint8",
+        fill_value=255,
+        compressors=None,
+    )
+    arr[...] = np.arange(12, dtype="uint8")
+
+    shard_buffer = sync(store.get("c/0", prototype=default_buffer_prototype()))
+    assert shard_buffer is not None
+    raw = shard_buffer.to_bytes()
+    index_size = 3 * 16 + 4  # ceil(12/5) = 3 index entries plus the crc32c checksum
+    assert len(raw) == 12 + index_size
+    offsets_and_lengths = np.frombuffer(raw[-index_size:-4], dtype="<u8").reshape(3, 2)
+    assert sorted(int(length) for length in offsets_and_lengths[:, 1]) == [2, 5, 5]
+
+
+def test_sharding_clipped_inner_chunks_rectilinear() -> None:
+    """A rectilinear outer grid whose shard edges are not divisible by the inner
+    chunk shape round-trips: the inner grid is clipped per shard.
+
+    Shard edges `[4, 6]` cannot be normalized to a regular grid (the larger edge
+    comes last), so this exercises `RectilinearChunkGridMetadata` for real —
+    unlike e.g. `[6, 4]`, which coincides with a clipped regular grid of size 6.
+    """
+    from zarr.core.metadata.v3 import RectilinearChunkGridMetadata
+
+    with zarr.config.set({"array.rectilinear_chunks": True}):
+        arr = zarr.create_array(
+            {}, shape=(10,), shards=[[4, 6]], chunks=(4,), dtype="uint8", fill_value=0
         )
+    assert isinstance(arr.metadata, ArrayV3Metadata)
+    assert isinstance(arr.metadata.chunk_grid, RectilinearChunkGridMetadata)
+    data = np.arange(10, dtype="uint8")
+    arr[...] = data
+    assert np.array_equal(arr[...], data)
+    assert arr.read_chunk_sizes == ((4, 4, 2),)
+
+    arr[3:9] = 42
+    data[3:9] = 42
+    assert np.array_equal(arr[...], data)
+
+
+def test_sharding_clipped_inner_chunks_nested() -> None:
+    """Clipping composes through nested sharding: 12-element shards split into
+    clipped 5-element inner chunks, each split into clipped 2-element chunks."""
+    inner = ShardingCodec(chunk_shape=(2,), codecs=(BytesCodec(),))
+    outer = ShardingCodec(chunk_shape=(5,), codecs=(inner,))
+    arr = zarr.create_array(
+        {},
+        shape=(24,),
+        chunks=(12,),
+        dtype="uint8",
+        fill_value=0,
+        serializer=outer,
+        compressors=None,
+        filters=None,
+    )
+    data = np.arange(24, dtype="uint8")
+    arr[...] = data
+    assert np.array_equal(arr[...], data)
 
+    arr[4:11] = 9
+    data[4:11] = 9
+    assert np.array_equal(arr[...], data)
 
-def test_invalid_shard_shape() -> None:
-    with pytest.raises(
-        ValueError,
-        match=(
-            f"Chunk edge length {16} in dimension {0} is not "
-            f"divisible by the shard's inner chunk size {9}\\."
-        ),
-    ):
-        zarr.create_array(
-            {},
-            shape=(16, 16),
-            shards=(16, 16),
-            chunks=(9, 9),
-            dtype=np.dtype("uint8"),
-            fill_value=0,
+
+def test_sharding_clipped_inner_chunk_sizes() -> None:
+    """Inner chunk grids restart at each shard boundary, so the reported inner
+    chunk sizes are clipped per shard (and by the array extent in the last
+    shard)."""
+    arr = zarr.create_array({}, shape=(23,), shards=(12,), chunks=(5,), dtype="uint8", fill_value=0)
+    assert arr.read_chunk_sizes == ((5, 5, 2, 5, 5, 1),)
+    assert arr.cdata_shape == (6,)
+    assert arr.nchunks == 6
+
+
+def test_sharding_clipped_nchunks_initialized() -> None:
+    """`nchunks_initialized` counts ceiling-division inner chunks per shard.
+
+    Regression test: it previously used floor division, which undercounts for
+    clipped grids and reports 0 when the inner chunk shape exceeds the shard
+    shape.
+    """
+    arr = zarr.create_array({}, shape=(16,), shards=(8,), chunks=(3,), dtype="uint8", fill_value=0)
+    arr[...] = np.arange(16, dtype="uint8")
+    # 2 shards, each holding ceil(8/3) = 3 inner chunks
+    assert arr.nchunks_initialized == 6
+
+    arr2 = zarr.create_array({}, shape=(11,), shards=(4,), chunks=(9,), dtype="uint8", fill_value=0)
+    arr2[...] = np.arange(11, dtype="uint8")
+    # 3 shards, each holding one inner chunk clipped to the shard shape
+    assert arr2.nchunks_initialized == 3
+
+
+def test_sharding_scalar_write_memoizes_complete_chunks(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A scalar broadcast partial write encodes the repeated complete inner
+    chunk once, not once per chunk.
+
+    Regression test: the memo gate must compare chunk specs by shape — an
+    object-identity gate never fires because `_get_chunk_spec` is uncached.
+    """
+    from zarr.core import chunk_utils
+
+    calls = 0
+    real_encode = chunk_utils.ChunkTransform.encode_chunk
+
+    def counting_encode(self: Any, *args: Any, **kwargs: Any) -> Any:
+        nonlocal calls
+        calls += 1
+        return real_encode(self, *args, **kwargs)
+
+    with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}):
+        arr = zarr.create_array(
+            {}, shape=(100,), shards=(100,), chunks=(10,), dtype="uint8", fill_value=0
         )
+        monkeypatch.setattr(chunk_utils.ChunkTransform, "encode_chunk", counting_encode)
+        arr[5:] = 7
+        expected = np.full(100, 7, dtype="uint8")
+        expected[:5] = 0
+        assert np.array_equal(arr[...], expected)
+    # one merge for the partial edge chunk, one memoized encode shared by the
+    # nine complete chunks, one shard-index encode
+    assert calls <= 4, f"scalar write encoded {calls} times; memo is not firing"
 
 
 @pytest.mark.parametrize("store", ["local"], indirect=["store"])