AttributeError using map_blocks with dask 2021.05.0

[pont-us]

What happened:

In an environment with xarray 0.18.0 and dask 2021.05.0 installed, I saved a dataset using to_zarr, opened it again using open_zarr, and called map_blocks on one of its variables. I got the following traceback:

Traceback (most recent call last):
  File "/home/pont/./dasktest2.py", line 12, in <module>
    ds2.myvar.map_blocks(lambda block: block)
  File "/home/pont/loc/envs/xcube-repos/lib/python3.9/site-packages/xarray/core/dataarray.py", line 3770, in map_blocks
    return map_blocks(func, self, args, kwargs, template)
  File "/home/pont/loc/envs/xcube-repos/lib/python3.9/site-packages/xarray/core/parallel.py", line 565, in map_blocks
    data = dask.array.Array(
  File "/home/pont/loc/envs/xcube-repos/lib/python3.9/site-packages/dask/array/core.py", line 1159, in __new__
    if layer.collection_annotations is None:
AttributeError: 'dict' object has no attribute 'collection_annotations'

What you expected to happen:

I expected map_blocks to complete successfully.

Minimal Complete Verifiable Example:

import xarray as xr
import numpy as np

ds1 = xr.Dataset({
    "myvar": (("x"), np.zeros(10)),
    "x": ("x", np.arange(10)),
})   
ds1.to_zarr("test.zarr", mode="w")
ds2 = xr.open_zarr("test.zarr")
ds2.myvar.map_blocks(lambda block: block)

Anything else we need to know?:

I wasn't sure whether to report this issue with dask or xcube. With dask 2021.04.1 the example runs without error, and it seems that dask PR 7309 introduced the breaking change. But my understanding of xarray's map_blocks implementation isn't sufficient to figure out where exactly the bug lies.

Environment:

Output of xr.show_versions()

INSTALLED VERSIONS ------------------ commit: None python: 3.9.4 | packaged by conda-forge | (default, May 10 2021, 22:13:33) [GCC 9.3.0] python-bits: 64 OS: Linux OS-release: 5.8.0-53-generic machine: x86_64 processor: x86_64 byteorder: little LC_ALL: None LANG: en_GB.UTF-8 LOCALE: en_GB.UTF-8 libhdf5: None libnetcdf: None

xarray: 0.18.0
pandas: 1.2.4
numpy: 1.20.2
scipy: None
netCDF4: None
pydap: None
h5netcdf: None
h5py: None
Nio: None
zarr: 2.8.1
cftime: None
nc_time_axis: None
PseudoNetCDF: None
rasterio: None
cfgrib: None
iris: None
bottleneck: None
dask: 2021.05.0
distributed: 2021.05.0
matplotlib: None
cartopy: None
seaborn: None
numbagg: None
pint: None
setuptools: 49.6.0.post20210108
pip: 21.1.1
conda: None
pytest: None
IPython: None
sphinx: None

[jrbourbeau]

Thanks for raising an issue @pont-us. I believe this may have been fixed by #5271 (which was merged just after the 0.18.0 release). Could you install the development version of xarray and try again to see if the problem persists?

[pont-us]

Thanks for raising an issue @pont-us. I believe this may have been fixed by #5271 (which was merged just after the 0.18.0 release). Could you install the development version of xarray and try again to see if the problem persists?

Thanks, I've just tested the current head (24c6152) with both the minimal example above and the xcube unit test that was failing with dask 2021.05.0, and can confirm that the error no longer occurs.

[max-sixty]

Thanks @pont-us . We just released a new version so you should be good to install that now.