Zarr to iris Cube: add tests and Zarr sample data (v2)

.github/workflows/ci-tests.yml

@@ -11,7 +11,6 @@ on:
   push:
     branches:
       - "main"
-      - "v*x"
     tags:
       - "v*"
   pull_request:
@@ -49,7 +48,7 @@ jobs:
 
       - name: "Install dependencies"
         run: |
-          conda install --yes numpy pytest pytest-mock iris xarray filelock requests
+          conda install --yes numpy pytest pytest-mock iris xarray filelock requests zarr aiohttp fsspec
 
       - name: "Install *latest* Iris"
         run: |

docs/changelog_fragments/145.dev.rst

@@ -0,0 +1 @@
+@valeriupredoi added test for Zarr conversion to Iris cubes.

tests/integration/test_zarr_to_iris.py

@@ -0,0 +1,111 @@
+"""Test conversion of remote and local Zarr store to iris Cube."""
+
+from importlib.resources import files as importlib_files
+from pathlib import Path
+
+import fsspec
+import iris
+import pytest
+import xarray as xr
+from ncdata.iris_xarray import cubes_from_xarray as conversion_func
+
+
+def _return_kwargs():
+    time_coder = xr.coders.CFDatetimeCoder(use_cftime=True)
+    return {
+        "consolidated": True,
+        "decode_times": time_coder,
+        "engine": "zarr",
+        "chunks": {},
+        "backend_kwargs": {},
+    }
+
+
+def _run_checks(cube):
+    """Run some standard checks."""
+    assert cube.var_name == "q"
+    assert cube.standard_name == "specific_humidity"
+    assert cube.long_name is None
+    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_zarr2_local():
+    """Test loading a Zarr2 store from local FS."""
+    zarr_path = (
+        Path(importlib_files("tests"))
+        / "testdata"
+        / "zarr-sample-data"
+        / "example_field_0.zarr2"
+    )
+
+    xr_kwargs = _return_kwargs()
+    zarr_xr = xr.open_dataset(zarr_path, **xr_kwargs)
+
+    cubes = conversion_func(zarr_xr)
+
+    assert len(cubes) == 1
+    cube = cubes[0]
+    _run_checks(cube)
+
+
+def test_load_zarr3_local():
+    """Test loading a Zarr3 store from local FS."""
+    zarr_path = (
+        Path(importlib_files("tests"))
+        / "testdata"
+        / "zarr-sample-data"
+        / "example_field_0.zarr3"
+    )
+
+    xr_kwargs = _return_kwargs()
+    zarr_xr = xr.open_dataset(zarr_path, **xr_kwargs)
+
+    cubes = conversion_func(zarr_xr)
+
+    assert len(cubes) == 1
+    cube = cubes[0]
+    _run_checks(cube)
+
+
+def _is_url_ok(url):
+    fs = fsspec.filesystem("http")
+    valid_zarr = True
+    try:
+        fs.open(str(url) + "/zarr.json", "rb")  # Zarr3
+    except Exception:  # noqa: BLE001
+        try:
+            fs.open(str(url) + "/.zmetadata", "rb")  # Zarr2
+        except Exception:  # noqa: BLE001
+            valid_zarr = False
+
+    return valid_zarr
+
+
+S3_TEST_PATH = (
+    "https://uor-aces-o.s3-ext.jc.rl.ac.uk/"
+    "esmvaltool-zarr/pr_Amon_CNRM-ESM2-1_02Kpd-11_r1i1p2f2_gr_200601-220112.zarr3"
+)
+_S3_accessible = _is_url_ok(S3_TEST_PATH)
+
+
+@pytest.mark.skipif(not _S3_accessible, reason="S3 url not accessible")
+def test_load_remote_zarr():
+    """Test loading a remote Zarr store.
+
+    This is a ~250MB compressed Zarr in an S3 bucket.
+    Conversion is done fully lazily, by passing chunks={}
+    to Xarray loader. Test takes ~3-4s and needs ~400MB res mem.
+    """
+    zarr_path = S3_TEST_PATH
+
+    xr_kwargs = _return_kwargs()
+    zarr_xr = xr.open_dataset(zarr_path, **xr_kwargs)
+
+    cubes = conversion_func(zarr_xr)
+
+    assert isinstance(cubes, iris.cube.CubeList)
+    assert len(cubes) == 1
+    assert cubes[0].has_lazy_data()

