Vectorized lazy indexing

doc/whats-new.rst

@@ -43,6 +43,10 @@ Documentation
 
 Enhancements
 ~~~~~~~~~~~~
+- Support lazy vectorized-indexing. After this change, flexible indexing such
+  as orthogonal/vectorized indexing, becomes possible for all the backend
+  arrays. (:issue:`1897`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - reduce methods such as :py:func:`DataArray.sum()` now accepts ``dtype``
   arguments. (:issue:`1838`)
   By `Keisuke Fujii <https://github.com/fujiisoup>`_.

xarray/backends/h5netcdf_.py

@@ -17,15 +17,20 @@
 
 class H5NetCDFArrayWrapper(BaseNetCDF4Array):
     def __getitem__(self, key):
-        key = indexing.unwrap_explicit_indexer(
-            key, self, allow=(indexing.BasicIndexer, indexing.OuterIndexer))
+        key, np_inds = indexing.decompose_indexer(key, self.shape,
+                                                  mode='outer_1vector')
+
         # h5py requires using lists for fancy indexing:
         # https://github.com/h5py/h5py/issues/992
-        # OuterIndexer only holds 1D integer ndarrays, so it's safe to convert
-        # them to lists.
-        key = tuple(list(k) if isinstance(k, np.ndarray) else k for k in key)
+        key = tuple(list(k) if isinstance(k, np.ndarray) else k for k in
+                    key.tuple)
         with self.datastore.ensure_open(autoclose=True):
-            return self.get_array()[key]
+            array = np.asarray(self.get_array()[key], dtype=self.dtype)
+
+        for ind in np_inds:
+            array = indexing.NumpyIndexingAdapter(array)[ind]
+
+        return array
 
 
 def maybe_decode_bytes(txt):

xarray/backends/netCDF4_.py

@@ -50,17 +50,16 @@ def get_array(self):
 
 class NetCDF4ArrayWrapper(BaseNetCDF4Array):
     def __getitem__(self, key):
-        key = indexing.unwrap_explicit_indexer(
-            key, self, allow=(indexing.BasicIndexer, indexing.OuterIndexer))
-
+        key, np_inds = indexing.decompose_indexer(key, self.shape,
+                                                  mode='outer')
         if self.datastore.is_remote:  # pragma: no cover
             getitem = functools.partial(robust_getitem, catch=RuntimeError)
         else:
             getitem = operator.getitem
 
         with self.datastore.ensure_open(autoclose=True):
             try:
-                data = getitem(self.get_array(), key)
+                array = getitem(self.get_array(), key.tuple)
             except IndexError:
                 # Catch IndexError in netCDF4 and return a more informative
                 # error message.  This is most often called when an unsorted
@@ -73,7 +72,10 @@ def __getitem__(self, key):
                     msg += '\n\nOriginal traceback:\n' + traceback.format_exc()
                 raise IndexError(msg)
 
-        return data
+        for ind in np_inds:
+            array = indexing.NumpyIndexingAdapter(array)[ind]
+
+        return array
 
 
 def _encode_nc4_variable(var):

xarray/backends/pydap_.py

@@ -24,18 +24,22 @@ def dtype(self):
         return self.array.dtype
 
     def __getitem__(self, key):
-        key = indexing.unwrap_explicit_indexer(
-            key, target=self, allow=indexing.BasicIndexer)
+        key, np_inds = indexing.decompose_indexer(key, self.shape,
+                                                  mode='basic')
 
         # pull the data from the array attribute if possible, to avoid
         # downloading coordinate data twice
         array = getattr(self.array, 'array', self.array)
-        result = robust_getitem(array, key, catch=ValueError)
+        result = robust_getitem(array, key.tuple, catch=ValueError)
         # pydap doesn't squeeze axes automatically like numpy
-        axis = tuple(n for n, k in enumerate(key)
+        axis = tuple(n for n, k in enumerate(key.tuple)
                      if isinstance(k, integer_types))
         if len(axis) > 0:
             result = np.squeeze(result, axis)
+
+        for ind in np_inds:
+            result = indexing.NumpyIndexingAdapter(result)[ind]
+
         return result
 
 

xarray/backends/pynio_.py

@@ -27,14 +27,18 @@ def get_array(self):
         return self.datastore.ds.variables[self.variable_name]
 
     def __getitem__(self, key):
-        key = indexing.unwrap_explicit_indexer(
-            key, target=self, allow=indexing.BasicIndexer)
+        key, np_inds = indexing.decompose_indexer(key, self.shape,
+                                                  mode='outer')
 
         with self.datastore.ensure_open(autoclose=True):
             array = self.get_array()
             if key == () and self.ndim == 0:
                 return array.get_value()
-            return array[key]
+
+            for ind in np_inds:
+                array = indexing.NumpyIndexingAdapter(array)[ind]
+
+            return array
 
 
 class NioDataStore(AbstractDataStore, DataStorePickleMixin):

xarray/backends/rasterio_.py

@@ -41,8 +41,10 @@ def shape(self):
         return self._shape
 
     def __getitem__(self, key):
-        key = indexing.unwrap_explicit_indexer(
-            key, self, allow=(indexing.BasicIndexer, indexing.OuterIndexer))
+        # TODO support indexing.decompose_indexer
+        if isinstance(key, indexing.VectorizedIndexer):
+            raise NotImplementedError
+        key = key.tuple
 
         # bands cannot be windowed but they can be listed
         band_key = key[0]