From 198f6caf2b5b9862e1aed4894fc6dc61f017cb3c Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Mon, 29 Sep 2025 22:08:49 +0100 Subject: [PATCH 1/3] improve memory for map_complete_blocks --- lib/iris/_lazy_data.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index 992789ed90..172c52f978 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -640,17 +640,39 @@ def map_complete_blocks(src, func, dims, out_sizes, dtype, *args, **kwargs): else: data = src.lazy_data() + shape = list(src.shape) + if result is None and data is not None: # Ensure dims are not chunked in_chunks = list(data.chunks) for dim in dims: - in_chunks[dim] = src.shape[dim] - data = data.rechunk(in_chunks) + in_chunks[dim] = (src.shape[dim],) # Determine output chunks - out_chunks = list(data.chunks) + out_chunks = in_chunks.copy() for dim, size in zip(dims, out_sizes): - out_chunks[dim] = size + out_chunks[dim] = (size,) + shape[dim] = size + + max_outchunks = [max(chunk) for chunk in out_chunks] + df = [False]*len(max_outchunks) + for dim in dims: + df[dim] = True + df = tuple(df) + opt_outchunks = _optimum_chunksize(max_outchunks, shape, dtype=dtype, dims_fixed=df) + for i, (chunk, max_out, opt_out) in enumerate(zip(out_chunks,max_outchunks,opt_outchunks)): + if opt_out < max_out: + new_chunks = [] + for c in chunk: + new_chunks.extend((c//opt_out)*[opt_out]) + if chunk_end := c % opt_out > 0: + new_chunks.append(chunk_end) + in_chunks[i] = tuple(new_chunks) + out_chunks[i] = tuple(new_chunks) + + + + data = data.rechunk(in_chunks) # Assume operation preserves mask. meta = da.utils.meta_from_array(data).astype(dtype) From acd68d62181010b2d206685a6ed5aeb3c6bc68d9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 29 Sep 2025 21:12:01 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- lib/iris/_lazy_data.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index 172c52f978..b12c8699ff 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -655,23 +655,25 @@ def map_complete_blocks(src, func, dims, out_sizes, dtype, *args, **kwargs): shape[dim] = size max_outchunks = [max(chunk) for chunk in out_chunks] - df = [False]*len(max_outchunks) + df = [False] * len(max_outchunks) for dim in dims: df[dim] = True df = tuple(df) - opt_outchunks = _optimum_chunksize(max_outchunks, shape, dtype=dtype, dims_fixed=df) - for i, (chunk, max_out, opt_out) in enumerate(zip(out_chunks,max_outchunks,opt_outchunks)): + opt_outchunks = _optimum_chunksize( + max_outchunks, shape, dtype=dtype, dims_fixed=df + ) + for i, (chunk, max_out, opt_out) in enumerate( + zip(out_chunks, max_outchunks, opt_outchunks) + ): if opt_out < max_out: new_chunks = [] for c in chunk: - new_chunks.extend((c//opt_out)*[opt_out]) + new_chunks.extend((c // opt_out) * [opt_out]) if chunk_end := c % opt_out > 0: new_chunks.append(chunk_end) in_chunks[i] = tuple(new_chunks) out_chunks[i] = tuple(new_chunks) - - data = data.rechunk(in_chunks) # Assume operation preserves mask. From bf4548ad774f5e22fd822a978dd597e5cfd7d160 Mon Sep 17 00:00:00 2001 From: "stephen.worsley" Date: Tue, 7 Oct 2025 09:38:16 +0100 Subject: [PATCH 3/3] fix test failures --- lib/iris/_lazy_data.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/iris/_lazy_data.py b/lib/iris/_lazy_data.py index 172c52f978..89d5eb9689 100644 --- a/lib/iris/_lazy_data.py +++ b/lib/iris/_lazy_data.py @@ -659,6 +659,7 @@ def map_complete_blocks(src, func, dims, out_sizes, dtype, *args, **kwargs): for dim in dims: df[dim] = True df = tuple(df) + dtype = np.dtype(dtype) opt_outchunks = _optimum_chunksize(max_outchunks, shape, dtype=dtype, dims_fixed=df) for i, (chunk, max_out, opt_out) in enumerate(zip(out_chunks,max_outchunks,opt_outchunks)): if opt_out < max_out: