Skip to content

Upgrade xarray-tensor and add tenorstore concat function - #429

Merged
dfulu merged 6 commits into
dev_feb2026_speedupsfrom
10_tensorstore_upgrade
Aug 7, 2026
Merged

Upgrade xarray-tensor and add tenorstore concat function#429
dfulu merged 6 commits into
dev_feb2026_speedupsfrom
10_tensorstore_upgrade

Conversation

@dfulu

@dfulu dfulu commented Jul 31, 2026

Copy link
Copy Markdown
Member

Pull Request

Description

This PR improves how we use xarray-tensorstore with multiple zarr by adding a concat function.

Currently we use a custom function which opens a bunch of zarrs using xarray and concatenates the data variables using tensorstore as a backend. We can never concatenate xarray-tensorstore objects after we have loaded them since xarray.concat doesn't support tensorstore backends and it tries to load the data eagerly to concat it.

The new concat_tensorstore() function in this PR introduces a way for us to lazily concat xarray-tensorstore objects. Now that we have this function, we no longer need the _tensorstore_open_zarrs() function which can only concat tensorstore objects whilst opening the zarrs. This simplifies our code and is more generally useful.

This has the additional benefit of allowing us to remove dask as a requirement. Previously we needed dask for one step in the original _tensorstore_open_zarrs() function. The dask step was also incredibly slow because it would build a dask graph over all the data variables. When opening 5 years of highly chunked satellite data this was taking ~25 seconds to open and concat using open_zarr_paths() which internally called _tensorstore_open_zarrs(). Now this takes under a second.

I've also taken the opportunity to upgrade us to the latest release of xarray-tensorstore which supports zarr3 and allows us to remove a lot of the custom code we had in the xr_tensorstore module.

I've also removed the Apache License and attribution from xarray-tensorstore since we are no longer redistributing any copied code from xarray-tensorstore, only importing it.

NOTE: This PR removes some functionality

Previously we allowed zarrs to be concatenated even if they didn't share the same coordinates. I think this might have been motivated by the Meteosat satellite position over the indian ocean shifting. This PR removes support for this and the coordinates now must match.

I strongly believe we should discourage ourselves from doing this kind of thing since it could affect model performance and I believe there should be other ways to solve satellite shifts better. But if absolutely required, we could put that functionality back into this new version later

Checklist:

  • My code follows OCF's coding style guidelines
  • I have performed a self-review of my own code
  • I have made corresponding changes to the documentation
  • I have added tests that prove my fix is effective or that my feature works
  • I have checked my code and corrected any misspellings

@dfulu
dfulu marked this pull request as ready for review July 31, 2026 10:52
@dfulu dfulu mentioned this pull request Jul 31, 2026
@dfulu
dfulu force-pushed the 10_tensorstore_upgrade branch from 0b311ac to 57b3882 Compare July 31, 2026 15:28
Comment thread src/ocf_data_sampler/common/xr_tensorstore.py Outdated

@Sukh-P Sukh-P left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great, really nice changes!

Base automatically changed from 8_split_generation to dev_feb2026_speedups August 7, 2026 13:16
@dfulu
dfulu force-pushed the 10_tensorstore_upgrade branch from e51aee4 to 2986d2f Compare August 7, 2026 13:17
@dfulu
dfulu merged commit d4bb173 into dev_feb2026_speedups Aug 7, 2026
4 checks passed
@dfulu
dfulu deleted the 10_tensorstore_upgrade branch August 7, 2026 13:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants