Zarr support (backend, in esmvalcore.preprocessor._io.py)

.github/workflows/run-tests.yml

@@ -21,6 +21,7 @@ on:
   push:
     branches:
       - main
+      - zarr_support
   # run the test only if the PR is to main
   # turn it on if required
   #pull_request:

environment.yml

@@ -5,6 +5,7 @@ channels:
   - nodefaults
 
 dependencies:
+  - aiohttp
   - cartopy
   - cf-units
   - cftime
@@ -18,6 +19,7 @@ dependencies:
   - fire
   - geopy
   - humanfriendly
+  - intake-esm
   - iris >=3.12.2  # https://github.com/SciTools/iris/issues/6417
   - iris-esmf-regrid >=0.11.0
   - iris-grib >=0.20.0  # github.com/ESMValGroup/ESMValCore/issues/2535
@@ -46,6 +48,7 @@ dependencies:
   - shapely >=2.0.0
   - xarray
   - yamale
+  - zarr >3
   # Python packages needed for building docs
   - autodocsumm >=0.2.2
   - ipython <9.0  # github.com/ESMValGroup/ESMValCore/issues/2680

esmvalcore/preprocessor/_io.py

@@ -9,7 +9,9 @@
 from itertools import groupby
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
 
+import fsspec
 import iris
 import ncdata
 import xarray as xr
@@ -75,6 +77,7 @@
 def load(
     file: str | Path | Cube | CubeList | xr.Dataset | ncdata.NcData,
     ignore_warnings: list[dict[str, Any]] | None = None,
+    backend_kwargs: dict[str, Any] | None = None,
 ) -> CubeList:
     """Load Iris cubes.
 
@@ -83,10 +86,18 @@
     file:
         File to be loaded. If ``file`` is already a loaded dataset, return it
         as a :class:`~iris.cube.CubeList`.
+        File as ``Path`` object could be a Zarr store.
     ignore_warnings:
         Keyword arguments passed to :func:`warnings.filterwarnings` used to
         ignore warnings issued by :func:`iris.load_raw`. Each list element
         corresponds to one call to :func:`warnings.filterwarnings`.
+    backend_kwargs:
+        Dict to hold info needed by storage backend e.g. to access
+        a PRIVATE S3 bucket containing object stores (e.g. netCDF4 files);
+        needed by ``fsspec`` and its extensions e.g. ``s3fs``, so
+        most of the times it is a ``storage_options`` dict. Note that Zarr
+        files are opened via ``http`` extension of ``fsspec``, so no need
+        for ``storage_options`` in that case (ie anon/anon).
 
     Returns
     -------
@@ -102,7 +113,19 @@
 
     """
     if isinstance(file, (str, Path)):
-        cubes = _load_from_file(file, ignore_warnings=ignore_warnings)
+        extension = (
+            file.suffix
+            if isinstance(file, Path)
+            else os.path.splitext(file)[1]
+        )
+        if "zarr" not in extension:
+            cubes = _load_from_file(file, ignore_warnings=ignore_warnings)
+        else:
+            cubes = _load_zarr(
+                file,
+                ignore_warnings=ignore_warnings,
+                backend_kwargs=backend_kwargs,
+            )
     elif isinstance(file, Cube):
         cubes = CubeList([file])
     elif isinstance(file, CubeList):
@@ -134,6 +157,50 @@
     return cubes
 
 
