feat(zarr-indexing): public part lowering, projection-carrying part views, result_into, asyncio example - #292
Open
d-v-b wants to merge 16 commits into
Open
feat(zarr-indexing): public part lowering, projection-carrying part views, result_into, asyncio example#292d-v-b wants to merge 16 commits into
d-v-b wants to merge 16 commits into
Conversation
…ns, add result_into Public API for consumers that plan reads here but perform them through their own I/O layer (the rio-tiler / zarr.AsyncArray pattern): - IndexTransform.as_basic_selection() lowers a box transform to the int/slice tuple that reads exactly its cells at exactly its domain shape, producing collapsed-constant singleton axes as length-1 slices and refusing queries, broadcasts, and axis permutations with ValueError. - Partition.source_selection / Partition.chunk_local_selection expose a part's read as the wrapped array's own basic selection and as the grid-cell-relative equivalent, so an async consumer's loop is out[part.out_selection] = await src.getitem(part.source_selection) and a decoded-chunk cache keyed on base_coords needs no coordinate arithmetic. - part.view.result() now passes the paired ChunkProjection to the reader instead of projection=None, matching the parent view's partitioned read; with_reader and repartitioning keep the pairing, a new selection drops it. - LazyArray.result_into(out, *, parts=None) is the non-allocating form of result(): the caller's validated buffer is filled in place and returned. ChainedIndexingStateMachine gained an invariant running both documented assembly loops literally under plain NumPy semantics, reader-free. Assisted-by: ClaudeCode:claude-fable-5
…ed I/O guide A peer to the Dask example: parts() driven by asyncio.gather against zarr.AsyncArray, with three loops — one fetch per part through Partition.source_selection, a decoded-chunk cache keyed on base_coords and placed with chunk_local_selection, and the ValueError fallback for query parts. The integrations guide gains a matching 'Consumer-owned I/O' section with a tested snippet. Assisted-by: ClaudeCode:claude-fable-5
A view that partitions a window plans over that window, so its parts named their cell in window coordinates while their transform addressed the source. Three consequences, all of them silent: - `chunk_domain` promises global storage coordinates, so the documented decoded-cell read `cell[part.chunk_local_selection]` fetched the wrong region for a part of a part. It is now translated back onto the source. - A part of a part carried that window-relative projection into its own `result()`, where a reader keyed on `chunk_coords` would fetch a chunk the part does not read. Before the pairing existed such a read refused loudly; it refuses again, and the partitioned path takes its context from the part view so both paths agree by construction rather than by coincidence. - `view.result(parts=view.parts())` on an unpartitioned view read through a freshly synthesized whole-base projection instead of the view's own. `base_coords` counts cells of the partitioned base, so it keys a decoded-chunk cache only alongside the grid that produced it; the asyncio example keys on the cell's global origin instead, and the docs say which is which. Assisted-by: ClaudeCode:claude-fable-5
…esult `result()` allocates, so two buffers it could never produce were unreachable until `result_into` let the caller supply one, and both failed silently: - A plain ndarray for a `numpy.ma` source passed the dtype check and came back holding the values beneath the mask, presented as data. - A buffer overlapping the wrapped array had each part overwrite cells the parts after it still had to read, so the result was wrong in a way that depended on the partitioning. Both are now refused. The overlap test asks the cheap bounds question first and only pays for an exact answer once memory is known to overlap. `result_into`'s own docstring recommended `final[part.out_selection]` as a destination with only a parenthetical hedge, though NumPy hands back a copy for a fancy `out_selection` — which this method would fill and return, leaving `final` untouched. It cannot tell a copy from a view, so the doc now branches on the selection instead of hedging. Translating a windowed cell's domain keeps its labels (roborev, edc35b5). Assisted-by: ClaudeCode:claude-fable-5
`LazyArray._select` admits NumPy's `None`, so `lazy[None, 2:8]` is a box view like any other — but its domain carries an axis no output map reads, and `as_basic_selection` had no spelling for one. Every part of such a view raised "the selectors must produce the domain's axes in increasing order", and since `is_box` is `True` throughout, a consumer following the documented precheck met an uncaught `ValueError` on a selection NumPy itself spells natively. An unreferenced axis is one nothing reads, which is what `None` does, so it now lowers back to `None` — except where a constant is due in its position, which keeps the axis as a length-1 slice, as before. Selections gain `None` alongside ints and slices, and the trailing-axis refusal becomes an assertion: a referenced axis is produced by the map that references it, and an unreferenced one is a proven singleton. `is_box` is documented as necessary but not sufficient — broadcasts and axes restored by repetition still refuse — with `ValueError` the decider. A `None` selector, like a negative step, is one a backend narrower than NumPy may reject. Assisted-by: ClaudeCode:claude-fable-5
…owering invariant The example runner guarded on the script's *name* — `if "dask" in stem` — so the asyncio example, which needs `zarr`, was the first to need a guard and not get one: in an environment without zarr (which the suite otherwise supports, and which pyproject deliberately leaves out of the test group) it failed instead of skipping. The guard now reads the imports out of the script, so it holds for whatever an example needs next. The lowering invariant excused any refusing box that contained a `ConstantMap`, which is most of them — an integer-indexed box like `[3, 4:6]` could start refusing and the invariant would stay green. It now names the two shapes that may refuse: an axis restored by repetition, and one broadcast from a single cell. That is also the whole set now that a fabricated axis lowers, so the check is tight in both directions rather than loose in one and wrong in the other. The paired-refusal check no longer re-evaluates the property that already raised on its way in. Assisted-by: ClaudeCode:claude-fable-5
…p reaches `as_basic_selection` computed the first and last coordinate a map addresses over its domain with the same four lines the reader's two slab-pushers each carry, so planning a read through the public lowering and executing one through `BasicReader` asked the same question of two implementations. A change to either — a clamping rule, an overflow guard — would have silently moved the planned request away from the executed read, with both files' doctests still passing. `DimensionMap.endpoints` is now that question, asked once. It returns `None` for an empty interval, which is the case with no coordinate to name and the reason the pair was never just two `checked_affine` calls. The Returns section no longer promises a `None` per unread axis, since a constant due in that position stands in for one (roborev, 413a520). Assisted-by: ClaudeCode:claude-fable-5
Truncating `dask.array` to `dask` asked a coarser question than the guard it replaced: a distribution can be installed while the submodule an example needs is not importable, and the example would then run and fail rather than skip. Keeping the dotted path restores that precision (roborev, a527676). Assisted-by: ClaudeCode:claude-fable-5
The state machine explored hard in the region where the code was already right. Its view was always the one built from the source — nothing ever set it to a part's view — so a view with a window, and every bug that needs one, sat outside the reachable state space. Its selections never carried `None`, so an axis no source axis backs was never drawn. And every reader it drew answers from `context.transform` alone, so a projection describing a different read than the transform beside it could not change a single value. Three additions, one per gap: - `descend_into_a_part` follows a part's view and boxes it again, reaching the part-of-a-part whose cell coordinates and transform count from different origins. The model follows by the documented assembly. - `fabricates_an_axis` draws a basic selection carrying `None`. - `ProjectionReader` reads each part out of the cell its `chunk_domain` names, the way a decoded-chunk cache does, and joins `basic_reader` as a reader every machine draws from — so both halves of a `ReadContext` now face the same NumPy model. Against this PR's first two commits each addition fails: the descent alone trips `box_parts_lower_to_basic_selections` with no newaxis involved, and resolving a nested part view returns the model's values through `basic_reader` while disagreeing through `ProjectionReader` — the value-level statement of the bug that check could not previously make. 1500 examples over 14 steps pass on three sources here. An explicit per-axis partitioning describes the source's extents, so it is skipped where a descent has narrowed the base it would have to sum to. Assisted-by: ClaudeCode:claude-fable-5
`partitionings` describes the source's extents, so a subclass declaring nothing but explicit per-axis sizes has none that fit a part's narrowed base — and `sampled_from` on the resulting empty list raises from inside Hypothesis rather than saying so. Falling back to leaving the view boxed as it already is states the same thing and keeps the run going (roborev, d1fb56f). Assisted-by: ClaudeCode:claude-fable-5
`as_basic_selection` is a partial function, and its refusal shared a type with everything else that can go wrong: the documented consumer fallback was `except ValueError`, which also catches a genuine defect in the lowering — silently degrading every part to the slow path with nothing reporting it. Refusals now raise `NoBasicSelectionError`, exported at the package root and subclassing `ValueError` so existing catch sites keep working. The asyncio example, the stateful invariant, and the error tests all catch the precise type now; the invariant in particular no longer excuses a bug that happens to raise `ValueError`. Assisted-by: ClaudeCode:claude-fable-5
`as_basic_selection` is partial by nature — integers and slices spell only diagonal, order-preserving reads — but its refusal was a dead end, and the reader privately owned the fact that it need not be: `_decompose` already factored every transform into an ascending cover plus a block-local residual. That machinery moves to transform.py and becomes public: `IndexTransform.decompose()` returns `(cover, residual)` with the value law (resolving the residual against `source[cover]` reads exactly the transform's cells) and the composition law (the cover, read as the transform it denotes, chained onto the residual, reads cell for cell as the original — the factorization is inverted by `compose`). Queries factor into their bounding interval plus a block-local gather, so a consumer with its own I/O layer can keep every part on it. `decompose_unit_step()` is the contiguous-cover variant `UnitStepReader` reads through; the readers now call the public methods, so the planned request and the executed read share one implementation. The shared walk also gained the negative-coordinate refusal, which previously produced a cover slice NumPy would have quietly wrapped. The stateful invariant now holds every part — query parts included — to the factorization law against the NumPy model, and the decompose tests run the composition law over every transform shape the lowering tests cover, including the ones as_basic_selection refuses. Assisted-by: ClaudeCode:claude-fable-5
`tuple[int | slice | None, ...]` was spelled out at six signature sites and had already drifted once (`source_selection` said `tuple[int | slice, ...]` until the newaxis lowering widened it). It is now the public alias `BasicSelection`, defined beside the lowering and exported at the root. A plain alias, deliberately not a NewType: the tuple's validity is relative to the array it is applied to — the same value can be a correct source_selection and a wrong chunk_local_selection — so a nominal brand would assert a provenance the type system cannot carry, and no signature inside or outside the package could enforce it. The name is the destination contract's: `zarr.AsyncArray.getitem` types its selection parameter `BasicSelection`, and NumPy calls the dialect basic indexing. The asyncio example's `AsyncSource` protocol now uses it in the one place the type appears in parameter position. Assisted-by: ClaudeCode:claude-fable-5
Reject destination views whose logical elements share storage while preserving valid reversed, transposed, and strided outputs. Add regressions for zero-stride and nonzero-stride overlap. Assisted-by: Codex:gpt-5
Normalize new-axis, negative-step, and query selections into an ascending cover plus an in-memory residual. Exercise the adapter with real AsyncArray reads and include the example in strict typechecking. Assisted-by: Codex:gpt-5
Rename every fragment to the PR-numbered convention required by the changelog validator and update the entries to describe the AsyncArray and result-buffer fixes. Assisted-by: Codex:gpt-5
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 🤖
API additions to
zarr-indexingidentified in an integration study for rio-tiler (an async raster tile server that wants to use this package as a chunk planner while doing I/O throughzarr.AsyncArray).Public lowering of parts to backend-native selections
IndexTransform.as_basic_selection()— lowers a box transform to a tuple of ints/slices such thatsource[selection]reads exactly the transform's cells, at exactly its domain shape. A single-coordinate gather (oindex[:, [2]]) collapses to aConstantMapat construction and leaves a singleton domain axis behind; those lower as length-1 slices so the shape contract holds. Queries, broadcasts (an axis restored by repetition — reachable by gathering a collapsed constant with duplicates), transposed/repeated axes, and negative coordinates raiseValueErrorinstead of guessing a slab.Partition.source_selection— that lowering of a part's global read, so an async consumer's whole loop isout[part.out_selection] = await arr.getitem(part.source_selection).Partition.chunk_local_selection— the same read relative toprojection.chunk_domain's origin, for decoded-chunk caches keyed onbase_coords.Fix:
part.view.result()dropped the projectionResolving a partition's view on its own called the reader with
ReadContext(projection=None), while the parentview.result(parts=...)passed the projection — a trap for custom readers keyed onchunk_coords, and a path the dask example itself uses. Part views now carry their pairedChunkProjectioninto their unpartitionedresult();with_readerand repartitioning keep the pairing, a new.lazyselection drops it.LazyArray.result_into(out, *, parts=None)The non-allocating sibling of
result()(a separate method rather than anout=kwarg, so each method's allocation contract is unconditional). The caller's writable buffer is validated against the view's shape/dtype, filled in place — every cell written exactly once — and returned. A view into a larger array qualifies, so a part's block can land directly in its final slot; anumpy.mabuffer keeps a masked source's mask.Async consumption story
A tested asyncio example (
examples/lazy_indexing_asyncio/, peer to the dask example) drivesparts()withasyncio.gatheragainstzarr.AsyncArray: one fetch per part viasource_selection, a decoded-chunk cache placed withchunk_local_selection, and theValueErrorfallback for query parts. The integrations guide gains a matching "Consumer-owned I/O" section with a tested snippet. NoAsyncReader/result_async— an internalgatherwould put a scheduler inside the package, against the design notes' "no internal scheduler" stance.Testing
tests/test_transform.py/tests/test_lazy_array.py.ChainedIndexingStateMachinegains a reader-free invariant that runs both documented assembly loops literally under plain NumPy semantics — it found the collapsed-constant and repeated-axis corner cases above.just lint/typecheck/test(1331 passed incl. Hypothesis + doctests) /test-tensorstore(88 passed) /docs-check(strict), plus the dask-dependent tests under a dask overlay.Changelog fragments use towncrier orphan names (
changes/+*.md); rename to this PR's number if preferred.🤖 Generated with Claude Code