Skip to content

Commit e118ab8

Browse files
committed
fix: make the chunk normalizer the one judge of a chunk specification
Chunk specifications were parsed by `normalize_chunks_nd`, but a separate duck-typed classifier, `_is_rectilinear_chunks`, ran on the raw input first at three sites to decide whether the spec was rectilinear. Two opinions on the same input is the shape of the bug in #4374, and they could disagree: a 0-d numpy array counted as rectilinear because it has `__iter__`. The classifier is gone. Each site normalizes first and asks the resulting `ChunkGrid` (`is_regular`); stored rectilinear metadata passed as `chunks=` counts as rectilinear even when its edges are uniform. The shard resolver sends regular and rectilinear shard specs through the same normalizer. Also fixed on the way: - The legacy v2 branch of `AsyncArray.create` tested `chunks or chunk_shape`, so `zarr.create(chunks=np.array([...]), zarr_format=2)` failed with "truth value of an array is ambiguous". It now uses the `is not None` form the v3 branch already had. - 0-d numpy arrays unwrap to their scalar in both normalizers instead of failing with "len() of unsized object". - A non-integer scalar spec (`2.0`, `np.float64`) raises the normalizer's own TypeError instead of "object has no len()". Assisted-by: ClaudeCode:claude-fable-5-1
1 parent ba883a5 commit e118ab8

3 files changed

Lines changed: 100 additions & 104 deletions

File tree

‎src/zarr/core/array.py‎

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,6 @@
4545
ChunkGrid,
4646
_is_auto,
4747
_is_keep,
48-
_is_rectilinear_chunks,
4948
guess_chunks,
5049
normalize_chunks_nd,
5150
resolve_outer_and_inner_chunks,
@@ -514,17 +513,15 @@ async def _create(
514513
)
515514
if dimension_names is not None:
516515
raise ValueError("dimension_names cannot be used for arrays with zarr_format 2.")
517-
if _is_rectilinear_chunks(_raw_chunks):
518-
raise ValueError("Zarr format 2 does not support rectilinear chunk grids.")
519-
520516
item_size = 1
521517
if isinstance(dtype_parsed, HasItemSize):
522518
item_size = dtype_parsed.item_size
523-
_raw = chunks or chunk_shape
524-
if _raw is None:
519+
if _raw_chunks is None:
525520
outer_chunks = guess_chunks(shape, item_size)
526521
else:
527-
outer_chunks = normalize_chunks_nd(_raw, shape)
522+
outer_chunks = normalize_chunks_nd(_raw_chunks, shape)
523+
if not outer_chunks.is_regular:
524+
raise ValueError("Zarr format 2 does not support rectilinear chunk grids.")
528525
_chunks = outer_chunks.chunk_shape
529526

