use get_partial_values(), remove explicit coallescing and concurrent_map

Author: aldenks

changes/3004.feature.md

@@ -1,7 +1,4 @@
-Optimizes reading multiple chunks from a shard. Reads of nearby chunks within
-the same shard are coalesced to reduce the number of calls to the store.
-After any coalescing, the resulting byte ranges are read in parallel.
-
-Coalescing respects two config options. Reads are coalesced if there are fewer
-than `sharding.read.coalesce_max_gap_bytes` bytes between chunks and the total
-size of the coalesced read is no more than `sharding.read.coalesce_max_bytes`.
+Optimizes reading multiple chunks from a shard.
+Serial calls to `.get()` in the sharding codec have been replaced with
+a single call to `.get_partial_values()` which stores may optimize by making
+concurrent requests and/or coalescing nearby requests to the same shard.

docs/user-guide/config.md

@@ -33,7 +33,6 @@ Configuration options include the following:
 - Async and threading options, e.g. `async.concurrency` and `threading.max_workers`
 - Selections of implementations of codecs, codec pipelines and buffers
 - Enabling GPU support with `zarr.config.enable_gpu()`. See GPU support for more.
-- Control request merging when reading multiple chunks from the same shard with `sharding.read.coalesce_max_gap_bytes` and `sharding.read.coalesce_max_bytes`
 
 For selecting custom implementations of codecs, pipelines, buffers and ndbuffers,
 first register the implementations in the registry and then select them in the config.

src/zarr/abc/store.py

@@ -16,10 +16,19 @@
 
     from zarr.core.buffer import Buffer, BufferPrototype
 
-__all__ = ["ByteGetter", "ByteSetter", "Store", "set_or_delete"]
-
-
-@dataclass
+__all__ = [
+    "ByteGetter",
+    "ByteSetter",
+    "Store",
+    "SupportsDeleteSync",
+    "SupportsGetSync",
+    "SupportsSetSync",
+    "SupportsSyncStore",
+    "set_or_delete",
+]
+
+
+@dataclass(frozen=True, slots=True)
 class RangeByteRequest:
     """Request a specific byte range"""
 
@@ -29,15 +38,15 @@ class RangeByteRequest:
     """The end of the byte range request (exclusive)."""
 
 
-@dataclass
+@dataclass(frozen=True, slots=True)
 class OffsetByteRequest:
     """Request all bytes starting from a given byte offset"""
 
     offset: int
     """The byte offset for the offset range request."""
 
 
-@dataclass
+@dataclass(frozen=True, slots=True)
 class SuffixByteRequest:
     """Request up to the last `n` bytes"""
 
@@ -686,20 +695,57 @@ async def get(
         self, prototype: BufferPrototype, byte_range: ByteRequest | None = None
     ) -> Buffer | None: ...
 
+    async def get_partial_values(
+        self,
+        prototype: BufferPrototype,
+        byte_ranges: Iterable[ByteRequest | None],
+    ) -> list[Buffer | None]: ...
+
 
 @runtime_checkable
 class ByteSetter(Protocol):
     async def get(
         self, prototype: BufferPrototype, byte_range: ByteRequest | None = None
     ) -> Buffer | None: ...
 
+    async def get_partial_values(
+        self,
+        prototype: BufferPrototype,
+        byte_ranges: Iterable[ByteRequest | None],
+    ) -> list[Buffer | None]: ...
+
     async def set(self, value: Buffer) -> None: ...
 
     async def delete(self) -> None: ...
 
     async def set_if_not_exists(self, default: Buffer) -> None: ...
 
 
