Upgrade xarray-tensor and add tenorstore concat function - #429
Merged
Conversation
dfulu
marked this pull request as ready for review
July 31, 2026 10:52
dfulu
force-pushed
the
10_tensorstore_upgrade
branch
from
July 31, 2026 15:28
0b311ac to
57b3882
Compare
Sukh-P
reviewed
Aug 4, 2026
Sukh-P
approved these changes
Aug 4, 2026
Sukh-P
left a comment
Member
There was a problem hiding this comment.
Looks great, really nice changes!
dfulu
force-pushed
the
10_tensorstore_upgrade
branch
from
August 7, 2026 13:17
e51aee4 to
2986d2f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.concatdoesn'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 usingopen_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-tensorstorewhich supports zarr3 and allows us to remove a lot of the custom code we had in thexr_tensorstoremodule.I've also removed the Apache License and attribution from
xarray-tensorstoresince we are no longer redistributing any copied code fromxarray-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: