diff --git a/docs/src/whatsnew/latest.rst b/docs/src/whatsnew/latest.rst index 74d090a006..68737d86cc 100644 --- a/docs/src/whatsnew/latest.rst +++ b/docs/src/whatsnew/latest.rst @@ -60,8 +60,8 @@ This document explains the changes made to Iris for this release 🚀 Performance Enhancements =========================== -#. `@bouweandela`_ made loading :class:`~iris.cube.Cube`s from small NetCDF - files faster. (:pull:`6229`) +#. `@bouweandela`_ made loading :class:`~iris.cube.Cube`s from NetCDF files + faster. (:pull:`6229` and :pull:`6252`) #. `@fnattino`_ enabled lazy cube interpolation using the linear and nearest-neighbour interpolators (:class:`iris.analysis.Linear` and diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index a3dfa1edb4..7546aaea04 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -19,6 +19,9 @@ import numpy as np import numpy.ma as ma +MAX_CACHE_SIZE = 100 +"""Maximum number of Dask arrays to cache.""" + def non_lazy(func): """Turn a lazy function into a function that returns a result immediately.""" @@ -202,6 +205,7 @@ def _optimum_chunksize_internals( dim = working[0] working = working[1:] result.append(dim) + result = tuple(result) return result @@ -227,6 +231,33 @@ def _optimum_chunksize( ) +class LRUCache: + def __init__(self, maxsize: int) -> None: + self._cache: dict = {} + self.maxsize = maxsize + + def __getitem__(self, key): + value = self._cache.pop(key) + self._cache[key] = value + return value + + def __setitem__(self, key, value): + self._cache[key] = value + if len(self._cache) > self.maxsize: + self._cache.pop(next(iter(self._cache))) + + def __contains__(self, key): + return key in self._cache + + def __repr__(self): + return ( + f"<{self.__class__.__name__} maxsize={self.maxsize} cache={self._cache!r} >" + ) + + +CACHE = LRUCache(MAX_CACHE_SIZE) + + def as_lazy_data(data, chunks=None, asarray=False, meta=None, dims_fixed=None): """Convert the input array `data` to a :class:`dask.array.Array`. @@ -264,6 +295,8 @@ def as_lazy_data(data, chunks=None, asarray=False, meta=None, dims_fixed=None): but reduced by a factor if that exceeds the dask default chunksize. """ + from iris.fileformats.netcdf._thread_safe_nc import NetCDFDataProxy + if isinstance(data, ma.core.MaskedConstant): data = ma.masked_array(data.data, mask=data.mask) @@ -277,7 +310,7 @@ def as_lazy_data(data, chunks=None, asarray=False, meta=None, dims_fixed=None): if chunks is None: # No existing chunks : Make a chunk the shape of the entire input array # (but we will subdivide it if too big). - chunks = list(data.shape) + chunks = tuple(data.shape) # Adjust chunk size for better dask performance, # NOTE: but only if no shape dimension is zero, so that we can handle the @@ -291,9 +324,24 @@ def as_lazy_data(data, chunks=None, asarray=False, meta=None, dims_fixed=None): dims_fixed=dims_fixed, ) - if not is_lazy_data(data): - data = da.from_array(data, chunks=chunks, asarray=asarray, meta=meta) - return data + # Define a cache key for caching arrays created from NetCDFDataProxy objects. + # Creating new Dask arrays is relatively slow, therefore caching is beneficial + # if many cubes in the same file share coordinate arrays. + if isinstance(data, NetCDFDataProxy): + key = (repr(data), chunks, asarray, meta.dtype, type(meta)) + else: + key = None + + if is_lazy_data(data): + result = data + elif key in CACHE: + result = CACHE[key].copy() + else: + result = da.from_array(data, chunks=chunks, asarray=asarray, meta=meta) + if key is not None: + CACHE[key] = result.copy() + + return result def _co_realise_lazy_arrays(arrays): diff --git a/lib/iris/tests/unit/lazy_data/test_as_lazy_data.py b/lib/iris/tests/unit/lazy_data/test_as_lazy_data.py index 821370ce6c..a0c13208ee 100644 --- a/lib/iris/tests/unit/lazy_data/test_as_lazy_data.py +++ b/lib/iris/tests/unit/lazy_data/test_as_lazy_data.py @@ -161,7 +161,7 @@ def test_default_chunks_limiting(self, mocker): as_lazy_data(data) assert limitcall_patch.call_args_list == [ mock.call( - list(test_shape), + tuple(test_shape), shape=test_shape, dtype=np.dtype("f4"), dims_fixed=None, diff --git a/lib/iris/tests/unit/lazy_data/test_lrucache.py b/lib/iris/tests/unit/lazy_data/test_lrucache.py new file mode 100644 index 0000000000..2a8abb0cde --- /dev/null +++ b/lib/iris/tests/unit/lazy_data/test_lrucache.py @@ -0,0 +1,25 @@ +# Copyright Iris contributors +# +# This file is part of Iris and is released under the BSD license. +# See LICENSE in the root of the repository for full licensing details. +"""Test function :func:`iris._lazy data.LRUCache`.""" + +from iris._lazy_data import LRUCache + + +def test_lrucache(): + cache = LRUCache(2) + + cache["a"] = 1 + + assert "a" in cache + assert cache["a"] == 1 + + cache["b"] = 2 + cache["c"] = 3 + + assert "a" not in cache + assert "b" in cache + assert "c" in cache + + assert str(cache) == ""