Use Store.get_many for whole-chunk reads in BatchedCodecPipeline

changes/1758.feature.md

@@ -0,0 +1,10 @@
+`BatchedCodecPipeline` now fetches the encoded bytes for a whole (non-sharded)
+read with a single `Store.get_many` call spanning the entire request, instead of
+issuing one `Store.get` per chunk. This lets a store batch or coalesce the
+underlying reads — for example `FsspecStore` coalesces nearby chunk reads via
+`cat_ranges`, and a custom store (such as virtualizarr's `ManifestStore` or
+icechunk's `IcechunkStore`) can override `get_many` to merge reads that resolve
+into the same underlying object — independently of `codec_pipeline.batch_size`,
+which still governs only decode batching. The sharding codec's partial-decode
+path is unchanged, and stores without a specialized `get_many` fall back to the
+previous concurrent per-chunk behavior.

changes/1806.feature.md

@@ -0,0 +1,11 @@
+Add `zarr.abc.store.Store.get_many`, a bulk counterpart to `Store.get` that
+retrieves many values — each a whole key or a `(key, byte_range)` pair — in a
+single call. It generalizes `Store.get_ranges` (many ranges of one key) to many
+keys, yielding `(request_index, Buffer | None)` batches in completion order so a
+store can coalesce reads that land in the same underlying object. The method is
+defined on the `Store` ABC with a default implementation that fetches the
+requests concurrently with `Store.get`, so every store inherits a working
+version; stores whose backend can retrieve many objects together should override
+it (`FsspecStore` does, coalescing via `fsspec`'s `cat_ranges`). Coalescing
+tuning is left to each store rather than exposed on the interface. This restores
+and generalizes the batched-fetch capability of the v2 `getitems` Store API.

src/zarr/abc/store.py

@@ -237,6 +237,71 @@ async def get_partial_values(
         """
         ...
 
+    async def get_many(
+        self,
+        requests: Sequence[tuple[str, ByteRequest | None] | str],
+        *,
+        prototype: BufferPrototype,
+    ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]:
+        """Retrieve many values, possibly from different keys, at once.
+
+        This is the bulk counterpart to :meth:`get`: the whole set of requests
+        is handed to the store in a single call, so an implementation can fetch
+        them together — for example by coalescing reads that land in the same
+        underlying object into fewer requests — rather than one at a time. It
+        generalizes :meth:`get_ranges` (which reads many ranges from a *single*
+        key) to many keys, each with an optional byte range.
+
+        Yields one batch per underlying I/O operation, each a sequence of
+        ``(request_index, Buffer | None)`` tuples where ``request_index`` is the
+        position of the request in ``requests``. Every request is reported
+        exactly once across all batches; a ``None`` buffer means that key is
+        absent. Batches arrive in completion order, not request order, so
+        callers use the indices to reassemble results.
+
+        The default implementation fetches each request concurrently with
+        :meth:`get`, so every store gets a working version for free; stores
+        whose backend can retrieve many objects together (e.g.
+        :class:`~zarr.storage.FsspecStore`, which coalesces nearby reads via
+        ``fsspec``) should override it. Anything specific to *how* a store
+        batches or coalesces (concurrency limits, gap thresholds, ...) is an
+        implementation concern of that store, not part of this interface.
+
+        Parameters
+        ----------
+        requests : Sequence[tuple[str, ByteRequest | None] | str]
+            The values to retrieve. Each request is either a bare key (the
+            whole value) or a ``(key, byte_range)`` tuple; a ``byte_range`` of
+            ``None`` also means the whole value. A key may appear more than
+            once with different ranges.
+        prototype : BufferPrototype
+            The prototype of the output buffers. Stores may support a default
+            buffer prototype.
+
+        Yields
+        ------
+        Sequence[tuple[int, Buffer | None]]
+            One batch per underlying I/O operation, each a sequence of
+            ``(request_index, Buffer | None)`` tuples.
+        """
+        # Local imports to avoid an import cycle at module load time.
+        from zarr.core.common import concurrent_map
+        from zarr.core.config import config
+
+        indexed = [
+            (i, req, None) if isinstance(req, str) else (i, req[0], req[1])
+            for i, req in enumerate(requests)
+        ]
+
+        async def _fetch(
+            index: int, key: str, byte_range: ByteRequest | None
+        ) -> tuple[int, Buffer | None]:
+            return index, await self.get(key, prototype, byte_range)
+
+        results = await concurrent_map(indexed, _fetch, config.get("async.concurrency"))
+        for result in results:
+            yield [result]
+
     @abstractmethod
     async def exists(self, key: str) -> bool:
         """Check if a key exists in the store.

src/zarr/core/codec_pipeline.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass, field
 from itertools import batched, pairwise
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 from warnings import warn
 
 from zarr.abc.codec import (
@@ -23,14 +23,61 @@
 from zarr.registry import register_pipeline
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator
+    from collections.abc import Iterable, Iterator, Mapping, Sequence
     from typing import Self
 
-    from zarr.abc.store import ByteGetter, ByteSetter
+    from zarr.abc.store import ByteGetter, ByteSetter, Store
     from zarr.core.array_spec import ArraySpec
     from zarr.core.buffer import Buffer, BufferPrototype, NDBuffer
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType
     from zarr.core.metadata.v3 import ChunkGridMetadata
+    from zarr.storage._common import StorePath
+
+
+def _bulk_store_keys(
+    batch_info_list: list[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
+) -> tuple[Store, list[str], BufferPrototype] | None:
+    """Return ``(store, keys, prototype)`` when a whole-chunk read batch can be
+    served by a single :meth:`~zarr.abc.store.Store.get_many` call, else ``None``.
+
+    Eligible when the batch is non-empty and every chunk is backed by a
+    ``StorePath`` sharing one ``Store`` (compared by identity) and one buffer
+    prototype. When it is not eligible (e.g. a mix of stores, or virtual byte
+    getters used by the sharding codec) the caller should fall back to
+    per-chunk fetching.
+    """
+    # Local import to avoid a module-load import cycle (storage imports core).
+    from zarr.storage._common import StorePath
+
+    if not batch_info_list:
+        return None
+    byte_getters = [byte_getter for byte_getter, *_ in batch_info_list]
+    if not all(isinstance(bg, StorePath) for bg in byte_getters):
+        return None
+    store_paths = cast("list[StorePath]", byte_getters)
+    store = store_paths[0].store
+    prototype = batch_info_list[0][1].prototype
+    if not all(sp.store is store for sp in store_paths):
+        return None
+    if not all(array_spec.prototype is prototype for _, array_spec, *_ in batch_info_list):
+        return None
+    return store, [sp.path for sp in store_paths], prototype
+
+
+async def _collect_get_many(
+    store: Store, keys: Sequence[str], prototype: BufferPrototype
+) -> list[Buffer | None]:
+    """Drive ``Store.get_many`` for a set of whole-key reads and return the
+    results as a positional list aligned to ``keys`` (``None`` for absent keys).
+
+    ``get_many`` yields ``(request_index, Buffer | None)`` batches in completion
+    order, so we scatter each result back to its input position.
+    """
+    out: list[Buffer | None] = [None] * len(keys)
+    async for batch in store.get_many(keys, prototype=prototype):
+        for index, buffer in batch:
+            out[index] = buffer
+    return out
 
 
 def _unzip2[T, U](iterable: Iterable[tuple[T, U]]) -> tuple[list[T], list[U]]:
@@ -353,11 +400,47 @@ async def encode_partial_batch(
         assert isinstance(self.array_bytes_codec, ArrayBytesCodecPartialEncodeMixin)
         await self.array_bytes_codec.encode_partial(batch_info)
 
+    async def _get_chunk_bytes_batch(
+        self,
+        batch_info_list: list[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
+        prefetched: Mapping[str, Buffer | None] | None = None,
+    ) -> list[Buffer | None]:
+        """Fetch the whole encoded bytes for each chunk in a batch.
+
+        If ``prefetched`` is supplied (a mapping of store key to already-read
+        ``Buffer``, as produced by :meth:`read`), the bytes are taken from it
+        directly. Otherwise, when every chunk is backed by a ``StorePath`` over
+        a single common ``Store`` and buffer prototype, the reads are handed to
+        the store as one :meth:`~zarr.abc.store.Store.get_many` call so a
+        backend that can batch or coalesce object reads gets the chance to do
+        so. Failing that, it falls back to fetching each chunk concurrently
+        with :meth:`~zarr.abc.store.ByteGetter.get`, matching prior behavior.
+        """
+        if prefetched is not None:
+            # read() already fetched these keys in one get_many call.
+            store_paths = cast("list[StorePath]", [bg for bg, *_ in batch_info_list])
+            return [prefetched.get(sp.path) for sp in store_paths]
+
+        plan = _bulk_store_keys(batch_info_list)
+        if plan is not None:
+            store, keys, prototype = plan
+            return await _collect_get_many(store, keys, prototype)
+
+        return await concurrent_map(
+            [
+                (byte_getter, array_spec.prototype)
+                for byte_getter, array_spec, *_ in batch_info_list
+            ],
+            lambda byte_getter, prototype: byte_getter.get(prototype),
+            config.get("async.concurrency"),
+        )
+
     async def read_batch(
         self,
         batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
+        prefetched: Mapping[str, Buffer | None] | None = None,
     ) -> tuple[GetResult, ...]:
         results: list[GetResult] = []
         if self.supports_partial_decode:
@@ -381,14 +464,7 @@ async def read_batch(
                     results.append(GetResult(status="missing"))
         else:
             batch_info_list = list(batch_info)
-            chunk_bytes_batch = await concurrent_map(
-                [
-                    (byte_getter, array_spec.prototype)
-                    for byte_getter, array_spec, *_ in batch_info_list
-                ],
-                lambda byte_getter, prototype: byte_getter.get(prototype),
-                config.get("async.concurrency"),
-            )
+            chunk_bytes_batch = await self._get_chunk_bytes_batch(batch_info_list, prefetched)
             chunk_array_batch = await self.decode_batch(
                 [
                     (chunk_bytes, chunk_spec)
@@ -590,9 +666,24 @@ async def read(
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
     ) -> tuple[GetResult, ...]:
+        batch_info = list(batch_info)
+        # For whole-chunk reads (not partial/sharded decode), fetch the encoded
+        # bytes for the entire request in a single Store.get_many call. This
+        # lets the store batch or coalesce the reads regardless of
+        # codec_pipeline.batch_size (which only governs decode batching), and
+        # restores the bulk-fetch behavior of the v2 getitems Store API. The
+        # per-batch read_batch calls then read their bytes from this mapping.
+        prefetched: Mapping[str, Buffer | None] | None = None
+        if not self.supports_partial_decode:
+            plan = _bulk_store_keys(batch_info)
+            if plan is not None:
+                store, keys, prototype = plan
+                values = await _collect_get_many(store, keys, prototype)
+                prefetched = dict(zip(keys, values, strict=True))
+
         batch_results = await concurrent_map(
             [
-                (single_batch_info, out, drop_axes)
+                (single_batch_info, out, drop_axes, prefetched)
                 for single_batch_info in batched(batch_info, self.batch_size)
             ],
             self.read_batch,

src/zarr/storage/_fsspec.py

@@ -22,7 +22,7 @@
 logger = getLogger(__name__)
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Iterable
+    from collections.abc import AsyncIterator, Iterable, Sequence
 
     from fsspec import AbstractFileSystem
     from fsspec.asyn import AsyncFileSystem
@@ -459,6 +459,23 @@ async def get_partial_values(
 
         return [None if isinstance(r, Exception) else prototype.buffer.from_bytes(r) for r in res]
 
+    async def get_many(
+        self,
+        requests: Sequence[tuple[str, ByteRequest | None] | str],
+        *,
+        prototype: BufferPrototype,
+    ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]:
+        # docstring inherited
+        # Hand the whole set of requests to fsspec in one call so it can
+        # coalesce nearby reads into fewer requests (via
+        # ``AbstractFileSystem._cat_ranges``), rather than issuing one request
+        # per key as the default Store.get_many does. get_partial_values
+        # returns results aligned to the input order, so we can index them
+        # directly and yield them as a single completed batch.
+        key_ranges = [(req, None) if isinstance(req, str) else req for req in requests]
+        values = await self.get_partial_values(prototype, key_ranges)
+        yield list(enumerate(values))
+
     async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         allfiles = await self.fs._find(self.path, detail=False, withdirs=False)

src/zarr/testing/store.py

@@ -262,7 +262,7 @@ async def test_get_raises(self, store: S) -> None:
         with pytest.raises((ValueError, TypeError), match=r"Unexpected byte_range, got.*"):
             await store.get("c/0", prototype=default_buffer_prototype(), byte_range=(0, 2))  # type: ignore[arg-type]
 
-    async def test_get_many(self, store: S) -> None:
+    async def test_get_many_streaming(self, store: S) -> None:
         """
         Ensure that multiple keys can be retrieved at once with the _get_many method.
         """
@@ -407,6 +407,41 @@ async def test_get_partial_values(
             obs.to_bytes() == exp.to_bytes() for obs, exp in zip(observed, expected, strict=True)
         )
 
+    async def test_get_many(self, store: S) -> None:
+        # put a handful of whole values
+        for key, data in {"c/0/0": b"aaaaa", "c/0/1": b"bb", "c/0/2": b"cccc"}.items():
+            await self.set(store, key, self.buffer_cls.from_bytes(data))
+
+        # mix bare keys, an explicit (key, None) tuple, a partial range, and a
+        # missing key. Each request must be reported exactly once, by index.
+        requests: list[tuple[str, ByteRequest | None] | str] = [
+            "c/0/0",
+            ("c/0/1", None),
+            ("c/0/0", RangeByteRequest(1, 3)),
+            "c/0/2",
+            "c/0/missing",
+        ]
+        collected: dict[int, Buffer | None] = {}
+        async for batch in store.get_many(requests, prototype=default_buffer_prototype()):
+            for index, value in batch:
+                assert index not in collected  # reported exactly once
+                collected[index] = value
+
+        assert set(collected) == set(range(len(requests)))
+        assert collected[4] is None  # missing key -> None (not omitted)
+        expected = {0: b"aaaaa", 1: b"bb", 2: b"aa", 3: b"cccc"}  # index 2 is "aaaaa"[1:3]
+        for index, want in expected.items():
+            buffer = collected[index]
+            assert buffer is not None
+            assert buffer.to_bytes() == want
+
+    async def test_get_many_empty(self, store: S) -> None:
+        # an empty request is valid and yields no results
+        batches = [
+            batch async for batch in store.get_many([], prototype=default_buffer_prototype())
+        ]
+        assert [pair for batch in batches for pair in batch] == []
+
     async def test_exists(self, store: S) -> None:
         assert not await store.exists("foo")
         await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar"))