From 88ab6319eae62f8eeaa8ef2d4488bf0628e88aaf Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 29 Oct 2026 11:53:31 +0100 Subject: [PATCH 1/8] 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 2/8] 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 3/8] 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 4/8] 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 5/8] 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 6/8] 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 7/8] 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 8/8] 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" }, +]