fix(array): clear the stored document whenever an array stores its metadata

Setting attributes of an array read from a document that had to be
upgraded stored the upgraded document with the new attributes, but left
the metadata marked with the document it was read from. A consolidated
group handle shares that metadata, so its next write stored the old
document in the consolidated metadata and the new attributes were lost
from it (zarr 3.4.0 kept them).

Every write of an array's own documents (creating, resizing, setting
attributes) now goes through `AsyncArray._save_metadata`, which clears
the mark afterwards, as the first chunk write does after storing the
upgrade. Both clear it through `_stored_document_replaced`, the one
place that declares it: the first chunk write clears it also when it
stores nothing (the store already holds a valid document, or none), so
it cannot be folded into the save.

The deep copy in `mark_upgraded` stays: the metadata's attributes share
objects with the document it was read from, which the caller also
holds. A test pins it.

Assisted-by: ClaudeCode:claude-opus-5-5
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: d-v-b

changes/4334.bugfix.md

@@ -4,4 +4,4 @@ A chunk edge length is now always at least 1, while an array extent may be 0. Ev
 
 Stored metadata with a regular chunk size of 0 or JSON `false` (Zarr format 2 `chunks`, Zarr format 3 `regular` `chunk_shape`) is now read as one chunk spanning the axis (a multiple of the inner chunk size for sharded arrays, which previously failed to open). zarr-python wrote such sizes for arrays created with a zero-length axis until 3.4, and, from 2.18.7 to 3.2.1, for an explicit chunk size of 0 or `False` on an axis of any length, as in `chunks=(0,)` (those arrays could store no data). On a zero-length axis this opens silently, as before; appending to such an axis then stores chunks of size 1. On an axis of positive length it opens with a `ZarrUserWarning`, which says that the axis holds only the fill value and how to store valid metadata: `array.update_attributes({})`, then `zarr.consolidate_metadata` if the metadata is consolidated. JSON `true`, as zarr-python 3.0 and 3.2 wrote for a chunk size of `True`, is read as 1, silently, in the chunk grid (regular, and the explicit edges and run-length encoded sizes of a rectilinear one) and in the inner chunk shape of every sharding codec, nested or not. A stored rectilinear chunk grid whose edge lengths are integral JSON floats (`[[4.0, 2]]`), as zarr-python 3.2 wrote for float edges, is read with those edges as integers, silently; a float anywhere no release wrote one (a regular chunk shape, the inner chunk shape of a sharding codec, a run-length repeat count, an edge below 1) is rejected.
 
-Writing data to an array read from such metadata (other than an empty selection) first stores the upgrade of the metadata the store then holds, unless another writer has stored valid metadata since, so that other readers find the chunks written; in a `ZipStore` this adds a second entry for the metadata document, as every metadata update does. If the metadata the store then holds lays out chunks differently (another writer resized the array keeping the chunk size of 0, say), the write raises a `ValueError` asking to reopen the array, and stores nothing. This is the only operation that stores the upgrade: a group's consolidated metadata keeps such an array's metadata as it was stored (`zarr.consolidate_metadata` copies it as the array's own document holds it), and changing a group reads and writes no metadata of its members.
+Writing data to an array read from such metadata (other than an empty selection) first stores the upgrade of the metadata the store then holds, unless another writer has stored valid metadata since, so that other readers find the chunks written; in a `ZipStore` this adds a second entry for the metadata document, as every metadata update does. If the metadata the store then holds lays out chunks differently (another writer resized the array keeping the chunk size of 0, say), the write raises a `ValueError` asking to reopen the array, and stores nothing. Apart from the array's own metadata writes (`update_attributes`, `resize`), which store the upgrade as before, this is the only operation that stores it: a group's consolidated metadata keeps such an array's metadata as it was stored (`zarr.consolidate_metadata` copies it as the array's own document holds it) until the array stores its upgrade through the group's handle, and changing a group reads and writes no metadata of its members.

src/zarr/core/array.py

