Remove get_many_ordered

Author: dgegen

changes/3207.bugfix.rst

@@ -2,17 +2,14 @@
   opening Zarr v3 arrays on synchronous fsspec filesystems, by implementing a
   fallback to sequential reads for non-concurrency-safe filesystems, ensuring
   robust metadata retrieval without sacrificing performance for safe
-  filesystems. Furthermore ``Store.get_many`` was modified to retrieve objects
+  filesystems. Furthermore ``Store._get_many`` was modified to retrieve objects
   concurrently from storage. The previous implementation was sequential,
-  awaiting each ``self.get(*req)`` before proceeding, contrary to the docstring.
-- 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. ``StorePath.get_many_ordered`` is used in
-  ``get_array_metadata``. ``Store._get_many_ordered`` is used in
-  ``_read_metadata_v2``.
-- Modified ``FsspecStore._get_many`` and ``FsspecStore._get_many_ordered``
-  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. Adding tests ``LockableFileSystem`` to test with
-  async/sync behavior.
+  awaiting each ``self.get(*req)`` before proceeding, contrary to the
+  docstring.
+- Introduced ``StorePath.get_many``, mimicing the behaviour of `StorePath.get`.
+- Use ``Store._get_many`` and ``StorePath.get_many`` in ``get_array_metadata``.
+- Implemented ``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. Adding tests ``LockableFileSystem`` to test
+  with async/sync behavior.

src/zarr/abc/store.py

