From c39b82a0d33ce351c71c57cdd04bff942399c714 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Fri, 29 Mar 2024 10:46:08 +0100 Subject: [PATCH 1/3] to_base_variable: convert multiindex -> np.ndarray --- xarray/core/variable.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/xarray/core/variable.py b/xarray/core/variable.py index ec284e411fc..717cfc623f3 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -22,6 +22,7 @@ BasicIndexer, OuterIndexer, PandasIndexingAdapter, + PandasMultiIndexingAdapter, VectorizedIndexer, as_indexable, ) @@ -530,8 +531,13 @@ def values(self, values): def to_base_variable(self) -> Variable: """Return this variable as a base xarray.Variable""" + if isinstance(self._data, PandasMultiIndexingAdapter): + data = np.asarray(self._data) + else: + data = self._data + return Variable( - self._dims, self._data, self._attrs, encoding=self._encoding, fastpath=True + self._dims, data, self._attrs, encoding=self._encoding, fastpath=True ) to_variable = utils.alias(to_base_variable, "to_variable") From 0fb6c8f10298c2f9dfdee16e85169f736e4080a7 Mon Sep 17 00:00:00 2001 From: Benoit Bovy Date: Fri, 29 Mar 2024 10:48:50 +0100 Subject: [PATCH 2/3] add reset_index test Ensure that reset index variable doesn't wrap any multiindex anymore. --- xarray/tests/test_dataset.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index e2a64964775..9cb1d88ad60 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -3560,6 +3560,9 @@ def test_reset_index_drop_convert( assert name not in reset.variables for name in converted: assert_identical(reset[name].variable, ds[name].variable.to_base_variable()) + # check that variable data is converted back to a numpy array + # (https://github.com/pydata/xarray/issues/8887) + assert isinstance(reset[name].variable._data, np.ndarray) for old_name, new_name in renamed.items(): assert_identical(ds[old_name].variable, reset[new_name].variable) From dd9c3b4ad88b6694b6e737e86e80ad1dcfa1527c Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Fri, 29 Mar 2024 08:46:27 -0600 Subject: [PATCH 3/3] Add consistency check --- xarray/tests/test_dataset.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 9cb1d88ad60..d01c0f604f4 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -3514,6 +3514,20 @@ def test_reset_index(self) -> None: with pytest.raises(ValueError, match=r".*not coordinates with an index"): ds.reset_index("y") + def test_reset_multi_index_resets_levels(self) -> None: + # ND DataArray that gets stacked along a multiindex + da = DataArray(np.ones((3, 3)), coords={"dim1": [1, 2, 3], "dim2": [4, 5, 6]}) + da = da.stack(feature=["dim1", "dim2"]) + + # Extract just the stacked coordinates for saving in a dataset + ds = Dataset(data_vars={"feature": da.feature}) + xr.testing.assertions._assert_internal_invariants( + ds.reset_index(["feature", "dim1", "dim2"]), check_default_indexes=False + ) # succeeds + xr.testing.assertions._assert_internal_invariants( + ds.reset_index(["feature"]), check_default_indexes=False + ) # fails, but no warning either + def test_reset_index_keep_attrs(self) -> None: coord_1 = DataArray([1, 2], dims=["coord_1"], attrs={"attrs": True}) ds = Dataset({}, {"coord_1": coord_1})