Generalize lazy backend indexing a little more (#10078)

doc/whats-new.rst

@@ -30,6 +30,8 @@ New Features
   By `Justus Magin <https://github.com/keewis>`_.
 - Added experimental support for coordinate transforms (not ready for public use yet!) (:pull:`9543`)
   By `Benoit Bovy <https://github.com/benbovy>`_.
+- Support reading to `GPU memory with Zarr <https://zarr.readthedocs.io/en/stable/user-guide/gpu.html>`_ (:pull:`10078`).
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/core/indexes.py

@@ -442,6 +442,7 @@ def safe_cast_to_index(array: Any) -> pd.Index:
     """
     from xarray.core.dataarray import DataArray
     from xarray.core.variable import Variable
+    from xarray.namedarray.pycompat import to_numpy
 
     if isinstance(array, pd.Index):
         index = array
@@ -468,7 +469,7 @@ def safe_cast_to_index(array: Any) -> pd.Index:
                 )
                 kwargs["dtype"] = "float64"
 
-        index = pd.Index(np.asarray(array), **kwargs)
+        index = pd.Index(to_numpy(array), **kwargs)
 
     return _maybe_cast_to_cftimeindex(index)
 

xarray/core/indexing.py

@@ -1013,8 +1013,8 @@ def explicit_indexing_adapter(
     raw_key, numpy_indices = decompose_indexer(key, shape, indexing_support)
     result = raw_indexing_method(raw_key.tuple)
     if numpy_indices.tuple:
-        # index the loaded np.ndarray
-        indexable = NumpyIndexingAdapter(result)
+        # index the loaded duck array
+        indexable = as_indexable(result)
         result = apply_indexer(indexable, numpy_indices)
     return result
 

xarray/tests/test_indexing.py

@@ -19,6 +19,7 @@
     raise_if_dask_computes,
     requires_dask,
 )
+from xarray.tests.arrays import DuckArrayWrapper
 
 B = IndexerMaker(indexing.BasicIndexer)
 
@@ -1052,3 +1053,15 @@ def test_advanced_indexing_dask_array() -> None:
     with raise_if_dask_computes():
         actual = ds.b.sel(x=ds.a.data)
     assert_identical(expected, actual)
+
+
+def test_backend_indexing_non_numpy() -> None:
+    """This model indexing of a Zarr store that reads to GPU memory."""
+    array = DuckArrayWrapper(np.array([1, 2, 3]))
+    indexed = indexing.explicit_indexing_adapter(
+        indexing.BasicIndexer((slice(1),)),
+        shape=array.shape,
+        indexing_support=indexing.IndexingSupport.BASIC,
+        raw_indexing_method=array.__getitem__,
+    )
+    np.testing.assert_array_equal(indexed.array, np.array([1]))