feat: return a useful value from CodecPipeline.read() (#3828)

* Add a new type `GetResult` to `src/zarr/abc/codec.py`.

This type is a typeddict that stores information about a completed
attempt to fetch a stored resource, e.g. a chunk. Currently, `GetResult`
only stores whether the resource was present or missing. We can add more capacity
to this type in the future.

* add tests for codec pipeline

* Remove whitespace

* changelog

Author: d-v-b

changes/3828.misc.md

@@ -0,0 +1,2 @@
+`CodecPipeline.read` and `CodecPipeline.read_batch` now return a tuple of typeddict objects
+that each carry information about the request for a chunk from storage.

src/zarr/abc/codec.py

@@ -2,7 +2,7 @@
 
 from abc import abstractmethod
 from collections.abc import Mapping
-from typing import TYPE_CHECKING, Generic, Protocol, TypeGuard, TypeVar, runtime_checkable
+from typing import TYPE_CHECKING, Generic, Literal, Protocol, TypeGuard, TypeVar, runtime_checkable
 
 from typing_extensions import ReadOnly, TypedDict
 
@@ -32,9 +32,17 @@
     "CodecInput",
     "CodecOutput",
     "CodecPipeline",
+    "GetResult",
     "SupportsSyncCodec",
 ]
 
+
+class GetResult(TypedDict):
+    """Metadata about a store get operation."""
+
+    status: Literal["present", "missing"]
+
+
 CodecInput = TypeVar("CodecInput", bound=NDBuffer | Buffer)
 CodecOutput = TypeVar("CodecOutput", bound=NDBuffer | Buffer)
 
@@ -433,13 +441,13 @@ async def read(
         batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
-    ) -> None:
+    ) -> tuple[GetResult, ...]:
         """Reads chunk data from the store, decodes it and writes it into an output array.
         Partial decoding may be utilized if the codecs and stores support it.
 
         Parameters
         ----------
-        batch_info : Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple]]
+        batch_info : Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]]
             Ordered set of information about the chunks.
             The first slice selection determines which parts of the chunk will be fetched.
             The second slice selection determines where in the output array the chunk data will be written.
@@ -451,6 +459,11 @@ async def read(
             ``out``) to the fill value for the array.
 
         out : NDBuffer
+
+        Returns
+        -------
+        tuple[GetResult, ...]
+            One result per chunk in ``batch_info``.
         """
         ...
 

src/zarr/core/codec_pipeline.py

@@ -13,6 +13,7 @@
     BytesBytesCodec,
     Codec,
     CodecPipeline,
+    GetResult,
 )
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
@@ -251,47 +252,58 @@ async def read_batch(
         batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
-    ) -> None:
+    ) -> tuple[GetResult, ...]:
+        results: list[GetResult] = []
         if self.supports_partial_decode:
+            batch_info_list = list(batch_info)
             chunk_array_batch = await self.decode_partial_batch(
                 [
                     (byte_getter, chunk_selection, chunk_spec)
-                    for byte_getter, chunk_spec, chunk_selection, *_ in batch_info
+                    for byte_getter, chunk_spec, chunk_selection, *_ in batch_info_list
                 ]
             )
             for chunk_array, (_, chunk_spec, _, out_selection, _) in zip(
-                chunk_array_batch, batch_info, strict=False
+                chunk_array_batch, batch_info_list, strict=False
             ):
                 if chunk_array is not None:
                     if drop_axes:
                         chunk_array = chunk_array.squeeze(axis=drop_axes)
                     out[out_selection] = chunk_array
+                    results.append(GetResult(status="present"))
                 else:
                     out[out_selection] = fill_value_or_default(chunk_spec)
+                    results.append(GetResult(status="missing"))
         else:
