From f9a50c05945594aeaa7a8ce9b528a8401216a371 Mon Sep 17 00:00:00 2001 From: FBumann <117816358+FBumann@users.noreply.github.com> Date: Thu, 16 Jul 2026 12:17:06 +0200 Subject: [PATCH] perf: make fill_missing_coords copy lazily (skip the no-op deep copy) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fill_missing_coords did `ds = ds.copy()` unconditionally, then filled a coord only for dims that lack one. When every dim already has a coord — the overwhelmingly common case, including every `var * array` op — it returned a byte-identical deep copy that is immediately discarded. Because `as_dataarray` ends in `fill_missing_coords`, the §8 check `coeff_da = as_dataarray(coefficient)` in Variable.to_linexpr copied the whole coefficient and held it live through reindex+fillna, inflating build peak on the exact-match path (the remaining half of the test_op[var_mul_array_match] regression after #838). Compute the missing dims first; copy (and mutate) only when there is something to fill, else return `ds` untouched. The copy still guards the one mutation it ever performs. memray, legacy default, GRID 3x4x1000: var_mul_array_match 217 -> 123 KiB, var_mul_array_bcast 998 -> 593 KiB. Full suite 7840 passed under both semantics. Co-Authored-By: Claude Opus 4.8 (1M context) --- linopy/alignment.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/linopy/alignment.py b/linopy/alignment.py index 76f566be..93eb0b60 100644 --- a/linopy/alignment.py +++ b/linopy/alignment.py @@ -354,16 +354,17 @@ def fill_missing_coords( Whether to fill in integer coordinates for helper dimensions, by default False. """ - ds = ds.copy() if not isinstance(ds, Dataset | DataArray): raise TypeError(f"Expected xarray.DataArray or xarray.Dataset, got {type(ds)}.") skip_dims = [] if fill_helper_dims else HELPER_DIMS + missing = [d for d in ds.dims if d not in ds.coords and d not in skip_dims] + if not missing: + return ds - # Fill in missing integer coordinates - for dim in ds.dims: - if dim not in ds.coords and dim not in skip_dims: - ds.coords[dim] = arange(ds.sizes[dim]) + ds = ds.copy() + for dim in missing: + ds.coords[dim] = arange(ds.sizes[dim]) return ds