From dc678ac100df6fa0716579ea50ea12b26c73e2ec Mon Sep 17 00:00:00 2001 From: Ben Dichter Date: Sun, 4 Oct 2026 16:02:07 -0400 Subject: [PATCH] Narrow partial reads of uncompressed chunks within a single row When a selection touches one index along the first axis, apply the same row window along the next axis, and so on. arr[t, y0:y1, x0:x1] then fetches rows y0 to y1 of frame t instead of the whole frame, and a single item is fetched for a point. The fetched block keeps the narrowed axes with length 1, so the shifted selection applies to it unchanged. Co-Authored-By: Claude Opus 5.5 --- changes/4458.feature.md | 2 +- src/zarr/codecs/bytes.py | 138 +++++++++++++++++++------------- tests/test_codecs/test_bytes.py | 100 +++++++++++++++++++---- 3 files changed, 171 insertions(+), 69 deletions(-) diff --git a/changes/4458.feature.md b/changes/4458.feature.md index eb4497c9b4..edca09ec76 100644 --- a/changes/4458.feature.md +++ b/changes/4458.feature.md @@ -1 +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. +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. A selection that touches a single row is narrowed in the same way along the next axis, so `arr[t, y0:y1, x0:x1]` fetches only rows `y0` to `y1` of frame `t`. 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 621dca78e9..9ebcaa4c02 100644 --- a/src/zarr/codecs/bytes.py +++ b/src/zarr/codecs/bytes.py @@ -1,5 +1,6 @@ from __future__ import annotations +import math import sys import warnings from dataclasses import dataclass, replace @@ -153,14 +154,16 @@ async def _decode_partial_single( 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. + selection is applied to them. When the selection touches a single row, + the same is done within that row along the next axis, and so on, so + that `arr[t, y0:y1, x0:x1]` fetches rows `y0` to `y1` of frame `t`. A + selection that touches every row reads the whole chunk, as before. """ - window = _row_window(selection, chunk_spec.shape) + window = _block_window(selection, chunk_spec.shape) chunk_bytes = await byte_getter.get( - prototype=chunk_spec.prototype, byte_range=_row_byte_range(window, chunk_spec) + prototype=chunk_spec.prototype, byte_range=_block_byte_range(window, chunk_spec) ) - return self._decode_rows(chunk_bytes, selection, window, chunk_spec) + return self._decode_block(chunk_bytes, selection, window, chunk_spec) def _decode_partial_sync( self, @@ -169,28 +172,29 @@ def _decode_partial_sync( chunk_spec: ArraySpec, ) -> NDBuffer | None: """Sync equivalent of `_decode_partial_single`.""" - window = _row_window(selection, chunk_spec.shape) + window = _block_window(selection, chunk_spec.shape) chunk_bytes = byte_getter.get_sync( - prototype=chunk_spec.prototype, byte_range=_row_byte_range(window, chunk_spec) + prototype=chunk_spec.prototype, byte_range=_block_byte_range(window, chunk_spec) ) - return self._decode_rows(chunk_bytes, selection, window, chunk_spec) + return self._decode_block(chunk_bytes, selection, window, chunk_spec) - def _decode_rows( + def _decode_block( self, chunk_bytes: Buffer | None, selection: SelectorTuple, - window: tuple[int, int, SelectorTuple] | None, + window: tuple[int, tuple[int, ...], SelectorTuple] | None, chunk_spec: ArraySpec, ) -> NDBuffer | None: - """Decode the bytes fetched for a row window and apply the selection to them.""" + """Decode the bytes fetched for a block 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 + offset, block_shape, block_selection = window + itemsize = _itemsize(chunk_spec) + start = offset * itemsize + end = start + math.prod(block_shape) * itemsize + chunk_size = math.prod(chunk_spec.shape) * itemsize # 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: @@ -206,8 +210,8 @@ def _decode_rows( 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] + block_spec = replace(chunk_spec, shape=block_shape) + return self._decode_sync(chunk_bytes, block_spec)[block_selection] def _encode_sync( self, @@ -239,58 +243,84 @@ def compute_encoded_size(self, input_byte_length: int, _chunk_spec: ArraySpec) - return input_byte_length -def _row_window( +def _block_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. +) -> tuple[int, tuple[int, ...], SelectorTuple] | None: + """The contiguous block of a C-order chunk that a selection touches. + + Along the first axis the block spans the first to the last row the + selection touches. If that is a single row, the block is narrowed in the + same way along the next axis, and so on. The axes after that are whole. + + Returns the offset of the block from the start of the chunk in items, the + shape of the block, and the selection shifted so that it indexes the block. + The block has as many axes as the chunk; the axes it was narrowed to a + single index along have length 1. Returns None when the block is the whole + chunk or cannot be determined, in which case the whole chunk is read. """ - if len(shape) == 0 or not isinstance(selection, tuple) or len(selection) == 0: + if not isinstance(selection, tuple): + return None + offset = 0 + block_shape = list(shape) + shifted = list(selection) + for axis, (selector, size) in enumerate(zip(selection, shape, strict=False)): + window = _axis_window(selector, size) + if window is None: + break + first, stop, shifted[axis] = window + offset += first * math.prod(shape[axis + 1 :]) + block_shape[axis] = stop - first + if stop - first > 1: + break + if tuple(block_shape) == shape: return None - first_axis = selection[0] + return offset, tuple(block_shape), tuple(shifted) + + +def _axis_window(selector: Selector, size: int) -> tuple[int, int, Selector] | None: + """The indices along one axis that a selector touches, and the selector relative to them. + + Returns the first index, one past the last index, and the selector shifted + so that it indexes an axis holding only those indices. Returns None when + the indices cannot be determined. + """ shifted: Selector - if isinstance(first_axis, slice): - rows = range(*first_axis.indices(shape[0])) - if len(rows) == 0 or rows.step < 0: + if isinstance(selector, slice): + indices = range(*selector.indices(size)) + if len(indices) == 0 or indices.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] + first, stop = indices[0], indices[-1] + 1 + shifted = slice(0, stop - first, indices.step) + elif isinstance(selector, int | np.integer): + first = int(selector) % size stop = first + 1 shifted = 0 - elif isinstance(first_axis, np.ndarray): - if first_axis.dtype == bool: - indices = np.nonzero(first_axis)[0] + elif isinstance(selector, np.ndarray): + if selector.dtype == bool: + positions = np.nonzero(selector)[0] else: - indices = first_axis % shape[0] - if indices.size == 0: + positions = selector % size + if positions.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 + first, stop = int(positions.min()), int(positions.max()) + 1 + shifted = selector[first:stop] if selector.dtype == bool else positions - first else: return None - if (first, stop) == (0, shape[0]): - return None - return first, stop, (shifted, *selection[1:]) + return first, stop, shifted -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 _itemsize(chunk_spec: ArraySpec) -> int: + """The number of bytes in one item of a chunk.""" + return chunk_spec.dtype.to_native_dtype().itemsize -def _row_byte_range( - window: tuple[int, int, SelectorTuple] | None, chunk_spec: ArraySpec +def _block_byte_range( + window: tuple[int, tuple[int, ...], SelectorTuple] | None, chunk_spec: ArraySpec ) -> RangeByteRequest | None: - """The byte range holding a row window, or None to read the whole chunk.""" + """The byte range holding a block 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) + offset, block_shape, _ = window + itemsize = _itemsize(chunk_spec) + start = offset * itemsize + return RangeByteRequest(start, start + math.prod(block_shape) * itemsize) diff --git a/tests/test_codecs/test_bytes.py b/tests/test_codecs/test_bytes.py index 8ac4df7ba1..de2cacd3c2 100644 --- a/tests/test_codecs/test_bytes.py +++ b/tests/test_codecs/test_bytes.py @@ -3,6 +3,7 @@ from __future__ import annotations import enum +import math import sys import warnings from typing import TYPE_CHECKING, Any, Literal, cast @@ -19,7 +20,7 @@ BytesCodec, Endian, EndianLiteral, - _row_window, + _block_window, ) from zarr.core.array_spec import ArrayConfig, ArraySpec from zarr.core.buffer import NDBuffer, default_buffer_prototype @@ -392,22 +393,52 @@ def _single_chunk_array( @pytest.mark.usefixtures("codec_pipeline") @pytest.mark.parametrize( - ("selection", "rows_read"), + ("selection", "items_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), + (np.s_[500:600, 3], 100 * 10), + (np.s_[500:600], 100 * 10), + (np.s_[777, 5], 1), # a single row is narrowed to the columns it needs + (np.s_[777, 2:6], 4), + (np.s_[777:778, 2:6], 4), + (np.s_[-1], 10), + (np.s_[10:20:3, ::2], 10 * 10), # rows 10 to 19 are fetched + (np.s_[...], 10_000 * 10), ], ) -def test_uncompressed_partial_read(selection: Any, rows_read: int) -> None: - """Reading part of an uncompressed chunk fetches only the rows it touches.""" +def test_uncompressed_partial_read(selection: Any, items_read: int) -> None: + """Reading part of an uncompressed chunk fetches only the block 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 + assert sum(n for *_, n in store.reads) == items_read * 2 + + +@pytest.mark.usefixtures("codec_pipeline") +@pytest.mark.parametrize( + ("selection", "byte_range"), + [ + ( + np.s_[7, 10:20, 5:9], + (7 * 1200 + 10 * 30, 7 * 1200 + 20 * 30), + ), # rows 10 to 19 of frame 7 + (np.s_[7, 10:20], (7 * 1200 + 10 * 30, 7 * 1200 + 20 * 30)), + (np.s_[7, 10, 5:9], (7 * 1200 + 10 * 30 + 5, 7 * 1200 + 10 * 30 + 9)), + (np.s_[7, 10, 5], (7 * 1200 + 10 * 30 + 5, 7 * 1200 + 10 * 30 + 6)), + (np.s_[7:8, 10:11, 5:9], (7 * 1200 + 10 * 30 + 5, 7 * 1200 + 10 * 30 + 9)), + (np.s_[7, :, 5:9], (7 * 1200, 8 * 1200)), # every row of frame 7 + (np.s_[7], (7 * 1200, 8 * 1200)), + (np.s_[7:9, 10:20, 5:9], (7 * 1200, 9 * 1200)), # two frames are fetched whole + ], +) +def test_uncompressed_partial_read_within_a_row( + selection: Any, byte_range: tuple[int, int] +) -> None: + """A selection that touches one index along the leading axes is narrowed along the next.""" + data = np.arange(50 * 40 * 30, dtype="uint8").reshape(50, 40, 30) + arr, store = _single_chunk_array(data, compressors=None) + np.testing.assert_array_equal(arr[selection], data[selection]) + ((_, served, _),) = store.reads + assert (served.start, served.end) == byte_range @pytest.mark.usefixtures("codec_pipeline") @@ -427,6 +458,19 @@ def test_uncompressed_partial_read_values(dtype: str, ndim: int) -> None: 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]) + # selections that touch a single index along the leading axes + single: list[Any] = [np.s_[4, 2:], np.s_[4, -1], np.s_[4:5, 1:3], np.s_[4, ..., 0]] + if ndim == 3: + single += [np.s_[4, 2, 1:], np.s_[4, 2:4, 1], np.s_[4, :, 1], np.s_[4, 3, ::2]] + for selection in single if ndim > 1 else []: + np.testing.assert_array_equal(arr[selection], data[selection]) + if ndim > 1: + np.testing.assert_array_equal(arr.oindex[[4], [3, 1]], data[[4]][:, [3, 1]]) + mask = np.zeros(shape[1], dtype=bool) + mask[[1, 3]] = True + np.testing.assert_array_equal(arr.oindex[4, mask], data[4][mask]) + point = tuple(np.array([n - 1, n - 1]) for n in shape) + np.testing.assert_array_equal(arr.vindex[point], data[point]) @pytest.mark.usefixtures("codec_pipeline") @@ -543,10 +587,38 @@ def test_uncompressed_partial_read_store_sends_unexpected_length() -> None: ((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 + ((0, slice(None)), (1, 4)), # the only row, and every column of it (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 +def test_block_window_reads_whole_chunk(selection: Any, shape: tuple[int, ...]) -> None: + """Selections with no narrower block than the whole chunk give None.""" + assert _block_window(selection, shape) is None + + +@pytest.mark.parametrize( + ("selection", "shape", "offset", "block_shape"), + [ + ((slice(2, 5), slice(None)), (10, 4), 8, (3, 4)), + ((3, slice(1, 3)), (10, 4), 13, (1, 2)), + ((3, 2), (10, 4), 14, (1, 1)), + ((0, 2), (1, 4), 2, (1, 1)), # an axis of length 1 is passed through + ((3, slice(2, 2)), (10, 4), 12, (1, 4)), # no columns: stop at the row + ((3, ..., 1), (10, 4, 5), 60, (1, 4, 5)), # stop at the ellipsis + ((3,), (10, 4), 12, (1, 4)), # fewer selectors than axes + ((np.array([3, 3]), np.array([1, 2])), (10, 4), 13, (1, 2)), + ((2, 1, slice(1, 4)), (5, 4, 6), 2 * 24 + 6 + 1, (1, 1, 3)), + ], +) +def test_block_window( + selection: Any, shape: tuple[int, ...], offset: int, block_shape: tuple[int, ...] +) -> None: + """The block starts at the first item the selection touches along the narrowed axes.""" + window = _block_window(selection, shape) + assert window is not None + assert window[:2] == (offset, block_shape) + # the shifted selection picks the same items from the block as the selection does from the chunk + data = np.arange(math.prod(shape)).reshape(shape) + block = data.reshape(-1)[offset : offset + math.prod(block_shape)].reshape(block_shape) + np.testing.assert_array_equal(block[window[2]], data[selection])