feat: narrow partial reads of uncompressed chunks within a single row

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.

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)

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])