Use Store.get_many for whole-chunk reads in BatchedCodecPipeline

BatchedCodecPipeline.read now fetches the encoded bytes for an entire
(non-sharded) read with a single Store.get_many call, instead of one
Store.get per chunk. It drives get_many over all chunk keys, scatters the
completion-ordered (index, buffer) results back into position, and feeds
them to the per-batch decode path.

This lets a store batch or coalesce the underlying reads (e.g. FsspecStore
via cat_ranges, or a custom store such as virtualizarr's ManifestStore /
icechunk's IcechunkStore that overrides get_many) regardless of
codec_pipeline.batch_size, which still governs only decode batching. The
sharding codec's partial-decode path is untouched, and stores without a
specialized get_many fall back to the previous concurrent per-chunk gets.

Author: TomNicholas

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.

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,

tests/test_codec_pipeline.py

@@ -19,9 +19,11 @@
 from zarr.storage import MemoryStore
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import AsyncIterator, Callable, Sequence
 
     from zarr.abc.codec import Codec
+    from zarr.abc.store import ByteRequest
+    from zarr.core.buffer import Buffer, BufferPrototype
 
 
 @pytest.mark.parametrize(
@@ -195,3 +197,48 @@ def test_codecs_from_list_outcome_matches_order_rules(labels: list[str]) -> None
         aa, _ab, bb = codecs_from_list(codecs)
         assert labels.count(_AA) == len(aa)
         assert labels.count(_BB) == len(bb)
+
+
+async def test_read_uses_bulk_get_many() -> None:
+    """The pipeline should fetch a whole multi-chunk read with a single
+    ``Store.get_many`` call (spanning the entire request, independent of
+    ``codec_pipeline.batch_size``), rather than one ``get`` per chunk."""
+    store = MemoryStore()
+    arr = zarr.create_array(store, shape=(20,), chunks=(5,), dtype="int64")  # 4 chunks
+    arr[:] = np.arange(20)
+
+    calls: dict[str, int] = {"get_many": 0, "requests": 0}
+    orig_get_many = store.get_many
+
+    # get_many is an async generator, so the spy is a sync function returning
+    # the underlying async iterator; count at call time.
+    def spy_get_many(
+        requests: Sequence[tuple[str, ByteRequest | None] | str],
+        *,
+        prototype: BufferPrototype,
+    ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]:
+        requests = list(requests)
+        calls["get_many"] += 1
+        calls["requests"] += len(requests)
+        return orig_get_many(requests, prototype=prototype)
+
+    store.get_many = spy_get_many  # type: ignore[method-assign]
+
+    result = arr[:]
+    np.testing.assert_array_equal(result, np.arange(20))
+    # one bulk call covering all four chunks
+    assert calls["get_many"] == 1
+    assert calls["requests"] == 4
+
+
+async def test_read_bulk_handles_missing_chunks() -> None:
+    """A bulk read where some chunks were never written must still fill those
+    positions with the fill value (get_many reports missing keys as None)."""
+    store = MemoryStore()
+    arr = zarr.open_array(store, mode="w", shape=(20,), chunks=(5,), dtype="int64", fill_value=-1)
+    arr[0:5] = 7  # write only the first chunk; the other three are missing
+
+    result = arr[:]
+    expected = np.full(20, -1, dtype="int64")
+    expected[0:5] = 7
+    np.testing.assert_array_equal(result, expected)