feat: Bytes partial decode

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.

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)