Skip to content

Add lazy reindexer that can operate on ManifestArrays#1031

Closed
sharkinsspatial wants to merge 3 commits into
mainfrom
align_nan
Closed

Add lazy reindexer that can operate on ManifestArrays#1031
sharkinsspatial wants to merge 3 commits into
mainfrom
align_nan

Conversation

@sharkinsspatial

@sharkinsspatial sharkinsspatial commented Jun 29, 2026

Copy link
Copy Markdown
Collaborator

What I did

Added lazy reindex and reindex_like accessor 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 as fill_value.

Acceptance criteria:

sharkinsspatial and others added 2 commits June 29, 2026 15:33
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

codecov Bot commented Jun 29, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 95.74468% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.89%. Comparing base (e1ebb4b) to head (8f42dcb).

Files with missing lines Patch % Lines
virtualizarr/manifests/array.py 87.50% 4 Missing ⚠️
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     
Files with missing lines Coverage Δ
virtualizarr/accessor.py 96.82% <100.00%> (+1.08%) ⬆️
virtualizarr/manifests/reindex.py 100.00% <100.00%> (ø)
virtualizarr/manifests/array.py 85.29% <87.50%> (+0.67%) ⬆️
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

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

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.

What would it take to support this type of reindexing in xarray?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

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

  1. xarray's reindex operates at the element level rather than the array level so there would be no duck array hook for reindex.
  2. 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.
  3. Internally reindex also uses an element level where call 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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I'll actually probably try to add an additional align accessor method which would be

  1. Calculate the union of the both ds indexes.
  2. Reindex each ds onto these unioned indexes

To make the the user experience a bit more seamless.

@dcherian dcherian Jun 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

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.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

@dcherian 🙇‍♂️ Thanks for this. I couldn't make the leap of logic for partitioning the indexer by chunk boundary. I'll try to investigate tomorrow, hopefully I can get this working for the simple case before we land #954 and then someone else will be on the hook for rectilinear partitioning 😆

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.

Support for sparsely tiled spatial grids?

3 participants