From 7a4b8b4c69fb091926fac268477dfc511e693715 Mon Sep 17 00:00:00 2001 From: Huashuo Lei <139048865+wpyaoxueros@users.noreply.github.com> Date: Fri, 23 Oct 2026 11:22:30 +0000 Subject: [PATCH 1/4] Prepare CuPy indexing, assignment and scalar truth fixes Prepared with Codex assistance. Public PR authorship and required human review remain pending. Co-authored-by: Codex --- xarray/core/common.py | 5 +- xarray/core/indexes.py | 7 +++ xarray/core/indexing.py | 17 ++++++- xarray/core/variable.py | 1 + xarray/tests/test_cupy.py | 83 ++++++++++++++++++++++++++++++++++ xarray/tests/test_dataarray.py | 10 ++++ 6 files changed, 121 insertions(+), 2 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index fbd3ebc57a1..466075c8f0c 100644 --- a/xarray/core/common.py +++ b/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) diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py index 2a513c5c938..f88c054c4ee 100644 --- a/xarray/core/indexes.py +++ b/xarray/core/indexes.py @@ -26,12 +26,14 @@ get_valid_numpy_dtype, is_allowed_extension_array_dtype, is_dict_like, + is_duck_array, is_scalar, ) 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 +838,11 @@ def isel( if isinstance(indxr, slice) and indxr == slice(None): return self + if is_duck_array(indxr) and not isinstance(indxr, np.ndarray): + from xarray.namedarray.pycompat import to_numpy + + indxr = to_numpy(cast("duckarray[Any, Any]", indxr)) + return self._replace(self.index[indxr]) # type: ignore[index,unused-ignore] @override diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 2de92f9bcdb..113d0850678 100644 --- a/xarray/core/indexing.py +++ b/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 diff --git a/xarray/core/variable.py b/xarray/core/variable.py index 02ee047f1e4..d7c0b23747e 100644 --- a/xarray/core/variable.py +++ b/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) diff --git a/xarray/tests/test_cupy.py b/xarray/tests/test_cupy.py index cb8945e8ae5..c75315e993c 100644 --- a/xarray/tests/test_cupy.py +++ b/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 diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 407cd08c6cd..c384a231fa4 100644 --- a/xarray/tests/test_dataarray.py +++ b/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) From 27a3f22c033650a3c83864e2802c7cd14809a427 Mon Sep 17 00:00:00 2001 From: Huashuo Lei <139048865+wpyaoxueros@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:04:45 +0000 Subject: [PATCH 2/4] Document CuPy indexing fixes in release notes Co-authored-by: Codex --- doc/whats-new.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 21f3671188a..bcd23e4d477 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -1486,6 +1486,10 @@ New features 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 `_. - Fix error raised when writing scalar variables to Zarr with ``region={}`` (:pull:`10796`). By `Stephan Hoyer `_. From 6cbb3002adfdc591c6245cba299203905a76797a Mon Sep 17 00:00:00 2001 From: Huashuo Lei <139048865+wpyaoxueros@users.noreply.github.com> Date: Fri, 23 Oct 2026 15:13:28 +0000 Subject: [PATCH 3/4] Keep index conversion imports at module scope Co-authored-by: Codex --- xarray/core/indexes.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/xarray/core/indexes.py b/xarray/core/indexes.py index f88c054c4ee..3fd1fbdfd56 100644 --- a/xarray/core/indexes.py +++ b/xarray/core/indexes.py @@ -29,6 +29,7 @@ is_duck_array, is_scalar, ) +from xarray.namedarray.pycompat import to_numpy if TYPE_CHECKING: from xarray.core.types import ErrorOptions, JoinOptions, Self @@ -839,8 +840,6 @@ def isel( return self if is_duck_array(indxr) and not isinstance(indxr, np.ndarray): - from xarray.namedarray.pycompat import to_numpy - indxr = to_numpy(cast("duckarray[Any, Any]", indxr)) return self._replace(self.index[indxr]) # type: ignore[index,unused-ignore] From e74437a1c1658020f610c15ad3f8e29d9f38d484 Mon Sep 17 00:00:00 2001 From: Huashuo Lei <139048865+wpyaoxueros@users.noreply.github.com> Date: Fri, 23 Oct 2026 23:46:17 +0800 Subject: [PATCH 4/4] Place CuPy release note under the unreleased version Co-authored-by: Codex --- doc/whats-new.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index bcd23e4d477..76db46bcc75 100644 --- a/doc/whats-new.rst +++ b/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 `_. - 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 @@ -1486,10 +1490,6 @@ New features 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 `_. - Fix error raised when writing scalar variables to Zarr with ``region={}`` (:pull:`10796`). By `Stephan Hoyer `_.