open_datatree() from zarr creates issues with `kwargs`

[aladinor]

• Python version: 3.12
• Operating System: Ubuntu 22.04

Description

HI everyone,

I have been working on creating a datatree with multiple datasets containing weather radar data, as shown here. The ultimate purpose is to store data in Analysis-Ready Cloud-Optimized format (Zarr).

When reading data back using the datatree.open_datatree method, it loads the data as it was stored, but the issue is that all Xarray.DataArrays within the tree are rechunked into one single chunk along all dimensions.

from datatree import open_datatree

path = '/media/alfonso/drive/Alfonso/python/raw2zarr/zarr/Guaviare.zarr/'
dtree = open_datatree(path, engine='zarr', decode_cf=False)

Then, the tree output looks like,

Digging around this issue, I found out that the xarray.open_dataset default variables, such as chunks and decode_cf, are set to None. Therefore, when passing an empty kwargs dictionary, all variables will be set to None.

https://github.com/xarray-contrib/datatree/blob/0afaa6cc1d6800987d8b9c37a604dc0a8c68aeaa/datatree/io.py#L88

Possible Solution

I think this can be resolved by using xarray.open_zarr instead of using xarray.open_dataset method in #L88

Please let me know your thoughts.

[TomNicholas]

Hi @aladinor - have you checked that this behaviour is actually different from what xr.open_dataset will do for one dataset? Because the aim here should be compliance with that behaviour, and clarity as to what options should be passed if you want different behaviour (e.g. passing chunks={}).

[aladinor]

Hi @TomNicholas. This is a really good question. So the main thing that I found out is that xarray.open_dataset has its default parameters. Thus, when passing the **kwargs dictionary from datatree.open_datatree to xarray.open_dataset, it will overwrite all the default parameters and set them None. At least in the test I performed, if none of those default arguments were passed, xr.open_dataset would behave as I mentioned in #328. I think that If we want to keep xr.open_dataset, all the parameters should be passed within the **kwargs dictionary or try to find a way to not overwrite the default args in xr.open_dataset. I could not find a way to do that without retyping/passing all of them from datatree.open_datatree to xr.open_dataset . That's why I felt using xr.open_zarr was more convenient. However, I will recheck it using one dataset

[aladinor]

Hi @TomNicholas. This could be a reproducible example for you to check.

import xarray as xr
import datatree
import numpy as np
import pandas as pd

def main():
    data = 283 + 5 * np.random.randn(5, 3, 4)
    times = pd.date_range("2018-01-01", periods=5)
    lons = np.linspace(-120, -90, 4)
    lats = np.linspace(25, 55, 3)
    temp = xr.DataArray(data, coords=[times, lats, lons], dims=["time", "lat", "lon"]).chunk({'time': 1,
                                                                                              'lat': -1,
                                                                                              'lon': -1})
    ds = xr.Dataset(data_vars={'Temp': temp})
    ds2 = xr.Dataset(data_vars={'Temp_C': temp - 273.15}).chunk({'time': 1, 'lat': 1, 'lon': 1})
    dtree = datatree.DataTree(data=ds)
    datatree.DataTree(ds2, parent=dtree, name='set2')
    print(dtree.Temp.chunksizes)
    print(dtree.set2.Temp_C.chunksizes)
    path_zarr = "test.zarr"
    dtree.to_zarr(path_zarr)

    dtree_back = datatree.open_datatree(path_zarr, engine='zarr')
    print(dtree_back.Temp.chunksizes)
    print(dtree_back.set2.Temp_C.chunksizes)

    ...

if __name__ == "__main__":
    main()

And the output

Frozen({'time': (1, 1, 1, 1, 1), 'lat': (3,), 'lon': (4,)})
Frozen({'time': (1, 1, 1, 1, 1), 'lat': (1, 1, 1), 'lon': (1, 1, 1, 1)})
Frozen({})
Frozen({})

The expected output without passing any parameters to open_datatree should be to keep the same chunk sizes. However, it didn't. Let me know your thoughts.

[aladinor]

Hi @TomNicholas. I looked again at xr.open_dataset documentation and I found this

In order to reproduce the default behavior of xr.open_zarr(...) use xr.open_dataset(..., engine='zarr', chunks={}). chunks='auto' will use dask auto chunking taking into account the engine preferred chunks. See dask chunking for more details.

I think based on I will close this issue and the PR. Thanks for your time!