diff --git a/xarray/coding/variables.py b/xarray/coding/variables.py index 96413da1e8f..662fec4b2c4 100644 --- a/xarray/coding/variables.py +++ b/xarray/coding/variables.py @@ -520,9 +520,9 @@ def decode(self, variable: Variable, name: T_Name = None) -> Variable: scale_factor = pop_to(attrs, encoding, "scale_factor", name=name) add_offset = pop_to(attrs, encoding, "add_offset", name=name) - if np.ndim(scale_factor) > 0: + if duck_array_ops.ndim(scale_factor) > 0: scale_factor = np.asarray(scale_factor).item() - if np.ndim(add_offset) > 0: + if duck_array_ops.ndim(add_offset) > 0: add_offset = np.asarray(add_offset).item() # if we have a _FillValue/masked_value in encoding we already have the wanted # floating point dtype here (via CFMaskCoder), so no check is necessary diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index 51ebfdd8177..e98ac0f36a1 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -796,6 +796,12 @@ def _nd_cum_func(cum_func, array, axis, **kwargs): return out +def ndim(array) -> int: + # Required part of the duck array and the array-api, but we fall back in case + # https://docs.xarray.dev/en/latest/internals/duck-arrays-integration.html#duck-array-requirements + return array.ndim if hasattr(array, "ndim") else np.ndim(array) + + def cumprod(array, axis=None, **kwargs): """N-dimensional version of cumprod.""" return _nd_cum_func(cumprod_1d, array, axis, **kwargs) diff --git a/xarray/core/extension_array.py b/xarray/core/extension_array.py index 9052f5ae0a0..7cc9db96d0d 100644 --- a/xarray/core/extension_array.py +++ b/xarray/core/extension_array.py @@ -66,6 +66,11 @@ def __extension_duck_array__where( return cast(T_ExtensionArray, pd.Series(x).where(condition, pd.Series(y)).array) +@implements(np.ndim) +def __extension_duck_array__ndim(x: PandasExtensionArray) -> int: + return x.ndim + + @implements(np.reshape) def __extension_duck_array__reshape( arr: T_ExtensionArray, shape: tuple diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 1813a25d7af..e14543e646f 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -769,14 +769,15 @@ def __repr__(self) -> str: def _wrap_numpy_scalars(array): """Wrap NumPy scalars in 0d arrays.""" - if np.ndim(array) == 0 and ( + ndim = duck_array_ops.ndim(array) + if ndim == 0 and ( isinstance(array, np.generic) or not (is_duck_array(array) or isinstance(array, NDArrayMixin)) ): return np.array(array) elif hasattr(array, "dtype"): return array - elif np.ndim(array) == 0: + elif ndim == 0: return np.array(array) else: return array diff --git a/xarray/indexes/range_index.py b/xarray/indexes/range_index.py index b04c86e1a5d..2b9a5e5071a 100644 --- a/xarray/indexes/range_index.py +++ b/xarray/indexes/range_index.py @@ -5,6 +5,7 @@ import numpy as np import pandas as pd +from xarray.core import duck_array_ops from xarray.core.coordinate_transform import CoordinateTransform from xarray.core.dataarray import DataArray from xarray.core.indexes import CoordinateTransformIndex, Index, PandasIndex @@ -320,7 +321,9 @@ def isel( if isinstance(idxer, slice): return RangeIndex(self.transform.slice(idxer)) - elif (isinstance(idxer, Variable) and idxer.ndim > 1) or np.ndim(idxer) == 0: + elif (isinstance(idxer, Variable) and idxer.ndim > 1) or duck_array_ops.ndim( + idxer + ) == 0: return None else: values = self.transform.forward({self.dim: np.asarray(idxer)})[ diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index 1e7c32dec1e..2f67e97522c 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -15,6 +15,7 @@ from xarray import DataArray, Dataset, IndexVariable, Variable, set_options from xarray.core import dtypes, duck_array_ops, indexing from xarray.core.common import full_like, ones_like, zeros_like +from xarray.core.extension_array import PandasExtensionArray from xarray.core.indexing import ( BasicIndexer, CopyOnWriteArray, @@ -2894,6 +2895,7 @@ class TestBackendIndexing: @pytest.fixture(autouse=True) def setUp(self): self.d = np.random.random((10, 3)).astype(np.float64) + self.cat = PandasExtensionArray(pd.Categorical(["a", "b"] * 5)) def check_orthogonal_indexing(self, v): assert np.allclose(v.isel(x=[8, 3], y=[2, 1]), self.d[[8, 3]][:, [2, 1]]) @@ -2913,6 +2915,14 @@ def test_NumpyIndexingAdapter(self): dims=("x", "y"), data=NumpyIndexingAdapter(NumpyIndexingAdapter(self.d)) ) + def test_extension_array_duck_array(self): + lazy = LazilyIndexedArray(self.cat) + assert (lazy.get_duck_array().array == self.cat).all() + + def test_extension_array_duck_indexed(self): + lazy = Variable(dims=("x"), data=LazilyIndexedArray(self.cat)) + assert (lazy[[0, 1, 5]] == ["a", "b", "b"]).all() + def test_LazilyIndexedArray(self): v = Variable(dims=("x", "y"), data=LazilyIndexedArray(self.d)) self.check_orthogonal_indexing(v) @@ -2951,12 +2961,14 @@ def test_MemoryCachedArray(self): def test_DaskIndexingAdapter(self): import dask.array as da - da = da.asarray(self.d) - v = Variable(dims=("x", "y"), data=DaskIndexingAdapter(da)) + dask_array = da.asarray(self.d) + v = Variable(dims=("x", "y"), data=DaskIndexingAdapter(dask_array)) self.check_orthogonal_indexing(v) self.check_vectorized_indexing(v) # doubly wrapping - v = Variable(dims=("x", "y"), data=CopyOnWriteArray(DaskIndexingAdapter(da))) + v = Variable( + dims=("x", "y"), data=CopyOnWriteArray(DaskIndexingAdapter(dask_array)) + ) self.check_orthogonal_indexing(v) self.check_vectorized_indexing(v)