Avoid calling np.asarray on lazy indexing classes

doc/whats-new.rst

@@ -47,6 +47,10 @@ 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>`_.
 
 
 .. _whats-new.2023.02.0:
@@ -210,9 +214,6 @@ Documentation
 - Add example of reading and writing individual groups to a single netCDF file to I/O docs page. (:pull:`7338`)
  By `Tom Nicholas <https://github.com/TomNicholas>`_.
 
-Internal Changes
-~~~~~~~~~~~~~~~~
-
 
 .. _whats-new.2022.11.0:
 

xarray/backends/common.py

@@ -84,9 +84,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=None):
  key = indexing.BasicIndexer((slice(None),) * self.ndim)
- return np.asarray(self[key], dtype=dtype)
+ return self[key]
 
 
 class AbstractDataStore:

xarray/backends/pydap_.py

@@ -46,6 +46,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

@@ -437,13 +437,25 @@ class ExplicitlyIndexed:
 
  __slots__ = ()
 
+ def __array__(self, dtype=None):
+ # 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=None):
+ # 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):
@@ -456,7 +468,10 @@ def __init__(self, array, indexer_cls=BasicIndexer):
  self.indexer_cls = indexer_cls
 
  def __array__(self, dtype=None):
- return np.asarray(self.array, dtype=dtype)
+ 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)
@@ -519,9 +534,16 @@ def shape(self) -> tuple[int, ...]:
  shape.append(k.size)
  return tuple(shape)
 
- def __array__(self, dtype=None):
+ def get_duck_array(self):
  array = as_indexable(self.array)
- return np.asarray(array[self.key], dtype=None)
+ array = array[self.key]
+ # 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)
@@ -572,8 +594,11 @@ 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]
+ 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)
@@ -619,8 +644,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]))
@@ -646,12 +671,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=None):
+ return np.asarray(self.get_duck_array(), dtype=dtype)
+
+ 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]))
@@ -815,7 +842,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
 
 
@@ -1425,6 +1452,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),)
@@ -1565,9 +1595,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

@@ -421,6 +421,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
 
@@ -528,6 +530,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

@@ -138,12 +138,37 @@ class UnexpectedDataAccess(Exception):
 
 
 class InaccessibleArray(utils.NDArrayMixin, ExplicitlyIndexed):
+ """Disallows any loading."""
+
  def __init__(self, array):
  self.array = array
 
+ def get_duck_array(self):
+ raise UnexpectedDataAccess("Tried accessing data")
+
+ def __array__(self, dtype=None):
+ raise UnexpectedDataAccess("Tried accessing data")
+
+ def __getitem__(self, key):
+ raise UnexpectedDataAccess("Tried accessing data")
+
+
+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=None):
  raise UnexpectedDataAccess("Tried accessing data")
 
+ def __array_namespace__(self):
+ """Present to satisfy is_duck_array test."""
+
 
 class ReturnItem:
  def __getitem__(self, key):

xarray/tests/test_dataset.py

@@ -36,6 +36,7 @@
 from xarray.core.pycompat import array_type, integer_types
 from xarray.core.utils import is_scalar
 from xarray.tests import (
+ DuckArrayWrapper,
  InaccessibleArray,
  UnexpectedDataAccess,
  assert_allclose,
@@ -202,6 +203,10 @@ def create_test_stacked_array() -> tuple[DataArray, DataArray]:
 
 
 class InaccessibleVariableDataStore(backends.InMemoryDataStore):
+ """
+ Store that does not allow any data access.
+ """
+
  def __init__(self):
  super().__init__()
  self._indexvars = set()
@@ -222,6 +227,47 @@ def lazy_inaccessible(k, v):
  return {k: lazy_inaccessible(k, v) for k, v in self._variables.items()}
 
 
+class DuckBackendArrayWrapper(backends.common.BackendArray):
+ """Mimic a BackendArray wrapper around DuckArrayWrapper"""
+
+ def __init__(self, array):
+ self.array = DuckArrayWrapper(array)
+ self.shape = array.shape
+ self.dtype = array.dtype
+
+ def get_array(self):
+ return self.array
+
+ def __getitem__(self, key):
+ return self.array[key.tuple]
+
+
+class AccessibleAsDuckArrayDataStore(backends.InMemoryDataStore):
+ """
+ Store that returns a duck array, not convertible to numpy array,
+ on read. Modeled after nVIDIA's kvikio.
+ """
+
+ def __init__(self):
+ super().__init__()
+ self._indexvars = set()
+
+ def store(self, variables, *args, **kwargs) -> None:
+ super().store(variables, *args, **kwargs)
+ for k, v in variables.items():
+ if isinstance(v, IndexVariable):
+ self._indexvars.add(k)
+
+ def get_variables(self):
+ def lazy_accessible(k, v):
+ if k in self._indexvars:
+ return v
+ data = indexing.LazilyIndexedArray(DuckBackendArrayWrapper(v.values))
+ return Variable(v.dims, data, v.attrs)
+
+ return {k: lazy_accessible(k, v) for k, v in self._variables.items()}
+
+
 class TestDataset:
  def test_repr(self) -> None:
  data = create_test_data(seed=123)
@@ -4678,6 +4724,29 @@ def test_lazy_load(self) -> None:
  ds.isel(time=10)
  ds.isel(time=slice(10), dim1=[0]).isel(dim1=0, dim2=-1)
 
+ def test_lazy_load_duck_array(self) -> None:
+ store = AccessibleAsDuckArrayDataStore()
+ create_test_data().dump_to_store(store)
+
+ for decode_cf in [True, False]:
+ ds = open_dataset(store, decode_cf=decode_cf)
+ with pytest.raises(UnexpectedDataAccess):
+ ds["var1"].values
+
+ # these should not raise UnexpectedDataAccess:
+ ds.var1.data
+ ds.isel(time=10)
+ ds.isel(time=slice(10), dim1=[0]).isel(dim1=0, dim2=-1)
+ repr(ds)
+
+ # preserve the duck array type and don't cast to array
+ assert isinstance(ds["var1"].load().data, DuckArrayWrapper)
+ assert isinstance(
+ ds["var1"].isel(dim2=0, dim1=0).load().data, DuckArrayWrapper
+ )
+
+ ds.close()
+
  def test_dropna(self) -> None:
  x = np.random.randn(4, 4)
  x[::2, 0] = np.nan