Replicate legacy dask behaviour pre dask/dask#9555.

trexfeathers · trexfeathers · commit 7363412a0693 · 2022-10-31T15:06:24.000Z
diff --git a/lib/iris/_data_manager.py b/lib/iris/_data_manager.py
@@ -13,7 +13,12 @@
 import numpy as np
 import numpy.ma as ma
 
-from iris._lazy_data import as_concrete_data, as_lazy_data, is_lazy_data
+from iris._lazy_data import (
+    as_concrete_data,
+    as_lazy_data,
+    copy_lazy_data,
+    is_lazy_data,
+)
 
 
 class DataManager:
@@ -174,7 +179,7 @@ def _deepcopy(self, memo, data=None):
             if data is None:
                 # Copy the managed data.
                 if self.has_lazy_data():
-                    data = copy.deepcopy(self._lazy_array, memo)
+                    data = copy_lazy_data(self._lazy_array)
                 else:
                     data = self._real_array.copy()
             else:
diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py
@@ -417,3 +417,27 @@ def map_complete_blocks(src, func, dims, out_sizes):
         out_chunks[dim] = size
 
     return data.map_blocks(func, chunks=out_chunks, dtype=src.dtype)
+
+
+def copy_lazy_data(data: da.Array) -> da.Array:
+    """
+    Copy a lazy data array, ensuring compute() will return a different array to the original.
+
+    Replicates legacy Dask behaviour removed in dask#9555.
+
+    Parameters
+    ----------
+    data : Dask Array
+        The lazy data array to be copied.
+
+    Returns
+    -------
+    Dask Array
+        A copy of `data`.
+
+    """
+    if data.npartitions == 1:
+        result = data.map_blocks(dask.utils.M.copy)
+    else:
+        result = data.copy()
+    return result
