diff --git a/esmvalcore/preprocessor/_time.py b/esmvalcore/preprocessor/_time.py index c6e3027414..635861b86c 100644 --- a/esmvalcore/preprocessor/_time.py +++ b/esmvalcore/preprocessor/_time.py @@ -412,7 +412,12 @@ def get_time_weights(cube: Cube) -> np.ndarray | da.core.Array: # Extract 1D time weights (= lengths of time intervals) time_weights = time.lazy_bounds()[:, 1] - time.lazy_bounds()[:, 0] - if not cube.has_lazy_data(): + if cube.has_lazy_data(): + # Align the weight chunks with the data chunks to avoid excessively + # large chunks as a result of broadcasting. + time_chunks = cube.lazy_data().chunks[coord_dims[0]] + time_weights = time_weights.rechunk(time_chunks) + else: time_weights = time_weights.compute() return time_weights diff --git a/tests/unit/preprocessor/_time/test_time.py b/tests/unit/preprocessor/_time/test_time.py index 6cb05a366d..a47da2a724 100644 --- a/tests/unit/preprocessor/_time/test_time.py +++ b/tests/unit/preprocessor/_time/test_time.py @@ -1777,10 +1777,11 @@ def test_get_time_weights(): def test_get_time_weights_lazy(): """Test ``get_time_weights`` for complex cube with lazy data.""" cube = _make_cube() - cube.data = cube.lazy_data() + cube.data = cube.lazy_data().rechunk((1, 1, 1, 3)) weights = get_time_weights(cube) assert isinstance(weights, da.Array) assert weights.shape == (2, ) + assert weights.chunks == ((1, 1), ) np.testing.assert_allclose(weights, [15.0, 30.0])