Skip to content

Commit baddb31

Browse files
d-v-bclaude
andcommitted
fix(sharding): parse the inner chunk shape with the chunk normalizer as a regular grid
ShardingCodec read its inner chunk_shape with parse_shapelike, which allows 0, so an inner size of 0 was accepted at construction and failed later with ZeroDivisionError. parse_regular_chunk_shape shares the normalizer's integer test (_chunk_int/_chunk_list) and requires every size to be at least 1; with no axis length to resolve against, -1 and False are rejected, and an explicit edge list (a rectilinear dimension) is rejected because the inner grid of a shard is regular. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent ad34ae8 commit baddb31

5 files changed

Lines changed: 92 additions & 21 deletions

File tree

‎changes/4376.bugfix.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
Chunk specifications are judged only by the chunk normalizer: the separate duck-typed check for rectilinear input is gone, so the Zarr format 2 and sharding restrictions apply to what the normalized grid says was declared. A chunk size is an integer by Python's integer protocol: numpy integer scalars and 0-d integer arrays are accepted wherever a chunk or shard shape is. The legacy `zarr.create(..., zarr_format=2)` no longer tests the truth value of `chunks`: only `None` means not given, as for Zarr format 3. `chunks=False` is now one chunk instead of an automatic chunk shape, numpy arrays work, and `0`, `[]` and `np.array([0])` raise instead of chunking automatically. A non-integer scalar chunk specification raises a `TypeError` naming the problem. An explicit list of chunk edges declares a rectilinear dimension everywhere, including a stored rectilinear chunk grid passed as `chunks=` whose edges happen to be uniform, so the Zarr format 2 and sharding restrictions treat it the same whether it is given as lists or as metadata, and `zarr.create` stores it as a rectilinear chunk grid, as `create_array` does. A `RectilinearChunkGridMetadata` made only of bare integers declares no explicit edge list, so it is read as the regular grid it describes: it is accepted for Zarr format 2 and as the chunk shape of a sharded array. The `chunks` and `chunk_shape` parameters of `zarr.create` and the `ShardsLike` type now declare every chunk specification the normalizer accepts, including numpy integers and arrays, explicit edge lists and stored chunk grid metadata.
1+
Chunk specifications are judged only by the chunk normalizer: the separate duck-typed check for rectilinear input is gone, so the Zarr format 2 and sharding restrictions apply to what the normalized grid says was declared. A chunk size is an integer by Python's integer protocol: numpy integer scalars and 0-d integer arrays are accepted wherever a chunk or shard shape is. The legacy `zarr.create(..., zarr_format=2)` no longer tests the truth value of `chunks`: only `None` means not given, as for Zarr format 3. `chunks=False` is now one chunk instead of an automatic chunk shape, numpy arrays work, and `0`, `[]` and `np.array([0])` raise instead of chunking automatically. A non-integer scalar chunk specification raises a `TypeError` naming the problem. An explicit list of chunk edges declares a rectilinear dimension everywhere, including a stored rectilinear chunk grid passed as `chunks=` whose edges happen to be uniform, so the Zarr format 2 and sharding restrictions treat it the same whether it is given as lists or as metadata, and `zarr.create` stores it as a rectilinear chunk grid, as `create_array` does. A `RectilinearChunkGridMetadata` made only of bare integers declares no explicit edge list, so it is read as the regular grid it describes: it is accepted for Zarr format 2 and as the chunk shape of a sharded array. The `chunks` and `chunk_shape` parameters of `zarr.create` and the `ShardsLike` type now declare every chunk specification the normalizer accepts, including numpy integers and arrays, explicit edge lists and stored chunk grid metadata. The inner chunk shape of `ShardingCodec` is parsed with the chunk normalizer's integer test as a regular chunk shape: every size must be at least 1, so a stored or given inner chunk size of 0 raises a `ValueError` instead of a `ZeroDivisionError` on read, and an explicit list of chunk edges is rejected with a `TypeError`.

