Merge branch 'main' into doc/threading-max-workers-docs

Author: d-v-b

changes/2720.doc.md

@@ -0,0 +1 @@
+Document removal of `zarr.storage.init_group` in v3 migration guide, with replacement using `zarr.open_group`/`zarr.create_group`.

changes/3748.feature.md

@@ -0,0 +1 @@
+Added `array.read_missing_chunks` configuration option. When set to `False`, reading missing chunks raises a `ChunkNotFoundError` instead of filling them with the array's fill value.

docs/user-guide/arrays.md

@@ -158,13 +158,25 @@ print(f"Shape after second append: {z.shape}")
 
 Zarr arrays are parametrized with a configuration that determines certain aspects of array behavior.
 
-We currently support two configuration options for arrays: `write_empty_chunks` and `order`.
+We currently support three configuration options for arrays: `write_empty_chunks`, `read_missing_chunks`, and `order`.
 
 | field | type | default | description |
 | - |     - | - | - |
 | `write_empty_chunks` | `bool` | `False` | Controls whether empty chunks are written to storage. See [Empty chunks](performance.md#empty-chunks).
+| `read_missing_chunks` | `bool` | `True` | Controls whether missing chunks are filled with the array's fill value on read. If `False`, reading missing chunks raises a [`ChunkNotFoundError`][zarr.errors.ChunkNotFoundError].
 | `order` | `Literal["C", "F"]` | `"C"` | The memory layout of arrays returned when reading data from the store.
 
+!!! info
+    The Zarr V3 spec states that readers should interpret an uninitialized chunk as containing the
+    array's `fill_value`. By default, Zarr-Python follows this behavior: a missing chunk is treated
+    as uninitialized and filled with the array's `fill_value`. However, if you know that all chunks
+    have been written (i.e., are initialized), you may want to treat a missing chunk as an error. Set
+    `read_missing_chunks=False` to raise a [`ChunkNotFoundError`][zarr.errors.ChunkNotFoundError] instead.
+
+!!! note
+    `write_empty_chunks=False` skips writing chunks that are entirely the array's fill value.
+    If `read_missing_chunks=False`, attempting to read these missing chunks will raise a [`ChunkNotFoundError`][zarr.errors.ChunkNotFoundError].
+
 You can specify the configuration when you create an array with the `config` keyword argument.
 `config` can be passed as either a `dict` or an `ArrayConfig` object.
 

docs/user-guide/config.md

@@ -30,6 +30,7 @@ Configuration options include the following:
 - Default Zarr format `default_zarr_version`
 - Default array order in memory `array.order`
 - Whether empty chunks are written to storage `array.write_empty_chunks`
+- Whether missing chunks are filled with the array's fill value on read `array.read_missing_chunks` (default `True`). Set to `False` to raise a [`ChunkNotFoundError`][zarr.errors.ChunkNotFoundError] instead.
 - Async and threading options, e.g. `async.concurrency` and `threading.max_workers`
 - Selections of implementations of codecs, codec pipelines and buffers
 - Enabling GPU support with `zarr.config.enable_gpu()`. See GPU support for more.

docs/user-guide/v3_migration.md

@@ -114,6 +114,15 @@ The following sections provide details on breaking changes in Zarr-Python 3.
    - Use [`zarr.Group.require_array`][] in place of `zarr.Group.require_dataset`
 3. Disallow "." syntax for getting group members. To get a member of a group named `foo`,
    use `group["foo"]` in place of `group.foo`.
+4. The `zarr.storage.init_group` low-level helper function has been removed. Use
+   [`zarr.open_group`][] or [`zarr.create_group`][] instead:
+
+   ```diff
+   - from zarr.storage import init_group
+   - init_group(store, overwrite=True, path="my/path")
+   + import zarr
+   + zarr.open_group(store, mode="w", path="my/path")
+   ```
 
 ### The Store class
 

src/zarr/abc/codec.py

@@ -67,20 +67,19 @@ def _check_codecjson_v2(data: object) -> TypeGuard[CodecJSON_V2[str]]:
 
 
 @runtime_checkable
-class SupportsSyncCodec(Protocol):
+class SupportsSyncCodec[CI: CodecInput, CO: CodecOutput](Protocol):
     """Protocol for codecs that support synchronous encode/decode.
 
-    Codecs implementing this protocol provide ``_decode_sync`` and ``_encode_sync``
+    Codecs implementing this protocol provide `_decode_sync` and `_encode_sync`
     methods that perform encoding/decoding without requiring an async event loop.
+
+    The type parameters mirror `BaseCodec`: `CI` is the decoded type and `CO` is
+    the encoded type.
     """
 
-    def _decode_sync(
-        self, chunk_data: NDBuffer | Buffer, chunk_spec: ArraySpec
-    ) -> NDBuffer | Buffer: ...
+    def _decode_sync(self, chunk_data: CO, chunk_spec: ArraySpec) -> CI: ...
 
-    def _encode_sync(
-        self, chunk_data: NDBuffer | Buffer, chunk_spec: ArraySpec
-    ) -> NDBuffer | Buffer | None: ...
+    def _encode_sync(self, chunk_data: CI, chunk_spec: ArraySpec) -> CO | None: ...
 
 
 class BaseCodec[CI: CodecInput, CO: CodecOutput](Metadata):

src/zarr/core/array.py

@@ -117,6 +117,7 @@
 from zarr.core.sync import sync
 from zarr.errors import (
     ArrayNotFoundError,
+    ChunkNotFoundError,
     MetadataValidationError,
     ZarrDeprecationWarning,
     ZarrUserWarning,
@@ -5610,7 +5611,8 @@ async def _get_selection(
             _config = replace(_config, order=order)
 
         # reading chunks and decoding them
-        await codec_pipeline.read(
+        indexed_chunks = list(indexer)
+        results = await codec_pipeline.read(
             [
                 (
                     store_path / metadata.encode_chunk_key(chunk_coords),
@@ -5619,11 +5621,26 @@ async def _get_selection(
                     out_selection,
                     is_complete_chunk,
                 )
-                for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexer
+                for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexed_chunks
             ],
             out_buffer,
             drop_axes=indexer.drop_axes,
         )
+        if _config.read_missing_chunks is False:
+            missing_info = []
+            for i, result in enumerate(results):
+                if result["status"] == "missing":
+                    coords = indexed_chunks[i][0]
+                    key = metadata.encode_chunk_key(coords)
+                    missing_info.append(f"  chunk '{key}' (grid position {coords})")
+            if missing_info:
+                chunks_str = "\n".join(missing_info)
+                raise ChunkNotFoundError(
+                    f"{len(missing_info)} chunk(s) not found in store '{store_path}'.\n"
+                    f"Set the 'array.read_missing_chunks' config to True to fill "
+                    f"missing chunks with the fill value.\n"
+                    f"Missing chunks:\n{chunks_str}"
+                )
     if isinstance(indexer, BasicIndexer) and indexer.shape == ():
         return out_buffer.as_scalar()
     return out_buffer.as_ndarray_like()

src/zarr/core/array_spec.py

@@ -28,6 +28,7 @@ class ArrayConfigParams(TypedDict):
 
     order: NotRequired[MemoryOrder]
     write_empty_chunks: NotRequired[bool]
+    read_missing_chunks: NotRequired[bool]
 
 
 @dataclass(frozen=True)
@@ -41,17 +42,25 @@ class ArrayConfig:
         The memory layout of the arrays returned when reading data from the store.
     write_empty_chunks : bool
         If True, empty chunks will be written to the store.
+    read_missing_chunks : bool
+        If True, missing chunks will be filled with the array's fill value on read.
+        If False, reading missing chunks will raise a ``ChunkNotFoundError``.
     """
 
     order: MemoryOrder
     write_empty_chunks: bool
+    read_missing_chunks: bool
 
-    def __init__(self, order: MemoryOrder, write_empty_chunks: bool) -> None:
+    def __init__(
+        self, order: MemoryOrder, write_empty_chunks: bool, *, read_missing_chunks: bool = True
+    ) -> None:
         order_parsed = parse_order(order)
         write_empty_chunks_parsed = parse_bool(write_empty_chunks)
+        read_missing_chunks_parsed = parse_bool(read_missing_chunks)
 
         object.__setattr__(self, "order", order_parsed)
         object.__setattr__(self, "write_empty_chunks", write_empty_chunks_parsed)
+        object.__setattr__(self, "read_missing_chunks", read_missing_chunks_parsed)
 
     @classmethod
     def from_dict(cls, data: ArrayConfigParams) -> Self:
@@ -62,7 +71,9 @@ def from_dict(cls, data: ArrayConfigParams) -> Self:
         """
         kwargs_out: ArrayConfigParams = {}
         for f in fields(ArrayConfig):
-            field_name = cast("Literal['order', 'write_empty_chunks']", f.name)
+            field_name = cast(
+                "Literal['order', 'write_empty_chunks', 'read_missing_chunks']", f.name
+            )
             if field_name not in data:
                 kwargs_out[field_name] = zarr_config.get(f"array.{field_name}")
             else:
@@ -73,7 +84,11 @@ def to_dict(self) -> ArrayConfigParams:
         """
         Serialize an instance of this class to a dict.
         """
-        return {"order": self.order, "write_empty_chunks": self.write_empty_chunks}
+        return {
+            "order": self.order,
+            "write_empty_chunks": self.write_empty_chunks,
+            "read_missing_chunks": self.read_missing_chunks,
+        }
 
 
 ArrayConfigLike = ArrayConfig | ArrayConfigParams

src/zarr/core/codec_pipeline.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from itertools import islice, pairwise
 from typing import TYPE_CHECKING, Any
 from warnings import warn
@@ -14,6 +14,7 @@
     Codec,
     CodecPipeline,
     GetResult,
+    SupportsSyncCodec,
 )
 from zarr.core.common import concurrent_map
 from zarr.core.config import config
@@ -66,6 +67,111 @@ def fill_value_or_default(chunk_spec: ArraySpec) -> Any:
         return fill_value
 
 
+@dataclass(slots=True, kw_only=True)
+class ChunkTransform:
+    """A synchronous codec chain bound to an ArraySpec.
+
+    Provides `encode` and `decode` for pure-compute codec operations
+    (no IO, no threading, no batching).
+
+    All codecs must implement `SupportsSyncCodec`. Construction will
+    raise `TypeError` if any codec does not.
+    """
+
+    codecs: tuple[Codec, ...]
+    array_spec: ArraySpec
+
+    # (sync codec, input_spec) pairs in pipeline order.
+    _aa_codecs: tuple[tuple[SupportsSyncCodec[NDBuffer, NDBuffer], ArraySpec], ...] = field(
+        init=False, repr=False, compare=False
+    )
+    _ab_codec: SupportsSyncCodec[NDBuffer, Buffer] = field(init=False, repr=False, compare=False)
+    _ab_spec: ArraySpec = field(init=False, repr=False, compare=False)
+    _bb_codecs: tuple[SupportsSyncCodec[Buffer, Buffer], ...] = field(
+        init=False, repr=False, compare=False
+    )
+
+    def __post_init__(self) -> None:
+        non_sync = [c for c in self.codecs if not isinstance(c, SupportsSyncCodec)]
+        if non_sync:
+            names = ", ".join(type(c).__name__ for c in non_sync)
+            raise TypeError(
+                f"All codecs must implement SupportsSyncCodec. The following do not: {names}"
+            )
+
+        aa, ab, bb = codecs_from_list(list(self.codecs))
+
+        aa_codecs: list[tuple[SupportsSyncCodec[NDBuffer, NDBuffer], ArraySpec]] = []
+        spec = self.array_spec
+        for aa_codec in aa:
+            assert isinstance(aa_codec, SupportsSyncCodec)
+            aa_codecs.append((aa_codec, spec))
+            spec = aa_codec.resolve_metadata(spec)
+
+        self._aa_codecs = tuple(aa_codecs)
+        assert isinstance(ab, SupportsSyncCodec)
+        self._ab_codec = ab
+        self._ab_spec = spec
+        bb_sync: list[SupportsSyncCodec[Buffer, Buffer]] = []
+        for bb_codec in bb:
+            assert isinstance(bb_codec, SupportsSyncCodec)
+            bb_sync.append(bb_codec)
+        self._bb_codecs = tuple(bb_sync)
+
+    def decode(
+        self,
+        chunk_bytes: Buffer,
+    ) -> NDBuffer:
+        """Decode a single chunk through the full codec chain, synchronously.
+
+        Pure compute -- no IO.
+        """
+        data: Buffer = chunk_bytes
+        for bb_codec in reversed(self._bb_codecs):
+            data = bb_codec._decode_sync(data, self._ab_spec)
+
+        chunk_array: NDBuffer = self._ab_codec._decode_sync(data, self._ab_spec)
+
+        for aa_codec, spec in reversed(self._aa_codecs):
+            chunk_array = aa_codec._decode_sync(chunk_array, spec)
+
+        return chunk_array
+
+    def encode(
+        self,
+        chunk_array: NDBuffer,
+    ) -> Buffer | None:
+        """Encode a single chunk through the full codec chain, synchronously.
+
+        Pure compute -- no IO.
+        """
+        aa_data: NDBuffer = chunk_array
+        for aa_codec, spec in self._aa_codecs:
+            aa_result = aa_codec._encode_sync(aa_data, spec)
+            if aa_result is None:
+                return None
+            aa_data = aa_result
+
+        ab_result = self._ab_codec._encode_sync(aa_data, self._ab_spec)
+        if ab_result is None:
+            return None
+
+        bb_data: Buffer = ab_result
+        for bb_codec in self._bb_codecs:
+            bb_result = bb_codec._encode_sync(bb_data, self._ab_spec)
+            if bb_result is None:
+                return None
+            bb_data = bb_result
+
+        return bb_data
+
+    def compute_encoded_size(self, byte_length: int, array_spec: ArraySpec) -> int:
+        for codec in self.codecs:
+            byte_length = codec.compute_encoded_size(byte_length, array_spec)
+            array_spec = codec.resolve_metadata(array_spec)
+        return byte_length
+
+
 @dataclass(frozen=True)
 class BatchedCodecPipeline(CodecPipeline):
     """Default codec pipeline.

src/zarr/core/config.py

@@ -96,6 +96,7 @@ def enable_gpu(self) -> ConfigSet:
             "array": {
                 "order": "C",
                 "write_empty_chunks": False,
+                "read_missing_chunks": True,
                 "target_shard_size_bytes": None,
             },
             "async": {"concurrency": 10, "timeout": None},