Skip to content

Commit df25933

Browse files
authored
Fix isel performance regression (#3319)
* Fix isel performance regression xref GH2227 Before: indexing.BooleanIndexing.time_indexing 898±0ms After indexing.BooleanIndexing.time_indexing 401±0ms * mypy fix
1 parent fddced0 commit df25933

2 files changed

Lines changed: 19 additions & 6 deletions

File tree

asv_bench/benchmarks/indexing.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,3 +125,16 @@ def setup(self, key):
125125
requires_dask()
126126
super().setup(key)
127127
self.ds = self.ds.chunk({"x": 100, "y": 50, "t": 50})
128+
129+
130+
class BooleanIndexing:
131+
# https://github.com/pydata/xarray/issues/2227
132+
def setup(self):
133+
self.ds = xr.Dataset(
134+
{"a": ("time", np.arange(10_000_000))},
135+
coords={"time": np.arange(10_000_000)},
136+
)
137+
self.time_filter = self.ds.time > 50_000
138+
139+
def time_indexing(self):
140+
self.ds.isel(time=self.time_filter)

xarray/core/dataset.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1781,7 +1781,7 @@ def _validate_indexers(
17811781
elif isinstance(v, Dataset):
17821782
raise TypeError("cannot use a Dataset as an indexer")
17831783
elif isinstance(v, Sequence) and len(v) == 0:
1784-
v = IndexVariable((k,), np.zeros((0,), dtype="int64"))
1784+
v = Variable((k,), np.zeros((0,), dtype="int64"))
17851785
else:
17861786
v = np.asarray(v)
17871787

@@ -1795,16 +1795,13 @@ def _validate_indexers(
17951795
if v.ndim == 0:
17961796
v = Variable((), v)
17971797
elif v.ndim == 1:
1798-
v = IndexVariable((k,), v)
1798+
v = Variable((k,), v)
17991799
else:
18001800
raise IndexError(
18011801
"Unlabeled multi-dimensional array cannot be "
18021802
"used for indexing: {}".format(k)
18031803
)
18041804

1805-
if v.ndim == 1:
1806-
v = v.to_index_variable()
1807-
18081805
indexers_list.append((k, v))
18091806

18101807
return indexers_list
@@ -2367,7 +2364,10 @@ def interp(
23672364
if kwargs is None:
23682365
kwargs = {}
23692366
coords = either_dict_or_kwargs(coords, coords_kwargs, "interp")
2370-
indexers = OrderedDict(self._validate_indexers(coords))
2367+
indexers = OrderedDict(
2368+
(k, v.to_index_variable() if isinstance(v, Variable) and v.ndim == 1 else v)
2369+
for k, v in self._validate_indexers(coords)
2370+
)
23712371

23722372
obj = self if assume_sorted else self.sortby([k for k in coords])
23732373

0 commit comments

Comments
 (0)