Skip to content

Commit

Permalink
Enable explicit use of key tuples (instead of *Indexer objects) in in…
Browse files Browse the repository at this point in the history
…dexing adapters and explicitly indexed arrays (#8870)

* pass key tuple to indexing adapters and explicitly indexed arrays

* update indexing in StackedBytesArray

* Update indexing in StackedBytesArray

* Add _IndexerKey type to _typing.py

* Update indexing in StackedBytesArray

* use tuple indexing in test_backend_array_deprecation_warning

* Add support for CompatIndexedTuple in explicit indexing adapter

This commit updates the `explicit_indexing_adapter` function to accept both
`ExplicitIndexer` and the new `CompatIndexedTuple`. The `CompatIndexedTuple` is
designed to facilitate the transition towards using raw tuples by carrying
additional metadata about the indexing type (basic, vectorized, or outer).

* remove unused code

* type hint fixes

* fix docstrings

* fix tests

* fix docstrings

* Apply suggestions from code review

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

* update docstrings and pass tuples directly

* Some test cleanup

* update docstring

* use `BasicIndexer` instead of `CompatIndexedTuple`

* support explicit indexing with tuples

* fix mypy errors

* remove unused IndexerMaker

* Update LazilyIndexedArray._updated_key to support explicit indexing with tuples

---------

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>
Co-authored-by: Deepak Cherian <deepak@cherian.net>
  • Loading branch information
3 people committed May 10, 2024
1 parent ddd4cdb commit 96ac4b7
Show file tree
Hide file tree
Showing 9 changed files with 212 additions and 199 deletions.
20 changes: 11 additions & 9 deletions xarray/coding/strings.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from xarray.core import indexing
from xarray.core.utils import module_available
from xarray.core.variable import Variable
from xarray.namedarray._typing import _IndexerKey
from xarray.namedarray.parallelcompat import get_chunked_array_type
from xarray.namedarray.pycompat import is_chunked_array

Expand Down Expand Up @@ -220,8 +221,7 @@ class StackedBytesArray(indexing.ExplicitlyIndexedNDArrayMixin):
"""Wrapper around array-like objects to create a new indexable object where
values, when accessed, are automatically stacked along the last dimension.
>>> indexer = indexing.BasicIndexer((slice(None),))
>>> StackedBytesArray(np.array(["a", "b", "c"], dtype="S1"))[indexer]
>>> StackedBytesArray(np.array(["a", "b", "c"], dtype="S1"))[(slice(None),)]
array(b'abc', dtype='|S3')
"""

Expand All @@ -240,7 +240,7 @@ def __init__(self, array):

@property
def dtype(self):
return np.dtype("S" + str(self.array.shape[-1]))
return np.dtype(f"S{str(self.array.shape[-1])}")

@property
def shape(self) -> tuple[int, ...]:
Expand All @@ -249,15 +249,17 @@ def shape(self) -> tuple[int, ...]:
def __repr__(self):
return f"{type(self).__name__}({self.array!r})"

def _vindex_get(self, key):
def _vindex_get(self, key: _IndexerKey):
return _numpy_char_to_bytes(self.array.vindex[key])

def _oindex_get(self, key):
def _oindex_get(self, key: _IndexerKey):
return _numpy_char_to_bytes(self.array.oindex[key])

def __getitem__(self, key):
def __getitem__(self, key: _IndexerKey):
from xarray.core.indexing import BasicIndexer

# require slicing the last dimension completely
key = type(key)(indexing.expanded_indexer(key.tuple, self.array.ndim))
if key.tuple[-1] != slice(None):
indexer = indexing.expanded_indexer(key, self.array.ndim)
if indexer[-1] != slice(None):
raise IndexError("too many indices")
return _numpy_char_to_bytes(self.array[key])
return _numpy_char_to_bytes(self.array[BasicIndexer(indexer)])
6 changes: 2 additions & 4 deletions xarray/coding/variables.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,7 @@ class NativeEndiannessArray(indexing.ExplicitlyIndexedNDArrayMixin):
>>> NativeEndiannessArray(x).dtype
dtype('int16')
>>> indexer = indexing.BasicIndexer((slice(None),))
>>> NativeEndiannessArray(x)[indexer].dtype
>>> NativeEndiannessArray(x)[(slice(None),)].dtype
dtype('int16')
"""

Expand Down Expand Up @@ -137,8 +136,7 @@ class BoolTypeArray(indexing.ExplicitlyIndexedNDArrayMixin):
>>> BoolTypeArray(x).dtype
dtype('bool')
>>> indexer = indexing.BasicIndexer((slice(None),))
>>> BoolTypeArray(x)[indexer].dtype
>>> BoolTypeArray(x)[(slice(None),)].dtype
dtype('bool')
"""

Expand Down
Loading

0 comments on commit 96ac4b7

Please sign in to comment.