Clean up backend indexing

Author: dcherian

xarray/backends/memory.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 from xarray.backends.common import AbstractWritableDataStore
+from xarray.core import indexing
 from xarray.core.variable import Variable
 
 
@@ -24,7 +25,12 @@ def get_attrs(self):
         return self._attributes
 
     def get_variables(self):
-        return self._variables
+        res = {}
+        for k, v in self._variables.items():
+            v = v.copy(deep=True)
+            res[k] = v
+            v._data = indexing.LazilyIndexedArray(v._data)
+        return res
 
     def get_dimensions(self):
         return {d: s for v in self._variables.values() for d, s in v.dims.items()}

xarray/backends/scipy_.py

@@ -190,7 +190,7 @@ def ds(self):
     def open_store_variable(self, name, var):
         return Variable(
             var.dimensions,
-            ScipyArrayWrapper(name, self),
+            indexing.LazilyIndexedArray(ScipyArrayWrapper(name, self)),
             _decode_attrs(var._attributes),
         )
 

xarray/conventions.py

@@ -18,7 +18,7 @@
 )
 from xarray.core.utils import emit_user_level_warning
 from xarray.core.variable import IndexVariable, Variable
-from xarray.namedarray.utils import is_duck_dask_array
+from xarray.namedarray.utils import is_duck_array
 
 CF_RELATED_DATA = (
     "bounds",
@@ -249,7 +249,20 @@ def decode_cf_variable(
 
     encoding.setdefault("dtype", original_dtype)
 
-    if not is_duck_dask_array(data):
+    if (
+        # we don't need to lazily index duck arrays
+        not is_duck_array(data)
+        # These arrays already support lazy indexing
+        # OR for IndexingAdapters, it makes no sense to wrap them
+        and not isinstance(data, indexing.ExplicitlyIndexedNDArrayMixin)
+    ):
+        # this path applies to bare BackendArray objects.
+        # It is not hit for any internal Xarray backend
+        emit_user_level_warning(
+            "The backend you are using has not explicitly supported lazy indexing with Xarray. "
+            "Please ask the backend developers to support lazy loading by wrapping with LazilyIndexedArray. "
+            "See https://docs.xarray.dev/en/stable/internals/how-to-add-new-backend.html#how-to-support-lazy-loading for more."
+        )
         data = indexing.LazilyIndexedArray(data)
 
     return Variable(dimensions, data, attributes, encoding=encoding, fastpath=True)

xarray/core/indexes.py

@@ -717,7 +717,7 @@ def from_variables(
 
         # preserve wrapped pd.Index (if any)
         # accessing `.data` can load data from disk, so we only access if needed
-        data = var._data.array if hasattr(var._data, "array") else var.data
+        data = var._data if isinstance(var._data, PandasIndexingAdapter) else var.data
         # multi-index level variable: get level index
         if isinstance(var._data, PandasMultiIndexingAdapter):
             level = var._data.level

xarray/core/indexing.py

@@ -658,17 +658,11 @@ def shape(self) -> _Shape:
     def get_duck_array(self):
         if isinstance(self.array, ExplicitlyIndexedNDArrayMixin):
             array = apply_indexer(self.array, self.key)
+            array = array.get_duck_array()
         else:
             # If the array is not an ExplicitlyIndexedNDArrayMixin,
-            # it may wrap a BackendArray so use its __getitem__
+            # it is a BackendArray so use its __getitem__
             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()
         return _wrap_numpy_scalars(array)
 
     def transpose(self, order):
@@ -734,16 +728,11 @@ def shape(self) -> _Shape:
     def get_duck_array(self):
         if isinstance(self.array, ExplicitlyIndexedNDArrayMixin):
             array = apply_indexer(self.array, self.key)
+            array = array.get_duck_array()
         else:
             # If the array is not an ExplicitlyIndexedNDArrayMixin,
-            # it may wrap a BackendArray so use its __getitem__
+            # it is a BackendArray so use its __getitem__
             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()
         return _wrap_numpy_scalars(array)
 
     def _updated_key(self, new_key: ExplicitIndexer):