Add lazy reindexer that can operate on ManifestArrays#1031
Add lazy reindexer that can operate on ManifestArrays#1031sharkinsspatial wants to merge 3 commits into
Conversation
Add ds.vz.reindex(...) and ds.vz.reindex_like(other) that conform a virtual dataset to target coordinate labels lazily, padding/reordering the chunk grid with null-path chunk entries instead of materializing fill values. Missing positions read back as the array's metadata.fill_value (no dtype promotion). Three layers: - chunk_index_map (manifests/reindex.py): pure label->chunk-grid translation; raises NotImplementedError on non-chunk-aligned targets, supports whole-chunk reorder (e.g. sort/reverse a per-timestep virtual dataset). - ManifestArray._reindex_axis: rebuilds the manifest along one axis with null slabs for missing chunks, re-keys inlined chunks, updates metadata shape. - vz.reindex / reindex_like: requires an in-memory index on the reindexed dim (else ValueError), routes ManifestArray-backed variables through the chunk-grid path and loaded vars/coords through xarray's own reindex, preserving coordinate status of virtual aux coords. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add API reference entries for VirtualiZarrDatasetAccessor.reindex and reindex_like, and a New Features release note. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1031 +/- ##
==========================================
+ Coverage 89.64% 89.89% +0.24%
==========================================
Files 36 37 +1
Lines 2202 2296 +94
==========================================
+ Hits 1974 2064 +90
- Misses 228 232 +4
🚀 New features to boost your workflow:
|
| reference and only be realized when a reader asks for those bytes (see | ||
| ``test_missing_chunk_reads_back_as_virtual_nan``). | ||
|
|
||
| What actually breaks today is xarray's *reindex machinery*, not the manifest |
There was a problem hiding this comment.
What would it take to support this type of reindexing in xarray?
There was a problem hiding this comment.
I did a few aborted investigations with Claude. Once I finalize this Draft I'll add some more context in the initial comment but I should start by saying my understanding of xarray internals is very limited. I'll have to look back through the results of my initial Claude session last week (I think I serialized as a memory) but IIUC
- xarray's reindex operates at the element level rather than the array level so there would be no duck array hook for reindex.
- Internally reindex is using a gather operation that operates at the element level. This element level gather could move elements within a chunk. We can only operate lazily at the chunk level without materializing the chunks so this would probably be undesirable.
- Internally reindex also uses an element level
wherecall that would require materializing.
My guess is that there may be some path to potentially updating xarray internals to handle this but I couldn't grok how that would be handled. Maybe you or @dcherian have insights here.
In the interim, users can call this accessor reindexer to mimic what would happen internally in xarray and then use xarray's native combination methods on the pre-aligned data.
There was a problem hiding this comment.
I'll actually probably try to add an additional align accessor method which would be
- Calculate the union of the both ds indexes.
- Reindex each ds onto these unioned indexes
To make the the user experience a bit more seamless.
There was a problem hiding this comment.
Xarray will provide the backing array library with an array indexer of ints with -1 indicating missing. What you can do is take that indexer, partition it by chunks and then for each partition assert either the partition has no -1 or is all -1 (that is the case you can support). if so, you should be able to translate that indexer to whatever approach is used here. If you support this, you will get Xarray's reindex & alignment machinery for free.
There was a problem hiding this comment.
What I did
Added lazy
reindexandreindex_likeaccessor methods that align a virtual dataset to new coordinate labels without materializing data, padding missing positions in the new target grid with null-path manifest entries so that the missing chunks are read asfill_value.Acceptance criteria:
docs/about/releases.md*.mdfile underdocs/api