Add complex dtype support to FillValueCoder for Zarr backend

xarray/backends/zarr.py

@@ -122,7 +122,9 @@ class FillValueCoder:
     """
 
     @classmethod
-    def encode(cls, value: int | float | str | bytes, dtype: np.dtype[Any]) -> Any:
+    def encode(
+        cls, value: int | float | complex | str | bytes, dtype: np.dtype[Any]
+    ) -> Any:
         if dtype.kind in "S":
             # byte string, this implies that 'value' must also be `bytes` dtype.
             assert isinstance(value, bytes)
@@ -132,16 +134,33 @@ def encode(cls, value: int | float | str | bytes, dtype: np.dtype[Any]) -> Any:
             return bool(value)
         elif dtype.kind in "iu":
             # todo: do we want to check for decimals?
+            assert isinstance(value, int | float)
             return int(value)
         elif dtype.kind in "f":
+            assert isinstance(value, int | float)
             return base64.standard_b64encode(struct.pack("<d", float(value))).decode()
+        elif dtype.kind in "c":
+            # complex - encode each component as base64, matching float encoding
+            assert isinstance(value, complex) or np.issubdtype(
+                type(value), np.complexfloating
+            )
+            return [
+                base64.standard_b64encode(
+                    struct.pack("<d", float(value.real))  # type: ignore[union-attr]
+                ).decode(),
+                base64.standard_b64encode(
+                    struct.pack("<d", float(value.imag))  # type: ignore[union-attr]
+                ).decode(),
+            ]
         elif dtype.kind in "U":
             return str(value)
         else:
             raise ValueError(f"Failed to encode fill_value. Unsupported dtype {dtype}")
 
     @classmethod
-    def decode(cls, value: int | float | str | bytes, dtype: str | np.dtype[Any]):
+    def decode(
+        cls, value: int | float | str | bytes | list, dtype: str | np.dtype[Any]
+    ):
         if dtype == "string":
             # zarr V3 string type
             return str(value)
@@ -153,9 +172,16 @@ def decode(cls, value: int | float | str | bytes, dtype: str | np.dtype[Any]):
         if np_dtype.kind in "f":
             assert isinstance(value, str | bytes)
             return struct.unpack("<d", base64.standard_b64decode(value))[0]
+        elif np_dtype.kind in "c":
+            # complex - decode each component from base64, matching float decoding
+            assert isinstance(value, list | tuple) and len(value) == 2
+            real = struct.unpack("<d", base64.standard_b64decode(value[0]))[0]
+            imag = struct.unpack("<d", base64.standard_b64decode(value[1]))[0]
+            return complex(real, imag)
         elif np_dtype.kind in "b":
             return bool(value)
         elif np_dtype.kind in "iu":
+            assert isinstance(value, int | float)
             return int(value)
         else:
             raise ValueError(f"Failed to decode fill_value. Unsupported dtype {dtype}")

xarray/tests/test_backends.py

@@ -1129,7 +1129,7 @@
         ):
             pytest.skip("uint8 data can't be written to non-NetCDF4 data")
 
         with self.roundtrip(decoded) as actual:
             for k in decoded.variables:
                 assert decoded.variables[k].dtype == actual.variables[k].dtype
                 # CF _FillValue is always on-disk type
@@ -1212,7 +1212,7 @@
         encoded = Dataset({"x": ("t", sb, attributes)})
         unsigned_dtype = np.dtype(f"u{sb.dtype.itemsize}")
 
         with _roundtrip_with_warnings(decoded) as actual:
             for k in decoded.variables:
                 assert decoded.variables[k].dtype == actual.variables[k].dtype
                 exp_fv = decoded.variables[k].encoding["_FillValue"]
@@ -1288,7 +1288,7 @@
 
     def test_grid_mapping_and_bounds_are_not_coordinates_in_file(self) -> None:
         original = self._create_cf_dataset()
         with self.roundtrip(original, open_kwargs={"decode_coords": False}) as ds:
             assert ds.coords["latitude"].attrs["bounds"] == "latitude_bnds"
             assert ds.coords["longitude"].attrs["bounds"] == "longitude_bnds"
             assert "coordinates" not in ds["variable"].attrs
@@ -1296,7 +1296,7 @@
 
     def test_coordinate_variables_after_dataset_roundtrip(self) -> None:
         original = self._create_cf_dataset()
         with self.roundtrip(original, open_kwargs={"decode_coords": "all"}) as actual:
             assert_identical(actual, original)
 
         with self.roundtrip(original) as actual:
@@ -1321,7 +1321,7 @@
         # xarray/tests/test_conventions.py::TestCFEncodedDataStore
         # needs the to_dataset. The other backends should be fine
         # without it.
         with pytest.warns(
             UserWarning,
             match=(
                 r"Variable\(s\) referenced in bounds not in variables: "
@@ -1423,7 +1423,7 @@
         ds = Dataset({"x": ("y", np.arange(10.0))})
 
         kwargs: dict[str, Any] = dict(encoding={"x": {"dtype": "f4"}})
         with self.roundtrip(ds, save_kwargs=kwargs) as actual:
             encoded_dtype = actual.x.encoding["dtype"]
             # On OS X, dtype sometimes switches endianness for unclear reasons
             assert encoded_dtype.kind == "f" and encoded_dtype.itemsize == 4
@@ -1506,7 +1506,7 @@
         # Test default encoding for float:
         ds = Dataset({"x": ("y", np.arange(10.0))})
         kwargs = dict(encoding={"x": {"dtype": "f4"}})
         with self.roundtrip(ds, save_kwargs=kwargs) as actual:
             assert math.isnan(actual.x.encoding["_FillValue"])
         assert ds.x.encoding == {}
 
@@ -7099,6 +7099,18 @@
     assert actual3 == expected3
 
 
+@requires_zarr
+@pytest.mark.parametrize("dtype", [complex, np.complex64, np.complex128])
+def test_fill_value_coder_complex(dtype) -> None:
+    """Test that FillValueCoder round-trips complex fill values."""
+    from xarray.backends.zarr import FillValueCoder
+
+    for value in [dtype(1 + 2j), dtype(-3.5 + 4.5j), dtype(complex("nan+nanj"))]:
+        encoded = FillValueCoder.encode(value, np.dtype(dtype))
+        decoded = FillValueCoder.decode(encoded, np.dtype(dtype))
+        np.testing.assert_equal(np.array(decoded, dtype=dtype), np.array(value))
+
+
 @requires_zarr
 def test_extract_zarr_variable_encoding() -> None:
     var = xr.Variable("x", [1, 2])
@@ -7141,7 +7153,7 @@
 
     m = fsspec.filesystem("memory")
     mm = m.get_mapper("out1.zarr")
     ds.to_zarr(mm)  # old interface
     ds0 = ds.copy()
     # pd.to_timedelta returns ns-precision, but the example data is in second precision
     # so we need to fix this