Feat/write empty chunks

src/zarr/codecs/pipeline.py

@@ -331,6 +331,7 @@ async def write_batch(
         value: NDBuffer,
         drop_axes: tuple[int, ...] = (),
     ) -> None:
+        write_empty_chunks = config.get("array.write_empty_chunks") == True  # noqa: E712
         if self.supports_partial_encode:
             await self.encode_partial_batch(
                 [
@@ -360,7 +361,7 @@ async def _read_key(
                 _read_key,
                 config.get("async.concurrency"),
             )
-            chunk_array_batch = await self.decode_batch(
+            chunk_array_decoded = await self.decode_batch(
                 [
                     (chunk_bytes, chunk_spec)
                     for chunk_bytes, (_, chunk_spec, _, _) in zip(
@@ -369,23 +370,25 @@ async def _read_key(
                 ],
             )
 
-            chunk_array_batch = [
+            chunk_array_merged = [
                 self._merge_chunk_array(
                     chunk_array, value, out_selection, chunk_spec, chunk_selection, drop_axes
                 )
                 for chunk_array, (_, chunk_spec, chunk_selection, out_selection) in zip(
-                    chunk_array_batch, batch_info, strict=False
-                )
-            ]
-
-            chunk_array_batch = [
-                None
-                if chunk_array is None or chunk_array.all_equal(chunk_spec.fill_value)
-                else chunk_array
-                for chunk_array, (_, chunk_spec, _, _) in zip(
-                    chunk_array_batch, batch_info, strict=False
+                    chunk_array_decoded, batch_info, strict=False
                 )
             ]
+            chunk_array_batch: list[NDBuffer | None] = []
+            for chunk_array, (_, chunk_spec, _, _) in zip(
+                chunk_array_merged, batch_info, strict=False
+            ):
+                if chunk_array is None:
+                    chunk_array_batch.append(None)  # type: ignore[unreachable]
+                else:
+                    if not write_empty_chunks and chunk_array.all_equal(chunk_spec.fill_value):
+                        chunk_array_batch.append(None)
+                    else:
+                        chunk_array_batch.append(chunk_array)
 
             chunk_bytes_batch = await self.encode_batch(
                 [

src/zarr/core/config.py

@@ -42,7 +42,7 @@ def reset(self) -> None:
     defaults=[
         {
             "default_zarr_version": 3,
-            "array": {"order": "C"},
+            "array": {"order": "C", "write_empty_chunks": False},
             "async": {"concurrency": 10, "timeout": None},
             "threading": {"max_workers": None},
             "json_indent": 2,

tests/test_array.py

@@ -11,6 +11,7 @@
 from zarr.core.array import chunks_initialized
 from zarr.core.buffer.cpu import NDBuffer
 from zarr.core.common import JSON, MemoryOrder, ZarrFormat
+from zarr.core.config import config
 from zarr.core.group import AsyncGroup
 from zarr.core.indexing import ceildiv
 from zarr.core.sync import sync
@@ -436,3 +437,39 @@ def test_array_create_order(
         assert vals.flags.f_contiguous
     else:
         raise AssertionError
+
+
+@pytest.mark.parametrize("store", ["memory"], indirect=True)
+@pytest.mark.parametrize("write_empty_chunks", [True, False])
+@pytest.mark.parametrize("fill_value", [0, 5])
+def test_write_empty_chunks(
+    zarr_format: ZarrFormat, store: MemoryStore, write_empty_chunks: bool, fill_value: int
+) -> None:
+    """
+    Check that the write_empty_chunks value of the config is applied correctly. We expect that
+    when write_empty_chunks is True, writing chunks equal to the fill value will result in
+    those chunks appearing in the store.
+
+    When write_empty_chunks is False, writing chunks that are equal to the fill value will result in
+    those chunks not being present in the store. In particular, they should be deleted if they were
+    already present.
+    """
+    arr = Array.create(
+        store=store,
+        shape=(2,),
+        zarr_format=zarr_format,
+        dtype="i4",
+        fill_value=fill_value,
+        chunk_shape=(1,),
+    )
+
+    # initialize the store with some non-fill value chunks
+    arr[:] = fill_value + 1
+    assert arr.nchunks_initialized == arr.nchunks
+
+    with config.set({"array.write_empty_chunks": write_empty_chunks}):
+        arr[:] = fill_value
+    if not write_empty_chunks:
+        assert arr.nchunks_initialized == 0
+    else:
+        assert arr.nchunks_initialized == arr.nchunks

tests/test_config.py

@@ -40,7 +40,7 @@ def test_config_defaults_set() -> None:
     assert config.defaults == [
         {
             "default_zarr_version": 3,
-            "array": {"order": "C"},
+            "array": {"order": "C", "write_empty_chunks": False},
             "async": {"concurrency": 10, "timeout": None},
             "threading": {"max_workers": None},
             "json_indent": 2,