+            batch_info_list = list(batch_info)
             chunk_bytes_batch = await concurrent_map(
-                [(byte_getter, array_spec.prototype) for byte_getter, array_spec, *_ in batch_info],
+                [
+                    (byte_getter, array_spec.prototype)
+                    for byte_getter, array_spec, *_ in batch_info_list
+                ],
                 lambda byte_getter, prototype: byte_getter.get(prototype),
                 config.get("async.concurrency"),
             )
             chunk_array_batch = await self.decode_batch(
                 [
                     (chunk_bytes, chunk_spec)
                     for chunk_bytes, (_, chunk_spec, *_) in zip(
-                        chunk_bytes_batch, batch_info, strict=False
+                        chunk_bytes_batch, batch_info_list, strict=False
                     )
                 ],
             )
             for chunk_array, (_, chunk_spec, chunk_selection, out_selection, _) in zip(
-                chunk_array_batch, batch_info, strict=False
+                chunk_array_batch, batch_info_list, strict=False
             ):
                 if chunk_array is not None:
                     tmp = chunk_array[chunk_selection]
                     if drop_axes:
                         tmp = tmp.squeeze(axis=drop_axes)
                     out[out_selection] = tmp
+                    results.append(GetResult(status="present"))
                 else:
                     out[out_selection] = fill_value_or_default(chunk_spec)
+                    results.append(GetResult(status="missing"))
+        return tuple(results)
 
     def _merge_chunk_array(
         self,
@@ -471,15 +483,19 @@ async def read(
         batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
         out: NDBuffer,
         drop_axes: tuple[int, ...] = (),
-    ) -> None:
-        await concurrent_map(
+    ) -> tuple[GetResult, ...]:
+        batch_results = await concurrent_map(
             [
                 (single_batch_info, out, drop_axes)
                 for single_batch_info in batched(batch_info, self.batch_size)
             ],
             self.read_batch,
             config.get("async.concurrency"),
         )
+        results: list[GetResult] = []
+        for batch in batch_results:
+            results.extend(batch)
+        return tuple(results)
 
     async def write(
         self,

tests/test_codec_pipeline.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import pytest
+
+import zarr
+from zarr.core.buffer.core import default_buffer_prototype
+from zarr.core.indexing import BasicIndexer
+from zarr.storage import MemoryStore
+
+
+@pytest.mark.parametrize(
+    ("write_slice", "read_slice", "expected_statuses"),
+    [
+        # Write all chunks, read all — all present
+        (slice(None), slice(None), ("present", "present", "present")),
+        # Write first chunk only, read all — first present, rest missing
+        (slice(0, 2), slice(None), ("present", "missing", "missing")),
+        # Write nothing, read all — all missing
+        (None, slice(None), ("missing", "missing", "missing")),
+    ],
+)
+async def test_read_returns_get_results(
+    write_slice: slice | None,
+    read_slice: slice,
+    expected_statuses: tuple[str, ...],
+) -> None:
+    """
+    Test that CodecPipeline.read returns a tuple of GetResult with correct statuses.
+    """
+    store = MemoryStore()
+    arr = zarr.open_array(store, mode="w", shape=(6,), chunks=(2,), dtype="int64", fill_value=-1)
+
+    if write_slice is not None:
+        arr[write_slice] = 0
+
+    async_arr = arr._async_array
+    pipeline = async_arr.codec_pipeline
+    metadata = async_arr.metadata
+
+    prototype = default_buffer_prototype()
+    config = async_arr.config
+    indexer = BasicIndexer(
+        read_slice,
+        shape=metadata.shape,
+        chunk_grid=metadata.chunk_grid,
+    )
+
+    out_buffer = prototype.nd_buffer.empty(
+        shape=indexer.shape,
+        dtype=metadata.dtype.to_native_dtype(),
+        order=config.order,
+    )
+
+    results = await pipeline.read(
+        [
+            (
+                async_arr.store_path / metadata.encode_chunk_key(chunk_coords),
+                metadata.get_chunk_spec(chunk_coords, config, prototype=prototype),
+                chunk_selection,
+                out_selection,
+                is_complete_chunk,
+            )
+            for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer
+        ],
+        out_buffer,
+        drop_axes=indexer.drop_axes,
+    )
+
+    assert len(results) == len(expected_statuses)
+    for result, expected_status in zip(results, expected_statuses, strict=True):
+        assert result["status"] == expected_status