support async orthogonal indexing via zarr-developers/zarr-python#3083

Author: TomNicholas

xarray/backends/zarr.py

@@ -218,6 +218,10 @@ async def _async_getitem(self, key):
         async_array = self._array._async_array
         return await async_array.getitem(key)
 
+    async def _async_oindex(self, key):
+        async_array = self._array._async_array
+        return await async_array.oindex.getitem(key)
+
     def __getitem__(self, key):
         array = self._array
         if isinstance(key, indexing.BasicIndexer):
@@ -241,12 +245,11 @@ async def async_getitem(self, key):
             # method = self._vindex
             raise NotImplementedError("async lazy vectorized indexing is not supported")
         elif isinstance(key, indexing.OuterIndexer):
-            # method = self._oindex
-            raise NotImplementedError("async lazy orthogonal indexing is not supported")
+            method = self._async_oindex
 
         print("did an async get")
         return await indexing.async_explicit_indexing_adapter(
-            key, array.shape, indexing.IndexingSupport.BASIC, method
+            key, array.shape, indexing.IndexingSupport.OUTER, method
         )
 
 

xarray/tests/test_async.py

@@ -25,6 +25,9 @@ class LatencyStore(WrapperStore[T_Store]):
 
         latency: float
 
+        # TODO only have to add this because of dumb behaviour in zarr where it raises with "ValueError: Store is not read-only but mode is 'r'"
+        read_only = True
+
         def __init__(
             self,
             store: T_Store,
@@ -175,7 +178,6 @@ async def test_concurrent_load_multiple_objects(self, xr_obj) -> None:
             total_time=timer.total_time, latency=self.LATENCY, n_loads=N_OBJECTS
         )
 
-    @pytest.mark.xfail(reason="not implemented")
     async def test_indexing(self, memorystore) -> None:
         latencystore = LatencyStore(memorystore, latency=self.LATENCY)
         ds = xr.open_zarr(latencystore, zarr_format=3, consolidated=False, chunks=None)