Zarr I/O PR 7: real-world test data and benchmarks

.github/workflows/benchmarks_run.yml

@@ -73,7 +73,7 @@ jobs:
     env:
       IRIS_TEST_DATA_LOC_PATH: benchmarks
       IRIS_TEST_DATA_PATH: benchmarks/iris-test-data
-      IRIS_TEST_DATA_VERSION: "2.30"
+      IRIS_TEST_DATA_VERSION: "2.32"
       # Lets us manually bump the cache to rebuild
       ENV_CACHE_BUILD: "1"
       TEST_DATA_CACHE_BUILD: "2"

.github/workflows/ci-tests.yml

@@ -54,7 +54,7 @@ jobs:
 
     env:
       # NOTE: IRIS_TEST_DATA_VERSION is also set in benchmarks_run.yml
-      IRIS_TEST_DATA_VERSION: "2.30"
+      IRIS_TEST_DATA_VERSION: "2.32"
       ENV_NAME: "ci-tests"
 
     steps:

benchmarks/benchmarks/generate_data/um_files.py

@@ -2,7 +2,7 @@
 #
 # This file is part of Iris and is released under the BSD license.
 # See LICENSE in the root of the repository for full licensing details.
-"""Generate FF, PP and NetCDF files based on a minimal synthetic FF file.
+"""Generate FF, PP, NetCDF and Zarr files based on a minimal synthetic FF file.
 
 NOTE: uses the Mule package, so depends on an environment with Mule installed.
 """
@@ -122,12 +122,13 @@ def add_field(level_: int, time_step_: int) -> None:
     ff_path = save_paths.get("FF", None)
     pp_path = save_paths.get("PP", None)
     nc_path = save_paths.get("NetCDF", None)
+    zarr_path = save_paths.get("Zarr", None)
 
     if ff_path:
         new_ff.to_file(ff_path)
     if pp_path:
         fields_to_pp_file(str(pp_path), new_ff.fields)
-    if nc_path:
+    if nc_path or zarr_path:
         temp_ff_path = None
         # Need an Iris Cube from the FF content.
         if ff_path:
@@ -139,12 +140,18 @@ def add_field(level_: int, time_step_: int) -> None:
             new_ff.to_file(temp_ff_path.name)
             ff_cube = load_cube(temp_ff_path.name)
 
-        save_cube(ff_cube, nc_path, zlib=compress)
+        if nc_path:
+            save_cube(ff_cube, nc_path, zlib=compress)
+        if zarr_path:
+            # Zarr's own spelling of compression. The Zarr saver rejects
+            #  netCDF keywords such as zlib= rather than translating them.
+            compressors = "auto" if compress else None
+            save_cube(ff_cube, zarr_path, mode="w", compressors=compressors)
         if temp_ff_path:
             temp_ff_path.close()
 
 
-FILE_EXTENSIONS = {"FF": "", "PP": ".pp", "NetCDF": ".nc"}
+FILE_EXTENSIONS = {"FF": "", "PP": ".pp", "NetCDF": ".nc", "Zarr": ".zarr"}
 
 
 def create_um_files(
@@ -177,7 +184,9 @@ def create_um_files(
     for file_type in file_types:
         file_ext = FILE_EXTENSIONS[file_type]
         save_path = (save_dir / f"{compress}").with_suffix(file_ext)
-        files_exist = files_exist and save_path.is_file()
+        # exists() not is_file(): a Zarr store is a directory, so is_file()
+        #  is always False for one and REUSE_DATA would never hold.
+        files_exist = files_exist and save_path.exists()
         save_paths[file_type] = str(save_path)
 
     if not REUSE_DATA or not files_exist:

benchmarks/benchmarks/load/__init__.py

@@ -18,7 +18,7 @@ class LoadAndRealise:
     params = (
         [(50, 50, 2), (1280, 960, 5), (2, 2, 1000)],
         [False, True],
-        ["FF", "PP", "NetCDF"],
+        ["FF", "PP", "NetCDF", "Zarr"],
     )
     param_names = ["xyz", "compressed", "file_format"]
 
@@ -78,7 +78,7 @@ def time_stash_constraint(self, _, __, ___) -> None:
 
 
 class TimeConstraint:
-    params = ([3, 20], ["FF", "PP", "NetCDF"])
+    params = ([3, 20], ["FF", "PP", "NetCDF", "Zarr"])
     param_names = ["time_dim_len", "file_format"]
 
     def setup_cache(self) -> dict:

benchmarks/benchmarks/save.py

@@ -9,35 +9,48 @@
 
 from .generate_data.ugrid import make_cube_like_2d_cubesphere
 
+#: Target name and save keywords per format. Zarr needs mode="w" because its
+#:  saver defaults to mode="w-" and raises FileExistsError on a target that
+#:  exists - and ASV calls a timing function many times between setup() calls.
+SAVE_TARGETS = {
+    "NetCDF": ("tmp.nc", {}),
+    "Zarr": ("tmp.zarr", {"mode": "w"}),
+}
 
-class NetcdfSave:
-    params = [[50, 600], [False, True]]
-    param_names = ["cubesphere-N", "is_unstructured"]
 
-    def setup(self, n_cubesphere, is_unstructured):
+class Save:
+    params = [[50, 600], [False, True], list(SAVE_TARGETS)]
+    param_names = ["cubesphere-N", "is_unstructured", "file_format"]
+
+    def setup(self, n_cubesphere, is_unstructured, file_format):
         self.cube = make_cube_like_2d_cubesphere(
             n_cube=n_cubesphere, with_mesh=is_unstructured
         )
+        self.save_target, self.save_kwargs = SAVE_TARGETS[file_format]
 
     def _save_data(self, cube, do_copy=True):
         if do_copy:
             # Copy the cube, to avoid distorting the results by changing it
             # Because we known that older Iris code realises lazy coords
             cube = cube.copy()
-        save(cube, "tmp.nc")
+        save(cube, self.save_target, **self.save_kwargs)
 
     def _save_mesh(self, cube):
         # In this case, we are happy that the mesh is *not* modified
         save_mesh(cube.mesh, "mesh.nc")
 
-    def time_netcdf_save_cube(self, n_cubesphere, is_unstructured):
+    def time_save_cube(self, n_cubesphere, is_unstructured, file_format):
         self._save_data(self.cube)
 
-    def time_netcdf_save_mesh(self, n_cubesphere, is_unstructured):
-        if is_unstructured:
+    def time_save_mesh(self, n_cubesphere, is_unstructured, file_format):
+        # iris.mesh.save_mesh is defined in iris.fileformats.netcdf.saver and
+        #  writes netCDF only, so there is nothing to time for Zarr. A cube
+        #  that *has* a mesh still saves it through time_save_cube, with
+        #  is_unstructured=True.
+        if is_unstructured and file_format == "NetCDF":
             self._save_mesh(self.cube)
 
-    def tracemalloc_netcdf_save(self, n_cubesphere, is_unstructured):
+    def tracemalloc_save(self, n_cubesphere, is_unstructured, file_format):
         # Don't need to copy the cube here since track_ benchmarks don't
         #  do repeats between self.setup() calls.
         self._save_data(self.cube, do_copy=False)