Merge branch 'main' of github.com:zarr-developers/zarr-python

Author: d-v-b

changes/3417.bugfix.md

@@ -0,0 +1,4 @@
+Fixed `BytesCodec.from_dict` so that `BytesCodec` instances roundtrip to / from
+their dict representation. `BytesCodec.from_dict` now interprets a missing
+`endian` configuration as `endian=None` (matching what `BytesCodec.to_dict`
+emits), instead of falling back to the system's native byte order.

changes/3987.feature.md

@@ -0,0 +1 @@
+Two new fields on `ArrayConfig` control how the sharding codec coalesces partial-shard reads: `sharding_coalesce_max_gap_bytes` (default 1 MiB) and `sharding_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 `sharding_coalesce_max_gap_bytes` and the merged read stays within `sharding_coalesce_max_bytes`. Defaults are seeded from the matching `array.sharding_coalesce_max_gap_bytes` / `array.sharding_coalesce_max_bytes` keys in [`zarr.config`][] at array-creation time, and can be overridden per array by passing `config={...}` to [`zarr.create_array`][].

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 `array.sharding_coalesce_max_gap_bytes` and `array.sharding_coalesce_max_bytes`. Reads of nearby chunks are coalesced into a single request to the store when separated by at most `sharding_coalesce_max_gap_bytes` and the resulting merged read is no larger than `sharding_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.

pyproject.toml

@@ -105,7 +105,7 @@ test = [
     "pytest-benchmark==5.2.3",
     "pytest-codspeed==5.0.3",
     "tomlkit==0.15.0",
-    "uv==0.11.19",
+    "uv==0.11.20",
 ]
 remote-tests = [
     {include-group = "test"},

src/zarr/codecs/bytes.py

@@ -62,6 +62,7 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
             data, "bytes", require_configuration=False
         )
         configuration_parsed = configuration_parsed or {}
+        configuration_parsed.setdefault("endian", None)
         return cls(**configuration_parsed)  # type: ignore[arg-type]
 
     def to_dict(self) -> dict[str, JSON]:

src/zarr/codecs/sharding.py

@@ -527,6 +527,8 @@ async def _decode_partial_single(
                 chunk_spec.prototype,
                 chunks_per_shard,
                 all_chunk_coords,
+                max_gap_bytes=shard_spec.config.sharding_coalesce_max_gap_bytes,
+                max_coalesced_bytes=shard_spec.config.sharding_coalesce_max_bytes,
             )
 
         if shard_dict_maybe is None:
@@ -846,10 +848,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:
@@ -873,7 +881,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:

src/zarr/core/array_spec.py

@@ -7,6 +7,7 @@
     MemoryOrder,
     parse_bool,
     parse_fill_value,
+    parse_int,
     parse_order,
     parse_shapelike,
 )
@@ -29,6 +30,8 @@ class ArrayConfigParams(TypedDict):
     order: NotRequired[MemoryOrder]
     write_empty_chunks: NotRequired[bool]
     read_missing_chunks: NotRequired[bool]
+    sharding_coalesce_max_gap_bytes: NotRequired[int]
+    sharding_coalesce_max_bytes: NotRequired[int]
 
 
 @dataclass(frozen=True)
@@ -45,22 +48,42 @@ class ArrayConfig:
     read_missing_chunks : bool
         If True, missing chunks will be filled with the array's fill value on read.
         If False, reading missing chunks will raise a ``ChunkNotFoundError``.
