Compatibility for zarr-python 3.x (#9552)

* Remove zarr pin

* Define zarr_v3 helper

* zarr-v3: filters / compressors -> codecs

* zarr-v3: update tests to avoid values equal to fillValue

* Various test fixes

* zarr_version fixes

* removed open_consolidated workarounds
* removed _store_version check
* pass through zarr_version

* fixup! zarr-v3: filters / compressors -> codecs

* fixup! fixup! zarr-v3: filters / compressors -> codecs

* fixup

* path / key normalization in set_variables

* fixes

* workaround nested consolidated metadata

* test: avoid fill_value

* test: Adjust call counts

* zarr-python 3.x Array.resize doesn't mutate

* test compatibility

- skip write_empty_chunks on 3.x
- update patch targets

* skip ZipStore with_mode

* test: more fill_value avoidance

* test: more fill_value avoidance

* v3 compat for instrumented test

* Handle zarr_version / zarr_format deprecation

* wip

* most Zarr tests passing

* unskip tests

* add custom Zarr _FillValue encoding / decoding

* relax dtype comparison in test_roundtrip_empty_vlen_string_array

* fix test_explicitly_omit_fill_value_via_encoding_kwarg

* fix test_append_string_length_mismatch_raises

* fix test_check_encoding_is_consistent_after_append for v3

* skip roundtrip_endian for zarr v3

* unskip datetimes and fix test_compressor_encoding

* unskip tests

* add back dtype skip

* point upstream to v3 branch

* Create temporary directory before using it

* Avoid zarr.storage.zip on v2

* fixed close_store_on_close bug

* Remove workaround, fixed upstream

* Restore original `w` mode.

* workaround for store closing with mode=w

* typing fixes

* compat

* Remove unnecessary pop

* fixed skip

* fixup types

* fixup types

* [test-upstream]

* Update install-upstream-wheels.sh

* set use_consolidated to false when user provides consolidated=False

* fix: import consolidated_metadata from package root

* fix: relax instrumented store checks for v3

* Adjust 2.18.3 thresholds

* skip datatree zarr tests w/ zarr 3 for now

* fixed kvstore usage

* typing fixes

* move zarr.codecs import

* fixup ignores

* storage options fix, skip

* fixed types

* Update ci/install-upstream-wheels.sh

* type fixes

* whats-new

* Update xarray/tests/test_backends_datatree.py

* fix type import

* set mapper, chunk_mapper

* Pass through zarr_format

* Fixup

* more cleanup

* revert test changes

* Update xarray/backends/zarr.py

* cleanup

* update docstring

* fix rtd

* tweak

---------

Co-authored-by: Ryan Abernathey <ryan.abernathey@gmail.com>
Co-authored-by: Joe Hamman <joe@earthmover.io>
Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>
Co-authored-by: Deepak Cherian <deepak@cherian.net>

Author: 4 people

ci/install-upstream-wheels.sh

@@ -45,15 +45,15 @@ python -m pip install \
     --pre \
     --upgrade \
     pyarrow
-# manually install `pint` to pull in new dependencies
-python -m pip install --upgrade pint
+# manually install `pint`, `donfig`, and `crc32c` to pull in new dependencies
+python -m pip install --upgrade pint donfig crc32c
 python -m pip install \
     --no-deps \
     --upgrade \
     git+https://github.com/dask/dask \
     git+https://github.com/dask/dask-expr \
     git+https://github.com/dask/distributed \
-    git+https://github.com/zarr-developers/zarr.git@main \
+    git+https://github.com/zarr-developers/zarr \
     git+https://github.com/Unidata/cftime \
     git+https://github.com/pypa/packaging \
     git+https://github.com/hgrecco/pint \

doc/user-guide/io.rst

@@ -823,8 +823,9 @@ For example:
 .. ipython:: python
 
     import zarr
+    from numcodecs.blosc import Blosc
 
-    compressor = zarr.Blosc(cname="zstd", clevel=3, shuffle=2)
+    compressor = Blosc(cname="zstd", clevel=3, shuffle=2)
     ds.to_zarr("foo.zarr", encoding={"foo": {"compressor": compressor}})
 
 .. note::

doc/whats-new.rst

@@ -39,6 +39,10 @@ New Features
   By `Holly Mandel <https://github.com/hollymandel>`_.
 - Implement handling of complex numbers (netcdf4/h5netcdf) and enums (h5netcdf) (:issue:`9246`, :issue:`3297`, :pull:`9509`).
   By `Kai Mühlbauer <https://github.com/kmuehlbauer>`_.
+- Support for Zarr-Python 3 (:issue:`95515`, :pull:`9552`).
+  By `Tom Augspurger <https://github.com/TomAugspurger>`_,
+  `Ryan Abernathey <https://github.com/rabernat>`_ and
+  `Joe Hamman <https://github.com/jhamman>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

pyproject.toml

@@ -49,7 +49,7 @@ dev = [
   "sphinx_autosummary_accessors",
   "xarray[complete]",
 ]
-io = ["netCDF4", "h5netcdf", "scipy", 'pydap; python_version<"3.10"', "zarr<3", "fsspec", "cftime", "pooch"]
+io = ["netCDF4", "h5netcdf", "scipy", 'pydap; python_version<"3.10"', "zarr", "fsspec", "cftime", "pooch"]
 etc = ["sparse"]
 parallel = ["dask[complete]"]
 viz = ["cartopy", "matplotlib", "nc-time-axis", "seaborn"]
@@ -124,6 +124,7 @@ module = [
   "nc_time_axis.*",
   "netCDF4.*",
   "netcdftime.*",
+  "numcodecs.*",
   "opt_einsum.*",
   "pint.*",
   "pooch.*",

xarray/backends/api.py

@@ -33,6 +33,7 @@
     _normalize_path,
 )
 from xarray.backends.locks import _get_scheduler
+from xarray.backends.zarr import _zarr_v3
 from xarray.core import indexing
 from xarray.core.combine import (
     _infer_concat_order_from_positions,
@@ -1685,6 +1686,7 @@ def to_zarr(
     safe_chunks: bool = True,
     storage_options: dict[str, str] | None = None,
     zarr_version: int | None = None,
+    zarr_format: int | None = None,
     write_empty_chunks: bool | None = None,
     chunkmanager_store_kwargs: dict[str, Any] | None = None,
 ) -> backends.ZarrStore | Delayed:
@@ -1703,21 +1705,28 @@ def to_zarr(
     store = _normalize_path(store)
     chunk_store = _normalize_path(chunk_store)
 
+    kwargs = {}
     if storage_options is None:
         mapper = store
         chunk_mapper = chunk_store
     else:
-        from fsspec import get_mapper
-
         if not isinstance(store, str):
             raise ValueError(
                 f"store must be a string to use storage_options. Got {type(store)}"
             )
-        mapper = get_mapper(store, **storage_options)
-        if chunk_store is not None:
-            chunk_mapper = get_mapper(chunk_store, **storage_options)
-        else:
+
+        if _zarr_v3():
+            kwargs["storage_options"] = storage_options
+            mapper = store
             chunk_mapper = chunk_store
+        else:
+            from fsspec import get_mapper
+
+            mapper = get_mapper(store, **storage_options)
+            if chunk_store is not None:
+                chunk_mapper = get_mapper(chunk_store, **storage_options)
+            else:
+                chunk_mapper = chunk_store
 
     if encoding is None:
         encoding = {}
@@ -1747,13 +1756,6 @@ def to_zarr(
     # validate Dataset keys, DataArray names
     _validate_dataset_names(dataset)
 
-    if zarr_version is None:
-        # default to 2 if store doesn't specify its version (e.g. a path)
-        zarr_version = int(getattr(store, "_store_version", 2))
-
-    if consolidated is None and zarr_version > 2:
-        consolidated = False
-
     if mode == "r+":
         already_consolidated = consolidated
         consolidate_on_close = False
@@ -1773,7 +1775,9 @@ def to_zarr(
         safe_chunks=safe_chunks,
         stacklevel=4,  # for Dataset.to_zarr()
         zarr_version=zarr_version,
+        zarr_format=zarr_format,
         write_empty=write_empty_chunks,
+        **kwargs,
     )
 
     if region is not None: