From f9023c05fd5391c6a17ab6368de7d58d35cccc40 Mon Sep 17 00:00:00 2001 From: Chirag Jangir <184008125+chiragjangiir@users.noreply.github.com> Date: Thu, 29 Oct 2026 16:06:35 +0530 Subject: [PATCH 1/2] BUG: restore vectorized indexing for arrays implementing both protocols Arrays such as sparse implement both __array_function__ and __array_namespace__, so as_indexable wraps them in an ArrayApiIndexingAdapter, whose _vindex_get unconditionally raised TypeError even though the numpy protocol supports paired indexing. Fall back to NdArrayLikeIndexingAdapter when __array_function__ is available, mirroring the protocol fallback used elsewhere in as_indexable, and keep raising for array-API-only arrays. This restores vectorized indexing (e.g. DataArray.sel with DataArray indexers) on sparse-backed arrays (GH9694). --- doc/whats-new.rst | 6 +++++ xarray/core/indexing.py | 6 +++++ xarray/tests/test_indexing.py | 51 +++++++++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 85e8bd13d0c..af3047b2357 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -441,6 +441,12 @@ Bug Fixes engine raises a ``TypeError`` that points to ``engine="zarr"`` (:issue:`11221`, :pull:`11226`). By `Carson Jones `_. +- Vectorized (paired) indexing, e.g. :py:meth:`DataArray.sel` with + :py:class:`DataArray` indexers, works again for arrays that implement both + ``__array_function__`` and ``__array_namespace__``, such as ``sparse`` + arrays, which previously raised ``TypeError: Vectorized indexing is not + supported`` (:issue:`9694`, :pull:`11828`). + By `Chirag Jangir `_. Documentation diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index 204567b76a3..e21291d2b77 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -1855,6 +1855,12 @@ def _oindex_get(self, indexer: OuterIndexer): @override def _vindex_get(self, indexer: VectorizedIndexer): + # Vectorized (paired) indexing is not part of the Array API standard, + # but arrays implementing both protocols (e.g. sparse) can still + # handle it through ``__array_function__``, mirroring the fallback + # used elsewhere in ``as_indexable`` (GH9694). + if hasattr(self.array, "__array_function__"): + return NdArrayLikeIndexingAdapter(self.array)._vindex_get(indexer) raise TypeError("Vectorized indexing is not supported") @override diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py index 047bf795914..73e13206251 100644 --- a/xarray/tests/test_indexing.py +++ b/xarray/tests/test_indexing.py @@ -1163,6 +1163,57 @@ def test_as_indexable(array, expected_type): assert isinstance(actual, expected_type) +class ArrayWithNamespaceNdArray(np.ndarray): + """ndarray subclass implementing both the Array API and numpy protocols.""" + + def __array_namespace__(self, /, *, api_version=None): + raise NotImplementedError + + +def test_array_api_adapter_vindex_falls_back_to_array_function() -> None: + # GH9694: vectorized indexing is not part of the Array API standard, but + # arrays implementing both protocols should still support it + array = np.arange(6, dtype=np.int64).reshape(2, 3).view(ArrayWithNamespaceNdArray) + adapter = indexing.ArrayApiIndexingAdapter(array) + indexer = indexing.VectorizedIndexer((np.array([1, 0]), np.array([2, 1]))) + + actual = adapter.vindex[indexer] + + expected = np.array([5, 1]) + np.testing.assert_array_equal(np.asarray(actual), expected) + + +def test_array_api_adapter_vindex_without_array_function_raises() -> None: + # GH9694: without the numpy protocol there is no way to vectorize + adapter = indexing.ArrayApiIndexingAdapter(ArrayWithNamespace()) + indexer = indexing.VectorizedIndexer((np.array([0]),)) + + with pytest.raises(TypeError, match="Vectorized indexing is not supported"): + adapter.vindex[indexer] + + +@requires_sparse +def test_vectorized_indexing_sparse_array() -> None: + # GH9694: DataArray.sel with DataArray indexers (vectorized indexing) + # raised TypeError for sparse-backed arrays + import sparse + + data = np.arange(15, dtype=np.float64).reshape(3, 5) + da = DataArray( + sparse.COO.from_numpy(data), + dims=("x", "y"), + coords={"x": ["x0", "x1", "x2"], "y": ["y0", "y1", "y2", "y3", "y4"]}, + ) + newdim = {"newdim": ["nd0", "nd1", "nd2"]} + x_idx = DataArray(["x2", "x1", "x2"], coords=newdim, dims="newdim") + y_idx = DataArray(["y4", "y2", "y0"], coords=newdim, dims="newdim") + + actual = da.sel(x=x_idx, y=y_idx) + + assert isinstance(actual.data, sparse.COO) + np.testing.assert_array_equal(actual.data.todense(), [14.0, 7.0, 10.0]) + + @requires_sparse @pytest.mark.parametrize( "key", From f2385dcc58b23b8e9eb2b4b9b5ada6abd657b6fc Mon Sep 17 00:00:00 2001 From: Chirag Jangir <184008125+chiragjangiir@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:54:30 +0530 Subject: [PATCH 2/2] BUG: support Array API integer-array indexing in ArrayApiIndexingAdapter._vindex_get (#11828) - Keep __array_function__ fallback first for dual-protocol arrays such as sparse - Implement Array API 2024.12+ integer-array indexing for array-API-only arrays in ArrayApiIndexingAdapter._vindex_get, converting slices to arange arrays and broadcasting indexers - Keep raising TypeError for assignment in ArrayApiIndexingAdapter._vindex_set - Add test coverage with array-api-strict for 1D, 2D indexers, interleaved slices, and unsupported assignment --- xarray/core/indexing.py | 41 +++++++++++++++++++++++++++++++---- xarray/tests/test_indexing.py | 35 ++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py index e21291d2b77..29da88de823 100644 --- a/xarray/core/indexing.py +++ b/xarray/core/indexing.py @@ -1855,13 +1855,46 @@ def _oindex_get(self, indexer: OuterIndexer): @override def _vindex_get(self, indexer: VectorizedIndexer): - # Vectorized (paired) indexing is not part of the Array API standard, - # but arrays implementing both protocols (e.g. sparse) can still - # handle it through ``__array_function__``, mirroring the fallback + # Arrays implementing both protocols (e.g. sparse) handle paired + # indexing through ``__array_function__``, mirroring the fallback # used elsewhere in ``as_indexable`` (GH9694). if hasattr(self.array, "__array_function__"): return NdArrayLikeIndexingAdapter(self.array)._vindex_get(indexer) - raise TypeError("Vectorized indexing is not supported") + + # For array-API-only arrays, integer array indexing is supported since + # the 2024.12 standard. Slices are converted to arange arrays broadcasting along + # their axis, yielding the vectorized dimension order directly. + try: + xp = self.array.__array_namespace__() + except (AttributeError, TypeError): + xp = None + + if xp is None or not hasattr(xp, "arange"): + raise TypeError("Vectorized indexing is not supported") + + device = getattr(self.array, "device", None) + key = indexer.tuple + shape = np.broadcast_shapes(*(k.shape for k in key if not isinstance(k, slice))) + n_slices = sum(isinstance(k, slice) for k in key) + new_key = [] + i_slice = 0 + for size, k in zip(self.array.shape, key, strict=True): + if isinstance(k, slice): + k = xp.arange(*k.indices(size), device=device) + new_shape = ( + (1,) * (len(shape) + i_slice) + + k.shape + + (1,) * (n_slices - i_slice - 1) + ) + i_slice += 1 + else: + k = xp.broadcast_to(xp.asarray(k, device=device), shape) + new_shape = shape + (1,) * n_slices + new_key.append(xp.reshape(k, new_shape)) + try: + return self.array[tuple(new_key)] + except TypeError as exc: + raise TypeError("Vectorized indexing is not supported") from exc @override def __getitem__(self, indexer: ExplicitIndexer): diff --git a/xarray/tests/test_indexing.py b/xarray/tests/test_indexing.py index 73e13206251..4bf8b652459 100644 --- a/xarray/tests/test_indexing.py +++ b/xarray/tests/test_indexing.py @@ -20,6 +20,7 @@ dask_array_type, has_dask, raise_if_dask_computes, + requires_array_api_strict, requires_dask, requires_pandas_3, requires_sparse, @@ -1192,6 +1193,40 @@ def test_array_api_adapter_vindex_without_array_function_raises() -> None: adapter.vindex[indexer] +@requires_array_api_strict +def test_array_api_adapter_vindex_array_api_strict() -> None: + import array_api_strict as xp + + # 3D array: test 1D indexers, 2D indexers, and slices interleaved between indexed axes + x_np = np.arange(24, dtype=np.int64).reshape(2, 3, 4) + x_xp = xp.asarray(x_np) + adapter = indexing.ArrayApiIndexingAdapter(x_xp) + + # 1D indexers with interleaved slice + idxr_1d = indexing.VectorizedIndexer( + (np.array([0, 1]), slice(1, 3), np.array([3, 2])) + ) + actual_1d = adapter.vindex[idxr_1d] + expected_1d = indexing.NdArrayLikeIndexingAdapter(x_np)._vindex_get(idxr_1d) + np.testing.assert_array_equal(np.asarray(actual_1d), expected_1d) + + # 2D indexers + idxr_2d = indexing.VectorizedIndexer( + ( + np.array([[0, 1], [1, 0]]), + slice(0, 2), + np.array([[2, 3], [1, 0]]), + ) + ) + actual_2d = adapter.vindex[idxr_2d] + expected_2d = indexing.NdArrayLikeIndexingAdapter(x_np)._vindex_get(idxr_2d) + np.testing.assert_array_equal(np.asarray(actual_2d), expected_2d) + + # Vectorized assignment is unsupported + with pytest.raises(TypeError, match="Vectorized indexing is not supported"): + adapter.vindex[idxr_1d] = 0 + + @requires_sparse def test_vectorized_indexing_sparse_array() -> None: # GH9694: DataArray.sel with DataArray indexers (vectorized indexing)