Creating datasets with dask xarray arrays issue

[miguelcarcamov]

When creating a dataset with xarray arrays with dask arrays inside then when doing a calculation there seems to be no dask array in the variables. However, if I force the dataset to use dask arrays using .data the variables are indeed dask arrays. I was wondering if xarray chooses to use dask array or numpy arrays when creating the dataset depending on the memory of the machine or if this is done automatically.

Example:

import xarray as xr
import dask.array as da
import numpy as np
nrows = 1000
nchan = 16
ncorr = 2
N = 256

_id = da.arange(0, nrows)
uvw_grid = xr.DataArray(da.random.uniform(size=(nrows, 3)).astype(np.int32) + N//2)
data = xr.DataArray(da.random.uniform(size=(nrows, nchan, ncorr)).astype(np.float32) + 1j*da.random.uniform(size=(nrows, nchan, ncorr)).astype(np.float32))
weight = xr.DataArray(da.random.uniform(size=(nrows, ncorr)))

xarr_ds = xr.Dataset(
                data_vars=dict(
                    U=(["row"], uvw_grid[:,0]),
                    V=(["row"], uvw_grid[:,1]),
                    W=(["row"], uvw_grid[:,2]),
                    DATA=(["row", "chan", "corr"], data),
                    WEIGHT=(["row","corr"], weight),
                ),
                coords=dict(
                     ROWID=(["row"], _id)
                )
            )

print(xarr_ds)

<xarray.Dataset>
Dimensions:  (chan: 16, corr: 2, row: 1000)
Coordinates:
    ROWID    (row) int64 dask.array<chunksize=(1000,), meta=np.ndarray>
Dimensions without coordinates: chan, corr, row
Data variables:
    U        (row) int32 128 128 128 128 128 128 128 ... 128 128 128 128 128 128
    V        (row) int32 128 128 128 128 128 128 128 ... 128 128 128 128 128 128
    W        (row) int32 128 128 128 128 128 128 128 ... 128 128 128 128 128 128
    DATA     (row, chan, corr) complex64 (0.81320506+0.025824035j) ... (0.186...
    WEIGHT   (row, corr) float64 0.3917 0.6287 0.2797 ... 0.5258 0.05505 0.03535

However, if I do:

xarr_ds = xr.Dataset(
                data_vars=dict(
                    U=(["row"], uvw_grid[:,0].data),
                    V=(["row"], uvw_grid[:,1].data),
                    W=(["row"], uvw_grid[:,2].data),
                    DATA=(["row", "chan", "corr"], data.data),
                    WEIGHT=(["row","corr"], weight.data),
                ),
                coords=dict(
                     ROWID=(["row"], _id)
                )
            )

print(xarr_ds)

I get:

<xarray.Dataset>
Dimensions:  (chan: 16, corr: 2, row: 1000)
Coordinates:
    ROWID    (row) int64 dask.array<chunksize=(1000,), meta=np.ndarray>
Dimensions without coordinates: chan, corr, row
Data variables:
    U        (row) int32 dask.array<chunksize=(1000,), meta=np.ndarray>
    V        (row) int32 dask.array<chunksize=(1000,), meta=np.ndarray>
    W        (row) int32 dask.array<chunksize=(1000,), meta=np.ndarray>
    DATA     (row, chan, corr) complex64 dask.array<chunksize=(1000, 16, 2), meta=np.ndarray>
    WEIGHT   (row, corr) float64 dask.array<chunksize=(1000, 2), meta=np.ndarray>

Cheers

[keewis]

The reason this fails is that xr.Dataset({var: ("dim", data_array)}) casts the DataArray to numpy (you're constructing a new variable from a DataArray and a new dimension name). With a recent version of xarray you should get a warning suggesting the use of .data, which will become a error in the future.

This means that the second version is the only one that is supported.

[miguelcarcamov]

However it is very slow when doing a groupby using dask arrays. It takes like 15 sec to do a groupby and then a map. When not using .data is much more faster. Does xarray dataset do a compute internally when doing groupby?

[keewis]

I'll let others confirm, but I suspect it is faster because the overhead introduced by dask is removed by converting to numpy (if your data indeed fits into memory).

[dcherian]

If you are using .groupby(dask_dataarray) then yes it will compute to find out the locations of groups. Can you provide your groupby statement?