+    sharding_coalesce_max_gap_bytes : int
+        When reading multiple chunks from the same shard, nearby byte ranges
+        separated by no more than this many bytes are coalesced into a single
+        request to the store.
+    sharding_coalesce_max_bytes : int
+        Requests will not be coalesced if doing so would exceed this byte size.
     """
 
     order: MemoryOrder
     write_empty_chunks: bool
     read_missing_chunks: bool
+    sharding_coalesce_max_gap_bytes: int
+    sharding_coalesce_max_bytes: int
 
     def __init__(
-        self, order: MemoryOrder, write_empty_chunks: bool, *, read_missing_chunks: bool = True
+        self,
+        order: MemoryOrder,
+        write_empty_chunks: bool,
+        *,
+        read_missing_chunks: bool = True,
+        sharding_coalesce_max_gap_bytes: int = 1 << 20,  # 1 MiB
+        sharding_coalesce_max_bytes: int = 16 << 20,  # 16 MiB
     ) -> None:
         order_parsed = parse_order(order)
         write_empty_chunks_parsed = parse_bool(write_empty_chunks)
         read_missing_chunks_parsed = parse_bool(read_missing_chunks)
+        sharding_coalesce_max_gap_bytes_parsed = parse_int(sharding_coalesce_max_gap_bytes)
+        sharding_coalesce_max_bytes_parsed = parse_int(sharding_coalesce_max_bytes)
 
         object.__setattr__(self, "order", order_parsed)
         object.__setattr__(self, "write_empty_chunks", write_empty_chunks_parsed)
         object.__setattr__(self, "read_missing_chunks", read_missing_chunks_parsed)
+        object.__setattr__(
+            self, "sharding_coalesce_max_gap_bytes", sharding_coalesce_max_gap_bytes_parsed
+        )
+        object.__setattr__(self, "sharding_coalesce_max_bytes", sharding_coalesce_max_bytes_parsed)
 
     @classmethod
     def from_dict(cls, data: ArrayConfigParams) -> Self:
@@ -72,7 +95,8 @@ def from_dict(cls, data: ArrayConfigParams) -> Self:
         kwargs_out: ArrayConfigParams = {}
         for f in fields(ArrayConfig):
             field_name = cast(
-                "Literal['order', 'write_empty_chunks', 'read_missing_chunks']", f.name
+                "Literal['order', 'write_empty_chunks', 'read_missing_chunks', 'sharding_coalesce_max_gap_bytes', 'sharding_coalesce_max_bytes']",
+                f.name,
             )
             if field_name not in data:
                 kwargs_out[field_name] = zarr_config.get(f"array.{field_name}")
@@ -88,6 +112,8 @@ def to_dict(self) -> ArrayConfigParams:
             "order": self.order,
             "write_empty_chunks": self.write_empty_chunks,
             "read_missing_chunks": self.read_missing_chunks,
+            "sharding_coalesce_max_gap_bytes": self.sharding_coalesce_max_gap_bytes,
+            "sharding_coalesce_max_bytes": self.sharding_coalesce_max_bytes,
         }
 
 

src/zarr/core/common.py

@@ -217,6 +217,12 @@ def parse_bool(data: Any) -> bool:
     raise ValueError(f"Expected bool, got {data} instead.")
 
 
+def parse_int(data: Any) -> int:
+    if isinstance(data, int) and not isinstance(data, bool):
+        return data
+    raise ValueError(f"Expected int, got {data} instead.")
+
+
 def _warn_write_empty_chunks_kwarg() -> None:
     # TODO: link to docs page on array configuration in this message
     msg = (

src/zarr/core/config.py

@@ -99,6 +99,8 @@ def enable_gpu(self) -> ConfigSet:
                 "read_missing_chunks": True,
                 "target_shard_size_bytes": None,
                 "rectilinear_chunks": False,
+                "sharding_coalesce_max_gap_bytes": 1 << 20,  # 1 MiB
+                "sharding_coalesce_max_bytes": 16 << 20,  # 16 MiB
             },
             "async": {"concurrency": 10, "timeout": None},
             "threading": {"max_workers": None},

tests/test_codecs/test_bytes.py

@@ -5,20 +5,98 @@
 import enum
 import sys
 import warnings
-from typing import Any, cast
+from typing import TYPE_CHECKING, Any, Literal, cast
 
+import numpy as np
 import pytest
 
+import zarr
+from tests.conftest import Expect, ExpectFail
+from zarr.abc.codec import SupportsSyncCodec
 from zarr.codecs.bytes import (
     ENDIAN,
     BytesCodec,
     Endian,
     EndianLiteral,
 )
 from zarr.core.array_spec import ArrayConfig, ArraySpec
-from zarr.core.buffer import default_buffer_prototype
+from zarr.core.buffer import NDBuffer, default_buffer_prototype
+from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.core.dtype.npy.int import Int8, Int32
 from zarr.core.dtype.npy.structured import Struct
+from zarr.storage import StorePath
+
+from .test_codecs import _AsyncArrayProxy
+
+if TYPE_CHECKING:
+    from zarr.abc.store import Store
+
+
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
+@pytest.mark.parametrize("input_dtype", [">u2", "<u2"])
+@pytest.mark.parametrize("store_endian", ["big", "little"])
+async def test_endian(
+    store: Store,
+    input_dtype: Literal[">u2", "<u2"],
+    store_endian: Literal["big", "little"],
+) -> None:
+    """
+    The `bytes` codec stores multi-byte data in the byte order configured on the
+    codec, regardless of the input array's byte order, and reads it back to the
+    original values. The input-dtype/store-endian cross-product exercises the
+    encode-side byteswap (input byte order != store byte order) and the no-op
+    case alike. Compression is disabled so the stored chunk is the codec's raw
+    output and its byte layout can be asserted directly.
+    """
+    data = np.arange(0, 256, dtype=input_dtype).reshape((16, 16))
+    path = "endian"
+    spath = StorePath(store, path)
+    a = await zarr.api.asynchronous.create_array(
+        spath,
+        shape=data.shape,
+        chunks=(16, 16),
+        dtype="uint16",
+        fill_value=0,
+        compressors=None,
+        serializer=BytesCodec(endian=store_endian),
+    )
+
+    await _AsyncArrayProxy(a)[:, :].set(data)
+
+    # The stored chunk is laid out in the byte order configured on the codec.
+    stored = await store.get(f"{path}/c/0/0", prototype=default_buffer_prototype())
+    assert stored is not None
+    expected_dtype = ">u2" if store_endian == "big" else "<u2"
+    assert stored.to_bytes() == data.astype(expected_dtype).tobytes()
+
+    # ... and the data reads back to the original values.
+    readback_data = await _AsyncArrayProxy(a)[:, :].get()
+    assert np.array_equal(data, readback_data)
+
+
+def test_bytes_codec_supports_sync() -> None:
+    assert isinstance(BytesCodec(), SupportsSyncCodec)
+
+
+def test_bytes_codec_sync_roundtrip() -> None:
+    codec = BytesCodec()
+    arr = np.arange(100, dtype="float64")
+    zdtype = get_data_type_from_native_dtype(arr.dtype)
+    spec = ArraySpec(
+        shape=arr.shape,
+        dtype=zdtype,
+        fill_value=zdtype.cast_scalar(0),
+        config=ArrayConfig(order="C", write_empty_chunks=True),
+        prototype=default_buffer_prototype(),
+    )
+    nd_buf: NDBuffer = default_buffer_prototype().nd_buffer.from_numpy_array(arr)
+
+    codec = codec.evolve_from_array_spec(spec)
+
+    encoded = codec._encode_sync(nd_buf, spec)
+    assert encoded is not None
+    decoded = codec._decode_sync(encoded, spec)
+    np.testing.assert_array_equal(arr, decoded.as_numpy_array())
 
 
 @pytest.mark.parametrize("endian", ENDIAN)
@@ -46,6 +124,43 @@ def test_bytes_codec_json_roundtrip(endian: EndianLiteral) -> None:
     assert restored == codec
 
 
+# to_dict and from_dict are inverses over this (endian setting, wire dict) mapping:
+# to_dict turns the endian setting into the dict; from_dict recovers it.
+_ENDIAN_DICT_CASES: list[Expect[EndianLiteral | None, dict[str, Any]]] = [
+    Expect(
+        input="little",
+        output={"name": "bytes", "configuration": {"endian": "little"}},
+        id="little",
+    ),
+    Expect(
+        input="big",
+        output={"name": "bytes", "configuration": {"endian": "big"}},
+        id="big",
+    ),
+    Expect(input=None, output={"name": "bytes"}, id="missing"),
+]
+
+
+@pytest.mark.parametrize("case", _ENDIAN_DICT_CASES, ids=lambda c: c.id)
+def test_to_dict(case: Expect[EndianLiteral | None, dict[str, Any]]) -> None:
+    assert BytesCodec(endian=case.input).to_dict() == case.output
+
+
+@pytest.mark.parametrize("case", _ENDIAN_DICT_CASES, ids=lambda c: c.id)
+def test_from_dict(case: Expect[EndianLiteral | None, dict[str, Any]]) -> None:
+    assert BytesCodec.from_dict(case.output).endian == case.input
+
+
+@pytest.mark.parametrize("endian", ["little", "big", pytest.param(None, id="missing")])
+def test_roundtrip(endian: EndianLiteral | None) -> None:
+    codec = BytesCodec(endian=endian)
+
+    encoded = codec.to_dict()
+    roundtripped = BytesCodec.from_dict(encoded)
+
+    assert codec == roundtripped
+
+
 @pytest.mark.parametrize(
     ("member", "expected"),
     [("little", "little"), ("big", "big")],
@@ -105,14 +220,25 @@ def test_bytes_codec_init_with_deprecated_class_member() -> None:
     assert codec.endian == "little"
 
 
-def test_bytes_codec_rejects_unknown_endian() -> None:
+@pytest.mark.parametrize(
+    "case",
+    [
+        ExpectFail(
+            input="north",
+            exception=ValueError,
+            id="unknown-string",
+            msg="endian must be one of",
+        ),
+    ],
+    ids=lambda c: c.id,
+)
+def test_bytes_codec_rejects_unknown_endian(case: ExpectFail[Any]) -> None:
     """
-    `BytesCodec.__init__` raises `ValueError` when given a string outside
+    `BytesCodec.__init__` raises `ValueError` when given a value outside
     `ENDIAN`, and the error message names the offending parameter.
     """
-    kwargs: dict[str, Any] = {"endian": "north"}
-    with pytest.raises(ValueError, match="endian must be one of"):
-        BytesCodec(**kwargs)
+    with case.raises():
+        BytesCodec(endian=case.input)
 
 
 def test_endian_attribute_error_for_unknown_member() -> None: