Reindex via indexer#1032
Open
sharkinsspatial wants to merge 7 commits into
Open
Conversation
Teach ManifestArray to handle xarray's reindex/alignment indexer at the array-protocol layer, so plain xr.reindex, xr.align, and xr.concat work over virtual arrays and stay lazy — missing positions become null-path manifest entries that read back as the array's fill_value, with no materialization and no dtype promotion. xarray's Variable._getitem_with_mask gathers with an integer indexer (-1 for missing) then masks with where(). The hooks that make this work lazily: - manifests/indexing.py: an integer-array indexer is routed to a chunk-grid remap (chunk_map_from_indexer -> ManifestArray._reindex_axis) instead of being rejected as fancy indexing. Each chunk-block must be all-missing (-> null chunk) or a chunk-aligned run (-> real chunk), else NotImplementedError; no chunk splitting. - manifests/reindex.py (new): chunk_map_from_indexer partitions the indexer by chunk and asserts the all-or-nothing condition. - array_api.py: np.where returns the gathered array unchanged when the requested fill mask matches its null chunks (the manifest's own fill_value governs); general masking still raises. np.result_type treats a lone ManifestArray's dtype as authoritative, so int arrays keep their dtype + declared sentinel fill rather than being promoted to float+NaN (also fixes a latent generator-consumption bug in the previous implementation). - array.py: __array_function__ now admits ndarray among the argument types (np.where's boolean condition arrives as one); adds ManifestArray._reindex_axis. 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 #1032 +/- ##
==========================================
+ Coverage 89.87% 89.99% +0.11%
==========================================
Files 36 37 +1
Lines 2202 2288 +86
==========================================
+ Hits 1979 2059 +80
- Misses 223 229 +6
🚀 New features to boost your workflow:
|
A 2D+ outer-join reindexes more than one dimension in a single call, and xarray expresses that as a broadcast (vectorized) indexer: one integer array per axis, shaped to the array's rank with the real labels on its own axis and length 1 everywhere else (e.g. (N,1) and (1,M)). _apply_reindex assumed a 1D indexer per axis, so it misread the array's rank — raising "axis 0 (chunk size 1)" or, on newer numpy, TypeError from int() on a non-scalar block. Add _collapse_outer_indexer: since reindex indexing is orthogonal, each broadcast component reshapes losslessly to its 1D per-axis indexer. A genuine pointwise/vectorized indexer (non-axis dim != 1) is rejected with a clear error. This makes the real ITS_LIVE x+y mosaic work via native reindex/concat. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
c435f3e to
180aee1
Compare
TomNicholas
reviewed
Jun 30, 2026
TomNicholas
left a comment
Member
There was a problem hiding this comment.
Excited to get this to work, but I think we'll have to go back and forth re-organizing and polishing before I feel confident about this.
Member
|
Also it would be quite useful to have an issue/docs example / both showing how this actually solves user's NaN-padding and alignment problems. |
This was referenced Jul 8, 2026
Collaborator
Author
|
Tests are failing due to an upstream GRIB parser issue. |
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.
What I did
Thanks to @dcherian for recommending this approach 🙇
Added support for ManifestArray to handle xarray's reindex/alignment indexer at the array layer, so vanilla
xr.reindex,xr.align(join="outer")that would introducenans work with virtual arrays and stay lazy. Missing positions become null-path manifest entries that Zarr returns as the array's fill_value — no materialization._collapse_outer_indexerreshapes broadcast (vectorized) indexers back to their 1D per-axis form.chunk_map_from_indexertranslates xarray's per-element reindex indexer (e.g. [0,1,2,-1,-1], where -1 marks a missing label) into a per-chunk plan. It slices the indexer into chunk-sized blocks and requires each block to be either all -1 → emit an empty null-path chunk, or a contiguous run that exactly fills one source chunk → reuse that chunk's reference. Anything else (a block mixing real and missing positions, or a partial/reordered chunk) would mean materializing, so it raises. The result is one source-chunk index or null chunk per target chunk slot.np.wherereturns the gathered array unchanged when the fill mask matches its null chunks.ManifestArray._reindex_axisand admits ndarray in__array_function__.This replaces the bespoke manifest surgery people need today to align virtual datasets onto a common grid when the alignment would result in
nans.By design there's no chunk splitting / within-chunk reordering — target labels must land on chunk boundaries. Because align sorts the union ascending, a descending coordinate (e.g. ITS_LIVE y coord) requires explicit reindex with a descending target.
docs/about/releases.md*.mdfile underdocs/api