Compatibility for zarr-python 3.x

doc/user-guide/io.rst

@@ -845,8 +845,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::

pyproject.toml

@@ -42,7 +42,7 @@ dev = [
   "ruff",
   "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"]
 parallel = ["dask[complete]"]
 viz = ["matplotlib", "seaborn", "nc-time-axis"]
 
@@ -116,6 +116,7 @@ module = [
   "nc_time_axis.*",
   "netCDF4.*",
   "netcdftime.*",
+  "numcodecs.*",
   "opt_einsum.*",
   "pint.*",
   "pooch.*",

xarray/backends/api.py

@@ -1732,13 +1732,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

xarray/backends/zarr.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import functools
 import json
 import os
 import warnings
 from collections.abc import Callable, Iterable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
+import packaging.version
 import pandas as pd
 
 from xarray import coding, conventions
@@ -40,8 +42,20 @@
     from xarray.core.dataset import Dataset
     from xarray.core.datatree import DataTree
 
+
+@functools.lru_cache
+def _zarr_v3() -> bool:
+    try:
+        import zarr
+    except ImportError:
+        return False
+    else:
+        return packaging.version.parse(zarr.__version__).major >= 3
+
+
 # need some special secret attributes to tell us the dimensions
 DIMENSION_KEY = "_ARRAY_DIMENSIONS"
+ZarrFormat = Literal[2, 3]
 
 
 def encode_zarr_attr_value(value):
@@ -75,8 +89,10 @@ def __init__(self, zarr_array):
         self.shape = self._array.shape
 
         # preserve vlen string object dtype (GH 7328)
-        if self._array.filters is not None and any(
-            [filt.codec_id == "vlen-utf8" for filt in self._array.filters]
+        if (
+            not _zarr_v3()
+            and self._array.filters is not None
+            and any([filt.codec_id == "vlen-utf8" for filt in self._array.filters])
         ):
             dtype = coding.strings.create_vlen_dtype(str)
         else:
@@ -317,6 +333,7 @@ def extract_zarr_variable_encoding(
 
     safe_to_drop = {"source", "original_shape"}
     valid_encodings = {
+        "codecs",
         "chunks",
         "compressor",
         "filters",
@@ -614,9 +631,25 @@ def open_store_variable(self, name, zarr_array=None):
         encoding = {
             "chunks": zarr_array.chunks,
             "preferred_chunks": dict(zip(dimensions, zarr_array.chunks, strict=True)),
-            "compressor": zarr_array.compressor,
-            "filters": zarr_array.filters,
         }
+
+        if _zarr_v3() and zarr_array.metadata.zarr_format == 3:
+            encoding["codecs"] = [x.to_dict() for x in zarr_array.metadata.codecs]
+        elif _zarr_v3():
+            encoding.update(
+                {
+                    "compressor": zarr_array.metadata.compressor,
+                    "filters": zarr_array.metadata.filters,
+                }
+            )
+        else:
+            encoding.update(
+                {
+                    "compressor": zarr_array.compressor,
+                    "filters": zarr_array.filters,
+                }
+            )
+
         # _FillValue needs to be in attributes, not encoding, so it will get
         # picked up by decode_cf
         if zarr_array.fill_value is not None:
@@ -786,7 +819,11 @@ def store(
             variables_to_set, check_encoding_set, writer, unlimited_dims=unlimited_dims
         )
         if self._consolidate_on_close:
-            zarr.consolidate_metadata(self.zarr_group.store)
+            kwargs = {}
+            if _zarr_v3():
+                # https://github.com/zarr-developers/zarr-python/pull/2113#issuecomment-2386718323
+                kwargs["path"] = self.zarr_group.name.lstrip("/")
+            zarr.consolidate_metadata(self.zarr_group.store, **kwargs)
 
     def sync(self):
         pass
@@ -850,9 +887,19 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
                     #     - Existing variables already have their attrs included in the consolidated metadata file.
                     #     - The size of dimensions can not be expanded, that would require a call using `append_dim`
                     #        which is mutually exclusive with `region`
+                    kwargs = {}
+                    if _zarr_v3():
+                        kwargs["store"] = self.zarr_group.store
+                    else:
+                        kwargs["store"] = self.zarr_group.chunk_store
+
+                    # TODO: see if zarr should normalize these strings.
                     zarr_array = zarr.open(
-                        store=self.zarr_group.chunk_store,
-                        path=f"{self.zarr_group.name}/{name}",
+                        **kwargs,
+                        # path=f"{self.zarr_group.name}/{name}",
+                        path="/".join([self.zarr_group.name.rstrip("/"), name]).lstrip(
+                            "/"
+                        ),
                         write_empty_chunks=self._write_empty,
                     )
                 else:
@@ -868,7 +915,10 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
 
                     new_shape = list(zarr_array.shape)
                     new_shape[append_axis] += v.shape[append_axis]
-                    zarr_array.resize(new_shape)
+                    if _zarr_v3():
+                        zarr_array = zarr_array.resize(new_shape)
+                    else:
+                        zarr_array.resize(new_shape)
 
                 zarr_shape = zarr_array.shape
 
@@ -913,6 +963,10 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=No
                     else:
                         encoding["write_empty_chunks"] = self._write_empty
 
+                if "codecs" in encoding:
+                    pipeline = encoding.pop("codecs")
+                    encoding["codecs"] = pipeline
+
                 zarr_array = self.zarr_group.create(
                     name,
                     shape=shape,
@@ -1396,40 +1450,34 @@ def _get_open_params(
     if isinstance(store, os.PathLike):
         store = os.fspath(store)
 
-    if zarr_version is None:
-        # default to 2 if store doesn't specify it's version (e.g. a path)
-        zarr_version = getattr(store, "_store_version", 2)
-
     open_kwargs = dict(
         # mode='a-' is a handcrafted xarray specialty
         mode="a" if mode == "a-" else mode,
         synchronizer=synchronizer,
         path=group,
     )
     open_kwargs["storage_options"] = storage_options
-    if zarr_version > 2:
-        open_kwargs["zarr_version"] = zarr_version
-
-        if consolidated or consolidate_on_close:
-            raise ValueError(
-                "consolidated metadata has not been implemented for zarr "
-                f"version {zarr_version} yet. Set consolidated=False for "
-                f"zarr version {zarr_version}. See also "
-                "https://github.com/zarr-developers/zarr-specs/issues/136"
-            )
 
-        if consolidated is None:
-            consolidated = False
+    if _zarr_v3():
+        open_kwargs["zarr_format"] = zarr_version
+    else:
+        open_kwargs["zarr_version"] = zarr_version
 
     if chunk_store is not None:
         open_kwargs["chunk_store"] = chunk_store
         if consolidated is None:
             consolidated = False
 
+    if _zarr_v3():
+        missing_exc: type[Exception] = ValueError
+    else:
+        missing_exc = zarr.errors.GroupNotFoundError
+
     if consolidated is None:
         try:
             zarr_group = zarr.open_consolidated(store, **open_kwargs)
-        except KeyError:
+        except (ValueError, KeyError):
+            # ValueError in zarr-python 3.x, KeyError in 2.x.
             try:
                 zarr_group = zarr.open_group(store, **open_kwargs)
                 warnings.warn(
@@ -1447,7 +1495,7 @@ def _get_open_params(
                     RuntimeWarning,
                     stacklevel=stacklevel,
                 )
-            except zarr.errors.GroupNotFoundError as err:
+            except missing_exc as err:
                 raise FileNotFoundError(
                     f"No such file or directory: '{store}'"
                 ) from err