+def _load_zarr(
+    file: str | Path | Cube | CubeList | xr.Dataset | ncdata.NcData,
+    ignore_warnings: list[dict[str, Any]] | None = None,
+    backend_kwargs: dict[str, Any] | None = None,
+) -> CubeList:
+    if isinstance(file, Path):
+        zarr_xr = xr.open_dataset(
+            file,
+            consolidated=False,
+            engine="zarr",
+        )
+    elif urlparse(file).scheme in ["http", "https"]:
+        # basic test that opens the Zarr/.zmetadata file for Zarr2
+        # or Zarr/zarr.json for Zarr3
+        fs = fsspec.filesystem("http")
+        zarr2 = zarr3 = True
+        try:
+            fs.open(file + "/.zmetadata", "rb")  # Zarr2
+        except Exception:  # noqa: BLE001
+            zarr2 = False
+        try:
+            fs.open(file + "/zarr.json", "rb")  # Zarr3
+        except Exception:  # noqa: BLE001
+            zarr3 = False
+        # we don't want to catch any specific aiohttp/fsspec exception
+        # bottom line is that that file has issues, so raise
+        if not zarr2 and not zarr3:
+            msg = f"File '{file}' can not be open as Zarr file at the moment."
+            raise ValueError(msg) from None
+
+        time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)
+        zarr_xr = xr.open_dataset(
+            file,
+            consolidated=True,
+            decode_times=time_coder,
+            engine="zarr",
+            backend_kwargs=backend_kwargs,
+        )
+    else:
+        return []
+
+    return dataset_to_iris(zarr_xr, ignore_warnings=ignore_warnings)
+
+
 def _load_from_file(
     file: str | Path,
     ignore_warnings: list[dict[str, Any]] | None = None,

pyproject.toml

@@ -32,6 +32,7 @@ dynamic = [
     "version",
 ]
 dependencies = [
+    "aiohttp",
     "cartopy",
     "cf-units",
     "dask[array,distributed]>=2025",  # Core/issues/2503
@@ -44,6 +45,7 @@ dependencies = [
     "fire",
     "geopy",
     "humanfriendly",
+    "intake-esm",
     "iris-grib>=0.20.0",  # github.com/ESMValGroup/ESMValCore/issues/2535
     "isodate>=0.7.0",
     "jinja2",
@@ -68,6 +70,7 @@ dependencies = [
     "stratify>=0.3",
     "xarray",
     "yamale",
+    "zarr>3",
 ]
 description = "A community tool for pre-processing data from Earth system models in CMIP and running analysis scripts"
 license = {text = "Apache License, Version 2.0"}

tests/integration/preprocessor/_io/test_load.py

@@ -4,6 +4,7 @@
 from importlib.resources import files as importlib_files
 from pathlib import Path
 
+import cf_units
 import iris
 import ncdata
 import numpy as np
@@ -120,6 +121,134 @@ def test_load_ncdata():
     assert not cube.coords()
 
 
+def test_load_zarr_local():
+    """Test loading a Zarr store as ncdata.NcData via Xarray."""
+    zarr_path = (
+        Path(importlib_files("tests"))
+        / "sample_data"
+        / "zarr-sample-data"
+        / "example_field_0.zarr2"
+    )
+
+    cubes = load(zarr_path)
+
+    assert len(cubes) == 1
+    cube = cubes[0]
+    assert cube.var_name == "q"
+    assert cube.standard_name == "specific_humidity"
+    assert cube.long_name is None
+    assert cube.units == cf_units.Unit("1")
+    coords = cube.coords()
+    coord_names = [coord.standard_name for coord in coords]
+    assert "longitude" in coord_names
+    assert "latitude" in coord_names
+
+
+def test_load_zarr_remote():
+    """
+    Test loading a Zarr store from a https Object Store.
+
+    We have a permanent bucket: esmvaltool-zarr at CEDA's object store
+    "url": "https://uor-aces-o.s3-ext.jc.rl.ac.uk",
+    where will host a number of test files.
+    """
+    zarr_path = (
+        "https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
+        "esmvaltool-zarr/example_field_0.zarr2"
+    )
+
+    # with "dummy" storage options
+    cubes = load(
+        zarr_path,
+        ignore_warnings=None,
+        backend_kwargs={"storage_options": {}},
+    )
+
+    assert len(cubes) == 1
+    cube = cubes[0]
+    assert cube.var_name == "q"
+    assert cube.standard_name == "specific_humidity"
+    assert cube.long_name is None
+    assert cube.units == cf_units.Unit("1")
+    coords = cube.coords()
+    coord_names = [coord.standard_name for coord in coords]
+    assert "longitude" in coord_names
+    assert "latitude" in coord_names
+
+    # without storage_options
+    cubes = load(zarr_path)
+
+    assert len(cubes) == 1
+    cube = cubes[0]
+    assert cube.var_name == "q"
+    assert cube.standard_name == "specific_humidity"
+    assert cube.long_name is None
+    assert cube.units == cf_units.Unit("1")
+    coords = cube.coords()
+    coord_names = [coord.standard_name for coord in coords]
+    assert "longitude" in coord_names
+    assert "latitude" in coord_names
+
+
+def test_load_zarr_remote_not_zarrfile():
+    """
+    Test loading a Zarr store from a https Object Store.
+
+    This fails due to the file being loaded not a Zarr file.
+    We have a permanent bucket: esmvaltool-zarr at CEDA's object store
+    "url": "https://uor-aces-o.s3-ext.jc.rl.ac.uk",
+    where will host a number of test files.
+    """
+    zarr_path = (
+        "https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
+        "esmvaltool-zarr/example_field_0.zarr17"
+    )
+
+    msg = (
+        "File 'https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
+        "esmvaltool-zarr/example_field_0.zarr17' can not "
+        "be open as Zarr file at the moment."
+    )
+    with pytest.raises(ValueError, match=msg):
+        load(zarr_path)
+
+
+def test_load_zarr_remote_not_file():
+    """
+    Test loading a Zarr store from a https Object Store.
+
+    This fails due to non-existing file.
+    We have a permanent bucket: esmvaltool-zarr at CEDA's object store
+    "url": "https://uor-aces-o.s3-ext.jc.rl.ac.uk",
+    where will host a number of test files.
+    """
+    zarr_path = (
+        "https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
+        "esmvaltool-zarr/example_field_0.zarr22"
+    )
+
+    msg = (
+        "File 'https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
+        "esmvaltool-zarr/example_field_0.zarr22' can not "
+        "be open as Zarr file at the moment."
+    )
+    with pytest.raises(ValueError, match=msg):
+        load(zarr_path)
+
+
+def test_load_not_zarr():
+    """
+    Test loadinng something that has a zarr extension.
+
+    But file doesn't hold any data / doesn't exist.
+    """
+    zarr_path = "esmvaltool-zarr/example_field_0.zarr22"
+
+    msg = "esmvaltool-zarr/example_field_0.zarr22 does not contain any data"
+    with pytest.raises(ValueError, match=msg):
+        load(zarr_path)
+
+
 def test_load_invalid_type_fail():
     """Test loading an invalid type."""
     with pytest.raises(TypeError):