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

docs/src/whatsnew/latest.rst

@@ -56,6 +56,12 @@ This document explains the changes made to Iris for this release
    always promoting the result to ``float64``. Integer inputs are still returned
    as ``float64``. (:issue:`4119`)
 
+#. :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

@@ -8,7 +8,7 @@
 
 """
 
-from functools import lru_cache, wraps
+from functools import lru_cache, partial, wraps
 from types import ModuleType
 from typing import Sequence
 
@@ -560,6 +560,13 @@ def co_realise_cubes(*cubes):
         cube.data = result
 
 
+def _as_array_wrapper(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(op(block))
+
+
 def lazy_elementwise(lazy_array, elementwise_op):
     """Apply a (numpy-style) elementwise array operation to a lazy array.
 
@@ -592,7 +599,12 @@ 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)
+    return da.map_blocks(
+        partial(_as_array_wrapper, elementwise_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