From 88ab6319eae62f8eeaa8ef2d4488bf0628e88aaf Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 11:53:31 +0100 Subject: [PATCH 01/48] feat(experimental): add ArrayIO and AsyncArrayIO, an explicit-IO array layer Declaring an array today eagerly stores its metadata, so chunks cannot be written first. ArrayIO and AsyncArrayIO bind an array metadata document to a store location without performing any IO on construction; storing metadata (store_metadata) and storing or reading chunks (setitem/getitem) are explicit method calls that run in whatever order the caller chooses, mirroring the array APIs of zarrs and zarrita. The normalization half of init_array is extracted into the pure build_array_metadata, which init_array and AsyncArrayIO.build share; init_array's behavior is unchanged. Refs zarr-developers/zarr-python#4217 Assisted-by: ClaudeCode:claude-fable-5 --- docs/api/zarr/experimental.md | 2 + src/zarr/core/array.py | 48 +++- src/zarr/experimental/__init__.py | 3 +- src/zarr/experimental/array_io.py | 318 +++++++++++++++++++++++ tests/test_experimental/test_array_io.py | 84 ++++++ 5 files changed, 451 insertions(+), 4 deletions(-) create mode 100644 src/zarr/experimental/array_io.py create mode 100644 tests/test_experimental/test_array_io.py diff --git a/docs/api/zarr/experimental.md b/docs/api/zarr/experimental.md index 60f1f987b5..030bb8b714 100644 --- a/docs/api/zarr/experimental.md +++ b/docs/api/zarr/experimental.md @@ -6,4 +6,6 @@ Experimental functionality is not stable and may change or be removed at any poi ## Classes +::: zarr.experimental.array_io + ::: zarr.experimental.cache_store diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index f650388f80..aba50ee14d 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -178,6 +178,7 @@ # Array and AsyncArray are defined in the base `zarr` namespace __all__ = [ "DEFAULT_FILL_VALUE", + "build_array_metadata", "create_codec_pipeline", "parse_array_metadata", ] @@ -4607,6 +4608,49 @@ async def init_array( AsyncArray The AsyncArray. """ + meta = build_array_metadata( + shape=shape, + dtype=dtype, + chunks=chunks, + shards=shards, + filters=filters, + compressors=compressors, + serializer=serializer, + fill_value=fill_value, + order=order, + zarr_format=zarr_format, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + dimension_names=dimension_names, + ) + arr = AsyncArray(metadata=meta, store_path=store_path, config=config) + await save_new_metadata(store_path, meta, overwrite=overwrite, ensure_parents=True) + return arr + + +def build_array_metadata( + *, + shape: ShapeLike, + dtype: ZDTypeLike, + chunks: ChunksLike | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + zarr_format: ZarrFormat | None = 3, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNamesLike = None, +) -> ArrayV2Metadata | ArrayV3Metadata: + """Build an array metadata document from creation-style keyword arguments. + + This is the pure half of `init_array`: it normalizes the same keyword + arguments into an in-memory metadata object without touching any store. + The parameters have the same meaning as in + [`zarr.api.asynchronous.create_array`][]. + """ if zarr_format is None: zarr_format = _default_zarr_format() @@ -4733,9 +4777,7 @@ async def init_array( attributes=attributes, ) - arr = AsyncArray(metadata=meta, store_path=store_path, config=config) - await save_new_metadata(store_path, meta, overwrite=overwrite, ensure_parents=True) - return arr + return meta async def create_array( diff --git a/src/zarr/experimental/__init__.py b/src/zarr/experimental/__init__.py index f7caaf96a1..9a7d50f61f 100644 --- a/src/zarr/experimental/__init__.py +++ b/src/zarr/experimental/__init__.py @@ -1,5 +1,6 @@ """The experimental module is a site for exporting new or experimental Zarr features.""" from zarr.core.chunk_grids import ChunkGrid, ChunkSpec +from zarr.experimental.array_io import ArrayIO, AsyncArrayIO -__all__ = ["ChunkGrid", "ChunkSpec"] +__all__ = ["ArrayIO", "AsyncArrayIO", "ChunkGrid", "ChunkSpec"] diff --git a/src/zarr/experimental/array_io.py b/src/zarr/experimental/array_io.py new file mode 100644 index 0000000000..30aaf46769 --- /dev/null +++ b/src/zarr/experimental/array_io.py @@ -0,0 +1,318 @@ +"""An explicit-IO layer for Zarr arrays. + +`AsyncArrayIO` and `ArrayIO` separate declaring an array from storing it. +Constructing either class performs no IO: an instance binds an array metadata +document to a location in a store, and every read or write of that store is an +explicit method call. In particular, metadata is stored only when +`store_metadata` is called, so chunks can be written before metadata — a +pattern the eager creation API (`zarr.create_array`) cannot express. This +mirrors the array APIs of `zarrs` and `zarrita`, where storing metadata is a +separate step from storing chunks. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +from zarr.core.array import DEFAULT_FILL_VALUE, Array, AsyncArray, build_array_metadata +from zarr.core.sync import sync +from zarr.storage._common import make_store_path + +if TYPE_CHECKING: + import numpy.typing as npt + + from zarr.core.array import ( + CompressorsLike, + FiltersLike, + SerializerLike, + ShardsLike, + ) + from zarr.core.array_spec import ArrayConfig, ArrayConfigLike + from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar + from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike + from zarr.core.common import ( + JSON, + ChunksLike, + DimensionNamesLike, + MemoryOrder, + ShapeLike, + ZarrFormat, + ) + from zarr.core.dtype import ZDTypeLike + from zarr.core.indexing import BasicSelection + from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata + from zarr.storage import StoreLike + from zarr.storage._common import StorePath + from zarr.types import AnyArray, AnyAsyncArray + +__all__ = ["ArrayIO", "AsyncArrayIO"] + + +class AsyncArrayIO: + """Explicit asynchronous IO for one Zarr array. + + An `AsyncArrayIO` binds an array metadata document to a location in a + store. Constructing one — directly, or with `build` — performs no IO: + nothing is read from or written to the store until an IO method is called. + `store_metadata` stores the metadata document, and `setitem` / `getitem` + store and read chunks, in whatever order the caller chooses. + """ + + _array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] + + def __init__( + self, + metadata: ArrayV2Metadata | ArrayV3Metadata, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: + self._array = AsyncArray(metadata=metadata, store_path=store_path, config=config) + + @classmethod + def _wrap(cls, array: AnyAsyncArray) -> AsyncArrayIO: + io = object.__new__(cls) + io._array = array + return io + + @classmethod + async def build( + cls, + store: StoreLike, + *, + name: str | None = None, + shape: ShapeLike, + dtype: ZDTypeLike, + chunks: ChunksLike | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + zarr_format: ZarrFormat | None = 3, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNamesLike = None, + config: ArrayConfigLike | None = None, + ) -> AsyncArrayIO: + """Declare an array bound to a location in a store, without storing anything. + + The keyword arguments have the same meaning as in + [`zarr.api.asynchronous.create_array`][]. The returned instance holds + the normalized metadata in memory; `store_metadata` stores it. + """ + metadata = build_array_metadata( + shape=shape, + dtype=dtype, + chunks=chunks, + shards=shards, + filters=filters, + compressors=compressors, + serializer=serializer, + fill_value=fill_value, + order=order, + zarr_format=zarr_format, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + dimension_names=dimension_names, + ) + store_path = await make_store_path(store, path=name) + return cls(metadata, store_path, config) + + @classmethod + async def open( + cls, + store: StoreLike, + *, + name: str | None = None, + zarr_format: ZarrFormat | None = 3, + config: ArrayConfigLike | None = None, + ) -> AsyncArrayIO: + """Read the metadata of an existing array and bind to it. + + This is the one classmethod constructor that performs IO: the array's + stored metadata document is read from the store. + """ + store_path = await make_store_path(store, path=name) + array = await AsyncArray.open(store_path, zarr_format, config=config) + return cls._wrap(array) + + @property + def metadata(self) -> ArrayMetadata: + """The in-memory metadata this instance is bound to.""" + return self._array.metadata + + @property + def store_path(self) -> StorePath: + """The location in a store this instance is bound to.""" + return self._array.store_path + + @property + def config(self) -> ArrayConfig: + """The runtime configuration of this instance.""" + return self._array.config + + def to_array(self) -> AnyAsyncArray: + """Return an `AsyncArray` handle over the same metadata and location.""" + return self._array + + async def store_metadata(self, *, ensure_parents: bool = False) -> None: + """Store this array's metadata document(s). + + The documents are written unconditionally, replacing the metadata of + any existing node at this location, as `store_metadata` does in + `zarrs`. Guarded creation, which raises when a node already exists, + is [`zarr.api.asynchronous.create_array`][]'s job. With + `ensure_parents`, missing parent groups are stored first. + """ + await self._array._save_metadata(self._array.metadata, ensure_parents=ensure_parents) + + async def getitem( + self, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """Read the selected region, reading absent chunks as the fill value.""" + return await self._array.getitem(selection, prototype=prototype) + + async def setitem( + self, + selection: BasicSelection, + value: npt.ArrayLike, + prototype: BufferPrototype | None = None, + ) -> None: + """Encode `value` and store the chunks covering `selection`.""" + await self._array.setitem(selection, value, prototype=prototype) + + def __repr__(self) -> str: + return f"" + + +class ArrayIO: + """Explicit synchronous IO for one Zarr array. See `AsyncArrayIO`.""" + + _async_io: AsyncArrayIO + + def __init__( + self, + metadata: ArrayV2Metadata | ArrayV3Metadata, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: + self._async_io = AsyncArrayIO(metadata, store_path, config) + + @classmethod + def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: + io = object.__new__(cls) + io._async_io = async_io + return io + + @classmethod + def build( + cls, + store: StoreLike, + *, + name: str | None = None, + shape: ShapeLike, + dtype: ZDTypeLike, + chunks: ChunksLike | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + zarr_format: ZarrFormat | None = 3, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNamesLike = None, + config: ArrayConfigLike | None = None, + ) -> ArrayIO: + """Declare an array bound to a location in a store, without storing anything. + + See `AsyncArrayIO.build`. + """ + return cls._wrap( + sync( + AsyncArrayIO.build( + store, + name=name, + shape=shape, + dtype=dtype, + chunks=chunks, + shards=shards, + filters=filters, + compressors=compressors, + serializer=serializer, + fill_value=fill_value, + order=order, + zarr_format=zarr_format, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + dimension_names=dimension_names, + config=config, + ) + ) + ) + + @classmethod + def open( + cls, + store: StoreLike, + *, + name: str | None = None, + zarr_format: ZarrFormat | None = 3, + config: ArrayConfigLike | None = None, + ) -> ArrayIO: + """Read the metadata of an existing array and bind to it. + + See `AsyncArrayIO.open`. + """ + return cls._wrap( + sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config)) + ) + + @property + def metadata(self) -> ArrayMetadata: + """The in-memory metadata this instance is bound to.""" + return self._async_io.metadata + + @property + def store_path(self) -> StorePath: + """The location in a store this instance is bound to.""" + return self._async_io.store_path + + @property + def config(self) -> ArrayConfig: + """The runtime configuration of this instance.""" + return self._async_io.config + + def to_array(self) -> AnyArray: + """Return an `Array` handle over the same metadata and location.""" + return Array(self._async_io.to_array()) + + def store_metadata(self, *, ensure_parents: bool = False) -> None: + """Store this array's metadata document(s). See `AsyncArrayIO.store_metadata`.""" + sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) + + def getitem( + self, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """Read the selected region, reading absent chunks as the fill value.""" + return sync(self._async_io.getitem(selection, prototype=prototype)) + + def setitem( + self, + selection: BasicSelection, + value: npt.ArrayLike, + prototype: BufferPrototype | None = None, + ) -> None: + """Encode `value` and store the chunks covering `selection`.""" + sync(self._async_io.setitem(selection, value, prototype=prototype)) + + def __repr__(self) -> str: + return f"" diff --git a/tests/test_experimental/test_array_io.py b/tests/test_experimental/test_array_io.py new file mode 100644 index 0000000000..5b570b20d8 --- /dev/null +++ b/tests/test_experimental/test_array_io.py @@ -0,0 +1,84 @@ +"""Tests for the explicit-IO array layer in `zarr.experimental.array_io`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from zarr.api.asynchronous import open_array +from zarr.errors import ArrayNotFoundError +from zarr.experimental import ArrayIO, AsyncArrayIO +from zarr.storage import MemoryStore + +if TYPE_CHECKING: + from zarr.core.common import ZarrFormat + +METADATA_DOCUMENTS = ("zarr.json", ".zarray", ".zattrs", ".zgroup") + + +async def _store_keys(store: MemoryStore) -> set[str]: + return {key async for key in store.list()} + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: + """Declaring an array stores nothing; chunks, then metadata, can be stored + in that order, and the result is an ordinary readable array.""" + store = MemoryStore() + data = np.arange(16, dtype="uint16").reshape(4, 4) + aio = await AsyncArrayIO.build( + store, + name="foo", + shape=data.shape, + dtype=data.dtype, + chunks=(2, 2), + zarr_format=zarr_format, + attributes={"answer": 42}, + ) + assert await _store_keys(store) == set() + + await aio.setitem(Ellipsis, data) + keys = await _store_keys(store) + assert len(keys) == 4 + assert not any(key.split("/")[-1] in METADATA_DOCUMENTS for key in keys) + np.testing.assert_array_equal(await aio.getitem(Ellipsis), data) + + await aio.store_metadata() + opened = await AsyncArrayIO.open(store, name="foo", zarr_format=zarr_format) + assert opened.metadata.attributes == {"answer": 42} + np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) + + arr = await open_array(store=store, path="foo", mode="r", zarr_format=zarr_format) + np.testing.assert_array_equal(await arr.getitem(Ellipsis), data) + + +def test_chunks_before_metadata_sync() -> None: + """The synchronous `ArrayIO` supports the same chunks-first flow.""" + store = MemoryStore() + data = np.arange(16, dtype="uint16").reshape(4, 4) + aio = ArrayIO.build(store, name="foo", shape=data.shape, dtype=data.dtype, chunks=(2, 2)) + aio.setitem(Ellipsis, data) + aio.store_metadata() + opened = ArrayIO.open(store, name="foo") + np.testing.assert_array_equal(opened.getitem(Ellipsis), data) + np.testing.assert_array_equal(opened.to_array()[:], data) + + +def test_build_v2_shards_raises() -> None: + with pytest.raises(ValueError, match="shard_shape"): + ArrayIO.build( + MemoryStore(), + name="foo", + shape=(4,), + dtype="uint8", + chunks=(2,), + shards=(4,), + zarr_format=2, + ) + + +async def test_open_missing_raises() -> None: + with pytest.raises(ArrayNotFoundError): + await AsyncArrayIO.open(MemoryStore(), name="foo") From c15b44a4dce317230bb7395b06f1afaf4d1dc353 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 11:56:37 +0100 Subject: [PATCH 02/48] docs: changelog fragment for #419 Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/419.feature.md diff --git a/changes/419.feature.md b/changes/419.feature.md new file mode 100644 index 0000000000..aa15ec44b5 --- /dev/null +++ b/changes/419.feature.md @@ -0,0 +1 @@ +Added `zarr.experimental.ArrayIO` and `zarr.experimental.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`) and storing or reading chunks (`setitem` / `getitem`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. From a51bdbe664cf96537fb5e19021f8eed2ac75e0c4 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 12:25:16 +0100 Subject: [PATCH 03/48] feat(experimental): give ArrayIO erasing inverses and a format-detecting open MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every write now has an inverse, copying zarrs: store_metadata ↔ erase_metadata, chunk writes ↔ erase_chunk / erase_chunks, each succeeding when the target is already absent. open's zarr_format now defaults to None (read format 3, fall back to format 2), matching zarrs' open, instead of requiring format 3. Adversarial review against zarrs and TensorStore also found contract gaps now documented: the bound metadata is assumed, never verified (TensorStore's assume_metadata hazard); store_metadata keeps any previous array's chunks; a repaired document is stored before the first chunk write; and the module docstring no longer cites zarrita, whose create is eager. Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 2 +- src/zarr/experimental/array_io.py | 106 +++++++++++++++++++++-- tests/test_experimental/test_array_io.py | 35 +++++++- 3 files changed, 133 insertions(+), 10 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index aa15ec44b5..74bac8cfba 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1 +1 @@ -Added `zarr.experimental.ArrayIO` and `zarr.experimental.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`) and storing or reading chunks (`setitem` / `getitem`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. +Added `zarr.experimental.ArrayIO` and `zarr.experimental.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `open` auto-detects the Zarr format by default, as `zarrs` does. diff --git a/src/zarr/experimental/array_io.py b/src/zarr/experimental/array_io.py index 30aaf46769..2f1cd7ddcc 100644 --- a/src/zarr/experimental/array_io.py +++ b/src/zarr/experimental/array_io.py @@ -6,15 +6,36 @@ explicit method call. In particular, metadata is stored only when `store_metadata` is called, so chunks can be written before metadata — a pattern the eager creation API (`zarr.create_array`) cannot express. This -mirrors the array APIs of `zarrs` and `zarrita`, where storing metadata is a -separate step from storing chunks. +mirrors the array API of `zarrs`, where storing, retrieving and erasing +metadata and chunks are each explicit operations. + +Every write has an erasing inverse: `store_metadata` ↔ `erase_metadata`, +and `setitem` ↔ `erase_chunk` / `erase_chunks`. `store_metadata` is the +publication point: until it is called, no Zarr reader resolves the array, so +a failed chunks-first run leaves no readable array. It does leave chunk keys +in the store, which occupy space and would be misread by a different array +stored later at the same location; `erase_chunks` removes them. + +One exception to "no implicit metadata IO": an array opened from a stored +document that needed a repair stores that repair before its first chunk +write, exactly as `AsyncArray` does. Call `store_metadata` first to make +that write explicit. """ from __future__ import annotations from typing import TYPE_CHECKING, Any, Literal -from zarr.core.array import DEFAULT_FILL_VALUE, Array, AsyncArray, build_array_metadata +from zarr.core.array import ( + DEFAULT_FILL_VALUE, + Array, + AsyncArray, + _iter_shard_keys, + build_array_metadata, +) +from zarr.core.common import concurrent_foreach, product +from zarr.core.config import config as zarr_config +from zarr.core.metadata.io import ARRAY_DOCUMENTS from zarr.core.sync import sync from zarr.storage._common import make_store_path @@ -48,6 +69,10 @@ __all__ = ["ArrayIO", "AsyncArrayIO"] +async def _delete_key(store_path: StorePath, key: str) -> None: + await (store_path / key).delete() + + class AsyncArrayIO: """Explicit asynchronous IO for one Zarr array. @@ -56,6 +81,12 @@ class AsyncArrayIO: nothing is read from or written to the store until an IO method is called. `store_metadata` stores the metadata document, and `setitem` / `getitem` store and read chunks, in whatever order the caller chooses. + + Warning: the bound metadata is assumed, never verified against the store. + Reads and partial writes decode stored chunks under it, so binding + metadata that disagrees with how existing chunks at this location were + encoded misreads or corrupts data — the hazard TensorStore documents for + `assume_metadata`. """ _array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] @@ -125,13 +156,15 @@ async def open( store: StoreLike, *, name: str | None = None, - zarr_format: ZarrFormat | None = 3, + zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> AsyncArrayIO: """Read the metadata of an existing array and bind to it. This is the one classmethod constructor that performs IO: the array's - stored metadata document is read from the store. + stored metadata document is read from the store. `zarr_format=None` + (the default) reads Zarr format 3 and falls back to format 2, as + `open` does in `zarrs`; pass 2 or 3 to require one format. """ store_path = await make_store_path(store, path=name) array = await AsyncArray.open(store_path, zarr_format, config=config) @@ -162,11 +195,56 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: The documents are written unconditionally, replacing the metadata of any existing node at this location, as `store_metadata` does in `zarrs`. Guarded creation, which raises when a node already exists, - is [`zarr.api.asynchronous.create_array`][]'s job. With - `ensure_parents`, missing parent groups are stored first. + is [`zarr.api.asynchronous.create_array`][]'s job. Chunks of any + previous array at this location are kept and would be read under the + new metadata; call `erase_chunks` first when repurposing a location. + With `ensure_parents`, missing parent groups are stored first. """ await self._array._save_metadata(self._array.metadata, ensure_parents=ensure_parents) + async def erase_metadata(self) -> None: + """Erase this array's metadata document(s), the inverse of `store_metadata`. + + Succeeds if the documents do not exist, as `erase_metadata` does in + `zarrs`. Stored chunks are kept; `erase_chunks` removes those. + """ + store_path = self._array.store_path + for document in ARRAY_DOCUMENTS[self.metadata.zarr_format]: + await (store_path / document).delete() + self._array._stored_document_replaced() + + async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: + """Erase one stored chunk by its chunk grid coordinates. + + When sharding is used, the stored chunk is the shard. Succeeds if the + chunk does not exist, as `erase_chunk` does in `zarrs`. + """ + grid_shape = self._array._shard_grid_shape + if len(chunk_coords) != len(grid_shape) or any( + coord < 0 or coord >= size for coord, size in zip(chunk_coords, grid_shape, strict=True) + ): + raise IndexError( + f"Chunk coordinates {chunk_coords} are outside the chunk grid " + f"of shape {grid_shape}." + ) + await _delete_key(self._array.store_path, self.metadata.encode_chunk_key(chunk_coords)) + + async def erase_chunks(self) -> None: + """Erase every stored chunk of this array, the inverse of its chunk writes. + + Chunks that do not exist are skipped; the metadata documents are kept. + """ + n_chunks = product(self._array._shard_grid_shape) + if n_chunks == 0: + return + limit = zarr_config.get("async.concurrency") or 1000 + store_path = self._array.store_path + await concurrent_foreach( + ((store_path, key) for key in _iter_shard_keys(self._array)), + _delete_key, + min(limit, n_chunks), + ) + async def getitem( self, selection: BasicSelection, @@ -262,7 +340,7 @@ def open( store: StoreLike, *, name: str | None = None, - zarr_format: ZarrFormat | None = 3, + zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> ArrayIO: """Read the metadata of an existing array and bind to it. @@ -296,6 +374,18 @@ def store_metadata(self, *, ensure_parents: bool = False) -> None: """Store this array's metadata document(s). See `AsyncArrayIO.store_metadata`.""" sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) + def erase_metadata(self) -> None: + """Erase this array's metadata document(s). See `AsyncArrayIO.erase_metadata`.""" + sync(self._async_io.erase_metadata()) + + def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: + """Erase one stored chunk by its chunk grid coordinates. See `AsyncArrayIO.erase_chunk`.""" + sync(self._async_io.erase_chunk(chunk_coords)) + + def erase_chunks(self) -> None: + """Erase every stored chunk of this array. See `AsyncArrayIO.erase_chunks`.""" + sync(self._async_io.erase_chunks()) + def getitem( self, selection: BasicSelection, diff --git a/tests/test_experimental/test_array_io.py b/tests/test_experimental/test_array_io.py index 5b570b20d8..615c697a77 100644 --- a/tests/test_experimental/test_array_io.py +++ b/tests/test_experimental/test_array_io.py @@ -46,13 +46,33 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: np.testing.assert_array_equal(await aio.getitem(Ellipsis), data) await aio.store_metadata() - opened = await AsyncArrayIO.open(store, name="foo", zarr_format=zarr_format) + # the default zarr_format=None auto-detects the format, like zarrs + opened = await AsyncArrayIO.open(store, name="foo") + assert opened.metadata.zarr_format == zarr_format assert opened.metadata.attributes == {"answer": 42} np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) arr = await open_array(store=store, path="foo", mode="r", zarr_format=zarr_format) np.testing.assert_array_equal(await arr.getitem(Ellipsis), data) + # every write has an erasing inverse: one chunk, all chunks, then metadata + await opened.erase_chunk((0, 0)) + partially_erased = np.asarray(await opened.getitem(Ellipsis)) + np.testing.assert_array_equal(partially_erased[:2, :2], np.zeros((2, 2), dtype="uint16")) + np.testing.assert_array_equal(partially_erased[2:, :], data[2:, :]) + + await opened.erase_chunks() + await opened.erase_chunks() # erasing absent chunks succeeds + np.testing.assert_array_equal(await opened.getitem(Ellipsis), np.zeros_like(data)) + keys = await _store_keys(store) + assert all(key.split("/")[-1] in METADATA_DOCUMENTS for key in keys) + + await opened.erase_metadata() + await opened.erase_metadata() # erasing absent metadata succeeds + assert await _store_keys(store) == set() + with pytest.raises(ArrayNotFoundError): + await AsyncArrayIO.open(store, name="foo") + def test_chunks_before_metadata_sync() -> None: """The synchronous `ArrayIO` supports the same chunks-first flow.""" @@ -64,6 +84,11 @@ def test_chunks_before_metadata_sync() -> None: opened = ArrayIO.open(store, name="foo") np.testing.assert_array_equal(opened.getitem(Ellipsis), data) np.testing.assert_array_equal(opened.to_array()[:], data) + opened.erase_chunk((0, 0)) + opened.erase_chunks() + opened.erase_metadata() + with pytest.raises(ArrayNotFoundError): + ArrayIO.open(store, name="foo") def test_build_v2_shards_raises() -> None: @@ -82,3 +107,11 @@ def test_build_v2_shards_raises() -> None: async def test_open_missing_raises() -> None: with pytest.raises(ArrayNotFoundError): await AsyncArrayIO.open(MemoryStore(), name="foo") + + +async def test_erase_chunk_out_of_grid_raises() -> None: + aio = await AsyncArrayIO.build(MemoryStore(), shape=(4, 4), dtype="uint8", chunks=(2, 2)) + with pytest.raises(IndexError, match="outside the chunk grid"): + await aio.erase_chunk((2, 0)) + with pytest.raises(IndexError, match="outside the chunk grid"): + await aio.erase_chunk((0,)) From c618bbef32b2173861606126885898b414e8e3be Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 12:29:19 +0100 Subject: [PATCH 04/48] feat(experimental): let ArrayIO verify its bound metadata against the store MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A writer that declared an array independently (the distributed chunks-first pattern: one coordinator stores metadata, N workers build the same declaration) can assert its assumption before writing chunks. verify() reads the stored document, favoring format 3 when both exist, and raises naming each differing field — the explicit counterpart of TensorStore's spec-constraint validation on open. Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 1 + src/zarr/experimental/array_io.py | 45 ++++++++++++++++++++++++ tests/test_experimental/test_array_io.py | 39 ++++++++++++++++++++ 3 files changed, 85 insertions(+) diff --git a/changes/419.feature.md b/changes/419.feature.md index 74bac8cfba..24f9990189 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1 +1,2 @@ Added `zarr.experimental.ArrayIO` and `zarr.experimental.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `open` auto-detects the Zarr format by default, as `zarrs` does. +`verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. diff --git a/src/zarr/experimental/array_io.py b/src/zarr/experimental/array_io.py index 2f1cd7ddcc..a0b30d4ee0 100644 --- a/src/zarr/experimental/array_io.py +++ b/src/zarr/experimental/array_io.py @@ -32,6 +32,8 @@ AsyncArray, _iter_shard_keys, build_array_metadata, + get_array_metadata, + parse_array_metadata, ) from zarr.core.common import concurrent_foreach, product from zarr.core.config import config as zarr_config @@ -202,6 +204,42 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: """ await self._array._save_metadata(self._array.metadata, ensure_parents=ensure_parents) + async def verify(self) -> None: + """Check that the stored metadata declares the array this instance is bound to. + + Reads the metadata stored at this location — favoring Zarr format 3 + when documents of both formats exist — and raises `ValueError` naming + each differing field when it does not equal the bound metadata. + Raises [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. + This is the explicit counterpart of the verification TensorStore + performs against spec constraints on open: a writer that declared its + own metadata with `build` can assert, before writing chunks, that it + agrees with what a coordinator stored. + """ + bound = self.metadata + stored = parse_array_metadata( + await get_array_metadata(self._array.store_path, zarr_format=None), + str(self._array.store_path), + ) + if stored == bound: + return + if type(stored) is not type(bound): + raise ValueError( + f"Array {str(self.store_path)!r}: the stored array is Zarr format " + f"{stored.zarr_format}; the bound metadata is Zarr format " + f"{bound.zarr_format}." + ) + stored_dict, bound_dict = stored.to_dict(), bound.to_dict() + differing = sorted( + key + for key in stored_dict.keys() | bound_dict.keys() + if stored_dict.get(key) != bound_dict.get(key) + ) + raise ValueError( + f"Array {str(self.store_path)!r}: the stored metadata does not match the " + f"bound metadata; differing fields: {', '.join(differing)}." + ) + async def erase_metadata(self) -> None: """Erase this array's metadata document(s), the inverse of `store_metadata`. @@ -374,6 +412,13 @@ def store_metadata(self, *, ensure_parents: bool = False) -> None: """Store this array's metadata document(s). See `AsyncArrayIO.store_metadata`.""" sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) + def verify(self) -> None: + """Check that the stored metadata declares the array this instance is bound to. + + See `AsyncArrayIO.verify`. + """ + sync(self._async_io.verify()) + def erase_metadata(self) -> None: """Erase this array's metadata document(s). See `AsyncArrayIO.erase_metadata`.""" sync(self._async_io.erase_metadata()) diff --git a/tests/test_experimental/test_array_io.py b/tests/test_experimental/test_array_io.py index 615c697a77..df1af45f2f 100644 --- a/tests/test_experimental/test_array_io.py +++ b/tests/test_experimental/test_array_io.py @@ -55,6 +55,18 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: arr = await open_array(store=store, path="foo", mode="r", zarr_format=zarr_format) np.testing.assert_array_equal(await arr.getitem(Ellipsis), data) + # a writer that declared the same array independently can verify its assumption + worker = await AsyncArrayIO.build( + store, + name="foo", + shape=data.shape, + dtype=data.dtype, + chunks=(2, 2), + zarr_format=zarr_format, + attributes={"answer": 42}, + ) + await worker.verify() + # every write has an erasing inverse: one chunk, all chunks, then metadata await opened.erase_chunk((0, 0)) partially_erased = np.asarray(await opened.getitem(Ellipsis)) @@ -109,6 +121,33 @@ async def test_open_missing_raises() -> None: await AsyncArrayIO.open(MemoryStore(), name="foo") +async def test_verify_mismatch_raises() -> None: + store = MemoryStore() + aio = await AsyncArrayIO.build(store, shape=(4, 4), dtype="uint8", chunks=(2, 2)) + await aio.store_metadata() + other = await AsyncArrayIO.build( + store, shape=(4, 4), dtype="uint8", chunks=(4, 4), attributes={"a": 1} + ) + with pytest.raises(ValueError, match="differing fields: attributes, chunk_grid"): + await other.verify() + + +async def test_verify_format_mismatch_raises() -> None: + store = MemoryStore() + await ( + await AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=2) + ).store_metadata() + bound_v3 = await AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=3) + with pytest.raises(ValueError, match="stored array is Zarr format 2"): + await bound_v3.verify() + + +async def test_verify_missing_raises() -> None: + aio = await AsyncArrayIO.build(MemoryStore(), shape=(4,), dtype="uint8") + with pytest.raises(ArrayNotFoundError): + await aio.verify() + + async def test_erase_chunk_out_of_grid_raises() -> None: aio = await AsyncArrayIO.build(MemoryStore(), shape=(4, 4), dtype="uint8", chunks=(2, 2)) with pytest.raises(IndexError, match="outside the chunk grid"): From 102a1addaa52732c32bf094d0f576222fbb36afc Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 13:04:28 +0100 Subject: [PATCH 05/48] refactor(array): make AsyncArray wrap AsyncArrayIO The explicit-IO layer is the lower layer, so it moves into zarr.core.array as a self-contained class holding its own (metadata, store path, config, codec pipeline, chunk grid) state, with zarr.experimental.array_io re-exporting it. AsyncArray now holds one field, the wrapped AsyncArrayIO, exposes metadata/store_path/config/ codec_pipeline/_chunk_grid as properties over it, and delegates its metadata and selection IO to it; resize rebinds the wrapped state, so the two layers cannot drift. AsyncArrayIO defines value equality over the full binding, preserving AsyncArray's previous dataclass equality, and to_array() returns an AsyncArray sharing the same binding. The repair-before-chunk-write hook and the stored-document mark reset become module functions used by the one remaining implementation. Refs zarr-developers/zarr-python#4217 Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 683 +++++++++++++++++++++++++----- src/zarr/experimental/array_io.py | 442 +------------------ 2 files changed, 580 insertions(+), 545 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index aba50ee14d..61f75e38e2 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -6,7 +6,7 @@ import warnings from asyncio import gather from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, replace from itertools import starmap from logging import getLogger from typing import ( @@ -178,6 +178,8 @@ # Array and AsyncArray are defined in the base `zarr` namespace __all__ = [ "DEFAULT_FILL_VALUE", + "ArrayIO", + "AsyncArrayIO", "build_array_metadata", "create_codec_pipeline", "parse_array_metadata", @@ -393,11 +395,7 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: The runtime configuration of the array. """ - metadata: T_ArrayMetadata - store_path: StorePath - codec_pipeline: CodecPipeline = field(init=False) - _chunk_grid: ChunkGrid = field(init=False) - config: ArrayConfig + _io: AsyncArrayIO @overload def __init__( @@ -421,18 +419,38 @@ def __init__( store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: - metadata_parsed = parse_array_metadata(metadata, str(store_path)) - config_parsed = parse_array_config(config) + object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config)) - object.__setattr__(self, "metadata", metadata_parsed) - object.__setattr__(self, "store_path", store_path) - object.__setattr__(self, "config", config_parsed) - object.__setattr__(self, "_chunk_grid", ChunkGrid.from_metadata(metadata_parsed)) - object.__setattr__( - self, - "codec_pipeline", - create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), - ) + @classmethod + def _from_io(cls, io: AsyncArrayIO) -> AnyAsyncArray: + """An array over an existing binding: the array and `io` share one state.""" + self = object.__new__(cls) + object.__setattr__(self, "_io", io) + return self + + @property + def metadata(self) -> T_ArrayMetadata: + """The in-memory metadata of this array.""" + return cast("T_ArrayMetadata", self._io.metadata) + + @property + def store_path(self) -> StorePath: + """The location in a store this array is bound to.""" + return self._io.store_path + + @property + def config(self) -> ArrayConfig: + """The runtime configuration of this array.""" + return self._io.config + + @property + def codec_pipeline(self) -> CodecPipeline: + """The codec pipeline used for encoding and decoding chunks.""" + return self._io._codec_pipeline + + @property + def _chunk_grid(self) -> ChunkGrid: + return self._io._chunk_grid @classmethod async def _create( @@ -1487,17 +1505,7 @@ async def _get_selection( out: NDBuffer | None = None, fields: Fields | None = None, ) -> NDArrayLikeOrScalar: - return await _get_selection( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - indexer, - prototype=prototype, - out=out, - fields=fields, - ) + return await self._io._get_selection(indexer, prototype=prototype, out=out, fields=fields) async def getitem( self, @@ -1541,15 +1549,7 @@ async def getitem( np.int32(0) """ - return await _getitem( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - selection, - prototype=prototype, - ) + return await self._io.getitem(selection, prototype=prototype) async def get_orthogonal_selection( self, @@ -1624,57 +1624,20 @@ async def get_coordinate_selection( out_array = cast("NDArrayLikeOrScalar", np.array(out_array).reshape(indexer.sel_shape)) return out_array - async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None: - """Store `metadata` as this array's own documents, then clear the - `_stored_document` mark (see `_stored_document_replaced`). `_resize` stores the - documents it encoded before deleting chunks directly, then clears the mark the - same way.""" - await save_metadata(self.store_path, metadata, ensure_parents=ensure_parents) - self._stored_document_replaced() + async def _save_metadata(self, ensure_parents: bool = False) -> None: + """Store this array's metadata documents, delegating to the wrapped + `AsyncArrayIO`, which clears the `_stored_document` mark (see + `_stored_document_replaced`). `_resize` stores the documents it encoded before + deleting chunks directly, then clears the mark the same way.""" + await self._io.store_metadata(ensure_parents=ensure_parents) def _stored_document_replaced(self) -> None: - """Record that the store no longer holds a document of this array that needs an - repair: it holds the repair, a valid document, or none. The metadata this handle - holds, which a consolidated group handle may share, then stops standing for the - document it was read from (see `mark_repaired`), so no later write through either - handle stores that document again.""" - object.__setattr__(self.metadata, "_stored_document", None) + """See the module-level `_stored_document_replaced`.""" + _stored_document_replaced(self.metadata) async def _store_repaired_document(self) -> None: - """Store the repair of this array's current stored document, if it needs one, - before chunks are written under this handle's metadata. - - Only for metadata read from a document whose repair moves chunks (see - `zarr.core.metadata.repair`). The document is read again, because the store may - hold a newer one than this handle's metadata. If that one lays out chunks - differently (another writer stored a different chunk size since), this handle - would write chunks no reader finds, so it raises and stores nothing. If it needs - no repair (the array was re-saved since, possibly by another implementation), it - is left as written; if there is none, there is nothing to repair. - - Storing the same repair twice is harmless, so handles that write chunks - concurrently need no coordination. The read and the store are not one atomic - step, though: a metadata write by another handle between them (a resize, an - attribute update) is replaced by the repair of the document read before it, as - with any two metadata writes that race. - """ - if self.metadata._stored_document is None: - return - zarr_format = self.metadata.zarr_format - documents = await read_documents(self.store_path, ARRAY_DOCUMENTS[zarr_format]) - try: - current = parse_stored_array(documents, zarr_format, str(self.store_path)) - except ArrayNotFoundError: - pass - else: - if _chunk_layout(current) != _chunk_layout(self.metadata): - raise ValueError( - f"Array {str(self.store_path)!r}: the metadata stored has changed since " - "this array was opened; reopen the array to write to it. Nothing was stored." - ) - if current._stored_document is not None: - await upsert_metadata(self.store_path, current, documents) - self._stored_document_replaced() + """See the module-level `_store_repaired_document`.""" + await _store_repaired_document(self.store_path, self.metadata) async def _set_selection( self, @@ -1684,21 +1647,7 @@ async def _set_selection( prototype: BufferPrototype, fields: Fields | None = None, ) -> None: - if product(indexer.shape) > 0: - # Chunks are about to be stored under the repaired metadata, so store it - # first: every reader of the store then agrees with them. - await self._store_repaired_document() - return await _set_selection( - self.store_path, - self.metadata, - self.codec_pipeline, - self.config, - self._chunk_grid, - indexer, - value, - prototype=prototype, - fields=fields, - ) + return await self._io._set_selection(indexer, value, prototype=prototype, fields=fields) async def setitem( self, @@ -1739,10 +1688,7 @@ async def setitem( - This method is asynchronous and should be awaited. - Supports basic indexing, where the selection is contiguous and does not involve advanced indexing. """ - if prototype is None: - prototype = default_buffer_prototype() - indexer = BasicIndexer(selection, shape=self.metadata.shape, chunk_grid=self._chunk_grid) - return await self._set_selection(indexer, value, prototype=prototype) + return await self._io.setitem(selection, value, prototype=prototype) @property def oindex(self) -> AsyncOIndex[T_ArrayMetadata]: @@ -5002,6 +4948,52 @@ def _sharding_codec(metadata: ArrayMetadata) -> ShardingCodec | None: return None +def _stored_document_replaced(metadata: ArrayMetadata) -> None: + """Record that the store no longer holds a document of this array that needs a + repair: it holds the repair, a valid document, or none. The metadata, which a + consolidated group handle may share with an array handle, then stops standing for + the document it was read from (see `mark_repaired`), so no later write through any + handle stores that document again.""" + object.__setattr__(metadata, "_stored_document", None) + + +async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadata) -> None: + """Store the repair of the array's current stored document, if it needs one, + before chunks are written under `metadata`. + + Only for metadata read from a document whose repair moves chunks (see + `zarr.core.metadata.repair`). The document is read again, because the store may + hold a newer one than `metadata`. If that one lays out chunks differently + (another writer stored a different chunk size since), the caller would write + chunks no reader finds, so it raises and stores nothing. If it needs no repair + (the array was re-saved since, possibly by another implementation), it is left as + written; if there is none, there is nothing to repair. + + Storing the same repair twice is harmless, so handles that write chunks + concurrently need no coordination. The read and the store are not one atomic + step, though: a metadata write by another handle between them (a resize, an + attribute update) is replaced by the repair of the document read before it, as + with any two metadata writes that race. + """ + if metadata._stored_document is None: + return + zarr_format = metadata.zarr_format + documents = await read_documents(store_path, ARRAY_DOCUMENTS[zarr_format]) + try: + current = parse_stored_array(documents, zarr_format, str(store_path)) + except ArrayNotFoundError: + pass + else: + if _chunk_layout(current) != _chunk_layout(metadata): + raise ValueError( + f"Array {str(store_path)!r}: the metadata stored has changed since " + "this array was opened; reopen the array to write to it. Nothing was stored." + ) + if current._stored_document is not None: + await upsert_metadata(store_path, current, documents) + _stored_document_replaced(metadata) + + def _stored_rectilinear_grid_or_none( metadata: ArrayMetadata, ) -> RectilinearChunkGridMetadata | None: @@ -6055,9 +6047,9 @@ async def _delete_chunk(coords: tuple[int, ...]) -> None: await store_documents(array.store_path, documents) array._stored_document_replaced() - # Update metadata and chunk_grid (in place) - object.__setattr__(array, "metadata", new_metadata) - object.__setattr__(array, "_chunk_grid", new_chunk_grid) + # Update the wrapped binding in place; the codec pipeline is kept, as before. + array._io._metadata = new_metadata + array._io._chunk_grid = new_chunk_grid async def _append( @@ -6145,7 +6137,7 @@ async def _update_attributes( array.metadata.attributes.update(new_attributes) # Write new metadata - await array._save_metadata(array.metadata) + await array._save_metadata() return array @@ -6172,3 +6164,480 @@ async def _info_complete( await _nshards_initialized(array), await array.store_path.store.getsize_prefix(array.store_path.path), ) + + +async def _delete_key(store_path: StorePath, key: str) -> None: + await (store_path / key).delete() + + +class AsyncArrayIO: + """Explicit asynchronous IO for one Zarr array. + + An `AsyncArrayIO` binds an array metadata document to a location in a + store. Constructing one — directly, or with `build` — performs no IO: + nothing is read from or written to the store until an IO method is called. + `store_metadata` stores the metadata document, and `setitem` / `getitem` + store and read chunks, in whatever order the caller chooses. Every write + has an erasing inverse: `store_metadata` ↔ `erase_metadata`, and chunk + writes ↔ `erase_chunk` / `erase_chunks`. `AsyncArray` is an array-like + layer over this one: it wraps an `AsyncArrayIO` and delegates its metadata + and selection IO to it. + + `store_metadata` is the publication point: until it is called, no Zarr + reader resolves the array, so a failed chunks-first run leaves no readable + array. It does leave chunk keys in the store, which occupy space and would + be misread by a different array stored later at the same location; + `erase_chunks` removes them. + + Warning: the bound metadata is assumed, never verified against the store. + Reads and partial writes decode stored chunks under it, so binding + metadata that disagrees with how existing chunks at this location were + encoded misreads or corrupts data — the hazard TensorStore documents for + `assume_metadata`. One exception to "no implicit metadata IO": an array + opened from a stored document that needed a repair stores that repair + before its first chunk write. Call `store_metadata` first to make that + write explicit. + """ + + _metadata: ArrayMetadata + _store_path: StorePath + _config: ArrayConfig + _codec_pipeline: CodecPipeline + _chunk_grid: ChunkGrid + + def __init__( + self, + metadata: ArrayMetadata | ArrayMetadataDict, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: + metadata_parsed = parse_array_metadata(metadata, str(store_path)) + self._metadata = metadata_parsed + self._store_path = store_path + self._config = parse_array_config(config) + self._chunk_grid = ChunkGrid.from_metadata(metadata_parsed) + self._codec_pipeline = create_codec_pipeline( + metadata=metadata_parsed, store=store_path.store + ) + + @classmethod + async def build( + cls, + store: StoreLike, + *, + name: str | None = None, + shape: ShapeLike, + dtype: ZDTypeLike, + chunks: ChunksLike | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + zarr_format: ZarrFormat | None = 3, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNamesLike = None, + config: ArrayConfigLike | None = None, + ) -> AsyncArrayIO: + """Declare an array bound to a location in a store, without storing anything. + + The keyword arguments have the same meaning as in + [`zarr.api.asynchronous.create_array`][]. The returned instance holds + the normalized metadata in memory; `store_metadata` stores it. + """ + metadata = build_array_metadata( + shape=shape, + dtype=dtype, + chunks=chunks, + shards=shards, + filters=filters, + compressors=compressors, + serializer=serializer, + fill_value=fill_value, + order=order, + zarr_format=zarr_format, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + dimension_names=dimension_names, + ) + store_path = await make_store_path(store, path=name) + return cls(metadata, store_path, config) + + @classmethod + async def open( + cls, + store: StoreLike, + *, + name: str | None = None, + zarr_format: ZarrFormat | None = None, + config: ArrayConfigLike | None = None, + ) -> AsyncArrayIO: + """Read the metadata of an existing array and bind to it. + + This is the one classmethod constructor that performs IO: the array's + stored metadata document is read from the store. `zarr_format=None` + (the default) reads Zarr format 3 and falls back to format 2, as + `open` does in `zarrs`; pass 2 or 3 to require one format. + """ + store_path = await make_store_path(store, path=name) + metadata = parse_array_metadata( + await get_array_metadata(store_path, zarr_format=zarr_format), str(store_path) + ) + return cls(metadata, store_path, config) + + @property + def metadata(self) -> ArrayMetadata: + """The in-memory metadata this instance is bound to.""" + return self._metadata + + @property + def store_path(self) -> StorePath: + """The location in a store this instance is bound to.""" + return self._store_path + + @property + def config(self) -> ArrayConfig: + """The runtime configuration of this instance.""" + return self._config + + @property + def _shard_grid_shape(self) -> tuple[int, ...]: + # The stored chunk grid is the shard grid when sharding is used, the + # chunk grid otherwise, so its keys are the stored chunk keys. + return self._chunk_grid.grid_shape + + def to_array(self) -> AnyAsyncArray: + """Return an `AsyncArray` over this binding; both share one state.""" + return AsyncArray._from_io(self) + + async def store_metadata(self, *, ensure_parents: bool = False) -> None: + """Store this array's metadata document(s). + + The documents are written unconditionally, replacing the metadata of + any existing node at this location, as `store_metadata` does in + `zarrs`. Guarded creation, which raises when a node already exists, + is [`zarr.api.asynchronous.create_array`][]'s job. Chunks of any + previous array at this location are kept and would be read under the + new metadata; call `erase_chunks` first when repurposing a location. + With `ensure_parents`, missing parent groups are stored first. + """ + await save_metadata(self._store_path, self._metadata, ensure_parents=ensure_parents) + _stored_document_replaced(self._metadata) + + async def verify(self) -> None: + """Check that the stored metadata declares the array this instance is bound to. + + Reads the metadata stored at this location — favoring Zarr format 3 + when documents of both formats exist — and raises `ValueError` naming + each differing field when it does not equal the bound metadata. + Raises [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. + This is the explicit counterpart of the verification TensorStore + performs against spec constraints on open: a writer that declared its + own metadata with `build` can assert, before writing chunks, that it + agrees with what a coordinator stored. + """ + bound = self._metadata + stored = parse_array_metadata( + await get_array_metadata(self._store_path, zarr_format=None), str(self._store_path) + ) + if stored == bound: + return + if type(stored) is not type(bound): + raise ValueError( + f"Array {str(self.store_path)!r}: the stored array is Zarr format " + f"{stored.zarr_format}; the bound metadata is Zarr format " + f"{bound.zarr_format}." + ) + stored_dict, bound_dict = stored.to_dict(), bound.to_dict() + differing = sorted( + key + for key in stored_dict.keys() | bound_dict.keys() + if stored_dict.get(key) != bound_dict.get(key) + ) + raise ValueError( + f"Array {str(self.store_path)!r}: the stored metadata does not match the " + f"bound metadata; differing fields: {', '.join(differing)}." + ) + + async def erase_metadata(self) -> None: + """Erase this array's metadata document(s), the inverse of `store_metadata`. + + Succeeds if the documents do not exist, as `erase_metadata` does in + `zarrs`. Stored chunks are kept; `erase_chunks` removes those. + """ + for document in ARRAY_DOCUMENTS[self._metadata.zarr_format]: + await (self._store_path / document).delete() + _stored_document_replaced(self._metadata) + + async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: + """Erase one stored chunk by its chunk grid coordinates. + + When sharding is used, the stored chunk is the shard. Succeeds if the + chunk does not exist, as `erase_chunk` does in `zarrs`. + """ + grid_shape = self._shard_grid_shape + if len(chunk_coords) != len(grid_shape) or any( + coord < 0 or coord >= size for coord, size in zip(chunk_coords, grid_shape, strict=True) + ): + raise IndexError( + f"Chunk coordinates {chunk_coords} are outside the chunk grid " + f"of shape {grid_shape}." + ) + await _delete_key(self._store_path, self._metadata.encode_chunk_key(chunk_coords)) + + async def erase_chunks(self) -> None: + """Erase every stored chunk of this array, the inverse of its chunk writes. + + Chunks that do not exist are skipped; the metadata documents are kept. + """ + grid_shape = self._shard_grid_shape + n_chunks = product(grid_shape) + if n_chunks == 0: + return + limit = zarr_config.get("async.concurrency") or 1000 + await concurrent_foreach( + ( + (self._store_path, self._metadata.encode_chunk_key(coords)) + for coords in _iter_grid(grid_shape) + ), + _delete_key, + min(limit, n_chunks), + ) + + async def _get_selection( + self, + indexer: Indexer, + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, + fields: Fields | None = None, + ) -> NDArrayLikeOrScalar: + return await _get_selection( + self._store_path, + self._metadata, + self._codec_pipeline, + self._config, + self._chunk_grid, + indexer, + prototype=prototype, + out=out, + fields=fields, + ) + + async def _set_selection( + self, + indexer: Indexer, + value: npt.ArrayLike, + *, + prototype: BufferPrototype, + fields: Fields | None = None, + ) -> None: + if product(indexer.shape) > 0: + # Chunks are about to be stored under the repaired metadata, so store it + # first: every reader of the store then agrees with them. + await _store_repaired_document(self._store_path, self._metadata) + return await _set_selection( + self._store_path, + self._metadata, + self._codec_pipeline, + self._config, + self._chunk_grid, + indexer, + value, + prototype=prototype, + fields=fields, + ) + + async def getitem( + self, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """Read the selected region, reading absent chunks as the fill value.""" + if prototype is None: + prototype = default_buffer_prototype() + indexer = BasicIndexer(selection, shape=self._metadata.shape, chunk_grid=self._chunk_grid) + return await self._get_selection(indexer, prototype=prototype) + + async def setitem( + self, + selection: BasicSelection, + value: npt.ArrayLike, + prototype: BufferPrototype | None = None, + ) -> None: + """Encode `value` and store the chunks covering `selection`.""" + if prototype is None: + prototype = default_buffer_prototype() + indexer = BasicIndexer(selection, shape=self._metadata.shape, chunk_grid=self._chunk_grid) + return await self._set_selection(indexer, value, prototype=prototype) + + def __repr__(self) -> str: + return f"" + + def __eq__(self, other: object) -> bool: + # Value equality over the full binding, matching what the dataclass-generated + # equality of AsyncArray compared when it held these as its own fields. + if not isinstance(other, AsyncArrayIO): + return NotImplemented + return ( + self._metadata == other._metadata + and self._store_path == other._store_path + and self._config == other._config + and self._codec_pipeline == other._codec_pipeline + and self._chunk_grid == other._chunk_grid + ) + + def __hash__(self) -> int: + # A subset of the equality fields: equal bindings hash equally. + return hash((type(self), self._metadata)) + + +class ArrayIO: + """Explicit synchronous IO for one Zarr array. See `AsyncArrayIO`.""" + + _async_io: AsyncArrayIO + + def __init__( + self, + metadata: ArrayMetadata | ArrayMetadataDict, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: + self._async_io = AsyncArrayIO(metadata, store_path, config) + + @classmethod + def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: + io = object.__new__(cls) + io._async_io = async_io + return io + + @classmethod + def build( + cls, + store: StoreLike, + *, + name: str | None = None, + shape: ShapeLike, + dtype: ZDTypeLike, + chunks: ChunksLike | Literal["auto"] = "auto", + shards: ShardsLike | None = None, + filters: FiltersLike = "auto", + compressors: CompressorsLike = "auto", + serializer: SerializerLike = "auto", + fill_value: Any | None = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + zarr_format: ZarrFormat | None = 3, + attributes: dict[str, JSON] | None = None, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + dimension_names: DimensionNamesLike = None, + config: ArrayConfigLike | None = None, + ) -> ArrayIO: + """Declare an array bound to a location in a store, without storing anything. + + See `AsyncArrayIO.build`. + """ + return cls._wrap( + sync( + AsyncArrayIO.build( + store, + name=name, + shape=shape, + dtype=dtype, + chunks=chunks, + shards=shards, + filters=filters, + compressors=compressors, + serializer=serializer, + fill_value=fill_value, + order=order, + zarr_format=zarr_format, + attributes=attributes, + chunk_key_encoding=chunk_key_encoding, + dimension_names=dimension_names, + config=config, + ) + ) + ) + + @classmethod + def open( + cls, + store: StoreLike, + *, + name: str | None = None, + zarr_format: ZarrFormat | None = None, + config: ArrayConfigLike | None = None, + ) -> ArrayIO: + """Read the metadata of an existing array and bind to it. + + See `AsyncArrayIO.open`. + """ + return cls._wrap( + sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config)) + ) + + @property + def metadata(self) -> ArrayMetadata: + """The in-memory metadata this instance is bound to.""" + return self._async_io.metadata + + @property + def store_path(self) -> StorePath: + """The location in a store this instance is bound to.""" + return self._async_io.store_path + + @property + def config(self) -> ArrayConfig: + """The runtime configuration of this instance.""" + return self._async_io.config + + def to_array(self) -> AnyArray: + """Return an `Array` over this binding; both share one state.""" + return Array(self._async_io.to_array()) + + def store_metadata(self, *, ensure_parents: bool = False) -> None: + """Store this array's metadata document(s). See `AsyncArrayIO.store_metadata`.""" + sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) + + def verify(self) -> None: + """Check that the stored metadata declares the array this instance is bound to. + + See `AsyncArrayIO.verify`. + """ + sync(self._async_io.verify()) + + def erase_metadata(self) -> None: + """Erase this array's metadata document(s). See `AsyncArrayIO.erase_metadata`.""" + sync(self._async_io.erase_metadata()) + + def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: + """Erase one stored chunk by its chunk grid coordinates. See `AsyncArrayIO.erase_chunk`.""" + sync(self._async_io.erase_chunk(chunk_coords)) + + def erase_chunks(self) -> None: + """Erase every stored chunk of this array. See `AsyncArrayIO.erase_chunks`.""" + sync(self._async_io.erase_chunks()) + + def getitem( + self, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """Read the selected region, reading absent chunks as the fill value.""" + return sync(self._async_io.getitem(selection, prototype=prototype)) + + def setitem( + self, + selection: BasicSelection, + value: npt.ArrayLike, + prototype: BufferPrototype | None = None, + ) -> None: + """Encode `value` and store the chunks covering `selection`.""" + sync(self._async_io.setitem(selection, value, prototype=prototype)) + + def __repr__(self) -> str: + return f"" diff --git a/src/zarr/experimental/array_io.py b/src/zarr/experimental/array_io.py index a0b30d4ee0..36705b695e 100644 --- a/src/zarr/experimental/array_io.py +++ b/src/zarr/experimental/array_io.py @@ -9,445 +9,11 @@ mirrors the array API of `zarrs`, where storing, retrieving and erasing metadata and chunks are each explicit operations. -Every write has an erasing inverse: `store_metadata` ↔ `erase_metadata`, -and `setitem` ↔ `erase_chunk` / `erase_chunks`. `store_metadata` is the -publication point: until it is called, no Zarr reader resolves the array, so -a failed chunks-first run leaves no readable array. It does leave chunk keys -in the store, which occupy space and would be misread by a different array -stored later at the same location; `erase_chunks` removes them. - -One exception to "no implicit metadata IO": an array opened from a stored -document that needed a repair stores that repair before its first chunk -write, exactly as `AsyncArray` does. Call `store_metadata` first to make -that write explicit. +The classes live in `zarr.core.array`, where `AsyncArray` wraps an +`AsyncArrayIO` and delegates its metadata and selection IO to it; this module +is their public, experimental home. """ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, Literal - -from zarr.core.array import ( - DEFAULT_FILL_VALUE, - Array, - AsyncArray, - _iter_shard_keys, - build_array_metadata, - get_array_metadata, - parse_array_metadata, -) -from zarr.core.common import concurrent_foreach, product -from zarr.core.config import config as zarr_config -from zarr.core.metadata.io import ARRAY_DOCUMENTS -from zarr.core.sync import sync -from zarr.storage._common import make_store_path - -if TYPE_CHECKING: - import numpy.typing as npt - - from zarr.core.array import ( - CompressorsLike, - FiltersLike, - SerializerLike, - ShardsLike, - ) - from zarr.core.array_spec import ArrayConfig, ArrayConfigLike - from zarr.core.buffer import BufferPrototype, NDArrayLikeOrScalar - from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike - from zarr.core.common import ( - JSON, - ChunksLike, - DimensionNamesLike, - MemoryOrder, - ShapeLike, - ZarrFormat, - ) - from zarr.core.dtype import ZDTypeLike - from zarr.core.indexing import BasicSelection - from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata - from zarr.storage import StoreLike - from zarr.storage._common import StorePath - from zarr.types import AnyArray, AnyAsyncArray +from zarr.core.array import ArrayIO, AsyncArrayIO __all__ = ["ArrayIO", "AsyncArrayIO"] - - -async def _delete_key(store_path: StorePath, key: str) -> None: - await (store_path / key).delete() - - -class AsyncArrayIO: - """Explicit asynchronous IO for one Zarr array. - - An `AsyncArrayIO` binds an array metadata document to a location in a - store. Constructing one — directly, or with `build` — performs no IO: - nothing is read from or written to the store until an IO method is called. - `store_metadata` stores the metadata document, and `setitem` / `getitem` - store and read chunks, in whatever order the caller chooses. - - Warning: the bound metadata is assumed, never verified against the store. - Reads and partial writes decode stored chunks under it, so binding - metadata that disagrees with how existing chunks at this location were - encoded misreads or corrupts data — the hazard TensorStore documents for - `assume_metadata`. - """ - - _array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] - - def __init__( - self, - metadata: ArrayV2Metadata | ArrayV3Metadata, - store_path: StorePath, - config: ArrayConfigLike | None = None, - ) -> None: - self._array = AsyncArray(metadata=metadata, store_path=store_path, config=config) - - @classmethod - def _wrap(cls, array: AnyAsyncArray) -> AsyncArrayIO: - io = object.__new__(cls) - io._array = array - return io - - @classmethod - async def build( - cls, - store: StoreLike, - *, - name: str | None = None, - shape: ShapeLike, - dtype: ZDTypeLike, - chunks: ChunksLike | Literal["auto"] = "auto", - shards: ShardsLike | None = None, - filters: FiltersLike = "auto", - compressors: CompressorsLike = "auto", - serializer: SerializerLike = "auto", - fill_value: Any | None = DEFAULT_FILL_VALUE, - order: MemoryOrder | None = None, - zarr_format: ZarrFormat | None = 3, - attributes: dict[str, JSON] | None = None, - chunk_key_encoding: ChunkKeyEncodingLike | None = None, - dimension_names: DimensionNamesLike = None, - config: ArrayConfigLike | None = None, - ) -> AsyncArrayIO: - """Declare an array bound to a location in a store, without storing anything. - - The keyword arguments have the same meaning as in - [`zarr.api.asynchronous.create_array`][]. The returned instance holds - the normalized metadata in memory; `store_metadata` stores it. - """ - metadata = build_array_metadata( - shape=shape, - dtype=dtype, - chunks=chunks, - shards=shards, - filters=filters, - compressors=compressors, - serializer=serializer, - fill_value=fill_value, - order=order, - zarr_format=zarr_format, - attributes=attributes, - chunk_key_encoding=chunk_key_encoding, - dimension_names=dimension_names, - ) - store_path = await make_store_path(store, path=name) - return cls(metadata, store_path, config) - - @classmethod - async def open( - cls, - store: StoreLike, - *, - name: str | None = None, - zarr_format: ZarrFormat | None = None, - config: ArrayConfigLike | None = None, - ) -> AsyncArrayIO: - """Read the metadata of an existing array and bind to it. - - This is the one classmethod constructor that performs IO: the array's - stored metadata document is read from the store. `zarr_format=None` - (the default) reads Zarr format 3 and falls back to format 2, as - `open` does in `zarrs`; pass 2 or 3 to require one format. - """ - store_path = await make_store_path(store, path=name) - array = await AsyncArray.open(store_path, zarr_format, config=config) - return cls._wrap(array) - - @property - def metadata(self) -> ArrayMetadata: - """The in-memory metadata this instance is bound to.""" - return self._array.metadata - - @property - def store_path(self) -> StorePath: - """The location in a store this instance is bound to.""" - return self._array.store_path - - @property - def config(self) -> ArrayConfig: - """The runtime configuration of this instance.""" - return self._array.config - - def to_array(self) -> AnyAsyncArray: - """Return an `AsyncArray` handle over the same metadata and location.""" - return self._array - - async def store_metadata(self, *, ensure_parents: bool = False) -> None: - """Store this array's metadata document(s). - - The documents are written unconditionally, replacing the metadata of - any existing node at this location, as `store_metadata` does in - `zarrs`. Guarded creation, which raises when a node already exists, - is [`zarr.api.asynchronous.create_array`][]'s job. Chunks of any - previous array at this location are kept and would be read under the - new metadata; call `erase_chunks` first when repurposing a location. - With `ensure_parents`, missing parent groups are stored first. - """ - await self._array._save_metadata(self._array.metadata, ensure_parents=ensure_parents) - - async def verify(self) -> None: - """Check that the stored metadata declares the array this instance is bound to. - - Reads the metadata stored at this location — favoring Zarr format 3 - when documents of both formats exist — and raises `ValueError` naming - each differing field when it does not equal the bound metadata. - Raises [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. - This is the explicit counterpart of the verification TensorStore - performs against spec constraints on open: a writer that declared its - own metadata with `build` can assert, before writing chunks, that it - agrees with what a coordinator stored. - """ - bound = self.metadata - stored = parse_array_metadata( - await get_array_metadata(self._array.store_path, zarr_format=None), - str(self._array.store_path), - ) - if stored == bound: - return - if type(stored) is not type(bound): - raise ValueError( - f"Array {str(self.store_path)!r}: the stored array is Zarr format " - f"{stored.zarr_format}; the bound metadata is Zarr format " - f"{bound.zarr_format}." - ) - stored_dict, bound_dict = stored.to_dict(), bound.to_dict() - differing = sorted( - key - for key in stored_dict.keys() | bound_dict.keys() - if stored_dict.get(key) != bound_dict.get(key) - ) - raise ValueError( - f"Array {str(self.store_path)!r}: the stored metadata does not match the " - f"bound metadata; differing fields: {', '.join(differing)}." - ) - - async def erase_metadata(self) -> None: - """Erase this array's metadata document(s), the inverse of `store_metadata`. - - Succeeds if the documents do not exist, as `erase_metadata` does in - `zarrs`. Stored chunks are kept; `erase_chunks` removes those. - """ - store_path = self._array.store_path - for document in ARRAY_DOCUMENTS[self.metadata.zarr_format]: - await (store_path / document).delete() - self._array._stored_document_replaced() - - async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: - """Erase one stored chunk by its chunk grid coordinates. - - When sharding is used, the stored chunk is the shard. Succeeds if the - chunk does not exist, as `erase_chunk` does in `zarrs`. - """ - grid_shape = self._array._shard_grid_shape - if len(chunk_coords) != len(grid_shape) or any( - coord < 0 or coord >= size for coord, size in zip(chunk_coords, grid_shape, strict=True) - ): - raise IndexError( - f"Chunk coordinates {chunk_coords} are outside the chunk grid " - f"of shape {grid_shape}." - ) - await _delete_key(self._array.store_path, self.metadata.encode_chunk_key(chunk_coords)) - - async def erase_chunks(self) -> None: - """Erase every stored chunk of this array, the inverse of its chunk writes. - - Chunks that do not exist are skipped; the metadata documents are kept. - """ - n_chunks = product(self._array._shard_grid_shape) - if n_chunks == 0: - return - limit = zarr_config.get("async.concurrency") or 1000 - store_path = self._array.store_path - await concurrent_foreach( - ((store_path, key) for key in _iter_shard_keys(self._array)), - _delete_key, - min(limit, n_chunks), - ) - - async def getitem( - self, - selection: BasicSelection, - *, - prototype: BufferPrototype | None = None, - ) -> NDArrayLikeOrScalar: - """Read the selected region, reading absent chunks as the fill value.""" - return await self._array.getitem(selection, prototype=prototype) - - async def setitem( - self, - selection: BasicSelection, - value: npt.ArrayLike, - prototype: BufferPrototype | None = None, - ) -> None: - """Encode `value` and store the chunks covering `selection`.""" - await self._array.setitem(selection, value, prototype=prototype) - - def __repr__(self) -> str: - return f"" - - -class ArrayIO: - """Explicit synchronous IO for one Zarr array. See `AsyncArrayIO`.""" - - _async_io: AsyncArrayIO - - def __init__( - self, - metadata: ArrayV2Metadata | ArrayV3Metadata, - store_path: StorePath, - config: ArrayConfigLike | None = None, - ) -> None: - self._async_io = AsyncArrayIO(metadata, store_path, config) - - @classmethod - def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: - io = object.__new__(cls) - io._async_io = async_io - return io - - @classmethod - def build( - cls, - store: StoreLike, - *, - name: str | None = None, - shape: ShapeLike, - dtype: ZDTypeLike, - chunks: ChunksLike | Literal["auto"] = "auto", - shards: ShardsLike | None = None, - filters: FiltersLike = "auto", - compressors: CompressorsLike = "auto", - serializer: SerializerLike = "auto", - fill_value: Any | None = DEFAULT_FILL_VALUE, - order: MemoryOrder | None = None, - zarr_format: ZarrFormat | None = 3, - attributes: dict[str, JSON] | None = None, - chunk_key_encoding: ChunkKeyEncodingLike | None = None, - dimension_names: DimensionNamesLike = None, - config: ArrayConfigLike | None = None, - ) -> ArrayIO: - """Declare an array bound to a location in a store, without storing anything. - - See `AsyncArrayIO.build`. - """ - return cls._wrap( - sync( - AsyncArrayIO.build( - store, - name=name, - shape=shape, - dtype=dtype, - chunks=chunks, - shards=shards, - filters=filters, - compressors=compressors, - serializer=serializer, - fill_value=fill_value, - order=order, - zarr_format=zarr_format, - attributes=attributes, - chunk_key_encoding=chunk_key_encoding, - dimension_names=dimension_names, - config=config, - ) - ) - ) - - @classmethod - def open( - cls, - store: StoreLike, - *, - name: str | None = None, - zarr_format: ZarrFormat | None = None, - config: ArrayConfigLike | None = None, - ) -> ArrayIO: - """Read the metadata of an existing array and bind to it. - - See `AsyncArrayIO.open`. - """ - return cls._wrap( - sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config)) - ) - - @property - def metadata(self) -> ArrayMetadata: - """The in-memory metadata this instance is bound to.""" - return self._async_io.metadata - - @property - def store_path(self) -> StorePath: - """The location in a store this instance is bound to.""" - return self._async_io.store_path - - @property - def config(self) -> ArrayConfig: - """The runtime configuration of this instance.""" - return self._async_io.config - - def to_array(self) -> AnyArray: - """Return an `Array` handle over the same metadata and location.""" - return Array(self._async_io.to_array()) - - def store_metadata(self, *, ensure_parents: bool = False) -> None: - """Store this array's metadata document(s). See `AsyncArrayIO.store_metadata`.""" - sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) - - def verify(self) -> None: - """Check that the stored metadata declares the array this instance is bound to. - - See `AsyncArrayIO.verify`. - """ - sync(self._async_io.verify()) - - def erase_metadata(self) -> None: - """Erase this array's metadata document(s). See `AsyncArrayIO.erase_metadata`.""" - sync(self._async_io.erase_metadata()) - - def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: - """Erase one stored chunk by its chunk grid coordinates. See `AsyncArrayIO.erase_chunk`.""" - sync(self._async_io.erase_chunk(chunk_coords)) - - def erase_chunks(self) -> None: - """Erase every stored chunk of this array. See `AsyncArrayIO.erase_chunks`.""" - sync(self._async_io.erase_chunks()) - - def getitem( - self, - selection: BasicSelection, - *, - prototype: BufferPrototype | None = None, - ) -> NDArrayLikeOrScalar: - """Read the selected region, reading absent chunks as the fill value.""" - return sync(self._async_io.getitem(selection, prototype=prototype)) - - def setitem( - self, - selection: BasicSelection, - value: npt.ArrayLike, - prototype: BufferPrototype | None = None, - ) -> None: - """Encode `value` and store the chunks covering `selection`.""" - sync(self._async_io.setitem(selection, value, prototype=prototype)) - - def __repr__(self) -> str: - return f"" From ea0c01b20332a648f8ce0f1d56a7ef99ef3a9b7f Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 13:05:03 +0100 Subject: [PATCH 06/48] docs: note the AsyncArray delegation in the changelog fragment Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 1 + 1 file changed, 1 insertion(+) diff --git a/changes/419.feature.md b/changes/419.feature.md index 24f9990189..7d3d33cf3e 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,2 +1,3 @@ Added `zarr.experimental.ArrayIO` and `zarr.experimental.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `open` auto-detects the Zarr format by default, as `zarrs` does. `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. +`AsyncArray` now wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it; the two layers share one state, and `to_array()` returns an array over the same binding. From 63c6ed75cb16feaf9412f9f5316e691bed3aad39 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 13:26:03 +0100 Subject: [PATCH 07/48] feat(array): make ArrayIO public and reachable from every array handle ArrayIO and AsyncArrayIO move out of zarr.experimental into the zarr namespace, next to Array and AsyncArray, which are built on them: the wrapped layer is reachable as Array.io / AsyncArray.io, the inverse of to_array(). The zarr.experimental.array_io shim is removed and the layer is documented on the array API page. Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 5 ++--- docs/api/zarr/array.md | 2 ++ docs/api/zarr/experimental.md | 2 -- src/zarr/__init__.py | 4 +++- src/zarr/core/array.py | 10 ++++++++++ src/zarr/experimental/__init__.py | 3 +-- src/zarr/experimental/array_io.py | 19 ------------------- .../{test_experimental => }/test_array_io.py | 9 +++++++-- 8 files changed, 25 insertions(+), 29 deletions(-) delete mode 100644 src/zarr/experimental/array_io.py rename tests/{test_experimental => }/test_array_io.py (92%) diff --git a/changes/419.feature.md b/changes/419.feature.md index 7d3d33cf3e..93515a4559 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,3 +1,2 @@ -Added `zarr.experimental.ArrayIO` and `zarr.experimental.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `open` auto-detects the Zarr format by default, as `zarrs` does. -`verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. -`AsyncArray` now wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it; the two layers share one state, and `to_array()` returns an array over the same binding. +Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. +`Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. diff --git a/docs/api/zarr/array.md b/docs/api/zarr/array.md index ff61cb1fe2..dc51e1913e 100644 --- a/docs/api/zarr/array.md +++ b/docs/api/zarr/array.md @@ -1,2 +1,4 @@ ::: zarr.Array ::: zarr.AsyncArray +::: zarr.ArrayIO +::: zarr.AsyncArrayIO diff --git a/docs/api/zarr/experimental.md b/docs/api/zarr/experimental.md index 030bb8b714..60f1f987b5 100644 --- a/docs/api/zarr/experimental.md +++ b/docs/api/zarr/experimental.md @@ -6,6 +6,4 @@ Experimental functionality is not stable and may change or be removed at any poi ## Classes -::: zarr.experimental.array_io - ::: zarr.experimental.cache_store diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index f305245f4e..c30b0cb4bf 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -34,7 +34,7 @@ zeros, zeros_like, ) -from zarr.core.array import Array, AsyncArray +from zarr.core.array import Array, ArrayIO, AsyncArray, AsyncArrayIO from zarr.core.config import config from zarr.core.group import AsyncGroup, Group @@ -147,7 +147,9 @@ def set_format(log_format: str) -> None: __all__ = [ "Array", + "ArrayIO", "AsyncArray", + "AsyncArrayIO", "AsyncGroup", "Group", "__version__", diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 61f75e38e2..4c8146d049 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -428,6 +428,11 @@ def _from_io(cls, io: AsyncArrayIO) -> AnyAsyncArray: object.__setattr__(self, "_io", io) return self + @property + def io(self) -> AsyncArrayIO: + """The explicit-IO layer this array delegates to; both share one state.""" + return self._io + @property def metadata(self) -> T_ArrayMetadata: """The in-memory metadata of this array.""" @@ -1898,6 +1903,11 @@ def async_array(self) -> AsyncArray[T_ArrayMetadata]: """ return self._async_array + @property + def io(self) -> ArrayIO: + """The explicit synchronous IO layer over this array's binding.""" + return ArrayIO._wrap(self._async_array._io) + @property def config(self) -> ArrayConfig: """ diff --git a/src/zarr/experimental/__init__.py b/src/zarr/experimental/__init__.py index 9a7d50f61f..f7caaf96a1 100644 --- a/src/zarr/experimental/__init__.py +++ b/src/zarr/experimental/__init__.py @@ -1,6 +1,5 @@ """The experimental module is a site for exporting new or experimental Zarr features.""" from zarr.core.chunk_grids import ChunkGrid, ChunkSpec -from zarr.experimental.array_io import ArrayIO, AsyncArrayIO -__all__ = ["ArrayIO", "AsyncArrayIO", "ChunkGrid", "ChunkSpec"] +__all__ = ["ChunkGrid", "ChunkSpec"] diff --git a/src/zarr/experimental/array_io.py b/src/zarr/experimental/array_io.py deleted file mode 100644 index 36705b695e..0000000000 --- a/src/zarr/experimental/array_io.py +++ /dev/null @@ -1,19 +0,0 @@ -"""An explicit-IO layer for Zarr arrays. - -`AsyncArrayIO` and `ArrayIO` separate declaring an array from storing it. -Constructing either class performs no IO: an instance binds an array metadata -document to a location in a store, and every read or write of that store is an -explicit method call. In particular, metadata is stored only when -`store_metadata` is called, so chunks can be written before metadata — a -pattern the eager creation API (`zarr.create_array`) cannot express. This -mirrors the array API of `zarrs`, where storing, retrieving and erasing -metadata and chunks are each explicit operations. - -The classes live in `zarr.core.array`, where `AsyncArray` wraps an -`AsyncArrayIO` and delegates its metadata and selection IO to it; this module -is their public, experimental home. -""" - -from zarr.core.array import ArrayIO, AsyncArrayIO - -__all__ = ["ArrayIO", "AsyncArrayIO"] diff --git a/tests/test_experimental/test_array_io.py b/tests/test_array_io.py similarity index 92% rename from tests/test_experimental/test_array_io.py rename to tests/test_array_io.py index df1af45f2f..8e19e518a6 100644 --- a/tests/test_experimental/test_array_io.py +++ b/tests/test_array_io.py @@ -1,4 +1,4 @@ -"""Tests for the explicit-IO array layer in `zarr.experimental.array_io`.""" +"""Tests for the explicit-IO array layer, `zarr.ArrayIO` and `zarr.AsyncArrayIO`.""" from __future__ import annotations @@ -7,9 +7,9 @@ import numpy as np import pytest +from zarr import ArrayIO, AsyncArrayIO from zarr.api.asynchronous import open_array from zarr.errors import ArrayNotFoundError -from zarr.experimental import ArrayIO, AsyncArrayIO from zarr.storage import MemoryStore if TYPE_CHECKING: @@ -54,6 +54,9 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: arr = await open_array(store=store, path="foo", mode="r", zarr_format=zarr_format) np.testing.assert_array_equal(await arr.getitem(Ellipsis), data) + # every AsyncArray is built on an AsyncArrayIO, reachable as .io + assert arr.io.metadata is arr.metadata + np.testing.assert_array_equal(await arr.io.getitem(Ellipsis), data) # a writer that declared the same array independently can verify its assumption worker = await AsyncArrayIO.build( @@ -96,6 +99,8 @@ def test_chunks_before_metadata_sync() -> None: opened = ArrayIO.open(store, name="foo") np.testing.assert_array_equal(opened.getitem(Ellipsis), data) np.testing.assert_array_equal(opened.to_array()[:], data) + # every Array is built on the IO layer, reachable as .io + np.testing.assert_array_equal(opened.to_array().io.getitem(Ellipsis), data) opened.erase_chunk((0, 0)) opened.erase_chunks() opened.erase_metadata() From adf48d22a384c31c7d111fc1ed8fa78c49d78814 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 14:18:24 +0100 Subject: [PATCH 08/48] feat(array): hold the ArrayIO document as a zarr-metadata model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit zarr-python now depends on zarr-metadata (>=0.6.0): AsyncArrayIO holds one source — the document, a zarr_metadata model, when constructed by open; the resolved metadata when constructed by build or by AsyncArray — and derives the other form, the chunk grid, the chunk key encoding and the codec pipeline as caches. Document operations (store_metadata, erase_metadata, erase_chunk, erase_chunks, verify, equality) never resolve codec or data type names, so an array whose codecs this environment lacks still opens for document work and refuses only when the data path is built; a document zarr can only read as repaired is refused by the strict model read, with AsyncArray.open remaining the repairing reader. Grid derivation goes through the same grid-metadata objects ChunkGrid.from_metadata dispatches on, so both derivations lay out chunks identically, and model JSON is normalized through JSON text because the model returns read-only forms. Resolved-seeded instances keep deriving their grid, chunk keys and equality from the resolved metadata: zarr-python writes some documents the strict model refuses (Zarr format 2 generic time units), and the array API's behavior must not depend on the model accepting them. Refs zarr-developers/zarr-python#4217 Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 1 + pyproject.toml | 1 + src/zarr/core/array.py | 356 ++++++++++++++++++++++++++++++++--------- tests/test_array_io.py | 45 +++++- uv.lock | 15 ++ 5 files changed, 345 insertions(+), 73 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index 93515a4559..9dd9f28247 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,2 +1,3 @@ Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. `Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. +The binding holds the metadata document as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.document`: document operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. diff --git a/pyproject.toml b/pyproject.toml index 0df173f805..067d2e7297 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ dependencies = [ 'typing_extensions>=4.14', 'donfig>=0.8', 'msgspec>=0.19', + 'zarr-metadata>=0.6.0', ] dynamic = [ diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 4c8146d049..33b7a7daef 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2,6 +2,7 @@ import contextlib import copy +import json import math import warnings from asyncio import gather @@ -21,6 +22,7 @@ import numpy as np from typing_extensions import Sentinel, deprecated +from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata import zarr from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec @@ -121,6 +123,7 @@ ) from zarr.core.metadata.io import ( ARRAY_DOCUMENTS, + _write_metadata, encode_documents, parse_stored_array, read_documents, @@ -138,7 +141,9 @@ from zarr.core.metadata.v3 import ( ChunkGridMetadata, RectilinearChunkGridMetadata, + RectilinearChunkGridMetadataJSON, RegularChunkGridMetadata, + RegularChunkGridMetadataJSON, create_chunk_grid_metadata, parse_node_type_array, ) @@ -420,6 +425,9 @@ def __init__( config: ArrayConfigLike | None = None, ) -> None: object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config)) + # An AsyncArray is the resolving layer: build its codec pipeline now, so + # pipeline-construction warnings fire at creation, as they always have. + _ = self._io._resolved_codec_pipeline @classmethod def _from_io(cls, io: AsyncArrayIO) -> AnyAsyncArray: @@ -451,7 +459,7 @@ def config(self) -> ArrayConfig: @property def codec_pipeline(self) -> CodecPipeline: """The codec pipeline used for encoding and decoding chunks.""" - return self._io._codec_pipeline + return self._io._resolved_codec_pipeline @property def _chunk_grid(self) -> ChunkGrid: @@ -6058,8 +6066,7 @@ async def _delete_chunk(coords: tuple[int, ...]) -> None: array._stored_document_replaced() # Update the wrapped binding in place; the codec pipeline is kept, as before. - array._io._metadata = new_metadata - array._io._chunk_grid = new_chunk_grid + array._io._rebind_resolved(new_metadata, new_chunk_grid) async def _append( @@ -6145,6 +6152,8 @@ async def _update_attributes( The array with the updated attributes. """ array.metadata.attributes.update(new_attributes) + # The document the binding derives must re-derive from the changed metadata. + array._io._rebind_resolved(array.metadata, array._chunk_grid) # Write new metadata await array._save_metadata() @@ -6180,12 +6189,50 @@ async def _delete_key(store_path: StorePath, key: str) -> None: await (store_path / key).delete() +type ArrayDocumentModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata +"""A Zarr array metadata document, held as a `zarr_metadata` model: the JSON as +written, validated, with extension names read but never resolved into classes.""" + + +def _document_json(document: ArrayDocumentModel) -> dict[str, Any]: + """The document as plain, mutable JSON. The model returns read-only forms — + arrays as tuples — while zarr's parsers take the stored forms, so the + document goes through JSON text.""" + return cast("dict[str, Any]", json.loads(json.dumps(document.to_json()))) + + +def _array_document_model(data: dict[str, JSON]) -> ArrayDocumentModel: + """The model of an array metadata document in the form `get_array_metadata` + returns: the Zarr format 2 form carries its attributes merged in. + + The read is strict: a document with problems — including one zarr can only + read as repaired — raises `zarr_metadata.MetadataValidationError` naming + each problem. A document whose extension names are unknown is not a + problem: the model reports them as unclaimed fields, and resolution + fails only when the data path is built. + """ + if data.get("zarr_format") == 2: + return ZarrV2ArrayMetadata(data) + return ZarrV3ArrayMetadata(data) + + class AsyncArrayIO: """Explicit asynchronous IO for one Zarr array. An `AsyncArrayIO` binds an array metadata document to a location in a - store. Constructing one — directly, or with `build` — performs no IO: - nothing is read from or written to the store until an IO method is called. + store. The document is held as a `zarr_metadata` model — the JSON as + written, with extension names read but not resolved into codec or data + type classes — so every document operation (`store_metadata`, + `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) works without + resolving anything: an array whose codecs this environment lacks can + still be opened, verified and erased. Reading and writing chunks + (`getitem` / `setitem`) resolves the document into classes on first use; + `metadata` is that resolved form, derived once from `document` — or, when + the instance was built from an already-resolved metadata object (as every + `AsyncArray` is), held directly, with `document` the derived one. + + Constructing one — directly, or with `build` — performs no IO: nothing is + read from or written to the store until an IO method is called. `store_metadata` stores the metadata document, and `setitem` / `getitem` store and read chunks, in whatever order the caller chooses. Every write has an erasing inverse: `store_metadata` ↔ `erase_metadata`, and chunk @@ -6199,9 +6246,9 @@ class AsyncArrayIO: be misread by a different array stored later at the same location; `erase_chunks` removes them. - Warning: the bound metadata is assumed, never verified against the store. - Reads and partial writes decode stored chunks under it, so binding - metadata that disagrees with how existing chunks at this location were + Warning: the bound document is assumed, never verified against the store. + Reads and partial writes decode stored chunks under it, so binding a + document that disagrees with how existing chunks at this location were encoded misreads or corrupts data — the hazard TensorStore documents for `assume_metadata`. One exception to "no implicit metadata IO": an array opened from a stored document that needed a repair stores that repair @@ -6209,26 +6256,35 @@ class AsyncArrayIO: write explicit. """ - _metadata: ArrayMetadata _store_path: StorePath _config: ArrayConfig - _codec_pipeline: CodecPipeline - _chunk_grid: ChunkGrid + _source: ArrayDocumentModel | ArrayMetadata + _document_cache: ArrayDocumentModel | None + _resolved_cache: ArrayMetadata | None + _chunk_grid_cache: ChunkGrid | None + _chunk_key_encoding_cache: ChunkKeyEncoding | None + _codec_pipeline_cache: CodecPipeline | None def __init__( self, - metadata: ArrayMetadata | ArrayMetadataDict, + document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: - metadata_parsed = parse_array_metadata(metadata, str(store_path)) - self._metadata = metadata_parsed self._store_path = store_path self._config = parse_array_config(config) - self._chunk_grid = ChunkGrid.from_metadata(metadata_parsed) - self._codec_pipeline = create_codec_pipeline( - metadata=metadata_parsed, store=store_path.store - ) + self._document_cache = None + self._resolved_cache = None + self._chunk_grid_cache = None + self._chunk_key_encoding_cache = None + self._codec_pipeline_cache = None + if isinstance(document, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): + self._source = document + else: + # A resolved metadata object passes through; a raw document dict is + # read eagerly, as repaired when it needs that, warning as `AsyncArray` + # always has. + self._source = parse_array_metadata(document, str(store_path)) @classmethod async def build( @@ -6284,23 +6340,50 @@ async def open( zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> AsyncArrayIO: - """Read the metadata of an existing array and bind to it. + """Read the metadata document of an existing array and bind to it. This is the one classmethod constructor that performs IO: the array's stored metadata document is read from the store. `zarr_format=None` (the default) reads Zarr format 3 and falls back to format 2, as - `open` does in `zarrs`; pass 2 or 3 to require one format. + `open` does in `zarrs`; pass 2 or 3 to require one format. The read is + strict and does not resolve extension names: an array with unknown + codecs opens for document work, while a stored document zarr can only + read as repaired raises `zarr_metadata.MetadataValidationError` — open + such an array with [`zarr.api.asynchronous.open_array`][], which reads + it as repaired. """ store_path = await make_store_path(store, path=name) - metadata = parse_array_metadata( - await get_array_metadata(store_path, zarr_format=zarr_format), str(store_path) - ) - return cls(metadata, store_path, config) + data = await get_array_metadata(store_path, zarr_format=zarr_format) + return cls(_array_document_model(data), store_path, config) + + @property + def document(self) -> ArrayDocumentModel: + """The metadata document this instance is bound to, as a `zarr_metadata` model.""" + source = self._source + if isinstance(source, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): + return source + if self._document_cache is None: + buffers = source.to_buffer_dict(default_buffer_prototype()) + key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} + if isinstance(source, ArrayV2Metadata): + self._document_cache = ZarrV2ArrayMetadata.from_key_value(key_value) + else: + self._document_cache = ZarrV3ArrayMetadata.from_key_value(key_value) + return self._document_cache @property def metadata(self) -> ArrayMetadata: - """The in-memory metadata this instance is bound to.""" - return self._metadata + """The resolved metadata: the document with its codec and data type names + resolved into classes, derived from `document` on first use.""" + source = self._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + return source + if self._resolved_cache is None: + data = _document_json(source) + if isinstance(source, ZarrV2ArrayMetadata): + data["attributes"] = dict(source.attributes) + self._resolved_cache = parse_array_metadata(data, str(self._store_path)) + return self._resolved_cache @property def store_path(self) -> StorePath: @@ -6312,12 +6395,110 @@ def config(self) -> ArrayConfig: """The runtime configuration of this instance.""" return self._config + @property + def _zarr_format(self) -> ZarrFormat: + if isinstance(self._source, ZarrV2ArrayMetadata | ArrayV2Metadata): + return 2 + return 3 + + @property + def _resolved_if_any(self) -> ArrayMetadata | None: + source = self._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + return source + return self._resolved_cache + + def _rebind_resolved( + self, resolved: ArrayMetadata, chunk_grid: ChunkGrid | None = None + ) -> None: + """Make `resolved` this binding's state; the document re-derives from it. + The codec pipeline is kept: rebinding is how `resize` swaps the shape, + which leaves the chunk encoding as it was.""" + self._source = resolved + self._document_cache = None + self._resolved_cache = None + self._chunk_grid_cache = chunk_grid + self._chunk_key_encoding_cache = None + + @property + def _chunk_grid(self) -> ChunkGrid: + """The stored chunk grid: from the resolved metadata when the instance + holds one, and otherwise from the document alone, so the storage-level + operations of a document-bound instance never resolve codecs.""" + if self._chunk_grid_cache is None: + resolved = self._resolved_if_any + if resolved is not None: + self._chunk_grid_cache = ChunkGrid.from_metadata(resolved) + return self._chunk_grid_cache + document = self.document + doc = _document_json(document) + if isinstance(document, ZarrV2ArrayMetadata): + self._chunk_grid_cache = ChunkGrid.from_sizes(document.shape, tuple(doc["chunks"])) + else: + grid = doc["chunk_grid"] + # Through the same grid-metadata objects `ChunkGrid.from_metadata` + # dispatches on, so both derivations lay out chunks identically. + # A grid needs its configuration to lay out keys, so the v3.1 + # bare-name shorthand cannot name a computable grid. + name = grid if isinstance(grid, str) else grid["name"] + sizes: Sequence[int | Sequence[int]] + if isinstance(grid, dict) and name == "regular": + sizes = tuple( + RegularChunkGridMetadata.from_dict( + cast("RegularChunkGridMetadataJSON", grid) + ).chunk_shape + ) + elif isinstance(grid, dict) and name == "rectilinear": + sizes = RectilinearChunkGridMetadata.from_dict( + cast("RectilinearChunkGridMetadataJSON", grid) + ).chunk_shapes + else: + raise ValueError( + f"Array {str(self._store_path)!r}: unknown chunk grid " + f"{name!r}; its stored chunk keys cannot be computed." + ) + self._chunk_grid_cache = ChunkGrid.from_sizes(document.shape, sizes) + return self._chunk_grid_cache + + def _encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: + """The store key of a stored chunk: from the resolved metadata when the + instance holds one, and otherwise from the document alone.""" + resolved = self._resolved_if_any + if resolved is not None: + return resolved.encode_chunk_key(chunk_coords) + return self._chunk_key_encoding.encode_chunk_key(chunk_coords) + + @property + def _chunk_key_encoding(self) -> ChunkKeyEncoding: + """The chunk key encoding, derived from the document alone.""" + if self._chunk_key_encoding_cache is None: + document = self.document + doc = _document_json(document) + if isinstance(document, ZarrV2ArrayMetadata): + self._chunk_key_encoding_cache = V2ChunkKeyEncoding( + separator=doc.get("dimension_separator", ".") + ) + else: + encoding = doc["chunk_key_encoding"] + if isinstance(encoding, str): # the v3.1 bare-name shorthand + encoding = {"name": encoding} + self._chunk_key_encoding_cache = parse_chunk_key_encoding(encoding) + return self._chunk_key_encoding_cache + @property def _shard_grid_shape(self) -> tuple[int, ...]: # The stored chunk grid is the shard grid when sharding is used, the # chunk grid otherwise, so its keys are the stored chunk keys. return self._chunk_grid.grid_shape + @property + def _resolved_codec_pipeline(self) -> CodecPipeline: + if self._codec_pipeline_cache is None: + self._codec_pipeline_cache = create_codec_pipeline( + metadata=self.metadata, store=self._store_path.store + ) + return self._codec_pipeline_cache + def to_array(self) -> AnyAsyncArray: """Return an `AsyncArray` over this binding; both share one state.""" return AsyncArray._from_io(self) @@ -6333,38 +6514,53 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: new metadata; call `erase_chunks` first when repurposing a location. With `ensure_parents`, missing parent groups are stored first. """ - await save_metadata(self._store_path, self._metadata, ensure_parents=ensure_parents) - _stored_document_replaced(self._metadata) + source = self._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + await save_metadata(self._store_path, source, ensure_parents=ensure_parents) + _stored_document_replaced(source) + else: + prototype = default_buffer_prototype() + to_save = { + str(key): prototype.buffer.from_bytes(value) + for key, value in source.to_key_value().items() + } + await _write_metadata( + self._store_path, to_save, self._zarr_format, ensure_parents=ensure_parents + ) async def verify(self) -> None: """Check that the stored metadata declares the array this instance is bound to. - Reads the metadata stored at this location — favoring Zarr format 3 - when documents of both formats exist — and raises `ValueError` naming - each differing field when it does not equal the bound metadata. - Raises [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. - This is the explicit counterpart of the verification TensorStore - performs against spec constraints on open: a writer that declared its - own metadata with `build` can assert, before writing chunks, that it + Reads the metadata document stored at this location — favoring Zarr + format 3 when documents of both formats exist — and raises + `ValueError` naming each differing field when it does not equal the + bound document. Document equality is semantic: a fill value is + compared in its canonical spelling. Raises + [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. This is + the explicit counterpart of the verification TensorStore performs + against spec constraints on open: a writer that declared its own + metadata with `build` can assert, before writing chunks, that it agrees with what a coordinator stored. """ - bound = self._metadata - stored = parse_array_metadata( - await get_array_metadata(self._store_path, zarr_format=None), str(self._store_path) - ) + bound = self.document + stored = _array_document_model(await get_array_metadata(self._store_path, zarr_format=None)) if stored == bound: return if type(stored) is not type(bound): + stored_format = 2 if isinstance(stored, ZarrV2ArrayMetadata) else 3 raise ValueError( f"Array {str(self.store_path)!r}: the stored array is Zarr format " - f"{stored.zarr_format}; the bound metadata is Zarr format " - f"{bound.zarr_format}." + f"{stored_format}; the bound document is Zarr format {self._zarr_format}." ) - stored_dict, bound_dict = stored.to_dict(), bound.to_dict() + stored_json = dict(stored.to_json()) + bound_json = dict(bound.to_json()) + if isinstance(stored, ZarrV2ArrayMetadata) and isinstance(bound, ZarrV2ArrayMetadata): + stored_json["attributes"] = dict(stored.attributes) + bound_json["attributes"] = dict(bound.attributes) differing = sorted( key - for key in stored_dict.keys() | bound_dict.keys() - if stored_dict.get(key) != bound_dict.get(key) + for key in stored_json.keys() | bound_json.keys() + if stored_json.get(key) != bound_json.get(key) ) raise ValueError( f"Array {str(self.store_path)!r}: the stored metadata does not match the " @@ -6377,9 +6573,11 @@ async def erase_metadata(self) -> None: Succeeds if the documents do not exist, as `erase_metadata` does in `zarrs`. Stored chunks are kept; `erase_chunks` removes those. """ - for document in ARRAY_DOCUMENTS[self._metadata.zarr_format]: + for document in ARRAY_DOCUMENTS[self._zarr_format]: await (self._store_path / document).delete() - _stored_document_replaced(self._metadata) + resolved = self._resolved_if_any + if resolved is not None: + _stored_document_replaced(resolved) async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: """Erase one stored chunk by its chunk grid coordinates. @@ -6395,7 +6593,7 @@ async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: f"Chunk coordinates {chunk_coords} are outside the chunk grid " f"of shape {grid_shape}." ) - await _delete_key(self._store_path, self._metadata.encode_chunk_key(chunk_coords)) + await _delete_key(self._store_path, self._encode_chunk_key(chunk_coords)) async def erase_chunks(self) -> None: """Erase every stored chunk of this array, the inverse of its chunk writes. @@ -6409,7 +6607,7 @@ async def erase_chunks(self) -> None: limit = zarr_config.get("async.concurrency") or 1000 await concurrent_foreach( ( - (self._store_path, self._metadata.encode_chunk_key(coords)) + (self._store_path, self._encode_chunk_key(coords)) for coords in _iter_grid(grid_shape) ), _delete_key, @@ -6426,8 +6624,8 @@ async def _get_selection( ) -> NDArrayLikeOrScalar: return await _get_selection( self._store_path, - self._metadata, - self._codec_pipeline, + self.metadata, + self._resolved_codec_pipeline, self._config, self._chunk_grid, indexer, @@ -6444,14 +6642,15 @@ async def _set_selection( prototype: BufferPrototype, fields: Fields | None = None, ) -> None: + metadata = self.metadata if product(indexer.shape) > 0: # Chunks are about to be stored under the repaired metadata, so store it # first: every reader of the store then agrees with them. - await _store_repaired_document(self._store_path, self._metadata) + await _store_repaired_document(self._store_path, metadata) return await _set_selection( self._store_path, - self._metadata, - self._codec_pipeline, + metadata, + self._resolved_codec_pipeline, self._config, self._chunk_grid, indexer, @@ -6469,7 +6668,7 @@ async def getitem( """Read the selected region, reading absent chunks as the fill value.""" if prototype is None: prototype = default_buffer_prototype() - indexer = BasicIndexer(selection, shape=self._metadata.shape, chunk_grid=self._chunk_grid) + indexer = BasicIndexer(selection, shape=self.metadata.shape, chunk_grid=self._chunk_grid) return await self._get_selection(indexer, prototype=prototype) async def setitem( @@ -6481,28 +6680,28 @@ async def setitem( """Encode `value` and store the chunks covering `selection`.""" if prototype is None: prototype = default_buffer_prototype() - indexer = BasicIndexer(selection, shape=self._metadata.shape, chunk_grid=self._chunk_grid) + indexer = BasicIndexer(selection, shape=self.metadata.shape, chunk_grid=self._chunk_grid) return await self._set_selection(indexer, value, prototype=prototype) def __repr__(self) -> str: - return f"" + return f"" def __eq__(self, other: object) -> bool: - # Value equality over the full binding, matching what the dataclass-generated - # equality of AsyncArray compared when it held these as its own fields. + # Semantic equality over the binding: the document (canonically, via the + # model's own equality), the location and the runtime configuration. if not isinstance(other, AsyncArrayIO): return NotImplemented - return ( - self._metadata == other._metadata - and self._store_path == other._store_path - and self._config == other._config - and self._codec_pipeline == other._codec_pipeline - and self._chunk_grid == other._chunk_grid - ) + if self._store_path != other._store_path or self._config != other._config: + return False + resolved, other_resolved = self._resolved_if_any, other._resolved_if_any + if resolved is not None and other_resolved is not None: + return resolved == other_resolved + return self.document == other.document def __hash__(self) -> int: - # A subset of the equality fields: equal bindings hash equally. - return hash((type(self), self._metadata)) + # A subset of the equality fields, available without deriving either + # form: equal bindings hash equally. + return hash((type(self), self._source.shape, self._store_path.path)) class ArrayIO: @@ -6512,11 +6711,11 @@ class ArrayIO: def __init__( self, - metadata: ArrayMetadata | ArrayMetadataDict, + document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: - self._async_io = AsyncArrayIO(metadata, store_path, config) + self._async_io = AsyncArrayIO(document, store_path, config) @classmethod def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: @@ -6581,7 +6780,7 @@ def open( zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> ArrayIO: - """Read the metadata of an existing array and bind to it. + """Read the metadata document of an existing array and bind to it. See `AsyncArrayIO.open`. """ @@ -6589,9 +6788,14 @@ def open( sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config)) ) + @property + def document(self) -> ArrayDocumentModel: + """The metadata document this instance is bound to, as a `zarr_metadata` model.""" + return self._async_io.document + @property def metadata(self) -> ArrayMetadata: - """The in-memory metadata this instance is bound to.""" + """The resolved metadata. See `AsyncArrayIO.metadata`.""" return self._async_io.metadata @property @@ -6650,4 +6854,12 @@ def setitem( sync(self._async_io.setitem(selection, value, prototype=prototype)) def __repr__(self) -> str: - return f"" + return f"" + + def __eq__(self, other: object) -> bool: + if not isinstance(other, ArrayIO): + return NotImplemented + return self._async_io == other._async_io + + def __hash__(self) -> int: + return hash((type(self), self._async_io)) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 8e19e518a6..68f60c2e5f 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -6,10 +6,12 @@ import numpy as np import pytest +from zarr_metadata import ZarrV3ArrayMetadata from zarr import ArrayIO, AsyncArrayIO from zarr.api.asynchronous import open_array -from zarr.errors import ArrayNotFoundError +from zarr.core.buffer import default_buffer_prototype +from zarr.errors import ArrayNotFoundError, UnknownCodecError from zarr.storage import MemoryStore if TYPE_CHECKING: @@ -22,6 +24,20 @@ async def _store_keys(store: MemoryStore) -> set[str]: return {key async for key in store.list()} +async def _store_foreign_codec_array(store: MemoryStore) -> None: + """Store a valid v3 array document whose codec this environment lacks.""" + document = ( + b'{"zarr_format": 3, "node_type": "array", "shape": [4], "data_type": "uint8",' + b' "fill_value": 0, "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [2]}},' + b' "chunk_key_encoding": {"name": "default"},' + b' "codecs": [{"name": "zarr-python-does-not-know-this-codec"}]}' + ) + prototype = default_buffer_prototype() + await store.set("foo/zarr.json", prototype.buffer.from_bytes(document)) + await store.set("foo/c/0", prototype.buffer.from_bytes(b"xx")) + await store.set("foo/c/1", prototype.buffer.from_bytes(b"yy")) + + @pytest.mark.parametrize("zarr_format", [2, 3]) async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: """Declaring an array stores nothing; chunks, then metadata, can be stored @@ -48,6 +64,9 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: await aio.store_metadata() # the default zarr_format=None auto-detects the format, like zarrs opened = await AsyncArrayIO.open(store, name="foo") + # the binding is a document; the stored and declared documents agree + assert opened.document == aio.document + assert opened.document.shape == data.shape assert opened.metadata.zarr_format == zarr_format assert opened.metadata.attributes == {"answer": 42} np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) @@ -126,6 +145,30 @@ async def test_open_missing_raises() -> None: await AsyncArrayIO.open(MemoryStore(), name="foo") +async def test_document_work_without_codec_resolution() -> None: + """An array whose codec this environment lacks still opens for document + work — reading the document, verifying it, erasing chunks and metadata — + and refuses only when the data path must be built.""" + store = MemoryStore() + await _store_foreign_codec_array(store) + + aio = await AsyncArrayIO.open(store, name="foo") + document = aio.document + assert isinstance(document, ZarrV3ArrayMetadata) + assert document.shape == (4,) + assert document.must_understand_fields == {} + await aio.verify() + + with pytest.raises(UnknownCodecError, match="zarr-python-does-not-know-this-codec"): + await aio.getitem(Ellipsis) + + await aio.erase_chunk((0,)) + await aio.erase_chunks() + await aio.store_metadata() + await aio.erase_metadata() + assert await _store_keys(store) == set() + + async def test_verify_mismatch_raises() -> None: store = MemoryStore() aio = await AsyncArrayIO.build(store, shape=(4, 4), dtype="uint8", chunks=(2, 2)) diff --git a/uv.lock b/uv.lock index 3b0dbe3049..23012125e9 100644 --- a/uv.lock +++ b/uv.lock @@ -3705,6 +3705,7 @@ dependencies = [ { name = "numpy" }, { name = "packaging" }, { name = "typing-extensions" }, + { name = "zarr-metadata" }, ] [package.optional-dependencies] @@ -3832,6 +3833,7 @@ requires-dist = [ { name = "typer", marker = "extra == 'cli'" }, { name = "typing-extensions", specifier = ">=4.14" }, { name = "universal-pathlib", marker = "extra == 'optional'" }, + { name = "zarr-metadata", specifier = ">=0.6.0" }, ] provides-extras = ["cast-value-rs", "cli", "gpu", "optional", "remote"] @@ -3924,3 +3926,16 @@ test = [ { name = "tomlkit", specifier = "==0.15.1" }, { name = "uv", specifier = "==0.13.4" }, ] + +[[package]] +name = "zarr-metadata" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5f/a3/4526430215c09ed45d47886a7c8b439218622a4f05e6ce1198147961ca64/zarr_metadata-0.6.0.tar.gz", hash = "sha256:4462dc91cdb8c19c370bcd548096125591359579586f42de789dc495a33119d0", size = 519797, upload-time = "2026-10-26T13:55:27.086Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/20/b2/c70c528d438bbf656873e8da51949c50efff866b85c6ee3335ba81bf1142/zarr_metadata-0.6.0-py3-none-any.whl", hash = "sha256:ab1f9f33d18094ffac020c2c975aa32f8737f36a0dfa8543b6c6d717a3469d1e", size = 274628, upload-time = "2026-10-26T13:55:25.798Z" }, +] From 1b5b4abb667d4785a859763f4f2ef7a065ba7e2b Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 10:38:37 +0100 Subject: [PATCH 09/48] refactor(array): make ArrayIO.build synchronous and declare its surface once MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Binding is pure, so build stops being awaitable: it takes a Store or a StorePath and binds with StorePath construction, never make_store_path, whose coercion of path and URL strings opens stores — the IO the method promises not to do. Stores open themselves on their first IO, so an unopened store binds safely. The declaration keyword surface is declared once, as the ArrayDeclarationKwargs TypedDict that both build methods unpack, instead of being copied into each; zarr_format on build no longer admits None, whose config-default meaning collided with open's auto-detect meaning one method away. The constructors return Self, and build_array_metadata's fill_value drops a redundant None from its type. Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 2 +- src/zarr/core/array.py | 125 ++++++++++++++++------------------------- tests/test_array_io.py | 18 +++--- 3 files changed, 57 insertions(+), 88 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index 9dd9f28247..27b60ce61c 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,3 +1,3 @@ -Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. +Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `build` is synchronous — binding is pure — and takes a `Store` or `StorePath`; `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. `Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. The binding holds the metadata document as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.document`: document operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 33b7a7daef..b0cda1db44 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -14,7 +14,9 @@ TYPE_CHECKING, Any, Literal, + Required, TypedDict, + Unpack, cast, overload, ) @@ -4601,7 +4603,7 @@ def build_array_metadata( filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", serializer: SerializerLike = "auto", - fill_value: Any | None = DEFAULT_FILL_VALUE, + fill_value: Any = DEFAULT_FILL_VALUE, order: MemoryOrder | None = None, zarr_format: ZarrFormat | None = 3, attributes: dict[str, JSON] | None = None, @@ -6216,6 +6218,27 @@ def _array_document_model(data: dict[str, JSON]) -> ArrayDocumentModel: return ZarrV3ArrayMetadata(data) +class ArrayDeclarationKwargs(TypedDict, total=False): + """The keyword arguments that declare an array: what `build_array_metadata` + takes, with the same meanings as in [`zarr.api.asynchronous.create_array`][]. + This is the one declaration of the surface; `ArrayIO.build` and + `AsyncArrayIO.build` take it as their keyword arguments.""" + + shape: Required[ShapeLike] + dtype: Required[ZDTypeLike] + chunks: ChunksLike | Literal["auto"] + shards: ShardsLike | None + filters: FiltersLike + compressors: CompressorsLike + serializer: SerializerLike + fill_value: Any + order: MemoryOrder | None + zarr_format: ZarrFormat + attributes: dict[str, JSON] | None + chunk_key_encoding: ChunkKeyEncodingLike | None + dimension_names: DimensionNamesLike + + class AsyncArrayIO: """Explicit asynchronous IO for one Zarr array. @@ -6287,48 +6310,29 @@ def __init__( self._source = parse_array_metadata(document, str(store_path)) @classmethod - async def build( + def build( cls, - store: StoreLike, + store: Store | StorePath, *, name: str | None = None, - shape: ShapeLike, - dtype: ZDTypeLike, - chunks: ChunksLike | Literal["auto"] = "auto", - shards: ShardsLike | None = None, - filters: FiltersLike = "auto", - compressors: CompressorsLike = "auto", - serializer: SerializerLike = "auto", - fill_value: Any | None = DEFAULT_FILL_VALUE, - order: MemoryOrder | None = None, - zarr_format: ZarrFormat | None = 3, - attributes: dict[str, JSON] | None = None, - chunk_key_encoding: ChunkKeyEncodingLike | None = None, - dimension_names: DimensionNamesLike = None, config: ArrayConfigLike | None = None, - ) -> AsyncArrayIO: + **kwargs: Unpack[ArrayDeclarationKwargs], + ) -> Self: """Declare an array bound to a location in a store, without storing anything. - The keyword arguments have the same meaning as in - [`zarr.api.asynchronous.create_array`][]. The returned instance holds - the normalized metadata in memory; `store_metadata` stores it. - """ - metadata = build_array_metadata( - shape=shape, - dtype=dtype, - chunks=chunks, - shards=shards, - filters=filters, - compressors=compressors, - serializer=serializer, - fill_value=fill_value, - order=order, - zarr_format=zarr_format, - attributes=attributes, - chunk_key_encoding=chunk_key_encoding, - dimension_names=dimension_names, - ) - store_path = await make_store_path(store, path=name) + Binding is pure, which is why this is the one synchronous constructor on + this class: `store` must already be a `Store` or `StorePath`, because + coercing a path or URL string opens a store, and that is + [`zarr.api.asynchronous.create_array`][]'s job. The declaration keyword + arguments have the same meaning as in `create_array`; unlike `zarrs`, + `fill_value` defaults to the data type's default fill value instead of + being required. The returned instance holds the normalized metadata in + memory; `store_metadata` stores it. + """ + metadata = build_array_metadata(**kwargs) + store_path = store if isinstance(store, StorePath) else StorePath(store) + if name is not None: + store_path = store_path / name return cls(metadata, store_path, config) @classmethod @@ -6339,7 +6343,7 @@ async def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, - ) -> AsyncArrayIO: + ) -> Self: """Read the metadata document of an existing array and bind to it. This is the one classmethod constructor that performs IO: the array's @@ -6718,7 +6722,7 @@ def __init__( self._async_io = AsyncArrayIO(document, store_path, config) @classmethod - def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: + def _wrap(cls, async_io: AsyncArrayIO) -> Self: io = object.__new__(cls) io._async_io = async_io return io @@ -6726,50 +6730,17 @@ def _wrap(cls, async_io: AsyncArrayIO) -> ArrayIO: @classmethod def build( cls, - store: StoreLike, + store: Store | StorePath, *, name: str | None = None, - shape: ShapeLike, - dtype: ZDTypeLike, - chunks: ChunksLike | Literal["auto"] = "auto", - shards: ShardsLike | None = None, - filters: FiltersLike = "auto", - compressors: CompressorsLike = "auto", - serializer: SerializerLike = "auto", - fill_value: Any | None = DEFAULT_FILL_VALUE, - order: MemoryOrder | None = None, - zarr_format: ZarrFormat | None = 3, - attributes: dict[str, JSON] | None = None, - chunk_key_encoding: ChunkKeyEncodingLike | None = None, - dimension_names: DimensionNamesLike = None, config: ArrayConfigLike | None = None, - ) -> ArrayIO: + **kwargs: Unpack[ArrayDeclarationKwargs], + ) -> Self: """Declare an array bound to a location in a store, without storing anything. See `AsyncArrayIO.build`. """ - return cls._wrap( - sync( - AsyncArrayIO.build( - store, - name=name, - shape=shape, - dtype=dtype, - chunks=chunks, - shards=shards, - filters=filters, - compressors=compressors, - serializer=serializer, - fill_value=fill_value, - order=order, - zarr_format=zarr_format, - attributes=attributes, - chunk_key_encoding=chunk_key_encoding, - dimension_names=dimension_names, - config=config, - ) - ) - ) + return cls._wrap(AsyncArrayIO.build(store, name=name, config=config, **kwargs)) @classmethod def open( @@ -6779,7 +6750,7 @@ def open( name: str | None = None, zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, - ) -> ArrayIO: + ) -> Self: """Read the metadata document of an existing array and bind to it. See `AsyncArrayIO.open`. diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 68f60c2e5f..2784e217d5 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -44,7 +44,7 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: in that order, and the result is an ordinary readable array.""" store = MemoryStore() data = np.arange(16, dtype="uint16").reshape(4, 4) - aio = await AsyncArrayIO.build( + aio = AsyncArrayIO.build( store, name="foo", shape=data.shape, @@ -78,7 +78,7 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: np.testing.assert_array_equal(await arr.io.getitem(Ellipsis), data) # a writer that declared the same array independently can verify its assumption - worker = await AsyncArrayIO.build( + worker = AsyncArrayIO.build( store, name="foo", shape=data.shape, @@ -171,9 +171,9 @@ async def test_document_work_without_codec_resolution() -> None: async def test_verify_mismatch_raises() -> None: store = MemoryStore() - aio = await AsyncArrayIO.build(store, shape=(4, 4), dtype="uint8", chunks=(2, 2)) + aio = AsyncArrayIO.build(store, shape=(4, 4), dtype="uint8", chunks=(2, 2)) await aio.store_metadata() - other = await AsyncArrayIO.build( + other = AsyncArrayIO.build( store, shape=(4, 4), dtype="uint8", chunks=(4, 4), attributes={"a": 1} ) with pytest.raises(ValueError, match="differing fields: attributes, chunk_grid"): @@ -182,22 +182,20 @@ async def test_verify_mismatch_raises() -> None: async def test_verify_format_mismatch_raises() -> None: store = MemoryStore() - await ( - await AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=2) - ).store_metadata() - bound_v3 = await AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=3) + await AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=2).store_metadata() + bound_v3 = AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=3) with pytest.raises(ValueError, match="stored array is Zarr format 2"): await bound_v3.verify() async def test_verify_missing_raises() -> None: - aio = await AsyncArrayIO.build(MemoryStore(), shape=(4,), dtype="uint8") + aio = AsyncArrayIO.build(MemoryStore(), shape=(4,), dtype="uint8") with pytest.raises(ArrayNotFoundError): await aio.verify() async def test_erase_chunk_out_of_grid_raises() -> None: - aio = await AsyncArrayIO.build(MemoryStore(), shape=(4, 4), dtype="uint8", chunks=(2, 2)) + aio = AsyncArrayIO.build(MemoryStore(), shape=(4, 4), dtype="uint8", chunks=(2, 2)) with pytest.raises(IndexError, match="outside the chunk grid"): await aio.erase_chunk((2, 0)) with pytest.raises(IndexError, match="outside the chunk grid"): From 8876ff83e73d9cfa4a584510b0718270329e39b8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 15:32:16 +0100 Subject: [PATCH 10/48] refactor(array): build each format's metadata document with its own routine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing hides the v2 / v3 difference any more: build_array_metadata is replaced by module-level build_array_metadata_v2 and build_array_metadata_v3, each taking the actual fields of its document (.zarray plus .zattrs, and zarr.json), with the bodies lifted from the AsyncArray._create_metadata_* staticmethods, which now delegate (the module-level home #4175 asks for). init_array regains its normalization and calls the per-format builders. ArrayIO and AsyncArrayIO declare with build_v2 and build_v3 over per-format TypedDicts, so a format mistake — shards on a v2 array, order on a v3 one — is inexpressible instead of guarded, and build_v3's chunk_grid accepts the grid metadata, its JSON, or a chunk shape declaring a regular grid, as ArrayBuilder does in zarrs. Refs zarr-developers/zarr-python#4175 Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 2 +- src/zarr/core/array.py | 357 ++++++++++++++++++++++++++--------------- tests/test_array_io.py | 71 ++++---- 3 files changed, 264 insertions(+), 166 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index 27b60ce61c..f249a6e75b 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,3 +1,3 @@ -Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. `build` is synchronous — binding is pure — and takes a `Store` or `StorePath`; `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. +Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. The declaring constructors, `build_v2` and `build_v3`, are synchronous — binding is pure — take a `Store` or `StorePath`, and take the fields of each format's metadata document, with nothing hiding the difference between the formats; the documents themselves are built by the module-level `build_array_metadata_v2` and `build_array_metadata_v3`. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. `Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. The binding holds the metadata document as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.document`: document operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index b0cda1db44..61ed7c8096 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -187,7 +187,8 @@ "DEFAULT_FILL_VALUE", "ArrayIO", "AsyncArrayIO", - "build_array_metadata", + "build_array_metadata_v2", + "build_array_metadata_v3", "create_codec_pipeline", "parse_array_metadata", ] @@ -618,43 +619,16 @@ def _create_metadata_v3( dimension_names: DimensionNamesLike = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV3Metadata: - """Create an instance of ArrayV3Metadata.""" - filters: tuple[ArrayArrayCodec, ...] - compressors: tuple[BytesBytesCodec, ...] - - shape = parse_shapelike(shape) - if codecs is None: - # no data types have default filters - filters = () - serializer = default_serializer_v3(dtype) - compressors = default_compressors_v3(dtype) - - codecs_parsed = (*filters, serializer, *compressors) - else: - codecs_parsed = tuple(codecs) - - chunk_key_encoding_parsed: ChunkKeyEncodingLike - if chunk_key_encoding is None: - chunk_key_encoding_parsed = {"name": "default", "separator": "/"} - else: - chunk_key_encoding_parsed = chunk_key_encoding - - if fill_value is DEFAULT_FILL_VALUE or fill_value is None: - # Use dtype's default scalar for the DEFAULT_FILL_VALUE sentinel - # For v3, None is converted to DEFAULT_FILL_VALUE behavior - fill_value_parsed = dtype.default_scalar() - else: - fill_value_parsed = fill_value - - return ArrayV3Metadata( + """See the module-level `build_array_metadata_v3`.""" + return build_array_metadata_v3( shape=shape, - data_type=dtype, + dtype=dtype, chunk_grid=chunk_grid, - chunk_key_encoding=chunk_key_encoding_parsed, - fill_value=fill_value_parsed, - codecs=codecs_parsed, # type: ignore[arg-type] - dimension_names=tuple(dimension_names) if dimension_names else None, - attributes=attributes or {}, + fill_value=fill_value, + chunk_key_encoding=chunk_key_encoding, + codecs=codecs, + dimension_names=dimension_names, + attributes=attributes, ) @classmethod @@ -711,25 +685,16 @@ def _create_metadata_v2( compressor: CompressorLikev2 = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV2Metadata: - if dimension_separator is None: - dimension_separator = "." - - # Handle the DEFAULT_FILL_VALUE sentinel - if fill_value is DEFAULT_FILL_VALUE: - fill_value_parsed: Any = dtype.default_scalar() - else: - # For v2, preserve None as-is (backward compatibility) - fill_value_parsed = fill_value - - return ArrayV2Metadata( + """See the module-level `build_array_metadata_v2`.""" + return build_array_metadata_v2( shape=shape, dtype=dtype, chunks=chunks, order=order, - dimension_separator=dimension_separator, - fill_value=fill_value_parsed, - compressor=compressor, + dimension_separator=dimension_separator if dimension_separator is not None else ".", + fill_value=fill_value, filters=filters, + compressor=compressor, attributes=attributes, ) @@ -4574,49 +4539,6 @@ async def init_array( AsyncArray The AsyncArray. """ - meta = build_array_metadata( - shape=shape, - dtype=dtype, - chunks=chunks, - shards=shards, - filters=filters, - compressors=compressors, - serializer=serializer, - fill_value=fill_value, - order=order, - zarr_format=zarr_format, - attributes=attributes, - chunk_key_encoding=chunk_key_encoding, - dimension_names=dimension_names, - ) - arr = AsyncArray(metadata=meta, store_path=store_path, config=config) - await save_new_metadata(store_path, meta, overwrite=overwrite, ensure_parents=True) - return arr - - -def build_array_metadata( - *, - shape: ShapeLike, - dtype: ZDTypeLike, - chunks: ChunksLike | Literal["auto"] = "auto", - shards: ShardsLike | None = None, - filters: FiltersLike = "auto", - compressors: CompressorsLike = "auto", - serializer: SerializerLike = "auto", - fill_value: Any = DEFAULT_FILL_VALUE, - order: MemoryOrder | None = None, - zarr_format: ZarrFormat | None = 3, - attributes: dict[str, JSON] | None = None, - chunk_key_encoding: ChunkKeyEncodingLike | None = None, - dimension_names: DimensionNamesLike = None, -) -> ArrayV2Metadata | ArrayV3Metadata: - """Build an array metadata document from creation-style keyword arguments. - - This is the pure half of `init_array`: it normalizes the same keyword - arguments into an in-memory metadata object without touching any store. - The parameters have the same meaning as in - [`zarr.api.asynchronous.create_array`][]. - """ if zarr_format is None: zarr_format = _default_zarr_format() @@ -4682,7 +4604,7 @@ def build_array_metadata( order_parsed = order chunk_key_encoding_parsed = cast("V2ChunkKeyEncoding", chunk_key_encoding_parsed) - meta = AsyncArray._create_metadata_v2( + meta = build_array_metadata_v2( shape=shape_parsed, dtype=zdtype, chunks=outer_chunks.chunk_shape, @@ -4732,7 +4654,7 @@ def build_array_metadata( if order is not None: _warn_order_kwarg() - meta = AsyncArray._create_metadata_v3( + meta = build_array_metadata_v3( shape=shape_parsed, dtype=zdtype, fill_value=fill_value, @@ -4743,7 +4665,118 @@ def build_array_metadata( attributes=attributes, ) - return meta + arr = AsyncArray(metadata=meta, store_path=store_path, config=config) + await save_new_metadata(store_path, meta, overwrite=overwrite, ensure_parents=True) + return arr + + +def build_array_metadata_v2( + *, + shape: ShapeLike, + dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + chunks: ShapeLike, + fill_value: Any = DEFAULT_FILL_VALUE, + order: MemoryOrder = "C", + dimension_separator: Literal[".", "/"] = ".", + filters: Iterable[dict[str, JSON] | Numcodec] | None = None, + compressor: CompressorLikev2 = None, + attributes: dict[str, JSON] | None = None, +) -> ArrayV2Metadata: + """Build a Zarr format 2 array metadata document from its fields. + + Each keyword is a field of the `.zarray` document, with `attributes` the + `.zattrs` document. `fill_value` defaults to the data type's default + scalar, and an explicit `None` is kept, because the field is nullable in + Zarr format 2. `filters` and `compressor` are numcodecs objects or their + JSON, exactly as the format stores them — nothing is guessed; the + guessing creation API is [`zarr.api.asynchronous.create_array`][]. + """ + zdtype = parse_dtype(dtype, zarr_format=2) + if fill_value is DEFAULT_FILL_VALUE: + fill_value_parsed: Any = zdtype.default_scalar() + else: + # An explicit None is kept: the field is nullable in Zarr format 2. + fill_value_parsed = fill_value + + return ArrayV2Metadata( + shape=parse_shapelike(shape), + dtype=zdtype, + chunks=parse_shapelike(chunks), + order=order, + dimension_separator=dimension_separator, + fill_value=fill_value_parsed, + compressor=compressor, + filters=filters, + attributes=attributes, + ) + + +def build_array_metadata_v3( + *, + shape: ShapeLike, + dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata | Mapping[str, JSON] | ShapeLike, + fill_value: Any = DEFAULT_FILL_VALUE, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + codecs: Iterable[Codec | dict[str, JSON]] | None = None, + dimension_names: DimensionNamesLike = None, + attributes: dict[str, JSON] | None = None, +) -> ArrayV3Metadata: + """Build a Zarr format 3 array metadata document from its fields. + + Each keyword is a field of the `zarr.json` document. `chunk_grid` takes + the grid's metadata or its JSON, or a chunk shape, which declares a + regular grid, as `ArrayBuilder` does in `zarrs`. `codecs` of `None` + declares the default chain for the data type (no filters, the default + serializer, the default compressors), and `chunk_key_encoding` of `None` + the default encoding with the `/` separator. `fill_value` defaults to the + data type's default scalar, and `None` also means that default, because + the field is required in Zarr format 3. + """ + zdtype = parse_dtype(dtype, zarr_format=3) + grid: ChunkGridMetadata | dict[str, JSON] + if isinstance(chunk_grid, RegularChunkGridMetadata | RectilinearChunkGridMetadata): + grid = chunk_grid + elif isinstance(chunk_grid, Mapping): + grid = dict(chunk_grid) + else: + grid = RegularChunkGridMetadata(chunk_shape=parse_shapelike(chunk_grid)) + + filters: tuple[ArrayArrayCodec, ...] + compressors: tuple[BytesBytesCodec, ...] + if codecs is None: + # no data types have default filters + filters = () + serializer = default_serializer_v3(zdtype) + compressors = default_compressors_v3(zdtype) + + codecs_parsed = (*filters, serializer, *compressors) + else: + codecs_parsed = tuple(codecs) + + chunk_key_encoding_parsed: ChunkKeyEncodingLike + if chunk_key_encoding is None: + chunk_key_encoding_parsed = {"name": "default", "separator": "/"} + else: + chunk_key_encoding_parsed = chunk_key_encoding + + if fill_value is DEFAULT_FILL_VALUE or fill_value is None: + # The dtype's default scalar: the field is required in Zarr format 3, + # so None also means the default. + fill_value_parsed = zdtype.default_scalar() + else: + fill_value_parsed = fill_value + + return ArrayV3Metadata( + shape=parse_shapelike(shape), + data_type=zdtype, + chunk_grid=grid, + chunk_key_encoding=chunk_key_encoding_parsed, + fill_value=fill_value_parsed, + codecs=codecs_parsed, # type: ignore[arg-type] + dimension_names=tuple(dimension_names) if dimension_names else None, + attributes=attributes if attributes is not None else {}, + ) async def create_array( @@ -6218,25 +6251,46 @@ def _array_document_model(data: dict[str, JSON]) -> ArrayDocumentModel: return ZarrV3ArrayMetadata(data) -class ArrayDeclarationKwargs(TypedDict, total=False): - """The keyword arguments that declare an array: what `build_array_metadata` - takes, with the same meanings as in [`zarr.api.asynchronous.create_array`][]. - This is the one declaration of the surface; `ArrayIO.build` and - `AsyncArrayIO.build` take it as their keyword arguments.""" +class ArrayV2DeclarationKwargs(TypedDict, total=False): + """The fields that declare a Zarr format 2 array: what + `build_array_metadata_v2` takes. This is the one declaration of the + surface; `ArrayIO.build_v2` and `AsyncArrayIO.build_v2` take it as their + keyword arguments.""" shape: Required[ShapeLike] - dtype: Required[ZDTypeLike] - chunks: ChunksLike | Literal["auto"] - shards: ShardsLike | None - filters: FiltersLike - compressors: CompressorsLike - serializer: SerializerLike + dtype: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] + chunks: Required[ShapeLike] fill_value: Any - order: MemoryOrder | None - zarr_format: ZarrFormat + order: MemoryOrder + dimension_separator: Literal[".", "/"] + filters: Iterable[dict[str, JSON] | Numcodec] | None + compressor: CompressorLikev2 attributes: dict[str, JSON] | None + + +class ArrayV3DeclarationKwargs(TypedDict, total=False): + """The fields that declare a Zarr format 3 array: what + `build_array_metadata_v3` takes. This is the one declaration of the + surface; `ArrayIO.build_v3` and `AsyncArrayIO.build_v3` take it as their + keyword arguments.""" + + shape: Required[ShapeLike] + dtype: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] + chunk_grid: Required[ChunkGridMetadata | Mapping[str, JSON] | ShapeLike] + fill_value: Any chunk_key_encoding: ChunkKeyEncodingLike | None + codecs: Iterable[Codec | dict[str, JSON]] | None dimension_names: DimensionNamesLike + attributes: dict[str, JSON] | None + + +def _bound_store_path(store: Store | StorePath, name: str | None) -> StorePath: + """The location `name` within `store`, bound without IO: a path or URL + string is not accepted, because coercing one opens a store.""" + store_path = store if isinstance(store, StorePath) else StorePath(store) + if name is not None: + store_path = store_path / name + return store_path class AsyncArrayIO: @@ -6310,30 +6364,52 @@ def __init__( self._source = parse_array_metadata(document, str(store_path)) @classmethod - def build( + def build_v2( + cls, + store: Store | StorePath, + *, + name: str | None = None, + config: ArrayConfigLike | None = None, + **kwargs: Unpack[ArrayV2DeclarationKwargs], + ) -> Self: + """Declare a Zarr format 2 array bound to a location in a store, without + storing anything. + + Binding is pure, which is why the declaring constructors are the + synchronous methods on this class: `store` must already be a `Store` + or `StorePath`, because coercing a path or URL string opens a store, + and that is [`zarr.api.asynchronous.create_array`][]'s job — as is + guessing: every keyword here is a field of the `.zarray` document (see + `build_array_metadata_v2`). Unlike `zarrs`, `fill_value` defaults to + the data type's default fill value instead of being required. The + returned instance holds the metadata in memory; `store_metadata` + stores it. + """ + return cls(build_array_metadata_v2(**kwargs), _bound_store_path(store, name), config) + + @classmethod + def build_v3( cls, store: Store | StorePath, *, name: str | None = None, config: ArrayConfigLike | None = None, - **kwargs: Unpack[ArrayDeclarationKwargs], + **kwargs: Unpack[ArrayV3DeclarationKwargs], ) -> Self: - """Declare an array bound to a location in a store, without storing anything. - - Binding is pure, which is why this is the one synchronous constructor on - this class: `store` must already be a `Store` or `StorePath`, because - coercing a path or URL string opens a store, and that is - [`zarr.api.asynchronous.create_array`][]'s job. The declaration keyword - arguments have the same meaning as in `create_array`; unlike `zarrs`, - `fill_value` defaults to the data type's default fill value instead of - being required. The returned instance holds the normalized metadata in - memory; `store_metadata` stores it. - """ - metadata = build_array_metadata(**kwargs) - store_path = store if isinstance(store, StorePath) else StorePath(store) - if name is not None: - store_path = store_path / name - return cls(metadata, store_path, config) + """Declare a Zarr format 3 array bound to a location in a store, without + storing anything. + + Binding is pure, which is why the declaring constructors are the + synchronous methods on this class: `store` must already be a `Store` + or `StorePath`, because coercing a path or URL string opens a store, + and that is [`zarr.api.asynchronous.create_array`][]'s job — as is + guessing: every keyword here is a field of the `zarr.json` document + (see `build_array_metadata_v3`). Unlike `zarrs`, `fill_value` defaults + to the data type's default fill value instead of being required. The + returned instance holds the metadata in memory; `store_metadata` + stores it. + """ + return cls(build_array_metadata_v3(**kwargs), _bound_store_path(store, name), config) @classmethod async def open( @@ -6728,19 +6804,36 @@ def _wrap(cls, async_io: AsyncArrayIO) -> Self: return io @classmethod - def build( + def build_v2( + cls, + store: Store | StorePath, + *, + name: str | None = None, + config: ArrayConfigLike | None = None, + **kwargs: Unpack[ArrayV2DeclarationKwargs], + ) -> Self: + """Declare a Zarr format 2 array bound to a location in a store, without + storing anything. + + See `AsyncArrayIO.build_v2`. + """ + return cls._wrap(AsyncArrayIO.build_v2(store, name=name, config=config, **kwargs)) + + @classmethod + def build_v3( cls, store: Store | StorePath, *, name: str | None = None, config: ArrayConfigLike | None = None, - **kwargs: Unpack[ArrayDeclarationKwargs], + **kwargs: Unpack[ArrayV3DeclarationKwargs], ) -> Self: - """Declare an array bound to a location in a store, without storing anything. + """Declare a Zarr format 3 array bound to a location in a store, without + storing anything. - See `AsyncArrayIO.build`. + See `AsyncArrayIO.build_v3`. """ - return cls._wrap(AsyncArrayIO.build(store, name=name, config=config, **kwargs)) + return cls._wrap(AsyncArrayIO.build_v3(store, name=name, config=config, **kwargs)) @classmethod def open( diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 2784e217d5..58271d2788 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import numpy as np import pytest @@ -24,6 +24,29 @@ async def _store_keys(store: MemoryStore) -> set[str]: return {key async for key in store.list()} +def _declare_answer_array( + zarr_format: ZarrFormat, store: MemoryStore, data: np.ndarray[Any, np.dtype[Any]] +) -> AsyncArrayIO: + """Declare the combo test's array with the builder of `zarr_format`.""" + if zarr_format == 2: + return AsyncArrayIO.build_v2( + store, + name="foo", + shape=data.shape, + dtype=data.dtype, + chunks=(2, 2), + attributes={"answer": 42}, + ) + return AsyncArrayIO.build_v3( + store, + name="foo", + shape=data.shape, + dtype=data.dtype, + chunk_grid=(2, 2), + attributes={"answer": 42}, + ) + + async def _store_foreign_codec_array(store: MemoryStore) -> None: """Store a valid v3 array document whose codec this environment lacks.""" document = ( @@ -44,15 +67,7 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: in that order, and the result is an ordinary readable array.""" store = MemoryStore() data = np.arange(16, dtype="uint16").reshape(4, 4) - aio = AsyncArrayIO.build( - store, - name="foo", - shape=data.shape, - dtype=data.dtype, - chunks=(2, 2), - zarr_format=zarr_format, - attributes={"answer": 42}, - ) + aio = _declare_answer_array(zarr_format, store, data) assert await _store_keys(store) == set() await aio.setitem(Ellipsis, data) @@ -78,15 +93,7 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: np.testing.assert_array_equal(await arr.io.getitem(Ellipsis), data) # a writer that declared the same array independently can verify its assumption - worker = AsyncArrayIO.build( - store, - name="foo", - shape=data.shape, - dtype=data.dtype, - chunks=(2, 2), - zarr_format=zarr_format, - attributes={"answer": 42}, - ) + worker = _declare_answer_array(zarr_format, store, data) await worker.verify() # every write has an erasing inverse: one chunk, all chunks, then metadata @@ -112,7 +119,7 @@ def test_chunks_before_metadata_sync() -> None: """The synchronous `ArrayIO` supports the same chunks-first flow.""" store = MemoryStore() data = np.arange(16, dtype="uint16").reshape(4, 4) - aio = ArrayIO.build(store, name="foo", shape=data.shape, dtype=data.dtype, chunks=(2, 2)) + aio = ArrayIO.build_v3(store, name="foo", shape=data.shape, dtype=data.dtype, chunk_grid=(2, 2)) aio.setitem(Ellipsis, data) aio.store_metadata() opened = ArrayIO.open(store, name="foo") @@ -127,16 +134,14 @@ def test_chunks_before_metadata_sync() -> None: ArrayIO.open(store, name="foo") -def test_build_v2_shards_raises() -> None: - with pytest.raises(ValueError, match="shard_shape"): - ArrayIO.build( +def test_build_v3_unknown_chunk_grid_raises() -> None: + with pytest.raises(ValueError, match="chunk grid"): + ArrayIO.build_v3( MemoryStore(), name="foo", shape=(4,), dtype="uint8", - chunks=(2,), - shards=(4,), - zarr_format=2, + chunk_grid={"name": "diagonal", "configuration": {}}, ) @@ -171,10 +176,10 @@ async def test_document_work_without_codec_resolution() -> None: async def test_verify_mismatch_raises() -> None: store = MemoryStore() - aio = AsyncArrayIO.build(store, shape=(4, 4), dtype="uint8", chunks=(2, 2)) + aio = AsyncArrayIO.build_v3(store, shape=(4, 4), dtype="uint8", chunk_grid=(2, 2)) await aio.store_metadata() - other = AsyncArrayIO.build( - store, shape=(4, 4), dtype="uint8", chunks=(4, 4), attributes={"a": 1} + other = AsyncArrayIO.build_v3( + store, shape=(4, 4), dtype="uint8", chunk_grid=(4, 4), attributes={"a": 1} ) with pytest.raises(ValueError, match="differing fields: attributes, chunk_grid"): await other.verify() @@ -182,20 +187,20 @@ async def test_verify_mismatch_raises() -> None: async def test_verify_format_mismatch_raises() -> None: store = MemoryStore() - await AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=2).store_metadata() - bound_v3 = AsyncArrayIO.build(store, shape=(4,), dtype="uint8", zarr_format=3) + await AsyncArrayIO.build_v2(store, shape=(4,), dtype="uint8", chunks=(4,)).store_metadata() + bound_v3 = AsyncArrayIO.build_v3(store, shape=(4,), dtype="uint8", chunk_grid=(4,)) with pytest.raises(ValueError, match="stored array is Zarr format 2"): await bound_v3.verify() async def test_verify_missing_raises() -> None: - aio = AsyncArrayIO.build(MemoryStore(), shape=(4,), dtype="uint8") + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4,), dtype="uint8", chunk_grid=(4,)) with pytest.raises(ArrayNotFoundError): await aio.verify() async def test_erase_chunk_out_of_grid_raises() -> None: - aio = AsyncArrayIO.build(MemoryStore(), shape=(4, 4), dtype="uint8", chunks=(2, 2)) + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4, 4), dtype="uint8", chunk_grid=(2, 2)) with pytest.raises(IndexError, match="outside the chunk grid"): await aio.erase_chunk((2, 0)) with pytest.raises(IndexError, match="outside the chunk grid"): From 8956b961d4d5ea9601e4f2ef18e63f3e6650269a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 15:57:16 +0100 Subject: [PATCH 11/48] refactor(array): name the model metadata and the resolved form metadata_legacy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ArrayIO's metadata is the zarr_metadata model, full stop: the property is named metadata, the resolved ArrayV2Metadata / ArrayV3Metadata view — which the data path runs on and AsyncArray.metadata still returns — is metadata_legacy, and the term document leaves the surface. The with_metadata, with_store and with_config methods return new bindings with one part replaced, with with_config merging partial configuration the way AsyncArray.with_config does. Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 3 +- src/zarr/core/array.py | 195 +++++++++++++++++++++++++---------------- tests/test_array_io.py | 48 +++++++--- 3 files changed, 160 insertions(+), 86 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index f249a6e75b..5b5bbb8c22 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,3 +1,4 @@ Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. The declaring constructors, `build_v2` and `build_v3`, are synchronous — binding is pure — take a `Store` or `StorePath`, and take the fields of each format's metadata document, with nothing hiding the difference between the formats; the documents themselves are built by the module-level `build_array_metadata_v2` and `build_array_metadata_v3`. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. `Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. -The binding holds the metadata document as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.document`: document operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. +The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. +`with_metadata`, `with_store` and `with_config` return new bindings with one part replaced. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 61ed7c8096..4eb1d7e639 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -447,7 +447,7 @@ def io(self) -> AsyncArrayIO: @property def metadata(self) -> T_ArrayMetadata: """The in-memory metadata of this array.""" - return cast("T_ArrayMetadata", self._io.metadata) + return cast("T_ArrayMetadata", self._io.metadata_legacy) @property def store_path(self) -> StorePath: @@ -6224,20 +6224,20 @@ async def _delete_key(store_path: StorePath, key: str) -> None: await (store_path / key).delete() -type ArrayDocumentModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata -"""A Zarr array metadata document, held as a `zarr_metadata` model: the JSON as -written, validated, with extension names read but never resolved into classes.""" +type ArrayMetadataModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata +"""Array metadata as a `zarr_metadata` model: the JSON as written, validated, +with extension names read but never resolved into classes.""" -def _document_json(document: ArrayDocumentModel) -> dict[str, Any]: - """The document as plain, mutable JSON. The model returns read-only forms — +def _metadata_model_json(document: ArrayMetadataModel) -> dict[str, Any]: + """The metadata as plain, mutable JSON. The model returns read-only forms — arrays as tuples — while zarr's parsers take the stored forms, so the - document goes through JSON text.""" + metadata goes through JSON text.""" return cast("dict[str, Any]", json.loads(json.dumps(document.to_json()))) -def _array_document_model(data: dict[str, JSON]) -> ArrayDocumentModel: - """The model of an array metadata document in the form `get_array_metadata` +def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: + """Array metadata as a `zarr_metadata` model, from the form `get_array_metadata` returns: the Zarr format 2 form carries its attributes merged in. The read is strict: a document with problems — including one zarr can only @@ -6296,26 +6296,27 @@ def _bound_store_path(store: Store | StorePath, name: str | None) -> StorePath: class AsyncArrayIO: """Explicit asynchronous IO for one Zarr array. - An `AsyncArrayIO` binds an array metadata document to a location in a - store. The document is held as a `zarr_metadata` model — the JSON as - written, with extension names read but not resolved into codec or data - type classes — so every document operation (`store_metadata`, - `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) works without - resolving anything: an array whose codecs this environment lacks can - still be opened, verified and erased. Reading and writing chunks - (`getitem` / `setitem`) resolves the document into classes on first use; - `metadata` is that resolved form, derived once from `document` — or, when - the instance was built from an already-resolved metadata object (as every - `AsyncArray` is), held directly, with `document` the derived one. - - Constructing one — directly, or with `build` — performs no IO: nothing is - read from or written to the store until an IO method is called. - `store_metadata` stores the metadata document, and `setitem` / `getitem` + An `AsyncArrayIO` binds array metadata to a location in a store. The + `metadata` is a `zarr_metadata` model — the JSON as written, with + extension names read but not resolved into codec or data type classes — + so every metadata operation (`store_metadata`, `erase_metadata`, + `erase_chunk`, `erase_chunks`, `verify`) works without resolving + anything: an array whose codecs this environment lacks can still be + opened, verified and erased. Reading and writing chunks (`getitem` / + `setitem`) resolves the metadata into the legacy classes on first use; + `metadata_legacy` is that resolved form, derived once — or, when the + instance was built from one of those objects (as every `AsyncArray` is), + carried as the source `metadata` derives from. + + Constructing one — directly, or with `build_v2` / `build_v3` — performs + no IO: nothing is read from or written to the store until an IO method is + called. `store_metadata` stores the metadata, and `setitem` / `getitem` store and read chunks, in whatever order the caller chooses. Every write has an erasing inverse: `store_metadata` ↔ `erase_metadata`, and chunk - writes ↔ `erase_chunk` / `erase_chunks`. `AsyncArray` is an array-like - layer over this one: it wraps an `AsyncArrayIO` and delegates its metadata - and selection IO to it. + writes ↔ `erase_chunk` / `erase_chunks`. The `with_metadata`, `with_store` + and `with_config` methods return new bindings with one part replaced. + `AsyncArray` is an array-like layer over this one: it wraps an + `AsyncArrayIO` and delegates its metadata and selection IO to it. `store_metadata` is the publication point: until it is called, no Zarr reader resolves the array, so a failed chunks-first run leaves no readable @@ -6323,20 +6324,20 @@ class AsyncArrayIO: be misread by a different array stored later at the same location; `erase_chunks` removes them. - Warning: the bound document is assumed, never verified against the store. - Reads and partial writes decode stored chunks under it, so binding a - document that disagrees with how existing chunks at this location were + Warning: the bound metadata is assumed, never verified against the store. + Reads and partial writes decode stored chunks under it, so binding + metadata that disagrees with how existing chunks at this location were encoded misreads or corrupts data — the hazard TensorStore documents for `assume_metadata`. One exception to "no implicit metadata IO": an array - opened from a stored document that needed a repair stores that repair + opened from stored metadata that needed a repair stores that repair before its first chunk write. Call `store_metadata` first to make that write explicit. """ _store_path: StorePath _config: ArrayConfig - _source: ArrayDocumentModel | ArrayMetadata - _document_cache: ArrayDocumentModel | None + _source: ArrayMetadataModel | ArrayMetadata + _metadata_model_cache: ArrayMetadataModel | None _resolved_cache: ArrayMetadata | None _chunk_grid_cache: ChunkGrid | None _chunk_key_encoding_cache: ChunkKeyEncoding | None @@ -6344,24 +6345,23 @@ class AsyncArrayIO: def __init__( self, - document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: self._store_path = store_path self._config = parse_array_config(config) - self._document_cache = None + self._metadata_model_cache = None self._resolved_cache = None self._chunk_grid_cache = None self._chunk_key_encoding_cache = None self._codec_pipeline_cache = None - if isinstance(document, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): - self._source = document + if isinstance(metadata, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): + self._source = metadata else: - # A resolved metadata object passes through; a raw document dict is - # read eagerly, as repaired when it needs that, warning as `AsyncArray` - # always has. - self._source = parse_array_metadata(document, str(store_path)) + # A legacy resolved object passes through; raw JSON is read eagerly, + # as repaired when it needs that, warning as `AsyncArray` always has. + self._source = parse_array_metadata(metadata, str(store_path)) @classmethod def build_v2( @@ -6434,32 +6434,34 @@ async def open( """ store_path = await make_store_path(store, path=name) data = await get_array_metadata(store_path, zarr_format=zarr_format) - return cls(_array_document_model(data), store_path, config) + return cls(_array_metadata_model(data), store_path, config) @property - def document(self) -> ArrayDocumentModel: - """The metadata document this instance is bound to, as a `zarr_metadata` model.""" + def metadata(self) -> ArrayMetadataModel: + """The metadata this instance is bound to, as a `zarr_metadata` model.""" source = self._source if isinstance(source, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): return source - if self._document_cache is None: + if self._metadata_model_cache is None: buffers = source.to_buffer_dict(default_buffer_prototype()) key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} if isinstance(source, ArrayV2Metadata): - self._document_cache = ZarrV2ArrayMetadata.from_key_value(key_value) + self._metadata_model_cache = ZarrV2ArrayMetadata.from_key_value(key_value) else: - self._document_cache = ZarrV3ArrayMetadata.from_key_value(key_value) - return self._document_cache + self._metadata_model_cache = ZarrV3ArrayMetadata.from_key_value(key_value) + return self._metadata_model_cache @property - def metadata(self) -> ArrayMetadata: - """The resolved metadata: the document with its codec and data type names - resolved into classes, derived from `document` on first use.""" + def metadata_legacy(self) -> ArrayMetadata: + """The metadata resolved into the legacy classes (`ArrayV2Metadata` / + `ArrayV3Metadata`), with codec and data type names resolved, derived + from `metadata` on first use. The data path runs on this form; these + classes are trending toward deprecation.""" source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): return source if self._resolved_cache is None: - data = _document_json(source) + data = _metadata_model_json(source) if isinstance(source, ZarrV2ArrayMetadata): data["attributes"] = dict(source.attributes) self._resolved_cache = parse_array_metadata(data, str(self._store_path)) @@ -6495,7 +6497,7 @@ def _rebind_resolved( The codec pipeline is kept: rebinding is how `resize` swaps the shape, which leaves the chunk encoding as it was.""" self._source = resolved - self._document_cache = None + self._metadata_model_cache = None self._resolved_cache = None self._chunk_grid_cache = chunk_grid self._chunk_key_encoding_cache = None @@ -6510,8 +6512,8 @@ def _chunk_grid(self) -> ChunkGrid: if resolved is not None: self._chunk_grid_cache = ChunkGrid.from_metadata(resolved) return self._chunk_grid_cache - document = self.document - doc = _document_json(document) + document = self.metadata + doc = _metadata_model_json(document) if isinstance(document, ZarrV2ArrayMetadata): self._chunk_grid_cache = ChunkGrid.from_sizes(document.shape, tuple(doc["chunks"])) else: @@ -6552,8 +6554,8 @@ def _encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: def _chunk_key_encoding(self) -> ChunkKeyEncoding: """The chunk key encoding, derived from the document alone.""" if self._chunk_key_encoding_cache is None: - document = self.document - doc = _document_json(document) + document = self.metadata + doc = _metadata_model_json(document) if isinstance(document, ZarrV2ArrayMetadata): self._chunk_key_encoding_cache = V2ChunkKeyEncoding( separator=doc.get("dimension_separator", ".") @@ -6575,7 +6577,7 @@ def _shard_grid_shape(self) -> tuple[int, ...]: def _resolved_codec_pipeline(self) -> CodecPipeline: if self._codec_pipeline_cache is None: self._codec_pipeline_cache = create_codec_pipeline( - metadata=self.metadata, store=self._store_path.store + metadata=self.metadata_legacy, store=self._store_path.store ) return self._codec_pipeline_cache @@ -6583,6 +6585,30 @@ def to_array(self) -> AnyAsyncArray: """Return an `AsyncArray` over this binding; both share one state.""" return AsyncArray._from_io(self) + def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: + """A new binding with `metadata` in place of this one's, at the same + location, with the same configuration. Nothing is stored; the data + path derives anew from the new metadata.""" + return type(self)(metadata, self._store_path, self._config) + + def with_store(self, store: Store | StorePath) -> Self: + """A new binding for the same metadata in `store`: at this binding's + path within a `Store`, or at a `StorePath` as given. Nothing is + stored; the data path derives anew against the new store.""" + store_path = ( + store if isinstance(store, StorePath) else StorePath(store, self._store_path.path) + ) + return type(self)(self._source, store_path, self._config) + + def with_config(self, config: ArrayConfigLike) -> Self: + """A new binding with `config` merged over this one's runtime + configuration, as `AsyncArray.with_config` merges it.""" + if isinstance(config, ArrayConfig): + new_config = config + else: + new_config = ArrayConfig(**{**self._config.to_dict(), **config}) # type: ignore[arg-type] + return type(self)(self._source, self._store_path, new_config) + async def store_metadata(self, *, ensure_parents: bool = False) -> None: """Store this array's metadata document(s). @@ -6622,8 +6648,8 @@ async def verify(self) -> None: metadata with `build` can assert, before writing chunks, that it agrees with what a coordinator stored. """ - bound = self.document - stored = _array_document_model(await get_array_metadata(self._store_path, zarr_format=None)) + bound = self.metadata + stored = _array_metadata_model(await get_array_metadata(self._store_path, zarr_format=None)) if stored == bound: return if type(stored) is not type(bound): @@ -6653,8 +6679,8 @@ async def erase_metadata(self) -> None: Succeeds if the documents do not exist, as `erase_metadata` does in `zarrs`. Stored chunks are kept; `erase_chunks` removes those. """ - for document in ARRAY_DOCUMENTS[self._zarr_format]: - await (self._store_path / document).delete() + for key in ARRAY_DOCUMENTS[self._zarr_format]: + await (self._store_path / key).delete() resolved = self._resolved_if_any if resolved is not None: _stored_document_replaced(resolved) @@ -6704,7 +6730,7 @@ async def _get_selection( ) -> NDArrayLikeOrScalar: return await _get_selection( self._store_path, - self.metadata, + self.metadata_legacy, self._resolved_codec_pipeline, self._config, self._chunk_grid, @@ -6722,7 +6748,7 @@ async def _set_selection( prototype: BufferPrototype, fields: Fields | None = None, ) -> None: - metadata = self.metadata + metadata = self.metadata_legacy if product(indexer.shape) > 0: # Chunks are about to be stored under the repaired metadata, so store it # first: every reader of the store then agrees with them. @@ -6748,7 +6774,9 @@ async def getitem( """Read the selected region, reading absent chunks as the fill value.""" if prototype is None: prototype = default_buffer_prototype() - indexer = BasicIndexer(selection, shape=self.metadata.shape, chunk_grid=self._chunk_grid) + indexer = BasicIndexer( + selection, shape=self.metadata_legacy.shape, chunk_grid=self._chunk_grid + ) return await self._get_selection(indexer, prototype=prototype) async def setitem( @@ -6760,7 +6788,9 @@ async def setitem( """Encode `value` and store the chunks covering `selection`.""" if prototype is None: prototype = default_buffer_prototype() - indexer = BasicIndexer(selection, shape=self.metadata.shape, chunk_grid=self._chunk_grid) + indexer = BasicIndexer( + selection, shape=self.metadata_legacy.shape, chunk_grid=self._chunk_grid + ) return await self._set_selection(indexer, value, prototype=prototype) def __repr__(self) -> str: @@ -6776,7 +6806,7 @@ def __eq__(self, other: object) -> bool: resolved, other_resolved = self._resolved_if_any, other._resolved_if_any if resolved is not None and other_resolved is not None: return resolved == other_resolved - return self.document == other.document + return self.metadata == other.metadata def __hash__(self) -> int: # A subset of the equality fields, available without deriving either @@ -6791,11 +6821,11 @@ class ArrayIO: def __init__( self, - document: ArrayDocumentModel | ArrayMetadata | ArrayMetadataDict, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: - self._async_io = AsyncArrayIO(document, store_path, config) + self._async_io = AsyncArrayIO(metadata, store_path, config) @classmethod def _wrap(cls, async_io: AsyncArrayIO) -> Self: @@ -6853,14 +6883,14 @@ def open( ) @property - def document(self) -> ArrayDocumentModel: - """The metadata document this instance is bound to, as a `zarr_metadata` model.""" - return self._async_io.document + def metadata(self) -> ArrayMetadataModel: + """The metadata this instance is bound to, as a `zarr_metadata` model.""" + return self._async_io.metadata @property - def metadata(self) -> ArrayMetadata: - """The resolved metadata. See `AsyncArrayIO.metadata`.""" - return self._async_io.metadata + def metadata_legacy(self) -> ArrayMetadata: + """The metadata resolved into the legacy classes. See `AsyncArrayIO.metadata_legacy`.""" + return self._async_io.metadata_legacy @property def store_path(self) -> StorePath: @@ -6876,6 +6906,21 @@ def to_array(self) -> AnyArray: """Return an `Array` over this binding; both share one state.""" return Array(self._async_io.to_array()) + def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: + """A new binding with `metadata` in place of this one's. See + `AsyncArrayIO.with_metadata`.""" + return self._wrap(self._async_io.with_metadata(metadata)) + + def with_store(self, store: Store | StorePath) -> Self: + """A new binding for the same metadata in `store`. See + `AsyncArrayIO.with_store`.""" + return self._wrap(self._async_io.with_store(store)) + + def with_config(self, config: ArrayConfigLike) -> Self: + """A new binding with `config` merged over this one's. See + `AsyncArrayIO.with_config`.""" + return self._wrap(self._async_io.with_config(config)) + def store_metadata(self, *, ensure_parents: bool = False) -> None: """Store this array's metadata document(s). See `AsyncArrayIO.store_metadata`.""" sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 58271d2788..7d1f16463e 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -79,17 +79,17 @@ async def test_chunks_before_metadata(zarr_format: ZarrFormat) -> None: await aio.store_metadata() # the default zarr_format=None auto-detects the format, like zarrs opened = await AsyncArrayIO.open(store, name="foo") - # the binding is a document; the stored and declared documents agree - assert opened.document == aio.document - assert opened.document.shape == data.shape - assert opened.metadata.zarr_format == zarr_format - assert opened.metadata.attributes == {"answer": 42} + # the binding is metadata; the stored and declared metadata agree + assert opened.metadata == aio.metadata + assert opened.metadata.shape == data.shape + assert opened.metadata_legacy.zarr_format == zarr_format + assert dict(opened.metadata.attributes) == {"answer": 42} np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) arr = await open_array(store=store, path="foo", mode="r", zarr_format=zarr_format) np.testing.assert_array_equal(await arr.getitem(Ellipsis), data) # every AsyncArray is built on an AsyncArrayIO, reachable as .io - assert arr.io.metadata is arr.metadata + assert arr.io.metadata_legacy is arr.metadata np.testing.assert_array_equal(await arr.io.getitem(Ellipsis), data) # a writer that declared the same array independently can verify its assumption @@ -158,10 +158,10 @@ async def test_document_work_without_codec_resolution() -> None: await _store_foreign_codec_array(store) aio = await AsyncArrayIO.open(store, name="foo") - document = aio.document - assert isinstance(document, ZarrV3ArrayMetadata) - assert document.shape == (4,) - assert document.must_understand_fields == {} + model = aio.metadata + assert isinstance(model, ZarrV3ArrayMetadata) + assert model.shape == (4,) + assert model.must_understand_fields == {} await aio.verify() with pytest.raises(UnknownCodecError, match="zarr-python-does-not-know-this-codec"): @@ -205,3 +205,31 @@ async def test_erase_chunk_out_of_grid_raises() -> None: await aio.erase_chunk((2, 0)) with pytest.raises(IndexError, match="outside the chunk grid"): await aio.erase_chunk((0,)) + + +async def test_with_methods_rebind() -> None: + """with_metadata, with_store and with_config return new bindings with one + part replaced and everything else kept.""" + store, other_store = MemoryStore(), MemoryStore() + data = np.arange(4, dtype="uint8") + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), dtype="uint8", chunk_grid=(2,)) + await aio.setitem(Ellipsis, data) + await aio.store_metadata() + + relabeled = aio.with_metadata(aio.metadata.update(attributes={"a": 1})) + assert relabeled.store_path == aio.store_path + assert dict(relabeled.metadata.attributes) == {"a": 1} + await relabeled.store_metadata() + reread = await AsyncArrayIO.open(store, name="foo") + assert reread.metadata == relabeled.metadata + + moved = aio.with_store(other_store) + assert moved.store_path.path == aio.store_path.path + await moved.setitem(Ellipsis, data) + await moved.store_metadata() + copied = await AsyncArrayIO.open(other_store, name="foo") + np.testing.assert_array_equal(await copied.getitem(Ellipsis), data) + + configured = aio.with_config({"order": "F"}) + assert configured.config.order == "F" + assert configured.metadata == aio.metadata From ee8b2d1e44dd9c7e873e0a188d75e6fe16c62be6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 16:01:47 +0100 Subject: [PATCH 12/48] docs(array): finish the metadata wording on ArrayIO Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 38 +++++++++++++++++++------------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 4eb1d7e639..56e8938a5f 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -6420,17 +6420,17 @@ async def open( zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> Self: - """Read the metadata document of an existing array and bind to it. + """Read the metadata of an existing array and bind to it. This is the one classmethod constructor that performs IO: the array's - stored metadata document is read from the store. `zarr_format=None` - (the default) reads Zarr format 3 and falls back to format 2, as - `open` does in `zarrs`; pass 2 or 3 to require one format. The read is - strict and does not resolve extension names: an array with unknown - codecs opens for document work, while a stored document zarr can only - read as repaired raises `zarr_metadata.MetadataValidationError` — open - such an array with [`zarr.api.asynchronous.open_array`][], which reads - it as repaired. + stored metadata is read from the store. `zarr_format=None` (the + default) reads Zarr format 3 and falls back to format 2, as `open` + does in `zarrs`; pass 2 or 3 to require one format. The read is strict + and does not resolve extension names: an array with unknown codecs + opens for metadata work, while stored metadata zarr can only read as + repaired raises `zarr_metadata.MetadataValidationError` — open such an + array with [`zarr.api.asynchronous.open_array`][], which reads it as + repaired. """ store_path = await make_store_path(store, path=name) data = await get_array_metadata(store_path, zarr_format=zarr_format) @@ -6610,7 +6610,7 @@ def with_config(self, config: ArrayConfigLike) -> Self: return type(self)(self._source, self._store_path, new_config) async def store_metadata(self, *, ensure_parents: bool = False) -> None: - """Store this array's metadata document(s). + """Store this array's metadata. The documents are written unconditionally, replacing the metadata of any existing node at this location, as `store_metadata` does in @@ -6637,10 +6637,10 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: async def verify(self) -> None: """Check that the stored metadata declares the array this instance is bound to. - Reads the metadata document stored at this location — favoring Zarr - format 3 when documents of both formats exist — and raises + Reads the metadata stored at this location — favoring Zarr format 3 + when both formats' metadata exists — and raises `ValueError` naming each differing field when it does not equal the - bound document. Document equality is semantic: a fill value is + bound metadata. Metadata equality is semantic: a fill value is compared in its canonical spelling. Raises [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. This is the explicit counterpart of the verification TensorStore performs @@ -6656,7 +6656,7 @@ async def verify(self) -> None: stored_format = 2 if isinstance(stored, ZarrV2ArrayMetadata) else 3 raise ValueError( f"Array {str(self.store_path)!r}: the stored array is Zarr format " - f"{stored_format}; the bound document is Zarr format {self._zarr_format}." + f"{stored_format}; the bound metadata is Zarr format {self._zarr_format}." ) stored_json = dict(stored.to_json()) bound_json = dict(bound.to_json()) @@ -6674,9 +6674,9 @@ async def verify(self) -> None: ) async def erase_metadata(self) -> None: - """Erase this array's metadata document(s), the inverse of `store_metadata`. + """Erase this array's stored metadata, the inverse of `store_metadata`. - Succeeds if the documents do not exist, as `erase_metadata` does in + Succeeds if the stored metadata does not exist, as `erase_metadata` does in `zarrs`. Stored chunks are kept; `erase_chunks` removes those. """ for key in ARRAY_DOCUMENTS[self._zarr_format]: @@ -6874,7 +6874,7 @@ def open( zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> Self: - """Read the metadata document of an existing array and bind to it. + """Read the metadata of an existing array and bind to it. See `AsyncArrayIO.open`. """ @@ -6922,7 +6922,7 @@ def with_config(self, config: ArrayConfigLike) -> Self: return self._wrap(self._async_io.with_config(config)) def store_metadata(self, *, ensure_parents: bool = False) -> None: - """Store this array's metadata document(s). See `AsyncArrayIO.store_metadata`.""" + """Store this array's metadata. See `AsyncArrayIO.store_metadata`.""" sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) def verify(self) -> None: @@ -6933,7 +6933,7 @@ def verify(self) -> None: sync(self._async_io.verify()) def erase_metadata(self) -> None: - """Erase this array's metadata document(s). See `AsyncArrayIO.erase_metadata`.""" + """Erase this array's stored metadata. See `AsyncArrayIO.erase_metadata`.""" sync(self._async_io.erase_metadata()) def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: From 5b0a08a9f563ed009833566d651f80330bf14a56 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 16:30:37 +0100 Subject: [PATCH 13/48] fix(array): make ArrayIO meet the global configuration the way Array does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The codec pipeline implementation is captured when the binding is made — the capture discipline AsyncArray has always had — through a new pipeline_class parameter on create_codec_pipeline, while the pipeline itself still builds on first use, so metadata-only work resolves nothing. build_v2's order defaults from the global array.order instead of a static C. Nothing about ArrayIO requires new configuration concepts; scoping resolution per binding stays in the follow-up that deprecates the global state. Assisted-by: ClaudeCode:claude-fable-5 --- changes/419.feature.md | 1 + src/zarr/core/array.py | 31 ++++++++++++++++++++++++------- tests/test_array_io.py | 22 ++++++++++++++++++++++ 3 files changed, 47 insertions(+), 7 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index 5b5bbb8c22..7e888b6a5e 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -2,3 +2,4 @@ Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arr `Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. `with_metadata`, `with_store` and `with_config` return new bindings with one part replaced. +`ArrayIO` responds to the global configuration the way `Array` always has: the codec pipeline implementation is captured when the binding is made, the runtime configuration defaults from `zarr.config`, and `build_v2`'s `order` defaults from `array.order`. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 56e8938a5f..018505b04c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -253,10 +253,17 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: raise TypeError # pragma: no cover -def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None) -> CodecPipeline: +def create_codec_pipeline( + metadata: ArrayMetadata, + *, + store: Store | None = None, + pipeline_class: type[CodecPipeline] | None = None, +) -> CodecPipeline: + if pipeline_class is None: + pipeline_class = get_pipeline_class() if store is not None: try: - return get_pipeline_class().from_array_metadata_and_store( + return pipeline_class.from_array_metadata_and_store( array_metadata=metadata, store=store ) except NotImplementedError: @@ -269,7 +276,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None # "disables partial reads" warning) fire here; `evolve_from_array_spec` # re-splits the same already-warned-about chain via # `codecs_from_list_unchecked`, so it does not re-emit them. - pipeline = get_pipeline_class().from_codecs(metadata.codecs) + pipeline = pipeline_class.from_codecs(metadata.codecs) # Use the regular chunk shape if available, otherwise use a # placeholder. The ChunkTransform is shape-agnostic — the actual @@ -288,7 +295,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None return pipeline.evolve_from_array_spec(chunk_spec) elif isinstance(metadata, ArrayV2Metadata): v2_codec = V2Codec(filters=metadata.filters, compressor=metadata.compressor) - pipeline = get_pipeline_class().from_codecs([v2_codec]) + pipeline = pipeline_class.from_codecs([v2_codec]) chunk_spec = ArraySpec( shape=metadata.chunks, dtype=metadata.dtype, @@ -4676,7 +4683,7 @@ def build_array_metadata_v2( dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], chunks: ShapeLike, fill_value: Any = DEFAULT_FILL_VALUE, - order: MemoryOrder = "C", + order: MemoryOrder | None = None, dimension_separator: Literal[".", "/"] = ".", filters: Iterable[dict[str, JSON] | Numcodec] | None = None, compressor: CompressorLikev2 = None, @@ -4687,11 +4694,14 @@ def build_array_metadata_v2( Each keyword is a field of the `.zarray` document, with `attributes` the `.zattrs` document. `fill_value` defaults to the data type's default scalar, and an explicit `None` is kept, because the field is nullable in - Zarr format 2. `filters` and `compressor` are numcodecs objects or their + Zarr format 2. `order` defaults to the global configuration's + `array.order`. `filters` and `compressor` are numcodecs objects or their JSON, exactly as the format stores them — nothing is guessed; the guessing creation API is [`zarr.api.asynchronous.create_array`][]. """ zdtype = parse_dtype(dtype, zarr_format=2) + if order is None: + order = zarr_config.get("array.order") if fill_value is DEFAULT_FILL_VALUE: fill_value_parsed: Any = zdtype.default_scalar() else: @@ -6336,6 +6346,7 @@ class AsyncArrayIO: _store_path: StorePath _config: ArrayConfig + _pipeline_class: type[CodecPipeline] _source: ArrayMetadataModel | ArrayMetadata _metadata_model_cache: ArrayMetadataModel | None _resolved_cache: ArrayMetadata | None @@ -6351,6 +6362,10 @@ def __init__( ) -> None: self._store_path = store_path self._config = parse_array_config(config) + # The pipeline implementation is chosen when the binding is made, from + # the global configuration, exactly as `AsyncArray` has always chosen + # it; the pipeline itself is built on first use. + self._pipeline_class = get_pipeline_class() self._metadata_model_cache = None self._resolved_cache = None self._chunk_grid_cache = None @@ -6577,7 +6592,9 @@ def _shard_grid_shape(self) -> tuple[int, ...]: def _resolved_codec_pipeline(self) -> CodecPipeline: if self._codec_pipeline_cache is None: self._codec_pipeline_cache = create_codec_pipeline( - metadata=self.metadata_legacy, store=self._store_path.store + metadata=self.metadata_legacy, + store=self._store_path.store, + pipeline_class=self._pipeline_class, ) return self._codec_pipeline_cache diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 7d1f16463e..3400c320f9 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -8,9 +8,12 @@ import pytest from zarr_metadata import ZarrV3ArrayMetadata +import zarr from zarr import ArrayIO, AsyncArrayIO from zarr.api.asynchronous import open_array from zarr.core.buffer import default_buffer_prototype +from zarr.core.codec_pipeline import FusedCodecPipeline +from zarr.core.metadata import ArrayV2Metadata from zarr.errors import ArrayNotFoundError, UnknownCodecError from zarr.storage import MemoryStore @@ -233,3 +236,22 @@ async def test_with_methods_rebind() -> None: configured = aio.with_config({"order": "F"}) assert configured.config.order == "F" assert configured.metadata == aio.metadata + + +def test_binding_follows_the_global_config() -> None: + """ArrayIO responds to the global configuration the way Array always has: + the pipeline implementation is captured when the binding is made, and + build_v2's order defaults from array.order.""" + with zarr.config.set( + { + "codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline", + "array.order": "F", + } + ): + v3 = ArrayIO.build_v3(MemoryStore(), shape=(2,), dtype="uint8", chunk_grid=(2,)) + v2 = ArrayIO.build_v2(MemoryStore(), shape=(2,), dtype="uint8", chunks=(2,)) + # outside the block, the choices made inside it hold + assert isinstance(v3.to_array().async_array.codec_pipeline, FusedCodecPipeline) + legacy = v2.metadata_legacy + assert isinstance(legacy, ArrayV2Metadata) + assert legacy.order == "F" From 19d2e634f4115dfa795f2fcb868273bbb4a3cde1 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 16:42:26 +0100 Subject: [PATCH 14/48] refactor(array): name the v3 builder's data type field data_type The zarr.json field is data_type, so build_array_metadata_v3 and the build_v3 declaration take data_type; the v2 builder keeps dtype, which is the .zarray field's own name. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 10 +++++----- tests/test_array_io.py | 22 ++++++++++++---------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 018505b04c..13afebd58d 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -629,7 +629,7 @@ def _create_metadata_v3( """See the module-level `build_array_metadata_v3`.""" return build_array_metadata_v3( shape=shape, - dtype=dtype, + data_type=dtype, chunk_grid=chunk_grid, fill_value=fill_value, chunk_key_encoding=chunk_key_encoding, @@ -4663,7 +4663,7 @@ async def init_array( _warn_order_kwarg() meta = build_array_metadata_v3( shape=shape_parsed, - dtype=zdtype, + data_type=zdtype, fill_value=fill_value, chunk_grid=grid, chunk_key_encoding=chunk_key_encoding_parsed, @@ -4724,7 +4724,7 @@ def build_array_metadata_v2( def build_array_metadata_v3( *, shape: ShapeLike, - dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + data_type: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], chunk_grid: ChunkGridMetadata | Mapping[str, JSON] | ShapeLike, fill_value: Any = DEFAULT_FILL_VALUE, chunk_key_encoding: ChunkKeyEncodingLike | None = None, @@ -4743,7 +4743,7 @@ def build_array_metadata_v3( data type's default scalar, and `None` also means that default, because the field is required in Zarr format 3. """ - zdtype = parse_dtype(dtype, zarr_format=3) + zdtype = parse_dtype(data_type, zarr_format=3) grid: ChunkGridMetadata | dict[str, JSON] if isinstance(chunk_grid, RegularChunkGridMetadata | RectilinearChunkGridMetadata): grid = chunk_grid @@ -6285,7 +6285,7 @@ class ArrayV3DeclarationKwargs(TypedDict, total=False): keyword arguments.""" shape: Required[ShapeLike] - dtype: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] + data_type: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] chunk_grid: Required[ChunkGridMetadata | Mapping[str, JSON] | ShapeLike] fill_value: Any chunk_key_encoding: ChunkKeyEncodingLike | None diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 3400c320f9..46c60205e3 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -44,7 +44,7 @@ def _declare_answer_array( store, name="foo", shape=data.shape, - dtype=data.dtype, + data_type=data.dtype, chunk_grid=(2, 2), attributes={"answer": 42}, ) @@ -122,7 +122,9 @@ def test_chunks_before_metadata_sync() -> None: """The synchronous `ArrayIO` supports the same chunks-first flow.""" store = MemoryStore() data = np.arange(16, dtype="uint16").reshape(4, 4) - aio = ArrayIO.build_v3(store, name="foo", shape=data.shape, dtype=data.dtype, chunk_grid=(2, 2)) + aio = ArrayIO.build_v3( + store, name="foo", shape=data.shape, data_type=data.dtype, chunk_grid=(2, 2) + ) aio.setitem(Ellipsis, data) aio.store_metadata() opened = ArrayIO.open(store, name="foo") @@ -143,7 +145,7 @@ def test_build_v3_unknown_chunk_grid_raises() -> None: MemoryStore(), name="foo", shape=(4,), - dtype="uint8", + data_type="uint8", chunk_grid={"name": "diagonal", "configuration": {}}, ) @@ -179,10 +181,10 @@ async def test_document_work_without_codec_resolution() -> None: async def test_verify_mismatch_raises() -> None: store = MemoryStore() - aio = AsyncArrayIO.build_v3(store, shape=(4, 4), dtype="uint8", chunk_grid=(2, 2)) + aio = AsyncArrayIO.build_v3(store, shape=(4, 4), data_type="uint8", chunk_grid=(2, 2)) await aio.store_metadata() other = AsyncArrayIO.build_v3( - store, shape=(4, 4), dtype="uint8", chunk_grid=(4, 4), attributes={"a": 1} + store, shape=(4, 4), data_type="uint8", chunk_grid=(4, 4), attributes={"a": 1} ) with pytest.raises(ValueError, match="differing fields: attributes, chunk_grid"): await other.verify() @@ -191,19 +193,19 @@ async def test_verify_mismatch_raises() -> None: async def test_verify_format_mismatch_raises() -> None: store = MemoryStore() await AsyncArrayIO.build_v2(store, shape=(4,), dtype="uint8", chunks=(4,)).store_metadata() - bound_v3 = AsyncArrayIO.build_v3(store, shape=(4,), dtype="uint8", chunk_grid=(4,)) + bound_v3 = AsyncArrayIO.build_v3(store, shape=(4,), data_type="uint8", chunk_grid=(4,)) with pytest.raises(ValueError, match="stored array is Zarr format 2"): await bound_v3.verify() async def test_verify_missing_raises() -> None: - aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4,), dtype="uint8", chunk_grid=(4,)) + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,)) with pytest.raises(ArrayNotFoundError): await aio.verify() async def test_erase_chunk_out_of_grid_raises() -> None: - aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4, 4), dtype="uint8", chunk_grid=(2, 2)) + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4, 4), data_type="uint8", chunk_grid=(2, 2)) with pytest.raises(IndexError, match="outside the chunk grid"): await aio.erase_chunk((2, 0)) with pytest.raises(IndexError, match="outside the chunk grid"): @@ -215,7 +217,7 @@ async def test_with_methods_rebind() -> None: part replaced and everything else kept.""" store, other_store = MemoryStore(), MemoryStore() data = np.arange(4, dtype="uint8") - aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), dtype="uint8", chunk_grid=(2,)) + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) await aio.setitem(Ellipsis, data) await aio.store_metadata() @@ -248,7 +250,7 @@ def test_binding_follows_the_global_config() -> None: "array.order": "F", } ): - v3 = ArrayIO.build_v3(MemoryStore(), shape=(2,), dtype="uint8", chunk_grid=(2,)) + v3 = ArrayIO.build_v3(MemoryStore(), shape=(2,), data_type="uint8", chunk_grid=(2,)) v2 = ArrayIO.build_v2(MemoryStore(), shape=(2,), dtype="uint8", chunks=(2,)) # outside the block, the choices made inside it hold assert isinstance(v3.to_array().async_array.codec_pipeline, FusedCodecPipeline) From 87c1d632b7d851ba59a8d9b5cd242a695becbbb1 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 17:01:22 +0100 Subject: [PATCH 15/48] docs(array): shorten the ArrayIO docstrings to what each function does Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 295 +++++++++++++---------------------------- 1 file changed, 90 insertions(+), 205 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 13afebd58d..cf36a46f0b 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4689,15 +4689,13 @@ def build_array_metadata_v2( compressor: CompressorLikev2 = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV2Metadata: - """Build a Zarr format 2 array metadata document from its fields. + """Build Zarr format 2 array metadata from its fields. - Each keyword is a field of the `.zarray` document, with `attributes` the - `.zattrs` document. `fill_value` defaults to the data type's default - scalar, and an explicit `None` is kept, because the field is nullable in - Zarr format 2. `order` defaults to the global configuration's + Each keyword is a `.zarray` field; `attributes` is the `.zattrs` one. + `fill_value` defaults to the data type's default scalar (an explicit + `None` is kept: the field is nullable); `order` defaults from the global `array.order`. `filters` and `compressor` are numcodecs objects or their - JSON, exactly as the format stores them — nothing is guessed; the - guessing creation API is [`zarr.api.asynchronous.create_array`][]. + JSON, as stored. Nothing is guessed. """ zdtype = parse_dtype(dtype, zarr_format=2) if order is None: @@ -4732,16 +4730,13 @@ def build_array_metadata_v3( dimension_names: DimensionNamesLike = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV3Metadata: - """Build a Zarr format 3 array metadata document from its fields. - - Each keyword is a field of the `zarr.json` document. `chunk_grid` takes - the grid's metadata or its JSON, or a chunk shape, which declares a - regular grid, as `ArrayBuilder` does in `zarrs`. `codecs` of `None` - declares the default chain for the data type (no filters, the default - serializer, the default compressors), and `chunk_key_encoding` of `None` - the default encoding with the `/` separator. `fill_value` defaults to the - data type's default scalar, and `None` also means that default, because - the field is required in Zarr format 3. + """Build Zarr format 3 array metadata from its fields. + + Each keyword is a `zarr.json` field. `chunk_grid` takes grid metadata, + its JSON, or a chunk shape, which declares a regular grid. `codecs` + `None` is the data type's default chain; `chunk_key_encoding` `None` is + the default encoding with `/`; `fill_value` defaults to the data type's + default scalar (`None` means the default too: the field is required). """ zdtype = parse_dtype(data_type, zarr_format=3) grid: ChunkGridMetadata | dict[str, JSON] @@ -6235,26 +6230,20 @@ async def _delete_key(store_path: StorePath, key: str) -> None: type ArrayMetadataModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata -"""Array metadata as a `zarr_metadata` model: the JSON as written, validated, -with extension names read but never resolved into classes.""" +"""Array metadata as a `zarr_metadata` model: validated JSON, extension names unresolved.""" def _metadata_model_json(document: ArrayMetadataModel) -> dict[str, Any]: - """The metadata as plain, mutable JSON. The model returns read-only forms — - arrays as tuples — while zarr's parsers take the stored forms, so the - metadata goes through JSON text.""" + """The metadata as plain, mutable JSON, via JSON text: the model returns read-only forms.""" return cast("dict[str, Any]", json.loads(json.dumps(document.to_json()))) def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: - """Array metadata as a `zarr_metadata` model, from the form `get_array_metadata` - returns: the Zarr format 2 form carries its attributes merged in. - - The read is strict: a document with problems — including one zarr can only - read as repaired — raises `zarr_metadata.MetadataValidationError` naming - each problem. A document whose extension names are unknown is not a - problem: the model reports them as unclaimed fields, and resolution - fails only when the data path is built. + """Model the metadata dict `get_array_metadata` returned. + + The read is strict: metadata zarr can only read as repaired raises + `zarr_metadata.MetadataValidationError`. Unknown extension names read as + unclaimed fields, not errors. """ if data.get("zarr_format") == 2: return ZarrV2ArrayMetadata(data) @@ -6262,10 +6251,7 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: class ArrayV2DeclarationKwargs(TypedDict, total=False): - """The fields that declare a Zarr format 2 array: what - `build_array_metadata_v2` takes. This is the one declaration of the - surface; `ArrayIO.build_v2` and `AsyncArrayIO.build_v2` take it as their - keyword arguments.""" + """The fields that declare a Zarr format 2 array: what `build_array_metadata_v2` and the `build_v2` methods take as keywords.""" shape: Required[ShapeLike] dtype: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] @@ -6279,10 +6265,7 @@ class ArrayV2DeclarationKwargs(TypedDict, total=False): class ArrayV3DeclarationKwargs(TypedDict, total=False): - """The fields that declare a Zarr format 3 array: what - `build_array_metadata_v3` takes. This is the one declaration of the - surface; `ArrayIO.build_v3` and `AsyncArrayIO.build_v3` take it as their - keyword arguments.""" + """The fields that declare a Zarr format 3 array: what `build_array_metadata_v3` and the `build_v3` methods take as keywords.""" shape: Required[ShapeLike] data_type: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] @@ -6295,8 +6278,7 @@ class ArrayV3DeclarationKwargs(TypedDict, total=False): def _bound_store_path(store: Store | StorePath, name: str | None) -> StorePath: - """The location `name` within `store`, bound without IO: a path or URL - string is not accepted, because coercing one opens a store.""" + """The location `name` within `store`, bound without IO.""" store_path = store if isinstance(store, StorePath) else StorePath(store) if name is not None: store_path = store_path / name @@ -6306,42 +6288,16 @@ def _bound_store_path(store: Store | StorePath, name: str | None) -> StorePath: class AsyncArrayIO: """Explicit asynchronous IO for one Zarr array. - An `AsyncArrayIO` binds array metadata to a location in a store. The - `metadata` is a `zarr_metadata` model — the JSON as written, with - extension names read but not resolved into codec or data type classes — - so every metadata operation (`store_metadata`, `erase_metadata`, - `erase_chunk`, `erase_chunks`, `verify`) works without resolving - anything: an array whose codecs this environment lacks can still be - opened, verified and erased. Reading and writing chunks (`getitem` / - `setitem`) resolves the metadata into the legacy classes on first use; - `metadata_legacy` is that resolved form, derived once — or, when the - instance was built from one of those objects (as every `AsyncArray` is), - carried as the source `metadata` derives from. - - Constructing one — directly, or with `build_v2` / `build_v3` — performs - no IO: nothing is read from or written to the store until an IO method is - called. `store_metadata` stores the metadata, and `setitem` / `getitem` - store and read chunks, in whatever order the caller chooses. Every write - has an erasing inverse: `store_metadata` ↔ `erase_metadata`, and chunk - writes ↔ `erase_chunk` / `erase_chunks`. The `with_metadata`, `with_store` - and `with_config` methods return new bindings with one part replaced. - `AsyncArray` is an array-like layer over this one: it wraps an - `AsyncArrayIO` and delegates its metadata and selection IO to it. - - `store_metadata` is the publication point: until it is called, no Zarr - reader resolves the array, so a failed chunks-first run leaves no readable - array. It does leave chunk keys in the store, which occupy space and would - be misread by a different array stored later at the same location; - `erase_chunks` removes them. - - Warning: the bound metadata is assumed, never verified against the store. - Reads and partial writes decode stored chunks under it, so binding - metadata that disagrees with how existing chunks at this location were - encoded misreads or corrupts data — the hazard TensorStore documents for - `assume_metadata`. One exception to "no implicit metadata IO": an array - opened from stored metadata that needed a repair stores that repair - before its first chunk write. Call `store_metadata` first to make that - write explicit. + Binds array metadata (`metadata`, a `zarr_metadata` model) to a location + in a store. Construction performs no IO; every read or write of the store + is an explicit method call, and every write has an erasing inverse. The + metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, + `erase_chunks`, `verify`) never resolve extension names; `getitem` and + `setitem` resolve them into the legacy classes (`metadata_legacy`) on + first use. The bound metadata is assumed, never checked against the + store, and an array opened from metadata that needed a repair stores that + repair before its first chunk write. `AsyncArray` wraps an `AsyncArrayIO` + and delegates its metadata and selection IO to it. """ _store_path: StorePath @@ -6387,18 +6343,11 @@ def build_v2( config: ArrayConfigLike | None = None, **kwargs: Unpack[ArrayV2DeclarationKwargs], ) -> Self: - """Declare a Zarr format 2 array bound to a location in a store, without - storing anything. - - Binding is pure, which is why the declaring constructors are the - synchronous methods on this class: `store` must already be a `Store` - or `StorePath`, because coercing a path or URL string opens a store, - and that is [`zarr.api.asynchronous.create_array`][]'s job — as is - guessing: every keyword here is a field of the `.zarray` document (see - `build_array_metadata_v2`). Unlike `zarrs`, `fill_value` defaults to - the data type's default fill value instead of being required. The - returned instance holds the metadata in memory; `store_metadata` - stores it. + """Declare a Zarr format 2 array bound to `store`, storing nothing. + + The keywords are the `.zarray` fields (see `build_array_metadata_v2`). + `store` must be a `Store` or `StorePath`: coercing strings opens + stores. `store_metadata` stores the result. """ return cls(build_array_metadata_v2(**kwargs), _bound_store_path(store, name), config) @@ -6411,18 +6360,11 @@ def build_v3( config: ArrayConfigLike | None = None, **kwargs: Unpack[ArrayV3DeclarationKwargs], ) -> Self: - """Declare a Zarr format 3 array bound to a location in a store, without - storing anything. - - Binding is pure, which is why the declaring constructors are the - synchronous methods on this class: `store` must already be a `Store` - or `StorePath`, because coercing a path or URL string opens a store, - and that is [`zarr.api.asynchronous.create_array`][]'s job — as is - guessing: every keyword here is a field of the `zarr.json` document - (see `build_array_metadata_v3`). Unlike `zarrs`, `fill_value` defaults - to the data type's default fill value instead of being required. The - returned instance holds the metadata in memory; `store_metadata` - stores it. + """Declare a Zarr format 3 array bound to `store`, storing nothing. + + The keywords are the `zarr.json` fields (see + `build_array_metadata_v3`). `store` must be a `Store` or `StorePath`: + coercing strings opens stores. `store_metadata` stores the result. """ return cls(build_array_metadata_v3(**kwargs), _bound_store_path(store, name), config) @@ -6435,17 +6377,12 @@ async def open( zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> Self: - """Read the metadata of an existing array and bind to it. - - This is the one classmethod constructor that performs IO: the array's - stored metadata is read from the store. `zarr_format=None` (the - default) reads Zarr format 3 and falls back to format 2, as `open` - does in `zarrs`; pass 2 or 3 to require one format. The read is strict - and does not resolve extension names: an array with unknown codecs - opens for metadata work, while stored metadata zarr can only read as - repaired raises `zarr_metadata.MetadataValidationError` — open such an - array with [`zarr.api.asynchronous.open_array`][], which reads it as - repaired. + """Read the stored metadata of an existing array and bind to it. + + `zarr_format=None` (the default) tries format 3, then format 2. The + read is strict and resolves nothing: metadata zarr can only read as + repaired raises `zarr_metadata.MetadataValidationError`; use + [`zarr.api.asynchronous.open_array`][] for the repairing read. """ store_path = await make_store_path(store, path=name) data = await get_array_metadata(store_path, zarr_format=zarr_format) @@ -6453,7 +6390,7 @@ async def open( @property def metadata(self) -> ArrayMetadataModel: - """The metadata this instance is bound to, as a `zarr_metadata` model.""" + """The bound metadata, as a `zarr_metadata` model.""" source = self._source if isinstance(source, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): return source @@ -6468,10 +6405,7 @@ def metadata(self) -> ArrayMetadataModel: @property def metadata_legacy(self) -> ArrayMetadata: - """The metadata resolved into the legacy classes (`ArrayV2Metadata` / - `ArrayV3Metadata`), with codec and data type names resolved, derived - from `metadata` on first use. The data path runs on this form; these - classes are trending toward deprecation.""" + """The metadata resolved into the legacy `ArrayV2Metadata` / `ArrayV3Metadata` classes, derived on first use.""" source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): return source @@ -6484,12 +6418,12 @@ def metadata_legacy(self) -> ArrayMetadata: @property def store_path(self) -> StorePath: - """The location in a store this instance is bound to.""" + """The bound location in a store.""" return self._store_path @property def config(self) -> ArrayConfig: - """The runtime configuration of this instance.""" + """The runtime configuration.""" return self._config @property @@ -6508,9 +6442,7 @@ def _resolved_if_any(self) -> ArrayMetadata | None: def _rebind_resolved( self, resolved: ArrayMetadata, chunk_grid: ChunkGrid | None = None ) -> None: - """Make `resolved` this binding's state; the document re-derives from it. - The codec pipeline is kept: rebinding is how `resize` swaps the shape, - which leaves the chunk encoding as it was.""" + """Make `resolved` this binding's state. The codec pipeline is kept: rebinding is how `resize` swaps the shape, which leaves the chunk encoding as it was.""" self._source = resolved self._metadata_model_cache = None self._resolved_cache = None @@ -6519,9 +6451,7 @@ def _rebind_resolved( @property def _chunk_grid(self) -> ChunkGrid: - """The stored chunk grid: from the resolved metadata when the instance - holds one, and otherwise from the document alone, so the storage-level - operations of a document-bound instance never resolve codecs.""" + """The stored chunk grid: from the resolved metadata when the instance holds one, otherwise from the metadata model alone.""" if self._chunk_grid_cache is None: resolved = self._resolved_if_any if resolved is not None: @@ -6558,8 +6488,7 @@ def _chunk_grid(self) -> ChunkGrid: return self._chunk_grid_cache def _encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: - """The store key of a stored chunk: from the resolved metadata when the - instance holds one, and otherwise from the document alone.""" + """The store key of a stored chunk.""" resolved = self._resolved_if_any if resolved is not None: return resolved.encode_chunk_key(chunk_coords) @@ -6567,7 +6496,7 @@ def _encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: @property def _chunk_key_encoding(self) -> ChunkKeyEncoding: - """The chunk key encoding, derived from the document alone.""" + """The chunk key encoding, from the metadata model alone.""" if self._chunk_key_encoding_cache is None: document = self.metadata doc = _metadata_model_json(document) @@ -6599,27 +6528,22 @@ def _resolved_codec_pipeline(self) -> CodecPipeline: return self._codec_pipeline_cache def to_array(self) -> AnyAsyncArray: - """Return an `AsyncArray` over this binding; both share one state.""" + """An `AsyncArray` over this binding; both share one state.""" return AsyncArray._from_io(self) def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: - """A new binding with `metadata` in place of this one's, at the same - location, with the same configuration. Nothing is stored; the data - path derives anew from the new metadata.""" + """This binding with `metadata` in place of its own. Nothing is stored.""" return type(self)(metadata, self._store_path, self._config) def with_store(self, store: Store | StorePath) -> Self: - """A new binding for the same metadata in `store`: at this binding's - path within a `Store`, or at a `StorePath` as given. Nothing is - stored; the data path derives anew against the new store.""" + """This binding at its path within `store`, or at a `StorePath` as given. Nothing is stored.""" store_path = ( store if isinstance(store, StorePath) else StorePath(store, self._store_path.path) ) return type(self)(self._source, store_path, self._config) def with_config(self, config: ArrayConfigLike) -> Self: - """A new binding with `config` merged over this one's runtime - configuration, as `AsyncArray.with_config` merges it.""" + """This binding with `config` merged over its runtime configuration, as `AsyncArray.with_config` merges it.""" if isinstance(config, ArrayConfig): new_config = config else: @@ -6627,15 +6551,11 @@ def with_config(self, config: ArrayConfigLike) -> Self: return type(self)(self._source, self._store_path, new_config) async def store_metadata(self, *, ensure_parents: bool = False) -> None: - """Store this array's metadata. + """Store the metadata, unconditionally replacing any stored metadata here. - The documents are written unconditionally, replacing the metadata of - any existing node at this location, as `store_metadata` does in - `zarrs`. Guarded creation, which raises when a node already exists, - is [`zarr.api.asynchronous.create_array`][]'s job. Chunks of any - previous array at this location are kept and would be read under the - new metadata; call `erase_chunks` first when repurposing a location. - With `ensure_parents`, missing parent groups are stored first. + Chunks of a previous array at this location are kept; `erase_chunks` + removes them. With `ensure_parents`, missing parent groups are stored + first. """ source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): @@ -6652,18 +6572,11 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: ) async def verify(self) -> None: - """Check that the stored metadata declares the array this instance is bound to. - - Reads the metadata stored at this location — favoring Zarr format 3 - when both formats' metadata exists — and raises - `ValueError` naming each differing field when it does not equal the - bound metadata. Metadata equality is semantic: a fill value is - compared in its canonical spelling. Raises - [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. This is - the explicit counterpart of the verification TensorStore performs - against spec constraints on open: a writer that declared its own - metadata with `build` can assert, before writing chunks, that it - agrees with what a coordinator stored. + """Check that the stored metadata equals the bound metadata. + + Favors format 3 when both formats exist, compares canonically, and + raises `ValueError` naming each differing field, or + [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. """ bound = self.metadata stored = _array_metadata_model(await get_array_metadata(self._store_path, zarr_format=None)) @@ -6691,11 +6604,7 @@ async def verify(self) -> None: ) async def erase_metadata(self) -> None: - """Erase this array's stored metadata, the inverse of `store_metadata`. - - Succeeds if the stored metadata does not exist, as `erase_metadata` does in - `zarrs`. Stored chunks are kept; `erase_chunks` removes those. - """ + """Delete the stored metadata; chunks are kept. Succeeds if absent.""" for key in ARRAY_DOCUMENTS[self._zarr_format]: await (self._store_path / key).delete() resolved = self._resolved_if_any @@ -6703,11 +6612,7 @@ async def erase_metadata(self) -> None: _stored_document_replaced(resolved) async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: - """Erase one stored chunk by its chunk grid coordinates. - - When sharding is used, the stored chunk is the shard. Succeeds if the - chunk does not exist, as `erase_chunk` does in `zarrs`. - """ + """Delete one stored chunk by grid coordinates (the shard, under sharding). Succeeds if absent.""" grid_shape = self._shard_grid_shape if len(chunk_coords) != len(grid_shape) or any( coord < 0 or coord >= size for coord, size in zip(chunk_coords, grid_shape, strict=True) @@ -6719,10 +6624,7 @@ async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: await _delete_key(self._store_path, self._encode_chunk_key(chunk_coords)) async def erase_chunks(self) -> None: - """Erase every stored chunk of this array, the inverse of its chunk writes. - - Chunks that do not exist are skipped; the metadata documents are kept. - """ + """Delete every stored chunk; the metadata is kept.""" grid_shape = self._shard_grid_shape n_chunks = product(grid_shape) if n_chunks == 0: @@ -6788,7 +6690,7 @@ async def getitem( *, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - """Read the selected region, reading absent chunks as the fill value.""" + """Read the selection, reading absent chunks as the fill value.""" if prototype is None: prototype = default_buffer_prototype() indexer = BasicIndexer( @@ -6859,11 +6761,7 @@ def build_v2( config: ArrayConfigLike | None = None, **kwargs: Unpack[ArrayV2DeclarationKwargs], ) -> Self: - """Declare a Zarr format 2 array bound to a location in a store, without - storing anything. - - See `AsyncArrayIO.build_v2`. - """ + """Declare a Zarr format 2 array bound to `store`, storing nothing. See `AsyncArrayIO.build_v2`.""" return cls._wrap(AsyncArrayIO.build_v2(store, name=name, config=config, **kwargs)) @classmethod @@ -6875,11 +6773,7 @@ def build_v3( config: ArrayConfigLike | None = None, **kwargs: Unpack[ArrayV3DeclarationKwargs], ) -> Self: - """Declare a Zarr format 3 array bound to a location in a store, without - storing anything. - - See `AsyncArrayIO.build_v3`. - """ + """Declare a Zarr format 3 array bound to `store`, storing nothing. See `AsyncArrayIO.build_v3`.""" return cls._wrap(AsyncArrayIO.build_v3(store, name=name, config=config, **kwargs)) @classmethod @@ -6891,17 +6785,14 @@ def open( zarr_format: ZarrFormat | None = None, config: ArrayConfigLike | None = None, ) -> Self: - """Read the metadata of an existing array and bind to it. - - See `AsyncArrayIO.open`. - """ + """Read the stored metadata of an existing array and bind to it. See `AsyncArrayIO.open`.""" return cls._wrap( sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config)) ) @property def metadata(self) -> ArrayMetadataModel: - """The metadata this instance is bound to, as a `zarr_metadata` model.""" + """The bound metadata, as a `zarr_metadata` model.""" return self._async_io.metadata @property @@ -6911,54 +6802,48 @@ def metadata_legacy(self) -> ArrayMetadata: @property def store_path(self) -> StorePath: - """The location in a store this instance is bound to.""" + """The bound location in a store.""" return self._async_io.store_path @property def config(self) -> ArrayConfig: - """The runtime configuration of this instance.""" + """The runtime configuration.""" return self._async_io.config def to_array(self) -> AnyArray: - """Return an `Array` over this binding; both share one state.""" + """An `Array` over this binding; both share one state.""" return Array(self._async_io.to_array()) def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: - """A new binding with `metadata` in place of this one's. See - `AsyncArrayIO.with_metadata`.""" + """This binding with `metadata` in place of its own. See `AsyncArrayIO.with_metadata`.""" return self._wrap(self._async_io.with_metadata(metadata)) def with_store(self, store: Store | StorePath) -> Self: - """A new binding for the same metadata in `store`. See - `AsyncArrayIO.with_store`.""" + """This binding in `store`. See `AsyncArrayIO.with_store`.""" return self._wrap(self._async_io.with_store(store)) def with_config(self, config: ArrayConfigLike) -> Self: - """A new binding with `config` merged over this one's. See - `AsyncArrayIO.with_config`.""" + """This binding with `config` merged over its own. See `AsyncArrayIO.with_config`.""" return self._wrap(self._async_io.with_config(config)) def store_metadata(self, *, ensure_parents: bool = False) -> None: - """Store this array's metadata. See `AsyncArrayIO.store_metadata`.""" + """Store the metadata. See `AsyncArrayIO.store_metadata`.""" sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) def verify(self) -> None: - """Check that the stored metadata declares the array this instance is bound to. - - See `AsyncArrayIO.verify`. - """ + """Check that the stored metadata equals the bound metadata. See `AsyncArrayIO.verify`.""" sync(self._async_io.verify()) def erase_metadata(self) -> None: - """Erase this array's stored metadata. See `AsyncArrayIO.erase_metadata`.""" + """Delete the stored metadata; chunks are kept. See `AsyncArrayIO.erase_metadata`.""" sync(self._async_io.erase_metadata()) def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: - """Erase one stored chunk by its chunk grid coordinates. See `AsyncArrayIO.erase_chunk`.""" + """Delete one stored chunk by grid coordinates. See `AsyncArrayIO.erase_chunk`.""" sync(self._async_io.erase_chunk(chunk_coords)) def erase_chunks(self) -> None: - """Erase every stored chunk of this array. See `AsyncArrayIO.erase_chunks`.""" + """Delete every stored chunk. See `AsyncArrayIO.erase_chunks`.""" sync(self._async_io.erase_chunks()) def getitem( @@ -6967,7 +6852,7 @@ def getitem( *, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: - """Read the selected region, reading absent chunks as the fill value.""" + """Read the selection, reading absent chunks as the fill value. See `AsyncArrayIO.getitem`.""" return sync(self._async_io.getitem(selection, prototype=prototype)) def setitem( @@ -6976,7 +6861,7 @@ def setitem( value: npt.ArrayLike, prototype: BufferPrototype | None = None, ) -> None: - """Encode `value` and store the chunks covering `selection`.""" + """Encode `value` and store the chunks. See `AsyncArrayIO.setitem`.""" sync(self._async_io.setitem(selection, value, prototype=prototype)) def __repr__(self) -> str: From 07d5aead215618660e9303b6c88183f55e4f5cf0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 17:28:12 +0100 Subject: [PATCH 16/48] fix(array): stop derived bindings aliasing their source's legacy metadata Addresses the roborev review of this branch (job 1359). with_metadata, with_store and with_config copy a legacy source object, so the _stored_document repair mark and in-place attribute updates never leak between bindings, and they carry the captured pipeline implementation instead of re-reading the global configuration at call time; AsyncArray.with_config routes through the binding and inherits all of that. Storing a v2 model without attributes now deletes a stale .zattrs, matching the unconditional-replacement contract, and the model branch honors json_indent. erase_chunks reads async.concurrency with an explicit None check instead of a falsy or-default. Behavioral tests cover array-to-binding state propagation, derivative independence, and the stale-.zattrs case. Assisted-by: ClaudeCode:claude-fable-5 --- src/zarr/core/array.py | 51 ++++++++++++++++++++++++++++++------------ tests/test_array_io.py | 45 ++++++++++++++++++++++++++++++++++++- 2 files changed, 81 insertions(+), 15 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index cf36a46f0b..bcf90dd5f0 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -1236,13 +1236,7 @@ def with_config(self, config: ArrayConfigLike) -> Self: ------- A new Array """ - if isinstance(config, ArrayConfig): - new_config = config - else: - # Merge new config with existing config, so missing keys are inherited - # from the current array rather than from global defaults - new_config = ArrayConfig(**{**self.config.to_dict(), **config}) # type: ignore[arg-type] - return type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) + return cast("Self", self._io.with_config(config).to_array()) async def nchunks_initialized(self) -> int: """ @@ -6531,16 +6525,40 @@ def to_array(self) -> AnyAsyncArray: """An `AsyncArray` over this binding; both share one state.""" return AsyncArray._from_io(self) + def _derived_source(self) -> ArrayMetadataModel | ArrayMetadata: + """The source for a derived binding: a legacy object is copied, so the + repair mark and in-place attribute updates never alias across bindings; + a model is immutable and shared.""" + source = self._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + return copy.copy(source) + return source + + def _derive( + self, + source: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict, + store_path: StorePath, + config: ArrayConfig, + ) -> Self: + """A new binding, keeping this one's captured pipeline implementation.""" + derived = type(self)(source, store_path, config) + derived._pipeline_class = self._pipeline_class + derived._codec_pipeline_cache = None + return derived + def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: - """This binding with `metadata` in place of its own. Nothing is stored.""" - return type(self)(metadata, self._store_path, self._config) + """This binding with `metadata` in place of its own. Nothing is stored; + a legacy `metadata` object is copied.""" + if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): + metadata = copy.copy(metadata) + return self._derive(metadata, self._store_path, self._config) def with_store(self, store: Store | StorePath) -> Self: """This binding at its path within `store`, or at a `StorePath` as given. Nothing is stored.""" store_path = ( store if isinstance(store, StorePath) else StorePath(store, self._store_path.path) ) - return type(self)(self._source, store_path, self._config) + return self._derive(self._derived_source(), store_path, self._config) def with_config(self, config: ArrayConfigLike) -> Self: """This binding with `config` merged over its runtime configuration, as `AsyncArray.with_config` merges it.""" @@ -6548,7 +6566,7 @@ def with_config(self, config: ArrayConfigLike) -> Self: new_config = config else: new_config = ArrayConfig(**{**self._config.to_dict(), **config}) # type: ignore[arg-type] - return type(self)(self._source, self._store_path, new_config) + return self._derive(self._derived_source(), self._store_path, new_config) async def store_metadata(self, *, ensure_parents: bool = False) -> None: """Store the metadata, unconditionally replacing any stored metadata here. @@ -6563,13 +6581,17 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: _stored_document_replaced(source) else: prototype = default_buffer_prototype() + key_value = source.to_key_value(indent=zarr_config.get("json_indent")) to_save = { - str(key): prototype.buffer.from_bytes(value) - for key, value in source.to_key_value().items() + str(key): prototype.buffer.from_bytes(value) for key, value in key_value.items() } await _write_metadata( self._store_path, to_save, self._zarr_format, ensure_parents=ensure_parents ) + if self._zarr_format == 2 and ZATTRS_JSON not in to_save: + # An attribute-less model emits no .zattrs; a stale one must not + # survive an unconditional metadata write. + await (self._store_path / ZATTRS_JSON).delete() async def verify(self) -> None: """Check that the stored metadata equals the bound metadata. @@ -6629,7 +6651,8 @@ async def erase_chunks(self) -> None: n_chunks = product(grid_shape) if n_chunks == 0: return - limit = zarr_config.get("async.concurrency") or 1000 + concurrency = zarr_config.get("async.concurrency") + limit = concurrency if concurrency is not None else 1000 await concurrent_foreach( ( (self._store_path, self._encode_chunk_key(coords)) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 46c60205e3..60711760de 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -6,7 +6,8 @@ import numpy as np import pytest -from zarr_metadata import ZarrV3ArrayMetadata +from zarr_metadata import UNSET, ZarrV3ArrayMetadata +from zarr_metadata import ZarrV2ArrayMetadata as ModelV2 import zarr from zarr import ArrayIO, AsyncArrayIO @@ -257,3 +258,45 @@ def test_binding_follows_the_global_config() -> None: legacy = v2.metadata_legacy assert isinstance(legacy, ArrayV2Metadata) assert legacy.order == "F" + + +async def test_array_state_propagates_to_the_shared_binding() -> None: + """A change through the array side (resize, update_attributes) is visible + on the shared binding's model, and a derived binding stays independent.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await aio.store_metadata() + arr = aio.to_array() + + derived = aio.with_config({"order": "F"}) + moved = aio.with_store(MemoryStore()) + + await arr.resize((6,)) + assert aio.metadata.shape == (6,) + + await arr.update_attributes({"a": 1}) + assert dict(aio.metadata.attributes) == {"a": 1} + + # the derivatives were made before the changes and do not alias them + assert derived.metadata_legacy.shape == (4,) + assert dict(derived.metadata_legacy.attributes) == {} + assert dict(moved.metadata_legacy.attributes) == {} + + +async def test_store_metadata_of_model_without_attributes_clears_zattrs() -> None: + """Storing a v2 model with no attributes removes a stale .zattrs.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v2( + store, name="foo", shape=(4,), dtype="uint8", chunks=(4,), attributes={"stale": 1} + ) + await aio.store_metadata() + assert "foo/.zattrs" in await _store_keys(store) + + opened = await AsyncArrayIO.open(store, name="foo") + model = opened.metadata + assert isinstance(model, ModelV2) + bare = opened.with_metadata(model.update(attributes=UNSET)) + await bare.store_metadata() + assert "foo/.zattrs" not in await _store_keys(store) + reread = await AsyncArrayIO.open(store, name="foo") + assert dict(reread.metadata.attributes) == {} From 34421c99da5687c8dd56a9bf28bc77641f2b9adf Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 31 Oct 2026 17:41:05 +0100 Subject: [PATCH 17/48] fix(array): copy a legacy source's attributes when deriving a binding `copy.copy` on a legacy metadata object shared its `attributes` dict, so an in-place attribute update still reached bindings derived before it. Derived sources now go through `_copied_legacy`, which copies the dict and keeps the repair mark. `_derive` fills the fields of a new binding directly instead of re-running `__init__` (which re-read the global pipeline choice), and `AsyncArray.with_config` builds the derived pipeline eagerly and constructs through `type(self)`, keeping subclass identity and warning timing. Assisted-by: ClaudeCode:claude-fable-5 Co-Authored-By: Claude Fable 5 --- src/zarr/core/array.py | 33 +++++++++++++++++++++++++++------ tests/test_array_io.py | 12 ++++++++++++ 2 files changed, 39 insertions(+), 6 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index bcf90dd5f0..5b7ffd3788 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -1236,7 +1236,11 @@ def with_config(self, config: ArrayConfigLike) -> Self: ------- A new Array """ - return cast("Self", self._io.with_config(config).to_array()) + derived = self._io.with_config(config) + # Build the pipeline now, so construction warnings fire here, as they + # do in __init__. + _ = derived._resolved_codec_pipeline + return cast("Self", type(self)._from_io(derived)) async def nchunks_initialized(self) -> int: """ @@ -6223,6 +6227,16 @@ async def _delete_key(store_path: StorePath, key: str) -> None: await (store_path / key).delete() +def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: + """A copy of a legacy metadata object that shares no mutable state with it: + the attributes dict is copied, and the `_stored_document` repair mark is + kept on the copy.""" + copied = copy.copy(source) + if source.attributes is not None: + object.__setattr__(copied, "attributes", dict(source.attributes)) + return copied + + type ArrayMetadataModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata """Array metadata as a `zarr_metadata` model: validated JSON, extension names unresolved.""" @@ -6531,18 +6545,25 @@ def _derived_source(self) -> ArrayMetadataModel | ArrayMetadata: a model is immutable and shared.""" source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): - return copy.copy(source) + return _copied_legacy(source) return source def _derive( self, - source: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict, + source: ArrayMetadataModel | ArrayMetadata, store_path: StorePath, config: ArrayConfig, ) -> Self: - """A new binding, keeping this one's captured pipeline implementation.""" - derived = type(self)(source, store_path, config) + """A new binding over `source`, keeping this one's captured pipeline.""" + derived = object.__new__(type(self)) + derived._store_path = store_path + derived._config = config derived._pipeline_class = self._pipeline_class + derived._source = source + derived._metadata_model_cache = None + derived._resolved_cache = None + derived._chunk_grid_cache = None + derived._chunk_key_encoding_cache = None derived._codec_pipeline_cache = None return derived @@ -6550,7 +6571,7 @@ def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: """This binding with `metadata` in place of its own. Nothing is stored; a legacy `metadata` object is copied.""" if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): - metadata = copy.copy(metadata) + metadata = _copied_legacy(metadata) return self._derive(metadata, self._store_path, self._config) def with_store(self, store: Store | StorePath) -> Self: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 60711760de..734fe60f09 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -282,6 +282,18 @@ async def test_array_state_propagates_to_the_shared_binding() -> None: assert dict(derived.metadata_legacy.attributes) == {} assert dict(moved.metadata_legacy.attributes) == {} + # the same independence holds for a legacy-sourced binding, whose mutable + # attributes dict must be copied, not shared + legacy_arr = await zarr.api.asynchronous.create_array( + MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,) + ) + legacy_derived = legacy_arr.io.with_config({"order": "F"}) + _ = legacy_derived.metadata # fill the model cache before the mutation + await legacy_arr.update_attributes({"b": 2}) + assert dict(legacy_arr.io.metadata.attributes) == {"b": 2} + assert dict(legacy_derived.metadata_legacy.attributes) == {} + assert dict(legacy_derived.metadata.attributes) == {} + async def test_store_metadata_of_model_without_attributes_clears_zattrs() -> None: """Storing a v2 model with no attributes removes a stale .zattrs.""" From 1546f2daf26ee19ef19976232156d62cc4b39fcc Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 12:09:23 +0100 Subject: [PATCH 18/48] refactor(metadata): record stored documents in a per-store cache, not on metadata Whether a stored array document needs a repair that moves chunks is a fact about the store, so it now lives in `zarr.core.metadata.stored`: a cache per store object, keyed by array path and Zarr format, filled by metadata reads (once the document has been read as metadata) and writes, and consulted before chunk writes, which read the store on a miss. Consolidated copies of such documents are recorded apart, so a group write keeps them as stored. This removes the `_stored_document` attribute from `ArrayV2Metadata` and `ArrayV3Metadata`, and every place that had to carry or clear it: derived bindings, re-resolution and resize no longer manage store state. `mark_repaired` becomes `warn_repairs`, and each `from_dict` is a private `_read_stored` reader plus its warning, so the cache can read stored documents without the rectilinear flag that gates storing them. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- changes/419.feature.md | 1 + src/zarr/api/asynchronous.py | 7 +- src/zarr/core/array.py | 122 ++++++++++--------- src/zarr/core/group.py | 81 +++++++++---- src/zarr/core/metadata/io.py | 85 +++++++++++-- src/zarr/core/metadata/repair.py | 38 ++---- src/zarr/core/metadata/stored.py | 137 +++++++++++++++++++++ src/zarr/core/metadata/v2.py | 18 +-- src/zarr/core/metadata/v3.py | 19 +-- tests/test_array_stateful.py | 7 +- tests/test_metadata/test_repair.py | 185 ++++++++++++++++++++++++++--- 11 files changed, 555 insertions(+), 145 deletions(-) create mode 100644 src/zarr/core/metadata/stored.py diff --git a/changes/419.feature.md b/changes/419.feature.md index 7e888b6a5e..597e3baf15 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -3,3 +3,4 @@ Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arr The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. `with_metadata`, `with_store` and `with_config` return new bindings with one part replaced. `ArrayIO` responds to the global configuration the way `Array` always has: the codec pipeline implementation is captured when the binding is made, the runtime configuration defaults from `zarr.config`, and `build_v2`'s `order` defaults from `array.order`. +Whether a stored array document needs a repair that moves chunks (such as a stored chunk size of 0) is now recorded per store object, in a cache of the metadata documents zarr has read and stored, instead of on the metadata read from it. Every array bound to one store object shares what any of them learned, derived bindings carry nothing, and the first chunk write through any binding stores the repair, including a pickled array or an `ArrayIO` declaration whose chunk layout matches the stored document; a store object that has not seen the document reads it first. diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 476276fd82..39d3a506f1 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -18,6 +18,7 @@ create_array, from_array, get_array_metadata, + observe_read, ) from zarr.core.array_spec import ArrayConfigLike, parse_array_config from zarr.core.buffer import NDArrayLike @@ -423,7 +424,11 @@ async def open( zarr_format = _metadata_dict["zarr_format"] is_v3_array = zarr_format == 3 and _metadata_dict.get("node_type") == "array" if is_v3_array or zarr_format == 2: - return AsyncArray(store_path=store_path, metadata=_metadata_dict, config=config) + array = AsyncArray( + store_path=store_path, metadata=_metadata_dict, config=config + ) + observe_read(store_path, metadata_dict) + return array except (FileNotFoundError, NodeTypeValidationError): pass _check_no_group_config(config, store_path) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 5b7ffd3788..08acae943b 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -129,11 +129,14 @@ encode_documents, parse_stored_array, read_documents, + read_repaired, save_metadata, save_new_metadata, store_documents, + stored_array_document, upsert_metadata, ) +from zarr.core.metadata.stored import moves_chunks, stored_documents from zarr.core.metadata.v2 import ( CompressorLikev2, get_object_codec_id, @@ -307,6 +310,13 @@ def create_codec_pipeline( raise TypeError # pragma: no cover +def observe_read(store_path: StorePath, document: dict[str, JSON]) -> None: + """Record the array document `get_array_metadata` read at `store_path`, once it has + been read as metadata, in the store's cache (see `zarr.core.metadata.stored`).""" + zarr_format: ZarrFormat = 2 if document.get("zarr_format") == 2 else 3 + stored_documents(store_path.store).observe(store_path.path, document, zarr_format) + + async def get_array_metadata( store_path: StorePath, zarr_format: ZarrFormat | None = 3 ) -> dict[str, JSON]: @@ -836,7 +846,9 @@ async def example(): metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) # TODO: remove this cast when we have better type hints _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict) - return cls(store_path=store_path, metadata=_metadata_dict, config=config) + array = cls(store_path=store_path, metadata=_metadata_dict, config=config) + observe_read(store_path, metadata_dict) + return array @property def store(self) -> Store: @@ -1610,20 +1622,9 @@ async def get_coordinate_selection( return out_array async def _save_metadata(self, ensure_parents: bool = False) -> None: - """Store this array's metadata documents, delegating to the wrapped - `AsyncArrayIO`, which clears the `_stored_document` mark (see - `_stored_document_replaced`). `_resize` stores the documents it encoded before - deleting chunks directly, then clears the mark the same way.""" + """Store this array's metadata documents, delegating to the wrapped `AsyncArrayIO`.""" await self._io.store_metadata(ensure_parents=ensure_parents) - def _stored_document_replaced(self) -> None: - """See the module-level `_stored_document_replaced`.""" - _stored_document_replaced(self.metadata) - - async def _store_repaired_document(self) -> None: - """See the module-level `_store_repaired_document`.""" - await _store_repaired_document(self.store_path, self.metadata) - async def _set_selection( self, indexer: Indexer, @@ -5004,26 +5005,23 @@ def _sharding_codec(metadata: ArrayMetadata) -> ShardingCodec | None: return None -def _stored_document_replaced(metadata: ArrayMetadata) -> None: - """Record that the store no longer holds a document of this array that needs a - repair: it holds the repair, a valid document, or none. The metadata, which a - consolidated group handle may share with an array handle, then stops standing for - the document it was read from (see `mark_repaired`), so no later write through any - handle stores that document again.""" - object.__setattr__(metadata, "_stored_document", None) - - async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadata) -> None: - """Store the repair of the array's current stored document, if it needs one, - before chunks are written under `metadata`. - - Only for metadata read from a document whose repair moves chunks (see - `zarr.core.metadata.repair`). The document is read again, because the store may - hold a newer one than `metadata`. If that one lays out chunks differently - (another writer stored a different chunk size since), the caller would write - chunks no reader finds, so it raises and stores nothing. If it needs no repair - (the array was re-saved since, possibly by another implementation), it is left as - written; if there is none, there is nothing to repair. + """Store the repair of the array document at `store_path`, if it needs one that + moves chunks, before chunks are written under `metadata`. + + The store's cache (see `zarr.core.metadata.stored`) says whether the document + needs such a repair. While it has not seen the document, a consolidated copy it + has seen stands for it, as consolidated metadata does for every reader; with + neither, the document is read. Only a document whose repair + lays out chunks as `metadata` does is acted on: one that lays them out otherwise + describes another array, and chunks written under `metadata` are its caller's + declaration, as over any document. The document is read again, because the store + may hold a newer one. If that one lays out chunks differently (another writer + stored a different chunk size since), the caller would write chunks no reader + finds, so it raises, stores nothing, and leaves the cache as it was, so every + later write raises too. If it needs no repair (the array was re-saved since, + possibly by another implementation), it is left as written; if there is none, + there is nothing to repair. Storing the same repair twice is harmless, so handles that write chunks concurrently need no coordination. The read and the store are not one atomic @@ -5031,23 +5029,38 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat attribute update) is replaced by the repair of the document read before it, as with any two metadata writes that race. """ - if metadata._stored_document is None: - return + cache = stored_documents(store_path.store) + path = store_path.path zarr_format = metadata.zarr_format - documents = await read_documents(store_path, ARRAY_DOCUMENTS[zarr_format]) - try: - current = parse_stored_array(documents, zarr_format, str(store_path)) - except ArrayNotFoundError: - pass + keys = ARRAY_DOCUMENTS[zarr_format] + documents: dict[str, Buffer] | None = None + if cache.knows(path, zarr_format): + stored = cache.needing_repair(path, zarr_format) else: - if _chunk_layout(current) != _chunk_layout(metadata): - raise ValueError( - f"Array {str(store_path)!r}: the metadata stored has changed since " - "this array was opened; reopen the array to write to it. Nothing was stored." - ) - if current._stored_document is not None: - await upsert_metadata(store_path, current, documents) - _stored_document_replaced(metadata) + stored = cache.copy_needing_repair(path, zarr_format) + if stored is None: + documents = await read_documents(store_path, keys) + cache.observe(path, stored_array_document(documents, zarr_format), zarr_format) + stored = cache.needing_repair(path, zarr_format) + if stored is None or _chunk_layout(read_repaired(stored, zarr_format)) != _chunk_layout( + metadata + ): + return + if documents is None: + documents = await read_documents(store_path, keys) + document = stored_array_document(documents, zarr_format) + if document is None: + cache.observe(path, None, zarr_format) + return + current = parse_stored_array(documents, zarr_format, str(store_path)) + if _chunk_layout(current) != _chunk_layout(metadata): + raise ValueError( + f"Array {str(store_path)!r}: the metadata stored has changed since " + "this array was opened; reopen the array to write to it. Nothing was stored." + ) + if moves_chunks(document, zarr_format): + await upsert_metadata(store_path, current, documents) + cache.observe(path, None, zarr_format) def _stored_rectilinear_grid_or_none( @@ -6101,7 +6114,6 @@ async def _delete_chunk(coords: tuple[int, ...]) -> None: # Write new metadata await store_documents(array.store_path, documents) - array._stored_document_replaced() # Update the wrapped binding in place; the codec pipeline is kept, as before. array._io._rebind_resolved(new_metadata, new_chunk_grid) @@ -6229,8 +6241,7 @@ async def _delete_key(store_path: StorePath, key: str) -> None: def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: """A copy of a legacy metadata object that shares no mutable state with it: - the attributes dict is copied, and the `_stored_document` repair mark is - kept on the copy.""" + the attributes dict is copied.""" copied = copy.copy(source) if source.attributes is not None: object.__setattr__(copied, "attributes", dict(source.attributes)) @@ -6394,7 +6405,9 @@ async def open( """ store_path = await make_store_path(store, path=name) data = await get_array_metadata(store_path, zarr_format=zarr_format) - return cls(_array_metadata_model(data), store_path, config) + binding = cls(_array_metadata_model(data), store_path, config) + observe_read(store_path, data) + return binding @property def metadata(self) -> ArrayMetadataModel: @@ -6599,7 +6612,6 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): await save_metadata(self._store_path, source, ensure_parents=ensure_parents) - _stored_document_replaced(source) else: prototype = default_buffer_prototype() key_value = source.to_key_value(indent=zarr_config.get("json_indent")) @@ -6650,9 +6662,9 @@ async def erase_metadata(self) -> None: """Delete the stored metadata; chunks are kept. Succeeds if absent.""" for key in ARRAY_DOCUMENTS[self._zarr_format]: await (self._store_path / key).delete() - resolved = self._resolved_if_any - if resolved is not None: - _stored_document_replaced(resolved) + stored_documents(self._store_path.store).observe( + self._store_path.path, None, self._zarr_format + ) async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: """Delete one stored chunk by grid coordinates (the shard, under sharding). Succeeds if absent.""" diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index f1b256c491..4ea3b36dd7 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -13,7 +13,6 @@ import zarr.api.asynchronous as async_api from zarr.abc.metadata import Metadata -from zarr.abc.store import Store, set_or_delete from zarr.core._info import GroupInfo from zarr.core._json import buffer_to_json_object, json_to_buffer from zarr.core.array import ( @@ -55,6 +54,7 @@ save_new_metadata, store_documents, ) +from zarr.core.metadata.stored import stored_documents from zarr.core.metadata.v3 import AllowedExtraField, check_storable, parse_extra_fields from zarr.core.sync import SyncMixin, sync from zarr.errors import ( @@ -82,11 +82,13 @@ ) from typing import Any + from zarr.abc.store import Store from zarr.core.array_spec import ArrayConfigLike from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike from zarr.core.common import MemoryOrder from zarr.core.dtype import ZDTypeLike + from zarr.core.metadata.repair import ArrayDocument from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 logger = logging.getLogger("zarr.group") @@ -153,17 +155,19 @@ class ConsolidatedMetadata: kind: Literal["inline"] = "inline" must_understand: Literal[False] = False - def to_dict(self) -> dict[str, JSON]: - """The consolidated metadata document. An array read from a stored document that - had to be repaired is written as that document was stored, so every reader of - the consolidated metadata reads it as repaired again (see `mark_repaired`).""" + def to_dict( + self, *, members_as_stored: Mapping[str, ArrayDocument] | None = None + ) -> dict[str, JSON]: + """The consolidated metadata document. A member in `members_as_stored` (by its + path here) is written as that stored document, so every reader of the + consolidated metadata reads it as repaired again (see + `zarr.core.metadata.io.members_as_stored`).""" + as_stored = {} if members_as_stored is None else members_as_stored return { "kind": self.kind, "must_understand": self.must_understand, "metadata": { - k: v.to_dict() - if isinstance(v, GroupMetadata) or v._stored_document is None - else dict(v._stored_document) + k: dict(as_stored[k]) if k in as_stored else v.to_dict() for k, v in sorted( self.flattened_metadata.items(), key=lambda item: ( @@ -384,12 +388,19 @@ class GroupMetadata(Metadata): node_type: Literal["group"] = field(default="group", init=False) extra_fields: dict[str, AllowedExtraField] = field(default_factory=dict) - def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: + def to_buffer_dict( + self, + prototype: BufferPrototype, + *, + members_as_stored: Mapping[str, ArrayDocument] | None = None, + ) -> dict[str, Buffer]: + """The group's metadata documents, by key. A consolidated member in + `members_as_stored` is written as that stored document (see + `ConsolidatedMetadata.to_dict`).""" + as_stored = {} if members_as_stored is None else members_as_stored if self.consolidated_metadata is not None: for path, member in self.consolidated_metadata.flattened_metadata.items(): - # A member read from a document that had to be repaired is stored as it - # was stored (see `ConsolidatedMetadata.to_dict`). - if isinstance(member, ArrayV3Metadata) and member._stored_document is None: + if isinstance(member, ArrayV3Metadata) and path not in as_stored: try: check_storable(member) except ValueError as e: @@ -397,7 +408,12 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: raise indent = config.get("json_indent") if self.zarr_format == 3: - return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)} + document = self.to_dict() + if self.consolidated_metadata is not None: + document["consolidated_metadata"] = self.consolidated_metadata.to_dict( + members_as_stored=as_stored + ) + return {ZARR_JSON: json_to_buffer(document, prototype=prototype, indent=indent)} else: items = { ZGROUP_JSON: json_to_buffer( @@ -410,7 +426,9 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: ZGROUP_JSON: {"zarr_format": self.zarr_format}, ZATTRS_JSON: self.attributes, } - consolidated_metadata = self.consolidated_metadata.to_dict()["metadata"] + consolidated_metadata = self.consolidated_metadata.to_dict( + members_as_stored=as_stored + )["metadata"] if not isinstance(consolidated_metadata, dict): raise TypeError( "Expected consolidated metadata to serialize to a dict, " @@ -744,10 +762,19 @@ def from_dict( elif node_type not in ("group", None): msg = f"Node type in metadata ({node_type}) is not 'group'" raise GroupNotFoundError(msg) - return cls( + group = cls( metadata=GroupMetadata.from_dict(data, path=str(store_path)), store_path=store_path, ) + consolidated = group.metadata.consolidated_metadata + raw_consolidated = data.get("consolidated_metadata") + if consolidated is not None and isinstance(raw_consolidated, dict): + members = raw_consolidated.get("metadata") + if isinstance(members, dict): + stored_documents(store_path.store).observe_consolidated( + store_path.path, members, group.metadata.zarr_format + ) + return group async def setitem(self, key: str, value: Any) -> None: """ @@ -2251,9 +2278,7 @@ async def update_attributes_async(self, new_attributes: dict[str, Any]) -> Group new_metadata = replace(self.metadata, attributes=new_attributes) # Write new metadata - to_save = new_metadata.to_buffer_dict(default_buffer_prototype()) - awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()] - await asyncio.gather(*awaitables) + await store_documents(self.store_path, encode_documents(self.store_path, new_metadata)) async_group = replace(self._async_group, metadata=new_metadata) return replace(self, _async_group=async_group) @@ -3422,6 +3447,11 @@ async def _store_nodes( ] created_object_keys = [] + node_documents: defaultdict[str, dict[str, Buffer]] = defaultdict(dict) + for key, value in documents.items(): + node, _, name = key.rpartition("/") + node_documents[node][name] = value + cache = stored_documents(store) for coro in asyncio.as_completed(create_tasks): created_key = await coro @@ -3439,6 +3469,7 @@ async def _store_nodes( node_name = created_key[: created_key.rfind("/")] meta_out = nodes[node_name] if meta_out.zarr_format == 3: + cache.observe_written(node_name, node_documents[node_name]) yield node_name, built[node_name] else: # For zarr v2 @@ -3454,6 +3485,7 @@ async def _store_nodes( meta_done = _join_paths([node_name, ZARRAY_JSON]) in created_object_keys if meta_done and attrs_done: + cache.observe_written(node_name, node_documents[node_name]) yield node_name, built[node_name] continue @@ -3709,9 +3741,11 @@ async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupM ) if zarr_json_bytes is None: raise FileNotFoundError(path) - return _build_metadata_v3( - buffer_to_json_object(zarr_json_bytes), path=str(StorePath(store, path)) - ) + document = buffer_to_json_object(zarr_json_bytes) + metadata = _build_metadata_v3(document, path=str(StorePath(store, path))) + if isinstance(metadata, ArrayV3Metadata): + stored_documents(store).observe(path, document, 3) + return metadata async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupMetadata: @@ -3746,7 +3780,10 @@ async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupM else: zmeta = buffer_to_json_object(zgroup_bytes) - return _build_metadata_v2(zmeta, zattrs, path=str(StorePath(store, path))) + metadata = _build_metadata_v2(zmeta, zattrs, path=str(StorePath(store, path))) + if isinstance(metadata, ArrayV2Metadata): + stored_documents(store).observe(path, {**zmeta, "attributes": zattrs}, 2) + return metadata async def _read_group_metadata_v2(store: Store, path: str) -> GroupMetadata: diff --git a/src/zarr/core/metadata/io.py b/src/zarr/core/metadata/io.py index c1cd927b94..0c0a58cede 100644 --- a/src/zarr/core/metadata/io.py +++ b/src/zarr/core/metadata/io.py @@ -10,7 +10,8 @@ from zarr.core.buffer.core import default_buffer_prototype from zarr.core.buffer.cpu import buffer_prototype as cpu_buffer_prototype from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON -from zarr.core.metadata.repair import mark_repaired, repair_array_document +from zarr.core.metadata.repair import repair_array_document +from zarr.core.metadata.stored import stored_documents from zarr.errors import ArrayNotFoundError, ContainsArrayError from zarr.storage._common import StorePath, ensure_no_existing_node @@ -21,6 +22,7 @@ from zarr.core.common import JSON, ZarrFormat from zarr.core.group import GroupMetadata from zarr.core.metadata import ArrayMetadata + from zarr.core.metadata.repair import ArrayDocument class _Absent(Enum): @@ -83,6 +85,11 @@ def encode_documents( from zarr.core.group import GroupMetadata try: + if isinstance(metadata, GroupMetadata): + return metadata.to_buffer_dict( + default_buffer_prototype(), + members_as_stored=members_as_stored(store_path, metadata), + ) return metadata.to_buffer_dict(default_buffer_prototype()) except ValueError as e: node = "Group" if isinstance(metadata, GroupMetadata) else "Array" @@ -91,10 +98,74 @@ def encode_documents( async def store_documents(store_path: StorePath, documents: Mapping[str, Buffer]) -> None: - """Store metadata documents encoded by `encode_documents` under `store_path`.""" + """Store metadata documents encoded by `encode_documents` under `store_path`, and + record them in the store's cache (see `stored_documents`).""" await asyncio.gather( *(set_or_delete(store_path / key, value) for key, value in documents.items()) ) + stored_documents(store_path.store).observe_written(store_path.path, documents) + + +def members_as_stored(store_path: StorePath, metadata: GroupMetadata) -> dict[str, ArrayDocument]: + """The members of the consolidated metadata of the group at `store_path` to store as + they are stored, by their paths in it: each array whose stored document, or else + whose stored consolidated copy, needs a repair that moves chunks and reads as the + member does. A group write then leaves those copies as they were stored, so every + reader of the consolidated metadata reads them as repaired again. An array whose + own document zarr knows needs no such repair is stored as it reads.""" + from zarr.core.group import GroupMetadata + from zarr.storage._utils import _join_paths + + consolidated = metadata.consolidated_metadata + if consolidated is None: + return {} + cache = stored_documents(store_path.store) + as_stored: dict[str, ArrayDocument] = {} + for key, member in consolidated.flattened_metadata.items(): + if isinstance(member, GroupMetadata): + continue + path = _join_paths([store_path.path, key]) + zarr_format = member.zarr_format + own = cache.needing_repair(path, zarr_format) + if cache.knows(path, zarr_format) and own is None: + continue + for document in (own, cache.copy_needing_repair(path, zarr_format)): + if document is not None and reads_as(document, member): + as_stored[key] = document + break + return as_stored + + +def reads_as(document: ArrayDocument, metadata: ArrayMetadata) -> bool: + """Whether the stored array document `document`, repaired, reads as `metadata`.""" + return read_repaired(document, metadata.zarr_format).to_dict() == metadata.to_dict() + + +def read_repaired(document: ArrayDocument, zarr_format: ZarrFormat) -> ArrayMetadata: + """The metadata the stored array document `document` reads as, without warnings.""" + from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata + + reader = ArrayV2Metadata if zarr_format == 2 else ArrayV3Metadata + metadata, _ = reader._read_stored(dict(document)) + return metadata + + +def stored_array_document( + documents: Mapping[str, Buffer], zarr_format: ZarrFormat +) -> ArrayDocument | None: + """The array document of `zarr_format` among `documents` (by store key, as + `read_documents` reads them), or `None` if there is none.""" + if zarr_format == 2: + if ZARRAY_JSON not in documents: + return None + document: dict[str, JSON] = buffer_to_json_object(documents[ZARRAY_JSON]) + zattrs = documents.get(ZATTRS_JSON) + document["attributes"] = {} if zattrs is None else buffer_to_json_object(zattrs) + return document + if ZARR_JSON not in documents: + return None + document = buffer_to_json_object(documents[ZARR_JSON]) + return document if document.get("node_type") == "array" else None async def read_documents(store_path: StorePath, keys: Iterable[str]) -> dict[str, Buffer]: @@ -139,8 +210,8 @@ def parse_stored_array( ) -> ArrayMetadata: """The metadata of an array from its documents (by store key, see `ARRAY_DOCUMENTS`), read with the repairs but without their warnings (whoever asks has warned, or reads - metadata built in code), and marked (see `mark_repaired`) if they had to be - repaired. Raises `ArrayNotFoundError` if there is no array document among them. + metadata built in code). Raises `ArrayNotFoundError` if there is no array document + among them. Only operations that store metadata read documents this way, so documents read as a rectilinear chunk grid require the rectilinear chunks flag, as storing it does; the @@ -157,9 +228,8 @@ def parse_stored_array( stored = _array_metadata_dict_v3(documents[ZARR_JSON]) else: raise ArrayNotFoundError(f"No Zarr format {zarr_format} array metadata document.") - repaired, readings = repair_array_document(stored, zarr_format) - metadata = parse_array_metadata(dict(repaired), path) - return mark_repaired(metadata, stored, readings, None, warn=False) + repaired, _ = repair_array_document(stored, zarr_format) + return parse_array_metadata(dict(repaired), path) def _build_parents(store_path: StorePath, zarr_format: ZarrFormat) -> dict[str, GroupMetadata]: @@ -281,3 +351,4 @@ async def _write_metadata( ) from e await asyncio.gather(*set_awaitables) + stored_documents(store_path.store).observe_written(store_path.path, to_save) diff --git a/src/zarr/core/metadata/repair.py b/src/zarr/core/metadata/repair.py index d8050b9f3d..2d6daddb30 100644 --- a/src/zarr/core/metadata/repair.py +++ b/src/zarr/core/metadata/repair.py @@ -9,18 +9,18 @@ read from the same document before these repairs existed is silent, so a document that opened without a warning still does. The warnings are given once per document, after the repaired document has passed the metadata constructor, so an invalid -document raises its own error, not a warning about how it was read. Metadata read -from a document whose repair moves chunks (a chunk size read as another size) is -marked (see `mark_repaired`), silent or not, so the array stores the repair before it -writes chunks under it. A repair that only respells a value zarr already read the -same way (`true` as 1, `4.0` as 4) moves no chunks, so it is not marked. +document raises its own error, not a warning about how it was read. A store whose +document needs a repair that moves chunks (a chunk size read as another size) gets +that repair stored before zarr writes chunks under it, silent or not; whether a +stored document needs one is recorded per store (see `zarr.core.metadata.stored`), +not on the metadata read from it. A repair that only respells a value zarr already +read the same way (`true` as 1, `4.0` as 4) moves no chunks. To read another kind of invalid document, add a repair to `ARRAY_REPAIRS`. """ from __future__ import annotations -import copy import json import warnings from collections.abc import Callable, Iterable, Mapping, Sequence @@ -71,30 +71,16 @@ class Reading(NamedTuple): and a re-save would keep it.""" -def mark_repaired[M]( - metadata: M, - stored: ArrayDocument, - readings: Sequence[Reading], - path: str | None, - *, - warn: bool = True, -) -> M: - """Record that `metadata` was read from the document `stored`, which needed the - repairs whose `readings` `repair_array_document` returned, if any. If a reading - moves chunks, keep a copy of `stored` as `_stored_document` on `metadata`, so the - array stores the repair before it writes chunks under it and consolidated metadata - stores it as it was stored. If `warn`, warn once with the readings' warnings (each - says what the user must act on and how), naming the array at `path` when the caller - knows it.""" - if any(reading.moves_chunks for reading in readings): - object.__setattr__(metadata, "_stored_document", copy.deepcopy(stored)) +def warn_repairs(readings: Sequence[Reading], path: str | None) -> None: + """Warn once with the warnings of the `readings` `repair_array_document` returned + (each says what the user must act on and how), naming the array at `path` when + the caller knows it.""" messages = [reading.warning for reading in readings if reading.warning is not None] - if warn and messages: + if messages: subject = "" if path is None else f"Array {path!r}: " # The synchronous API parses metadata on zarr's IO thread, whose stack holds no # user code, so the warning points at the `from_dict` that read the document. warnings.warn(f"{subject}{' '.join(messages)}", ZarrUserWarning, stacklevel=2) - return metadata def _is_int_list(value: object) -> TypeGuard[list[int]]: @@ -375,7 +361,7 @@ def repair_array_document( """Apply the repairs for `zarr_format` to a stored array metadata document. Returns the repaired document and the reading of each repair that changed it, for - `mark_repaired` once the document has been validated. + `warn_repairs` once the document has been validated. """ readings: list[Reading] = [] for repair in ARRAY_REPAIRS[zarr_format]: diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py new file mode 100644 index 0000000000..9569a141e0 --- /dev/null +++ b/src/zarr/core/metadata/stored.py @@ -0,0 +1,137 @@ +"""What zarr knows of the array metadata documents a store holds: a cache over the store. + +Whether an array's stored document needs a repair that moves chunks (see +`zarr.core.metadata.repair`) is a fact about the store, not about metadata read from +it. zarr records it here, per store object, from the array documents it reads and +stores, and asks here before it writes chunks. Every binding of one store object +shares what any of them learned, and metadata stays a value that copies and +re-resolves without carrying store state. + +The cache observes; it never serves a read. A path it knows nothing about is read +from the store when the answer is needed. +""" + +from __future__ import annotations + +import copy +import weakref +from functools import partial +from typing import TYPE_CHECKING, Final + +from zarr.core._json import buffer_to_json_object +from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON +from zarr.core.metadata.repair import repair_array_document + +if TYPE_CHECKING: + from collections.abc import Mapping + + from zarr.abc.store import Store + from zarr.core.buffer import Buffer + from zarr.core.common import JSON, ZarrFormat + from zarr.core.metadata.repair import ArrayDocument + + +def moves_chunks(document: ArrayDocument, zarr_format: ZarrFormat) -> bool: + """Whether reading `document` needs a repair that moves chunks.""" + _, readings = repair_array_document(document, zarr_format) + return any(reading.moves_chunks for reading in readings) + + +class StoredDocuments: + """The array metadata documents zarr has read from or stored to one store. + + Each array document zarr has seen, by its path and Zarr format (a path may hold a + document of each), maps to the document as it was stored when reading it needs a + repair that moves chunks, and to `None` when it needs none or there is no + document; a document zarr has not seen is unknown. Copies of array documents in + consolidated metadata are recorded apart, by the array's path and format, only + when reading them needs such a repair, so a group write can store them as they + were stored. zarr records a document it read only once it has read it as + metadata, so a document it could not read is unknown. + """ + + def __init__(self) -> None: + self._arrays: dict[tuple[str, ZarrFormat], ArrayDocument | None] = {} + self._copies: dict[tuple[str, ZarrFormat], ArrayDocument] = {} + + def knows(self, path: str, zarr_format: ZarrFormat) -> bool: + """Whether zarr has read or stored the array document of `zarr_format` at `path`.""" + return (path, zarr_format) in self._arrays + + def needing_repair(self, path: str, zarr_format: ZarrFormat) -> ArrayDocument | None: + """The document of `zarr_format` stored at `path` as it was stored, if reading it + needs a repair that moves chunks; `None` if it needs none, there is none, or it + is unknown.""" + return self._arrays.get((path, zarr_format)) + + def copy_needing_repair(self, path: str, zarr_format: ZarrFormat) -> ArrayDocument | None: + """The consolidated copy of the array document of `zarr_format` at `path` as it + was stored, if reading it needs a repair that moves chunks.""" + return self._copies.get((path, zarr_format)) + + def observe(self, path: str, document: ArrayDocument | None, zarr_format: ZarrFormat) -> None: + """Record that the array document of `zarr_format` at `path` is `document` + (`None`: none).""" + if document is not None and moves_chunks(document, zarr_format): + self._arrays[path, zarr_format] = copy.deepcopy(document) + else: + self._arrays[path, zarr_format] = None + + def observe_written(self, path: str, documents: Mapping[str, Buffer | None]) -> None: + """Record the metadata documents just stored under `path`, by key (`None`: + deleted), as `store_documents` takes them.""" + if ZARR_JSON in documents: + zarr_json = documents[ZARR_JSON] + document = None if zarr_json is None else buffer_to_json_object(zarr_json) + if document is not None and document.get("node_type") != "array": + document = None + self.observe(path, document, 3) + if ZARRAY_JSON in documents: + zarray = documents[ZARRAY_JSON] + if zarray is None: + self.observe(path, None, 2) + else: + zattrs = documents.get(ZATTRS_JSON) + v2_document: dict[str, JSON] = buffer_to_json_object(zarray) + v2_document["attributes"] = {} if zattrs is None else buffer_to_json_object(zattrs) + self.observe(path, v2_document, 2) + + def observe_consolidated( + self, group_path: str, members: Mapping[str, JSON], zarr_format: ZarrFormat + ) -> None: + """Record the member documents of consolidated metadata read from the group at + `group_path`, by their paths relative to it.""" + prefix = f"{group_path}/" if group_path else "" + for key, member in members.items(): + if not isinstance(member, dict) or member.get("node_type", "array") != "array": + continue + if zarr_format == 2 and "shape" not in member: + continue # a v2 group: it has no `shape` + path = f"{prefix}{key}" + if moves_chunks(member, zarr_format): + self._copies[path, zarr_format] = copy.deepcopy(member) + else: + self._copies.pop((path, zarr_format), None) + + +def _forget(key: int, ref: weakref.ref[Store]) -> None: + entry = _CACHES.get(key) + if entry is not None and entry[0] is ref: + del _CACHES[key] + + +_CACHES: Final[dict[int, tuple[weakref.ref[Store], StoredDocuments]]] = {} +"""The cache of each live store object, by its `id`: stores define `__eq__`, so they +are not hashable, and two equal stores are not one store object.""" + + +def stored_documents(store: Store) -> StoredDocuments: + """The cache of `store`, shared by every binding of this store object; it is dropped + with the store.""" + key = id(store) + entry = _CACHES.get(key) + if entry is not None and entry[0]() is store: + return entry[1] + documents = StoredDocuments() + _CACHES[key] = (weakref.ref(store, partial(_forget, key)), documents) + return documents diff --git a/src/zarr/core/metadata/v2.py b/src/zarr/core/metadata/v2.py index 136267fc58..6b32b41676 100644 --- a/src/zarr/core/metadata/v2.py +++ b/src/zarr/core/metadata/v2.py @@ -4,7 +4,7 @@ import warnings from collections.abc import Iterable, Sequence from functools import cached_property -from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast +from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast from zarr.abc.metadata import Metadata from zarr.abc.numcodec import Numcodec, _is_numcodec @@ -25,7 +25,7 @@ TBaseScalar, ZDType, ) - from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.repair import Reading from dataclasses import dataclass, field, fields, replace @@ -45,7 +45,7 @@ from zarr.core.config import config, parse_indexing_order from zarr.core.json_parse import parse_field from zarr.core.metadata.common import parse_attributes -from zarr.core.metadata.repair import mark_repaired, repair_array_document +from zarr.core.metadata.repair import repair_array_document, warn_repairs class ArrayV2MetadataDict(TypedDict): @@ -73,9 +73,6 @@ class ArrayV2Metadata(Metadata): compressor: Numcodec | None attributes: dict[str, JSON] = field(default_factory=dict) zarr_format: Literal[2] = field(init=False, default=2) - _stored_document: ClassVar[ArrayDocument | None] = None - """The stored document `from_dict` read this metadata from, if it had to repair it - (set on the instance by `mark_repaired`): the store may still hold it.""" def __init__( self, @@ -156,6 +153,13 @@ def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2M """Read a stored `.zarray` document (with its attributes). An invalid document that `zarr.core.metadata.repair` can read is read as repaired; a reading the user must act on warns, naming the array at `path`.""" + metadata, readings = cls._read_stored(data) + warn_repairs(readings, path) + return metadata + + @classmethod + def _read_stored(cls, data: dict[str, Any]) -> tuple[ArrayV2Metadata, list[Reading]]: + """`from_dict` without its warnings: the metadata, and how the repairs read it.""" repaired, readings = repair_array_document(data, 2) # a new dict, because we are modifying it _data: dict[str, Any] = dict(repaired) @@ -209,7 +213,7 @@ def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2M _data = {k: v for k, v in _data.items() if k in expected} - return mark_repaired(cls(**_data), data, readings, path) + return cls(**_data), readings def to_dict(self) -> dict[str, JSON]: zarray_dict = super().to_dict() diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2d5d169e6b..186e68f1c3 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -4,7 +4,7 @@ from collections.abc import Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, TypeGuard, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast from typing_extensions import TypedDict @@ -41,8 +41,8 @@ from zarr.core.json_parse import parse_field from zarr.core.metadata.common import parse_attributes from zarr.core.metadata.repair import ( - mark_repaired, repair_array_document, + warn_repairs, ) from zarr.errors import MetadataValidationError, NodeTypeValidationError from zarr.registry import get_codec_class @@ -54,7 +54,7 @@ from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_grids import ChunkGrid from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar - from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.repair import Reading def parse_zarr_format(data: object) -> Literal[3]: @@ -519,9 +519,6 @@ class ArrayV3Metadata(Metadata): node_type: Literal["array"] = field(default="array", init=False) storage_transformers: tuple[dict[str, JSON], ...] extra_fields: dict[str, AllowedExtraField] - _stored_document: ClassVar[ArrayDocument | None] = None - """The stored document `from_dict` read this metadata from, if it had to repair it - (set on the instance by `mark_repaired`): the store may still hold it.""" def __init__( self, @@ -679,6 +676,14 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: if path is not None: e.add_note(f"Array {path!r}: nothing was read.") raise + metadata, readings = cls._read_stored(data) + warn_repairs(readings, path) + return metadata + + @classmethod + def _read_stored(cls, data: dict[str, JSON]) -> tuple[Self, list[Reading]]: + """`from_dict` without its warnings or the rectilinear chunks flag: the + metadata, and how the repairs read it.""" repaired, readings = repair_array_document(data, 3) # a new dict, because we are modifying it _data = dict(repaired) @@ -723,7 +728,7 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: extra_fields=allowed_extra_fields, storage_transformers=_data_typed.get("storage_transformers", ()), # type: ignore[arg-type] ) - return mark_repaired(metadata, data, readings, path) + return metadata, readings def to_dict(self) -> dict[str, JSON]: out_dict = super().to_dict() diff --git a/tests/test_array_stateful.py b/tests/test_array_stateful.py index ca9b1d3ad8..2bd69df822 100644 --- a/tests/test_array_stateful.py +++ b/tests/test_array_stateful.py @@ -37,6 +37,7 @@ import zarr from zarr.core.buffer import cpu, default_buffer_prototype from zarr.core.chunk_grids import ChunkGrid +from zarr.core.metadata.stored import stored_documents from zarr.core.sync import sync from zarr.errors import ZarrUserWarning from zarr.storage import MemoryStore @@ -157,7 +158,11 @@ def _open(self) -> zarr.Array[Any]: warned = any(issubclass(w.category, ZarrUserWarning) for w in record) must_warn = any(self.shape[axis] > 0 for axis in self.legacy_axes) assert warned is must_warn, [str(w.message) for w in record] - assert (arr.metadata._stored_document is not None) is bool(self.legacy_axes) + cache = stored_documents(arr.store_path.store) + needs_repair = ( + cache.needing_repair(arr.store_path.path, arr.metadata.zarr_format) is not None + ) + assert needs_repair is bool(self.legacy_axes) return arr # ----------------------------------------------------------------- model diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index 5247754533..bf9311299d 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -4,7 +4,9 @@ import asyncio import dataclasses +import gc import json +import pickle import re import warnings import zipfile @@ -24,6 +26,7 @@ RESAVE_HINT, repair_array_document, ) +from zarr.core.metadata.stored import _CACHES, StoredDocuments, moves_chunks, stored_documents from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata from zarr.core.sync import sync from zarr.dtype import Int16 @@ -39,6 +42,13 @@ from zarr.types import AnyArray +def _store_needs_repair(array: AnyArray) -> bool: + """Whether the cache of the array's store records that its stored document needs a + repair that moves chunks.""" + cache = stored_documents(array.store_path.store) + return cache.needing_repair(array.store_path.path, array.metadata.zarr_format) is not None + + def _v2_doc(shape: list[int], chunks: list[Any]) -> dict[str, JSON]: return { "zarr_format": 2, @@ -203,7 +213,7 @@ def test_repair_array_document( warnings.simplefilter("always") metadata = metadata_cls.from_dict(dict(doc), path="group/array") assert _chunk_shapes(metadata) == expected - assert metadata._stored_document == (doc if repaired == "moved" else None) + assert moves_chunks(doc, cast("ZarrFormat", doc["zarr_format"])) is (repaired == "moved") messages = [str(w.message) for w in record] if warning is None: assert messages == [] @@ -314,14 +324,14 @@ def test_read_invalid_edges_in_rectilinear_grid( """A stored rectilinear chunk grid whose explicit edges or run-length encoded sizes are integral floats or JSON `true`, as zarr-python wrote them when given float or `True` edges, is read with those edges as the `int`s they equal, silently: zarr read - these edges so before. The reading moves no chunks, so the metadata is not marked - for re-saving.""" + these edges so before. The reading moves no chunks, so no repair is stored before + chunk writes.""" shape = [sum(edges) for edges in expected] doc = _rectilinear_doc(shape, chunk_shapes) with zarr.config.set({"array.rectilinear_chunks": True}): metadata = _read_strictly(doc) assert metadata.chunk_grid == RectilinearChunkGridMetadata(chunk_shapes=expected) - assert metadata._stored_document is None + assert not moves_chunks(doc, 3) @pytest.mark.parametrize( @@ -675,7 +685,7 @@ async def write_twice() -> None: sync(write_twice()) - assert arr.metadata._stored_document is None + assert not _store_needs_repair(arr) with warnings.catch_warnings(): warnings.simplefilter("error", ZarrUserWarning) reopened = zarr.open_array(store=path, mode="r") @@ -715,7 +725,7 @@ def test_stale_handle_write_keeps_newer_metadata( stale[0] = 9 - assert stale.metadata._stored_document is None + assert not _store_needs_repair(stale) reopened = _open_strictly(path) assert reopened.shape == (9,) assert reopened.attrs.asdict() == {"x": 1} @@ -843,7 +853,7 @@ def test_stale_handle_write_after_chunk_grid_change_raises( with pytest.raises(ValueError, match="has changed since this array was opened; reopen"): stale[0:3] = [7, 8, 9] - assert stale.metadata._stored_document is not None + assert _store_needs_repair(stale) assert {p.name: p.read_bytes() for p in path.iterdir()} == documents @@ -903,7 +913,7 @@ async def failing_set(self: LocalStore, key: str, *args: Any, **kwargs: Any) -> with pytest.raises(OSError, match=f"cannot store {re.escape(doc_name)}"): operation(arr) - assert arr.metadata._stored_document is not None + assert _store_needs_repair(arr) assert (path / doc_name).read_bytes() == stored @@ -916,11 +926,11 @@ def test_write_without_stored_document(zarr_format: Literal[2, 3]) -> None: doc = _v2_doc([3], [0]) if zarr_format == 2 else _v3_doc([3], [0]) with pytest.warns(ZarrUserWarning, match="is read as"): array = zarr.Array(AsyncArray.from_dict(StorePath(store), doc)) - repaired = array.metadata._stored_document is not None + repaired = moves_chunks(doc, zarr_format) array[:] = [1, 2, 3] - assert (repaired, array.metadata._stored_document) == (True, None) + assert (repaired, _store_needs_repair(array)) == (True, False) np.testing.assert_array_equal(array[:], [1, 2, 3]) assert not [key for key in store._store_dict if key.endswith((".zarray", "zarr.json"))] @@ -979,7 +989,7 @@ def test_from_dict_keeps_values_that_are_not_json(codec: Any) -> None: doc: dict[str, Any] = _v3_doc([4], [4]) doc["codecs"] = [codec] metadata = ArrayV3Metadata.from_dict(doc) - assert metadata._stored_document is None + assert not moves_chunks(doc, 3) assert isinstance(metadata.codecs[0], ShardingCodec) @@ -1219,20 +1229,26 @@ def test_consolidated_repaired_member_attributes_kept_by_group_write( @pytest.mark.parametrize("zarr_format", [2, 3]) -def test_repaired_metadata_keeps_the_document_it_read(zarr_format: Literal[2, 3]) -> None: - """Metadata read from a document that had to be repaired keeps that document as it - was read, whatever becomes of the caller's dict (or the objects in it, which the - metadata's attributes may hold): consolidated metadata stores it as read.""" +def test_cache_keeps_the_document_it_read(zarr_format: Literal[2, 3]) -> None: + """A store's cache keeps a stored document that needs a repair as it was read, + whatever becomes of the caller's dict (or the objects in it), and consolidated + metadata stores a member given as stored exactly so.""" doc = _v2_doc([0], [0]) if zarr_format == 2 else _v3_doc([0], [0]) doc["attributes"] = {"k": [1]} read = json.loads(json.dumps(doc)) metadata_cls = ArrayV2Metadata if zarr_format == 2 else ArrayV3Metadata - metadata = metadata_cls.from_dict(doc) + metadata = metadata_cls.from_dict(json.loads(json.dumps(doc))) + cache = StoredDocuments() + cache.observe("a", doc, zarr_format) - cast("list[int]", metadata.attributes["k"]).append(2) + cast("list[int]", cast("dict[str, Any]", doc["attributes"])["k"]).append(2) doc["shape"] = [4] - assert ConsolidatedMetadata(metadata={"a": metadata}).to_dict()["metadata"] == {"a": read} + stored = cache.needing_repair("a", zarr_format) + assert stored is not None + assert stored == read + consolidated = ConsolidatedMetadata(metadata={"a": metadata}) + assert consolidated.to_dict(members_as_stored={"a": stored})["metadata"] == {"a": read} @pytest.mark.parametrize("zarr_format", [2, 3]) @@ -1246,7 +1262,7 @@ def test_empty_write_stores_no_metadata(tmp_path: Path, zarr_format: Literal[2, array[0:0] = np.empty(0, dtype="int16") - assert array.metadata._stored_document is not None + assert _store_needs_repair(array) assert {p.name: p.read_bytes() for p in path.iterdir()} == documents @@ -1653,3 +1669,134 @@ def test_delete_member_without_flag_keeps_group(tmp_path: Path) -> None: assert group.metadata.consolidated_metadata is not None assert sorted(group.metadata.consolidated_metadata.metadata) == ["n", "r"] assert {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} == stored + + +def _bind_opened(store: LocalStore) -> AnyArray: + return zarr.open_array(store=store, mode="r+") + + +def _bind_pickled(store: LocalStore) -> AnyArray: + revived: AnyArray = pickle.loads(pickle.dumps(zarr.open_array(store=store, mode="r+"))) + return revived + + +def _bind_derived(store: LocalStore) -> AnyArray: + return zarr.open_array(store=store, mode="r+").io.with_config({"order": "F"}).to_array() + + +def _bind_declared_same_layout(store: LocalStore) -> AnyArray: + return zarr.ArrayIO.build_v2(store, shape=(3,), dtype="int16", chunks=(1,)).to_array() + + +def _bind_declared_other_layout(store: LocalStore) -> AnyArray: + return zarr.ArrayIO.build_v2(store, shape=(3,), dtype="int16", chunks=(3,)).to_array() + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize( + ("bind", "stored_after"), + [ + (_bind_opened, [1]), + (_bind_pickled, [1]), + (_bind_derived, [1]), + (_bind_declared_same_layout, [1]), + (_bind_declared_other_layout, [0]), + ], + ids=["opened", "pickled", "derived", "declared-same-layout", "declared-other-layout"], +) +def test_write_stores_repair_however_bound( + tmp_path: Path, bind: Callable[[LocalStore], AnyArray], stored_after: list[int] +) -> None: + """A chunk write stores the repair of a stored document that needs one, however the + writing array was bound to the store: the store's cache answers, and a store + object that has not seen the document reads it. A declaration that lays out + chunks otherwise describes another array, so the document is left as stored.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, 2) + array = bind(LocalStore(path)) + + array[:] = [7, 8, 9] + + assert json.loads((path / ".zarray").read_text())["chunks"] == stored_after + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +def test_bindings_of_one_store_share_what_is_stored( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Arrays bound to one store object share its cache: once one has stored the repair, + another writes its chunks without reading a metadata document first.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, 2) + store = LocalStore(path) + first = zarr.open_array(store=store, mode="r+") + second = zarr.open_array(store=store, mode="r+") + first[:] = [7, 8, 9] + accesses = _record_store_access(monkeypatch) + + second[:] = [1, 2, 3] + + assert not [key for _, key in accesses if key.endswith((".zarray", ".zattrs"))] + np.testing.assert_array_equal(_open_strictly(path)[...], [1, 2, 3]) + + +def test_stored_documents_belongs_to_one_store_object() -> None: + """Each store object has its own cache, also when it equals another store, and the + cache is dropped with its store.""" + store, equal = MemoryStore(), MemoryStore() + assert store == equal + cache = stored_documents(store) + cache.observe("a", None, 3) + + assert stored_documents(store) is cache + assert not stored_documents(equal).knows("a", 3) + key = id(store) + del store, cache + gc.collect() + assert key not in _CACHES + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +def test_documents_of_each_format_are_recorded_apart(tmp_path: Path) -> None: + """A path may hold a document of each Zarr format (as `migrate_v3` leaves it before + `remove_metadata`): reading one does not stand for the other, so a v2 write still + stores the repair of its `.zarray` after the store object read the `zarr.json`.""" + path = tmp_path / "both.zarr" + _legacy_array(path, 2) + zarr.create_array(store=tmp_path / "v3.zarr", shape=(3,), chunks=(1,), dtype="int16") + (path / "zarr.json").write_bytes((tmp_path / "v3.zarr" / "zarr.json").read_bytes()) + store = LocalStore(path) + v2 = zarr.open_array(store=store, mode="r+", zarr_format=2) + zarr.open_array(store=store, mode="r+", zarr_format=3) + + v2[:] = [7, 8, 9] + + assert json.loads((path / ".zarray").read_text())["chunks"] == [1] + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +def test_unreadable_member_document_leaves_its_copy_as_stored(tmp_path: Path) -> None: + """A member document zarr fails to read (here a rectilinear grid while the flag is + off) is not recorded, so a group write keeps that member's consolidated copy as it + was stored instead of encoding the member anew.""" + path = tmp_path / "group.zarr" + zarr.open_group(path, mode="w") + zarr.create_array(store=path / "a", shape=(3, 6), chunks=(3, 3), dtype="int8") + _rewrite_doc( + path / "a", + 3, + lambda doc: doc["chunk_grid"]["configuration"].update(chunk_shape=[[1, 2], 3]), + ) + zarr.consolidate_metadata(path) + with zarr.config.set({"array.rectilinear_chunks": True}): + zarr.open_array(store=path / "a", mode="r+")[0, 0] = 1 + stored = json.dumps(_consolidated_member(path, 3, "a")) + store = LocalStore(path) + group = zarr.open_group(store, mode="r+", use_consolidated=True) + with pytest.raises(ValueError, match="experimental and disabled"): + list(zarr.open_group(store, mode="r+", use_consolidated=False).members()) + + group.attrs["y"] = 2 + + assert json.dumps(_consolidated_member(path, 3, "a")) == stored From a56367333783ded70903240d990d9c4b89b3e3ae Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 12:29:28 +0100 Subject: [PATCH 19/48] build: raise the typing-extensions floor to 4.16, as zarr-metadata requires zarr-metadata 0.6.0 requires typing_extensions>=4.16, so zarr's effective minimum was already 4.16 while it declared >=4.14, and the min_deps test environment's typing_extensions==4.14.* pin could not resolve. Declare the real floor and pin min_deps to 4.16.*. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- changes/419.feature.md | 2 +- pyproject.toml | 4 ++-- uv.lock | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index 597e3baf15..1a3df83588 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,6 +1,6 @@ Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. The declaring constructors, `build_v2` and `build_v3`, are synchronous — binding is pure — take a `Store` or `StorePath`, and take the fields of each format's metadata document, with nothing hiding the difference between the formats; the documents themselves are built by the module-level `build_array_metadata_v2` and `build_array_metadata_v3`. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. `Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. -The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata` — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. +The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata`, which raises its minimum `typing-extensions` to 4.16 — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. `with_metadata`, `with_store` and `with_config` return new bindings with one part replaced. `ArrayIO` responds to the global configuration the way `Array` always has: the codec pipeline implementation is captured when the binding is made, the runtime configuration defaults from `zarr.config`, and `build_v2`'s `order` defaults from `array.order`. Whether a stored array document needs a repair that moves chunks (such as a stored chunk size of 0) is now recorded per store object, in a cache of the metadata documents zarr has read and stored, instead of on the metadata read from it. Every array bound to one store object shares what any of them learned, derived bindings carry nothing, and the first chunk write through any binding stores the repair, including a pickled array or an `ArrayIO` declaration whose chunk layout matches the stored document; a store object that has not seen the document reads it first. diff --git a/pyproject.toml b/pyproject.toml index 067d2e7297..7007b2979c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ dependencies = [ # 0.16.6 and 0.17.1 refuse to unpickle untrusted data (GHSA-jqc8-qjfv-5m6r); 0.17.0 does not. 'numcodecs>=0.16.6,!=0.17.0', 'google-crc32c>=1.5', - 'typing_extensions>=4.14', + 'typing_extensions>=4.16', 'donfig>=0.8', 'msgspec>=0.19', 'zarr-metadata>=0.6.0', @@ -258,7 +258,7 @@ extra-dependencies = [ 'fsspec==2023.10.0', 's3fs==2023.10.0', 'universal_pathlib==0.2.0', - 'typing_extensions==4.14.*', + 'typing_extensions==4.16.*', 'donfig==0.8.*', 'obstore==0.5.*', 'msgspec==0.19.*', diff --git a/uv.lock b/uv.lock index 23012125e9..6064c46e87 100644 --- a/uv.lock +++ b/uv.lock @@ -3831,7 +3831,7 @@ requires-dist = [ { name = "obstore", marker = "extra == 'remote'", specifier = ">=0.5.1" }, { name = "packaging", specifier = ">=22.0" }, { name = "typer", marker = "extra == 'cli'" }, - { name = "typing-extensions", specifier = ">=4.14" }, + { name = "typing-extensions", specifier = ">=4.16" }, { name = "universal-pathlib", marker = "extra == 'optional'" }, { name = "zarr-metadata", specifier = ">=0.6.0" }, ] From f39623440a5da26552c7d1e7296046d7639bd69f Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 12:31:42 +0100 Subject: [PATCH 20/48] fix(metadata): carry a binding's stored-document entry through pickling A store's cache does not pickle with the store, so every unpickled binding (each task of a distributed write) read its metadata document again before its first chunk write. A pickled AsyncArrayIO now carries what its store's cache knew of its own document, restored into the new store object's cache unless that one already knows it, as the removed mark used to pickle with the metadata. The StoredDocuments docstring now says that a chunk write's own read is recorded as read. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 26 +++++++++++++++++++++++--- src/zarr/core/metadata/stored.py | 11 +++++++++-- tests/test_metadata/test_repair.py | 18 ++++++++++++++++++ 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 08acae943b..b858f63d00 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -181,6 +181,7 @@ from zarr.codecs.sharding import IndexLocation, ShardingCodec from zarr.core.buffer import Buffer from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar + from zarr.core.metadata.repair import ArrayDocument from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -6552,10 +6553,29 @@ def to_array(self) -> AnyAsyncArray: """An `AsyncArray` over this binding; both share one state.""" return AsyncArray._from_io(self) + def __getstate__(self) -> dict[str, object]: + # A store's cache does not pickle with the store: carry what it knows of this + # binding's document, so the unpickled binding does not read it again. + state = dict(self.__dict__) + cache = stored_documents(self._store_path.store) + path, zarr_format = self._store_path.path, self._zarr_format + if cache.knows(path, zarr_format): + state["_stored_entry"] = cache.needing_repair(path, zarr_format) + return state + + def __setstate__(self, state: dict[str, object]) -> None: + known = "_stored_entry" in state + entry = cast("ArrayDocument | None", state.pop("_stored_entry", None)) + self.__dict__.update(state) + if known: + stored_documents(self._store_path.store).restore( + self._store_path.path, self._zarr_format, entry + ) + def _derived_source(self) -> ArrayMetadataModel | ArrayMetadata: - """The source for a derived binding: a legacy object is copied, so the - repair mark and in-place attribute updates never alias across bindings; - a model is immutable and shared.""" + """The source for a derived binding: a legacy object is copied, so in-place + attribute updates never alias across bindings; a model is immutable and + shared.""" source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): return _copied_legacy(source) diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index 9569a141e0..6e8b93e96f 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -46,8 +46,9 @@ class StoredDocuments: document; a document zarr has not seen is unknown. Copies of array documents in consolidated metadata are recorded apart, by the array's path and format, only when reading them needs such a repair, so a group write can store them as they - were stored. zarr records a document it read only once it has read it as - metadata, so a document it could not read is unknown. + were stored. A document zarr reads to open an array is recorded only once it has + been read as metadata, so one it could not open is unknown; one a chunk write reads + is recorded as read. """ def __init__(self) -> None: @@ -69,6 +70,12 @@ def copy_needing_repair(self, path: str, zarr_format: ZarrFormat) -> ArrayDocume was stored, if reading it needs a repair that moves chunks.""" return self._copies.get((path, zarr_format)) + def restore(self, path: str, zarr_format: ZarrFormat, document: ArrayDocument | None) -> None: + """Record `document` as `needing_repair(path, zarr_format)` returned it from + another store object's cache (a binding carries it when pickled), unless this + cache already knows the document: what it learned itself is newer.""" + self._arrays.setdefault((path, zarr_format), document) + def observe(self, path: str, document: ArrayDocument | None, zarr_format: ZarrFormat) -> None: """Record that the array document of `zarr_format` at `path` is `document` (`None`: none).""" diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index bf9311299d..4c90d23efe 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -1800,3 +1800,21 @@ def test_unreadable_member_document_leaves_its_copy_as_stored(tmp_path: Path) -> group.attrs["y"] = 2 assert json.dumps(_consolidated_member(path, 3, "a")) == stored + + +def test_pickled_array_carries_what_its_store_knows( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A pickled array carries what its store's cache knows of its document, so the + unpickled array writes its chunks without reading a metadata document first, as + each task of a distributed write unpickles a store of its own.""" + path = tmp_path / "valid.zarr" + store = LocalStore(path) + zarr.create_array(store=store, shape=(3,), chunks=(1,), dtype="int16") + revived: AnyArray = pickle.loads(pickle.dumps(zarr.open_array(store=store, mode="r+"))) + accesses = _record_store_access(monkeypatch) + + revived[:] = [1, 2, 3] + + assert not [key for _, key in accesses if key.endswith("zarr.json")] + np.testing.assert_array_equal(_open_strictly(path)[...], [1, 2, 3]) From 6a397374777d06097c4d287e14caa8cc41b534e6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 12:40:28 +0100 Subject: [PATCH 21/48] refactor(array): set a binding's fields in one place `__init__` and `_derive` each listed every field of an AsyncArrayIO; both now go through `_bind`, so a field added later cannot be missed by one of them. `_copied_legacy` drops a guard that never fired and says what its copy isolates: top-level attribute updates. Tests cover with_store and with_metadata from a legacy source, a derived binding keeping its captured pipeline while the global configuration names one that cannot load, and StoredDocuments.restore keeping what the cache learned itself. The installation guide gives the new typing_extensions floor. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- docs/user-guide/installation.md | 2 +- src/zarr/core/array.py | 47 ++++++++++++++++-------------- src/zarr/core/metadata/stored.py | 2 +- tests/test_array_io.py | 20 +++++++++---- tests/test_metadata/test_repair.py | 14 +++++++++ 5 files changed, 56 insertions(+), 29 deletions(-) diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index a7487e83c8..fc5e526504 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -9,7 +9,7 @@ Required dependencies include: - [numpy](https://numpy.org) (2.0 or later) - [numcodecs](https://numcodecs.readthedocs.io/en/stable/) (0.14 or later) - [google-crc32c](https://github.com/googleapis/python-crc32c) (1.5 or later) -- [typing_extensions](https://typing-extensions.readthedocs.io/en/latest/) (4.14 or later) +- [typing_extensions](https://typing-extensions.readthedocs.io/en/latest/) (4.16 or later) - [donfig](https://donfig.readthedocs.io/en/latest/) (0.8 or later) ## pip diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index b858f63d00..b7e02e5671 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -6241,11 +6241,11 @@ async def _delete_key(store_path: StorePath, key: str) -> None: def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: - """A copy of a legacy metadata object that shares no mutable state with it: - the attributes dict is copied.""" + """A copy of a legacy metadata object with its own attributes dict, so an + attribute update through one binding (which replaces top-level keys in + place) never reaches another.""" copied = copy.copy(source) - if source.attributes is not None: - object.__setattr__(copied, "attributes", dict(source.attributes)) + object.__setattr__(copied, "attributes", dict(source.attributes)) return copied @@ -6336,23 +6336,34 @@ def __init__( store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: - self._store_path = store_path - self._config = parse_array_config(config) + source: ArrayMetadataModel | ArrayMetadata + if isinstance(metadata, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): + source = metadata + else: + # A legacy resolved object passes through; raw JSON is read eagerly, + # as repaired when it needs that, warning as `AsyncArray` always has. + source = parse_array_metadata(metadata, str(store_path)) # The pipeline implementation is chosen when the binding is made, from # the global configuration, exactly as `AsyncArray` has always chosen # it; the pipeline itself is built on first use. - self._pipeline_class = get_pipeline_class() + self._bind(source, store_path, parse_array_config(config), get_pipeline_class()) + + def _bind( + self, + source: ArrayMetadataModel | ArrayMetadata, + store_path: StorePath, + config: ArrayConfig, + pipeline_class: type[CodecPipeline], + ) -> None: + self._store_path = store_path + self._config = config + self._pipeline_class = pipeline_class + self._source = source self._metadata_model_cache = None self._resolved_cache = None self._chunk_grid_cache = None self._chunk_key_encoding_cache = None self._codec_pipeline_cache = None - if isinstance(metadata, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): - self._source = metadata - else: - # A legacy resolved object passes through; raw JSON is read eagerly, - # as repaired when it needs that, warning as `AsyncArray` always has. - self._source = parse_array_metadata(metadata, str(store_path)) @classmethod def build_v2( @@ -6589,15 +6600,7 @@ def _derive( ) -> Self: """A new binding over `source`, keeping this one's captured pipeline.""" derived = object.__new__(type(self)) - derived._store_path = store_path - derived._config = config - derived._pipeline_class = self._pipeline_class - derived._source = source - derived._metadata_model_cache = None - derived._resolved_cache = None - derived._chunk_grid_cache = None - derived._chunk_key_encoding_cache = None - derived._codec_pipeline_cache = None + derived._bind(source, store_path, config, self._pipeline_class) return derived def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index 6e8b93e96f..a4dfddf476 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -73,7 +73,7 @@ def copy_needing_repair(self, path: str, zarr_format: ZarrFormat) -> ArrayDocume def restore(self, path: str, zarr_format: ZarrFormat, document: ArrayDocument | None) -> None: """Record `document` as `needing_repair(path, zarr_format)` returned it from another store object's cache (a binding carries it when pickled), unless this - cache already knows the document: what it learned itself is newer.""" + cache already knows the document: it keeps what it learned itself.""" self._arrays.setdefault((path, zarr_format), document) def observe(self, path: str, document: ArrayDocument | None, zarr_format: ZarrFormat) -> None: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 734fe60f09..e4f98f43ed 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -253,8 +253,12 @@ def test_binding_follows_the_global_config() -> None: ): v3 = ArrayIO.build_v3(MemoryStore(), shape=(2,), data_type="uint8", chunk_grid=(2,)) v2 = ArrayIO.build_v2(MemoryStore(), shape=(2,), dtype="uint8", chunks=(2,)) - # outside the block, the choices made inside it hold + # outside the block, the choices made inside it hold, also for a binding + # derived while the global configuration names a pipeline that cannot load assert isinstance(v3.to_array().async_array.codec_pipeline, FusedCodecPipeline) + with zarr.config.set({"codec_pipeline.path": "nonexistent.module.Pipeline"}): + derived = v3.with_config({"order": "C"}) + assert isinstance(derived.to_array().async_array.codec_pipeline, FusedCodecPipeline) legacy = v2.metadata_legacy assert isinstance(legacy, ArrayV2Metadata) assert legacy.order == "F" @@ -287,12 +291,18 @@ async def test_array_state_propagates_to_the_shared_binding() -> None: legacy_arr = await zarr.api.asynchronous.create_array( MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,) ) - legacy_derived = legacy_arr.io.with_config({"order": "F"}) - _ = legacy_derived.metadata # fill the model cache before the mutation + legacy_derivatives = [ + legacy_arr.io.with_config({"order": "F"}), + legacy_arr.io.with_store(MemoryStore()), + legacy_arr.io.with_metadata(legacy_arr.io.metadata_legacy), + ] + for legacy_derived in legacy_derivatives: + _ = legacy_derived.metadata # fill the model cache before the mutation await legacy_arr.update_attributes({"b": 2}) assert dict(legacy_arr.io.metadata.attributes) == {"b": 2} - assert dict(legacy_derived.metadata_legacy.attributes) == {} - assert dict(legacy_derived.metadata.attributes) == {} + for legacy_derived in legacy_derivatives: + assert dict(legacy_derived.metadata_legacy.attributes) == {} + assert dict(legacy_derived.metadata.attributes) == {} async def test_store_metadata_of_model_without_attributes_clears_zattrs() -> None: diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index 4c90d23efe..dff254f211 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -1818,3 +1818,17 @@ def test_pickled_array_carries_what_its_store_knows( assert not [key for _, key in accesses if key.endswith("zarr.json")] np.testing.assert_array_equal(_open_strictly(path)[...], [1, 2, 3]) + + +def test_restore_keeps_what_the_cache_learned_itself() -> None: + """A carried entry fills only a document the cache does not know: what the + cache recorded itself, such as a repair it stored, stays.""" + cache = StoredDocuments() + broken = _v2_doc([3], [0]) + cache.observe("a", None, 2) + + cache.restore("a", 2, broken) + cache.restore("b", 2, broken) + + assert cache.needing_repair("a", 2) is None + assert cache.needing_repair("b", 2) == broken From 203d21be8e16e9e37f14a7d018ad2f07f9cdb823 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 12:49:30 +0100 Subject: [PATCH 22/48] fix(array): read a with_metadata legacy object as the constructor reads it with_metadata copied a legacy object without the reading parse_array_metadata gives it, so ArrayV2Metadata with a chunk size of 0 bound unrepaired there while the constructor bound it repaired. A test binds such metadata through the constructor, build_v2 and with_metadata. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 5 +++-- tests/test_array_io.py | 28 +++++++++++++++++++++++++++- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index b7e02e5671..8e0056fdb3 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -6605,9 +6605,10 @@ def _derive( def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: """This binding with `metadata` in place of its own. Nothing is stored; - a legacy `metadata` object is copied.""" + a legacy `metadata` object is read as the constructor reads it, then + copied.""" if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): - metadata = _copied_legacy(metadata) + metadata = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) return self._derive(metadata, self._store_path, self._config) def with_store(self, store: Store | StorePath) -> Self: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index e4f98f43ed..08418fcbfd 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -15,8 +15,9 @@ from zarr.core.buffer import default_buffer_prototype from zarr.core.codec_pipeline import FusedCodecPipeline from zarr.core.metadata import ArrayV2Metadata +from zarr.dtype import Int16 from zarr.errors import ArrayNotFoundError, UnknownCodecError -from zarr.storage import MemoryStore +from zarr.storage import MemoryStore, StorePath if TYPE_CHECKING: from zarr.core.common import ZarrFormat @@ -322,3 +323,28 @@ async def test_store_metadata_of_model_without_attributes_clears_zattrs() -> Non assert "foo/.zattrs" not in await _store_keys(store) reread = await AsyncArrayIO.open(store, name="foo") assert dict(reread.metadata.attributes) == {} + + +def _bind_zero_chunk_by_constructor(metadata: ArrayV2Metadata) -> AsyncArrayIO: + return AsyncArrayIO(metadata, StorePath(MemoryStore())) + + +def _bind_zero_chunk_by_build(metadata: ArrayV2Metadata) -> AsyncArrayIO: + return AsyncArrayIO.build_v2(MemoryStore(), shape=metadata.shape, dtype="int16", chunks=(0,)) + + +def _bind_zero_chunk_by_with_metadata(metadata: ArrayV2Metadata) -> AsyncArrayIO: + declared = AsyncArrayIO.build_v2(MemoryStore(), shape=(3,), dtype="int16", chunks=(3,)) + return declared.with_metadata(metadata) + + +@pytest.mark.parametrize( + "bind", + [_bind_zero_chunk_by_constructor, _bind_zero_chunk_by_build, _bind_zero_chunk_by_with_metadata], + ids=["constructor", "build_v2", "with_metadata"], +) +def test_zero_chunk_v2_metadata_binds_repaired(bind: Any) -> None: + """`ArrayV2Metadata` with a chunk size of 0 is read as repaired however it + reaches a binding, as `parse_array_metadata` reads it.""" + metadata = ArrayV2Metadata(shape=(3,), chunks=(0,), dtype=Int16(), fill_value=0, order="C") + assert bind(metadata).metadata_legacy.chunks == (1,) From e886f9ec95ff7b6eefaa3d5da404f862c3194e68 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 13:27:00 +0100 Subject: [PATCH 23/48] docs(array): say how fresh the stored-document repair check is The repair check before a chunk write is only as fresh as the store object's cache: once another binding of that store object has read or stored a newer document, an array opened before then writes its chunks under its own metadata without the check, as any stale array handle does. This is the accepted behavior of the per-store cache; the removed metadata mark had re-read the store for such handles. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 8e0056fdb3..a06283d52b 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -5024,6 +5024,13 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat possibly by another implementation), it is left as written; if there is none, there is nothing to repair. + The check is only as fresh as the cache. Once a binding of the same store object + has read or stored a newer document (for example, `create_array(..., + overwrite=True)` on that store), the cache no longer records the repair, so an + array opened before that writes its chunks under its own metadata unchecked, as + any array whose document another writer replaced does: zarr does not detect a + stale array handle. + Storing the same repair twice is harmless, so handles that write chunks concurrently need no coordination. The read and the store are not one atomic step, though: a metadata write by another handle between them (a resize, an From c42f8159e6b082856fa38a469f2ff1538405796b Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 14:15:39 +0100 Subject: [PATCH 24/48] fix(array): resolve v2 model attributes as plain JSON and check v2 object codecs A v2 binding read as a zarr-metadata model copied only the top level of the model's read-only attributes into its legacy metadata, so nested lists and objects stayed tuples and mapping proxies and every metadata write through it (update_attributes, resize, append) failed to encode. The attributes now come from the model's own `.zattrs` document. build_array_metadata_v2 (and build_v2) accepted a variable-length string data type with no object codec, storing metadata nothing can read; it now makes the check create_array makes, shared as `_check_object_codec_v2`. store_metadata's docstring says what it replaces: the documents of its own format, leaving the other format's, a group's and created parents in place. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 45 ++++++++++++++++++++++++++++++------------ tests/test_array_io.py | 27 ++++++++++++++++++++++++- 2 files changed, 58 insertions(+), 14 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index a06283d52b..fb98057470 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -4706,7 +4706,7 @@ def build_array_metadata_v2( # An explicit None is kept: the field is nullable in Zarr format 2. fill_value_parsed = fill_value - return ArrayV2Metadata( + metadata = ArrayV2Metadata( shape=parse_shapelike(shape), dtype=zdtype, chunks=parse_shapelike(chunks), @@ -4717,6 +4717,8 @@ def build_array_metadata_v2( filters=filters, attributes=attributes, ) + _check_object_codec_v2(zdtype, metadata.filters, metadata.compressor) + return metadata def build_array_metadata_v3( @@ -5343,18 +5345,28 @@ def _parse_chunk_encoding_v2( ) raise TypeError(msg) _filters = parse_filters(filters) + _check_object_codec_v2(dtype, _filters, _compressor) + return _filters, _compressor + + +def _check_object_codec_v2( + dtype: ZDType[TBaseDType, TBaseScalar], + filters: tuple[Numcodec, ...] | None, + compressor: Numcodec | None, +) -> None: + """Raise if `dtype` needs an object codec that neither `filters` nor `compressor` holds.""" if isinstance(dtype, HasObjectCodec): # check the filters and the compressor for the object codec required for this data type - if _filters is None: - if _compressor is None: + if filters is None: + if compressor is None: object_codec_id = None else: - object_codec_id = get_object_codec_id((_compressor.get_config(),)) + object_codec_id = get_object_codec_id((compressor.get_config(),)) else: object_codec_id = get_object_codec_id( ( - *[f.get_config() for f in _filters], - _compressor.get_config() if _compressor is not None else None, + *[f.get_config() for f in filters], + compressor.get_config() if compressor is not None else None, ) ) if object_codec_id is None: @@ -5370,7 +5382,6 @@ def _parse_chunk_encoding_v2( "this array. " ) raise ValueError(msg) - return _filters, _compressor def _parse_chunk_encoding_v3( @@ -6452,7 +6463,11 @@ def metadata_legacy(self) -> ArrayMetadata: if self._resolved_cache is None: data = _metadata_model_json(source) if isinstance(source, ZarrV2ArrayMetadata): - data["attributes"] = dict(source.attributes) + # The model's attributes are read-only JSON; its `.zattrs` is plain. + zattrs = {str(key): value for key, value in source.to_key_value().items()} + data["attributes"] = ( + json.loads(zattrs[ZATTRS_JSON]) if ZATTRS_JSON in zattrs else {} + ) self._resolved_cache = parse_array_metadata(data, str(self._store_path)) return self._resolved_cache @@ -6634,11 +6649,15 @@ def with_config(self, config: ArrayConfigLike) -> Self: return self._derive(self._derived_source(), self._store_path, new_config) async def store_metadata(self, *, ensure_parents: bool = False) -> None: - """Store the metadata, unconditionally replacing any stored metadata here. - - Chunks of a previous array at this location are kept; `erase_chunks` - removes them. With `ensure_parents`, missing parent groups are stored - first. + """Store the metadata documents of this array's Zarr format, unconditionally + replacing those stored here. + + Documents of the other format, or of a group, at this location are left + as they are, as `erase_metadata` leaves them; `verify` reports when a + format 3 document would be opened instead. Chunks of a previous array at + this location are kept; `erase_chunks` removes them. With + `ensure_parents`, missing parent groups are stored first; they are nodes + of their own, which `erase_metadata` leaves in place. """ source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 08418fcbfd..4a0c6c5f21 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -20,7 +20,7 @@ from zarr.storage import MemoryStore, StorePath if TYPE_CHECKING: - from zarr.core.common import ZarrFormat + from zarr.core.common import JSON, ZarrFormat METADATA_DOCUMENTS = ("zarr.json", ".zarray", ".zattrs", ".zgroup") @@ -348,3 +348,28 @@ def test_zero_chunk_v2_metadata_binds_repaired(bind: Any) -> None: reaches a binding, as `parse_array_metadata` reads it.""" metadata = ArrayV2Metadata(shape=(3,), chunks=(0,), dtype=Int16(), fill_value=0, order="C") assert bind(metadata).metadata_legacy.chunks == (1,) + + +async def test_v2_model_binding_keeps_attributes_writable() -> None: + """A v2 binding read as a model resolves its attributes as plain, mutable + JSON, so writes through it that re-encode the attributes succeed.""" + store = MemoryStore() + attributes: dict[str, JSON] = {"k": [1, 2], "d": {"n": None}} + aio = AsyncArrayIO.build_v2( + store, name="foo", shape=(4,), dtype="uint8", chunks=(2,), attributes=attributes + ) + await aio.store_metadata() + arr = (await AsyncArrayIO.open(store, name="foo")).to_array() + + assert arr.metadata.attributes == attributes + await arr.update_attributes({"new": 1}) + await arr.resize((6,)) + + reread = await open_array(store=store, path="foo") + assert reread.metadata.attributes == {**attributes, "new": 1} + assert reread.shape == (6,) + + +def test_build_v2_vlen_without_object_codec_raises() -> None: + with pytest.raises(ValueError, match="requires the numcodecs.VLenUTF8 codec"): + AsyncArrayIO.build_v2(MemoryStore(), shape=(4,), dtype=str, chunks=(2,)) From 416c5a9a15d2b690bbe963ddb1459f9dc037034e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 14:30:58 +0100 Subject: [PATCH 25/48] fix(array): keep AsyncArray's dataclass surface and read nothing new AsyncArray's only dataclass field had become the wrapped binding, so dataclasses.fields/asdict/replace, the generated equality and hashing, the pickled state and dask tokens (dask tokenizes dataclasses by their fields) all changed, and copy.copy shared one binding. AsyncArray now declares the fields it always had (metadata, store_path, codec_pipeline, _chunk_grid, config), each a property over the binding that object.__setattr__ can still set, and pickles as that field state, so pickles load across versions in both directions and copies are independent. _save_metadata keeps its signature, and with_config rebuilds the array as it always did, re-reading the global pipeline choice. A chunk write no longer reads a document the store's cache has not seen: unknown means nothing is known to repair, as unmarked metadata always did, so writes through consolidated members, code-built arrays and their unpickled copies make no new metadata reads. A pickled binding carries a consolidated copy's entry as well as its document's. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- changes/419.feature.md | 2 +- src/zarr/core/array.py | 198 +++++++++++++++++++++-------- src/zarr/core/metadata/stored.py | 6 + tests/test_array.py | 80 ++++++++++++ tests/test_metadata/test_repair.py | 11 +- 5 files changed, 240 insertions(+), 57 deletions(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index 1a3df83588..9a76e2b323 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -3,4 +3,4 @@ Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arr The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata`, which raises its minimum `typing-extensions` to 4.16 — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. `with_metadata`, `with_store` and `with_config` return new bindings with one part replaced. `ArrayIO` responds to the global configuration the way `Array` always has: the codec pipeline implementation is captured when the binding is made, the runtime configuration defaults from `zarr.config`, and `build_v2`'s `order` defaults from `array.order`. -Whether a stored array document needs a repair that moves chunks (such as a stored chunk size of 0) is now recorded per store object, in a cache of the metadata documents zarr has read and stored, instead of on the metadata read from it. Every array bound to one store object shares what any of them learned, derived bindings carry nothing, and the first chunk write through any binding stores the repair, including a pickled array or an `ArrayIO` declaration whose chunk layout matches the stored document; a store object that has not seen the document reads it first. +Whether a stored array document needs a repair that moves chunks (such as a stored chunk size of 0) is now recorded per store object, in a cache of the metadata documents zarr has read and stored, instead of on the metadata read from it. Every array bound to one store object shares what any of them learned, derived bindings carry nothing, and the first chunk write through any binding that knows the document, including a pickled one, stores the repair; nothing is read that was not read before. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index fb98057470..67c1d130a5 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -7,12 +7,14 @@ import warnings from asyncio import gather from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from itertools import starmap from logging import getLogger from typing import ( TYPE_CHECKING, Any, + ClassVar, + Final, Literal, Required, TypedDict, @@ -421,7 +423,18 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: The runtime configuration of the array. """ - _io: AsyncArrayIO + # The dataclass fields are the ones this class has always had, so equality, + # hashing, `dataclasses.replace`/`fields`/`asdict`, pickling and dask tokens + # see the same array they always did. Each reads (and, under + # `object.__setattr__`, writes) the wrapped binding; see the properties + # attached after the class. + metadata: T_ArrayMetadata + store_path: StorePath + codec_pipeline: CodecPipeline = field(init=False) + _chunk_grid: ChunkGrid = field(init=False) + config: ArrayConfig + # Not a dataclass field: the binding every field reads. + _io: ClassVar[AsyncArrayIO] @overload def __init__( @@ -462,29 +475,32 @@ def io(self) -> AsyncArrayIO: """The explicit-IO layer this array delegates to; both share one state.""" return self._io - @property - def metadata(self) -> T_ArrayMetadata: - """The in-memory metadata of this array.""" - return cast("T_ArrayMetadata", self._io.metadata_legacy) - - @property - def store_path(self) -> StorePath: - """The location in a store this array is bound to.""" - return self._io.store_path - - @property - def config(self) -> ArrayConfig: - """The runtime configuration of this array.""" - return self._io.config - - @property - def codec_pipeline(self) -> CodecPipeline: - """The codec pipeline used for encoding and decoding chunks.""" - return self._io._resolved_codec_pipeline + def __getstate__(self) -> dict[str, object]: + # The pickled state is the fields', as it has always been, so a pickle + # loads in either version; what the store's cache knows rides along. + state: dict[str, object] = { + "metadata": self.metadata, + "store_path": self.store_path, + "codec_pipeline": self.codec_pipeline, + "_chunk_grid": self._chunk_grid, + "config": self.config, + } + state.update(self._io._carried_cache_entry()) + return state - @property - def _chunk_grid(self) -> ChunkGrid: - return self._io._chunk_grid + def __setstate__(self, state: dict[str, Any]) -> None: + metadata = state["metadata"] + pipeline = state["codec_pipeline"] + io = object.__new__(AsyncArrayIO) + io._bind(metadata, state["store_path"], state["config"], type(pipeline)) + io._codec_pipeline_cache = pipeline + io._chunk_grid_cache = state["_chunk_grid"] + # A pickle from before the store's cache held this carries it on the metadata. + legacy_mark = metadata.__dict__.pop("_stored_document", None) + if legacy_mark is not None: + state = {**state, "_stored_entry": legacy_mark} + io._restore_cache_entry(state) + object.__setattr__(self, "_io", io) @classmethod async def _create( @@ -1249,11 +1265,13 @@ def with_config(self, config: ArrayConfigLike) -> Self: ------- A new Array """ - derived = self._io.with_config(config) - # Build the pipeline now, so construction warnings fire here, as they - # do in __init__. - _ = derived._resolved_codec_pipeline - return cast("Self", type(self)._from_io(derived)) + if isinstance(config, ArrayConfig): + new_config = config + else: + # Merge new config with existing config, so missing keys are inherited + # from the current array rather than from global defaults + new_config = ArrayConfig(**{**self.config.to_dict(), **config}) # type: ignore[arg-type] + return type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) async def nchunks_initialized(self) -> int: """ @@ -1622,9 +1640,13 @@ async def get_coordinate_selection( out_array = cast("NDArrayLikeOrScalar", np.array(out_array).reshape(indexer.sel_shape)) return out_array - async def _save_metadata(self, ensure_parents: bool = False) -> None: - """Store this array's metadata documents, delegating to the wrapped `AsyncArrayIO`.""" - await self._io.store_metadata(ensure_parents=ensure_parents) + async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None: + """Store `metadata` as this array's own documents: through the wrapped + `AsyncArrayIO` when it is this array's metadata.""" + if metadata is self.metadata: + await self._io.store_metadata(ensure_parents=ensure_parents) + else: + await save_metadata(self.store_path, metadata, ensure_parents=ensure_parents) async def _set_selection( self, @@ -1866,6 +1888,60 @@ def _info( ) +def _get_metadata(self: AsyncArray[Any]) -> ArrayMetadata: + """The in-memory metadata of this array.""" + return self._io.metadata_legacy + + +def _set_metadata(self: AsyncArray[Any], value: ArrayMetadata) -> None: + self._io._rebind_resolved(value) + + +def _get_store_path(self: AsyncArray[Any]) -> StorePath: + """The location in a store this array is bound to.""" + return self._io.store_path + + +def _set_store_path(self: AsyncArray[Any], value: StorePath) -> None: + self._io._store_path = value + + +def _get_config(self: AsyncArray[Any]) -> ArrayConfig: + """The runtime configuration of this array.""" + return self._io.config + + +def _set_config(self: AsyncArray[Any], value: ArrayConfig) -> None: + self._io._config = value + + +def _get_codec_pipeline(self: AsyncArray[Any]) -> CodecPipeline: + """The codec pipeline used for encoding and decoding chunks.""" + return self._io._resolved_codec_pipeline + + +def _set_codec_pipeline(self: AsyncArray[Any], value: CodecPipeline) -> None: + self._io._codec_pipeline_cache = value + + +def _get_chunk_grid(self: AsyncArray[Any]) -> ChunkGrid: + return self._io._chunk_grid + + +def _set_chunk_grid(self: AsyncArray[Any], value: ChunkGrid) -> None: + self._io._chunk_grid_cache = value + + +# The dataclass fields of `AsyncArray` read the wrapped binding. They are +# attached once the dataclass is made, so `@dataclass` sees fields, not +# defaults; the setters keep `object.__setattr__(array, name, value)` working. +setattr(AsyncArray, "metadata", property(_get_metadata, _set_metadata)) # noqa: B010 +setattr(AsyncArray, "store_path", property(_get_store_path, _set_store_path)) # noqa: B010 +setattr(AsyncArray, "config", property(_get_config, _set_config)) # noqa: B010 +setattr(AsyncArray, "codec_pipeline", property(_get_codec_pipeline, _set_codec_pipeline)) # noqa: B010 +setattr(AsyncArray, "_chunk_grid", property(_get_chunk_grid, _set_chunk_grid)) # noqa: B010 + + # TODO: Array can be a frozen data class again once property setters (e.g. shape) are removed @dataclass(frozen=False) class Array[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: @@ -5015,7 +5091,9 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat The store's cache (see `zarr.core.metadata.stored`) says whether the document needs such a repair. While it has not seen the document, a consolidated copy it has seen stands for it, as consolidated metadata does for every reader; with - neither, the document is read. Only a document whose repair + neither, nothing is known to need a repair, and nothing is read: metadata built + in code, or bound to a store that never read the document, writes its chunks + as any array does. Only a document whose repair lays out chunks as `metadata` does is acted on: one that lays them out otherwise describes another array, and chunks written under `metadata` are its caller's declaration, as over any document. The document is read again, because the store @@ -5048,10 +5126,6 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat stored = cache.needing_repair(path, zarr_format) else: stored = cache.copy_needing_repair(path, zarr_format) - if stored is None: - documents = await read_documents(store_path, keys) - cache.observe(path, stored_array_document(documents, zarr_format), zarr_format) - stored = cache.needing_repair(path, zarr_format) if stored is None or _chunk_layout(read_repaired(stored, zarr_format)) != _chunk_layout( metadata ): @@ -6225,7 +6299,7 @@ async def _update_attributes( array._io._rebind_resolved(array.metadata, array._chunk_grid) # Write new metadata - await array._save_metadata() + await array._save_metadata(array.metadata) return array @@ -6258,6 +6332,10 @@ async def _delete_key(store_path: StorePath, key: str) -> None: await (store_path / key).delete() +_CARRIED_KEYS: Final = ("_stored_entry", "_stored_copy") +"""The keys a pickled binding carries its store cache's entry under.""" + + def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: """A copy of a legacy metadata object with its own attributes dict, so an attribute update through one binding (which replaces top-level keys in @@ -6334,7 +6412,10 @@ class AsyncArrayIO: `setitem` resolve them into the legacy classes (`metadata_legacy`) on first use. The bound metadata is assumed, never checked against the store, and an array opened from metadata that needed a repair stores that - repair before its first chunk write. `AsyncArray` wraps an `AsyncArrayIO` + repair before its first chunk write. The codec pipeline implementation is + chosen from the global configuration when the binding is made; the + pipeline itself, with any options it reads from the configuration, is + built at the first data operation. `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it. """ @@ -6587,23 +6668,38 @@ def to_array(self) -> AnyAsyncArray: return AsyncArray._from_io(self) def __getstate__(self) -> dict[str, object]: - # A store's cache does not pickle with the store: carry what it knows of this - # binding's document, so the unpickled binding does not read it again. state = dict(self.__dict__) - cache = stored_documents(self._store_path.store) - path, zarr_format = self._store_path.path, self._zarr_format - if cache.knows(path, zarr_format): - state["_stored_entry"] = cache.needing_repair(path, zarr_format) + state.update(self._carried_cache_entry()) return state def __setstate__(self, state: dict[str, object]) -> None: - known = "_stored_entry" in state - entry = cast("ArrayDocument | None", state.pop("_stored_entry", None)) + state = dict(state) + carried: dict[str, object] = {k: state.pop(k) for k in _CARRIED_KEYS if k in state} self.__dict__.update(state) - if known: - stored_documents(self._store_path.store).restore( - self._store_path.path, self._zarr_format, entry - ) + self._restore_cache_entry(carried) + + def _carried_cache_entry(self) -> dict[str, object]: + """What this binding's store cache knows of its document, to carry in a + pickle: a store's cache does not pickle with the store, and the unpickled + binding's first chunk write must know whether to store a repair first.""" + cache = stored_documents(self._store_path.store) + path, zarr_format = self._store_path.path, self._zarr_format + if cache.knows(path, zarr_format): + return {"_stored_entry": cache.needing_repair(path, zarr_format)} + copy_entry = cache.copy_needing_repair(path, zarr_format) + if copy_entry is not None: + return {"_stored_copy": copy_entry} + return {} + + def _restore_cache_entry(self, state: Mapping[str, object]) -> None: + """Record what `_carried_cache_entry` carried in this binding's store cache.""" + cache = stored_documents(self._store_path.store) + path, zarr_format = self._store_path.path, self._zarr_format + if "_stored_entry" in state: + entry = cast("ArrayDocument | None", state["_stored_entry"]) + cache.restore(path, zarr_format, entry) + elif "_stored_copy" in state: + cache.restore_copy(path, zarr_format, cast("ArrayDocument", state["_stored_copy"])) def _derived_source(self) -> ArrayMetadataModel | ArrayMetadata: """The source for a derived binding: a legacy object is copied, so in-place diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index a4dfddf476..3906680b95 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -76,6 +76,12 @@ def restore(self, path: str, zarr_format: ZarrFormat, document: ArrayDocument | cache already knows the document: it keeps what it learned itself.""" self._arrays.setdefault((path, zarr_format), document) + def restore_copy(self, path: str, zarr_format: ZarrFormat, document: ArrayDocument) -> None: + """Record a consolidated copy as `copy_needing_repair` returned it from another + store object's cache, unless this cache already knows the array's document.""" + if not self.knows(path, zarr_format): + self._copies.setdefault((path, zarr_format), document) + def observe(self, path: str, document: ArrayDocument | None, zarr_format: ZarrFormat) -> None: """Record that the array document of `zarr_format` at `path` is `document` (`None`: none).""" diff --git a/tests/test_array.py b/tests/test_array.py index 3e4835e686..0d2345922a 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2637,3 +2637,83 @@ async def test_create_array_huge_chunk_count() -> None: arr = await create_array(store={}, shape=(2**62,), chunks=(1,), dtype="int32") assert arr.shape == (2**62,) assert arr.chunks == (1,) + + +def test_async_array_keeps_its_dataclass_surface() -> None: + """`AsyncArray` wraps an `AsyncArrayIO` but keeps the dataclass surface it has + always had: its fields, `replace`, `asdict`, equality, (un)hashability, pickled + state, swapping a field with `object.__setattr__`, and independent copies.""" + import copy + import pickle + + from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline + + arr = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + aa = arr.async_array + + assert [f.name for f in dataclasses.fields(aa)] == [ + "metadata", + "store_path", + "codec_pipeline", + "_chunk_grid", + "config", + ] + assert set(dataclasses.asdict(aa)) == { + "metadata", + "store_path", + "codec_pipeline", + "_chunk_grid", + "config", + } + resized = dataclasses.replace(aa, metadata=dataclasses.replace(aa.metadata, shape=(8,))) + assert isinstance(resized, type(aa)) + assert (resized.shape, aa.shape) == ((8,), (4,)) + + with pytest.raises(TypeError, match="unhashable"): + hash(aa) + with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + fused = aa.with_config({}) + assert isinstance(fused.codec_pipeline, FusedCodecPipeline) + assert fused != aa + + swapped = copy.copy(aa) + object.__setattr__(swapped, "codec_pipeline", fused.codec_pipeline) + assert isinstance(swapped.codec_pipeline, FusedCodecPipeline) + assert isinstance(aa.codec_pipeline, BatchedCodecPipeline) + + copied = copy.copy(aa) + sync(copied.resize((6,))) + assert (copied.shape, aa.shape) == ((6,), (4,)) + + state = aa.__getstate__() + assert {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config"} <= set(state) + revived = pickle.loads(pickle.dumps(aa)) + assert revived == aa + + +def test_async_array_loads_a_pickled_field_state() -> None: + """A pickle of an `AsyncArray` holds its field state, as it always has, so an + array pickled before it wrapped an `AsyncArrayIO` loads.""" + arr = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + aa = arr.async_array + old_state = { + "metadata": aa.metadata, + "store_path": aa.store_path, + "codec_pipeline": aa.codec_pipeline, + "_chunk_grid": aa._chunk_grid, + "config": aa.config, + } + revived = object.__new__(type(aa)) + revived.__setstate__(old_state) + assert revived == aa + np.testing.assert_array_equal(sync(revived.getitem(slice(None))), np.zeros(4, dtype="uint8")) + + +def test_dask_token_of_an_array_is_deterministic() -> None: + dask_base = pytest.importorskip("dask.base") + store = MemoryStore() + arr = zarr.create_array(store, shape=(4,), dtype="uint8", chunks=(2,)) + before = dask_base.tokenize(arr) + _ = arr[:] + assert dask_base.tokenize(arr) == before + assert dask_base.tokenize(zarr.open_array(store)) == before diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index dff254f211..edc4aec728 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -1699,7 +1699,7 @@ def _bind_declared_other_layout(store: LocalStore) -> AnyArray: (_bind_opened, [1]), (_bind_pickled, [1]), (_bind_derived, [1]), - (_bind_declared_same_layout, [1]), + (_bind_declared_same_layout, [0]), (_bind_declared_other_layout, [0]), ], ids=["opened", "pickled", "derived", "declared-same-layout", "declared-other-layout"], @@ -1707,10 +1707,11 @@ def _bind_declared_other_layout(store: LocalStore) -> AnyArray: def test_write_stores_repair_however_bound( tmp_path: Path, bind: Callable[[LocalStore], AnyArray], stored_after: list[int] ) -> None: - """A chunk write stores the repair of a stored document that needs one, however the - writing array was bound to the store: the store's cache answers, and a store - object that has not seen the document reads it. A declaration that lays out - chunks otherwise describes another array, so the document is left as stored.""" + """A chunk write stores the repair of a stored document that needs one when the + store's cache knows it: an array opened, pickled or derived from one that read it. + A declaration over a document no binding of its store read knows nothing to repair, + and writes its chunks as any array does, as does one that lays out chunks + otherwise.""" path = tmp_path / "legacy.zarr" _legacy_array(path, 2) array = bind(LocalStore(path)) From a9de1777d624721919c474c40cdfc9c010c0c9d5 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 15:01:25 +0100 Subject: [PATCH 26/48] fix(array): bind metadata dicts as models and settle ArrayIO's IO signatures AsyncArrayIO and ArrayIO read a metadata dict as a zarr-metadata model when it is valid JSON metadata, as zarrista's from_metadata does, so an extension zarr-python lacks fails only when chunks are read or written; metadata built in code, or an invalid document zarr reads as repaired, is still read into the legacy classes when bound. AsyncArray keeps reading dicts the legacy way. erase_chunk takes grid coordinates positionally, as any sequence, and setitem takes its prototype by keyword, matching zarrista before either signature is released. AsyncArray's pickled state carries the pipeline class its binding chose, which may differ from its pipeline instance's class. The object-codec check applies to build_array_metadata_v2 only, so zarr.create(zarr_format=2) behaves as before. Also: drop an unpickling shim for a state no release wrote, a dead read guard, and a duplicate attribute encode. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 115 ++++++++++++++++++++--------- tests/test_array.py | 45 ++++++++--- tests/test_array_io.py | 72 +++++++++++++++++- tests/test_metadata/test_repair.py | 20 +++++ 4 files changed, 204 insertions(+), 48 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 67c1d130a5..621bf497cc 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -26,6 +26,7 @@ import numpy as np from typing_extensions import Sentinel, deprecated +from zarr_metadata import MetadataValidationError as ModelValidationError from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata import zarr @@ -458,7 +459,9 @@ def __init__( store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: - object.__setattr__(self, "_io", AsyncArrayIO(metadata, store_path, config)) + # An AsyncArray reads metadata into the legacy classes, as it always has. + metadata_parsed = parse_array_metadata(metadata, str(store_path)) + object.__setattr__(self, "_io", AsyncArrayIO(metadata_parsed, store_path, config)) # An AsyncArray is the resolving layer: build its codec pipeline now, so # pipeline-construction warnings fire at creation, as they always have. _ = self._io._resolved_codec_pipeline @@ -484,6 +487,7 @@ def __getstate__(self) -> dict[str, object]: "codec_pipeline": self.codec_pipeline, "_chunk_grid": self._chunk_grid, "config": self.config, + "_pipeline_class": self._io._pipeline_class, } state.update(self._io._carried_cache_entry()) return state @@ -492,13 +496,12 @@ def __setstate__(self, state: dict[str, Any]) -> None: metadata = state["metadata"] pipeline = state["codec_pipeline"] io = object.__new__(AsyncArrayIO) - io._bind(metadata, state["store_path"], state["config"], type(pipeline)) + # The class the binding chose: a pipeline may build an instance of + # another one (zarrs-python falls back for stores it cannot use). + pipeline_class = state.get("_pipeline_class", type(pipeline)) + io._bind(metadata, state["store_path"], state["config"], pipeline_class) io._codec_pipeline_cache = pipeline io._chunk_grid_cache = state["_chunk_grid"] - # A pickle from before the store's cache held this carries it on the metadata. - legacy_mark = metadata.__dict__.pop("_stored_document", None) - if legacy_mark is not None: - state = {**state, "_stored_entry": legacy_mark} io._restore_cache_entry(state) object.__setattr__(self, "_io", io) @@ -720,7 +723,7 @@ def _create_metadata_v2( attributes: dict[str, JSON] | None = None, ) -> ArrayV2Metadata: """See the module-level `build_array_metadata_v2`.""" - return build_array_metadata_v2( + return _array_metadata_v2( shape=shape, dtype=dtype, chunks=chunks, @@ -4687,7 +4690,7 @@ async def init_array( order_parsed = order chunk_key_encoding_parsed = cast("V2ChunkKeyEncoding", chunk_key_encoding_parsed) - meta = build_array_metadata_v2( + meta = _array_metadata_v2( shape=shape_parsed, dtype=zdtype, chunks=outer_chunks.chunk_shape, @@ -4753,7 +4756,7 @@ async def init_array( return arr -def build_array_metadata_v2( +def _array_metadata_v2( *, shape: ShapeLike, dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], @@ -4765,14 +4768,9 @@ def build_array_metadata_v2( compressor: CompressorLikev2 = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV2Metadata: - """Build Zarr format 2 array metadata from its fields. - - Each keyword is a `.zarray` field; `attributes` is the `.zattrs` one. - `fill_value` defaults to the data type's default scalar (an explicit - `None` is kept: the field is nullable); `order` defaults from the global - `array.order`. `filters` and `compressor` are numcodecs objects or their - JSON, as stored. Nothing is guessed. - """ + """`build_array_metadata_v2` without its object codec check, which the + callers that create arrays from user arguments make (or, for the legacy + `create` path, have never made) themselves.""" zdtype = parse_dtype(dtype, zarr_format=2) if order is None: order = zarr_config.get("array.order") @@ -4782,7 +4780,7 @@ def build_array_metadata_v2( # An explicit None is kept: the field is nullable in Zarr format 2. fill_value_parsed = fill_value - metadata = ArrayV2Metadata( + return ArrayV2Metadata( shape=parse_shapelike(shape), dtype=zdtype, chunks=parse_shapelike(chunks), @@ -4793,7 +4791,40 @@ def build_array_metadata_v2( filters=filters, attributes=attributes, ) - _check_object_codec_v2(zdtype, metadata.filters, metadata.compressor) + + +def build_array_metadata_v2( + *, + shape: ShapeLike, + dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + chunks: ShapeLike, + fill_value: Any = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + dimension_separator: Literal[".", "/"] = ".", + filters: Iterable[dict[str, JSON] | Numcodec] | None = None, + compressor: CompressorLikev2 = None, + attributes: dict[str, JSON] | None = None, +) -> ArrayV2Metadata: + """Build Zarr format 2 array metadata from its fields. + + Each keyword is a `.zarray` field; `attributes` is the `.zattrs` one. + `fill_value` defaults to the data type's default scalar (an explicit + `None` is kept: the field is nullable); `order` defaults from the global + `array.order`. `filters` and `compressor` are numcodecs objects or their + JSON, as stored. Nothing is guessed. + """ + metadata = _array_metadata_v2( + shape=shape, + dtype=dtype, + chunks=chunks, + fill_value=fill_value, + order=order, + dimension_separator=dimension_separator, + filters=filters, + compressor=compressor, + attributes=attributes, + ) + _check_object_codec_v2(metadata.dtype, metadata.filters, metadata.compressor) return metadata @@ -5099,8 +5130,9 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat declaration, as over any document. The document is read again, because the store may hold a newer one. If that one lays out chunks differently (another writer stored a different chunk size since), the caller would write chunks no reader - finds, so it raises, stores nothing, and leaves the cache as it was, so every - later write raises too. If it needs no repair (the array was re-saved since, + finds, so it raises, stores nothing, and leaves the cache as it was, so later + writes raise too, until a binding of this store object reads or stores a newer + document. If it needs no repair (the array was re-saved since, possibly by another implementation), it is left as written; if there is none, there is nothing to repair. @@ -5108,8 +5140,8 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat has read or stored a newer document (for example, `create_array(..., overwrite=True)` on that store), the cache no longer records the repair, so an array opened before that writes its chunks under its own metadata unchecked, as - any array whose document another writer replaced does: zarr does not detect a - stale array handle. + any array whose document another writer replaced does: outside the check above, + zarr does not detect a stale array handle. Storing the same repair twice is harmless, so handles that write chunks concurrently need no coordination. The read and the store are not one atomic @@ -5121,7 +5153,6 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat path = store_path.path zarr_format = metadata.zarr_format keys = ARRAY_DOCUMENTS[zarr_format] - documents: dict[str, Buffer] | None = None if cache.knows(path, zarr_format): stored = cache.needing_repair(path, zarr_format) else: @@ -5130,8 +5161,7 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat metadata ): return - if documents is None: - documents = await read_documents(store_path, keys) + documents = await read_documents(store_path, keys) document = stored_array_document(documents, zarr_format) if document is None: cache.observe(path, None, zarr_format) @@ -6336,6 +6366,18 @@ async def _delete_key(store_path: StorePath, key: str) -> None: """The keys a pickled binding carries its store cache's entry under.""" +def _bound_metadata_dict(data: dict[str, Any], path: str) -> ArrayMetadataModel | ArrayMetadata: + """A metadata dict given to a binding: read as a model when it is valid + JSON metadata, so names resolve only when chunks are read or written; + otherwise (metadata built in code, or an invalid document zarr reads as + repaired) read into the legacy classes now, warning as `AsyncArray` always + has.""" + try: + return _array_metadata_model(data) + except ModelValidationError: + return parse_array_metadata(data, path) + + def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: """A copy of a legacy metadata object with its own attributes dict, so an attribute update through one binding (which replaces top-level keys in @@ -6438,10 +6480,12 @@ def __init__( source: ArrayMetadataModel | ArrayMetadata if isinstance(metadata, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): source = metadata - else: - # A legacy resolved object passes through; raw JSON is read eagerly, - # as repaired when it needs that, warning as `AsyncArray` always has. + elif isinstance(metadata, ArrayMetadata): + # A resolved object passes through (v2 with a chunk size of 0 is read + # as repaired). source = parse_array_metadata(metadata, str(store_path)) + else: + source = _bound_metadata_dict(cast("dict[str, Any]", metadata), str(store_path)) # The pipeline implementation is chosen when the binding is made, from # the global configuration, exactly as `AsyncArray` has always chosen # it; the pipeline itself is built on first use. @@ -6544,11 +6588,7 @@ def metadata_legacy(self) -> ArrayMetadata: if self._resolved_cache is None: data = _metadata_model_json(source) if isinstance(source, ZarrV2ArrayMetadata): - # The model's attributes are read-only JSON; its `.zattrs` is plain. - zattrs = {str(key): value for key, value in source.to_key_value().items()} - data["attributes"] = ( - json.loads(zattrs[ZATTRS_JSON]) if ZATTRS_JSON in zattrs else {} - ) + data.setdefault("attributes", {}) self._resolved_cache = parse_array_metadata(data, str(self._store_path)) return self._resolved_cache @@ -6812,8 +6852,9 @@ async def erase_metadata(self) -> None: self._store_path.path, None, self._zarr_format ) - async def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: + async def erase_chunk(self, chunk_coords: Sequence[int], /) -> None: """Delete one stored chunk by grid coordinates (the shard, under sharding). Succeeds if absent.""" + chunk_coords = tuple(chunk_coords) grid_shape = self._shard_grid_shape if len(chunk_coords) != len(grid_shape) or any( coord < 0 or coord >= size for coord, size in zip(chunk_coords, grid_shape, strict=True) @@ -6904,6 +6945,7 @@ async def setitem( self, selection: BasicSelection, value: npt.ArrayLike, + *, prototype: BufferPrototype | None = None, ) -> None: """Encode `value` and store the chunks covering `selection`.""" @@ -7040,7 +7082,7 @@ def erase_metadata(self) -> None: """Delete the stored metadata; chunks are kept. See `AsyncArrayIO.erase_metadata`.""" sync(self._async_io.erase_metadata()) - def erase_chunk(self, chunk_coords: tuple[int, ...]) -> None: + def erase_chunk(self, chunk_coords: Sequence[int], /) -> None: """Delete one stored chunk by grid coordinates. See `AsyncArrayIO.erase_chunk`.""" sync(self._async_io.erase_chunk(chunk_coords)) @@ -7061,6 +7103,7 @@ def setitem( self, selection: BasicSelection, value: npt.ArrayLike, + *, prototype: BufferPrototype | None = None, ) -> None: """Encode `value` and store the chunks. See `AsyncArrayIO.setitem`.""" diff --git a/tests/test_array.py b/tests/test_array.py index 0d2345922a..0d46f2e73b 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -56,6 +56,7 @@ resolve_outer_and_inner_chunks, ) from zarr.core.chunk_key_encodings import ChunkKeyEncodingParams +from zarr.core.codec_pipeline import BatchedCodecPipeline from zarr.core.common import JSON, ChunksLike, ZarrFormat, ceildiv from zarr.core.dtype import ( DateTime64, @@ -2642,7 +2643,9 @@ async def test_create_array_huge_chunk_count() -> None: def test_async_array_keeps_its_dataclass_surface() -> None: """`AsyncArray` wraps an `AsyncArrayIO` but keeps the dataclass surface it has always had: its fields, `replace`, `asdict`, equality, (un)hashability, pickled - state, swapping a field with `object.__setattr__`, and independent copies.""" + state, swapping a field with `object.__setattr__`, and copies whose binding + (shape, pipeline) is their own; like before, a copy shares the metadata + object, so attribute updates made in place are seen through both.""" import copy import pickle @@ -2696,17 +2699,24 @@ def test_async_array_loads_a_pickled_field_state() -> None: array pickled before it wrapped an `AsyncArrayIO` loads.""" arr = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) aa = arr.async_array - old_state = { - "metadata": aa.metadata, - "store_path": aa.store_path, - "codec_pipeline": aa.codec_pipeline, - "_chunk_grid": aa._chunk_grid, - "config": aa.config, - } + sync(aa.setitem(slice(None), np.arange(4, dtype="uint8"))) + old_state = pickle.loads( + pickle.dumps( + { + "metadata": aa.metadata, + "store_path": aa.store_path, + "codec_pipeline": aa.codec_pipeline, + "_chunk_grid": aa._chunk_grid, + "config": aa.config, + } + ) + ) revived = object.__new__(type(aa)) revived.__setstate__(old_state) assert revived == aa - np.testing.assert_array_equal(sync(revived.getitem(slice(None))), np.zeros(4, dtype="uint8")) + np.testing.assert_array_equal(sync(revived.getitem(slice(None))), np.arange(4)) + sync(revived.setitem(slice(None), np.full(4, 7, dtype="uint8"))) + np.testing.assert_array_equal(sync(revived.getitem(slice(None))), np.full(4, 7)) def test_dask_token_of_an_array_is_deterministic() -> None: @@ -2717,3 +2727,20 @@ def test_dask_token_of_an_array_is_deterministic() -> None: _ = arr[:] assert dask_base.tokenize(arr) == before assert dask_base.tokenize(zarr.open_array(store)) == before + + +class _ChosenPipeline(BatchedCodecPipeline): + """A pipeline class whose instances may be another class's.""" + + +def test_pickled_array_keeps_the_pipeline_class_it_chose() -> None: + """A pickled array keeps the pipeline class its binding chose, also when that + class built an instance of another one (as zarrs-python does for stores it + cannot use).""" + from zarr.core.codec_pipeline import BatchedCodecPipeline + + aa = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)).async_array + aa.io._pipeline_class = _ChosenPipeline + revived = pickle.loads(pickle.dumps(aa)) + assert revived.io._pipeline_class is _ChosenPipeline + assert type(revived.codec_pipeline) is BatchedCodecPipeline diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 4a0c6c5f21..1f1a301a2c 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import numpy as np import pytest @@ -20,6 +20,8 @@ from zarr.storage import MemoryStore, StorePath if TYPE_CHECKING: + from collections.abc import Callable + from zarr.core.common import JSON, ZarrFormat METADATA_DOCUMENTS = ("zarr.json", ".zarray", ".zattrs", ".zgroup") @@ -330,7 +332,14 @@ def _bind_zero_chunk_by_constructor(metadata: ArrayV2Metadata) -> AsyncArrayIO: def _bind_zero_chunk_by_build(metadata: ArrayV2Metadata) -> AsyncArrayIO: - return AsyncArrayIO.build_v2(MemoryStore(), shape=metadata.shape, dtype="int16", chunks=(0,)) + return AsyncArrayIO.build_v2( + MemoryStore(), + shape=metadata.shape, + dtype=metadata.dtype, + chunks=metadata.chunks, + fill_value=metadata.fill_value, + order=metadata.order, + ) def _bind_zero_chunk_by_with_metadata(metadata: ArrayV2Metadata) -> AsyncArrayIO: @@ -343,7 +352,9 @@ def _bind_zero_chunk_by_with_metadata(metadata: ArrayV2Metadata) -> AsyncArrayIO [_bind_zero_chunk_by_constructor, _bind_zero_chunk_by_build, _bind_zero_chunk_by_with_metadata], ids=["constructor", "build_v2", "with_metadata"], ) -def test_zero_chunk_v2_metadata_binds_repaired(bind: Any) -> None: +def test_zero_chunk_v2_metadata_binds_repaired( + bind: Callable[[ArrayV2Metadata], AsyncArrayIO], +) -> None: """`ArrayV2Metadata` with a chunk size of 0 is read as repaired however it reaches a binding, as `parse_array_metadata` reads it.""" metadata = ArrayV2Metadata(shape=(3,), chunks=(0,), dtype=Int16(), fill_value=0, order="C") @@ -373,3 +384,58 @@ async def test_v2_model_binding_keeps_attributes_writable() -> None: def test_build_v2_vlen_without_object_codec_raises() -> None: with pytest.raises(ValueError, match="requires the numcodecs.VLenUTF8 codec"): AsyncArrayIO.build_v2(MemoryStore(), shape=(4,), dtype=str, chunks=(2,)) + + +async def test_metadata_dict_binds_as_a_model_until_chunks_move() -> None: + """A valid metadata dict binds as a model, so an extension zarr-python + lacks fails only when chunks are read or written; metadata zarr can only + read as repaired is read, repaired, when bound.""" + store = MemoryStore() + doc: dict[str, JSON] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "bfloat16", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "zarrs.zfp", "configuration": {"mode": "reversible"}}], + } + aio = AsyncArrayIO(cast("Any", doc), StorePath(store, "foo")) + assert isinstance(aio.metadata, ZarrV3ArrayMetadata) + await aio.store_metadata() + await aio.verify() + await aio.erase_chunks() + with pytest.raises(ValueError, match="bfloat16"): + await aio.getitem(Ellipsis) + + zero_chunk = {**doc, "data_type": "uint8", "codecs": [{"name": "bytes"}]} + zero_chunk["chunk_grid"] = {"name": "regular", "configuration": {"chunk_shape": [0]}} + with pytest.warns(zarr.errors.ZarrUserWarning, match="is read as"): + repaired = AsyncArrayIO(cast("Any", zero_chunk), StorePath(store, "bar")) + assert repaired.metadata_legacy.chunks == (1,) + + +async def test_chunk_and_selection_io_signatures() -> None: + """erase_chunk takes grid coordinates positionally, as any sequence; + setitem takes its prototype by keyword.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8")) + await aio.erase_chunk([0]) + assert "foo/c/0" not in await _store_keys(store) + with pytest.raises(TypeError): + await aio.erase_chunk(chunk_coords=(1,)) # type: ignore[call-arg] + with pytest.raises(TypeError): + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8"), default_buffer_prototype()) # type: ignore[call-arg] + + +async def test_v2_model_binding_without_attributes() -> None: + store = MemoryStore() + aio = AsyncArrayIO.build_v2(store, name="foo", shape=(4,), dtype="uint8", chunks=(2,)) + await aio.store_metadata() + await store.delete("foo/.zattrs") + arr = (await AsyncArrayIO.open(store, name="foo")).to_array() + assert arr.metadata.attributes == {} + await arr.update_attributes({"x": 1}) + assert (await open_array(store=store, path="foo")).metadata.attributes == {"x": 1} diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index edc4aec728..b92a28fdd2 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -1833,3 +1833,23 @@ def test_restore_keeps_what_the_cache_learned_itself() -> None: assert cache.needing_repair("a", 2) is None assert cache.needing_repair("b", 2) == broken + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_pickled_consolidated_member_stores_its_repair( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A member read from consolidated metadata carries its copy's entry when + pickled, so the unpickled member's first chunk write stores the repair.""" + path = tmp_path / "group.zarr" + document = _consolidated_legacy_member(path, zarr_format) + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + member = group["a"] + assert isinstance(member, zarr.Array) + revived: AnyArray = pickle.loads(pickle.dumps(member)) + + revived[:] = [7, 8, 9] + + assert _stored_chunks(json.loads(document.read_text())) == [1] From 0d28b55771f8bcfb46f140018ce1d711e283a9b5 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 15:14:42 +0100 Subject: [PATCH 27/48] fix(metadata): carry group cache entries, bound the cache, import models lazily A pickled or deep-copied consolidated group lost what its store's cache knew of its members, so a group write through the copy re-encoded a member read from a repaired copy instead of keeping it as stored (and raised for a rectilinear one). AsyncGroup now carries those entries in its pickled state. An AsyncArray built from a stored document given as a dict records it, so its first chunk write stores the repair, as before. The cache records a valid document only where it replaces one that needs a repair or a consolidated copy that does, so it no longer grows with every array a store object sees; group writes are not re-parsed to record them, and zarr_metadata is imported when a binding first holds a model, so `import zarr` costs what it did. A legacy-sourced binding derives its model on each access, so an in-place change through a handle sharing its metadata is never stale. ArrayIO.open opens a path writable without creating it. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 94 ++++++++++++++++++++---------- src/zarr/core/group.py | 41 +++++++++++-- src/zarr/core/metadata/io.py | 43 +++++++++++--- src/zarr/core/metadata/stored.py | 47 ++++++++++++--- tests/test_array_io.py | 8 +++ tests/test_metadata/test_repair.py | 59 ++++++++++++++++++- 6 files changed, 240 insertions(+), 52 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 621bf497cc..745bab1b43 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2,8 +2,10 @@ import contextlib import copy +import functools import json import math +import os import warnings from asyncio import gather from collections.abc import Iterable, Mapping, Sequence @@ -26,8 +28,6 @@ import numpy as np from typing_extensions import Sentinel, deprecated -from zarr_metadata import MetadataValidationError as ModelValidationError -from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata import zarr from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec @@ -178,11 +178,13 @@ from typing import Self import numpy.typing as npt + from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata from zarr.abc.codec import CodecPipeline from zarr.abc.store import Store from zarr.codecs.sharding import IndexLocation, ShardingCodec from zarr.core.buffer import Buffer + from zarr.core.common import AccessModeLiteral from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.core.metadata.repair import ArrayDocument from zarr.storage import StoreLike @@ -461,6 +463,10 @@ def __init__( ) -> None: # An AsyncArray reads metadata into the legacy classes, as it always has. metadata_parsed = parse_array_metadata(metadata, str(store_path)) + if isinstance(metadata, dict): + # A document given as a dict stands for what the store holds, as the + # document an array is opened from does. + _record_given_document(store_path, metadata, metadata_parsed.zarr_format) object.__setattr__(self, "_io", AsyncArrayIO(metadata_parsed, store_path, config)) # An AsyncArray is the resolving layer: build its codec pipeline now, so # pipeline-construction warnings fire at creation, as they always have. @@ -6236,7 +6242,7 @@ async def _delete_chunk(coords: tuple[int, ...]) -> None: ) # Write new metadata - await store_documents(array.store_path, documents) + await store_documents(array.store_path, documents, array=True) # Update the wrapped binding in place; the codec pipeline is kept, as before. array._io._rebind_resolved(new_metadata, new_chunk_grid) @@ -6372,12 +6378,25 @@ def _bound_metadata_dict(data: dict[str, Any], path: str) -> ArrayMetadataModel otherwise (metadata built in code, or an invalid document zarr reads as repaired) read into the legacy classes now, warning as `AsyncArray` always has.""" + from zarr_metadata import MetadataValidationError + try: return _array_metadata_model(data) - except ModelValidationError: + except MetadataValidationError: return parse_array_metadata(data, path) +def _record_given_document( + store_path: StorePath, document: Mapping[str, Any], zarr_format: ZarrFormat +) -> None: + """Record `document` as the one stored at `store_path` when reading it needs a + repair that moves chunks, unless the store's cache already knows that document.""" + if moves_chunks(document, zarr_format): + stored_documents(store_path.store).restore( + store_path.path, zarr_format, copy.deepcopy(dict(document)) + ) + + def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: """A copy of a legacy metadata object with its own attributes dict, so an attribute update through one binding (which replaces top-level keys in @@ -6403,9 +6422,19 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: `zarr_metadata.MetadataValidationError`. Unknown extension names read as unclaimed fields, not errors. """ + model_v2, model_v3 = _model_classes() if data.get("zarr_format") == 2: - return ZarrV2ArrayMetadata(data) - return ZarrV3ArrayMetadata(data) + return model_v2(data) + return model_v3(data) + + +@functools.cache +def _model_classes() -> tuple[type[ZarrV2ArrayMetadata], type[ZarrV3ArrayMetadata]]: + """The `zarr_metadata` model classes, imported on first use: the import is + slow, and only bindings that hold a model need it.""" + from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata + + return ZarrV2ArrayMetadata, ZarrV3ArrayMetadata class ArrayV2DeclarationKwargs(TypedDict, total=False): @@ -6465,7 +6494,6 @@ class AsyncArrayIO: _config: ArrayConfig _pipeline_class: type[CodecPipeline] _source: ArrayMetadataModel | ArrayMetadata - _metadata_model_cache: ArrayMetadataModel | None _resolved_cache: ArrayMetadata | None _chunk_grid_cache: ChunkGrid | None _chunk_key_encoding_cache: ChunkKeyEncoding | None @@ -6478,12 +6506,12 @@ def __init__( config: ArrayConfigLike | None = None, ) -> None: source: ArrayMetadataModel | ArrayMetadata - if isinstance(metadata, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): - source = metadata - elif isinstance(metadata, ArrayMetadata): + if isinstance(metadata, ArrayMetadata): # A resolved object passes through (v2 with a chunk size of 0 is read # as repaired). source = parse_array_metadata(metadata, str(store_path)) + elif isinstance(metadata, _model_classes()): + source = metadata else: source = _bound_metadata_dict(cast("dict[str, Any]", metadata), str(store_path)) # The pipeline implementation is chosen when the binding is made, from @@ -6502,7 +6530,6 @@ def _bind( self._config = config self._pipeline_class = pipeline_class self._source = source - self._metadata_model_cache = None self._resolved_cache = None self._chunk_grid_cache = None self._chunk_key_encoding_cache = None @@ -6558,7 +6585,9 @@ async def open( repaired raises `zarr_metadata.MetadataValidationError`; use [`zarr.api.asynchronous.open_array`][] for the repairing read. """ - store_path = await make_store_path(store, path=name) + # A path opens writable without creating anything when no array is there. + mode: AccessModeLiteral | None = "r+" if isinstance(store, str | os.PathLike) else None + store_path = await make_store_path(store, path=name, mode=mode) data = await get_array_metadata(store_path, zarr_format=zarr_format) binding = cls(_array_metadata_model(data), store_path, config) observe_read(store_path, data) @@ -6568,16 +6597,15 @@ async def open( def metadata(self) -> ArrayMetadataModel: """The bound metadata, as a `zarr_metadata` model.""" source = self._source - if isinstance(source, ZarrV2ArrayMetadata | ZarrV3ArrayMetadata): + if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): return source - if self._metadata_model_cache is None: - buffers = source.to_buffer_dict(default_buffer_prototype()) - key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} - if isinstance(source, ArrayV2Metadata): - self._metadata_model_cache = ZarrV2ArrayMetadata.from_key_value(key_value) - else: - self._metadata_model_cache = ZarrV3ArrayMetadata.from_key_value(key_value) - return self._metadata_model_cache + # Derived on every access, not cached: a legacy object may be changed in + # place through another handle that shares it (as `AsyncArray` copies do). + buffers = source.to_buffer_dict(default_buffer_prototype()) + key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} + model_v2, model_v3 = _model_classes() + model_class = model_v2 if isinstance(source, ArrayV2Metadata) else model_v3 + return model_class.from_key_value(key_value) @property def metadata_legacy(self) -> ArrayMetadata: @@ -6587,7 +6615,7 @@ def metadata_legacy(self) -> ArrayMetadata: return source if self._resolved_cache is None: data = _metadata_model_json(source) - if isinstance(source, ZarrV2ArrayMetadata): + if isinstance(source, _model_classes()[0]): data.setdefault("attributes", {}) self._resolved_cache = parse_array_metadata(data, str(self._store_path)) return self._resolved_cache @@ -6604,9 +6632,12 @@ def config(self) -> ArrayConfig: @property def _zarr_format(self) -> ZarrFormat: - if isinstance(self._source, ZarrV2ArrayMetadata | ArrayV2Metadata): + source = self._source + if isinstance(source, ArrayV2Metadata): return 2 - return 3 + if isinstance(source, ArrayV3Metadata): + return 3 + return 2 if isinstance(source, _model_classes()[0]) else 3 @property def _resolved_if_any(self) -> ArrayMetadata | None: @@ -6620,7 +6651,6 @@ def _rebind_resolved( ) -> None: """Make `resolved` this binding's state. The codec pipeline is kept: rebinding is how `resize` swaps the shape, which leaves the chunk encoding as it was.""" self._source = resolved - self._metadata_model_cache = None self._resolved_cache = None self._chunk_grid_cache = chunk_grid self._chunk_key_encoding_cache = None @@ -6635,7 +6665,7 @@ def _chunk_grid(self) -> ChunkGrid: return self._chunk_grid_cache document = self.metadata doc = _metadata_model_json(document) - if isinstance(document, ZarrV2ArrayMetadata): + if isinstance(document, _model_classes()[0]): self._chunk_grid_cache = ChunkGrid.from_sizes(document.shape, tuple(doc["chunks"])) else: grid = doc["chunk_grid"] @@ -6676,7 +6706,7 @@ def _chunk_key_encoding(self) -> ChunkKeyEncoding: if self._chunk_key_encoding_cache is None: document = self.metadata doc = _metadata_model_json(document) - if isinstance(document, ZarrV2ArrayMetadata): + if isinstance(document, _model_classes()[0]): self._chunk_key_encoding_cache = V2ChunkKeyEncoding( separator=doc.get("dimension_separator", ".") ) @@ -6805,7 +6835,11 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: str(key): prototype.buffer.from_bytes(value) for key, value in key_value.items() } await _write_metadata( - self._store_path, to_save, self._zarr_format, ensure_parents=ensure_parents + self._store_path, + to_save, + self._zarr_format, + ensure_parents=ensure_parents, + array=True, ) if self._zarr_format == 2 and ZATTRS_JSON not in to_save: # An attribute-less model emits no .zattrs; a stale one must not @@ -6824,14 +6858,14 @@ async def verify(self) -> None: if stored == bound: return if type(stored) is not type(bound): - stored_format = 2 if isinstance(stored, ZarrV2ArrayMetadata) else 3 + stored_format = 2 if isinstance(stored, _model_classes()[0]) else 3 raise ValueError( f"Array {str(self.store_path)!r}: the stored array is Zarr format " f"{stored_format}; the bound metadata is Zarr format {self._zarr_format}." ) stored_json = dict(stored.to_json()) bound_json = dict(bound.to_json()) - if isinstance(stored, ZarrV2ArrayMetadata) and isinstance(bound, ZarrV2ArrayMetadata): + if isinstance(stored, _model_classes()[0]) and isinstance(bound, _model_classes()[0]): stored_json["attributes"] = dict(stored.attributes) bound_json["attributes"] = dict(bound.attributes) differing = sorted( diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 4ea3b36dd7..a65cf0ec52 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -89,6 +89,7 @@ from zarr.core.common import MemoryOrder from zarr.core.dtype import ZDTypeLike from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.stored import CarriedEntry from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 logger = logging.getLogger("zarr.group") @@ -688,6 +689,32 @@ async def open( use_consolidated=use_consolidated, ) + def __getstate__(self) -> dict[str, object]: + # A store's cache does not pickle with the store: carry what it knows of the + # consolidated members, so a write through the unpickled group keeps the + # members it would have kept as stored. + state = dict(self.__dict__) + consolidated = self.metadata.consolidated_metadata + if consolidated is not None: + paths = [ + _join_paths([self.store_path.path, key]) for key in consolidated.flattened_metadata + ] + entries = stored_documents(self.store_path.store).entries( + paths, self.metadata.zarr_format + ) + if entries: + state["_stored_members"] = entries + return state + + def __setstate__(self, state: dict[str, object]) -> None: + state = dict(state) + entries = cast("dict[str, CarriedEntry] | None", state.pop("_stored_members", None)) + self.__dict__.update(state) + if entries is not None: + stored_documents(self.store_path.store).restore_entries( + entries, self.metadata.zarr_format + ) + @classmethod def _from_bytes_v2( cls, @@ -893,7 +920,9 @@ async def delitem(self, key: str) -> None: # In place, so every handle sharing this consolidated metadata (a parent's or a # subgroup's) sees the deletion. consolidated.metadata.pop(key, None) - await store_documents(self.store_path, encode_documents(self.store_path, self.metadata)) + await store_documents( + self.store_path, encode_documents(self.store_path, self.metadata), array=False + ) async def get[DefaultT]( self, key: str, default: DefaultT | None = None @@ -2278,7 +2307,9 @@ async def update_attributes_async(self, new_attributes: dict[str, Any]) -> Group new_metadata = replace(self.metadata, attributes=new_attributes) # Write new metadata - await store_documents(self.store_path, encode_documents(self.store_path, new_metadata)) + await store_documents( + self.store_path, encode_documents(self.store_path, new_metadata), array=False + ) async_group = replace(self._async_group, metadata=new_metadata) return replace(self, _async_group=async_group) @@ -3469,7 +3500,8 @@ async def _store_nodes( node_name = created_key[: created_key.rfind("/")] meta_out = nodes[node_name] if meta_out.zarr_format == 3: - cache.observe_written(node_name, node_documents[node_name]) + if not isinstance(meta_out, GroupMetadata): + cache.observe_written(node_name, node_documents[node_name]) yield node_name, built[node_name] else: # For zarr v2 @@ -3485,7 +3517,8 @@ async def _store_nodes( meta_done = _join_paths([node_name, ZARRAY_JSON]) in created_object_keys if meta_done and attrs_done: - cache.observe_written(node_name, node_documents[node_name]) + if not isinstance(meta_out, GroupMetadata): + cache.observe_written(node_name, node_documents[node_name]) yield node_name, built[node_name] continue diff --git a/src/zarr/core/metadata/io.py b/src/zarr/core/metadata/io.py index 0c0a58cede..5242ae1d3a 100644 --- a/src/zarr/core/metadata/io.py +++ b/src/zarr/core/metadata/io.py @@ -97,13 +97,16 @@ def encode_documents( raise -async def store_documents(store_path: StorePath, documents: Mapping[str, Buffer]) -> None: - """Store metadata documents encoded by `encode_documents` under `store_path`, and - record them in the store's cache (see `stored_documents`).""" +async def store_documents( + store_path: StorePath, documents: Mapping[str, Buffer], *, array: bool +) -> None: + """Store metadata documents encoded by `encode_documents` under `store_path`; an + array's are recorded in the store's cache (see `stored_documents`).""" await asyncio.gather( *(set_or_delete(store_path / key, value) for key, value in documents.items()) ) - stored_documents(store_path.store).observe_written(store_path.path, documents) + if array: + stored_documents(store_path.store).observe_written(store_path.path, documents) def members_as_stored(store_path: StorePath, metadata: GroupMetadata) -> dict[str, ArrayDocument]: @@ -194,7 +197,11 @@ async def upsert_metadata( {key: buffer_to_json_object(buf) for key, buf in documents.items()}, ) changed = {change.path[0] for change in changes} - await store_documents(store_path, {k: v for k, v in documents.items() if k in changed}) + await store_documents( + store_path, + {k: v for k, v in documents.items() if k in changed}, + array=not _is_group(metadata), + ) return changes @@ -270,7 +277,13 @@ async def save_metadata( ValueError """ to_save = encode_documents(store_path, metadata) - await _write_metadata(store_path, to_save, metadata.zarr_format, ensure_parents=ensure_parents) + await _write_metadata( + store_path, + to_save, + metadata.zarr_format, + ensure_parents=ensure_parents, + array=not _is_group(metadata), + ) async def save_new_metadata( @@ -303,7 +316,19 @@ async def save_new_metadata( await store_path.delete_dir() else: await ensure_no_existing_node(store_path, zarr_format=metadata.zarr_format) - await _write_metadata(store_path, to_save, metadata.zarr_format, ensure_parents=ensure_parents) + await _write_metadata( + store_path, + to_save, + metadata.zarr_format, + ensure_parents=ensure_parents, + array=not _is_group(metadata), + ) + + +def _is_group(metadata: ArrayMetadata | GroupMetadata) -> bool: + from zarr.core.group import GroupMetadata + + return isinstance(metadata, GroupMetadata) async def _write_metadata( @@ -312,6 +337,7 @@ async def _write_metadata( zarr_format: ZarrFormat, *, ensure_parents: bool, + array: bool, ) -> None: set_awaitables = [set_or_delete(store_path / key, value) for key, value in to_save.items()] @@ -351,4 +377,5 @@ async def _write_metadata( ) from e await asyncio.gather(*set_awaitables) - stored_documents(store_path.store).observe_written(store_path.path, to_save) + if array: + stored_documents(store_path.store).observe_written(store_path.path, to_save) diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index 3906680b95..f5017c105b 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -16,14 +16,14 @@ import copy import weakref from functools import partial -from typing import TYPE_CHECKING, Final +from typing import TYPE_CHECKING, Final, Literal from zarr.core._json import buffer_to_json_object from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON from zarr.core.metadata.repair import repair_array_document if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Iterable, Mapping from zarr.abc.store import Store from zarr.core.buffer import Buffer @@ -31,6 +31,13 @@ from zarr.core.metadata.repair import ArrayDocument +type CarriedEntry = ( + tuple[Literal["document"], ArrayDocument | None] | tuple[Literal["copy"], ArrayDocument] +) +"""What a cache records of one array document, carried in a pickle: the document's +entry, or else a consolidated copy's.""" + + def moves_chunks(document: ArrayDocument, zarr_format: ZarrFormat) -> bool: """Whether reading `document` needs a repair that moves chunks.""" _, readings = repair_array_document(document, zarr_format) @@ -42,8 +49,10 @@ class StoredDocuments: Each array document zarr has seen, by its path and Zarr format (a path may hold a document of each), maps to the document as it was stored when reading it needs a - repair that moves chunks, and to `None` when it needs none or there is no - document; a document zarr has not seen is unknown. Copies of array documents in + repair that moves chunks. A document that needs none (or its absence) is recorded, + as `None`, only where it replaces such a document or a consolidated copy that + needs one; otherwise it is unknown, which means the same, so the cache holds only + what stands for a repair. Copies of array documents in consolidated metadata are recorded apart, by the array's path and format, only when reading them needs such a repair, so a group write can store them as they were stored. A document zarr reads to open an array is recorded only once it has @@ -85,14 +94,34 @@ def restore_copy(self, path: str, zarr_format: ZarrFormat, document: ArrayDocume def observe(self, path: str, document: ArrayDocument | None, zarr_format: ZarrFormat) -> None: """Record that the array document of `zarr_format` at `path` is `document` (`None`: none).""" + key = (path, zarr_format) if document is not None and moves_chunks(document, zarr_format): - self._arrays[path, zarr_format] = copy.deepcopy(document) - else: - self._arrays[path, zarr_format] = None + self._arrays[key] = copy.deepcopy(document) + elif key in self._arrays or key in self._copies: + self._arrays[key] = None + + def entries(self, paths: Iterable[str], zarr_format: ZarrFormat) -> dict[str, CarriedEntry]: + """What this cache records of the array documents at `paths`, to carry to + another store object's cache (see `restore_entries`).""" + carried: dict[str, CarriedEntry] = {} + for path in paths: + if self.knows(path, zarr_format): + carried[path] = ("document", self.needing_repair(path, zarr_format)) + elif (copy_entry := self.copy_needing_repair(path, zarr_format)) is not None: + carried[path] = ("copy", copy_entry) + return carried + + def restore_entries(self, entries: Mapping[str, CarriedEntry], zarr_format: ZarrFormat) -> None: + """Record what `entries` carried, keeping what this cache already knows.""" + for path, entry in entries.items(): + if entry[0] == "document": + self.restore(path, zarr_format, entry[1]) + else: + self.restore_copy(path, zarr_format, entry[1]) def observe_written(self, path: str, documents: Mapping[str, Buffer | None]) -> None: - """Record the metadata documents just stored under `path`, by key (`None`: - deleted), as `store_documents` takes them.""" + """Record the array metadata documents just stored under `path`, by key + (`None`: deleted), as `store_documents` takes them.""" if ZARR_JSON in documents: zarr_json = documents[ZARR_JSON] document = None if zarr_json is None else buffer_to_json_object(zarr_json) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 1f1a301a2c..afc5ff9646 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: from collections.abc import Callable + from pathlib import Path from zarr.core.common import JSON, ZarrFormat @@ -439,3 +440,10 @@ async def test_v2_model_binding_without_attributes() -> None: assert arr.metadata.attributes == {} await arr.update_attributes({"x": 1}) assert (await open_array(store=store, path="foo")).metadata.attributes == {"x": 1} + + +def test_open_a_missing_path_creates_nothing(tmp_path: Path) -> None: + missing = tmp_path / "missing.zarr" + with pytest.raises(FileNotFoundError): + ArrayIO.open(str(missing)) + assert not missing.exists() diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index b92a28fdd2..73a53064e1 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import copy import dataclasses import gc import json @@ -1823,9 +1824,11 @@ def test_pickled_array_carries_what_its_store_knows( def test_restore_keeps_what_the_cache_learned_itself() -> None: """A carried entry fills only a document the cache does not know: what the - cache recorded itself, such as a repair it stored, stays.""" + cache recorded itself, such as a repair it stored over a broken document, + stays.""" cache = StoredDocuments() broken = _v2_doc([3], [0]) + cache.observe("a", broken, 2) cache.observe("a", None, 2) cache.restore("a", 2, broken) @@ -1853,3 +1856,57 @@ def test_pickled_consolidated_member_stores_its_repair( revived[:] = [7, 8, 9] assert _stored_chunks(json.loads(document.read_text())) == [1] + + +def test_cache_records_only_what_stands_for_a_repair() -> None: + """A valid document is recorded only where it replaces one that needs a repair + (or a consolidated copy that does), so the cache does not grow with every + array it sees.""" + cache = StoredDocuments() + valid, broken = _v2_doc([3], [3]), _v2_doc([3], [0]) + for path in ("a", "b", "c"): + cache.observe(path, valid, 2) + cache.observe("b", broken, 2) + cache.observe("b", valid, 2) + cache.observe_consolidated("", {"c": broken}, 2) + cache.observe("c", valid, 2) + assert (cache.knows("a", 2), cache.knows("b", 2), cache.knows("c", 2)) == (False, True, True) + assert cache.needing_repair("b", 2) is None + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_pickled_group_keeps_consolidated_members_as_stored( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A pickled (or deep-copied) consolidated group carries what its store's cache + knew of its members, so a group write through the copy keeps a member read + from a repaired copy as it was stored, as the original would.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "a") + stored = json.dumps(_consolidated_member(path, zarr_format, "a")) + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + for revived in (pickle.loads(pickle.dumps(group)), copy.deepcopy(group)): + revived.attrs["x"] = 1 + assert json.dumps(_consolidated_member(path, zarr_format, "a")) == stored + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_array_from_a_stored_document_dict_stores_its_repair( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """An `AsyncArray` built from a stored document given as a dict stores that + document's repair before its first chunk write, as one opened from it does.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, zarr_format) + doc_name = ".zarray" if zarr_format == 2 else "zarr.json" + doc = json.loads((path / doc_name).read_text()) + if zarr_format == 2: + doc["attributes"] = {} + array = zarr.Array(AsyncArray(doc, StorePath(LocalStore(path)))) + + array[:] = [7, 8, 9] + + assert _stored_chunks(json.loads((path / doc_name).read_text())) == [1] From 3c7b3dc68bc5480b8d7d7b0aacb93fe8f78744e3 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 15:27:00 +0100 Subject: [PATCH 28/48] fix(array): reject storage transformers in bound dicts; pin cache semantics A metadata dict with storage transformers bound as a model and failed only at the first chunk operation; it is now rejected when bound, as AsyncArray always rejected it. The cache's `knows` and the repair check's docstring say that a valid document is recorded only where it replaces a broken one, and a test pins that a member read valid and then from a broken consolidated copy is kept as the copy stored, as before. Tests also cover a pickled subgroup of a consolidated parent, a binding derived from an unpickled array building the pipeline class it chose, and `import zarr` leaving zarr_metadata unimported. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 22 +++++++++----- src/zarr/core/metadata/stored.py | 6 +++- tests/test_array.py | 16 ++++++++-- tests/test_array_io.py | 16 ++++++++++ tests/test_metadata/test_repair.py | 49 ++++++++++++++++++++++++++++++ 5 files changed, 97 insertions(+), 12 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 745bab1b43..303d9ff7cc 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -728,7 +728,8 @@ def _create_metadata_v2( compressor: CompressorLikev2 = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV2Metadata: - """See the module-level `build_array_metadata_v2`.""" + """See the module-level `build_array_metadata_v2`; this legacy path makes + no object codec check, as it never has (see `_array_metadata_v2`).""" return _array_metadata_v2( shape=shape, dtype=dtype, @@ -5126,11 +5127,12 @@ async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadat moves chunks, before chunks are written under `metadata`. The store's cache (see `zarr.core.metadata.stored`) says whether the document - needs such a repair. While it has not seen the document, a consolidated copy it - has seen stands for it, as consolidated metadata does for every reader; with - neither, nothing is known to need a repair, and nothing is read: metadata built - in code, or bound to a store that never read the document, writes its chunks - as any array does. Only a document whose repair + needs such a repair. Where it records nothing for the document (it records a + valid one only where that replaces a broken one), a consolidated copy it has + seen that needs a repair stands for it, as consolidated metadata does for every + reader, and the document is read to check; with neither, nothing is known to + need a repair, and nothing is read: metadata built in code, or bound to a store + that never read the document, writes its chunks as any array does. Only a document whose repair lays out chunks as `metadata` does is acted on: one that lays them out otherwise describes another array, and chunks written under `metadata` are its caller's declaration, as over any document. The document is read again, because the store @@ -6380,6 +6382,9 @@ def _bound_metadata_dict(data: dict[str, Any], path: str) -> ArrayMetadataModel has.""" from zarr_metadata import MetadataValidationError + if data.get("storage_transformers"): + # Bound as legacy metadata, which rejects them, as AsyncArray always has. + return parse_array_metadata(data, path) try: return _array_metadata_model(data) except MetadataValidationError: @@ -6599,8 +6604,9 @@ def metadata(self) -> ArrayMetadataModel: source = self._source if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): return source - # Derived on every access, not cached: a legacy object may be changed in - # place through another handle that shares it (as `AsyncArray` copies do). + # Derived on every access, not cached: a legacy object's attributes may be + # changed in place through another handle that shares it (as `AsyncArray` + # copies do); the chunk grid and key encoding such changes cannot touch. buffers = source.to_buffer_dict(default_buffer_prototype()) key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} model_v2, model_v3 = _model_classes() diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index f5017c105b..551e8378f7 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -65,7 +65,11 @@ def __init__(self) -> None: self._copies: dict[tuple[str, ZarrFormat], ArrayDocument] = {} def knows(self, path: str, zarr_format: ZarrFormat) -> bool: - """Whether zarr has read or stored the array document of `zarr_format` at `path`.""" + """Whether this cache records the array document of `zarr_format` at `path`: + one that needs a repair that moves chunks, or a valid one that replaced such + a document or stands against a consolidated copy that needs one. A valid + document read with nothing to replace is not recorded, so whether it is + known depends on what was read first; either way it means no repair.""" return (path, zarr_format) in self._arrays def needing_repair(self, path: str, zarr_format: ZarrFormat) -> ArrayDocument | None: diff --git a/tests/test_array.py b/tests/test_array.py index 0d46f2e73b..8fd0ecd3f3 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2737,10 +2737,20 @@ def test_pickled_array_keeps_the_pipeline_class_it_chose() -> None: """A pickled array keeps the pipeline class its binding chose, also when that class built an instance of another one (as zarrs-python does for stores it cannot use).""" - from zarr.core.codec_pipeline import BatchedCodecPipeline - aa = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)).async_array aa.io._pipeline_class = _ChosenPipeline revived = pickle.loads(pickle.dumps(aa)) - assert revived.io._pipeline_class is _ChosenPipeline assert type(revived.codec_pipeline) is BatchedCodecPipeline + # A binding derived from the revived one builds its pipeline from the class. + derived = revived.io.with_store(MemoryStore()) + assert type(derived.to_array().codec_pipeline) is _ChosenPipeline + + +def test_import_zarr_leaves_zarr_metadata_unimported() -> None: + """`zarr_metadata` is imported when a binding first holds a model, not by + `import zarr`, whose cost stays what it was.""" + import subprocess + import sys + + code = "import sys, zarr; assert 'zarr_metadata' not in sys.modules" + subprocess.run([sys.executable, "-c", code], check=True) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index afc5ff9646..7c9c229fe3 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -447,3 +447,19 @@ def test_open_a_missing_path_creates_nothing(tmp_path: Path) -> None: with pytest.raises(FileNotFoundError): ArrayIO.open(str(missing)) assert not missing.exists() + + +def test_metadata_dict_with_storage_transformers_raises() -> None: + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "uint8", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes"}], + "storage_transformers": [{"name": "nope"}], + } + with pytest.raises(ValueError, match="storage transformers"): + AsyncArrayIO(cast("Any", doc), StorePath(MemoryStore())) diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index 73a53064e1..d969b0c5d2 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -1910,3 +1910,52 @@ def test_array_from_a_stored_document_dict_stores_its_repair( array[:] = [7, 8, 9] assert _stored_chunks(json.loads((path / doc_name).read_text())) == [1] + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_member_read_valid_then_from_a_broken_copy_keeps_the_copy( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A member whose own document was read valid, then read again from a broken + consolidated copy, is kept as that copy stored by a group write, as a member + read from such a copy always was.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "a") + stored = json.dumps(_consolidated_member(path, zarr_format, "a")) + store = LocalStore(path) + zarr.open_array(store=store, path="a", mode="r") + group = zarr.open_group(store, mode="r+", use_consolidated=True) + + group.attrs["x"] = 1 + + assert json.dumps(_consolidated_member(path, zarr_format, "a")) == stored + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_pickled_subgroup_keeps_nested_members_as_stored( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A subgroup read from a consolidated parent carries its members' entries under + its own path when pickled, so a write through the copy stores what a write + through the subgroup itself stores, nested members kept as stored.""" + + def write_through(sub_of: Callable[[zarr.Group], zarr.Group], name: str) -> dict[str, bytes]: + path = tmp_path / name + _flagged_consolidated_group(path, zarr_format, "g/a") + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + sub = group["g"] + assert isinstance(sub, zarr.Group) + sub_of(sub).attrs["x"] = 1 + return { + str(p.relative_to(path)): p.read_bytes() + for p in sorted(path.rglob("*")) + if p.is_file() and p.name in ("zarr.json", ".zmetadata", ".zattrs", ".zgroup") + } + + assert write_through(lambda sub: pickle.loads(pickle.dumps(sub)), "pickled") == write_through( + lambda sub: sub, "original" + ) From 8af2f87ee05ef6e8b2286c89da1a4cad2576dd56 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 15:34:14 +0100 Subject: [PATCH 29/48] fix(array): take zarr-metadata JSON in bindings and rebuild pickled pipelines ArrayIO and AsyncArrayIO type their metadata argument to take zarr_metadata's JSON types (what zarrista's arrays hold), and with_metadata reads a metadata dict as the constructor does instead of binding the dict itself. A pickled AsyncArrayIO no longer carries its built codec pipeline, which is bound to the pickling process's store object and may not pickle; it is rebuilt on first use. The class docstring words name resolution in the data path as what today's pipelines do. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 43 ++++++++++++++++++++++++++++++------------ tests/test_array_io.py | 25 ++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 12 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 303d9ff7cc..ebc9b2d7f0 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -178,7 +178,12 @@ from typing import Self import numpy.typing as npt - from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata + from zarr_metadata import ( + ZarrV2ArrayMetadata, + ZarrV2ArrayMetadataJSON, + ZarrV3ArrayMetadata, + ZarrV3ArrayMetadataJSON, + ) from zarr.abc.codec import CodecPipeline from zarr.abc.store import Store @@ -6412,6 +6417,8 @@ def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: type ArrayMetadataModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata +type ArrayMetadataModelJSON = ZarrV2ArrayMetadataJSON | ZarrV3ArrayMetadataJSON +"""Array metadata as `zarr_metadata` types it in JSON (zarrista's metadata type).""" """Array metadata as a `zarr_metadata` model: validated JSON, extension names unresolved.""" @@ -6484,9 +6491,9 @@ class AsyncArrayIO: in a store. Construction performs no IO; every read or write of the store is an explicit method call, and every write has an erasing inverse. The metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, - `erase_chunks`, `verify`) never resolve extension names; `getitem` and - `setitem` resolve them into the legacy classes (`metadata_legacy`) on - first use. The bound metadata is assumed, never checked against the + `erase_chunks`, `verify`) never resolve extension names; with the codec + pipelines zarr-python has today, `getitem` and `setitem` resolve them into + the legacy classes (`metadata_legacy`) on first use. The bound metadata is assumed, never checked against the store, and an array opened from metadata that needed a repair stores that repair before its first chunk write. The codec pipeline implementation is chosen from the global configuration when the binding is made; the @@ -6506,7 +6513,7 @@ class AsyncArrayIO: def __init__( self, - metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: @@ -6745,6 +6752,9 @@ def to_array(self) -> AnyAsyncArray: def __getstate__(self) -> dict[str, object]: state = dict(self.__dict__) + # The pipeline is rebuilt on first use: a built one is bound to this + # process's store object, and may not pickle. + state["_codec_pipeline_cache"] = None state.update(self._carried_cache_entry()) return state @@ -6797,13 +6807,20 @@ def _derive( derived._bind(source, store_path, config, self._pipeline_class) return derived - def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: + def with_metadata( + self, metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataModelJSON + ) -> Self: """This binding with `metadata` in place of its own. Nothing is stored; - a legacy `metadata` object is read as the constructor reads it, then - copied.""" + a legacy `metadata` object or a metadata dict is read as the + constructor reads it, and a legacy object is copied.""" + source: ArrayMetadataModel | ArrayMetadata if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): - metadata = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) - return self._derive(metadata, self._store_path, self._config) + source = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) + elif isinstance(metadata, dict): + source = _bound_metadata_dict(cast("dict[str, Any]", metadata), str(self._store_path)) + else: + source = metadata + return self._derive(source, self._store_path, self._config) def with_store(self, store: Store | StorePath) -> Self: """This binding at its path within `store`, or at a `StorePath` as given. Nothing is stored.""" @@ -7024,7 +7041,7 @@ class ArrayIO: def __init__( self, - metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: @@ -7098,7 +7115,9 @@ def to_array(self) -> AnyArray: """An `Array` over this binding; both share one state.""" return Array(self._async_io.to_array()) - def with_metadata(self, metadata: ArrayMetadataModel | ArrayMetadata) -> Self: + def with_metadata( + self, metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataModelJSON + ) -> Self: """This binding with `metadata` in place of its own. See `AsyncArrayIO.with_metadata`.""" return self._wrap(self._async_io.with_metadata(metadata)) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 7c9c229fe3..40a91c2ca0 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -2,6 +2,7 @@ from __future__ import annotations +import pickle from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -463,3 +464,27 @@ def test_metadata_dict_with_storage_transformers_raises() -> None: } with pytest.raises(ValueError, match="storage transformers"): AsyncArrayIO(cast("Any", doc), StorePath(MemoryStore())) + + +async def test_bindings_take_metadata_json() -> None: + """A binding takes array metadata as `zarr_metadata` types it in JSON (what + zarrista's arrays hold), through the constructor and with_metadata.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + json_metadata = aio.metadata.to_json() + rebound = aio.with_metadata(json_metadata) + assert AsyncArrayIO(json_metadata, aio.store_path) == aio + assert rebound == aio + await rebound.store_metadata() + await rebound.setitem(Ellipsis, np.arange(4, dtype="uint8")) + np.testing.assert_array_equal(await aio.getitem(Ellipsis), np.arange(4)) + + +async def test_pickled_binding_rebuilds_its_pipeline() -> None: + aio = AsyncArrayIO.build_v3( + MemoryStore(), name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,) + ) + _ = await aio.getitem(Ellipsis) + revived = pickle.loads(pickle.dumps(aio)) + assert revived._codec_pipeline_cache is None + assert revived == aio From 0f4bf4576316313cc45e79805cb5af40a6270e41 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 15:53:23 +0100 Subject: [PATCH 30/48] fix(array): leave groups to erase, copy given metadata, name verify's fields erase_metadata deleted the array documents of its format at its location even when they were a group's (a v3 group's zarr.json, a v2 group's .zattrs); it now deletes only an array's, and storing an attribute-less v2 model leaves a group's .zattrs in place. The AsyncArrayIO constructor copies a legacy object it is given, as with_metadata and the derivations do, while AsyncArray keeps binding the object itself through _from_resolved. A metadata dict whose reading moves chunks binds as repaired legacy metadata and is recorded as the stored document, also where the model accepts it (a v2 chunk size of 0). verify names the fields that differ as the model compares them. build_array_metadata_v3 keeps an empty dimension_names list (the creation paths still read it as none, as they always have), erase_chunk takes integer coordinates only, and binding equality no longer depends on what has been resolved. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 129 +++++++++++++++++++++++++++++++---------- tests/test_array_io.py | 101 ++++++++++++++++++++++++++++++++ 2 files changed, 200 insertions(+), 30 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index ebc9b2d7f0..253d7392f1 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -5,6 +5,7 @@ import functools import json import math +import operator import os import warnings from asyncio import gather @@ -69,6 +70,7 @@ ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON, + ZGROUP_JSON, ChunksLike, DimensionNamesLike, MemoryOrder, @@ -472,7 +474,9 @@ def __init__( # A document given as a dict stands for what the store holds, as the # document an array is opened from does. _record_given_document(store_path, metadata, metadata_parsed.zarr_format) - object.__setattr__(self, "_io", AsyncArrayIO(metadata_parsed, store_path, config)) + object.__setattr__( + self, "_io", AsyncArrayIO._from_resolved(metadata_parsed, store_path, config) + ) # An AsyncArray is the resolving layer: build its codec pipeline now, so # pipeline-construction warnings fire at creation, as they always have. _ = self._io._resolved_codec_pipeline @@ -675,7 +679,8 @@ def _create_metadata_v3( fill_value=fill_value, chunk_key_encoding=chunk_key_encoding, codecs=codecs, - dimension_names=dimension_names, + # As these paths always have, an empty list declares no names. + dimension_names=dimension_names or None, attributes=attributes, ) @@ -4759,7 +4764,8 @@ async def init_array( chunk_grid=grid, chunk_key_encoding=chunk_key_encoding_parsed, codecs=codecs_out, - dimension_names=dimension_names, + # As these paths always have, an empty list declares no names. + dimension_names=dimension_names or None, attributes=attributes, ) @@ -4900,7 +4906,7 @@ def build_array_metadata_v3( chunk_key_encoding=chunk_key_encoding_parsed, fill_value=fill_value_parsed, codecs=codecs_parsed, # type: ignore[arg-type] - dimension_names=tuple(dimension_names) if dimension_names else None, + dimension_names=tuple(dimension_names) if dimension_names is not None else None, attributes=attributes if attributes is not None else {}, ) @@ -6379,7 +6385,9 @@ async def _delete_key(store_path: StorePath, key: str) -> None: """The keys a pickled binding carries its store cache's entry under.""" -def _bound_metadata_dict(data: dict[str, Any], path: str) -> ArrayMetadataModel | ArrayMetadata: +def _bound_metadata_dict( + data: dict[str, Any], store_path: StorePath +) -> ArrayMetadataModel | ArrayMetadata: """A metadata dict given to a binding: read as a model when it is valid JSON metadata, so names resolve only when chunks are read or written; otherwise (metadata built in code, or an invalid document zarr reads as @@ -6387,13 +6395,23 @@ def _bound_metadata_dict(data: dict[str, Any], path: str) -> ArrayMetadataModel has.""" from zarr_metadata import MetadataValidationError + path = str(store_path) + zarr_format: ZarrFormat = 2 if data.get("zarr_format") == 2 else 3 if data.get("storage_transformers"): # Bound as legacy metadata, which rejects them, as AsyncArray always has. return parse_array_metadata(data, path) - try: - return _array_metadata_model(data) - except MetadataValidationError: - return parse_array_metadata(data, path) + if not moves_chunks(data, zarr_format): + # A document zarr reads as repaired, moving chunks, is never bound as a + # model, also where the model accepts it (a v2 chunk size of 0). + try: + return _array_metadata_model(data) + except MetadataValidationError: + pass + legacy = parse_array_metadata(data, path) + # An invalid document given as a dict stands for what the store holds, as + # one given to AsyncArray does. + _record_given_document(store_path, data, legacy.zarr_format) + return legacy def _record_given_document( @@ -6407,6 +6425,37 @@ def _record_given_document( ) +def _differing_fields( + model_class: type[ArrayMetadataModel], + stored: ArrayMetadataModel, + stored_json: dict[str, Any], + bound_json: dict[str, Any], +) -> list[str]: + """The fields of `stored` that differ from `bound_json`'s as the model compares + them: a field whose JSON differs (so `-0.0` from `0.0`) counts unless putting + the bound value in `stored` leaves it equal (an absent field against its + default).""" + from zarr_metadata import MetadataValidationError + + differing = [] + for key in sorted(stored_json.keys() | bound_json.keys()): + stored_value, bound_value = stored_json.get(key), bound_json.get(key) + if json.dumps(stored_value, sort_keys=True) == json.dumps(bound_value, sort_keys=True): + continue + patched = dict(stored_json) + if key in bound_json: + patched[key] = bound_value + else: + patched.pop(key) + try: + unchanged = model_class(patched) == stored + except MetadataValidationError: + unchanged = False + if not unchanged: + differing.append(key) + return differing + + def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: """A copy of a legacy metadata object with its own attributes dict, so an attribute update through one binding (which replaces top-level keys in @@ -6519,13 +6568,15 @@ def __init__( ) -> None: source: ArrayMetadataModel | ArrayMetadata if isinstance(metadata, ArrayMetadata): - # A resolved object passes through (v2 with a chunk size of 0 is read - # as repaired). + # A resolved object is read (v2 with a chunk size of 0 as repaired) + # and never shared with the caller. source = parse_array_metadata(metadata, str(store_path)) + if source is metadata: + source = _copied_legacy(source) elif isinstance(metadata, _model_classes()): source = metadata else: - source = _bound_metadata_dict(cast("dict[str, Any]", metadata), str(store_path)) + source = _bound_metadata_dict(cast("dict[str, Any]", metadata), store_path) # The pipeline implementation is chosen when the binding is made, from # the global configuration, exactly as `AsyncArray` has always chosen # it; the pipeline itself is built on first use. @@ -6547,6 +6598,22 @@ def _bind( self._chunk_key_encoding_cache = None self._codec_pipeline_cache = None + @classmethod + def _from_resolved( + cls, source: ArrayMetadata, store_path: StorePath, config: ArrayConfigLike | None + ) -> Self: + """A binding over `source` itself, already resolved (as `AsyncArray` binds the + metadata it is given): read as `parse_array_metadata` reads a metadata + object, but not copied.""" + binding = object.__new__(cls) + binding._bind( + parse_array_metadata(source, str(store_path)), + store_path, + parse_array_config(config), + get_pipeline_class(), + ) + return binding + @classmethod def build_v2( cls, @@ -6817,7 +6884,7 @@ def with_metadata( if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): source = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) elif isinstance(metadata, dict): - source = _bound_metadata_dict(cast("dict[str, Any]", metadata), str(self._store_path)) + source = _bound_metadata_dict(cast("dict[str, Any]", metadata), self._store_path) else: source = metadata return self._derive(source, self._store_path, self._config) @@ -6864,9 +6931,13 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: ensure_parents=ensure_parents, array=True, ) - if self._zarr_format == 2 and ZATTRS_JSON not in to_save: + if ( + self._zarr_format == 2 + and ZATTRS_JSON not in to_save + and not await (self._store_path / ZGROUP_JSON).exists() + ): # An attribute-less model emits no .zattrs; a stale one must not - # survive an unconditional metadata write. + # survive an unconditional metadata write. A group's is its own. await (self._store_path / ZATTRS_JSON).delete() async def verify(self) -> None: @@ -6891,27 +6962,25 @@ async def verify(self) -> None: if isinstance(stored, _model_classes()[0]) and isinstance(bound, _model_classes()[0]): stored_json["attributes"] = dict(stored.attributes) bound_json["attributes"] = dict(bound.attributes) - differing = sorted( - key - for key in stored_json.keys() | bound_json.keys() - if stored_json.get(key) != bound_json.get(key) - ) + differing = _differing_fields(type(stored), stored, stored_json, bound_json) raise ValueError( f"Array {str(self.store_path)!r}: the stored metadata does not match the " f"bound metadata; differing fields: {', '.join(differing)}." ) async def erase_metadata(self) -> None: - """Delete the stored metadata; chunks are kept. Succeeds if absent.""" - for key in ARRAY_DOCUMENTS[self._zarr_format]: - await (self._store_path / key).delete() - stored_documents(self._store_path.store).observe( - self._store_path.path, None, self._zarr_format - ) + """Delete the stored array metadata of this array's format; chunks are + kept. Succeeds if absent; a group's documents here are left in place.""" + zarr_format = self._zarr_format + keys = ARRAY_DOCUMENTS[zarr_format] + if stored_array_document(await read_documents(self._store_path, keys), zarr_format): + for key in keys: + await (self._store_path / key).delete() + stored_documents(self._store_path.store).observe(self._store_path.path, None, zarr_format) async def erase_chunk(self, chunk_coords: Sequence[int], /) -> None: """Delete one stored chunk by grid coordinates (the shard, under sharding). Succeeds if absent.""" - chunk_coords = tuple(chunk_coords) + chunk_coords = tuple(operator.index(coord) for coord in chunk_coords) grid_shape = self._shard_grid_shape if len(chunk_coords) != len(grid_shape) or any( coord < 0 or coord >= size for coord, size in zip(chunk_coords, grid_shape, strict=True) @@ -7023,9 +7092,9 @@ def __eq__(self, other: object) -> bool: return NotImplemented if self._store_path != other._store_path or self._config != other._config: return False - resolved, other_resolved = self._resolved_if_any, other._resolved_if_any - if resolved is not None and other_resolved is not None: - return resolved == other_resolved + legacy = (ArrayV2Metadata, ArrayV3Metadata) + if isinstance(self._source, legacy) and isinstance(other._source, legacy): + return self._source == other._source return self.metadata == other.metadata def __hash__(self) -> int: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 40a91c2ca0..8b05def2c3 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import pickle from typing import TYPE_CHECKING, Any, cast @@ -488,3 +489,103 @@ async def test_pickled_binding_rebuilds_its_pipeline() -> None: revived = pickle.loads(pickle.dumps(aio)) assert revived._codec_pipeline_cache is None assert revived == aio + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_erase_and_store_metadata_leave_a_group_in_place(zarr_format: ZarrFormat) -> None: + """erase_metadata deletes only an array's documents, and storing a v2 model + without attributes keeps a group's `.zattrs`.""" + store = MemoryStore() + group = await zarr.api.asynchronous.open_group( + store, path="g", zarr_format=zarr_format, mode="w" + ) + await group.update_attributes({"owner": "group"}) + before = await _store_keys(store) + declared = ( + AsyncArrayIO.build_v2(store, name="g", shape=(4,), dtype="uint8", chunks=(4,)) + if zarr_format == 2 + else AsyncArrayIO.build_v3(store, name="g", shape=(4,), data_type="uint8", chunk_grid=(4,)) + ) + await declared.erase_metadata() + assert await _store_keys(store) == before + if zarr_format == 2: + model = declared.metadata + assert isinstance(model, ModelV2) + await declared.with_metadata(model.update(attributes=UNSET)).store_metadata() + reopened = await zarr.api.asynchronous.open_group(store, path="g", mode="r", zarr_format=2) + assert dict(reopened.attrs) == {"owner": "group"} + + +def test_constructor_never_shares_the_callers_metadata() -> None: + mine = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,)).metadata + aio = AsyncArrayIO(mine, StorePath(MemoryStore())) + assert aio.metadata_legacy is not mine + assert aio.metadata_legacy.attributes is not mine.attributes + + +async def test_metadata_dict_of_a_repaired_document_stores_its_repair() -> None: + """An invalid stored document given to AsyncArrayIO as a dict is recorded as + what the store holds, as AsyncArray records it, so the first chunk write + stores the repair.""" + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 2, + "shape": [3], + "chunks": [0], + "dtype": " None: + store = MemoryStore() + aio = AsyncArrayIO.build_v3( + store, name="foo", shape=(4,), data_type="float32", chunk_grid=(4,), fill_value=0.0 + ) + await aio.store_metadata() + other = aio.with_metadata(aio.metadata.update(fill_value=-0.0)) + with pytest.raises(ValueError, match=r"differing fields: fill_value\.$"): + await other.verify() + + +def test_build_v3_keeps_empty_dimension_names_on_a_0d_array() -> None: + """The builder keeps an empty name list for a 0-d array; the creation paths + read it as no names, as they always have.""" + built = ArrayIO.build_v3( + MemoryStore(), shape=(), data_type="int8", chunk_grid=(), dimension_names=[] + ) + assert built.metadata_legacy.dimension_names == () # type: ignore[union-attr] + created = zarr.create_array(MemoryStore(), shape=(), dtype="int8", dimension_names=[]) + assert created.metadata.dimension_names is None # type: ignore[union-attr] + + +async def test_erase_chunk_takes_integer_coordinates() -> None: + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8")) + with pytest.raises(TypeError): + await aio.erase_chunk([cast("int", np.float64(1))]) + await aio.erase_chunk([cast("int", np.int64(1))]) + assert "foo/c/1" not in await _store_keys(store) + + +def test_equality_does_not_depend_on_resolution() -> None: + doc = AsyncArrayIO.build_v3( + MemoryStore(), name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).metadata.to_json() + a, b = (AsyncArrayIO(doc, StorePath(MemoryStore(), "foo")) for _ in "ab") + a = b.with_store(a.store_path) + before = a == b + _ = a.metadata_legacy + assert (a == b) is before From 4be3936fcc98683f258d4d56cf7d46e801975b36 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 16:34:24 +0100 Subject: [PATCH 31/48] fix(array): keep 3.4.x pickles repairing, and group writes as fast as before zarr 3.4.1 and 3.4.2 released the repair mark on metadata objects (`_stored_document`), so pickles cross versions: an array pickled by them now has its mark read into the store's cache, a consolidated group's members likewise, and an array pickled here marks a copy of its metadata the same way for those versions to read. A v3 group write no longer encodes its consolidated metadata twice (30% of every write on a large consolidated group), and the check for whether a document's repair moves chunks returns early for documents whose regular chunk sizes are all positive, so reads of consolidated metadata cost what they did. AsyncArray and AsyncArrayIO raise TypeError for metadata that is not a mapping, as AsyncArray always did. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 22 +++++++++++++++++++--- src/zarr/core/group.py | 20 ++++++++++++++++++-- src/zarr/core/metadata/repair.py | 19 +++++++++++++++++++ src/zarr/core/metadata/stored.py | 4 +++- tests/test_array.py | 9 ++++++++- tests/test_metadata/test_repair.py | 30 ++++++++++++++++++++++++++++++ 6 files changed, 97 insertions(+), 7 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 253d7392f1..f004febdcb 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -496,15 +496,23 @@ def io(self) -> AsyncArrayIO: def __getstate__(self) -> dict[str, object]: # The pickled state is the fields', as it has always been, so a pickle # loads in either version; what the store's cache knows rides along. + carried = self._io._carried_cache_entry() + metadata = self.metadata + mark = carried.get("_stored_entry", carried.get("_stored_copy")) + if mark is not None: + # zarr 3.4.1 and 3.4.2 read the repair a stored document needs from + # this mark on the metadata itself; this version reads the carried entry. + metadata = copy.copy(metadata) + object.__setattr__(metadata, "_stored_document", copy.deepcopy(mark)) state: dict[str, object] = { - "metadata": self.metadata, + "metadata": metadata, "store_path": self.store_path, "codec_pipeline": self.codec_pipeline, "_chunk_grid": self._chunk_grid, "config": self.config, "_pipeline_class": self._io._pipeline_class, } - state.update(self._io._carried_cache_entry()) + state.update(carried) return state def __setstate__(self, state: dict[str, Any]) -> None: @@ -517,6 +525,11 @@ def __setstate__(self, state: dict[str, Any]) -> None: io._bind(metadata, state["store_path"], state["config"], pipeline_class) io._codec_pipeline_cache = pipeline io._chunk_grid_cache = state["_chunk_grid"] + if not any(key in state for key in _CARRIED_KEYS): + # A pickle from zarr 3.4.1 or 3.4.2 marks the metadata itself. + legacy_mark = vars(metadata).get("_stored_document") + if legacy_mark is not None: + state = {**state, "_stored_entry": legacy_mark} io._restore_cache_entry(state) object.__setattr__(self, "_io", io) @@ -6575,8 +6588,11 @@ def __init__( source = _copied_legacy(source) elif isinstance(metadata, _model_classes()): source = metadata - else: + elif isinstance(metadata, Mapping): source = _bound_metadata_dict(cast("dict[str, Any]", metadata), store_path) + else: + # Raises the TypeError AsyncArray always raised for such input. + source = parse_array_metadata(metadata, str(store_path)) # The pipeline implementation is chosen when the binding is made, from # the global configuration, exactly as `AsyncArray` has always chosen # it; the pipeline itself is built on first use. diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index a65cf0ec52..3fa7f33d3a 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -410,7 +410,8 @@ def to_buffer_dict( indent = config.get("json_indent") if self.zarr_format == 3: document = self.to_dict() - if self.consolidated_metadata is not None: + if self.consolidated_metadata is not None and as_stored: + # Encoded again only where a member is kept as stored. document["consolidated_metadata"] = self.consolidated_metadata.to_dict( members_as_stored=as_stored ) @@ -710,11 +711,26 @@ def __setstate__(self, state: dict[str, object]) -> None: state = dict(state) entries = cast("dict[str, CarriedEntry] | None", state.pop("_stored_members", None)) self.__dict__.update(state) - if entries is not None: + if entries is None: + entries = self._legacy_member_marks() + if entries: stored_documents(self.store_path.store).restore_entries( entries, self.metadata.zarr_format ) + def _legacy_member_marks(self) -> dict[str, CarriedEntry]: + """The consolidated members a pickle from zarr 3.4.1 or 3.4.2 marks as read + from a document that needs a repair, as the entries this version carries.""" + consolidated = self.metadata.consolidated_metadata + if consolidated is None: + return {} + marks: dict[str, CarriedEntry] = {} + for key, member in consolidated.flattened_metadata.items(): + mark = vars(member).get("_stored_document") + if mark is not None: + marks[_join_paths([self.store_path.path, key])] = ("copy", mark) + return marks + @classmethod def _from_bytes_v2( cls, diff --git a/src/zarr/core/metadata/repair.py b/src/zarr/core/metadata/repair.py index 2d6daddb30..e574da38a4 100644 --- a/src/zarr/core/metadata/repair.py +++ b/src/zarr/core/metadata/repair.py @@ -355,6 +355,25 @@ def _invalid_edge_lengths_v3(doc: ArrayDocument) -> tuple[ArrayDocument, Reading """The repairs of an array document of each Zarr format, applied in order.""" +def may_move_chunks(doc: ArrayDocument, zarr_format: ZarrFormat) -> bool: + """Whether some repair might move chunks in `doc`: `False` only when every + stored regular chunk size is a JSON int >= 1 or `true`, for only a stored 0 (or + `false`) or an edge list in a regular grid is read so that it moves chunks (see + `_read_chunk_size` and `_invalid_chunk_sizes_v3`). A cheap check that answers + for most documents before `repair_array_document` runs.""" + if zarr_format == 2: + sizes = doc.get("chunks") + else: + grid = doc.get("chunk_grid") + if not isinstance(grid, Mapping) or grid.get("name") != "regular": + return True + configuration = grid.get("configuration") + sizes = configuration.get("chunk_shape") if isinstance(configuration, Mapping) else None + if not isinstance(sizes, list | tuple): + return True + return not all(size is True or (type(size) is int and size >= 1) for size in sizes) + + def repair_array_document( doc: ArrayDocument, zarr_format: ZarrFormat ) -> tuple[ArrayDocument, list[Reading]]: diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index 551e8378f7..20051fde4a 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -20,7 +20,7 @@ from zarr.core._json import buffer_to_json_object from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON -from zarr.core.metadata.repair import repair_array_document +from zarr.core.metadata.repair import may_move_chunks, repair_array_document if TYPE_CHECKING: from collections.abc import Iterable, Mapping @@ -40,6 +40,8 @@ def moves_chunks(document: ArrayDocument, zarr_format: ZarrFormat) -> bool: """Whether reading `document` needs a repair that moves chunks.""" + if not may_move_chunks(document, zarr_format): + return False _, readings = repair_array_document(document, zarr_format) return any(reading.moves_chunks for reading in readings) diff --git a/tests/test_array.py b/tests/test_array.py index 8fd0ecd3f3..f47a4c171d 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -7,7 +7,7 @@ import re import sys from itertools import accumulate, starmap -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, cast from unittest import mock import numcodecs @@ -2754,3 +2754,10 @@ def test_import_zarr_leaves_zarr_metadata_unimported() -> None: code = "import sys, zarr; assert 'zarr_metadata' not in sys.modules" subprocess.run([sys.executable, "-c", code], check=True) + + +def test_async_array_rejects_metadata_that_is_not_a_mapping() -> None: + with pytest.raises(TypeError): + AsyncArray(cast("Any", 42), StorePath(MemoryStore())) + with pytest.raises(TypeError): + zarr.AsyncArrayIO(cast("Any", 42), StorePath(MemoryStore())) diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index d969b0c5d2..e775e9156e 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -1959,3 +1959,33 @@ def write_through(sub_of: Callable[[zarr.Group], zarr.Group], name: str) -> dict assert write_through(lambda sub: pickle.loads(pickle.dumps(sub)), "pickled") == write_through( lambda sub: sub, "original" ) + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_pickles_cross_versions_with_the_repair(tmp_path: Path, zarr_format: Literal[2, 3]) -> None: + """An array pickled by zarr 3.4.1 or 3.4.2 marks its metadata with the stored + document it was read from; this version reads that mark, and marks the metadata + it pickles the same way, so a write after unpickling stores the repair in + either direction.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, zarr_format) + doc_name = ".zarray" if zarr_format == 2 else "zarr.json" + stored = json.loads((path / doc_name).read_text()) + array = zarr.open_array(store=LocalStore(path), mode="r+") + + state = array.async_array.__getstate__() + assert vars(state["metadata"])["_stored_document"] is not None + + old_state = { + key: state[key] for key in ("store_path", "codec_pipeline", "_chunk_grid", "config") + } + old_metadata = copy.copy(array.metadata) + object.__setattr__(old_metadata, "_stored_document", stored) + old_state["metadata"] = old_metadata + revived = object.__new__(type(array.async_array)) + revived.__setstate__( + pickle.loads(pickle.dumps({**old_state, "store_path": StorePath(LocalStore(path))})) + ) + sync(revived.setitem(slice(None), np.array([7, 8, 9], dtype="int16"))) + assert _stored_chunks(json.loads((path / doc_name).read_text())) == [1] From b6da26f0516497530c27f76ca6cfb5da23b01a12 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 16:46:43 +0100 Subject: [PATCH 32/48] fix(array): reject storage transformers on every bind, erase corrupt documents ArrayIO.open and a model given to the constructor bound metadata declaring storage transformers and failed only at the first chunk operation; every model is now checked when bound, as dicts and AsyncArray are. erase_metadata no longer parses documents before deleting them, so a corrupt or node_type-less array document is removed rather than raising or surviving; only a group's documents are left (a v2 `.zattrs` beside a `.zgroup`, a v3 group's `zarr.json`). Pickles no longer carry a stale 3.4.x mark, and a consolidated group pickled here marks its members for 3.4.1 and 3.4.2 to read. Binding equality is always the models' canonical equality, the rectilinear grid skips the repair scan it never needed, AsyncArrayIO names the type it was given when it cannot bind it, and with_metadata takes any mapping. Tests pin these and check may_move_chunks against the repairs. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 84 +++++++++++++++++------ src/zarr/core/group.py | 47 +++++++++++++ src/zarr/core/metadata/repair.py | 11 +-- tests/test_array.py | 9 ++- tests/test_array_io.py | 105 +++++++++++++++++++++++++---- tests/test_metadata/test_repair.py | 99 +++++++++++++++++++++------ 6 files changed, 293 insertions(+), 62 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index f004febdcb..4248b295fe 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -266,7 +266,7 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: return ArrayV2Metadata.from_dict(data, path=path) else: raise ValueError(f"Invalid zarr_format: {zarr_format}. Expected 2 or 3") - raise TypeError # pragma: no cover + raise TypeError def create_codec_pipeline( @@ -497,13 +497,12 @@ def __getstate__(self) -> dict[str, object]: # The pickled state is the fields', as it has always been, so a pickle # loads in either version; what the store's cache knows rides along. carried = self._io._carried_cache_entry() - metadata = self.metadata + metadata: ArrayMetadata = self.metadata mark = carried.get("_stored_entry", carried.get("_stored_copy")) - if mark is not None: + if mark is not None or "_stored_document" in vars(metadata): # zarr 3.4.1 and 3.4.2 read the repair a stored document needs from # this mark on the metadata itself; this version reads the carried entry. - metadata = copy.copy(metadata) - object.__setattr__(metadata, "_stored_document", copy.deepcopy(mark)) + metadata = _with_legacy_mark(metadata, cast("ArrayDocument | None", mark)) state: dict[str, object] = { "metadata": metadata, "store_path": self.store_path, @@ -530,6 +529,9 @@ def __setstate__(self, state: dict[str, Any]) -> None: legacy_mark = vars(metadata).get("_stored_document") if legacy_mark is not None: state = {**state, "_stored_entry": legacy_mark} + if "_stored_document" in vars(metadata): + # Read once; the cache holds it from here on. + object.__delattr__(metadata, "_stored_document") io._restore_cache_entry(state) object.__setattr__(self, "_io", io) @@ -6439,7 +6441,6 @@ def _record_given_document( def _differing_fields( - model_class: type[ArrayMetadataModel], stored: ArrayMetadataModel, stored_json: dict[str, Any], bound_json: dict[str, Any], @@ -6461,7 +6462,7 @@ def _differing_fields( else: patched.pop(key) try: - unchanged = model_class(patched) == stored + unchanged = type(stored)(patched) == stored except MetadataValidationError: unchanged = False if not unchanged: @@ -6469,6 +6470,38 @@ def _differing_fields( return differing +def _reject_storage_transformers(model: ArrayMetadataModel) -> None: + """Raise, as `parse_array_metadata` does, for a model that declares storage + transformers: zarr-python cannot read or write chunks through them.""" + transformers = getattr(model, "storage_transformers", ()) + if transformers: + raise ValueError( + f"Array metadata contains storage transformers: {transformers}." + "Arrays with storage transformers are not supported in zarr-python at this time." + ) + + +def _is_group_document(document: Buffer) -> bool: + """Whether a stored `zarr.json` is a group's; one that is not JSON is not.""" + try: + decoded = buffer_to_json_object(document) + except (ValueError, TypeError): + return False + return decoded.get("node_type") == "group" + + +def _with_legacy_mark(metadata: ArrayMetadata, mark: ArrayDocument | None) -> ArrayMetadata: + """A copy of `metadata` carrying `mark` as the `_stored_document` attribute zarr + 3.4.1 and 3.4.2 read the repair a stored document needs from (none: no mark).""" + marked = copy.copy(metadata) + if mark is None: + if "_stored_document" in vars(marked): + object.__delattr__(marked, "_stored_document") + else: + object.__setattr__(marked, "_stored_document", copy.deepcopy(mark)) + return marked + + def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: """A copy of a legacy metadata object with its own attributes dict, so an attribute update through one binding (which replaces top-level keys in @@ -6479,9 +6512,9 @@ def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: type ArrayMetadataModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata +"""Array metadata as a `zarr_metadata` model: validated JSON, extension names unresolved.""" type ArrayMetadataModelJSON = ZarrV2ArrayMetadataJSON | ZarrV3ArrayMetadataJSON """Array metadata as `zarr_metadata` types it in JSON (zarrista's metadata type).""" -"""Array metadata as a `zarr_metadata` model: validated JSON, extension names unresolved.""" def _metadata_model_json(document: ArrayMetadataModel) -> dict[str, Any]: @@ -6591,8 +6624,12 @@ def __init__( elif isinstance(metadata, Mapping): source = _bound_metadata_dict(cast("dict[str, Any]", metadata), store_path) else: - # Raises the TypeError AsyncArray always raised for such input. - source = parse_array_metadata(metadata, str(store_path)) + raise TypeError( + "Expected array metadata (a zarr_metadata model, ArrayV2Metadata or " + f"ArrayV3Metadata) or a metadata mapping, got {type(metadata).__name__}." + ) + if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + _reject_storage_transformers(source) # The pipeline implementation is chosen when the binding is made, from # the global configuration, exactly as `AsyncArray` has always chosen # it; the pipeline itself is built on first use. @@ -6891,7 +6928,8 @@ def _derive( return derived def with_metadata( - self, metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataModelJSON + self, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, ) -> Self: """This binding with `metadata` in place of its own. Nothing is stored; a legacy `metadata` object or a metadata dict is read as the @@ -6899,8 +6937,8 @@ def with_metadata( source: ArrayMetadataModel | ArrayMetadata if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): source = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) - elif isinstance(metadata, dict): - source = _bound_metadata_dict(cast("dict[str, Any]", metadata), self._store_path) + elif isinstance(metadata, Mapping) and not isinstance(metadata, _model_classes()): + source = _bound_metadata_dict(dict(metadata), self._store_path) else: source = metadata return self._derive(source, self._store_path, self._config) @@ -6978,7 +7016,7 @@ async def verify(self) -> None: if isinstance(stored, _model_classes()[0]) and isinstance(bound, _model_classes()[0]): stored_json["attributes"] = dict(stored.attributes) bound_json["attributes"] = dict(bound.attributes) - differing = _differing_fields(type(stored), stored, stored_json, bound_json) + differing = _differing_fields(stored, stored_json, bound_json) raise ValueError( f"Array {str(self.store_path)!r}: the stored metadata does not match the " f"bound metadata; differing fields: {', '.join(differing)}." @@ -6988,10 +7026,14 @@ async def erase_metadata(self) -> None: """Delete the stored array metadata of this array's format; chunks are kept. Succeeds if absent; a group's documents here are left in place.""" zarr_format = self._zarr_format - keys = ARRAY_DOCUMENTS[zarr_format] - if stored_array_document(await read_documents(self._store_path, keys), zarr_format): - for key in keys: - await (self._store_path / key).delete() + if zarr_format == 2: + await (self._store_path / ZARRAY_JSON).delete() + if not await (self._store_path / ZGROUP_JSON).exists(): + await (self._store_path / ZATTRS_JSON).delete() + else: + document = await (self._store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype) + if document is not None and not _is_group_document(document): + await (self._store_path / ZARR_JSON).delete() stored_documents(self._store_path.store).observe(self._store_path.path, None, zarr_format) async def erase_chunk(self, chunk_coords: Sequence[int], /) -> None: @@ -7108,9 +7150,6 @@ def __eq__(self, other: object) -> bool: return NotImplemented if self._store_path != other._store_path or self._config != other._config: return False - legacy = (ArrayV2Metadata, ArrayV3Metadata) - if isinstance(self._source, legacy) and isinstance(other._source, legacy): - return self._source == other._source return self.metadata == other.metadata def __hash__(self) -> int: @@ -7201,7 +7240,8 @@ def to_array(self) -> AnyArray: return Array(self._async_io.to_array()) def with_metadata( - self, metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataModelJSON + self, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, ) -> Self: """This binding with `metadata` in place of its own. See `AsyncArrayIO.with_metadata`.""" return self._wrap(self._async_io.with_metadata(metadata)) diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 3fa7f33d3a..99238f0eaa 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import copy import logging import unicodedata import warnings @@ -143,6 +144,28 @@ def _parse_async_node( raise TypeError(f"Unknown node type, got {type(node)}") +def _with_member_marks( + consolidated: ConsolidatedMetadata, prefix: str, marks: Mapping[str, ArrayDocument] +) -> ConsolidatedMetadata: + """A copy of `consolidated` whose array members at the paths in `marks` (relative + to the group, below `prefix`) are copies carrying their mark as the + `_stored_document` attribute zarr 3.4.1 and 3.4.2 read.""" + members: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata] = {} + for name, member in consolidated.metadata.items(): + key = f"{prefix}{name}" + if isinstance(member, GroupMetadata): + nested = member.consolidated_metadata + if nested is not None: + member = replace( + member, consolidated_metadata=_with_member_marks(nested, f"{key}/", marks) + ) + elif key in marks: + member = copy.copy(member) + object.__setattr__(member, "_stored_document", copy.deepcopy(marks[key])) + members[name] = member + return replace(consolidated, metadata=members) + + @dataclass(frozen=True) class ConsolidatedMetadata: """ @@ -705,6 +728,19 @@ def __getstate__(self) -> dict[str, object]: ) if entries: state["_stored_members"] = entries + # zarr 3.4.1 and 3.4.2 read a member's repair from a mark on the + # member itself: pickle marked copies for them to read. + prefix = f"{self.store_path.path}/" if self.store_path.path else "" + marks = { + path.removeprefix(prefix): entry[1] + for path, entry in entries.items() + if entry[1] is not None + } + if marks: + state["metadata"] = replace( + self.metadata, + consolidated_metadata=_with_member_marks(consolidated, "", marks), + ) return state def __setstate__(self, state: dict[str, object]) -> None: @@ -713,6 +749,7 @@ def __setstate__(self, state: dict[str, object]) -> None: self.__dict__.update(state) if entries is None: entries = self._legacy_member_marks() + self._clear_member_marks() if entries: stored_documents(self.store_path.store).restore_entries( entries, self.metadata.zarr_format @@ -731,6 +768,16 @@ def _legacy_member_marks(self) -> dict[str, CarriedEntry]: marks[_join_paths([self.store_path.path, key])] = ("copy", mark) return marks + def _clear_member_marks(self) -> None: + """Drop the marks a pickle leaves on consolidated members: once read, the + store's cache holds them.""" + consolidated = self.metadata.consolidated_metadata + if consolidated is None: + return + for member in consolidated.flattened_metadata.values(): + if "_stored_document" in vars(member): + object.__delattr__(member, "_stored_document") + @classmethod def _from_bytes_v2( cls, diff --git a/src/zarr/core/metadata/repair.py b/src/zarr/core/metadata/repair.py index e574da38a4..839df71ea0 100644 --- a/src/zarr/core/metadata/repair.py +++ b/src/zarr/core/metadata/repair.py @@ -356,15 +356,18 @@ def _invalid_edge_lengths_v3(doc: ArrayDocument) -> tuple[ArrayDocument, Reading def may_move_chunks(doc: ArrayDocument, zarr_format: ZarrFormat) -> bool: - """Whether some repair might move chunks in `doc`: `False` only when every - stored regular chunk size is a JSON int >= 1 or `true`, for only a stored 0 (or - `false`) or an edge list in a regular grid is read so that it moves chunks (see - `_read_chunk_size` and `_invalid_chunk_sizes_v3`). A cheap check that answers + """Whether some repair might move chunks in `doc`: `False` for a rectilinear + grid, and when every stored regular chunk size is a JSON int >= 1 or `true`, + for only a stored 0 (or `false`) or an edge list in a regular grid is read so + that it moves chunks (see `_read_chunk_size` and `_invalid_chunk_sizes_v3`). A cheap check that answers for most documents before `repair_array_document` runs.""" if zarr_format == 2: sizes = doc.get("chunks") else: grid = doc.get("chunk_grid") + if isinstance(grid, Mapping) and grid.get("name") == "rectilinear": + # Its only repair (`_invalid_edge_lengths_v3`) respells edges in place. + return False if not isinstance(grid, Mapping) or grid.get("name") != "regular": return True configuration = grid.get("configuration") diff --git a/tests/test_array.py b/tests/test_array.py index f47a4c171d..40bde50783 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -2750,10 +2750,13 @@ def test_import_zarr_leaves_zarr_metadata_unimported() -> None: """`zarr_metadata` is imported when a binding first holds a model, not by `import zarr`, whose cost stays what it was.""" import subprocess - import sys - code = "import sys, zarr; assert 'zarr_metadata' not in sys.modules" - subprocess.run([sys.executable, "-c", code], check=True) + code = "import sys, zarr; print('zarr_metadata' in sys.modules)" + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=False + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False" def test_async_array_rejects_metadata_that_is_not_a_mapping() -> None: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 8b05def2c3..978c0226cf 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -4,6 +4,7 @@ import json import pickle +import tempfile from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -19,7 +20,7 @@ from zarr.core.metadata import ArrayV2Metadata from zarr.dtype import Int16 from zarr.errors import ArrayNotFoundError, UnknownCodecError -from zarr.storage import MemoryStore, StorePath +from zarr.storage import LocalStore, MemoryStore, StorePath if TYPE_CHECKING: from collections.abc import Callable @@ -482,13 +483,16 @@ async def test_bindings_take_metadata_json() -> None: async def test_pickled_binding_rebuilds_its_pipeline() -> None: + """A pickled binding drops its built pipeline (a pipeline assigned to it is not + carried either) and builds its class's pipeline on first use.""" aio = AsyncArrayIO.build_v3( - MemoryStore(), name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,) + LocalStore(tempfile.mkdtemp()), name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,) ) - _ = await aio.getitem(Ellipsis) + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8")) revived = pickle.loads(pickle.dumps(aio)) assert revived._codec_pipeline_cache is None - assert revived == aio + np.testing.assert_array_equal(await revived.getitem(Ellipsis), np.arange(4)) + assert isinstance(revived._resolved_codec_pipeline, revived._pipeline_class) @pytest.mark.parametrize("zarr_format", [2, 3]) @@ -580,12 +584,89 @@ async def test_erase_chunk_takes_integer_coordinates() -> None: assert "foo/c/1" not in await _store_keys(store) -def test_equality_does_not_depend_on_resolution() -> None: - doc = AsyncArrayIO.build_v3( +def test_equality_is_canonical_whatever_is_resolved() -> None: + """Bindings compare by canonical metadata, before and after resolving, and + whatever they were built from: a fill value of -0.0 differs from 0.0.""" + store_path = StorePath(MemoryStore(), "foo") + positive = AsyncArrayIO.build_v3( + store_path, shape=(4,), data_type="float32", chunk_grid=(4,), fill_value=0.0 + ) + negative = positive.with_metadata(positive.metadata.update(fill_value=-0.0)) + legacy_negative = AsyncArrayIO(negative.metadata_legacy, store_path) + assert positive != negative + assert positive != legacy_negative + assert negative == legacy_negative + _ = (positive.metadata_legacy, negative.metadata_legacy) + assert positive != negative + assert negative == legacy_negative + + +@pytest.mark.parametrize("bind", ["open", "model"]) +async def test_storage_transformers_are_rejected_when_bound(bind: str) -> None: + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "uint8", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes"}], + "storage_transformers": [{"name": "nope"}], + } + if bind == "open": + prototype = default_buffer_prototype() + await store.set("foo/zarr.json", prototype.buffer.from_bytes(json.dumps(doc).encode())) + with pytest.raises(ValueError, match="storage transformers"): + await AsyncArrayIO.open(store, name="foo") + else: + model = ZarrV3ArrayMetadata(doc) + with pytest.raises(ValueError, match="storage transformers"): + AsyncArrayIO(model, StorePath(store, "foo")) + + +@pytest.mark.parametrize( + ("documents", "left"), + [ + ({"foo/zarr.json": b"not json"}, set()), + ({"foo/zarr.json": b'{"zarr_format": 3}'}, set()), + ({"foo/zarr.json": b'{"zarr_format": 3, "node_type": "group"}'}, {"foo/zarr.json"}), + ({"foo/.zattrs": b"{}"}, set()), + ( + {"foo/.zarray": b"{}", "foo/.zattrs": b"{}", "foo/.zgroup": b"{}"}, + {"foo/.zattrs", "foo/.zgroup"}, + ), + ], + ids=["corrupt-v3", "v3-without-node-type", "v3-group", "stray-v2-zattrs", "v2-array-and-group"], +) +async def test_erase_metadata_removes_array_documents_however_written( + documents: dict[str, bytes], left: set[str] +) -> None: + """erase_metadata removes array documents even when they are corrupt, and + leaves only a group's documents in place.""" + store = MemoryStore() + prototype = default_buffer_prototype() + for key, value in documents.items(): + await store.set(key, prototype.buffer.from_bytes(value)) + zarr_format: ZarrFormat = 3 if "foo/zarr.json" in documents else 2 + declared = ( + AsyncArrayIO.build_v2(store, name="foo", shape=(4,), dtype="uint8", chunks=(4,)) + if zarr_format == 2 + else AsyncArrayIO.build_v3( + store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + ) + await declared.erase_metadata() + assert await _store_keys(store) == left + + +async def test_with_metadata_takes_any_mapping() -> None: + from types import MappingProxyType + + aio = AsyncArrayIO.build_v3( MemoryStore(), name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,) - ).metadata.to_json() - a, b = (AsyncArrayIO(doc, StorePath(MemoryStore(), "foo")) for _ in "ab") - a = b.with_store(a.store_path) - before = a == b - _ = a.metadata_legacy - assert (a == b) is before + ) + rebound = aio.with_metadata(cast("Any", MappingProxyType(dict(aio.metadata.to_json())))) + assert rebound == aio + await rebound.store_metadata() diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index e775e9156e..cab7f07544 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -13,8 +13,10 @@ import zipfile from typing import TYPE_CHECKING, Any, Literal, cast +import hypothesis.strategies as st import numpy as np import pytest +from hypothesis import given, settings import zarr from zarr.codecs import ShardingCodec @@ -25,6 +27,7 @@ from zarr.core.metadata.repair import ( RECREATE_HINT, RESAVE_HINT, + may_move_chunks, repair_array_document, ) from zarr.core.metadata.stored import _CACHES, StoredDocuments, moves_chunks, stored_documents @@ -1936,29 +1939,24 @@ def test_member_read_valid_then_from_a_broken_copy_keeps_the_copy( @pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") @pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") @pytest.mark.parametrize("zarr_format", [2, 3]) -def test_pickled_subgroup_keeps_nested_members_as_stored( - tmp_path: Path, zarr_format: Literal[2, 3] +@pytest.mark.parametrize("pickled", [False, True]) +def test_subgroup_write_keeps_nested_members_as_stored( + tmp_path: Path, zarr_format: Literal[2, 3], pickled: bool ) -> None: - """A subgroup read from a consolidated parent carries its members' entries under - its own path when pickled, so a write through the copy stores what a write - through the subgroup itself stores, nested members kept as stored.""" - - def write_through(sub_of: Callable[[zarr.Group], zarr.Group], name: str) -> dict[str, bytes]: - path = tmp_path / name - _flagged_consolidated_group(path, zarr_format, "g/a") - group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) - sub = group["g"] - assert isinstance(sub, zarr.Group) - sub_of(sub).attrs["x"] = 1 - return { - str(p.relative_to(path)): p.read_bytes() - for p in sorted(path.rglob("*")) - if p.is_file() and p.name in ("zarr.json", ".zmetadata", ".zattrs", ".zgroup") - } + """A subgroup read from a consolidated parent keeps a nested member read from a + repaired copy as it was stored when it writes its own documents, also once + pickled: the copy carries the members' entries under the subgroup's path.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "g/a") + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + sub = group["g"] + assert isinstance(sub, zarr.Group) + if pickled: + sub = pickle.loads(pickle.dumps(sub)) - assert write_through(lambda sub: pickle.loads(pickle.dumps(sub)), "pickled") == write_through( - lambda sub: sub, "original" - ) + sub.attrs["x"] = 1 + + assert _stored_chunks(_consolidated_member(path / "g", zarr_format, "a")) == [0] @pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") @@ -1989,3 +1987,62 @@ def test_pickles_cross_versions_with_the_repair(tmp_path: Path, zarr_format: Lit ) sync(revived.setitem(slice(None), np.array([7, 8, 9], dtype="int16"))) assert _stored_chunks(json.loads((path / doc_name).read_text())) == [1] + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_group_pickles_cross_versions_with_member_marks( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A consolidated group pickled here marks its repaired members as zarr 3.4.1 + and 3.4.2 read them, and a group pickled by them, with marked members, keeps + those members as stored when it writes.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "a") + stored = _consolidated_member(path, zarr_format, "a") + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + + state = group._async_group.__getstate__() + member = state["metadata"].consolidated_metadata.metadata["a"] # type: ignore[attr-defined] + assert _stored_chunks(dict(vars(member)["_stored_document"])) == [0] + live = group.metadata.consolidated_metadata + assert live is not None + assert "_stored_document" not in vars(live.metadata["a"]) + + old_state = {k: v for k, v in state.items() if k != "_stored_members"} + old_state["store_path"] = StorePath(LocalStore(path)) + revived = object.__new__(type(group._async_group)) + revived.__setstate__(pickle.loads(pickle.dumps(old_state))) + sync(revived.update_attributes({"x": 1})) + assert json.dumps(_consolidated_member(path, zarr_format, "a")) == json.dumps(stored) + + +chunk_sizes = st.one_of( + st.integers(min_value=-2, max_value=4), + st.booleans(), + st.floats(allow_nan=False), + st.just([2, 2]), +) + + +@settings(deadline=None, max_examples=300) +@given( + zarr_format=st.sampled_from([2, 3]), + sizes=st.lists(chunk_sizes, max_size=3), + grid_name=st.sampled_from(["regular", "rectilinear", "other"]), +) +def test_may_move_chunks_agrees_with_the_repairs( + zarr_format: Literal[2, 3], sizes: list[Any], grid_name: str +) -> None: + """`may_move_chunks` answers no only where no repair moves chunks.""" + shape = [3] * len(sizes) + if zarr_format == 2: + doc: dict[str, Any] = {**_v2_doc(shape, sizes)} + else: + doc = {**_v3_doc(shape, sizes)} + if grid_name != "regular": + doc["chunk_grid"] = {"name": grid_name, "configuration": {"chunk_shapes": sizes}} + _, readings = repair_array_document(doc, zarr_format) + if any(reading.moves_chunks for reading in readings): + assert may_move_chunks(doc, zarr_format) From 378fcf64e6438571b3cdeb849997448abca9a962 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 17:00:43 +0100 Subject: [PATCH 33/48] fix(array): name verify's fields over nested v2 attributes, reject v2 chunk size 0 on strict reads - verify builds both sides as plain JSON, so nested attributes (read-only mapping proxies in the model) no longer raise TypeError for v2 arrays. - The strict read (`open`, `verify`) rejects a v2 chunk size zarr reads as repaired, as it does for v3, though the v2 model accepts it. - with_metadata raises TypeError for anything that is not array metadata. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 46 +++++++++++++++++++++++++++-------- tests/test_array_io.py | 55 +++++++++++++++++++++++++++++++++++++++--- 2 files changed, 87 insertions(+), 14 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 4248b295fe..847f7f98b1 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -6519,7 +6519,12 @@ def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: def _metadata_model_json(document: ArrayMetadataModel) -> dict[str, Any]: """The metadata as plain, mutable JSON, via JSON text: the model returns read-only forms.""" - return cast("dict[str, Any]", json.loads(json.dumps(document.to_json()))) + return cast("dict[str, Any]", _plain_json(document.to_json())) + + +def _plain_json(value: object) -> Any: + """`value`, a model's read-only JSON (mapping proxies at any depth), as plain JSON.""" + return json.loads(json.dumps(value, default=dict)) def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: @@ -6529,10 +6534,26 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: `zarr_metadata.MetadataValidationError`. Unknown extension names read as unclaimed fields, not errors. """ + from zarr_metadata import MetadataValidationError, ValidationProblem + model_v2, model_v3 = _model_classes() - if data.get("zarr_format") == 2: - return model_v2(data) - return model_v3(data) + if data.get("zarr_format") != 2: + return model_v3(data) + if moves_chunks(data, 2): + # The v2 model accepts a chunk size of 0, which zarr reads as repaired. + raise MetadataValidationError( + [ + ValidationProblem( + ("chunks", index), + f"expected an integer >= 1, got {size!r}", + "invalid_value", + input=size, + ) + for index, size in enumerate(cast("list[JSON]", data["chunks"])) + if not (type(size) is int and size >= 1) + ] + ) + return model_v2(data) @functools.cache @@ -6937,10 +6958,15 @@ def with_metadata( source: ArrayMetadataModel | ArrayMetadata if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): source = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) - elif isinstance(metadata, Mapping) and not isinstance(metadata, _model_classes()): + elif isinstance(metadata, _model_classes()): + source = metadata + elif isinstance(metadata, Mapping): source = _bound_metadata_dict(dict(metadata), self._store_path) else: - source = metadata + raise TypeError( + "Expected array metadata (a zarr_metadata model, ArrayV2Metadata or " + f"ArrayV3Metadata) or a metadata mapping, got {type(metadata).__name__}." + ) return self._derive(source, self._store_path, self._config) def with_store(self, store: Store | StorePath) -> Self: @@ -7011,11 +7037,11 @@ async def verify(self) -> None: f"Array {str(self.store_path)!r}: the stored array is Zarr format " f"{stored_format}; the bound metadata is Zarr format {self._zarr_format}." ) - stored_json = dict(stored.to_json()) - bound_json = dict(bound.to_json()) + stored_json = _metadata_model_json(stored) + bound_json = _metadata_model_json(bound) if isinstance(stored, _model_classes()[0]) and isinstance(bound, _model_classes()[0]): - stored_json["attributes"] = dict(stored.attributes) - bound_json["attributes"] = dict(bound.attributes) + stored_json["attributes"] = _plain_json(stored.attributes) + bound_json["attributes"] = _plain_json(bound.attributes) differing = _differing_fields(stored, stored_json, bound_json) raise ValueError( f"Array {str(self.store_path)!r}: the stored metadata does not match the " diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 978c0226cf..372b1c0439 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -9,7 +9,7 @@ import numpy as np import pytest -from zarr_metadata import UNSET, ZarrV3ArrayMetadata +from zarr_metadata import UNSET, MetadataValidationError, ZarrV3ArrayMetadata from zarr_metadata import ZarrV2ArrayMetadata as ModelV2 import zarr @@ -445,6 +445,32 @@ async def test_v2_model_binding_without_attributes() -> None: assert (await open_array(store=store, path="foo")).metadata.attributes == {"x": 1} +async def test_open_rejects_a_v2_chunk_size_zarr_reads_as_repaired() -> None: + """The strict read rejects a v2 chunk size of 0, as it does a v3 one, though + the v2 model accepts it.""" + store = MemoryStore() + document = { + "zarr_format": 2, + "shape": [4], + "chunks": [0], + "dtype": " None: + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,)) + with pytest.raises(TypeError, match="got str"): + aio.with_metadata(cast("Any", "nonsense")) + + def test_open_a_missing_path_creates_nothing(tmp_path: Path) -> None: missing = tmp_path / "missing.zarr" with pytest.raises(FileNotFoundError): @@ -552,10 +578,31 @@ async def test_metadata_dict_of_a_repaired_document_stores_its_repair() -> None: assert json.loads(buffer.to_bytes())["chunks"] == [1] -async def test_verify_names_the_fields_that_differ() -> None: +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_verify_names_the_fields_that_differ(zarr_format: ZarrFormat) -> None: + """Nested attributes (read-only in the model) do not stop the field naming.""" store = MemoryStore() - aio = AsyncArrayIO.build_v3( - store, name="foo", shape=(4,), data_type="float32", chunk_grid=(4,), fill_value=0.0 + attributes: dict[str, JSON] = {"units": {"name": "m", "scale": [1, 2]}} + aio = ( + AsyncArrayIO.build_v2( + store, + name="foo", + shape=(4,), + dtype="float32", + chunks=(4,), + fill_value=0.0, + attributes=attributes, + ) + if zarr_format == 2 + else AsyncArrayIO.build_v3( + store, + name="foo", + shape=(4,), + data_type="float32", + chunk_grid=(4,), + fill_value=0.0, + attributes=attributes, + ) ) await aio.store_metadata() other = aio.with_metadata(aio.metadata.update(fill_value=-0.0)) From 756cc2aafec754a522dd35e5dfe1e2bcbc273031 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 17:08:47 +0100 Subject: [PATCH 34/48] fix(array): parse no written document the cache does not need, keep subclass state in pickles - A metadata write records its documents in the store's cache only when the cache knows the path or the written metadata might move chunks, judged from its chunk sizes; writes of large attributes cost what they did on main. store_documents takes the written array metadata (default: none, a group's), so its two-argument form works again. - AsyncArray pickles keep a subclass's own instance attributes, and a shallow copy shares the metadata object again, also for a repaired array. - Two threads first touching one store object share one cache. - zarr.core.group re-exports Store and set_or_delete, as 3.4.2 did. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 29 +++++++++++-- src/zarr/core/group.py | 16 ++++--- src/zarr/core/metadata/io.py | 34 +++++++++------ src/zarr/core/metadata/stored.py | 49 ++++++++++++++++++---- tests/test_array_io.py | 71 +++++++++++++++++++++++++++++++- 5 files changed, 166 insertions(+), 33 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 847f7f98b1..3020bfd68a 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -503,7 +503,10 @@ def __getstate__(self) -> dict[str, object]: # zarr 3.4.1 and 3.4.2 read the repair a stored document needs from # this mark on the metadata itself; this version reads the carried entry. metadata = _with_legacy_mark(metadata, cast("ArrayDocument | None", mark)) - state: dict[str, object] = { + # Attributes of the instance's own (a subclass's) pickle with it, as + # they did when the fields were its only state. + state: dict[str, object] = {key: v for key, v in vars(self).items() if key != "_io"} + state |= { "metadata": metadata, "store_path": self.store_path, "codec_pipeline": self.codec_pipeline, @@ -534,6 +537,19 @@ def __setstate__(self, state: dict[str, Any]) -> None: object.__delattr__(metadata, "_stored_document") io._restore_cache_entry(state) object.__setattr__(self, "_io", io) + for key, value in state.items(): + if key not in _ASYNC_ARRAY_STATE_KEYS: + object.__setattr__(self, key, value) + + def __copy__(self) -> Self: + # A shallow copy shares the field values, as a dataclass copy does; the + # binding is copied so that rebinding one (resize) leaves the other. + io = object.__new__(AsyncArrayIO) + io.__dict__.update(vars(self._io)) + copied = object.__new__(type(self)) + copied.__dict__.update(vars(self)) + object.__setattr__(copied, "_io", io) + return copied @classmethod async def _create( @@ -6270,7 +6286,7 @@ async def _delete_chunk(coords: tuple[int, ...]) -> None: ) # Write new metadata - await store_documents(array.store_path, documents, array=True) + await store_documents(array.store_path, documents, written=new_metadata) # Update the wrapped binding in place; the codec pipeline is kept, as before. array._io._rebind_resolved(new_metadata, new_chunk_grid) @@ -6399,6 +6415,13 @@ async def _delete_key(store_path: StorePath, key: str) -> None: _CARRIED_KEYS: Final = ("_stored_entry", "_stored_copy") """The keys a pickled binding carries its store cache's entry under.""" +_ASYNC_ARRAY_STATE_KEYS: Final = frozenset( + {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config", "_pipeline_class"} + | set(_CARRIED_KEYS) +) +"""The keys of a pickled `AsyncArray`'s state that rebuild its binding; any other key is +an attribute of the instance's own.""" + def _bound_metadata_dict( data: dict[str, Any], store_path: StorePath @@ -7009,7 +7032,7 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: to_save, self._zarr_format, ensure_parents=ensure_parents, - array=True, + written=source, ) if ( self._zarr_format == 2 diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 99238f0eaa..f674eb6ac9 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -14,6 +14,9 @@ import zarr.api.asynchronous as async_api from zarr.abc.metadata import Metadata + +# Importable from here, as they have been, though this module no longer uses them. +from zarr.abc.store import Store, set_or_delete # noqa: F401, TC001 from zarr.core._info import GroupInfo from zarr.core._json import buffer_to_json_object, json_to_buffer from zarr.core.array import ( @@ -83,7 +86,6 @@ ) from typing import Any - from zarr.abc.store import Store from zarr.core.array_spec import ArrayConfigLike from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike @@ -983,9 +985,7 @@ async def delitem(self, key: str) -> None: # In place, so every handle sharing this consolidated metadata (a parent's or a # subgroup's) sees the deletion. consolidated.metadata.pop(key, None) - await store_documents( - self.store_path, encode_documents(self.store_path, self.metadata), array=False - ) + await store_documents(self.store_path, encode_documents(self.store_path, self.metadata)) async def get[DefaultT]( self, key: str, default: DefaultT | None = None @@ -2370,9 +2370,7 @@ async def update_attributes_async(self, new_attributes: dict[str, Any]) -> Group new_metadata = replace(self.metadata, attributes=new_attributes) # Write new metadata - await store_documents( - self.store_path, encode_documents(self.store_path, new_metadata), array=False - ) + await store_documents(self.store_path, encode_documents(self.store_path, new_metadata)) async_group = replace(self._async_group, metadata=new_metadata) return replace(self, _async_group=async_group) @@ -3564,7 +3562,7 @@ async def _store_nodes( meta_out = nodes[node_name] if meta_out.zarr_format == 3: if not isinstance(meta_out, GroupMetadata): - cache.observe_written(node_name, node_documents[node_name]) + cache.observe_written(node_name, node_documents[node_name], meta_out) yield node_name, built[node_name] else: # For zarr v2 @@ -3581,7 +3579,7 @@ async def _store_nodes( if meta_done and attrs_done: if not isinstance(meta_out, GroupMetadata): - cache.observe_written(node_name, node_documents[node_name]) + cache.observe_written(node_name, node_documents[node_name], meta_out) yield node_name, built[node_name] continue diff --git a/src/zarr/core/metadata/io.py b/src/zarr/core/metadata/io.py index 5242ae1d3a..024c9a57bb 100644 --- a/src/zarr/core/metadata/io.py +++ b/src/zarr/core/metadata/io.py @@ -3,7 +3,7 @@ import asyncio from enum import Enum from itertools import zip_longest -from typing import TYPE_CHECKING, Final, NamedTuple +from typing import TYPE_CHECKING, Final, NamedTuple, cast from zarr.abc.store import set_or_delete from zarr.core._json import buffer_to_json_object, json_equal @@ -23,6 +23,7 @@ from zarr.core.group import GroupMetadata from zarr.core.metadata import ArrayMetadata from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.stored import WrittenArrayMetadata class _Absent(Enum): @@ -98,15 +99,19 @@ def encode_documents( async def store_documents( - store_path: StorePath, documents: Mapping[str, Buffer], *, array: bool + store_path: StorePath, + documents: Mapping[str, Buffer], + *, + written: WrittenArrayMetadata | None = None, ) -> None: - """Store metadata documents encoded by `encode_documents` under `store_path`; an - array's are recorded in the store's cache (see `stored_documents`).""" + """Store metadata documents encoded by `encode_documents` under `store_path`; those + of the array metadata `written` (none: a group's) are recorded in the store's cache + (see `stored_documents`).""" await asyncio.gather( *(set_or_delete(store_path / key, value) for key, value in documents.items()) ) - if array: - stored_documents(store_path.store).observe_written(store_path.path, documents) + if written is not None: + stored_documents(store_path.store).observe_written(store_path.path, documents, written) def members_as_stored(store_path: StorePath, metadata: GroupMetadata) -> dict[str, ArrayDocument]: @@ -200,7 +205,7 @@ async def upsert_metadata( await store_documents( store_path, {k: v for k, v in documents.items() if k in changed}, - array=not _is_group(metadata), + written=_written_array(metadata), ) return changes @@ -282,7 +287,7 @@ async def save_metadata( to_save, metadata.zarr_format, ensure_parents=ensure_parents, - array=not _is_group(metadata), + written=_written_array(metadata), ) @@ -321,10 +326,15 @@ async def save_new_metadata( to_save, metadata.zarr_format, ensure_parents=ensure_parents, - array=not _is_group(metadata), + written=_written_array(metadata), ) +def _written_array(metadata: ArrayMetadata | GroupMetadata) -> ArrayMetadata | None: + """`metadata` as the array metadata a write records, or `None` for a group.""" + return None if _is_group(metadata) else cast("ArrayMetadata", metadata) + + def _is_group(metadata: ArrayMetadata | GroupMetadata) -> bool: from zarr.core.group import GroupMetadata @@ -337,7 +347,7 @@ async def _write_metadata( zarr_format: ZarrFormat, *, ensure_parents: bool, - array: bool, + written: WrittenArrayMetadata | None, ) -> None: set_awaitables = [set_or_delete(store_path / key, value) for key, value in to_save.items()] @@ -377,5 +387,5 @@ async def _write_metadata( ) from e await asyncio.gather(*set_awaitables) - if array: - stored_documents(store_path.store).observe_written(store_path.path, to_save) + if written is not None: + stored_documents(store_path.store).observe_written(store_path.path, to_save, written) diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index 20051fde4a..558c398c54 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -14,9 +14,10 @@ from __future__ import annotations import copy +import threading import weakref from functools import partial -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final, Literal, cast from zarr.core._json import buffer_to_json_object from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON @@ -25,11 +26,16 @@ if TYPE_CHECKING: from collections.abc import Iterable, Mapping + from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata + from zarr.abc.store import Store from zarr.core.buffer import Buffer from zarr.core.common import JSON, ZarrFormat + from zarr.core.metadata import ArrayMetadata from zarr.core.metadata.repair import ArrayDocument + type WrittenArrayMetadata = ArrayMetadata | ZarrV2ArrayMetadata | ZarrV3ArrayMetadata + type CarriedEntry = ( tuple[Literal["document"], ArrayDocument | None] | tuple[Literal["copy"], ArrayDocument] @@ -125,9 +131,16 @@ def restore_entries(self, entries: Mapping[str, CarriedEntry], zarr_format: Zarr else: self.restore_copy(path, zarr_format, entry[1]) - def observe_written(self, path: str, documents: Mapping[str, Buffer | None]) -> None: - """Record the array metadata documents just stored under `path`, by key - (`None`: deleted), as `store_documents` takes them.""" + def observe_written( + self, path: str, documents: Mapping[str, Buffer | None], written: WrittenArrayMetadata + ) -> None: + """Record the array metadata documents of `written` just stored under `path`, by + key (`None`: deleted), as `store_documents` takes them. They are parsed only + when this cache knows the path or `written` might move chunks, so a write of + valid metadata to a path the cache knows nothing about costs no parse.""" + key = (path, written.zarr_format) + if key not in self._arrays and key not in self._copies and not _may_move(written): + return if ZARR_JSON in documents: zarr_json = documents[ZARR_JSON] document = None if zarr_json is None else buffer_to_json_object(zarr_json) @@ -162,12 +175,30 @@ def observe_consolidated( self._copies.pop((path, zarr_format), None) +def _may_move(written: WrittenArrayMetadata) -> bool: + """`may_move_chunks` for the document `written` encodes, read from its chunk sizes.""" + from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata + + if isinstance(written, ArrayV2Metadata): + return may_move_chunks({"chunks": list(written.chunks)}, 2) + if isinstance(written, ArrayV3Metadata): + return may_move_chunks( + cast("ArrayDocument", {"chunk_grid": written.chunk_grid.to_dict()}), 3 + ) + if written.zarr_format == 2: + return may_move_chunks({"chunks": list(written.chunks)}, 2) + return may_move_chunks(cast("ArrayDocument", {"chunk_grid": written.chunk_grid.json}), 3) + + def _forget(key: int, ref: weakref.ref[Store]) -> None: entry = _CACHES.get(key) if entry is not None and entry[0] is ref: del _CACHES[key] +_CREATING: Final = threading.Lock() +"""Held while a cache is created, so two threads first touching one store object share one.""" + _CACHES: Final[dict[int, tuple[weakref.ref[Store], StoredDocuments]]] = {} """The cache of each live store object, by its `id`: stores define `__eq__`, so they are not hashable, and two equal stores are not one store object.""" @@ -180,6 +211,10 @@ def stored_documents(store: Store) -> StoredDocuments: entry = _CACHES.get(key) if entry is not None and entry[0]() is store: return entry[1] - documents = StoredDocuments() - _CACHES[key] = (weakref.ref(store, partial(_forget, key)), documents) - return documents + with _CREATING: + entry = _CACHES.get(key) + if entry is not None and entry[0]() is store: + return entry[1] + documents = StoredDocuments() + _CACHES[key] = (weakref.ref(store, partial(_forget, key)), documents) + return documents diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 372b1c0439..215b60c39e 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -2,6 +2,7 @@ from __future__ import annotations +import copy import json import pickle import tempfile @@ -13,13 +14,14 @@ from zarr_metadata import ZarrV2ArrayMetadata as ModelV2 import zarr -from zarr import ArrayIO, AsyncArrayIO +from zarr import ArrayIO, AsyncArray, AsyncArrayIO from zarr.api.asynchronous import open_array from zarr.core.buffer import default_buffer_prototype from zarr.core.codec_pipeline import FusedCodecPipeline from zarr.core.metadata import ArrayV2Metadata +from zarr.core.metadata.stored import stored_documents from zarr.dtype import Int16 -from zarr.errors import ArrayNotFoundError, UnknownCodecError +from zarr.errors import ArrayNotFoundError, UnknownCodecError, ZarrUserWarning from zarr.storage import LocalStore, MemoryStore, StorePath if TYPE_CHECKING: @@ -717,3 +719,68 @@ async def test_with_metadata_takes_any_mapping() -> None: rebound = aio.with_metadata(cast("Any", MappingProxyType(dict(aio.metadata.to_json())))) assert rebound == aio await rebound.store_metadata() + + +class _TaggedAsyncArray(AsyncArray[Any]): + """An `AsyncArray` subclass with an instance attribute of its own.""" + + tag: str + + def __init__(self, metadata: Any, store_path: StorePath, tag: str = "default") -> None: + super().__init__(metadata, store_path) + object.__setattr__(self, "tag", tag) + + +def test_async_array_subclass_state_survives_pickle_and_copy() -> None: + """A subclass's own instance attributes pickle and copy with it, as they did + when the dataclass fields were an array's only state.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + tagged = _TaggedAsyncArray(created.metadata, created.store_path, tag="mine") + for restored in (pickle.loads(pickle.dumps(tagged)), copy.copy(tagged), copy.deepcopy(tagged)): + assert type(restored) is _TaggedAsyncArray + assert restored.tag == "mine" + assert restored == tagged + + +async def test_copying_an_array_shares_its_metadata() -> None: + """A shallow copy shares the metadata object, also for an array opened from a + document that needs a repair; rebinding one (resize) leaves the other.""" + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "int32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + } + prototype = default_buffer_prototype() + await store.set("r/zarr.json", prototype.buffer.from_bytes(json.dumps(doc).encode())) + with pytest.warns(ZarrUserWarning): + opened = await open_array(store=store, path="r") + copied = copy.copy(opened) + assert copied.metadata is opened.metadata + assert "_stored_document" not in vars(copied.metadata) + await opened.resize((8,)) + assert copied.shape == (4,) + assert opened.shape == (8,) + + +async def test_writing_valid_metadata_parses_nothing(monkeypatch: pytest.MonkeyPatch) -> None: + """Storing metadata that cannot move chunks, at a path the store's cache knows + nothing about, records nothing and parses no document.""" + import zarr.core.metadata.stored as stored_module + + def refuse(_: object) -> None: + raise AssertionError("a written document was parsed") + + monkeypatch.setattr(stored_module, "buffer_to_json_object", refuse) + store = MemoryStore() + arr = zarr.create_array(store, name="a", shape=(4,), dtype="uint8", attributes={"x": 1}) + arr.attrs["y"] = 2 + arr.resize((8,)) + v2 = zarr.create_array(store, name="b", shape=(4,), dtype="uint8", zarr_format=2) + v2.attrs["y"] = 2 + assert not stored_documents(store).knows("a", 3) From a41599c5cb0fb2195c7d4611252f8ca8207d50bd Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 17:44:58 +0100 Subject: [PATCH 35/48] fix(array): give each AsyncArray its own state, check every bound v2 model's chunks - A change to an AsyncArray's state (resize, append, attribute updates, a field set with object.__setattr__) gives it a binding of its own, so an array copied by its __dict__, or a binding taken from `io` before the change, keeps the state it had, as when the fields were its only state. - An AsyncArray over a model-backed binding makes the metadata it hands out the binding's source, so the binding's model never disagrees with it. - A zarr_metadata v2 model bound through the constructor or with_metadata is checked as the strict read checks it (a chunk size of 0 is rejected), and a model with storage transformers is rejected by with_metadata too; a tuple-valued mapping with chunk size 0 binds as repaired, as AsyncArray reads it. - Bindings of different classes that compare equal hash equally. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 117 +++++++++++++++++++++++++++-------------- tests/test_array_io.py | 100 +++++++++++++++++++++++++++++++++-- 2 files changed, 173 insertions(+), 44 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 3020bfd68a..33907b9580 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -483,14 +483,17 @@ def __init__( @classmethod def _from_io(cls, io: AsyncArrayIO) -> AnyAsyncArray: - """An array over an existing binding: the array and `io` share one state.""" + """An array over an existing binding, until a change to the array's state + gives it a binding of its own (see `_own_binding`).""" self = object.__new__(cls) object.__setattr__(self, "_io", io) return self @property def io(self) -> AsyncArrayIO: - """The explicit-IO layer this array delegates to; both share one state.""" + """The explicit-IO binding this array delegates to, as of now: a change to + the array's state (resize, append, an attribute update) gives the array a + binding of its own, and a binding taken before keeps the state it had.""" return self._io def __getstate__(self) -> dict[str, object]: @@ -542,13 +545,10 @@ def __setstate__(self, state: dict[str, Any]) -> None: object.__setattr__(self, key, value) def __copy__(self) -> Self: - # A shallow copy shares the field values, as a dataclass copy does; the - # binding is copied so that rebinding one (resize) leaves the other. - io = object.__new__(AsyncArrayIO) - io.__dict__.update(vars(self._io)) + # A shallow copy shares the field values, as a dataclass copy does (a + # change to either takes a binding of its own; see `_own_binding`). copied = object.__new__(type(self)) copied.__dict__.update(vars(self)) - object.__setattr__(copied, "_io", io) return copied @classmethod @@ -1941,11 +1941,24 @@ def _info( def _get_metadata(self: AsyncArray[Any]) -> ArrayMetadata: """The in-memory metadata of this array.""" - return self._io.metadata_legacy + return self._io._legacy_as_source() + + +def _own_binding(self: AsyncArray[Any]) -> AsyncArrayIO: + """A copy of this array's binding, made the array's own, for a change to its state. + + An array's state is its own, as when its fields were its only state: another array + made by copying this one (its `__dict__`), or a binding taken from `io`, keeps the + state it had. + """ + io = object.__new__(AsyncArrayIO) + io.__dict__.update(vars(self._io)) + object.__setattr__(self, "_io", io) + return io def _set_metadata(self: AsyncArray[Any], value: ArrayMetadata) -> None: - self._io._rebind_resolved(value) + _own_binding(self)._rebind_resolved(value) def _get_store_path(self: AsyncArray[Any]) -> StorePath: @@ -1954,7 +1967,7 @@ def _get_store_path(self: AsyncArray[Any]) -> StorePath: def _set_store_path(self: AsyncArray[Any], value: StorePath) -> None: - self._io._store_path = value + _own_binding(self)._store_path = value def _get_config(self: AsyncArray[Any]) -> ArrayConfig: @@ -1963,7 +1976,7 @@ def _get_config(self: AsyncArray[Any]) -> ArrayConfig: def _set_config(self: AsyncArray[Any], value: ArrayConfig) -> None: - self._io._config = value + _own_binding(self)._config = value def _get_codec_pipeline(self: AsyncArray[Any]) -> CodecPipeline: @@ -1972,7 +1985,7 @@ def _get_codec_pipeline(self: AsyncArray[Any]) -> CodecPipeline: def _set_codec_pipeline(self: AsyncArray[Any], value: CodecPipeline) -> None: - self._io._codec_pipeline_cache = value + _own_binding(self)._codec_pipeline_cache = value def _get_chunk_grid(self: AsyncArray[Any]) -> ChunkGrid: @@ -1980,7 +1993,7 @@ def _get_chunk_grid(self: AsyncArray[Any]) -> ChunkGrid: def _set_chunk_grid(self: AsyncArray[Any], value: ChunkGrid) -> None: - self._io._chunk_grid_cache = value + _own_binding(self)._chunk_grid_cache = value # The dataclass fields of `AsyncArray` read the wrapped binding. They are @@ -6289,7 +6302,7 @@ async def _delete_chunk(coords: tuple[int, ...]) -> None: await store_documents(array.store_path, documents, written=new_metadata) # Update the wrapped binding in place; the codec pipeline is kept, as before. - array._io._rebind_resolved(new_metadata, new_chunk_grid) + _own_binding(array)._rebind_resolved(new_metadata, new_chunk_grid) async def _append( @@ -6376,7 +6389,7 @@ async def _update_attributes( """ array.metadata.attributes.update(new_attributes) # The document the binding derives must re-derive from the changed metadata. - array._io._rebind_resolved(array.metadata, array._chunk_grid) + _own_binding(array)._rebind_resolved(array.metadata, array._chunk_grid) # Write new metadata await array._save_metadata(array.metadata) @@ -6445,7 +6458,9 @@ def _bound_metadata_dict( return _array_metadata_model(data) except MetadataValidationError: pass - legacy = parse_array_metadata(data, path) + # Read again as an object, as `AsyncArray` binds it: that read repairs a v2 chunk + # size of 0 the document repairs leave (a tuple-valued shape). + legacy = parse_array_metadata(parse_array_metadata(data, path), path) # An invalid document given as a dict stands for what the store holds, as # one given to AsyncArray does. _record_given_document(store_path, data, legacy.zarr_format) @@ -6557,26 +6572,41 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: `zarr_metadata.MetadataValidationError`. Unknown extension names read as unclaimed fields, not errors. """ - from zarr_metadata import MetadataValidationError, ValidationProblem - model_v2, model_v3 = _model_classes() if data.get("zarr_format") != 2: return model_v3(data) - if moves_chunks(data, 2): - # The v2 model accepts a chunk size of 0, which zarr reads as repaired. - raise MetadataValidationError( - [ - ValidationProblem( - ("chunks", index), - f"expected an integer >= 1, got {size!r}", - "invalid_value", - input=size, - ) - for index, size in enumerate(cast("list[JSON]", data["chunks"])) - if not (type(size) is int and size >= 1) - ] + chunks = data.get("chunks") + if isinstance(chunks, list | tuple): + _reject_v2_chunk_sizes(chunks) + return _checked_model(model_v2(data)) + + +def _checked_model(model: ArrayMetadataModel) -> ArrayMetadataModel: + """`model`, once checked as zarr checks a model it binds: it declares no storage + transformers, and a Zarr format 2 model no chunk size zarr reads as repaired.""" + _reject_storage_transformers(model) + if isinstance(model, _model_classes()[0]): + _reject_v2_chunk_sizes(model.chunks) + return model + + +def _reject_v2_chunk_sizes(chunks: Sequence[object]) -> None: + """Raise `zarr_metadata.MetadataValidationError` for each Zarr format 2 chunk size + that is not an integer >= 1: the v2 model accepts a 0, which zarr reads as repaired.""" + from zarr_metadata import MetadataValidationError, ValidationProblem + + problems = [ + ValidationProblem( + ("chunks", index), + f"expected an integer >= 1, got {size!r}", + "invalid_value", + input=cast("JSON", size), ) - return model_v2(data) + for index, size in enumerate(chunks) + if not (type(size) is int and size >= 1) + ] + if problems: + raise MetadataValidationError(problems) @functools.cache @@ -6664,7 +6694,7 @@ def __init__( if source is metadata: source = _copied_legacy(source) elif isinstance(metadata, _model_classes()): - source = metadata + source = _checked_model(metadata) elif isinstance(metadata, Mapping): source = _bound_metadata_dict(cast("dict[str, Any]", metadata), store_path) else: @@ -6672,8 +6702,6 @@ def __init__( "Expected array metadata (a zarr_metadata model, ArrayV2Metadata or " f"ArrayV3Metadata) or a metadata mapping, got {type(metadata).__name__}." ) - if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): - _reject_storage_transformers(source) # The pipeline implementation is chosen when the binding is made, from # the global configuration, exactly as `AsyncArray` has always chosen # it; the pipeline itself is built on first use. @@ -6797,6 +6825,15 @@ def metadata_legacy(self) -> ArrayMetadata: self._resolved_cache = parse_array_metadata(data, str(self._store_path)) return self._resolved_cache + def _legacy_as_source(self) -> ArrayMetadata: + """`metadata_legacy`, made this binding's source. An `AsyncArray` hands it out as + its `metadata`, which callers change in place (its attributes) and share with + other arrays, so this binding's model must be read from it from then on.""" + resolved = self.metadata_legacy + if self._source is not resolved: + self._source = resolved + return resolved + @property def store_path(self) -> StorePath: """The bound location in a store.""" @@ -6911,7 +6948,9 @@ def _resolved_codec_pipeline(self) -> CodecPipeline: return self._codec_pipeline_cache def to_array(self) -> AnyAsyncArray: - """An `AsyncArray` over this binding; both share one state.""" + """An `AsyncArray` over this binding. A change to the array's state (resize, + append, an attribute update) gives the array a binding of its own, and + this binding keeps the state it had.""" return AsyncArray._from_io(self) def __getstate__(self) -> dict[str, object]: @@ -6982,7 +7021,7 @@ def with_metadata( if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): source = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) elif isinstance(metadata, _model_classes()): - source = metadata + source = _checked_model(metadata) elif isinstance(metadata, Mapping): source = _bound_metadata_dict(dict(metadata), self._store_path) else: @@ -7204,7 +7243,7 @@ def __eq__(self, other: object) -> bool: def __hash__(self) -> int: # A subset of the equality fields, available without deriving either # form: equal bindings hash equally. - return hash((type(self), self._source.shape, self._store_path.path)) + return hash((self._source.shape, self._store_path.path)) class ArrayIO: @@ -7351,4 +7390,4 @@ def __eq__(self, other: object) -> bool: return self._async_io == other._async_io def __hash__(self) -> int: - return hash((type(self), self._async_io)) + return hash(self._async_io) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 215b60c39e..4ebff44223 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -273,9 +273,10 @@ def test_binding_follows_the_global_config() -> None: assert legacy.order == "F" -async def test_array_state_propagates_to_the_shared_binding() -> None: - """A change through the array side (resize, update_attributes) is visible - on the shared binding's model, and a derived binding stays independent.""" +async def test_array_state_changes_rebind_only_the_array() -> None: + """A change through the array side (resize, update_attributes) shows on the + array's current binding (`io`), not on a binding taken before it, and a + derived binding stays independent.""" store = MemoryStore() aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) await aio.store_metadata() @@ -285,10 +286,11 @@ async def test_array_state_propagates_to_the_shared_binding() -> None: moved = aio.with_store(MemoryStore()) await arr.resize((6,)) - assert aio.metadata.shape == (6,) + assert arr.io.metadata.shape == (6,) + assert aio.metadata.shape == (4,) await arr.update_attributes({"a": 1}) - assert dict(aio.metadata.attributes) == {"a": 1} + assert dict(arr.io.metadata.attributes) == {"a": 1} # the derivatives were made before the changes and do not alias them assert derived.metadata_legacy.shape == (4,) @@ -784,3 +786,91 @@ def refuse(_: object) -> None: v2 = zarr.create_array(store, name="b", shape=(4,), dtype="uint8", zarr_format=2) v2.attrs["y"] = 2 assert not stored_documents(store).knows("a", 3) + + +async def test_an_array_copied_by_its_dict_keeps_its_own_state() -> None: + """An array made by copying another's `__dict__` (as subclasses do) is changed + apart from it: a resize, an attribute update or a field set through either + leaves the other, as when its fields were its only state; so does a binding + taken from `io` before the change.""" + original = zarr.create_array(MemoryStore(), shape=(4,), dtype="int32", chunks=(2,))._async_array + taken = original.io + clone = object.__new__(AsyncArray) + clone.__dict__.update(original.__dict__) + await clone.resize((6,)) + assert (original.shape, clone.shape, taken.metadata_legacy.shape) == ((4,), (6,), (4,)) + await clone.update_attributes({"x": 1}) + assert original.attrs == {} + object.__setattr__(clone, "metadata", clone.metadata.update_shape((3,))) + assert original.shape == (4,) + + +_V2_CHUNK_ZERO: dict[str, Any] = { + "zarr_format": 2, + "shape": [0], + "chunks": [0], + "dtype": " AsyncArrayIO: + return AsyncArrayIO(model, StorePath(MemoryStore(), "a")) + + +def _bind_by_with_metadata(model: ModelV2) -> AsyncArrayIO: + declared = AsyncArrayIO.build_v2(MemoryStore(), shape=(4,), dtype=" None: + """The v2 model accepts a chunk size of 0, which zarr reads as repaired; bound + as a model, it is rejected as the strict read rejects it.""" + model = ModelV2(_V2_CHUNK_ZERO) + with pytest.raises(MetadataValidationError, match="chunks"): + bind(model) + + +def test_a_tuple_valued_v2_chunk_size_zero_binds_as_repaired() -> None: + """A metadata mapping written in Python (tuples, not lists) binds as `AsyncArray` + reads it: a chunk size of 0 as repaired.""" + document = {**_V2_CHUNK_ZERO, "shape": (0,), "chunks": (0,)} + binding = ArrayIO(cast("Any", document), StorePath(MemoryStore(), "a")) + assert binding.metadata_legacy.chunks == (1,) + binding.store_metadata() + binding.verify() + + +async def test_an_array_over_a_model_shows_one_document() -> None: + """An array over a model-backed binding hands out its metadata, which callers + change in place (attribute updates through a copy); the binding's model then + reads from it, so storing through the binding keeps the change.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,), attributes={"v": 0} + ).store_metadata() + original = (await AsyncArrayIO.open(store, name="a")).to_array() + _ = original.metadata + for copied in (copy.copy(original), original.with_config({"order": "F"})): + await copied.update_attributes({"v": id(copied)}) + assert dict(original.io.metadata.attributes) == original.attrs + await original.io.store_metadata() + await original.io.verify() + + +def test_bindings_of_different_classes_hash_as_they_compare() -> None: + class _Binding(AsyncArrayIO): + pass + + base = AsyncArrayIO.build_v3( + MemoryStore(), name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + other = _Binding(base.metadata, base.store_path) + assert other == base + assert hash(other) == hash(base) From ebd8f2c8f1bdaf993ca0cde57485390fde6c571d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 18:19:29 +0100 Subject: [PATCH 36/48] fix(array): keep a binding's document as stored when an array reads it - An array made from a binding (`to_array`) works on a copy of it, so nothing done through the array (an attribute update, a resize) changes the binding it came from, and reading the array no longer replaces the binding's model with its legacy round trip, which lost NaN payloads, codec configuration and v2 float fill values. - A model-backed binding whose resolved metadata an array changed in place (attribute updates through copies sharing it) reports those attributes in its model, and stores them. - Copies and pickles of an array keep the class of its binding. - The dataclass fields of AsyncArray are no longer attributes of the class, so a @dataclass subclass may annotate one again or add a field without a default, as before. - The strict v2 read reports the model's problems together with its chunk sizes; binding is the one place a model is checked; a metadata mapping holding tuples is read once, as its JSON; parse_array_metadata's TypeError names what it got, and the storage-transformer error has its space back. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 180 +++++++++++++++++++++---------- src/zarr/core/metadata/repair.py | 5 +- src/zarr/core/metadata/stored.py | 15 ++- tests/test_array_io.py | 159 +++++++++++++++++++++++---- 4 files changed, 274 insertions(+), 85 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 33907b9580..9f607163ca 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -176,11 +176,12 @@ from zarr.storage._utils import _relativize_path if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Callable, Iterator from typing import Self import numpy.typing as npt from zarr_metadata import ( + ValidationProblem, ZarrV2ArrayMetadata, ZarrV2ArrayMetadataJSON, ZarrV3ArrayMetadata, @@ -255,18 +256,16 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: zarr_format = data.get("zarr_format") if zarr_format == 3: meta_out = ArrayV3Metadata.from_dict(data, path=path) - if len(meta_out.storage_transformers) > 0: - msg = ( - f"Array metadata contains storage transformers: {meta_out.storage_transformers}." - "Arrays with storage transformers are not supported in zarr-python at this time." - ) - raise ValueError(msg) + _reject_storage_transformers(meta_out) return meta_out elif zarr_format == 2: return ArrayV2Metadata.from_dict(data, path=path) else: raise ValueError(f"Invalid zarr_format: {zarr_format}. Expected 2 or 3") - raise TypeError + raise TypeError( + "Expected array metadata (ArrayV2Metadata or ArrayV3Metadata) or a metadata " + f"mapping, got {type(data).__name__}." + ) def create_codec_pipeline( @@ -483,17 +482,19 @@ def __init__( @classmethod def _from_io(cls, io: AsyncArrayIO) -> AnyAsyncArray: - """An array over an existing binding, until a change to the array's state - gives it a binding of its own (see `_own_binding`).""" + """An array over a copy of `io` (see `AsyncArrayIO._detached`): nothing done + through the array changes `io`.""" self = object.__new__(cls) - object.__setattr__(self, "_io", io) + object.__setattr__(self, "_io", io._detached()) return self @property def io(self) -> AsyncArrayIO: - """The explicit-IO binding this array delegates to, as of now: a change to - the array's state (resize, append, an attribute update) gives the array a - binding of its own, and a binding taken before keeps the state it had.""" + """The explicit-IO binding this array delegates to, as of now. A resize or an + append gives the array a binding of its own, so a binding taken before keeps + the shape it had; an attribute update changes the metadata in place, as it + always has, so it shows through a binding taken before, and through copies + of the array that share its metadata.""" return self._io def __getstate__(self) -> dict[str, object]: @@ -517,13 +518,17 @@ def __getstate__(self) -> dict[str, object]: "config": self.config, "_pipeline_class": self._io._pipeline_class, } + if type(self._io) is not AsyncArrayIO: + # Carried only when it is a subclass, so the pickle stays loadable where + # that class is not importable only if it is used. + state["_binding_class"] = type(self._io) state.update(carried) return state def __setstate__(self, state: dict[str, Any]) -> None: metadata = state["metadata"] pipeline = state["codec_pipeline"] - io = object.__new__(AsyncArrayIO) + io = object.__new__(state.get("_binding_class", AsyncArrayIO)) # The class the binding chose: a pipeline may build an instance of # another one (zarrs-python falls back for stores it cannot use). pipeline_class = state.get("_pipeline_class", type(pipeline)) @@ -1939,9 +1944,32 @@ def _info( ) +class _BindingField: + """A dataclass field of `AsyncArray`, read from and written to its binding. Like a + field without a default, it is not an attribute of the class, so a `@dataclass` + subclass that annotates the field again finds no default.""" + + def __init__( + self, + get: Callable[[AsyncArray[Any]], Any], + set_: Callable[[AsyncArray[Any], Any], None], + ) -> None: + self._get = get + self._set = set_ + self.__doc__ = get.__doc__ + + def __get__(self, instance: AsyncArray[Any] | None, owner: type | None = None) -> Any: + if instance is None: + raise AttributeError(f"{self._get.__name__.removeprefix('_get_')} is an instance field") + return self._get(instance) + + def __set__(self, instance: AsyncArray[Any], value: object) -> None: + self._set(instance, value) + + def _get_metadata(self: AsyncArray[Any]) -> ArrayMetadata: """The in-memory metadata of this array.""" - return self._io._legacy_as_source() + return self._io.metadata_legacy def _own_binding(self: AsyncArray[Any]) -> AsyncArrayIO: @@ -1951,7 +1979,7 @@ def _own_binding(self: AsyncArray[Any]) -> AsyncArrayIO: made by copying this one (its `__dict__`), or a binding taken from `io`, keeps the state it had. """ - io = object.__new__(AsyncArrayIO) + io = object.__new__(type(self._io)) io.__dict__.update(vars(self._io)) object.__setattr__(self, "_io", io) return io @@ -1999,11 +2027,13 @@ def _set_chunk_grid(self: AsyncArray[Any], value: ChunkGrid) -> None: # The dataclass fields of `AsyncArray` read the wrapped binding. They are # attached once the dataclass is made, so `@dataclass` sees fields, not # defaults; the setters keep `object.__setattr__(array, name, value)` working. -setattr(AsyncArray, "metadata", property(_get_metadata, _set_metadata)) # noqa: B010 -setattr(AsyncArray, "store_path", property(_get_store_path, _set_store_path)) # noqa: B010 -setattr(AsyncArray, "config", property(_get_config, _set_config)) # noqa: B010 -setattr(AsyncArray, "codec_pipeline", property(_get_codec_pipeline, _set_codec_pipeline)) # noqa: B010 -setattr(AsyncArray, "_chunk_grid", property(_get_chunk_grid, _set_chunk_grid)) # noqa: B010 +setattr(AsyncArray, "metadata", _BindingField(_get_metadata, _set_metadata)) # noqa: B010 +setattr(AsyncArray, "store_path", _BindingField(_get_store_path, _set_store_path)) # noqa: B010 +setattr(AsyncArray, "config", _BindingField(_get_config, _set_config)) # noqa: B010 +setattr( # noqa: B010 + AsyncArray, "codec_pipeline", _BindingField(_get_codec_pipeline, _set_codec_pipeline) +) +setattr(AsyncArray, "_chunk_grid", _BindingField(_get_chunk_grid, _set_chunk_grid)) # noqa: B010 # TODO: Array can be a frozen data class again once property setters (e.g. shape) are removed @@ -6429,7 +6459,15 @@ async def _delete_key(store_path: StorePath, key: str) -> None: """The keys a pickled binding carries its store cache's entry under.""" _ASYNC_ARRAY_STATE_KEYS: Final = frozenset( - {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config", "_pipeline_class"} + { + "metadata", + "store_path", + "codec_pipeline", + "_chunk_grid", + "config", + "_pipeline_class", + "_binding_class", + } | set(_CARRIED_KEYS) ) """The keys of a pickled `AsyncArray`'s state that rebuild its binding; any other key is @@ -6447,26 +6485,36 @@ def _bound_metadata_dict( from zarr_metadata import MetadataValidationError path = str(store_path) + # Metadata written in Python may hold tuples where its JSON holds lists; the + # repairs read JSON. + data = cast("dict[str, Any]", _lists_for_tuples(data)) zarr_format: ZarrFormat = 2 if data.get("zarr_format") == 2 else 3 if data.get("storage_transformers"): # Bound as legacy metadata, which rejects them, as AsyncArray always has. return parse_array_metadata(data, path) if not moves_chunks(data, zarr_format): # A document zarr reads as repaired, moving chunks, is never bound as a - # model, also where the model accepts it (a v2 chunk size of 0). + # model, also where the model accepts it. try: return _array_metadata_model(data) except MetadataValidationError: pass - # Read again as an object, as `AsyncArray` binds it: that read repairs a v2 chunk - # size of 0 the document repairs leave (a tuple-valued shape). - legacy = parse_array_metadata(parse_array_metadata(data, path), path) + legacy = parse_array_metadata(data, path) # An invalid document given as a dict stands for what the store holds, as # one given to AsyncArray does. _record_given_document(store_path, data, legacy.zarr_format) return legacy +def _lists_for_tuples(value: object) -> object: + """`value` with every tuple, at any depth in its dicts and lists, made a list.""" + if isinstance(value, dict): + return {key: _lists_for_tuples(item) for key, item in value.items()} + if isinstance(value, list | tuple): + return [_lists_for_tuples(item) for item in value] + return value + + def _record_given_document( store_path: StorePath, document: Mapping[str, Any], zarr_format: ZarrFormat ) -> None: @@ -6508,13 +6556,13 @@ def _differing_fields( return differing -def _reject_storage_transformers(model: ArrayMetadataModel) -> None: - """Raise, as `parse_array_metadata` does, for a model that declares storage +def _reject_storage_transformers(metadata: ArrayMetadataModel | ArrayMetadata) -> None: + """Raise for metadata (a model or a legacy object) that declares storage transformers: zarr-python cannot read or write chunks through them.""" - transformers = getattr(model, "storage_transformers", ()) + transformers = getattr(metadata, "storage_transformers", ()) if transformers: raise ValueError( - f"Array metadata contains storage transformers: {transformers}." + f"Array metadata contains storage transformers: {transformers}. " "Arrays with storage transformers are not supported in zarr-python at this time." ) @@ -6572,30 +6620,41 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: `zarr_metadata.MetadataValidationError`. Unknown extension names read as unclaimed fields, not errors. """ + from zarr_metadata import MetadataValidationError + model_v2, model_v3 = _model_classes() if data.get("zarr_format") != 2: return model_v3(data) chunks = data.get("chunks") - if isinstance(chunks, list | tuple): - _reject_v2_chunk_sizes(chunks) - return _checked_model(model_v2(data)) + chunk_problems = _v2_chunk_problems(chunks) if isinstance(chunks, list | tuple) else [] + try: + model = model_v2(data) + except MetadataValidationError as error: + if not chunk_problems: + raise + raise MetadataValidationError([*error.problems, *chunk_problems]) from None + if chunk_problems: + raise MetadataValidationError(chunk_problems) + return model def _checked_model(model: ArrayMetadataModel) -> ArrayMetadataModel: """`model`, once checked as zarr checks a model it binds: it declares no storage transformers, and a Zarr format 2 model no chunk size zarr reads as repaired.""" + from zarr_metadata import MetadataValidationError + _reject_storage_transformers(model) - if isinstance(model, _model_classes()[0]): - _reject_v2_chunk_sizes(model.chunks) + if isinstance(model, _model_classes()[0]) and (problems := _v2_chunk_problems(model.chunks)): + raise MetadataValidationError(problems) return model -def _reject_v2_chunk_sizes(chunks: Sequence[object]) -> None: - """Raise `zarr_metadata.MetadataValidationError` for each Zarr format 2 chunk size - that is not an integer >= 1: the v2 model accepts a 0, which zarr reads as repaired.""" - from zarr_metadata import MetadataValidationError, ValidationProblem +def _v2_chunk_problems(chunks: Sequence[object]) -> list[ValidationProblem]: + """A problem for each Zarr format 2 chunk size that is not an integer >= 1: the v2 + model accepts a 0, which zarr reads as repaired.""" + from zarr_metadata import ValidationProblem - problems = [ + return [ ValidationProblem( ("chunks", index), f"expected an integer >= 1, got {size!r}", @@ -6605,8 +6664,6 @@ def _reject_v2_chunk_sizes(chunks: Sequence[object]) -> None: for index, size in enumerate(chunks) if not (type(size) is int and size >= 1) ] - if problems: - raise MetadataValidationError(problems) @functools.cache @@ -6802,6 +6859,13 @@ def metadata(self) -> ArrayMetadataModel: """The bound metadata, as a `zarr_metadata` model.""" source = self._source if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + resolved = self._resolved_cache + if resolved is not None: + # An `AsyncArray` over this binding hands out the resolved object, + # whose attributes callers change in place: the model reads them. + attributes = _plain_json(resolved.attributes) + if attributes != _plain_json(source.attributes): + return source.update(attributes=attributes) return source # Derived on every access, not cached: a legacy object's attributes may be # changed in place through another handle that shares it (as `AsyncArray` @@ -6825,14 +6889,17 @@ def metadata_legacy(self) -> ArrayMetadata: self._resolved_cache = parse_array_metadata(data, str(self._store_path)) return self._resolved_cache - def _legacy_as_source(self) -> ArrayMetadata: - """`metadata_legacy`, made this binding's source. An `AsyncArray` hands it out as - its `metadata`, which callers change in place (its attributes) and share with - other arrays, so this binding's model must be read from it from then on.""" - resolved = self.metadata_legacy - if self._source is not resolved: - self._source = resolved - return resolved + def _detached(self) -> Self: + """A copy of this binding that shares nothing that is changed in place: its + legacy metadata (whose attributes an `AsyncArray` updates in place) is copied, + and its caches of immutable values are shared.""" + copied = object.__new__(type(self)) + copied.__dict__.update(self.__dict__) + if isinstance(self._source, ArrayV2Metadata | ArrayV3Metadata): + copied._source = _copied_legacy(self._source) + elif self._resolved_cache is not None: + copied._resolved_cache = _copied_legacy(self._resolved_cache) + return copied @property def store_path(self) -> StorePath: @@ -6948,9 +7015,8 @@ def _resolved_codec_pipeline(self) -> CodecPipeline: return self._codec_pipeline_cache def to_array(self) -> AnyAsyncArray: - """An `AsyncArray` over this binding. A change to the array's state (resize, - append, an attribute update) gives the array a binding of its own, and - this binding keeps the state it had.""" + """An `AsyncArray` over a copy of this binding: nothing done through the array + (a resize, an attribute update) changes this binding.""" return AsyncArray._from_io(self) def __getstate__(self) -> dict[str, object]: @@ -6997,7 +7063,7 @@ def _derived_source(self) -> ArrayMetadataModel | ArrayMetadata: source = self._source if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): return _copied_legacy(source) - return source + return self.metadata def _derive( self, @@ -7061,6 +7127,7 @@ async def store_metadata(self, *, ensure_parents: bool = False) -> None: if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): await save_metadata(self._store_path, source, ensure_parents=ensure_parents) else: + source = self.metadata prototype = default_buffer_prototype() key_value = source.to_key_value(indent=zarr_config.get("json_indent")) to_save = { @@ -7324,7 +7391,8 @@ def config(self) -> ArrayConfig: return self._async_io.config def to_array(self) -> AnyArray: - """An `Array` over this binding; both share one state.""" + """An `Array` over a copy of this binding: nothing done through the array (a + resize, an attribute update) changes this binding.""" return Array(self._async_io.to_array()) def with_metadata( diff --git a/src/zarr/core/metadata/repair.py b/src/zarr/core/metadata/repair.py index 839df71ea0..682a3eaa90 100644 --- a/src/zarr/core/metadata/repair.py +++ b/src/zarr/core/metadata/repair.py @@ -359,8 +359,9 @@ def may_move_chunks(doc: ArrayDocument, zarr_format: ZarrFormat) -> bool: """Whether some repair might move chunks in `doc`: `False` for a rectilinear grid, and when every stored regular chunk size is a JSON int >= 1 or `true`, for only a stored 0 (or `false`) or an edge list in a regular grid is read so - that it moves chunks (see `_read_chunk_size` and `_invalid_chunk_sizes_v3`). A cheap check that answers - for most documents before `repair_array_document` runs.""" + that it moves chunks (see `_read_chunk_size` and `_invalid_chunk_sizes_v3`). + A cheap check that answers for most documents before `repair_array_document` + runs.""" if zarr_format == 2: sizes = doc.get("chunks") else: diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py index 558c398c54..d71f660a01 100644 --- a/src/zarr/core/metadata/stored.py +++ b/src/zarr/core/metadata/stored.py @@ -177,17 +177,16 @@ def observe_consolidated( def _may_move(written: WrittenArrayMetadata) -> bool: """`may_move_chunks` for the document `written` encodes, read from its chunk sizes.""" - from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata + from zarr.core.metadata import ArrayV3Metadata - if isinstance(written, ArrayV2Metadata): - return may_move_chunks({"chunks": list(written.chunks)}, 2) - if isinstance(written, ArrayV3Metadata): - return may_move_chunks( - cast("ArrayDocument", {"chunk_grid": written.chunk_grid.to_dict()}), 3 - ) if written.zarr_format == 2: return may_move_chunks({"chunks": list(written.chunks)}, 2) - return may_move_chunks(cast("ArrayDocument", {"chunk_grid": written.chunk_grid.json}), 3) + grid: object = ( + written.chunk_grid.to_dict() + if isinstance(written, ArrayV3Metadata) + else written.chunk_grid.json + ) + return may_move_chunks(cast("ArrayDocument", {"chunk_grid": grid}), 3) def _forget(key: int, ref: weakref.ref[Store]) -> None: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 4ebff44223..2fb6f6e3a9 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -3,9 +3,10 @@ from __future__ import annotations import copy +import dataclasses import json import pickle -import tempfile +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast import numpy as np @@ -18,8 +19,9 @@ from zarr.api.asynchronous import open_array from zarr.core.buffer import default_buffer_prototype from zarr.core.codec_pipeline import FusedCodecPipeline -from zarr.core.metadata import ArrayV2Metadata +from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.stored import stored_documents +from zarr.core.sync import sync from zarr.dtype import Int16 from zarr.errors import ArrayNotFoundError, UnknownCodecError, ZarrUserWarning from zarr.storage import LocalStore, MemoryStore, StorePath @@ -449,24 +451,45 @@ async def test_v2_model_binding_without_attributes() -> None: assert (await open_array(store=store, path="foo")).metadata.attributes == {"x": 1} -async def test_open_rejects_a_v2_chunk_size_zarr_reads_as_repaired() -> None: - """The strict read rejects a v2 chunk size of 0, as it does a v3 one, though - the v2 model accepts it.""" - store = MemoryStore() - document = { - "zarr_format": 2, - "shape": [4], - "chunks": [0], - "dtype": " None: prototype = default_buffer_prototype() await store.set("a/.zarray", prototype.buffer.from_bytes(json.dumps(document).encode())) - with pytest.raises(MetadataValidationError, match="chunks"): + + +_V2_DOCUMENT: dict[str, Any] = { + "zarr_format": 2, + "shape": [4], + "chunks": [4], + "dtype": " None: + """The strict reads (`open`, `verify`) reject a v2 chunk size of 0, as they do a + v3 one, though the v2 model accepts it, naming where and what kind.""" + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "chunks": [0]}) + with pytest.raises(MetadataValidationError) as opened: + await AsyncArrayIO.open(store, name="a") + binding = AsyncArrayIO.build_v2(store, name="a", shape=(4,), dtype=" None: + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "chunks": [0], "dtype": "nope"}) + with pytest.raises(MetadataValidationError) as error: await AsyncArrayIO.open(store, name="a") + locations = [p.loc for p in error.value.problems] + assert ("chunks", 0) in locations + assert len(locations) > 1 def test_with_metadata_rejects_what_is_not_metadata() -> None: @@ -512,11 +535,11 @@ async def test_bindings_take_metadata_json() -> None: np.testing.assert_array_equal(await aio.getitem(Ellipsis), np.arange(4)) -async def test_pickled_binding_rebuilds_its_pipeline() -> None: +async def test_pickled_binding_rebuilds_its_pipeline(tmp_path: Path) -> None: """A pickled binding drops its built pipeline (a pipeline assigned to it is not carried either) and builds its class's pipeline on first use.""" aio = AsyncArrayIO.build_v3( - LocalStore(tempfile.mkdtemp()), name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,) + LocalStore(tmp_path), name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,) ) await aio.setitem(Ellipsis, np.arange(4, dtype="uint8")) revived = pickle.loads(pickle.dumps(aio)) @@ -742,6 +765,7 @@ def test_async_array_subclass_state_survives_pickle_and_copy() -> None: assert type(restored) is _TaggedAsyncArray assert restored.tag == "mine" assert restored == tagged + assert set(vars(restored)) == {"_io", "tag"} async def test_copying_an_array_shares_its_metadata() -> None: @@ -874,3 +898,100 @@ class _Binding(AsyncArrayIO): other = _Binding(base.metadata, base.store_path) assert other == base assert hash(other) == hash(base) + + +async def test_attribute_updates_through_an_array() -> None: + """An attribute update through an array never reaches the binding the array was + made from (`to_array`), shows through a binding taken from the array (`io`) + consistently in both forms, and through an array copied by its `__dict__`, as + it did when the fields were an array's only state.""" + store = MemoryStore() + origin = AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,), attributes={"v": 0} + ) + await origin.store_metadata() + array = origin.to_array() + taken = array.io + clone = object.__new__(AsyncArray) + clone.__dict__.update(array.__dict__) + await array.update_attributes({"v": 1}) + assert dict(origin.metadata.attributes) == {"v": 0} + assert dict(origin.metadata_legacy.attributes) == {"v": 0} + assert dict(taken.metadata.attributes) == taken.metadata_legacy.attributes == {"v": 1} + assert clone.attrs == {"v": 1} + await taken.store_metadata() + await taken.verify() + + +@pytest.mark.parametrize( + ("data_type", "fill_value"), [("float32", "0x7fc00001"), ("float64", "0x7ff8000000000001")] +) +async def test_reading_an_array_leaves_its_binding_as_stored( + data_type: str, fill_value: str +) -> None: + """Reading an array made from a binding (its shape, its data) leaves the binding + holding the document as stored: it still verifies, compares equal to a fresh + open, and stores that document unchanged (a NaN payload is kept).""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": data_type, + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": fill_value, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + stored = default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + await store.set("a/zarr.json", stored) + binding = await AsyncArrayIO.open(store, name="a") + array = binding.to_array() + _ = array.shape + await array.getitem(slice(None)) + assert binding == await AsyncArrayIO.open(store, name="a") + await binding.verify() + await binding.store_metadata() + rewritten = await store.get("a/zarr.json", prototype=default_buffer_prototype()) + assert rewritten is not None + assert json.loads(rewritten.to_bytes()) == document + + +class _LoggingBinding(AsyncArrayIO): + """A binding subclass, as a backend would define one.""" + + +def test_an_array_keeps_its_binding_class() -> None: + """A change to an array's state, and pickling, keep the class of the binding + it was made from.""" + binding = _LoggingBinding.build_v3( + MemoryStore(), name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + array = binding.to_array() + for kept in (array, pickle.loads(pickle.dumps(array))): + sync(kept.resize((6,))) + assert type(kept.io) is _LoggingBinding + + +@dataclass(frozen=True) +class _Typed(AsyncArray[Any]): + """A `@dataclass` subclass that annotates a field again and adds one with no default.""" + + metadata: ArrayV3Metadata + label: str + + def __init__( + self, metadata: Any, store_path: StorePath, config: Any = None, label: str = "x" + ) -> None: + super().__init__(metadata, store_path, config) + object.__setattr__(self, "label", label) + + +def test_a_dataclass_subclass_may_annotate_fields_again() -> None: + """A `@dataclass` subclass that annotates a field again, or adds one without a + default, defines and works, as when the fields had no class attribute.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + typed = _Typed(created.metadata, created.store_path) + assert typed.shape == (4,) + assert dataclasses.replace(typed, label="y").label == "y" From 5c141489ed70a16282ce8798a3c3660d83a218f8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 18:26:54 +0100 Subject: [PATCH 37/48] docs(changes): say to_array binds a copy of the binding Assisted-by: ClaudeCode:claude-opus-5-5 --- changes/419.feature.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changes/419.feature.md b/changes/419.feature.md index 9a76e2b323..15da292245 100644 --- a/changes/419.feature.md +++ b/changes/419.feature.md @@ -1,5 +1,5 @@ Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. The declaring constructors, `build_v2` and `build_v3`, are synchronous — binding is pure — take a `Store` or `StorePath`, and take the fields of each format's metadata document, with nothing hiding the difference between the formats; the documents themselves are built by the module-level `build_array_metadata_v2` and `build_array_metadata_v3`. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. -`Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over the same binding. +`Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over a copy of the binding, so nothing done through the array (an attribute update, a resize) changes the binding it came from. The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata`, which raises its minimum `typing-extensions` to 4.16 — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. `with_metadata`, `with_store` and `with_config` return new bindings with one part replaced. `ArrayIO` responds to the global configuration the way `Array` always has: the codec pipeline implementation is captured when the binding is made, the runtime configuration defaults from `zarr.config`, and `build_v2`'s `order` defaults from `array.order`. From d3d0656caf31bb58b0a3ff021e55ed2826cc01d6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 19:14:44 +0100 Subject: [PATCH 38/48] fix(array): rebind every binding through one hook, keep the model through rebuilds - Every way a binding is made or changed goes through `_bind`, including the in-place rebind of a resize or attribute update, so a binding subclass that derives state of its own there (a native handle) never goes stale. - An array over a model binding keeps the model when rebuilt by `with_config`, pickling or a deep copy, so its binding still verifies against an unchanged store and stores the document as it was (NaN payloads, codec configuration); `with_config` also keeps the binding's class. - Binding equality no longer raises when attributes were changed in place to values JSON cannot hold, and reading the model names such a value. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 75 +++++++++++++++++++++++++++++++++----- tests/test_array_io.py | 82 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 149 insertions(+), 8 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 9f607163ca..cec89ddbde 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -522,6 +522,11 @@ def __getstate__(self) -> dict[str, object]: # Carried only when it is a subclass, so the pickle stays loadable where # that class is not importable only if it is used. state["_binding_class"] = type(self._io) + source = self._io._source + if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + # The document as stored, which the legacy metadata cannot always + # reproduce (a NaN payload, a codec's configuration). + state["_model"] = source state.update(carried) return state @@ -532,7 +537,15 @@ def __setstate__(self, state: dict[str, Any]) -> None: # The class the binding chose: a pipeline may build an instance of # another one (zarrs-python falls back for stores it cannot use). pipeline_class = state.get("_pipeline_class", type(pipeline)) - io._bind(metadata, state["store_path"], state["config"], pipeline_class) + model = state.get("_model") + io._bind( + metadata if model is None else model, + state["store_path"], + state["config"], + pipeline_class, + ) + if model is not None: + io._resolved_cache = metadata io._codec_pipeline_cache = pipeline io._chunk_grid_cache = state["_chunk_grid"] if not any(key in state for key in _CARRIED_KEYS): @@ -1327,7 +1340,12 @@ def with_config(self, config: ArrayConfigLike) -> Self: # Merge new config with existing config, so missing keys are inherited # from the current array rather than from global defaults new_config = ArrayConfig(**{**self.config.to_dict(), **config}) # type: ignore[arg-type] - return type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) + # Built through the constructor, as it always was, so a subclass's own + # __init__ runs; then over this array's kind of binding (its class, and + # its model, so the document stays as stored). + rebuilt = type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) + object.__setattr__(rebuilt, "_io", self._io._rebound_like(rebuilt._io)) + return rebuilt async def nchunks_initialized(self) -> int: """ @@ -6467,6 +6485,7 @@ async def _delete_key(store_path: StorePath, key: str) -> None: "config", "_pipeline_class", "_binding_class", + "_model", } | set(_CARRIED_KEYS) ) @@ -6610,7 +6629,15 @@ def _metadata_model_json(document: ArrayMetadataModel) -> dict[str, Any]: def _plain_json(value: object) -> Any: """`value`, a model's read-only JSON (mapping proxies at any depth), as plain JSON.""" - return json.loads(json.dumps(value, default=dict)) + return json.loads(json.dumps(value, default=_mapping_as_dict)) + + +def _mapping_as_dict(value: object) -> dict[str, Any]: + """`json.dumps`' fallback for `_plain_json`: a mapping as a dict; anything else is + not JSON.""" + if isinstance(value, Mapping): + return dict(value) + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: @@ -6726,6 +6753,11 @@ class AsyncArrayIO: pipeline itself, with any options it reads from the configuration, is built at the first data operation. `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it. + + Every way a binding is made or changed (construction, the builders, + `open`, the `with_*` methods, pickling, and the resize or attribute update + of an array over it) goes through `_bind`, so a subclass that derives + state of its own from the binding (a native handle) derives it there. """ _store_path: StorePath @@ -6901,6 +6933,26 @@ def _detached(self) -> Self: copied._resolved_cache = _copied_legacy(self._resolved_cache) return copied + def _rebound_like(self, built: AsyncArrayIO) -> Self: + """A binding of this binding's class, over this binding's model when it holds + one (sharing `built`'s resolved metadata), at `built`'s location, with its + configuration, pipeline class and caches: how an array rebuilt through its + constructor keeps the binding it had.""" + rebound = object.__new__(type(self)) + source = self._source + model_backed = not isinstance(source, ArrayV2Metadata | ArrayV3Metadata) + rebound._bind( + source if model_backed else built._source, + built._store_path, + built._config, + built._pipeline_class, + ) + if model_backed: + rebound._resolved_cache = built.metadata_legacy + rebound._codec_pipeline_cache = built._codec_pipeline_cache + rebound._chunk_grid_cache = built._chunk_grid_cache + return rebound + @property def store_path(self) -> StorePath: """The bound location in a store.""" @@ -6930,11 +6982,13 @@ def _resolved_if_any(self) -> ArrayMetadata | None: def _rebind_resolved( self, resolved: ArrayMetadata, chunk_grid: ChunkGrid | None = None ) -> None: - """Make `resolved` this binding's state. The codec pipeline is kept: rebinding is how `resize` swaps the shape, which leaves the chunk encoding as it was.""" - self._source = resolved - self._resolved_cache = None + """Make `resolved` this binding's state, through `_bind`. The codec pipeline is + kept: rebinding is how `resize` swaps the shape, which leaves the chunk + encoding as it was.""" + pipeline = self._codec_pipeline_cache + self._bind(resolved, self._store_path, self._config, self._pipeline_class) + self._codec_pipeline_cache = pipeline self._chunk_grid_cache = chunk_grid - self._chunk_key_encoding_cache = None @property def _chunk_grid(self) -> ChunkGrid: @@ -7305,7 +7359,12 @@ def __eq__(self, other: object) -> bool: return NotImplemented if self._store_path != other._store_path or self._config != other._config: return False - return self.metadata == other.metadata + try: + return self.metadata == other.metadata + except TypeError: + # Attributes changed in place to values JSON cannot hold: compare the + # resolved forms, which hold them. + return self.metadata_legacy == other.metadata_legacy def __hash__(self) -> int: # A subset of the equality fields, available without deriving either diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 2fb6f6e3a9..165b84abc5 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -995,3 +995,85 @@ def test_a_dataclass_subclass_may_annotate_fields_again() -> None: typed = _Typed(created.metadata, created.store_path) assert typed.shape == (4,) assert dataclasses.replace(typed, label="y").label == "y" + + +class _NativeBinding(AsyncArrayIO): + """A binding that derives state of its own (a native handle's shape) in `_bind`.""" + + native_shape: tuple[int, ...] + + def _bind(self, source: Any, *args: Any) -> None: + super()._bind(source, *args) + self.native_shape = tuple(source.shape) + + +async def test_every_route_rebinds_through_bind() -> None: + """Every way a binding is made or changed goes through `_bind`, so state a + subclass derives there follows the metadata, also after a resize.""" + store = MemoryStore() + built = _NativeBinding.build_v3(store, name="a", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await built.store_metadata() + array = built.to_array() + await array.resize((6,)) + bindings = [ + built, + await _NativeBinding.open(store, name="a"), + built.with_metadata(built.metadata), + built.with_store(MemoryStore()), + built.with_config({"order": "F"}), + pickle.loads(pickle.dumps(built)), + array.io, + array.with_config({"order": "F"}).io, + pickle.loads(pickle.dumps(array)).io, + copy.copy(array).io, + ] + for binding in bindings: + assert isinstance(binding, _NativeBinding) + assert binding.native_shape == binding.metadata.shape + + +@pytest.mark.parametrize("derive", ["with_config", "pickle", "deepcopy"]) +async def test_an_array_over_a_model_keeps_the_document_as_stored(derive: str) -> None: + """An array rebuilt from one over a model binding (`with_config`, pickling, a + deep copy) keeps the model, so its binding still verifies against the store + and stores the document unchanged (a NaN payload is kept).""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "float32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": "0x7fc00001", + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + await store.set( + "a/zarr.json", default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + ) + array = (await AsyncArrayIO.open(store, name="a")).to_array() + if derive == "with_config": + rebuilt = array.with_config({"order": "F"}) + elif derive == "pickle": + rebuilt = pickle.loads(pickle.dumps(array)) + else: + rebuilt = copy.deepcopy(array) + await rebuilt.io.verify() + await rebuilt.io.store_metadata() + stored = await store.get("a/zarr.json", prototype=default_buffer_prototype()) + assert stored is not None + assert json.loads(stored.to_bytes()) == document + + +async def test_bindings_compare_with_attributes_json_cannot_hold() -> None: + """Equality falls back to the resolved metadata when attributes were changed in + place to values JSON cannot hold; reading the model names the value.""" + array = AsyncArrayIO.build_v3( + MemoryStore(), name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).to_array() + binding = array.io + array.metadata.attributes["n"] = np.int64(3) + assert binding == copy.copy(binding) + with pytest.raises(TypeError, match="int64 is not JSON serializable"): + _ = binding.metadata From e88c1a7456d87b63620c77134a4d064ff21262a0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 19:29:11 +0100 Subject: [PATCH 39/48] fix(array): honor subclass slots and reductions, read fields as MISSING on the class - AsyncArray pickles and copies carry the values of slots a subclass declares, and a subclass that defines its own reduction (`__reduce_ex__`, `__reduce__`, `__getstate__`, `__setstate__`) is copied its way. - The dataclass fields read as `dataclasses.MISSING` on the class instead of raising, so a loop over `dir(AsyncArray)` works while a `@dataclass` subclass still finds no default. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 36 +++++++++++++++++++++++++++++++----- tests/test_array_io.py | 41 +++++++++++++++++++++++++++++++++++++++-- 2 files changed, 70 insertions(+), 7 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index cec89ddbde..21555037a5 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -10,7 +10,7 @@ import warnings from asyncio import gather from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field, replace +from dataclasses import MISSING, dataclass, field, replace from itertools import starmap from logging import getLogger from typing import ( @@ -510,6 +510,7 @@ def __getstate__(self) -> dict[str, object]: # Attributes of the instance's own (a subclass's) pickle with it, as # they did when the fields were its only state. state: dict[str, object] = {key: v for key, v in vars(self).items() if key != "_io"} + state |= _slot_values(self) state |= { "metadata": metadata, "store_path": self.store_path, @@ -563,10 +564,15 @@ def __setstate__(self, state: dict[str, Any]) -> None: object.__setattr__(self, key, value) def __copy__(self) -> Self: + if _overrides_reduction(type(self)): + # A subclass that decides how it is copied is copied its way. + return cast("Self", copy._reconstruct(self, None, *self.__reduce_ex__(4))) # type: ignore[attr-defined] # A shallow copy shares the field values, as a dataclass copy does (a # change to either takes a binding of its own; see `_own_binding`). copied = object.__new__(type(self)) copied.__dict__.update(vars(self)) + for name, value in _slot_values(self).items(): + object.__setattr__(copied, name, value) return copied @classmethod @@ -1963,9 +1969,9 @@ def _info( class _BindingField: - """A dataclass field of `AsyncArray`, read from and written to its binding. Like a - field without a default, it is not an attribute of the class, so a `@dataclass` - subclass that annotates the field again finds no default.""" + """A dataclass field of `AsyncArray`, read from and written to its binding. On the + class it reads as `dataclasses.MISSING`, so a `@dataclass` subclass that annotates + the field again finds no default, as for a field without one.""" def __init__( self, @@ -1978,7 +1984,7 @@ def __init__( def __get__(self, instance: AsyncArray[Any] | None, owner: type | None = None) -> Any: if instance is None: - raise AttributeError(f"{self._get.__name__.removeprefix('_get_')} is an instance field") + return MISSING return self._get(instance) def __set__(self, instance: AsyncArray[Any], value: object) -> None: @@ -1990,6 +1996,26 @@ def _get_metadata(self: AsyncArray[Any]) -> ArrayMetadata: return self._io.metadata_legacy +def _slot_values(array: AsyncArray[Any]) -> dict[str, object]: + """The values in the `__slots__` a subclass of `AsyncArray` declares, which + `vars` leaves out, by name; unset slots are left out.""" + values: dict[str, object] = {} + for klass in type(array).__mro__: + slots = vars(klass).get("__slots__", ()) + for name in (slots,) if isinstance(slots, str) else slots: + if name not in ("__dict__", "__weakref__") and hasattr(array, name): + values[name] = getattr(array, name) + return values + + +_REDUCTION_HOOKS: Final = ("__reduce_ex__", "__reduce__", "__getstate__", "__setstate__") + + +def _overrides_reduction(cls: type) -> bool: + """Whether a subclass of `AsyncArray` defines how it is pickled or copied.""" + return any(getattr(cls, hook) is not getattr(AsyncArray, hook) for hook in _REDUCTION_HOOKS) + + def _own_binding(self: AsyncArray[Any]) -> AsyncArrayIO: """A copy of this array's binding, made the array's own, for a change to its state. diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 165b84abc5..b83470a8f5 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -6,8 +6,8 @@ import dataclasses import json import pickle -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, cast +from dataclasses import MISSING, dataclass +from typing import TYPE_CHECKING, Any, SupportsIndex, cast import numpy as np import pytest @@ -1077,3 +1077,40 @@ async def test_bindings_compare_with_attributes_json_cannot_hold() -> None: assert binding == copy.copy(binding) with pytest.raises(TypeError, match="int64 is not JSON serializable"): _ = binding.metadata + + +class _SlottedAsyncArray(AsyncArray[Any]): + """An `AsyncArray` subclass with a slot of its own.""" + + __slots__ = ("extra",) + extra: int + + +class _ReducingAsyncArray(AsyncArray[Any]): + """An `AsyncArray` subclass that decides how it is copied.""" + + def __reduce_ex__(self, protocol: SupportsIndex) -> Any: + return (_rebuilt_marked, (self.metadata, self.store_path)) + + +def _rebuilt_marked(metadata: Any, store_path: StorePath) -> _ReducingAsyncArray: + rebuilt = _ReducingAsyncArray(metadata, store_path) + object.__setattr__(rebuilt, "marked", True) + return rebuilt + + +def test_subclass_slots_and_reductions_are_honored() -> None: + """A subclass's slots pickle and copy with it, and a subclass that defines its + own reduction is copied its way, as when the fields were an array's only state.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + slotted = _SlottedAsyncArray(created.metadata, created.store_path) + object.__setattr__(slotted, "extra", 5) + for copied in (pickle.loads(pickle.dumps(slotted)), copy.copy(slotted), copy.deepcopy(slotted)): + assert copied.extra == 5 + reducing = _ReducingAsyncArray(created.metadata, created.store_path) + assert vars(copy.copy(reducing))["marked"] is True + + +def test_fields_read_as_missing_on_the_class() -> None: + """The fields have no class value, as fields without defaults have none.""" + assert all(getattr(AsyncArray, name) is MISSING for name in ("metadata", "config")) From 0c7863766a87f0e668b8a5914bbdaf4427f0c9ab Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 19:54:06 +0100 Subject: [PATCH 40/48] fix(array): keep what a subclass __init__ built, pickle slots and models portably - AsyncArray.with_config keeps the metadata a subclass __init__ built, binding the model only when the constructor bound this array's own metadata. - Pickles carry slots as the default reduction does (the instance dict and the slots' values, name-mangled ones included), so a slotted subclass's pickle loads in either version, and carry a model binding's documents as bytes, so a zarr without zarr-metadata loads it. - Unpickling a binding, and setting an array's store_path or config, go through `_bind`; the docstring says attributes alone change in place. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 97 ++++++++++++++++++++++++++++++------------ tests/test_array_io.py | 52 ++++++++++++++++++++-- 2 files changed, 119 insertions(+), 30 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 21555037a5..92ebe32838 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2,6 +2,7 @@ import contextlib import copy +import copyreg import functools import json import math @@ -510,7 +511,6 @@ def __getstate__(self) -> dict[str, object]: # Attributes of the instance's own (a subclass's) pickle with it, as # they did when the fields were its only state. state: dict[str, object] = {key: v for key, v in vars(self).items() if key != "_io"} - state |= _slot_values(self) state |= { "metadata": metadata, "store_path": self.store_path, @@ -526,19 +526,26 @@ def __getstate__(self) -> dict[str, object]: source = self._io._source if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): # The document as stored, which the legacy metadata cannot always - # reproduce (a NaN payload, a codec's configuration). - state["_model"] = source + # reproduce (a NaN payload, a codec's configuration), as its bytes, + # so a zarr without zarr-metadata loads the pickle too. + state["_model_documents"] = dict(source.to_key_value()) state.update(carried) - return state - - def __setstate__(self, state: dict[str, Any]) -> None: + # Slots as the default reduction carries them, so either version loads them. + slots = _slot_values(self) + return (state, slots) if slots else state # type: ignore[return-value] + + def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: + if isinstance(state, tuple): + # A subclass with slots: the instance's dict and its slots' values. + state = {**state[0], **(state[1] or {})} metadata = state["metadata"] pipeline = state["codec_pipeline"] io = object.__new__(state.get("_binding_class", AsyncArrayIO)) # The class the binding chose: a pipeline may build an instance of # another one (zarrs-python falls back for stores it cannot use). pipeline_class = state.get("_pipeline_class", type(pipeline)) - model = state.get("_model") + documents = state.get("_model_documents") + model = None if documents is None else _model_from_documents(documents, metadata) io._bind( metadata if model is None else model, state["store_path"], @@ -1350,7 +1357,7 @@ def with_config(self, config: ArrayConfigLike) -> Self: # __init__ runs; then over this array's kind of binding (its class, and # its model, so the document stays as stored). rebuilt = type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) - object.__setattr__(rebuilt, "_io", self._io._rebound_like(rebuilt._io)) + object.__setattr__(rebuilt, "_io", self._io._rebound_like(rebuilt._io, rebuilt.metadata)) return rebuilt async def nchunks_initialized(self) -> int: @@ -1996,16 +2003,23 @@ def _get_metadata(self: AsyncArray[Any]) -> ArrayMetadata: return self._io.metadata_legacy +def _model_from_documents( + documents: Mapping[str, bytes], metadata: ArrayMetadata +) -> ArrayMetadataModel: + """The model a pickled `AsyncArray` carried as its documents' bytes.""" + model_v2, model_v3 = _model_classes() + model_class = model_v2 if metadata.zarr_format == 2 else model_v3 + return model_class.from_key_value(cast("Any", documents)) + + def _slot_values(array: AsyncArray[Any]) -> dict[str, object]: """The values in the `__slots__` a subclass of `AsyncArray` declares, which `vars` leaves out, by name; unset slots are left out.""" - values: dict[str, object] = {} - for klass in type(array).__mro__: - slots = vars(klass).get("__slots__", ()) - for name in (slots,) if isinstance(slots, str) else slots: - if name not in ("__dict__", "__weakref__") and hasattr(array, name): - values[name] = getattr(array, name) - return values + return { + name: getattr(array, name) + for name in copyreg._slotnames(type(array)) # type: ignore[attr-defined] + if hasattr(array, name) + } _REDUCTION_HOOKS: Final = ("__reduce_ex__", "__reduce__", "__getstate__", "__setstate__") @@ -2039,7 +2053,8 @@ def _get_store_path(self: AsyncArray[Any]) -> StorePath: def _set_store_path(self: AsyncArray[Any], value: StorePath) -> None: - _own_binding(self)._store_path = value + binding = _own_binding(self) + binding._rebind_at(value, binding._config) def _get_config(self: AsyncArray[Any]) -> ArrayConfig: @@ -2048,7 +2063,8 @@ def _get_config(self: AsyncArray[Any]) -> ArrayConfig: def _set_config(self: AsyncArray[Any], value: ArrayConfig) -> None: - _own_binding(self)._config = value + binding = _own_binding(self) + binding._rebind_at(binding._store_path, value) def _get_codec_pipeline(self: AsyncArray[Any]) -> CodecPipeline: @@ -6511,7 +6527,7 @@ async def _delete_key(store_path: StorePath, key: str) -> None: "config", "_pipeline_class", "_binding_class", - "_model", + "_model_documents", } | set(_CARRIED_KEYS) ) @@ -6784,6 +6800,9 @@ class AsyncArrayIO: `open`, the `with_*` methods, pickling, and the resize or attribute update of an array over it) goes through `_bind`, so a subclass that derives state of its own from the binding (a native handle) derives it there. + Attributes alone may change in place, without a `_bind` (an `AsyncArray` + over the binding updates them so, as it always has): derived state must not + capture them, and writing metadata reads `metadata`. """ _store_path: StorePath @@ -6959,22 +6978,26 @@ def _detached(self) -> Self: copied._resolved_cache = _copied_legacy(self._resolved_cache) return copied - def _rebound_like(self, built: AsyncArrayIO) -> Self: - """A binding of this binding's class, over this binding's model when it holds - one (sharing `built`'s resolved metadata), at `built`'s location, with its - configuration, pipeline class and caches: how an array rebuilt through its - constructor keeps the binding it had.""" + def _rebound_like(self, built: AsyncArrayIO, metadata: ArrayMetadata) -> Self: + """How an array rebuilt through its constructor (`built`, which bound + `metadata`) keeps the binding it had: a binding of this binding's class, at + `built`'s location, with its configuration, pipeline class and caches, over + this binding's model when `metadata` is still this binding's resolved form + (shared), else over `metadata` itself (a subclass `__init__` substituted it).""" rebound = object.__new__(type(self)) source = self._source - model_backed = not isinstance(source, ArrayV2Metadata | ArrayV3Metadata) + keeps_model = ( + not isinstance(source, ArrayV2Metadata | ArrayV3Metadata) + and metadata is self._resolved_if_any + ) rebound._bind( - source if model_backed else built._source, + source if keeps_model else metadata, built._store_path, built._config, built._pipeline_class, ) - if model_backed: - rebound._resolved_cache = built.metadata_legacy + if keeps_model: + rebound._resolved_cache = metadata rebound._codec_pipeline_cache = built._codec_pipeline_cache rebound._chunk_grid_cache = built._chunk_grid_cache return rebound @@ -7005,6 +7028,24 @@ def _resolved_if_any(self) -> ArrayMetadata | None: return source return self._resolved_cache + def _rebind_at(self, store_path: StorePath, config: ArrayConfig) -> None: + """Bind this binding again, through `_bind`, at `store_path` with `config`, + keeping its caches: setting an array's location or configuration never + rebuilt its pipeline.""" + caches = ( + self._resolved_cache, + self._chunk_grid_cache, + self._chunk_key_encoding_cache, + self._codec_pipeline_cache, + ) + self._bind(self._source, store_path, config, self._pipeline_class) + ( + self._resolved_cache, + self._chunk_grid_cache, + self._chunk_key_encoding_cache, + self._codec_pipeline_cache, + ) = caches + def _rebind_resolved( self, resolved: ArrayMetadata, chunk_grid: ChunkGrid | None = None ) -> None: @@ -7111,6 +7152,8 @@ def __setstate__(self, state: dict[str, object]) -> None: state = dict(state) carried: dict[str, object] = {k: state.pop(k) for k in _CARRIED_KEYS if k in state} self.__dict__.update(state) + # Through `_bind`, so state a subclass derives there is derived again. + self._rebind_at(self._store_path, self._config) self._restore_cache_entry(carried) def _carried_cache_entry(self) -> dict[str, object]: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index b83470a8f5..ac097b7dd9 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -1001,10 +1001,12 @@ class _NativeBinding(AsyncArrayIO): """A binding that derives state of its own (a native handle's shape) in `_bind`.""" native_shape: tuple[int, ...] + native_where: str - def _bind(self, source: Any, *args: Any) -> None: - super()._bind(source, *args) + def _bind(self, source: Any, store_path: StorePath, *args: Any) -> None: + super()._bind(source, store_path, *args) self.native_shape = tuple(source.shape) + self.native_where = str(store_path) async def test_every_route_rebinds_through_bind() -> None: @@ -1027,9 +1029,13 @@ async def test_every_route_rebinds_through_bind() -> None: pickle.loads(pickle.dumps(array)).io, copy.copy(array).io, ] + moved = copy.copy(array) + object.__setattr__(moved, "store_path", StorePath(MemoryStore(), "b")) + bindings.append(moved.io) for binding in bindings: assert isinstance(binding, _NativeBinding) assert binding.native_shape == binding.metadata.shape + assert binding.native_where == str(binding.store_path) @pytest.mark.parametrize("derive", ["with_config", "pickle", "deepcopy"]) @@ -1082,7 +1088,7 @@ async def test_bindings_compare_with_attributes_json_cannot_hold() -> None: class _SlottedAsyncArray(AsyncArray[Any]): """An `AsyncArray` subclass with a slot of its own.""" - __slots__ = ("extra",) + __slots__ = ("__private", "extra") extra: int @@ -1105,8 +1111,10 @@ def test_subclass_slots_and_reductions_are_honored() -> None: created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) slotted = _SlottedAsyncArray(created.metadata, created.store_path) object.__setattr__(slotted, "extra", 5) + object.__setattr__(slotted, "_SlottedAsyncArray__private", 6) for copied in (pickle.loads(pickle.dumps(slotted)), copy.copy(slotted), copy.deepcopy(slotted)): assert copied.extra == 5 + assert vars(type(copied))["_SlottedAsyncArray__private"].__get__(copied) == 6 reducing = _ReducingAsyncArray(created.metadata, created.store_path) assert vars(copy.copy(reducing))["marked"] is True @@ -1114,3 +1122,41 @@ def test_subclass_slots_and_reductions_are_honored() -> None: def test_fields_read_as_missing_on_the_class() -> None: """The fields have no class value, as fields without defaults have none.""" assert all(getattr(AsyncArray, name) is MISSING for name in ("metadata", "config")) + + +class _StampingAsyncArray(AsyncArray[Any]): + """A subclass whose `__init__` changes the metadata it is given.""" + + def __init__(self, metadata: Any, store_path: StorePath, config: Any = None) -> None: + stamped = dataclasses.replace( + metadata, attributes={**metadata.attributes, "n": metadata.attributes.get("n", 0) + 1} + ) + super().__init__(stamped, store_path, config) + + +async def test_with_config_keeps_what_a_subclass_init_built() -> None: + """`with_config` rebuilds through the constructor and keeps the metadata a + subclass `__init__` built, as it always has, also over a model binding.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + for array in ( + _StampingAsyncArray(zarr.open_array(store, path="a").metadata, StorePath(store, "a")), + _StampingAsyncArray._from_io(await AsyncArrayIO.open(store, name="a")), + ): + before = cast("int", array.attrs.get("n", 0)) + assert array.with_config({"order": "F"}).attrs["n"] == before + 1 + + +async def test_a_pickled_array_carries_its_model_as_bytes() -> None: + """The model a pickled array carries is its documents' bytes, so a zarr without + zarr-metadata loads the pickle too.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + pickled = pickle.dumps(array) + assert b"zarr_metadata" not in pickled + assert pickle.loads(pickled).io == array.io From d072ad329d94f2af89085a2716f3904181f514aa Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 20:20:55 +0100 Subject: [PATCH 41/48] fix(array): keep a binding's pickled state under its own key, report each problem once - AsyncArray pickles carry the binding's own state (pipeline class, binding class, a model's documents, the store cache's entry) under one key, so a subclass attribute of any other name pickles as it always did. - The strict v2 read reports a chunk size the model already rejected once. - A model binding compares changed attributes as JSON text and keeps the model it made for them, so NaN attributes are stable and reads are cheap; equality without JSON attributes compares documents without them, then the resolved attributes. - A subclass that defines its own reduction is copied without private `copy` helpers; `__getstate__` is annotated with the tuple it may return; the docstring says a binding whose derived state does not pickle leaves it out. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 133 ++++++++++++++++++++--------- tests/test_array_io.py | 94 +++++++++++++++++--- tests/test_metadata/test_repair.py | 2 +- 3 files changed, 174 insertions(+), 55 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 92ebe32838..63627eeef6 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -498,9 +498,10 @@ def io(self) -> AsyncArrayIO: of the array that share its metadata.""" return self._io - def __getstate__(self) -> dict[str, object]: + def __getstate__(self) -> dict[str, object] | tuple[dict[str, object], dict[str, object]]: # The pickled state is the fields', as it has always been, so a pickle - # loads in either version; what the store's cache knows rides along. + # loads in either version; the binding's own state rides along under one + # key of its own, so no name an instance may hold is taken. carried = self._io._carried_cache_entry() metadata: ArrayMetadata = self.metadata mark = carried.get("_stored_entry", carried.get("_stored_copy")) @@ -517,22 +518,23 @@ def __getstate__(self) -> dict[str, object]: "codec_pipeline": self.codec_pipeline, "_chunk_grid": self._chunk_grid, "config": self.config, - "_pipeline_class": self._io._pipeline_class, } + binding_state: dict[str, object] = {"pipeline_class": self._io._pipeline_class} if type(self._io) is not AsyncArrayIO: # Carried only when it is a subclass, so the pickle stays loadable where # that class is not importable only if it is used. - state["_binding_class"] = type(self._io) + binding_state["binding_class"] = type(self._io) source = self._io._source if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): # The document as stored, which the legacy metadata cannot always # reproduce (a NaN payload, a codec's configuration), as its bytes, # so a zarr without zarr-metadata loads the pickle too. - state["_model_documents"] = dict(source.to_key_value()) - state.update(carried) + binding_state["model_documents"] = dict(source.to_key_value()) + binding_state |= carried + state[_BINDING_STATE] = binding_state # Slots as the default reduction carries them, so either version loads them. slots = _slot_values(self) - return (state, slots) if slots else state # type: ignore[return-value] + return (state, slots) if slots else state def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: if isinstance(state, tuple): @@ -540,11 +542,12 @@ def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, A state = {**state[0], **(state[1] or {})} metadata = state["metadata"] pipeline = state["codec_pipeline"] - io = object.__new__(state.get("_binding_class", AsyncArrayIO)) + binding_state = cast("dict[str, Any]", state.get(_BINDING_STATE, {})) + io = object.__new__(binding_state.get("binding_class", AsyncArrayIO)) # The class the binding chose: a pipeline may build an instance of # another one (zarrs-python falls back for stores it cannot use). - pipeline_class = state.get("_pipeline_class", type(pipeline)) - documents = state.get("_model_documents") + pipeline_class = binding_state.get("pipeline_class", type(pipeline)) + documents = binding_state.get("model_documents") model = None if documents is None else _model_from_documents(documents, metadata) io._bind( metadata if model is None else model, @@ -556,15 +559,16 @@ def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, A io._resolved_cache = metadata io._codec_pipeline_cache = pipeline io._chunk_grid_cache = state["_chunk_grid"] - if not any(key in state for key in _CARRIED_KEYS): + carried: dict[str, Any] = {k: binding_state[k] for k in _CARRIED_KEYS if k in binding_state} + if _BINDING_STATE not in state: # A pickle from zarr 3.4.1 or 3.4.2 marks the metadata itself. legacy_mark = vars(metadata).get("_stored_document") if legacy_mark is not None: - state = {**state, "_stored_entry": legacy_mark} + carried = {"_stored_entry": legacy_mark} if "_stored_document" in vars(metadata): # Read once; the cache holds it from here on. object.__delattr__(metadata, "_stored_document") - io._restore_cache_entry(state) + io._restore_cache_entry(carried) object.__setattr__(self, "_io", io) for key, value in state.items(): if key not in _ASYNC_ARRAY_STATE_KEYS: @@ -573,7 +577,7 @@ def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, A def __copy__(self) -> Self: if _overrides_reduction(type(self)): # A subclass that decides how it is copied is copied its way. - return cast("Self", copy._reconstruct(self, None, *self.__reduce_ex__(4))) # type: ignore[attr-defined] + return cast("Self", _copied_by_reduction(self)) # A shallow copy shares the field values, as a dataclass copy does (a # change to either takes a binding of its own; see `_own_binding`). copied = object.__new__(type(self)) @@ -2003,6 +2007,19 @@ def _get_metadata(self: AsyncArray[Any]) -> ArrayMetadata: return self._io.metadata_legacy +def _copied_by_reduction(array: AsyncArray[Any]) -> object: + """A shallow copy of `array` made from its reduction, as `copy.copy` makes one.""" + reduced = array.__reduce_ex__(4) + if isinstance(reduced, str): + return array + rebuild, args, *rest = reduced + copied = rebuild(*args) + state = rest[0] if rest else None + if state is not None: + copied.__setstate__(state) + return copied + + def _model_from_documents( documents: Mapping[str, bytes], metadata: ArrayMetadata ) -> ArrayMetadataModel: @@ -6518,18 +6535,12 @@ async def _delete_key(store_path: StorePath, key: str) -> None: _CARRIED_KEYS: Final = ("_stored_entry", "_stored_copy") """The keys a pickled binding carries its store cache's entry under.""" +_BINDING_STATE: Final = "_zarr_binding_state" +"""The key a pickled `AsyncArray` carries its binding's own state under: its pipeline +class, its class when a subclass, a model's documents, its store cache's entry.""" + _ASYNC_ARRAY_STATE_KEYS: Final = frozenset( - { - "metadata", - "store_path", - "codec_pipeline", - "_chunk_grid", - "config", - "_pipeline_class", - "_binding_class", - "_model_documents", - } - | set(_CARRIED_KEYS) + {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config", _BINDING_STATE} ) """The keys of a pickled `AsyncArray`'s state that rebuild its binding; any other key is an attribute of the instance's own.""" @@ -6674,6 +6685,28 @@ def _plain_json(value: object) -> Any: return json.loads(json.dumps(value, default=_mapping_as_dict)) +def _document_without_attributes(binding: AsyncArrayIO) -> ArrayMetadataModel: + """The document `binding` holds, as a model, with no attributes.""" + source = binding._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + return _model_of_legacy(replace(source, attributes={})) + return source.update(attributes={}) + + +def _model_of_legacy(metadata: ArrayMetadata) -> ArrayMetadataModel: + """The document `metadata` encodes, as a model.""" + buffers = metadata.to_buffer_dict(default_buffer_prototype()) + key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} + model_v2, model_v3 = _model_classes() + model_class = model_v2 if isinstance(metadata, ArrayV2Metadata) else model_v3 + return model_class.from_key_value(key_value) + + +def _json_text(value: object) -> str: + """`value`, JSON with mapping proxies at any depth, as canonical JSON text.""" + return json.dumps(value, sort_keys=True, default=_mapping_as_dict) + + def _mapping_as_dict(value: object) -> dict[str, Any]: """`json.dumps`' fallback for `_plain_json`: a mapping as a dict; anything else is not JSON.""" @@ -6699,9 +6732,17 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: try: model = model_v2(data) except MetadataValidationError as error: - if not chunk_problems: + # The model rejects chunk sizes below 0 itself; add only what it accepts + # (what no problem it reports at the chunk size or above covers). + reported = {problem.loc for problem in error.problems} + unreported = [ + problem + for problem in chunk_problems + if not any(problem.loc[:depth] in reported for depth in range(1, len(problem.loc) + 1)) + ] + if not unreported: raise - raise MetadataValidationError([*error.problems, *chunk_problems]) from None + raise MetadataValidationError([*error.problems, *unreported]) from None if chunk_problems: raise MetadataValidationError(chunk_problems) return model @@ -6802,7 +6843,9 @@ class AsyncArrayIO: state of its own from the binding (a native handle) derives it there. Attributes alone may change in place, without a `_bind` (an `AsyncArray` over the binding updates them so, as it always has): derived state must not - capture them, and writing metadata reads `metadata`. + capture them, and writing metadata reads `metadata`. Pickling a binding + pickles its attributes and unpickling binds it again; a subclass whose + derived state does not pickle leaves it out in `__getstate__`. """ _store_path: StorePath @@ -6813,6 +6856,7 @@ class AsyncArrayIO: _chunk_grid_cache: ChunkGrid | None _chunk_key_encoding_cache: ChunkKeyEncoding | None _codec_pipeline_cache: CodecPipeline | None + _attributes_cache: tuple[str, ArrayMetadataModel] | None def __init__( self, @@ -6856,6 +6900,7 @@ def _bind( self._chunk_grid_cache = None self._chunk_key_encoding_cache = None self._codec_pipeline_cache = None + self._attributes_cache = None @classmethod def _from_resolved( @@ -6937,21 +6982,22 @@ def metadata(self) -> ArrayMetadataModel: source = self._source if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): resolved = self._resolved_cache - if resolved is not None: - # An `AsyncArray` over this binding hands out the resolved object, - # whose attributes callers change in place: the model reads them. - attributes = _plain_json(resolved.attributes) - if attributes != _plain_json(source.attributes): - return source.update(attributes=attributes) - return source + if resolved is None: + return source + # An `AsyncArray` over this binding hands out the resolved object, whose + # attributes callers change in place: the model reads them. Compared as + # JSON text, in which NaN equals itself, and kept until they change. + text = _json_text(resolved.attributes) + cached = self._attributes_cache + if cached is None or cached[0] != text: + same = text == _json_text(source.attributes) + cached = (text, source if same else source.update(attributes=json.loads(text))) + self._attributes_cache = cached + return cached[1] # Derived on every access, not cached: a legacy object's attributes may be # changed in place through another handle that shares it (as `AsyncArray` # copies do); the chunk grid and key encoding such changes cannot touch. - buffers = source.to_buffer_dict(default_buffer_prototype()) - key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} - model_v2, model_v3 = _model_classes() - model_class = model_v2 if isinstance(source, ArrayV2Metadata) else model_v3 - return model_class.from_key_value(key_value) + return _model_of_legacy(source) @property def metadata_legacy(self) -> ArrayMetadata: @@ -7432,8 +7478,11 @@ def __eq__(self, other: object) -> bool: return self.metadata == other.metadata except TypeError: # Attributes changed in place to values JSON cannot hold: compare the - # resolved forms, which hold them. - return self.metadata_legacy == other.metadata_legacy + # documents without them, then the resolved attributes, which hold them. + return ( + _document_without_attributes(self) == _document_without_attributes(other) + and self.metadata_legacy.attributes == other.metadata_legacy.attributes + ) def __hash__(self) -> int: # A subset of the equality fields, available without deriving either diff --git a/tests/test_array_io.py b/tests/test_array_io.py index ac097b7dd9..7bfbfa6c26 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -5,7 +5,9 @@ import copy import dataclasses import json +import math import pickle +import threading from dataclasses import MISSING, dataclass from typing import TYPE_CHECKING, Any, SupportsIndex, cast @@ -488,8 +490,17 @@ async def test_a_strict_v2_read_reports_every_problem_with_its_chunk_sizes() -> with pytest.raises(MetadataValidationError) as error: await AsyncArrayIO.open(store, name="a") locations = [p.loc for p in error.value.problems] - assert ("chunks", 0) in locations - assert len(locations) > 1 + assert locations.count(("chunks", 0)) == 1 + assert any(loc[0] == "dtype" for loc in locations) + + +async def test_a_strict_v2_read_reports_a_chunk_size_once() -> None: + """A chunk size the model rejects itself is reported once.""" + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "chunks": [-1]}) + with pytest.raises(MetadataValidationError) as error: + await AsyncArrayIO.open(store, name="a") + assert [p.loc for p in error.value.problems] == [("chunks",)] def test_with_metadata_rejects_what_is_not_metadata() -> None: @@ -1073,16 +1084,35 @@ async def test_an_array_over_a_model_keeps_the_document_as_stored(derive: str) - async def test_bindings_compare_with_attributes_json_cannot_hold() -> None: - """Equality falls back to the resolved metadata when attributes were changed in - place to values JSON cannot hold; reading the model names the value.""" - array = AsyncArrayIO.build_v3( - MemoryStore(), name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) - ).to_array() - binding = array.io - array.metadata.attributes["n"] = np.int64(3) - assert binding == copy.copy(binding) + """Equality of model bindings whose attributes were changed in place to values + JSON cannot hold compares their documents without the attributes, then the + attributes; reading the model names the value.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="float32", chunk_grid=(4,), fill_value=0.0 + ).store_metadata() + first, second = [(await AsyncArrayIO.open(store, name="a")).to_array() for _ in range(2)] + for array in (first, second): + array.metadata.attributes["n"] = np.int64(3) + assert first.io == second.io + second.metadata.attributes["n"] = np.int64(4) + assert first.io != second.io with pytest.raises(TypeError, match="int64 is not JSON serializable"): - _ = binding.metadata + _ = first.io.metadata + + +async def test_a_model_binding_reads_changed_attributes_once() -> None: + """A model binding whose attributes were changed in place gives one model for + them until they change again, also for attributes holding NaN.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + array.metadata.attributes["x"] = float("nan") + binding = array.io + assert binding.metadata is binding.metadata + assert math.isnan(cast("float", binding.metadata.attributes["x"])) class _SlottedAsyncArray(AsyncArray[Any]): @@ -1159,4 +1189,44 @@ async def test_a_pickled_array_carries_its_model_as_bytes() -> None: array = (await AsyncArrayIO.open(store, name="a")).to_array() pickled = pickle.dumps(array) assert b"zarr_metadata" not in pickled - assert pickle.loads(pickled).io == array.io + restored = pickle.loads(pickled).io + assert restored == array.io + assert not isinstance(restored._source, ArrayV2Metadata | ArrayV3Metadata) + + +class _UnpicklableBinding(AsyncArrayIO): + """A binding deriving state that does not pickle, which it leaves out of its own.""" + + lock: Any + + def _bind(self, *args: Any) -> None: + super()._bind(*args) + self.lock = threading.Lock() + + def __getstate__(self) -> dict[str, object]: + state = super().__getstate__() + state.pop("lock") + return state + + +def test_a_binding_derives_its_state_again_when_unpickled() -> None: + binding = _UnpicklableBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + for restored in (pickle.loads(pickle.dumps(binding)), copy.deepcopy(binding)): + assert restored == binding + assert restored.lock is not binding.lock + + +class _NamedAsyncArray(AsyncArray[Any]): + """A subclass with attributes named as a binding's own state might be.""" + + +def test_subclass_attributes_keep_names_the_binding_uses() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + named = _NamedAsyncArray(created.metadata, created.store_path) + names = ("_context", "_pipeline_class", "_binding_class", "_model_documents", "_stored_entry") + for name in names: + object.__setattr__(named, name, f"mine {name}") + for restored in (pickle.loads(pickle.dumps(named)), copy.deepcopy(named)): + assert all(vars(restored)[name] == f"mine {name}" for name in names) diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index cab7f07544..84df6f6c4e 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -1972,7 +1972,7 @@ def test_pickles_cross_versions_with_the_repair(tmp_path: Path, zarr_format: Lit stored = json.loads((path / doc_name).read_text()) array = zarr.open_array(store=LocalStore(path), mode="r+") - state = array.async_array.__getstate__() + state = cast("dict[str, Any]", array.async_array.__getstate__()) assert vars(state["metadata"])["_stored_document"] is not None old_state = { From 5529888654cf189c25df29a8f86b0ca6d23a59a5 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 20:44:32 +0100 Subject: [PATCH 42/48] fix(array): bind a pickle's model only with its own metadata, honor subclass reductions - A pickled array's model is bound again only when the pickle still holds the metadata it was pickled with: a zarr that keeps the binding's state but changes the metadata (3.4.2 resizing an unpickled array) no longer brings back a stale model. - A subclass whose own reduction carries a state of its own, with no `__setstate__`, copies and pickles as the default protocol would. - Attributes changed in place are written in their own order; the strict v2 read reports a chunk size of 0 beside the model's other problems. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 30 ++++++++++++++-------- tests/test_array_io.py | 58 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 77 insertions(+), 11 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 63627eeef6..b8c365b52d 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -530,6 +530,10 @@ def __getstate__(self) -> dict[str, object] | tuple[dict[str, object], dict[str, # reproduce (a NaN payload, a codec's configuration), as its bytes, # so a zarr without zarr-metadata loads the pickle too. binding_state["model_documents"] = dict(source.to_key_value()) + # The metadata the model belongs with (one object in the pickle): a + # zarr that does not read the binding's state keeps it as it was while + # changing the metadata (a resize), and the model is then stale. + binding_state["metadata"] = metadata binding_state |= carried state[_BINDING_STATE] = binding_state # Slots as the default reduction carries them, so either version loads them. @@ -540,6 +544,12 @@ def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, A if isinstance(state, tuple): # A subclass with slots: the instance's dict and its slots' values. state = {**state[0], **(state[1] or {})} + if "metadata" not in state: + # A subclass's own state, from its own reduction, whose factory made + # the array: set it as the default protocol would. + for key, value in state.items(): + object.__setattr__(self, key, value) + return metadata = state["metadata"] pipeline = state["codec_pipeline"] binding_state = cast("dict[str, Any]", state.get(_BINDING_STATE, {})) @@ -548,7 +558,10 @@ def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, A # another one (zarrs-python falls back for stores it cannot use). pipeline_class = binding_state.get("pipeline_class", type(pipeline)) documents = binding_state.get("model_documents") - model = None if documents is None else _model_from_documents(documents, metadata) + current = binding_state.get("metadata") is metadata + model = ( + None if documents is None or not current else _model_from_documents(documents, metadata) + ) io._bind( metadata if model is None else model, state["store_path"], @@ -6732,14 +6745,8 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: try: model = model_v2(data) except MetadataValidationError as error: - # The model rejects chunk sizes below 0 itself; add only what it accepts - # (what no problem it reports at the chunk size or above covers). - reported = {problem.loc for problem in error.problems} - unreported = [ - problem - for problem in chunk_problems - if not any(problem.loc[:depth] in reported for depth in range(1, len(problem.loc) + 1)) - ] + # The model rejects other chunk sizes itself: add those it accepts, a 0. + unreported = [problem for problem in chunk_problems if problem.input == 0] if not unreported: raise raise MetadataValidationError([*error.problems, *unreported]) from None @@ -6991,7 +6998,10 @@ def metadata(self) -> ArrayMetadataModel: cached = self._attributes_cache if cached is None or cached[0] != text: same = text == _json_text(source.attributes) - cached = (text, source if same else source.update(attributes=json.loads(text))) + # Built from the attributes in their own order; the sorted text is + # only the key they are compared by. + attributes = _plain_json(resolved.attributes) + cached = (text, source if same else source.update(attributes=attributes)) self._attributes_cache = cached return cached[1] # Derived on every access, not cached: a legacy object's attributes may be diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 7bfbfa6c26..4bdf0572f1 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -29,7 +29,7 @@ from zarr.storage import LocalStore, MemoryStore, StorePath if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Callable, Mapping from pathlib import Path from zarr.core.common import JSON, ZarrFormat @@ -1230,3 +1230,59 @@ def test_subclass_attributes_keep_names_the_binding_uses() -> None: object.__setattr__(named, name, f"mine {name}") for restored in (pickle.loads(pickle.dumps(named)), copy.deepcopy(named)): assert all(vars(restored)[name] == f"mine {name}" for name in names) + + +class _OwnReductionAsyncArray(AsyncArray[Any]): + """A subclass whose reduction carries a state of its own, with no `__setstate__`.""" + + def __reduce__(self) -> Any: + return (_own_reduction_rebuilt, (self.metadata, self.store_path), {"note": "hi"}) + + +def _own_reduction_rebuilt(metadata: Any, store_path: StorePath) -> _OwnReductionAsyncArray: + return _OwnReductionAsyncArray(metadata, store_path) + + +def test_a_subclass_reduction_with_its_own_state_copies_and_pickles() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + array = _OwnReductionAsyncArray(created.metadata, created.store_path) + for copied in (copy.copy(array), copy.deepcopy(array), pickle.loads(pickle.dumps(array))): + assert vars(copied)["note"] == "hi" + assert copied.shape == (4,) + + +async def test_attributes_changed_in_place_keep_their_order() -> None: + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + array.metadata.attributes.update({"zeta": 1, "alpha": {"y": 1, "b": 2}}) + attributes = array.io.metadata.attributes + assert list(attributes) == ["zeta", "alpha"] + assert list(cast("Mapping[str, Any]", attributes["alpha"])) == ["y", "b"] + + +async def test_a_strict_v2_read_reports_a_zero_chunk_size_beside_others() -> None: + """A chunk size of 0, which the model accepts, is reported with the problems + the model reports for the other chunk sizes.""" + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "shape": [4, 4], "chunks": [-1, 0]}) + with pytest.raises(MetadataValidationError) as error: + await AsyncArrayIO.open(store, name="a") + assert ("chunks", 1) in [p.loc for p in error.value.problems] + + +async def test_a_pickle_whose_metadata_was_replaced_binds_the_new_metadata() -> None: + """A pickle that another zarr loaded, resized and pickled again (keeping the + binding's state as it was) binds the metadata it now holds, not the stale model.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + state = cast("dict[str, Any]", array.__getstate__()) + resized = dataclasses.replace(state["metadata"], shape=(8,)) + revived = object.__new__(AsyncArray) + revived.__setstate__({**state, "metadata": resized}) + assert revived.io.metadata.shape == (8,) From bd9e2714fc9abcb0c98e23157bdc3907b66a2d33 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 21:29:32 +0100 Subject: [PATCH 43/48] fix(array): compare attributes with mixed key types, report a rejected chunk size once - A model binding compares attributes changed in place as JSON makes them (keys as strings), so attributes holding an int key no longer raise. - The strict v2 read adds only an int chunk size of 0 to the model's problems; `False` and `0.0`, which the model rejects itself, are reported once. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 11 ++++++++--- tests/test_array_io.py | 20 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index b8c365b52d..eeedb280ba 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -6716,8 +6716,9 @@ def _model_of_legacy(metadata: ArrayMetadata) -> ArrayMetadataModel: def _json_text(value: object) -> str: - """`value`, JSON with mapping proxies at any depth, as canonical JSON text.""" - return json.dumps(value, sort_keys=True, default=_mapping_as_dict) + """`value`, JSON with mapping proxies at any depth, as canonical JSON text (keys + made strings first, as JSON makes them, so mixed key types sort).""" + return json.dumps(_plain_json(value), sort_keys=True) def _mapping_as_dict(value: object) -> dict[str, Any]: @@ -6746,7 +6747,11 @@ def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: model = model_v2(data) except MetadataValidationError as error: # The model rejects other chunk sizes itself: add those it accepts, a 0. - unreported = [problem for problem in chunk_problems if problem.input == 0] + unreported = [ + problem + for problem in chunk_problems + if type(problem.input) is int and problem.input == 0 + ] if not unreported: raise raise MetadataValidationError([*error.problems, *unreported]) from None diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 4bdf0572f1..356a2eefeb 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -1286,3 +1286,23 @@ async def test_a_pickle_whose_metadata_was_replaced_binds_the_new_metadata() -> revived = object.__new__(AsyncArray) revived.__setstate__({**state, "metadata": resized}) assert revived.io.metadata.shape == (8,) + + +@pytest.mark.parametrize("chunks", [[False, 2], [0.0, 2]]) +async def test_a_strict_v2_read_reports_a_rejected_chunk_size_once(chunks: list[Any]) -> None: + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "shape": [4, 4], "chunks": chunks}) + with pytest.raises(MetadataValidationError) as error: + await AsyncArrayIO.open(store, name="a") + assert [p.loc for p in error.value.problems] == [("chunks",)] + + +async def test_a_model_binding_reads_attributes_with_mixed_key_types() -> None: + """Attributes changed in place to hold an int key read as JSON writes them.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,), attributes={"b": 1} + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + cast("dict[Any, Any]", array.metadata.attributes)[1] = "x" + assert dict(array.io.metadata.attributes) == {"b": 1, "1": "x"} From b983870d318c0414a4059ec50f7a8a0d67012dd1 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 1 Nov 2026 23:21:12 +0100 Subject: [PATCH 44/48] fix(array): run the pipeline a rebound binding chose after with_config `AsyncArray.with_config` carries the pipeline its constructor built into the rebound binding only when that binding runs the same pipeline class, so a binding subclass whose `_bind` chooses its own pipeline (a backend) runs it. Assisted-by: ClaudeCode:claude-opus-5-5 --- src/zarr/core/array.py | 5 ++++- tests/test_array_io.py | 24 +++++++++++++++++++++++- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index eeedb280ba..6e05dbf7ff 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -7059,7 +7059,10 @@ def _rebound_like(self, built: AsyncArrayIO, metadata: ArrayMetadata) -> Self: ) if keeps_model: rebound._resolved_cache = metadata - rebound._codec_pipeline_cache = built._codec_pipeline_cache + if built._pipeline_class is rebound._pipeline_class: + # The pipeline the constructor built is the one this binding runs; a + # `_bind` that chose another pipeline class builds its own. + rebound._codec_pipeline_cache = built._codec_pipeline_cache rebound._chunk_grid_cache = built._chunk_grid_cache return rebound diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 356a2eefeb..0acc7ee6aa 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -20,7 +20,7 @@ from zarr import ArrayIO, AsyncArray, AsyncArrayIO from zarr.api.asynchronous import open_array from zarr.core.buffer import default_buffer_prototype -from zarr.core.codec_pipeline import FusedCodecPipeline +from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.stored import stored_documents from zarr.core.sync import sync @@ -1306,3 +1306,25 @@ async def test_a_model_binding_reads_attributes_with_mixed_key_types() -> None: array = (await AsyncArrayIO.open(store, name="a")).to_array() cast("dict[Any, Any]", array.metadata.attributes)[1] = "x" assert dict(array.io.metadata.attributes) == {"b": 1, "1": "x"} + + +class _OwnPipeline(BatchedCodecPipeline): + """A pipeline class a backend binding chooses for itself.""" + + +class _OwnPipelineBinding(AsyncArrayIO): + """A binding whose `_bind` runs a pipeline class of its own, as a backend might.""" + + def _bind(self, *args: Any) -> None: + *rest, _ = args + bind: Any = super()._bind + bind(*rest, _OwnPipeline) + + +def test_with_config_runs_the_pipeline_a_binding_chose() -> None: + array = _OwnPipelineBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,) + ).to_array() + configured = array.with_config({"order": "F"}) + assert configured.io._pipeline_class is _OwnPipeline + assert isinstance(configured.codec_pipeline, _OwnPipeline) From 022ed23075dd01e22e990ba491434271a30479f8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 2 Nov 2026 01:18:02 +0100 Subject: [PATCH 45/48] test(array): refuse a parse with a Mock instead of a nested function Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- tests/test_array_io.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 0acc7ee6aa..cb2b66b1c8 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -10,6 +10,7 @@ import threading from dataclasses import MISSING, dataclass from typing import TYPE_CHECKING, Any, SupportsIndex, cast +from unittest import mock import numpy as np import pytest @@ -810,10 +811,11 @@ async def test_writing_valid_metadata_parses_nothing(monkeypatch: pytest.MonkeyP nothing about, records nothing and parses no document.""" import zarr.core.metadata.stored as stored_module - def refuse(_: object) -> None: - raise AssertionError("a written document was parsed") - - monkeypatch.setattr(stored_module, "buffer_to_json_object", refuse) + monkeypatch.setattr( + stored_module, + "buffer_to_json_object", + mock.Mock(side_effect=AssertionError("a written document was parsed")), + ) store = MemoryStore() arr = zarr.create_array(store, name="a", shape=(4,), dtype="uint8", attributes={"x": 1}) arr.attrs["y"] = 2 From 596194181ed730fcd72a44eee258a997e0eb8f9c Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 2 Nov 2026 01:56:24 +0100 Subject: [PATCH 46/48] fix(array): share a binding's first metadata resolution across threads A binding over a model resolves its legacy metadata lazily. Threads that resolved it at once each kept their own object, so an attribute update made through one was lost when another stored its resolution afterwards. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 14 +++++++++++++- tests/test_array_io.py | 32 +++++++++++++++++++++++++++++++- 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 6e05dbf7ff..5640860cb5 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -8,6 +8,7 @@ import math import operator import os +import threading import warnings from asyncio import gather from collections.abc import Iterable, Mapping, Sequence @@ -6832,6 +6833,10 @@ def _bound_store_path(store: Store | StorePath, name: str | None) -> StorePath: return store_path +_FIRST_RESOLUTION_LOCK: Final = threading.Lock() +"""Held while a binding stores the legacy metadata it resolved first.""" + + class AsyncArrayIO: """Explicit asynchronous IO for one Zarr array. @@ -7024,7 +7029,14 @@ def metadata_legacy(self) -> ArrayMetadata: data = _metadata_model_json(source) if isinstance(source, _model_classes()[0]): data.setdefault("attributes", {}) - self._resolved_cache = parse_array_metadata(data, str(self._store_path)) + resolved = parse_array_metadata(data, str(self._store_path)) + # Threads resolving at once all keep the first object stored: an + # `AsyncArray` changes its attributes in place, so they must share it. + # (A rebind of this binding racing the read is a write racing a read, + # which no binding orders.) + with _FIRST_RESOLUTION_LOCK: + if self._resolved_cache is None: + self._resolved_cache = resolved return self._resolved_cache def _detached(self) -> Self: diff --git a/tests/test_array_io.py b/tests/test_array_io.py index cb2b66b1c8..b04d90a77a 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -20,9 +20,10 @@ import zarr from zarr import ArrayIO, AsyncArray, AsyncArrayIO from zarr.api.asynchronous import open_array +from zarr.core.array import parse_array_metadata from zarr.core.buffer import default_buffer_prototype from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline -from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.stored import stored_documents from zarr.core.sync import sync from zarr.dtype import Int16 @@ -1330,3 +1331,32 @@ def test_with_config_runs_the_pipeline_a_binding_chose() -> None: configured = array.with_config({"order": "F"}) assert configured.io._pipeline_class is _OwnPipeline assert isinstance(configured.codec_pipeline, _OwnPipeline) + + +class _ResolvedFirstElsewhere: + """Stands in for `parse_array_metadata` while another thread stores its own + resolution of `binding` first.""" + + def __init__(self, binding: AsyncArrayIO, elsewhere: ArrayMetadata) -> None: + self.binding = binding + self.elsewhere = elsewhere + + def __call__(self, *args: Any, **kwargs: Any) -> ArrayMetadata: + resolved = parse_array_metadata(*args, **kwargs) + self.binding._resolved_cache = self.elsewhere + return resolved + + +async def test_bindings_resolving_at_once_share_the_first_resolution() -> None: + """A binding over a model keeps the legacy metadata resolved first, so threads + resolving it at once share one object, whose attributes an array changes in place.""" + store = MemoryStore() + built = AsyncArrayIO.build_v3(store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,)) + await built.store_metadata() + binding = await AsyncArrayIO.open(store, name="a") + elsewhere = parse_array_metadata(built.metadata_legacy.to_dict(), "a") + with mock.patch( + "zarr.core.array.parse_array_metadata", _ResolvedFirstElsewhere(binding, elsewhere) + ): + assert binding.metadata_legacy is elsewhere + assert binding.metadata_legacy is elsewhere From 2dea537c396d82fcb8b819842c2a650bbf483783 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 2 Nov 2026 13:23:37 +0100 Subject: [PATCH 47/48] fix(array): read a subclass's pickled state in every form, keep attribute order AsyncArray.__setstate__ took any state holding a `metadata` entry for its own field state, and raised on a subclass reduction's `(None, slots)` state; it now recognizes its own state by both `metadata` and `codec_pipeline`. AsyncArrayIO's pickle carries its store-cache entry under one key of its own, so a subclass attribute of the same name survives, and leaves out the attributes' derived model. Attributes reordered in place now count as changed, so the written order reaches the model. _rebound_like says a binding's pipeline class wins over one a subclass __init__ assigned, and a test pins it. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 27 ++++++----- tests/test_array_io.py | 101 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 115 insertions(+), 13 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 5640860cb5..8a24bfa60b 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -543,9 +543,9 @@ def __getstate__(self) -> dict[str, object] | tuple[dict[str, object], dict[str, def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: if isinstance(state, tuple): - # A subclass with slots: the instance's dict and its slots' values. - state = {**state[0], **(state[1] or {})} - if "metadata" not in state: + # A subclass with slots: the instance's dict, if any, and its slots' values. + state = {**(state[0] or {}), **(state[1] or {})} + if "metadata" not in state or "codec_pipeline" not in state: # A subclass's own state, from its own reduction, whose factory made # the array: set it as the default protocol would. for key, value in state.items(): @@ -6717,9 +6717,9 @@ def _model_of_legacy(metadata: ArrayMetadata) -> ArrayMetadataModel: def _json_text(value: object) -> str: - """`value`, JSON with mapping proxies at any depth, as canonical JSON text (keys - made strings first, as JSON makes them, so mixed key types sort).""" - return json.dumps(_plain_json(value), sort_keys=True) + """`value`, JSON with mapping proxies at any depth, as JSON text in its own order + (keys made strings, as JSON makes them), so a change of order is a change.""" + return json.dumps(_plain_json(value)) def _mapping_as_dict(value: object) -> dict[str, Any]: @@ -7008,8 +7008,6 @@ def metadata(self) -> ArrayMetadataModel: cached = self._attributes_cache if cached is None or cached[0] != text: same = text == _json_text(source.attributes) - # Built from the attributes in their own order; the sorted text is - # only the key they are compared by. attributes = _plain_json(resolved.attributes) cached = (text, source if same else source.update(attributes=attributes)) self._attributes_cache = cached @@ -7056,7 +7054,9 @@ def _rebound_like(self, built: AsyncArrayIO, metadata: ArrayMetadata) -> Self: `metadata`) keeps the binding it had: a binding of this binding's class, at `built`'s location, with its configuration, pipeline class and caches, over this binding's model when `metadata` is still this binding's resolved form - (shared), else over `metadata` itself (a subclass `__init__` substituted it).""" + (shared), else over `metadata` itself (a subclass `__init__` substituted it). + A `_bind` that runs another pipeline class wins over any pipeline `built` + holds, also one a subclass `__init__` assigned.""" rebound = object.__new__(type(self)) source = self._source keeps_model = ( @@ -7219,14 +7219,17 @@ def to_array(self) -> AnyAsyncArray: def __getstate__(self) -> dict[str, object]: state = dict(self.__dict__) # The pipeline is rebuilt on first use: a built one is bound to this - # process's store object, and may not pickle. + # process's store object, and may not pickle. The attributes' model is + # derived again. state["_codec_pipeline_cache"] = None - state.update(self._carried_cache_entry()) + state["_attributes_cache"] = None + # Under one key of its own, so no name a subclass's attribute may hold is taken. + state[_BINDING_STATE] = self._carried_cache_entry() return state def __setstate__(self, state: dict[str, object]) -> None: state = dict(state) - carried: dict[str, object] = {k: state.pop(k) for k in _CARRIED_KEYS if k in state} + carried = cast("dict[str, object]", state.pop(_BINDING_STATE, {})) self.__dict__.update(state) # Through `_bind`, so state a subclass derives there is derived again. self._rebind_at(self._store_path, self._config) diff --git a/tests/test_array_io.py b/tests/test_array_io.py index b04d90a77a..40d56ea603 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -20,7 +20,7 @@ import zarr from zarr import ArrayIO, AsyncArray, AsyncArrayIO from zarr.api.asynchronous import open_array -from zarr.core.array import parse_array_metadata +from zarr.core.array import create_codec_pipeline, parse_array_metadata from zarr.core.buffer import default_buffer_prototype from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata @@ -1254,6 +1254,72 @@ def test_a_subclass_reduction_with_its_own_state_copies_and_pickles() -> None: assert copied.shape == (4,) +class _SlotStateReductionAsyncArray(AsyncArray[Any]): + """A subclass whose reduction carries a state in the default protocol's + `(None, slots)` form.""" + + def __reduce__(self) -> Any: + return (_slot_state_rebuilt, (self.metadata, self.store_path), (None, {"note": "hi"})) + + +def _slot_state_rebuilt(metadata: Any, store_path: StorePath) -> _SlotStateReductionAsyncArray: + return _SlotStateReductionAsyncArray(metadata, store_path) + + +class _MetadataStateReductionAsyncArray(AsyncArray[Any]): + """A subclass whose own state holds a `metadata` entry, which is not an array's + field state: it has no `codec_pipeline`.""" + + def __reduce__(self) -> Any: + return ( + _metadata_state_rebuilt, + (self.metadata, self.store_path), + {"metadata": self.metadata}, + ) + + +def _metadata_state_rebuilt( + metadata: Any, store_path: StorePath +) -> _MetadataStateReductionAsyncArray: + return _MetadataStateReductionAsyncArray(metadata, store_path) + + +def test_a_subclass_reduction_state_is_read_in_every_form() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + slotted = _SlotStateReductionAsyncArray(created.metadata, created.store_path) + for copied in (copy.copy(slotted), copy.deepcopy(slotted), pickle.loads(pickle.dumps(slotted))): + assert vars(copied)["note"] == "hi" + holding = _MetadataStateReductionAsyncArray(created.metadata, created.store_path) + for copied in (copy.copy(holding), copy.deepcopy(holding), pickle.loads(pickle.dumps(holding))): + assert copied.shape == (4,) + + +class _AttributeNamedLikeTheCacheEntry(AsyncArrayIO): + """A binding subclass with an attribute of the name a carried cache entry had.""" + + _stored_entry: str + + +async def test_a_pickled_binding_keeps_its_attributes_apart_from_what_it_carries() -> None: + """What a binding carries in a pickle takes no name of a subclass's, and the + model derived from attributes changed in place is not pickled.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + binding = _AttributeNamedLikeTheCacheEntry.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + binding._stored_entry = "mine" + assert pickle.loads(pickle.dumps(binding))._stored_entry == "mine" + + array = (await AsyncArrayIO.open(store, name="a")).to_array() + array.metadata.attributes["changed"] = 1 + assert dict(array.io.metadata.attributes) == {"changed": 1} + assert array.io.__getstate__()["_attributes_cache"] is None + assert dict(pickle.loads(pickle.dumps(array.io)).metadata.attributes) == {"changed": 1} + + async def test_attributes_changed_in_place_keep_their_order() -> None: store = MemoryStore() await AsyncArrayIO.build_v3( @@ -1265,6 +1331,12 @@ async def test_attributes_changed_in_place_keep_their_order() -> None: assert list(attributes) == ["zeta", "alpha"] assert list(cast("Mapping[str, Any]", attributes["alpha"])) == ["y", "b"] + # a change of order alone is a change + stored = array.metadata.attributes + alpha = stored.pop("zeta") + stored["zeta"] = alpha + assert list(array.io.metadata.attributes) == ["alpha", "zeta"] + async def test_a_strict_v2_read_reports_a_zero_chunk_size_beside_others() -> None: """A chunk size of 0, which the model accepts, is reported with the problems @@ -1333,6 +1405,33 @@ def test_with_config_runs_the_pipeline_a_binding_chose() -> None: assert isinstance(configured.codec_pipeline, _OwnPipeline) +class _AssignedPipeline(BatchedCodecPipeline): + """A pipeline class a subclass `__init__` assigns.""" + + +class _PipelineAssigningAsyncArray(AsyncArray[Any]): + """A subclass whose `__init__` assigns a pipeline of its own.""" + + def __init__(self, metadata: Any, store_path: StorePath, config: Any = None) -> None: + super().__init__(metadata, store_path, config) + pipeline = create_codec_pipeline(self.metadata, pipeline_class=_AssignedPipeline) + object.__setattr__(self, "codec_pipeline", pipeline) + + +def test_with_config_prefers_a_binding_pipeline_over_an_assigned_one() -> None: + """A binding that chose a pipeline class runs it, over a pipeline a subclass + `__init__` assigned; without such a binding, the assigned pipeline runs.""" + binding = _OwnPipelineBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + chosen = _PipelineAssigningAsyncArray._from_io(binding).with_config({"order": "F"}) + assert isinstance(chosen.codec_pipeline, _OwnPipeline) + + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + assigned = _PipelineAssigningAsyncArray(created.metadata, created.store_path) + assert isinstance(assigned.with_config({"order": "F"}).codec_pipeline, _AssignedPipeline) + + class _ResolvedFirstElsewhere: """Stands in for `parse_array_metadata` while another thread stores its own resolution of `binding` first.""" From a79a325810778698356dbefc5c36e74522dc26ad Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Mon, 2 Nov 2026 13:28:37 +0100 Subject: [PATCH 48/48] fix(array): treat a pickled state as an array's only when it holds every field AsyncArray.__setstate__ took a state with `metadata` and `codec_pipeline` for its own field state and then raised KeyError on the fields it lacked, so a subclass reduction holding just those two could not load or copy. It now requires all five fields. Assisted-by: ClaudeCode:claude-opus-5-5 Co-Authored-By: Claude Opus 5.5 --- src/zarr/core/array.py | 10 +++++++--- tests/test_array_io.py | 6 +++--- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 8a24bfa60b..867a2bdd6b 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -545,7 +545,7 @@ def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, A if isinstance(state, tuple): # A subclass with slots: the instance's dict, if any, and its slots' values. state = {**(state[0] or {}), **(state[1] or {})} - if "metadata" not in state or "codec_pipeline" not in state: + if not state.keys() >= _ASYNC_ARRAY_FIELDS: # A subclass's own state, from its own reduction, whose factory made # the array: set it as the default protocol would. for key, value in state.items(): @@ -6553,9 +6553,13 @@ async def _delete_key(store_path: StorePath, key: str) -> None: """The key a pickled `AsyncArray` carries its binding's own state under: its pipeline class, its class when a subclass, a model's documents, its store cache's entry.""" -_ASYNC_ARRAY_STATE_KEYS: Final = frozenset( - {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config", _BINDING_STATE} +_ASYNC_ARRAY_FIELDS: Final = frozenset( + {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config"} ) +"""The fields a pickled `AsyncArray`'s state holds; a state without them all is a +subclass's own.""" + +_ASYNC_ARRAY_STATE_KEYS: Final = _ASYNC_ARRAY_FIELDS | {_BINDING_STATE} """The keys of a pickled `AsyncArray`'s state that rebuild its binding; any other key is an attribute of the instance's own.""" diff --git a/tests/test_array_io.py b/tests/test_array_io.py index 40d56ea603..c7f357fc68 100644 --- a/tests/test_array_io.py +++ b/tests/test_array_io.py @@ -1267,14 +1267,14 @@ def _slot_state_rebuilt(metadata: Any, store_path: StorePath) -> _SlotStateReduc class _MetadataStateReductionAsyncArray(AsyncArray[Any]): - """A subclass whose own state holds a `metadata` entry, which is not an array's - field state: it has no `codec_pipeline`.""" + """A subclass whose own state holds some of an array's fields, which is not an + array's field state: it does not hold them all.""" def __reduce__(self) -> Any: return ( _metadata_state_rebuilt, (self.metadata, self.store_path), - {"metadata": self.metadata}, + {"metadata": self.metadata, "codec_pipeline": self.codec_pipeline}, )