Adding a warning when changing data with zarr backend

src/parcels/_core/model.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import warnings
 from abc import ABC, abstractmethod
 from collections.abc import Hashable, Sequence
 from typing import Any, Self
 
 import cf_xarray  # noqa: F401
 import uxarray as ux
 import xarray as xr
+import zarr
 from dask import is_dask_collection
 
 import parcels._sgrid as sgrid
@@ -131,6 +133,17 @@ def preprocess_sgrid_model_data(ds: xr.Dataset) -> xr.Dataset:
     return ds
 
 
+def validate_field_data(ds: xr.Dataset) -> xr.Dataset:
+    if any(isinstance(da.variable._data, zarr.Array) for da in ds.data_vars.values()):
+        warnings.warn(
+            "Changing a Zarr-backed dataset. This may convert the Parcels backend to NumPy. "
+            "If you want to keep the Zarr backend, please use `skip_field_data_validation=True` when creating the FieldSet.",
+            UserWarning,
+            stacklevel=2,
+        )
+    return ds.fillna(0)
+
+
 class StructuredModelData(ModelData):
     def __init__(
         self,
@@ -144,7 +157,7 @@ def __init__(
 
         data = preprocess_sgrid_model_data(data)
         if not skip_field_data_validation:
-            data = data.fillna(0)
+            data = validate_field_data(data)
         grid = XGrid(data, mesh)
 
         self.data = data

tests/test_fieldset.py

@@ -397,6 +397,21 @@ def test_fieldset_add_error_on_duplicate_context_values():
         fset1 + fset2
 
 
+@pytest.mark.parametrize("skip", [True, False])
+def test_zarr_warning_on_fieldset_creation(skip, tmp_path):
+    """Test that creating a FieldSet from a Zarr-backed dataset raises a warning about potential backend changes."""
+    ds = parcels.tutorial.open_dataset("CopernicusMarine_data_for_Argo_tutorial/data")
+    ds = convert.copernicusmarine_to_sgrid(fields={"U": ds["uo"], "V": ds["vo"]})
+    path = tmp_path / "ds.zarr"
+    ds.to_zarr(path)
+    ds_zarr = open_raw_zarr(path)
+    if not skip:
+        with pytest.warns(UserWarning, match="Changing a Zarr-backed dataset"):
+            FieldSet.from_sgrid_conventions(ds_zarr, skip_field_data_validation=skip)
+    else:
+        FieldSet.from_sgrid_conventions(ds_zarr, skip_field_data_validation=skip)
+
+
 def test_fieldset_add_context_values():
     """Test that context values from both FieldSets are present in the combined FieldSet."""
     ds1 = datasets_structured["ds_2d_left"][["U_A_grid", "grid"]].rename({"U_A_grid": "U1"})
@@ -494,15 +509,15 @@ def test_fieldset_describe_backends(tmp_path):
     path = tmp_path / "ds.zarr"
     ds_fset.to_zarr(path)
     ds_zarr = open_raw_zarr(path)
-    fieldset = FieldSet.from_sgrid_conventions(ds_zarr)
+    fieldset = FieldSet.from_sgrid_conventions(ds_zarr, skip_field_data_validation=True)
 
     io = StringIO()
     expected = """\
 | Name   | Type        |   Grid number | Interp method / value   | Parcels backend   |
 |:-------|:------------|--------------:|:------------------------|:------------------|
-| U      | Field       |             0 | XLinear(...)            | NumPy             |
-| V      | Field       |             0 | XLinear(...)            | NumPy             |
-| W      | Field       |             0 | XLinear(...)            | NumPy             |
+| U      | Field       |             0 | XLinear(...)            | Zarr              |
+| V      | Field       |             0 | XLinear(...)            | Zarr              |
+| W      | Field       |             0 | XLinear(...)            | Zarr              |
 | UV     | VectorField |             0 | CGrid_Velocity(...)     | -                 |
 | UVW    | VectorField |             0 | CGrid_Velocity(...)     | -                 |