Fix CuPy indexing, assignment and scalar truth conversion

doc/whats-new.rst

@@ -165,6 +165,10 @@ Deprecations
 
 Bug Fixes
 ~~~~~~~~~
+- Fix CuPy indexers with pandas-backed coordinates, basic and label-based
+  assignment to CuPy arrays, and scalar truth conversion while preserving
+  pandas extension-array behavior (:issue:`11739`, :pull:`11797`).
+  By `Huashuo Lei <https://github.com/wpyaoxueros>`_.
 - Results of the string accessor methods ``strip``, ``pad``, ``center``,
   ``ljust``, ``rjust``, ``zfill`` and ``replace`` on dask arrays are no longer
   truncated to a single character, and ``repeat`` no longer fails on dask arrays

xarray/core/common.py

@@ -166,7 +166,10 @@ class AbstractArray(Generic[DimType_co]):  # noqa: UP046
     __slots__ = ()
 
     def __bool__(self: Any) -> bool:
-        return bool(self.values)
+        data = self.data
+        if isinstance(data, pd.api.extensions.ExtensionArray):
+            return bool(self.values)
+        return bool(data)
 
     def __float__(self: Any) -> float:
         return float(self.values)

xarray/core/indexes.py

@@ -26,12 +26,15 @@
     get_valid_numpy_dtype,
     is_allowed_extension_array_dtype,
     is_dict_like,
+    is_duck_array,
     is_scalar,
 )
+from xarray.namedarray.pycompat import to_numpy
 
 if TYPE_CHECKING:
     from xarray.core.types import ErrorOptions, JoinOptions, Self
     from xarray.core.variable import Variable
+    from xarray.namedarray._typing import duckarray
 
 
 IndexVars = dict[Any, "Variable"]
@@ -836,6 +839,9 @@ def isel(
         if isinstance(indxr, slice) and indxr == slice(None):
             return self
 
+        if is_duck_array(indxr) and not isinstance(indxr, np.ndarray):
+            indxr = to_numpy(cast("duckarray[Any, Any]", indxr))
+
         return self._replace(self.index[indxr])  # type: ignore[index,unused-ignore]
 
     @override

xarray/core/indexing.py

@@ -2090,16 +2090,31 @@ def _convert_scalar(self, item) -> np.ndarray:
     def _index_get(
         self, indexer: ExplicitIndexer, func_name: str
     ) -> PandasIndexingAdapter | np.ndarray:
-        key = indexer.tuple
+        key: Any = indexer.tuple
 
         if len(key) == 1:
             # unpack key so it can index a pandas.Index object (pandas.Index
             # objects don't like tuples)
             (key,) = key
 
+        cupy_array_type = array_type("cupy")
+        if isinstance(key, cupy_array_type):
+            key = to_numpy(key)
+        elif isinstance(key, tuple):
+            key = tuple(
+                to_numpy(item) if isinstance(item, cupy_array_type) else item
+                for item in key
+            )
+
         # if multidimensional key, convert the index to numpy array and index the latter
         if getattr(key, "ndim", 0) > 1:
             indexable = NumpyIndexingAdapter(np.asarray(self))
+            indexer = type(indexer)(
+                tuple(
+                    to_numpy(item) if isinstance(item, cupy_array_type) else item
+                    for item in indexer.tuple
+                )
+            )
             return getattr(indexable, func_name)(indexer)
 
         # otherwise index the pandas index then re-wrap or convert the result

xarray/core/variable.py

@@ -1129,6 +1129,7 @@ def __setitem__(self, key, value):
                 value = Variable(dims[-value.ndim :], value)
         # broadcast to become assignable
         value = value.set_dims(dims).data
+        value = to_like_array(value, self._data)
 
         if new_order:
             value = duck_array_ops.asarray(value)

xarray/tests/test_cupy.py

@@ -90,3 +90,86 @@ def test_isel_0d_cupy_indexer(
     assert actual.dims == expected_dims
     assert isinstance(actual.data, cp.ndarray)
     np.testing.assert_array_equal(actual.data.get(), expected)
+
+
+def test_isel_cupy_vector_indexer_with_pandas_index() -> None:
+    da = xr.DataArray(
+        cp.arange(24).reshape(4, 6),
+        dims=("x", "y"),
+        coords={"x": [10, 20, 30, 40]},
+    )
+
+    actual = da.isel(x=cp.array([1, 2]))
+    expected_data = cp.arange(24).reshape(4, 6)[cp.array([1, 2])]
+
+    assert actual.get_index("x").equals(pd.Index([20, 30]))
+    np.testing.assert_array_equal(actual.data.get(), expected_data.get())
+
+
+def test_isel_cupy_multidimensional_vector_indexer_with_pandas_index() -> None:
+    da = xr.DataArray(
+        cp.arange(4),
+        dims="x",
+        coords={"x": [10, 20, 30, 40]},
+    )
+    indexer = xr.DataArray(cp.array([[0, 2], [3, 1]]), dims=("row", "col"))
+
+    actual = da.isel(x=indexer)
+
+    assert actual.dims == ("row", "col")
+    np.testing.assert_array_equal(actual.data.get(), [[0, 2], [3, 1]])
+    np.testing.assert_array_equal(actual.coords["x"].data, [[10, 30], [40, 20]])
+
+
+def test_isel_cupy_argmax_indexer_with_pandas_index() -> None:
+    da = xr.DataArray(
+        cp.arange(24).reshape(4, 6),
+        dims=("x", "y"),
+        coords={"x": [10, 20, 30, 40]},
+    )
+    indexer = da.argmax("x")
+
+    actual = da.isel(x=indexer)
+    expected_indexer = xr.DataArray(indexer.data.get(), dims=indexer.dims)
+    expected = da.isel(x=expected_indexer)
+
+    assert actual.dims == expected.dims
+    np.testing.assert_array_equal(actual.data.get(), expected.data.get())
+    np.testing.assert_array_equal(actual.coords["x"].data, expected.coords["x"].data)
+
+
+@pytest.mark.parametrize("indexer", [2, slice(1, 3)])
+def test_cupy_basic_indexer_assignment(indexer: int | slice) -> None:
+    data = cp.arange(12).reshape(4, 3)
+    da = xr.DataArray(data.copy(), dims=("x", "y"))
+    expected = data.copy()
+    expected[indexer] = -1
+
+    da[{"x": indexer}] = -1
+
+    np.testing.assert_array_equal(da.data.get(), expected.get())
+
+
+def test_cupy_loc_assignment() -> None:
+    data = cp.arange(12).reshape(4, 3)
+    da = xr.DataArray(data.copy(), dims=("x", "y"), coords={"x": [0, 1, 2, 3]})
+    expected = data.copy()
+    expected[2] = -1
+
+    da.loc[{"x": 2}] = -1
+
+    np.testing.assert_array_equal(da.data.get(), expected.get())
+
+
+@pytest.mark.parametrize("value", [True, False])
+def test_bool_zero_dim_cupy_dataarray(value: bool) -> None:
+    assert bool(xr.DataArray(cp.array(value))) is value
+
+
+@pytest.mark.parametrize(("method", "expected"), [("idxmax", 1), ("idxmin", 0)])
+def test_cupy_idxminmax_skipna(method: str, expected: int) -> None:
+    da = xr.DataArray(cp.array([1, 3, 2]), dims="x", coords={"x": [0, 1, 2]})
+
+    result = getattr(da, method)(skipna=True)
+
+    assert result.item() == expected

xarray/tests/test_dataarray.py

@@ -86,6 +86,16 @@ def setup(self):
             np.uint64
         )
 
+    def test_bool_pandas_extension_array(self) -> None:
+        false_value = DataArray(pd.array([False], dtype="boolean"), dims="x")
+        true_value = DataArray(pd.array([True], dtype="boolean"), dims="x")
+        missing_value = DataArray(pd.array([pd.NA], dtype="boolean"), dims="x")
+
+        assert not bool(false_value)
+        assert bool(true_value)
+        with pytest.raises(TypeError, match="boolean value of NA is ambiguous"):
+            bool(missing_value)
+
     def test_repr(self) -> None:
         v = Variable(["time", "x"], [[1, 2, 3], [4, 5, 6]], {"foo": "bar"})
         v = v.astype(np.uint64)