Skip to content

Commit ff0d056

Browse files
andersy005dcherian
andauthored
refactor DaskIndexingAdapter __getitem__ method (#8758)
* remove the unnecessary check for `VectorizedIndexer` * remvove test * update whats-new * update test * Trim --------- Co-authored-by: Deepak Cherian <[email protected]>
1 parent 5d93331 commit ff0d056

File tree

3 files changed

+4
-24
lines changed

3 files changed

+4
-24
lines changed

doc/whats-new.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,9 @@ Internal Changes
102102
- Adds :py:func:`open_datatree` into ``xarray/backends`` (:pull:`8697`) By `Matt
103103
Savoie <https://github.com/flamingbear>`_.
104104

105+
- Refactor :py:meth:`xarray.core.indexing.DaskIndexingAdapter.__getitem__` to remove an unnecessary rewrite of the indexer key
106+
(:issue: `8377`, :pull:`8758`) By `Anderson Banihirwe <https://github.com/andersy005>`
107+
105108
.. _whats-new.2024.01.1:
106109

107110
v2024.01.1 (23 Jan, 2024)

xarray/core/indexing.py

Lines changed: 1 addition & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import functools
55
import operator
66
from collections import Counter, defaultdict
7-
from collections.abc import Hashable, Iterable, Mapping
7+
from collections.abc import Hashable, Mapping
88
from contextlib import suppress
99
from dataclasses import dataclass, field
1010
from datetime import timedelta
@@ -1418,23 +1418,6 @@ def __init__(self, array):
14181418
self.array = array
14191419

14201420
def __getitem__(self, key):
1421-
if not isinstance(key, VectorizedIndexer):
1422-
# if possible, short-circuit when keys are effectively slice(None)
1423-
# This preserves dask name and passes lazy array equivalence checks
1424-
# (see duck_array_ops.lazy_array_equiv)
1425-
rewritten_indexer = False
1426-
new_indexer = []
1427-
for idim, k in enumerate(key.tuple):
1428-
if isinstance(k, Iterable) and (
1429-
not is_duck_dask_array(k)
1430-
and duck_array_ops.array_equiv(k, np.arange(self.array.shape[idim]))
1431-
):
1432-
new_indexer.append(slice(None))
1433-
rewritten_indexer = True
1434-
else:
1435-
new_indexer.append(k)
1436-
if rewritten_indexer:
1437-
key = type(key)(tuple(new_indexer))
14381421

14391422
if isinstance(key, BasicIndexer):
14401423
return self.array[key.tuple]

xarray/tests/test_dask.py

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1662,16 +1662,10 @@ def test_lazy_array_equiv_merge(compat):
16621662
lambda a: a.assign_attrs(new_attr="anew"),
16631663
lambda a: a.assign_coords(cxy=a.cxy),
16641664
lambda a: a.copy(),
1665-
lambda a: a.isel(x=np.arange(a.sizes["x"])),
16661665
lambda a: a.isel(x=slice(None)),
16671666
lambda a: a.loc[dict(x=slice(None))],
1668-
lambda a: a.loc[dict(x=np.arange(a.sizes["x"]))],
1669-
lambda a: a.loc[dict(x=a.x)],
1670-
lambda a: a.sel(x=a.x),
1671-
lambda a: a.sel(x=a.x.values),
16721667
lambda a: a.transpose(...),
16731668
lambda a: a.squeeze(), # no dimensions to squeeze
1674-
lambda a: a.sortby("x"), # "x" is already sorted
16751669
lambda a: a.reindex(x=a.x),
16761670
lambda a: a.reindex_like(a),
16771671
lambda a: a.rename({"cxy": "cnew"}).rename({"cnew": "cxy"}),

0 commit comments

Comments
 (0)