diff --git a/changes/202.bugfix.md b/changes/202.bugfix.md new file mode 100644 index 0000000000..9c9bd40f21 --- /dev/null +++ b/changes/202.bugfix.md @@ -0,0 +1 @@ +Stop emitting an `UnstableSpecificationWarning` when serializing the `struct` data type to Zarr V3 metadata. The `struct` data type now has a stable Zarr V3 specification. The legacy `structured` alias and the unspecified `null_terminated_bytes`, `raw_bytes`, and `variable_length_bytes` data types continue to warn. diff --git a/changes/251.feature.md b/changes/251.feature.md new file mode 100644 index 0000000000..82143d742b --- /dev/null +++ b/changes/251.feature.md @@ -0,0 +1,10 @@ +The sharding codec no longer requires the inner `chunk_shape` to evenly divide +the shard shape. Inner chunks are laid out on a semi-regular grid: they are +spaced at `chunk_shape` intervals and clipped by the shard boundary, so the +number of inner chunks per dimension is the ceiling division of the shard shape +by `chunk_shape`, and clipped chunks are stored at their clipped shape. This +implements version 1.1 of the `sharding_indexed` specification (see +[zarr-specs #370](https://github.com/zarr-developers/zarr-specs/pull/370)) and +works with regular and rectilinear chunk grids, including nested sharding. +Note that arrays written with a non-divisible `chunk_shape` are not readable by +implementations that only support version 1.0 of the sharding specification. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index 51b2fa1a17..23774f74c6 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -594,6 +594,14 @@ In this example a shard shape of (1000, 1000) and a chunk shape of (100, 100) is This means that `10*10` chunks are stored in each shard, and there are `10*10` shards in total. Without the `shards` argument, there would be 10,000 chunks stored as individual files. +The chunk shape does not need to evenly divide the shard shape. Chunks are laid +out on a regular grid within each shard and clipped by the shard boundary, so a +shard shape of `(1000,)` with a chunk shape of `(300,)` stores four chunks per +shard with sizes `300, 300, 300, 100`. Note that arrays that rely on this +clipping (version 1.1 of the `sharding_indexed` specification) cannot be read by +implementations that only support version 1.0, which requires the chunk shape +to evenly divide the shard shape. + ## Rectilinear (variable) chunk grids !!! warning "Experimental" @@ -728,8 +736,9 @@ print("Roundtrip OK") Rectilinear chunk grids can also be used for shard boundaries when combined with sharding. In this case, the outer grid (shards) is rectilinear while the -inner chunks remain regular. Each shard dimension must be divisible by the -corresponding inner chunk size: +inner chunks remain regular. The inner chunk size does not need to divide the +shard sizes evenly: inner chunks are clipped by each shard's boundary, so every +shard holds a semi-regular grid of inner chunks: ```python exec="true" session="arrays" source="above" result="ansi" z = zarr.create_array( diff --git a/packages/zarr-metadata/src/zarr_metadata/v3/codec/sharding_indexed.py b/packages/zarr-metadata/src/zarr_metadata/v3/codec/sharding_indexed.py index a8c9247ec4..572a7f0871 100644 --- a/packages/zarr-metadata/src/zarr_metadata/v3/codec/sharding_indexed.py +++ b/packages/zarr-metadata/src/zarr_metadata/v3/codec/sharding_indexed.py @@ -27,8 +27,10 @@ class ShardingIndexedCodecConfiguration(TypedDict): """ Configuration for the Zarr v3 `sharding_indexed` codec. - `chunk_shape` is the shape of inner chunks along each dimension; - it must evenly divide the shard shape. + `chunk_shape` is the shape of inner chunks along each dimension. It does + not need to evenly divide the shard shape: inner chunks are clipped by the + shard shape (`sharding_indexed` spec version 1.1; version 1.0 required + `chunk_shape` to evenly divide the shard shape). `codecs` is the codec pipeline applied to each inner chunk; exactly one array-to-bytes codec is required. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 8f23606011..1fa0624c16 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -1,6 +1,6 @@ from __future__ import annotations -from collections.abc import Iterable, Mapping, MutableMapping, Sequence +from collections.abc import Callable, Iterable, Mapping, MutableMapping, Sequence from dataclasses import dataclass, replace from functools import lru_cache from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple @@ -370,6 +370,15 @@ class ShardingCodec( ): """Sharding codec. + Inner chunks are laid out on a semi-regular grid within each shard: + `chunk_shape` gives the nominal spacing and does not need to evenly divide + the shard shape. The number of inner chunks per dimension is the ceiling + division of the shard shape by `chunk_shape`, and chunks that straddle the + shard boundary are clipped — they are encoded and stored at their clipped + shape (`sharding_indexed` spec version 1.1). Version 1.0 implementations + reject a `chunk_shape` that does not evenly divide the shard shape, so + arrays relying on clipping are not readable by them. + `subchunk_write_order` controls the physical order of subchunks within a shard. It is a write-time setting only: it is not stored in array metadata, so reopening a sharded array does not recover it (the setting reverts to the `morton` default per codec instance). @@ -415,6 +424,7 @@ def __init__( object.__setattr__(self, "_get_index_chunk_spec", lru_cache()(self._get_index_chunk_spec)) object.__setattr__(self, "_get_chunks_per_shard", lru_cache()(self._get_chunks_per_shard)) object.__setattr__(self, "_shard_index_size", lru_cache()(self._shard_index_size)) + object.__setattr__(self, "_inner_grid", lru_cache()(self._inner_grid)) object.__setattr__( self, "_get_inner_chunk_transform", lru_cache()(self._get_inner_chunk_transform) ) @@ -441,6 +451,7 @@ def __setstate__(self, state: dict[str, Any]) -> None: object.__setattr__(self, "_get_index_chunk_spec", lru_cache()(self._get_index_chunk_spec)) object.__setattr__(self, "_get_chunks_per_shard", lru_cache()(self._get_chunks_per_shard)) object.__setattr__(self, "_shard_index_size", lru_cache()(self._shard_index_size)) + object.__setattr__(self, "_inner_grid", lru_cache()(self._inner_grid)) object.__setattr__( self, "_get_inner_chunk_transform", lru_cache()(self._get_inner_chunk_transform) ) @@ -540,24 +551,13 @@ def validate( raise ValueError( "The shard's `chunk_shape` and array's `shape` need to have the same number of dimensions." ) - if isinstance(chunk_grid, RegularChunkGridMetadata): - edges_per_dim: tuple[tuple[int, ...], ...] = tuple((s,) for s in chunk_grid.chunk_shape) - elif isinstance(chunk_grid, RectilinearChunkGridMetadata): - edges_per_dim = tuple( - (s,) if isinstance(s, int) else s for s in chunk_grid.chunk_shapes - ) - else: + if not isinstance(chunk_grid, (RegularChunkGridMetadata, RectilinearChunkGridMetadata)): raise TypeError( f"Sharding is only compatible with regular and rectilinear chunk grids, " f"got {type(chunk_grid)}" ) - for i, (edges, inner) in enumerate(zip(edges_per_dim, self.chunk_shape, strict=False)): - for edge in set(edges): - if edge % inner != 0: - raise ValueError( - f"Chunk edge length {edge} in dimension {i} is not " - f"divisible by the shard's inner chunk size {inner}." - ) + # `chunk_shape` does not need to evenly divide the shard shape: inner + # chunks are clipped by the shard shape (sharding_indexed v1.1). def _get_inner_chunk_transform(self, shard_spec: ArraySpec) -> Any: """The synchronous transform for the inner codec chain. @@ -639,18 +639,17 @@ def _decode_sync( See TODO: make issue for handling subchunk parallelism """ shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) inner_transform = self._get_inner_chunk_transform(shard_spec) indexer = BasicIndexer( tuple(slice(0, s) for s in shard_shape), shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) - out = chunk_spec.prototype.nd_buffer.empty( + out = shard_spec.prototype.nd_buffer.empty( shape=shard_shape, dtype=shard_spec.dtype.to_native_dtype(), order=shard_spec.order, @@ -669,7 +668,7 @@ def _decode_sync( decode_and_scatter_chunk( shard_dict.get(chunk_coords), out, - chunk_spec=chunk_spec, + chunk_spec=get_chunk_spec(chunk_coords), chunk_selection=chunk_selection, out_selection=out_selection, drop_axes=(), @@ -709,13 +708,13 @@ def _encode_sync( """ shard_shape = shard_spec.shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) inner_transform = self._get_inner_chunk_transform(shard_spec) indexer = BasicIndexer( tuple(slice(0, s) for s in shard_shape), shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) # Key order here is immaterial; _encode_shard_dict_sync lays the present @@ -727,7 +726,9 @@ def _encode_sync( for chunk_coords, _chunk_selection, out_selection, _ in indexer: # None = chunk normalized to missing (see encode_or_elide_chunk) shard_builder[chunk_coords] = encode_or_elide_chunk( - shard_array[out_selection], chunk_spec, inner_transform.encode_chunk + shard_array[out_selection], + get_chunk_spec(chunk_coords), + inner_transform.encode_chunk, ) return self._encode_shard_dict_sync( @@ -758,14 +759,14 @@ def _encode_partial_sync( """ shard_shape = shard_spec.shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) inner_transform = self._get_inner_chunk_transform(shard_spec) indexer = list( get_indexer( selection, shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) ) @@ -798,18 +799,21 @@ def _encode_partial_sync( # the canonical merge_and_encode_chunk (None = normalized to missing). # # Scalar fast path: when the written value is a scalar broadcast, every - # *complete* inner chunk is byte-for-byte identical — same fill, same - # empty-check, same encoded bytes. Compute that outcome once and reuse it - # for all complete chunks instead of re-merging, re-checking, and - # re-encoding tens of thousands of identical chunks. Incomplete (edge) - # chunks still merge against their own existing data individually. + # *complete* inner chunk of the nominal shape is byte-for-byte identical + # — same fill, same empty-check, same encoded bytes. Compute that + # outcome once and reuse it for all such chunks instead of re-merging, + # re-checking, and re-encoding tens of thousands of identical chunks. + # Chunks clipped by the shard boundary have a different shape (and thus + # different encoded bytes), so they are excluded from the memo, as are + # incomplete chunks, which merge against their own existing data. # `_sentinel` distinguishes "not computed yet" from a memoized `None` # (an empty chunk). _sentinel = object() scalar_complete_result: Buffer | None | object = _sentinel for chunk_coords, chunk_sel, out_sel, is_complete_chunk in indexer: - if is_scalar and is_complete_chunk: + chunk_spec = get_chunk_spec(chunk_coords) + if is_scalar and is_complete_chunk and chunk_spec.shape == self.chunk_shape: if scalar_complete_result is _sentinel: scalar_complete_result = merge_and_encode_chunk( None, @@ -942,18 +946,17 @@ async def _decode_single( shard_spec: ArraySpec, ) -> NDBuffer: shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) indexer = BasicIndexer( tuple(slice(0, s) for s in shard_shape), shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) # setup output array - out = chunk_spec.prototype.nd_buffer.empty( + out = shard_spec.prototype.nd_buffer.empty( shape=shard_shape, dtype=shard_spec.dtype.to_native_dtype(), order=shard_spec.order, @@ -969,7 +972,7 @@ async def _decode_single( [ ( _ShardingByteGetter(shard_dict, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -988,14 +991,13 @@ async def _decode_partial_single( shard_spec: ArraySpec, ) -> NDBuffer | None: shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) indexer = get_indexer( selection, shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) # setup output array @@ -1014,14 +1016,14 @@ async def _decode_partial_single( # read entire shard shard_dict_maybe = await self._load_full_shard_maybe( byte_getter=byte_getter, - prototype=chunk_spec.prototype, + prototype=shard_spec.prototype, chunks_per_shard=chunks_per_shard, ) else: # read some chunks within the shard shard_dict_maybe = await self._load_partial_shard_maybe( byte_getter, - chunk_spec.prototype, + shard_spec.prototype, chunks_per_shard, all_chunk_coords, max_gap_bytes=shard_spec.config.sharding_coalesce_max_gap_bytes, @@ -1037,7 +1039,7 @@ async def _decode_partial_single( [ ( _ShardingByteGetter(shard_dict, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -1107,6 +1109,14 @@ def _decode_full_shard_bulk_if_uncompressed( return None ab_codec = self.codecs[0] + # Inner chunks clipped by the shard boundary (chunk_shape not evenly + # dividing the shard shape) have variable payload sizes; the bulk path + # assumes a uniform payload per chunk, so those shards take the + # per-chunk path. (The blob-length and is_dense checks below would also + # reject them, but the gate makes the precondition explicit.) + if not self._is_evenly_divided(shard_spec.shape): + return None + chunks_per_shard = self._get_chunks_per_shard(shard_spec) chunk_spec = self._get_chunk_spec(shard_spec) n_chunks = product(chunks_per_shard) @@ -1200,15 +1210,14 @@ def _decode_partial_sync( the inner chunks the selection touches, fetch those, and decode. """ shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) inner_transform = self._get_inner_chunk_transform(shard_spec) indexer = get_indexer( selection, shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) out = shard_spec.prototype.nd_buffer.empty( @@ -1222,7 +1231,7 @@ def _decode_partial_sync( # Read just the inner chunks we need. if self._is_total_shard(all_chunk_coords, chunks_per_shard): - shard_bytes = byte_getter.get_sync(prototype=chunk_spec.prototype) + shard_bytes = byte_getter.get_sync(prototype=shard_spec.prototype) if shard_bytes is None: return None bulk = self._decode_full_shard_bulk_if_uncompressed(shard_bytes, shard_spec, indexer) @@ -1239,7 +1248,7 @@ def _decode_partial_sync( # / #3004). Returns None if the shard is absent. partial = self._load_partial_shard_maybe_sync( byte_getter, - chunk_spec.prototype, + shard_spec.prototype, chunks_per_shard, all_chunk_coords, max_gap_bytes=shard_spec.config.sharding_coalesce_max_gap_bytes, @@ -1255,7 +1264,7 @@ def _decode_partial_sync( decode_and_scatter_chunk( shard_dict.get(chunk_coords), out, - chunk_spec=chunk_spec, + chunk_spec=get_chunk_spec(chunk_coords), chunk_selection=chunk_selection, out_selection=out_selection, drop_axes=(), @@ -1272,15 +1281,14 @@ async def _encode_single( shard_spec: ArraySpec, ) -> Buffer | None: shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) indexer = list( BasicIndexer( tuple(slice(0, s) for s in shard_shape), shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) ) shard_builder = dict.fromkeys(lexicographic_order_coords(chunks_per_shard)) @@ -1289,7 +1297,7 @@ async def _encode_single( [ ( _ShardingByteSetter(shard_builder, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -1313,15 +1321,14 @@ async def _encode_partial_single( shard_spec: ArraySpec, ) -> None: shard_shape = shard_spec.shape - chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - chunk_spec = self._get_chunk_spec(shard_spec) + get_chunk_spec = self._make_chunk_spec_getter(shard_spec) indexer = list( get_indexer( selection, shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) ) @@ -1330,7 +1337,7 @@ async def _encode_partial_single( else: shard_reader = await self._load_full_shard_maybe( byte_getter=byte_setter, - prototype=chunk_spec.prototype, + prototype=shard_spec.prototype, chunks_per_shard=chunks_per_shard, ) shard_reader = shard_reader or _ShardReader.create_empty(chunks_per_shard) @@ -1343,7 +1350,7 @@ async def _encode_partial_single( [ ( _ShardingByteSetter(shard_dict, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -1384,12 +1391,12 @@ async def _encode_shard_dict( def _is_total_shard( self, all_chunk_coords: set[tuple[int, ...]], chunks_per_shard: tuple[int, ...] ) -> bool: - # `all_chunk_coords` comes from an indexer over this shard's chunk grid, so - # it is always a subset of that grid (`validate` requires the shard shape to - # be divisible by the inner chunk shape, so the indexer cannot produce an - # out-of-grid coordinate). A subset whose size equals the grid's is the - # whole grid, so the count check alone proves totality — no need to build - # and membership-test the full coordinate set on this hot path. + # `all_chunk_coords` comes from an indexer over this shard's inner chunk + # grid (`_inner_grid`, which covers the whole shard via ceiling + # division), so it is always a subset of that grid. A subset whose size + # equals the grid's is the whole grid, so the count check alone proves + # totality — no need to build and membership-test the full coordinate + # set on this hot path. return len(all_chunk_coords) == product(chunks_per_shard) def _is_complete_shard_write( @@ -1499,15 +1506,54 @@ def _get_chunk_spec(self, shard_spec: ArraySpec) -> ArraySpec: prototype=shard_spec.prototype, ) + def _inner_grid(self, shard_shape: tuple[int, ...]) -> ChunkGrid: + """The semi-regular grid of inner chunks within one shard. + + Inner chunks are spaced at `chunk_shape` intervals and clipped by the + shard shape: the grid shape is the ceiling division of the shard shape + by `chunk_shape`, and chunks straddling the shard boundary are stored + at their clipped shape. Memoized via instance-local `lru_cache`. + """ + return ChunkGrid.from_sizes(shard_shape, self.chunk_shape) + + def _is_evenly_divided(self, shard_shape: tuple[int, ...]) -> bool: + """True when every inner chunk of a shard has the full `chunk_shape`.""" + return all(s % c == 0 for s, c in zip(shard_shape, self.chunk_shape, strict=False)) + + def _make_chunk_spec_getter( + self, shard_spec: ArraySpec + ) -> Callable[[tuple[int, ...]], ArraySpec]: + """Per-inner-chunk `ArraySpec` lookup for one shard. + + Inner chunks clipped by the shard boundary are encoded at their clipped + shape, so their spec differs from the nominal `chunk_shape` spec. + Within one getter, chunks with the nominal shape all share one spec + object, which keeps per-spec caches downstream (e.g. + `ChunkTransform._resolve_specs`) effective. Note that `_get_chunk_spec` + is not cached (see #3054), so the nominal spec here is NOT the same + object as one obtained from a separate `_get_chunk_spec` call — callers + must compare specs by shape, not identity. + """ + nominal = self._get_chunk_spec(shard_spec) + if self._is_evenly_divided(shard_spec.shape): + return lambda chunk_coords: nominal + grid = self._inner_grid(shard_spec.shape) + specs: dict[tuple[int, ...], ArraySpec] = {nominal.shape: nominal} + + def get_spec(chunk_coords: tuple[int, ...]) -> ArraySpec: + chunk = grid[chunk_coords] + assert chunk is not None, f"chunk coords {chunk_coords} outside the inner grid" + shape = chunk.shape + spec = specs.get(shape) + if spec is None: + spec = replace(nominal, shape=shape) + specs[shape] = spec + return spec + + return get_spec + def _get_chunks_per_shard(self, shard_spec: ArraySpec) -> tuple[int, ...]: - return tuple( - s // c - for s, c in zip( - shard_spec.shape, - self.chunk_shape, - strict=False, - ) - ) + return self._inner_grid(shard_spec.shape).grid_shape def _shard_index_byte_range( self, chunks_per_shard: tuple[int, ...] diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index f75ef72415..d8be45be8c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -197,6 +197,25 @@ def _chunk_sizes_from_shape( return tuple(result) +def _inner_chunk_sizes_from_outer( + outer_chunk_sizes: tuple[tuple[int, ...], ...], inner_chunk_shape: tuple[int, ...] +) -> tuple[tuple[int, ...], ...]: + """Compute dask-style inner (sub-shard) chunk sizes from outer chunk sizes. + + The inner chunk grid restarts at every shard boundary and inner chunks are + clipped by the shard shape, so along each dimension the sizes are the + ceiling partition of each outer chunk's data size by the inner chunk size, + concatenated across the outer chunks. + """ + result: list[tuple[int, ...]] = [] + for outer_sizes, c in zip(outer_chunk_sizes, inner_chunk_shape, strict=True): + sizes: list[int] = [] + for outer_size in outer_sizes: + sizes.extend(min(c, outer_size - i * c) for i in range(ceildiv(outer_size, c))) + result.append(tuple(sizes)) + return tuple(result) + + def parse_array_metadata(data: Any) -> ArrayMetadata: if isinstance(data, ArrayMetadata): return data @@ -883,8 +902,11 @@ def read_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) if len(codecs) == 1 and isinstance(codecs[0], ShardingCodec): + # The inner chunk grid restarts at every shard boundary (inner + # chunks are clipped by the shard shape), so derive the sizes from + # the outer chunk sizes rather than the array shape alone. inner_chunk_shape = codecs[0].chunk_shape - return _chunk_sizes_from_shape(self.shape, inner_chunk_shape) + return _inner_chunk_sizes_from_outer(self._chunk_grid.chunk_sizes, inner_chunk_shape) return self._chunk_grid.chunk_sizes @property @@ -1121,9 +1143,14 @@ def _chunk_grid_shape(self) -> tuple[int, ...]: codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) if len(codecs) == 1 and isinstance(codecs[0], ShardingCodec): - # When sharding, count inner chunks across the whole array + # When sharding, count inner chunks across the whole array. The + # inner chunk grid restarts at every shard boundary (inner chunks + # are clipped by the shard shape), so count per outer chunk. chunk_shape = codecs[0].chunk_shape - return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True))) + return tuple( + sum(ceildiv(outer_size, c) for outer_size in outer_sizes) + for outer_sizes, c in zip(self._chunk_grid.chunk_sizes, chunk_shape, strict=True) + ) return self._chunk_grid.grid_shape @property @@ -5318,8 +5345,10 @@ async def _nchunks_initialized( if array.shards is None: chunks_per_shard = 1 else: + # Inner chunks are clipped by the shard shape, so the count per shard + # is the ceiling division of the shard shape by the chunk shape. chunks_per_shard = product( - tuple(a // b for a, b in zip(array.shards, array.chunks, strict=True)) + tuple(ceildiv(a, b) for a, b in zip(array.shards, array.chunks, strict=True)) ) return (await _nshards_initialized(array)) * chunks_per_shard diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index de576dbef5..19e88bea7d 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -32,6 +32,7 @@ from zarr.core.buffer import NDArrayLike, default_buffer_prototype from zarr.core.indexing import lexicographic_order_coords from zarr.core.metadata.v3 import ArrayV3Metadata +from zarr.core.sync import sync from zarr.storage import MemoryStore, StorePath, ZipStore from ..conftest import ArrayRequest @@ -831,34 +832,189 @@ def test_invalid_metadata(store: Store) -> None: dtype=np.dtype("uint8"), fill_value=0, ) - spath2 = StorePath(store, "invalid_inner_chunk_shape") - with pytest.raises(ValueError): - zarr.create_array( - spath2, - shape=(16, 16), - shards=(16, 16), - chunks=(8, 7), - dtype=np.dtype("uint8"), - fill_value=0, + + +@pytest.mark.parametrize( + ("array_shape", "shard_shape", "chunk_shape"), + [ + ((24,), (12,), (5,)), # trailing inner chunk clipped to 2 elements + ((10,), (4,), (6,)), # inner chunk larger than the shard: one clipped chunk per shard + ((23, 17), (12, 12), (5, 5)), # array shape not divisible by the shard shape either + ((16, 16), (16, 16), (9, 9)), # single shard, clipped in both dimensions + ((16, 16), (8, 16), (4, 4)), # evenly divisible control case + ], +) +@pytest.mark.parametrize("index_location", ["start", "end"]) +def test_sharding_clipped_inner_chunks( + array_shape: tuple[int, ...], + shard_shape: tuple[int, ...], + chunk_shape: tuple[int, ...], + index_location: IndexLocation, +) -> None: + """`chunk_shape` does not need to evenly divide the shard shape + (`sharding_indexed` v1.1): inner chunks are clipped by the shard boundary. + + Full and partial reads and writes round-trip for reasonable combinations of + array, shard, and inner chunk shape. + """ + data = np.arange(np.prod(array_shape), dtype="uint16").reshape(array_shape) + arr = zarr.create_array( + {}, + shape=array_shape, + shards={"shape": shard_shape, "index_location": index_location}, + chunks=chunk_shape, + dtype=data.dtype, + fill_value=0, + ) + arr[...] = data + assert np.array_equal(arr[...], data) + + # partial read crossing clipped inner chunk boundaries + sel = tuple(slice(1, s - 1) for s in array_shape) + assert np.array_equal(arr[sel], data[sel]) + + # partial write crossing clipped inner chunk boundaries + arr[sel] = 0 + expected = data.copy() + expected[sel] = 0 + assert np.array_equal(arr[...], expected) + + +def test_sharding_clipped_inner_chunks_byte_layout() -> None: + """Spec example: a shard of shape [12] with inner chunk shape [5] holds 3 + inner chunks stored at their clipped shapes ([5], [5], and [2]). + + With an uncompressed inner codec chain the shard blob holds exactly 12 data + bytes (not 15) plus a 3-entry shard index. + """ + store = MemoryStore() + arr = zarr.create_array( + store, + shape=(12,), + shards=(12,), + chunks=(5,), + dtype="uint8", + fill_value=255, + compressors=None, + ) + arr[...] = np.arange(12, dtype="uint8") + + shard_buffer = sync(store.get("c/0", prototype=default_buffer_prototype())) + assert shard_buffer is not None + raw = shard_buffer.to_bytes() + index_size = 3 * 16 + 4 # ceil(12/5) = 3 index entries plus the crc32c checksum + assert len(raw) == 12 + index_size + offsets_and_lengths = np.frombuffer(raw[-index_size:-4], dtype=" None: + """A rectilinear outer grid whose shard edges are not divisible by the inner + chunk shape round-trips: the inner grid is clipped per shard. + + Shard edges `[4, 6]` cannot be normalized to a regular grid (the larger edge + comes last), so this exercises `RectilinearChunkGridMetadata` for real — + unlike e.g. `[6, 4]`, which coincides with a clipped regular grid of size 6. + """ + from zarr.core.metadata.v3 import RectilinearChunkGridMetadata + + with zarr.config.set({"array.rectilinear_chunks": True}): + arr = zarr.create_array( + {}, shape=(10,), shards=[[4, 6]], chunks=(4,), dtype="uint8", fill_value=0 ) + assert isinstance(arr.metadata, ArrayV3Metadata) + assert isinstance(arr.metadata.chunk_grid, RectilinearChunkGridMetadata) + data = np.arange(10, dtype="uint8") + arr[...] = data + assert np.array_equal(arr[...], data) + assert arr.read_chunk_sizes == ((4, 4, 2),) + + arr[3:9] = 42 + data[3:9] = 42 + assert np.array_equal(arr[...], data) + + +def test_sharding_clipped_inner_chunks_nested() -> None: + """Clipping composes through nested sharding: 12-element shards split into + clipped 5-element inner chunks, each split into clipped 2-element chunks.""" + inner = ShardingCodec(chunk_shape=(2,), codecs=(BytesCodec(),)) + outer = ShardingCodec(chunk_shape=(5,), codecs=(inner,)) + arr = zarr.create_array( + {}, + shape=(24,), + chunks=(12,), + dtype="uint8", + fill_value=0, + serializer=outer, + compressors=None, + filters=None, + ) + data = np.arange(24, dtype="uint8") + arr[...] = data + assert np.array_equal(arr[...], data) + arr[4:11] = 9 + data[4:11] = 9 + assert np.array_equal(arr[...], data) -def test_invalid_shard_shape() -> None: - with pytest.raises( - ValueError, - match=( - f"Chunk edge length {16} in dimension {0} is not " - f"divisible by the shard's inner chunk size {9}\\." - ), - ): - zarr.create_array( - {}, - shape=(16, 16), - shards=(16, 16), - chunks=(9, 9), - dtype=np.dtype("uint8"), - fill_value=0, + +def test_sharding_clipped_inner_chunk_sizes() -> None: + """Inner chunk grids restart at each shard boundary, so the reported inner + chunk sizes are clipped per shard (and by the array extent in the last + shard).""" + arr = zarr.create_array({}, shape=(23,), shards=(12,), chunks=(5,), dtype="uint8", fill_value=0) + assert arr.read_chunk_sizes == ((5, 5, 2, 5, 5, 1),) + assert arr.cdata_shape == (6,) + assert arr.nchunks == 6 + + +def test_sharding_clipped_nchunks_initialized() -> None: + """`nchunks_initialized` counts ceiling-division inner chunks per shard. + + Regression test: it previously used floor division, which undercounts for + clipped grids and reports 0 when the inner chunk shape exceeds the shard + shape. + """ + arr = zarr.create_array({}, shape=(16,), shards=(8,), chunks=(3,), dtype="uint8", fill_value=0) + arr[...] = np.arange(16, dtype="uint8") + # 2 shards, each holding ceil(8/3) = 3 inner chunks + assert arr.nchunks_initialized == 6 + + arr2 = zarr.create_array({}, shape=(11,), shards=(4,), chunks=(9,), dtype="uint8", fill_value=0) + arr2[...] = np.arange(11, dtype="uint8") + # 3 shards, each holding one inner chunk clipped to the shard shape + assert arr2.nchunks_initialized == 3 + + +def test_sharding_scalar_write_memoizes_complete_chunks(monkeypatch: pytest.MonkeyPatch) -> None: + """A scalar broadcast partial write encodes the repeated complete inner + chunk once, not once per chunk. + + Regression test: the memo gate must compare chunk specs by shape — an + object-identity gate never fires because `_get_chunk_spec` is uncached. + """ + from zarr.core import chunk_utils + + calls = 0 + real_encode = chunk_utils.ChunkTransform.encode_chunk + + def counting_encode(self: Any, *args: Any, **kwargs: Any) -> Any: + nonlocal calls + calls += 1 + return real_encode(self, *args, **kwargs) + + with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + arr = zarr.create_array( + {}, shape=(100,), shards=(100,), chunks=(10,), dtype="uint8", fill_value=0 ) + monkeypatch.setattr(chunk_utils.ChunkTransform, "encode_chunk", counting_encode) + arr[5:] = 7 + expected = np.full(100, 7, dtype="uint8") + expected[:5] = 0 + assert np.array_equal(arr[...], expected) + # one merge for the partial edge chunk, one memoized encode shared by the + # nine complete chunks, one shard-index encode + assert calls <= 4, f"scalar write encoded {calls} times; memo is not firing" @pytest.mark.parametrize("store", ["local"], indirect=["store"]) diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index 0df5a5d9fd..26209b828e 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -1048,8 +1048,9 @@ def test_sharding_accepts_rectilinear_outer_grid() -> None: ) -def test_sharding_rejects_non_divisible_rectilinear() -> None: - """Rectilinear shard sizes not divisible by inner chunk_shape should raise.""" +def test_sharding_accepts_non_divisible_rectilinear() -> None: + """Rectilinear shard sizes not divisible by inner chunk_shape are accepted: + inner chunks are clipped by the shard shape (sharding_indexed v1.1).""" from zarr.codecs.sharding import ShardingCodec from zarr.core.dtype import Float32 from zarr.core.metadata.v3 import RectilinearChunkGridMetadata @@ -1057,12 +1058,11 @@ def test_sharding_rejects_non_divisible_rectilinear() -> None: codec = ShardingCodec(chunk_shape=(5, 5)) grid_meta = RectilinearChunkGridMetadata(chunk_shapes=((10, 20, 17), (50, 50))) - with pytest.raises(ValueError, match="divisible"): - codec.validate( - shape=(47, 100), - dtype=Float32(), - chunk_grid=grid_meta, - ) + codec.validate( + shape=(47, 100), + dtype=Float32(), + chunk_grid=grid_meta, + ) def test_sharding_accepts_divisible_rectilinear() -> None: @@ -1081,17 +1081,18 @@ def test_sharding_accepts_divisible_rectilinear() -> None: ) -def test_sharding_rejects_non_divisible_among_repeated_edges() -> None: - """Shard validation catches a non-divisible edge even among many repeated valid ones.""" +def test_sharding_accepts_non_divisible_among_repeated_edges() -> None: + """Shard validation accepts a non-divisible edge among repeated divisible + ones: inner chunks are clipped per shard (sharding_indexed v1.1).""" from zarr.codecs.sharding import ShardingCodec from zarr.core.dtype import Float32 from zarr.core.metadata.v3 import RectilinearChunkGridMetadata - # edges (10, 10, 7) — 7 is not divisible by 5 + # edges (10, 10, 7) — 7 is not divisible by 5; the last shard holds + # inner chunks of sizes (5, 2) codec = ShardingCodec(chunk_shape=(5,)) grid_meta = RectilinearChunkGridMetadata(chunk_shapes=((10, 10, 7),)) - with pytest.raises(ValueError, match="divisible"): - codec.validate(shape=(27,), dtype=Float32(), chunk_grid=grid_meta) + codec.validate(shape=(27,), dtype=Float32(), chunk_grid=grid_meta) def test_sharding_accepts_all_repeated_divisible_edges() -> None: @@ -1761,16 +1762,20 @@ def test_pipeline_rectilinear_shards_partial_read(tmp_path: Path) -> None: np.testing.assert_array_equal(result, data[50:70, 40:60]) -def test_pipeline_rectilinear_shards_validates_divisibility(tmp_path: Path) -> None: - """Inner chunk_shape must divide every shard's dimensions.""" - with pytest.raises(ValueError, match="divisible"): - zarr.create_array( - store=tmp_path / "bad.zarr", - shape=(120, 100), - chunks=(10, 10), - shards=[[60, 45, 15], [50, 50]], - dtype="int32", - ) +def test_pipeline_rectilinear_shards_non_divisible(tmp_path: Path) -> None: + """Inner chunk_shape need not divide the shard dimensions: inner chunks are + clipped per shard (sharding_indexed v1.1) and data round-trips.""" + arr = zarr.create_array( + store=tmp_path / "clipped.zarr", + shape=(120, 100), + chunks=(10, 10), + shards=[[60, 45, 15], [50, 50]], + dtype="int32", + ) + data = np.arange(120 * 100, dtype="int32").reshape(120, 100) + arr[:] = data + np.testing.assert_array_equal(arr[:], data) + np.testing.assert_array_equal(arr[55:70, 40:60], data[55:70, 40:60]) def test_pipeline_nchunks(tmp_path: Path) -> None: