Optimize partial shard reads (zarr-developers#3004)

* Add performance test of partial shard reads

* WIP Consolidate reads of multiple chunks in the same shard

Add test and make max gap and max coalesce size config options

Code clarity and comments

Test that chunk request coalescing reduces calls to store

Profile a few values for coalesce_max_gap

Update [doc]tests to include new sharding.read.* values

document sharded read config options in user-guide/config.rst

tweak logic: start new coalesced group if coalescing would exceed `coalesce_max_bytes`

previous logic only started a new group if existing group
was size already exceeded coalesce_max_bytes.

set `mypy_path = "src"` to help pre-commit mypy find imported classes

Reorder methods in sharding.py, add docstring + commenting

wording docs fix

docstring clarification

trigger precommit on all python files changed in this pull request

trying to get the ruff format that's happening locally during pre-commit
to match the pre-commit run that is failing on CI.

revert trigger for pre-commit ruff format

* Add changes/3004.feature.rst

* Consistently return None on failure and test partial shard read failure modes

Use range of integers as out_selection not slice in CoordinateIndexer
To fix issue when using vindex with repeated indexes in indexer

test: improve formatting and add debugging breakpoint in array property tests

test: disable hypothesis deadline for test_array_roundtrip to prevent timeout

fix: initialize decode buffers with shard_spec.fill_value instead of 0 to fix partial shard holes

style: reformat code for improved readability and consistency in sharding.py

fix: revert incorrect RangeByteRequest length fix in sharding byte retrieval

* Fix and test for case where some chunks in shard are all fill

* Self review

* Removing profiling code masquerading as a skipped test

* revert change to indexing.py, not needed

* Add test for duplicate integer indexing into a coalesced group

* Undo change to fill value when initializing shard arrays

* Undo change to set mypy_path = "src"

* Commenting and revert uncessary changes to files for smaller diff

* remove now redundant cast

* Document runtime config keys

* Improve changelog entry and .rst -> .md

* .coords -> .chunk_coords in _ChunkCoordsByteSlice dataclass

* Update test env in docs/contributing.md

* Move `config.get` calls up into `_decode_partial_single`

* Ensure no change in behavior when ByteGetter.get returns None + comment

* Add test_sharing_unit.py, focusing on coallesce behavior, but with basic tests for other components

* Fix typing errors in test_sharing_unit.py

* Only use concurrent_map over chunks within a shard if > 1 groups after coalescing

* no-op test change to retry CI after unavailable runner failure

* Another no-op test change to retry CI after unavailable runner failure

* use get_partial_values(), remove explicit coallescing and concurrent_map

* cleanup

* self review

* Unit tests for new get_partial_values implementations

* tests: work around mypy not seeing equivalent protocols as equivalent

* Re-work to use Store.get_ranges. Simplify significantly.

* Add missing assert in test_sharding.py

* Remove unit tests that tested now-removed _ShardIndex.is_dense

---------

Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>

Author: aldenks

changes/3004.feature.md

@@ -0,0 +1,4 @@
+Optimizes reading multiple chunks from a shard. Serial calls to `Store.get()`
+in the sharding codec have been replaced with a single call to
+`Store.get_ranges()`, which coalesces nearby byte ranges and fetches them
+concurrently.

src/zarr/codecs/sharding.py

@@ -59,6 +59,7 @@
     parse_codecs,
 )
 from zarr.registry import get_ndbuffer_class, get_pipeline_class
+from zarr.storage._common import StorePath
 from zarr.storage._utils import _normalize_byte_range_index
 
 if TYPE_CHECKING:
@@ -467,32 +468,26 @@ async def _decode_partial_single(
         all_chunk_coords = {chunk_coords for chunk_coords, *_ in indexed_chunks}
 
         # reading bytes of all requested chunks
-        shard_dict: ShardMapping = {}
+        shard_dict_maybe: ShardMapping | None
         if self._is_total_shard(all_chunk_coords, chunks_per_shard):
             # read entire shard
             shard_dict_maybe = await self._load_full_shard_maybe(
                 byte_getter=byte_getter,
                 prototype=chunk_spec.prototype,
                 chunks_per_shard=chunks_per_shard,
             )
-            if shard_dict_maybe is None:
-                return None
-            shard_dict = shard_dict_maybe
         else:
             # read some chunks within the shard
-            shard_index = await self._load_shard_index_maybe(byte_getter, chunks_per_shard)
-            if shard_index is None:
-                return None
-            shard_dict = {}
-            for chunk_coords in all_chunk_coords:
-                chunk_byte_slice = shard_index.get_chunk_slice(chunk_coords)
-                if chunk_byte_slice:
-                    chunk_bytes = await byte_getter.get(
-                        prototype=chunk_spec.prototype,
-                        byte_range=RangeByteRequest(chunk_byte_slice[0], chunk_byte_slice[1]),
-                    )
-                    if chunk_bytes:
-                        shard_dict[chunk_coords] = chunk_bytes
+            shard_dict_maybe = await self._load_partial_shard_maybe(
+                byte_getter,
+                chunk_spec.prototype,
+                chunks_per_shard,
+                all_chunk_coords,
+            )
+
+        if shard_dict_maybe is None:
+            return None
+        shard_dict = shard_dict_maybe
 
         # decoding chunks and writing them into the output buffer
         await self.codec_pipeline.read(
@@ -779,6 +774,66 @@ async def _load_full_shard_maybe(
             else None
         )
 
+    async def _load_partial_shard_maybe(
+        self,
+        byte_getter: ByteGetter,
+        prototype: BufferPrototype,
+        chunks_per_shard: tuple[int, ...],
+        all_chunk_coords: set[tuple[int, ...]],
+    ) -> ShardMapping | None:
+        """
+        Read chunks from `byte_getter` for the case where the read is less than a full shard.
+        Returns a mapping of chunk coordinates to bytes or None.
+        """
+        shard_index = await self._load_shard_index_maybe(byte_getter, chunks_per_shard)
+        if shard_index is None:
+            return None
+
+        # Pair up chunks and their byte ranges as list[tuple[chunk_coord, byte_range]]
+        chunk_coord_byte_ranges: list[tuple[tuple[int, ...], RangeByteRequest]] = []
+        for chunk_coord in all_chunk_coords:
+            chunk_byte_slice = shard_index.get_chunk_slice(chunk_coord)
+            if chunk_byte_slice is not None:
+                chunk_coord_byte_ranges.append(
+                    (chunk_coord, RangeByteRequest(chunk_byte_slice[0], chunk_byte_slice[1]))
+                )
+
+        if not chunk_coord_byte_ranges:
+            return {}
+
+        shard_dict: ShardMutableMapping = {}
+        if isinstance(byte_getter, StorePath):
+            # External store: use Store.get_ranges for coalescing + concurrency.
+            byte_ranges = [byte_range for _, byte_range in chunk_coord_byte_ranges]
+            try:
+                async for group in byte_getter.store.get_ranges(
+                    byte_getter.path, byte_ranges, prototype=prototype
+                ):
+                    for idx, buf in group:
+                        if buf is not None:
+                            chunk_coord, _ = chunk_coord_byte_ranges[idx]
+                            shard_dict[chunk_coord] = buf
+            except BaseExceptionGroup as eg:
+                # `Store.get_ranges` raises FileNotFoundError (wrapped in a
+                # BaseExceptionGroup) if any underlying fetch indicates the key is
+                # absent. The shard index loaded above, so this typically means a
+                # race where the shard was deleted mid-read; treat it as "shard
+                # gone" to match the index-missing branch (return None). Anything
+                # else in the group (e.g. IO errors) is re-raised.
+                _, rest = eg.split(FileNotFoundError)
+                if rest is not None:
+                    raise rest from None
+                return None
+        else:
+            # Any other ByteGetter. In practice only `_ShardingByteGetter` for
+            # nested sharding, which slices an in-memory buffer (no I/O to coalesce).
+            for chunk_coord, byte_range in chunk_coord_byte_ranges:
+                buf = await byte_getter.get(prototype, byte_range)
+                if buf is not None:
+                    shard_dict[chunk_coord] = buf
+
+        return shard_dict
+
     def compute_encoded_size(self, input_byte_length: int, shard_spec: ArraySpec) -> int:
         chunks_per_shard = self._get_chunks_per_shard(shard_spec)
         return input_byte_length + self._shard_index_size(chunks_per_shard)