Add open_raw_zarr helper (#2668)

* Add open_raw_zarr helper

* Make open_raw_zarr read metadata

And add test

* Add check that data_vars are Zarr arrays

Author: VeckoTheGecko

src/parcels/__init__.py

@@ -10,6 +10,7 @@
 import warnings as _stdlib_warnings
 
 from parcels._core.fieldset import FieldSet
+from parcels._xarray import open_raw_zarr
 from parcels._core.particleset import ParticleSet
 from parcels._core.particlefile import ParticleFile, read_particlefile
 from parcels._core.particle import (
@@ -42,6 +43,7 @@
 __all__ = [  # noqa: RUF022
     # Core classes
     "FieldSet",
+    "open_raw_zarr",
     "ParticleSet",
     "ParticleFile",
     "Variable",

src/parcels/_xarray.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+import xarray as xr
+import zarr
+import zarr.storage
+from zarr.abc.store import Store
+
+
+def _not_implemented(*args, **kwargs):
+    raise NotImplementedError("This function it not implemented")
+
+
+def open_raw_zarr(store: Store):
+    """Open a Zarr dataset in an Xarray dataset, bypassing Dask."""
+    with xr.open_zarr(store) as ds:
+        var_to_dims = {name: var.dims for name, var in ds.variables.items()}
+        var_to_attrs = {name: var.attrs for name, var in ds.variables.items()}
+        coords = {name: ds[name].variable.load() for name in ds.coords}
+        ds_attrs = ds.attrs
+
+    group = zarr.open(store, mode="r")
+    assert isinstance(group, zarr.Group)
+
+    data_vars = {}
+    for name, array in group.members():
+        if not isinstance(array, zarr.Array):
+            raise ValueError("Discovered a zarr.Group in the root group. open_raw_zarr doesn't work with nested groups")
+        if name in coords:
+            continue
+
+        # trick xarray to prevent coersion to a numpy array
+        array.__array_function__ = _not_implemented  # type: ignore[attr-defined]
+        data_vars[name] = xr.Variable(var_to_dims[name], array, attrs=var_to_attrs[name])
+
+    return xr.Dataset(data_vars, coords, attrs=ds_attrs)

tests/test_xarray.py

@@ -0,0 +1,20 @@
+import pytest
+import xarray as xr
+import zarr
+
+from parcels import open_raw_zarr
+from parcels._datasets.structured.generic import datasets
+
+
+@pytest.mark.parametrize("ds", [pytest.param(v, id=k) for k, v in datasets.items()])
+def test_open_raw_zarr(ds, tmp_path):
+    path = tmp_path / "ds.zarr"
+    ds.to_zarr(path)
+
+    result = open_raw_zarr(path)
+
+    for k in result.data_vars:
+        # tests that the internal representation within Xarray isn't coerced into a numpy array
+        assert isinstance(result[k]._variable._data, zarr.Array)
+
+    xr.testing.assert_identical(result.load(), ds)