feat: zarr3

Author: floriankrb

src/anemoi/datasets/__init__.py

@@ -9,6 +9,7 @@
 
 from typing import List
 
+from .add_zarr_support import ZarrSupport
 from .data import MissingDateError
 from .data import add_dataset_path
 from .data import add_named_dataset
@@ -30,4 +31,5 @@
     "MissingDateError",
     "open_dataset",
     "__version__",
+    "ZarrSupport",
 ]

src/anemoi/datasets/add_zarr_support.py

@@ -0,0 +1,129 @@
+# (C) Copyright 2025 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+import logging
+
+import zarr
+
+LOG = logging.getLogger(__name__)
+
+
+class ZarrSupportV2:
+    @classmethod
+    def base_store(cls):
+        return zarr.storage.BaseStore
+
+    @classmethod
+    def is_zarr_group(cls, obj):
+        return isinstance(obj, zarr.hierarchy.Group)
+
+    @classmethod
+    def create_array(cls, zarr_root, *args, **kwargs):
+        return zarr_root.create_dataset(*args, **kwargs)
+
+    @classmethod
+    def change_dtype_datetime64(cls, dtype):
+        return dtype
+
+    @classmethod
+    def cast_dtype_datetime64(cls, array, dtype):
+        return array, dtype
+
+    @classmethod
+    def get_not_found_exception(cls):
+        return zarr.errors.PathNotFoundError
+
+    @classmethod
+    def zarr_open_mode_append(cls):
+        return "w+"
+
+    @classmethod
+    def zarr_open_to_patch_in_tests(cls):
+        return "zarr.convenience.open"
+
+    @classmethod
+    def get_read_only_store_class(cls):
+        class ReadOnlyStore(zarr.storage.BaseStore):
+            """A base class for read-only stores."""
+
+            def __delitem__(self, key: str) -> None:
+                """Prevent deletion of items."""
+                raise NotImplementedError()
+
+            def __setitem__(self, key: str, value: bytes) -> None:
+                """Prevent setting of items."""
+                raise NotImplementedError()
+
+            def __len__(self) -> int:
+                """Return the number of items in the store."""
+                raise NotImplementedError()
+
+            def __iter__(self) -> iter:
+                """Return an iterator over the store."""
+                raise NotImplementedError()
+
+        return ReadOnlyStore
+
+
+class ZarrSupportV3:
+    @classmethod
+    def base_store(cls):
+        return zarr.abc.store.Store
+
+    @classmethod
+    def is_zarr_group(cls, obj):
+        return isinstance(obj, zarr.Group)
+
+    @classmethod
+    def create_array(cls, zarr_root, *args, **kwargs):
+        return zarr_root.create_array(*args, **kwargs)
+
+    @classmethod
+    def get_not_found_exception(cls):
+        return FileNotFoundError
+
+    @classmethod
+    def zarr_open_mode_append(cls):
+        return "a"
+
+    @classmethod
+    def change_dtype_datetime64(cls, dtype):
+        # remove this flag (and the relevant code) when Zarr 3 supports datetime64
+        # https://github.com/zarr-developers/zarr-python/issues/2616
+        import numpy as np
+
+        if dtype == "datetime64[s]":
+            dtype = np.dtype("int64")
+        return dtype
+
+    @classmethod
+    def cast_dtype_datetime64(cls, array, dtype):
+        # remove this flag (and the relevant code) when Zarr 3 supports datetime64
+        # https://github.com/zarr-developers/zarr-python/issues/2616
+        import numpy as np
+
+        if dtype == np.dtype("datetime64[s]"):
+            dtype = "int64"
+            array = array.astype(dtype)
+
+        return array, dtype
+
+    @classmethod
+    def zarr_open_to_patch_in_tests(cls):
+        return "zarr.open"
+
+    @classmethod
+    def get_read_only_store_class(cls):
+        raise NotImplementedError("TODO")
+
+
+if zarr.__version__.startswith("3"):
+    ZarrSupport = ZarrSupportV3
+else:
+    LOG.warning("Using Zarr 2 : only zarr datasets build with zarr 2 are supported")
+    ZarrSupport = ZarrSupportV2

src/anemoi/datasets/commands/copy.py

@@ -20,6 +20,8 @@
 from anemoi.utils.remote import Transfer
 from anemoi.utils.remote import TransferMethodNotImplementedError
 
+from anemoi.datasets import ZarrSupport
+
 from . import Command
 
 LOG = logging.getLogger(__name__)
@@ -49,8 +51,6 @@ class ZarrCopier:
         Flag to resume copying an existing dataset.
     verbosity : int
         Verbosity level of logging.
-    nested : bool
-        Flag to use ZARR's nested directory backend.
     rechunk : str
         Rechunk size for the target data array.
     """
@@ -64,7 +64,6 @@ def __init__(
         overwrite: bool,
         resume: bool,
         verbosity: int,
-        nested: bool,
         rechunk: str,
         **kwargs: Any,
     ) -> None:
@@ -86,8 +85,6 @@ def __init__(
             Flag to resume copying an existing dataset.
         verbosity : int
             Verbosity level of logging.
-        nested : bool
-            Flag to use ZARR's nested directory backend.
         rechunk : str
             Rechunk size for the target data array.
         **kwargs : Any
@@ -100,7 +97,6 @@ def __init__(
         self.overwrite = overwrite
         self.resume = resume
         self.verbosity = verbosity
-        self.nested = nested
         self.rechunk = rechunk
 
         self.rechunking = rechunk.split(",") if rechunk else []
@@ -113,27 +109,6 @@ def __init__(
                 raise NotImplementedError("Rechunking with SSH not implemented.")
             assert NotImplementedError("SSH not implemented.")
 
-    def _store(self, path: str, nested: bool = False) -> Any:
-        """Get the storage path.
-
-        Parameters
-        ----------
-        path : str
-            Path to the storage.
-        nested : bool, optional
-            Flag to use nested directory storage.
-
-        Returns
-        -------
-        Any
-            Storage path.
-        """
-        if nested:
-            import zarr
-
-            return zarr.storage.NestedDirectoryStore(path)
-        return path
-
     def copy_chunk(self, n: int, m: int, source: Any, target: Any, _copy: Any, verbosity: int) -> Optional[slice]:
         """Copy a chunk of data from source to target.
 
@@ -237,7 +212,8 @@ def copy_data(self, source: Any, target: Any, _copy: Any, verbosity: int) -> Non
         target_data = (
             target["data"]
             if "data" in target
-            else target.create_dataset(
+            else ZarrSupport.create_array(
+                target,
                 "data",
                 shape=source_data.shape,
                 chunks=self.data_chunks,
@@ -317,13 +293,12 @@ def copy_group(self, source: Any, target: Any, _copy: Any, verbosity: int) -> No
         verbosity : int
             Verbosity level of logging.
         """
-        import zarr
 
         for k, v in source.attrs.items():
             target.attrs[k] = v
 
         for name in sorted(source.keys()):
-            if isinstance(source[name], zarr.hierarchy.Group):
+            if ZarrSupport.is_zarr_group(source[name]):
                 group = target[name] if name in target else target.create_group(name)
                 self.copy_group(
                     source[name],
@@ -376,13 +351,13 @@ def run(self) -> None:
 
         def target_exists() -> bool:
             try:
-                zarr.open(self._store(self.target), mode="r")
+                zarr.open(self.target, mode="r")
                 return True
             except ValueError:
                 return False
 
         def target_finished() -> bool:
-            target = zarr.open(self._store(self.target), mode="r")
+            target = zarr.open(self.target, mode="r")
             if "_copy" in target:
                 done = sum(1 if x else 0 for x in target["_copy"])
                 todo = len(target["_copy"])
@@ -400,19 +375,19 @@ def target_finished() -> bool:
         def open_target() -> Any:
 
             if not target_exists():
-                return zarr.open(self._store(self.target, self.nested), mode="w")
+                return zarr.open(self.target, mode="w")
 
             if self.overwrite:
                 LOG.error("Target already exists, overwriting.")
-                return zarr.open(self._store(self.target, self.nested), mode="w")
+                return zarr.open(self.target, mode="w")
 
             if self.resume:
                 if target_finished():
                     LOG.error("Target already exists and is finished.")
                     sys.exit(0)
 
                 LOG.error("Target already exists, resuming copy.")
-                return zarr.open(self._store(self.target, self.nested), mode="w+")
+                return zarr.open(self.target, mode=ZarrSupport.zarr_open_mode_append())
 
             LOG.error("Target already exists, use either --overwrite or --resume.")
             sys.exit(1)
@@ -421,7 +396,7 @@ def open_target() -> Any:
 
         assert target is not None, target
 
-        source = zarr.open(self._store(self.source), mode="r")
+        source = zarr.open(self.source, mode="r")
         self.copy(source, target, self.verbosity)
 
 
@@ -455,7 +430,6 @@ def add_arguments(self, command_parser: Any) -> None:
             help="Verbosity level. 0 is silent, 1 is normal, 2 is verbose.",
             default=1,
         )
-        command_parser.add_argument("--nested", action="store_true", help="Use ZARR's nested directpry backend.")
         command_parser.add_argument(
             "--rechunk", help="Rechunk the target data array. Rechunk size should be a diviser of the block size."
         )

src/anemoi/datasets/create/__init__.py

@@ -32,6 +32,7 @@
 from earthkit.data.core.order import build_remapping
 
 from anemoi.datasets import MissingDateError
+from anemoi.datasets import ZarrSupport
 from anemoi.datasets import open_dataset
 from anemoi.datasets.create.input.trace import enable_trace
 from anemoi.datasets.create.persistent import build_storage
@@ -154,7 +155,7 @@ def _path_readable(path: str) -> bool:
     try:
         zarr.open(path, "r")
         return True
-    except zarr.errors.PathNotFoundError:
+    except ZarrSupport.get_not_found_exception():
         return False
 
 
@@ -208,7 +209,7 @@ def update_metadata(self, **kwargs: Any) -> None:
         import zarr
 
         LOG.debug(f"Updating metadata {kwargs}")
-        z = zarr.open(self.path, mode="w+")
+        z = zarr.open(self.path, mode=ZarrSupport.zarr_open_mode_append())
         for k, v in kwargs.items():
             if isinstance(v, np.datetime64):
                 v = v.astype(datetime.datetime)
@@ -1520,7 +1521,7 @@ def run(self) -> None:
 
         LOG.info(stats)
 
-        if not all(self.registry.get_flags(sync=False)):
+        if not all(self.registry.get_flags()):
             raise Exception(f"❗Zarr {self.path} is not fully built, not writing statistics into dataset.")
 
         for k in ["mean", "stdev", "minimum", "maximum", "sums", "squares", "count", "has_nans"]:

src/anemoi/datasets/create/patch.py

@@ -14,6 +14,8 @@
 
 import zarr
 
+from anemoi.datasets import ZarrSupport
+
 LOG = logging.getLogger(__name__)
 
 
@@ -134,7 +136,7 @@ def apply_patch(path: str, verbose: bool = True, dry_run: bool = False) -> None:
 
     try:
         attrs = zarr.open(path, mode="r").attrs.asdict()
-    except zarr.errors.PathNotFoundError as e:
+    except ZarrSupport.get_not_found_exception() as e:
         LOG.error(f"Failed to open {path}")
         LOG.error(e)
         exit(0)