Merge pull request #390 from davidhassell/persist-environment

Persist data after computation

Author: davidhassell

Changelog.rst

@@ -3,6 +3,10 @@ Version NEXTVERSION
 
 **2026-??-??**
 
+* New keyword parameter to `cfdm.Data.compute`: ``persist``
+  (https://github.com/NCAS-CMS/cfdm/issues/389)
+* New function to control the persistence of computed data:
+  `cfdm.persist_data` (https://github.com/NCAS-CMS/cfdm/issues/389)
 * Support for HEALPix grids
   (https://github.com/NCAS-CMS/cfdm/issues/370)
 * New default backend for netCDF-4 in `cfdm.write`: ``h5netcdf-h5py``,

cfdm/__init__.py

@@ -65,6 +65,7 @@
     integer_dtype,
     log_level,
     parse_indices,
+    persist_data,
     rtol,
     unique_constructs,
     _disable_logging,

cfdm/core/data/data.py

@@ -834,7 +834,7 @@ def source(self, default=ValueError()):
         >>> f = {{package}}.read('file.nc')[0]
         >>> d = f.data
         >>> d.source()
-        <{{repr}}NetCDF4Array(149, 182): file=file.nc variable=latitude>
+        <{{repr}}PyfiveArray(149, 182): file.nc latitude(149, 182)>
 
         """
         return self._get_component("array", default=default)

cfdm/data/data.py

@@ -22,6 +22,7 @@
     display_data,
     is_log_level_info,
     parse_indices,
+    persist_data,
 )
 from ..mixin.container import Container
 from ..mixin.files import Files
@@ -140,6 +141,9 @@ def __new__(cls, *args, **kwargs):
         # Function for determining whether or not to display data
         # elements during `__str__`
         instance._display_data = display_data
+        # Function for determining whether or not to persist computed
+        # data
+        instance._persist_data = persist_data
         return instance
 
     def __init__(
@@ -2721,11 +2725,17 @@ def array(self):
         """A numpy array copy of the data.
 
         In-place changes to the returned numpy array do not affect the
-        underlying dask array.
+        underlying Dask array.
 
         The returned numpy array has the same mask hardness and fill
         values as the data.
 
+        .. note:: If the `{{package}}.persist_data` function returns
+                  True then calling `array` will persist the
+                  underlying lazy Dask array into an equivalent
+                  chunked Dask array, but now with the results fully
+                  computed and cached memory.
+
         Compare with `compute`.
 
         **Performance**
@@ -2734,7 +2744,8 @@ def array(self):
         returned `numpy` array is a deep copy of that returned by
         created `compute`.
 
-        .. seealso:: `datetime_array`, `compute`, `persist`
+        .. seealso:: `datetime_array`, `compute`, `persist`,
+                     `{{package}}.persist_data`
 
         **Examples**
 
@@ -4102,7 +4113,9 @@ def compressed(self, inplace=False):
         d._set_dask(dx, clear=self._ALL, in_memory=True)
         return d
 
-    def compute(self, _force_to_memory=True, _cache_elements=True):
+    def compute(
+        self, persist=None, _force_to_memory=True, _cache_elements=True
+    ):
         """A view of the computed data.
 
         In-place changes to the returned array *might* affect the
@@ -4120,11 +4133,30 @@ def compute(self, _force_to_memory=True, _cache_elements=True):
 
         .. versionadded:: (cfdm) 1.11.2.0
 
-        .. seealso:: `persist`, `array`, `datetime_array`,
-                     `sparse_array`
+        .. seealso:: `array`, `datetime_array`, `sparse_array`,
+                     `persist`, `{{package}}.persist_data`
 
         :Parameters:
 
+            persist: `None` or `bool`, optional
+                Control the persistence of computed data. Persisting
+                turns the underlying lazy dask array into an
+                equivalent chunked dask array, but now with the
+                results fully computed and cached memory. This can
+                avoid the expense of re-reading the data from disk, or
+                re-computing it, when the data is accessed on multiple
+                occasions.
+
+                If *persist* is `None` (the default) then the value of
+                *persist* will be taken from the
+                `{{package}}.persist_data` function. If *persist* is
+                True then the data is persisted, regardless of value
+                returned by `{{package}}.persist_data`. If *persist*
+                is False then the data is not persisted, regardless of
+                value returned by `{{package}}.persist_data`.
+
+                .. versionadded:: (cfdm) NEXTVERSION
+
             _force_to_memory: `bool`, optional
                 If True (the default) then force the data resulting
                 from computing the returned Dask graph to be in
@@ -4172,12 +4204,28 @@ def compute(self, _force_to_memory=True, _cache_elements=True):
          [0.029 0.059 0.039 0.07  0.058 0.072 0.009 0.017]
          [0.006 0.036 0.019 0.035 0.018 0.037 0.034 0.013]]
         >>> f.data.compute(_force_to_memory=False)
-        <{{repr}}NetCDF4Array(5, 8): file.nc, q(5, 8)>
+        <{{repr}}PyfiveArray(5, 8): file.nc, q(5, 8)>
 
         """
         dx = self.to_dask_array(
             _force_mask_hardness=False, _force_to_memory=_force_to_memory
         )
+
+        if persist is None:
+            persist = self._persist_data()
+
+        if persist:
+            dx = dx.persist()
+
+            # Note to developers: If the following `_set_dask` call is
+            #                     changed, consider making the same
+            #                     changes in `persist`.
+            self._set_dask(
+                dx,
+                clear=self._ALL ^ self._ARRAY ^ self._CACHE,
+                in_memory=True,
+            )
+
         a = dx.compute()
 
         if np.ma.isMA(a) and a is not np.ma.masked:
@@ -4189,11 +4237,16 @@ def compute(self, _force_to_memory=True, _cache_elements=True):
             a.set_fill_value(self.get_fill_value(None))
 
         if _cache_elements:
-            from scipy.sparse import issparse
+            ok = False
+            if isinstance(a, (np.ndarray, int, float, bool, str)):
+                ok = True
+            else:
+                from scipy.sparse import issparse
 
-            if isinstance(a, (np.ndarray, int, float, bool, str)) or issparse(
-                a
-            ):
+                if issparse(a):
+                    ok = True
+
+            if ok:
                 self.cache_elements(_array=a)
 
         return a
@@ -6584,6 +6637,10 @@ def persist(self, inplace=False):
             _force_mask_hardness=False, _force_to_memory=True
         )
         dx = dx.persist()
+
+        # Note to developers: If the following `_set_dask` call is
+        #                     changed, consider making the same
+        #                     changes in `compute`.
         d._set_dask(
             dx, clear=self._ALL ^ self._ARRAY ^ self._CACHE, in_memory=True
         )

cfdm/docstring/docstring.py

@@ -834,7 +834,7 @@
     # persist
     "{{persist description}}": """Persisting turns an underlying lazy dask array into an
         equivalent chunked dask array, but now with the results fully
-        computed and in memory. This can avoid the expense of
+        computed and cached in memory. This can avoid the expense of
         re-reading the data from disk, or re-computing it, when the
         data is accessed on multiple occasions.""",
     # ----------------------------------------------------------------