diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 12720cecae..99fc16f626 100644 --- a/docs/src/whatsnew/latest.rst +++ b/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 ======================= diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index 2de7f8c5ac..26efb106fc 100644 --- a/lib/iris/_lazy_data.py +++ b/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): diff --git a/lib/iris/tests/unit/lazy_data/test_lazy_elementwise.py b/lib/iris/tests/unit/lazy_data/test_lazy_elementwise.py index 1600067d79..629ffa34c3 100644 --- a/lib/iris/tests/unit/lazy_data/test_lazy_elementwise.py +++ b/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