Keep lazy_elementwise blocks as arrays for scalar-returning ops (#6965)

docs/src/whatsnew/latest.rst

@@ -36,7 +36,11 @@ This document explains the changes made to Iris for this release
 🐛 Bugs Fixed
 =============
 
-#. N/A
+#. :user:`gaoflow` fixed an error when computing (e.g. saving) a scalar lazy
+   cube whose units had been converted with
+   :meth:`~iris.cube.Cube.convert_units`. The unit conversion could yield a
+   plain Python scalar for the 0-dimensional block, which Dask was then unable
+   to store. (:issue:`6965`)
 
 
 💣 Incompatible Changes

lib/iris/_lazy_data.py

@@ -592,7 +592,13 @@ def lazy_elementwise(lazy_array, elementwise_op):
     dtype = elementwise_op(np.zeros(1, lazy_array.dtype)).dtype
     meta = da.utils.meta_from_array(lazy_array).astype(dtype)
 
-    return da.map_blocks(elementwise_op, lazy_array, dtype=dtype, meta=meta)
+    def wrapped_op(block):
+        # Some operations return a Python scalar for a 0-dimensional block
+        # (e.g. cf_units.Unit.convert on a scalar array), which Dask cannot
+        # store. Ensure each block remains an array. See #6965.
+        return np.asanyarray(elementwise_op(block))
+
+    return da.map_blocks(wrapped_op, lazy_array, dtype=dtype, meta=meta)
 
 
 def map_complete_blocks(src, func, dims, out_sizes, dtype, *args, **kwargs):

lib/iris/tests/unit/lazy_data/test_lazy_elementwise.py

@@ -38,3 +38,20 @@ def test_dtype_change(self):
         assert is_lazy_data(wrapped)
         assert wrapped.dtype == np.int_
         assert wrapped.compute().dtype == wrapped.dtype
+
+    def test_scalar_returning_op_on_0d(self):
+        # An op that returns a Python scalar for a 0-d block (e.g.
+        # cf_units.Unit.convert on a scalar array) must still yield an array,
+        # so that the result can be stored by Dask (#6965).
+        def scalar_op(array):
+            # Mimics returning a Python float for a scalar input.
+            return float(array) if array.ndim == 0 else array + 1.0
+
+        lazy_array = as_lazy_data(np.array(3.0))
+        assert lazy_array.ndim == 0
+        wrapped = lazy_elementwise(lazy_array, scalar_op)
+        assert is_lazy_data(wrapped)
+        result = wrapped.compute()
+        assert isinstance(result, np.ndarray)
+        assert result.shape == ()
+        assert result[()] == 3.0