diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index f48c9e72af1..d587ee3a02b 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -4,7 +4,7 @@ from collections import defaultdict from contextlib import suppress from datetime import timedelta -from typing import Any, Callable, Sequence, Tuple, Union +from typing import Any, Callable, Iterable, Sequence, Tuple, Union import numpy as np import pandas as pd @@ -1291,6 +1291,24 @@ def __init__(self, array): self.array = array def __getitem__(self, key): + + if not isinstance(key, VectorizedIndexer): + # if possible, short-circuit when keys are effectively slice(None) + # This preserves dask name and passes lazy array equivalence checks + # (see duck_array_ops.lazy_array_equiv) + rewritten_indexer = False + new_indexer = [] + for idim, k in enumerate(key.tuple): + if isinstance(k, Iterable) and duck_array_ops.array_equiv( + k, np.arange(self.array.shape[idim]) + ): + new_indexer.append(slice(None)) + rewritten_indexer = True + else: + new_indexer.append(k) + if rewritten_indexer: + key = type(key)(tuple(new_indexer)) + if isinstance(key, BasicIndexer): return self.array[key.tuple] elif isinstance(key, VectorizedIndexer): diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 1c3ab2ec229..8e62341d5ee 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -5,7 +5,7 @@ from collections import defaultdict from datetime import timedelta from distutils.version import LooseVersion -from typing import Any, Dict, Hashable, Iterable, Mapping, TypeVar, Union +from typing import Any, Dict, Hashable, Mapping, TypeVar, Union import numpy as np import pandas as pd @@ -540,20 +540,6 @@ def _broadcast_indexes(self, key): k.item() if isinstance(k, np.ndarray) and k.ndim == 0 else k for k in key ) - key_dict = dict(zip(self.dims, key)) - for dim, k in key_dict.items(): - if isinstance(k, Iterable): - # let da.sel(x=da.x) pass but skip if Variable has different dimensions - # e.g. da.sel(x=Variable(("points",), [0, 1, 2])) - if isinstance(k, Variable) and k.dims != (dim,): - continue - if duck_array_ops.array_equiv(k, np.arange(self.sizes[dim])): - # short-circuit when keys are effectively slice(None) - # This preserves dask name and passes lazy array equivalence checks - # (see duck_array_ops.lazy_array_equiv) - key_dict[dim] = slice(None) - key = tuple(key_dict.values()) - if all(isinstance(k, BASIC_INDEXING_TYPES) for k in key): return self._broadcast_indexes_basic(key)