From eab28a3f9dab50cee62a3156da9644e38d9e5612 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 11:01:35 +0200 Subject: [PATCH 1/2] 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") From 943f9b47fcc691e58493a6db74435605671069ca Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 11:07:23 +0200 Subject: [PATCH 2/2] feat(zarr): store the source array metadata document on read When a variable is read from a Zarr store, keep the array's complete metadata document on encoding["zarr_array_metadata"] for provenance and introspection. The key is read-only: it is listed in ZARR_READ_ONLY_ENCODING_KEYS, dropped on write, and never forwarded to zarr's array-creation call. The document is emitted by private per-format helpers (_v2_array_metadata, _v3_array_metadata) that read the metadata keys individually and return locally-defined TypedDicts mirroring the key names of the zarr-metadata package's ZarrV2ArrayMetadataJSON / ZarrV3ArrayMetadataJSON, with deliberately wide value types. These helpers are the only places xarray touches zarr-python's metadata API, so future changes to zarr-python's array class shape are absorbed there and surface loudly. Assisted-by: ClaudeCode:claude-fable-5 Co-authored-by: Claude --- doc/user-guide/io.rst | 9 +++ xarray/backends/zarr.py | 113 +++++++++++++++++++++++++++++++++- xarray/tests/test_backends.py | 23 +++++++ 3 files changed, 142 insertions(+), 3 deletions(-) diff --git a/doc/user-guide/io.rst b/doc/user-guide/io.rst index 84cd1d97323..2f6b433c665 100644 --- a/doc/user-guide/io.rst +++ b/doc/user-guide/io.rst @@ -1110,6 +1110,15 @@ For example: Not all native zarr compression and filtering options have been tested with xarray. +.. note:: + + When a variable is read from a Zarr store, xarray also stores the array's + complete Zarr metadata document (as returned by + ``zarr.Array.metadata.to_dict()``) on ``encoding["zarr_array_metadata"]``. + This is provided for provenance and introspection only: it is read-only, so + it is dropped when writing and is not used to control how the array is + written. + .. _io.zarr.appending: Modifying existing Zarr stores diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 0401f9e7145..74bda41b8c8 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -4,8 +4,8 @@ import json import os import struct -from collections.abc import Hashable, Iterable, Mapping -from typing import TYPE_CHECKING, Any, Literal, Self, cast +from collections.abc import Hashable, Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Literal, NotRequired, Self, TypedDict, cast import numpy as np import pandas as pd @@ -122,10 +122,110 @@ def _choose_default_mode( # 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"} + {"source", "original_shape", "preferred_chunks", "zarr_array_metadata"} ) +class _ZarrV2ArrayMetadata(TypedDict): + """A Zarr format 2 array metadata document. + + The keys mirror ``zarr_metadata.ZarrV2ArrayMetadataJSON``; the value types + are deliberately wide, since xarray only relies on the key names. + """ + + zarr_format: Literal[2] + shape: Sequence[int] + chunks: Sequence[int] + dtype: object + compressor: object + fill_value: object + order: object + filters: object + dimension_separator: NotRequired[object] + attributes: NotRequired[Mapping[str, object]] + + +class _ZarrV3ArrayMetadata(TypedDict): + """A Zarr format 3 array metadata document. + + The keys mirror ``zarr_metadata.ZarrV3ArrayMetadataJSON``; the value types + are deliberately wide, since xarray only relies on the key names. + """ + + zarr_format: Literal[3] + node_type: Literal["array"] + data_type: object + shape: Sequence[int] + chunk_grid: object + chunk_key_encoding: object + fill_value: object + codecs: Sequence[object] + attributes: NotRequired[Mapping[str, object]] + storage_transformers: NotRequired[Sequence[object]] + dimension_names: NotRequired[Sequence[str | None]] + + +def _array_zarr_format(zarr_array: ZarrArray) -> ZarrFormat: + """Return the Zarr format of an array. + + This is the single place where xarray reads the format from a zarr-python + array object, so a change to zarr-python's array or metadata API only + needs to be accommodated here. + """ + zarr_format = zarr_array.metadata.zarr_format + if zarr_format not in (2, 3): + raise ValueError(f"unsupported Zarr format: {zarr_format!r}") + return zarr_format + + +def _v2_array_metadata(zarr_array: ZarrArray) -> _ZarrV2ArrayMetadata: + """Emit the format 2 metadata document for a zarr array. + + This function and ``_v3_array_metadata`` are the only places where xarray + reads a zarr-python array's metadata document, so a change to zarr-python's + array or metadata API only needs to be accommodated here. The keys are read + individually so that the result always has the declared shape and a missing + key fails loudly. + """ + raw = zarr_array.metadata.to_dict() + document = { + "zarr_format": raw["zarr_format"], + "shape": raw["shape"], + "chunks": raw["chunks"], + "dtype": raw["dtype"], + "compressor": raw["compressor"], + "fill_value": raw["fill_value"], + "order": raw["order"], + "filters": raw["filters"], + } + for key in ("dimension_separator", "attributes"): + if key in raw: + document[key] = raw[key] + return cast("_ZarrV2ArrayMetadata", document) + + +def _v3_array_metadata(zarr_array: ZarrArray) -> _ZarrV3ArrayMetadata: + """Emit the format 3 metadata document for a zarr array. + + See ``_v2_array_metadata`` for the role these two functions play. + """ + raw = zarr_array.metadata.to_dict() + document = { + "zarr_format": raw["zarr_format"], + "node_type": raw["node_type"], + "data_type": raw["data_type"], + "shape": raw["shape"], + "chunk_grid": raw["chunk_grid"], + "chunk_key_encoding": raw["chunk_key_encoding"], + "fill_value": raw["fill_value"], + "codecs": raw["codecs"], + } + for key in ("attributes", "storage_transformers", "dimension_names"): + if key in raw: + document[key] = raw[key] + return cast("_ZarrV3ArrayMetadata", document) + + class FillValueCoder: """Handle custom logic to safely encode and decode fill values in Zarr. Possibly redundant with logic in xarray/coding/variables.py but needs to be @@ -903,6 +1003,13 @@ def open_store_variable(self, name): ) if self.zarr_group.metadata.zarr_format == 3: encoding.update({"serializer": zarr_array.serializer}) + # Store the source array's full metadata document for provenance and + # introspection. This is read-only: it is dropped on write (see + # ``ZARR_READ_ONLY_ENCODING_KEYS``) and is not consumed by the writer. + if _array_zarr_format(zarr_array) == 3: + encoding["zarr_array_metadata"] = _v3_array_metadata(zarr_array) + else: + encoding["zarr_array_metadata"] = _v2_array_metadata(zarr_array) if self._use_zarr_fill_value_as_mask: # Setting this attribute triggers CF decoding for missing values diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 5240804ed25..7db233d7975 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -4503,6 +4503,29 @@ def test_avoid_excess_metadata_calls(self) -> None: assert mock.call_count == call_count +@requires_zarr +@pytest.mark.parametrize( + "zarr_format, layout_key", + [(2, "chunks"), (3, "codecs")], +) +def test_zarr_array_metadata_stored_on_read(tmp_path, zarr_format, layout_key) -> None: + # The source array's full metadata document is stored on read for + # provenance / introspection. It is a read-only encoding key: dropped on + # write and never forwarded to the zarr array-creation call. + ds = xr.Dataset({"a": ("x", [1.0, 2.0, 3.0, 4.0])}) + ds.to_zarr(tmp_path / "s.zarr", zarr_format=zarr_format, mode="w") + + opened = xr.open_zarr(tmp_path / "s.zarr") + frag = opened["a"].encoding["zarr_array_metadata"] + assert frag["zarr_format"] == zarr_format + assert layout_key in frag + + # rewriting the opened dataset must succeed (the read-only key is dropped, + # not passed to zarr's create) and reproduce the original data + opened.to_zarr(tmp_path / "s2.zarr", zarr_format=zarr_format, mode="w") + assert_identical(xr.open_zarr(tmp_path / "s2.zarr"), ds) + + @requires_scipy class TestScipyInMemoryData(CFEncodedBase, NetCDF3Only, InMemoryNetCDF): engine: T_NetcdfEngine = "scipy"