Skip to content

Commit 2475d49

Browse files
authored
Restrict fastpath isel indexes to the case of all PandasIndex (#10066)
1 parent df33a9a commit 2475d49

File tree

3 files changed

+41
-4
lines changed

3 files changed

+41
-4
lines changed

Diff for: doc/whats-new.rst

+2
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ Bug fixes
5555
- Fix DataArray().drop_attrs(deep=False) and add support for attrs to
5656
DataArray()._replace(). (:issue:`10027`, :pull:`10030`). By `Jan
5757
Haacker <https://github.com/j-haacker>`_.
58+
- Fix ``isel`` for multi-coordinate Xarray indexes (:issue:`10063`, :pull:`10066`).
59+
By `Benoit Bovy <https://github.com/benbovy>`_.
5860

5961

6062
Documentation

Diff for: xarray/core/indexes.py

+5-4
Original file line numberDiff line numberDiff line change
@@ -1991,10 +1991,11 @@ def isel_indexes(
19911991
indexes: Indexes[Index],
19921992
indexers: Mapping[Any, Any],
19931993
) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]:
1994-
# TODO: remove if clause in the future. It should be unnecessary.
1995-
# See failure introduced when removed
1996-
# https://github.com/pydata/xarray/pull/9002#discussion_r1590443756
1997-
if any(isinstance(v, PandasMultiIndex) for v in indexes._indexes.values()):
1994+
# Fast path function _apply_indexes_fast does not work with multi-coordinate
1995+
# Xarray indexes (see https://github.com/pydata/xarray/issues/10063).
1996+
# -> call it only in the most common case where all indexes are default
1997+
# PandasIndex each associated to a single 1-dimensional coordinate.
1998+
if any(type(idx) is not PandasIndex for idx in indexes._indexes.values()):
19981999
return _apply_indexes(indexes, indexers, "isel")
19992000
else:
20002001
return _apply_indexes_fast(indexes, indexers, "isel")

Diff for: xarray/tests/test_dataset.py

+34
Original file line numberDiff line numberDiff line change
@@ -1593,6 +1593,40 @@ def test_isel_fancy_convert_index_variable(self) -> None:
15931593
assert "x" not in actual.xindexes
15941594
assert not isinstance(actual.x.variable, IndexVariable)
15951595

1596+
def test_isel_multicoord_index(self) -> None:
1597+
# regression test https://github.com/pydata/xarray/issues/10063
1598+
# isel on a multi-coordinate index should return a unique index associated
1599+
# to each coordinate
1600+
class MultiCoordIndex(xr.Index):
1601+
def __init__(self, idx1, idx2):
1602+
self.idx1 = idx1
1603+
self.idx2 = idx2
1604+
1605+
@classmethod
1606+
def from_variables(cls, variables, *, options=None):
1607+
idx1 = PandasIndex.from_variables(
1608+
{"x": variables["x"]}, options=options
1609+
)
1610+
idx2 = PandasIndex.from_variables(
1611+
{"y": variables["y"]}, options=options
1612+
)
1613+
1614+
return cls(idx1, idx2)
1615+
1616+
def create_variables(self, variables=None):
1617+
return {**self.idx1.create_variables(), **self.idx2.create_variables()}
1618+
1619+
def isel(self, indexers):
1620+
idx1 = self.idx1.isel({"x": indexers.get("x", slice(None))})
1621+
idx2 = self.idx2.isel({"y": indexers.get("y", slice(None))})
1622+
return MultiCoordIndex(idx1, idx2)
1623+
1624+
coords = xr.Coordinates(coords={"x": [0, 1], "y": [1, 2]}, indexes={})
1625+
ds = xr.Dataset(coords=coords).set_xindex(["x", "y"], MultiCoordIndex)
1626+
1627+
ds2 = ds.isel(x=slice(None), y=slice(None))
1628+
assert ds2.xindexes["x"] is ds2.xindexes["y"]
1629+
15961630
def test_sel(self) -> None:
15971631
data = create_test_data()
15981632
int_slicers = {"dim1": slice(None, None, 2), "dim2": slice(2), "dim3": slice(3)}

0 commit comments

Comments
 (0)