From 8df355ef1c6948ac21fd2fc735be36cd39492652 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Mon, 15 Dec 2025 21:39:26 -0800 Subject: [PATCH 1/6] Refactor Dataset.map to merge attrs instead of copying Update the `keep_attrs` behavior in `Dataset.map()` and `DataTree.map()` to merge attributes from the original and function results using the `drop_conflicts` strategy, rather than unconditionally copying original attrs. When `keep_attrs=True`, matching attrs are kept and conflicting attrs are dropped. When `keep_attrs=False`, only attrs set by the function are retained. Add comprehensive tests for the new attr merging behavior. --- xarray/core/dataset.py | 23 ++++++++++++++--------- xarray/core/datatree.py | 13 ++++++++++--- xarray/tests/test_dataset.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 53 insertions(+), 12 deletions(-) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index bce048048da..9a7b88263ef 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -6910,8 +6910,10 @@ def map( DataArray. keep_attrs : bool or None, optional If True, both the dataset's and variables' attributes (`attrs`) will be - copied from the original objects to the new ones. If False, the new dataset - and variables will be returned without copying the attributes. + combined from the original objects and the function results using the + ``drop_conflicts`` strategy: matching attrs are kept, conflicting attrs + are dropped. If False, the new dataset and variables will have only + the attributes set by the function. args : iterable, optional Positional arguments passed on to `func`. **kwargs : Any @@ -6960,16 +6962,19 @@ def map( coords = Coordinates._construct_direct(coords=coord_vars, indexes=indexes) if keep_attrs: + # Merge attrs from function result and original, dropping conflicts + from xarray.structure.merge import merge_attrs + for k, v in variables.items(): - v._copy_attrs_from(self.data_vars[k]) + v.attrs = merge_attrs( + [v.attrs, self.data_vars[k].attrs], "drop_conflicts" + ) for k, v in coords.items(): if k in self.coords: - v._copy_attrs_from(self.coords[k]) - else: - for v in variables.values(): - v.attrs = {} - for v in coords.values(): - v.attrs = {} + v.attrs = merge_attrs( + [v.attrs, self.coords[k].attrs], "drop_conflicts" + ) + # When keep_attrs=False, leave attrs as the function returned them attrs = self.attrs if keep_attrs else None return type(self)(variables, coords=coords, attrs=attrs) diff --git a/xarray/core/datatree.py b/xarray/core/datatree.py index e079332780c..a64ceefb207 100644 --- a/xarray/core/datatree.py +++ b/xarray/core/datatree.py @@ -397,8 +397,10 @@ def map( # type: ignore[override] DataArray. keep_attrs : bool | None, optional If True, both the dataset's and variables' attributes (`attrs`) will be - copied from the original objects to the new ones. If False, the new dataset - and variables will be returned without copying the attributes. + combined from the original objects and the function results using the + ``drop_conflicts`` strategy: matching attrs are kept, conflicting attrs + are dropped. If False, the new dataset and variables will have only + the attributes set by the function. args : iterable, optional Positional arguments passed on to `func`. **kwargs : Any @@ -438,8 +440,13 @@ def map( # type: ignore[override] for k, v in self.data_vars.items() } if keep_attrs: + # Merge attrs from function result and original, dropping conflicts + from xarray.structure.merge import merge_attrs + for k, v in variables.items(): - v._copy_attrs_from(self.data_vars[k]) + v.attrs = merge_attrs( + [v.attrs, self.data_vars[k].attrs], "drop_conflicts" + ) attrs = self.attrs if keep_attrs else None # return type(self)(variables, attrs=attrs) return Dataset(variables, attrs=attrs) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 6dce32aeb5c..83ce11269c5 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6452,6 +6452,35 @@ def mixed_func(x): expected = xr.Dataset({"foo": 42, "bar": ("y", [4, 5])}) assert_identical(result, expected) + def test_map_preserves_function_attrs(self) -> None: + # Regression test for GH11019 + # Attrs added by function should be preserved in result + ds = xr.Dataset({"test": ("x", [1, 2, 3], {"original": "value"})}) + + def add_attr(da): + return da.assign_attrs(new_attr="foobar") + + # With keep_attrs=True: merge using drop_conflicts (no conflict here) + result = ds.map(add_attr, keep_attrs=True) + assert result["test"].attrs == {"original": "value", "new_attr": "foobar"} + + # With keep_attrs=False: function's attrs preserved + result = ds.map(add_attr, keep_attrs=False) + assert result["test"].attrs == {"original": "value", "new_attr": "foobar"} + + # When function modifies existing attr with keep_attrs=True, conflict is dropped + def modify_attr(da): + return da.assign_attrs(original="modified", extra="added") + + result = ds.map(modify_attr, keep_attrs=True) + assert result["test"].attrs == { + "extra": "added" + } # "original" dropped due to conflict + + # When function modifies existing attr with keep_attrs=False, function wins + result = ds.map(modify_attr, keep_attrs=False) + assert result["test"].attrs == {"original": "modified", "extra": "added"} + def test_apply_pending_deprecated_map(self) -> None: data = create_test_data() data.attrs["foo"] = "bar" From b558513131d05e3f5fe0a9b627711d694ab79927 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Mon, 15 Dec 2025 22:19:22 -0800 Subject: [PATCH 2/6] Fix weighted operations to respect keep_attrs=False MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Weighted operations internally propagate attrs from weights through computations like dot(). When keep_attrs=False is passed, users expect no attrs on the result, but attrs from weights were leaking through. Clear attrs explicitly in _implementation when keep_attrs is False. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- xarray/computation/weighted.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index b311290aabf..12d61cedc7d 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -544,13 +544,29 @@ def _implementation(self, func, dim, **kwargs) -> DataArray: dataset = self.obj._to_temp_dataset() dataset = dataset.map(func, dim=dim, **kwargs) - return self.obj._from_temp_dataset(dataset) + result = self.obj._from_temp_dataset(dataset) + + # Clear attrs when keep_attrs is explicitly False + # (weighted operations can propagate attrs from weights through internal computations) + if kwargs.get("keep_attrs") is False: + result.attrs = {} + + return result class DatasetWeighted(Weighted["Dataset"]): def _implementation(self, func, dim, **kwargs) -> Dataset: self._check_dim(dim) - return self.obj.map(func, dim=dim, **kwargs) + result = self.obj.map(func, dim=dim, **kwargs) + + # Clear attrs when keep_attrs is explicitly False + # (weighted operations can propagate attrs from weights through internal computations) + if kwargs.get("keep_attrs") is False: + result.attrs = {} + for var in result.data_vars.values(): + var.attrs = {} + + return result def _inject_docstring(cls, cls_name): From d243e16d97a4952cbfa9da5cc5d6ee321d288473 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Mon, 15 Dec 2025 23:01:35 -0800 Subject: [PATCH 3/6] Fix RTD builds by restricting test-nightly to linux/macOS platforms The `test-nightly` environment uses pandas nightly wheels from PyPI, which currently don't have win-64 builds available. This causes `pixi lock` to fail when solving for all platforms. RTD builds fail because they have no lock file cache (unlike GitHub Actions CI which caches pixi.lock). When RTD runs `pixi install -e doc`, pixi must generate the lock file from scratch, which fails on the unsolvable test-nightly/win-64 combination. This restriction can be removed once pandas nightly provides win-64 wheels again. Co-authored-by: Claude --- pixi.toml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pixi.toml b/pixi.toml index bfa51c3fb54..443c8c8f211 100644 --- a/pixi.toml +++ b/pixi.toml @@ -143,6 +143,12 @@ sparse = "0.15.*" toolz = "0.12.*" zarr = "2.18.*" +# TODO: Remove `platforms` restriction once pandas nightly has win-64 wheels again. +# Without this, `pixi lock` fails because it can't solve the nightly feature for win-64, +# which breaks RTD builds (RTD has no lock file cache, unlike GitHub Actions CI). +[feature.nightly] +platforms = ["linux-64", "osx-arm64"] + [feature.nightly.dependencies] python = "*" From f9a56696b267569c6d14c9a895e4604200985761 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Tue, 16 Dec 2025 00:00:34 -0800 Subject: [PATCH 4/6] Add whats-new entry for Dataset.map attrs behavior MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- doc/whats-new.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 7e3badc7143..00c88f19239 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -26,6 +26,10 @@ Deprecations Bug Fixes ~~~~~~~~~ +- :py:meth:`Dataset.map` now merges attrs from the function result and the original + using the ``drop_conflicts`` strategy when ``keep_attrs=True``, preserving attrs + set by the function (:issue:`11019`, :pull:`11020`). + By `Maximilian Roos `_. - Ensure that ``keep_attrs='drop'`` and ``keep_attrs=False`` remove attrs from result, even when there is only one xarray object given to ``apply_ufunc`` (:issue:`10982` :pull:`10997`). By `Julia Signell `_. From 966b49fcd67d5baaaae433213a6ecd063c5416b0 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Fri, 16 Jan 2026 11:36:43 -0800 Subject: [PATCH 5/6] Clear coord attrs in weighted operations when keep_attrs=False Address PR feedback to also clear coordinate attrs (not just data_vars attrs) when keep_attrs=False in both DataArrayWeighted and DatasetWeighted. Added test to verify coord attrs are cleared for both DataArray and Dataset. Co-authored-by: Claude --- xarray/computation/weighted.py | 4 ++++ xarray/tests/test_weighted.py | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/xarray/computation/weighted.py b/xarray/computation/weighted.py index 12d61cedc7d..d19cc4fea90 100644 --- a/xarray/computation/weighted.py +++ b/xarray/computation/weighted.py @@ -550,6 +550,8 @@ def _implementation(self, func, dim, **kwargs) -> DataArray: # (weighted operations can propagate attrs from weights through internal computations) if kwargs.get("keep_attrs") is False: result.attrs = {} + for var in result.coords.values(): + var.attrs = {} return result @@ -565,6 +567,8 @@ def _implementation(self, func, dim, **kwargs) -> Dataset: result.attrs = {} for var in result.data_vars.values(): var.attrs = {} + for var in result.coords.values(): + var.attrs = {} return result diff --git a/xarray/tests/test_weighted.py b/xarray/tests/test_weighted.py index 5e913c00629..3cdec727cdd 100644 --- a/xarray/tests/test_weighted.py +++ b/xarray/tests/test_weighted.py @@ -786,6 +786,25 @@ def test_weighted_mean_keep_attrs_ds(): assert data.coords["dim_1"].attrs == result.coords["dim_1"].attrs +@pytest.mark.parametrize("as_dataset", (True, False)) +def test_weighted_operations_drop_coord_attrs(as_dataset): + # Test that coord attrs are cleared when keep_attrs=False + weights = DataArray(np.random.randn(2)) + data = Dataset( + {"a": (["dim_0", "dim_1"], np.random.randn(2, 2), {"attr": "data"})}, + coords={"dim_1": ("dim_1", ["a", "b"], {"coord_attr": "value"})}, + ) + + if not as_dataset: + data = data["a"] + + result = data.weighted(weights).mean(dim="dim_0", keep_attrs=False) + + # All attrs should be cleared + assert result.attrs == {} + assert result.coords["dim_1"].attrs == {} + + @pytest.mark.parametrize("operation", ("sum_of_weights", "sum", "mean", "quantile")) @pytest.mark.parametrize("as_dataset", (True, False)) def test_weighted_bad_dim(operation, as_dataset): From ec5d6fc041b258e82253d3582d61c5cc796ad019 Mon Sep 17 00:00:00 2001 From: Maximilian Roos Date: Fri, 16 Jan 2026 14:37:50 -0800 Subject: [PATCH 6/6] Fix mypy type error in test_weighted_operations_drop_coord_attrs Use proper type annotation with DataArray | Dataset union type to avoid incompatible assignment error. Co-authored-by: Claude --- xarray/tests/test_weighted.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/xarray/tests/test_weighted.py b/xarray/tests/test_weighted.py index 3cdec727cdd..3d860b5b17a 100644 --- a/xarray/tests/test_weighted.py +++ b/xarray/tests/test_weighted.py @@ -790,13 +790,12 @@ def test_weighted_mean_keep_attrs_ds(): def test_weighted_operations_drop_coord_attrs(as_dataset): # Test that coord attrs are cleared when keep_attrs=False weights = DataArray(np.random.randn(2)) - data = Dataset( + ds = Dataset( {"a": (["dim_0", "dim_1"], np.random.randn(2, 2), {"attr": "data"})}, coords={"dim_1": ("dim_1", ["a", "b"], {"coord_attr": "value"})}, ) - if not as_dataset: - data = data["a"] + data: DataArray | Dataset = ds if as_dataset else ds["a"] result = data.weighted(weights).mean(dim="dim_0", keep_attrs=False)