get open_mfdataset to return numpy arrays

[raybellwaves]

I'm using open_mfdataset on many zarr stores which each have a large number of chunks.

I would like to return a Dataset where the data_var is a numpy array i.e. bring it all into memory in the open_mfdataset call but i'm having difficulties. I've played around the chunks arg and parallel arg but they all seem return a dask.array. It could be the chunking on the underlying zarr stores which require loading. Here's some code to test:

Generate zarr stores

import xarray as xr
from tqdm.auto import tqdm


ds = xr.tutorial.open_dataset("air_temperature")
for i in tqdm(range(30)):
    _ds = ds.isel(time=i).expand_dims("time")
    _ds = _ds.chunk({"lat": 4, "lon": 8, "time": 1})
    _ds["air"].encoding = {}
    _ds.to_zarr(f"/tmp/_{i}.zarr", consolidated=True, mode="w")

Read zarr stores

expected = xr.open_mfdataset("/tmp/_*.zarr", engine="zarr", chunks={}, parallel=True).load()
expected

actual = xr.open_mfdataset("/tmp/_*.zarr", engine="zarr", chunks=None, parallel=False)
actual

[Illviljan]

To open without dask you have to pass chunks=None to open_dataset. But in this line in open_mfdataset chunks is changed to dict:

xarray/xarray/backends/api.py

Line 899 in 4bb9d9c

open_kwargs = dict(engine=engine, chunks=chunks or {}, **kwargs)

Not sure why that is.

Maybe it's possible to get around this using preprocess to trigger a .compute()?

[mathause]

or call .load()