From 9bfe855a6cfd6985073fd5e4e5b5be027628eca8 Mon Sep 17 00:00:00 2001 From: Davis Bennett Date: Tue, 14 Jul 2026 13:25:16 +0200 Subject: [PATCH 1/5] fix: byte-order handling for structured dtypes in the bytes codec (#220) * fix: byte-order handling for structured dtypes in the bytes codec The bytes codec neither byte-swapped structured-dtype fields to its configured endian on encode (numpy reports byteorder '|' for void dtypes, so the top-level byteorder comparison never detected a mismatch) nor honored its endian when decoding, silently corrupting any structured data whose field byte order differed from the stored one (e.g. virtual references to external big-endian data). Encode now detects byte-order mismatches by comparing full dtypes via newbyteorder, and decode reinterprets raw bytes in the stored byte order before converting to the data type's declared byte order, so the stored layout (codec state) and the in-memory layout (array data type) are independent. Closes #4141 Assisted-by: ClaudeCode:claude-fable-5 * test: fold structured byte-order cases into existing bytes codec tests Extend test_endian's parametrization with structured dtypes and test_bytes_codec_sync_roundtrip with endian/dtype parametrization plus stored-layout and decoded-dtype assertions, instead of adding parallel test functions for the same properties. Assisted-by: ClaudeCode:claude-fable-5 * refactor: rename stored_dtype to view_dtype in BytesCodec decode The variable is the dtype used to view the raw chunk bytes (byte order from the codec's endian configuration), not a property of the stored data or of the returned buffer, which always carries the array's declared dtype. Assisted-by: ClaudeCode:claude-fable-5 * docs: note that the decode-side byte-order conversion copies the chunk Assisted-by: ClaudeCode:claude-fable-5 --- changes/4141.bugfix.md | 1 + src/zarr/codecs/bytes.py | 33 +++++++++++----- tests/test_codecs/test_bytes.py | 69 ++++++++++++++++++++++++++------- 3 files changed, 80 insertions(+), 23 deletions(-) create mode 100644 changes/4141.bugfix.md diff --git a/changes/4141.bugfix.md b/changes/4141.bugfix.md new file mode 100644 index 0000000000..6a132da3f5 --- /dev/null +++ b/changes/4141.bugfix.md @@ -0,0 +1 @@ +Fix silent byte-order corruption for structured dtypes with the `bytes` codec: multi-byte fields are now byte-swapped to the codec's configured `endian` on write and decoded honoring it on read, so non-native-endian structured data (e.g. big-endian fields, as produced by virtual references to external data) round-trips correctly. diff --git a/src/zarr/codecs/bytes.py b/src/zarr/codecs/bytes.py index 240c077627..fae762fd08 100644 --- a/src/zarr/codecs/bytes.py +++ b/src/zarr/codecs/bytes.py @@ -100,14 +100,28 @@ def _decode_sync( chunk_spec: ArraySpec, ) -> NDBuffer: endian_str = self.endian + dtype = chunk_spec.dtype.to_native_dtype() + # The byte order of the stored data is set by this codec's `endian` + # configuration; the byte order of the decoded array is set by the array's + # data type. The two are independent: the raw bytes are viewed with a dtype + # in the stored byte order, then converted to the declared dtype if needed. if isinstance(chunk_spec.dtype, HasEndianness): - dtype = replace(chunk_spec.dtype, endianness=endian_str).to_native_dtype() # type: ignore[call-arg] + view_dtype = replace(chunk_spec.dtype, endianness=endian_str).to_native_dtype() # type: ignore[call-arg] + elif isinstance(chunk_spec.dtype, Struct) and endian_str is not None: + # Per the struct data type spec, all multi-byte fields are stored in the + # byte order configured on this codec. + view_dtype = dtype.newbyteorder(endian_str) else: - dtype = chunk_spec.dtype.to_native_dtype() + view_dtype = dtype as_array_like = chunk_bytes.as_array_like() chunk_array = chunk_spec.prototype.nd_buffer.from_ndarray_like( - as_array_like.view(dtype=dtype) # type: ignore[attr-defined] + as_array_like.view(dtype=view_dtype) # type: ignore[attr-defined] ) + if view_dtype != dtype: + # This byte-swapping conversion copies the chunk. The dtype inequality + # guard keeps the common case, where the stored and declared byte orders + # already match, on the zero-copy view path above. + chunk_array = chunk_array.astype(dtype) # ensure correct chunk shape if chunk_array.shape != chunk_spec.shape: @@ -129,13 +143,14 @@ def _encode_sync( chunk_spec: ArraySpec, ) -> Buffer | None: assert isinstance(chunk_array, NDBuffer) - if ( - chunk_array.dtype.itemsize > 1 - and self.endian is not None - and self.endian != chunk_array.byteorder - ): + if chunk_array.dtype.itemsize > 1 and self.endian is not None: + # Compare full dtypes rather than the top-level byteorder: numpy reports + # byteorder '|' for structured dtypes even when their fields are + # byte-order-sensitive, so newbyteorder is the only reliable way to + # detect (and normalize) a byte-order mismatch. new_dtype = chunk_array.dtype.newbyteorder(self.endian) - chunk_array = chunk_array.astype(new_dtype) + if new_dtype != chunk_array.dtype: + chunk_array = chunk_array.astype(new_dtype) nd_array = chunk_array.as_ndarray_like() # Flatten the nd-array (only copy if needed) and reinterpret as bytes diff --git a/tests/test_codecs/test_bytes.py b/tests/test_codecs/test_bytes.py index 03dd0b40c6..ead778f526 100644 --- a/tests/test_codecs/test_bytes.py +++ b/tests/test_codecs/test_bytes.py @@ -33,29 +33,50 @@ @pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"]) -@pytest.mark.parametrize("input_dtype", [">u2", "u2", + "f4"), ("mask", ">i4")], + [("flux", "u2", " None: """ The `bytes` codec stores multi-byte data in the byte order configured on the codec, regardless of the input array's byte order, and reads it back to the - original values. The input-dtype/store-endian cross-product exercises the - encode-side byteswap (input byte order != store byte order) and the no-op - case alike. Compression is disabled so the stored chunk is the codec's raw - output and its byte layout can be asserted directly. + original values. For structured dtypes this applies to every multi-byte + field, per the `struct` data type spec; the struct cases guard against the + endianness bugs from + https://github.com/zarr-developers/zarr-python/issues/4141, where the + encode path never byte-swapped struct fields (numpy reports byteorder '|' + for void dtypes) and the decode path ignored the codec's endian entirely. + The input-dtype/store-endian cross-product exercises the encode-side + byteswap (input byte order != store byte order) and the no-op case alike. + Compression is disabled so the stored chunk is the codec's raw output and + its byte layout can be asserted directly. """ - data = np.arange(0, 256, dtype=input_dtype).reshape((16, 16)) + dtype = np.dtype(input_dtype) + if dtype.fields is None: + data = np.arange(0, 256, dtype=dtype).reshape((16, 16)) + else: + data = np.zeros((16, 16), dtype=dtype) + data["flux"] = np.arange(0, 256).reshape((16, 16)) + data["mask"] = np.arange(256, 512).reshape((16, 16)) path = "endian" spath = StorePath(store, path) a = await zarr.api.asynchronous.create_array( spath, shape=data.shape, chunks=(16, 16), - dtype="uint16", + dtype=dtype, fill_value=0, compressors=None, serializer=BytesCodec(endian=store_endian), @@ -66,8 +87,7 @@ async def test_endian( # The stored chunk is laid out in the byte order configured on the codec. stored = await store.get(f"{path}/c/0/0", prototype=default_buffer_prototype()) assert stored is not None - expected_dtype = ">u2" if store_endian == "big" else " None: assert isinstance(BytesCodec(), SupportsSyncCodec) -def test_bytes_codec_sync_roundtrip() -> None: - codec = BytesCodec() - arr = np.arange(100, dtype="float64") +@pytest.mark.parametrize("endian", ENDIAN) +@pytest.mark.parametrize( + "native_dtype", + [np.dtype("float64"), np.dtype(">u2"), np.dtype([("a", ">f4"), ("b", " None: + """ + The synchronous encode/decode path round-trips data, and the two byte + orders involved are independent: the codec's `endian` configuration governs + only the stored byte layout (every multi-byte value, including struct + fields, is laid out in the codec's byte order regardless of the input + array's byte order), while the decoded buffer's byte order is governed by + the array's data type regardless of the codec's. The mixed-endian struct + case pins that per-field byte order of the in-memory dtype survives a + roundtrip through a single stored byte order. + """ + if native_dtype.fields is None: + arr = np.arange(100, dtype=native_dtype) + else: + arr = np.array([(1.5, 2), (3.5, 4), (5.5, 6), (7.5, 8)], dtype=native_dtype) zdtype = get_data_type_from_native_dtype(arr.dtype) spec = ArraySpec( shape=arr.shape, @@ -91,11 +129,14 @@ def test_bytes_codec_sync_roundtrip() -> None: ) nd_buf: NDBuffer = default_buffer_prototype().nd_buffer.from_numpy_array(arr) - codec = codec.evolve_from_array_spec(spec) + codec = BytesCodec(endian=endian).evolve_from_array_spec(spec) encoded = codec._encode_sync(nd_buf, spec) assert encoded is not None + assert encoded.to_bytes() == arr.astype(native_dtype.newbyteorder(endian)).tobytes() + decoded = codec._decode_sync(encoded, spec) + assert decoded.dtype == zdtype.to_native_dtype() np.testing.assert_array_equal(arr, decoded.as_numpy_array()) From a024bb2e5d400ba50d5a0d33c7ec9a0239e443b5 Mon Sep 17 00:00:00 2001 From: Davis Bennett Date: Tue, 28 Jul 2026 22:05:12 +0200 Subject: [PATCH 2/5] fix: stop emitting UnstableSpecificationWarning for the spec'd "struct" data type (#202) * chore(deps): bump the actions group across 1 directory with 8 updates (#176) Bumps the actions group with 8 updates in the / directory: | Package | From | To | | --- | --- | --- | | [prefix-dev/setup-pixi](https://github.com/prefix-dev/setup-pixi) | `0.9.5` | `0.9.6` | | [codecov/codecov-action](https://github.com/codecov/codecov-action) | `6.0.0` | `6.0.1` | | [github/issue-metrics](https://github.com/github/issue-metrics) | `4.2.2` | `4.2.7` | | [j178/prek-action](https://github.com/j178/prek-action) | `2.0.3` | `2.0.4` | | [actions/upload-artifact](https://github.com/actions/upload-artifact) | `7.0.0` | `7.0.1` | | [actions/download-artifact](https://github.com/actions/download-artifact) | `7.0.0` | `8.0.1` | | [pypa/gh-action-pypi-publish](https://github.com/pypa/gh-action-pypi-publish) | `1.13.0` | `1.14.0` | | [zizmorcore/zizmor-action](https://github.com/zizmorcore/zizmor-action) | `0.5.3` | `0.5.6` | Updates `prefix-dev/setup-pixi` from 0.9.5 to 0.9.6 - [Release notes](https://github.com/prefix-dev/setup-pixi/releases) - [Commits](https://github.com/prefix-dev/setup-pixi/compare/1b2de7f3351f171c8b4dfeb558c639cb58ed4ec0...5185adfbffb4bd703da3010310260805d89ebb11) Updates `codecov/codecov-action` from 6.0.0 to 6.0.1 - [Release notes](https://github.com/codecov/codecov-action/releases) - [Changelog](https://github.com/codecov/codecov-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/codecov/codecov-action/compare/57e3a136b779b570ffcdbf80b3bdc90e7fab3de2...e79a6962e0d4c0c17b229090214935d2e33f8354) Updates `github/issue-metrics` from 4.2.2 to 4.2.7 - [Release notes](https://github.com/github/issue-metrics/releases) - [Commits](https://github.com/github/issue-metrics/compare/c9e9838147fd355dace335ba787f01b6641a400a...1e38d5e62363e14db8019ed7d106b9855bdba6cc) Updates `j178/prek-action` from 2.0.3 to 2.0.4 - [Release notes](https://github.com/j178/prek-action/releases) - [Commits](https://github.com/j178/prek-action/compare/6ad80277337ad479fe43bd70701c3f7f8aa74db3...bdca6f102f98e2b4c7029491a53dfd366469e33d) Updates `actions/upload-artifact` from 7.0.0 to 7.0.1 - [Release notes](https://github.com/actions/upload-artifact/releases) - [Commits](https://github.com/actions/upload-artifact/compare/v7...043fb46d1a93c77aae656e7c1c64a875d1fc6a0a) Updates `actions/download-artifact` from 7.0.0 to 8.0.1 - [Release notes](https://github.com/actions/download-artifact/releases) - [Commits](https://github.com/actions/download-artifact/compare/v7...3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c) Updates `pypa/gh-action-pypi-publish` from 1.13.0 to 1.14.0 - [Release notes](https://github.com/pypa/gh-action-pypi-publish/releases) - [Commits](https://github.com/pypa/gh-action-pypi-publish/compare/v1.13.0...cef221092ed1bacb1cc03d23a2d87d1d172e277b) Updates `zizmorcore/zizmor-action` from 0.5.3 to 0.5.6 - [Release notes](https://github.com/zizmorcore/zizmor-action/releases) - [Commits](https://github.com/zizmorcore/zizmor-action/compare/b1d7e1fb5de872772f31590499237e7cce841e8e...5f14fd08f7cf1cb1609c1e344975f152c7ee938d) --- updated-dependencies: - dependency-name: prefix-dev/setup-pixi dependency-version: 0.9.6 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: codecov/codecov-action dependency-version: 6.0.1 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: github/issue-metrics dependency-version: 4.2.7 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: j178/prek-action dependency-version: 2.0.4 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: actions/upload-artifact dependency-version: 7.0.1 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions - dependency-name: actions/download-artifact dependency-version: 8.0.1 dependency-type: direct:production update-type: version-update:semver-major dependency-group: actions - dependency-name: pypa/gh-action-pypi-publish dependency-version: 1.14.0 dependency-type: direct:production update-type: version-update:semver-minor dependency-group: actions - dependency-name: zizmorcore/zizmor-action dependency-version: 0.5.6 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: actions ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> * fix: don't warn on the spec'd "struct" data type in v3 The `struct` data type now has a stable Zarr V3 specification (zarr-extensions/data-types/struct), so serializing it no longer warrants an UnstableSpecificationWarning. The legacy `structured` alias and the unspecified bytes data types (null_terminated_bytes, raw_bytes, variable_length_bytes) continue to warn. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_019rRWaxXfZ3ZmbiZYWZoDBP * docs: add changelog fragment for struct warning fix Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_019rRWaxXfZ3ZmbiZYWZoDBP --------- Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) --- changes/202.bugfix.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/202.bugfix.md 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. From 45c67bc88441fb6a15cb5fbb348b77518d612c18 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 10:28:32 +0200 Subject: [PATCH 3/5] feat(sharding): manage inner chunks as a semi-regular grid (spec v1.1) Implement zarr-specs #370: the sharding codec's chunk_shape no longer needs to evenly divide the shard shape. Inner chunks are laid out on a semi-regular grid managed by ChunkGrid: chunks per shard is the ceiling division of the shard shape by chunk_shape, and chunks straddling the shard boundary are encoded and stored at their clipped shape. - validate() drops the divisibility requirement (dimension-count and grid-type checks remain) - all encode/decode paths (sync and async, full and partial) use per-chunk ArraySpecs derived from the inner grid - the uncompressed whole-shard bulk-decode fast path gates on an evenly divided grid; the scalar-broadcast write memo only reuses encoded bytes across nominal-shape chunks - Array.read_chunk_sizes and cdata_shape account for the inner grid restarting at every shard boundary - composes with rectilinear outer chunk grids and nested sharding Arrays written with a non-divisible chunk_shape require a v1.1-aware implementation to read; v1.0 implementations must reject them. Assisted-by: ClaudeCode:claude-fable-5 --- changes/4180.feature.md | 10 + docs/user-guide/arrays.md | 13 +- .../v3/codec/sharding_indexed.py | 6 +- src/zarr/codecs/sharding.py | 186 +++++++++++------- src/zarr/core/array.py | 33 +++- tests/test_codecs/test_sharding.py | 149 +++++++++++--- tests/test_unified_chunk_grid.py | 51 ++--- 7 files changed, 321 insertions(+), 127 deletions(-) create mode 100644 changes/4180.feature.md diff --git a/changes/4180.feature.md b/changes/4180.feature.md new file mode 100644 index 0000000000..82143d742b --- /dev/null +++ b/changes/4180.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 a1488f7c30..1ff3f99f5c 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 2d4d63d400..959d46abdc 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). @@ -419,6 +428,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) ) @@ -445,6 +455,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) ) @@ -544,24 +555,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. @@ -643,18 +643,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, @@ -673,7 +672,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=(), @@ -713,13 +712,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 @@ -731,7 +730,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( @@ -762,14 +763,15 @@ 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) + nominal_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), ) ) @@ -802,18 +804,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 is nominal_chunk_spec: if scalar_complete_result is _sentinel: scalar_complete_result = merge_and_encode_chunk( None, @@ -946,18 +951,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, @@ -973,7 +977,7 @@ async def _decode_single( [ ( _ShardingByteGetter(shard_dict, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -992,14 +996,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 @@ -1018,14 +1021,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, @@ -1041,7 +1044,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, @@ -1111,6 +1114,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) @@ -1204,15 +1215,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( @@ -1226,7 +1236,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) @@ -1243,7 +1253,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, @@ -1259,7 +1269,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=(), @@ -1276,15 +1286,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)) @@ -1293,7 +1302,7 @@ async def _encode_single( [ ( _ShardingByteSetter(shard_builder, chunk_coords), - chunk_spec, + get_chunk_spec(chunk_coords), chunk_selection, out_selection, is_complete_shard, @@ -1317,15 +1326,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), ) ) @@ -1334,7 +1342,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) @@ -1347,7 +1355,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, @@ -1388,12 +1396,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( @@ -1481,15 +1489,51 @@ 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. + Chunks with the nominal shape all share one spec object, which keeps + per-spec caches downstream (e.g. `ChunkTransform._resolve_specs`) + effective and makes `spec is nominal_spec` a valid uniformity check. + """ + 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..4b233ec151 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 diff --git a/tests/test_codecs/test_sharding.py b/tests/test_codecs/test_sharding.py index 9e6bebd8df..3086ec2be7 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 @@ -808,34 +809,130 @@ 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, - ) -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, - ) +@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.""" + arr = zarr.create_array( + {}, shape=(10,), shards=[[6, 4]], chunks=(4,), dtype="uint8", fill_value=0 + ) + data = np.arange(10, dtype="uint8") + arr[...] = data + assert np.array_equal(arr[...], data) + assert arr.read_chunk_sizes == ((4, 2, 4),) + + arr[5:9] = 42 + data[5: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_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 @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: From 944dbd66aaf7e04a2b7a96e90975d9e63e537159 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 11:14:13 +0200 Subject: [PATCH 4/5] fix(sharding): address adversarial-review findings on clipped grids Fixes from a multi-agent adversarial review of the semi-regular grid change: - nchunks_initialized used floor division for chunks-per-shard, which undercounts on clipped grids and reports 0 when chunk_shape exceeds the shard shape; use ceiling division to match the inner grid - the scalar-broadcast write memo in _encode_partial_sync gated on ArraySpec object identity, which never holds because _get_chunk_spec is uncached (#3054); gate on shape equality so the memo fires again (one encode instead of one per complete chunk) - the rectilinear clipping test used shard edges [[6, 4]], which normalize to a regular grid; use [[4, 6]] with the rectilinear config flag so RectilinearChunkGridMetadata is actually exercised Adds regression tests for the first two. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/codecs/sharding.py | 12 +++-- src/zarr/core/array.py | 4 +- tests/test_codecs/test_sharding.py | 73 +++++++++++++++++++++++++++--- 3 files changed, 76 insertions(+), 13 deletions(-) diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 959d46abdc..becce96ff2 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -763,7 +763,6 @@ def _encode_partial_sync( """ shard_shape = shard_spec.shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) - nominal_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) @@ -818,7 +817,7 @@ def _encode_partial_sync( for chunk_coords, chunk_sel, out_sel, is_complete_chunk in indexer: chunk_spec = get_chunk_spec(chunk_coords) - if is_scalar and is_complete_chunk and chunk_spec is nominal_chunk_spec: + 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, @@ -1510,9 +1509,12 @@ def _make_chunk_spec_getter( Inner chunks clipped by the shard boundary are encoded at their clipped shape, so their spec differs from the nominal `chunk_shape` spec. - Chunks with the nominal shape all share one spec object, which keeps - per-spec caches downstream (e.g. `ChunkTransform._resolve_specs`) - effective and makes `spec is nominal_spec` a valid uniformity check. + 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): diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 4b233ec151..d8be45be8c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -5345,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 3086ec2be7..e7c3615809 100644 --- a/tests/test_codecs/test_sharding.py +++ b/tests/test_codecs/test_sharding.py @@ -887,17 +887,27 @@ def test_sharding_clipped_inner_chunks_byte_layout() -> None: def test_sharding_clipped_inner_chunks_rectilinear() -> 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.""" - arr = zarr.create_array( - {}, shape=(10,), shards=[[6, 4]], chunks=(4,), dtype="uint8", fill_value=0 - ) + 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, 2, 4),) + assert arr.read_chunk_sizes == ((4, 4, 2),) - arr[5:9] = 42 - data[5:9] = 42 + arr[3:9] = 42 + data[3:9] = 42 assert np.array_equal(arr[...], data) @@ -935,6 +945,55 @@ def test_sharding_clipped_inner_chunk_sizes() -> None: 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"]) def test_sharding_mixed_integer_list_indexing(store: Store) -> None: """Regression test for https://github.com/zarr-developers/zarr-python/issues/3691. From 0ea51c9eaf81b6b0debab87e4886da75eb93f417 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 29 Jul 2026 11:15:21 +0200 Subject: [PATCH 5/5] chore: rename changelog fragment to PR number Assisted-by: ClaudeCode:claude-fable-5 --- changes/{4180.feature.md => 251.feature.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename changes/{4180.feature.md => 251.feature.md} (100%) diff --git a/changes/4180.feature.md b/changes/251.feature.md similarity index 100% rename from changes/4180.feature.md rename to changes/251.feature.md