fix(chunk-grids): read mixed regular/rectilinear chunk grids written by zarr 3.2.x

changes/4375.bugfix.md

@@ -0,0 +1,3 @@
+Arrays whose stored `regular` chunk grid mixes chunk sizes with lists of chunk edge lengths, such as `"chunk_shape": [2, [5, 10, 5]]` (written by zarr 3.2.0 and 3.2.1 for `chunks=(2, (5, 10, 5))`, and as `[2, [5.0, 10.0, 5.0]]` for float edges), can be read again, without enabling `array.rectilinear_chunks`. The grid is read as the rectilinear chunk grid it describes, with a `ZarrUserWarning`; re-saving the array's metadata (or writing chunks to it) stores that rectilinear grid, which requires the flag. A group's consolidated metadata keeps such an array's metadata as it was stored.
+
+The `array.rectilinear_chunks` flag now gates reading and storing array metadata that declares a rectilinear chunk grid, including in a group's consolidated metadata, instead of constructing `RectilinearChunkGridMetadata`. An operation that would store such metadata with the flag off (creating an array, `Array.resize`, deleting a member of a group with consolidated metadata, `create_hierarchy`) raises before it deletes or writes anything, and the error names the array, also when it is a member of a group's consolidated metadata.

src/zarr/core/array.py

