Avoid duplicate metadata fetches when zarr.open() falls back from array to group

src/zarr/api/asynchronous.py

@@ -13,11 +13,12 @@
 from zarr.core.array import (
     DEFAULT_FILL_VALUE,
     Array,
+    ArrayMetadataProbe,
     AsyncArray,
     CompressorLike,
+    _probe_array_metadata,
     create_array,
     from_array,
-    get_array_metadata,
 )
 from zarr.core.array_spec import ArrayConfigLike, parse_array_config
 from zarr.core.buffer import NDArrayLike
@@ -35,6 +36,7 @@
     AsyncGroup,
     ConsolidatedMetadata,
     GroupMetadata,
+    _PreFetchedGroupMetadata,
     create_hierarchy,
 )
 from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata
@@ -108,6 +110,24 @@ def _infer_overwrite(mode: AccessModeLiteral) -> bool:
     return mode in _OVERWRITE_MODES
 
 
+def _array_probe_to_group_prefetch(
+    probe: ArrayMetadataProbe,
+) -> _PreFetchedGroupMetadata | None:
+    """Translate array-probe buffers into group-open pre-fetched metadata.
+
+    The ``zarr.json`` and ``.zattrs`` documents fetched while probing for an array
+    are the same documents ``AsyncGroup.open`` would otherwise fetch, so they can
+    be reused to avoid duplicate requests. Returns None if nothing reusable was
+    fetched (e.g. an explicit ``zarr_format`` probe, which does not overlap).
+    """
+    pre_fetched: _PreFetchedGroupMetadata = {}
+    if "zarr_json" in probe:
+        pre_fetched["zarr_json"] = probe["zarr_json"]
+    if "zattrs" in probe:
+        pre_fetched["zattrs"] = probe["zattrs"]
+    return pre_fetched or None
+
+
 def _get_shape_chunks(a: ArrayLike | Any) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]:
     """Helper function to get the shape and chunks from an array-like object"""
     shape = None
@@ -358,20 +378,33 @@ async def open(
 
     # TODO: the mode check below seems wrong!
     if "shape" not in kwargs and mode in {"a", "r", "r+", "w"}:
-        try:
-            metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format)
+        # Probe for an array once. The probe also returns the zarr.json/.zattrs
+        # buffers it fetched so the group open below can reuse them instead of
+        # re-fetching (only populated for zarr_format=None).
+        result = await _probe_array_metadata(store_path, zarr_format=zarr_format)
+        if result.metadata is not None:
             # TODO: remove this cast when we fix typing for array metadata dicts
-            _metadata_dict = cast("ArrayMetadataDict", metadata_dict)
-            # for v2, the above would already have raised an exception if not an array
+            _metadata_dict = cast("ArrayMetadataDict", result.metadata)
             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=kwargs.get("config")
                 )
-        except (AssertionError, FileNotFoundError, NodeTypeValidationError):
-            pass
-        return await open_group(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs)
+        elif not isinstance(result.error, (FileNotFoundError, NodeTypeValidationError)):
+            # Only "not an array" errors trigger the group fallback; anything else
+            # (e.g. an invalid zarr_format) is propagated, as it was before.
+            assert result.error is not None
+            raise result.error
+        # Fall back to opening a group, reusing the buffers already fetched during
+        # the array probe (zarr.json, .zattrs) so they are not fetched again.
+        return await open_group(
+            store=store_path,
+            zarr_format=zarr_format,
+            mode=mode,
+            _pre_fetched_metadata=_array_probe_to_group_prefetch(result.probe),
+            **kwargs,
+        )
 
     try:
         return await open_array(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs)
