Zarr-v3 Consolidated Metadata

src/zarr/api/asynchronous.py

@@ -1,15 +1,24 @@
 from __future__ import annotations
 
 import asyncio
+import dataclasses
 import warnings
 from typing import TYPE_CHECKING, Any, Literal, Union, cast
 
 import numpy as np
 import numpy.typing as npt
 
 from zarr.core.array import Array, AsyncArray
-from zarr.core.common import JSON, AccessModeLiteral, ChunkCoords, MemoryOrder, ZarrFormat
-from zarr.core.group import AsyncGroup
+from zarr.core.buffer import NDArrayLike
+from zarr.core.chunk_key_encodings import ChunkKeyEncoding
+from zarr.core.common import (
+    JSON,
+    AccessModeLiteral,
+    ChunkCoords,
+    MemoryOrder,
+    ZarrFormat,
+)
+from zarr.core.group import AsyncGroup, ConsolidatedMetadata
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
 from zarr.store import (
     StoreLike,
@@ -129,8 +138,38 @@ def _default_zarr_version() -> ZarrFormat:
     return 3
 
 
-async def consolidate_metadata(*args: Any, **kwargs: Any) -> AsyncGroup:
-    raise NotImplementedError
+async def consolidate_metadata(store: StoreLike) -> AsyncGroup:
+    """
+    Consolidate the metadata of all nodes in a hierarchy.
+
+    Upon completion, the metadata of the root node in the Zarr hierarchy will be
+    updated to include all the metadata of child nodes.
+
+    Parameters
+    ----------
+    store: StoreLike
+        The store-like object whose metadata you wish to consolidate.
+
+    Returns
+    -------
+    group: AsyncGroup
+        The group, with the ``consolidated_metadata`` field set to include
+        the metadata of each child node.
+    """
+    group = await AsyncGroup.open(store)
+    members = dict([x async for x in group.members(max_depth=None)])
+    members_metadata = {}
+
+    members_metadata = {k: v.metadata for k, v in members.items()}
+
+    consolidated_metadata = ConsolidatedMetadata(metadata=members_metadata)
+    metadata = dataclasses.replace(group.metadata, consolidated_metadata=consolidated_metadata)
+    group = dataclasses.replace(
+        group,
+        metadata=metadata,
+    )
+    await group._save_metadata()
+    return group
 
 
 async def copy(*args: Any, **kwargs: Any) -> tuple[int, int, int]:
@@ -232,7 +271,8 @@ async def open(
 
 
 async def open_consolidated(*args: Any, **kwargs: Any) -> AsyncGroup:
-    raise NotImplementedError
+    kwargs.setdefault("open_consolidated", True)
+    return await open_group(*args, **kwargs)
 
 
 async def save(
@@ -492,6 +532,7 @@ async def open_group(
     zarr_format: ZarrFormat | None = None,
     meta_array: Any | None = None,  # not used
     attributes: dict[str, JSON] | None = None,
+    open_consolidated: bool = False,
 ) -> AsyncGroup:
     """Open a group using file-mode-like semantics.
 
@@ -551,7 +592,9 @@ async def open_group(
         attributes = {}
 
     try:
-        return await AsyncGroup.open(store_path, zarr_format=zarr_format)
+        return await AsyncGroup.open(
+            store_path, zarr_format=zarr_format, open_consolidated=open_consolidated
+        )
     except (KeyError, FileNotFoundError):
         return await AsyncGroup.create(
             store_path, zarr_format=zarr_format, exists_ok=True, attributes=attributes
@@ -706,7 +749,9 @@ async def create(
             )
         else:
             warnings.warn(
-                "dimension_separator is not yet implemented", RuntimeWarning, stacklevel=2
+                "dimension_separator is not yet implemented",
+                RuntimeWarning,
+                stacklevel=2,
             )
     if write_empty_chunks:
         warnings.warn("write_empty_chunks is not yet implemented", RuntimeWarning, stacklevel=2)

src/zarr/core/array.py

@@ -386,9 +386,15 @@ async def open(
         else:
             # V3 arrays are comprised of a zarr.json object
             assert zarr_json_bytes is not None
+            zarr_metadata = json.loads(zarr_json_bytes.to_bytes())
+            if zarr_metadata.get("node_type") != "array":
+                # This KeyError is load bearing for `open`. That currently tries
+                # to open the node as an `array` and then falls back to opening
+                # as a group.
+                raise KeyError
             return cls(
                 store_path=store_path,
-                metadata=ArrayV3Metadata.from_dict(json.loads(zarr_json_bytes.to_bytes())),
+                metadata=ArrayV3Metadata.from_dict(zarr_metadata),
             )
 
     @property

src/zarr/core/common.py

@@ -16,6 +16,8 @@
     overload,
 )
 
+import numcodecs
+
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable, Iterator
 
@@ -26,6 +28,7 @@
 ZARRAY_JSON = ".zarray"
 ZGROUP_JSON = ".zgroup"
 ZATTRS_JSON = ".zattrs"
+ZMETADATA_v2_JSON = ".zmetadata"
 
 BytesLike = bytes | bytearray | memoryview
 ShapeLike = tuple[int, ...] | int
@@ -168,3 +171,31 @@ def parse_order(data: Any) -> Literal["C", "F"]:
     if data in ("C", "F"):
         return cast(Literal["C", "F"], data)
     raise ValueError(f"Expected one of ('C', 'F'), got {data} instead.")
+
+
+def _json_convert(o: Any) -> Any:
+    if isinstance(o, np.dtype):
+        return str(o)
+    if np.isscalar(o):
+        # convert numpy scalar to python type, and pass
+        # python types through
+        if hasattr(o, "dtype") and o.dtype.kind == "M" and hasattr(o, "view"):
+            # https://github.com/zarr-developers/zarr-python/issues/2119
+            # `.item()` on a datetime type might or might not return an
+            # integer, depending on the value.
+            # Explicitly cast to an int first, and then grab .item()
+            out = o.view("i8").item()
+        else:
+            out = getattr(o, "item", lambda: o)()
+        if isinstance(out, complex):
+            # python complex types are not JSON serializable, so we use the
+            # serialization defined in the zarr v3 spec
+            return [out.real, out.imag]
+        return out
+    if isinstance(o, Enum):
+        return o.name
+    # this serializes numcodecs compressors
+    # todo: implement to_dict for codecs
+    elif isinstance(o, numcodecs.abc.Codec):
+        config: dict[str, Any] = o.get_config()
+        return config