Add IndexingAdapter mixin

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/coding/variables.py

@@ -105,6 +105,7 @@ def __getitem__(self, key) -> np.ndarray:
         return np.asarray(self.array[key], dtype=self.dtype)
 
 
+
 def _apply_mask(
     data: np.ndarray,
     encoded_fill_values: list,

xarray/core/indexing.py

@@ -516,16 +516,30 @@ def get_duck_array(self):
         return self.array
 
 
-class ExplicitlyIndexedNDArrayMixin(NDArrayMixin, ExplicitlyIndexed):
-    __slots__ = ()
+class IndexingAdapter:
+    """Marker class for indexing adapters.
+
+    These classes translate between Xarray's indexing semantics and the underlying array's
+    indexing semantics.
+    """
 
     def get_duck_array(self):
         key = BasicIndexer((slice(None),) * self.ndim)
         return self[key]
 
     async def async_get_duck_array(self):
-        key = BasicIndexer((slice(None),) * self.ndim)
-        return await self.async_getitem(key)
+        """These classes are applied to in-memory arrays, so specific async support isn't needed."""
+        return self.get_duck_array()
+
+
+class ExplicitlyIndexedNDArrayMixin(NDArrayMixin, ExplicitlyIndexed):
+    __slots__ = ()
+
+    def get_duck_array(self):
+        raise NotImplementedError
+
+    async def async_get_duck_array(self):
+        raise NotImplementedError
 
     def _oindex_get(self, indexer: OuterIndexer):
         raise NotImplementedError(
@@ -650,37 +664,25 @@ 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):
-            array = array.get_duck_array()
+        if isinstance(self.array, BackendArray):
+            array = self.array[self.key]
+        else:
+            array = apply_indexer(self.array, self.key)
+            if isinstance(array, ExplicitlyIndexed):
+                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)
+            if isinstance(array, ExplicitlyIndexed):
+                array = await array.async_get_duck_array()
         return _wrap_numpy_scalars(array)
 
     def transpose(self, order):
@@ -744,36 +746,26 @@ 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):
-            array = array.get_duck_array()
+        else:
+            array = apply_indexer(self.array, self.key)
+            if isinstance(array, ExplicitlyIndexed):
+                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)
+            if isinstance(array, ExplicitlyIndexed):
+                array = await array.async_get_duck_array()
         return _wrap_numpy_scalars(array)
 
     def _updated_key(self, new_key: ExplicitIndexer):
@@ -1589,7 +1581,7 @@ def is_fancy_indexer(indexer: Any) -> bool:
     return True
 
 
-class NumpyIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
+class NumpyIndexingAdapter(IndexingAdapter, ExplicitlyIndexedNDArrayMixin):
     """Wrap a NumPy array to use explicit indexing."""
 
     __slots__ = ("array",)
@@ -1668,7 +1660,7 @@ def __init__(self, array):
         self.array = array
 
 
-class ArrayApiIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
+class ArrayApiIndexingAdapter(IndexingAdapter, ExplicitlyIndexedNDArrayMixin):
     """Wrap an array API array to use explicit indexing."""
 
     __slots__ = ("array",)
@@ -1733,7 +1725,7 @@ def _assert_not_chunked_indexer(idxr: tuple[Any, ...]) -> None:
         )
 
 
-class DaskIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
+class DaskIndexingAdapter(IndexingAdapter, ExplicitlyIndexedNDArrayMixin):
     """Wrap a dask array to support explicit indexing."""
 
     __slots__ = ("array",)
@@ -1809,7 +1801,7 @@ def transpose(self, order):
         return self.array.transpose(order)
 
 
-class PandasIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
+class PandasIndexingAdapter(IndexingAdapter, ExplicitlyIndexedNDArrayMixin):
     """Wrap a pandas.Index to preserve dtypes and handle explicit indexing."""
 
     __slots__ = ("_dtype", "array")
@@ -1872,15 +1864,6 @@ 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),)
@@ -2135,7 +2118,9 @@ def copy(self, deep: bool = True) -> Self:
         return type(self)(array, self._dtype, self.level)
 
 
-class CoordinateTransformIndexingAdapter(ExplicitlyIndexedNDArrayMixin):
+class CoordinateTransformIndexingAdapter(
+    IndexingAdapter, ExplicitlyIndexedNDArrayMixin
+):
     """Wrap a CoordinateTransform as a lazy coordinate array.
 
     Supports explicit indexing (both outer and vectorized).

xarray/namedarray/pycompat.py

@@ -153,20 +153,15 @@ async def async_to_duck_array(
     from xarray.core.indexing import (
         ExplicitlyIndexed,
         ImplicitToExplicitIndexingAdapter,
+        IndexingAdapter,
     )
-    from xarray.namedarray.parallelcompat import get_chunked_array_type
 
     print(type(data))
-
-    if is_chunked_array(data):
-        chunkmanager = get_chunked_array_type(data)
-        loaded_data, *_ = chunkmanager.compute(data, **kwargs)  # type: ignore[var-annotated]
-        return loaded_data
-
-    if isinstance(data, ExplicitlyIndexed | ImplicitToExplicitIndexingAdapter):
+    if isinstance(data, IndexingAdapter):
+        # These wrap in-memory arrays, and async isn't needed
+        return data.get_duck_array()
+    elif isinstance(data, ExplicitlyIndexed | ImplicitToExplicitIndexingAdapter):
         print("async inside to_duck_array")
         return await data.async_get_duck_array()  # type: ignore[no-untyped-call, no-any-return]
-    elif is_duck_array(data):
-        return data
     else:
-        return np.asarray(data)  # type: ignore[return-value]
+        return to_duck_array(data, **kwargs)