diff --git a/xarray/coding/strings.py b/xarray/coding/strings.py index db95286f6aa..5612cdd11e9 100644 --- a/xarray/coding/strings.py +++ b/xarray/coding/strings.py @@ -249,15 +249,20 @@ def shape(self) -> tuple[int, ...]: def __repr__(self): return f"{type(self).__name__}({self.array!r})" - def _vindex_get(self, key): - return _numpy_char_to_bytes(self.array.vindex[key]) - - def _oindex_get(self, key): - return _numpy_char_to_bytes(self.array.oindex[key]) + def _check_and_raise_if_non_basic_indexer(self, indexer) -> None: + ... + # TODO: this is a temporary fix until BackendArray supports vindex and oindex def __getitem__(self, key): # require slicing the last dimension completely key = type(key)(indexing.expanded_indexer(key.tuple, self.array.ndim)) if key.tuple[-1] != slice(None): raise IndexError("too many indices") - return _numpy_char_to_bytes(self.array[key]) + # TODO: this is a temporary fix until BackendArray supports vindex and oindex + if isinstance(key, indexing.OuterIndexer): + data = self.array.oindex[key] + elif isinstance(key, indexing.VectorizedIndexer): + data = self.array.vindex[key] + else: + data = self.array[key] + return _numpy_char_to_bytes(data) diff --git a/xarray/coding/variables.py b/xarray/coding/variables.py index d31cb6e626a..6fb00031bcc 100644 --- a/xarray/coding/variables.py +++ b/xarray/coding/variables.py @@ -68,14 +68,19 @@ def __init__(self, array, func: Callable, dtype: np.typing.DTypeLike): def dtype(self) -> np.dtype: return np.dtype(self._dtype) - def _oindex_get(self, key): - return type(self)(self.array.oindex[key], self.func, self.dtype) - - def _vindex_get(self, key): - return type(self)(self.array.vindex[key], self.func, self.dtype) + def _check_and_raise_if_non_basic_indexer(self, indexer) -> None: + ... + # TODO: this is a temporary fix until BackendArray supports vindex and oindex def __getitem__(self, key): - return type(self)(self.array[key], self.func, self.dtype) + # TODO: this is a temporary fix until BackendArray supports vindex and oindex + if isinstance(key, indexing.OuterIndexer): + data = self.array.oindex[key] + elif isinstance(key, indexing.VectorizedIndexer): + data = self.array.vindex[key] + else: + data = self.array[key] + return type(self)(data, self.func, self.dtype) def get_duck_array(self): return self.func(self.array.get_duck_array()) @@ -113,14 +118,19 @@ def __init__(self, array) -> None: def dtype(self) -> np.dtype: return np.dtype(self.array.dtype.kind + str(self.array.dtype.itemsize)) - def _oindex_get(self, key): - return np.asarray(self.array.oindex[key], dtype=self.dtype) - - def _vindex_get(self, key): - return np.asarray(self.array.vindex[key], dtype=self.dtype) + def _check_and_raise_if_non_basic_indexer(self, indexer) -> None: + ... + # TODO: this is a temporary fix until BackendArray supports vindex and oindex def __getitem__(self, key) -> np.ndarray: - return np.asarray(self.array[key], dtype=self.dtype) + # TODO: this is a temporary fix until BackendArray supports vindex and oindex + if isinstance(key, indexing.OuterIndexer): + data = self.array.oindex[key] + elif isinstance(key, indexing.VectorizedIndexer): + data = self.array.vindex[key] + else: + data = self.array[key] + return np.asarray(data, dtype=self.dtype) class BoolTypeArray(indexing.ExplicitlyIndexedNDArrayMixin): @@ -151,14 +161,19 @@ def __init__(self, array) -> None: def dtype(self) -> np.dtype: return np.dtype("bool") - def _oindex_get(self, key): - return np.asarray(self.array.oindex[key], dtype=self.dtype) - - def _vindex_get(self, key): - return np.asarray(self.array.vindex[key], dtype=self.dtype) + def _check_and_raise_if_non_basic_indexer(self, indexer) -> None: + ... + # TODO: this is a temporary fix until BackendArray supports vindex and oindex def __getitem__(self, key) -> np.ndarray: - return np.asarray(self.array[key], dtype=self.dtype) + # TODO: this is a temporary fix until BackendArray supports vindex and oindex + if isinstance(key, indexing.OuterIndexer): + data = self.array.oindex[key] + elif isinstance(key, indexing.VectorizedIndexer): + data = self.array.vindex[key] + else: + data = self.array[key] + return np.asarray(data, dtype=self.dtype) def lazy_elemwise_func(array, func: Callable, dtype: np.typing.DTypeLike): diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 0926da6fd80..aecfadf9fb0 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -639,7 +639,18 @@ def shape(self) -> _Shape: return tuple(shape) def get_duck_array(self): - if isinstance(self.array, ExplicitlyIndexedNDArrayMixin): + from xarray.coding import strings, variables + + if isinstance(self.array, ExplicitlyIndexedNDArrayMixin) and not isinstance( + self.array, + ( + strings.StackedBytesArray, + variables._ElementwiseFunctionArray, + variables.BoolTypeArray, + variables.NativeEndiannessArray, + ), + ): + # TODO: Remove the isinstance check for variables.BoolTypeArray and variables.NativeEndiannessArray once the BackendArrray is updated with oindex and vindex properties array = apply_indexer(self.array, self.key) else: # If the array is not an ExplicitlyIndexedNDArrayMixin, @@ -715,7 +726,18 @@ def shape(self) -> _Shape: return np.broadcast(*self.key.tuple).shape def get_duck_array(self): - if isinstance(self.array, ExplicitlyIndexedNDArrayMixin): + from xarray.coding import strings, variables + + if isinstance(self.array, ExplicitlyIndexedNDArrayMixin) and not isinstance( + self.array, + ( + strings.StackedBytesArray, + variables._ElementwiseFunctionArray, + variables.BoolTypeArray, + variables.NativeEndiannessArray, + ), + ): + # TODO: Remove the isinstance check for variables.BoolTypeArray and variables.NativeEndiannessArray once the BackendArrray is updated with oindex and vindex properties array = apply_indexer(self.array, self.key) else: # If the array is not an ExplicitlyIndexedNDArrayMixin, diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index fdf181b583a..43c27d88d9e 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -3211,6 +3211,15 @@ def test_bytes_pickle(self) -> None: unpickled = pickle.loads(pickle.dumps(ds)) assert_identical(unpickled, data) + def test_scipy_wrapper_array_oindex_vindex(self) -> None: + ds = xr.Dataset() + ds["A"] = xr.DataArray([[1, "a"], [2, "b"]], dims=["x", "y"]) + with create_tmp_file(allow_cleanup_failure=False) as path: + ds.to_netcdf(path, engine="scipy") + with xr.open_dataset(path, engine="scipy") as ds2: + with create_tmp_file(allow_cleanup_failure=False) as path2: + ds2.sel(y=[1]).to_netcdf(path2) + @requires_scipy class TestScipyFileObject(CFEncodedBase, NetCDF3Only): diff --git a/xarray/tests/test_coding_strings.py b/xarray/tests/test_coding_strings.py index 51f63ea72dd..f1eca00f9a1 100644 --- a/xarray/tests/test_coding_strings.py +++ b/xarray/tests/test_coding_strings.py @@ -181,7 +181,7 @@ def test_StackedBytesArray_vectorized_indexing() -> None: V = IndexerMaker(indexing.VectorizedIndexer) indexer = V[np.array([[0, 1], [1, 0]])] - actual = stacked.vindex[indexer] + actual = stacked[indexer] assert_array_equal(actual, expected)