Plan conflict-free write batches for mismatched chunk grids - #324
Open
d-v-b wants to merge 20 commits into
Open
Conversation
Restricting a transform to a chunk box distributes over output dimensions whenever each output map reads its own input axis, which is every basic and orthogonal selection. Chunk resolution therefore no longer intersects the whole transform with every candidate chunk; it resolves each axis once against its grid into a table (StridedSet / IndexedSet), sorts correlated (vindex) index arrays into chunks once into a JointSet, and derives each ChunkProjection as one row of each table. ChunkPlan.partition() and partition_transform() expose the factored form, so a consumer can read the tables directly instead of materializing an object graph per chunk. The projections a plan yields are unchanged; the general whole-transform walk remains for hand-built diagonals, which have no factored form. Along the way: _intersect_general reuses a precomputed _CorrelatedBlock and accepts survivor positions; checked_affine has identity and dtype-bounded fast paths; ArrayMap._with_affine shares frozen index arrays on translate; IndexDomain._unchecked / IndexTransform._unchecked skip validation for objects derived from an already-valid transform. Assisted-by: ClaudeCode:claude-fable-5-1
Assisted-by: ClaudeCode:claude-fable-5-1
The package and its tests import nothing from zarr; the old comments claimed the chunk-resolution tests needed zarr's ChunkGrid, which stopped being true once the package grew its own grids. The real reason is the shared pinned test toolchain. Assisted-by: ClaudeCode:claude-fable-5-1
…hunk narrative The module docstring described intersecting the whole transform with every candidate chunk as "the algorithm"; that walk is now the fallback for hand-built diagonals only. It now explains the factored form and its three tables, and why they cost the sum of the touched chunks per axis. The visual guide gains a final integrator section, "A plan is a product of per-axis tables", with an executable snippet that reads the StridedSet, IndexedSet and JointSet tables off real plans and checks the plan's projections against the partition's rows. Integration boundaries gains "Reading the tables directly", a consumer that assembles a strided box from the tables with no projection materialized. The API index, landing page and design notes (TensorStore lineage, the performance caveat, and the box/query split) point at the new section. Assisted-by: ClaudeCode:claude-fable-5-1
…w fixes Adversarial review (roborev, a correctness reviewer, a complexity reviewer, and ~24k differential examples against main) of the grid partition. Cuts. The whole-transform walk that remained for hand-built diagonals is gone: it was unreachable for every index-array shape, its key builder was duplicated verbatim in _chunk_keys, and for the one shape it served it produced wrong projections (a three-point diagonal yielded four projections covering six cells, on main too). A DimensionMap diagonal is now rejected with ValueError. With it go the sorted-1-D fast path, the three cell-transform helpers, the block/positions parameters of _intersect_general, the correlated-residual check that admitted a diagonal and then crashed, GridPartition.__getitem__, partition_transform as public API, the object-dtype column fallback (StridedSet.origin is now a position along the request axis, so every column is intp), checked_affine's dtype-bound shortcut (measured at noise; the identity shortcut stays and now accepts bool via np.can_cast, as main did), and StridedSet.chunk_map/cell_map. Fixes. GridPartition.n_rows is an exact integer and len raises OverflowError instead of wrapping to zero; table columns are read-only, so a memoized partition cannot drift under a consumer; the documented table consumer now handles reversed axes, inserted axes and transposed transforms, and the snippet checks all three. Docs. Corrected the diagonal statement everywhere it appeared, the memoized "fresh walk" wording, the "vectorized per axis" claim, and the TensorStore correspondence (its strided sets are per input dimension; it keeps one index array set per connected component). The guide no longer restates the class docstrings. Assisted-by: ClaudeCode:claude-fable-5-1
A zero-stride DimensionMap over a domain wider than np.intp is valid and touches one storage cell; coercing every StridedSet column to intp made it raise OverflowError where main returned one projection. `extent` and `origin` are the two columns measured along the request axis, whose bounds are arbitrary Python ints, so they now fall back to exact-int (object) columns when a value does not fit. Chunk-local columns stay intp. Also corrects the design note that said both affine-diagonal cases raise NotImplementedError: two slice maps sharing an axis now raise ValueError. Assisted-by: ClaudeCode:claude-fable-5-1
…d the minimal grid protocol Every varying grid in the partition cases summed exactly to its extent, so the boundary where a chunk's data extent is shorter than its declared size, the rectilinear-specific case, was unpinned; so was a grid without data_size. Both now run through the evaluation oracle for strided, orthogonal and correlated selections. Assisted-by: ClaudeCode:claude-fable-5-1
Assisted-by: Codex:GPT-6
Assisted-by: Codex:GPT-6
Assisted-by: Codex:GPT-6
Assisted-by: Codex:GPT-6
…cies Assisted-by: Codex:GPT-6
Assisted-by: Codex:GPT-6
Retain public selection-flow documentation and fix singleton data-extent coverage. Move the execution prototype to a follow-up review. Assisted-by: Codex:GPT-6
Preserve source tasks, schedule destination write-unit conflicts, and expose explicit preserve/reorder policies. Include adversarial property tests, real concurrent codec copies, documentation, and comparisons with Dask and Rechunker planners. Assisted-by: Codex:GPT-6
Assisted-by: Codex:GPT-6
Merge current main, use the retained chunk_coords API, and match the factored planner diagonal rejection contract. Verify multidimensional gather schedules against independently enumerated write units. Assisted-by: Codex:GPT-6
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
1 participant
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.
🤖 AI text below 🤖
Incoming chunks can write disjoint array slices yet race when the destination performs read-modify-write on shared chunks or shards.
plan_rechunkkeeps each source chunk as one task and schedules batches with disjoint destination write units. Chunks of 3 copied into chunks of 4 over length 12 produce((0, 2), (1, 3)).plan_write_batcheshandles existing request-to-destination transforms. Its default preserves input order between conflicting tasks;order="reorder"uses deterministic first-fit coloring. Plans contain task/piece metadata and perform no I/O. This branch now targets main, where the grid-planning prerequisite has landed.Use shards as the destination units when writes replace shard objects. Complete all destination I/O in a batch before starting the next; prefetched destination snapshots invalidate the guarantee. Source and destination must not alias, and unrelated writers must be excluded. Task-internal duplicate assignments remain the caller's responsibility.
Planning is eager, without a byte-memory budget or optimal coloring guarantee. Affine diagonals and shared-axis mixed transforms raise ValueError; unsupported correlated mixtures raise NotImplementedError. The guide documents these limits and compares scheduling with Rechunker and Dask.
Review added a multidimensional property test comparing scheduling against enumerated write footprints, and updated planning calls for current main. Validation: 1,426 tests passed, 4 existing skips; Ruff passed and Pyright reported zero errors.