Basic functional lazy saving.

lib/iris/fileformats/netcdf/saver.py

@@ -23,6 +23,7 @@
 import warnings
 
 import cf_units
+import dask
 import dask.array as da
 import netCDF4
 import numpy as np
@@ -490,10 +491,51 @@ def __setitem__(self, keys, arr):
 MESH_ELEMENTS = ("node", "edge", "face")
 
 
+class DeferredSaveWrapper:
+    """
+    An object which mimics the data access of a netCDF4.Variable, and can be written to.
+    It encapsulates the netcdf file and variable which are actually to be written to.
+    This opens the file each time, to enable writing the data chunk, then closes it.
+    TODO: could be improved with a caching scheme, but this just about works.
+    """
+
+    def __init__(self, filepath, cf_var):
+        # Grab useful properties of the variable, including the identifying 'name'.
+        self.path = filepath
+        for key in ("shape", "dtype", "ndim", "name"):
+            setattr(self, key, getattr(cf_var, key))
+
+    def __setitem__(self, keys, array_data):
+        # Write to the variable.
+        # Re-open the file for writing.
+        dataset = netCDF4.Dataset(self.path, "r+")
+        try:
+            var = dataset.variables[self.name]
+            var[keys] = array_data
+        finally:
+            dataset.close()
+
+    def __repr__(self):
+        fmt = (
+            "<{self.__class__.__name__} shape={self.shape}"
+            " dtype={self.dtype!r} path={self.path!r}"
+            " name={self.name!r}>"
+        )
+        return fmt.format(self=self)
+
+
+@dask.delayed
+def combined_delayeds(*args):
+    """A delayed function which simply computes all its arguments."""
+    # Dask computes the lazy args before passing them here.
+    # So job done -- we don't need to do anything with them.
+    pass
+
+
 class Saver:
     """A manager for saving netcdf files."""
 
-    def __init__(self, filename, netcdf_format):
+    def __init__(self, filename, netcdf_format, compute=True):
         """
         A manager for saving netcdf files.
 
@@ -506,6 +548,15 @@ def __init__(self, filename, netcdf_format):
             Underlying netCDF file format, one of 'NETCDF4', 'NETCDF4_CLASSIC',
             'NETCDF3_CLASSIC' or 'NETCDF3_64BIT'. Default is 'NETCDF4' format.
 
+        * compute (bool):
+            If True, the Saver performs normal 'synchronous' data writes, where data
+            is streamed directly into file variables during the save operation.
+            If False, the file is created as normal, but computation and streaming of
+            any lazy array content is instead deferred to :class:`dask.delayed` objects,
+            which are held in a list in the saver 'delayed_writes' property.
+            The relavant file variables are created empty, and the write can
+            subsequently be completed by computing the 'save.deferred_writes'.
+
         Returns:
             None.
 
@@ -542,7 +593,14 @@ def __init__(self, filename, netcdf_format):
         self._mesh_dims = {}
         #: A dictionary, mapping formula terms to owner cf variable name
         self._formula_terms_cache = {}
+        #: Whether lazy saving.
+        self.lazy_saves = not compute
+        #: A list of deferred writes (if lazy saving)
+        self.deferred_writes = []
+        #: Target filepath
+        self.filepath = filename
         #: NetCDF dataset
+        self._dataset = None
         try:
             self._dataset = netCDF4.Dataset(
                 filename, mode="w", format=netcdf_format
@@ -2442,8 +2500,7 @@ def _increment_name(self, varname):
 
         return "{}_{}".format(varname, num)
 
-    @staticmethod
-    def _lazy_stream_data(data, fill_value, fill_warn, cf_var):
+    def _lazy_stream_data(self, data, fill_value, fill_warn, cf_var):
         if hasattr(data, "shape") and data.shape == (1,) + cf_var.shape:
             # (Don't do this check for string data).
             # Reduce dimensionality where the data array has an extra dimension
@@ -2453,13 +2510,36 @@ def _lazy_stream_data(data, fill_value, fill_warn, cf_var):
             data = data.squeeze(axis=0)
 
         if is_lazy_data(data):
+            if self.lazy_saves:
+                # deferred lazy streaming
+                def store(data, cf_var, fill_value):
+                    # Create a data-writeable object that we can stream into, which
+                    # encapsulates the file to be opened + variable to be written.
+                    writeable_var_wrapper = DeferredSaveWrapper(
+                        self.filepath, cf_var
+                    )
+                    # Add a delayed save to our 'deferred_writes' list.
+                    self.deferred_writes.append(
+                        da.store(
+                            [data], [writeable_var_wrapper], compute=False
+                        )
+                    )
+                    # NOTE: in this case, no checking of fill-value violations so just
+                    # return dummy values for this.
+                    # TODO: just for now -- can probably make this work later
+                    is_masked, contains_value = False, False
+                    return is_masked, contains_value
 
-            def store(data, cf_var, fill_value):
-                # Store lazy data and check whether it is masked and contains
-                # the fill value
-                target = _FillValueMaskCheckAndStoreTarget(cf_var, fill_value)
-                da.store([data], [target])
-                return target.is_masked, target.contains_value
+            else:
+                # Immediate streaming store : check mask+fill as we go.
+                def store(data, cf_var, fill_value):
+                    # Store lazy data and check whether it is masked and contains
+                    # the fill value
+                    target = _FillValueMaskCheckAndStoreTarget(
+                        cf_var, fill_value
+                    )
+                    da.store([data], [target])
+                    return target.is_masked, target.contains_value
 
         else:
 
@@ -2526,6 +2606,7 @@ def save(
     least_significant_digit=None,
     packing=None,
     fill_value=None,
+    compute=True,
 ):
     """
     Save cube(s) to a netCDF file, given the cube and the filename.
