search

xarray-contrib / datatree

WIP implementation of a tree-like hierarchical data structure for xarray.
https://xarray-datatree.readthedocs.io
Apache License 2.0
161 stars 43 forks source link

Allow opening selected groups only #338

Open mraspaud opened 1 week ago

mraspaud commented 1 week ago

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.

mraspaud commented 1 week ago

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 commented 1 week ago

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 commented 1 week ago

@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 commented 1 week ago

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 commented 1 week ago

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

TomNicholas commented 1 week ago

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.