Skip to content

Commit

Permalink
Restrict fastpath isel indexes to the case of all PandasIndex (#10066)
Browse files Browse the repository at this point in the history
  • Loading branch information
benbovy authored Feb 24, 2025
1 parent df33a9a commit 2475d49
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 4 deletions.
2 changes: 2 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ Bug fixes
- Fix DataArray().drop_attrs(deep=False) and add support for attrs to
DataArray()._replace(). (:issue:`10027`, :pull:`10030`). By `Jan
Haacker <https://github.com/j-haacker>`_.
- Fix ``isel`` for multi-coordinate Xarray indexes (:issue:`10063`, :pull:`10066`).
By `Benoit Bovy <https://github.com/benbovy>`_.


Documentation
Expand Down
9 changes: 5 additions & 4 deletions xarray/core/indexes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1991,10 +1991,11 @@ def isel_indexes(
indexes: Indexes[Index],
indexers: Mapping[Any, Any],
) -> tuple[dict[Hashable, Index], dict[Hashable, Variable]]:
# TODO: remove if clause in the future. It should be unnecessary.
# See failure introduced when removed
# https://github.com/pydata/xarray/pull/9002#discussion_r1590443756
if any(isinstance(v, PandasMultiIndex) for v in indexes._indexes.values()):
# Fast path function _apply_indexes_fast does not work with multi-coordinate
# Xarray indexes (see https://github.com/pydata/xarray/issues/10063).
# -> call it only in the most common case where all indexes are default
# PandasIndex each associated to a single 1-dimensional coordinate.
if any(type(idx) is not PandasIndex for idx in indexes._indexes.values()):
return _apply_indexes(indexes, indexers, "isel")
else:
return _apply_indexes_fast(indexes, indexers, "isel")
Expand Down
34 changes: 34 additions & 0 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -1593,6 +1593,40 @@ def test_isel_fancy_convert_index_variable(self) -> None:
assert "x" not in actual.xindexes
assert not isinstance(actual.x.variable, IndexVariable)

def test_isel_multicoord_index(self) -> None:
# regression test https://github.com/pydata/xarray/issues/10063
# isel on a multi-coordinate index should return a unique index associated
# to each coordinate
class MultiCoordIndex(xr.Index):
def __init__(self, idx1, idx2):
self.idx1 = idx1
self.idx2 = idx2

@classmethod
def from_variables(cls, variables, *, options=None):
idx1 = PandasIndex.from_variables(
{"x": variables["x"]}, options=options
)
idx2 = PandasIndex.from_variables(
{"y": variables["y"]}, options=options
)

return cls(idx1, idx2)

def create_variables(self, variables=None):
return {**self.idx1.create_variables(), **self.idx2.create_variables()}

def isel(self, indexers):
idx1 = self.idx1.isel({"x": indexers.get("x", slice(None))})
idx2 = self.idx2.isel({"y": indexers.get("y", slice(None))})
return MultiCoordIndex(idx1, idx2)

coords = xr.Coordinates(coords={"x": [0, 1], "y": [1, 2]}, indexes={})
ds = xr.Dataset(coords=coords).set_xindex(["x", "y"], MultiCoordIndex)

ds2 = ds.isel(x=slice(None), y=slice(None))
assert ds2.xindexes["x"] is ds2.xindexes["y"]

def test_sel(self) -> None:
data = create_test_data()
int_slicers = {"dim1": slice(None, None, 2), "dim2": slice(2), "dim3": slice(3)}
Expand Down

0 comments on commit 2475d49

Please sign in to comment.