Skip to content

Commit

Permalink
pandas 3 MultiIndex fixes (pydata#8847)
Browse files Browse the repository at this point in the history
* Fix dropping of muiltiindexes

xref pydata#8844
Closes xarray-contrib/flox#342

* More fixes
  • Loading branch information
dcherian authored Mar 20, 2024
1 parent 0b6716d commit 4a0bb2e
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 6 deletions.
2 changes: 1 addition & 1 deletion xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -5867,7 +5867,7 @@ def drop_vars(
for var in names_set:
maybe_midx = self._indexes.get(var, None)
if isinstance(maybe_midx, PandasMultiIndex):
idx_coord_names = set(maybe_midx.index.names + [maybe_midx.dim])
idx_coord_names = set(list(maybe_midx.index.names) + [maybe_midx.dim])
idx_other_names = idx_coord_names - set(names_set)
other_names.update(idx_other_names)
if other_names:
Expand Down
11 changes: 6 additions & 5 deletions xarray/tests/test_indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -352,7 +352,7 @@ def test_constructor(self) -> None:
# default level names
pd_idx = pd.MultiIndex.from_arrays([foo_data, bar_data])
index = PandasMultiIndex(pd_idx, "x")
assert index.index.names == ("x_level_0", "x_level_1")
assert list(index.index.names) == ["x_level_0", "x_level_1"]

def test_from_variables(self) -> None:
v_level1 = xr.Variable(
Expand All @@ -370,7 +370,7 @@ def test_from_variables(self) -> None:
assert index.dim == "x"
assert index.index.equals(expected_idx)
assert index.index.name == "x"
assert index.index.names == ["level1", "level2"]
assert list(index.index.names) == ["level1", "level2"]

var = xr.Variable(("x", "y"), [[1, 2, 3], [4, 5, 6]])
with pytest.raises(
Expand Down Expand Up @@ -413,7 +413,8 @@ def test_stack(self) -> None:
index = PandasMultiIndex.stack(prod_vars, "z")

assert index.dim == "z"
assert index.index.names == ["x", "y"]
# TODO: change to tuple when pandas 3 is minimum
assert list(index.index.names) == ["x", "y"]
np.testing.assert_array_equal(
index.index.codes, [[0, 0, 0, 1, 1, 1], [0, 1, 2, 0, 1, 2]]
)
Expand Down Expand Up @@ -531,12 +532,12 @@ def test_rename(self) -> None:
assert new_index is index

new_index = index.rename({"two": "three"}, {})
assert new_index.index.names == ["one", "three"]
assert list(new_index.index.names) == ["one", "three"]
assert new_index.dim == "x"
assert new_index.level_coords_dtype == {"one": "<U1", "three": np.int32}

new_index = index.rename({}, {"x": "y"})
assert new_index.index.names == ["one", "two"]
assert list(new_index.index.names) == ["one", "two"]
assert new_index.dim == "y"
assert new_index.level_coords_dtype == level_coords_dtype

Expand Down

0 comments on commit 4a0bb2e

Please sign in to comment.