design doc: rectilinear reindex - #2179
Conversation
| For a chunk grid of `[1, 1, 2, 3]` shifted right by 1 chunk, there are four options that preserve the array shape. | ||
|
|
||
| 1. Periodic boundary - `[3, 1, 1, 2]` | ||
| 2. One fill chunk that takes up all empty space `[3, 1, 1, 2]` |
There was a problem hiding this comment.
in would be most helpful to start with a shift > 1 to see the difference between 1 and 2.
| 4. Rechunk to include the empty part in the existing chunk. So the first chunk would become (3+1) `4,1,2` | ||
|
|
||
|
|
||
| Option 3 will not work because a chunk with size 1 is far too small for real-world applications. Option 4 has the potential to be computationally intensive as it might end up rechunking real data to combine it with empty chunks. |
There was a problem hiding this comment.
I think rejecting option 3 for this reason is right, but only because Zarr doesn't allow to modify the grid. Icechunk could allow you. Right?
There was a problem hiding this comment.
could allow you
Could, but does not currently. Would be a ncie feature to be able to re-chunk empty chunks as that would be a very cheap metadata op.
|
|
||
| Option 3 will not work because a chunk with size 1 is far too small for real-world applications. Option 4 has the potential to be computationally intensive as it might end up rechunking real data to combine it with empty chunks. | ||
|
|
||
| That leaves Options 1 and 2, which are always equivalent for a shift of magnitude 1. They differ with a shift of magnitude 2 or greater. For a shift by 2 with our example grid of `[1, 1, 2, 3]` they give: |
There was a problem hiding this comment.
let's remind the reader that the shift is in chunk space.
|
|
||
| For any periodic data, like chunks of months, Option 1 is more sensible than creating a large empty chunk that may not fit well with writing workflows. | ||
|
|
||
| That leaves Option 1, which does have the downside that it may be slightly confusing that chunk sizes are shifted around the periodic boundary while values are **not** shifted to the other side of the boundary. However, it seems to be the only strategy that gives generally usable results. |
There was a problem hiding this comment.
How about the lowest level option: ask the user, "hey, after shifting I have to fill holes of size [Nx, Ny, ...] array elements, how do you want to chunk those holes?". Here Ni is the sum of the chunk sizes shifted in that dimension . This could be, for example, in the shape of a function argument:
// the sizes across one dimension of the chunks that overflow
type DimensionOverFlow = Vec<NonZero<usize>>
fn shift_array(...., rechunk_hole: &[DimensionOverflow] -> Vec<DimensionOverflow>)
rechunk_hole indicates, for each dimension the new grid for the holes.
Of course, we need to place some constraints and checks on the returned values (input and output size must be same size, sum of elements for each dim must be equal, etc.
Then we can offer "pre-made" functions that implement rotation, single chunk, fixed_grid_of_size(n), etc.
|
|
||
| Option 2 splits the computation of chunk grid and new location and would require 2 calls across the Rust/Python boundary for each chunk in the array. | ||
|
|
||
| So Option 1 is cleaner and easier for the user. However, it has the downside that it is an API change to `reindex_array`. This is an acceptable risk because there are likely very few users of `reindex_array` today. |
There was a problem hiding this comment.
Is there an option here to not break the API by allowing the function return
list[int] |
tuple[list[int] | None, list[int] | None] |
None
?
There was a problem hiding this comment.
yes that is possible, just kinda drifts farther into one python's most menacing features. variable return sizing. but I guess explicitly requiring a tuple or even a custom object could get around that. Might be the cleanest. Though I do think (no evidence) that reindex_array has zero real world user today.
| - Per-axis rectilinear consistency (chunks sharing an axis index agree on that axis's size) | ||
| - Per-axis sums equal dimension length | ||
|
|
||
| Option 2, doing validation, has a lot of potential for complexity and we may accidentally miss some edge cases. So an incomplete validation is likely worse than none at all as it would give false confidence. Given that users of `reindex_array` are doing an advanced operation it is acceptable for them to accept the risk of writing an incorrect chunk grid. I.e. they assume the responsibility for carefully validating their final chunk grid. Additionally this is safe because they can always roll back the operation using Icechunk time travel. |
There was a problem hiding this comment.
I think I disagree here:
- Are there really many edge cases?
- If it that's so, I think an effort at validation is better than no validation. We don't need to advertise full validation exists.
|
|
||
| Options: | ||
|
|
||
| 1. Accept things as they are and just import both directions. |
|
|
||
| ### Array Metadata implementation location | ||
|
|
||
| Any fix for the chunk grid necessarily will need access to `ArrayMetadata` and related structs which currently live in `store.rs`. However `shift_array` and `reindex_array` are in `session.rs`. So to fix the chunk grid issues we will need to import from store into session which inverts their current import relationship. While Rust can handle this, it's a sign that a new abstraction around metadata would be helpful. |
There was a problem hiding this comment.
Can we make the session functions lower level and have store call them, and then update the metadata based on output from those functions?
| forward(idx: list[int]) -> list[int] | None | ||
| backward(idx: list[int]) -> list[int] | None | ||
| forward_chunk_shape(idx: list[int]) -> list[int] | ||
| backward_chunk_shape(idx: list[int]) -> list[int] |
There was a problem hiding this comment.
Help me think about this, do we also need two functions here? Can we have only regrid_vacuumed_chunk that is called on every chunk that was cleared, doesn't matter if it was cleared because of forward or backward?
There was a problem hiding this comment.
ah very good point. and I like the term "vacuumed"
There was a problem hiding this comment.
this comment sent me into a deep rabbit hole. very impactful :) there's actually several restrictions about what chunks we even need user input for. in many cases the chunk size si fully constrained by other chunks. putting in draft till i can sort it all out a bit better.
|
Ok after considering the impact of vacuuming I spent awhile considering how to state the problem more precisely and it turns out there are both more constraints and more ambiguities than I had considered at first. I've re-written a big portion to better state the problem. Be curious to have your take on re-reading the first few sections @paraseba prior to pushing forward on the API design. I want to make sure I fully captured the actual problem space first. |
|
Another thing I hadn't thought that I suspect is also broken. When you have a regular chunk grid where the array elements aren't evenly divisible by the number of chunks the final chunk along an axis is truncated. I'd be that that is broken today as well edit: tried this out, zarr-python ends up handling this by figuring out the fall off behavior for the final chunk correctly. but maybe that is somehow in conflict with the other semantics of shift. this will require more thought |
…ver#2257) ## Real human summary earth-mover#2151 documents a bug with `shift_array` and `reindex_array` for non-regular chunk grids. We've decided the complete fix is necessary but likely to take longer than we're willing to wait before addressing the risk of this bug hitting a real use case. So the stop-gap we've decided on is to reject shift_array/reindex_array for non-regular chunk grids. Then we'll come back to implement the more complete fix in a later PR. 👇 below is Claude's summary of the PR. ---- Part of earth-mover#2151 (intentionally not closing it — this is the stopgap, not the fix). ## Problem `shift_array` / `reindex_array` relabel chunk indices in the manifest without ever looking at chunk sizes. That's safe on regular grids, where every chunk is encoded at the full uniform chunk shape, but on a rectilinear grid slot `i` expects exactly `chunk_shapes[i]` elements. A shift commits cleanly and leaves payloads in slots that decode to different sizes — the array can no longer be read back (`ValueError: cannot reshape array of size 1 into shape (2,)`). `reindex_array` is exposed to Python with arbitrary index mappings, so it has the same failure mode. ## Change - `reindex_array` (and therefore `shift_array`) now fails up front with a new `UnsupportedChunkGridForReindex` error unless the node's `user_data` parses as zarr array metadata positively declaring a `regular` chunk grid. The check fails closed: metadata we can't parse is rejected too. Parsing is intentionally minimal — only the chunk grid name is read. - The rejected operation stages nothing in the change set. - Docs: warning in the moving-chunks guide, notes on both docstrings, and `design-docs/018-shift-array-rectilinear-guard.md` recording the decision. ## Relationship to earth-mover#2179 earth-mover#2179 designs the real fix (rewriting the chunk grid during reindex, gap ambiguity, settled/open edges in N dimensions). This PR is the immediate data-corruption stopper; the design doc here defers to earth-mover#2179 for the follow-up and takes number 018 since earth-mover#2179 claims 017. For plain `shift_array` the follow-up is a metadata-only right-rotation of `chunk_shapes`; arbitrary reindex mappings need the user-facing grid API from earth-mover#2179. ## Tests - Rust: rejection of both operations on both spec versions, plus the regular-grid happy path. - Python: repro from the issue asserting the error and that the session stays clean. Works on zarr 3.2.x (tuple-of-tuples `chunks`) as well as zarr main (`RectilinearChunkGrid`). - Stateful suite: the `shift_array` rule now expects the rejection when hypothesis draws a rectilinear array — it was previously exposed to this corruption since the package conftest enables `array.rectilinear_chunks`. The grid-type check matches both the `RegularChunkGrid` (zarr main) and `RegularChunkGridMetadata` (zarr 3.2.x) class names. Changelog entry deliberately omitted per `RELEASE.md` (written at release time). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Sebastian Galkin <code@amisdelabc.com>
2 participants
Design doc for fixing: #2151
There were far more options than I considered at first!
All prose is hand written. I used AI only for grammar/spelling and a small amount of making the func signatures.