+@runtime_checkable
+class SupportsGetSync(Protocol):
+    def get_sync(
+        self,
+        key: str,
+        *,
+        prototype: BufferPrototype | None = None,
+        byte_range: ByteRequest | None = None,
+    ) -> Buffer | None: ...
+
+
+@runtime_checkable
+class SupportsSetSync(Protocol):
+    def set_sync(self, key: str, value: Buffer) -> None: ...
+
+
+@runtime_checkable
+class SupportsDeleteSync(Protocol):
+    def delete_sync(self, key: str) -> None: ...
+
+
+@runtime_checkable
+class SupportsSyncStore(SupportsGetSync, SupportsSetSync, SupportsDeleteSync, Protocol): ...
+
+
 async def set_or_delete(byte_setter: ByteSetter, value: Buffer | None) -> None:
     """Set or delete a value in a byte setter
 

src/zarr/codecs/sharding.py

@@ -5,7 +5,7 @@
 from enum import Enum
 from functools import lru_cache
 from operator import itemgetter
-from typing import TYPE_CHECKING, Any, NamedTuple
+from typing import TYPE_CHECKING, Any, NamedTuple, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -37,13 +37,11 @@
 from zarr.core.chunk_grids import ChunkGrid, RegularChunkGrid
 from zarr.core.common import (
     ShapeLike,
-    concurrent_map,
     parse_enum,
     parse_named_configuration,
     parse_shapelike,
     product,
 )
-from zarr.core.config import config
 from zarr.core.dtype.npy.int import UInt64
 from zarr.core.indexing import (
     BasicIndexer,
@@ -102,6 +100,13 @@ async def get(
         start, stop = _normalize_byte_range_index(value, byte_range)
         return value[start:stop]
 
+    async def get_partial_values(
+        self,
+        prototype: BufferPrototype,
+        byte_ranges: Iterable[ByteRequest | None],
+    ) -> list[Buffer | None]:
+        return [await self.get(prototype, br) for br in byte_ranges]
+
 
 @dataclass(frozen=True)
 class _ShardingByteSetter(_ShardingByteGetter, ByteSetter):
@@ -295,14 +300,6 @@ def to_dict_vectorized(
         return result
 
 
-@dataclass(frozen=True)
-class _ChunkCoordsByteSlice:
-    """Holds a core.indexing.ChunkProjection.chunk_coords and its byte range in a serialized shard."""
-
-    chunk_coords: tuple[int, ...]
-    byte_slice: slice
-
-
 @dataclass(frozen=True)
 class ShardingCodec(
     ArrayBytesCodec,
@@ -485,26 +482,19 @@ async def _decode_partial_single(
         all_chunk_coords = {chunk_coords for chunk_coords, *_ in indexed_chunks}
 
         # reading bytes of all requested chunks
-        shard_dict_maybe: ShardMapping | None = {}
+        shard_dict_maybe: ShardMapping | None = 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, chunk_spec.prototype, chunks_per_shard
             )
         else:
             # read some chunks within the shard
-            max_gap_bytes = config.get("sharding.read.coalesce_max_gap_bytes")
-            coalesce_max_bytes = config.get("sharding.read.coalesce_max_bytes")
-            async_concurrency = config.get("async.concurrency")
-
             shard_dict_maybe = await self._load_partial_shard_maybe(
                 byte_getter,
                 chunk_spec.prototype,
                 chunks_per_shard,
                 all_chunk_coords,
-                max_gap_bytes,
-                coalesce_max_bytes,
-                async_concurrency,
             )
 
         if shard_dict_maybe is None:
@@ -789,112 +779,34 @@ async def _load_partial_shard_maybe(
         prototype: BufferPrototype,
         chunks_per_shard: tuple[int, ...],
         all_chunk_coords: set[tuple[int, ...]],
-        max_gap_bytes: int,
-        coalesce_max_bytes: int,
-        async_concurrency: 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.
-
-        Reads are coalesced if there are fewer than `max_gap_bytes` bytes between chunks
-        and the total size of the coalesced read is no more than `coalesce_max_bytes`.
         """
         shard_index = await self._load_shard_index_maybe(byte_getter, chunks_per_shard)
         if shard_index is None:
             return None  # shard index read failure, the ByteGetter returned None
 
-        chunks = [
-            _ChunkCoordsByteSlice(chunk_coords, slice(*chunk_byte_slice))
-            for chunk_coords in all_chunk_coords
-            # Drop chunks where index lookup fails
-            # e.g. empty chunks when write_empty_chunks = False
-            if (chunk_byte_slice := shard_index.get_chunk_slice(chunk_coords))
-        ]
+        # Build parallel lists of chunk coordinates and byte ranges for non-empty chunks
+        chunk_coords_list: list[tuple[int, ...]] = []
+        byte_ranges: list[RangeByteRequest] = []
+        for chunk_coords in all_chunk_coords:
+            chunk_byte_slice = shard_index.get_chunk_slice(chunk_coords)
+            if chunk_byte_slice is not None:
+                chunk_coords_list.append(chunk_coords)
+                byte_ranges.append(RangeByteRequest(chunk_byte_slice[0], chunk_byte_slice[1]))
 
