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.

Author: d-v-b

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,

src/zarr/core/common.py

@@ -50,6 +50,7 @@
 DimensionNamesLike = Iterable[str | None] | None
 DimensionNames = DimensionNamesLike  # for backwards compatibility
 
+
 TName = TypeVar("TName", bound=str)
 TConfig = TypeVar("TConfig", bound=Mapping[str, object])