Allow opening selected groups only

[mraspaud]

This PR allows opening selected groups only in open_datatree.

The use case is speeding up loading of files with many groups, in our case netcdf, where we actually need a handful of groups to be loaded.

• [x] Tests added
• [x] Passes pre-commit run --all-files
• [x] Changes are summarized in docs/source/whats-new.rst

[mraspaud]

There seems to be failing tests that I don't think is our doing, as we could reproduct them on the main branch (before our changes where added), is that to be expected?

[keewis]

before you spend more time here: could you check if the version that was integrated into xarray does this already? And if not, open the PR there?

Edit: but yes, the failing tests seem unrelated, that's because of a change in the Dataset / DataArray repr. Edit2: also, the version of open_datatree is much faster now, so we might not even need the manual optimization

[mraspaud]

@keewis thanks for the heads up. We have checked the latest DataTree for the xarray integration, and while it indeed is much faster, it's still to slow for our need.

We need to read batches of 80 files, which have around 70 groups each, on my laptop that takes now around 2 second per file, so almost three minutes to generate the datatrees. As this is for a process that needs to run in realtime, with a new batch every 10 minutes, we are looking for all the performance gains we can get. The optimisation we are looking for with this PR comes from the fact that there are groups which are duplicated across the 80 files (so we can just read them from one file and reuse them for the other files), and that some data from the files we don't need at all.

[keewis]

okay, sure. I'd still recommend checking the version in xarray (which is not public API yet so may still change – though this is pretty unlikely at this point) to see if the group parameter already does what you need it to.

[mraspaud]

From what I understand, the groupparameter just sets the root group, so different purpose.

[TomNicholas]

Hi @mraspaud - thanks for this contribution! I can see how this might be useful. I apologise for the indeterminate state of datatree right now.

From what I understand, the group parameter just sets the root group, so different purpose.

This repository will soon be archived, so if you want this feature then your PR here will need to be reconciled with what's now in xarray main.

The recent PR's that @keewis mentioned are especially pertinent - they speed up opening DataTree objects by multiple orders of magnitude!

We should think about whether your use of the groups kwarg here can be made compatible with the interpretation of group upstream to mean "the root group". e.g. could the type of group be str | Iterable[str] | None?

Another idea you might want to think about is whether the suggested open_dict_of_datasets function might be better suited for your use case (see https://github.com/pydata/xarray/issues/9137). That's already "lower-level", so might be a more natural place to accept an argument that means you only open certain groups.

[TomNicholas]

As there have been multiple relevant changes upstream, it's not clear if this issue is still relevant. @mraspaud can you please try it with the upstream code (including the open_groups function instead of open_datatree) and if that still doesn't work for you raise a new issue upstream?

[mraspaud]

Thanks for coming back to me on this, and sorry for the long silence. I will check with the novelties upstream and report back there if needed.