diff --git a/changes/4458.feature.md b/changes/4458.feature.md new file mode 100644 index 0000000000..eb4497c9b4 --- /dev/null +++ b/changes/4458.feature.md @@ -0,0 +1 @@ +Reading part of an uncompressed chunk now fetches only the bytes it needs. When `BytesCodec` is the only codec, a selection is read with a single byte range request covering the rows along the first axis that it touches, instead of the whole chunk. Chunks with a compressor or any other codec are read whole, as before. diff --git a/src/zarr/codecs/bytes.py b/src/zarr/codecs/bytes.py index 622161c458..621dca78e9 100644 --- a/src/zarr/codecs/bytes.py +++ b/src/zarr/codecs/bytes.py @@ -3,9 +3,12 @@ import sys import warnings from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, ClassVar, Final, Literal +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal -from zarr.abc.codec import ArrayBytesCodec +import numpy as np + +from zarr.abc.codec import ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin +from zarr.abc.store import RangeByteRequest from zarr.codecs._deprecated_enum import _coerce_enum_input, _DeprecatedStrEnumMeta from zarr.core.common import JSON, parse_named_configuration from zarr.core.dtype.common import HasEndianness @@ -14,8 +17,10 @@ if TYPE_CHECKING: from typing import Self + from zarr.abc.store import ByteGetter from zarr.core.array_spec import ArraySpec from zarr.core.buffer import Buffer, NDBuffer + from zarr.core.indexing import Selector, SelectorTuple EndianLiteral = Literal["little", "big"] @@ -40,7 +45,7 @@ def _parse_endian(data: object) -> EndianLiteral: @dataclass(frozen=True) -class BytesCodec(ArrayBytesCodec): +class BytesCodec(ArrayBytesCodec, ArrayBytesCodecPartialDecodeMixin): """bytes codec""" is_fixed_size = True @@ -137,6 +142,73 @@ async def _decode_single( ) -> NDBuffer: return self._decode_sync(chunk_bytes, chunk_spec) + async def _decode_partial_single( + self, + byte_getter: ByteGetter, + selection: SelectorTuple, + chunk_spec: ArraySpec, + ) -> NDBuffer | None: + """Read only the part of an uncompressed chunk that a selection needs. + + The chunk is stored in C order, so each row along its first axis is a + contiguous run of bytes. The rows from the first to the last one the + selection touches are fetched with a single range request, and the + selection is applied to them. A selection that touches every row reads + the whole chunk, as before. + """ + window = _row_window(selection, chunk_spec.shape) + chunk_bytes = await byte_getter.get( + prototype=chunk_spec.prototype, byte_range=_row_byte_range(window, chunk_spec) + ) + return self._decode_rows(chunk_bytes, selection, window, chunk_spec) + + def _decode_partial_sync( + self, + byte_getter: Any, + selection: SelectorTuple, + chunk_spec: ArraySpec, + ) -> NDBuffer | None: + """Sync equivalent of `_decode_partial_single`.""" + window = _row_window(selection, chunk_spec.shape) + chunk_bytes = byte_getter.get_sync( + prototype=chunk_spec.prototype, byte_range=_row_byte_range(window, chunk_spec) + ) + return self._decode_rows(chunk_bytes, selection, window, chunk_spec) + + def _decode_rows( + self, + chunk_bytes: Buffer | None, + selection: SelectorTuple, + window: tuple[int, int, SelectorTuple] | None, + chunk_spec: ArraySpec, + ) -> NDBuffer | None: + """Decode the bytes fetched for a row window and apply the selection to them.""" + if chunk_bytes is None: + return None + if window is None: + return self._decode_sync(chunk_bytes, chunk_spec)[selection] + first, stop, rows_selection = window + row_bytes = _row_bytes(chunk_spec) + start, end = first * row_bytes, stop * row_bytes + chunk_size = chunk_spec.shape[0] * row_bytes + # A store may not honor the byte range it was given. The length of what + # it sent tells which bytes those are. + if len(chunk_bytes) == end - start: + pass + elif len(chunk_bytes) == chunk_size: + # The whole chunk, as an HTTP server that ignores the Range header sends. + chunk_bytes = chunk_bytes[start:end] + elif len(chunk_bytes) == chunk_size - start: + # Everything from the start of the range to the end of the chunk. + chunk_bytes = chunk_bytes[: end - start] + else: + raise ValueError( + f"Requested bytes {start} to {end} of a {chunk_size} byte chunk, " + f"but the store returned {len(chunk_bytes)} bytes." + ) + rows_spec = replace(chunk_spec, shape=(stop - first, *chunk_spec.shape[1:])) + return self._decode_sync(chunk_bytes, rows_spec)[rows_selection] + def _encode_sync( self, chunk_array: NDBuffer, @@ -165,3 +237,60 @@ async def _encode_single( def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) -> int: return input_byte_length + + +def _row_window( + selection: SelectorTuple, shape: tuple[int, ...] +) -> tuple[int, int, SelectorTuple] | None: + """The rows along axis 0 that a selection touches, and the selection relative to them. + + Returns the first row, one past the last row, and the selection shifted so + that it indexes an array holding only those rows. Returns None when the + selection touches every row or the rows cannot be determined, in which case + the whole chunk is read. + """ + if len(shape) == 0 or not isinstance(selection, tuple) or len(selection) == 0: + return None + first_axis = selection[0] + shifted: Selector + if isinstance(first_axis, slice): + rows = range(*first_axis.indices(shape[0])) + if len(rows) == 0 or rows.step < 0: + return None + first, stop = rows[0], rows[-1] + 1 + shifted = slice(0, stop - first, rows.step) + elif isinstance(first_axis, int | np.integer): + first = int(first_axis) % shape[0] + stop = first + 1 + shifted = 0 + elif isinstance(first_axis, np.ndarray): + if first_axis.dtype == bool: + indices = np.nonzero(first_axis)[0] + else: + indices = first_axis % shape[0] + if indices.size == 0: + return None + first, stop = int(indices.min()), int(indices.max()) + 1 + shifted = first_axis[first:stop] if first_axis.dtype == bool else indices - first + else: + return None + if (first, stop) == (0, shape[0]): + return None + return first, stop, (shifted, *selection[1:]) + + +def _row_bytes(chunk_spec: ArraySpec) -> int: + """The number of bytes in one row along axis 0 of a chunk.""" + row_items = int(np.prod(chunk_spec.shape[1:])) + return chunk_spec.dtype.to_native_dtype().itemsize * row_items + + +def _row_byte_range( + window: tuple[int, int, SelectorTuple] | None, chunk_spec: ArraySpec +) -> RangeByteRequest | None: + """The byte range holding a row window, or None to read the whole chunk.""" + if window is None: + return None + first, stop, _ = window + row_bytes = _row_bytes(chunk_spec) + return RangeByteRequest(first * row_bytes, stop * row_bytes) diff --git a/tests/test_codecs/test_bytes.py b/tests/test_codecs/test_bytes.py index ead778f526..8ac4df7ba1 100644 --- a/tests/test_codecs/test_bytes.py +++ b/tests/test_codecs/test_bytes.py @@ -13,11 +13,13 @@ import zarr from tests.conftest import Expect, ExpectFail from zarr.abc.codec import SupportsSyncCodec +from zarr.abc.store import OffsetByteRequest, RangeByteRequest from zarr.codecs.bytes import ( ENDIAN, BytesCodec, Endian, EndianLiteral, + _row_window, ) from zarr.core.array_spec import ArrayConfig, ArraySpec from zarr.core.buffer import NDBuffer, default_buffer_prototype @@ -29,6 +31,8 @@ from .test_codecs import _AsyncArrayProxy if TYPE_CHECKING: + from collections.abc import Callable, Iterator + from zarr.abc.store import Store @@ -338,3 +342,211 @@ def test_bytes_codec_evolve_structured_single_byte_fields_clears_endian() -> Non spec = _make_array_spec(dtype) evolved = codec.evolve_from_array_spec(spec) assert evolved.endian is None + + +@pytest.fixture( + params=[ + "zarr.core.codec_pipeline.BatchedCodecPipeline", + "zarr.core.codec_pipeline.FusedCodecPipeline", + ] +) +def codec_pipeline(request: pytest.FixtureRequest) -> Iterator[None]: + """Run a test once with each codec pipeline.""" + with zarr.config.set({"codec_pipeline.path": request.param}): + yield + + +class _RangeLoggingStore(zarr.storage.MemoryStore): + """A memory store that records the byte ranges it serves for chunk keys.""" + + def __init__(self) -> None: + super().__init__() + self.reads: list[tuple[str, Any, int]] = [] + + def _log(self, key: str, byte_range: Any, buf: Any) -> None: + if buf is not None and not key.endswith("zarr.json"): + self.reads.append((key, byte_range, len(buf))) + + async def get(self, key: str, prototype: Any = None, byte_range: Any = None) -> Any: + buf = await super().get(key, prototype, byte_range) + self._log(key, byte_range, buf) + return buf + + def get_sync(self, key: str, *, prototype: Any = None, byte_range: Any = None) -> Any: + buf = super().get_sync(key, prototype=prototype, byte_range=byte_range) + self._log(key, byte_range, buf) + return buf + + +def _single_chunk_array( + data: np.ndarray[Any, Any], **kwargs: Any +) -> tuple[Any, _RangeLoggingStore]: + store = _RangeLoggingStore() + arr = zarr.create_array( + store, shape=data.shape, dtype=data.dtype, chunks=data.shape, fill_value=0, **kwargs + ) + arr[...] = data + store.reads.clear() + return arr, store + + +@pytest.mark.usefixtures("codec_pipeline") +@pytest.mark.parametrize( + ("selection", "rows_read"), + [ + (np.s_[500:600, 3], 100), + (np.s_[500:600], 100), + (np.s_[777, 5], 1), + (np.s_[-1], 1), + (np.s_[10:20:3, ::2], 10), # rows 10 to 19 are fetched + (np.s_[...], 10_000), + ], +) +def test_uncompressed_partial_read(selection: Any, rows_read: int) -> None: + """Reading part of an uncompressed chunk fetches only the rows it touches.""" + data = np.arange(100_000, dtype="int16").reshape(10_000, 10) + arr, store = _single_chunk_array(data, compressors=None) + np.testing.assert_array_equal(arr[selection], data[selection]) + assert sum(n for *_, n in store.reads) == rows_read * 10 * 2 + + +@pytest.mark.usefixtures("codec_pipeline") +@pytest.mark.parametrize("dtype", [">u2", " None: + """Partial reads return the same values as numpy for several dtypes, byte orders, and shapes.""" + shape = {1: (1000,), 2: (100, 7), 3: (20, 5, 3)}[ndim] + data = (np.arange(int(np.prod(shape))) % 200).astype(dtype).reshape(shape) + arr, _ = _single_chunk_array(data, compressors=None) + selections: list[Any] = [np.s_[3:9], np.s_[4], np.s_[-3:], np.s_[1:15:4]] + for selection in selections: + np.testing.assert_array_equal(arr[selection], data[selection]) + np.testing.assert_array_equal(arr.oindex[[1, 5, 2]], data[[1, 5, 2]]) + mask = np.zeros(shape[0], dtype=bool) + mask[[2, 5, 11]] = True + np.testing.assert_array_equal(arr.oindex[mask], data[mask]) + coords = tuple(np.array([0, 6, 3]) % n for n in shape) + np.testing.assert_array_equal(arr.vindex[coords], data[coords]) + + +@pytest.mark.usefixtures("codec_pipeline") +def test_uncompressed_partial_read_across_chunks() -> None: + """A selection spanning several chunks reads only the touched rows of each.""" + data = np.arange(40_000, dtype="int32").reshape(4000, 10) + store = _RangeLoggingStore() + arr = zarr.create_array( + store, shape=data.shape, dtype="int32", chunks=(1000, 10), compressors=None + ) + arr[...] = data + store.reads.clear() + np.testing.assert_array_equal(arr[990:1010, 2], data[990:1010, 2]) + assert sorted((key, n) for key, _, n in store.reads) == [("c/0/0", 400), ("c/1/0", 400)] + + +@pytest.mark.usefixtures("codec_pipeline") +def test_compressed_chunks_are_read_whole() -> None: + """With a compressor the chunk bytes cannot be split, so the whole chunk is fetched.""" + data = np.arange(100_000, dtype="int16").reshape(10_000, 10) + arr, store = _single_chunk_array(data) # default compressor + np.testing.assert_array_equal(arr[500:600, 3], data[500:600, 3]) + ((_, byte_range, _),) = store.reads + assert byte_range is None + + +@pytest.mark.usefixtures("codec_pipeline") +@pytest.mark.parametrize("selection", [np.s_[10:20], np.s_[...]]) +def test_uncompressed_partial_read_missing_chunk(selection: Any) -> None: + """An unwritten chunk reads as the fill value, for a partial and a whole-chunk read.""" + store = _RangeLoggingStore() + arr = zarr.create_array( + store, shape=(100, 4), dtype="int16", chunks=(100, 4), compressors=None, fill_value=7 + ) + expected = np.full((100, 4), 7, dtype="int16")[selection] + np.testing.assert_array_equal(arr[selection], expected) + + +class _RewritesRangeStore(zarr.storage.MemoryStore): + """A memory store that does not serve the byte range it is asked for. + + `rewrite` maps the requested range to the one that is served instead. + """ + + def __init__(self, rewrite: Callable[[RangeByteRequest], Any]) -> None: + super().__init__() + self.rewrite = rewrite + + def _served(self, byte_range: Any) -> Any: + if isinstance(byte_range, RangeByteRequest): + return self.rewrite(byte_range) + return byte_range + + async def get(self, key: str, prototype: Any = None, byte_range: Any = None) -> Any: + return await super().get(key, prototype, self._served(byte_range)) + + def get_sync(self, key: str, *, prototype: Any = None, byte_range: Any = None) -> Any: + return super().get_sync(key, prototype=prototype, byte_range=self._served(byte_range)) + + +def _ignore_range(byte_range: RangeByteRequest) -> None: + """Send the whole value, as an HTTP server that ignores the Range header does.""" + + +def _ignore_range_end(byte_range: RangeByteRequest) -> OffsetByteRequest: + """Send everything from the start of the range to the end of the value.""" + return OffsetByteRequest(byte_range.start) + + +def _one_byte_short(byte_range: RangeByteRequest) -> RangeByteRequest: + return RangeByteRequest(byte_range.start, byte_range.end - 1) + + +def _rewriting_array( + rewrite: Callable[[RangeByteRequest], Any], +) -> tuple[Any, np.ndarray[Any, Any]]: + data = np.arange(400, dtype=" None: + """A store that sends more than the requested range still reads correctly.""" + arr, data = _rewriting_array(rewrite) + np.testing.assert_array_equal(arr[selection], data[selection]) + + +@pytest.mark.usefixtures("codec_pipeline") +def test_uncompressed_partial_read_store_sends_unexpected_length() -> None: + """A response whose length matches no known byte range behavior is an error.""" + arr, _ = _rewriting_array(_one_byte_short) + with pytest.raises(ValueError, match="the store returned 15 bytes"): + arr[50:52, 0] + + +@pytest.mark.parametrize( + ("selection", "shape"), + [ + ((slice(5, 5), slice(None)), (10, 4)), # no rows + ((slice(8, 2, -1), slice(None)), (10, 4)), # negative step + ((np.array([], dtype=np.intp), slice(None)), (10, 4)), # no rows + ((slice(None), 2), (10, 4)), # every row + ((np.array([0, 9]), slice(None)), (10, 4)), # first and last row + ((..., 2), (10, 4)), # rows not determined + (slice(2, 4), (10,)), # not a tuple + ((), ()), # zero-dimensional chunk + ], +) +def test_row_window_reads_whole_chunk(selection: Any, shape: tuple[int, ...]) -> None: + """Selections with no narrower row window than the whole chunk give None.""" + assert _row_window(selection, shape) is None