Dask now supports item assignment. Shouldn't xarray support it too?

[tammasloughran]

As of a few weeks ago, dask now supports item assignment. Documentation here.

At the moment, xarray prevents item assignment if a DataArray uses dask arrays.

>>> import xarray as xr
>>> foo = xr.DataArray([1,2,3,4,5,6], dims=('x'), coords={'x':[1,2,3,4,5,6]}).chunk(2)
>>> foo.loc[1] = 44
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/tammas/miniconda3/envs/da/lib/python3.8/site-packages/xarray/core/dataarray.py", line 213, in __setitem__
    self.data_array[pos_indexers] = value
  File "/home/tammas/miniconda3/envs/da/lib/python3.8/site-packages/xarray/core/dataarray.py", line 724, in __setitem__
    self.variable[key] = value
  File "/home/tammas/miniconda3/envs/da/lib/python3.8/site-packages/xarray/core/variable.py", line 889, in __setitem__
    indexable[index_tuple] = value
  File "/home/tammas/miniconda3/envs/da/lib/python3.8/site-packages/xarray/core/indexing.py", line 1379, in __setitem__
    raise TypeError(
TypeError: this variable's data is stored in a dask array, which does not support item assignment. To assign to this variable, you must first load it into memory explicitly using the .load() method or accessing its .values attribute.

Now that this is possible in dask, shouldn't xarray allow it too if the installed dask version supports it?

Kind regards,
Tam

[max-sixty]

Yes! We'd definitely take a PR for this.

Related — we can also speed up .stack for dask arrays, which we currently have a workaround for because of this.

[tammasloughran]

PR #5174 . This is a start. I'm not too familiar with the inner workings of xarray, so it might need more work.