@@ -756,6 +789,7 @@ async def open_group(
     meta_array: Any | None = None,  # not used
     attributes: dict[str, JSON] | None = None,
     use_consolidated: bool | str | None = None,
+    _pre_fetched_metadata: _PreFetchedGroupMetadata | None = None,
 ) -> AsyncGroup:
     """Open a group using file-mode-like semantics.
 
@@ -806,6 +840,10 @@ async def open_group(
         Zarr format 2 allowed configuring the key storing the consolidated metadata
         (``.zmetadata`` by default). Specify the custom key as ``use_consolidated``
         to load consolidated metadata from a non-default key.
+    _pre_fetched_metadata : _PreFetchedGroupMetadata or None, default None
+        Private. Already-fetched metadata buffers that ``AsyncGroup.open`` may
+        reuse instead of re-fetching. Used by ``zarr.open`` to avoid duplicate
+        fetches when falling back from an array probe to a group open.
 
     Returns
     -------
@@ -829,7 +867,10 @@ async def open_group(
     try:
         if mode in _READ_MODES:
             return await AsyncGroup.open(
-                store_path, zarr_format=zarr_format, use_consolidated=use_consolidated
+                store_path,
+                zarr_format=zarr_format,
+                use_consolidated=use_consolidated,
+                _pre_fetched_metadata=_pre_fetched_metadata,
             )
     except (KeyError, FileNotFoundError):
         pass

src/zarr/api/synchronous.py

@@ -40,6 +40,7 @@
         ZarrFormat,
     )
     from zarr.core.dtype import ZDTypeLike
+    from zarr.core.group import _PreFetchedGroupMetadata
     from zarr.storage import StoreLike
     from zarr.types import AnyArray
 
@@ -473,6 +474,7 @@ def open_group(
     meta_array: Any | None = None,  # not used in async api
     attributes: dict[str, JSON] | None = None,
     use_consolidated: bool | str | None = None,
+    _pre_fetched_metadata: _PreFetchedGroupMetadata | None = None,
 ) -> Group:
     """Open a group using file-mode-like semantics.
 
@@ -523,6 +525,10 @@ def open_group(
         Zarr format 2 allowed configuring the key storing the consolidated metadata
         (``.zmetadata`` by default). Specify the custom key as ``use_consolidated``
         to load consolidated metadata from a non-default key.
+    _pre_fetched_metadata : _PreFetchedGroupMetadata or None, default None
+        Private. Already-fetched metadata buffers that ``AsyncGroup.open`` may
+        reuse instead of re-fetching. Used by ``zarr.open`` to avoid duplicate
+        fetches when falling back from an array probe to a group open.
 
     Returns
     -------
@@ -543,6 +549,7 @@ def open_group(
                 meta_array=meta_array,
                 attributes=attributes,
                 use_consolidated=use_consolidated,
+                _pre_fetched_metadata=_pre_fetched_metadata,
             )
         )
     )

src/zarr/core/array.py

@@ -10,6 +10,7 @@
     TYPE_CHECKING,
     Any,
     Literal,
+    NotRequired,
     TypedDict,
     cast,
     overload,
@@ -133,6 +134,7 @@
     ArrayNotFoundError,
     ChunkNotFoundError,
     MetadataValidationError,
+    NodeTypeValidationError,
     ZarrDeprecationWarning,
     ZarrUserWarning,
 )
@@ -154,6 +156,7 @@
     from zarr.abc.codec import CodecPipeline
     from zarr.abc.store import Store
     from zarr.codecs.sharding import ShardingCodecIndexLocation
+    from zarr.core.buffer import Buffer
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
     from zarr.storage import StoreLike
     from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3
@@ -234,9 +237,48 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None
     raise TypeError  # pragma: no cover
 
 
-async def get_array_metadata(
+class ArrayMetadataProbe(TypedDict):
+    """Buffers fetched while probing for array metadata.
+
+    Private. Carried on ``ArrayProbeResult`` so callers (e.g. ``zarr.open``) can
+    reuse the ``zarr.json`` and ``.zattrs`` documents when falling back to a group
+    open, instead of fetching them a second time. A key is present only if that
+    document was actually fetched (a ``None`` value means fetched-but-absent). Only
+    the ``zarr_format=None`` probe overlaps with a group probe, so only that case
+    populates these keys.
+    """
+
+    zarr_json: NotRequired[Buffer | None]
+    zattrs: NotRequired[Buffer | None]
+
+
+@dataclass(frozen=True, kw_only=True)
+class ArrayProbeResult:
+    """Outcome of probing a path for array metadata.
+
+    Private. Either ``metadata`` is set (the path holds an array) or ``error`` is
+    set (it does not, or the request was invalid). ``probe`` carries any reusable
+    buffers fetched during a ``zarr_format=None`` probe so a caller can hand them
+    to ``AsyncGroup.open`` without re-fetching.
+    """
+
+    metadata: dict[str, JSON] | None = None
+    error: Exception | None = None
+    probe: ArrayMetadataProbe = field(default_factory=ArrayMetadataProbe)
+
+
+async def _probe_array_metadata(
     store_path: StorePath, zarr_format: ZarrFormat | None = 3
-) -> dict[str, JSON]:
+) -> ArrayProbeResult:
+    """Fetch and parse array metadata without raising on "not an array".
+
+    Returns an ``ArrayProbeResult`` whose ``metadata`` is set on success and whose
+    ``error`` is set (but not raised) when the path does not hold an array or the
+    request is invalid. ``probe`` always carries the reusable buffers fetched
+    during a ``zarr_format=None`` probe, even when ``error`` is set, which is
+    exactly when ``zarr.open`` needs to fall back to a group open.
+    """
+    probe: ArrayMetadataProbe = {}
     if zarr_format == 2:
         zarray_bytes, zattrs_bytes = await gather(
             (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype),
@@ -247,21 +289,24 @@ async def get_array_metadata(
                 "A Zarr V2 array metadata document was not found in store "
                 f"{store_path.store!r} at path {store_path.path!r}."
             )
-            raise ArrayNotFoundError(msg)
+            return ArrayProbeResult(error=ArrayNotFoundError(msg), probe=probe)
     elif zarr_format == 3:
         zarr_json_bytes = await (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype)
         if zarr_json_bytes is None:
             msg = (
                 "A Zarr V3 array metadata document was not found in store "
                 f"{store_path.store!r} at path {store_path.path!r}."
             )
-            raise ArrayNotFoundError(msg)
+            return ArrayProbeResult(error=ArrayNotFoundError(msg), probe=probe)
     elif zarr_format is None:
         zarr_json_bytes, zarray_bytes, zattrs_bytes = await gather(
             (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype),
             (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype),
             (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype),
         )
+        # Record the buffers that overlap with a group probe so the caller may
+        # reuse them. These are recorded even when array detection fails below.
+        probe = {"zarr_json": zarr_json_bytes, "zattrs": zattrs_bytes}
         if zarr_json_bytes is not None and zarray_bytes is not None:
             # warn and favor v3
             msg = f"Both zarr.json (Zarr format 3) and .zarray (Zarr format 2) metadata objects exist at {store_path}. Zarr v3 will be used."
@@ -271,15 +316,15 @@ async def get_array_metadata(
                 f"Neither Zarr V3 nor Zarr V2 array metadata documents "
                 f"were found in store {store_path.store!r} at path {store_path.path!r}."
             )
-            raise ArrayNotFoundError(msg)
+            return ArrayProbeResult(error=ArrayNotFoundError(msg), probe=probe)
         # set zarr_format based on which keys were found
         if zarr_json_bytes is not None:
             zarr_format = 3
         else:
             zarr_format = 2
     else:
         msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'."  # type: ignore[unreachable]
-        raise MetadataValidationError(msg)
+        return ArrayProbeResult(error=MetadataValidationError(msg), probe=probe)
 
     metadata_dict: dict[str, JSON]
     if zarr_format == 2:
@@ -293,9 +338,22 @@ async def get_array_metadata(
         assert zarr_json_bytes is not None
         metadata_dict = buffer_to_json_object(zarr_json_bytes)
 
-        parse_node_type_array(metadata_dict.get("node_type"))
+        try:
+            parse_node_type_array(metadata_dict.get("node_type"))
+        except NodeTypeValidationError as exc:
+            return ArrayProbeResult(error=exc, probe=probe)
+
+    return ArrayProbeResult(metadata=metadata_dict, probe=probe)
 
-    return metadata_dict
+
+async def get_array_metadata(
+    store_path: StorePath, zarr_format: ZarrFormat | None = 3
+) -> dict[str, JSON]:
+    result = await _probe_array_metadata(store_path, zarr_format=zarr_format)
+    if result.error is not None:
+        raise result.error
+    assert result.metadata is not None
+    return result.metadata
 
 
 @dataclass(frozen=True)

src/zarr/core/group.py

@@ -8,7 +8,7 @@
 from collections import defaultdict
 from dataclasses import asdict, dataclass, field, fields, replace
 from itertools import accumulate
-from typing import TYPE_CHECKING, Literal, assert_never, cast, overload
+from typing import TYPE_CHECKING, Literal, NotRequired, TypedDict, assert_never, cast, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -83,6 +83,21 @@
 logger = logging.getLogger("zarr.group")
 
 
+class _PreFetchedGroupMetadata(TypedDict):
+    """Already-fetched metadata buffers that ``AsyncGroup.open`` may reuse.
+
+    This is a private optimization used by ``zarr.open`` to avoid re-fetching
+    the ``zarr.json`` and ``.zattrs`` documents it already read while probing
+    for an array. Only the format-detection documents are cached here; the
+    consolidated-metadata document is still fetched by ``AsyncGroup.open``
+    because the key depends on ``use_consolidated``.
+    """
+
+    zarr_json: NotRequired[Buffer | None]
+    zgroup: NotRequired[Buffer | None]
+    zattrs: NotRequired[Buffer | None]
+
+
 def parse_zarr_format(data: Any) -> ZarrFormat:
     """Parse the zarr_format field from metadata."""
     if data in (2, 3):
@@ -490,6 +505,7 @@ async def open(
         store: StoreLike,
         zarr_format: ZarrFormat | None = 3,
         use_consolidated: bool | str | None = None,
+        _pre_fetched_metadata: _PreFetchedGroupMetadata | None = None,
     ) -> AsyncGroup:
         """Open a new AsyncGroup
 
@@ -514,6 +530,11 @@ async def open(
             Zarr format 2 allowed configuring the key storing the consolidated metadata
             (``.zmetadata`` by default). Specify the custom key as ``use_consolidated``
             to load consolidated metadata from a non-default key.
+        _pre_fetched_metadata : _PreFetchedGroupMetadata or None, default None
+            Private. Already-fetched metadata buffers (``zarr.json``, ``.zgroup``,
+            ``.zattrs``) that this method may reuse instead of re-fetching. Only
+            consulted when ``zarr_format`` is None. Used by ``zarr.open`` to avoid
+            duplicate fetches when falling back from an array probe to a group open.
         """
         store_path = await make_store_path(store)
         if not store_path.store.supports_consolidated_metadata:
@@ -554,16 +575,38 @@ async def open(
             if zarr_json_bytes is None:
                 raise FileNotFoundError(store_path)
         elif zarr_format is None:
+            # Reuse any buffers that a caller (e.g. ``zarr.open``) already
+            # fetched while probing for an array, only fetching what is missing.
+            pre_fetched = _pre_fetched_metadata or {}
+            keys = (ZARR_JSON, ZGROUP_JSON, ZATTRS_JSON, str(consolidated_key))
+            cached = (
+                pre_fetched.get("zarr_json"),
+                pre_fetched.get("zgroup"),
+                pre_fetched.get("zattrs"),
+                None,  # consolidated metadata is never pre-fetched
+            )
+            present = (
+                "zarr_json" in pre_fetched,
+                "zgroup" in pre_fetched,
+                "zattrs" in pre_fetched,
+                False,
+            )
+            results = await asyncio.gather(
+                *(
+                    (store_path / key).get()
+                    for key, is_present in zip(keys, present, strict=True)
+                    if not is_present
+                )
+            )
+            results_iter = iter(results)
             (
                 zarr_json_bytes,
                 zgroup_bytes,
                 zattrs_bytes,
                 maybe_consolidated_metadata_bytes,
-            ) = await asyncio.gather(
-                (store_path / ZARR_JSON).get(),
-                (store_path / ZGROUP_JSON).get(),
-                (store_path / ZATTRS_JSON).get(),
-                (store_path / str(consolidated_key)).get(),
+            ) = tuple(
+                value if is_present else next(results_iter)
+                for value, is_present in zip(cached, present, strict=True)
             )
             if zarr_json_bytes is not None and zgroup_bytes is not None:
                 # warn and favor v3