Fill missing chunks (#3748)

* Add `codec_pipeline.fill_missing_chunks` config

* Set default for `fill_missing_chunks` in config.py. Add test replicating
example in zarr-python #486.

* Add fill_missing_chunks to examples of config options.

* Add to /changes

* Parameterize tests to make sure we hit both branches of `if
self.supports_partial_decode`.

* Fix lint errors: remove parentheses, type kwargs.

* Move config from codec_pipeline -> array. Update docs, tests.

* Delegate missing-shard detection away from _get_chunk_spec. Codify
expected behaviour of fill_missing_chunks for both sharding and
write_empty_chunks via tests. Use elif to make control flow slightly
clearer.

* Define ChunkNotFoundError; expose chunk key and chunk index in ChunkNotFoundError

* update docs

* fix links

* cleanup

* Pass chunk indexes up

* fill_missing_chunks -> read_missing_chunks

* Resolve behavioural differences between main and maxrjones/zarr-python@37a40e3.
Update docstrings to match current behaviour. Move description of
sharding behaviour to test, now that it has no dedicated codepath.

---------

Co-authored-by: Tom White <tom.e.white@gmail.com>
Co-authored-by: Davis Bennett <davis.v.bennett@gmail.com>
Co-authored-by: Max Jones <14077947+maxrjones@users.noreply.github.com>

Author: 4 people

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.

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/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},

src/zarr/errors.py

@@ -3,6 +3,7 @@
     "ArrayNotFoundError",
     "BaseZarrError",
     "BoundsCheckError",
+    "ChunkNotFoundError",
     "ContainsArrayAndGroupError",
     "ContainsArrayError",
     "ContainsGroupError",
@@ -144,3 +145,9 @@ class BoundsCheckError(IndexError): ...
 
 
 class ArrayIndexError(IndexError): ...
+
+
+class ChunkNotFoundError(BaseZarrError):
+    """
+    Raised when a chunk that was expected to exist in storage was not retrieved successfully.
+    """

tests/test_config.py

@@ -23,7 +23,7 @@
 from zarr.core.codec_pipeline import BatchedCodecPipeline
 from zarr.core.config import BadConfigError, config
 from zarr.core.indexing import SelectorTuple
-from zarr.errors import ZarrUserWarning
+from zarr.errors import ChunkNotFoundError, ZarrUserWarning
 from zarr.registry import (
     fully_qualified_name,
     get_buffer_class,
@@ -53,6 +53,7 @@ def test_config_defaults_set() -> None:
                 "array": {
                     "order": "C",
                     "write_empty_chunks": False,
+                    "read_missing_chunks": True,
                     "target_shard_size_bytes": None,
                 },
                 "async": {"concurrency": 10, "timeout": None},
@@ -319,6 +320,108 @@ class NewCodec2(BytesCodec):
         get_codec_class("new_codec")
 
 
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"shards": (4, 4)},
+        {"compressors": None},
+    ],
+    ids=["partial_decode", "full_decode"],
+)
+def test_config_read_missing_chunks(store: Store, kwargs: dict[str, Any]) -> None:
+    arr = zarr.create_array(
+        store=store,
+        shape=(4, 4),
+        chunks=(2, 2),
+        dtype="int32",
+        fill_value=42,
+        **kwargs,
+    )
+
+    # default behavior: missing chunks are filled with the fill value
+    result = zarr.open_array(store)[:]
+    assert np.array_equal(result, np.full((4, 4), 42, dtype="int32"))
+
+    # with read_missing_chunks=False, reading missing chunks raises an error
+    with config.set({"array.read_missing_chunks": False}):
+        with pytest.raises(ChunkNotFoundError):
+            zarr.open_array(store)[:]
+
+    # after writing data, all chunks exist and no error is raised
+    arr[:] = np.arange(16, dtype="int32").reshape(4, 4)
+    with config.set({"array.read_missing_chunks": False}):
+        result = zarr.open_array(store)[:]
+        assert np.array_equal(result, np.arange(16, dtype="int32").reshape(4, 4))
+
+
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
+def test_config_read_missing_chunks_sharded_inner(store: Store) -> None:
+    """Because the shard index and inner chunks should be stored
+    together in a single storage object (read: a file or blob),
+    we delegate to the shard index the responsibility of determining
+    what chunks should be present.
+
+    Thus, `read_missing_chunks` raises an error only if the entire *shard*
+    is missing. Missing inner chunks are filled with the array's fill value
+    and do not raise an error, even if `read_missing_chunks=False` at the
+    array level.
+    """
+    arr = zarr.create_array(
+        store=store,
+        shape=(8, 4),
+        chunks=(2, 2),
+        shards=(4, 4),
+        dtype="int32",
+        fill_value=42,
+    )
+
+    # write only one inner chunk in the first shard, leaving the second shard empty
+    arr[0:2, 0:2] = np.ones((2, 2), dtype="int32")
+
+    with config.set({"array.read_missing_chunks": False}):
+        a = zarr.open_array(store)
+
+        # first shard exists: missing inner chunks are filled, no error
+        result = a[:4]
+        expected = np.full((4, 4), 42, dtype="int32")
+        expected[0:2, 0:2] = 1
+        assert np.array_equal(result, expected)
+
+        # second shard is entirely missing: raises an error
+        with pytest.raises(ChunkNotFoundError):
+            a[4:]
+
+
+@pytest.mark.parametrize("store", ["local", "memory"], indirect=["store"])
+def test_config_read_missing_chunks_write_empty_chunks(store: Store) -> None:
+    """write_empty_chunks=False drops chunks equal to fill_value, which then
+    appear missing to read_missing_chunks=False."""
+    arr = zarr.create_array(
+        store=store,
+        shape=(4,),
+        chunks=(2,),
+        dtype="int32",
+        fill_value=0,
+        config={"write_empty_chunks": False, "read_missing_chunks": False},
+    )
+
+    # write non-fill-value data: chunks are stored
+    arr[:] = [1, 2, 3, 4]
+    assert np.array_equal(arr[:], [1, 2, 3, 4])
+
+    # overwrite with fill_value: chunks are dropped by write_empty_chunks=False
+    arr[:] = 0
+    with pytest.raises(ChunkNotFoundError):
+        arr[:]
+
+    # with write_empty_chunks=True, chunks are kept and no error is raised
+    with config.set({"array.write_empty_chunks": True}):
+        arr = zarr.open_array(store)
+        arr[:] = 0
+        assert np.array_equal(arr[:], [0, 0, 0, 0])
+
+
 @pytest.mark.parametrize(
     "key",
     [