-        groups = self._coalesce_chunks(chunks, max_gap_bytes, coalesce_max_bytes)
+        if not byte_ranges:
+            return {}
 
-        shard_dict: ShardMutableMapping = {}
-        if len(groups) == 1:
-            # Avoid thread start overhead when there's only one group
-            shard_dict_result = await self._get_group_bytes(groups[0], byte_getter, prototype)
-            # can be None if the ByteGetter returned None when reading chunk data
-            if shard_dict_result is not None:
-                shard_dict.update(shard_dict_result)
-        else:
-            shard_dicts = await concurrent_map(
-                [(group, byte_getter, prototype) for group in groups],
-                self._get_group_bytes,
-                async_concurrency,
-            )
+        # Fetch all chunk byte ranges via get_partial_values
+        buffers = await byte_getter.get_partial_values(prototype, byte_ranges)
 
-            for shard_dict_result in shard_dicts:
-                if shard_dict_result is not None:
-                    shard_dict.update(shard_dict_result)
-
-        return shard_dict
-
-    def _coalesce_chunks(
-        self,
-        chunks: list[_ChunkCoordsByteSlice],
-        max_gap_bytes: int,
-        coalesce_max_bytes: int,
-    ) -> list[list[_ChunkCoordsByteSlice]]:
-        """
-        Combine chunks from a single shard into groups that should be read together
-        in a single request to the store.
-        """
-        sorted_chunks = sorted(chunks, key=lambda c: c.byte_slice.start)
-
-        if len(sorted_chunks) == 0:
-            return []
-
-        groups = []
-        current_group = [sorted_chunks[0]]
-
-        for chunk in sorted_chunks[1:]:
-            gap_to_chunk = chunk.byte_slice.start - current_group[-1].byte_slice.stop
-            size_if_coalesced = chunk.byte_slice.stop - current_group[0].byte_slice.start
-            if gap_to_chunk < max_gap_bytes and size_if_coalesced < coalesce_max_bytes:
-                current_group.append(chunk)
-            else:
-                groups.append(current_group)
-                current_group = [chunk]
-
-        groups.append(current_group)
-
-        return groups
-
-    async def _get_group_bytes(
-        self,
-        group: list[_ChunkCoordsByteSlice],
-        byte_getter: ByteGetter,
-        prototype: BufferPrototype,
-    ) -> ShardMapping | None:
-        """
-        Reads a possibly coalesced group of one or more chunks from a shard.
-        Returns a mapping of chunk coordinates to bytes.
-        """
-        # _coalesce_chunks ensures that the group is not empty.
-        group_start = group[0].byte_slice.start
-        group_end = group[-1].byte_slice.stop
-
-        # A single call to retrieve the bytes for the entire group.
-        group_bytes = await byte_getter.get(
-            prototype=prototype,
-            byte_range=RangeByteRequest(group_start, group_end),
-        )
-        if group_bytes is None:
-            return None
-
-        # Extract the bytes corresponding to each chunk in group from group_bytes.
-        shard_dict = {}
-        for chunk in group:
-            chunk_slice = slice(
-                chunk.byte_slice.start - group_start,
-                chunk.byte_slice.stop - group_start,
-            )
-            shard_dict[chunk.chunk_coords] = group_bytes[chunk_slice]
+        shard_dict: ShardMutableMapping = {}
+        for chunk_coords, buf in zip(chunk_coords_list, buffers, strict=True):
+            if buf is not None:
+                shard_dict[chunk_coords] = buf
 
         return shard_dict
 

src/zarr/core/config.py

@@ -100,12 +100,6 @@ def enable_gpu(self) -> ConfigSet:
             },
             "async": {"concurrency": 10, "timeout": None},
             "threading": {"max_workers": None},
-            "sharding": {
-                "read": {
-                    "coalesce_max_bytes": 100 * 2**20,  # 100MiB
-                    "coalesce_max_gap_bytes": 2**20,  # 1MiB
-                }
-            },
             "json_indent": 2,
             "codec_pipeline": {
                 "path": "zarr.core.codec_pipeline.BatchedCodecPipeline",