Add support for handling zarr urls

changelog/7113.breaking.rst

@@ -0,0 +1,4 @@
+:user:`trexfeathers` altered :func:`iris.io.decode_uri` to separate out URI
+**fragments** (the section following the ``#`` character) -
+:func:`~iris.io.decode_uri` now returns a 3-part tuple: ``scheme``, ``path``,
+and ``fragment``.

changelog/7113.feature.rst

@@ -0,0 +1,4 @@
+:user:`trexfeathers` and :user:`stephenworsley` added support for loading and 
+saving of Zarr files formatted in Zarr Storage Spec Version 2 via the
+`NetCDF NCZarr <https://docs.unidata.ucar.edu/nug/current/nczarr_head.html>`_
+implementation. (:issue:`6967`)

changelog/7113.internal.rst

@@ -0,0 +1,4 @@
+:user:`trexfeathers` and :user:`stephenworsley` added functionality to the NetCDF
+loader and saver to handle Zarr file I/O using the
+`NetCDF NCZarr <https://docs.unidata.ucar.edu/nug/current/nczarr_head.html>`_
+implementation. (:issue:`6967`)

docs/src/common_links.inc

@@ -27,6 +27,7 @@
 .. _matplotlib: https://matplotlib.org/stable/
 .. _napolean: https://sphinxcontrib-napoleon.readthedocs.io/en/latest/sphinxcontrib.napoleon.html
 .. _nc-time-axis: https://github.com/SciTools/nc-time-axis
+.. _NcZarr: https://docs.unidata.ucar.edu/nug/current/nczarr_head.html
 .. _nox: https://nox.thea.codes/en/stable/
 .. _New Issue: https://github.com/scitools/iris/issues/new/choose
 .. _pre-commit: https://pre-commit.com/

docs/src/user_manual/tutorial/loading_iris_cubes.rst

@@ -3,6 +3,8 @@
 
    A lesson in how Iris loading works.
 
+.. include:: ../../common_links.inc
+
 .. _loading_iris_cubes:
 
 ===================
@@ -34,9 +36,21 @@ of the given files and attempts to produce Iris Cubes from their contents.
 
 .. note::
 
-    Currently there is support for CF NetCDF, GRIB 1 & 2, PP and FieldsFiles
+    Currently there is support for CF NetCDF, GRIB 1 & 2, PP, FieldsFiles and
+    Zarr (via `NcZarr`_)
     file formats with a framework for this to be extended to custom formats.
 
+.. admonition:: NcZarr loading
+
+   NcZarr provides I/O for Zarr files via the NetCDF API, allowing Iris to use
+   its existing NetCDF loading code to read Zarr files. As discussed in the
+   `NcZarr`_ docs, NcZarr receives filepaths as URLs, with the appropriate
+   ``mode`` fragments. Iris will load via NcZarr if it detects a URL with these
+   fragments. For example:
+
+   .. code-block:: python
+
+       cubes = iris.load("file:///path/to/file.zarr#mode=nczarr,file")
 
 In order to find out what has been loaded, the result can be printed:
 

docs/src/user_manual/tutorial/saving_iris_cubes.rst

@@ -3,6 +3,8 @@
 
    A lesson on writing Iris Cubes to file in various formats.
 
+.. include:: ../../common_links.inc
+
 .. _saving_iris_cubes:
 
 ==================
@@ -16,6 +18,7 @@ Iris supports the saving of cubes and cube lists to:
 * CF netCDF (version 1.7)
 * GRIB edition 2  (if `iris-grib  <https://github.com/SciTools/iris-grib>`_ is installed)
 * Met Office PP
+* Zarr (via `NcZarr`_)
 
 
 The :py:func:`iris.save` function saves one or more cubes to a file.
@@ -45,6 +48,18 @@ and the keyword argument `saver` is not required.
     data can result in corruption. Users should proceed with caution when
     attempting to overwrite an existing file.
 
+.. admonition:: NcZarr saving
+
+   NcZarr provides I/O for Zarr files via the NetCDF API, allowing Iris to use
+   its existing NetCDF saving code to read Zarr files. As discussed in the
+   `NcZarr`_ docs, NcZarr receives filepaths as URLs, with the appropriate
+   ``mode`` fragments. Iris will save via NcZarr if it detects a URL with these
+   fragments. For example:
+
+   .. code-block:: python
+
+       cubes = iris.save(my_cube, "file:///path/to/file.zarr#mode=nczarr,file")
+
 
 Controlling the Save Process
 ----------------------------

lib/iris/__init__.py

@@ -55,7 +55,7 @@
 
     .. warning::
 
-        If supplying a URL, only OPeNDAP Data Sources are supported.
+        If supplying a URL, only OPeNDAP and NcZarr data sources are supported.
 
 * constraints:
     Either a single constraint, or an iterable of constraints.

lib/iris/fileformats/__init__.py

@@ -10,6 +10,7 @@
    API reference
 """
 
+from iris.io import _is_nczarr_fragment
 from iris.io.format_picker import (
     DataSourceObjectProtocol,
     FileExtension,
@@ -33,6 +34,16 @@
 )
 
 
+def _uri_is_http_not_nczarr(uri_parts):
+    scheme, _part, fragment = uri_parts
+    return scheme in ["http", "https"] and not _is_nczarr_fragment(fragment)
+
+
+def _uri_is_nczarr(uri_parts):
+    _scheme, _part, fragment = uri_parts
+    return _is_nczarr_fragment(fragment)
+
+
 #
 # PP files.
 #
@@ -131,7 +142,7 @@ def _load_grib(*args, **kwargs):
     FormatSpecification(
         "NetCDF OPeNDAP",
         UriProtocol(),
-        lambda protocol: protocol in ["http", "https"],
+        _uri_is_http_not_nczarr,
         netcdf.load_cubes,
         priority=6,
         constraint_aware_handler=True,
@@ -155,6 +166,20 @@ def _load_grib(*args, **kwargs):
 )
 
 
+#
+# Zarr files.
+#
+FORMAT_AGENT.add_spec(
+    FormatSpecification(
+        "NcZarr",
+        UriProtocol(),
+        _uri_is_nczarr,
+        netcdf.load_cubes,
+        priority=3,
+    )
+)
+
+
 #
 # UM Fieldsfiles.
 #

lib/iris/fileformats/cf.py

@@ -52,6 +52,9 @@
     re.VERBOSE,
 )
 
+# NCZarr stores scalar variables as size-1 arrays on this pseudo-dimension.
+_NCZARR_SCALAR_DIMENSION = "_scalar_"
+
 # NetCDF variable attributes handled by the netCDF4 module and
 # therefore automatically classed as "used" attributes.
 _CF_ATTRS_IGNORE = set(["_FillValue", "add_offset", "missing_value", "scale_factor"])
@@ -171,7 +174,11 @@ def spans(self, cf_variable):
         bool
 
         """
-        result = set(self.dimensions).issubset(cf_variable.dimensions)
+        dimensions = tuple(self.dimensions)
+        if dimensions == (_NCZARR_SCALAR_DIMENSION,):
+            return True
+
+        result = set(dimensions).issubset(cf_variable.dimensions)
         return result
 
     def __eq__(self, other):

lib/iris/fileformats/netcdf/loader.py

@@ -650,18 +650,24 @@ def inner(cf_datavar):
 def load_cubes(file_sources, callback=None, constraints=None):
     """Load cubes from a list of NetCDF filenames/OPeNDAP URLs.
 
+    Also supports Zarr files in the NcZarr URL format, e.g.
+    ``file:///path/to/file#mode=nczarr,file``. Note that NcZarr is limited to
+    Zarr Storage Specification version 2. See the
+    `NcZarr docs <https://docs.unidata.ucar.edu/nug/current/nczarr_head.html>`
+    for more.
+
     Parameters
     ----------
     file_sources : str or list
-        One or more NetCDF filenames/OPeNDAP URLs to load from.
+        One or more NetCDF filenames/OPeNDAP URLs/NcZarr URLs to load from.
         OR open datasets.
     callback : function, optional
         Function which can be passed on to :func:`iris.io.run_callback`.
     constraints : optional
 
     Returns
     -------
-    Generator of loaded NetCDF :class:`iris.cube.Cube`.
+    Generator of loaded NetCDF/NcZarr :class:`iris.cube.Cube`.
 
     """
     # Deferred import to avoid circular imports.

lib/iris/fileformats/netcdf/saver.py

@@ -27,6 +27,7 @@
 import string
 import typing
 from typing import TYPE_CHECKING
+from urllib.parse import urlsplit
 import warnings
 
 import cf_units
@@ -317,15 +318,21 @@ class VariableEmulator(typing.Protocol):
 
 
 class Saver:
-    """A manager for saving netcdf files."""
+    """A manager for saving NetCDF/NcZarr files."""
 
     def __init__(self, filename, netcdf_format, compute=True):
         """Manage saving netcdf files.
 
+        Also supports Zarr files in the NcZarr URL format, e.g.
+        ``file:///path/to/file#mode=nczarr,file``. Note that NcZarr is limited to
+        Zarr Storage Specification version 2. See the
+        `NcZarr docs <https://docs.unidata.ucar.edu/nug/current/nczarr_head.html>`
+        for more.
+
         Parameters
         ----------
         filename : str or netCDF4.Dataset
-            Name of the netCDF file to save the cube.
+            Name of the NetCDF file, or an NcZarr URL, to save the cube.
             OR a writeable object supporting the :class:`netCF4.Dataset` api.
         netcdf_format : str
             Underlying netCDF file format, one of 'NETCDF4', 'NETCDF4_CLASSIC',
@@ -396,9 +403,15 @@ def __init__(self, filename, netcdf_format, compute=True):
             None  # this line just for the API page -- value is set later
         )
 
+        #: Whether the save target is an NCZarr URL.
+        self._is_nczarr = False
         # A list of delayed writes for lazy saving
         # a list of couples (source, target).
         self._delayed_writes = []
+        # A list of pre-close writes for NCZarr lazy saving.
+        # NCZarr cannot be reopened for deferred writes, so lazy data is
+        # written before __exit__ closes the file.
+        self._nczarr_writes = []
 
         # Detect if we were passed a pre-opened dataset (or something like one)
         self._to_open_dataset = hasattr(filename, "createVariable")
@@ -426,12 +439,26 @@ def __init__(self, filename, netcdf_format, compute=True):
         else:
             # Given a filepath string/path : create a dataset from that
             try:
-                self.filepath = Path(filename).absolute()
+                # Lazy import to avoid circular import overhead at module import-time.
+                from iris.io import _is_nczarr_fragment
+
+                self._is_nczarr = _is_nczarr_fragment(
+                    urlsplit(str(filename)).fragment or None
+                )
+                if self._is_nczarr:
+                    # NCZarr URLs contain a #mode= fragment; Path() strips it.
+                    # Keep as a plain string and pass directly to DatasetWrapper.
+                    self.filepath = str(filename)
+                else:
+                    filepath = Path(filename)
+                    self.filepath = filepath.absolute()
                 self._dataset = _thread_safe_nc.DatasetWrapper(
                     self.filepath, mode="w", format=netcdf_format
                 )
             except RuntimeError:
-                dir_name = self.filepath.parent
+                if self._is_nczarr:
+                    raise
+                dir_name = Path(self.filepath).parent
                 if not dir_name.is_dir():
                     msg = "No such file or directory: {}".format(dir_name)
                     raise IOError(msg)
@@ -447,7 +474,12 @@ def __enter__(self):
         return self
 
     def __exit__(self, type, value, traceback):
-        """Flush any buffered data to the CF-netCDF file before closing."""
+        """Flush any buffered data to the CF-NetCDF/NcZarr file before closing."""
+        if self._nczarr_writes:
+            # NCZarr lazy writes must be computed while the file is still open;
+            # the deferred reopen-write pattern used for netCDF is not supported.
+            sources, targets = zip(*self._nczarr_writes)
+            da.store(list(sources), list(targets))
         self._dataset.sync()
         if not self._to_open_dataset:
             # Only close if the Saver created it.
@@ -474,10 +506,16 @@ def write(
     ):
         """Wrap for saving cubes to a NetCDF file.
 
+        Also supports Zarr files in the NcZarr URL format, e.g.
+        ``file:///path/to/file#mode=nczarr,file``. Note that NcZarr is limited to
+        Zarr Storage Specification version 2. See the
+        `NcZarr docs <https://docs.unidata.ucar.edu/nug/current/nczarr_head.html>`
+        for more.
+
         Parameters
         ----------
         cube : :class:`iris.cube.Cube`
-            A :class:`iris.cube.Cube` to be saved to a netCDF file.
+            A :class:`iris.cube.Cube` to be saved to a NetCDF/NcZarr file.
         local_keys : iterable of str, optional
             An iterable of cube attribute keys. Any cube attributes with
             matching keys will become attributes on the data variable rather
@@ -2544,6 +2582,12 @@ def _lazy_stream_data(
         else:
             doing_delayed_save = is_lazy_data(data)
             if doing_delayed_save:
+                if self._is_nczarr:
+                    # NCZarr cannot be reopened for deferred writes.
+                    # Collect here; da.store() will be called in __exit__.
+                    self._nczarr_writes.append((data, cf_var))
+                    return
+
                 # save lazy data with a delayed operation.  For now, we just record the
                 # necessary information -- a single, complete delayed action is constructed
                 # later by a call to delayed_completion().
@@ -2635,7 +2679,13 @@ def save(
     fill_value=None,
     compute=True,
 ):
-    r"""Save cube(s) to a netCDF file, given the cube and the filename.
+    r"""Save cube(s) to a NetCDF file, given the cube and the filename.
+
+    Also supports Zarr files in the NcZarr URL format, e.g.
+    ``file:///path/to/file#mode=nczarr,file``. Note that NcZarr is limited to
+    Zarr Storage Specification Version 2. See the
+    `NcZarr docs <https://docs.unidata.ucar.edu/nug/current/nczarr_head.html>`
+    for more.
 
     * Iris will write CF 1.7 compliant NetCDF files.
     * **If split-attribute saving is disabled**, i.e.
@@ -2661,7 +2711,7 @@ def save(
         A :class:`iris.cube.Cube`, :class:`iris.cube.CubeList` or other
         iterable of cubes to be saved to a netCDF file.
     filename : str
-        Name of the netCDF file to save the cube(s).
+        Name of the NetCDF file or NcZarr URL to save the cube(s).
         **Or** an open, writeable :class:`netCDF4.Dataset`, or compatible object.
 
         .. note::