Better impl

Author: dcherian

xarray/backends/zarr.py

@@ -242,6 +242,7 @@ def __getitem__(self, key):
         # could possibly have a work-around for 0d data here
 
     async def async_getitem(self, key):
+        print("async getting")
         array = self._array
         if isinstance(key, indexing.BasicIndexer):
             method = self._async_getitem

xarray/core/indexing.py

@@ -523,11 +523,16 @@ class IndexingAdapter:
     indexing semantics.
     """
 
+    # FIXME: @abstractmethod here
+    def get_duck_array():
+        raise NotImplementedError
+
     async def async_get_duck_array():
-        """These classes are applied to in-memory arrays, so async get isn't needed."""
-        raise RuntimeError(
-            f"{type(self).__name__} async_get_duck_array was called. This should not happen."
-        )
+        """These classes are applied to in-memory arrays, so specific async support isn't needed."""
+        return self.get_duck_array()
+        # raise RuntimeError(
+        #     f"{type(self).__name__} async_get_duck_array was called. This should not happen."
+        # )
 
 
 class ExplicitlyIndexedNDArrayMixin(NDArrayMixin, ExplicitlyIndexed):
@@ -664,37 +669,23 @@ def shape(self) -> _Shape:
         return self._shape
 
     def get_duck_array(self):
-        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 = self.array[self.key]
+        from xarray.backends.common import BackendArray
 
-        # 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):
+        if isinstance(self.array, BackendArray):
+            array = self.array[self.key]
+        else:
+            array = apply_indexer(self.array, self.key)
             array = array.get_duck_array()
         return _wrap_numpy_scalars(array)
 
     async def async_get_duck_array(self):
-        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 (async) getitem
-            array = await self.array.async_getitem(self.key)
+        from xarray.backends.common import BackendArray
 
-        # 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):
-            # At this point, we have issued completed the possible async load from disk
-            # and array is in-memory. So use the sync get
-            array = array.get_duck_array()
+        if isinstance(self.array, BackendArray):
+            array = await self.array.async_getitem(self.key)
+        else:
+            array = apply_indexer(self.array, self.key)
+            array = await array.async_get_duck_array()
         return _wrap_numpy_scalars(array)
 
     def transpose(self, order):
@@ -758,36 +749,24 @@ def shape(self) -> _Shape:
         return np.broadcast(*self.key.tuple).shape
 
     def get_duck_array(self):
-        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__
+        from xarray.backends.common import BackendArray
+
+        if isinstance(self.array, BackendArray):
             array = self.array[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):
+        else:
+            array = apply_indexer(self.array, self.key)
             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__
+        from xarray.backends.common import BackendArray
+
+        if isinstance(self.array, BackendArray):
             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):
-            # At this point, we have issued completed the possible async load from disk
-            # and array is in-memory. So use the sync get
-            array = array.get_duck_array()
+        else:
+            array = apply_indexer(self.array, self.key)
+            array = await array.async_get_duck_array()
         return _wrap_numpy_scalars(array)
 
     def _updated_key(self, new_key: ExplicitIndexer):