Allow setting fill_value on Zarr format 3 arrays (#10161)

* Allow setting `fill_value` on Zarr format 3 arrays

Closes #10064

* fix

* fix format detection

* fix

* Set use_zarr_fill_value_as_mask=False

Author: dcherian

doc/user-guide/io.rst

@@ -1021,6 +1021,22 @@ reads. Because this fall-back option is so much slower, xarray issues a
        instead of falling back to try reading non-consolidated metadata.
 
 
+Fill Values
+~~~~~~~~~~~
+
+Zarr arrays have a ``fill_value`` that is used for chunks that were never written to disk.
+For the Zarr version 2 format, Xarray will set ``fill_value`` to be equal to the CF/NetCDF ``"_FillValue"``.
+This is ``np.nan`` by default for floats, and unset otherwise. Note that the Zarr library will set a
+default ``fill_value`` if not specified (usually ``0``).
+
+For the Zarr version 3 format, ``_FillValue`` and ```fill_value`` are decoupled.
+So you can set ``fill_value`` in ``encoding`` as usual.
+
+Note that at read-time, you can control whether ``_FillValue`` is masked using the
+``mask_and_scale`` kwarg; and whether Zarr's ``fill_value`` is treated as synonymous
+with ``_FillValue`` using the ``use_zarr_fill_value_as_mask`` kwarg to :py:func:`xarray.open_zarr`.
+
+
 .. _io.kerchunk:
 
 Kerchunk

doc/whats-new.rst

@@ -22,6 +22,8 @@ v2025.03.1 (unreleased)
 New Features
 ~~~~~~~~~~~~
 
+- Allow setting a ``fill_value`` for Zarr format 3 arrays. Specify ``fill_value`` in ``encoding`` as usual.
+  (:issue:`10064`). By `Deepak Cherian <https://github.com/dcherian>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/backends/zarr.py

@@ -426,6 +426,7 @@ def extract_zarr_variable_encoding(
     raise_on_invalid=False,
     name=None,
     *,
+    zarr_format: ZarrFormat,
     safe_chunks=True,
     region=None,
     mode=None,
@@ -443,6 +444,7 @@ def extract_zarr_variable_encoding(
     region: tuple[slice, ...], optional
     mode: str, optional
     shape: tuple[int, ...], optional
+    zarr_format: Literal[2,3]
     Returns
     -------
     encoding : dict
@@ -463,16 +465,23 @@ def extract_zarr_variable_encoding(
         "cache_metadata",
         "write_empty_chunks",
     }
+    if zarr_format == 3:
+        valid_encodings.add("fill_value")
 
     for k in safe_to_drop:
         if k in encoding:
             del encoding[k]
 
     if raise_on_invalid:
         invalid = [k for k in encoding if k not in valid_encodings]
+        if "fill_value" in invalid and zarr_format == 2:
+            msg = " Use `_FillValue` to set the Zarr array `fill_value`"
+        else:
+            msg = ""
+
         if invalid:
             raise ValueError(
-                f"unexpected encoding parameters for zarr backend:  {invalid!r}"
+                f"unexpected encoding parameters for zarr backend:  {invalid!r}." + msg
             )
     else:
         for k in list(encoding):
@@ -1147,7 +1156,7 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
             if self._use_zarr_fill_value_as_mask:
                 fill_value = attrs.pop("_FillValue", None)
             else:
-                fill_value = None
+                fill_value = v.encoding.pop("fill_value", None)
                 if "_FillValue" in attrs:
                     # replace with encoded fill value
                     fv = attrs.pop("_FillValue")
@@ -1198,6 +1207,7 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
                 region=region,
                 mode=self._mode,
                 shape=zarr_shape,
+                zarr_format=3 if is_zarr_v3_format else 2,
             )
 
             if self._mode == "w" or name not in existing_keys:

xarray/tests/test_backends.py

@@ -13,6 +13,7 @@
 import tempfile
 import uuid
 import warnings
+from collections import ChainMap
 from collections.abc import Generator, Iterator, Mapping
 from contextlib import ExitStack
 from io import BytesIO
@@ -3343,6 +3344,67 @@ def test_cache_members(self) -> None:
             observed_keys_2 = sorted(zstore_mut.array_keys())
             assert observed_keys_2 == sorted(array_keys + [new_key])
 
+    @requires_dask
+    @pytest.mark.parametrize("dtype", [int, float])
+    def test_zarr_fill_value_setting(self, dtype):
+        # When zarr_format=2, _FillValue sets fill_value
+        # When zarr_format=3, fill_value is set independently
+        # We test this by writing a dask array with compute=False,
+        # on read we should receive chunks filled with `fill_value`
+        fv = -1
+        ds = xr.Dataset(
+            {"foo": ("x", dask.array.from_array(np.array([0, 0, 0], dtype=dtype)))}
+        )
+        expected = xr.Dataset({"foo": ("x", [fv] * 3)})
+
+        zarr_format_2 = (
+            has_zarr_v3 and zarr.config.get("default_zarr_format") == 2
+        ) or not has_zarr_v3
+        if zarr_format_2:
+            attr = "_FillValue"
+            expected.foo.attrs[attr] = fv
+        else:
+            attr = "fill_value"
+            if dtype is float:
+                # for floats, Xarray inserts a default `np.nan`
+                expected.foo.attrs["_FillValue"] = np.nan
+
+        # turn off all decoding so we see what Zarr returns to us.
+        # Since chunks, are not written, we should receive on `fill_value`
+        open_kwargs = {
+            "mask_and_scale": False,
+            "consolidated": False,
+            "use_zarr_fill_value_as_mask": False,
+        }
+        save_kwargs = dict(compute=False, consolidated=False)
+        with self.roundtrip(
+            ds,
+            save_kwargs=ChainMap(save_kwargs, dict(encoding={"foo": {attr: fv}})),
+            open_kwargs=open_kwargs,
+        ) as actual:
+            assert_identical(actual, expected)
+
+        ds.foo.encoding[attr] = fv
+        with self.roundtrip(
+            ds, save_kwargs=save_kwargs, open_kwargs=open_kwargs
+        ) as actual:
+            assert_identical(actual, expected)
+
+        if zarr_format_2:
+            ds = ds.drop_encoding()
+            with pytest.raises(ValueError, match="_FillValue"):
+                with self.roundtrip(
+                    ds,
+                    save_kwargs=ChainMap(
+                        save_kwargs, dict(encoding={"foo": {"fill_value": fv}})
+                    ),
+                    open_kwargs=open_kwargs,
+                ):
+                    pass
+            # TODO: this doesn't fail because of the
+            # ``raise_on_invalid=vn in check_encoding_set`` line in zarr.py
+            # ds.foo.encoding["fill_value"] = fv
+
 
 @requires_zarr
 @pytest.mark.skipif(
@@ -5827,23 +5889,23 @@ def test_encode_zarr_attr_value() -> None:
 @requires_zarr
 def test_extract_zarr_variable_encoding() -> None:
     var = xr.Variable("x", [1, 2])
-    actual = backends.zarr.extract_zarr_variable_encoding(var)
+    actual = backends.zarr.extract_zarr_variable_encoding(var, zarr_format=3)
     assert "chunks" in actual
     assert actual["chunks"] == ("auto" if has_zarr_v3 else None)
 
     var = xr.Variable("x", [1, 2], encoding={"chunks": (1,)})
-    actual = backends.zarr.extract_zarr_variable_encoding(var)
+    actual = backends.zarr.extract_zarr_variable_encoding(var, zarr_format=3)
     assert actual["chunks"] == (1,)
 
     # does not raise on invalid
     var = xr.Variable("x", [1, 2], encoding={"foo": (1,)})
-    actual = backends.zarr.extract_zarr_variable_encoding(var)
+    actual = backends.zarr.extract_zarr_variable_encoding(var, zarr_format=3)
 
     # raises on invalid
     var = xr.Variable("x", [1, 2], encoding={"foo": (1,)})
     with pytest.raises(ValueError, match=r"unexpected encoding parameters"):
         actual = backends.zarr.extract_zarr_variable_encoding(
-            var, raise_on_invalid=True
+            var, raise_on_invalid=True, zarr_format=3
         )