(fix): ndim accessible as np.ndim on PandasExtensionArray (#10414)

Author: ilan-gold

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

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)

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

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

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)})[

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)