Enhanced Metadata Retrieval

- Introduced `StorePath.get_many` to retrieve multiple metadata files in a single
  call, optimizing the retrieval process and reducing overhead based on
  concurrency safety. `get_many` is used in `get_array_metadata`.
- Introduced `FsspecStore._get_many` to retrieve multiple metadata files in a
  single call. `_get_many` is used in `_read_metadata_v2`.
- Introduced `StorePath._is_concurrency_save` to determine if the file system
  supports concurrent access.
- Updated logic 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.

Author: dgegen

changes/3207.bugfix.rst

@@ -0,0 +1,11 @@
+- Introduced ``StorePath.get_many`` to retrieve multiple metadata files in a single
+  call, optimizing the retrieval process and reducing overhead based on
+  concurrency safety. ``get_many`` is used in ``get_array_metadata``.
+- Introduced ``FsspecStore._get_many`` to retrieve multiple metadata files in a
+  single call. ``_get_many`` is used in ``_read_metadata_v2``.
+- Introduced ``StorePath._is_concurrency_save`` to determine if the file system
+  supports concurrent access.
+- Updated logic 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.

src/zarr/core/array.py

@@ -200,9 +200,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(
+            ZARRAY_JSON,
+            ZATTRS_JSON,
+            prototype=cpu_buffer_prototype,
         )
         if zarray_bytes is None:
             raise FileNotFoundError(store_path)
@@ -211,10 +212,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(
+            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
@@ -1429,7 +1431,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(
+                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,11 +3476,20 @@ 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()),
-    )
+    requests = [
+        (_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),
+    ]
+
+    # Use the _get_many method to retrieve the data
+    results = {}
+    async for key, buffer in store._get_many(requests):
+        results[key] = buffer
+
+    zarray_bytes = results.get(_join_paths([path, ZARRAY_JSON]))
+    zgroup_bytes = results.get(_join_paths([path, ZGROUP_JSON]))
+    zattrs_bytes = results.get(_join_paths([path, ZATTRS_JSON]))
 
     if zattrs_bytes is None:
         zattrs = {}

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,24 @@ async def get(
             prototype = default_buffer_prototype()
         return await self.store.get(self.path, prototype=prototype, byte_range=byte_range)
 
+    async def get_many(
+        self,
+        *suffixes : str,
+        prototype: BufferPrototype | None = None,
+        byte_range: ByteRequest | None = None,
+    ):
+        tasks = [
+            (self / suffix).get(prototype=prototype, byte_range=byte_range) for suffix in suffixes
+        ]
+        if await self._is_concurrency_save():
+            return await gather(*tasks)
+        else:
+            results = []
+            for task in tasks:
+                result = await task
+                results.append(result)
+            return results
+
     async def set(self, value: Buffer, byte_range: ByteRequest | None = None) -> None:
         """
         Write bytes to the store.
@@ -263,6 +282,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,