From 9e9ee4eba85ca33bc410029f9ce70d5ace12562b Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 23:53:00 +0200 Subject: [PATCH 1/8] fix(chunk-grids): make array creation O(1) in per-dimension chunk count Chunk normalization now returns a ChunkGrid whose uniform dimensions are stored as FixedDimension (size + extent) instead of one array entry per chunk, so create_array(shape=(2**62,), chunks=(1,)) succeeds instantly instead of raising "array is too big", and chunks=(1, 1) on a (2**31, 2**31) array no longer allocates ~17 GB per dimension. Explicit per-chunk lists collapse to FixedDimension when they describe a regular grid; genuinely irregular lists become VaryingDimension. Both variants bind chunk sizes to their extent, so the two forms carry the same invariants. The intermediate ChunksTuple type and as_regular_shape helper are removed; create_chunk_grid_metadata consumes the grid directly and serializes uniform dimensions of mixed rectilinear grids as the spec's bare-int step-size shorthand. Creation-time counterpart of the gh-4174 indexing fix. Assisted-by: ClaudeCode:claude-fable-5 --- changes/4174.bugfix.md | 8 ++ src/zarr/core/array.py | 9 +- src/zarr/core/chunk_grids.py | 118 ++++++++++++------------- src/zarr/core/metadata/v3.py | 38 ++++---- tests/conftest.py | 5 +- tests/test_array.py | 18 +++- tests/test_chunk_grids.py | 166 ++++++++++++++++++++++------------- 7 files changed, 213 insertions(+), 149 deletions(-) create mode 100644 changes/4174.bugfix.md diff --git a/changes/4174.bugfix.md b/changes/4174.bugfix.md new file mode 100644 index 0000000000..672a87f807 --- /dev/null +++ b/changes/4174.bugfix.md @@ -0,0 +1,8 @@ +Array creation is now O(1) in the number of chunks per dimension. Chunk +normalization returns a `ChunkGrid` whose uniform dimensions are stored as a +size + extent pair (`FixedDimension`) instead of being expanded to one entry +per chunk, so creating arrays like +`zarr.create_array(store, shape=(2**62,), chunks=(1,), dtype='int32')` succeeds +instantly instead of raising `ValueError` or allocating gigabytes of memory. +The intermediate `ChunksTuple` representation was removed in the process. +This is the creation-time counterpart of the indexing fix in #4172. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 903f7b8f26..9c4ec072eb 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -43,7 +43,6 @@ SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, _is_rectilinear_chunks, - as_regular_shape, guess_chunks, normalize_chunks_nd, resolve_outer_and_inner_chunks, @@ -523,7 +522,7 @@ async def _create( outer_chunks = guess_chunks(shape, item_size) else: outer_chunks = normalize_chunks_nd(_raw, shape) - _chunks = as_regular_shape(outer_chunks) + _chunks = outer_chunks.chunk_shape if order is None: order_parsed = config_parsed.order @@ -4472,7 +4471,7 @@ async def init_array( "chunks=(inner_size, ...), shards=[[shard_sizes], ...]" ) - # Normalize the user's chunks into canonical ChunksTuple form + # Normalize the user's chunks into a canonical ChunkGrid if chunks == "auto": max_bytes = None if shards is None else SHARDED_INNER_CHUNK_MAX_BYTES @@ -4515,7 +4514,7 @@ async def init_array( meta = AsyncArray._create_metadata_v2( shape=shape_parsed, dtype=zdtype, - chunks=as_regular_shape(outer_chunks), + chunks=outer_chunks.chunk_shape, dimension_separator=chunk_key_encoding_parsed.separator, fill_value=fill_value, order=order_parsed, @@ -4534,7 +4533,7 @@ async def init_array( grid = create_chunk_grid_metadata(outer_chunks) codecs_out: tuple[Codec, ...] if inner is not None: - inner_chunks_flat = as_regular_shape(inner.outer_chunks) + inner_chunks_flat = inner.outer_chunks.chunk_shape index_location: IndexLocation = "end" if isinstance(shards, dict): index_location = cast("IndexLocation", shards.get("index_location", "end")) diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 584829bc6c..e1486b8694 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -12,7 +12,6 @@ TYPE_CHECKING, Any, NamedTuple, - NewType, Protocol, TypeGuard, cast, @@ -43,31 +42,6 @@ is not `None`. Explicit chunk sizes are not affected by this value. """ -ChunksTuple = NewType("ChunksTuple", tuple[np.ndarray[tuple[int], np.dtype[np.int64]], ...]) -"""Normalized chunk specification: one 1D int64 array of chunk sizes per dimension. - -Produced exclusively by `normalize_chunks_nd` and `guess_chunks`. -Consumers should use this type to ensure they receive validated, -canonical chunk specifications rather than raw user input. -""" - - -class ChunkLayout(NamedTuple): - """Result of resolving user `chunks`/`shards` into grid metadata inputs. - - outer_chunks - Chunk sizes for the chunk grid metadata. When sharding is active - these are the shard sizes; otherwise they are the user's chunk sizes. - inner - Recursive sub-structure inside each chunk. `None` means the chunk is - opaque (no sharding). When present, `inner.outer_chunks` gives the - sub-chunk sizes passed to `ShardingCodec`, and `inner.inner` gives - the next level of nesting (for nested sharding), or `None`. - """ - - outer_chunks: ChunksTuple - inner: ChunkLayout | None = None - @dataclass(frozen=True) class FixedDimension: @@ -389,12 +363,6 @@ def is_regular_nd( return all(is_regular_1d(d) for d in chunks) -def as_regular_shape(chunks: ChunksTuple) -> tuple[int, ...]: - """Flatten a regular ChunksTuple to one int per dimension.""" - assert is_regular_nd(chunks), f"expected regular chunks, got {chunks}" - return tuple(int(dim[0]) for dim in chunks) - - @dataclass(frozen=True) class ChunkGrid: """ @@ -489,6 +457,11 @@ def from_sizes( # -- Properties -- + @property + def dimensions(self) -> tuple[DimensionGrid, ...]: + """The per-dimension grids (`FixedDimension` or `VaryingDimension`).""" + return self._dimensions + @property def ndim(self) -> int: return len(self._dimensions) @@ -641,6 +614,23 @@ def update_shape(self, new_shape: tuple[int, ...]) -> ChunkGrid: return ChunkGrid(dimensions=dims) +class ChunkLayout(NamedTuple): + """Result of resolving user `chunks`/`shards` into grid metadata inputs. + + outer_chunks + Chunk grid for the chunk grid metadata. When sharding is active + this holds the shard sizes; otherwise it holds the user's chunk sizes. + inner + Recursive sub-structure inside each chunk. `None` means the chunk is + opaque (no sharding). When present, `inner.outer_chunks` gives the + sub-chunk sizes passed to `ShardingCodec`, and `inner.inner` gives + the next level of nesting (for nested sharding), or `None`. + """ + + outer_chunks: ChunkGrid + inner: ChunkLayout | None = None + + def _guess_regular_chunks( shape: tuple[int, ...] | int, typesize: int, @@ -717,30 +707,31 @@ def _guess_regular_chunks( return tuple(int(x) for x in chunks) -def normalize_chunks_1d( - chunks: int | Iterable[object], span: int -) -> np.ndarray[tuple[int], np.dtype[np.int64]]: +def normalize_chunks_1d(chunks: int | Iterable[object], span: int) -> DimensionGrid: """ - Normalize a one-dimensional chunk specification into a 1D int64 array of - chunk sizes that cover the span. + Normalize a one-dimensional chunk specification into a dimension grid: + `FixedDimension` for uniform chunk sizes, `VaryingDimension` for explicit + per-chunk sizes that genuinely vary. Both variants bind the chunk sizes to + the span, and the uniform form is O(1) in the number of chunks — a + dimension with `2**62` chunks must not materialize one entry per chunk. `-1` means "one chunk covering the entire span." - For an integer chunk size, all chunks are uniform — the last chunk may - overhang the span. The actual data extent of each chunk is determined - by the chunk grid at runtime, not by this function. + Explicit chunk size lists must sum to the span exactly; lists that describe + a regular grid (all sizes equal, or equal with a smaller boundary chunk) + collapse to `FixedDimension`. For uniform sizes the last chunk may overhang + the span. """ # `numbers.Integral` rather than `int` so that numpy integer scalars (which are not # `int` subclasses) take the uniform-chunk path instead of being treated as a sequence. + # The `-1` sentinel check lives inside this branch so that numpy-array chunk + # specifications never hit an ambiguous-truth-value error on `chunks == -1`. if isinstance(chunks, numbers.Integral): chunk_size = int(chunks) if chunk_size == -1: - return np.array([span], dtype=np.int64) + return FixedDimension(size=span, extent=span) if chunk_size <= 0: raise ValueError(f"Chunk size must be positive, got {chunk_size}") - if span == 0: - return np.array([chunk_size], dtype=np.int64) - n = ceildiv(span, chunk_size) - return np.full(n, chunk_size, dtype=np.int64) + return FixedDimension(size=chunk_size, extent=span) else: try: chunk_list = list(chunks) # type: ignore[arg-type] @@ -766,23 +757,30 @@ def normalize_chunks_1d( raise ValueError(f"All chunk sizes must be positive, got {ints}") if sum(ints) != span: raise ValueError(f"Chunk sizes {ints} do not sum to span {span}") - return np.asarray(ints, dtype=np.int64) + if is_regular_1d(ints): + return FixedDimension(size=ints[0], extent=span) + return VaryingDimension(ints, extent=span) def normalize_chunks_nd( chunks: Any, shape: tuple[int, ...], -) -> ChunksTuple: +) -> ChunkGrid: """ - Normalize a chunk specification into a `ChunksTuple`. + Normalize a chunk specification into a `ChunkGrid`. This is a mechanical transformation — no heuristics, no guessing. Handles `False` ("all data in one chunk"), scalar ints, `-1` sentinels (one chunk per dimension covering the full span), and explicit per-dimension lists of chunk sizes (regular or rectilinear). + This is the strict parser for user-supplied chunk specifications; use + `ChunkGrid.from_sizes` / `ChunkGrid.from_metadata` for stored metadata, + which is validated under more tolerant rules (e.g. trailing edges beyond + the array extent). + For auto-chunking, use `guess_chunks` which returns a - `ChunksTuple` directly. `chunks=None` and `chunks=True` are rejected + `ChunkGrid` directly. `chunks=None` and `chunks=True` are rejected here — the caller is responsible for choosing between explicit sizes and auto-chunking. """ @@ -793,7 +791,9 @@ def normalize_chunks_nd( # handle no chunking if chunks is False: - return ChunksTuple(tuple(np.array([s], dtype=np.int64) for s in shape)) + return ChunkGrid( + dimensions=tuple(FixedDimension(size=int(s), extent=int(s)) for s in shape) + ) # handle 1D convenience form. bool is excluded above so this only catches actual ints. if isinstance(chunks, numbers.Integral): @@ -805,20 +805,20 @@ def normalize_chunks_nd( f"chunks has {len(chunks)} dimensions but shape has {len(shape)} dimensions" ) - return ChunksTuple( - tuple(normalize_chunks_1d(c, span=s) for c, s in zip(chunks, shape, strict=True)) + return ChunkGrid( + dimensions=tuple(normalize_chunks_1d(c, span=s) for c, s in zip(chunks, shape, strict=True)) ) def guess_chunks( shape: tuple[int, ...], typesize: int, *, max_bytes: int | None = None -) -> ChunksTuple: +) -> ChunkGrid: """ Heuristically determine chunk sizes for an array. This is the policy function — it makes opinionated choices about chunk sizes based on array shape and element size, and returns a - normalized `ChunksTuple`. + normalized `ChunkGrid`. Parameters ---------- @@ -877,7 +877,7 @@ def _guess_num_chunks_per_axis_shard( def resolve_outer_and_inner_chunks( *, array_shape: tuple[int, ...], - chunks: ChunksTuple, + chunks: ChunkGrid, shard_shape: ShardsLike | None, item_size: int, ) -> ChunkLayout: @@ -888,7 +888,7 @@ def resolve_outer_and_inner_chunks( array_shape The array shape. chunks - Normalized chunk specification (the user's `chunks=`). + Normalized chunk grid (the user's `chunks=`). shard_shape Raw shard specification (the user's `shards=`). `None` means no sharding, `"auto"` triggers heuristic inference, @@ -900,7 +900,7 @@ def resolve_outer_and_inner_chunks( Returns ------- ChunkLayout - `outer_chunks` is the `ChunksTuple` for chunk grid + `outer_chunks` is the `ChunkGrid` for chunk grid metadata. `inner` holds the sub-chunk structure for `ShardingCodec`, or is `None` when sharding is not active. """ @@ -912,8 +912,8 @@ def resolve_outer_and_inner_chunks( outer = normalize_chunks_nd(shard_shape, array_shape) return ChunkLayout(outer_chunks=outer, inner=ChunkLayout(outer_chunks=chunks)) - # Extract the flat chunk shape (first size per dimension) for arithmetic. - chunk_shape_flat = as_regular_shape(chunks) + # Extract the flat chunk shape (uniform size per dimension) for arithmetic. + chunk_shape_flat = chunks.chunk_shape if shard_shape == "auto": warnings.warn( diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index fc47f8fc95..fe63b11d33 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -12,7 +12,7 @@ from zarr.core._json import json_to_buffer from zarr.core.array_spec import ArrayConfig, ArraySpec from zarr.core.buffer.core import default_buffer_prototype -from zarr.core.chunk_grids import is_regular_nd +from zarr.core.chunk_grids import FixedDimension, VaryingDimension from zarr.core.chunk_key_encodings import ( ChunkKeyEncoding, ChunkKeyEncodingLike, @@ -43,7 +43,7 @@ from typing import Self from zarr.core.buffer import Buffer, BufferPrototype - from zarr.core.chunk_grids import ChunksTuple + from zarr.core.chunk_grids import ChunkGrid from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar @@ -373,32 +373,36 @@ def from_dict(cls, data: RectilinearChunkGridMetadataJSON) -> Self: # type: ign def create_chunk_grid_metadata( - chunks: ChunksTuple, + chunks: ChunkGrid, ) -> ChunkGridMetadata: - """Construct a chunk grid metadata object from a normalized `ChunksTuple`. + """Construct a chunk grid metadata object from a normalized `ChunkGrid`. - Regular chunks produce a `RegularChunkGridMetadata`. - Rectilinear chunks produce a `RectilinearChunkGridMetadata`. + Regular grids produce a `RegularChunkGridMetadata`. + Rectilinear grids produce a `RectilinearChunkGridMetadata`. Parameters ---------- - chunks : ChunksTuple - Normalized chunk specification, as returned by + chunks : ChunkGrid + Normalized chunk grid, as returned by `normalize_chunks_nd` or `guess_chunks`. See Also -------- parse_chunk_grid : Deserialize a chunk grid from stored JSON metadata. """ - if is_regular_nd(chunks): - # If we know the chunks specification is regular, then we can take the first - # chunk size for each dimension as the chunk shape. - chunk_shape = tuple(int(dim_chunks[0]) for dim_chunks in chunks) - return RegularChunkGridMetadata(chunk_shape=chunk_shape) - else: - return RectilinearChunkGridMetadata( - chunk_shapes=tuple(tuple(int(x) for x in d) for d in chunks) - ) + if chunks.is_regular: + return RegularChunkGridMetadata(chunk_shape=chunks.chunk_shape) + # Uniform dimensions stay bare ints — the rectilinear grid spec treats + # a bare int as a step size repeating to cover the axis. + chunk_shapes: list[int | tuple[int, ...]] = [] + for dim in chunks.dimensions: + if isinstance(dim, FixedDimension): + chunk_shapes.append(dim.size) + elif isinstance(dim, VaryingDimension): + chunk_shapes.append(dim.edges) + else: + raise TypeError(f"Unknown dimension grid type: {type(dim)}") + return RectilinearChunkGridMetadata(chunk_shapes=tuple(chunk_shapes)) def parse_chunk_grid( diff --git a/tests/conftest.py b/tests/conftest.py index 7ccf9958e7..6fc96237af 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -25,7 +25,6 @@ ) from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, - as_regular_shape, guess_chunks, normalize_chunks_nd, resolve_outer_and_inner_chunks, @@ -393,7 +392,7 @@ def create_array_metadata( return ArrayV2Metadata( shape=shape_parsed, dtype=dtype_parsed, - chunks=as_regular_shape(outer_chunks), + chunks=outer_chunks.chunk_shape, order=order_parsed, dimension_separator=chunk_key_encoding_parsed.separator, fill_value=fill_value, @@ -412,7 +411,7 @@ def create_array_metadata( sub_codecs: tuple[Codec, ...] = (*array_array, array_bytes, *bytes_bytes) codecs_out: tuple[Codec, ...] if inner is not None: - inner_chunks_flat = as_regular_shape(inner.outer_chunks) + inner_chunks_flat = inner.outer_chunks.chunk_shape index_location: IndexLocation = "end" if isinstance(shards, dict): index_location = cast("IndexLocation", shards.get("index_location", "end")) diff --git a/tests/test_array.py b/tests/test_array.py index e18ce3cfb1..46890244ec 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -1091,7 +1091,7 @@ def test_auto_partition_auto_shards( shard_shape="auto", item_size=dtype.itemsize, ) - auto_shards = tuple(dim[0] for dim in outer_chunks) + auto_shards = outer_chunks.chunk_shape assert auto_shards == expected_shards @@ -1106,7 +1106,7 @@ def test_auto_partition_auto_shards_with_auto_chunks_should_be_close_to_1MiB() - chunks_normalized = guess_chunks( array_shape, item_size, max_bytes=SHARDED_INNER_CHUNK_MAX_BYTES ) - chunk_shape = tuple(dim[0] for dim in chunks_normalized) + chunk_shape = chunks_normalized.chunk_shape chunk_bytes = np.prod(chunk_shape) * item_size assert chunk_bytes <= SHARDED_INNER_CHUNK_MAX_BYTES assert chunk_bytes > SHARDED_INNER_CHUNK_MAX_BYTES // 4 # should be in the right ballpark @@ -1123,7 +1123,7 @@ def test_auto_partition_auto_shards_with_auto_chunks_should_be_close_to_1MiB() - item_size=item_size, ) assert inner is not None - shard_shape = tuple(dim[0] for dim in outer_chunks) + shard_shape = outer_chunks.chunk_shape # Shard dimensions must be multiples of chunk dimensions assert all(s % c == 0 for s, c in zip(shard_shape, chunk_shape, strict=True)) @@ -2449,3 +2449,15 @@ async def test_create_array_chunks_3d( shape = (10, 12, 15) arr = await create_array(store={}, shape=shape, chunks=chunk_input, dtype="float64") assert arr.write_chunk_sizes == expected + + +async def test_create_array_huge_chunk_count() -> None: + """Array creation must be O(1) in the number of chunks per dimension. + + With `shape=(2**62,)` and `chunks=(1,)` this dimension has 2**62 chunks; + materializing one entry per chunk would raise ("array is too big"). + Companion to the indexing-time fix from gh-4174. + """ + arr = await create_array(store={}, shape=(2**62,), chunks=(1,), dtype="int32") + assert arr.shape == (2**62,) + assert arr.chunks == (1,) diff --git a/tests/test_chunk_grids.py b/tests/test_chunk_grids.py index 4640c43d1c..eec1a928b5 100644 --- a/tests/test_chunk_grids.py +++ b/tests/test_chunk_grids.py @@ -5,7 +5,10 @@ from tests.conftest import Expect, ExpectFail from zarr.core.chunk_grids import ( + ChunkGrid, ChunkLayout, + FixedDimension, + VaryingDimension, _guess_regular_chunks, normalize_chunks_1d, normalize_chunks_nd, @@ -14,15 +17,30 @@ def _assert_chunks_equal( - actual: tuple[Any, ...], - expected: tuple[tuple[int, ...], ...], + actual: ChunkGrid, + expected: tuple[int | tuple[int, ...], ...], + shape: tuple[int, ...], ) -> None: - """Compare a ChunksTuple (tuple of np.int64 arrays) against a tuple of int tuples.""" - assert len(actual) == len(expected), f"axis count mismatch: {len(actual)} vs {len(expected)}" - for axis, (a, e) in enumerate(zip(actual, expected, strict=True)): - assert np.array_equal(a, np.asarray(e, dtype=np.int64)), ( - f"axis {axis}: {list(a)} != {list(e)}" - ) + """Compare a normalized ChunkGrid against per-dimension expectations. + + An expected bare `int` requires a `FixedDimension` with that uniform size; + an expected tuple requires a `VaryingDimension` with those exact edges. + Every dimension must carry the corresponding extent from `shape`. + """ + dims = actual.dimensions + assert len(dims) == len(expected) == len(shape), ( + f"axis count mismatch: {len(dims)} vs {len(expected)} vs {len(shape)}" + ) + for axis, (a, e, span) in enumerate(zip(dims, expected, shape, strict=True)): + if isinstance(e, int): + assert isinstance(a, FixedDimension), f"axis {axis}: expected FixedDimension, got {a!r}" + assert a.size == e, f"axis {axis}: size {a.size} != {e}" + else: + assert isinstance(a, VaryingDimension), ( + f"axis {axis}: expected VaryingDimension, got {a!r}" + ) + assert a.edges == tuple(e), f"axis {axis}: edges {a.edges} != {tuple(e)}" + assert a.extent == span, f"axis {axis}: extent {a.extent} != {span}" @pytest.mark.parametrize( @@ -42,79 +60,83 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None: @pytest.mark.parametrize( ("chunks", "shape", "expected"), [ - # 1D cases - ((10,), (100,), ((10,) * 10,)), - ([10], (100,), ((10,) * 10,)), - (10, (100,), ((10,) * 10,)), + # 1D cases (uniform sizes stay bare ints) + ((10,), (100,), (10,)), + ([10], (100,), (10,)), + (10, (100,), (10,)), # 2D cases - ((10, 10), (100, 10), ((10,) * 10, (10,))), - (10, (100, 10), ((10,) * 10, (10,))), - ((10, -1), (100, 10), ((10,) * 10, (10,))), + ((10, 10), (100, 10), (10, 10)), + (10, (100, 10), (10, 10)), + ((10, -1), (100, 10), (10, 10)), # 3D cases - (30, (100, 20, 10), ((30, 30, 30, 30), (30,), (30,))), - ((30, -1, -1), (100, 20, 10), ((30, 30, 30, 30), (20,), (10,))), - ((30, 20, -1), (100, 20, 10), ((30, 30, 30, 30), (20,), (10,))), - ((30, 20, 10), (100, 20, 10), ((30, 30, 30, 30), (20,), (10,))), - # dask-style chunks (explicit per-chunk sizes) - (((100, 100, 100), (50, 50)), (300, 100), ((100, 100, 100), (50, 50))), - (((100, 100, 50),), (250,), ((100, 100, 50),)), - (((100,),), (100,), ((100,),)), + (30, (100, 20, 10), (30, 30, 30)), + ((30, -1, -1), (100, 20, 10), (30, 20, 10)), + ((30, 20, -1), (100, 20, 10), (30, 20, 10)), + ((30, 20, 10), (100, 20, 10), (30, 20, 10)), + # dask-style chunks describing a regular grid collapse to the uniform form + (((100, 100, 100), (50, 50)), (300, 100), (100, 50)), + (((100, 100, 50),), (250,), (100,)), + (((100,),), (100,), (100,)), + # genuinely irregular explicit chunks keep the per-chunk form + (((10, 20, 70), (50, 50)), (100, 100), ((10, 20, 70), 50)), # no chunking (False means each dimension is one chunk spanning the full extent) - (False, (100,), ((100,),)), - (False, (100, 50), ((100,), (50,))), + (False, (100,), (100,)), + (False, (100, 50), (100, 50)), # sentinel values - (-1, (100,), ((100,),)), + (-1, (100,), (100,)), # zero-length dimensions preserve the declared chunk size - (10, (0,), ((10,),)), - ((5, 10), (0, 100), ((5,), (10,) * 10)), - ((5, 10), (20, 0), ((5, 5, 5, 5), (10,))), + (10, (0,), (10,)), + ((5, 10), (0, 100), (5, 10)), + ((5, 10), (20, 0), (5, 10)), # numpy integers are accepted anywhere a python int is, whether as the scalar # convenience form, as per-dimension entries, or as the `-1` sentinel. - (np.int64(10), (100,), ((10,) * 10,)), - ((np.int64(2), np.int64(2)), (4, 4), ((2, 2), (2, 2))), - ((1, 3, np.int64(16), np.int64(16)), (1, 3, 32, 32), ((1,), (3,), (16, 16), (16, 16))), - ((np.int32(30), np.int64(-1)), (100, 20), ((30, 30, 30, 30), (20,))), - (np.array([10, 10]), (100, 100), ((10,) * 10, (10,) * 10)), + (np.int64(10), (100,), (10,)), + ((np.int64(2), np.int64(2)), (4, 4), (2, 2)), + ((1, 3, np.int64(16), np.int64(16)), (1, 3, 32, 32), (1, 3, 16, 16)), + ((np.int32(30), np.int64(-1)), (100, 20), (30, 20)), + (np.array([10, 10]), (100, 100), (10, 10)), # rectilinear chunks given as numpy arrays - ((np.array([60, 40]), np.array([50, 50])), (100, 100), ((60, 40), (50, 50))), + ((np.array([60, 40]), np.array([50, 50])), (100, 100), (60, 50)), ], ) def test_normalize_chunks( - chunks: Any, shape: tuple[int, ...], expected: tuple[tuple[int, ...], ...] + chunks: Any, shape: tuple[int, ...], expected: tuple[int | tuple[int, ...], ...] ) -> None: - _assert_chunks_equal(normalize_chunks_nd(chunks, shape), expected) + _assert_chunks_equal(normalize_chunks_nd(chunks, shape), expected, shape) @pytest.mark.parametrize( ("array_shape", "chunks_input", "shard_shape", "expected_outer", "expected_inner_outer"), [ # no sharding: outer = chunks, inner = None - ((100,), (10,), None, ((10,) * 10,), None), + ((100,), (10,), None, (10,), None), # explicit regular shards - ((100,), (10,), (50,), ((50, 50),), ((10,) * 10,)), - # rectilinear shards - ((100,), (10,), ((60, 40),), ((60, 40),), ((10,) * 10,)), + ((100,), (10,), (50,), (50,), (10,)), + # rectilinear shards describing a regular-with-boundary grid collapse + ((100,), (10,), ((60, 40),), (60,), (10,)), + # genuinely irregular rectilinear shards keep the per-chunk form + ((100,), (10,), ((30, 60, 10),), ((30, 60, 10),), (10,)), # dict-style shards - ((100, 100), (10, 10), {"shape": (50, 50)}, ((50, 50), (50, 50)), ((10,) * 10, (10,) * 10)), + ((100, 100), (10, 10), {"shape": (50, 50)}, (50, 50), (10, 10)), ], ) def test_resolve_outer_and_inner_chunks( array_shape: tuple[int, ...], chunks_input: tuple[int, ...], shard_shape: Any, - expected_outer: tuple[tuple[int, ...], ...], - expected_inner_outer: tuple[tuple[int, ...], ...] | None, + expected_outer: tuple[int | tuple[int, ...], ...], + expected_inner_outer: tuple[int | tuple[int, ...], ...] | None, ) -> None: chunks = normalize_chunks_nd(chunks_input, array_shape) outer_chunks, inner = resolve_outer_and_inner_chunks( array_shape=array_shape, chunks=chunks, shard_shape=shard_shape, item_size=1 ) - _assert_chunks_equal(outer_chunks, expected_outer) + _assert_chunks_equal(outer_chunks, expected_outer, array_shape) if expected_inner_outer is None: assert inner is None else: assert inner is not None - _assert_chunks_equal(inner.outer_chunks, expected_inner_outer) + _assert_chunks_equal(inner.outer_chunks, expected_inner_outer, array_shape) assert inner.inner is None @@ -128,11 +150,11 @@ def test_chunk_layout_nested() -> None: top = ChunkLayout(outer_chunks=normalize_chunks_nd((50, 50), (100, 100)), inner=mid) # Three levels: top -> mid -> leaf - _assert_chunks_equal(top.outer_chunks, ((50, 50), (50, 50))) + _assert_chunks_equal(top.outer_chunks, (50, 50), (100, 100)) assert top.inner is not None - _assert_chunks_equal(top.inner.outer_chunks, ((25,) * 4, (25,) * 4)) + _assert_chunks_equal(top.inner.outer_chunks, (25, 25), (100, 100)) assert top.inner.inner is not None - _assert_chunks_equal(top.inner.inner.outer_chunks, ((5,) * 20, (5,) * 20)) + _assert_chunks_equal(top.inner.inner.outer_chunks, (5, 5), (100, 100)) assert top.inner.inner.inner is None @@ -258,22 +280,42 @@ def test_normalize_chunks_nd_errors(case: ExpectFail[tuple[Any, tuple[int, ...]] @pytest.mark.parametrize( "case", [ - # uniform-chunks branch: one int → broadcast across span via np.full. - Expect(input=(1000, 100_000), output=[1000] * 100, id="uniform"), - # explicit-per-chunk branch. - Expect(input=([10, 20, 30, 40], 100), output=[10, 20, 30, 40], id="explicit-list"), + # uniform-chunks branch: O(1) size+extent record, never one entry per chunk. + Expect( + input=(1000, 100_000), output=FixedDimension(size=1000, extent=100_000), id="uniform" + ), + # uniform chunks on a span too large to expand per-chunk (creation-time + # counterpart of the gh-4174 indexing fix). + Expect(input=(1, 2**62), output=FixedDimension(size=1, extent=2**62), id="uniform-huge"), # -1 sentinel branch: one chunk covering the full span. - Expect(input=(-1, 100), output=[100], id="full-span-sentinel"), + Expect(input=(-1, 100), output=FixedDimension(size=100, extent=100), id="full-span"), + # zero-length span preserves the declared chunk size. + Expect(input=(10, 0), output=FixedDimension(size=10, extent=0), id="uniform-zero-span"), + # explicit lists that describe a regular grid collapse to the uniform form. + Expect( + input=([10, 10, 10], 30), + output=FixedDimension(size=10, extent=30), + id="explicit-regular", + ), + Expect( + input=([10, 10, 5], 25), + output=FixedDimension(size=10, extent=25), + id="explicit-boundary", + ), + # genuinely irregular edges keep the explicit per-chunk form. + Expect( + input=([10, 20, 70], 100), + output=VaryingDimension([10, 20, 70], extent=100), + id="explicit-irregular", + ), ], ids=lambda c: c.id, ) -def test_normalize_chunks_1d_returns_int64_array( - case: Expect[tuple[Any, int], list[int]], +def test_normalize_chunks_1d( + case: Expect[tuple[Any, int], FixedDimension | VaryingDimension], ) -> None: - """Every branch of normalize_chunks_1d must produce a 1D int64 array.""" + """Both output variants bind chunk sizes to the span: uniform specs become + `FixedDimension` (O(1) regardless of chunk count), irregular explicit + lists become `VaryingDimension`.""" chunks, span = case.input - result = normalize_chunks_1d(chunks, span) - assert isinstance(result, np.ndarray) - assert result.dtype == np.int64 - assert result.ndim == 1 - assert result.tolist() == case.output + assert normalize_chunks_1d(chunks, span) == case.output From bf1a243803909e7a0b73ca4af8b6be2606b2f819 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 12:32:03 +0200 Subject: [PATCH 2/8] refactor(chunk-grids): drop dead ndarray branch from is_regular_1d Since normalize_chunks_nd returns ChunkGrid dimensions instead of int64 arrays, is_regular_1d only ever receives plain Python sequences; the vectorized numpy path was unreachable in production code. Remove it and narrow the signatures of is_regular_1d / is_regular_nd to Sequence[int]. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/chunk_grids.py | 11 ++--------- tests/test_metadata/test_v3.py | 16 +--------------- 2 files changed, 3 insertions(+), 24 deletions(-) diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index e1486b8694..bd68f301eb 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -334,9 +334,7 @@ def _is_rectilinear_chunks(chunks: Any) -> TypeGuard[Sequence[Sequence[int]]]: return False -def is_regular_1d( - dim_chunks: Sequence[int] | np.ndarray[tuple[int], np.dtype[np.int64]], -) -> bool: +def is_regular_1d(dim_chunks: Sequence[int]) -> bool: """Check if a single dimension's chunk sizes represent a regular grid. A regular dimension has either all chunks the same size, or all @@ -346,9 +344,6 @@ def is_regular_1d( if len(dim_chunks) <= 1: return True first = dim_chunks[0] - if isinstance(dim_chunks, np.ndarray): - # Vectorized comparison avoids per-element Python iteration over int64 arrays. - return bool((dim_chunks[1:-1] == first).all() and dim_chunks[-1] <= first) for c in dim_chunks[1:-1]: if c != first: return False @@ -356,9 +351,7 @@ def is_regular_1d( return dim_chunks[-1] <= first -def is_regular_nd( - chunks: Iterable[Sequence[int] | np.ndarray[tuple[int], np.dtype[np.int64]]], -) -> bool: +def is_regular_nd(chunks: Iterable[Sequence[int]]) -> bool: """Check if an N-dimensional chunk specification represents a regular grid.""" return all(is_regular_1d(d) for d in chunks) diff --git a/tests/test_metadata/test_v3.py b/tests/test_metadata/test_v3.py index d1e156e500..77a0bad2e3 100644 --- a/tests/test_metadata/test_v3.py +++ b/tests/test_metadata/test_v3.py @@ -5,7 +5,6 @@ import json from typing import TYPE_CHECKING -import numpy as np import pytest from tests.conftest import Expect, ExpectFail @@ -110,9 +109,6 @@ def test_parse_codecs_unknown_raises(monkeypatch: pytest.MonkeyPatch) -> None: # Chunk-grid regularity helpers # --------------------------------------------------------------------------- -# Cases used for both list/tuple (Python-sequence path) and ndarray (vectorized -# path) of `is_regular_1d`. Parametrizing the input form ensures both branches -# are exercised by the same suite of edge cases. _REGULAR_1D_CASES: list[Expect[list[int], bool]] = [ Expect(input=[], output=True, id="empty"), Expect(input=[10], output=True, id="single-chunk"), @@ -129,19 +125,11 @@ def test_parse_codecs_unknown_raises(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.parametrize("case", _REGULAR_1D_CASES, ids=lambda c: c.id) def test_is_regular_1d_sequence(case: Expect[list[int], bool]) -> None: - """`is_regular_1d` accepts plain Python sequences and uses the iterative path.""" - # list and tuple both go through the non-ndarray branch. + """`is_regular_1d` accepts plain Python sequences.""" assert is_regular_1d(case.input) is case.output assert is_regular_1d(tuple(case.input)) is case.output -@pytest.mark.parametrize("case", _REGULAR_1D_CASES, ids=lambda c: c.id) -def test_is_regular_1d_ndarray(case: Expect[list[int], bool]) -> None: - """`is_regular_1d` accepts int64 ndarrays and uses the vectorized path.""" - arr = np.asarray(case.input, dtype=np.int64) - assert is_regular_1d(arr) is case.output - - @pytest.mark.parametrize( "case", [ @@ -156,8 +144,6 @@ def test_is_regular_1d_ndarray(case: Expect[list[int], bool]) -> None: def test_is_regular_nd_sequence(case: Expect[list[list[int]], bool]) -> None: """`is_regular_nd` returns True iff every per-dim spec is regular.""" assert is_regular_nd(case.input) is case.output - # Same result via ndarray inputs. - assert is_regular_nd([np.asarray(d, dtype=np.int64) for d in case.input]) is case.output # --------------------------------------------------------------------------- From 3b277af98dbc6bf3cfc95d6582ed5c4260c25e75 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 13:26:08 +0200 Subject: [PATCH 3/8] test(metadata): cover the unknown-dimension-type guard in create_chunk_grid_metadata The defensive TypeError branch is unreachable through the public API, so exercise it directly with a stub dimension object. This was the only genuinely uncovered patch line in #4218; the other lines codecov flagged came from a partial coverage upload. Assisted-by: ClaudeCode:claude-fable-5 --- tests/test_metadata/test_v3.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_metadata/test_v3.py b/tests/test_metadata/test_v3.py index 77a0bad2e3..9f78ed7b70 100644 --- a/tests/test_metadata/test_v3.py +++ b/tests/test_metadata/test_v3.py @@ -10,7 +10,7 @@ from tests.conftest import Expect, ExpectFail from tests.test_metadata.conftest import minimal_metadata_dict_v3 from zarr.core.buffer import default_buffer_prototype -from zarr.core.chunk_grids import is_regular_1d, is_regular_nd +from zarr.core.chunk_grids import ChunkGrid, is_regular_1d, is_regular_nd from zarr.core.config import config from zarr.core.dtype import Float64, UInt8 from zarr.core.group import GroupMetadata, parse_node_type @@ -19,6 +19,7 @@ ARRAY_METADATA_KEYS, ArrayMetadataJSON_V3, ArrayV3Metadata, + create_chunk_grid_metadata, parse_codecs, parse_dimension_names, parse_node_type_array, @@ -146,6 +147,13 @@ def test_is_regular_nd_sequence(case: Expect[list[list[int]], bool]) -> None: assert is_regular_nd(case.input) is case.output +def test_create_chunk_grid_metadata_unknown_dimension_type() -> None: + """`create_chunk_grid_metadata` rejects dimension grids it does not recognize.""" + grid = ChunkGrid(dimensions=(object(),)) # type: ignore[arg-type] + with pytest.raises(TypeError, match="Unknown dimension grid type"): + create_chunk_grid_metadata(grid) + + # --------------------------------------------------------------------------- # Types # --------------------------------------------------------------------------- From d3c5c17d5cadeb3098bfcb5f6f34c8a3cdd3dd9e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 13:35:23 +0200 Subject: [PATCH 4/8] fix(testing): pass bare-int rectilinear dims through the chunks= conversion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stateful hypothesis tests convert generated chunk grid metadata back into a create_array chunks= argument. Bare-int dimensions — the spec's step-size shorthand, now produced when a uniform dimension of a mixed grid collapses — were wrapped as single-element lists, turning "repeat to cover the axis" into "exactly one chunk" and failing the sum-to-span check (e.g. chunks=[1] for span 3). Extract the conversion into chunks_param_from_rectilinear, pass bare ints through unchanged, and widen ChunksLike to admit mixed int | sequence per-dimension specs, which the normalizer already accepted. Assisted-by: ClaudeCode:claude-fable-5 --- changes/4174.bugfix.md | 5 ++++- src/zarr/core/common.py | 4 +++- src/zarr/testing/strategies.py | 20 +++++++++++++++----- tests/test_properties.py | 26 ++++++++++++++++++++++++++ 4 files changed, 48 insertions(+), 7 deletions(-) diff --git a/changes/4174.bugfix.md b/changes/4174.bugfix.md index 672a87f807..2e4fcb2eec 100644 --- a/changes/4174.bugfix.md +++ b/changes/4174.bugfix.md @@ -4,5 +4,8 @@ size + extent pair (`FixedDimension`) instead of being expanded to one entry per chunk, so creating arrays like `zarr.create_array(store, shape=(2**62,), chunks=(1,), dtype='int32')` succeeds instantly instead of raising `ValueError` or allocating gigabytes of memory. -The intermediate `ChunksTuple` representation was removed in the process. +The intermediate `ChunksTuple` representation was removed in the process, and +`ChunksLike` now admits per-dimension specs that mix a bare int (uniform chunk +size) with explicit edge-length sequences, matching what the normalizer and +the rectilinear grid spec already accepted. This is the creation-time counterpart of the indexing fix in #4172. diff --git a/src/zarr/core/common.py b/src/zarr/core/common.py index 3da5c108b6..3008b1fef7 100644 --- a/src/zarr/core/common.py +++ b/src/zarr/core/common.py @@ -35,7 +35,9 @@ BytesLike = bytes | bytearray | memoryview ShapeLike = Iterable[int | np.integer[Any]] | int | np.integer[Any] -ChunksLike = ShapeLike | Iterable[Iterable[int]] +# Per-dimension chunk specs may mix a bare int (uniform chunk size, the +# rectilinear spec's step-size shorthand) with explicit edge-length sequences. +ChunksLike = ShapeLike | Iterable[int | Iterable[int]] # For backwards compatibility ChunkCoords = tuple[int, ...] ZarrFormat = Literal[2, 3] diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 6679dbcee4..db14d5f29b 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -312,17 +312,14 @@ def arrays( # - RegularChunkGridMetadata -> flat tuple of ints # - RectilinearChunkGridMetadata -> nested list of ints (triggers rectilinear path) # - v2 -> flat tuple of ints - chunks_param: tuple[int, ...] | list[list[int]] + chunks_param: tuple[int, ...] | list[int | list[int]] shard_shape = None dim_names = None if zarr_format == 3: chunk_grid_meta = draw(st.none() | chunk_grids(shape=nparray.shape), label="chunk grid") dim_names = draw(dimension_names(ndim=nparray.ndim), label="dimension names") if isinstance(chunk_grid_meta, RectilinearChunkGridMetadata): - chunks_param = [ - list(dim) if isinstance(dim, tuple) else [dim] - for dim in chunk_grid_meta.chunk_shapes - ] + chunks_param = chunks_param_from_rectilinear(chunk_grid_meta) elif isinstance(chunk_grid_meta, RegularChunkGridMetadata): chunks_param = chunk_grid_meta.chunk_shape else: @@ -404,6 +401,19 @@ def simple_arrays( ) +def chunks_param_from_rectilinear( + meta: RectilinearChunkGridMetadata, +) -> list[int | list[int]]: + """Convert rectilinear chunk grid metadata into a `chunks=` argument. + + Explicit edge tuples become lists. Bare ints — the spec's step-size + shorthand meaning "repeat to cover the axis" — pass through unchanged; + wrapping one in a single-element list would instead declare exactly one + chunk, which fails normalization whenever the axis needs more than one. + """ + return [list(dim) if isinstance(dim, tuple) else dim for dim in meta.chunk_shapes] + + @st.composite def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[int]]: """Generate valid rectilinear chunk shapes for a given array shape. diff --git a/tests/test_properties.py b/tests/test_properties.py index 33888bfd4e..a73a6d8a7c 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -448,3 +448,29 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N assert serialized_complex_float_is_valid(asdict_dict["fill_value"]) elif dtype_native.kind in ("M", "m") and np.isnat(meta.fill_value): assert asdict_dict["fill_value"] == -9223372036854775808 + + +def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: + """Bare-int dims in rectilinear metadata (the spec's step-size shorthand, + produced when a uniform dimension of a mixed grid is collapsed) must pass + through the `chunks=` conversion unchanged. Wrapping one in a + single-element list turns "repeat to cover the axis" into "exactly one + chunk" and re-creation fails the sum-to-span check.""" + from zarr.core.metadata.v3 import RectilinearChunkGridMetadata + from zarr.storage import MemoryStore + from zarr.testing.strategies import chunks_param_from_rectilinear + + with zarr.config.set({"array.rectilinear_chunks": True}): + src = zarr.create_array( + MemoryStore(), shape=(3, 3), chunks=[[1, 2], [1, 1, 1]], dtype="uint8" + ) + grid = src.metadata.chunk_grid # type: ignore[union-attr] + assert isinstance(grid, RectilinearChunkGridMetadata) + assert grid.chunk_shapes == ((1, 2), 1) + dst = zarr.create_array( + MemoryStore(), + shape=src.shape, + chunks=chunks_param_from_rectilinear(grid), + dtype="uint8", + ) + assert dst.metadata.chunk_grid == grid # type: ignore[union-attr] From 535fa0ab948a68a373e2aa1e47a094d3c7bcdf93 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 26 Aug 2026 14:38:25 +0200 Subject: [PATCH 5/8] fix(chunk-grids): keep explicit per-chunk lists rectilinear MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the regular-grid collapse for explicit list input in normalize_chunks_1d: a per-chunk size list now always produces VaryingDimension, even when the sizes are uniform or uniform plus a short tail. Scalar specs (ints, numpy integers, and the -1 sentinel) still produce FixedDimension, which is the path that makes array creation O(1) in per-dimension chunk count; explicit lists are already O(n) in the input, so nothing is lost. The grid kind now follows the input syntax, matching 3.2.x behavior: previously an explicitly rectilinear spec whose edges happened to look regular was silently stored as RegularChunkGridMetadata, which changes resize semantics — a regular grid grows by extending the uniform pattern while a rectilinear grid appends an edge chunk, breaking append-oriented layouts like (168,) * 13 + (24,). This is resolution option 1 from gh-4272. Uniform dimensions of mixed scalar/list specs still serialize as the spec's bare-int step-size shorthand; explicit lists serialize as edge lists. The touched-chunk-keys assertions in the resize regression test follow the approach from gh-4290, which fixes the same issue on main via a requested_rectilinear flag. Fixes #4272 Co-authored-by: Shurong Cao Assisted-by: ClaudeCode:claude-fable-5 --- changes/4272.bugfix.md | 11 ++++ src/zarr/core/chunk_grids.py | 16 +++--- tests/test_chunk_grids.py | 32 ++++++------ tests/test_properties.py | 6 +-- tests/test_unified_chunk_grid.py | 90 ++++++++++++++++++++++++++++++++ 5 files changed, 127 insertions(+), 28 deletions(-) create mode 100644 changes/4272.bugfix.md diff --git a/changes/4272.bugfix.md b/changes/4272.bugfix.md new file mode 100644 index 0000000000..0df9c41960 --- /dev/null +++ b/changes/4272.bugfix.md @@ -0,0 +1,11 @@ +Explicit per-chunk size lists now always produce a rectilinear chunk grid, +even when the sizes happen to describe a regular grid (all equal, or all equal +with a smaller trailing chunk). Previously such input was silently collapsed to +a regular grid, which changed resize semantics: a regular grid grows by +extending the uniform pattern, while a rectilinear grid appends a new edge +chunk — the behavior an append-oriented layout like `(168,) * 13 + (24,)` +relies on. The grid kind now follows the input syntax, matching 3.2.x: +scalar chunk sizes (including numpy integers and the `-1` sentinel) produce a +regular grid, nested sequences produce a rectilinear grid. Rectilinear grids +remain gated behind `zarr.config.set({"array.rectilinear_chunks": True})`. +See #4174 for the accompanying O(1) chunk normalization change. diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index bd68f301eb..f23011cbdb 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -703,16 +703,18 @@ def _guess_regular_chunks( def normalize_chunks_1d(chunks: int | Iterable[object], span: int) -> DimensionGrid: """ Normalize a one-dimensional chunk specification into a dimension grid: - `FixedDimension` for uniform chunk sizes, `VaryingDimension` for explicit - per-chunk sizes that genuinely vary. Both variants bind the chunk sizes to + `FixedDimension` for scalar chunk sizes, `VaryingDimension` for explicit + per-chunk size lists. Both variants bind the chunk sizes to the span, and the uniform form is O(1) in the number of chunks — a dimension with `2**62` chunks must not materialize one entry per chunk. `-1` means "one chunk covering the entire span." - Explicit chunk size lists must sum to the span exactly; lists that describe - a regular grid (all sizes equal, or equal with a smaller boundary chunk) - collapse to `FixedDimension`. For uniform sizes the last chunk may overhang - the span. + Explicit chunk size lists must sum to the span exactly and always produce + `VaryingDimension`, even when the sizes happen to be uniform: the input + syntax declares the grid kind, so a per-chunk list is preserved as a + rectilinear dimension rather than silently collapsed to a regular one, + which would change how the dimension grows on resize. For scalar sizes + the last chunk may overhang the span. """ # `numbers.Integral` rather than `int` so that numpy integer scalars (which are not # `int` subclasses) take the uniform-chunk path instead of being treated as a sequence. @@ -750,8 +752,6 @@ def normalize_chunks_1d(chunks: int | Iterable[object], span: int) -> DimensionG raise ValueError(f"All chunk sizes must be positive, got {ints}") if sum(ints) != span: raise ValueError(f"Chunk sizes {ints} do not sum to span {span}") - if is_regular_1d(ints): - return FixedDimension(size=ints[0], extent=span) return VaryingDimension(ints, extent=span) diff --git a/tests/test_chunk_grids.py b/tests/test_chunk_grids.py index eec1a928b5..90ec67fe9d 100644 --- a/tests/test_chunk_grids.py +++ b/tests/test_chunk_grids.py @@ -73,12 +73,12 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None: ((30, -1, -1), (100, 20, 10), (30, 20, 10)), ((30, 20, -1), (100, 20, 10), (30, 20, 10)), ((30, 20, 10), (100, 20, 10), (30, 20, 10)), - # dask-style chunks describing a regular grid collapse to the uniform form - (((100, 100, 100), (50, 50)), (300, 100), (100, 50)), - (((100, 100, 50),), (250,), (100,)), - (((100,),), (100,), (100,)), - # genuinely irregular explicit chunks keep the per-chunk form - (((10, 20, 70), (50, 50)), (100, 100), ((10, 20, 70), 50)), + # dask-style explicit lists always keep the per-chunk rectilinear form, + # even when the sizes describe a regular grid (gh-4272) + (((100, 100, 100), (50, 50)), (300, 100), ((100, 100, 100), (50, 50))), + (((100, 100, 50),), (250,), ((100, 100, 50),)), + (((100,),), (100,), ((100,),)), + (((10, 20, 70), (50, 50)), (100, 100), ((10, 20, 70), (50, 50))), # no chunking (False means each dimension is one chunk spanning the full extent) (False, (100,), (100,)), (False, (100, 50), (100, 50)), @@ -96,7 +96,7 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None: ((np.int32(30), np.int64(-1)), (100, 20), (30, 20)), (np.array([10, 10]), (100, 100), (10, 10)), # rectilinear chunks given as numpy arrays - ((np.array([60, 40]), np.array([50, 50])), (100, 100), (60, 50)), + ((np.array([60, 40]), np.array([50, 50])), (100, 100), ((60, 40), (50, 50))), ], ) def test_normalize_chunks( @@ -112,9 +112,9 @@ def test_normalize_chunks( ((100,), (10,), None, (10,), None), # explicit regular shards ((100,), (10,), (50,), (50,), (10,)), - # rectilinear shards describing a regular-with-boundary grid collapse - ((100,), (10,), ((60, 40),), (60,), (10,)), - # genuinely irregular rectilinear shards keep the per-chunk form + # rectilinear shards keep the per-chunk form even when the sizes + # describe a regular-with-boundary grid (gh-4272) + ((100,), (10,), ((60, 40),), ((60, 40),), (10,)), ((100,), (10,), ((30, 60, 10),), ((30, 60, 10),), (10,)), # dict-style shards ((100, 100), (10, 10), {"shape": (50, 50)}, (50, 50), (10, 10)), @@ -291,18 +291,18 @@ def test_normalize_chunks_nd_errors(case: ExpectFail[tuple[Any, tuple[int, ...]] Expect(input=(-1, 100), output=FixedDimension(size=100, extent=100), id="full-span"), # zero-length span preserves the declared chunk size. Expect(input=(10, 0), output=FixedDimension(size=10, extent=0), id="uniform-zero-span"), - # explicit lists that describe a regular grid collapse to the uniform form. + # explicit lists always keep the per-chunk form, even when the sizes + # describe a regular grid (gh-4272). Expect( input=([10, 10, 10], 30), - output=FixedDimension(size=10, extent=30), + output=VaryingDimension([10, 10, 10], extent=30), id="explicit-regular", ), Expect( input=([10, 10, 5], 25), - output=FixedDimension(size=10, extent=25), + output=VaryingDimension([10, 10, 5], extent=25), id="explicit-boundary", ), - # genuinely irregular edges keep the explicit per-chunk form. Expect( input=([10, 20, 70], 100), output=VaryingDimension([10, 20, 70], extent=100), @@ -314,8 +314,8 @@ def test_normalize_chunks_nd_errors(case: ExpectFail[tuple[Any, tuple[int, ...]] def test_normalize_chunks_1d( case: Expect[tuple[Any, int], FixedDimension | VaryingDimension], ) -> None: - """Both output variants bind chunk sizes to the span: uniform specs become - `FixedDimension` (O(1) regardless of chunk count), irregular explicit + """Both output variants bind chunk sizes to the span: scalar specs become + `FixedDimension` (O(1) regardless of chunk count), explicit per-chunk lists become `VaryingDimension`.""" chunks, span = case.input assert normalize_chunks_1d(chunks, span) == case.output diff --git a/tests/test_properties.py b/tests/test_properties.py index a73a6d8a7c..651c61e264 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -452,7 +452,7 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: """Bare-int dims in rectilinear metadata (the spec's step-size shorthand, - produced when a uniform dimension of a mixed grid is collapsed) must pass + produced by a scalar dimension of a mixed chunk spec) must pass through the `chunks=` conversion unchanged. Wrapping one in a single-element list turns "repeat to cover the axis" into "exactly one chunk" and re-creation fails the sum-to-span check.""" @@ -461,9 +461,7 @@ def test_chunks_param_from_rectilinear_bare_int_roundtrip() -> None: from zarr.testing.strategies import chunks_param_from_rectilinear with zarr.config.set({"array.rectilinear_chunks": True}): - src = zarr.create_array( - MemoryStore(), shape=(3, 3), chunks=[[1, 2], [1, 1, 1]], dtype="uint8" - ) + src = zarr.create_array(MemoryStore(), shape=(3, 3), chunks=([1, 2], 1), dtype="uint8") grid = src.metadata.chunk_grid # type: ignore[union-attr] assert isinstance(grid, RectilinearChunkGridMetadata) assert grid.chunk_shapes == ((1, 2), 1) diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index f0b54519ab..50def1a81e 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -1027,6 +1027,96 @@ def test_e2e_chunk_grid_name_regular_from_dict(tmp_path: Path) -> None: assert chunk_grid_dict["name"] == "regular" +# --------------------------------------------------------------------------- +# Input syntax determines grid kind (gh-4272) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("shape", "chunks", "expected_type"), + [ + # scalar specs (including numpy integers and the -1 sentinel) produce a regular grid + ((30,), (10,), RegularChunkGridMetadata), + ((30,), 10, RegularChunkGridMetadata), + ((30,), (np.int64(10),), RegularChunkGridMetadata), + ((30,), -1, RegularChunkGridMetadata), + # explicit per-chunk lists stay rectilinear even when the sizes are uniform + ((30,), [[10, 10, 10]], RectilinearChunkGridMetadata), + # ... or uniform with a short tail (the append-oriented time-series layout of gh-4272) + ((168 * 13 + 24,), [[168] * 13 + [24]], RectilinearChunkGridMetadata), + # genuinely varying edges are rectilinear too + ((60,), [[10, 20, 30]], RectilinearChunkGridMetadata), + ], + ids=[ + "scalar-tuple", + "scalar-int", + "scalar-numpy-int", + "full-span-sentinel", + "explicit-uniform", + "explicit-uniform-short-tail", + "explicit-varying", + ], +) +def test_chunk_grid_kind_follows_input_syntax( + tmp_path: Path, shape: tuple[int, ...], chunks: Any, expected_type: type +) -> None: + """The stored grid kind is decided by the input syntax, not the chunk sizes: + scalar specs produce a regular grid, while explicit per-chunk lists produce + a rectilinear grid even when the sizes happen to describe a regular one + (gh-4272).""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + arr = zarr.create_array(store=tmp_path / "arr.zarr", shape=shape, chunks=chunks, dtype="int32") + assert isinstance(arr.metadata, ArrayV3Metadata) + assert isinstance(arr.metadata.chunk_grid, expected_type) + + +def test_mixed_scalar_and_list_dims_keep_shorthand(tmp_path: Path) -> None: + """In a mixed spec, a scalar dimension serializes as the bare-int step-size + shorthand while an explicit list keeps its per-chunk edges.""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + arr = zarr.create_array( + store=tmp_path / "arr.zarr", shape=(30, 100), chunks=(5, [10, 20, 70]), dtype="int32" + ) + assert isinstance(arr.metadata, ArrayV3Metadata) + grid = arr.metadata.chunk_grid + assert isinstance(grid, RectilinearChunkGridMetadata) + assert grid.chunk_shapes == (5, (10, 20, 70)) + + +def test_resize_uniform_rectilinear_appends_edge() -> None: + """Growing an explicitly rectilinear array whose edges look regular appends + a new edge chunk, while the same sizes declared as a scalar extend the + uniform pattern instead (gh-4272), so an append-only workload writing the + grown region touches exactly one new chunk.""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + rect_store: dict[str, Any] = {} + rect = zarr.create_array(store=rect_store, shape=(30,), chunks=[[10, 10, 10]], dtype="int32") + rect.resize((45,)) + assert isinstance(rect.metadata, ArrayV3Metadata) + rect_grid = rect.metadata.chunk_grid + assert isinstance(rect_grid, RectilinearChunkGridMetadata) + assert rect_grid.chunk_shapes == ((10, 10, 10, 15),) + assert rect.nchunks == 4 + # the appended region is exactly the new edge chunk + rect[30:45] = 1 + assert sorted(k for k in rect_store if k.startswith("c/")) == ["c/3"] + + reg_store: dict[str, Any] = {} + reg = zarr.create_array(store=reg_store, shape=(30,), chunks=(10,), dtype="int32") + reg.resize((45,)) + assert isinstance(reg.metadata, ArrayV3Metadata) + reg_grid = reg.metadata.chunk_grid + assert isinstance(reg_grid, RegularChunkGridMetadata) + assert reg_grid.chunk_shape == (10,) + assert reg.nchunks == 5 + # the same write straddles two chunks of the extended uniform pattern + reg[30:45] = 1 + assert sorted(k for k in reg_store if k.startswith("c/")) == ["c/3", "c/4"] + + # --------------------------------------------------------------------------- # Sharding compatibility tests # --------------------------------------------------------------------------- From ed26f627ff2e2d11906869870708f7df6a194533 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 26 Aug 2026 15:49:02 +0200 Subject: [PATCH 6/8] fix(array): guard chunks/shards access by stored grid kind, not runtime grid The stateful hypothesis tests in CI failed in from_array on a rectilinear source whose edges happen to be uniform (e.g. chunk_shapes=((3,),)): _parse_keep_array_attr and _info guarded data.chunks / data.shards with the runtime ChunkGrid.is_regular, but the runtime grid collapses uniform rectilinear dimensions to FixedDimension as an optimization, while metadata.chunks raises for any array whose stored grid is rectilinear. Before explicit per-chunk lists stopped collapsing at creation, the two notions of regularity always agreed for created arrays, so the mismatch was unreachable. Add _stored_chunk_grid_is_regular, which dispatches on the metadata kind, and use it at the three sites that read .chunks/.shards. A uniform-edged rectilinear source now round-trips through from_array with its grid kind intact. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 21 ++++++++++++++++++--- tests/test_unified_chunk_grid.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 9c4ec072eb..fbe152f2c7 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -125,6 +125,7 @@ ) from zarr.core.metadata.v3 import ( ChunkGridMetadata, + RegularChunkGridMetadata, create_chunk_grid_metadata, parse_node_type_array, ) @@ -1821,7 +1822,7 @@ async def info_complete(self) -> Any: def _info( self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None ) -> Any: - chunk_shape = self.chunks if self._chunk_grid.is_regular else None + chunk_shape = self.chunks if _stored_chunk_grid_is_regular(self.metadata) else None return ArrayInfo( _zarr_format=self.metadata.zarr_format, _data_type=self._zdtype, @@ -4764,6 +4765,20 @@ async def create_array( ) +def _stored_chunk_grid_is_regular(metadata: ArrayMetadata) -> bool: + """Whether the *stored* chunk grid is regular, i.e. `.chunks` and `.shards` + are defined. + + Distinct from the runtime ``ChunkGrid.is_regular``: the runtime grid + collapses a rectilinear dimension whose edges happen to be uniform to a + ``FixedDimension`` as an optimization, so it can report regular for an + array whose stored metadata — and therefore `.chunks` — is rectilinear. + """ + if isinstance(metadata, ArrayV2Metadata): + return True + return isinstance(metadata.chunk_grid, RegularChunkGridMetadata) + + def _parse_keep_array_attr( data: AnyArray | npt.ArrayLike, chunks: ChunksLike | Literal["auto", "keep"], @@ -4792,12 +4807,12 @@ def _parse_keep_array_attr( ]: if isinstance(data, Array): if chunks == "keep": - if data._chunk_grid.is_regular: + if _stored_chunk_grid_is_regular(data.metadata): chunks = data.chunks else: chunks = data.write_chunk_sizes if shards == "keep": - shards = data.shards if data._chunk_grid.is_regular else None + shards = data.shards if _stored_chunk_grid_is_regular(data.metadata) else None if zarr_format is None: zarr_format = data.metadata.zarr_format if filters == "keep": diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 50def1a81e..591d262b7e 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -1085,6 +1085,24 @@ def test_mixed_scalar_and_list_dims_keep_shorthand(tmp_path: Path) -> None: assert grid.chunk_shapes == (5, (10, 20, 70)) +def test_from_array_keeps_uniform_rectilinear_grid() -> None: + """`from_array` on a rectilinear source whose edges happen to be uniform + keeps the rectilinear grid kind instead of raising. The runtime grid + collapses uniform edges to `FixedDimension` as an optimization, but + `.chunks` is only defined by the stored metadata kind, so the "keep" + logic must dispatch on the metadata, not the runtime grid (gh-4272).""" + from zarr.core.metadata.v3 import ArrayV3Metadata + + src = zarr.create_array(MemoryStore(), shape=(3,), chunks=[[3]], dtype="uint8") + src[:] = 1 + assert src.info is not None # _info must not raise either + dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=False) + assert isinstance(dst.metadata, ArrayV3Metadata) + assert isinstance(src.metadata, ArrayV3Metadata) + assert isinstance(dst.metadata.chunk_grid, RectilinearChunkGridMetadata) + assert dst.metadata.chunk_grid == src.metadata.chunk_grid + + def test_resize_uniform_rectilinear_appends_edge() -> None: """Growing an explicitly rectilinear array whose edges look regular appends a new edge chunk, while the same sizes declared as a scalar extend the From 79873b87688bd085193d7b4d56fcc2d2edf581de Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 29 Aug 2026 00:04:43 +0200 Subject: [PATCH 7/8] fix(array): make from_array 'keep' faithful for every stored chunk grid Fixes the three failures maxrjones identified in review of #4218, whose shared root cause was that from_array is really create_array plus an inverse mapping from stored grid metadata back into the chunks=/shards= parameter space, and that inverse mapping only handled regular grids: - The write_data=True copy path iterated shard regions via the .chunks/.shards accessors, which raise for rectilinear grids. _iter_shard_regions (and _shard_grid_shape) now iterate the stored chunk grid directly, which describes the write regions for every grid kind. - chunks='keep' expanded uniform dimensions of a rectilinear grid to one entry per chunk via write_chunk_sizes, resurrecting the O(nchunks) path and changing resize semantics. The stored chunk_shapes now pass through as-is, preserving the bare-int shorthand in O(ndim). - shards='keep' silently dropped sharding when the shard grid was rectilinear, and .info raised on such arrays through the .shards accessor. The shard grid's chunk_shapes now round-trip through the shards= parameter, and _info guards both accessors by stored grid kind. metadata.chunks now returns the inner chunk shape for any sharded array (inner chunks are always regular, whatever the shard grid), and the repeated sharding-codec isinstance dance is extracted into ArrayV3Metadata.sharding_codec per the existing TODO. ShardsLike admits mixed bare-int/sequence specs, matching what the normalizer accepts. Tests: the ad-hoc from_array regression test is replaced by a spec matrix (FROM_ARRAY_KEEP_CASES) covering every kind of chunk/shard spec create_array accepts, round-tripped through from_array at both write_data settings with grid, codec, and data equality asserted, plus an O(1) pin at shape=(2**62, 30) and error tests for the two accessors that still raise. The stateful hypothesis machine drops its write_data=False workaround, restoring property coverage of the default copy path. Assisted-by: ClaudeCode:claude-fable-5 --- changes/4272.bugfix.md | 9 +++ src/zarr/core/array.py | 127 +++++++++++++++++++------------ src/zarr/core/metadata/v3.py | 29 ++++--- src/zarr/testing/stateful.py | 8 +- tests/test_unified_chunk_grid.py | 76 +++++++++++++++--- 5 files changed, 170 insertions(+), 79 deletions(-) diff --git a/changes/4272.bugfix.md b/changes/4272.bugfix.md index 0df9c41960..bb303e3063 100644 --- a/changes/4272.bugfix.md +++ b/changes/4272.bugfix.md @@ -9,3 +9,12 @@ scalar chunk sizes (including numpy integers and the `-1` sentinel) produce a regular grid, nested sequences produce a rectilinear grid. Rectilinear grids remain gated behind `zarr.config.set({"array.rectilinear_chunks": True})`. See #4174 for the accompanying O(1) chunk normalization change. + +`zarr.from_array` with the default `chunks="keep"` / `shards="keep"` now +reproduces the source's stored grid exactly: a rectilinear grid is passed +through in O(number of dimensions), with uniform dimensions keeping their +bare-int shorthand; sharding under a rectilinear shard grid is preserved +instead of being silently dropped; and the default `write_data=True` copy +works for every grid kind. `Array.info` no longer raises for sharded arrays +with a rectilinear shard grid, and `Array.chunks` is now defined for any +sharded array (the inner chunks of a shard are always regular). diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index fbe152f2c7..c7e429f262 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -100,7 +100,6 @@ Selection, VIndex, _iter_grid, - _iter_regions, check_fields, check_no_multi_fields, is_pure_fancy_indexing, @@ -125,7 +124,7 @@ ) from zarr.core.metadata.v3 import ( ChunkGridMetadata, - RegularChunkGridMetadata, + RectilinearChunkGridMetadata, create_chunk_grid_metadata, parse_node_type_array, ) @@ -154,7 +153,7 @@ from zarr.abc.codec import CodecPipeline from zarr.abc.store import Store - from zarr.codecs.sharding import IndexLocation + from zarr.codecs.sharding import IndexLocation, ShardingCodec from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -848,10 +847,12 @@ def shape(self) -> tuple[int, ...]: @property def chunks(self) -> tuple[int, ...]: """Returns the chunk shape of the Array. - If sharding is used the inner chunk shape is returned. + If sharding is used the inner chunk shape is returned, which is defined + for any chunk grid (inner chunks are always regular). - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Otherwise, only defined for arrays using a regular chunk grid: for a + non-sharded array with a rectilinear chunk grid, `NotImplementedError` + is raised. Use `read_chunk_sizes` for general access. Returns ------- @@ -923,8 +924,9 @@ def shards(self) -> tuple[int, ...] | None: """Returns the shard shape of the Array. Returns None if sharding is not used. - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Only defined when the shard grid is regular: for a sharded array with a + rectilinear chunk grid, `NotImplementedError` is raised. Use + `write_chunk_sizes` for general access. Returns ------- @@ -1121,14 +1123,9 @@ def _chunk_grid_shape(self) -> tuple[int, ...]: tuple[int, ...] The number of chunks along each dimension. """ - # TODO: refactor — extract a sharding_codec property on ArrayV3Metadata - # to replace the repeated `len == 1 and isinstance` pattern. - from zarr.codecs.sharding import ShardingCodec - - codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) - if len(codecs) == 1 and isinstance(codecs[0], ShardingCodec): + if (sharding_codec := _sharding_codec(self.metadata)) is not None: # When sharding, count inner chunks across the whole array - chunk_shape = codecs[0].chunk_shape + chunk_shape = sharding_codec.chunk_shape return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True))) return self._chunk_grid.grid_shape @@ -1144,11 +1141,9 @@ def _shard_grid_shape(self) -> tuple[int, ...]: tuple[int, ...] The shape of the shard grid for this array. """ - if self.shards is None: - shard_shape = self.chunks - else: - shard_shape = self.shards - return tuple(starmap(ceildiv, zip(self.shape, shard_shape, strict=True))) + # The stored chunk grid is the shard grid when sharding is used, the + # chunk grid otherwise. Works for regular and rectilinear grids alike. + return self._chunk_grid.grid_shape @property def nchunks(self) -> int: @@ -1822,14 +1817,20 @@ async def info_complete(self) -> Any: def _info( self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None ) -> Any: - chunk_shape = self.chunks if _stored_chunk_grid_is_regular(self.metadata) else None + rectilinear_grid = _stored_rectilinear_grid(self.metadata) + sharded = _sharding_codec(self.metadata) is not None + # `.chunks` (the inner chunk shape when sharded) is undefined only for a + # non-sharded rectilinear grid; `.shards` is undefined for a rectilinear + # shard grid. ArrayInfo renders None as "" / omits the line. + chunk_shape = self.chunks if (rectilinear_grid is None or sharded) else None + shard_shape = self.shards if rectilinear_grid is None else None return ArrayInfo( _zarr_format=self.metadata.zarr_format, _data_type=self._zdtype, _fill_value=self.metadata.fill_value, _shape=self.shape, _order=self.order, - _shard_shape=self.shards, + _shard_shape=shard_shape, _chunk_shape=chunk_shape, _read_only=self.read_only, _compressors=self.compressors, @@ -2023,10 +2024,12 @@ def shape(self, value: tuple[int, ...]) -> None: @property def chunks(self) -> tuple[int, ...]: """Returns a tuple of integers describing the length of each dimension of a chunk of the array. - If sharding is used the inner chunk shape is returned. + If sharding is used the inner chunk shape is returned, which is defined + for any chunk grid (inner chunks are always regular). - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Otherwise, only defined for arrays using a regular chunk grid: for a + non-sharded array with a rectilinear chunk grid, `NotImplementedError` + is raised. Use `read_chunk_sizes` for general access. Returns ------- @@ -2090,8 +2093,9 @@ def shards(self) -> tuple[int, ...] | None: """Returns a tuple of integers describing the length of each dimension of a shard of the array. Returns None if sharding is not used. - Only defined for arrays using a regular chunk grid. - If array uses a rectilinear chunk grid, `NotImplementedError` is raised. + Only defined when the shard grid is regular: for a sharded array with a + rectilinear chunk grid, `NotImplementedError` is raised. Use + `write_chunk_sizes` for general access. Returns ------- @@ -4062,7 +4066,9 @@ class ShardsConfigParam(TypedDict): index_location: IndexLocation | None -type ShardsLike = tuple[int, ...] | Sequence[Sequence[int]] | ShardsConfigParam | Literal["auto"] +type ShardsLike = ( + tuple[int, ...] | Sequence[int | Sequence[int]] | ShardsConfigParam | Literal["auto"] +) async def from_array( @@ -4765,18 +4771,31 @@ async def create_array( ) -def _stored_chunk_grid_is_regular(metadata: ArrayMetadata) -> bool: - """Whether the *stored* chunk grid is regular, i.e. `.chunks` and `.shards` - are defined. +def _sharding_codec(metadata: ArrayMetadata) -> ShardingCodec | None: + """The array's sharding codec, or None if the array is not sharded. - Distinct from the runtime ``ChunkGrid.is_regular``: the runtime grid - collapses a rectilinear dimension whose edges happen to be uniform to a - ``FixedDimension`` as an optimization, so it can report regular for an - array whose stored metadata — and therefore `.chunks` — is rectilinear. + Zarr format 2 arrays are never sharded. """ - if isinstance(metadata, ArrayV2Metadata): - return True - return isinstance(metadata.chunk_grid, RegularChunkGridMetadata) + if isinstance(metadata, ArrayV3Metadata): + return metadata.sharding_codec + return None + + +def _stored_rectilinear_grid(metadata: ArrayMetadata) -> RectilinearChunkGridMetadata | None: + """The *stored* rectilinear chunk grid, or None if the stored grid is regular + (in which case `.chunks` and `.shards` are defined). + + Dispatches on the stored metadata, not the runtime ``ChunkGrid``: the + runtime grid collapses a rectilinear dimension whose edges happen to be + uniform to a ``FixedDimension`` as an optimization, so it can report regular + for an array whose stored metadata — and therefore `.chunks` — is + rectilinear. Zarr format 2 grids are always regular. + """ + if isinstance(metadata, ArrayV3Metadata) and isinstance( + metadata.chunk_grid, RectilinearChunkGridMetadata + ): + return metadata.chunk_grid + return None def _parse_keep_array_attr( @@ -4806,13 +4825,27 @@ def _parse_keep_array_attr( dict[str, JSON] | None, ]: if isinstance(data, Array): + rectilinear_grid = _stored_rectilinear_grid(data.metadata) + sharded = _sharding_codec(data.metadata) is not None if chunks == "keep": - if _stored_chunk_grid_is_regular(data.metadata): + if rectilinear_grid is None or sharded: + # `.chunks` is the inner chunk shape when sharding is used, and + # inner chunks are regular whatever the shape of the shard grid. chunks = data.chunks else: - chunks = data.write_chunk_sizes + # Pass the stored spec through as-is: it is O(ndim), not + # O(nchunks), and keeps the bare-int shorthand of uniform + # dimensions (which also preserves resize semantics). + chunks = rectilinear_grid.chunk_shapes if shards == "keep": - shards = data.shards if _stored_chunk_grid_is_regular(data.metadata) else None + if rectilinear_grid is None: + shards = data.shards + elif sharded: + # The stored grid is the shard grid; its spec round-trips + # through the `shards=` parameter of array creation. + shards = rectilinear_grid.chunk_shapes + else: + shards = None if zarr_format is None: zarr_format = data.metadata.zarr_format if filters == "keep": @@ -5285,14 +5318,10 @@ def _iter_shard_regions( A tuple of slice objects representing the region spanned by each shard in the selection or chunk when no shards are present. """ - if array.shards is None: - shard_shape = array.chunks - else: - shard_shape = array.shards - - return _iter_regions( - array.shape, shard_shape, origin=origin, selection_shape=selection_shape, trim_excess=True - ) + # The stored chunk grid always describes the write regions: the shard grid + # when sharding is used, the chunk grid otherwise. Iterating it directly + # works for regular and rectilinear grids alike. + return array._chunk_grid.iter_chunk_regions(origin=origin, selection_shape=selection_shape) def _iter_chunk_regions( diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index fe63b11d33..2ab3ff56fb 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -42,6 +42,7 @@ if TYPE_CHECKING: from typing import Self + from zarr.codecs.sharding import ShardingCodec from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_grids import ChunkGrid from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar @@ -569,26 +570,32 @@ def dtype(self) -> ZDType[TBaseDType, TBaseScalar]: # TODO: move these properties to the Array class. # They require knowledge of codecs (ShardingCodec) and don't belong on a metadata DTO. + @property + def sharding_codec(self) -> ShardingCodec | None: + """The array's sharding codec, or None if the array is not sharded.""" + from zarr.codecs.sharding import ShardingCodec + + if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): + return self.codecs[0] + return None + @property def chunks(self) -> tuple[int, ...]: + if (sharding_codec := self.sharding_codec) is not None: + # Inner chunks are always regular, whatever the shape of the outer + # (shard) grid. + return sharding_codec.chunk_shape if not isinstance(self.chunk_grid, RegularChunkGridMetadata): msg = ( "The `chunks` attribute is only defined for arrays using regular chunk grids. " "This array has a rectilinear chunk grid. Use `read_chunk_sizes` for general access." ) raise NotImplementedError(msg) - - from zarr.codecs.sharding import ShardingCodec - - if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): - return self.codecs[0].chunk_shape return self.chunk_grid.chunk_shape @property def shards(self) -> tuple[int, ...] | None: - from zarr.codecs.sharding import ShardingCodec - - if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): + if self.sharding_codec is not None: if not isinstance(self.chunk_grid, RegularChunkGridMetadata): msg = ( "The `shards` attribute is only defined for arrays using regular chunk grids. " @@ -600,10 +607,8 @@ def shards(self) -> tuple[int, ...] | None: @property def inner_codecs(self) -> tuple[Codec, ...]: - from zarr.codecs.sharding import ShardingCodec - - if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec): - return self.codecs[0].codecs + if (sharding_codec := self.sharding_codec) is not None: + return sharding_codec.codecs return self.codecs def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: diff --git a/src/zarr/testing/stateful.py b/src/zarr/testing/stateful.py index 652e90e60a..59c5b72e62 100644 --- a/src/zarr/testing/stateful.py +++ b/src/zarr/testing/stateful.py @@ -152,16 +152,12 @@ def add_array(self, data: DataObject, name: str) -> None: ) note(f"Adding array: path='{path}' shape={a.shape} chunks={a.metadata.chunk_grid}") - # Recreate the same array in the store under test. - # The data is copied here rather than by `write_data=True`, - # whose shard-wise copy does not support rectilinear chunk grids. - arr = zarr.from_array( + # Recreate the same array, including its data, in the store under test. + zarr.from_array( self.store, data=a, name=path, - write_data=False, ) - arr[:] = a[:] self.all_arrays.add(path) @rule() diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 591d262b7e..b44b5d9177 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -1085,22 +1085,74 @@ def test_mixed_scalar_and_list_dims_keep_shorthand(tmp_path: Path) -> None: assert grid.chunk_shapes == (5, (10, 20, 70)) -def test_from_array_keeps_uniform_rectilinear_grid() -> None: - """`from_array` on a rectilinear source whose edges happen to be uniform - keeps the rectilinear grid kind instead of raising. The runtime grid - collapses uniform edges to `FixedDimension` as an optimization, but - `.chunks` is only defined by the stored metadata kind, so the "keep" - logic must dispatch on the metadata, not the runtime grid (gh-4272).""" +# The chunk/shard spec matrix for `from_array` round-trip tests. Every kind of +# spec `create_array` accepts should appear here, so that any change to how +# specs are normalized or stored is automatically checked against the "keep" +# path of `from_array` as well. +FROM_ARRAY_KEEP_CASES = [ + pytest.param((30, 40), (10, 20), None, id="regular"), + pytest.param((40, 40), (5, 10), (10, 20), id="regular-sharded"), + pytest.param((60, 100), [[10, 20, 30], [50, 50]], None, id="rectilinear"), + pytest.param((9, 30), (3, [10, 20]), None, id="rectilinear-mixed-bare-int"), + pytest.param((3,), [[3]], None, id="rectilinear-uniform-list"), + pytest.param((100,), (10,), [[50, 50]], id="rectilinear-sharded"), +] + + +@pytest.mark.parametrize("write_data", [True, False]) +@pytest.mark.parametrize(("shape", "chunks", "shards"), FROM_ARRAY_KEEP_CASES) +def test_from_array_keep_roundtrips_chunk_grid( + shape: tuple[int, ...], chunks: Any, shards: Any, write_data: bool +) -> None: + """For every chunk/shard spec accepted by `create_array`, `from_array` with + the default "keep" parameters reproduces the source's stored chunk grid and + codecs exactly: the grid kind is preserved (gh-4272), uniform dimensions + keep their bare-int shorthand instead of being expanded per chunk, and + sharding survives — including under a rectilinear shard grid. `.info` is + defined for every kind of source, and the default data copy works for + every kind of grid.""" from zarr.core.metadata.v3 import ArrayV3Metadata - src = zarr.create_array(MemoryStore(), shape=(3,), chunks=[[3]], dtype="uint8") - src[:] = 1 - assert src.info is not None # _info must not raise either - dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=False) - assert isinstance(dst.metadata, ArrayV3Metadata) + src = zarr.create_array( + MemoryStore(), shape=shape, chunks=chunks, shards=shards, dtype="uint16" + ) + src[:] = np.arange(np.prod(shape), dtype="uint16").reshape(shape) + assert src.info is not None + dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=write_data) assert isinstance(src.metadata, ArrayV3Metadata) - assert isinstance(dst.metadata.chunk_grid, RectilinearChunkGridMetadata) + assert isinstance(dst.metadata, ArrayV3Metadata) assert dst.metadata.chunk_grid == src.metadata.chunk_grid + assert dst.metadata.codecs == src.metadata.codecs + if write_data: + np.testing.assert_array_equal(dst[:], src[:]) + + +def test_chunks_raises_for_nonsharded_rectilinear_grid() -> None: + """`.chunks` has no uniform value for a non-sharded rectilinear grid.""" + arr = zarr.create_array(MemoryStore(), shape=(30,), chunks=[[10, 20]], dtype="int32") + with pytest.raises(NotImplementedError, match="regular chunk grids"): + _ = arr.chunks + + +def test_shards_raises_for_rectilinear_shard_grid() -> None: + """`.shards` has no uniform value when the shard grid is rectilinear.""" + arr = zarr.create_array( + MemoryStore(), shape=(100,), chunks=(10,), shards=[[50, 50]], dtype="int32" + ) + with pytest.raises(NotImplementedError, match="regular chunk grids"): + _ = arr.shards + + +def test_from_array_keep_is_o1_in_chunk_count() -> None: + """`chunks="keep"` passes uniform dimensions through as bare-int shorthand + rather than expanding one entry per chunk, so copying array metadata is + O(ndim), not O(nchunks): this completes instantly despite ~2**60 chunks + along the first dimension (and would hang if the shorthand were expanded).""" + src = zarr.create_array(MemoryStore(), shape=(2**62, 30), chunks=(3, [10, 20]), dtype="uint8") + dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=False) + grid = dst.metadata.chunk_grid + assert isinstance(grid, RectilinearChunkGridMetadata) + assert grid.chunk_shapes == (3, (10, 20)) def test_resize_uniform_rectilinear_appends_edge() -> None: From ad3a9226fd82c10b79881e1d244fa2431f6c620a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 29 Aug 2026 10:04:38 +0200 Subject: [PATCH 8/8] fix(array): coherent introspection for rectilinear-sharded arrays Follow-up to the from_array 'keep' fixes, addressing the remaining inconsistencies in rectilinear-sharded introspection: - .info / info_complete rendered such arrays as unsharded: no shard line, and the stored-shard count labeled 'Chunks Initialized'. ArrayInfo._shard_shape now admits a '' sentinel, so the shard shape renders as and the count is labeled 'Shards Initialized'. - nchunks_initialized raised through the .shards accessor. It now sums the chunk count of each initialized shard individually (shard sizes vary under a rectilinear shard grid; the per-shard sum reduces to the old uniform multiplier for regular shard grids). - The user guide still said .chunks is only available for regular grids; it now documents that sharded arrays always have .chunks, and the rectilinear-shard section shows .chunks/.write_chunk_sizes/.info for such arrays. The round-trip matrix gains an initialized-count assertion, plus tests for the info rendering and the per-shard count. Assisted-by: ClaudeCode:claude-fable-5 --- changes/4272.bugfix.md | 8 +++-- docs/user-guide/arrays.md | 18 ++++++++-- src/zarr/core/_info.py | 4 ++- src/zarr/core/array.py | 60 ++++++++++++++++++++------------ tests/test_unified_chunk_grid.py | 31 +++++++++++++++++ 5 files changed, 93 insertions(+), 28 deletions(-) diff --git a/changes/4272.bugfix.md b/changes/4272.bugfix.md index bb303e3063..417128d092 100644 --- a/changes/4272.bugfix.md +++ b/changes/4272.bugfix.md @@ -15,6 +15,8 @@ reproduces the source's stored grid exactly: a rectilinear grid is passed through in O(number of dimensions), with uniform dimensions keeping their bare-int shorthand; sharding under a rectilinear shard grid is preserved instead of being silently dropped; and the default `write_data=True` copy -works for every grid kind. `Array.info` no longer raises for sharded arrays -with a rectilinear shard grid, and `Array.chunks` is now defined for any -sharded array (the inner chunks of a shard are always regular). +works for every grid kind. `Array.chunks` is now defined for any sharded array +(the inner chunks of a shard are always regular), and for sharded arrays with +a rectilinear shard grid `Array.info` no longer raises — it reports the shard +shape as `` — while `Array.nchunks_initialized` counts the chunks of +each initialized shard individually instead of raising. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index a192845f9e..6e31b2e9f9 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -674,8 +674,11 @@ z_regular = zarr.create_array( print(z_regular.write_chunk_sizes) ``` -Note that the `.chunks` property is only available for regular chunk grids. For -rectilinear arrays, use `.write_chunk_sizes` (or `.read_chunk_sizes`) instead. +Note that the `.chunks` property is not available for non-sharded rectilinear +arrays, since there is no single uniform chunk shape — use `.write_chunk_sizes` +(or `.read_chunk_sizes`) instead. Sharded arrays always have `.chunks`: it +returns the inner chunk shape, which is regular even when the shard grid is +rectilinear (see [Rectilinear shard boundaries](#rectilinear-shard-boundaries)). ### Resizing and appending @@ -746,6 +749,17 @@ print(z[50:70, 40:60]) Note that rectilinear inner chunks with sharding are not supported — only the shard boundaries can be rectilinear. +For such arrays, `.chunks` returns the (regular) inner chunk shape, while +`.shards` raises `NotImplementedError` since there is no single uniform shard +shape — use `.write_chunk_sizes` for the per-dimension shard sizes. `.info` +reports the shard shape as ``: + +```python exec="true" session="arrays" source="above" result="ansi" +print(f"chunks={z.chunks}") +print(f"shard sizes={z.write_chunk_sizes}") +print(z.info) +``` + ### Metadata format Rectilinear chunk grid metadata uses run-length encoding (RLE) for compact diff --git a/src/zarr/core/_info.py b/src/zarr/core/_info.py index 1503f05b26..58dcbdefb3 100644 --- a/src/zarr/core/_info.py +++ b/src/zarr/core/_info.py @@ -82,7 +82,9 @@ class ArrayInfo: _data_type: ZDType[TBaseDType, TBaseScalar] _fill_value: object _shape: tuple[int, ...] - _shard_shape: tuple[int, ...] | None = None + # "" marks a sharded array whose shard grid is rectilinear, so + # there is no uniform shard shape; None means the array is not sharded. + _shard_shape: tuple[int, ...] | Literal[""] | None = None _chunk_shape: tuple[int, ...] | None = None _order: Literal["C", "F"] _read_only: bool diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index c7e429f262..80b356ad81 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -1207,10 +1207,10 @@ async def nchunks_initialized(self) -> int: """ Calculate the number of chunks that have been initialized in storage. - This value is calculated as the product of the number of initialized shards and the number - of chunks per shard. For arrays that do not use sharding, the number of chunks per shard is - effectively 1, and in that case the number of chunks initialized is the same as the number - of stored objects associated with an array. + This value is calculated as the sum of the number of chunks in every initialized shard + (shard sizes can vary when the shard grid is rectilinear). For arrays that do not use + sharding, each stored object holds one chunk, so the number of chunks initialized is the + same as the number of stored objects associated with an array. Returns ------- @@ -1820,10 +1820,17 @@ def _info( rectilinear_grid = _stored_rectilinear_grid(self.metadata) sharded = _sharding_codec(self.metadata) is not None # `.chunks` (the inner chunk shape when sharded) is undefined only for a - # non-sharded rectilinear grid; `.shards` is undefined for a rectilinear - # shard grid. ArrayInfo renders None as "" / omits the line. + # non-sharded rectilinear grid, which ArrayInfo renders as ""; + # `.shards` is undefined for a rectilinear shard grid, where the + # "" sentinel keeps the array rendered as sharded. chunk_shape = self.chunks if (rectilinear_grid is None or sharded) else None - shard_shape = self.shards if rectilinear_grid is None else None + shard_shape: tuple[int, ...] | Literal[""] | None + if rectilinear_grid is None: + shard_shape = self.shards + elif sharded: + shard_shape = "" + else: + shard_shape = None return ArrayInfo( _zarr_format=self.metadata.zarr_format, _data_type=self._zdtype, @@ -2301,10 +2308,10 @@ def nchunks_initialized(self) -> int: """ Calculate the number of chunks that have been initialized in storage. - This value is calculated as the product of the number of initialized shards and the number of - chunks per shard. For arrays that do not use sharding, the number of chunks per shard is effectively 1, - and in that case the number of chunks initialized is the same as the number of stored objects associated with an - array. For a direct count of the number of initialized stored objects, see `nshards_initialized`. + This value is calculated as the sum of the number of chunks in every initialized shard + (shard sizes can vary when the shard grid is rectilinear). For arrays that do not use sharding, + each stored object holds one chunk, so the number of chunks initialized is the same as the number + of stored objects associated with an array. For a direct count of the number of initialized stored objects, see `nshards_initialized`. Returns ------- @@ -5359,10 +5366,10 @@ async def _nchunks_initialized( """ Calculate the number of chunks that have been initialized in storage. - This value is calculated as the product of the number of initialized shards and the number - of chunks per shard. For arrays that do not use sharding, the number of chunks per shard is - effectively 1, and in that case the number of chunks initialized is the same as the number - of stored objects associated with an array. + This value is calculated as the sum of the number of chunks in every initialized shard + (shard sizes can vary when the shard grid is rectilinear). For arrays that do not use + sharding, each stored object holds one chunk, so the number of chunks initialized is the + same as the number of stored objects associated with an array. Parameters ---------- @@ -5374,13 +5381,22 @@ async def _nchunks_initialized( nchunks_initialized : int The number of chunks that have been initialized. """ - if array.shards is None: - chunks_per_shard = 1 - else: - chunks_per_shard = product( - tuple(a // b for a, b in zip(array.shards, array.chunks, strict=True)) - ) - return (await _nshards_initialized(array)) * chunks_per_shard + if _sharding_codec(array.metadata) is None: + return await _nshards_initialized(array) + # Count the inner chunks of each initialized shard individually: shard + # sizes vary when the shard grid is rectilinear. `.chunks` is the inner + # chunk shape, which is defined for any sharded array. + inner_chunks = array.chunks + grid = array._chunk_grid + initialized = set(await _shards_initialized(array)) + total = 0 + for coords in grid.all_chunk_coords(): + spec = grid[coords] + if spec is not None and array.metadata.encode_chunk_key(coords) in initialized: + total += product( + tuple(s // c for s, c in zip(spec.codec_shape, inner_chunks, strict=True)) + ) + return total async def _nshards_initialized( diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index b44b5d9177..963bd95973 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -1118,6 +1118,9 @@ def test_from_array_keep_roundtrips_chunk_grid( ) src[:] = np.arange(np.prod(shape), dtype="uint16").reshape(shape) assert src.info is not None + # The array is fully written, so every chunk is initialized, whatever the + # grid kind and sharding layout. + assert src.nchunks_initialized == src.nchunks dst = zarr.from_array(MemoryStore(), data=src, name="0", write_data=write_data) assert isinstance(src.metadata, ArrayV3Metadata) assert isinstance(dst.metadata, ArrayV3Metadata) @@ -1127,6 +1130,34 @@ def test_from_array_keep_roundtrips_chunk_grid( np.testing.assert_array_equal(dst[:], src[:]) +def test_info_reports_variable_shard_shape_for_rectilinear_shard_grid() -> None: + """A sharded array whose shard grid is rectilinear has no uniform shard + shape. `.info` must still present it as sharded — shard shape rendered as + ``, and the initialized count labeled as shards, not chunks.""" + arr = zarr.create_array( + MemoryStore(), shape=(100,), chunks=(10,), shards=[[50, 50]], dtype="int32" + ) + arr[:] = 1 + assert "Shard shape : " in repr(arr.info) + info = arr.info_complete() + assert "Shards Initialized : 2" in repr(info) + assert arr.nchunks_initialized == 10 + + +def test_nchunks_initialized_counts_per_shard_for_rectilinear_shard_grid() -> None: + """With a rectilinear shard grid the chunk count per shard varies, so the + initialized-chunk count is summed per stored shard rather than derived + from a uniform multiplier.""" + arr = zarr.create_array( + MemoryStore(), shape=(90,), chunks=(10,), shards=[[60, 30]], dtype="int32" + ) + assert arr.nchunks_initialized == 0 + arr[60:90] = 1 # initializes only the second shard, which holds 3 chunks + assert arr.nchunks_initialized == 3 + arr[0:10] = 1 # initializes the first shard, which holds 6 chunks + assert arr.nchunks_initialized == 9 + + def test_chunks_raises_for_nonsharded_rectilinear_grid() -> None: """`.chunks` has no uniform value for a non-sharded rectilinear grid.""" arr = zarr.create_array(MemoryStore(), shape=(30,), chunks=[[10, 20]], dtype="int32")