implement async vectorized indexing

Author: TomNicholas

xarray/backends/zarr.py

@@ -222,6 +222,10 @@ async def _async_oindex(self, key):
         async_array = self._array._async_array
         return await async_array.oindex.getitem(key)
 
+    async def _async_vindex(self, key):
+        async_array = self._array._async_array
+        return await async_array.vindex.getitem(key)
+
     def __getitem__(self, key):
         array = self._array
         if isinstance(key, indexing.BasicIndexer):
@@ -242,14 +246,11 @@ async def async_getitem(self, key):
         if isinstance(key, indexing.BasicIndexer):
             method = self._async_getitem
         elif isinstance(key, indexing.VectorizedIndexer):
-            # method = self._vindex
-            raise NotImplementedError("async lazy vectorized indexing is not supported")
+            method = self._async_vindex
         elif isinstance(key, indexing.OuterIndexer):
             method = self._async_oindex
-
-        print("did an async get")
         return await indexing.async_explicit_indexing_adapter(
-            key, array.shape, indexing.IndexingSupport.OUTER, method
+            key, array.shape, indexing.IndexingSupport.VECTORIZED, method
         )
 
 

xarray/core/indexing.py

@@ -525,7 +525,7 @@ def get_duck_array(self):
 
     async def async_get_duck_array(self):
         key = BasicIndexer((slice(None),) * self.ndim)
-        return self[key]
+        return await self.async_getitem(key)
 
     def _oindex_get(self, indexer: OuterIndexer):
         raise NotImplementedError(
@@ -756,6 +756,22 @@ def get_duck_array(self):
             array = array.get_duck_array()
         return _wrap_numpy_scalars(array)
 
+    async def async_get_duck_array(self):
+        print("inside LazilyVectorizedIndexedArray.async_get_duck_array")
+        if isinstance(self.array, ExplicitlyIndexedNDArrayMixin):
+            array = apply_indexer(self.array, self.key)
+        else:
+            # If the array is not an ExplicitlyIndexedNDArrayMixin,
+            # it may wrap a BackendArray so use its __getitem__
+            array = await self.array.async_getitem(self.key)
+        # self.array[self.key] is now a numpy array when
+        # self.array is a BackendArray subclass
+        # and self.key is BasicIndexer((slice(None, None, None),))
+        # so we need the explicit check for ExplicitlyIndexed
+        if isinstance(array, ExplicitlyIndexed):
+            array = await array.async_get_duck_array()
+        return _wrap_numpy_scalars(array)
+
     def _updated_key(self, new_key: ExplicitIndexer):
         return _combine_indexers(self.key, self.shape, new_key)
 
@@ -1608,6 +1624,16 @@ def __getitem__(self, indexer: ExplicitIndexer):
         key = indexer.tuple + (Ellipsis,)
         return array[key]
 
+    async def async_getitem(self, indexer: ExplicitIndexer):
+        self._check_and_raise_if_non_basic_indexer(indexer)
+
+        array = self.array
+        # We want 0d slices rather than scalars. This is achieved by
+        # appending an ellipsis (see
+        # https://numpy.org/doc/stable/reference/arrays.indexing.html#detailed-notes).
+        key = indexer.tuple + (Ellipsis,)
+        return array[key]
+
     def _safe_setitem(self, array, key: tuple[Any, ...], value: Any) -> None:
         try:
             array[key] = value
@@ -1855,6 +1881,15 @@ def get_duck_array(self) -> np.ndarray | PandasExtensionArray:
             return PandasExtensionArray(self.array.array)
         return np.asarray(self)
 
+    async def async_get_duck_array(self) -> np.ndarray | PandasExtensionArray:
+        # TODO this must surely be wrong - it's not async yet
+        print("in PandasIndexingAdapter")
+        if pd.api.types.is_extension_array_dtype(self.array):
+            from xarray.core.extension_array import PandasExtensionArray
+
+            return PandasExtensionArray(self.array.array)
+        return np.asarray(self)
+
     @property
     def shape(self) -> _Shape:
         return (len(self.array),)

xarray/core/variable.py

@@ -959,7 +959,6 @@ def load(self, **kwargs):
         return self
 
     async def load_async(self, **kwargs):
-        print("async inside Variable")
         self._data = await async_to_duck_array(self._data, **kwargs)
         return self
 

xarray/tests/test_async.py

@@ -184,7 +184,7 @@ async def test_concurrent_load_multiple_objects(self, xr_obj) -> None:
             ("sel", {"x": 2}),
             ("sel", {"x": [2, 3]}),
             (
-                "isel",
+                "sel",
                 {
                     "x": xr.DataArray([2, 3], dims="points"),
                     "y": xr.DataArray([2, 3], dims="points"),
@@ -198,6 +198,7 @@ async def test_indexing(self, memorystore, method, indexer) -> None:
         latencystore = LatencyStore(memorystore, latency=0.0)
         ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)
 
+        # TODO we're not actually testing that these indexing methods are not blocking...
         result = await getattr(ds, method)(**indexer).load_async()
         expected = getattr(ds, method)(**indexer).load()
         xrt.assert_identical(result, expected)