From eab28a3f9dab50cee62a3156da9644e38d9e5612 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 11:01:35 +0200 Subject: [PATCH] refactor(zarr): centralize and split encoding validation on the write path Define the encoding keys the Zarr backend accepts as module-level constants (ZARR_V2_ENCODING_KEYS, ZARR_V3_ENCODING_KEYS, ZARR_READ_ONLY_ENCODING_KEYS) and extract a pure _validate_zarr_variable_encoding helper from extract_zarr_variable_encoding. _create_new_array now assembles zarr create() arguments in its own dict instead of mutating the variable's encoding, takes the dimension names explicitly, and records them in one place for both zarr formats (native dimension_names metadata for format 3, the hidden _ARRAY_DIMENSIONS attribute for format 2). Remove unreachable "order" handling from array creation, and correct the formatting of the conflicting-write_empty_chunks error message. Behavior is otherwise unchanged; tests are split one case per function and new tests pin the constants and the conflict error. Assisted-by: ClaudeCode:claude-fable-5 Co-authored-by: Claude --- doc/whats-new.rst | 7 +- xarray/backends/zarr.py | 156 +++++++++++++++++++++------------- xarray/tests/test_backends.py | 91 +++++++++++++++++--- 3 files changed, 186 insertions(+), 68 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 63177c82cca..0d77da9e5cc 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -61,6 +61,12 @@ Documentation Internal Changes ~~~~~~~~~~~~~~~~ +- Refactor how the Zarr backend validates the variable ``encoding`` dict on + write: the accepted keys are now defined by the module-level constants + ``ZARR_V2_ENCODING_KEYS`` and ``ZARR_V3_ENCODING_KEYS`` in + ``xarray.backends.zarr``. Behavior is unchanged, except that the error raised + for conflicting ``write_empty_chunks`` settings now has a correctly formatted + message. By `Davis Bennett `_. .. _whats-new.2026.07.0: @@ -163,7 +169,6 @@ Documentation Internal Changes ~~~~~~~~~~~~~~~~ - .. _whats-new.2026.04.0: v2026.04.0 (Apr 13, 2026) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 6abb7fef07e..0401f9e7145 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -97,6 +97,34 @@ def _choose_default_mode( DIMENSION_KEY = "_ARRAY_DIMENSIONS" ZarrFormat = Literal[2, 3] +# --- Zarr encoding keys ----------------------------------------------------- +# The variable ``.encoding`` dict is xarray's channel for storage-level +# metadata. The sets below are the single source of truth for the keys the +# Zarr backend accepts on write and forwards to zarr's array-creation routine. +ZARR_V2_ENCODING_KEYS: frozenset[str] = frozenset( + { + "chunks", + "shards", + "compressors", + "filters", + "serializer", + "cache_metadata", + "write_empty_chunks", + "chunk_key_encoding", + } +) + +# Format 3 additionally accepts ``fill_value``; in format 2 the array fill +# value is carried by the ``_FillValue`` attribute instead. +ZARR_V3_ENCODING_KEYS: frozenset[str] = ZARR_V2_ENCODING_KEYS | {"fill_value"} + +# Informational keys that xarray populates in ``.encoding`` on read (or that +# originate from other backends) but that must never be forwarded to zarr on +# write. These are dropped silently. +ZARR_READ_ONLY_ENCODING_KEYS: frozenset[str] = frozenset( + {"source", "original_shape", "preferred_chunks"} +) + class FillValueCoder: """Handle custom logic to safely encode and decode fill values in Zarr. @@ -438,6 +466,39 @@ def _get_zarr_dims_and_attrs(zarr_obj, dimension_key, try_nczarr): return dimensions, attributes +def _validate_zarr_variable_encoding( + encoding: Mapping[str, object], + *, + raise_on_invalid: bool, + zarr_format: ZarrFormat, +) -> dict[str, object]: + """Filter an encoding mapping down to the keys the Zarr backend accepts. + + Read-only/informational keys (``ZARR_READ_ONLY_ENCODING_KEYS``) are always + dropped. Remaining keys are checked against the writable key set for + ``zarr_format``: when ``raise_on_invalid`` is True any unrecognized key + raises ``ValueError``; otherwise unrecognized keys are dropped silently. + + Returns a new dict; the input mapping is never mutated. + """ + valid_keys = ZARR_V3_ENCODING_KEYS if zarr_format == 3 else ZARR_V2_ENCODING_KEYS + encoding = { + k: v for k, v in encoding.items() if k not in ZARR_READ_ONLY_ENCODING_KEYS + } + + invalid = [k for k in encoding if k not in valid_keys] + if len(invalid) > 0 and raise_on_invalid: + msg = ( + " Use `_FillValue` to set the Zarr array `fill_value`" + if "fill_value" in invalid and zarr_format == 2 + else "" + ) + raise ValueError( + f"unexpected encoding parameters for zarr backend: {invalid!r}." + msg + ) + return {k: v for k, v in encoding.items() if k not in invalid} + + def extract_zarr_variable_encoding( variable, raise_on_invalid=False, @@ -460,41 +521,9 @@ def extract_zarr_variable_encoding( Zarr encoding for `variable` """ - encoding = variable.encoding.copy() - - safe_to_drop = {"source", "original_shape", "preferred_chunks"} - valid_encodings = { - "chunks", - "shards", - "compressors", - "filters", - "serializer", - "cache_metadata", - "write_empty_chunks", - "chunk_key_encoding", - } - 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}." + msg - ) - else: - for k in list(encoding): - if k not in valid_encodings: - del encoding[k] + encoding = _validate_zarr_variable_encoding( + variable.encoding, raise_on_invalid=raise_on_invalid, zarr_format=zarr_format + ) chunks = _determine_zarr_chunks( enc_chunks=encoding.get("chunks"), @@ -1104,38 +1133,58 @@ def _open_existing_array(self, *, name) -> ZarrArray: return cast(ZarrArray, zarr_array) def _create_new_array( - self, *, name, shape, dtype, fill_value, encoding, attrs + self, *, name, dims, shape, dtype, fill_value, encoding, attrs ) -> ZarrArray: if coding.strings.check_vlen_dtype(dtype) is str: dtype = str + # Zarr array creation takes the variable's storage `encoding` together + # with store-level parameters (overwrite, dimension names, write-empty + # policy). Collect them in their own dict so that `encoding` remains a + # plain description of the variable's storage, separate from the + # arguments accepted by zarr's `create()`. + create_kwargs = dict(encoding) + create_kwargs["overwrite"] = self._mode == "w" + + # Zarr format 3 stores dimension names natively in the array metadata. + # Format 2 has no such field, so xarray records them in the hidden + # _ARRAY_DIMENSIONS attribute instead. + if self.zarr_group.metadata.zarr_format == 3: + create_kwargs["dimension_names"] = dims + else: + attrs = dict(attrs) + attrs[DIMENSION_KEY] = dims + if self._write_empty is not None: if ( - "write_empty_chunks" in encoding - and encoding["write_empty_chunks"] != self._write_empty + "write_empty_chunks" in create_kwargs + and create_kwargs["write_empty_chunks"] != self._write_empty ): raise ValueError( - 'Differing "write_empty_chunks" values in encoding and parameters' - f'Got {encoding["write_empty_chunks"] = } and {self._write_empty = }' + 'Differing "write_empty_chunks" values in encoding and parameters. ' + f"Got write_empty_chunks={create_kwargs['write_empty_chunks']!r} in " + f"encoding and write_empty_chunks={self._write_empty!r} as a parameter." ) else: - encoding["write_empty_chunks"] = self._write_empty - - # zarr v3 passes write_empty_chunks and order via the config argument - encoding["config"] = {} - for c in ("write_empty_chunks", "order"): - if c in encoding: - encoding["config"][c] = encoding.pop(c) + create_kwargs["write_empty_chunks"] = self._write_empty + + # zarr-python 3 expects write_empty_chunks in the config argument + # rather than as a top-level parameter + create_kwargs["config"] = {} + if "write_empty_chunks" in create_kwargs: + create_kwargs["config"]["write_empty_chunks"] = create_kwargs.pop( + "write_empty_chunks" + ) # fill_value is passed explicitly; remove from encoding to avoid duplicates - encoding.pop("fill_value", None) + create_kwargs.pop("fill_value", None) zarr_array = self.zarr_group.create( name, shape=shape, dtype=dtype, fill_value=fill_value, - **encoding, + **create_kwargs, ) zarr_array = _put_attrs(zarr_array, attrs) return zarr_array @@ -1271,16 +1320,9 @@ def set_variables( if self._mode == "w" or name not in existing_keys: # new variable encoded_attrs = {k: self.encode_attribute(v) for k, v in attrs.items()} - # the magic for storing the hidden dimension data - if is_zarr_v3_format: - encoding["dimension_names"] = dims - else: - encoded_attrs[DIMENSION_KEY] = dims - - encoding["overwrite"] = self._mode == "w" - zarr_array = self._create_new_array( name=name, + dims=dims, dtype=dtype, shape=shape, fill_value=fill_value, diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 94dab2a6a59..5240804ed25 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -7117,28 +7117,99 @@ def test_fill_value_coder_inf_nan(value, dtype) -> None: @requires_zarr -def test_extract_zarr_variable_encoding() -> None: - var = xr.Variable("x", [1, 2]) +@pytest.mark.parametrize( + "encoding, expected_chunks", + [({}, "auto"), ({"chunks": (1,)}, (1,))], +) +def test_extract_zarr_encoding_resolves_chunks(encoding, expected_chunks) -> None: + var = xr.Variable("x", [1, 2], encoding=encoding) actual = backends.zarr.extract_zarr_variable_encoding(var, zarr_format=3) - assert "chunks" in actual - assert actual["chunks"] == "auto" + assert actual["chunks"] == expected_chunks - var = xr.Variable("x", [1, 2], encoding={"chunks": (1,)}) - actual = backends.zarr.extract_zarr_variable_encoding(var, zarr_format=3) - assert actual["chunks"] == (1,) - # does not raise on invalid +@requires_zarr +def test_extract_zarr_encoding_drops_invalid_key() -> None: var = xr.Variable("x", [1, 2], encoding={"foo": (1,)}) actual = backends.zarr.extract_zarr_variable_encoding(var, zarr_format=3) + assert "foo" not in actual + - # raises on invalid +@requires_zarr +def test_extract_zarr_encoding_raises_on_invalid_key() -> None: 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( + backends.zarr.extract_zarr_variable_encoding( var, raise_on_invalid=True, zarr_format=3 ) +@requires_zarr +@pytest.mark.parametrize( + "read_only_key", ["preferred_chunks", "source", "original_shape"] +) +def test_validate_zarr_encoding_drops_read_only_key(read_only_key) -> None: + validate = backends.zarr._validate_zarr_variable_encoding + src = {"chunks": (1,), read_only_key: "value"} + out = validate(src, raise_on_invalid=False, zarr_format=3) + assert out == {"chunks": (1,)} + assert read_only_key in src # input is not mutated + + +@requires_zarr +def test_validate_zarr_encoding_drops_unknown_key_when_not_raising() -> None: + validate = backends.zarr._validate_zarr_variable_encoding + out = validate({"foo": 1, "chunks": (1,)}, raise_on_invalid=False, zarr_format=3) + assert out == {"chunks": (1,)} + + +@requires_zarr +def test_validate_zarr_encoding_raises_on_unknown_key() -> None: + validate = backends.zarr._validate_zarr_variable_encoding + with pytest.raises(ValueError, match=r"unexpected encoding parameters"): + validate({"foo": 1}, raise_on_invalid=True, zarr_format=3) + + +@requires_zarr +def test_validate_zarr_encoding_accepts_fill_value_for_v3() -> None: + validate = backends.zarr._validate_zarr_variable_encoding + out = validate({"fill_value": 0}, raise_on_invalid=True, zarr_format=3) + assert out == {"fill_value": 0} + + +@requires_zarr +def test_validate_zarr_encoding_rejects_fill_value_for_v2() -> None: + validate = backends.zarr._validate_zarr_variable_encoding + with pytest.raises(ValueError, match=r"Use `_FillValue`"): + validate({"fill_value": 0}, raise_on_invalid=True, zarr_format=2) + + +@requires_zarr +def test_valid_zarr_encoding_keys_fill_value_is_v3_only() -> None: + # fill_value is the only format-specific encoding key: valid for v3 only, + # since in format 2 the fill value is carried by the _FillValue attribute. + v3_only = backends.zarr.ZARR_V3_ENCODING_KEYS - backends.zarr.ZARR_V2_ENCODING_KEYS + assert v3_only == {"fill_value"} + + +@requires_zarr +def test_read_only_encoding_keys_are_not_writable() -> None: + read_only = backends.zarr.ZARR_READ_ONLY_ENCODING_KEYS + assert read_only.isdisjoint(backends.zarr.ZARR_V3_ENCODING_KEYS) + + +@requires_zarr +def test_write_empty_chunks_conflict_raises(tmp_path) -> None: + # conflicting write_empty_chunks settings via encoding and via parameter + ds = xr.Dataset({"a": ("x", [1.0, 2.0])}) + with pytest.raises(ValueError, match=r'Differing "write_empty_chunks"'): + ds.to_zarr( + tmp_path / "s.zarr", + mode="w", + write_empty_chunks=True, + encoding={"a": {"write_empty_chunks": False}}, + ) + + @requires_zarr @requires_fsspec @pytest.mark.filterwarnings("ignore:deallocating CachingFileManager")