Avoid calling np.asarray on lazy indexing classes (#6874)

* Add get_array to lazy indexing array types.

This returns the underlying array type instead of always casting
to np.array. This is necessary for Zarr stores where the
Zarr Array wraps a cupy array (for example kvikio.zarr.GDSStoree).
In that case, we cannot call np.asarray because __array__ is
expected to always return a numpy array.

We use get_array in Variable.data to make sure we don't load arrays
from such GDSStores.

* Rename to short_array_repr; use Variable.data

instead of always casting to np.asarray

* Fix Variable.load

* Make get_array recursive.

* Some cleanups

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add get_array to PandasIndexingAdaptor

* Finish short_array_repr refactoring

* Rename to get_duck_array

* Try without hasattr check

* Return bare array from LazilyIndexedArray.get_duck_array

* Add get_duck_array to AbstractArray

Clean up short_array_repr.

* Fix zerodim test

* Fix LazilyVectorizedIndexedArray

* Inherit __array__ from ExplicitlyIndexed

* Fix InaccessibleArray in tests

* Fix BackendArray

* reprs Use .data on AbstractArray

* Force netCDF and h5netCDF to return arrays

* Add whats-new

* Add comments; review feedback

* Fix merge.

* Remove another np.asarray

* Avoid np.asarray on __getitem__.

for BoolTypeArray, NativeEndiannessArray

* [WIP] ExplicitlyIndexedBackendArray

* Handle Indexing Adapter classes explicitly.

* Revert "Handle Indexing Adapter classes explicitly."

This reverts commit 46d98ec.

* Revert "[WIP] ExplicitlyIndexedBackendArray"

This reverts commit 9b727e6.

* Fix pydap now that NumpyIndexingAdapter does not automatically cast to array

* Update xarray/backends/pydap_.py

* Add test store.

* [skip-ci] Update whats-new

* Fix test

* fix mypy?

* Fix Zarr test

* test the repr too

* Guard np.asarray for scalars.

* Revert casting to arrays in backend

* Wrap numpy scalars in Explicitly*Indexed*.get_duck_aray

* Apply suggestions from code review

Co-authored-by: Illviljan <[email protected]>
Co-authored-by: Stephan Hoyer <[email protected]>

* Update xarray/tests/__init__.py

Co-authored-by: Illviljan <[email protected]>

* Update xarray/core/indexing.py

* Apply suggestions from code review

Co-authored-by: Illviljan <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Bring back the ugly check

* Update whats-new

* Fix pre-commit

* silence mypy error

* minimize diff

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Illviljan <[email protected]>
Co-authored-by: Stephan Hoyer <[email protected]>

Author: 4 people

doc/whats-new.rst

@@ -57,6 +57,11 @@ Documentation
 Internal Changes
 ~~~~~~~~~~~~~~~~
 
+- Don't assume that arrays read from disk will be Numpy arrays. This is a step toward
+  enabling reads from a Zarr store using the `Kvikio <https://docs.rapids.ai/api/kvikio/stable/api.html#zarr>`_
+  or `TensorStore <https://google.github.io/tensorstore/>`_ libraries.
+  (:pull:`6874`). By `Deepak Cherian <https://github.com/dcherian>`_.
+
 - Remove internal support for reading GRIB files through the ``cfgrib`` backend. ``cfgrib`` now uses the external
   backend interface, so no existing code should break.
   By `Deepak Cherian <https://github.com/dcherian>`_.

xarray/backends/common.py

@@ -86,9 +86,9 @@ def robust_getitem(array, key, catch=Exception, max_retries=6, initial_delay=500
 class BackendArray(NdimSizeLenMixin, indexing.ExplicitlyIndexed):
     __slots__ = ()
 
-    def __array__(self, dtype=None):
+    def get_duck_array(self, dtype: np.typing.DTypeLike = None):
         key = indexing.BasicIndexer((slice(None),) * self.ndim)
-        return np.asarray(self[key], dtype=dtype)
+        return self[key]  # type: ignore [index]
 
 
 class AbstractDataStore:

xarray/backends/pydap_.py

@@ -54,6 +54,7 @@ def _getitem(self, key):
         # downloading coordinate data twice
         array = getattr(self.array, "array", self.array)
         result = robust_getitem(array, key, catch=ValueError)
+        result = np.asarray(result)
         # in some cases, pydap doesn't squeeze axes automatically like numpy
         axis = tuple(n for n, k in enumerate(key) if isinstance(k, integer_types))
         if result.ndim + len(axis) != array.ndim and axis:

xarray/coding/variables.py

@@ -69,8 +69,8 @@ def dtype(self) -> np.dtype:
     def __getitem__(self, key):
         return type(self)(self.array[key], self.func, self.dtype)
 
-    def __array__(self, dtype=None):
-        return self.func(self.array)
+    def get_duck_array(self):
+        return self.func(self.array.get_duck_array())
 
     def __repr__(self) -> str:
         return "{}({!r}, func={!r}, dtype={!r})".format(
@@ -224,7 +224,7 @@ def decode(self, variable: Variable, name: T_Name = None):
 
 
 def _scale_offset_decoding(data, scale_factor, add_offset, dtype: np.typing.DTypeLike):
-    data = np.array(data, dtype=dtype, copy=True)
+    data = data.astype(dtype=dtype, copy=True)
     if scale_factor is not None:
         data *= scale_factor
     if add_offset is not None:

xarray/core/formatting.py

@@ -16,7 +16,7 @@
 from pandas.errors import OutOfBoundsDatetime
 
 from xarray.core.duck_array_ops import array_equiv
-from xarray.core.indexing import MemoryCachedArray
+from xarray.core.indexing import ExplicitlyIndexed, MemoryCachedArray
 from xarray.core.options import OPTIONS, _get_boolean_with_default
 from xarray.core.pycompat import array_type
 from xarray.core.utils import is_duck_array
@@ -557,8 +557,15 @@ def limit_lines(string: str, *, limit: int):
     return string
 
 
-def short_numpy_repr(array):
-    array = np.asarray(array)
+def short_array_repr(array):
+    from xarray.core.common import AbstractArray
+
+    if isinstance(array, ExplicitlyIndexed):
+        array = array.get_duck_array()
+    elif isinstance(array, AbstractArray):
+        array = array.data
+    if not is_duck_array(array):
+        array = np.asarray(array)
 
     # default to lower precision so a full (abbreviated) line can fit on
     # one line with the default display_width
@@ -582,11 +589,11 @@ def short_data_repr(array):
     """Format "data" for DataArray and Variable."""
     internal_data = getattr(array, "variable", array)._data
     if isinstance(array, np.ndarray):
-        return short_numpy_repr(array)
+        return short_array_repr(array)
     elif is_duck_array(internal_data):
         return limit_lines(repr(array.data), limit=40)
     elif array._in_memory:
-        return short_numpy_repr(array)
+        return short_array_repr(array)
     else:
         # internal xarray array type
         return f"[{array.size} values with dtype={array.dtype}]"
@@ -831,7 +838,7 @@ def diff_array_repr(a, b, compat):
         equiv = array_equiv
 
     if not equiv(a.data, b.data):
-        temp = [wrap_indent(short_numpy_repr(obj), start="    ") for obj in (a, b)]
+        temp = [wrap_indent(short_array_repr(obj), start="    ") for obj in (a, b)]
         diff_data_repr = [
             ab_side + "\n" + ab_data_repr
             for ab_side, ab_data_repr in zip(("L", "R"), temp)

xarray/core/indexing.py

@@ -449,13 +449,25 @@ class ExplicitlyIndexed:
 
     __slots__ = ()
 
+    def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray:
+        # Leave casting to an array up to the underlying array type.
+        return np.asarray(self.get_duck_array(), dtype=dtype)
+
+    def get_duck_array(self):
+        return self.array
+
 
 class ExplicitlyIndexedNDArrayMixin(NDArrayMixin, ExplicitlyIndexed):
     __slots__ = ()
 
-    def __array__(self, dtype=None):
+    def get_duck_array(self):
         key = BasicIndexer((slice(None),) * self.ndim)
-        return np.asarray(self[key], dtype=dtype)
+        return self[key]
+
+    def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray:
+        # This is necessary because we apply the indexing key in self.get_duck_array()
+        # Note this is the base class for all lazy indexing classes
+        return np.asarray(self.get_duck_array(), dtype=dtype)
 
 
 class ImplicitToExplicitIndexingAdapter(NDArrayMixin):
@@ -467,8 +479,11 @@ def __init__(self, array, indexer_cls=BasicIndexer):
         self.array = as_indexable(array)
         self.indexer_cls = indexer_cls
 
-    def __array__(self, dtype=None):
-        return np.asarray(self.array, dtype=dtype)
+    def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray:
+        return np.asarray(self.get_duck_array(), dtype=dtype)
+
+    def get_duck_array(self):
+        return self.array.get_duck_array()
 
     def __getitem__(self, key):
         key = expanded_indexer(key, self.ndim)
@@ -531,9 +546,15 @@ def shape(self) -> tuple[int, ...]:
                 shape.append(k.size)
         return tuple(shape)
 
-    def __array__(self, dtype=None):
-        array = as_indexable(self.array)
-        return np.asarray(array[self.key], dtype=None)
+    def get_duck_array(self):
+        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):
         return LazilyVectorizedIndexedArray(self.array, self.key).transpose(order)
@@ -584,8 +605,15 @@ def __init__(self, array, key):
     def shape(self) -> tuple[int, ...]:
         return np.broadcast(*self.key.tuple).shape
 
-    def __array__(self, dtype=None):
-        return np.asarray(self.array[self.key], dtype=None)
+    def get_duck_array(self):
+        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):
         return _combine_indexers(self.key, self.shape, new_key)
@@ -631,8 +659,8 @@ def _ensure_copied(self):
             self.array = as_indexable(np.array(self.array))
             self._copied = True
 
-    def __array__(self, dtype=None):
-        return np.asarray(self.array, dtype=dtype)
+    def get_duck_array(self):
+        return self.array.get_duck_array()
 
     def __getitem__(self, key):
         return type(self)(_wrap_numpy_scalars(self.array[key]))
@@ -658,12 +686,14 @@ def __init__(self, array):
         self.array = _wrap_numpy_scalars(as_indexable(array))
 
     def _ensure_cached(self):
-        if not isinstance(self.array, NumpyIndexingAdapter):
-            self.array = NumpyIndexingAdapter(np.asarray(self.array))
+        self.array = as_indexable(self.array.get_duck_array())
+
+    def __array__(self, dtype: np.typing.DTypeLike = None) -> np.ndarray:
+        return np.asarray(self.get_duck_array(), dtype=dtype)
 
-    def __array__(self, dtype=None):
+    def get_duck_array(self):
         self._ensure_cached()
-        return np.asarray(self.array, dtype=dtype)
+        return self.array.get_duck_array()
 
     def __getitem__(self, key):
         return type(self)(_wrap_numpy_scalars(self.array[key]))
@@ -827,7 +857,7 @@ def explicit_indexing_adapter(
     result = raw_indexing_method(raw_key.tuple)
     if numpy_indices.tuple:
         # index the loaded np.ndarray
-        result = NumpyIndexingAdapter(np.asarray(result))[numpy_indices]
+        result = NumpyIndexingAdapter(result)[numpy_indices]
     return result
 
 
@@ -1463,6 +1493,9 @@ def __array__(self, dtype: DTypeLike = None) -> np.ndarray:
                 array = array.astype("object")
         return np.asarray(array.values, dtype=dtype)
 
+    def get_duck_array(self) -> np.ndarray:
+        return np.asarray(self)
+
     @property
     def shape(self) -> tuple[int, ...]:
         return (len(self.array),)
@@ -1603,9 +1636,9 @@ def _repr_inline_(self, max_width: int) -> str:
             return format_array_flat(self._get_array_subset(), max_width)
 
     def _repr_html_(self) -> str:
-        from xarray.core.formatting import short_numpy_repr
+        from xarray.core.formatting import short_array_repr
 
-        array_repr = short_numpy_repr(self._get_array_subset())
+        array_repr = short_array_repr(self._get_array_subset())
         return f"<pre>{escape(array_repr)}</pre>"
 
     def copy(self, deep: bool = True) -> PandasMultiIndexingAdapter:

xarray/core/variable.py

@@ -426,6 +426,8 @@ def data(self) -> Any:
         """
         if is_duck_array(self._data):
             return self._data
+        elif isinstance(self._data, indexing.ExplicitlyIndexed):
+            return self._data.get_duck_array()
         else:
             return self.values
 
@@ -533,6 +535,8 @@ def load(self, **kwargs):
         """
         if is_duck_dask_array(self._data):
             self._data = as_compatible_data(self._data.compute(**kwargs))
+        elif isinstance(self._data, indexing.ExplicitlyIndexed):
+            self._data = self._data.get_duck_array()
         elif not is_duck_array(self._data):
             self._data = np.asarray(self._data)
         return self

xarray/tests/__init__.py

@@ -139,13 +139,18 @@ class UnexpectedDataAccess(Exception):
 
 
 class InaccessibleArray(utils.NDArrayMixin, ExplicitlyIndexed):
+    """Disallows any loading."""
+
     def __init__(self, array):
         self.array = array
 
-    def __getitem__(self, key):
-        raise UnexpectedDataAccess("Tried accessing data.")
+    def get_duck_array(self):
+        raise UnexpectedDataAccess("Tried accessing data")
+
+    def __array__(self, dtype: np.typing.DTypeLike = None):
+        raise UnexpectedDataAccess("Tried accessing data")
 
-    def __array__(self):
+    def __getitem__(self, key):
         raise UnexpectedDataAccess("Tried accessing data.")
 
 
@@ -157,6 +162,23 @@ def __getitem__(self, key):
         return self.array[tuple_idxr]
 
 
+class DuckArrayWrapper(utils.NDArrayMixin):
+    """Array-like that prevents casting to array.
+    Modeled after cupy."""
+
+    def __init__(self, array: np.ndarray):
+        self.array = array
+
+    def __getitem__(self, key):
+        return type(self)(self.array[key])
+
+    def __array__(self, dtype: np.typing.DTypeLike = None):
+        raise UnexpectedDataAccess("Tried accessing data")
+
+    def __array_namespace__(self):
+        """Present to satisfy is_duck_array test."""
+
+
 class ReturnItem:
     def __getitem__(self, key):
         return key