@@ -426,16 +426,6 @@ async def _get_with_name(
             task = await completed
             yield task
 
-    async def _get_many_ordered(
-        self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
-    ) -> tuple[Buffer | None, ...]:
-        """
-        Retrieve a collection of objects from storage in the order they were requested.
-        """
-        tasks = [self.get(*req) for req in requests]
-
-        return tuple(await gather(*tasks))
-
     async def getsize(self, key: str) -> int:
         """
         Return the size, in bytes, of a value in a Store.

src/zarr/core/array.py

@@ -211,24 +211,25 @@ 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 store_path.get_many_ordered(
-            ZARRAY_JSON,
-            ZATTRS_JSON,
-            prototype=cpu_buffer_prototype,
-        )
+        requests = [(key, default_buffer_prototype(), None) for key in [ZARRAY_JSON, ZATTRS_JSON]]
+        retrieved_buffers = {key: value async for key, value in store_path.get_many(requests)}
+        zarray_bytes, zattrs_bytes = tuple(retrieved_buffers.get(req[0]) for req in requests)
+
         if zarray_bytes is None:
             raise FileNotFoundError(store_path)
     elif zarr_format == 3:
         zarr_json_bytes = await (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype)
         if zarr_json_bytes is None:
             raise FileNotFoundError(store_path)
     elif zarr_format is None:
-        zarr_json_bytes, zarray_bytes, zattrs_bytes = await store_path.get_many_ordered(
-            ZARR_JSON,
-            ZARRAY_JSON,
-            ZATTRS_JSON,
-            prototype=cpu_buffer_prototype,
+        requests = [
+            (key, default_buffer_prototype(), None) for key in [ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON]
+        ]
+        retrieved_buffers = {key: value async for key, value in store_path.get_many(requests)}
+        zarr_json_bytes, zarray_bytes, zattrs_bytes = tuple(
+            retrieved_buffers.get(req[0]) for req in requests
         )
+
         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."

src/zarr/core/group.py

@@ -513,19 +513,23 @@ async def open(
             consolidated_key = use_consolidated
 
         if zarr_format == 2:
-            paths = [store_path / ZGROUP_JSON, store_path / ZATTRS_JSON]
+            requests = [
+                (key, default_buffer_prototype(), None) for key in [ZGROUP_JSON, ZATTRS_JSON]
+            ]
             if use_consolidated or use_consolidated is None:
-                paths.append(store_path / consolidated_key)
+                requests.append((consolidated_key, default_buffer_prototype(), None))
 
-            zgroup_bytes, zattrs_bytes, *rest = await asyncio.gather(
-                *[path.get() for path in paths]
+            retrieved_buffers = {key: value async for key, value in store_path.get_many(requests)}
+            zgroup_bytes, zattrs_bytes = (
+                retrieved_buffers[ZGROUP_JSON],
+                retrieved_buffers[ZATTRS_JSON],
             )
+
             if zgroup_bytes is None:
                 raise FileNotFoundError(store_path)
 
             if use_consolidated or use_consolidated is None:
-                maybe_consolidated_metadata_bytes = rest[0]
-
+                maybe_consolidated_metadata_bytes = retrieved_buffers[consolidated_key]
             else:
                 maybe_consolidated_metadata_bytes = None
 
@@ -534,17 +538,18 @@ async def open(
             if zarr_json_bytes is None:
                 raise FileNotFoundError(store_path)
         elif zarr_format is None:
+            requests = [
+                (key, default_buffer_prototype(), None)
+                for key in [ZARR_JSON, ZGROUP_JSON, ZATTRS_JSON, consolidated_key]
+            ]
+            retrieved_buffers = {key: value async for key, value in store_path.get_many(requests)}
             (
                 zarr_json_bytes,
                 zgroup_bytes,
                 zattrs_bytes,
                 maybe_consolidated_metadata_bytes,
-            ) = await store_path.get_many_ordered(
-                ZARR_JSON,
-                ZGROUP_JSON,
-                ZATTRS_JSON,
-                consolidated_key,
-            )
+            ) = tuple(retrieved_buffers.get(req[0]) for req in requests)
+
             if zarr_json_bytes is not None and zgroup_bytes is not None:
                 # warn and favor v3
                 msg = f"Both zarr.json (Zarr format 3) and .zgroup (Zarr format 2) metadata objects exist at {store_path}. Zarr format 3 will be used."
@@ -3476,12 +3481,14 @@ 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 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),
-        ]
+    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),
+    ]
+    retrieved_buffers = {key: value async for key, value in store._get_many(requests)}
+    zarray_bytes, zgroup_bytes, zattrs_bytes = tuple(
+        retrieved_buffers.get(req[0]) for req in requests
     )
 
     if zattrs_bytes is None:

src/zarr/storage/_common.py

@@ -28,6 +28,8 @@
     FSMap = None
 
 if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Iterable
+
     from zarr.core.buffer import BufferPrototype
 
 
@@ -164,34 +166,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_ordered(
-        self,
-        *path_components: str,
-        prototype: BufferPrototype | None = None,
-        byte_range: ByteRequest | None = None,
-    ) -> tuple[Buffer | None, ...]:
+    async def get_many(
+        self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
+    ) -> AsyncGenerator[tuple[str, Buffer | None], 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.
+        requests : Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
 
-        Returns
+        Yields
         -------
-        tuple[Buffer | None, ...]
-            A tuple of buffers read from the store, in the order of the provided path_components.
+        tuple[str, Buffer | None]
         """
-        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])
+        path_component_dict = {(self / req[0]).path: req[0] for req in requests}
+        complete_requests = [((self / req[0]).path, *req[1:]) for req in requests]
+        async for result in self.store._get_many(complete_requests):
+            yield (path_component_dict[result[0]], *result[1:])
 
     async def set(self, value: Buffer, byte_range: ByteRequest | None = None) -> None:
         """

src/zarr/storage/_fsspec.py

@@ -331,24 +331,12 @@ 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 key, value in super()._get_many(requests):
-                yield (key, value)
+            async for result in super()._get_many(requests):
+                yield result
         else:
             for key, prototype, byte_range in requests:
                 value = await self.get(key, prototype, byte_range)
                 yield (key, value)
-    
-    async def _get_many_ordered(
-        self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]]
-    ) -> tuple[Buffer | None, ...]:
-        if getattr(self.fs, "asynchronous", True):
-            return await super()._get_many_ordered(requests)
-        else:
-            results = []
-            for key, prototype, byte_range in requests:
-                value = await self.get(key, prototype, byte_range)
-                results.append(value)
-            return tuple(results)
 
     async def set(
         self,

src/zarr/testing/store.py

@@ -265,30 +265,6 @@ 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:

tests/test_store/test_fsspec.py

@@ -532,25 +532,10 @@ async def test_get_many_synchronous_fs(self):
         results = []
         async for k, v in self.store_sync._get_many(requests):
             results.append((k, v.to_bytes() if v else None))
-        # Results should already be in the same order as requests
+        # In the synchronous case, results should be in the same order as requests
 
         assert results == true_results
 
-    async def test_get_many_ordered_synchronous_fs(self):
-        requests, true_results = self.get_requests_and_true_results()
-
-        results = await self.store_sync._get_many_ordered(requests)
-        results = [value.to_bytes() if value else None for value in results]
-
-        assert results == [value[1] for value in true_results]
-
-    async def test_get_many_ordered_asynchronous_fs(self):
-        requests, true_results = self.get_requests_and_true_results()
-
-        results = await self.store_async._get_many_ordered(requests)
-        results = [value.to_bytes() if value else None for value in results]
-
-        assert results == [value[1] for value in true_results]
 
     async def test_asynchronous_locked_fs_raises(self):
         store = LockableFileSystem(asynchronous=True, lock=True).get_store(path="root")
@@ -559,6 +544,3 @@ async def test_asynchronous_locked_fs_raises(self):
         with pytest.raises(RuntimeError, match="Concurrent requests!"):
             async for _, _ in store._get_many(requests):
                 pass
-
-        with pytest.raises(RuntimeError, match="Concurrent requests!"):
-            await store._get_many_ordered(requests)