‎src/zarr/codecs/sharding.py‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@
3636
default_buffer_prototype,
3737
numpy_buffer_prototype,
3838
)
39-
from zarr.core.chunk_grids import ChunkGrid
39+
from zarr.core.chunk_grids import ChunkGrid, parse_regular_chunk_shape
4040
from zarr.core.chunk_utils import (
4141
ChunkTransform,
4242
decode_and_scatter_chunk,
@@ -47,7 +47,6 @@
4747
from zarr.core.common import (
4848
ShapeLike,
4949
parse_named_configuration,
50-
parse_shapelike,
5150
product,
5251
)
5352
from zarr.core.config import config as zarr_config
@@ -468,7 +467,7 @@ def __init__(
468467
index_location: ShardingCodecIndexLocation | IndexLocation = "end",
469468
subchunk_write_order: SubchunkWriteOrder = "morton",
470469
) -> None:
471-
chunk_shape_parsed = parse_shapelike(chunk_shape)
470+
chunk_shape_parsed = parse_regular_chunk_shape(chunk_shape)
472471
codecs_parsed = parse_codecs(codecs)
473472
index_codecs_parsed = parse_codecs(index_codecs)
474473
_check_index_codecs_fixed_size(index_codecs_parsed)
@@ -507,7 +506,7 @@ def __getstate__(self) -> dict[str, Any]:
507506

508507
def __setstate__(self, state: dict[str, Any]) -> None:
509508
config = state["configuration"]
510-
object.__setattr__(self, "chunk_shape", parse_shapelike(config["chunk_shape"]))
509+
object.__setattr__(self, "chunk_shape", parse_regular_chunk_shape(config["chunk_shape"]))
511510
object.__setattr__(self, "codecs", parse_codecs(config["codecs"]))
512511
object.__setattr__(self, "index_codecs", parse_codecs(config["index_codecs"]))
513512
object.__setattr__(self, "index_location", _parse_index_location(config["index_location"]))

‎src/zarr/core/chunk_grids.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -851,6 +851,35 @@ def normalize_chunks_nd(
851851
)
852852

853853

854+
def parse_regular_chunk_shape(chunk_shape: ShapeLike) -> tuple[int, ...]:
855+
"""Parse the chunk shape of a regular chunk grid that has no array shape to be
856+
normalized against, such as the inner chunk shape of a shard.
857+
858+
The chunk normalizer's integer test (`_chunk_int`) decides what a chunk size is, and
859+
every chunk size is at least 1. A single integer is a one-dimensional chunk shape.
860+
Without the lengths of the axes `-1` and `False` cannot be resolved, so they are
861+
rejected as non-positive; an explicit list of chunk edges declares a rectilinear
862+
dimension, which a regular chunk grid cannot have, and is rejected as a non-integer.
863+
"""
864+
size = _chunk_int(chunk_shape)
865+
elements = [size] if size is not None else _chunk_list(chunk_shape)
866+
as_ints = [_chunk_int(e) for e in elements]
867+
non_int = [
868+
(idx, e) for idx, (e, i) in enumerate(zip(elements, as_ints, strict=True)) if i is None
869+
]
870+
if non_int:
871+
non_int_idxs, non_int_vals = [*zip(*non_int, strict=False)]
872+
raise TypeError(
873+
f"Each chunk size of a regular chunk shape must be an integer; got non-integer "
874+
f"element(s) {non_int_vals!r} at indices {non_int_idxs!r}. A regular chunk shape "
875+
f"is a flat sequence of positive integers, one per dimension (e.g. [10, 10])."
876+
)
877+
sizes = tuple(i for i in as_ints if i is not None)
878+
if any(c < 1 for c in sizes):
879+
raise ValueError(f"All chunk sizes must be positive, got {list(sizes)}")
880+
return sizes
881+
882+
854883
def guess_chunks(
855884
shape: tuple[int, ...], typesize: int, *, max_bytes: int | None = None
856885
) -> ChunkGrid:

‎tests/test_codecs/test_sharding.py‎

Lines changed: 54 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import math
33
import pickle
44
import warnings
5-
from typing import Any, cast, get_args
5+
from typing import TYPE_CHECKING, Any, cast, get_args
66
from unittest.mock import AsyncMock
77

88
import numpy as np
@@ -40,6 +40,9 @@
4040
from ..conftest import ArrayRequest
4141
from .test_codecs import _AsyncArrayProxy, order_from_dim
4242

43+
if TYPE_CHECKING:
44+
from zarr.core.common import JSON
45+
4346

4447
def _reads_are_sync(store_mock: AsyncMock) -> bool:
4548
"""True when the partial-shard read for this store+pipeline goes through the
@@ -765,15 +768,13 @@ def test_structured_dtype_fill_value() -> None:
765768
[
766769
ShardingCodec(chunk_shape=(8, 8)),
767770
ShardingCodec(chunk_shape=(8, 8), subchunk_write_order="lexicographic"),
768-
ShardingCodec(chunk_shape=(0,)),
769771
],
770-
ids=["default", "lexicographic", "zero-chunk-size"],
772+
ids=["default", "lexicographic"],
771773
)
772774
def test_pickle(codec: ShardingCodec) -> None:
773775
"""ShardingCodec round-trips through pickle, including the non-serialized
774776
``subchunk_write_order`` (which ``to_dict`` omits and which must not silently
775-
revert to the ``morton`` default), and an inner chunk size of 0, which the
776-
constructor accepts."""
777+
revert to the ``morton`` default)."""
777778
restored = pickle.loads(pickle.dumps(codec))
778779
assert restored == codec
779780
assert restored.chunk_shape == codec.chunk_shape
@@ -1395,3 +1396,51 @@ def test_sharding_orthogonal_set_multiple_array_dims(
13951396
expected[ix] = value.reshape(expected[ix].shape)
13961397
assert np.array_equal(a[:], expected)
13971398
assert np.array_equal(a.oindex[selection], value)
1399+
1400+
1401+
@pytest.mark.parametrize(
1402+
("chunk_shape", "expected"),
1403+
[
1404+
((4,), (4,)),
1405+
([4, 2], (4, 2)),
1406+
(4, (4,)),
1407+
(np.int64(4), (4,)),
1408+
((np.int64(4),), (4,)),
1409+
(np.array([4]), (4,)),
1410+
((True,), (1,)),
1411+
(range(4, 5), (4,)),
1412+
((), ()),
1413+
],
1414+
ids=repr,
1415+
)
1416+
def test_sharding_codec_chunk_shape(chunk_shape: Any, expected: tuple[int, ...]) -> None:
1417+
"""The inner chunk shape of a shard is parsed by the chunk normalizer's integer test as
1418+
a regular chunk shape: one integer of at least 1 per dimension, written back as
1419+
given."""
1420+
codec = ShardingCodec(chunk_shape=chunk_shape)
1421+
assert codec.chunk_shape == expected
1422+
assert all(type(size) is int for size in codec.chunk_shape)
1423+
configuration = cast("dict[str, JSON]", codec.to_dict()["configuration"])
1424+
assert configuration["chunk_shape"] == expected
1425+
assert ShardingCodec.from_dict(codec.to_dict()).chunk_shape == expected
1426+
1427+
1428+
@pytest.mark.parametrize("chunk_shape", [(0,), (False,), (-1,), 0, [8, 0]], ids=repr)
1429+
def test_sharding_codec_chunk_shape_rejects_non_positive(chunk_shape: Any) -> None:
1430+
"""An inner chunk size below 1 is rejected, including `-1` and `False`, which have no
1431+
axis length to be resolved against."""
1432+
with pytest.raises(ValueError, match="All chunk sizes must be positive"):
1433+
ShardingCodec(chunk_shape=chunk_shape)
1434+
1435+
1436+
@pytest.mark.parametrize("chunk_shape", [[[4, 4]], [8, [4, 4]], (4.0,), "4"], ids=repr)
1437+
def test_sharding_codec_chunk_shape_rejects_non_integer(chunk_shape: Any) -> None:
1438+
"""An inner chunk shape is regular: an explicit list of chunk edges, which declares a
1439+
rectilinear dimension, is rejected like any other non-integer chunk size."""
1440+
with pytest.raises(TypeError, match="Each chunk size of a regular chunk shape must be"):
1441+
ShardingCodec(chunk_shape=chunk_shape)
1442+
1443+
1444+
def test_sharding_codec_chunk_shape_rejects_non_iterable() -> None:
1445+
with pytest.raises(TypeError, match="must be an integer or an iterable of integers"):
1446+
ShardingCodec(chunk_shape=None) # type: ignore[arg-type]

‎tests/test_metadata/test_repair.py‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ def test_read_invalid_edges_in_rectilinear_grid(
331331
(_v3_doc([4], [4.5]), "Dimension 0: chunk edge length must be an int, got 4.5"),
332332
(
333333
_v3_doc([8], [4], inner=[2.0]),
334-
"Expected an iterable of integers. Got [2.0] instead.",
334+
"got non-integer element(s) (2.0,) at indices (0,)",
335335
),
336336
(_v2_doc([20], [10.0]), "Expected an iterable of integers. Got [10.0] instead."),
337337
(
@@ -494,18 +494,12 @@ def test_regular_chunk_grid_rejects_chunk_shape_not_a_sequence(chunk_shape: Any)
494494
RegularChunkGridMetadata(chunk_shape=chunk_shape)
495495

496496

497-
def _sharding_chunk_shape(chunks: Any) -> tuple[tuple[int, ...], object]:
498-
codec = ShardingCodec(chunk_shape=chunks)
499-
configuration = cast("dict[str, JSON]", codec.to_dict()["configuration"])
500-
return codec.chunk_shape, configuration["chunk_shape"]
501-
502-
503497
CHUNK_SHAPE_SITES: dict[str, Callable[[Any], tuple[tuple[int, ...], object]]] = {
504498
"v2": lambda chunks: ((md := _v2_metadata(chunks)).chunks, md.to_dict()["chunks"]),
505-
"sharding-inner": _sharding_chunk_shape,
506499
}
507-
"""`ArrayV2Metadata` and `ShardingCodec` read a chunk shape as an array shape, returning
508-
the chunk shape and the value `to_dict` writes for it."""
500+
"""`ArrayV2Metadata` reads a chunk shape as an array shape, returning the chunk shape
501+
and the value `to_dict` writes for it. The inner chunk shape of `ShardingCodec` is a
502+
regular chunk shape instead (see `test_sharding_codec_chunk_shape`)."""
509503

510504

511505
@pytest.mark.parametrize("site", CHUNK_SHAPE_SITES)
@@ -527,7 +521,7 @@ def _sharding_chunk_shape(chunks: Any) -> tuple[tuple[int, ...], object]:
527521
def test_chunk_shape_read_as_array_shape(
528522
site: str, chunks: object, expected: tuple[int, ...]
529523
) -> None:
530-
"""`ArrayV2Metadata` and `ShardingCodec` read their chunk shape as `parse_shapelike`
524+
"""`ArrayV2Metadata` reads its chunk shape as `parse_shapelike`
531525
reads an array shape: an integer or an iterable of non-negative integers, including
532526
NumPy integers and bools. A chunk size of 0 is written back as given; reading a
533527
stored 0 is `zarr.core.metadata.repair`' business."""

0 commit comments

Comments
 (0)