@@ -2648,6 +2729,14 @@ def save(
         `:class:`iris.cube.CubeList`, or a single element, and each element of
         this argument will be applied to each cube separately.
 
+    * compute (bool):
+        When False, create the output file but defer writing any lazy array content to
+        its variables, such as (lazy) data and aux-coords points and bounds.
+        Instead return a class:`dask.delayed` which, when computed, will compute all
+        the lazy content and stream it to complete the file.
+        Several such data saves can be performed in parallel, by passing a list of them
+        into a :func:`dask.compute` call.
+
     Returns:
         None.
 
@@ -2748,7 +2837,7 @@ def is_valid_packspec(p):
                 raise ValueError(msg)
 
     # Initialise Manager for saving
-    with Saver(filename, netcdf_format) as sman:
+    with Saver(filename, netcdf_format, compute=compute) as sman:
         # Iterate through the cubelist.
         for cube, packspec, fill_value in zip(cubes, packspecs, fill_values):
             sman.write(
@@ -2793,3 +2882,10 @@ def is_valid_packspec(p):
 
         # Add conventions attribute.
         sman.update_global_attributes(Conventions=conventions)
+
+        if compute:
+            result = None
+        else:
+            # For lazy save, return a single 'delayed' representing all lazy writes.
+            result = combined_delayeds(sman.deferred_writes)
+        return result

lib/iris/io/__init__.py

@@ -444,7 +444,7 @@ def save(source, target, saver=None, **kwargs):
 
     # Single cube?
     if isinstance(source, Cube):
-        saver(source, target, **kwargs)
+        result = saver(source, target, **kwargs)
 
     # CubeList or sequence of cubes?
     elif isinstance(source, CubeList) or (
@@ -463,13 +463,18 @@ def save(source, target, saver=None, **kwargs):
             # Force append=True for the tail cubes. Don't modify the incoming
             # kwargs.
             kwargs = kwargs.copy()
+            result = []
             for i, cube in enumerate(source):
                 if i != 0:
                     kwargs["append"] = True
                 saver(cube, target, **kwargs)
+
+            result = None
         # Netcdf saver.
         else:
-            saver(source, target, **kwargs)
+            result = saver(source, target, **kwargs)
 
     else:
         raise ValueError("Cannot save; non Cube found in source")
+
+    return result