tests/testdata/zarr-sample-data/example_field_0.zarr2/.zattrs

@@ -0,0 +1,3 @@
+{
+    "Conventions": "CF-1.12"
+}

tests/testdata/zarr-sample-data/example_field_0.zarr2/.zgroup

@@ -0,0 +1,3 @@
+{
+    "zarr_format": 2
+}

tests/testdata/zarr-sample-data/example_field_0.zarr2/.zmetadata

@@ -0,0 +1,171 @@
+{
+    "metadata": {
+        ".zattrs": {
+            "Conventions": "CF-1.12"
+        },
+        ".zgroup": {
+            "zarr_format": 2
+        },
+        "lat/.zarray": {
+            "chunks": [
+                5
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f8",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                5
+            ],
+            "zarr_format": 2
+        },
+        "lat/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "lat"
+            ],
+            "bounds": "lat_bnds",
+            "standard_name": "latitude",
+            "units": "degrees_north"
+        },
+        "lat_bnds/.zarray": {
+            "chunks": [
+                3,
+                2
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f8",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                5,
+                2
+            ],
+            "zarr_format": 2
+        },
+        "lat_bnds/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "lat",
+                "bounds2"
+            ]
+        },
+        "lon/.zarray": {
+            "chunks": [
+                8
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f8",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                8
+            ],
+            "zarr_format": 2
+        },
+        "lon/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "lon"
+            ],
+            "bounds": "lon_bnds",
+            "standard_name": "longitude",
+            "units": "degrees_east"
+        },
+        "lon_bnds/.zarray": {
+            "chunks": [
+                4,
+                2
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f8",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                8,
+                2
+            ],
+            "zarr_format": 2
+        },
+        "lon_bnds/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "lon",
+                "bounds2"
+            ]
+        },
+        "q/.zarray": {
+            "chunks": [
+                3,
+                4
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f8",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                5,
+                8
+            ],
+            "zarr_format": 2
+        },
+        "q/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "lat",
+                "lon"
+            ],
+            "cell_methods": "area: mean",
+            "coordinates": "time",
+            "project": "research",
+            "standard_name": "specific_humidity",
+            "units": "1"
+        },
+        "time/.zarray": {
+            "chunks": [],
+            "compressor": null,
+            "dtype": "<f8",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [],
+            "zarr_format": 2
+        },
+        "time/.zattrs": {
+            "_ARRAY_DIMENSIONS": [],
+            "standard_name": "time",
+            "units": "days since 2018-12-01"
+        }
+    },
+    "zarr_consolidated_format": 1
+}

tests/testdata/zarr-sample-data/example_field_0.zarr2/lat/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        5
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f8",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        5
+    ],
+    "zarr_format": 2
+}

tests/testdata/zarr-sample-data/example_field_0.zarr2/lat/.zattrs

@@ -0,0 +1,8 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "lat"
+    ],
+    "bounds": "lat_bnds",
+    "standard_name": "latitude",
+    "units": "degrees_north"
+}

tests/testdata/zarr-sample-data/example_field_0.zarr2/lat_bnds/.zarray

@@ -0,0 +1,22 @@
+{
+    "chunks": [
+        3,
+        2
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f8",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        5,
+        2
+    ],
+    "zarr_format": 2
+}

tests/testdata/zarr-sample-data/example_field_0.zarr2/lat_bnds/.zattrs

@@ -0,0 +1,6 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "lat",
+        "bounds2"
+    ]
+}

tests/testdata/zarr-sample-data/example_field_0.zarr2/lon/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        8
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f8",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        8
+    ],
+    "zarr_format": 2
+}