@@ -120,10 +120,12 @@
 )
 from zarr.core.metadata.io import (
     ARRAY_DOCUMENTS,
+    encode_documents,
     parse_stored_array,
     read_documents,
     save_metadata,
     save_new_metadata,
+    store_documents,
     upsert_metadata,
 )
 from zarr.core.metadata.v2 import (
@@ -646,7 +648,6 @@ async def _create_v3(
             dimension_names=dimension_names,
             attributes=attributes,
         )
-
         array = cls(metadata=metadata, store_path=store_path, config=config)
         await save_new_metadata(store_path, metadata, overwrite=overwrite, ensure_parents=True)
         return array
@@ -727,7 +728,6 @@ async def _create_v2(
             compressor=compressor_parsed,
             attributes=attributes,
         )
-
         array = cls(metadata=metadata, store_path=store_path, config=config)
         await save_new_metadata(store_path, metadata, overwrite=overwrite, ensure_parents=True)
         return array
@@ -1607,7 +1607,9 @@ async def get_coordinate_selection(
 
     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`)."""
+        `_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()
 
@@ -1642,14 +1644,14 @@ async def _store_repaired_document(self) -> None:
         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)
+            current = parse_stored_array(documents, zarr_format, str(self.store_path))
         except ArrayNotFoundError:
             pass
         else:
             if _chunk_layout(current) != _chunk_layout(self.metadata):
                 raise ValueError(
-                    f"The metadata stored for the array at {str(self.store_path)!r} has "
-                    "changed since this array was opened: reopen the array to write to it."
+                    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)
@@ -5904,6 +5906,10 @@ async def _resize(
     # ensure deletion is only run if array is shrinking as the delete_outside_chunks path is unbounded in memory
     only_growing = all(new >= old for new, old in zip(new_shape, array.metadata.shape, strict=True))
 
+    # Encode the new metadata before deleting any chunk: metadata that cannot be stored
+    # then fails with the store untouched.
+    documents = encode_documents(array.store_path, new_metadata)
+
     if delete_outside_chunks and not only_growing:
         # Remove all chunks outside of the new shape
         old_chunk_coords = set(array._chunk_grid.all_chunk_coords())
@@ -5922,7 +5928,8 @@ async def _delete_key(key: str) -> None:
         )
 
     # Write new metadata
-    await array._save_metadata(new_metadata)
+    await store_documents(array.store_path, documents)
+    array._stored_document_replaced()
 
     # Update metadata and chunk_grid (in place)
     object.__setattr__(array, "metadata", new_metadata)

src/zarr/core/group.py

@@ -49,8 +49,13 @@
 from zarr.core.dtype import parse_data_type
 from zarr.core.json_parse import parse_field
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
-from zarr.core.metadata.io import save_metadata, save_new_metadata
-from zarr.core.metadata.v3 import AllowedExtraField, parse_extra_fields
+from zarr.core.metadata.io import (
+    encode_documents,
+    save_metadata,
+    save_new_metadata,
+    store_documents,
+)
+from zarr.core.metadata.v3 import AllowedExtraField, check_storable, parse_extra_fields
 from zarr.core.sync import SyncMixin, sync
 from zarr.errors import (
     ArrayNotFoundError,
@@ -380,6 +385,16 @@ class GroupMetadata(Metadata):
     extra_fields: dict[str, AllowedExtraField] = field(default_factory=dict)
 
     def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]:
+        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:
+                    try:
+                        check_storable(member)
+                    except ValueError as e:
+                        e.add_note(f"Array {path!r} in the consolidated metadata.")
+                        raise
         indent = config.get("json_indent")
         if self.zarr_format == 3:
             return {ZARR_JSON: json_to_buffer(self.to_dict(), prototype=prototype, indent=indent)}
@@ -834,11 +849,24 @@ async def delitem(self, key: str) -> None:
             Array or group name
         """
         store_path = self.store_path / key
-
+        consolidated = self.metadata.consolidated_metadata
+        if consolidated is None:
+            await store_path.delete_dir()
+            return
+        # Encode the group metadata without the member before deleting it: metadata that
+        # cannot be stored then fails with the store and this group untouched. What is
+        # stored is encoded after the deletion, from the metadata as it then is, so
+        # concurrent deletions each store the deletions made before them.
+        members = {name: node for name, node in consolidated.metadata.items() if name != key}
+        encode_documents(
+            self.store_path,
+            replace(self.metadata, consolidated_metadata=replace(consolidated, metadata=members)),
+        )
         await store_path.delete_dir()
-        if self.metadata.consolidated_metadata:
-            self.metadata.consolidated_metadata.metadata.pop(key, None)
-            await self._save_metadata()
+        # In place, so every handle sharing this consolidated metadata (a parent's or a
+        # subgroup's) sees the deletion.
+        consolidated.metadata.pop(key, None)
+        await store_documents(self.store_path, encode_documents(self.store_path, self.metadata))
 
     async def get[DefaultT](
         self, key: str, default: DefaultT | None = None
@@ -3307,11 +3335,11 @@ async def create_hierarchy(
             else:
                 nodes_explicit[k] = v
 
-    # Build every node before deleting or storing anything: metadata that no array or
-    # group can be built from then fails with the store untouched.
-    built = _build_nodes(store, nodes_explicit)
+    # Build and encode every node before deleting anything: a node that cannot be built
+    # or whose metadata cannot be stored then fails with the store untouched.
+    built, documents = _prepare_nodes(store, nodes_explicit)
     await asyncio.gather(*(store.delete_dir(key) for key in to_delete_keys))
-    async for key, node in _store_nodes(store, nodes_explicit, built):
+    async for key, node in _store_nodes(store, nodes_explicit, built, documents):
         yield key, node
 
 
@@ -3341,34 +3369,42 @@ async def create_nodes(
     AsyncGroup | AsyncArray
         The created nodes in the order they are created.
     """
-    async for key, node in _store_nodes(store, nodes, _build_nodes(store, nodes)):
+    async for key, node in _store_nodes(store, nodes, *_prepare_nodes(store, nodes)):
         yield key, node
 
 
-def _build_nodes(
+def _prepare_nodes(
     store: Store, nodes: Mapping[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata]
-) -> dict[str, AsyncGroup | AnyAsyncArray]:
-    """The array or group each of `nodes` describes, at its path in `store`."""
-    return {
+) -> tuple[dict[str, AsyncGroup | AnyAsyncArray], dict[str, Buffer]]:
+    """The array or group each of `nodes` describes, at its path in `store`, and the
+    metadata documents of `nodes`, by their keys in the store: every node is built and
+    encoded before anything is stored."""
+    built = {
         path: _build_node(store=store, path=path, metadata=meta) for path, meta in nodes.items()
     }
+    documents = {
+        _join_paths([path, key]): value
+        for path, metadata in nodes.items()
+        for key, value in encode_documents(StorePath(store, path), metadata).items()
+    }
+    return built, documents
 
 
 async def _store_nodes(
     store: Store,
     nodes: Mapping[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata],
     built: Mapping[str, AsyncGroup | AnyAsyncArray],
+    documents: Mapping[str, Buffer],
 ) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]:
-    """Store the metadata of `nodes` and yield the nodes `_build_nodes` built from them
-    (see `create_nodes`)."""
+    """Store the `documents` encoded from `nodes` and yield the nodes built from them, as
+    `_prepare_nodes` returns them (see `create_nodes`)."""
     # Note: the only way to alter this value is via the config. If that's undesirable for some reason,
     # then we should consider adding a keyword argument to this function
     semaphore = asyncio.Semaphore(config.get("async.concurrency"))
-    create_tasks: list[Coroutine[None, None, str]] = []
-
-    for key, value in nodes.items():
-        # make the key absolute
-        create_tasks.extend(_persist_metadata(store, key, value, semaphore=semaphore))
+    create_tasks = [
+        _set_return_key(store=store, key=key, value=value, semaphore=semaphore)
+        for key, value in documents.items()
+    ]
 
     created_object_keys = []
 
@@ -3881,23 +3917,6 @@ async def _set_return_key(
     return key
 
 
-def _persist_metadata(
-    store: Store,
-    path: str,
-    metadata: ArrayV2Metadata | ArrayV3Metadata | GroupMetadata,
-    semaphore: asyncio.Semaphore | None = None,
-) -> tuple[Coroutine[None, None, str], ...]:
-    """
-    Prepare to save a metadata document to storage, returning a tuple of coroutines that must be awaited.
-    """
-
-    to_save = metadata.to_buffer_dict(default_buffer_prototype())
-    return tuple(
-        _set_return_key(store=store, key=_join_paths([path, key]), value=value, semaphore=semaphore)
-        for key, value in to_save.items()
-    )
-
-
 async def create_rooted_hierarchy(
     *,
     store: Store,

src/zarr/core/metadata/io.py

@@ -70,8 +70,28 @@ def _diff(
             yield DocumentChange(path, stored, new)
 
 
+def encode_documents(
+    store_path: StorePath, metadata: ArrayMetadata | GroupMetadata
+) -> dict[str, Buffer]:
+    """The metadata documents `metadata` stores under `store_path`, by key (see
+    `to_buffer_dict`).
+
+    An operation that deletes or writes store content encodes its documents first, so
+    metadata that cannot be stored fails with the store untouched; the error then names
+    the node at `store_path`, as the warnings about stored documents do.
+    """
+    from zarr.core.group import GroupMetadata
+
+    try:
+        return metadata.to_buffer_dict(default_buffer_prototype())
+    except ValueError as e:
+        node = "Group" if isinstance(metadata, GroupMetadata) else "Array"
+        e.add_note(f"{node} {str(store_path)!r}: nothing was stored.")
+        raise
+
+
 async def store_documents(store_path: StorePath, documents: Mapping[str, Buffer]) -> None:
-    """Store metadata documents encoded by `to_buffer_dict` under `store_path`."""
+    """Store metadata documents encoded by `encode_documents` under `store_path`."""
     await asyncio.gather(
         *(set_or_delete(store_path / key, value) for key, value in documents.items())
     )
@@ -97,7 +117,7 @@ async def upsert_metadata(
     The documents are encoded before any is stored, so metadata that cannot be stored
     fails with the store untouched.
     """
-    documents = metadata.to_buffer_dict(default_buffer_prototype())
+    documents = encode_documents(store_path, metadata)
     changes = diff_documents(
         {key: buffer_to_json_object(buf) for key, buf in stored.items() if key in documents},
         {key: buffer_to_json_object(buf) for key, buf in documents.items()},
@@ -114,11 +134,17 @@ async def upsert_metadata(
 """The store keys of the metadata documents of an array of each Zarr format."""
 
 
-def parse_stored_array(documents: Mapping[str, Buffer], zarr_format: ZarrFormat) -> ArrayMetadata:
+def parse_stored_array(
+    documents: Mapping[str, Buffer], zarr_format: ZarrFormat, path: str | None = None
+) -> 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."""
+    repaired. 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
+    error names the array at `path`."""
     from zarr.core.array import (
         _array_metadata_dict_v2,
         _array_metadata_dict_v3,
@@ -132,7 +158,8 @@ def parse_stored_array(documents: Mapping[str, Buffer], zarr_format: ZarrFormat)
     else:
         raise ArrayNotFoundError(f"No Zarr format {zarr_format} array metadata document.")
     repaired, readings = repair_array_document(stored, zarr_format)
-    return mark_repaired(parse_array_metadata(dict(repaired)), stored, readings, None, warn=False)
+    metadata = parse_array_metadata(dict(repaired), path)
+    return mark_repaired(metadata, stored, readings, None, warn=False)
 
 
 def _build_parents(store_path: StorePath, zarr_format: ZarrFormat) -> dict[str, GroupMetadata]:
@@ -172,7 +199,7 @@ async def save_metadata(
     ------
     ValueError
     """
-    to_save = metadata.to_buffer_dict(default_buffer_prototype())
+    to_save = encode_documents(store_path, metadata)
     await _write_metadata(store_path, to_save, metadata.zarr_format, ensure_parents=ensure_parents)
 
 
@@ -185,8 +212,8 @@ async def save_new_metadata(
 ) -> None:
     """Save the metadata of a new array or group, replacing any existing node if requested.
 
-    The metadata is encoded before the store is modified, so metadata that cannot be
-    encoded raises without deleting an existing node.
+    The metadata is encoded first (see `encode_documents`), so metadata that cannot be
+    stored raises, naming the node, without deleting an existing node.
 
     Parameters
     ----------
@@ -201,7 +228,7 @@ async def save_new_metadata(
     ensure_parents : bool, optional
         Create any missing parent groups, and check no existing parents are arrays.
     """
-    to_save = metadata.to_buffer_dict(default_buffer_prototype())
+    to_save = encode_documents(store_path, metadata)
     if overwrite and store_path.store.supports_deletes:
         await store_path.delete_dir()
     else: