Array indexing with dask arrays

[ulijh]

Code example

da = xr.DataArray(np.ones((10, 10))).chunk(2)
indc = xr.DataArray(np.random.randint(0, 9, 10)).chunk(2)

# This fails:
da[{'dim_1' : indc}].values

Problem description

Indexing with chunked arrays fails, whereas it's fine with "normal" arrays. In case the indices are the result of a lazy calculation, I would like to continue lazily.

Expected Output

I would expect an output just like in the "un-chunked" case:

da[{'dim_1' : indc.compute()}].values
# Returns: array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])

Output of xr.show_versions()

INSTALLED VERSIONS ------------------ commit: None python: 3.7.0.final.0 python-bits: 64 OS: Linux OS-release: 4.18.14-arch1-1-ARCH machine: x86_64 processor: byteorder: little LC_ALL: None LANG: de_DE.utf8 LOCALE: de_DE.UTF-8

xarray: 0.10.9
pandas: 0.23.4
numpy: 1.15.2
scipy: 1.1.0
netCDF4: None
h5netcdf: 0.6.2
h5py: 2.8.0
Nio: None
zarr: None
cftime: None
PseudonetCDF: None
rasterio: None
iris: None
bottleneck: 1.2.1
cyordereddict: None
dask: 0.19.4
distributed: None
matplotlib: 2.2.3
cartopy: 0.16.0
seaborn: None
setuptools: 40.4.3
pip: 18.0
conda: None
pytest: 3.8.2
IPython: 6.5.0
sphinx: 1.8.0

[shoyer]

For reference, here's the current stacktrace/error message:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-74fe4ba70f9d> in <module>()
----> 1 da[{'dim_1' : indc}]

/usr/local/lib/python3.6/dist-packages/xarray/core/dataarray.py in __getitem__(self, key)
    472         else:
    473             # xarray-style array indexing
--> 474             return self.isel(indexers=self._item_key_to_dict(key))
    475 
    476     def __setitem__(self, key, value):

/usr/local/lib/python3.6/dist-packages/xarray/core/dataarray.py in isel(self, indexers, drop, **indexers_kwargs)
    817         """
    818         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'isel')
--> 819         ds = self._to_temp_dataset().isel(drop=drop, indexers=indexers)
    820         return self._from_temp_dataset(ds)
    821 

/usr/local/lib/python3.6/dist-packages/xarray/core/dataset.py in isel(self, indexers, drop, **indexers_kwargs)
   1537         for name, var in iteritems(self._variables):
   1538             var_indexers = {k: v for k, v in indexers_list if k in var.dims}
-> 1539             new_var = var.isel(indexers=var_indexers)
   1540             if not (drop and name in var_indexers):
   1541                 variables[name] = new_var

/usr/local/lib/python3.6/dist-packages/xarray/core/variable.py in isel(self, indexers, drop, **indexers_kwargs)
    905             if dim in indexers:
    906                 key[i] = indexers[dim]
--> 907         return self[tuple(key)]
    908 
    909     def squeeze(self, dim=None):

/usr/local/lib/python3.6/dist-packages/xarray/core/variable.py in __getitem__(self, key)
    614         array `x.values` directly.
    615         """
--> 616         dims, indexer, new_order = self._broadcast_indexes(key)
    617         data = as_indexable(self._data)[indexer]
    618         if new_order:

/usr/local/lib/python3.6/dist-packages/xarray/core/variable.py in _broadcast_indexes(self, key)
    487             return self._broadcast_indexes_outer(key)
    488 
--> 489         return self._broadcast_indexes_vectorized(key)
    490 
    491     def _broadcast_indexes_basic(self, key):

/usr/local/lib/python3.6/dist-packages/xarray/core/variable.py in _broadcast_indexes_vectorized(self, key)
    599             new_order = None
    600 
--> 601         return out_dims, VectorizedIndexer(tuple(out_key)), new_order
    602 
    603     def __getitem__(self, key):

/usr/local/lib/python3.6/dist-packages/xarray/core/indexing.py in __init__(self, key)
    423             else:
    424                 raise TypeError('unexpected indexer type for {}: {!r}'
--> 425                                 .format(type(self).__name__, k))
    426             new_key.append(k)
    427 

TypeError: unexpected indexer type for VectorizedIndexer: dask.array<xarray-<this-array>, shape=(10,), dtype=int64, chunksize=(2,)>

It looks like we could support this relatively easily since dask.array supports indexing with dask arrays now. This would be a welcome enhancement!

[ulijh]

It seem's working fine with the following change but it has a lot of dublicated code...

diff --git a/xarray/core/indexing.py b/xarray/core/indexing.py
index d51da471..9fe93581 100644
--- a/xarray/core/indexing.py
+++ b/xarray/core/indexing.py
@@ -7,6 +7,7 @@ from datetime import timedelta
 
 import numpy as np
 import pandas as pd
+import dask.array as da
 
 from . import duck_array_ops, nputils, utils
 from .pycompat import (
@@ -420,6 +421,19 @@ class VectorizedIndexer(ExplicitIndexer):
                                      'have different numbers of dimensions: {}'
                                      .format(ndims))
                 k = np.asarray(k, dtype=np.int64)
+            elif isinstance(k, dask_array_type):
+                if not np.issubdtype(k.dtype, np.integer):
+                    raise TypeError('invalid indexer array, does not have '
+                                    'integer dtype: {!r}'.format(k))
+                if ndim is None:
+                    ndim = k.ndim
+                elif ndim != k.ndim:
+                    ndims = [k.ndim for k in key
+                             if isinstance(k, (np.ndarray) + dask_array_type)]
+                    raise ValueError('invalid indexer key: ndarray arguments '
+                                     'have different numbers of dimensions: {}'
+                                     .format(ndims))
+                k = da.array(k, dtype=np.int64)
             else:
                 raise TypeError('unexpected indexer type for {}: {!r}'
                                 .format(type(self).__name__, k))

[ulijh]

As of version 0.12 indexing with dask arrays works out of the box... I think this can be closed now.

[ulijh]

Even though the example from above does work, sadly, the following does not:

import xarray as xr
import dask.array as da
import numpy as np
da = xr.DataArray(np.random.rand(3*4*5).reshape((3,4,5))).chunk(dict(dim_0=1))
idcs = da.argmax('dim_2')
da[dict(dim_2=idcs)]

results in

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-4-3542cdd6d61c> in <module>
----> 1 da[dict(dim_2=idcs)]

~/src/xarray/xarray/core/dataarray.py in __getitem__(self, key)
    604         else:
    605             # xarray-style array indexing
--> 606             return self.isel(indexers=self._item_key_to_dict(key))
    607 
    608     def __setitem__(self, key: Any, value: Any) -> None:

~/src/xarray/xarray/core/dataarray.py in isel(self, indexers, drop, **indexers_kwargs)
    986         """
    987         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "isel")
--> 988         ds = self._to_temp_dataset().isel(drop=drop, indexers=indexers)
    989         return self._from_temp_dataset(ds)
    990 

~/src/xarray/xarray/core/dataset.py in isel(self, indexers, drop, **indexers_kwargs)
   1901                     indexes[name] = new_index
   1902             else:
-> 1903                 new_var = var.isel(indexers=var_indexers)
   1904 
   1905             variables[name] = new_var

~/src/xarray/xarray/core/variable.py in isel(self, indexers, drop, **indexers_kwargs)
    984             if dim in indexers:
    985                 key[i] = indexers[dim]
--> 986         return self[tuple(key)]
    987 
    988     def squeeze(self, dim=None):

~/src/xarray/xarray/core/variable.py in __getitem__(self, key)
    675         array `x.values` directly.
    676         """
--> 677         dims, indexer, new_order = self._broadcast_indexes(key)
    678         data = as_indexable(self._data)[indexer]
    679         if new_order:

~/src/xarray/xarray/core/variable.py in _broadcast_indexes(self, key)
    532             if isinstance(k, Variable):
    533                 if len(k.dims) > 1:
--> 534                     return self._broadcast_indexes_vectorized(key)
    535                 dims.append(k.dims[0])
    536             elif not isinstance(k, integer_types):

~/src/xarray/xarray/core/variable.py in _broadcast_indexes_vectorized(self, key)
    660             new_order = None
    661 
--> 662         return out_dims, VectorizedIndexer(tuple(out_key)), new_order
    663 
    664     def __getitem__(self, key):

~/src/xarray/xarray/core/indexing.py in __init__(self, key)
    460                 raise TypeError(
    461                     "unexpected indexer type for {}: {!r}".format(
--> 462                         type(self).__name__, k
    463                     )
    464                 )

TypeError: unexpected indexer type for VectorizedIndexer: dask.array<arg_agg-aggregate, shape=(3, 4), dtype=int64, chunksize=(1, 4)>

[shoyer]

Yes, something seems to be going wrong here...

[ulijh]

I think the problem is somewhere here:

xarray/xarray/core/utils.py

Lines 85 to 103 in aaeea62

	def safe_cast_to_index(array: Any) -> pd.Index:
	"""Given an array, safely cast it to a pandas.Index.
	If it is already a pandas.Index, return it unchanged.
	Unlike pandas.Index, if the array has dtype=object or dtype=timedelta64,
	this function will not attempt to do automatic type conversion but will
	always return an index with dtype=object.
	"""
	if isinstance(array, pd.Index):
	index = array
	elif hasattr(array, "to_index"):
	index = array.to_index()
	else:
	kwargs = {}
	if hasattr(array, "dtype") and array.dtype.kind == "O":
	kwargs["dtype"] = object
	index = pd.Index(np.asarray(array), **kwargs)
	return _maybe_cast_to_cftimeindex(index)

I don't think pandas.Index can hold lazy arrays. Could there be a way around exploiting dask.dataframe indexing methods?

[rafa-guedes]

I'm having similar issue, here is an example:

import numpy as np
import dask.array as da
import xarray as xr

darr = xr.DataArray(data=[0.2, 0.4, 0.6], coords={"z": range(3)}, dims=("z",))
good_indexer = xr.DataArray(
    data=np.random.randint(0, 3, 8).reshape(4, 2).astype(int),
    coords={"y": range(4), "x": range(2)},
    dims=("y", "x")
)
bad_indexer = xr.DataArray(
    data=da.random.randint(0, 3, 8).reshape(4, 2).astype(int),
    coords={"y": range(4), "x": range(2)},
    dims=("y", "x")
)

In [5]: darr                                                                                                                                                            
Out[5]: 
<xarray.DataArray (z: 3)>
array([0.2, 0.4, 0.6])
Coordinates:
  * z        (z) int64 0 1 2

In [6]: good_indexer                                                     
Out[6]: 
<xarray.DataArray (y: 4, x: 2)>
array([[0, 1],
       [2, 2],
       [1, 2],
       [1, 0]])
Coordinates:
  * y        (y) int64 0 1 2 3
  * x        (x) int64 0 1

In [7]: bad_indexer                                                      
Out[7]: 
<xarray.DataArray 'reshape-417766b2035dcb1227ddde8505297039' (y: 4, x: 2)>
dask.array<reshape, shape=(4, 2), dtype=int64, chunksize=(4, 2), chunktype=numpy.ndarray>
Coordinates:
  * y        (y) int64 0 1 2 3
  * x        (x) int64 0 1

In [8]: darr[good_indexer]                                                                                                                                              
Out[8]: 
<xarray.DataArray (y: 4, x: 2)>
array([[0.2, 0.4],
       [0.6, 0.6],
       [0.4, 0.6],
       [0.4, 0.2]])
Coordinates:
    z        (y, x) int64 0 1 2 2 1 2 1 0
  * y        (y) int64 0 1 2 3
  * x        (x) int64 0 1

In [9]: darr[bad_indexer]                                                                                                                                               
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-8-2a57c1a2eade> in <module>
----> 1 darr[bad_indexer]

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/dataarray.py in __getitem__(self, key)
    638         else:
    639             # xarray-style array indexing
--> 640             return self.isel(indexers=self._item_key_to_dict(key))
    641 
    642     def __setitem__(self, key: Any, value: Any) -> None:

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/dataarray.py in isel(self, indexers, drop, **indexers_kwargs)
   1012         """
   1013         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "isel")
-> 1014         ds = self._to_temp_dataset().isel(drop=drop, indexers=indexers)
   1015         return self._from_temp_dataset(ds)
   1016 

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/dataset.py in isel(self, indexers, drop, **indexers_kwargs)
   1920             if name in self.indexes:
   1921                 new_var, new_index = isel_variable_and_index(
-> 1922                     name, var, self.indexes[name], var_indexers
   1923                 )
   1924                 if new_index is not None:

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/indexes.py in isel_variable_and_index(name, variable, index, indexers)
     79         )
     80 
---> 81     new_variable = variable.isel(indexers)
     82 
     83     if new_variable.dims != (name,):

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/variable.py in isel(self, indexers, **indexers_kwargs)
   1052 
   1053         key = tuple(indexers.get(dim, slice(None)) for dim in self.dims)
-> 1054         return self[key]
   1055 
   1056     def squeeze(self, dim=None):

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/variable.py in __getitem__(self, key)
    700         array `x.values` directly.
    701         """
--> 702         dims, indexer, new_order = self._broadcast_indexes(key)
    703         data = as_indexable(self._data)[indexer]
    704         if new_order:

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/variable.py in _broadcast_indexes(self, key)
    557             if isinstance(k, Variable):
    558                 if len(k.dims) > 1:
--> 559                     return self._broadcast_indexes_vectorized(key)
    560                 dims.append(k.dims[0])
    561             elif not isinstance(k, integer_types):

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/variable.py in _broadcast_indexes_vectorized(self, key)
    685             new_order = None
    686 
--> 687         return out_dims, VectorizedIndexer(tuple(out_key)), new_order
    688 
    689     def __getitem__(self: VariableType, key) -> VariableType:

~/.virtualenvs/py3/local/lib/python3.7/site-packages/xarray/core/indexing.py in __init__(self, key)
    447             else:
    448                 raise TypeError(
--> 449                     f"unexpected indexer type for {type(self).__name__}: {k!r}"
    450                 )
    451             new_key.append(k)

TypeError: unexpected indexer type for VectorizedIndexer: dask.array<reshape, shape=(4, 2), dtype=int64, chunksize=(4, 2), chunktype=numpy.ndarray>

In [10]: xr.__version__                                                   
Out[10]: '0.14.1'

In [11]: import dask; dask.__version__                                    
Out[11]: '2.9.0'

[roxyboy]

I'm just curious if there's been any progress on this issue. I'm also getting the same error: TypeError: unexpected indexer type for VectorizedIndexer and I would greatly benefit from lazy vectorized indexing.

[dcherian]

I don't think any one is working on it. We would appreciate it if you could try to fix it.