feat(zarr-indexing): LazyArray — generic lazy indexing over array-API arrays - #4222
Draft
d-v-b wants to merge 20 commits into
Draft
feat(zarr-indexing): LazyArray — generic lazy indexing over array-API arrays#4222d-v-b wants to merge 20 commits into
d-v-b wants to merge 20 commits into
Conversation
… arrays A generic wrapper for any array-API-like source (numpy, zarr, cupy, ...) adding TensorStore-style lazy indexing with a positional NumPy dialect: eager __getitem__, .lazy/.oindex/.vindex composing transforms without data access, result()/__array__ materializing. Resolution is partition-based: parts() iterates the base array's partitions projected through the view as resolvable sub-LazyArrays (Partition carries global box coordinates, out placement, and completeness); with_parts() re-partitions the same base explicitly; a single lowering engine serves both partitioned and whole-array sources. Partitioning is discovered from the source (read_chunk_sizes, .chunks) or declared, and never surfaced as a chunks vocabulary. Box selections (no index arrays; affine, interval-representable) are first-class: is_box, bounding_box() (exact hull up to stride), and strides() complete the slab-read story; the design-notes page records the box-vs-query taxonomy and the relationship to TensorStore. Dunders: __dask_tokenize__ (deterministic, canonical-ndsel-body-based), __len__, __iter__, 0-d conversions, pickling. Degenerate all-singleton index-array maps now collapse to constant maps in the transform algebra, and NumPy advanced-index placement rules are implemented faithfully. Assisted-by: ClaudeCode:claude-fable-5
`arr[::-1]` reverses. One desugaring rule covers both signs, as TensorStore 0.1.84 does and as ndsel PR #2 now specifies: omitted bounds resolve on the side the traversal starts and stops (`hi-1` and `lo-1` going down), the source interval is [start, stop) going up and [stop+1, start+1) going down, an empty interval is legal at any coordinate, an interval running the wrong way is an error rather than a silent empty, and the origin is trunc(start/step) for either sign. The corpus is re-vendored from ndsel 92d6a32 in this same commit, because it is the definition of correct here: `slice.json` gains ten negative-step fixtures, `errors.json` retires `negative_step_unsupported` for three `bounds_out_of_order` fixtures, and the message layer is changed to satisfy them. The retired reason code is documented as such rather than removed. A latent bug that only negative steps could reach: `_reindex_array` built `slice(pos, pos + size*step, step)`, and a downward walk reaching the front of the array computes a negative stop, which NumPy reads as counting from the end — `slice(6, -1, -1)` selects nothing where `slice(6, None, -1)` selects seven elements reversed. Both reindex helpers now go through `_positional_slice`. The stride<0 branches of `_intersect_dimension_map` and `iter_chunk_transforms` were written defensively and had never been reachable. They are now, and they were right: the parts-coverage test gains two reversing views, and the seeded sweep generates downward slices (4,352 of 7,200 chains carry one) across every partitioning. At the wrapper boundary the dialect stays NumPy's, which differs in one place: a reversed *positional* interval like `lazy[2:5:-1]` is empty, not an error, because that is what `x[2:5:-1]` means. Only literal coordinates call it a direction error. Tests: the study's recorded TensorStore corpus lands in `test_tensorstore_parity.py` — fifteen desugarings with their domains, offsets and strides, the three rows that discriminate trunc from floor and ceil, both error families, empty-outside-the-domain, five recorded compositions, and the recorded index-array reversal (a negative step over a gathered axis reverses the array rather than attaching a stride). Assisted-by: ClaudeCode:claude-fable-5
…ly Partition, strides docs - slice step zero now raises ValueError, matching NumPy in the wrapper's positional dialect (was IndexError) - Partition is keyword-only: box was inserted mid-field-list, so positional construction would silently misbind - strides() documents the empty-box case (bounding_box None, strides still defined) Assisted-by: ClaudeCode:claude-fable-5
github-actions
Bot
added
the
needs release notes
Rewrite the package's documentation surfaces — docs/index.md, docs/design-notes.md, docs/ndsel.md, docs/api/index.md, the LazyArray, boundary, and transform docstrings, and the 267 changelog fragment — in plain declarative English. Metaphor, personification, rhetorical framing, and emphasis used for effect are replaced with statements of the same technical content. No technical claim, API name, example, or example output changes. Assisted-by: ClaudeCode:claude-fable-5
Assisted-by: ClaudeCode:claude-fable-5
A new ruff release (the CI job floats via uvx) flags BLE001/S110 at the chunk-discovery tolerance and tokenize-fallback sites. Both catches are intentional contracts: discovery must degrade to no-information on any foreign-object failure, and a token call must never raise. Configured as per-file-ignores rather than noqa comments because the pinned pre-commit ruff strips the comments as unused (RUF100) while the floating CI ruff requires them. Assisted-by: ClaudeCode:claude-fable-5
Mirrors the main-branch pin (#271) so this PR's workflow runs the same ruff version; bump together with the pyproject pin. Assisted-by: ClaudeCode:claude-fable-5
Two runnable examples in the house style: wrapping a NumPy array in LazyArray (attribute forwarding, composing selections, box vs query selections, partitions), and using a LazyArray with Dask (from_array, one task per partition, deterministic tokens). Assisted-by: ClaudeCode:claude-fable-5
Adds a timed comparison of chained selections through dask.array against the same selections composed into one transform, and a section on when each is the right tool: dask's graph earns its cost when there is computation across chunks, and is overhead when it only defers indexing. Assisted-by: ClaudeCode:claude-fable-5
The example runner rewrote only the `zarr` dependency to the local checkout, so an example depending on an in-repo package resolved it from git main instead — and the lazy-indexing examples failed in CI, since LazyArray is not on main yet. Rewrite every package this repository ships, leaving dependencies an example does not declare alone. Assisted-by: ClaudeCode:claude-fable-5
`LazyArray` assumed every wrapped array could do basic slicing and did all fancy work itself, reading a block and post-indexing it with NumPy. That over-reads when the source could gather natively, and it walks a source axis by axis with `take` where one request would do. Each wrapper now carries an `IndexingSupport` level — BASIC, OUTER, OUTER_1VECTOR, VECTORIZED, the taxonomy and member names of xarray's `IndexingSupport` — and every read is split into the largest part of the selection that level can express, asked of the source in one call through `oindex`/`vindex` when it has them, and a residual transform applied to the block that comes back. The split is applied per partition as well as per whole-array read, so a part costs one request. The level is resolved at construction: an explicit `with_indexing_support` wins, then the source's own `__zarr_indexing_support__` (read defensively), then conservative inference — a NumPy array or zarr's `oindex`/`vindex` pair reads as VECTORIZED, everything else as BASIC, the only assumption that is always correct. A multi-array outer request only ever goes to an `oindex` accessor, because a bare `__getitem__` key with two arrays means an outer product to HDF5 and a correlated gather to NumPy. The level decides how much data crosses the boundary, never what `result()` returns. Tests hold that invariant directly: every selection case, at all four levels, against NumPy and zarr sources, partitioned and not; plus test doubles that raise when handed a key their declared level forbids, so exceeding a declaration fails loudly rather than working by accident. Assisted-by: ClaudeCode:claude-fable-5
The token included the partitioning while deliberately excluding the indexing-support level, though both are read strategies that leave the values unchanged. Excluding both means two wrappers that describe the same data token alike, so a consumer caching on tokens reuses one result across partitionings and support levels. Assisted-by: ClaudeCode:claude-fable-5
A slice with a negative step whose start lies before the front of the axis selects nothing, but the positional slice was written as `slice(start, stop, step)` with a negative stop, which NumPy reads as counting from the end: an empty selection of a fancy axis returned the whole axis reversed, and the correlated form raised a broadcast error. Write an empty selection out explicitly. The randomized basic-selection generator drew negative-step starts from `[0, size)` only, so a start before the front of the axis was unreachable and the suite could not see this. It now draws from below `-size` as well, and three cases pin the behavior directly. Assisted-by: ClaudeCode:claude-fable-5
An `oindex`/`vindex` step whose entries are all slices carries no coordinates: it narrows the view's own axes and must compose like basic indexing. `_reindex_array_oindex` instead applied each entry positionally to the corresponding axis of the existing index array, without asking whether that axis is one the array varies over or a singleton it merely broadcasts along — the distinction its basic-indexing sibling `_reindex_array` has always made. A slice starting past 0 therefore indexed a size-1 broadcast axis out of range and truncated the whole index array to size 0. A view with no coordinates left resolves to no parts, and `result()` handed back its unwritten `np.empty` buffer: live, on the default path, for any source that advertises chunks. `_reindex_array_oindex` now takes the `ArrayMap` and applies an entry only along its dependency axes (plus the `input_dimension` that breaks the tie for a degenerate length-1 orthogonal selection), preserving a broadcast singleton whatever the slice says. Coordinates never reach a broadcast axis — `_guard_fancy_after_fancy` still rejects genuine fancy-after-fancy with `NotImplementedError`, now under test. Four defects from the same review ride along: - `_array_map_dependency_axes` counted a length-**0** axis as an axis the array varies over, so an empty orthogonal selection classified as correlated and `array_map_dependent_axis` rejected it — a raise on the unpartitioned path where every partitioned path returned the right empty answer. An axis of size 0 carries no dependency any more than a singleton does. - `parts()` raised on a view emptied by a slice over an axis of extent 1. A correlated selection of one point normalizes to an all-singleton index array, so emptying the domain leaves the array at size 1 and the resolver went looking for a chunk. An empty input domain now yields no parts and meets no output domain, matching `result()`. - `sub_transform_to_selections` built `slice(stop + 1, start + 1, stride)` for a negative stride — endpoints swapped, step still negative, so it selected nothing where the reversed axis was meant. Both branches now lower through `_positional_slice`, the same walk an `ArrayMap` axis is reindexed by, which knows a downward walk reaching the front must stop at `None`. - `compose()` evaluated an inner index array over `range(size)` rather than over the outer domain's own range, and addressed it from 0 rather than from the inner domain's origin. Every coordinate resolved to the wrong cell whenever a domain did not start at 0 — which a step-1 slice and a negative-step slice both produce routinely here. - `transform_from_canonical` now rejects a non-integer `index_array` with an `NdselError` carrying `invalid_json`, instead of silently truncating `[0.9, 1.9]` to cells 0 and 1, coercing booleans, or leaking NumPy's own conversion error for strings. The fuzzer missed the first two because `_random_chain` drew at most one fancy step and had no way to spell a step that goes through a fancy accessor while carrying only slices. It now draws such a step separately, and the seeded sweep gained a `parts()` counterpart: `result()` can absorb a defect that the iteration contract cannot, since an empty view assembles correctly from no parts at all. Both sweeps fail on the pre-fix source, as does the new exhaustive stride/extent sweep over the chunk-selection bridge. Assisted-by: ClaudeCode:claude-fable-5
A `vindex` coordinate array with a singleton broadcast axis contributes an axis it does not vary over. A later basic index that consumes the axis it *does* vary over collapses the map to a `ConstantMap` and leaves the broadcast axis in the domain, referenced by nothing. Three places assumed that could not happen: - `sub_transform_to_selections` built `out_selection` with one entry per output map, so a view with such an axis got an index tuple of lower rank than the buffer. `out[out_selection] = value` then placed the part against the leading axes and broadcast the rest — silently wrong data on a partitioned read, and a `parts()` walk that left cells unwritten. - `_restore_domain_axis_order` put an unreferenced axis back as a singleton whatever the domain said. At extent 0 that fabricated a row for a selection whose own `shape` reported it empty. - `_lower_correlated` built its flat gather index from the domain's broadcast shape but added coordinates straight off the stored index array, which is singleton on the axes it does not vary over. The two disagree exactly when a correlated map is constant along a shared broadcast axis. `out_selection` is now built per domain dimension throughout, an unreferenced axis is restored at its own extent, and a correlated map's coordinates are broadcast to the block before being combined. The randomized chain sweep never generated the shape at fault: `_random_vindex` only produced `(length,)` and `(length, 1)` coordinate arrays, neither of which leaves a singleton axis for a later step to strand. It now draws a broadcast rank and places each array's varying axis within it, which reproduces all three failures on the unfixed code. Assisted-by: ClaudeCode:claude-fable-5
Five fixes to the wrapper's edges, none of which changes what a selection means. `result()` and `__array__` no longer alias the wrapped array. An unpartitioned read of a basic selection lowers to plain slicing, so it came back as a *view* of the source; NumPy 2 hands whatever `__array__` returns straight to the caller, so `numpy.array(view, copy=True)` aliased it and a write reached through. Under any partitioning the same read allocates, so this also made the answer depend on how the read was divided. The result is now detached whenever it may share memory with the wrapped array, and the `copy=False` refusal no longer justifies itself with a claim the other branch violated. `result()` verifies that the partition walk covered the output before returning it. The buffer is deliberately uninitialized, so any defect in the walk was reported as plausible-looking numbers rather than as an error. The cells each part addresses are counted from the selectors' own shapes — nothing is read — and a walk that does not add up to the view's size raises. Measured on a 16 MiB read: 51 us of accounting against 6.5 ms of read for 64 parts, within noise end to end, and +1.7% at 512 parts. The `BASIC` floor is a promise about the *source*, not about the blocks it returns. The residual is finished with `take`, `reshape` and `transpose`, which were applied to the block unconverted — so a source meeting exactly the documented floor crashed on `oindex[[4, 0, 0], :, :]`. A block that is neither a NumPy array nor an array-API namespace of its own is now coerced, which leaves a device array where it is. `numpy.matrix` is refused at construction: it never reduces rank, so a view's shape and its result disagree on every rank-reducing selection. A `numpy.ma` source keeps its mask through a partitioned read, which allocates a masked buffer. A declaration holding a *foreign* enum member that names one of these four levels — xarray's `IndexingSupport`, whose members these are borrowed from — is honored rather than discarded, since discarding it fell through to inference and answered with a *more* permissive level than the source asked for. `is_complete` is true for a reversing view, which reads every cell of its box back to front; the stride-1 test it failed was about direction, not coverage. `with_parts` accepts `(0,)` and `(0, 0)` for a zero-length axis, which said the same thing as the `()` and uniform spellings it already took, and the positivity error names the working form. Above the token digest limit and without dask, `__dask_tokenize__` returns a value that matches nothing rather than a shape-and-dtype description that two different 4 MiB arrays shared. A cache keyed on it misses instead of lying. Assisted-by: ClaudeCode:claude-fable-5
Every statement below was executed before being rewritten, and the replacement was executed too. - "`view + 1`" / "arithmetic materializes through `__array__`" is false. `LazyArray` defines no arithmetic dunders, so `view + 1` raises `TypeError`. What does work is a NumPy *function* — `numpy.add(view, 1)`, `numpy.sum(view)`, `numpy.stack([view, view])` — and an ndarray on the left of the operator. Corrected in the module docstring, `docs/index.md` and the changelog fragment. - "An empty selection returns `None` from both" is false: an empty *box* reports `strides()` and only `bounding_box()` is `None`. The `strides()` docstring already said so; the design notes now agree with it. - "A box touches a contiguous run of parts" is false for a strided box — `[::4]` over 2-wide parts visits every other part. The true property, and the one a partition-walk optimizer would want, is a regularly-spaced run in increasing order, each part at most once. - "TensorStore permits a lower-rank index array" is backwards. Checked against tensorstore 0.1.84: its JSON parser rejects a rank-1 array over a rank-2 domain and accepts full rank with singletons, which is what we emit. *Our* loader is the permissive one. The passage now says both models want full rank, keeps the real rationale (the singletons are what makes the orthogonal/vectorized distinction derivable), and describes our lower-rank acceptance as the compatibility affordance it is. - "Two limits remain" omitted fancy-after-fancy, which is a live `NotImplementedError` reachable from the documented surface, while `index.md` invited chaining fancy steps "anywhere in the chain". Current scope now lists five limits, including the diagonal-view and mixed correlated/orthogonal ones, and both prose pages point at it. - "A single whole-array part stays in the wrapped array's namespace" is only true with *no* partitioning: `result()` branches on whether a partitioning is in force, not on how many boxes it has, so `with_parts((4, 6))` on a 4x6 array returns a plain ndarray. - The changelog stated the support-detection precedence backwards (declaration wins, not inference); `index.md` had a sentence missing its noun; the module docstring's one-line `bounding_box()` summary dropped the stride caveat the three other locations keep; and the package README, the PyPI long description, never mentioned `LazyArray`. Assisted-by: ClaudeCode:claude-fable-5
The index-array rank was checked only from above, so a lower-rank array could exist inside the engine and be read for dependency axes it did not have. Nothing produced one: the tolerance was there for a test asserting compatibility with a body TensorStore itself rejects (verified against 0.1.84 — a rank-1 array over a rank-3 domain is an error in its JSON parser). Require the full input rank in the type, widen a lower-rank array at the JSON boundary where external input arrives, and give the test the shape TensorStore accepts. Assisted-by: ClaudeCode:claude-fable-5
An index array axis must be the domain's extent or a singleton it broadcasts over. Any other size leaves input coordinates with no entry, which read as a smaller selection rather than as the error it is: the truncated array behind one of this review's silent-corruption bugs was a (3, 0) array over a (3, 2) domain, which this rejects at construction. Two fixtures carried the inconsistency they were meant to exercise — an empty array over a domain with room for two coordinates, and a widening case whose array covered three of four positions — and now describe domains their arrays span. Assisted-by: ClaudeCode:claude-fable-5
…nk-0 part it found Adds `zarr_indexing.testing`, behind a `testing` extra: a Hypothesis state machine that composes indexing steps onto a LazyArray and checks each step's shape, `result()`, and `parts()` assembly against NumPy, plus the selection strategies on their own. A project can point it at its own array by overriding one method. The machine asserts the documented assembly literally — a part's values must arrive at the shape its out_selection addresses — which is how it found the defect it also fixes: intersecting a correlated transform with a part's bounds collapsed the surviving broadcast block into one axis even when the block was already rank 0, so a view narrowed to a single point produced parts of rank 1. A rank-0 block now stays rank 0, and `result()` drops the reshape that was absorbing the mismatch. Merged from the branch that produced it, which predates the coverage guard in `result()`; the guard stays and the reshape it compensated with goes. Assisted-by: ClaudeCode:claude-fable-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.
Summary
Adds a generic
LazyArrayadapter class tozarr-indexing. This is a class that wraps an eagerly indexed array with a lazy indexing accessor field.lazy. The idea is to support grafting lazy indexing semantics on top of existing eager arrays. Maybe this is useful for zarr-python, we will see.written by claude, origin PR at d-v-b#267
Author attestation
TODO
docs/user-guide/*.mdchanges/