diff --git a/changes/251.feature.md b/changes/251.feature.md new file mode 100644 index 0000000000..c5fb03437e --- /dev/null +++ b/changes/251.feature.md @@ -0,0 +1,11 @@ +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. +Chunk counts retain constant-space arithmetic for uniform shard dimensions, including very large sparse arrays. diff --git a/docs/user-guide/arrays.md b/docs/user-guide/arrays.md index 140bc4bd09..7f2f54870e 100644 --- a/docs/user-guide/arrays.md +++ b/docs/user-guide/arrays.md @@ -593,6 +593,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" @@ -731,8 +739,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 bd0760f7e3..8cf472526b 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 @@ -425,6 +425,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). @@ -470,6 +479,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) ) @@ -496,6 +506,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) ) @@ -595,24 +606,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. @@ -694,18 +694,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, @@ -724,7 +723,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=(), @@ -764,13 +763,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 @@ -782,7 +781,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( @@ -812,7 +813,7 @@ def _encode_partial_sync( inner chunks, and rewrites the whole shard. """ 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, value = self._get_shard_indexer_and_value(selection, shard_spec, value) @@ -846,18 +847,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 | object | None = _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, @@ -990,18 +994,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, @@ -1017,7 +1020,7 @@ async def _decode_single( [ ( _ShardingByteGetter(shard_dict, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -1036,14 +1039,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 @@ -1062,14 +1064,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, @@ -1085,7 +1087,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, @@ -1162,7 +1164,7 @@ def _decode_full_shard_bulk_if_uncompressed( # The byte-order handling below lacks the structured-dtype branch of # `BytesCodec._decode_sync` (which applies `newbyteorder` to multi-byte # struct fields), so structured dtypes must take the per-chunk path. - if isinstance(shard_spec.dtype, Struct): + if isinstance(shard_spec.dtype, Struct) or not self._is_evenly_divided(shard_spec.shape): return None chunks_per_shard = self._get_chunks_per_shard(shard_spec) @@ -1255,15 +1257,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( @@ -1277,7 +1278,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) @@ -1294,7 +1295,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, @@ -1310,7 +1311,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=(), @@ -1327,15 +1328,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)) @@ -1344,7 +1344,7 @@ async def _encode_single( [ ( _ShardingByteSetter(shard_builder, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -1368,7 +1368,7 @@ async def _encode_partial_single( shard_spec: ArraySpec, ) -> None: 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, shard_array = self._get_shard_indexer_and_value(selection, shard_spec, shard_array) @@ -1377,7 +1377,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) @@ -1390,7 +1390,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, @@ -1457,7 +1457,7 @@ def _get_shard_indexer_and_value( indexer = get_indexer( selection, shape=shard_shape, - chunk_grid=ChunkGrid.from_sizes(shard_shape, self.chunk_shape), + chunk_grid=self._inner_grid(shard_shape), ) if ( isinstance(indexer, CoordinateIndexer) @@ -1470,12 +1470,12 @@ def _get_shard_indexer_and_value( 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( @@ -1585,15 +1585,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 047e7bb3b3..011e35a3b2 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -6,7 +6,6 @@ from asyncio import gather from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace -from itertools import starmap from logging import getLogger from typing import ( TYPE_CHECKING, @@ -43,6 +42,7 @@ from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, + FixedDimension, _is_auto, _is_keep, _is_rectilinear_chunks, @@ -200,6 +200,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 @@ -900,7 +919,9 @@ def read_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: ((20, 20),) """ if (sharding_codec := _sharding_codec(self.metadata)) is not None: - return _chunk_sizes_from_shape(self.shape, sharding_codec.chunk_shape) + return _inner_chunk_sizes_from_outer( + self._chunk_grid.chunk_sizes, sharding_codec.chunk_shape + ) return self._chunk_grid.chunk_sizes @property @@ -1143,9 +1164,20 @@ def _chunk_grid_shape(self) -> tuple[int, ...]: The number of chunks along each dimension. """ if (sharding_codec := _sharding_codec(self.metadata)) is not None: - # When sharding, count inner chunks across the whole array + # The inner grid restarts at every shard boundary. chunk_shape = sharding_codec.chunk_shape - return tuple(starmap(ceildiv, zip(self.shape, chunk_shape, strict=True))) + counts: list[int] = [] + for dimension, inner in zip(self._chunk_grid.dimensions, chunk_shape, strict=True): + if isinstance(dimension, FixedDimension): + full, remainder = divmod(dimension.extent, dimension.size) + counts.append(full * ceildiv(dimension.size, inner) + ceildiv(remainder, inner)) + else: + counts.append( + sum( + ceildiv(dimension.data_size(i), inner) for i in range(dimension.nchunks) + ) + ) + return tuple(counts) return self._chunk_grid.grid_shape @property @@ -4117,7 +4149,7 @@ async def _shards_initialized( class ShardsConfigParam(TypedDict): - shape: tuple[int, ...] + shape: Sequence[int | Sequence[int]] index_location: IndexLocation | None @@ -5467,7 +5499,10 @@ async def _nchunks_initialized( # Uniform shard shape: the exact chunk count per shard is a single # multiply, O(1) after the storage listing. chunks_per_shard = product( - tuple(a // b for a, b in zip(meta.chunk_grid.chunk_shape, inner_chunks, strict=True)) + tuple( + ceildiv(a, b) + for a, b in zip(meta.chunk_grid.chunk_shape, inner_chunks, strict=True) + ) ) return (await _nshards_initialized(array)) * chunks_per_shard # Rectilinear shard grid: the chunk count varies per shard, so decode the @@ -5478,7 +5513,7 @@ async def _nchunks_initialized( spec = grid[meta.chunk_key_encoding.decode_chunk_key(key)] if spec is not None: total += product( - tuple(s // c for s, c in zip(spec.codec_shape, inner_chunks, strict=True)) + tuple(ceildiv(s, c) for s, c in zip(spec.codec_shape, inner_chunks, strict=True)) ) return total diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index cc27366027..33c1662dcb 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -960,6 +960,7 @@ def resolve_outer_and_inner_chunks( # Extract the flat chunk shape (uniform size per dimension) for arithmetic. chunk_shape_flat = chunks.chunk_shape + shard_flat: tuple[int | Sequence[int], ...] if _is_auto(shard_shape): warnings.warn( "Automatic shard shape inference is experimental and may change without notice.", diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 42a88ddc7a..323c84bfd0 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -33,6 +33,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 @@ -832,34 +833,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"]) @@ -1337,3 +1493,16 @@ def test_sharding_orthogonal_set_multiple_array_dims( expected[ix] = value.reshape(expected[ix].shape) assert np.array_equal(a[:], expected) assert np.array_equal(a.oindex[selection], value) + + +def test_sharding_huge_grid_cardinality(monkeypatch: pytest.MonkeyPatch) -> None: + """Counting chunks in sparse large arrays must not allocate a per-shard table.""" + from zarr.core.chunk_grids import ChunkGrid + + array = zarr.create_array({}, shape=(10**12,), chunks=(2,), shards=(3,), dtype="u1") + + def materialized_sizes(grid: ChunkGrid) -> tuple[tuple[int, ...], ...]: + raise AssertionError("cardinality tried to materialize every shard") + + monkeypatch.setattr(ChunkGrid, "chunk_sizes", property(materialized_sizes)) + assert array.cdata_shape == (666666666667,) diff --git a/tests/test_codecs/test_sharding_clipped_properties.py b/tests/test_codecs/test_sharding_clipped_properties.py new file mode 100644 index 0000000000..1cdd4d3837 --- /dev/null +++ b/tests/test_codecs/test_sharding_clipped_properties.py @@ -0,0 +1,75 @@ +"""NumPy oracle for clipped inner chunks across both pipeline implementations.""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import numpy as np +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +import zarr + +if TYPE_CHECKING: + from zarr.codecs.sharding import IndexLocation + + +@st.composite +def clipped_grids( + draw: st.DrawFn, +) -> tuple[tuple[int, ...], tuple[int | tuple[int, ...], ...], tuple[int, ...]]: + shape: list[int] = [] + shards: list[int | tuple[int, ...]] = [] + chunks: list[int] = [] + for _ in range(draw(st.integers(1, 3))): + if draw(st.booleans()): + edges = tuple(draw(st.lists(st.integers(1, 6), min_size=1, max_size=3))) + shards.append(edges) + shape.append(sum(edges)) + else: + shards.append(draw(st.integers(1, 6))) + shape.append(draw(st.integers(1, 12))) + chunks.append(draw(st.integers(1, 8))) + return tuple(shape), tuple(shards), tuple(chunks) + + +@pytest.mark.parametrize("pipeline", ["BatchedCodecPipeline", "FusedCodecPipeline"]) +@pytest.mark.parametrize("index_location", ["start", "end"]) +@settings(max_examples=30, deadline=None) +@given(case=clipped_grids()) +def test_clipped_sharding_matches_numpy( + pipeline: str, + index_location: IndexLocation, + case: tuple[tuple[int, ...], tuple[int | tuple[int, ...], ...], tuple[int, ...]], +) -> None: + shape, shards, chunks = case + expected = np.arange(math.prod(shape), dtype="i4").reshape(shape) + selection = tuple(slice(1, None, 2) for _ in shape) + with zarr.config.set( + { + "array.rectilinear_chunks": True, + "codec_pipeline.path": f"zarr.core.codec_pipeline.{pipeline}", + } + ): + array = zarr.create_array( + {}, + shape=shape, + shards={"shape": shards, "index_location": index_location}, + chunks=chunks, + dtype="i4", + compressors=None, + fill_value=0, + ) + array[:] = expected + np.testing.assert_array_equal(array[selection], expected[selection]) + array[selection] = -7 + expected[selection] = -7 + np.testing.assert_array_equal(array[:], expected) + array[selection] = 0 + expected[selection] = 0 + reopened = zarr.open_array(array.store, mode="r") + np.testing.assert_array_equal(reopened[:], expected) + assert array.cdata_shape == tuple(len(axis) for axis in array.read_chunk_sizes) + assert tuple(sum(axis) for axis in array.read_chunk_sizes) == shape diff --git a/tests/test_unified_chunk_grid.py b/tests/test_unified_chunk_grid.py index b8289d2135..31bad92d21 100644 --- a/tests/test_unified_chunk_grid.py +++ b/tests/test_unified_chunk_grid.py @@ -1327,8 +1327,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 @@ -1336,12 +1337,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: @@ -1360,17 +1360,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: @@ -2040,16 +2041,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: