Enhanced Metadata Retrieval

- Introduced `Store.get_many_ordered` and `StorePath.get_many_ordered` to
  retrieve multiple metadata files in a single call, optimizing the retrieval
  process and reducing overhead. `get_many_ordered` is used in
  `get_array_metadata`.
- Modified `FsspecStore._get_many` to conditionally use `asyncio.gather`
  based on the concurrency safety of the underlying file system, enhancing
  compatibility with synchronous file systems by avoiding deadlocks when
  accessing metadata concurrently.
- Modified `_read_metadata_v2` to use `Store._get_many`

Author: dgegen

changes/3207.bugfix.rst

@@ -0,0 +1,9 @@
+- Introduced ``Store.get_many_ordered`` and ``StorePath.get_many_ordered`` to
+  retrieve multiple metadata files in a single call, optimizing the retrieval
+  process and reducing overhead. ``get_many_ordered`` is used in
+  ``get_array_metadata``.
+- Modified ``FsspecStore._get_many`` to conditionally use ``asyncio.gather``
+  based on the concurrency safety of the underlying file system, enhancing
+  compatibility with synchronous file systems by avoiding deadlocks when
+  accessing metadata concurrently.
+- Modified ``_read_metadata_v2`` to use ``Store._get_many``

src/zarr/abc/store.py

@@ -11,7 +11,7 @@
 from zarr.core.config import config
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, AsyncIterator, Iterable
+    from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Sequence
     from types import TracebackType
     from typing import Any, Self, TypeAlias
 
@@ -417,6 +417,18 @@ async def _get_many(
         for req in requests:
             yield (req[0], await self.get(*req))
 
+    async def _get_many_ordered(
+        self, requests: Sequence[tuple[str, BufferPrototype, ByteRequest | None]]
+    ) -> tuple[Buffer | None, ...]:
+        """
+        Retrieve a collection of objects from storage in the order they were requested.
+        """
+        key_to_index = {req[0]: i for i, req in enumerate(requests)}
+        results: list[Buffer | None] = [None] * len(requests)
+        async for key, value in self._get_many(requests):
+            results[key_to_index[key]] = value
+        return tuple(results)
+
     async def getsize(self, key: str) -> int:
         """
         Return the size, in bytes, of a value in a Store.

src/zarr/core/array.py

@@ -205,9 +205,10 @@ async def get_array_metadata(
     store_path: StorePath, zarr_format: ZarrFormat | None = 3
 ) -> dict[str, JSON]:
     if zarr_format == 2:
-        zarray_bytes, zattrs_bytes = await gather(
-            (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype),
-            (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype),
+        zarray_bytes, zattrs_bytes = await store_path.get_many_ordered(
+            ZARRAY_JSON,
+            ZATTRS_JSON,
+            prototype=cpu_buffer_prototype,
         )
         if zarray_bytes is None:
             raise FileNotFoundError(store_path)
@@ -216,10 +217,11 @@ async def get_array_metadata(
         if zarr_json_bytes is None:
             raise FileNotFoundError(store_path)
     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),
+        zarr_json_bytes, zarray_bytes, zattrs_bytes = await store_path.get_many_ordered(
+            ZARR_JSON,
+            ZARRAY_JSON,
+            ZATTRS_JSON,
+            prototype=cpu_buffer_prototype,
         )
         if zarr_json_bytes is not None and zarray_bytes is not None:
             # warn and favor v3
@@ -1436,7 +1438,6 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F
                         ).items()
                     ]
                 )
-
         await gather(*awaitables)
 
     async def _set_selection(

src/zarr/core/group.py

@@ -540,11 +540,11 @@ async def open(
                 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(),
+            ) = await store_path.get_many_ordered(
+                ZARR_JSON,
+                ZGROUP_JSON,
+                ZATTRS_JSON,
+                consolidated_key,
             )
             if zarr_json_bytes is not None and zgroup_bytes is not None:
                 # warn and favor v3
@@ -3476,10 +3476,12 @@ async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupM
     """
     # TODO: consider first fetching array metadata, and only fetching group metadata when we don't
     # find an array
-    zarray_bytes, zgroup_bytes, zattrs_bytes = await asyncio.gather(
-        store.get(_join_paths([path, ZARRAY_JSON]), prototype=default_buffer_prototype()),
-        store.get(_join_paths([path, ZGROUP_JSON]), prototype=default_buffer_prototype()),
-        store.get(_join_paths([path, ZATTRS_JSON]), prototype=default_buffer_prototype()),
+    zarray_bytes, zgroup_bytes, zattrs_bytes = await store._get_many_ordered(
+        [
+            (_join_paths([path, ZARRAY_JSON]), default_buffer_prototype(), None),
+            (_join_paths([path, ZGROUP_JSON]), default_buffer_prototype(), None),
+            (_join_paths([path, ZATTRS_JSON]), default_buffer_prototype(), None),
+        ]
     )
 
     if zattrs_bytes is None:

src/zarr/storage/_common.py

@@ -2,6 +2,7 @@
 
 import importlib.util
 import json
+from asyncio import gather
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias
 
@@ -163,6 +164,35 @@ async def get(
             prototype = default_buffer_prototype()
         return await self.store.get(self.path, prototype=prototype, byte_range=byte_range)
 
+    async def get_many_ordered(
+        self,
+        *path_components: str,
+        prototype: BufferPrototype | None = None,
+        byte_range: ByteRequest | None = None,
+    ) -> tuple[Buffer | None, ...]:
+        """
+        Read multiple bytes from the store in order of the provided path_components.
+
+        Parameters
+        ----------
+        path_components : str
+            Components to append to the store path.
+        prototype : BufferPrototype, optional
+            The buffer prototype to use when reading the bytes.
+        byte_range : ByteRequest, optional
+            The range of bytes to read.
+
+        Returns
+        -------
+        tuple[Buffer | None, ...]
+            A tuple of buffers read from the store, in the order of the provided path_components.
+        """
+        if prototype is None:
+            prototype = default_buffer_prototype()
+
+        tasks = [(self / component).path for component in path_components]
+        return await self.store._get_many_ordered([(task, prototype, byte_range) for task in tasks])
+
     async def set(self, value: Buffer, byte_range: ByteRequest | None = None) -> None:
         """
         Write bytes to the store.
@@ -263,6 +293,10 @@ def __eq__(self, other: object) -> bool:
             pass
         return False
 
+    async def _is_concurrency_save(self):
+        fs = getattr(self.store, "fs", None)
+        return getattr(fs, "asynchronous", True)
+
 
 StoreLike: TypeAlias = Store | StorePath | FSMap | Path | str | dict[str, Buffer]
 

src/zarr/storage/_fsspec.py

@@ -18,7 +18,7 @@
 from zarr.storage._common import _dereference_path
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Iterable
+    from collections.abc import AsyncGenerator, AsyncIterator, Iterable
 
     from fsspec import AbstractFileSystem
     from fsspec.asyn import AsyncFileSystem
@@ -326,6 +326,17 @@ async def get(
         else:
             return value
 
+    async def _get_many(
+        self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
+    ) -> AsyncGenerator[tuple[str, Buffer | None], None]:
+        if getattr(self.fs, "asynchronous", True):
+            async for result in super()._get_many(requests=requests):
+                yield result
+        else:
+            for key, prototype, byte_range in requests:
+                value = await self.get(key, prototype, byte_range)
+                yield (key, value)
+
     async def set(
         self,
         key: str,

src/zarr/testing/store.py

@@ -265,6 +265,30 @@ async def test_get_many(self, store: S) -> None:
         expected_kvs = sorted(((k, b) for k, b in zip(keys, values, strict=False)))
         assert observed_kvs == expected_kvs
 
+    async def test_get_many_ordered(self, store: S) -> None:
+        """
+        Ensure that multiple keys can be retrieved at once with the _get_many method,
+        preserving the order of keys.
+        """
+        keys = tuple(map(str, range(10)))
+        values = tuple(f"{k}".encode() for k in keys)
+
+        for k, v in zip(keys, values, strict=False):
+            await self.set(store, k, self.buffer_cls.from_bytes(v))
+
+        observed_buffers = await store._get_many_ordered(
+            tuple(
+                zip(
+                    keys,
+                    (default_buffer_prototype(),) * len(keys),
+                    (None,) * len(keys),
+                    strict=False,
+                )
+            )
+        )
+        observed_bytes = tuple(b.to_bytes() for b in observed_buffers if b is not None)
+        assert observed_bytes == values, f"Expected {values}, got {observed_bytes}"
+
     @pytest.mark.parametrize("key", ["c/0", "foo/c/0.0", "foo/0/0"])
     @pytest.mark.parametrize("data", [b"\x01\x02\x03\x04", b""])
     async def test_getsize(self, store: S, key: str, data: bytes) -> None: