Multidimensional dask coordinates unexpectedly computed

[djhoese]

MCVE Code Sample

from dask.diagnostics import ProgressBar
import xarray as xr
import numpy as np
import dask.array as da

a = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'y': np.arange(10), 'x': np.arange(10), 'lons': (('y', 'x'), da.zeros((10, 10), chunks=2))}) 
b = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'y': np.arange(10), 'x': np.arange(10), 'lons': (('y', 'x'), da.zeros((10, 10), chunks=2))}) 

with ProgressBar():
    c = a + b

Output:

[########################################] | 100% Completed |  0.1s

Problem Description

Using arrays with 2D dask array coordinates results in the coordinates being computed for any binary operations (anything combining two or more DataArrays). I use ProgressBar in the above example to show when coordinates are being computed.

In my own work, when I learned that 2D dask coordinates were possible, I started adding longitude and latitude coordinates. These are rather large and can take a while to load/compute so I was surprised that simple operations (ex. a.fillna(b)) were causing things to be computed and taking a long time.

Is this computation by design or a possible bug?

Expected Output

No output from the ProgressBar, hoping that no coordinates would be computed/loaded.

Output of xr.show_versions()

INSTALLED VERSIONS

commit: None
python: 3.6.7 | packaged by conda-forge | (default, Feb 28 2019, 02:16:08)
[GCC 4.2.1 Compatible Clang 4.0.1 (tags/RELEASE_401/final)]
python-bits: 64
OS: Darwin
OS-release: 18.6.0
machine: x86_64
processor: i386
byteorder: little
LC_ALL: None
LANG: en_US.UTF-8
LOCALE: en_US.UTF-8
libhdf5: 1.10.4
libnetcdf: 4.6.2

xarray: 0.12.1
pandas: 0.24.2
numpy: 1.14.3
scipy: 1.3.0
netCDF4: 1.5.1.2
pydap: None
h5netcdf: 0.7.4
h5py: 2.9.0
Nio: None
zarr: 2.3.2
cftime: 1.0.3.4
nc_time_axis: None
PseudonetCDF: None
rasterio: 1.0.22
cfgrib: None
iris: None
bottleneck: 1.2.1
dask: 2.0.0
distributed: 2.0.0
matplotlib: 3.1.0
cartopy: 0.17.1.dev147+HEAD.detached.at.5e624fe
seaborn: None
setuptools: 41.0.1
pip: 19.1.1
conda: None
pytest: 4.6.3
IPython: 7.5.0
sphinx: 2.1.2

[djhoese]

Ok I'm getting a little more of an understanding on this. The main issue is that the dask array is not literally considered the same object because I'm creating the object twice. If I create a single dask array and pass it:

lons = da.zeros((10, 10), chunks=2)
a = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'y': np.arange(10), 'x': np.arange(10), 'lons': (('y', 'x'), lons)})
b = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'y': np.arange(10), 'x': np.arange(10), 'lons': (('y', 'x'), lons)})

I still get the progress bar because xarray is creating two new DataArray objects for this lons coordinate. So lons_data_arr.variable._data is not lons_data_arr2.variable._data causing the equivalency check here to fail.

If I make a single DataArray that becomes the coordinate variable then it seems to work:

lons2 = xr.DataArray(lons, dims=('y', 'x'))
a = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'y': np.arange(10), 'x': np.arange(10), 'lons': (('y', 'x'), lons2)})
b = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'y': np.arange(10), 'x': np.arange(10), 'lons': (('y', 'x'), lons2)})

I get no progress bar.

[djhoese]

Ok another update. In the previous example I accidentally added the lons coordinate DataArray with the dimensions redefined (('y', 'x'), lons2) which is technically redundant but it worked (no progress bar).

However, if I fix this redundancy and do:

a = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'lons': lons2})
b = xr.DataArray(da.zeros((10, 10), chunks=2), dims=('y', 'x'), coords={'lons': lons2})
with ProgressBar():
    c = a + b

I do get a progress bar again (lons2 is being computed). I've tracked it down to this transpose which is transposing when it doesn't need to which is causing the dask array to change:

https://github.com/pydata/xarray/blob/master/xarray/core/variable.py#L1223

I'm not sure if this would be considered a bug in dask or xarray. Also, not sure why the redundant version of the example worked.

[djhoese]

Modifying this line to be:

if dims == expanded_vars.sizes:
    return expanded_vars
return expanded_var.transpose(*dims)

Then this issue is avoided for at least the + case.

[dhirschfeld]

FYI: @djhoese, you can inline code snippits using the permanent link to the source:

xarray/xarray/core/variable.py

Line 1223 in e5bb647

return expanded_var.transpose(*dims)

[shoyer]

The source of the problem here is that when combining objects, xarray needs to decide what coordinates should remain. Our current heuristic, which pre-dates dask support, was really designed for array in memory: we keep around coordinates if they are equal on both arguments, and remove them otherwise. In some cases we can avoid the computation, if we know that the coordinates are the same object.

I am open to ideas on how to make this work better.

[djhoese]

@shoyer Understood. That explains why something like this wasn't caught before, but what would be the best solution for a short term fix?

For the long term, I also understand that there isn't really a good way to check equality of two dask arrays. I wonder if dask's graph optimization could be used to "simplify" two dask arrays' graph separately and check the graph equality. For example, two dask arrays created by doing da.zeros((10, 10), chunks=2) + 5 should be theoretically equal because their dask graphs are made up of the same tasks.

Edit: "short term fix": What is the best way to avoid the unnecessary transpose? Or is this not even the right way to approach this? Change dask to avoid the unnecessary transpose or change xarray to not do the tranpose or something else?

[shoyer]

For the long term, I also understand that there isn't really a good way to check equality of two dask arrays. I wonder if dask's graph optimization could be used to "simplify" two dask arrays' graph separately and check the graph equality. For example, two dask arrays created by doing da.zeros((10, 10), chunks=2) + 5 should be theoretically equal because their dask graphs are made up of the same tasks.

Dask actually already does this canonicalization. If two arrays have the same name, they use the same dask graph, e.g.,

In [5]: x = da.zeros((10, 10), chunks=2) + 5

In [6]: y = da.zeros((10, 10), chunks=2) + 5

In [7]: x.name
Out[7]: 'add-f7441a0f46f5cf40458391cd08406c23'

In [8]: y.name
Out[8]: 'add-f7441a0f46f5cf40458391cd08406c23'

So xarray could safely look at .name on dask arrays (e.g., inside Variable.equals or duck_array_ops.array_equiv) for determining that two dask arrays are the same, rather than merely using is to check if they are the same objects.

[djhoese]

Ah, good call. The transpose currently in xarray would still be a problem though.