@@ -1625,10 +1625,18 @@ async def get_coordinate_selection(
         return out_array
 
     async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None:
-        """
-        Asynchronously save the array metadata.
-        """
+        """Store `metadata` as this array's own documents: every write of them
+        (creating, resizing, setting attributes) goes through here."""
         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
+        upgrade: it holds the upgrade, 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_upgraded`), so no later write through either
+        handle stores that document again."""
+        object.__setattr__(self.metadata, "_stored_document", None)
 
     async def _store_upgraded_document(self) -> None:
         """Store the upgrade of this array's current stored document, if it needs one,
@@ -1660,7 +1668,7 @@ async def _store_upgraded_document(self) -> None:
                 )
             if current._stored_document is not None:
                 await upsert_metadata(self.store_path, current, documents)
-        object.__setattr__(self.metadata, "_stored_document", None)
+        self._stored_document_replaced()
 
     async def _set_selection(
         self,
@@ -5931,7 +5939,7 @@ async def _delete_key(key: str) -> None:
         )
 
     # Write new metadata
-    await save_metadata(array.store_path, new_metadata)
+    await array._save_metadata(new_metadata)
 
     # Update metadata and chunk_grid (in place)
     object.__setattr__(array, "metadata", new_metadata)
@@ -6023,7 +6031,7 @@ async def _update_attributes(
     array.metadata.attributes.update(new_attributes)
 
     # Write new metadata
-    await save_metadata(array.store_path, array.metadata)
+    await array._save_metadata(array.metadata)
 
     return array
 

tests/test_metadata/test_upgrades.py

@@ -16,6 +16,7 @@
 from zarr.codecs import ShardingCodec
 from zarr.codecs.numcodecs import Quantize
 from zarr.core.array import AsyncArray
+from zarr.core.group import ConsolidatedMetadata
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
 from zarr.core.metadata.upgrades import (
     RESAVE_HINT,
@@ -1007,6 +1008,48 @@ def test_consolidated_upgraded_member_write_after_chunk_grid_change_raises(
     assert {p: p.read_bytes() for p in path.rglob("*") if p.is_file()} == documents
 
 
+@pytest.mark.filterwarnings("ignore:Consolidated metadata is currently not part:UserWarning")
+@pytest.mark.parametrize("zarr_format", [2, 3])
+def test_consolidated_upgraded_member_attributes_kept_by_group_write(
+    tmp_path: Path, zarr_format: Literal[2, 3]
+) -> None:
+    """Setting attributes of a member read from consolidated metadata stores the member's
+    upgraded document with them. A later write of the group, whose consolidated metadata
+    shares the member's metadata, then stores that upgrade: the new attributes, not the
+    legacy document as it was stored before."""
+    path = tmp_path / "group.zarr"
+    _consolidated_legacy_member(path, zarr_format)
+    with pytest.warns(ZarrUserWarning, match="is read as"):
+        group = zarr.open_group(path, mode="r+", use_consolidated=True)
+
+    group["a"].attrs["x"] = 1
+    group.attrs["y"] = 2
+
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        attributes = dict(zarr.open_group(path, mode="r", use_consolidated=True)["a"].attrs)
+    assert attributes == {"x": 1}
+    assert _stored_chunks(_consolidated_member(path, zarr_format, "a")) == [3]
+    assert not [w for w in record if "is read as" in str(w.message)]
+
+
+@pytest.mark.parametrize("zarr_format", [2, 3])
+def test_upgraded_metadata_keeps_the_document_it_read(zarr_format: Literal[2, 3]) -> None:
+    """Metadata read from a document that had to be upgraded 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."""
+    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)
+
+    cast("list[int]", metadata.attributes["k"]).append(2)
+    doc["shape"] = [4]
+
+    assert ConsolidatedMetadata(metadata={"a": metadata}).to_dict()["metadata"] == {"a": read}
+
+
 @pytest.mark.parametrize("zarr_format", [2, 3])
 def test_empty_write_stores_no_metadata(tmp_path: Path, zarr_format: Literal[2, 3]) -> None:
     """A write of an empty selection stores no chunks, so it stores no metadata either."""