Skip to content

Commit 10c133b

Browse files
andersy005dcherian
andauthored
add .oindex and .vindex to BackendArray (#8885)
* add .oindex and .vindex to BackendArray * Add support for .oindex and .vindex in H5NetCDFArrayWrapper * Add support for .oindex and .vindex in NetCDF4ArrayWrapper, PydapArrayWrapper, NioArrayWrapper, and ZarrArrayWrapper * add deprecation warning * Fix deprecation warning message formatting * add tests * Update xarray/core/indexing.py Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com> * Update ZarrArrayWrapper class in xarray/backends/zarr.py Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com> --------- Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>
1 parent b81b451 commit 10c133b

File tree

9 files changed

+193
-37
lines changed

9 files changed

+193
-37
lines changed

xarray/backends/common.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,24 @@ def get_duck_array(self, dtype: np.typing.DTypeLike = None):
210210
key = indexing.BasicIndexer((slice(None),) * self.ndim)
211211
return self[key] # type: ignore [index]
212212

213+
def _oindex_get(self, key: indexing.OuterIndexer):
214+
raise NotImplementedError(
215+
f"{self.__class__.__name__}._oindex_get method should be overridden"
216+
)
217+
218+
def _vindex_get(self, key: indexing.VectorizedIndexer):
219+
raise NotImplementedError(
220+
f"{self.__class__.__name__}._vindex_get method should be overridden"
221+
)
222+
223+
@property
224+
def oindex(self) -> indexing.IndexCallable:
225+
return indexing.IndexCallable(self._oindex_get)
226+
227+
@property
228+
def vindex(self) -> indexing.IndexCallable:
229+
return indexing.IndexCallable(self._vindex_get)
230+
213231

214232
class AbstractDataStore:
215233
__slots__ = ()

xarray/backends/h5netcdf_.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,17 @@ def get_array(self, needs_lock=True):
4848
ds = self.datastore._acquire(needs_lock)
4949
return ds.variables[self.variable_name]
5050

51-
def __getitem__(self, key):
51+
def _oindex_get(self, key: indexing.OuterIndexer):
52+
return indexing.explicit_indexing_adapter(
53+
key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
54+
)
55+
56+
def _vindex_get(self, key: indexing.VectorizedIndexer):
57+
return indexing.explicit_indexing_adapter(
58+
key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
59+
)
60+
61+
def __getitem__(self, key: indexing.BasicIndexer):
5262
return indexing.explicit_indexing_adapter(
5363
key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
5464
)

xarray/backends/netCDF4_.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,17 @@ def get_array(self, needs_lock=True):
9797
variable.set_auto_chartostring(False)
9898
return variable
9999

100-
def __getitem__(self, key):
100+
def _oindex_get(self, key: indexing.OuterIndexer):
101+
return indexing.explicit_indexing_adapter(
102+
key, self.shape, indexing.IndexingSupport.OUTER, self._getitem
103+
)
104+
105+
def _vindex_get(self, key: indexing.VectorizedIndexer):
106+
return indexing.explicit_indexing_adapter(
107+
key, self.shape, indexing.IndexingSupport.OUTER, self._getitem
108+
)
109+
110+
def __getitem__(self, key: indexing.BasicIndexer):
101111
return indexing.explicit_indexing_adapter(
102112
key, self.shape, indexing.IndexingSupport.OUTER, self._getitem
103113
)

xarray/backends/pydap_.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,17 @@ def shape(self) -> tuple[int, ...]:
4343
def dtype(self):
4444
return self.array.dtype
4545

46-
def __getitem__(self, key):
46+
def _oindex_get(self, key: indexing.OuterIndexer):
47+
return indexing.explicit_indexing_adapter(
48+
key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
49+
)
50+
51+
def _vindex_get(self, key: indexing.VectorizedIndexer):
52+
return indexing.explicit_indexing_adapter(
53+
key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
54+
)
55+
56+
def __getitem__(self, key: indexing.BasicIndexer):
4757
return indexing.explicit_indexing_adapter(
4858
key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
4959
)

xarray/backends/pynio_.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,17 @@ def get_array(self, needs_lock=True):
5050
ds = self.datastore._manager.acquire(needs_lock)
5151
return ds.variables[self.variable_name]
5252

53-
def __getitem__(self, key):
53+
def _oindex_get(self, key: indexing.OuterIndexer):
54+
return indexing.explicit_indexing_adapter(
55+
key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
56+
)
57+
58+
def _vindex_get(self, key: indexing.VectorizedIndexer):
59+
return indexing.explicit_indexing_adapter(
60+
key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
61+
)
62+
63+
def __getitem__(self, key: indexing.BasicIndexer):
5464
return indexing.explicit_indexing_adapter(
5565
key, self.shape, indexing.IndexingSupport.BASIC, self._getitem
5666
)

xarray/backends/scipy_.py

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -67,15 +67,7 @@ def get_variable(self, needs_lock=True):
6767
ds = self.datastore._manager.acquire(needs_lock)
6868
return ds.variables[self.variable_name]
6969

70-
def _getitem(self, key):
71-
with self.datastore.lock:
72-
data = self.get_variable(needs_lock=False).data
73-
return data[key]
74-
75-
def __getitem__(self, key):
76-
data = indexing.explicit_indexing_adapter(
77-
key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
78-
)
70+
def _finalize_result(self, data):
7971
# Copy data if the source file is mmapped. This makes things consistent
8072
# with the netCDF4 library by ensuring we can safely read arrays even
8173
# after closing associated files.
@@ -88,6 +80,29 @@ def __getitem__(self, key):
8880

8981
return np.array(data, dtype=self.dtype, copy=copy)
9082

83+
def _getitem(self, key):
84+
with self.datastore.lock:
85+
data = self.get_variable(needs_lock=False).data
86+
return data[key]
87+
88+
def _vindex_get(self, key: indexing.VectorizedIndexer):
89+
data = indexing.explicit_indexing_adapter(
90+
key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
91+
)
92+
return self._finalize_result(data)
93+
94+
def _oindex_get(self, key: indexing.OuterIndexer):
95+
data = indexing.explicit_indexing_adapter(
96+
key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
97+
)
98+
return self._finalize_result(data)
99+
100+
def __getitem__(self, key):
101+
data = indexing.explicit_indexing_adapter(
102+
key, self.shape, indexing.IndexingSupport.OUTER_1VECTOR, self._getitem
103+
)
104+
return self._finalize_result(data)
105+
91106
def __setitem__(self, key, value):
92107
with self.datastore.lock:
93108
data = self.get_variable(needs_lock=False)

xarray/backends/zarr.py

Lines changed: 31 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -84,25 +84,38 @@ def __init__(self, zarr_array):
8484
def get_array(self):
8585
return self._array
8686

87-
def _oindex(self, key):
88-
return self._array.oindex[key]
89-
90-
def _vindex(self, key):
91-
return self._array.vindex[key]
92-
93-
def _getitem(self, key):
94-
return self._array[key]
95-
96-
def __getitem__(self, key):
97-
array = self._array
98-
if isinstance(key, indexing.BasicIndexer):
99-
method = self._getitem
100-
elif isinstance(key, indexing.VectorizedIndexer):
101-
method = self._vindex
102-
elif isinstance(key, indexing.OuterIndexer):
103-
method = self._oindex
87+
def _oindex_get(self, key: indexing.OuterIndexer):
88+
def raw_indexing_method(key):
89+
return self._array.oindex[key]
90+
91+
return indexing.explicit_indexing_adapter(
92+
key,
93+
self._array.shape,
94+
indexing.IndexingSupport.VECTORIZED,
95+
raw_indexing_method,
96+
)
97+
98+
def _vindex_get(self, key: indexing.VectorizedIndexer):
99+
100+
def raw_indexing_method(key):
101+
return self._array.vindex[key]
102+
103+
return indexing.explicit_indexing_adapter(
104+
key,
105+
self._array.shape,
106+
indexing.IndexingSupport.VECTORIZED,
107+
raw_indexing_method,
108+
)
109+
110+
def __getitem__(self, key: indexing.BasicIndexer):
111+
def raw_indexing_method(key):
112+
return self._array[key]
113+
104114
return indexing.explicit_indexing_adapter(
105-
key, array.shape, indexing.IndexingSupport.VECTORIZED, method
115+
key,
116+
self._array.shape,
117+
indexing.IndexingSupport.VECTORIZED,
118+
raw_indexing_method,
106119
)
107120

108121
# if self.ndim == 0:

xarray/core/indexing.py

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import enum
44
import functools
55
import operator
6+
import warnings
67
from collections import Counter, defaultdict
78
from collections.abc import Hashable, Iterable, Mapping
89
from contextlib import suppress
@@ -564,6 +565,14 @@ def __getitem__(self, key: Any):
564565
return result
565566

566567

568+
BackendArray_fallback_warning_message = (
569+
"The array `{0}` does not support indexing using the .vindex and .oindex properties. "
570+
"The __getitem__ method is being used instead. This fallback behavior will be "
571+
"removed in a future version. Please ensure that the backend array `{1}` implements "
572+
"support for the .vindex and .oindex properties to avoid potential issues."
573+
)
574+
575+
567576
class LazilyIndexedArray(ExplicitlyIndexedNDArrayMixin):
568577
"""Wrap an array to make basic and outer indexing lazy."""
569578

@@ -615,11 +624,18 @@ def shape(self) -> _Shape:
615624
return tuple(shape)
616625

617626
def get_duck_array(self):
618-
if isinstance(self.array, ExplicitlyIndexedNDArrayMixin):
627+
try:
619628
array = apply_indexer(self.array, self.key)
620-
else:
629+
except NotImplementedError as _:
621630
# If the array is not an ExplicitlyIndexedNDArrayMixin,
622-
# it may wrap a BackendArray so use its __getitem__
631+
# it may wrap a BackendArray subclass that doesn't implement .oindex and .vindex. so use its __getitem__
632+
warnings.warn(
633+
BackendArray_fallback_warning_message.format(
634+
self.array.__class__.__name__, self.array.__class__.__name__
635+
),
636+
category=DeprecationWarning,
637+
stacklevel=2,
638+
)
623639
array = self.array[self.key]
624640

625641
# self.array[self.key] is now a numpy array when
@@ -691,12 +707,20 @@ def shape(self) -> _Shape:
691707
return np.broadcast(*self.key.tuple).shape
692708

693709
def get_duck_array(self):
694-
if isinstance(self.array, ExplicitlyIndexedNDArrayMixin):
710+
try:
695711
array = apply_indexer(self.array, self.key)
696-
else:
712+
except NotImplementedError as _:
697713
# If the array is not an ExplicitlyIndexedNDArrayMixin,
698-
# it may wrap a BackendArray so use its __getitem__
714+
# it may wrap a BackendArray subclass that doesn't implement .oindex and .vindex. so use its __getitem__
715+
warnings.warn(
716+
BackendArray_fallback_warning_message.format(
717+
self.array.__class__.__name__, self.array.__class__.__name__
718+
),
719+
category=PendingDeprecationWarning,
720+
stacklevel=2,
721+
)
699722
array = self.array[self.key]
723+
700724
# self.array[self.key] is now a numpy array when
701725
# self.array is a BackendArray subclass
702726
# and self.key is BasicIndexer((slice(None, None, None),))

xarray/tests/test_backends.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5828,3 +5828,49 @@ def test_zarr_region_chunk_partial_offset(tmp_path):
58285828
# This write is unsafe, and should raise an error, but does not.
58295829
# with pytest.raises(ValueError):
58305830
# da.isel(x=slice(5, 25)).chunk(x=(10, 10)).to_zarr(store, region="auto")
5831+
5832+
5833+
def test_backend_array_deprecation_warning(capsys):
5834+
class CustomBackendArray(xr.backends.common.BackendArray):
5835+
def __init__(self):
5836+
array = self.get_array()
5837+
self.shape = array.shape
5838+
self.dtype = array.dtype
5839+
5840+
def get_array(self):
5841+
return np.arange(10)
5842+
5843+
def __getitem__(self, key):
5844+
return xr.core.indexing.explicit_indexing_adapter(
5845+
key, self.shape, xr.core.indexing.IndexingSupport.BASIC, self._getitem
5846+
)
5847+
5848+
def _getitem(self, key):
5849+
array = self.get_array()
5850+
return array[key]
5851+
5852+
cba = CustomBackendArray()
5853+
indexer = xr.core.indexing.VectorizedIndexer(key=(np.array([0]),))
5854+
5855+
la = xr.core.indexing.LazilyIndexedArray(cba, indexer)
5856+
5857+
with warnings.catch_warnings(record=True) as w:
5858+
warnings.simplefilter("always")
5859+
la.vindex[indexer].get_duck_array()
5860+
5861+
captured = capsys.readouterr()
5862+
assert len(w) == 1
5863+
assert issubclass(w[-1].category, PendingDeprecationWarning)
5864+
assert (
5865+
"The array `CustomBackendArray` does not support indexing using the .vindex and .oindex properties."
5866+
in str(w[-1].message)
5867+
)
5868+
assert "The __getitem__ method is being used instead." in str(w[-1].message)
5869+
assert "This fallback behavior will be removed in a future version." in str(
5870+
w[-1].message
5871+
)
5872+
assert (
5873+
"Please ensure that the backend array `CustomBackendArray` implements support for the .vindex and .oindex properties to avoid potential issues."
5874+
in str(w[-1].message)
5875+
)
5876+
assert captured.out == ""

0 commit comments

Comments
 (0)