Skip to content

Commit

Permalink
Handle .oindex and .vindex for the PandasMultiIndexingAdapter and Pan…
Browse files Browse the repository at this point in the history
…dasIndexingAdapter (#8869)
  • Loading branch information
andersy005 authored Mar 23, 2024
1 parent 7c3d2dd commit 6af547c
Showing 1 changed file with 89 additions and 14 deletions.
103 changes: 89 additions & 14 deletions xarray/core/indexing.py
Original file line number Diff line number Diff line change
Expand Up @@ -1680,11 +1680,65 @@ def _convert_scalar(self, item):
# a NumPy array.
return to_0d_array(item)

def _oindex_get(self, indexer: OuterIndexer):
return self.__getitem__(indexer)
def _prepare_key(self, key: tuple[Any, ...]) -> tuple[Any, ...]:
if isinstance(key, tuple) and len(key) == 1:
# unpack key so it can index a pandas.Index object (pandas.Index
# objects don't like tuples)
(key,) = key

def _vindex_get(self, indexer: VectorizedIndexer):
return self.__getitem__(indexer)
return key

def _handle_result(
self, result: Any
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
if isinstance(result, pd.Index):
return type(self)(result, dtype=self.dtype)
else:
return self._convert_scalar(result)

def _oindex_get(
self, indexer: OuterIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
key = self._prepare_key(indexer.tuple)

if getattr(key, "ndim", 0) > 1: # Return np-array if multidimensional
indexable = NumpyIndexingAdapter(np.asarray(self))
return indexable.oindex[indexer]

result = self.array[key]

return self._handle_result(result)

def _vindex_get(
self, indexer: VectorizedIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
key = self._prepare_key(indexer.tuple)

if getattr(key, "ndim", 0) > 1: # Return np-array if multidimensional
indexable = NumpyIndexingAdapter(np.asarray(self))
return indexable.vindex[indexer]

result = self.array[key]

return self._handle_result(result)

def __getitem__(
self, indexer: ExplicitIndexer
Expand All @@ -1695,22 +1749,15 @@ def __getitem__(
| np.datetime64
| np.timedelta64
):
key = indexer.tuple
if isinstance(key, tuple) and len(key) == 1:
# unpack key so it can index a pandas.Index object (pandas.Index
# objects don't like tuples)
(key,) = key
key = self._prepare_key(indexer.tuple)

if getattr(key, "ndim", 0) > 1: # Return np-array if multidimensional
indexable = NumpyIndexingAdapter(np.asarray(self))
return apply_indexer(indexable, indexer)
return indexable[indexer]

result = self.array[key]

if isinstance(result, pd.Index):
return type(self)(result, dtype=self.dtype)
else:
return self._convert_scalar(result)
return self._handle_result(result)

def transpose(self, order) -> pd.Index:
return self.array # self.array should be always one-dimensional
Expand Down Expand Up @@ -1766,6 +1813,34 @@ def _convert_scalar(self, item):
item = item[idx]
return super()._convert_scalar(item)

def _oindex_get(
self, indexer: OuterIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
result = super()._oindex_get(indexer)
if isinstance(result, type(self)):
result.level = self.level
return result

def _vindex_get(
self, indexer: VectorizedIndexer
) -> (
PandasIndexingAdapter
| NumpyIndexingAdapter
| np.ndarray
| np.datetime64
| np.timedelta64
):
result = super()._vindex_get(indexer)
if isinstance(result, type(self)):
result.level = self.level
return result

def __getitem__(self, indexer: ExplicitIndexer):
result = super().__getitem__(indexer)
if isinstance(result, type(self)):
Expand Down

0 comments on commit 6af547c

Please sign in to comment.