diff --git a/changes/419.feature.md b/changes/419.feature.md new file mode 100644 index 0000000000..15da292245 --- /dev/null +++ b/changes/419.feature.md @@ -0,0 +1,6 @@ +Added `zarr.ArrayIO` and `zarr.AsyncArrayIO`, an explicit-IO layer over Zarr arrays. Constructing one performs no IO against storage: storing metadata (`store_metadata`), storing or reading chunks (`setitem` / `getitem`), and their erasing inverses (`erase_metadata`, `erase_chunk`, `erase_chunks`) are explicit method calls that run in whatever order the caller chooses, so chunks can be written before metadata. The declaring constructors, `build_v2` and `build_v3`, are synchronous — binding is pure — take a `Store` or `StorePath`, and take the fields of each format's metadata document, with nothing hiding the difference between the formats; the documents themselves are built by the module-level `build_array_metadata_v2` and `build_array_metadata_v3`. `open` auto-detects the Zarr format by default, as `zarrs` does, and `verify` reads the stored metadata and raises, naming each differing field, when it does not match the metadata the instance is bound to. +`Array` and `AsyncArray` are built on this layer: every `AsyncArray` wraps an `AsyncArrayIO` and delegates its metadata and selection IO to it, the wrapped layer is reachable as `Array.io` / `AsyncArray.io`, and `to_array()` returns an array over a copy of the binding, so nothing done through the array (an attribute update, a resize) changes the binding it came from. +The binding holds its metadata as a `zarr-metadata` model — `zarr-python` now depends on `zarr-metadata`, which raises its minimum `typing-extensions` to 4.16 — reachable as `ArrayIO.metadata`, with `ArrayIO.metadata_legacy` the same metadata resolved into the legacy classes: metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, `erase_chunks`, `verify`) never resolve codec or data type names, so an array whose codecs this environment lacks can still be opened, verified and erased, and names resolve into classes only when chunks are read or written. +`with_metadata`, `with_store` and `with_config` return new bindings with one part replaced. +`ArrayIO` responds to the global configuration the way `Array` always has: the codec pipeline implementation is captured when the binding is made, the runtime configuration defaults from `zarr.config`, and `build_v2`'s `order` defaults from `array.order`. +Whether a stored array document needs a repair that moves chunks (such as a stored chunk size of 0) is now recorded per store object, in a cache of the metadata documents zarr has read and stored, instead of on the metadata read from it. Every array bound to one store object shares what any of them learned, derived bindings carry nothing, and the first chunk write through any binding that knows the document, including a pickled one, stores the repair; nothing is read that was not read before. diff --git a/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/user-guide/installation.md b/docs/user-guide/installation.md index a7487e83c8..fc5e526504 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -9,7 +9,7 @@ Required dependencies include: - [numpy](https://numpy.org) (2.0 or later) - [numcodecs](https://numcodecs.readthedocs.io/en/stable/) (0.14 or later) - [google-crc32c](https://github.com/googleapis/python-crc32c) (1.5 or later) -- [typing_extensions](https://typing-extensions.readthedocs.io/en/latest/) (4.14 or later) +- [typing_extensions](https://typing-extensions.readthedocs.io/en/latest/) (4.16 or later) - [donfig](https://donfig.readthedocs.io/en/latest/) (0.8 or later) ## pip diff --git a/pyproject.toml b/pyproject.toml index 0df173f805..7007b2979c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,9 +50,10 @@ dependencies = [ # 0.16.6 and 0.17.1 refuse to unpickle untrusted data (GHSA-jqc8-qjfv-5m6r); 0.17.0 does not. 'numcodecs>=0.16.6,!=0.17.0', 'google-crc32c>=1.5', - 'typing_extensions>=4.14', + 'typing_extensions>=4.16', 'donfig>=0.8', 'msgspec>=0.19', + 'zarr-metadata>=0.6.0', ] dynamic = [ @@ -257,7 +258,7 @@ extra-dependencies = [ 'fsspec==2023.10.0', 's3fs==2023.10.0', 'universal_pathlib==0.2.0', - 'typing_extensions==4.14.*', + 'typing_extensions==4.16.*', 'donfig==0.8.*', 'obstore==0.5.*', 'msgspec==0.19.*', diff --git a/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/api/asynchronous.py b/src/zarr/api/asynchronous.py index 476276fd82..39d3a506f1 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -18,6 +18,7 @@ create_array, from_array, get_array_metadata, + observe_read, ) from zarr.core.array_spec import ArrayConfigLike, parse_array_config from zarr.core.buffer import NDArrayLike @@ -423,7 +424,11 @@ async def open( zarr_format = _metadata_dict["zarr_format"] is_v3_array = zarr_format == 3 and _metadata_dict.get("node_type") == "array" if is_v3_array or zarr_format == 2: - return AsyncArray(store_path=store_path, metadata=_metadata_dict, config=config) + array = AsyncArray( + store_path=store_path, metadata=_metadata_dict, config=config + ) + observe_read(store_path, metadata_dict) + return array except (FileNotFoundError, NodeTypeValidationError): pass _check_no_group_config(config, store_path) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index f650388f80..867a2bdd6b 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -2,18 +2,28 @@ import contextlib import copy +import copyreg +import functools +import json import math +import operator +import os +import threading import warnings from asyncio import gather from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass, field, replace +from dataclasses import MISSING, dataclass, field, replace from itertools import starmap from logging import getLogger from typing import ( TYPE_CHECKING, Any, + ClassVar, + Final, Literal, + Required, TypedDict, + Unpack, cast, overload, ) @@ -62,6 +72,7 @@ ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON, + ZGROUP_JSON, ChunksLike, DimensionNamesLike, MemoryOrder, @@ -121,14 +132,18 @@ ) from zarr.core.metadata.io import ( ARRAY_DOCUMENTS, + _write_metadata, encode_documents, parse_stored_array, read_documents, + read_repaired, save_metadata, save_new_metadata, store_documents, + stored_array_document, upsert_metadata, ) +from zarr.core.metadata.stored import moves_chunks, stored_documents from zarr.core.metadata.v2 import ( CompressorLikev2, get_object_codec_id, @@ -138,7 +153,9 @@ from zarr.core.metadata.v3 import ( ChunkGridMetadata, RectilinearChunkGridMetadata, + RectilinearChunkGridMetadataJSON, RegularChunkGridMetadata, + RegularChunkGridMetadataJSON, create_chunk_grid_metadata, parse_node_type_array, ) @@ -161,16 +178,25 @@ from zarr.storage._utils import _relativize_path if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import Callable, Iterator from typing import Self import numpy.typing as npt + from zarr_metadata import ( + ValidationProblem, + ZarrV2ArrayMetadata, + ZarrV2ArrayMetadataJSON, + ZarrV3ArrayMetadata, + ZarrV3ArrayMetadataJSON, + ) from zarr.abc.codec import CodecPipeline from zarr.abc.store import Store from zarr.codecs.sharding import IndexLocation, ShardingCodec from zarr.core.buffer import Buffer + from zarr.core.common import AccessModeLiteral from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar + from zarr.core.metadata.repair import ArrayDocument from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -178,6 +204,10 @@ # Array and AsyncArray are defined in the base `zarr` namespace __all__ = [ "DEFAULT_FILL_VALUE", + "ArrayIO", + "AsyncArrayIO", + "build_array_metadata_v2", + "build_array_metadata_v3", "create_codec_pipeline", "parse_array_metadata", ] @@ -228,24 +258,29 @@ def parse_array_metadata(data: Any, path: str | None = None) -> ArrayMetadata: zarr_format = data.get("zarr_format") if zarr_format == 3: meta_out = ArrayV3Metadata.from_dict(data, path=path) - if len(meta_out.storage_transformers) > 0: - msg = ( - f"Array metadata contains storage transformers: {meta_out.storage_transformers}." - "Arrays with storage transformers are not supported in zarr-python at this time." - ) - raise ValueError(msg) + _reject_storage_transformers(meta_out) return meta_out elif zarr_format == 2: return ArrayV2Metadata.from_dict(data, path=path) else: raise ValueError(f"Invalid zarr_format: {zarr_format}. Expected 2 or 3") - raise TypeError # pragma: no cover + raise TypeError( + "Expected array metadata (ArrayV2Metadata or ArrayV3Metadata) or a metadata " + f"mapping, got {type(data).__name__}." + ) -def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None) -> CodecPipeline: +def create_codec_pipeline( + metadata: ArrayMetadata, + *, + store: Store | None = None, + pipeline_class: type[CodecPipeline] | None = None, +) -> CodecPipeline: + if pipeline_class is None: + pipeline_class = get_pipeline_class() if store is not None: try: - return get_pipeline_class().from_array_metadata_and_store( + return pipeline_class.from_array_metadata_and_store( array_metadata=metadata, store=store ) except NotImplementedError: @@ -258,7 +293,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None # "disables partial reads" warning) fire here; `evolve_from_array_spec` # re-splits the same already-warned-about chain via # `codecs_from_list_unchecked`, so it does not re-emit them. - pipeline = get_pipeline_class().from_codecs(metadata.codecs) + pipeline = pipeline_class.from_codecs(metadata.codecs) # Use the regular chunk shape if available, otherwise use a # placeholder. The ChunkTransform is shape-agnostic — the actual @@ -277,7 +312,7 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None return pipeline.evolve_from_array_spec(chunk_spec) elif isinstance(metadata, ArrayV2Metadata): v2_codec = V2Codec(filters=metadata.filters, compressor=metadata.compressor) - pipeline = get_pipeline_class().from_codecs([v2_codec]) + pipeline = pipeline_class.from_codecs([v2_codec]) chunk_spec = ArraySpec( shape=metadata.chunks, dtype=metadata.dtype, @@ -289,6 +324,13 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None raise TypeError # pragma: no cover +def observe_read(store_path: StorePath, document: dict[str, JSON]) -> None: + """Record the array document `get_array_metadata` read at `store_path`, once it has + been read as metadata, in the store's cache (see `zarr.core.metadata.stored`).""" + zarr_format: ZarrFormat = 2 if document.get("zarr_format") == 2 else 3 + stored_documents(store_path.store).observe(store_path.path, document, zarr_format) + + async def get_array_metadata( store_path: StorePath, zarr_format: ZarrFormat | None = 3 ) -> dict[str, JSON]: @@ -392,11 +434,18 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: The runtime configuration of the array. """ + # The dataclass fields are the ones this class has always had, so equality, + # hashing, `dataclasses.replace`/`fields`/`asdict`, pickling and dask tokens + # see the same array they always did. Each reads (and, under + # `object.__setattr__`, writes) the wrapped binding; see the properties + # attached after the class. metadata: T_ArrayMetadata store_path: StorePath codec_pipeline: CodecPipeline = field(init=False) _chunk_grid: ChunkGrid = field(init=False) config: ArrayConfig + # Not a dataclass field: the binding every field reads. + _io: ClassVar[AsyncArrayIO] @overload def __init__( @@ -420,18 +469,136 @@ def __init__( store_path: StorePath, config: ArrayConfigLike | None = None, ) -> None: + # An AsyncArray reads metadata into the legacy classes, as it always has. metadata_parsed = parse_array_metadata(metadata, str(store_path)) - config_parsed = parse_array_config(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)) + if isinstance(metadata, dict): + # A document given as a dict stands for what the store holds, as the + # document an array is opened from does. + _record_given_document(store_path, metadata, metadata_parsed.zarr_format) object.__setattr__( - self, - "codec_pipeline", - create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), + self, "_io", AsyncArrayIO._from_resolved(metadata_parsed, store_path, config) ) + # An AsyncArray is the resolving layer: build its codec pipeline now, so + # pipeline-construction warnings fire at creation, as they always have. + _ = self._io._resolved_codec_pipeline + + @classmethod + def _from_io(cls, io: AsyncArrayIO) -> AnyAsyncArray: + """An array over a copy of `io` (see `AsyncArrayIO._detached`): nothing done + through the array changes `io`.""" + self = object.__new__(cls) + object.__setattr__(self, "_io", io._detached()) + return self + + @property + def io(self) -> AsyncArrayIO: + """The explicit-IO binding this array delegates to, as of now. A resize or an + append gives the array a binding of its own, so a binding taken before keeps + the shape it had; an attribute update changes the metadata in place, as it + always has, so it shows through a binding taken before, and through copies + of the array that share its metadata.""" + return self._io + + def __getstate__(self) -> dict[str, object] | tuple[dict[str, object], dict[str, object]]: + # The pickled state is the fields', as it has always been, so a pickle + # loads in either version; the binding's own state rides along under one + # key of its own, so no name an instance may hold is taken. + carried = self._io._carried_cache_entry() + metadata: ArrayMetadata = self.metadata + mark = carried.get("_stored_entry", carried.get("_stored_copy")) + if mark is not None or "_stored_document" in vars(metadata): + # zarr 3.4.1 and 3.4.2 read the repair a stored document needs from + # this mark on the metadata itself; this version reads the carried entry. + metadata = _with_legacy_mark(metadata, cast("ArrayDocument | None", mark)) + # Attributes of the instance's own (a subclass's) pickle with it, as + # they did when the fields were its only state. + state: dict[str, object] = {key: v for key, v in vars(self).items() if key != "_io"} + state |= { + "metadata": metadata, + "store_path": self.store_path, + "codec_pipeline": self.codec_pipeline, + "_chunk_grid": self._chunk_grid, + "config": self.config, + } + binding_state: dict[str, object] = {"pipeline_class": self._io._pipeline_class} + if type(self._io) is not AsyncArrayIO: + # Carried only when it is a subclass, so the pickle stays loadable where + # that class is not importable only if it is used. + binding_state["binding_class"] = type(self._io) + source = self._io._source + if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + # The document as stored, which the legacy metadata cannot always + # reproduce (a NaN payload, a codec's configuration), as its bytes, + # so a zarr without zarr-metadata loads the pickle too. + binding_state["model_documents"] = dict(source.to_key_value()) + # The metadata the model belongs with (one object in the pickle): a + # zarr that does not read the binding's state keeps it as it was while + # changing the metadata (a resize), and the model is then stale. + binding_state["metadata"] = metadata + binding_state |= carried + state[_BINDING_STATE] = binding_state + # Slots as the default reduction carries them, so either version loads them. + slots = _slot_values(self) + return (state, slots) if slots else state + + def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: + if isinstance(state, tuple): + # A subclass with slots: the instance's dict, if any, and its slots' values. + state = {**(state[0] or {}), **(state[1] or {})} + if not state.keys() >= _ASYNC_ARRAY_FIELDS: + # A subclass's own state, from its own reduction, whose factory made + # the array: set it as the default protocol would. + for key, value in state.items(): + object.__setattr__(self, key, value) + return + metadata = state["metadata"] + pipeline = state["codec_pipeline"] + binding_state = cast("dict[str, Any]", state.get(_BINDING_STATE, {})) + io = object.__new__(binding_state.get("binding_class", AsyncArrayIO)) + # The class the binding chose: a pipeline may build an instance of + # another one (zarrs-python falls back for stores it cannot use). + pipeline_class = binding_state.get("pipeline_class", type(pipeline)) + documents = binding_state.get("model_documents") + current = binding_state.get("metadata") is metadata + model = ( + None if documents is None or not current else _model_from_documents(documents, metadata) + ) + io._bind( + metadata if model is None else model, + state["store_path"], + state["config"], + pipeline_class, + ) + if model is not None: + io._resolved_cache = metadata + io._codec_pipeline_cache = pipeline + io._chunk_grid_cache = state["_chunk_grid"] + carried: dict[str, Any] = {k: binding_state[k] for k in _CARRIED_KEYS if k in binding_state} + if _BINDING_STATE not in state: + # A pickle from zarr 3.4.1 or 3.4.2 marks the metadata itself. + legacy_mark = vars(metadata).get("_stored_document") + if legacy_mark is not None: + carried = {"_stored_entry": legacy_mark} + if "_stored_document" in vars(metadata): + # Read once; the cache holds it from here on. + object.__delattr__(metadata, "_stored_document") + io._restore_cache_entry(carried) + object.__setattr__(self, "_io", io) + for key, value in state.items(): + if key not in _ASYNC_ARRAY_STATE_KEYS: + object.__setattr__(self, key, value) + + def __copy__(self) -> Self: + if _overrides_reduction(type(self)): + # A subclass that decides how it is copied is copied its way. + return cast("Self", _copied_by_reduction(self)) + # A shallow copy shares the field values, as a dataclass copy does (a + # change to either takes a binding of its own; see `_own_binding`). + copied = object.__new__(type(self)) + copied.__dict__.update(vars(self)) + for name, value in _slot_values(self).items(): + object.__setattr__(copied, name, value) + return copied @classmethod async def _create( @@ -584,43 +751,17 @@ def _create_metadata_v3( dimension_names: DimensionNamesLike = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV3Metadata: - """Create an instance of ArrayV3Metadata.""" - filters: tuple[ArrayArrayCodec, ...] - compressors: tuple[BytesBytesCodec, ...] - - shape = parse_shapelike(shape) - if codecs is None: - # no data types have default filters - filters = () - serializer = default_serializer_v3(dtype) - compressors = default_compressors_v3(dtype) - - codecs_parsed = (*filters, serializer, *compressors) - else: - codecs_parsed = tuple(codecs) - - chunk_key_encoding_parsed: ChunkKeyEncodingLike - if chunk_key_encoding is None: - chunk_key_encoding_parsed = {"name": "default", "separator": "/"} - else: - chunk_key_encoding_parsed = chunk_key_encoding - - if fill_value is DEFAULT_FILL_VALUE or fill_value is None: - # Use dtype's default scalar for the DEFAULT_FILL_VALUE sentinel - # For v3, None is converted to DEFAULT_FILL_VALUE behavior - fill_value_parsed = dtype.default_scalar() - else: - fill_value_parsed = fill_value - - return ArrayV3Metadata( + """See the module-level `build_array_metadata_v3`.""" + return build_array_metadata_v3( shape=shape, data_type=dtype, chunk_grid=chunk_grid, - chunk_key_encoding=chunk_key_encoding_parsed, - fill_value=fill_value_parsed, - codecs=codecs_parsed, # type: ignore[arg-type] - dimension_names=tuple(dimension_names) if dimension_names else None, - attributes=attributes or {}, + fill_value=fill_value, + chunk_key_encoding=chunk_key_encoding, + codecs=codecs, + # As these paths always have, an empty list declares no names. + dimension_names=dimension_names or None, + attributes=attributes, ) @classmethod @@ -677,25 +818,17 @@ def _create_metadata_v2( compressor: CompressorLikev2 = None, attributes: dict[str, JSON] | None = None, ) -> ArrayV2Metadata: - if dimension_separator is None: - dimension_separator = "." - - # Handle the DEFAULT_FILL_VALUE sentinel - if fill_value is DEFAULT_FILL_VALUE: - fill_value_parsed: Any = dtype.default_scalar() - else: - # For v2, preserve None as-is (backward compatibility) - fill_value_parsed = fill_value - - return ArrayV2Metadata( + """See the module-level `build_array_metadata_v2`; this legacy path makes + no object codec check, as it never has (see `_array_metadata_v2`).""" + return _array_metadata_v2( shape=shape, dtype=dtype, chunks=chunks, order=order, - dimension_separator=dimension_separator, - fill_value=fill_value_parsed, - compressor=compressor, + dimension_separator=dimension_separator if dimension_separator is not None else ".", + fill_value=fill_value, filters=filters, + compressor=compressor, attributes=attributes, ) @@ -830,7 +963,9 @@ async def example(): metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) # TODO: remove this cast when we have better type hints _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict) - return cls(store_path=store_path, metadata=_metadata_dict, config=config) + array = cls(store_path=store_path, metadata=_metadata_dict, config=config) + observe_read(store_path, metadata_dict) + return array @property def store(self) -> Store: @@ -1236,7 +1371,12 @@ def with_config(self, config: ArrayConfigLike) -> Self: # Merge new config with existing config, so missing keys are inherited # from the current array rather than from global defaults new_config = ArrayConfig(**{**self.config.to_dict(), **config}) # type: ignore[arg-type] - return type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) + # Built through the constructor, as it always was, so a subclass's own + # __init__ runs; then over this array's kind of binding (its class, and + # its model, so the document stays as stored). + rebuilt = type(self)(metadata=self.metadata, store_path=self.store_path, config=new_config) + object.__setattr__(rebuilt, "_io", self._io._rebound_like(rebuilt._io, rebuilt.metadata)) + return rebuilt async def nchunks_initialized(self) -> int: """ @@ -1486,17 +1626,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, @@ -1540,15 +1670,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,56 +1746,12 @@ async def get_coordinate_selection( 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() - - 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) - - 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 + """Store `metadata` as this array's own documents: through the wrapped + `AsyncArrayIO` when it is this array's metadata.""" + if metadata is self.metadata: + await self._io.store_metadata(ensure_parents=ensure_parents) else: - 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() + await save_metadata(self.store_path, metadata, ensure_parents=ensure_parents) async def _set_selection( self, @@ -1683,21 +1761,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, @@ -1738,10 +1802,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]: @@ -1932,6 +1993,140 @@ def _info( ) +class _BindingField: + """A dataclass field of `AsyncArray`, read from and written to its binding. On the + class it reads as `dataclasses.MISSING`, so a `@dataclass` subclass that annotates + the field again finds no default, as for a field without one.""" + + def __init__( + self, + get: Callable[[AsyncArray[Any]], Any], + set_: Callable[[AsyncArray[Any], Any], None], + ) -> None: + self._get = get + self._set = set_ + self.__doc__ = get.__doc__ + + def __get__(self, instance: AsyncArray[Any] | None, owner: type | None = None) -> Any: + if instance is None: + return MISSING + return self._get(instance) + + def __set__(self, instance: AsyncArray[Any], value: object) -> None: + self._set(instance, value) + + +def _get_metadata(self: AsyncArray[Any]) -> ArrayMetadata: + """The in-memory metadata of this array.""" + return self._io.metadata_legacy + + +def _copied_by_reduction(array: AsyncArray[Any]) -> object: + """A shallow copy of `array` made from its reduction, as `copy.copy` makes one.""" + reduced = array.__reduce_ex__(4) + if isinstance(reduced, str): + return array + rebuild, args, *rest = reduced + copied = rebuild(*args) + state = rest[0] if rest else None + if state is not None: + copied.__setstate__(state) + return copied + + +def _model_from_documents( + documents: Mapping[str, bytes], metadata: ArrayMetadata +) -> ArrayMetadataModel: + """The model a pickled `AsyncArray` carried as its documents' bytes.""" + model_v2, model_v3 = _model_classes() + model_class = model_v2 if metadata.zarr_format == 2 else model_v3 + return model_class.from_key_value(cast("Any", documents)) + + +def _slot_values(array: AsyncArray[Any]) -> dict[str, object]: + """The values in the `__slots__` a subclass of `AsyncArray` declares, which + `vars` leaves out, by name; unset slots are left out.""" + return { + name: getattr(array, name) + for name in copyreg._slotnames(type(array)) # type: ignore[attr-defined] + if hasattr(array, name) + } + + +_REDUCTION_HOOKS: Final = ("__reduce_ex__", "__reduce__", "__getstate__", "__setstate__") + + +def _overrides_reduction(cls: type) -> bool: + """Whether a subclass of `AsyncArray` defines how it is pickled or copied.""" + return any(getattr(cls, hook) is not getattr(AsyncArray, hook) for hook in _REDUCTION_HOOKS) + + +def _own_binding(self: AsyncArray[Any]) -> AsyncArrayIO: + """A copy of this array's binding, made the array's own, for a change to its state. + + An array's state is its own, as when its fields were its only state: another array + made by copying this one (its `__dict__`), or a binding taken from `io`, keeps the + state it had. + """ + io = object.__new__(type(self._io)) + io.__dict__.update(vars(self._io)) + object.__setattr__(self, "_io", io) + return io + + +def _set_metadata(self: AsyncArray[Any], value: ArrayMetadata) -> None: + _own_binding(self)._rebind_resolved(value) + + +def _get_store_path(self: AsyncArray[Any]) -> StorePath: + """The location in a store this array is bound to.""" + return self._io.store_path + + +def _set_store_path(self: AsyncArray[Any], value: StorePath) -> None: + binding = _own_binding(self) + binding._rebind_at(value, binding._config) + + +def _get_config(self: AsyncArray[Any]) -> ArrayConfig: + """The runtime configuration of this array.""" + return self._io.config + + +def _set_config(self: AsyncArray[Any], value: ArrayConfig) -> None: + binding = _own_binding(self) + binding._rebind_at(binding._store_path, value) + + +def _get_codec_pipeline(self: AsyncArray[Any]) -> CodecPipeline: + """The codec pipeline used for encoding and decoding chunks.""" + return self._io._resolved_codec_pipeline + + +def _set_codec_pipeline(self: AsyncArray[Any], value: CodecPipeline) -> None: + _own_binding(self)._codec_pipeline_cache = value + + +def _get_chunk_grid(self: AsyncArray[Any]) -> ChunkGrid: + return self._io._chunk_grid + + +def _set_chunk_grid(self: AsyncArray[Any], value: ChunkGrid) -> None: + _own_binding(self)._chunk_grid_cache = value + + +# The dataclass fields of `AsyncArray` read the wrapped binding. They are +# attached once the dataclass is made, so `@dataclass` sees fields, not +# defaults; the setters keep `object.__setattr__(array, name, value)` working. +setattr(AsyncArray, "metadata", _BindingField(_get_metadata, _set_metadata)) # noqa: B010 +setattr(AsyncArray, "store_path", _BindingField(_get_store_path, _set_store_path)) # noqa: B010 +setattr(AsyncArray, "config", _BindingField(_get_config, _set_config)) # noqa: B010 +setattr( # noqa: B010 + AsyncArray, "codec_pipeline", _BindingField(_get_codec_pipeline, _set_codec_pipeline) +) +setattr(AsyncArray, "_chunk_grid", _BindingField(_get_chunk_grid, _set_chunk_grid)) # noqa: B010 + + # TODO: Array can be a frozen data class again once property setters (e.g. shape) are removed @dataclass(frozen=False) class Array[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: @@ -1951,6 +2146,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: """ @@ -4672,7 +4872,7 @@ async def init_array( order_parsed = order chunk_key_encoding_parsed = cast("V2ChunkKeyEncoding", chunk_key_encoding_parsed) - meta = AsyncArray._create_metadata_v2( + meta = _array_metadata_v2( shape=shape_parsed, dtype=zdtype, chunks=outer_chunks.chunk_shape, @@ -4722,14 +4922,15 @@ async def init_array( if order is not None: _warn_order_kwarg() - meta = AsyncArray._create_metadata_v3( + meta = build_array_metadata_v3( shape=shape_parsed, - dtype=zdtype, + data_type=zdtype, fill_value=fill_value, chunk_grid=grid, chunk_key_encoding=chunk_key_encoding_parsed, codecs=codecs_out, - dimension_names=dimension_names, + # As these paths always have, an empty list declares no names. + dimension_names=dimension_names or None, attributes=attributes, ) @@ -4738,6 +4939,143 @@ async def init_array( return arr +def _array_metadata_v2( + *, + shape: ShapeLike, + dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + chunks: ShapeLike, + fill_value: Any = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + dimension_separator: Literal[".", "/"] = ".", + filters: Iterable[dict[str, JSON] | Numcodec] | None = None, + compressor: CompressorLikev2 = None, + attributes: dict[str, JSON] | None = None, +) -> ArrayV2Metadata: + """`build_array_metadata_v2` without its object codec check, which the + callers that create arrays from user arguments make (or, for the legacy + `create` path, have never made) themselves.""" + zdtype = parse_dtype(dtype, zarr_format=2) + if order is None: + order = zarr_config.get("array.order") + if fill_value is DEFAULT_FILL_VALUE: + fill_value_parsed: Any = zdtype.default_scalar() + else: + # An explicit None is kept: the field is nullable in Zarr format 2. + fill_value_parsed = fill_value + + return ArrayV2Metadata( + shape=parse_shapelike(shape), + dtype=zdtype, + chunks=parse_shapelike(chunks), + order=order, + dimension_separator=dimension_separator, + fill_value=fill_value_parsed, + compressor=compressor, + filters=filters, + attributes=attributes, + ) + + +def build_array_metadata_v2( + *, + shape: ShapeLike, + dtype: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + chunks: ShapeLike, + fill_value: Any = DEFAULT_FILL_VALUE, + order: MemoryOrder | None = None, + dimension_separator: Literal[".", "/"] = ".", + filters: Iterable[dict[str, JSON] | Numcodec] | None = None, + compressor: CompressorLikev2 = None, + attributes: dict[str, JSON] | None = None, +) -> ArrayV2Metadata: + """Build Zarr format 2 array metadata from its fields. + + Each keyword is a `.zarray` field; `attributes` is the `.zattrs` one. + `fill_value` defaults to the data type's default scalar (an explicit + `None` is kept: the field is nullable); `order` defaults from the global + `array.order`. `filters` and `compressor` are numcodecs objects or their + JSON, as stored. Nothing is guessed. + """ + metadata = _array_metadata_v2( + shape=shape, + dtype=dtype, + chunks=chunks, + fill_value=fill_value, + order=order, + dimension_separator=dimension_separator, + filters=filters, + compressor=compressor, + attributes=attributes, + ) + _check_object_codec_v2(metadata.dtype, metadata.filters, metadata.compressor) + return metadata + + +def build_array_metadata_v3( + *, + shape: ShapeLike, + data_type: ZDTypeLike | ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGridMetadata | Mapping[str, JSON] | ShapeLike, + fill_value: Any = DEFAULT_FILL_VALUE, + chunk_key_encoding: ChunkKeyEncodingLike | None = None, + codecs: Iterable[Codec | dict[str, JSON]] | None = None, + dimension_names: DimensionNamesLike = None, + attributes: dict[str, JSON] | None = None, +) -> ArrayV3Metadata: + """Build Zarr format 3 array metadata from its fields. + + Each keyword is a `zarr.json` field. `chunk_grid` takes grid metadata, + its JSON, or a chunk shape, which declares a regular grid. `codecs` + `None` is the data type's default chain; `chunk_key_encoding` `None` is + the default encoding with `/`; `fill_value` defaults to the data type's + default scalar (`None` means the default too: the field is required). + """ + zdtype = parse_dtype(data_type, zarr_format=3) + grid: ChunkGridMetadata | dict[str, JSON] + if isinstance(chunk_grid, RegularChunkGridMetadata | RectilinearChunkGridMetadata): + grid = chunk_grid + elif isinstance(chunk_grid, Mapping): + grid = dict(chunk_grid) + else: + grid = RegularChunkGridMetadata(chunk_shape=parse_shapelike(chunk_grid)) + + filters: tuple[ArrayArrayCodec, ...] + compressors: tuple[BytesBytesCodec, ...] + if codecs is None: + # no data types have default filters + filters = () + serializer = default_serializer_v3(zdtype) + compressors = default_compressors_v3(zdtype) + + codecs_parsed = (*filters, serializer, *compressors) + else: + codecs_parsed = tuple(codecs) + + chunk_key_encoding_parsed: ChunkKeyEncodingLike + if chunk_key_encoding is None: + chunk_key_encoding_parsed = {"name": "default", "separator": "/"} + else: + chunk_key_encoding_parsed = chunk_key_encoding + + if fill_value is DEFAULT_FILL_VALUE or fill_value is None: + # The dtype's default scalar: the field is required in Zarr format 3, + # so None also means the default. + fill_value_parsed = zdtype.default_scalar() + else: + fill_value_parsed = fill_value + + return ArrayV3Metadata( + shape=parse_shapelike(shape), + data_type=zdtype, + chunk_grid=grid, + chunk_key_encoding=chunk_key_encoding_parsed, + fill_value=fill_value_parsed, + codecs=codecs_parsed, # type: ignore[arg-type] + dimension_names=tuple(dimension_names) if dimension_names is not None else None, + attributes=attributes if attributes is not None else {}, + ) + + async def create_array( store: StoreLike, *, @@ -4960,6 +5298,69 @@ def _sharding_codec(metadata: ArrayMetadata) -> ShardingCodec | None: return None +async def _store_repaired_document(store_path: StorePath, metadata: ArrayMetadata) -> None: + """Store the repair of the array document at `store_path`, if it needs one that + moves chunks, before chunks are written under `metadata`. + + The store's cache (see `zarr.core.metadata.stored`) says whether the document + needs such a repair. Where it records nothing for the document (it records a + valid one only where that replaces a broken one), a consolidated copy it has + seen that needs a repair stands for it, as consolidated metadata does for every + reader, and the document is read to check; with neither, nothing is known to + need a repair, and nothing is read: metadata built in code, or bound to a store + that never read the document, writes its chunks as any array does. Only a document whose repair + lays out chunks as `metadata` does is acted on: one that lays them out otherwise + describes another array, and chunks written under `metadata` are its caller's + declaration, as over any document. The document is read again, because the store + may hold a newer one. If that one lays out chunks differently (another writer + stored a different chunk size since), the caller would write chunks no reader + finds, so it raises, stores nothing, and leaves the cache as it was, so later + writes raise too, until a binding of this store object reads or stores a newer + document. If it needs no repair (the array was re-saved since, + possibly by another implementation), it is left as written; if there is none, + there is nothing to repair. + + The check is only as fresh as the cache. Once a binding of the same store object + has read or stored a newer document (for example, `create_array(..., + overwrite=True)` on that store), the cache no longer records the repair, so an + array opened before that writes its chunks under its own metadata unchecked, as + any array whose document another writer replaced does: outside the check above, + zarr does not detect a stale array handle. + + Storing the same repair twice is harmless, so handles that write chunks + concurrently need no coordination. The read and the store are not one atomic + 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. + """ + cache = stored_documents(store_path.store) + path = store_path.path + zarr_format = metadata.zarr_format + keys = ARRAY_DOCUMENTS[zarr_format] + if cache.knows(path, zarr_format): + stored = cache.needing_repair(path, zarr_format) + else: + stored = cache.copy_needing_repair(path, zarr_format) + if stored is None or _chunk_layout(read_repaired(stored, zarr_format)) != _chunk_layout( + metadata + ): + return + documents = await read_documents(store_path, keys) + document = stored_array_document(documents, zarr_format) + if document is None: + cache.observe(path, None, zarr_format) + return + current = parse_stored_array(documents, zarr_format, str(store_path)) + if _chunk_layout(current) != _chunk_layout(metadata): + raise ValueError( + f"Array {str(store_path)!r}: the metadata stored has changed since " + "this array was opened; reopen the array to write to it. Nothing was stored." + ) + if moves_chunks(document, zarr_format): + await upsert_metadata(store_path, current, documents) + cache.observe(path, None, zarr_format) + + def _stored_rectilinear_grid_or_none( metadata: ArrayMetadata, ) -> RectilinearChunkGridMetadata | None: @@ -5232,18 +5633,28 @@ def _parse_chunk_encoding_v2( ) raise TypeError(msg) _filters = parse_filters(filters) + _check_object_codec_v2(dtype, _filters, _compressor) + return _filters, _compressor + + +def _check_object_codec_v2( + dtype: ZDType[TBaseDType, TBaseScalar], + filters: tuple[Numcodec, ...] | None, + compressor: Numcodec | None, +) -> None: + """Raise if `dtype` needs an object codec that neither `filters` nor `compressor` holds.""" if isinstance(dtype, HasObjectCodec): # check the filters and the compressor for the object codec required for this data type - if _filters is None: - if _compressor is None: + if filters is None: + if compressor is None: object_codec_id = None else: - object_codec_id = get_object_codec_id((_compressor.get_config(),)) + object_codec_id = get_object_codec_id((compressor.get_config(),)) else: object_codec_id = get_object_codec_id( ( - *[f.get_config() for f in _filters], - _compressor.get_config() if _compressor is not None else None, + *[f.get_config() for f in filters], + compressor.get_config() if compressor is not None else None, ) ) if object_codec_id is None: @@ -5259,7 +5670,6 @@ def _parse_chunk_encoding_v2( "this array. " ) raise ValueError(msg) - return _filters, _compressor def _parse_chunk_encoding_v3( @@ -6010,12 +6420,10 @@ async def _delete_chunk(coords: tuple[int, ...]) -> None: ) # Write new metadata - await store_documents(array.store_path, documents) - array._stored_document_replaced() + await store_documents(array.store_path, documents, written=new_metadata) - # 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. + _own_binding(array)._rebind_resolved(new_metadata, new_chunk_grid) async def _append( @@ -6101,6 +6509,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. + _own_binding(array)._rebind_resolved(array.metadata, array._chunk_grid) # Write new metadata await array._save_metadata(array.metadata) @@ -6130,3 +6540,1136 @@ 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() + + +_CARRIED_KEYS: Final = ("_stored_entry", "_stored_copy") +"""The keys a pickled binding carries its store cache's entry under.""" + +_BINDING_STATE: Final = "_zarr_binding_state" +"""The key a pickled `AsyncArray` carries its binding's own state under: its pipeline +class, its class when a subclass, a model's documents, its store cache's entry.""" + +_ASYNC_ARRAY_FIELDS: Final = frozenset( + {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config"} +) +"""The fields a pickled `AsyncArray`'s state holds; a state without them all is a +subclass's own.""" + +_ASYNC_ARRAY_STATE_KEYS: Final = _ASYNC_ARRAY_FIELDS | {_BINDING_STATE} +"""The keys of a pickled `AsyncArray`'s state that rebuild its binding; any other key is +an attribute of the instance's own.""" + + +def _bound_metadata_dict( + data: dict[str, Any], store_path: StorePath +) -> ArrayMetadataModel | ArrayMetadata: + """A metadata dict given to a binding: read as a model when it is valid + JSON metadata, so names resolve only when chunks are read or written; + otherwise (metadata built in code, or an invalid document zarr reads as + repaired) read into the legacy classes now, warning as `AsyncArray` always + has.""" + from zarr_metadata import MetadataValidationError + + path = str(store_path) + # Metadata written in Python may hold tuples where its JSON holds lists; the + # repairs read JSON. + data = cast("dict[str, Any]", _lists_for_tuples(data)) + zarr_format: ZarrFormat = 2 if data.get("zarr_format") == 2 else 3 + if data.get("storage_transformers"): + # Bound as legacy metadata, which rejects them, as AsyncArray always has. + return parse_array_metadata(data, path) + if not moves_chunks(data, zarr_format): + # A document zarr reads as repaired, moving chunks, is never bound as a + # model, also where the model accepts it. + try: + return _array_metadata_model(data) + except MetadataValidationError: + pass + legacy = parse_array_metadata(data, path) + # An invalid document given as a dict stands for what the store holds, as + # one given to AsyncArray does. + _record_given_document(store_path, data, legacy.zarr_format) + return legacy + + +def _lists_for_tuples(value: object) -> object: + """`value` with every tuple, at any depth in its dicts and lists, made a list.""" + if isinstance(value, dict): + return {key: _lists_for_tuples(item) for key, item in value.items()} + if isinstance(value, list | tuple): + return [_lists_for_tuples(item) for item in value] + return value + + +def _record_given_document( + store_path: StorePath, document: Mapping[str, Any], zarr_format: ZarrFormat +) -> None: + """Record `document` as the one stored at `store_path` when reading it needs a + repair that moves chunks, unless the store's cache already knows that document.""" + if moves_chunks(document, zarr_format): + stored_documents(store_path.store).restore( + store_path.path, zarr_format, copy.deepcopy(dict(document)) + ) + + +def _differing_fields( + stored: ArrayMetadataModel, + stored_json: dict[str, Any], + bound_json: dict[str, Any], +) -> list[str]: + """The fields of `stored` that differ from `bound_json`'s as the model compares + them: a field whose JSON differs (so `-0.0` from `0.0`) counts unless putting + the bound value in `stored` leaves it equal (an absent field against its + default).""" + from zarr_metadata import MetadataValidationError + + differing = [] + for key in sorted(stored_json.keys() | bound_json.keys()): + stored_value, bound_value = stored_json.get(key), bound_json.get(key) + if json.dumps(stored_value, sort_keys=True) == json.dumps(bound_value, sort_keys=True): + continue + patched = dict(stored_json) + if key in bound_json: + patched[key] = bound_value + else: + patched.pop(key) + try: + unchanged = type(stored)(patched) == stored + except MetadataValidationError: + unchanged = False + if not unchanged: + differing.append(key) + return differing + + +def _reject_storage_transformers(metadata: ArrayMetadataModel | ArrayMetadata) -> None: + """Raise for metadata (a model or a legacy object) that declares storage + transformers: zarr-python cannot read or write chunks through them.""" + transformers = getattr(metadata, "storage_transformers", ()) + if transformers: + raise ValueError( + f"Array metadata contains storage transformers: {transformers}. " + "Arrays with storage transformers are not supported in zarr-python at this time." + ) + + +def _is_group_document(document: Buffer) -> bool: + """Whether a stored `zarr.json` is a group's; one that is not JSON is not.""" + try: + decoded = buffer_to_json_object(document) + except (ValueError, TypeError): + return False + return decoded.get("node_type") == "group" + + +def _with_legacy_mark(metadata: ArrayMetadata, mark: ArrayDocument | None) -> ArrayMetadata: + """A copy of `metadata` carrying `mark` as the `_stored_document` attribute zarr + 3.4.1 and 3.4.2 read the repair a stored document needs from (none: no mark).""" + marked = copy.copy(metadata) + if mark is None: + if "_stored_document" in vars(marked): + object.__delattr__(marked, "_stored_document") + else: + object.__setattr__(marked, "_stored_document", copy.deepcopy(mark)) + return marked + + +def _copied_legacy(source: ArrayMetadata) -> ArrayMetadata: + """A copy of a legacy metadata object with its own attributes dict, so an + attribute update through one binding (which replaces top-level keys in + place) never reaches another.""" + copied = copy.copy(source) + object.__setattr__(copied, "attributes", dict(source.attributes)) + return copied + + +type ArrayMetadataModel = ZarrV2ArrayMetadata | ZarrV3ArrayMetadata +"""Array metadata as a `zarr_metadata` model: validated JSON, extension names unresolved.""" +type ArrayMetadataModelJSON = ZarrV2ArrayMetadataJSON | ZarrV3ArrayMetadataJSON +"""Array metadata as `zarr_metadata` types it in JSON (zarrista's metadata type).""" + + +def _metadata_model_json(document: ArrayMetadataModel) -> dict[str, Any]: + """The metadata as plain, mutable JSON, via JSON text: the model returns read-only forms.""" + return cast("dict[str, Any]", _plain_json(document.to_json())) + + +def _plain_json(value: object) -> Any: + """`value`, a model's read-only JSON (mapping proxies at any depth), as plain JSON.""" + return json.loads(json.dumps(value, default=_mapping_as_dict)) + + +def _document_without_attributes(binding: AsyncArrayIO) -> ArrayMetadataModel: + """The document `binding` holds, as a model, with no attributes.""" + source = binding._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + return _model_of_legacy(replace(source, attributes={})) + return source.update(attributes={}) + + +def _model_of_legacy(metadata: ArrayMetadata) -> ArrayMetadataModel: + """The document `metadata` encodes, as a model.""" + buffers = metadata.to_buffer_dict(default_buffer_prototype()) + key_value = {key: buffer.to_bytes() for key, buffer in buffers.items()} + model_v2, model_v3 = _model_classes() + model_class = model_v2 if isinstance(metadata, ArrayV2Metadata) else model_v3 + return model_class.from_key_value(key_value) + + +def _json_text(value: object) -> str: + """`value`, JSON with mapping proxies at any depth, as JSON text in its own order + (keys made strings, as JSON makes them), so a change of order is a change.""" + return json.dumps(_plain_json(value)) + + +def _mapping_as_dict(value: object) -> dict[str, Any]: + """`json.dumps`' fallback for `_plain_json`: a mapping as a dict; anything else is + not JSON.""" + if isinstance(value, Mapping): + return dict(value) + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def _array_metadata_model(data: dict[str, JSON]) -> ArrayMetadataModel: + """Model the metadata dict `get_array_metadata` returned. + + The read is strict: metadata zarr can only read as repaired raises + `zarr_metadata.MetadataValidationError`. Unknown extension names read as + unclaimed fields, not errors. + """ + from zarr_metadata import MetadataValidationError + + model_v2, model_v3 = _model_classes() + if data.get("zarr_format") != 2: + return model_v3(data) + chunks = data.get("chunks") + chunk_problems = _v2_chunk_problems(chunks) if isinstance(chunks, list | tuple) else [] + try: + model = model_v2(data) + except MetadataValidationError as error: + # The model rejects other chunk sizes itself: add those it accepts, a 0. + unreported = [ + problem + for problem in chunk_problems + if type(problem.input) is int and problem.input == 0 + ] + if not unreported: + raise + raise MetadataValidationError([*error.problems, *unreported]) from None + if chunk_problems: + raise MetadataValidationError(chunk_problems) + return model + + +def _checked_model(model: ArrayMetadataModel) -> ArrayMetadataModel: + """`model`, once checked as zarr checks a model it binds: it declares no storage + transformers, and a Zarr format 2 model no chunk size zarr reads as repaired.""" + from zarr_metadata import MetadataValidationError + + _reject_storage_transformers(model) + if isinstance(model, _model_classes()[0]) and (problems := _v2_chunk_problems(model.chunks)): + raise MetadataValidationError(problems) + return model + + +def _v2_chunk_problems(chunks: Sequence[object]) -> list[ValidationProblem]: + """A problem for each Zarr format 2 chunk size that is not an integer >= 1: the v2 + model accepts a 0, which zarr reads as repaired.""" + from zarr_metadata import ValidationProblem + + return [ + ValidationProblem( + ("chunks", index), + f"expected an integer >= 1, got {size!r}", + "invalid_value", + input=cast("JSON", size), + ) + for index, size in enumerate(chunks) + if not (type(size) is int and size >= 1) + ] + + +@functools.cache +def _model_classes() -> tuple[type[ZarrV2ArrayMetadata], type[ZarrV3ArrayMetadata]]: + """The `zarr_metadata` model classes, imported on first use: the import is + slow, and only bindings that hold a model need it.""" + from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata + + return ZarrV2ArrayMetadata, ZarrV3ArrayMetadata + + +class ArrayV2DeclarationKwargs(TypedDict, total=False): + """The fields that declare a Zarr format 2 array: what `build_array_metadata_v2` and the `build_v2` methods take as keywords.""" + + shape: Required[ShapeLike] + dtype: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] + chunks: Required[ShapeLike] + fill_value: Any + order: MemoryOrder + dimension_separator: Literal[".", "/"] + filters: Iterable[dict[str, JSON] | Numcodec] | None + compressor: CompressorLikev2 + attributes: dict[str, JSON] | None + + +class ArrayV3DeclarationKwargs(TypedDict, total=False): + """The fields that declare a Zarr format 3 array: what `build_array_metadata_v3` and the `build_v3` methods take as keywords.""" + + shape: Required[ShapeLike] + data_type: Required[ZDTypeLike | ZDType[TBaseDType, TBaseScalar]] + chunk_grid: Required[ChunkGridMetadata | Mapping[str, JSON] | ShapeLike] + fill_value: Any + chunk_key_encoding: ChunkKeyEncodingLike | None + codecs: Iterable[Codec | dict[str, JSON]] | None + dimension_names: DimensionNamesLike + attributes: dict[str, JSON] | None + + +def _bound_store_path(store: Store | StorePath, name: str | None) -> StorePath: + """The location `name` within `store`, bound without IO.""" + store_path = store if isinstance(store, StorePath) else StorePath(store) + if name is not None: + store_path = store_path / name + return store_path + + +_FIRST_RESOLUTION_LOCK: Final = threading.Lock() +"""Held while a binding stores the legacy metadata it resolved first.""" + + +class AsyncArrayIO: + """Explicit asynchronous IO for one Zarr array. + + Binds array metadata (`metadata`, a `zarr_metadata` model) to a location + in a store. Construction performs no IO; every read or write of the store + is an explicit method call, and every write has an erasing inverse. The + metadata operations (`store_metadata`, `erase_metadata`, `erase_chunk`, + `erase_chunks`, `verify`) never resolve extension names; with the codec + pipelines zarr-python has today, `getitem` and `setitem` resolve them into + the legacy classes (`metadata_legacy`) on first use. The bound metadata is assumed, never checked against the + store, and an array opened from metadata that needed a repair stores that + repair before its first chunk write. The codec pipeline implementation is + chosen from the global configuration when the binding is made; the + pipeline itself, with any options it reads from the configuration, is + built at the first data operation. `AsyncArray` wraps an `AsyncArrayIO` + and delegates its metadata and selection IO to it. + + Every way a binding is made or changed (construction, the builders, + `open`, the `with_*` methods, pickling, and the resize or attribute update + of an array over it) goes through `_bind`, so a subclass that derives + state of its own from the binding (a native handle) derives it there. + Attributes alone may change in place, without a `_bind` (an `AsyncArray` + over the binding updates them so, as it always has): derived state must not + capture them, and writing metadata reads `metadata`. Pickling a binding + pickles its attributes and unpickling binds it again; a subclass whose + derived state does not pickle leaves it out in `__getstate__`. + """ + + _store_path: StorePath + _config: ArrayConfig + _pipeline_class: type[CodecPipeline] + _source: ArrayMetadataModel | ArrayMetadata + _resolved_cache: ArrayMetadata | None + _chunk_grid_cache: ChunkGrid | None + _chunk_key_encoding_cache: ChunkKeyEncoding | None + _codec_pipeline_cache: CodecPipeline | None + _attributes_cache: tuple[str, ArrayMetadataModel] | None + + def __init__( + self, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: + source: ArrayMetadataModel | ArrayMetadata + if isinstance(metadata, ArrayMetadata): + # A resolved object is read (v2 with a chunk size of 0 as repaired) + # and never shared with the caller. + source = parse_array_metadata(metadata, str(store_path)) + if source is metadata: + source = _copied_legacy(source) + elif isinstance(metadata, _model_classes()): + source = _checked_model(metadata) + elif isinstance(metadata, Mapping): + source = _bound_metadata_dict(cast("dict[str, Any]", metadata), store_path) + else: + raise TypeError( + "Expected array metadata (a zarr_metadata model, ArrayV2Metadata or " + f"ArrayV3Metadata) or a metadata mapping, got {type(metadata).__name__}." + ) + # The pipeline implementation is chosen when the binding is made, from + # the global configuration, exactly as `AsyncArray` has always chosen + # it; the pipeline itself is built on first use. + self._bind(source, store_path, parse_array_config(config), get_pipeline_class()) + + def _bind( + self, + source: ArrayMetadataModel | ArrayMetadata, + store_path: StorePath, + config: ArrayConfig, + pipeline_class: type[CodecPipeline], + ) -> None: + self._store_path = store_path + self._config = config + self._pipeline_class = pipeline_class + self._source = source + self._resolved_cache = None + self._chunk_grid_cache = None + self._chunk_key_encoding_cache = None + self._codec_pipeline_cache = None + self._attributes_cache = None + + @classmethod + def _from_resolved( + cls, source: ArrayMetadata, store_path: StorePath, config: ArrayConfigLike | None + ) -> Self: + """A binding over `source` itself, already resolved (as `AsyncArray` binds the + metadata it is given): read as `parse_array_metadata` reads a metadata + object, but not copied.""" + binding = object.__new__(cls) + binding._bind( + parse_array_metadata(source, str(store_path)), + store_path, + parse_array_config(config), + get_pipeline_class(), + ) + return binding + + @classmethod + def build_v2( + cls, + store: Store | StorePath, + *, + name: str | None = None, + config: ArrayConfigLike | None = None, + **kwargs: Unpack[ArrayV2DeclarationKwargs], + ) -> Self: + """Declare a Zarr format 2 array bound to `store`, storing nothing. + + The keywords are the `.zarray` fields (see `build_array_metadata_v2`). + `store` must be a `Store` or `StorePath`: coercing strings opens + stores. `store_metadata` stores the result. + """ + return cls(build_array_metadata_v2(**kwargs), _bound_store_path(store, name), config) + + @classmethod + def build_v3( + cls, + store: Store | StorePath, + *, + name: str | None = None, + config: ArrayConfigLike | None = None, + **kwargs: Unpack[ArrayV3DeclarationKwargs], + ) -> Self: + """Declare a Zarr format 3 array bound to `store`, storing nothing. + + The keywords are the `zarr.json` fields (see + `build_array_metadata_v3`). `store` must be a `Store` or `StorePath`: + coercing strings opens stores. `store_metadata` stores the result. + """ + return cls(build_array_metadata_v3(**kwargs), _bound_store_path(store, name), config) + + @classmethod + async def open( + cls, + store: StoreLike, + *, + name: str | None = None, + zarr_format: ZarrFormat | None = None, + config: ArrayConfigLike | None = None, + ) -> Self: + """Read the stored metadata of an existing array and bind to it. + + `zarr_format=None` (the default) tries format 3, then format 2. The + read is strict and resolves nothing: metadata zarr can only read as + repaired raises `zarr_metadata.MetadataValidationError`; use + [`zarr.api.asynchronous.open_array`][] for the repairing read. + """ + # A path opens writable without creating anything when no array is there. + mode: AccessModeLiteral | None = "r+" if isinstance(store, str | os.PathLike) else None + store_path = await make_store_path(store, path=name, mode=mode) + data = await get_array_metadata(store_path, zarr_format=zarr_format) + binding = cls(_array_metadata_model(data), store_path, config) + observe_read(store_path, data) + return binding + + @property + def metadata(self) -> ArrayMetadataModel: + """The bound metadata, as a `zarr_metadata` model.""" + source = self._source + if not isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + resolved = self._resolved_cache + if resolved is None: + return source + # An `AsyncArray` over this binding hands out the resolved object, whose + # attributes callers change in place: the model reads them. Compared as + # JSON text, in which NaN equals itself, and kept until they change. + text = _json_text(resolved.attributes) + cached = self._attributes_cache + if cached is None or cached[0] != text: + same = text == _json_text(source.attributes) + attributes = _plain_json(resolved.attributes) + cached = (text, source if same else source.update(attributes=attributes)) + self._attributes_cache = cached + return cached[1] + # Derived on every access, not cached: a legacy object's attributes may be + # changed in place through another handle that shares it (as `AsyncArray` + # copies do); the chunk grid and key encoding such changes cannot touch. + return _model_of_legacy(source) + + @property + def metadata_legacy(self) -> ArrayMetadata: + """The metadata resolved into the legacy `ArrayV2Metadata` / `ArrayV3Metadata` classes, derived on first use.""" + source = self._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + return source + if self._resolved_cache is None: + data = _metadata_model_json(source) + if isinstance(source, _model_classes()[0]): + data.setdefault("attributes", {}) + resolved = parse_array_metadata(data, str(self._store_path)) + # Threads resolving at once all keep the first object stored: an + # `AsyncArray` changes its attributes in place, so they must share it. + # (A rebind of this binding racing the read is a write racing a read, + # which no binding orders.) + with _FIRST_RESOLUTION_LOCK: + if self._resolved_cache is None: + self._resolved_cache = resolved + return self._resolved_cache + + def _detached(self) -> Self: + """A copy of this binding that shares nothing that is changed in place: its + legacy metadata (whose attributes an `AsyncArray` updates in place) is copied, + and its caches of immutable values are shared.""" + copied = object.__new__(type(self)) + copied.__dict__.update(self.__dict__) + if isinstance(self._source, ArrayV2Metadata | ArrayV3Metadata): + copied._source = _copied_legacy(self._source) + elif self._resolved_cache is not None: + copied._resolved_cache = _copied_legacy(self._resolved_cache) + return copied + + def _rebound_like(self, built: AsyncArrayIO, metadata: ArrayMetadata) -> Self: + """How an array rebuilt through its constructor (`built`, which bound + `metadata`) keeps the binding it had: a binding of this binding's class, at + `built`'s location, with its configuration, pipeline class and caches, over + this binding's model when `metadata` is still this binding's resolved form + (shared), else over `metadata` itself (a subclass `__init__` substituted it). + A `_bind` that runs another pipeline class wins over any pipeline `built` + holds, also one a subclass `__init__` assigned.""" + rebound = object.__new__(type(self)) + source = self._source + keeps_model = ( + not isinstance(source, ArrayV2Metadata | ArrayV3Metadata) + and metadata is self._resolved_if_any + ) + rebound._bind( + source if keeps_model else metadata, + built._store_path, + built._config, + built._pipeline_class, + ) + if keeps_model: + rebound._resolved_cache = metadata + if built._pipeline_class is rebound._pipeline_class: + # The pipeline the constructor built is the one this binding runs; a + # `_bind` that chose another pipeline class builds its own. + rebound._codec_pipeline_cache = built._codec_pipeline_cache + rebound._chunk_grid_cache = built._chunk_grid_cache + return rebound + + @property + def store_path(self) -> StorePath: + """The bound location in a store.""" + return self._store_path + + @property + def config(self) -> ArrayConfig: + """The runtime configuration.""" + return self._config + + @property + def _zarr_format(self) -> ZarrFormat: + source = self._source + if isinstance(source, ArrayV2Metadata): + return 2 + if isinstance(source, ArrayV3Metadata): + return 3 + return 2 if isinstance(source, _model_classes()[0]) else 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_at(self, store_path: StorePath, config: ArrayConfig) -> None: + """Bind this binding again, through `_bind`, at `store_path` with `config`, + keeping its caches: setting an array's location or configuration never + rebuilt its pipeline.""" + caches = ( + self._resolved_cache, + self._chunk_grid_cache, + self._chunk_key_encoding_cache, + self._codec_pipeline_cache, + ) + self._bind(self._source, store_path, config, self._pipeline_class) + ( + self._resolved_cache, + self._chunk_grid_cache, + self._chunk_key_encoding_cache, + self._codec_pipeline_cache, + ) = caches + + def _rebind_resolved( + self, resolved: ArrayMetadata, chunk_grid: ChunkGrid | None = None + ) -> None: + """Make `resolved` this binding's state, through `_bind`. The codec pipeline is + kept: rebinding is how `resize` swaps the shape, which leaves the chunk + encoding as it was.""" + pipeline = self._codec_pipeline_cache + self._bind(resolved, self._store_path, self._config, self._pipeline_class) + self._codec_pipeline_cache = pipeline + self._chunk_grid_cache = chunk_grid + + @property + def _chunk_grid(self) -> ChunkGrid: + """The stored chunk grid: from the resolved metadata when the instance holds one, otherwise from the metadata model alone.""" + if self._chunk_grid_cache is None: + resolved = self._resolved_if_any + if resolved is not None: + self._chunk_grid_cache = ChunkGrid.from_metadata(resolved) + return self._chunk_grid_cache + document = self.metadata + doc = _metadata_model_json(document) + if isinstance(document, _model_classes()[0]): + 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.""" + 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, from the metadata model alone.""" + if self._chunk_key_encoding_cache is None: + document = self.metadata + doc = _metadata_model_json(document) + if isinstance(document, _model_classes()[0]): + 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_legacy, + store=self._store_path.store, + pipeline_class=self._pipeline_class, + ) + return self._codec_pipeline_cache + + def to_array(self) -> AnyAsyncArray: + """An `AsyncArray` over a copy of this binding: nothing done through the array + (a resize, an attribute update) changes this binding.""" + return AsyncArray._from_io(self) + + def __getstate__(self) -> dict[str, object]: + state = dict(self.__dict__) + # The pipeline is rebuilt on first use: a built one is bound to this + # process's store object, and may not pickle. The attributes' model is + # derived again. + state["_codec_pipeline_cache"] = None + state["_attributes_cache"] = None + # Under one key of its own, so no name a subclass's attribute may hold is taken. + state[_BINDING_STATE] = self._carried_cache_entry() + return state + + def __setstate__(self, state: dict[str, object]) -> None: + state = dict(state) + carried = cast("dict[str, object]", state.pop(_BINDING_STATE, {})) + self.__dict__.update(state) + # Through `_bind`, so state a subclass derives there is derived again. + self._rebind_at(self._store_path, self._config) + self._restore_cache_entry(carried) + + def _carried_cache_entry(self) -> dict[str, object]: + """What this binding's store cache knows of its document, to carry in a + pickle: a store's cache does not pickle with the store, and the unpickled + binding's first chunk write must know whether to store a repair first.""" + cache = stored_documents(self._store_path.store) + path, zarr_format = self._store_path.path, self._zarr_format + if cache.knows(path, zarr_format): + return {"_stored_entry": cache.needing_repair(path, zarr_format)} + copy_entry = cache.copy_needing_repair(path, zarr_format) + if copy_entry is not None: + return {"_stored_copy": copy_entry} + return {} + + def _restore_cache_entry(self, state: Mapping[str, object]) -> None: + """Record what `_carried_cache_entry` carried in this binding's store cache.""" + cache = stored_documents(self._store_path.store) + path, zarr_format = self._store_path.path, self._zarr_format + if "_stored_entry" in state: + entry = cast("ArrayDocument | None", state["_stored_entry"]) + cache.restore(path, zarr_format, entry) + elif "_stored_copy" in state: + cache.restore_copy(path, zarr_format, cast("ArrayDocument", state["_stored_copy"])) + + def _derived_source(self) -> ArrayMetadataModel | ArrayMetadata: + """The source for a derived binding: a legacy object is copied, so in-place + attribute updates never alias across bindings; a model is immutable and + shared.""" + source = self._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + return _copied_legacy(source) + return self.metadata + + def _derive( + self, + source: ArrayMetadataModel | ArrayMetadata, + store_path: StorePath, + config: ArrayConfig, + ) -> Self: + """A new binding over `source`, keeping this one's captured pipeline.""" + derived = object.__new__(type(self)) + derived._bind(source, store_path, config, self._pipeline_class) + return derived + + def with_metadata( + self, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, + ) -> Self: + """This binding with `metadata` in place of its own. Nothing is stored; + a legacy `metadata` object or a metadata dict is read as the + constructor reads it, and a legacy object is copied.""" + source: ArrayMetadataModel | ArrayMetadata + if isinstance(metadata, ArrayV2Metadata | ArrayV3Metadata): + source = _copied_legacy(parse_array_metadata(metadata, str(self._store_path))) + elif isinstance(metadata, _model_classes()): + source = _checked_model(metadata) + elif isinstance(metadata, Mapping): + source = _bound_metadata_dict(dict(metadata), self._store_path) + else: + raise TypeError( + "Expected array metadata (a zarr_metadata model, ArrayV2Metadata or " + f"ArrayV3Metadata) or a metadata mapping, got {type(metadata).__name__}." + ) + return self._derive(source, self._store_path, self._config) + + def with_store(self, store: Store | StorePath) -> Self: + """This binding at its path within `store`, or at a `StorePath` as given. Nothing is stored.""" + store_path = ( + store if isinstance(store, StorePath) else StorePath(store, self._store_path.path) + ) + return self._derive(self._derived_source(), store_path, self._config) + + def with_config(self, config: ArrayConfigLike) -> Self: + """This binding with `config` merged over its runtime configuration, as `AsyncArray.with_config` merges it.""" + if isinstance(config, ArrayConfig): + new_config = config + else: + new_config = ArrayConfig(**{**self._config.to_dict(), **config}) # type: ignore[arg-type] + return self._derive(self._derived_source(), self._store_path, new_config) + + async def store_metadata(self, *, ensure_parents: bool = False) -> None: + """Store the metadata documents of this array's Zarr format, unconditionally + replacing those stored here. + + Documents of the other format, or of a group, at this location are left + as they are, as `erase_metadata` leaves them; `verify` reports when a + format 3 document would be opened instead. Chunks of a previous array at + this location are kept; `erase_chunks` removes them. With + `ensure_parents`, missing parent groups are stored first; they are nodes + of their own, which `erase_metadata` leaves in place. + """ + source = self._source + if isinstance(source, ArrayV2Metadata | ArrayV3Metadata): + await save_metadata(self._store_path, source, ensure_parents=ensure_parents) + else: + source = self.metadata + prototype = default_buffer_prototype() + key_value = source.to_key_value(indent=zarr_config.get("json_indent")) + to_save = { + str(key): prototype.buffer.from_bytes(value) for key, value in key_value.items() + } + await _write_metadata( + self._store_path, + to_save, + self._zarr_format, + ensure_parents=ensure_parents, + written=source, + ) + if ( + self._zarr_format == 2 + and ZATTRS_JSON not in to_save + and not await (self._store_path / ZGROUP_JSON).exists() + ): + # An attribute-less model emits no .zattrs; a stale one must not + # survive an unconditional metadata write. A group's is its own. + await (self._store_path / ZATTRS_JSON).delete() + + async def verify(self) -> None: + """Check that the stored metadata equals the bound metadata. + + Favors format 3 when both formats exist, compares canonically, and + raises `ValueError` naming each differing field, or + [`zarr.errors.ArrayNotFoundError`][] when nothing is stored. + """ + bound = self.metadata + stored = _array_metadata_model(await get_array_metadata(self._store_path, zarr_format=None)) + if stored == bound: + return + if type(stored) is not type(bound): + stored_format = 2 if isinstance(stored, _model_classes()[0]) else 3 + raise ValueError( + f"Array {str(self.store_path)!r}: the stored array is Zarr format " + f"{stored_format}; the bound metadata is Zarr format {self._zarr_format}." + ) + stored_json = _metadata_model_json(stored) + bound_json = _metadata_model_json(bound) + if isinstance(stored, _model_classes()[0]) and isinstance(bound, _model_classes()[0]): + stored_json["attributes"] = _plain_json(stored.attributes) + bound_json["attributes"] = _plain_json(bound.attributes) + differing = _differing_fields(stored, stored_json, bound_json) + raise ValueError( + f"Array {str(self.store_path)!r}: the stored metadata does not match the " + f"bound metadata; differing fields: {', '.join(differing)}." + ) + + async def erase_metadata(self) -> None: + """Delete the stored array metadata of this array's format; chunks are + kept. Succeeds if absent; a group's documents here are left in place.""" + zarr_format = self._zarr_format + if zarr_format == 2: + await (self._store_path / ZARRAY_JSON).delete() + if not await (self._store_path / ZGROUP_JSON).exists(): + await (self._store_path / ZATTRS_JSON).delete() + else: + document = await (self._store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype) + if document is not None and not _is_group_document(document): + await (self._store_path / ZARR_JSON).delete() + stored_documents(self._store_path.store).observe(self._store_path.path, None, zarr_format) + + async def erase_chunk(self, chunk_coords: Sequence[int], /) -> None: + """Delete one stored chunk by grid coordinates (the shard, under sharding). Succeeds if absent.""" + chunk_coords = tuple(operator.index(coord) for coord in chunk_coords) + grid_shape = self._shard_grid_shape + if len(chunk_coords) != len(grid_shape) or any( + coord < 0 or coord >= size for coord, size in zip(chunk_coords, grid_shape, strict=True) + ): + raise IndexError( + f"Chunk coordinates {chunk_coords} are outside the chunk grid " + f"of shape {grid_shape}." + ) + await _delete_key(self._store_path, self._encode_chunk_key(chunk_coords)) + + async def erase_chunks(self) -> None: + """Delete every stored chunk; the metadata is kept.""" + grid_shape = self._shard_grid_shape + n_chunks = product(grid_shape) + if n_chunks == 0: + return + concurrency = zarr_config.get("async.concurrency") + limit = concurrency if concurrency is not None else 1000 + await concurrent_foreach( + ( + (self._store_path, self._encode_chunk_key(coords)) + 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_legacy, + self._resolved_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: + metadata = self.metadata_legacy + if product(indexer.shape) > 0: + # Chunks are about to be stored under the repaired metadata, so store it + # first: every reader of the store then agrees with them. + await _store_repaired_document(self._store_path, metadata) + return await _set_selection( + self._store_path, + metadata, + self._resolved_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 selection, reading absent chunks as the fill value.""" + if prototype is None: + prototype = default_buffer_prototype() + indexer = BasicIndexer( + selection, shape=self.metadata_legacy.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_legacy.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: + # 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 + if self._store_path != other._store_path or self._config != other._config: + return False + try: + return self.metadata == other.metadata + except TypeError: + # Attributes changed in place to values JSON cannot hold: compare the + # documents without them, then the resolved attributes, which hold them. + return ( + _document_without_attributes(self) == _document_without_attributes(other) + and self.metadata_legacy.attributes == other.metadata_legacy.attributes + ) + + def __hash__(self) -> int: + # A subset of the equality fields, available without deriving either + # form: equal bindings hash equally. + return hash((self._source.shape, self._store_path.path)) + + +class ArrayIO: + """Explicit synchronous IO for one Zarr array. See `AsyncArrayIO`.""" + + _async_io: AsyncArrayIO + + def __init__( + self, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, + store_path: StorePath, + config: ArrayConfigLike | None = None, + ) -> None: + self._async_io = AsyncArrayIO(metadata, store_path, config) + + @classmethod + def _wrap(cls, async_io: AsyncArrayIO) -> Self: + io = object.__new__(cls) + io._async_io = async_io + return io + + @classmethod + def build_v2( + cls, + store: Store | StorePath, + *, + name: str | None = None, + config: ArrayConfigLike | None = None, + **kwargs: Unpack[ArrayV2DeclarationKwargs], + ) -> Self: + """Declare a Zarr format 2 array bound to `store`, storing nothing. See `AsyncArrayIO.build_v2`.""" + return cls._wrap(AsyncArrayIO.build_v2(store, name=name, config=config, **kwargs)) + + @classmethod + def build_v3( + cls, + store: Store | StorePath, + *, + name: str | None = None, + config: ArrayConfigLike | None = None, + **kwargs: Unpack[ArrayV3DeclarationKwargs], + ) -> Self: + """Declare a Zarr format 3 array bound to `store`, storing nothing. See `AsyncArrayIO.build_v3`.""" + return cls._wrap(AsyncArrayIO.build_v3(store, name=name, config=config, **kwargs)) + + @classmethod + def open( + cls, + store: StoreLike, + *, + name: str | None = None, + zarr_format: ZarrFormat | None = None, + config: ArrayConfigLike | None = None, + ) -> Self: + """Read the stored metadata of an existing array and bind to it. See `AsyncArrayIO.open`.""" + return cls._wrap( + sync(AsyncArrayIO.open(store, name=name, zarr_format=zarr_format, config=config)) + ) + + @property + def metadata(self) -> ArrayMetadataModel: + """The bound metadata, as a `zarr_metadata` model.""" + return self._async_io.metadata + + @property + def metadata_legacy(self) -> ArrayMetadata: + """The metadata resolved into the legacy classes. See `AsyncArrayIO.metadata_legacy`.""" + return self._async_io.metadata_legacy + + @property + def store_path(self) -> StorePath: + """The bound location in a store.""" + return self._async_io.store_path + + @property + def config(self) -> ArrayConfig: + """The runtime configuration.""" + return self._async_io.config + + def to_array(self) -> AnyArray: + """An `Array` over a copy of this binding: nothing done through the array (a + resize, an attribute update) changes this binding.""" + return Array(self._async_io.to_array()) + + def with_metadata( + self, + metadata: ArrayMetadataModel | ArrayMetadata | ArrayMetadataDict | ArrayMetadataModelJSON, + ) -> Self: + """This binding with `metadata` in place of its own. See `AsyncArrayIO.with_metadata`.""" + return self._wrap(self._async_io.with_metadata(metadata)) + + def with_store(self, store: Store | StorePath) -> Self: + """This binding in `store`. See `AsyncArrayIO.with_store`.""" + return self._wrap(self._async_io.with_store(store)) + + def with_config(self, config: ArrayConfigLike) -> Self: + """This binding with `config` merged over its own. See `AsyncArrayIO.with_config`.""" + return self._wrap(self._async_io.with_config(config)) + + def store_metadata(self, *, ensure_parents: bool = False) -> None: + """Store the metadata. See `AsyncArrayIO.store_metadata`.""" + sync(self._async_io.store_metadata(ensure_parents=ensure_parents)) + + def verify(self) -> None: + """Check that the stored metadata equals the bound metadata. See `AsyncArrayIO.verify`.""" + sync(self._async_io.verify()) + + def erase_metadata(self) -> None: + """Delete the stored metadata; chunks are kept. See `AsyncArrayIO.erase_metadata`.""" + sync(self._async_io.erase_metadata()) + + def erase_chunk(self, chunk_coords: Sequence[int], /) -> None: + """Delete one stored chunk by grid coordinates. See `AsyncArrayIO.erase_chunk`.""" + sync(self._async_io.erase_chunk(chunk_coords)) + + def erase_chunks(self) -> None: + """Delete every stored chunk. See `AsyncArrayIO.erase_chunks`.""" + sync(self._async_io.erase_chunks()) + + def getitem( + self, + selection: BasicSelection, + *, + prototype: BufferPrototype | None = None, + ) -> NDArrayLikeOrScalar: + """Read the selection, reading absent chunks as the fill value. See `AsyncArrayIO.getitem`.""" + 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. See `AsyncArrayIO.setitem`.""" + sync(self._async_io.setitem(selection, value, prototype=prototype)) + + def __repr__(self) -> str: + 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(self._async_io) diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index f1b256c491..f674eb6ac9 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import copy import logging import unicodedata import warnings @@ -13,7 +14,9 @@ import zarr.api.asynchronous as async_api from zarr.abc.metadata import Metadata -from zarr.abc.store import Store, set_or_delete + +# Importable from here, as they have been, though this module no longer uses them. +from zarr.abc.store import Store, set_or_delete # noqa: F401, TC001 from zarr.core._info import GroupInfo from zarr.core._json import buffer_to_json_object, json_to_buffer from zarr.core.array import ( @@ -55,6 +58,7 @@ save_new_metadata, store_documents, ) +from zarr.core.metadata.stored import stored_documents from zarr.core.metadata.v3 import AllowedExtraField, check_storable, parse_extra_fields from zarr.core.sync import SyncMixin, sync from zarr.errors import ( @@ -87,6 +91,8 @@ from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike from zarr.core.common import MemoryOrder from zarr.core.dtype import ZDTypeLike + from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.stored import CarriedEntry from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 logger = logging.getLogger("zarr.group") @@ -140,6 +146,28 @@ def _parse_async_node( raise TypeError(f"Unknown node type, got {type(node)}") +def _with_member_marks( + consolidated: ConsolidatedMetadata, prefix: str, marks: Mapping[str, ArrayDocument] +) -> ConsolidatedMetadata: + """A copy of `consolidated` whose array members at the paths in `marks` (relative + to the group, below `prefix`) are copies carrying their mark as the + `_stored_document` attribute zarr 3.4.1 and 3.4.2 read.""" + members: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata] = {} + for name, member in consolidated.metadata.items(): + key = f"{prefix}{name}" + if isinstance(member, GroupMetadata): + nested = member.consolidated_metadata + if nested is not None: + member = replace( + member, consolidated_metadata=_with_member_marks(nested, f"{key}/", marks) + ) + elif key in marks: + member = copy.copy(member) + object.__setattr__(member, "_stored_document", copy.deepcopy(marks[key])) + members[name] = member + return replace(consolidated, metadata=members) + + @dataclass(frozen=True) class ConsolidatedMetadata: """ @@ -153,17 +181,19 @@ class ConsolidatedMetadata: kind: Literal["inline"] = "inline" must_understand: Literal[False] = False - def to_dict(self) -> dict[str, JSON]: - """The consolidated metadata document. An array read from a stored document that - had to be repaired is written as that document was stored, so every reader of - the consolidated metadata reads it as repaired again (see `mark_repaired`).""" + def to_dict( + self, *, members_as_stored: Mapping[str, ArrayDocument] | None = None + ) -> dict[str, JSON]: + """The consolidated metadata document. A member in `members_as_stored` (by its + path here) is written as that stored document, so every reader of the + consolidated metadata reads it as repaired again (see + `zarr.core.metadata.io.members_as_stored`).""" + as_stored = {} if members_as_stored is None else members_as_stored return { "kind": self.kind, "must_understand": self.must_understand, "metadata": { - k: v.to_dict() - if isinstance(v, GroupMetadata) or v._stored_document is None - else dict(v._stored_document) + k: dict(as_stored[k]) if k in as_stored else v.to_dict() for k, v in sorted( self.flattened_metadata.items(), key=lambda item: ( @@ -384,12 +414,19 @@ class GroupMetadata(Metadata): node_type: Literal["group"] = field(default="group", init=False) extra_fields: dict[str, AllowedExtraField] = field(default_factory=dict) - def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: + def to_buffer_dict( + self, + prototype: BufferPrototype, + *, + members_as_stored: Mapping[str, ArrayDocument] | None = None, + ) -> dict[str, Buffer]: + """The group's metadata documents, by key. A consolidated member in + `members_as_stored` is written as that stored document (see + `ConsolidatedMetadata.to_dict`).""" + as_stored = {} if members_as_stored is None else members_as_stored if self.consolidated_metadata is not None: for path, member in self.consolidated_metadata.flattened_metadata.items(): - # A member read from a document that had to be repaired is stored as it - # was stored (see `ConsolidatedMetadata.to_dict`). - if isinstance(member, ArrayV3Metadata) and member._stored_document is None: + if isinstance(member, ArrayV3Metadata) and path not in as_stored: try: check_storable(member) except ValueError as e: @@ -397,7 +434,13 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: raise indent = config.get("json_indent") if self.zarr_format == 3: - return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)} + document = self.to_dict() + if self.consolidated_metadata is not None and as_stored: + # Encoded again only where a member is kept as stored. + document["consolidated_metadata"] = self.consolidated_metadata.to_dict( + members_as_stored=as_stored + ) + return {ZARR_JSON: json_to_buffer(document, prototype=prototype, indent=indent)} else: items = { ZGROUP_JSON: json_to_buffer( @@ -410,7 +453,9 @@ def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: ZGROUP_JSON: {"zarr_format": self.zarr_format}, ZATTRS_JSON: self.attributes, } - consolidated_metadata = self.consolidated_metadata.to_dict()["metadata"] + consolidated_metadata = self.consolidated_metadata.to_dict( + members_as_stored=as_stored + )["metadata"] if not isinstance(consolidated_metadata, dict): raise TypeError( "Expected consolidated metadata to serialize to a dict, " @@ -670,6 +715,71 @@ async def open( use_consolidated=use_consolidated, ) + def __getstate__(self) -> dict[str, object]: + # A store's cache does not pickle with the store: carry what it knows of the + # consolidated members, so a write through the unpickled group keeps the + # members it would have kept as stored. + state = dict(self.__dict__) + consolidated = self.metadata.consolidated_metadata + if consolidated is not None: + paths = [ + _join_paths([self.store_path.path, key]) for key in consolidated.flattened_metadata + ] + entries = stored_documents(self.store_path.store).entries( + paths, self.metadata.zarr_format + ) + if entries: + state["_stored_members"] = entries + # zarr 3.4.1 and 3.4.2 read a member's repair from a mark on the + # member itself: pickle marked copies for them to read. + prefix = f"{self.store_path.path}/" if self.store_path.path else "" + marks = { + path.removeprefix(prefix): entry[1] + for path, entry in entries.items() + if entry[1] is not None + } + if marks: + state["metadata"] = replace( + self.metadata, + consolidated_metadata=_with_member_marks(consolidated, "", marks), + ) + return state + + def __setstate__(self, state: dict[str, object]) -> None: + state = dict(state) + entries = cast("dict[str, CarriedEntry] | None", state.pop("_stored_members", None)) + self.__dict__.update(state) + if entries is None: + entries = self._legacy_member_marks() + self._clear_member_marks() + if entries: + stored_documents(self.store_path.store).restore_entries( + entries, self.metadata.zarr_format + ) + + def _legacy_member_marks(self) -> dict[str, CarriedEntry]: + """The consolidated members a pickle from zarr 3.4.1 or 3.4.2 marks as read + from a document that needs a repair, as the entries this version carries.""" + consolidated = self.metadata.consolidated_metadata + if consolidated is None: + return {} + marks: dict[str, CarriedEntry] = {} + for key, member in consolidated.flattened_metadata.items(): + mark = vars(member).get("_stored_document") + if mark is not None: + marks[_join_paths([self.store_path.path, key])] = ("copy", mark) + return marks + + def _clear_member_marks(self) -> None: + """Drop the marks a pickle leaves on consolidated members: once read, the + store's cache holds them.""" + consolidated = self.metadata.consolidated_metadata + if consolidated is None: + return + for member in consolidated.flattened_metadata.values(): + if "_stored_document" in vars(member): + object.__delattr__(member, "_stored_document") + @classmethod def _from_bytes_v2( cls, @@ -744,10 +854,19 @@ def from_dict( elif node_type not in ("group", None): msg = f"Node type in metadata ({node_type}) is not 'group'" raise GroupNotFoundError(msg) - return cls( + group = cls( metadata=GroupMetadata.from_dict(data, path=str(store_path)), store_path=store_path, ) + consolidated = group.metadata.consolidated_metadata + raw_consolidated = data.get("consolidated_metadata") + if consolidated is not None and isinstance(raw_consolidated, dict): + members = raw_consolidated.get("metadata") + if isinstance(members, dict): + stored_documents(store_path.store).observe_consolidated( + store_path.path, members, group.metadata.zarr_format + ) + return group async def setitem(self, key: str, value: Any) -> None: """ @@ -2251,9 +2370,7 @@ async def update_attributes_async(self, new_attributes: dict[str, Any]) -> Group new_metadata = replace(self.metadata, attributes=new_attributes) # Write new metadata - to_save = new_metadata.to_buffer_dict(default_buffer_prototype()) - awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()] - await asyncio.gather(*awaitables) + await store_documents(self.store_path, encode_documents(self.store_path, new_metadata)) async_group = replace(self._async_group, metadata=new_metadata) return replace(self, _async_group=async_group) @@ -3422,6 +3539,11 @@ async def _store_nodes( ] created_object_keys = [] + node_documents: defaultdict[str, dict[str, Buffer]] = defaultdict(dict) + for key, value in documents.items(): + node, _, name = key.rpartition("/") + node_documents[node][name] = value + cache = stored_documents(store) for coro in asyncio.as_completed(create_tasks): created_key = await coro @@ -3439,6 +3561,8 @@ async def _store_nodes( node_name = created_key[: created_key.rfind("/")] meta_out = nodes[node_name] if meta_out.zarr_format == 3: + if not isinstance(meta_out, GroupMetadata): + cache.observe_written(node_name, node_documents[node_name], meta_out) yield node_name, built[node_name] else: # For zarr v2 @@ -3454,6 +3578,8 @@ async def _store_nodes( meta_done = _join_paths([node_name, ZARRAY_JSON]) in created_object_keys if meta_done and attrs_done: + if not isinstance(meta_out, GroupMetadata): + cache.observe_written(node_name, node_documents[node_name], meta_out) yield node_name, built[node_name] continue @@ -3709,9 +3835,11 @@ async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupM ) if zarr_json_bytes is None: raise FileNotFoundError(path) - return _build_metadata_v3( - buffer_to_json_object(zarr_json_bytes), path=str(StorePath(store, path)) - ) + document = buffer_to_json_object(zarr_json_bytes) + metadata = _build_metadata_v3(document, path=str(StorePath(store, path))) + if isinstance(metadata, ArrayV3Metadata): + stored_documents(store).observe(path, document, 3) + return metadata async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupMetadata: @@ -3746,7 +3874,10 @@ async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupM else: zmeta = buffer_to_json_object(zgroup_bytes) - return _build_metadata_v2(zmeta, zattrs, path=str(StorePath(store, path))) + metadata = _build_metadata_v2(zmeta, zattrs, path=str(StorePath(store, path))) + if isinstance(metadata, ArrayV2Metadata): + stored_documents(store).observe(path, {**zmeta, "attributes": zattrs}, 2) + return metadata async def _read_group_metadata_v2(store: Store, path: str) -> GroupMetadata: diff --git a/src/zarr/core/metadata/io.py b/src/zarr/core/metadata/io.py index c1cd927b94..024c9a57bb 100644 --- a/src/zarr/core/metadata/io.py +++ b/src/zarr/core/metadata/io.py @@ -3,14 +3,15 @@ import asyncio from enum import Enum from itertools import zip_longest -from typing import TYPE_CHECKING, Final, NamedTuple +from typing import TYPE_CHECKING, Final, NamedTuple, cast from zarr.abc.store import set_or_delete from zarr.core._json import buffer_to_json_object, json_equal from zarr.core.buffer.core import default_buffer_prototype from zarr.core.buffer.cpu import buffer_prototype as cpu_buffer_prototype from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON -from zarr.core.metadata.repair import mark_repaired, repair_array_document +from zarr.core.metadata.repair import repair_array_document +from zarr.core.metadata.stored import stored_documents from zarr.errors import ArrayNotFoundError, ContainsArrayError from zarr.storage._common import StorePath, ensure_no_existing_node @@ -21,6 +22,8 @@ from zarr.core.common import JSON, ZarrFormat from zarr.core.group import GroupMetadata from zarr.core.metadata import ArrayMetadata + from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.stored import WrittenArrayMetadata class _Absent(Enum): @@ -83,6 +86,11 @@ def encode_documents( from zarr.core.group import GroupMetadata try: + if isinstance(metadata, GroupMetadata): + return metadata.to_buffer_dict( + default_buffer_prototype(), + members_as_stored=members_as_stored(store_path, metadata), + ) return metadata.to_buffer_dict(default_buffer_prototype()) except ValueError as e: node = "Group" if isinstance(metadata, GroupMetadata) else "Array" @@ -90,11 +98,82 @@ def encode_documents( raise -async def store_documents(store_path: StorePath, documents: Mapping[str, Buffer]) -> None: - """Store metadata documents encoded by `encode_documents` under `store_path`.""" +async def store_documents( + store_path: StorePath, + documents: Mapping[str, Buffer], + *, + written: WrittenArrayMetadata | None = None, +) -> None: + """Store metadata documents encoded by `encode_documents` under `store_path`; those + of the array metadata `written` (none: a group's) are recorded in the store's cache + (see `stored_documents`).""" await asyncio.gather( *(set_or_delete(store_path / key, value) for key, value in documents.items()) ) + if written is not None: + stored_documents(store_path.store).observe_written(store_path.path, documents, written) + + +def members_as_stored(store_path: StorePath, metadata: GroupMetadata) -> dict[str, ArrayDocument]: + """The members of the consolidated metadata of the group at `store_path` to store as + they are stored, by their paths in it: each array whose stored document, or else + whose stored consolidated copy, needs a repair that moves chunks and reads as the + member does. A group write then leaves those copies as they were stored, so every + reader of the consolidated metadata reads them as repaired again. An array whose + own document zarr knows needs no such repair is stored as it reads.""" + from zarr.core.group import GroupMetadata + from zarr.storage._utils import _join_paths + + consolidated = metadata.consolidated_metadata + if consolidated is None: + return {} + cache = stored_documents(store_path.store) + as_stored: dict[str, ArrayDocument] = {} + for key, member in consolidated.flattened_metadata.items(): + if isinstance(member, GroupMetadata): + continue + path = _join_paths([store_path.path, key]) + zarr_format = member.zarr_format + own = cache.needing_repair(path, zarr_format) + if cache.knows(path, zarr_format) and own is None: + continue + for document in (own, cache.copy_needing_repair(path, zarr_format)): + if document is not None and reads_as(document, member): + as_stored[key] = document + break + return as_stored + + +def reads_as(document: ArrayDocument, metadata: ArrayMetadata) -> bool: + """Whether the stored array document `document`, repaired, reads as `metadata`.""" + return read_repaired(document, metadata.zarr_format).to_dict() == metadata.to_dict() + + +def read_repaired(document: ArrayDocument, zarr_format: ZarrFormat) -> ArrayMetadata: + """The metadata the stored array document `document` reads as, without warnings.""" + from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata + + reader = ArrayV2Metadata if zarr_format == 2 else ArrayV3Metadata + metadata, _ = reader._read_stored(dict(document)) + return metadata + + +def stored_array_document( + documents: Mapping[str, Buffer], zarr_format: ZarrFormat +) -> ArrayDocument | None: + """The array document of `zarr_format` among `documents` (by store key, as + `read_documents` reads them), or `None` if there is none.""" + if zarr_format == 2: + if ZARRAY_JSON not in documents: + return None + document: dict[str, JSON] = buffer_to_json_object(documents[ZARRAY_JSON]) + zattrs = documents.get(ZATTRS_JSON) + document["attributes"] = {} if zattrs is None else buffer_to_json_object(zattrs) + return document + if ZARR_JSON not in documents: + return None + document = buffer_to_json_object(documents[ZARR_JSON]) + return document if document.get("node_type") == "array" else None async def read_documents(store_path: StorePath, keys: Iterable[str]) -> dict[str, Buffer]: @@ -123,7 +202,11 @@ async def upsert_metadata( {key: buffer_to_json_object(buf) for key, buf in documents.items()}, ) changed = {change.path[0] for change in changes} - await store_documents(store_path, {k: v for k, v in documents.items() if k in changed}) + await store_documents( + store_path, + {k: v for k, v in documents.items() if k in changed}, + written=_written_array(metadata), + ) return changes @@ -139,8 +222,8 @@ def parse_stored_array( ) -> ArrayMetadata: """The metadata of an array from its documents (by store key, see `ARRAY_DOCUMENTS`), read with the repairs but without their warnings (whoever asks has warned, or reads - metadata built in code), and marked (see `mark_repaired`) if they had to be - repaired. Raises `ArrayNotFoundError` if there is no array document among them. + metadata built in code). Raises `ArrayNotFoundError` if there is no array document + among them. Only operations that store metadata read documents this way, so documents read as a rectilinear chunk grid require the rectilinear chunks flag, as storing it does; the @@ -157,9 +240,8 @@ def parse_stored_array( stored = _array_metadata_dict_v3(documents[ZARR_JSON]) else: raise ArrayNotFoundError(f"No Zarr format {zarr_format} array metadata document.") - repaired, readings = repair_array_document(stored, zarr_format) - metadata = parse_array_metadata(dict(repaired), path) - return mark_repaired(metadata, stored, readings, None, warn=False) + repaired, _ = repair_array_document(stored, zarr_format) + return parse_array_metadata(dict(repaired), path) def _build_parents(store_path: StorePath, zarr_format: ZarrFormat) -> dict[str, GroupMetadata]: @@ -200,7 +282,13 @@ async def save_metadata( ValueError """ to_save = encode_documents(store_path, metadata) - await _write_metadata(store_path, to_save, metadata.zarr_format, ensure_parents=ensure_parents) + await _write_metadata( + store_path, + to_save, + metadata.zarr_format, + ensure_parents=ensure_parents, + written=_written_array(metadata), + ) async def save_new_metadata( @@ -233,7 +321,24 @@ async def save_new_metadata( await store_path.delete_dir() else: await ensure_no_existing_node(store_path, zarr_format=metadata.zarr_format) - await _write_metadata(store_path, to_save, metadata.zarr_format, ensure_parents=ensure_parents) + await _write_metadata( + store_path, + to_save, + metadata.zarr_format, + ensure_parents=ensure_parents, + written=_written_array(metadata), + ) + + +def _written_array(metadata: ArrayMetadata | GroupMetadata) -> ArrayMetadata | None: + """`metadata` as the array metadata a write records, or `None` for a group.""" + return None if _is_group(metadata) else cast("ArrayMetadata", metadata) + + +def _is_group(metadata: ArrayMetadata | GroupMetadata) -> bool: + from zarr.core.group import GroupMetadata + + return isinstance(metadata, GroupMetadata) async def _write_metadata( @@ -242,6 +347,7 @@ async def _write_metadata( zarr_format: ZarrFormat, *, ensure_parents: bool, + written: WrittenArrayMetadata | None, ) -> None: set_awaitables = [set_or_delete(store_path / key, value) for key, value in to_save.items()] @@ -281,3 +387,5 @@ async def _write_metadata( ) from e await asyncio.gather(*set_awaitables) + if written is not None: + stored_documents(store_path.store).observe_written(store_path.path, to_save, written) diff --git a/src/zarr/core/metadata/repair.py b/src/zarr/core/metadata/repair.py index d8050b9f3d..682a3eaa90 100644 --- a/src/zarr/core/metadata/repair.py +++ b/src/zarr/core/metadata/repair.py @@ -9,18 +9,18 @@ read from the same document before these repairs existed is silent, so a document that opened without a warning still does. The warnings are given once per document, after the repaired document has passed the metadata constructor, so an invalid -document raises its own error, not a warning about how it was read. Metadata read -from a document whose repair moves chunks (a chunk size read as another size) is -marked (see `mark_repaired`), silent or not, so the array stores the repair before it -writes chunks under it. A repair that only respells a value zarr already read the -same way (`true` as 1, `4.0` as 4) moves no chunks, so it is not marked. +document raises its own error, not a warning about how it was read. A store whose +document needs a repair that moves chunks (a chunk size read as another size) gets +that repair stored before zarr writes chunks under it, silent or not; whether a +stored document needs one is recorded per store (see `zarr.core.metadata.stored`), +not on the metadata read from it. A repair that only respells a value zarr already +read the same way (`true` as 1, `4.0` as 4) moves no chunks. To read another kind of invalid document, add a repair to `ARRAY_REPAIRS`. """ from __future__ import annotations -import copy import json import warnings from collections.abc import Callable, Iterable, Mapping, Sequence @@ -71,30 +71,16 @@ class Reading(NamedTuple): and a re-save would keep it.""" -def mark_repaired[M]( - metadata: M, - stored: ArrayDocument, - readings: Sequence[Reading], - path: str | None, - *, - warn: bool = True, -) -> M: - """Record that `metadata` was read from the document `stored`, which needed the - repairs whose `readings` `repair_array_document` returned, if any. If a reading - moves chunks, keep a copy of `stored` as `_stored_document` on `metadata`, so the - array stores the repair before it writes chunks under it and consolidated metadata - stores it as it was stored. If `warn`, warn once with the readings' warnings (each - says what the user must act on and how), naming the array at `path` when the caller - knows it.""" - if any(reading.moves_chunks for reading in readings): - object.__setattr__(metadata, "_stored_document", copy.deepcopy(stored)) +def warn_repairs(readings: Sequence[Reading], path: str | None) -> None: + """Warn once with the warnings of the `readings` `repair_array_document` returned + (each says what the user must act on and how), naming the array at `path` when + the caller knows it.""" messages = [reading.warning for reading in readings if reading.warning is not None] - if warn and messages: + if messages: subject = "" if path is None else f"Array {path!r}: " # The synchronous API parses metadata on zarr's IO thread, whose stack holds no # user code, so the warning points at the `from_dict` that read the document. warnings.warn(f"{subject}{' '.join(messages)}", ZarrUserWarning, stacklevel=2) - return metadata def _is_int_list(value: object) -> TypeGuard[list[int]]: @@ -369,13 +355,36 @@ def _invalid_edge_lengths_v3(doc: ArrayDocument) -> tuple[ArrayDocument, Reading """The repairs of an array document of each Zarr format, applied in order.""" +def may_move_chunks(doc: ArrayDocument, zarr_format: ZarrFormat) -> bool: + """Whether some repair might move chunks in `doc`: `False` for a rectilinear + grid, and when every stored regular chunk size is a JSON int >= 1 or `true`, + for only a stored 0 (or `false`) or an edge list in a regular grid is read so + that it moves chunks (see `_read_chunk_size` and `_invalid_chunk_sizes_v3`). + A cheap check that answers for most documents before `repair_array_document` + runs.""" + if zarr_format == 2: + sizes = doc.get("chunks") + else: + grid = doc.get("chunk_grid") + if isinstance(grid, Mapping) and grid.get("name") == "rectilinear": + # Its only repair (`_invalid_edge_lengths_v3`) respells edges in place. + return False + if not isinstance(grid, Mapping) or grid.get("name") != "regular": + return True + configuration = grid.get("configuration") + sizes = configuration.get("chunk_shape") if isinstance(configuration, Mapping) else None + if not isinstance(sizes, list | tuple): + return True + return not all(size is True or (type(size) is int and size >= 1) for size in sizes) + + def repair_array_document( doc: ArrayDocument, zarr_format: ZarrFormat ) -> tuple[ArrayDocument, list[Reading]]: """Apply the repairs for `zarr_format` to a stored array metadata document. Returns the repaired document and the reading of each repair that changed it, for - `mark_repaired` once the document has been validated. + `warn_repairs` once the document has been validated. """ readings: list[Reading] = [] for repair in ARRAY_REPAIRS[zarr_format]: diff --git a/src/zarr/core/metadata/stored.py b/src/zarr/core/metadata/stored.py new file mode 100644 index 0000000000..d71f660a01 --- /dev/null +++ b/src/zarr/core/metadata/stored.py @@ -0,0 +1,219 @@ +"""What zarr knows of the array metadata documents a store holds: a cache over the store. + +Whether an array's stored document needs a repair that moves chunks (see +`zarr.core.metadata.repair`) is a fact about the store, not about metadata read from +it. zarr records it here, per store object, from the array documents it reads and +stores, and asks here before it writes chunks. Every binding of one store object +shares what any of them learned, and metadata stays a value that copies and +re-resolves without carrying store state. + +The cache observes; it never serves a read. A path it knows nothing about is read +from the store when the answer is needed. +""" + +from __future__ import annotations + +import copy +import threading +import weakref +from functools import partial +from typing import TYPE_CHECKING, Final, Literal, cast + +from zarr.core._json import buffer_to_json_object +from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON +from zarr.core.metadata.repair import may_move_chunks, repair_array_document + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping + + from zarr_metadata import ZarrV2ArrayMetadata, ZarrV3ArrayMetadata + + from zarr.abc.store import Store + from zarr.core.buffer import Buffer + from zarr.core.common import JSON, ZarrFormat + from zarr.core.metadata import ArrayMetadata + from zarr.core.metadata.repair import ArrayDocument + + type WrittenArrayMetadata = ArrayMetadata | ZarrV2ArrayMetadata | ZarrV3ArrayMetadata + + +type CarriedEntry = ( + tuple[Literal["document"], ArrayDocument | None] | tuple[Literal["copy"], ArrayDocument] +) +"""What a cache records of one array document, carried in a pickle: the document's +entry, or else a consolidated copy's.""" + + +def moves_chunks(document: ArrayDocument, zarr_format: ZarrFormat) -> bool: + """Whether reading `document` needs a repair that moves chunks.""" + if not may_move_chunks(document, zarr_format): + return False + _, readings = repair_array_document(document, zarr_format) + return any(reading.moves_chunks for reading in readings) + + +class StoredDocuments: + """The array metadata documents zarr has read from or stored to one store. + + Each array document zarr has seen, by its path and Zarr format (a path may hold a + document of each), maps to the document as it was stored when reading it needs a + repair that moves chunks. A document that needs none (or its absence) is recorded, + as `None`, only where it replaces such a document or a consolidated copy that + needs one; otherwise it is unknown, which means the same, so the cache holds only + what stands for a repair. Copies of array documents in + consolidated metadata are recorded apart, by the array's path and format, only + when reading them needs such a repair, so a group write can store them as they + were stored. A document zarr reads to open an array is recorded only once it has + been read as metadata, so one it could not open is unknown; one a chunk write reads + is recorded as read. + """ + + def __init__(self) -> None: + self._arrays: dict[tuple[str, ZarrFormat], ArrayDocument | None] = {} + self._copies: dict[tuple[str, ZarrFormat], ArrayDocument] = {} + + def knows(self, path: str, zarr_format: ZarrFormat) -> bool: + """Whether this cache records the array document of `zarr_format` at `path`: + one that needs a repair that moves chunks, or a valid one that replaced such + a document or stands against a consolidated copy that needs one. A valid + document read with nothing to replace is not recorded, so whether it is + known depends on what was read first; either way it means no repair.""" + return (path, zarr_format) in self._arrays + + def needing_repair(self, path: str, zarr_format: ZarrFormat) -> ArrayDocument | None: + """The document of `zarr_format` stored at `path` as it was stored, if reading it + needs a repair that moves chunks; `None` if it needs none, there is none, or it + is unknown.""" + return self._arrays.get((path, zarr_format)) + + def copy_needing_repair(self, path: str, zarr_format: ZarrFormat) -> ArrayDocument | None: + """The consolidated copy of the array document of `zarr_format` at `path` as it + was stored, if reading it needs a repair that moves chunks.""" + return self._copies.get((path, zarr_format)) + + def restore(self, path: str, zarr_format: ZarrFormat, document: ArrayDocument | None) -> None: + """Record `document` as `needing_repair(path, zarr_format)` returned it from + another store object's cache (a binding carries it when pickled), unless this + cache already knows the document: it keeps what it learned itself.""" + self._arrays.setdefault((path, zarr_format), document) + + def restore_copy(self, path: str, zarr_format: ZarrFormat, document: ArrayDocument) -> None: + """Record a consolidated copy as `copy_needing_repair` returned it from another + store object's cache, unless this cache already knows the array's document.""" + if not self.knows(path, zarr_format): + self._copies.setdefault((path, zarr_format), document) + + def observe(self, path: str, document: ArrayDocument | None, zarr_format: ZarrFormat) -> None: + """Record that the array document of `zarr_format` at `path` is `document` + (`None`: none).""" + key = (path, zarr_format) + if document is not None and moves_chunks(document, zarr_format): + self._arrays[key] = copy.deepcopy(document) + elif key in self._arrays or key in self._copies: + self._arrays[key] = None + + def entries(self, paths: Iterable[str], zarr_format: ZarrFormat) -> dict[str, CarriedEntry]: + """What this cache records of the array documents at `paths`, to carry to + another store object's cache (see `restore_entries`).""" + carried: dict[str, CarriedEntry] = {} + for path in paths: + if self.knows(path, zarr_format): + carried[path] = ("document", self.needing_repair(path, zarr_format)) + elif (copy_entry := self.copy_needing_repair(path, zarr_format)) is not None: + carried[path] = ("copy", copy_entry) + return carried + + def restore_entries(self, entries: Mapping[str, CarriedEntry], zarr_format: ZarrFormat) -> None: + """Record what `entries` carried, keeping what this cache already knows.""" + for path, entry in entries.items(): + if entry[0] == "document": + self.restore(path, zarr_format, entry[1]) + else: + self.restore_copy(path, zarr_format, entry[1]) + + def observe_written( + self, path: str, documents: Mapping[str, Buffer | None], written: WrittenArrayMetadata + ) -> None: + """Record the array metadata documents of `written` just stored under `path`, by + key (`None`: deleted), as `store_documents` takes them. They are parsed only + when this cache knows the path or `written` might move chunks, so a write of + valid metadata to a path the cache knows nothing about costs no parse.""" + key = (path, written.zarr_format) + if key not in self._arrays and key not in self._copies and not _may_move(written): + return + if ZARR_JSON in documents: + zarr_json = documents[ZARR_JSON] + document = None if zarr_json is None else buffer_to_json_object(zarr_json) + if document is not None and document.get("node_type") != "array": + document = None + self.observe(path, document, 3) + if ZARRAY_JSON in documents: + zarray = documents[ZARRAY_JSON] + if zarray is None: + self.observe(path, None, 2) + else: + zattrs = documents.get(ZATTRS_JSON) + v2_document: dict[str, JSON] = buffer_to_json_object(zarray) + v2_document["attributes"] = {} if zattrs is None else buffer_to_json_object(zattrs) + self.observe(path, v2_document, 2) + + def observe_consolidated( + self, group_path: str, members: Mapping[str, JSON], zarr_format: ZarrFormat + ) -> None: + """Record the member documents of consolidated metadata read from the group at + `group_path`, by their paths relative to it.""" + prefix = f"{group_path}/" if group_path else "" + for key, member in members.items(): + if not isinstance(member, dict) or member.get("node_type", "array") != "array": + continue + if zarr_format == 2 and "shape" not in member: + continue # a v2 group: it has no `shape` + path = f"{prefix}{key}" + if moves_chunks(member, zarr_format): + self._copies[path, zarr_format] = copy.deepcopy(member) + else: + self._copies.pop((path, zarr_format), None) + + +def _may_move(written: WrittenArrayMetadata) -> bool: + """`may_move_chunks` for the document `written` encodes, read from its chunk sizes.""" + from zarr.core.metadata import ArrayV3Metadata + + if written.zarr_format == 2: + return may_move_chunks({"chunks": list(written.chunks)}, 2) + grid: object = ( + written.chunk_grid.to_dict() + if isinstance(written, ArrayV3Metadata) + else written.chunk_grid.json + ) + return may_move_chunks(cast("ArrayDocument", {"chunk_grid": grid}), 3) + + +def _forget(key: int, ref: weakref.ref[Store]) -> None: + entry = _CACHES.get(key) + if entry is not None and entry[0] is ref: + del _CACHES[key] + + +_CREATING: Final = threading.Lock() +"""Held while a cache is created, so two threads first touching one store object share one.""" + +_CACHES: Final[dict[int, tuple[weakref.ref[Store], StoredDocuments]]] = {} +"""The cache of each live store object, by its `id`: stores define `__eq__`, so they +are not hashable, and two equal stores are not one store object.""" + + +def stored_documents(store: Store) -> StoredDocuments: + """The cache of `store`, shared by every binding of this store object; it is dropped + with the store.""" + key = id(store) + entry = _CACHES.get(key) + if entry is not None and entry[0]() is store: + return entry[1] + with _CREATING: + entry = _CACHES.get(key) + if entry is not None and entry[0]() is store: + return entry[1] + documents = StoredDocuments() + _CACHES[key] = (weakref.ref(store, partial(_forget, key)), documents) + return documents diff --git a/src/zarr/core/metadata/v2.py b/src/zarr/core/metadata/v2.py index 136267fc58..6b32b41676 100644 --- a/src/zarr/core/metadata/v2.py +++ b/src/zarr/core/metadata/v2.py @@ -4,7 +4,7 @@ import warnings from collections.abc import Iterable, Sequence from functools import cached_property -from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, cast +from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast from zarr.abc.metadata import Metadata from zarr.abc.numcodec import Numcodec, _is_numcodec @@ -25,7 +25,7 @@ TBaseScalar, ZDType, ) - from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.repair import Reading from dataclasses import dataclass, field, fields, replace @@ -45,7 +45,7 @@ from zarr.core.config import config, parse_indexing_order from zarr.core.json_parse import parse_field from zarr.core.metadata.common import parse_attributes -from zarr.core.metadata.repair import mark_repaired, repair_array_document +from zarr.core.metadata.repair import repair_array_document, warn_repairs class ArrayV2MetadataDict(TypedDict): @@ -73,9 +73,6 @@ class ArrayV2Metadata(Metadata): compressor: Numcodec | None attributes: dict[str, JSON] = field(default_factory=dict) zarr_format: Literal[2] = field(init=False, default=2) - _stored_document: ClassVar[ArrayDocument | None] = None - """The stored document `from_dict` read this metadata from, if it had to repair it - (set on the instance by `mark_repaired`): the store may still hold it.""" def __init__( self, @@ -156,6 +153,13 @@ def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2M """Read a stored `.zarray` document (with its attributes). An invalid document that `zarr.core.metadata.repair` can read is read as repaired; a reading the user must act on warns, naming the array at `path`.""" + metadata, readings = cls._read_stored(data) + warn_repairs(readings, path) + return metadata + + @classmethod + def _read_stored(cls, data: dict[str, Any]) -> tuple[ArrayV2Metadata, list[Reading]]: + """`from_dict` without its warnings: the metadata, and how the repairs read it.""" repaired, readings = repair_array_document(data, 2) # a new dict, because we are modifying it _data: dict[str, Any] = dict(repaired) @@ -209,7 +213,7 @@ def from_dict(cls, data: dict[str, Any], *, path: str | None = None) -> ArrayV2M _data = {k: v for k, v in _data.items() if k in expected} - return mark_repaired(cls(**_data), data, readings, path) + return cls(**_data), readings def to_dict(self) -> dict[str, JSON]: zarray_dict = super().to_dict() diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 2d5d169e6b..186e68f1c3 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -4,7 +4,7 @@ from collections.abc import Iterable, Mapping, Sequence from collections.abc import Set as AbstractSet from dataclasses import dataclass, field, replace -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NotRequired, TypeGuard, cast +from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast from typing_extensions import TypedDict @@ -41,8 +41,8 @@ from zarr.core.json_parse import parse_field from zarr.core.metadata.common import parse_attributes from zarr.core.metadata.repair import ( - mark_repaired, repair_array_document, + warn_repairs, ) from zarr.errors import MetadataValidationError, NodeTypeValidationError from zarr.registry import get_codec_class @@ -54,7 +54,7 @@ from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_grids import ChunkGrid from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar - from zarr.core.metadata.repair import ArrayDocument + from zarr.core.metadata.repair import Reading def parse_zarr_format(data: object) -> Literal[3]: @@ -519,9 +519,6 @@ class ArrayV3Metadata(Metadata): node_type: Literal["array"] = field(default="array", init=False) storage_transformers: tuple[dict[str, JSON], ...] extra_fields: dict[str, AllowedExtraField] - _stored_document: ClassVar[ArrayDocument | None] = None - """The stored document `from_dict` read this metadata from, if it had to repair it - (set on the instance by `mark_repaired`): the store may still hold it.""" def __init__( self, @@ -679,6 +676,14 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: if path is not None: e.add_note(f"Array {path!r}: nothing was read.") raise + metadata, readings = cls._read_stored(data) + warn_repairs(readings, path) + return metadata + + @classmethod + def _read_stored(cls, data: dict[str, JSON]) -> tuple[Self, list[Reading]]: + """`from_dict` without its warnings or the rectilinear chunks flag: the + metadata, and how the repairs read it.""" repaired, readings = repair_array_document(data, 3) # a new dict, because we are modifying it _data = dict(repaired) @@ -723,7 +728,7 @@ def from_dict(cls, data: dict[str, JSON], *, path: str | None = None) -> Self: extra_fields=allowed_extra_fields, storage_transformers=_data_typed.get("storage_transformers", ()), # type: ignore[arg-type] ) - return mark_repaired(metadata, data, readings, path) + return metadata, readings def to_dict(self) -> dict[str, JSON]: out_dict = super().to_dict() diff --git a/tests/test_array.py b/tests/test_array.py index 3e4835e686..40bde50783 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -7,7 +7,7 @@ import re import sys from itertools import accumulate, starmap -from typing import TYPE_CHECKING, Any, Literal +from typing import TYPE_CHECKING, Any, Literal, cast from unittest import mock import numcodecs @@ -56,6 +56,7 @@ resolve_outer_and_inner_chunks, ) from zarr.core.chunk_key_encodings import ChunkKeyEncodingParams +from zarr.core.codec_pipeline import BatchedCodecPipeline from zarr.core.common import JSON, ChunksLike, ZarrFormat, ceildiv from zarr.core.dtype import ( DateTime64, @@ -2637,3 +2638,129 @@ async def test_create_array_huge_chunk_count() -> None: arr = await create_array(store={}, shape=(2**62,), chunks=(1,), dtype="int32") assert arr.shape == (2**62,) assert arr.chunks == (1,) + + +def test_async_array_keeps_its_dataclass_surface() -> None: + """`AsyncArray` wraps an `AsyncArrayIO` but keeps the dataclass surface it has + always had: its fields, `replace`, `asdict`, equality, (un)hashability, pickled + state, swapping a field with `object.__setattr__`, and copies whose binding + (shape, pipeline) is their own; like before, a copy shares the metadata + object, so attribute updates made in place are seen through both.""" + import copy + import pickle + + from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline + + arr = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + aa = arr.async_array + + assert [f.name for f in dataclasses.fields(aa)] == [ + "metadata", + "store_path", + "codec_pipeline", + "_chunk_grid", + "config", + ] + assert set(dataclasses.asdict(aa)) == { + "metadata", + "store_path", + "codec_pipeline", + "_chunk_grid", + "config", + } + resized = dataclasses.replace(aa, metadata=dataclasses.replace(aa.metadata, shape=(8,))) + assert isinstance(resized, type(aa)) + assert (resized.shape, aa.shape) == ((8,), (4,)) + + with pytest.raises(TypeError, match="unhashable"): + hash(aa) + with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + fused = aa.with_config({}) + assert isinstance(fused.codec_pipeline, FusedCodecPipeline) + assert fused != aa + + swapped = copy.copy(aa) + object.__setattr__(swapped, "codec_pipeline", fused.codec_pipeline) + assert isinstance(swapped.codec_pipeline, FusedCodecPipeline) + assert isinstance(aa.codec_pipeline, BatchedCodecPipeline) + + copied = copy.copy(aa) + sync(copied.resize((6,))) + assert (copied.shape, aa.shape) == ((6,), (4,)) + + state = aa.__getstate__() + assert {"metadata", "store_path", "codec_pipeline", "_chunk_grid", "config"} <= set(state) + revived = pickle.loads(pickle.dumps(aa)) + assert revived == aa + + +def test_async_array_loads_a_pickled_field_state() -> None: + """A pickle of an `AsyncArray` holds its field state, as it always has, so an + array pickled before it wrapped an `AsyncArrayIO` loads.""" + arr = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + aa = arr.async_array + sync(aa.setitem(slice(None), np.arange(4, dtype="uint8"))) + old_state = pickle.loads( + pickle.dumps( + { + "metadata": aa.metadata, + "store_path": aa.store_path, + "codec_pipeline": aa.codec_pipeline, + "_chunk_grid": aa._chunk_grid, + "config": aa.config, + } + ) + ) + revived = object.__new__(type(aa)) + revived.__setstate__(old_state) + assert revived == aa + np.testing.assert_array_equal(sync(revived.getitem(slice(None))), np.arange(4)) + sync(revived.setitem(slice(None), np.full(4, 7, dtype="uint8"))) + np.testing.assert_array_equal(sync(revived.getitem(slice(None))), np.full(4, 7)) + + +def test_dask_token_of_an_array_is_deterministic() -> None: + dask_base = pytest.importorskip("dask.base") + store = MemoryStore() + arr = zarr.create_array(store, shape=(4,), dtype="uint8", chunks=(2,)) + before = dask_base.tokenize(arr) + _ = arr[:] + assert dask_base.tokenize(arr) == before + assert dask_base.tokenize(zarr.open_array(store)) == before + + +class _ChosenPipeline(BatchedCodecPipeline): + """A pipeline class whose instances may be another class's.""" + + +def test_pickled_array_keeps_the_pipeline_class_it_chose() -> None: + """A pickled array keeps the pipeline class its binding chose, also when that + class built an instance of another one (as zarrs-python does for stores it + cannot use).""" + aa = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)).async_array + aa.io._pipeline_class = _ChosenPipeline + revived = pickle.loads(pickle.dumps(aa)) + assert type(revived.codec_pipeline) is BatchedCodecPipeline + # A binding derived from the revived one builds its pipeline from the class. + derived = revived.io.with_store(MemoryStore()) + assert type(derived.to_array().codec_pipeline) is _ChosenPipeline + + +def test_import_zarr_leaves_zarr_metadata_unimported() -> None: + """`zarr_metadata` is imported when a binding first holds a model, not by + `import zarr`, whose cost stays what it was.""" + import subprocess + + code = "import sys, zarr; print('zarr_metadata' in sys.modules)" + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=False + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False" + + +def test_async_array_rejects_metadata_that_is_not_a_mapping() -> None: + with pytest.raises(TypeError): + AsyncArray(cast("Any", 42), StorePath(MemoryStore())) + with pytest.raises(TypeError): + zarr.AsyncArrayIO(cast("Any", 42), StorePath(MemoryStore())) diff --git a/tests/test_array_io.py b/tests/test_array_io.py new file mode 100644 index 0000000000..c7f357fc68 --- /dev/null +++ b/tests/test_array_io.py @@ -0,0 +1,1461 @@ +"""Tests for the explicit-IO array layer, `zarr.ArrayIO` and `zarr.AsyncArrayIO`.""" + +from __future__ import annotations + +import copy +import dataclasses +import json +import math +import pickle +import threading +from dataclasses import MISSING, dataclass +from typing import TYPE_CHECKING, Any, SupportsIndex, cast +from unittest import mock + +import numpy as np +import pytest +from zarr_metadata import UNSET, MetadataValidationError, ZarrV3ArrayMetadata +from zarr_metadata import ZarrV2ArrayMetadata as ModelV2 + +import zarr +from zarr import ArrayIO, AsyncArray, AsyncArrayIO +from zarr.api.asynchronous import open_array +from zarr.core.array import create_codec_pipeline, parse_array_metadata +from zarr.core.buffer import default_buffer_prototype +from zarr.core.codec_pipeline import BatchedCodecPipeline, FusedCodecPipeline +from zarr.core.metadata import ArrayMetadata, ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata.stored import stored_documents +from zarr.core.sync import sync +from zarr.dtype import Int16 +from zarr.errors import ArrayNotFoundError, UnknownCodecError, ZarrUserWarning +from zarr.storage import LocalStore, MemoryStore, StorePath + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + from pathlib import Path + + from zarr.core.common import JSON, ZarrFormat + +METADATA_DOCUMENTS = ("zarr.json", ".zarray", ".zattrs", ".zgroup") + + +async def _store_keys(store: MemoryStore) -> set[str]: + return {key async for key in store.list()} + + +def _declare_answer_array( + zarr_format: ZarrFormat, store: MemoryStore, data: np.ndarray[Any, np.dtype[Any]] +) -> AsyncArrayIO: + """Declare the combo test's array with the builder of `zarr_format`.""" + if zarr_format == 2: + return AsyncArrayIO.build_v2( + store, + name="foo", + shape=data.shape, + dtype=data.dtype, + chunks=(2, 2), + attributes={"answer": 42}, + ) + return AsyncArrayIO.build_v3( + store, + name="foo", + shape=data.shape, + data_type=data.dtype, + chunk_grid=(2, 2), + attributes={"answer": 42}, + ) + + +async def _store_foreign_codec_array(store: MemoryStore) -> None: + """Store a valid v3 array document whose codec this environment lacks.""" + document = ( + 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 + in that order, and the result is an ordinary readable array.""" + store = MemoryStore() + data = np.arange(16, dtype="uint16").reshape(4, 4) + aio = _declare_answer_array(zarr_format, store, data) + 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() + # the default zarr_format=None auto-detects the format, like zarrs + opened = await AsyncArrayIO.open(store, name="foo") + # the binding is metadata; the stored and declared metadata agree + assert opened.metadata == aio.metadata + assert opened.metadata.shape == data.shape + assert opened.metadata_legacy.zarr_format == zarr_format + assert dict(opened.metadata.attributes) == {"answer": 42} + np.testing.assert_array_equal(await opened.getitem(Ellipsis), data) + + arr = await open_array(store=store, path="foo", mode="r", zarr_format=zarr_format) + np.testing.assert_array_equal(await arr.getitem(Ellipsis), data) + # every AsyncArray is built on an AsyncArrayIO, reachable as .io + assert arr.io.metadata_legacy is arr.metadata + np.testing.assert_array_equal(await arr.io.getitem(Ellipsis), data) + + # a writer that declared the same array independently can verify its assumption + worker = _declare_answer_array(zarr_format, store, data) + 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)) + 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.""" + store = MemoryStore() + data = np.arange(16, dtype="uint16").reshape(4, 4) + aio = ArrayIO.build_v3( + store, name="foo", shape=data.shape, data_type=data.dtype, chunk_grid=(2, 2) + ) + aio.setitem(Ellipsis, data) + aio.store_metadata() + opened = ArrayIO.open(store, name="foo") + 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() + with pytest.raises(ArrayNotFoundError): + ArrayIO.open(store, name="foo") + + +def test_build_v3_unknown_chunk_grid_raises() -> None: + with pytest.raises(ValueError, match="chunk grid"): + ArrayIO.build_v3( + MemoryStore(), + name="foo", + shape=(4,), + data_type="uint8", + chunk_grid={"name": "diagonal", "configuration": {}}, + ) + + +async def test_open_missing_raises() -> None: + with pytest.raises(ArrayNotFoundError): + 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") + model = aio.metadata + assert isinstance(model, ZarrV3ArrayMetadata) + assert model.shape == (4,) + assert model.must_understand_fields == {} + await aio.verify() + + with pytest.raises(UnknownCodecError, match="zarr-python-does-not-know-this-codec"): + 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 = AsyncArrayIO.build_v3(store, shape=(4, 4), data_type="uint8", chunk_grid=(2, 2)) + await aio.store_metadata() + other = AsyncArrayIO.build_v3( + store, shape=(4, 4), data_type="uint8", chunk_grid=(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 AsyncArrayIO.build_v2(store, shape=(4,), dtype="uint8", chunks=(4,)).store_metadata() + bound_v3 = AsyncArrayIO.build_v3(store, shape=(4,), data_type="uint8", chunk_grid=(4,)) + with pytest.raises(ValueError, match="stored array is Zarr format 2"): + await bound_v3.verify() + + +async def test_verify_missing_raises() -> None: + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,)) + with pytest.raises(ArrayNotFoundError): + await aio.verify() + + +async def test_erase_chunk_out_of_grid_raises() -> None: + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4, 4), data_type="uint8", chunk_grid=(2, 2)) + with pytest.raises(IndexError, match="outside the chunk grid"): + await aio.erase_chunk((2, 0)) + with pytest.raises(IndexError, match="outside the chunk grid"): + await aio.erase_chunk((0,)) + + +async def test_with_methods_rebind() -> None: + """with_metadata, with_store and with_config return new bindings with one + part replaced and everything else kept.""" + store, other_store = MemoryStore(), MemoryStore() + data = np.arange(4, dtype="uint8") + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await aio.setitem(Ellipsis, data) + await aio.store_metadata() + + relabeled = aio.with_metadata(aio.metadata.update(attributes={"a": 1})) + assert relabeled.store_path == aio.store_path + assert dict(relabeled.metadata.attributes) == {"a": 1} + await relabeled.store_metadata() + reread = await AsyncArrayIO.open(store, name="foo") + assert reread.metadata == relabeled.metadata + + moved = aio.with_store(other_store) + assert moved.store_path.path == aio.store_path.path + await moved.setitem(Ellipsis, data) + await moved.store_metadata() + copied = await AsyncArrayIO.open(other_store, name="foo") + np.testing.assert_array_equal(await copied.getitem(Ellipsis), data) + + configured = aio.with_config({"order": "F"}) + assert configured.config.order == "F" + assert configured.metadata == aio.metadata + + +def test_binding_follows_the_global_config() -> None: + """ArrayIO responds to the global configuration the way Array always has: + the pipeline implementation is captured when the binding is made, and + build_v2's order defaults from array.order.""" + with zarr.config.set( + { + "codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline", + "array.order": "F", + } + ): + v3 = ArrayIO.build_v3(MemoryStore(), shape=(2,), data_type="uint8", chunk_grid=(2,)) + v2 = ArrayIO.build_v2(MemoryStore(), shape=(2,), dtype="uint8", chunks=(2,)) + # outside the block, the choices made inside it hold, also for a binding + # derived while the global configuration names a pipeline that cannot load + assert isinstance(v3.to_array().async_array.codec_pipeline, FusedCodecPipeline) + with zarr.config.set({"codec_pipeline.path": "nonexistent.module.Pipeline"}): + derived = v3.with_config({"order": "C"}) + assert isinstance(derived.to_array().async_array.codec_pipeline, FusedCodecPipeline) + legacy = v2.metadata_legacy + assert isinstance(legacy, ArrayV2Metadata) + assert legacy.order == "F" + + +async def test_array_state_changes_rebind_only_the_array() -> None: + """A change through the array side (resize, update_attributes) shows on the + array's current binding (`io`), not on a binding taken before it, and a + derived binding stays independent.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await aio.store_metadata() + arr = aio.to_array() + + derived = aio.with_config({"order": "F"}) + moved = aio.with_store(MemoryStore()) + + await arr.resize((6,)) + assert arr.io.metadata.shape == (6,) + assert aio.metadata.shape == (4,) + + await arr.update_attributes({"a": 1}) + assert dict(arr.io.metadata.attributes) == {"a": 1} + + # the derivatives were made before the changes and do not alias them + assert derived.metadata_legacy.shape == (4,) + assert dict(derived.metadata_legacy.attributes) == {} + assert dict(moved.metadata_legacy.attributes) == {} + + # the same independence holds for a legacy-sourced binding, whose mutable + # attributes dict must be copied, not shared + legacy_arr = await zarr.api.asynchronous.create_array( + MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,) + ) + legacy_derivatives = [ + legacy_arr.io.with_config({"order": "F"}), + legacy_arr.io.with_store(MemoryStore()), + legacy_arr.io.with_metadata(legacy_arr.io.metadata_legacy), + ] + for legacy_derived in legacy_derivatives: + _ = legacy_derived.metadata # fill the model cache before the mutation + await legacy_arr.update_attributes({"b": 2}) + assert dict(legacy_arr.io.metadata.attributes) == {"b": 2} + for legacy_derived in legacy_derivatives: + assert dict(legacy_derived.metadata_legacy.attributes) == {} + assert dict(legacy_derived.metadata.attributes) == {} + + +async def test_store_metadata_of_model_without_attributes_clears_zattrs() -> None: + """Storing a v2 model with no attributes removes a stale .zattrs.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v2( + store, name="foo", shape=(4,), dtype="uint8", chunks=(4,), attributes={"stale": 1} + ) + await aio.store_metadata() + assert "foo/.zattrs" in await _store_keys(store) + + opened = await AsyncArrayIO.open(store, name="foo") + model = opened.metadata + assert isinstance(model, ModelV2) + bare = opened.with_metadata(model.update(attributes=UNSET)) + await bare.store_metadata() + assert "foo/.zattrs" not in await _store_keys(store) + reread = await AsyncArrayIO.open(store, name="foo") + assert dict(reread.metadata.attributes) == {} + + +def _bind_zero_chunk_by_constructor(metadata: ArrayV2Metadata) -> AsyncArrayIO: + return AsyncArrayIO(metadata, StorePath(MemoryStore())) + + +def _bind_zero_chunk_by_build(metadata: ArrayV2Metadata) -> AsyncArrayIO: + return AsyncArrayIO.build_v2( + MemoryStore(), + shape=metadata.shape, + dtype=metadata.dtype, + chunks=metadata.chunks, + fill_value=metadata.fill_value, + order=metadata.order, + ) + + +def _bind_zero_chunk_by_with_metadata(metadata: ArrayV2Metadata) -> AsyncArrayIO: + declared = AsyncArrayIO.build_v2(MemoryStore(), shape=(3,), dtype="int16", chunks=(3,)) + return declared.with_metadata(metadata) + + +@pytest.mark.parametrize( + "bind", + [_bind_zero_chunk_by_constructor, _bind_zero_chunk_by_build, _bind_zero_chunk_by_with_metadata], + ids=["constructor", "build_v2", "with_metadata"], +) +def test_zero_chunk_v2_metadata_binds_repaired( + bind: Callable[[ArrayV2Metadata], AsyncArrayIO], +) -> None: + """`ArrayV2Metadata` with a chunk size of 0 is read as repaired however it + reaches a binding, as `parse_array_metadata` reads it.""" + metadata = ArrayV2Metadata(shape=(3,), chunks=(0,), dtype=Int16(), fill_value=0, order="C") + assert bind(metadata).metadata_legacy.chunks == (1,) + + +async def test_v2_model_binding_keeps_attributes_writable() -> None: + """A v2 binding read as a model resolves its attributes as plain, mutable + JSON, so writes through it that re-encode the attributes succeed.""" + store = MemoryStore() + attributes: dict[str, JSON] = {"k": [1, 2], "d": {"n": None}} + aio = AsyncArrayIO.build_v2( + store, name="foo", shape=(4,), dtype="uint8", chunks=(2,), attributes=attributes + ) + await aio.store_metadata() + arr = (await AsyncArrayIO.open(store, name="foo")).to_array() + + assert arr.metadata.attributes == attributes + await arr.update_attributes({"new": 1}) + await arr.resize((6,)) + + reread = await open_array(store=store, path="foo") + assert reread.metadata.attributes == {**attributes, "new": 1} + assert reread.shape == (6,) + + +def test_build_v2_vlen_without_object_codec_raises() -> None: + with pytest.raises(ValueError, match="requires the numcodecs.VLenUTF8 codec"): + AsyncArrayIO.build_v2(MemoryStore(), shape=(4,), dtype=str, chunks=(2,)) + + +async def test_metadata_dict_binds_as_a_model_until_chunks_move() -> None: + """A valid metadata dict binds as a model, so an extension zarr-python + lacks fails only when chunks are read or written; metadata zarr can only + read as repaired is read, repaired, when bound.""" + store = MemoryStore() + doc: dict[str, JSON] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "bfloat16", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "zarrs.zfp", "configuration": {"mode": "reversible"}}], + } + aio = AsyncArrayIO(cast("Any", doc), StorePath(store, "foo")) + assert isinstance(aio.metadata, ZarrV3ArrayMetadata) + await aio.store_metadata() + await aio.verify() + await aio.erase_chunks() + with pytest.raises(ValueError, match="bfloat16"): + await aio.getitem(Ellipsis) + + zero_chunk = {**doc, "data_type": "uint8", "codecs": [{"name": "bytes"}]} + zero_chunk["chunk_grid"] = {"name": "regular", "configuration": {"chunk_shape": [0]}} + with pytest.warns(zarr.errors.ZarrUserWarning, match="is read as"): + repaired = AsyncArrayIO(cast("Any", zero_chunk), StorePath(store, "bar")) + assert repaired.metadata_legacy.chunks == (1,) + + +async def test_chunk_and_selection_io_signatures() -> None: + """erase_chunk takes grid coordinates positionally, as any sequence; + setitem takes its prototype by keyword.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8")) + await aio.erase_chunk([0]) + assert "foo/c/0" not in await _store_keys(store) + with pytest.raises(TypeError): + await aio.erase_chunk(chunk_coords=(1,)) # type: ignore[call-arg] + with pytest.raises(TypeError): + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8"), default_buffer_prototype()) # type: ignore[call-arg] + + +async def test_v2_model_binding_without_attributes() -> None: + store = MemoryStore() + aio = AsyncArrayIO.build_v2(store, name="foo", shape=(4,), dtype="uint8", chunks=(2,)) + await aio.store_metadata() + await store.delete("foo/.zattrs") + arr = (await AsyncArrayIO.open(store, name="foo")).to_array() + assert arr.metadata.attributes == {} + await arr.update_attributes({"x": 1}) + assert (await open_array(store=store, path="foo")).metadata.attributes == {"x": 1} + + +async def _store_v2_document(store: MemoryStore, document: dict[str, Any]) -> None: + prototype = default_buffer_prototype() + await store.set("a/.zarray", prototype.buffer.from_bytes(json.dumps(document).encode())) + + +_V2_DOCUMENT: dict[str, Any] = { + "zarr_format": 2, + "shape": [4], + "chunks": [4], + "dtype": " None: + """The strict reads (`open`, `verify`) reject a v2 chunk size of 0, as they do a + v3 one, though the v2 model accepts it, naming where and what kind.""" + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "chunks": [0]}) + with pytest.raises(MetadataValidationError) as opened: + await AsyncArrayIO.open(store, name="a") + binding = AsyncArrayIO.build_v2(store, name="a", shape=(4,), dtype=" None: + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "chunks": [0], "dtype": "nope"}) + with pytest.raises(MetadataValidationError) as error: + await AsyncArrayIO.open(store, name="a") + locations = [p.loc for p in error.value.problems] + assert locations.count(("chunks", 0)) == 1 + assert any(loc[0] == "dtype" for loc in locations) + + +async def test_a_strict_v2_read_reports_a_chunk_size_once() -> None: + """A chunk size the model rejects itself is reported once.""" + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "chunks": [-1]}) + with pytest.raises(MetadataValidationError) as error: + await AsyncArrayIO.open(store, name="a") + assert [p.loc for p in error.value.problems] == [("chunks",)] + + +def test_with_metadata_rejects_what_is_not_metadata() -> None: + aio = AsyncArrayIO.build_v3(MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,)) + with pytest.raises(TypeError, match="got str"): + aio.with_metadata(cast("Any", "nonsense")) + + +def test_open_a_missing_path_creates_nothing(tmp_path: Path) -> None: + missing = tmp_path / "missing.zarr" + with pytest.raises(FileNotFoundError): + ArrayIO.open(str(missing)) + assert not missing.exists() + + +def test_metadata_dict_with_storage_transformers_raises() -> None: + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "uint8", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes"}], + "storage_transformers": [{"name": "nope"}], + } + with pytest.raises(ValueError, match="storage transformers"): + AsyncArrayIO(cast("Any", doc), StorePath(MemoryStore())) + + +async def test_bindings_take_metadata_json() -> None: + """A binding takes array metadata as `zarr_metadata` types it in JSON (what + zarrista's arrays hold), through the constructor and with_metadata.""" + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + json_metadata = aio.metadata.to_json() + rebound = aio.with_metadata(json_metadata) + assert AsyncArrayIO(json_metadata, aio.store_path) == aio + assert rebound == aio + await rebound.store_metadata() + await rebound.setitem(Ellipsis, np.arange(4, dtype="uint8")) + np.testing.assert_array_equal(await aio.getitem(Ellipsis), np.arange(4)) + + +async def test_pickled_binding_rebuilds_its_pipeline(tmp_path: Path) -> None: + """A pickled binding drops its built pipeline (a pipeline assigned to it is not + carried either) and builds its class's pipeline on first use.""" + aio = AsyncArrayIO.build_v3( + LocalStore(tmp_path), name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,) + ) + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8")) + revived = pickle.loads(pickle.dumps(aio)) + assert revived._codec_pipeline_cache is None + np.testing.assert_array_equal(await revived.getitem(Ellipsis), np.arange(4)) + assert isinstance(revived._resolved_codec_pipeline, revived._pipeline_class) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_erase_and_store_metadata_leave_a_group_in_place(zarr_format: ZarrFormat) -> None: + """erase_metadata deletes only an array's documents, and storing a v2 model + without attributes keeps a group's `.zattrs`.""" + store = MemoryStore() + group = await zarr.api.asynchronous.open_group( + store, path="g", zarr_format=zarr_format, mode="w" + ) + await group.update_attributes({"owner": "group"}) + before = await _store_keys(store) + declared = ( + AsyncArrayIO.build_v2(store, name="g", shape=(4,), dtype="uint8", chunks=(4,)) + if zarr_format == 2 + else AsyncArrayIO.build_v3(store, name="g", shape=(4,), data_type="uint8", chunk_grid=(4,)) + ) + await declared.erase_metadata() + assert await _store_keys(store) == before + if zarr_format == 2: + model = declared.metadata + assert isinstance(model, ModelV2) + await declared.with_metadata(model.update(attributes=UNSET)).store_metadata() + reopened = await zarr.api.asynchronous.open_group(store, path="g", mode="r", zarr_format=2) + assert dict(reopened.attrs) == {"owner": "group"} + + +def test_constructor_never_shares_the_callers_metadata() -> None: + mine = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(4,)).metadata + aio = AsyncArrayIO(mine, StorePath(MemoryStore())) + assert aio.metadata_legacy is not mine + assert aio.metadata_legacy.attributes is not mine.attributes + + +async def test_metadata_dict_of_a_repaired_document_stores_its_repair() -> None: + """An invalid stored document given to AsyncArrayIO as a dict is recorded as + what the store holds, as AsyncArray records it, so the first chunk write + stores the repair.""" + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 2, + "shape": [3], + "chunks": [0], + "dtype": " None: + """Nested attributes (read-only in the model) do not stop the field naming.""" + store = MemoryStore() + attributes: dict[str, JSON] = {"units": {"name": "m", "scale": [1, 2]}} + aio = ( + AsyncArrayIO.build_v2( + store, + name="foo", + shape=(4,), + dtype="float32", + chunks=(4,), + fill_value=0.0, + attributes=attributes, + ) + if zarr_format == 2 + else AsyncArrayIO.build_v3( + store, + name="foo", + shape=(4,), + data_type="float32", + chunk_grid=(4,), + fill_value=0.0, + attributes=attributes, + ) + ) + await aio.store_metadata() + other = aio.with_metadata(aio.metadata.update(fill_value=-0.0)) + with pytest.raises(ValueError, match=r"differing fields: fill_value\.$"): + await other.verify() + + +def test_build_v3_keeps_empty_dimension_names_on_a_0d_array() -> None: + """The builder keeps an empty name list for a 0-d array; the creation paths + read it as no names, as they always have.""" + built = ArrayIO.build_v3( + MemoryStore(), shape=(), data_type="int8", chunk_grid=(), dimension_names=[] + ) + assert built.metadata_legacy.dimension_names == () # type: ignore[union-attr] + created = zarr.create_array(MemoryStore(), shape=(), dtype="int8", dimension_names=[]) + assert created.metadata.dimension_names is None # type: ignore[union-attr] + + +async def test_erase_chunk_takes_integer_coordinates() -> None: + store = MemoryStore() + aio = AsyncArrayIO.build_v3(store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await aio.setitem(Ellipsis, np.arange(4, dtype="uint8")) + with pytest.raises(TypeError): + await aio.erase_chunk([cast("int", np.float64(1))]) + await aio.erase_chunk([cast("int", np.int64(1))]) + assert "foo/c/1" not in await _store_keys(store) + + +def test_equality_is_canonical_whatever_is_resolved() -> None: + """Bindings compare by canonical metadata, before and after resolving, and + whatever they were built from: a fill value of -0.0 differs from 0.0.""" + store_path = StorePath(MemoryStore(), "foo") + positive = AsyncArrayIO.build_v3( + store_path, shape=(4,), data_type="float32", chunk_grid=(4,), fill_value=0.0 + ) + negative = positive.with_metadata(positive.metadata.update(fill_value=-0.0)) + legacy_negative = AsyncArrayIO(negative.metadata_legacy, store_path) + assert positive != negative + assert positive != legacy_negative + assert negative == legacy_negative + _ = (positive.metadata_legacy, negative.metadata_legacy) + assert positive != negative + assert negative == legacy_negative + + +@pytest.mark.parametrize("bind", ["open", "model"]) +async def test_storage_transformers_are_rejected_when_bound(bind: str) -> None: + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "uint8", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes"}], + "storage_transformers": [{"name": "nope"}], + } + if bind == "open": + prototype = default_buffer_prototype() + await store.set("foo/zarr.json", prototype.buffer.from_bytes(json.dumps(doc).encode())) + with pytest.raises(ValueError, match="storage transformers"): + await AsyncArrayIO.open(store, name="foo") + else: + model = ZarrV3ArrayMetadata(doc) + with pytest.raises(ValueError, match="storage transformers"): + AsyncArrayIO(model, StorePath(store, "foo")) + + +@pytest.mark.parametrize( + ("documents", "left"), + [ + ({"foo/zarr.json": b"not json"}, set()), + ({"foo/zarr.json": b'{"zarr_format": 3}'}, set()), + ({"foo/zarr.json": b'{"zarr_format": 3, "node_type": "group"}'}, {"foo/zarr.json"}), + ({"foo/.zattrs": b"{}"}, set()), + ( + {"foo/.zarray": b"{}", "foo/.zattrs": b"{}", "foo/.zgroup": b"{}"}, + {"foo/.zattrs", "foo/.zgroup"}, + ), + ], + ids=["corrupt-v3", "v3-without-node-type", "v3-group", "stray-v2-zattrs", "v2-array-and-group"], +) +async def test_erase_metadata_removes_array_documents_however_written( + documents: dict[str, bytes], left: set[str] +) -> None: + """erase_metadata removes array documents even when they are corrupt, and + leaves only a group's documents in place.""" + store = MemoryStore() + prototype = default_buffer_prototype() + for key, value in documents.items(): + await store.set(key, prototype.buffer.from_bytes(value)) + zarr_format: ZarrFormat = 3 if "foo/zarr.json" in documents else 2 + declared = ( + AsyncArrayIO.build_v2(store, name="foo", shape=(4,), dtype="uint8", chunks=(4,)) + if zarr_format == 2 + else AsyncArrayIO.build_v3( + store, name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + ) + await declared.erase_metadata() + assert await _store_keys(store) == left + + +async def test_with_metadata_takes_any_mapping() -> None: + from types import MappingProxyType + + aio = AsyncArrayIO.build_v3( + MemoryStore(), name="foo", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + rebound = aio.with_metadata(cast("Any", MappingProxyType(dict(aio.metadata.to_json())))) + assert rebound == aio + await rebound.store_metadata() + + +class _TaggedAsyncArray(AsyncArray[Any]): + """An `AsyncArray` subclass with an instance attribute of its own.""" + + tag: str + + def __init__(self, metadata: Any, store_path: StorePath, tag: str = "default") -> None: + super().__init__(metadata, store_path) + object.__setattr__(self, "tag", tag) + + +def test_async_array_subclass_state_survives_pickle_and_copy() -> None: + """A subclass's own instance attributes pickle and copy with it, as they did + when the dataclass fields were an array's only state.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + tagged = _TaggedAsyncArray(created.metadata, created.store_path, tag="mine") + for restored in (pickle.loads(pickle.dumps(tagged)), copy.copy(tagged), copy.deepcopy(tagged)): + assert type(restored) is _TaggedAsyncArray + assert restored.tag == "mine" + assert restored == tagged + assert set(vars(restored)) == {"_io", "tag"} + + +async def test_copying_an_array_shares_its_metadata() -> None: + """A shallow copy shares the metadata object, also for an array opened from a + document that needs a repair; rebinding one (resize) leaves the other.""" + store = MemoryStore() + doc: dict[str, Any] = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "int32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [0]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + } + prototype = default_buffer_prototype() + await store.set("r/zarr.json", prototype.buffer.from_bytes(json.dumps(doc).encode())) + with pytest.warns(ZarrUserWarning): + opened = await open_array(store=store, path="r") + copied = copy.copy(opened) + assert copied.metadata is opened.metadata + assert "_stored_document" not in vars(copied.metadata) + await opened.resize((8,)) + assert copied.shape == (4,) + assert opened.shape == (8,) + + +async def test_writing_valid_metadata_parses_nothing(monkeypatch: pytest.MonkeyPatch) -> None: + """Storing metadata that cannot move chunks, at a path the store's cache knows + nothing about, records nothing and parses no document.""" + import zarr.core.metadata.stored as stored_module + + monkeypatch.setattr( + stored_module, + "buffer_to_json_object", + mock.Mock(side_effect=AssertionError("a written document was parsed")), + ) + store = MemoryStore() + arr = zarr.create_array(store, name="a", shape=(4,), dtype="uint8", attributes={"x": 1}) + arr.attrs["y"] = 2 + arr.resize((8,)) + v2 = zarr.create_array(store, name="b", shape=(4,), dtype="uint8", zarr_format=2) + v2.attrs["y"] = 2 + assert not stored_documents(store).knows("a", 3) + + +async def test_an_array_copied_by_its_dict_keeps_its_own_state() -> None: + """An array made by copying another's `__dict__` (as subclasses do) is changed + apart from it: a resize, an attribute update or a field set through either + leaves the other, as when its fields were its only state; so does a binding + taken from `io` before the change.""" + original = zarr.create_array(MemoryStore(), shape=(4,), dtype="int32", chunks=(2,))._async_array + taken = original.io + clone = object.__new__(AsyncArray) + clone.__dict__.update(original.__dict__) + await clone.resize((6,)) + assert (original.shape, clone.shape, taken.metadata_legacy.shape) == ((4,), (6,), (4,)) + await clone.update_attributes({"x": 1}) + assert original.attrs == {} + object.__setattr__(clone, "metadata", clone.metadata.update_shape((3,))) + assert original.shape == (4,) + + +_V2_CHUNK_ZERO: dict[str, Any] = { + "zarr_format": 2, + "shape": [0], + "chunks": [0], + "dtype": " AsyncArrayIO: + return AsyncArrayIO(model, StorePath(MemoryStore(), "a")) + + +def _bind_by_with_metadata(model: ModelV2) -> AsyncArrayIO: + declared = AsyncArrayIO.build_v2(MemoryStore(), shape=(4,), dtype=" None: + """The v2 model accepts a chunk size of 0, which zarr reads as repaired; bound + as a model, it is rejected as the strict read rejects it.""" + model = ModelV2(_V2_CHUNK_ZERO) + with pytest.raises(MetadataValidationError, match="chunks"): + bind(model) + + +def test_a_tuple_valued_v2_chunk_size_zero_binds_as_repaired() -> None: + """A metadata mapping written in Python (tuples, not lists) binds as `AsyncArray` + reads it: a chunk size of 0 as repaired.""" + document = {**_V2_CHUNK_ZERO, "shape": (0,), "chunks": (0,)} + binding = ArrayIO(cast("Any", document), StorePath(MemoryStore(), "a")) + assert binding.metadata_legacy.chunks == (1,) + binding.store_metadata() + binding.verify() + + +async def test_an_array_over_a_model_shows_one_document() -> None: + """An array over a model-backed binding hands out its metadata, which callers + change in place (attribute updates through a copy); the binding's model then + reads from it, so storing through the binding keeps the change.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,), attributes={"v": 0} + ).store_metadata() + original = (await AsyncArrayIO.open(store, name="a")).to_array() + _ = original.metadata + for copied in (copy.copy(original), original.with_config({"order": "F"})): + await copied.update_attributes({"v": id(copied)}) + assert dict(original.io.metadata.attributes) == original.attrs + await original.io.store_metadata() + await original.io.verify() + + +def test_bindings_of_different_classes_hash_as_they_compare() -> None: + class _Binding(AsyncArrayIO): + pass + + base = AsyncArrayIO.build_v3( + MemoryStore(), name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + other = _Binding(base.metadata, base.store_path) + assert other == base + assert hash(other) == hash(base) + + +async def test_attribute_updates_through_an_array() -> None: + """An attribute update through an array never reaches the binding the array was + made from (`to_array`), shows through a binding taken from the array (`io`) + consistently in both forms, and through an array copied by its `__dict__`, as + it did when the fields were an array's only state.""" + store = MemoryStore() + origin = AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,), attributes={"v": 0} + ) + await origin.store_metadata() + array = origin.to_array() + taken = array.io + clone = object.__new__(AsyncArray) + clone.__dict__.update(array.__dict__) + await array.update_attributes({"v": 1}) + assert dict(origin.metadata.attributes) == {"v": 0} + assert dict(origin.metadata_legacy.attributes) == {"v": 0} + assert dict(taken.metadata.attributes) == taken.metadata_legacy.attributes == {"v": 1} + assert clone.attrs == {"v": 1} + await taken.store_metadata() + await taken.verify() + + +@pytest.mark.parametrize( + ("data_type", "fill_value"), [("float32", "0x7fc00001"), ("float64", "0x7ff8000000000001")] +) +async def test_reading_an_array_leaves_its_binding_as_stored( + data_type: str, fill_value: str +) -> None: + """Reading an array made from a binding (its shape, its data) leaves the binding + holding the document as stored: it still verifies, compares equal to a fresh + open, and stores that document unchanged (a NaN payload is kept).""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": data_type, + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": fill_value, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + stored = default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + await store.set("a/zarr.json", stored) + binding = await AsyncArrayIO.open(store, name="a") + array = binding.to_array() + _ = array.shape + await array.getitem(slice(None)) + assert binding == await AsyncArrayIO.open(store, name="a") + await binding.verify() + await binding.store_metadata() + rewritten = await store.get("a/zarr.json", prototype=default_buffer_prototype()) + assert rewritten is not None + assert json.loads(rewritten.to_bytes()) == document + + +class _LoggingBinding(AsyncArrayIO): + """A binding subclass, as a backend would define one.""" + + +def test_an_array_keeps_its_binding_class() -> None: + """A change to an array's state, and pickling, keep the class of the binding + it was made from.""" + binding = _LoggingBinding.build_v3( + MemoryStore(), name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + array = binding.to_array() + for kept in (array, pickle.loads(pickle.dumps(array))): + sync(kept.resize((6,))) + assert type(kept.io) is _LoggingBinding + + +@dataclass(frozen=True) +class _Typed(AsyncArray[Any]): + """A `@dataclass` subclass that annotates a field again and adds one with no default.""" + + metadata: ArrayV3Metadata + label: str + + def __init__( + self, metadata: Any, store_path: StorePath, config: Any = None, label: str = "x" + ) -> None: + super().__init__(metadata, store_path, config) + object.__setattr__(self, "label", label) + + +def test_a_dataclass_subclass_may_annotate_fields_again() -> None: + """A `@dataclass` subclass that annotates a field again, or adds one without a + default, defines and works, as when the fields had no class attribute.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + typed = _Typed(created.metadata, created.store_path) + assert typed.shape == (4,) + assert dataclasses.replace(typed, label="y").label == "y" + + +class _NativeBinding(AsyncArrayIO): + """A binding that derives state of its own (a native handle's shape) in `_bind`.""" + + native_shape: tuple[int, ...] + native_where: str + + def _bind(self, source: Any, store_path: StorePath, *args: Any) -> None: + super()._bind(source, store_path, *args) + self.native_shape = tuple(source.shape) + self.native_where = str(store_path) + + +async def test_every_route_rebinds_through_bind() -> None: + """Every way a binding is made or changed goes through `_bind`, so state a + subclass derives there follows the metadata, also after a resize.""" + store = MemoryStore() + built = _NativeBinding.build_v3(store, name="a", shape=(4,), data_type="uint8", chunk_grid=(2,)) + await built.store_metadata() + array = built.to_array() + await array.resize((6,)) + bindings = [ + built, + await _NativeBinding.open(store, name="a"), + built.with_metadata(built.metadata), + built.with_store(MemoryStore()), + built.with_config({"order": "F"}), + pickle.loads(pickle.dumps(built)), + array.io, + array.with_config({"order": "F"}).io, + pickle.loads(pickle.dumps(array)).io, + copy.copy(array).io, + ] + moved = copy.copy(array) + object.__setattr__(moved, "store_path", StorePath(MemoryStore(), "b")) + bindings.append(moved.io) + for binding in bindings: + assert isinstance(binding, _NativeBinding) + assert binding.native_shape == binding.metadata.shape + assert binding.native_where == str(binding.store_path) + + +@pytest.mark.parametrize("derive", ["with_config", "pickle", "deepcopy"]) +async def test_an_array_over_a_model_keeps_the_document_as_stored(derive: str) -> None: + """An array rebuilt from one over a model binding (`with_config`, pickling, a + deep copy) keeps the model, so its binding still verifies against the store + and stores the document unchanged (a NaN payload is kept).""" + store = MemoryStore() + document = { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "data_type": "float32", + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "chunk_key_encoding": {"name": "default"}, + "fill_value": "0x7fc00001", + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + await store.set( + "a/zarr.json", default_buffer_prototype().buffer.from_bytes(json.dumps(document).encode()) + ) + array = (await AsyncArrayIO.open(store, name="a")).to_array() + if derive == "with_config": + rebuilt = array.with_config({"order": "F"}) + elif derive == "pickle": + rebuilt = pickle.loads(pickle.dumps(array)) + else: + rebuilt = copy.deepcopy(array) + await rebuilt.io.verify() + await rebuilt.io.store_metadata() + stored = await store.get("a/zarr.json", prototype=default_buffer_prototype()) + assert stored is not None + assert json.loads(stored.to_bytes()) == document + + +async def test_bindings_compare_with_attributes_json_cannot_hold() -> None: + """Equality of model bindings whose attributes were changed in place to values + JSON cannot hold compares their documents without the attributes, then the + attributes; reading the model names the value.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="float32", chunk_grid=(4,), fill_value=0.0 + ).store_metadata() + first, second = [(await AsyncArrayIO.open(store, name="a")).to_array() for _ in range(2)] + for array in (first, second): + array.metadata.attributes["n"] = np.int64(3) + assert first.io == second.io + second.metadata.attributes["n"] = np.int64(4) + assert first.io != second.io + with pytest.raises(TypeError, match="int64 is not JSON serializable"): + _ = first.io.metadata + + +async def test_a_model_binding_reads_changed_attributes_once() -> None: + """A model binding whose attributes were changed in place gives one model for + them until they change again, also for attributes holding NaN.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + array.metadata.attributes["x"] = float("nan") + binding = array.io + assert binding.metadata is binding.metadata + assert math.isnan(cast("float", binding.metadata.attributes["x"])) + + +class _SlottedAsyncArray(AsyncArray[Any]): + """An `AsyncArray` subclass with a slot of its own.""" + + __slots__ = ("__private", "extra") + extra: int + + +class _ReducingAsyncArray(AsyncArray[Any]): + """An `AsyncArray` subclass that decides how it is copied.""" + + def __reduce_ex__(self, protocol: SupportsIndex) -> Any: + return (_rebuilt_marked, (self.metadata, self.store_path)) + + +def _rebuilt_marked(metadata: Any, store_path: StorePath) -> _ReducingAsyncArray: + rebuilt = _ReducingAsyncArray(metadata, store_path) + object.__setattr__(rebuilt, "marked", True) + return rebuilt + + +def test_subclass_slots_and_reductions_are_honored() -> None: + """A subclass's slots pickle and copy with it, and a subclass that defines its + own reduction is copied its way, as when the fields were an array's only state.""" + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + slotted = _SlottedAsyncArray(created.metadata, created.store_path) + object.__setattr__(slotted, "extra", 5) + object.__setattr__(slotted, "_SlottedAsyncArray__private", 6) + for copied in (pickle.loads(pickle.dumps(slotted)), copy.copy(slotted), copy.deepcopy(slotted)): + assert copied.extra == 5 + assert vars(type(copied))["_SlottedAsyncArray__private"].__get__(copied) == 6 + reducing = _ReducingAsyncArray(created.metadata, created.store_path) + assert vars(copy.copy(reducing))["marked"] is True + + +def test_fields_read_as_missing_on_the_class() -> None: + """The fields have no class value, as fields without defaults have none.""" + assert all(getattr(AsyncArray, name) is MISSING for name in ("metadata", "config")) + + +class _StampingAsyncArray(AsyncArray[Any]): + """A subclass whose `__init__` changes the metadata it is given.""" + + def __init__(self, metadata: Any, store_path: StorePath, config: Any = None) -> None: + stamped = dataclasses.replace( + metadata, attributes={**metadata.attributes, "n": metadata.attributes.get("n", 0) + 1} + ) + super().__init__(stamped, store_path, config) + + +async def test_with_config_keeps_what_a_subclass_init_built() -> None: + """`with_config` rebuilds through the constructor and keeps the metadata a + subclass `__init__` built, as it always has, also over a model binding.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + for array in ( + _StampingAsyncArray(zarr.open_array(store, path="a").metadata, StorePath(store, "a")), + _StampingAsyncArray._from_io(await AsyncArrayIO.open(store, name="a")), + ): + before = cast("int", array.attrs.get("n", 0)) + assert array.with_config({"order": "F"}).attrs["n"] == before + 1 + + +async def test_a_pickled_array_carries_its_model_as_bytes() -> None: + """The model a pickled array carries is its documents' bytes, so a zarr without + zarr-metadata loads the pickle too.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + pickled = pickle.dumps(array) + assert b"zarr_metadata" not in pickled + restored = pickle.loads(pickled).io + assert restored == array.io + assert not isinstance(restored._source, ArrayV2Metadata | ArrayV3Metadata) + + +class _UnpicklableBinding(AsyncArrayIO): + """A binding deriving state that does not pickle, which it leaves out of its own.""" + + lock: Any + + def _bind(self, *args: Any) -> None: + super()._bind(*args) + self.lock = threading.Lock() + + def __getstate__(self) -> dict[str, object]: + state = super().__getstate__() + state.pop("lock") + return state + + +def test_a_binding_derives_its_state_again_when_unpickled() -> None: + binding = _UnpicklableBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + for restored in (pickle.loads(pickle.dumps(binding)), copy.deepcopy(binding)): + assert restored == binding + assert restored.lock is not binding.lock + + +class _NamedAsyncArray(AsyncArray[Any]): + """A subclass with attributes named as a binding's own state might be.""" + + +def test_subclass_attributes_keep_names_the_binding_uses() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + named = _NamedAsyncArray(created.metadata, created.store_path) + names = ("_context", "_pipeline_class", "_binding_class", "_model_documents", "_stored_entry") + for name in names: + object.__setattr__(named, name, f"mine {name}") + for restored in (pickle.loads(pickle.dumps(named)), copy.deepcopy(named)): + assert all(vars(restored)[name] == f"mine {name}" for name in names) + + +class _OwnReductionAsyncArray(AsyncArray[Any]): + """A subclass whose reduction carries a state of its own, with no `__setstate__`.""" + + def __reduce__(self) -> Any: + return (_own_reduction_rebuilt, (self.metadata, self.store_path), {"note": "hi"}) + + +def _own_reduction_rebuilt(metadata: Any, store_path: StorePath) -> _OwnReductionAsyncArray: + return _OwnReductionAsyncArray(metadata, store_path) + + +def test_a_subclass_reduction_with_its_own_state_copies_and_pickles() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + array = _OwnReductionAsyncArray(created.metadata, created.store_path) + for copied in (copy.copy(array), copy.deepcopy(array), pickle.loads(pickle.dumps(array))): + assert vars(copied)["note"] == "hi" + assert copied.shape == (4,) + + +class _SlotStateReductionAsyncArray(AsyncArray[Any]): + """A subclass whose reduction carries a state in the default protocol's + `(None, slots)` form.""" + + def __reduce__(self) -> Any: + return (_slot_state_rebuilt, (self.metadata, self.store_path), (None, {"note": "hi"})) + + +def _slot_state_rebuilt(metadata: Any, store_path: StorePath) -> _SlotStateReductionAsyncArray: + return _SlotStateReductionAsyncArray(metadata, store_path) + + +class _MetadataStateReductionAsyncArray(AsyncArray[Any]): + """A subclass whose own state holds some of an array's fields, which is not an + array's field state: it does not hold them all.""" + + def __reduce__(self) -> Any: + return ( + _metadata_state_rebuilt, + (self.metadata, self.store_path), + {"metadata": self.metadata, "codec_pipeline": self.codec_pipeline}, + ) + + +def _metadata_state_rebuilt( + metadata: Any, store_path: StorePath +) -> _MetadataStateReductionAsyncArray: + return _MetadataStateReductionAsyncArray(metadata, store_path) + + +def test_a_subclass_reduction_state_is_read_in_every_form() -> None: + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + slotted = _SlotStateReductionAsyncArray(created.metadata, created.store_path) + for copied in (copy.copy(slotted), copy.deepcopy(slotted), pickle.loads(pickle.dumps(slotted))): + assert vars(copied)["note"] == "hi" + holding = _MetadataStateReductionAsyncArray(created.metadata, created.store_path) + for copied in (copy.copy(holding), copy.deepcopy(holding), pickle.loads(pickle.dumps(holding))): + assert copied.shape == (4,) + + +class _AttributeNamedLikeTheCacheEntry(AsyncArrayIO): + """A binding subclass with an attribute of the name a carried cache entry had.""" + + _stored_entry: str + + +async def test_a_pickled_binding_keeps_its_attributes_apart_from_what_it_carries() -> None: + """What a binding carries in a pickle takes no name of a subclass's, and the + model derived from attributes changed in place is not pickled.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + binding = _AttributeNamedLikeTheCacheEntry.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + binding._stored_entry = "mine" + assert pickle.loads(pickle.dumps(binding))._stored_entry == "mine" + + array = (await AsyncArrayIO.open(store, name="a")).to_array() + array.metadata.attributes["changed"] = 1 + assert dict(array.io.metadata.attributes) == {"changed": 1} + assert array.io.__getstate__()["_attributes_cache"] is None + assert dict(pickle.loads(pickle.dumps(array.io)).metadata.attributes) == {"changed": 1} + + +async def test_attributes_changed_in_place_keep_their_order() -> None: + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + array.metadata.attributes.update({"zeta": 1, "alpha": {"y": 1, "b": 2}}) + attributes = array.io.metadata.attributes + assert list(attributes) == ["zeta", "alpha"] + assert list(cast("Mapping[str, Any]", attributes["alpha"])) == ["y", "b"] + + # a change of order alone is a change + stored = array.metadata.attributes + alpha = stored.pop("zeta") + stored["zeta"] = alpha + assert list(array.io.metadata.attributes) == ["alpha", "zeta"] + + +async def test_a_strict_v2_read_reports_a_zero_chunk_size_beside_others() -> None: + """A chunk size of 0, which the model accepts, is reported with the problems + the model reports for the other chunk sizes.""" + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "shape": [4, 4], "chunks": [-1, 0]}) + with pytest.raises(MetadataValidationError) as error: + await AsyncArrayIO.open(store, name="a") + assert ("chunks", 1) in [p.loc for p in error.value.problems] + + +async def test_a_pickle_whose_metadata_was_replaced_binds_the_new_metadata() -> None: + """A pickle that another zarr loaded, resized and pickled again (keeping the + binding's state as it was) binds the metadata it now holds, not the stale model.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,) + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + state = cast("dict[str, Any]", array.__getstate__()) + resized = dataclasses.replace(state["metadata"], shape=(8,)) + revived = object.__new__(AsyncArray) + revived.__setstate__({**state, "metadata": resized}) + assert revived.io.metadata.shape == (8,) + + +@pytest.mark.parametrize("chunks", [[False, 2], [0.0, 2]]) +async def test_a_strict_v2_read_reports_a_rejected_chunk_size_once(chunks: list[Any]) -> None: + store = MemoryStore() + await _store_v2_document(store, {**_V2_DOCUMENT, "shape": [4, 4], "chunks": chunks}) + with pytest.raises(MetadataValidationError) as error: + await AsyncArrayIO.open(store, name="a") + assert [p.loc for p in error.value.problems] == [("chunks",)] + + +async def test_a_model_binding_reads_attributes_with_mixed_key_types() -> None: + """Attributes changed in place to hold an int key read as JSON writes them.""" + store = MemoryStore() + await AsyncArrayIO.build_v3( + store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,), attributes={"b": 1} + ).store_metadata() + array = (await AsyncArrayIO.open(store, name="a")).to_array() + cast("dict[Any, Any]", array.metadata.attributes)[1] = "x" + assert dict(array.io.metadata.attributes) == {"b": 1, "1": "x"} + + +class _OwnPipeline(BatchedCodecPipeline): + """A pipeline class a backend binding chooses for itself.""" + + +class _OwnPipelineBinding(AsyncArrayIO): + """A binding whose `_bind` runs a pipeline class of its own, as a backend might.""" + + def _bind(self, *args: Any) -> None: + *rest, _ = args + bind: Any = super()._bind + bind(*rest, _OwnPipeline) + + +def test_with_config_runs_the_pipeline_a_binding_chose() -> None: + array = _OwnPipelineBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,) + ).to_array() + configured = array.with_config({"order": "F"}) + assert configured.io._pipeline_class is _OwnPipeline + assert isinstance(configured.codec_pipeline, _OwnPipeline) + + +class _AssignedPipeline(BatchedCodecPipeline): + """A pipeline class a subclass `__init__` assigns.""" + + +class _PipelineAssigningAsyncArray(AsyncArray[Any]): + """A subclass whose `__init__` assigns a pipeline of its own.""" + + def __init__(self, metadata: Any, store_path: StorePath, config: Any = None) -> None: + super().__init__(metadata, store_path, config) + pipeline = create_codec_pipeline(self.metadata, pipeline_class=_AssignedPipeline) + object.__setattr__(self, "codec_pipeline", pipeline) + + +def test_with_config_prefers_a_binding_pipeline_over_an_assigned_one() -> None: + """A binding that chose a pipeline class runs it, over a pipeline a subclass + `__init__` assigned; without such a binding, the assigned pipeline runs.""" + binding = _OwnPipelineBinding.build_v3( + MemoryStore(), shape=(4,), data_type="uint8", chunk_grid=(4,) + ) + chosen = _PipelineAssigningAsyncArray._from_io(binding).with_config({"order": "F"}) + assert isinstance(chosen.codec_pipeline, _OwnPipeline) + + created = zarr.create_array(MemoryStore(), shape=(4,), dtype="uint8", chunks=(2,)) + assigned = _PipelineAssigningAsyncArray(created.metadata, created.store_path) + assert isinstance(assigned.with_config({"order": "F"}).codec_pipeline, _AssignedPipeline) + + +class _ResolvedFirstElsewhere: + """Stands in for `parse_array_metadata` while another thread stores its own + resolution of `binding` first.""" + + def __init__(self, binding: AsyncArrayIO, elsewhere: ArrayMetadata) -> None: + self.binding = binding + self.elsewhere = elsewhere + + def __call__(self, *args: Any, **kwargs: Any) -> ArrayMetadata: + resolved = parse_array_metadata(*args, **kwargs) + self.binding._resolved_cache = self.elsewhere + return resolved + + +async def test_bindings_resolving_at_once_share_the_first_resolution() -> None: + """A binding over a model keeps the legacy metadata resolved first, so threads + resolving it at once share one object, whose attributes an array changes in place.""" + store = MemoryStore() + built = AsyncArrayIO.build_v3(store, name="a", shape=(4,), data_type="uint8", chunk_grid=(4,)) + await built.store_metadata() + binding = await AsyncArrayIO.open(store, name="a") + elsewhere = parse_array_metadata(built.metadata_legacy.to_dict(), "a") + with mock.patch( + "zarr.core.array.parse_array_metadata", _ResolvedFirstElsewhere(binding, elsewhere) + ): + assert binding.metadata_legacy is elsewhere + assert binding.metadata_legacy is elsewhere diff --git a/tests/test_array_stateful.py b/tests/test_array_stateful.py index ca9b1d3ad8..2bd69df822 100644 --- a/tests/test_array_stateful.py +++ b/tests/test_array_stateful.py @@ -37,6 +37,7 @@ import zarr from zarr.core.buffer import cpu, default_buffer_prototype from zarr.core.chunk_grids import ChunkGrid +from zarr.core.metadata.stored import stored_documents from zarr.core.sync import sync from zarr.errors import ZarrUserWarning from zarr.storage import MemoryStore @@ -157,7 +158,11 @@ def _open(self) -> zarr.Array[Any]: warned = any(issubclass(w.category, ZarrUserWarning) for w in record) must_warn = any(self.shape[axis] > 0 for axis in self.legacy_axes) assert warned is must_warn, [str(w.message) for w in record] - assert (arr.metadata._stored_document is not None) is bool(self.legacy_axes) + cache = stored_documents(arr.store_path.store) + needs_repair = ( + cache.needing_repair(arr.store_path.path, arr.metadata.zarr_format) is not None + ) + assert needs_repair is bool(self.legacy_axes) return arr # ----------------------------------------------------------------- model diff --git a/tests/test_metadata/test_repair.py b/tests/test_metadata/test_repair.py index 5247754533..84df6f6c4e 100644 --- a/tests/test_metadata/test_repair.py +++ b/tests/test_metadata/test_repair.py @@ -3,15 +3,20 @@ from __future__ import annotations import asyncio +import copy import dataclasses +import gc import json +import pickle import re import warnings import zipfile from typing import TYPE_CHECKING, Any, Literal, cast +import hypothesis.strategies as st import numpy as np import pytest +from hypothesis import given, settings import zarr from zarr.codecs import ShardingCodec @@ -22,8 +27,10 @@ from zarr.core.metadata.repair import ( RECREATE_HINT, RESAVE_HINT, + may_move_chunks, repair_array_document, ) +from zarr.core.metadata.stored import _CACHES, StoredDocuments, moves_chunks, stored_documents from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata from zarr.core.sync import sync from zarr.dtype import Int16 @@ -39,6 +46,13 @@ from zarr.types import AnyArray +def _store_needs_repair(array: AnyArray) -> bool: + """Whether the cache of the array's store records that its stored document needs a + repair that moves chunks.""" + cache = stored_documents(array.store_path.store) + return cache.needing_repair(array.store_path.path, array.metadata.zarr_format) is not None + + def _v2_doc(shape: list[int], chunks: list[Any]) -> dict[str, JSON]: return { "zarr_format": 2, @@ -203,7 +217,7 @@ def test_repair_array_document( warnings.simplefilter("always") metadata = metadata_cls.from_dict(dict(doc), path="group/array") assert _chunk_shapes(metadata) == expected - assert metadata._stored_document == (doc if repaired == "moved" else None) + assert moves_chunks(doc, cast("ZarrFormat", doc["zarr_format"])) is (repaired == "moved") messages = [str(w.message) for w in record] if warning is None: assert messages == [] @@ -314,14 +328,14 @@ def test_read_invalid_edges_in_rectilinear_grid( """A stored rectilinear chunk grid whose explicit edges or run-length encoded sizes are integral floats or JSON `true`, as zarr-python wrote them when given float or `True` edges, is read with those edges as the `int`s they equal, silently: zarr read - these edges so before. The reading moves no chunks, so the metadata is not marked - for re-saving.""" + these edges so before. The reading moves no chunks, so no repair is stored before + chunk writes.""" shape = [sum(edges) for edges in expected] doc = _rectilinear_doc(shape, chunk_shapes) with zarr.config.set({"array.rectilinear_chunks": True}): metadata = _read_strictly(doc) assert metadata.chunk_grid == RectilinearChunkGridMetadata(chunk_shapes=expected) - assert metadata._stored_document is None + assert not moves_chunks(doc, 3) @pytest.mark.parametrize( @@ -675,7 +689,7 @@ async def write_twice() -> None: sync(write_twice()) - assert arr.metadata._stored_document is None + assert not _store_needs_repair(arr) with warnings.catch_warnings(): warnings.simplefilter("error", ZarrUserWarning) reopened = zarr.open_array(store=path, mode="r") @@ -715,7 +729,7 @@ def test_stale_handle_write_keeps_newer_metadata( stale[0] = 9 - assert stale.metadata._stored_document is None + assert not _store_needs_repair(stale) reopened = _open_strictly(path) assert reopened.shape == (9,) assert reopened.attrs.asdict() == {"x": 1} @@ -843,7 +857,7 @@ def test_stale_handle_write_after_chunk_grid_change_raises( with pytest.raises(ValueError, match="has changed since this array was opened; reopen"): stale[0:3] = [7, 8, 9] - assert stale.metadata._stored_document is not None + assert _store_needs_repair(stale) assert {p.name: p.read_bytes() for p in path.iterdir()} == documents @@ -903,7 +917,7 @@ async def failing_set(self: LocalStore, key: str, *args: Any, **kwargs: Any) -> with pytest.raises(OSError, match=f"cannot store {re.escape(doc_name)}"): operation(arr) - assert arr.metadata._stored_document is not None + assert _store_needs_repair(arr) assert (path / doc_name).read_bytes() == stored @@ -916,11 +930,11 @@ def test_write_without_stored_document(zarr_format: Literal[2, 3]) -> None: doc = _v2_doc([3], [0]) if zarr_format == 2 else _v3_doc([3], [0]) with pytest.warns(ZarrUserWarning, match="is read as"): array = zarr.Array(AsyncArray.from_dict(StorePath(store), doc)) - repaired = array.metadata._stored_document is not None + repaired = moves_chunks(doc, zarr_format) array[:] = [1, 2, 3] - assert (repaired, array.metadata._stored_document) == (True, None) + assert (repaired, _store_needs_repair(array)) == (True, False) np.testing.assert_array_equal(array[:], [1, 2, 3]) assert not [key for key in store._store_dict if key.endswith((".zarray", "zarr.json"))] @@ -979,7 +993,7 @@ def test_from_dict_keeps_values_that_are_not_json(codec: Any) -> None: doc: dict[str, Any] = _v3_doc([4], [4]) doc["codecs"] = [codec] metadata = ArrayV3Metadata.from_dict(doc) - assert metadata._stored_document is None + assert not moves_chunks(doc, 3) assert isinstance(metadata.codecs[0], ShardingCodec) @@ -1219,20 +1233,26 @@ def test_consolidated_repaired_member_attributes_kept_by_group_write( @pytest.mark.parametrize("zarr_format", [2, 3]) -def test_repaired_metadata_keeps_the_document_it_read(zarr_format: Literal[2, 3]) -> None: - """Metadata read from a document that had to be repaired keeps that document as it - was read, whatever becomes of the caller's dict (or the objects in it, which the - metadata's attributes may hold): consolidated metadata stores it as read.""" +def test_cache_keeps_the_document_it_read(zarr_format: Literal[2, 3]) -> None: + """A store's cache keeps a stored document that needs a repair as it was read, + whatever becomes of the caller's dict (or the objects in it), and consolidated + metadata stores a member given as stored exactly so.""" doc = _v2_doc([0], [0]) if zarr_format == 2 else _v3_doc([0], [0]) doc["attributes"] = {"k": [1]} read = json.loads(json.dumps(doc)) metadata_cls = ArrayV2Metadata if zarr_format == 2 else ArrayV3Metadata - metadata = metadata_cls.from_dict(doc) + metadata = metadata_cls.from_dict(json.loads(json.dumps(doc))) + cache = StoredDocuments() + cache.observe("a", doc, zarr_format) - cast("list[int]", metadata.attributes["k"]).append(2) + cast("list[int]", cast("dict[str, Any]", doc["attributes"])["k"]).append(2) doc["shape"] = [4] - assert ConsolidatedMetadata(metadata={"a": metadata}).to_dict()["metadata"] == {"a": read} + stored = cache.needing_repair("a", zarr_format) + assert stored is not None + assert stored == read + consolidated = ConsolidatedMetadata(metadata={"a": metadata}) + assert consolidated.to_dict(members_as_stored={"a": stored})["metadata"] == {"a": read} @pytest.mark.parametrize("zarr_format", [2, 3]) @@ -1246,7 +1266,7 @@ def test_empty_write_stores_no_metadata(tmp_path: Path, zarr_format: Literal[2, array[0:0] = np.empty(0, dtype="int16") - assert array.metadata._stored_document is not None + assert _store_needs_repair(array) assert {p.name: p.read_bytes() for p in path.iterdir()} == documents @@ -1653,3 +1673,376 @@ def test_delete_member_without_flag_keeps_group(tmp_path: Path) -> None: assert group.metadata.consolidated_metadata is not None assert sorted(group.metadata.consolidated_metadata.metadata) == ["n", "r"] assert {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} == stored + + +def _bind_opened(store: LocalStore) -> AnyArray: + return zarr.open_array(store=store, mode="r+") + + +def _bind_pickled(store: LocalStore) -> AnyArray: + revived: AnyArray = pickle.loads(pickle.dumps(zarr.open_array(store=store, mode="r+"))) + return revived + + +def _bind_derived(store: LocalStore) -> AnyArray: + return zarr.open_array(store=store, mode="r+").io.with_config({"order": "F"}).to_array() + + +def _bind_declared_same_layout(store: LocalStore) -> AnyArray: + return zarr.ArrayIO.build_v2(store, shape=(3,), dtype="int16", chunks=(1,)).to_array() + + +def _bind_declared_other_layout(store: LocalStore) -> AnyArray: + return zarr.ArrayIO.build_v2(store, shape=(3,), dtype="int16", chunks=(3,)).to_array() + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize( + ("bind", "stored_after"), + [ + (_bind_opened, [1]), + (_bind_pickled, [1]), + (_bind_derived, [1]), + (_bind_declared_same_layout, [0]), + (_bind_declared_other_layout, [0]), + ], + ids=["opened", "pickled", "derived", "declared-same-layout", "declared-other-layout"], +) +def test_write_stores_repair_however_bound( + tmp_path: Path, bind: Callable[[LocalStore], AnyArray], stored_after: list[int] +) -> None: + """A chunk write stores the repair of a stored document that needs one when the + store's cache knows it: an array opened, pickled or derived from one that read it. + A declaration over a document no binding of its store read knows nothing to repair, + and writes its chunks as any array does, as does one that lays out chunks + otherwise.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, 2) + array = bind(LocalStore(path)) + + array[:] = [7, 8, 9] + + assert json.loads((path / ".zarray").read_text())["chunks"] == stored_after + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +def test_bindings_of_one_store_share_what_is_stored( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Arrays bound to one store object share its cache: once one has stored the repair, + another writes its chunks without reading a metadata document first.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, 2) + store = LocalStore(path) + first = zarr.open_array(store=store, mode="r+") + second = zarr.open_array(store=store, mode="r+") + first[:] = [7, 8, 9] + accesses = _record_store_access(monkeypatch) + + second[:] = [1, 2, 3] + + assert not [key for _, key in accesses if key.endswith((".zarray", ".zattrs"))] + np.testing.assert_array_equal(_open_strictly(path)[...], [1, 2, 3]) + + +def test_stored_documents_belongs_to_one_store_object() -> None: + """Each store object has its own cache, also when it equals another store, and the + cache is dropped with its store.""" + store, equal = MemoryStore(), MemoryStore() + assert store == equal + cache = stored_documents(store) + cache.observe("a", None, 3) + + assert stored_documents(store) is cache + assert not stored_documents(equal).knows("a", 3) + key = id(store) + del store, cache + gc.collect() + assert key not in _CACHES + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +def test_documents_of_each_format_are_recorded_apart(tmp_path: Path) -> None: + """A path may hold a document of each Zarr format (as `migrate_v3` leaves it before + `remove_metadata`): reading one does not stand for the other, so a v2 write still + stores the repair of its `.zarray` after the store object read the `zarr.json`.""" + path = tmp_path / "both.zarr" + _legacy_array(path, 2) + zarr.create_array(store=tmp_path / "v3.zarr", shape=(3,), chunks=(1,), dtype="int16") + (path / "zarr.json").write_bytes((tmp_path / "v3.zarr" / "zarr.json").read_bytes()) + store = LocalStore(path) + v2 = zarr.open_array(store=store, mode="r+", zarr_format=2) + zarr.open_array(store=store, mode="r+", zarr_format=3) + + v2[:] = [7, 8, 9] + + assert json.loads((path / ".zarray").read_text())["chunks"] == [1] + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +def test_unreadable_member_document_leaves_its_copy_as_stored(tmp_path: Path) -> None: + """A member document zarr fails to read (here a rectilinear grid while the flag is + off) is not recorded, so a group write keeps that member's consolidated copy as it + was stored instead of encoding the member anew.""" + path = tmp_path / "group.zarr" + zarr.open_group(path, mode="w") + zarr.create_array(store=path / "a", shape=(3, 6), chunks=(3, 3), dtype="int8") + _rewrite_doc( + path / "a", + 3, + lambda doc: doc["chunk_grid"]["configuration"].update(chunk_shape=[[1, 2], 3]), + ) + zarr.consolidate_metadata(path) + with zarr.config.set({"array.rectilinear_chunks": True}): + zarr.open_array(store=path / "a", mode="r+")[0, 0] = 1 + stored = json.dumps(_consolidated_member(path, 3, "a")) + store = LocalStore(path) + group = zarr.open_group(store, mode="r+", use_consolidated=True) + with pytest.raises(ValueError, match="experimental and disabled"): + list(zarr.open_group(store, mode="r+", use_consolidated=False).members()) + + group.attrs["y"] = 2 + + assert json.dumps(_consolidated_member(path, 3, "a")) == stored + + +def test_pickled_array_carries_what_its_store_knows( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A pickled array carries what its store's cache knows of its document, so the + unpickled array writes its chunks without reading a metadata document first, as + each task of a distributed write unpickles a store of its own.""" + path = tmp_path / "valid.zarr" + store = LocalStore(path) + zarr.create_array(store=store, shape=(3,), chunks=(1,), dtype="int16") + revived: AnyArray = pickle.loads(pickle.dumps(zarr.open_array(store=store, mode="r+"))) + accesses = _record_store_access(monkeypatch) + + revived[:] = [1, 2, 3] + + assert not [key for _, key in accesses if key.endswith("zarr.json")] + np.testing.assert_array_equal(_open_strictly(path)[...], [1, 2, 3]) + + +def test_restore_keeps_what_the_cache_learned_itself() -> None: + """A carried entry fills only a document the cache does not know: what the + cache recorded itself, such as a repair it stored over a broken document, + stays.""" + cache = StoredDocuments() + broken = _v2_doc([3], [0]) + cache.observe("a", broken, 2) + cache.observe("a", None, 2) + + cache.restore("a", 2, broken) + cache.restore("b", 2, broken) + + assert cache.needing_repair("a", 2) is None + assert cache.needing_repair("b", 2) == broken + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_pickled_consolidated_member_stores_its_repair( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A member read from consolidated metadata carries its copy's entry when + pickled, so the unpickled member's first chunk write stores the repair.""" + path = tmp_path / "group.zarr" + document = _consolidated_legacy_member(path, zarr_format) + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + member = group["a"] + assert isinstance(member, zarr.Array) + revived: AnyArray = pickle.loads(pickle.dumps(member)) + + revived[:] = [7, 8, 9] + + assert _stored_chunks(json.loads(document.read_text())) == [1] + + +def test_cache_records_only_what_stands_for_a_repair() -> None: + """A valid document is recorded only where it replaces one that needs a repair + (or a consolidated copy that does), so the cache does not grow with every + array it sees.""" + cache = StoredDocuments() + valid, broken = _v2_doc([3], [3]), _v2_doc([3], [0]) + for path in ("a", "b", "c"): + cache.observe(path, valid, 2) + cache.observe("b", broken, 2) + cache.observe("b", valid, 2) + cache.observe_consolidated("", {"c": broken}, 2) + cache.observe("c", valid, 2) + assert (cache.knows("a", 2), cache.knows("b", 2), cache.knows("c", 2)) == (False, True, True) + assert cache.needing_repair("b", 2) is None + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_pickled_group_keeps_consolidated_members_as_stored( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A pickled (or deep-copied) consolidated group carries what its store's cache + knew of its members, so a group write through the copy keeps a member read + from a repaired copy as it was stored, as the original would.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "a") + stored = json.dumps(_consolidated_member(path, zarr_format, "a")) + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + for revived in (pickle.loads(pickle.dumps(group)), copy.deepcopy(group)): + revived.attrs["x"] = 1 + assert json.dumps(_consolidated_member(path, zarr_format, "a")) == stored + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_array_from_a_stored_document_dict_stores_its_repair( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """An `AsyncArray` built from a stored document given as a dict stores that + document's repair before its first chunk write, as one opened from it does.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, zarr_format) + doc_name = ".zarray" if zarr_format == 2 else "zarr.json" + doc = json.loads((path / doc_name).read_text()) + if zarr_format == 2: + doc["attributes"] = {} + array = zarr.Array(AsyncArray(doc, StorePath(LocalStore(path)))) + + array[:] = [7, 8, 9] + + assert _stored_chunks(json.loads((path / doc_name).read_text())) == [1] + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_member_read_valid_then_from_a_broken_copy_keeps_the_copy( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A member whose own document was read valid, then read again from a broken + consolidated copy, is kept as that copy stored by a group write, as a member + read from such a copy always was.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "a") + stored = json.dumps(_consolidated_member(path, zarr_format, "a")) + store = LocalStore(path) + zarr.open_array(store=store, path="a", mode="r") + group = zarr.open_group(store, mode="r+", use_consolidated=True) + + group.attrs["x"] = 1 + + assert json.dumps(_consolidated_member(path, zarr_format, "a")) == stored + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("pickled", [False, True]) +def test_subgroup_write_keeps_nested_members_as_stored( + tmp_path: Path, zarr_format: Literal[2, 3], pickled: bool +) -> None: + """A subgroup read from a consolidated parent keeps a nested member read from a + repaired copy as it was stored when it writes its own documents, also once + pickled: the copy carries the members' entries under the subgroup's path.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "g/a") + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + sub = group["g"] + assert isinstance(sub, zarr.Group) + if pickled: + sub = pickle.loads(pickle.dumps(sub)) + + sub.attrs["x"] = 1 + + assert _stored_chunks(_consolidated_member(path / "g", zarr_format, "a")) == [0] + + +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_pickles_cross_versions_with_the_repair(tmp_path: Path, zarr_format: Literal[2, 3]) -> None: + """An array pickled by zarr 3.4.1 or 3.4.2 marks its metadata with the stored + document it was read from; this version reads that mark, and marks the metadata + it pickles the same way, so a write after unpickling stores the repair in + either direction.""" + path = tmp_path / "legacy.zarr" + _legacy_array(path, zarr_format) + doc_name = ".zarray" if zarr_format == 2 else "zarr.json" + stored = json.loads((path / doc_name).read_text()) + array = zarr.open_array(store=LocalStore(path), mode="r+") + + state = cast("dict[str, Any]", array.async_array.__getstate__()) + assert vars(state["metadata"])["_stored_document"] is not None + + old_state = { + key: state[key] for key in ("store_path", "codec_pipeline", "_chunk_grid", "config") + } + old_metadata = copy.copy(array.metadata) + object.__setattr__(old_metadata, "_stored_document", stored) + old_state["metadata"] = old_metadata + revived = object.__new__(type(array.async_array)) + revived.__setstate__( + pickle.loads(pickle.dumps({**old_state, "store_path": StorePath(LocalStore(path))})) + ) + sync(revived.setitem(slice(None), np.array([7, 8, 9], dtype="int16"))) + assert _stored_chunks(json.loads((path / doc_name).read_text())) == [1] + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning") +@pytest.mark.filterwarnings("ignore:.*is read as:zarr.errors.ZarrUserWarning") +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_group_pickles_cross_versions_with_member_marks( + tmp_path: Path, zarr_format: Literal[2, 3] +) -> None: + """A consolidated group pickled here marks its repaired members as zarr 3.4.1 + and 3.4.2 read them, and a group pickled by them, with marked members, keeps + those members as stored when it writes.""" + path = tmp_path / "group.zarr" + _flagged_consolidated_group(path, zarr_format, "a") + stored = _consolidated_member(path, zarr_format, "a") + group = zarr.open_group(LocalStore(path), mode="r+", use_consolidated=True) + + state = group._async_group.__getstate__() + member = state["metadata"].consolidated_metadata.metadata["a"] # type: ignore[attr-defined] + assert _stored_chunks(dict(vars(member)["_stored_document"])) == [0] + live = group.metadata.consolidated_metadata + assert live is not None + assert "_stored_document" not in vars(live.metadata["a"]) + + old_state = {k: v for k, v in state.items() if k != "_stored_members"} + old_state["store_path"] = StorePath(LocalStore(path)) + revived = object.__new__(type(group._async_group)) + revived.__setstate__(pickle.loads(pickle.dumps(old_state))) + sync(revived.update_attributes({"x": 1})) + assert json.dumps(_consolidated_member(path, zarr_format, "a")) == json.dumps(stored) + + +chunk_sizes = st.one_of( + st.integers(min_value=-2, max_value=4), + st.booleans(), + st.floats(allow_nan=False), + st.just([2, 2]), +) + + +@settings(deadline=None, max_examples=300) +@given( + zarr_format=st.sampled_from([2, 3]), + sizes=st.lists(chunk_sizes, max_size=3), + grid_name=st.sampled_from(["regular", "rectilinear", "other"]), +) +def test_may_move_chunks_agrees_with_the_repairs( + zarr_format: Literal[2, 3], sizes: list[Any], grid_name: str +) -> None: + """`may_move_chunks` answers no only where no repair moves chunks.""" + shape = [3] * len(sizes) + if zarr_format == 2: + doc: dict[str, Any] = {**_v2_doc(shape, sizes)} + else: + doc = {**_v3_doc(shape, sizes)} + if grid_name != "regular": + doc["chunk_grid"] = {"name": grid_name, "configuration": {"chunk_shapes": sizes}} + _, readings = repair_array_document(doc, zarr_format) + if any(reading.moves_chunks for reading in readings): + assert may_move_chunks(doc, zarr_format) diff --git a/uv.lock b/uv.lock index 3b0dbe3049..6064c46e87 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] @@ -3830,8 +3831,9 @@ requires-dist = [ { name = "obstore", marker = "extra == 'remote'", specifier = ">=0.5.1" }, { name = "packaging", specifier = ">=22.0" }, { name = "typer", marker = "extra == 'cli'" }, - { name = "typing-extensions", specifier = ">=4.14" }, + { name = "typing-extensions", specifier = ">=4.16" }, { name = "universal-pathlib", marker = "extra == 'optional'" }, + { name = "zarr-metadata", specifier = ">=0.6.0" }, ] 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" }, +]