diff --git a/changes/3987.feature.md b/changes/3987.feature.md new file mode 100644 index 0000000000..bc6030e201 --- /dev/null +++ b/changes/3987.feature.md @@ -0,0 +1 @@ +Two new runtime config options control how the sharding codec coalesces partial-shard reads: `sharding.read.coalesce_max_gap_bytes` (default 1 MiB) and `sharding.read.coalesce_max_bytes` (default 16 MiB). When reading multiple chunks from the same shard, nearby byte ranges are merged into a single request to the store if separated by no more than `coalesce_max_gap_bytes` and the merged read stays within `coalesce_max_bytes`. diff --git a/docs/user-guide/config.md b/docs/user-guide/config.md index 8a8fa94c3d..1063b6f35a 100644 --- a/docs/user-guide/config.md +++ b/docs/user-guide/config.md @@ -35,6 +35,7 @@ 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`. Reads of nearby chunks are coalesced into a single request to the store when separated by at most `coalesce_max_gap_bytes` and the resulting merged read is no larger than `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. diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 41f9b7e9c6..37ee69b512 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -41,6 +41,7 @@ parse_shapelike, product, ) +from zarr.core.config import config from zarr.core.dtype.npy.int import UInt64 from zarr.core.indexing import ( BasicIndexer, @@ -483,6 +484,8 @@ async def _decode_partial_single( chunk_spec.prototype, chunks_per_shard, all_chunk_coords, + max_gap_bytes=config.get("sharding.read.coalesce_max_gap_bytes"), + max_coalesced_bytes=config.get("sharding.read.coalesce_max_bytes"), ) if shard_dict_maybe is None: @@ -780,10 +783,16 @@ async def _load_partial_shard_maybe( prototype: BufferPrototype, chunks_per_shard: tuple[int, ...], all_chunk_coords: set[tuple[int, ...]], + max_gap_bytes: int, + max_coalesced_bytes: 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. + + `max_gap_bytes` and `max_coalesced_bytes` are forwarded to + `Store.get_ranges` to control byte-range coalescing across the requested + chunks. """ shard_index = await self._load_shard_index_maybe(byte_getter, chunks_per_shard) if shard_index is None: @@ -807,7 +816,11 @@ async def _load_partial_shard_maybe( 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 + byte_getter.path, + byte_ranges, + prototype=prototype, + max_gap_bytes=max_gap_bytes, + max_coalesced_bytes=max_coalesced_bytes, ): for idx, buf in group: if buf is not None: diff --git a/src/zarr/core/config.py b/src/zarr/core/config.py index 7dcbc78e31..3dd9fc18a7 100644 --- a/src/zarr/core/config.py +++ b/src/zarr/core/config.py @@ -102,6 +102,12 @@ def enable_gpu(self) -> ConfigSet: }, "async": {"concurrency": 10, "timeout": None}, "threading": {"max_workers": None}, + "sharding": { + "read": { + "coalesce_max_gap_bytes": 1 << 20, # 1 MiB + "coalesce_max_bytes": 16 << 20, # 16 MiB + }, + }, "json_indent": 2, "codec_pipeline": { "path": "zarr.core.codec_pipeline.BatchedCodecPipeline", diff --git a/tests/test_codecs/test_sharding_unit.py b/tests/test_codecs/test_sharding_unit.py index 6e022ed9fa..7b07da7749 100644 --- a/tests/test_codecs/test_sharding_unit.py +++ b/tests/test_codecs/test_sharding_unit.py @@ -1,3 +1,6 @@ +from typing import TYPE_CHECKING, cast +from unittest.mock import AsyncMock + import numpy as np import pytest @@ -10,9 +13,13 @@ ) from zarr.core.buffer import default_buffer_prototype from zarr.core.buffer.cpu import Buffer +from zarr.core.config import config from zarr.storage._common import StorePath from zarr.storage._memory import MemoryStore +if TYPE_CHECKING: + from zarr.core.array import ShardsConfigParam + # ============================================================================ # _ShardIndex tests # ============================================================================ @@ -155,6 +162,8 @@ async def test_load_partial_shard_maybe_index_load_fails() -> None: prototype=default_buffer_prototype(), chunks_per_shard=(2,), all_chunk_coords={(0,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) assert result is None @@ -187,6 +196,8 @@ async def mock_load_index( prototype=default_buffer_prototype(), chunks_per_shard=chunks_per_shard, all_chunk_coords={(0,), (1,), (2,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) assert result is not None @@ -220,6 +231,8 @@ async def mock_load_index( prototype=default_buffer_prototype(), chunks_per_shard=chunks_per_shard, all_chunk_coords={(0,), (1,), (2,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) assert result == {} @@ -251,6 +264,8 @@ async def mock_load_index( prototype=default_buffer_prototype(), chunks_per_shard=chunks_per_shard, all_chunk_coords={(0,), (1,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) assert result is not None @@ -292,6 +307,8 @@ async def mock_load_index( prototype=default_buffer_prototype(), chunks_per_shard=chunks_per_shard, all_chunk_coords={(0,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) assert result is None @@ -336,6 +353,8 @@ async def boom(*args: object, **kwargs: object) -> Buffer | None: prototype=default_buffer_prototype(), chunks_per_shard=chunks_per_shard, all_chunk_coords={(0,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) @@ -368,6 +387,8 @@ async def mock_load_index( prototype=default_buffer_prototype(), chunks_per_shard=chunks_per_shard, all_chunk_coords={(0,), (1,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) assert result is not None @@ -405,6 +426,8 @@ async def mock_load_index( prototype=default_buffer_prototype(), chunks_per_shard=chunks_per_shard, all_chunk_coords={(0,)}, + max_gap_bytes=1 << 20, + max_coalesced_bytes=16 << 20, ) assert result == {} @@ -486,3 +509,114 @@ def test_is_total_shard_1d() -> None: # Partial partial_coords: set[tuple[int, ...]] = {(0,), (2,)} assert codec._is_total_shard(partial_coords, chunks_per_shard) is False + + +# ============================================================================ +# Coalescing config option tests +# +# Asserts that the `sharding.read.coalesce_max_gap_bytes` and +# `sharding.read.coalesce_max_bytes` config options flow through to +# `Store.get_ranges` as `max_gap_bytes` / `max_coalesced_bytes` kwargs. +# ============================================================================ + + +def _trigger_partial_shard_read() -> AsyncMock: + """Build a sharded array on a mocked `MemoryStore`, trigger a partial-shard + read via the public read path, and return the `get_ranges` mock. + """ + import zarr + from zarr.codecs.sharding import ShardingCodecIndexLocation + + chunk_shape = (2,) + shard_shape = (8,) + data = np.arange(8, dtype="int32") + + store = MemoryStore() + store_mock = AsyncMock(wraps=store, spec=store.__class__) + + shards: ShardsConfigParam = { + "shape": shard_shape, + "index_location": ShardingCodecIndexLocation.end, + } + a = zarr.create_array( + StorePath(store_mock), + shape=(8,), + chunks=chunk_shape, + shards=shards, + dtype=data.dtype, + fill_value=-1, + ) + a[:] = data + + store_mock.reset_mock() + + # Read a strict subset of chunks to take the partial-shard read path. + _ = a[0:4] + + return cast(AsyncMock, store_mock.get_ranges) + + +def test_load_partial_shard_forwards_config_to_get_ranges() -> None: + """`sharding.read.*` config values are forwarded to `Store.get_ranges`.""" + with config.set( + { + "sharding.read.coalesce_max_gap_bytes": 4242, + "sharding.read.coalesce_max_bytes": 424242, + } + ): + get_ranges_mock = _trigger_partial_shard_read() + + assert get_ranges_mock.call_count >= 1 + for call in get_ranges_mock.call_args_list: + kwargs = call.kwargs + assert kwargs["max_gap_bytes"] == 4242 + assert kwargs["max_coalesced_bytes"] == 424242 + + +def test_load_partial_shard_uses_config_defaults() -> None: + """Without explicit config, defaults from `zarr.config` are forwarded.""" + get_ranges_mock = _trigger_partial_shard_read() + + assert get_ranges_mock.call_count >= 1 + for call in get_ranges_mock.call_args_list: + kwargs = call.kwargs + assert kwargs["max_gap_bytes"] == config.get("sharding.read.coalesce_max_gap_bytes") + assert kwargs["max_coalesced_bytes"] == config.get("sharding.read.coalesce_max_bytes") + + +async def test_load_partial_shard_explicit_kwargs_passthrough( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """`_load_partial_shard_maybe` forwards its explicit kwargs to `get_ranges`.""" + codec = ShardingCodec(chunk_shape=(2,)) + chunks_per_shard = (4,) + + index = _ShardIndex.create_empty(chunks_per_shard) + index.set_chunk_slice((0,), slice(0, 100)) + index.set_chunk_slice((2,), slice(200, 300)) + + store = MemoryStore() + await store.set("shard", Buffer.from_bytes(b"x" * 300)) + store_mock = AsyncMock(wraps=store, spec=store.__class__) + byte_getter = StorePath(store_mock, "shard") + + async def mock_load_index( + self: ShardingCodec, byte_getter: StorePath, cps: tuple[int, ...] + ) -> _ShardIndex: + return index + + monkeypatch.setattr(ShardingCodec, "_load_shard_index_maybe", mock_load_index) + + await codec._load_partial_shard_maybe( + byte_getter=byte_getter, + prototype=default_buffer_prototype(), + chunks_per_shard=chunks_per_shard, + all_chunk_coords={(0,), (2,)}, + max_gap_bytes=12345, + max_coalesced_bytes=67890, + ) + + store_mock.get_ranges.assert_called_once() + kwargs = store_mock.get_ranges.call_args.kwargs + assert kwargs["max_gap_bytes"] == 12345 + assert kwargs["max_coalesced_bytes"] == 67890 diff --git a/tests/test_config.py b/tests/test_config.py index 4e293e968f..622305880e 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -59,6 +59,12 @@ def test_config_defaults_set() -> None: }, "async": {"concurrency": 10, "timeout": None}, "threading": {"max_workers": None}, + "sharding": { + "read": { + "coalesce_max_gap_bytes": 1 << 20, + "coalesce_max_bytes": 16 << 20, + }, + }, "json_indent": 2, "codec_pipeline": { "path": "zarr.core.codec_pipeline.BatchedCodecPipeline",