fix(sharding): validate nested sharding codecs and reject an inner chunk size of 0

changes/4480.bugfix.md

@@ -0,0 +1 @@
+A sharding codec nested in the codecs of another sharding codec is now validated as the outer one is: its inner chunk shape must divide the chunk it encodes. Such an array used to be created, and opened, without error, and then read back wrong data or its fill value. `ShardingCodec` now rejects an inner chunk size of 0 (or `False`) when it is constructed, with a `ValueError` that names the dimension, as other Zarr implementations do; it used to accept one, which raised a `ZeroDivisionError` when an array was created or opened with the codec or, where the codec was nested in another sharding codec, on the first write to the array.

src/zarr/codecs/sharding.py

@@ -46,6 +46,7 @@
 )
 from zarr.core.common import (
     ShapeLike,
+    parse_chunk_shape,
     parse_named_configuration,
     parse_shapelike,
     product,
@@ -437,6 +438,12 @@ def _check_index_codecs_fixed_size(index_codecs: tuple[Codec, ...]) -> None:
         )
 
 
+def _parse_inner_chunk_shape(data: ShapeLike) -> tuple[int, ...]:
+    """Parse the inner chunk shape of a sharding codec: a shape (see `parse_shapelike`)
+    whose every size is a chunk edge length, so at least 1 (see `parse_chunk_shape`)."""
+    return parse_chunk_shape(parse_shapelike(data))
+
+
 @dataclass(frozen=True)
 class ShardingCodec(
     ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin, ArrayBytesCodecPartialEncodeMixin
@@ -468,7 +475,7 @@ def __init__(
         index_location: ShardingCodecIndexLocation | IndexLocation = "end",
         subchunk_write_order: SubchunkWriteOrder = "morton",
     ) -> None:
-        chunk_shape_parsed = parse_shapelike(chunk_shape)
+        chunk_shape_parsed = _parse_inner_chunk_shape(chunk_shape)
         codecs_parsed = parse_codecs(codecs)
         index_codecs_parsed = parse_codecs(index_codecs)
         _check_index_codecs_fixed_size(index_codecs_parsed)
@@ -507,7 +514,7 @@ def __getstate__(self) -> dict[str, Any]:
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         config = state["configuration"]
-        object.__setattr__(self, "chunk_shape", parse_shapelike(config["chunk_shape"]))
+        object.__setattr__(self, "chunk_shape", _parse_inner_chunk_shape(config["chunk_shape"]))
         object.__setattr__(self, "codecs", parse_codecs(config["codecs"]))
         object.__setattr__(self, "index_codecs", parse_codecs(config["index_codecs"]))
         object.__setattr__(self, "index_location", _parse_index_location(config["index_location"]))
@@ -635,6 +642,28 @@ def validate(
                         f"Chunk edge length {edge} in dimension {i} is not "
                         f"divisible by the shard's inner chunk size {inner}."
                     )
+        self._validate_inner_codecs(dtype)
+
+    def _validate_inner_codecs(self, dtype: ZDType[TBaseDType, TBaseScalar]) -> None:
+        """Validate the codecs that encode each inner chunk, as array metadata validates
+        its codecs: each against the chunk the codecs before it leave it, a regular
+        grid of one inner chunk. A sharding codec among them is thereby checked against
+        the chunk it splits, which `validate` of the outer codec alone does not see.
+        """
+        spec = ArraySpec(
+            shape=self.chunk_shape,
+            dtype=dtype,
+            fill_value=dtype.default_scalar(),
+            config=ArrayConfig.from_dict({}),
+            prototype=default_buffer_prototype(),
+        )
+        for codec in self.codecs:
+            codec.validate(
+                shape=spec.shape,
+                dtype=spec.dtype,
+                chunk_grid=RegularChunkGridMetadata(chunk_shape=spec.shape),
+            )
+            spec = codec.resolve_metadata(spec)
 
     def _get_inner_chunk_transform(self, shard_spec: ArraySpec) -> Any:
         """The synchronous transform for the inner codec chain.

tests/test_codecs/test_sharding.py

@@ -33,6 +33,7 @@
     _ShardReader,
 )
 from zarr.core.buffer import NDArrayLike, default_buffer_prototype
+from zarr.core.dtype import Int32
 from zarr.core.indexing import lexicographic_order_coords
 from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.storage import MemoryStore, StorePath, ZipStore
@@ -768,15 +769,13 @@ def test_structured_dtype_fill_value() -> None:
     [
         ShardingCodec(chunk_shape=(8, 8)),
         ShardingCodec(chunk_shape=(8, 8), subchunk_write_order="lexicographic"),
-        ShardingCodec(chunk_shape=(0,)),
     ],
-    ids=["default", "lexicographic", "zero-chunk-size"],
+    ids=["default", "lexicographic"],
 )
 def test_pickle(codec: ShardingCodec) -> None:
     """ShardingCodec round-trips through pickle, including the non-serialized
     ``subchunk_write_order`` (which ``to_dict`` omits and which must not silently
-    revert to the ``morton`` default), and an inner chunk size of 0, which the
-    constructor accepts."""
+    revert to the ``morton`` default)."""
     restored = pickle.loads(pickle.dumps(codec))
     assert restored == codec
     assert restored.chunk_shape == codec.chunk_shape
@@ -1436,3 +1435,90 @@ def test_sharding_truncated_shard_raises(
         }
         with pytest.raises(ValueError, match="too short to hold its index"):
             ops[op]()
+
+
+@pytest.mark.filterwarnings(
+    "ignore:Combining a `sharding_indexed` codec:zarr.errors.ZarrUserWarning"
+)
+@pytest.mark.parametrize(
+    "inner_codecs",
+    [
+        (ShardingCodec(chunk_shape=(3, 10)),),
+        (ShardingCodec(chunk_shape=(6, 20)),),
+        (TransposeCodec(order=(1, 0)), ShardingCodec(chunk_shape=(10, 3))),
+        (ShardingCodec(chunk_shape=(3, 10), codecs=(ShardingCodec(chunk_shape=(1, 5)),)),),
+    ],
+    ids=["divides", "equal", "after-transpose", "three-levels"],
+)
+def test_nested_sharding_roundtrip(inner_codecs: tuple[Any, ...]) -> None:
+    """A sharding codec nested in a sharding codec is accepted when its inner chunk
+    shape divides the chunk it encodes (the outer codec's inner chunk, as the codecs
+    before it leave it, so permuted after a transpose), and data round-trips."""
+    data = np.arange(1, 121, dtype="int32").reshape(6, 20)
+    store = MemoryStore()
+    arr = zarr.create_array(
+        store,
+        shape=data.shape,
+        dtype=data.dtype,
+        chunks=(6, 20),
+        serializer=ShardingCodec(chunk_shape=(6, 20), codecs=inner_codecs),
+    )
+    arr[:] = data
+    np.testing.assert_array_equal(zarr.open_array(store, mode="r")[:], data)
+
+
+@pytest.mark.filterwarnings(
+    "ignore:Combining a `sharding_indexed` codec:zarr.errors.ZarrUserWarning"
+)
+@pytest.mark.parametrize("nested_chunk_shape", [(4, 7), (12, 40), (3, 20, 1)])
+def test_nested_sharding_rejects_indivisible_inner_chunk_shape(
+    nested_chunk_shape: tuple[int, ...],
+) -> None:
+    """A sharding codec nested in a sharding codec is checked as the outer one is: an
+    inner chunk shape that does not divide the chunk it encodes, or has another number
+    of dimensions, is rejected when the array is created, and nothing is stored. Such
+    arrays used to be created and to read back wrong data."""
+    store = MemoryStore()
+    with pytest.raises(ValueError, match="not divisible|same number of dimensions"):
+        zarr.create_array(
+            store,
+            shape=(6, 20),
+            dtype="int32",
+            chunks=(6, 20),
+            serializer=ShardingCodec(
+                chunk_shape=(6, 20), codecs=(ShardingCodec(chunk_shape=nested_chunk_shape),)
+            ),
+        )
+    assert store._store_dict == {}
+
+
+@pytest.mark.parametrize("chunk_shape", [(0, 5), (5, 0), (False, 5), [0], 0])
+def test_sharding_codec_rejects_inner_chunk_size_zero(chunk_shape: Any) -> None:
+    """`ShardingCodec` rejects an inner chunk size of 0 (or `False`) when it is
+    constructed, with a `ValueError` naming the dimension, so no array metadata, nested
+    codec or stored document can hold one."""
+    with pytest.raises(ValueError, match=r"Dimension \d: chunk edge length must be >= 1, got"):
+        ShardingCodec(chunk_shape=chunk_shape)
+
+
+@pytest.mark.filterwarnings(
+    "ignore:Combining a `sharding_indexed` codec:zarr.errors.ZarrUserWarning"
+)
+def test_open_rejects_stored_nested_sharding_with_indivisible_inner_chunk_shape() -> None:
+    """A stored array document whose nested sharding codec has an inner chunk shape
+    that does not divide the chunk it encodes is rejected when the metadata is parsed,
+    not when chunks are read or written."""
+    valid = ArrayV3Metadata(
+        shape=(6, 20),
+        data_type=Int32(),
+        chunk_grid={"name": "regular", "configuration": {"chunk_shape": (6, 20)}},
+        chunk_key_encoding={"name": "default"},
+        fill_value=0,
+        codecs=(ShardingCodec(chunk_shape=(6, 20), codecs=(ShardingCodec(chunk_shape=(3, 10)),)),),
+        attributes={},
+        dimension_names=None,
+    )
+    doc: Any = valid.to_dict()
+    doc["codecs"][0]["configuration"]["codecs"][0]["configuration"]["chunk_shape"] = (4, 7)
+    with pytest.raises(ValueError, match="not divisible"):
+        ArrayV3Metadata.from_dict(doc)

tests/test_metadata/test_repair.py

@@ -519,24 +519,30 @@ def _sharding_chunk_shape(chunks: Any) -> tuple[tuple[int, ...], object]:
         ((np.int64(4),), (4,)),
         (np.array([4]), (4,)),
         ((True,), (1,)),
-        ((0,), (0,)),
-        ((False,), (0,)),
         (range(4, 5), (4,)),
     ],
 )
 def test_chunk_shape_read_as_array_shape(
     site: str, chunks: object, expected: tuple[int, ...]
 ) -> None:
     """`ArrayV2Metadata` and `ShardingCodec` read their chunk shape as `parse_shapelike`
-    reads an array shape: an integer or an iterable of non-negative integers, including
-    NumPy integers and bools. A chunk size of 0 is written back as given; reading a
-    stored 0 is `zarr.core.metadata.repair`' business."""
+    reads an array shape: an integer or an iterable of integers, including NumPy
+    integers and bools."""
     parsed, written = CHUNK_SHAPE_SITES[site](chunks)
     assert parsed == expected
     assert all(type(size) is int for size in parsed)
     assert written == expected
 
 
+@pytest.mark.parametrize("chunks", [(0,), (False,)])
+def test_v2_chunk_size_zero_written_as_given(chunks: tuple[Any, ...]) -> None:
+    """`ArrayV2Metadata` accepts a chunk size of 0 and writes it back as given; reading
+    a stored 0 is `zarr.core.metadata.repair`' business. `ShardingCodec` rejects one."""
+    metadata = _v2_metadata(chunks)
+    assert metadata.chunks == (0,)
+    assert metadata.to_dict()["chunks"] == (0,)
+
+
 @pytest.mark.parametrize("site", CHUNK_SHAPE_SITES)
 def test_chunk_shape_read_as_array_shape_rejects_negative(site: str) -> None:
     with pytest.raises(ValueError, match="Expected all values to be non-negative"):