530527
if order is None:
@@ -4530,8 +4527,18 @@ async def init_array(
45304527

45314528
await _prepare_overwrite(store_path, zarr_format=zarr_format, overwrite=overwrite)
45324529

4533-
# Validate rectilinear chunks constraints
4534-
if _is_rectilinear_chunks(chunks):
4530+
# Normalize the user's chunks into a canonical ChunkGrid
4531+
4532+
if _is_auto(chunks):
4533+
max_bytes = None if shards is None else SHARDED_INNER_CHUNK_MAX_BYTES
4534+
chunks_normalized = guess_chunks(shape_parsed, item_size, max_bytes=max_bytes)
4535+
else:
4536+
chunks_normalized = normalize_chunks_nd(chunks, shape_parsed)
4537+
4538+
# Validate rectilinear chunks constraints. The normalized grid is the one
4539+
# judge of what the user declared; stored rectilinear metadata counts as
4540+
# rectilinear even when its edges happen to be uniform.
4541+
if isinstance(chunks, RectilinearChunkGridMetadata) or not chunks_normalized.is_regular:
45354542
if zarr_format == 2:
45364543
raise ValueError("Zarr format 2 does not support rectilinear chunk grids.")
45374544
if shards is not None:
@@ -4541,14 +4548,6 @@ async def init_array(
45414548
"chunks=(inner_size, ...), shards=[[shard_sizes], ...]"
45424549
)
45434550

4544-
# Normalize the user's chunks into a canonical ChunkGrid
4545-
4546-
if _is_auto(chunks):
4547-
max_bytes = None if shards is None else SHARDED_INNER_CHUNK_MAX_BYTES
4548-
chunks_normalized = guess_chunks(shape_parsed, item_size, max_bytes=max_bytes)
4549-
else:
4550-
chunks_normalized = normalize_chunks_nd(chunks, shape_parsed)
4551-
45524551
# Resolve chunks + shards into outer_chunks (grid metadata) and
45534552
# inner (sub-chunk structure for ShardingCodec, None if no sharding)
45544553
outer_chunks, inner = resolve_outer_and_inner_chunks(

‎src/zarr/core/chunk_grids.py‎

Lines changed: 17 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@
1414
Literal,
1515
NamedTuple,
1616
Protocol,
17-
cast,
1817
runtime_checkable,
1918
)
2019

@@ -334,30 +333,6 @@ def _is_keep(spec: object) -> TypeIs[Literal["keep"]]:
334333
return isinstance(spec, str) and spec == "keep"
335334

336335

337-
def _is_rectilinear_chunks(chunks: Any) -> bool:
338-
"""Check if chunks specifies a rectilinear grid along any dimension.
339-
340-
Returns True for nested sequences like [[10, 20], [5, 5]], for mixed
341-
per-dimension specs like (5, [10, 20]) regardless of which dimension
342-
carries the sequence, and for stored rectilinear grid metadata.
343-
Returns False for flat sequences like (10, 10) or [10, 10].
344-
"""
345-
from zarr.core.metadata.v3 import RectilinearChunkGridMetadata
346-
347-
if isinstance(chunks, RectilinearChunkGridMetadata):
348-
return True
349-
if isinstance(chunks, (str, int, ChunkGrid)):
350-
return False
351-
if not hasattr(chunks, "__iter__"):
352-
return False
353-
try:
354-
return any(
355-
hasattr(elem, "__iter__") and not isinstance(elem, (str, bytes, int)) for elem in chunks
356-
)
357-
except TypeError:
358-
return False
359-
360-
361336
def is_regular_1d(dim_chunks: Sequence[int]) -> bool:
362337
"""Check if a single dimension's chunk sizes represent a regular grid.
363338
@@ -756,6 +731,8 @@ def normalize_chunks_1d(chunks: int | Iterable[object], span: int) -> DimensionG
756731
# `int` subclasses) take the uniform-chunk path instead of being treated as a sequence.
757732
# The `-1` sentinel check lives inside this branch so that numpy-array chunk
758733
# specifications never hit an ambiguous-truth-value error on `chunks == -1`.
734+
if isinstance(chunks, np.ndarray) and chunks.ndim == 0:
735+
chunks = chunks[()]
759736
if isinstance(chunks, numbers.Integral):
760737
chunk_size = int(chunks)
761738
if chunk_size < -1 or chunk_size == 0:
@@ -834,9 +811,16 @@ def normalize_chunks_nd(
834811
chunks = -1
835812

836813
# handle 1D convenience form. bool is excluded above so this only catches actual ints.
814+
if isinstance(chunks, np.ndarray) and chunks.ndim == 0:
815+
chunks = chunks[()]
837816
if isinstance(chunks, numbers.Integral):
838817
chunks = tuple(int(chunks) for _ in shape)
839818

819+
if not hasattr(chunks, "__len__"):
820+
raise TypeError(
821+
f"Chunk specification must be an integer or an iterable of integers; got "
822+
f"{chunks!r} of type {type(chunks).__name__}."
823+
)
840824
# handle bad dimensionality
841825
if len(chunks) != len(shape):
842826
raise ValueError(
@@ -952,15 +936,10 @@ def resolve_outer_and_inner_chunks(
952936
if shard_shape is None:
953937
return ChunkLayout(outer_chunks=chunks)
954938

955-
# Rectilinear shards: normalize the nested sequence directly.
956-
if _is_rectilinear_chunks(shard_shape):
957-
outer = normalize_chunks_nd(shard_shape, array_shape)
958-
return ChunkLayout(outer_chunks=outer, inner=ChunkLayout(outer_chunks=chunks))
959-
960-
# Extract the flat chunk shape (uniform size per dimension) for arithmetic.
961-
chunk_shape_flat = chunks.chunk_shape
962-
939+
shard_spec: Any
963940
if _is_auto(shard_shape):
941+
# Extract the flat chunk shape (uniform size per dimension) for arithmetic.
942+
chunk_shape_flat = chunks.chunk_shape
964943
warnings.warn(
965944
"Automatic shard shape inference is experimental and may change without notice.",
966945
ZarrUserWarning,
@@ -984,11 +963,12 @@ def resolve_outer_and_inner_chunks(
984963
_shards_out += (c_shape * num_chunks_per_shard_axis,)
985964
else:
986965
_shards_out += (c_shape,)
987-
shard_flat = _shards_out
966+
shard_spec = _shards_out
988967
elif isinstance(shard_shape, dict):
989-
shard_flat = tuple(shard_shape["shape"])
968+
shard_spec = shard_shape["shape"]
990969
else:
991-
shard_flat = cast("tuple[int, ...]", shard_shape)
970+
shard_spec = shard_shape
992971

993-
outer = normalize_chunks_nd(shard_flat, array_shape)
972+
# Regular and rectilinear shard specifications go through the one normalizer.
973+
outer = normalize_chunks_nd(shard_spec, array_shape)
994974
return ChunkLayout(outer_chunks=outer, inner=ChunkLayout(outer_chunks=chunks))

‎tests/test_unified_chunk_grid.py‎

Lines changed: 67 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
ChunkSpec,
2020
FixedDimension,
2121
VaryingDimension,
22-
_is_rectilinear_chunks,
2322
)
2423
from zarr.core.common import compress_rle, expand_rle
2524
from zarr.core.metadata.v3 import (
@@ -557,55 +556,6 @@ def test_expand_rle_pair_with_float_count() -> None:
557556
assert result == [10, 10, 10]
558557

559558

560-
# ---------------------------------------------------------------------------
561-
# _is_rectilinear_chunks tests
562-
# ---------------------------------------------------------------------------
563-
564-
565-
@pytest.mark.parametrize(
566-
("value", "expected"),
567-
[
568-
([[10, 20], [5, 5]], True),
569-
(((10, 20), (5, 5)), True),
570-
((10, 20), False),
571-
([10, 20], False),
572-
(10, False),
573-
("auto", False),
574-
([], False),
575-
([[]], True),
576-
(ChunkGrid.from_sizes((10,), (5,)), False),
577-
(None, False),
578-
(3.14, False),
579-
],
580-
ids=[
581-
"nested-lists",
582-
"nested-tuples",
583-
"flat-tuple",
584-
"flat-list",
585-
"single-int",
586-
"string",
587-
"empty-list",
588-
"empty-nested-list",
589-
"chunk-grid-instance",
590-
"none",
591-
"float",
592-
],
593-
)
594-
def test_is_rectilinear_chunks(value: Any, expected: bool) -> None:
595-
"""_is_rectilinear_chunks correctly identifies nested sequences as rectilinear"""
596-
assert _is_rectilinear_chunks(value) is expected
597-
598-
599-
def test_is_rectilinear_chunks_handles_broken_iterable() -> None:
600-
"""_is_rectilinear_chunks returns False for objects that raise on iteration."""
601-
602-
class BrokenIter:
603-
def __iter__(self) -> Any:
604-
raise TypeError("cannot iterate")
605-
606-
assert _is_rectilinear_chunks(BrokenIter()) is False
607-
608-
609559
# ---------------------------------------------------------------------------
610560
# Serialization tests
611561
# ---------------------------------------------------------------------------
@@ -1211,6 +1161,73 @@ def test_mixed_chunks_gates_are_order_independent(chunks: Any) -> None:
12111161
zarr.create_array(MemoryStore(), shape=shape, chunks=chunks, zarr_format=2, dtype="uint8")
12121162

12131163

1164+
@pytest.mark.parametrize(
1165+
("chunks", "expected"),
1166+
[
1167+
pytest.param(np.int64(2), (2, 2), id="numpy-scalar"),
1168+
pytest.param(np.array(2), (2, 2), id="0-d-array"),
1169+
pytest.param((np.int64(5), np.int32(3)), (5, 3), id="tuple-of-numpy-ints"),
1170+
pytest.param(np.array([5, 3], dtype=np.uint8), (5, 3), id="1-d-array"),
1171+
pytest.param([np.int64(5), np.int64(3)], (5, 3), id="list-of-numpy-ints"),
1172+
pytest.param((np.array(5), 3), (5, 3), id="0-d-array-element"),
1173+
pytest.param((np.int64(-1), 3), (10, 3), id="numpy-minus-one"),
1174+
],
1175+
)
1176+
@pytest.mark.parametrize("creator", ["create_array-v3", "create_array-v2", "create-v2"])
1177+
def test_numpy_chunk_specs_normalize(chunks: Any, expected: tuple[int, ...], creator: str) -> None:
1178+
"""numpy scalars, 0-d arrays, and integer arrays are accepted wherever a
1179+
chunk shape is, in every zarr format and creation API, and are stored as
1180+
plain ints. The legacy `zarr.create` v2 path used to test the spec's truth
1181+
value, which a numpy array does not have."""
1182+
store = MemoryStore()
1183+
if creator == "create_array-v3":
1184+
arr = zarr.create_array(store, shape=(10, 6), chunks=chunks, dtype="uint8")
1185+
elif creator == "create_array-v2":
1186+
arr = zarr.create_array(store, shape=(10, 6), chunks=chunks, dtype="uint8", zarr_format=2)
1187+
else:
1188+
arr = zarr.create(store=store, shape=(10, 6), chunks=chunks, dtype="uint8", zarr_format=2)
1189+
assert arr.chunks == expected
1190+
assert all(type(c) is int for c in arr.chunks)
1191+
1192+
1193+
@pytest.mark.parametrize(
1194+
("chunks", "expected"),
1195+
[
1196+
pytest.param((2, np.array([3, 3])), (2, (3, 3)), id="int-and-array-edges"),
1197+
pytest.param((np.int64(2), [3, 3]), (2, (3, 3)), id="numpy-int-and-list"),
1198+
pytest.param(np.array([[5, 5], [3, 3]]), ((5, 5), (3, 3)), id="2-d-array"),
1199+
pytest.param(([np.int64(5), np.int64(5)], 3), ((5, 5), 3), id="edges-of-numpy-ints"),
1200+
],
1201+
)
1202+
def test_numpy_rectilinear_chunk_specs_normalize(
1203+
chunks: Any, expected: tuple[int | tuple[int, ...], ...]
1204+
) -> None:
1205+
"""numpy values inside a rectilinear `chunks=` spec are stored as plain ints
1206+
and the grid kind is decided by the normalized grid, not by duck typing."""
1207+
arr = zarr.create_array(MemoryStore(), shape=(10, 6), chunks=chunks, dtype="uint8")
1208+
assert arr.metadata.chunk_grid == RectilinearChunkGridMetadata(chunk_shapes=expected)
1209+
1210+
1211+
@pytest.mark.parametrize(
1212+
"shards",
1213+
[np.array([10, 6]), (np.int64(10), np.int32(6)), [np.int64(10), np.int64(6)]],
1214+
ids=["array", "tuple-of-numpy-ints", "list-of-numpy-ints"],
1215+
)
1216+
def test_numpy_shard_specs_normalize(shards: Any) -> None:
1217+
"""numpy values in `shards=` normalize like those in `chunks=`."""
1218+
arr = zarr.create_array(
1219+
MemoryStore(), shape=(10, 6), chunks=(2, 3), shards=shards, dtype="uint8"
1220+
)
1221+
assert arr.shards == (10, 6)
1222+
assert arr.chunks == (2, 3)
1223+
1224+
1225+
def test_zero_dim_float_array_chunks_rejected() -> None:
1226+
"""A 0-d array unwraps to its scalar, which must still be an integer."""
1227+
with pytest.raises(TypeError, match="must be an integer or an iterable of integers"):
1228+
zarr.create_array(MemoryStore(), shape=(10,), chunks=np.array(2.0), dtype="uint8")
1229+
1230+
12141231
def test_from_array_keep_preserves_all_bare_int_rectilinear_grid() -> None:
12151232
"""A rectilinear grid using the bare-int shorthand on every dimension can
12161233
only arrive from externally written metadata — `create_array` never

0 commit comments

Comments
 (0)