feat(zarr-indexing): LazyArray — generic lazy indexing over array-API arrays - #267
Draft
d-v-b wants to merge 17 commits into
Draft
feat(zarr-indexing): LazyArray — generic lazy indexing over array-API arrays#267d-v-b wants to merge 17 commits into
d-v-b wants to merge 17 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
d-v-b
force-pushed
the
feat/lazy-array-wrapper
branch
from
July 31, 2026 09:54
691bb02 to
02e4c6f
Compare
7 tasks
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
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 🤖
Adds a generic
LazyArraywrapper topackages/zarr-indexing. Nothing insrc/zarris touched.What it is
LazyArraywraps any array-API-like array — NumPy, zarr, CuPy, anything withshape,dtype, and__getitem__— forwards the array-API attributes, and adds a.lazyaccessor for TensorStore-style deferred indexing:.lazy[...],.lazy.oindex[...], and.lazy.vindex[...]compose anIndexTransformand return a newLazyArray.__getitem__(no.lazy) is eager — mirroringzarr.Array's design, and what makes the wrapper a drop-in duck array fordask.array.from_array.Design
Positional dialect, and why it differs from zarr's. The transform algebra speaks literal domain coordinates: after
v = arr.lazy[10:50],v's first element is at coordinate 10, and a negative index is genuinely negative rather than counted from the end (TensorStore's convention).zarr.Array.lazydeliberately exposes that dialect, because a zarr view's coordinates are meant to stay comparable with the parent array's.LazyArraycannot afford that: it is a duck array, so it has to behave like the thing it wraps. Every view re-zeroes its domain and selections are positional NumPy semantics — index 0 is the first element of the current view, negatives wrap, boolean scalars are rejected (mindingisinstance(True, int)), masks must match the view shape exactly, and integer arrays are bounds-checked positionally. The newzarr_indexing.boundarymodule is the generic (zarr-free) translation from positions to literal coordinates beforeselection_to_transform. It duplicates whatzarr/core/array.pydoes today for its own boundary; consolidating zarr onto it is left to a follow-up, noted in the module docstring.Chunk discovery, in order: explicit
chunks=(either convention) >array.read_chunk_sizes(zarr's dask-convention property, sharding-aware; read through a guardedgetattrsince zarr raises anAttributeError-flavouredLazyViewErroron views) >array.chunks, disambiguated by element type (tuple-of-tuples = dask sizes, tuple-of-ints = uniform chunk shape) > none, treated as a single whole-array chunk.Two resolution strategies.
iter_chunk_transforms(transform, grids); read each touched chunk once with plain basic slicing, convert the local transform viasub_transform_to_selections, and scatter into the output buffer (flat scatter for correlated maps). This mirrorszarr.core.array._get_selection_via_transform, reimplemented minimally for in-memory blocks. The buffer is NumPy for v1 because the resolver's scatter indices are host-side; the device caveat is documented.__getitem__, orthogonalArrayMaps become successive per-axistake(via__array_namespace__when available), and correlated maps flatten the correlated axes, gather once, and reshape back. Constant maps become integer selections.The two strategies must agree, which the test matrix enforces directly.
New in
zarr_indexing.grid:EdgeDimensionGrid, a concreteDimensionGridLikebuilt from an explicit tuple of per-chunk sizes (offsets by cumsum, lookups bysearchsorted, so the grid need not be regular), plusdimension_grids_from_chunks(chunks, shape), which normalizes either chunk convention — validating that dask-convention sizes sum to the shape, and expanding a uniform chunk shape with a clipped tail chunk.Tests
packages/zarr-indexing/tests/test_lazy_array.py— one parametrized oracle plus one test per error case, per the repo's testing philosophy.oindex(unsorted + duplicates + multi-axis, negative, boolean axis),vindex(coordinates, broadcast pair, negatives, full-shape mask), and composed chains (basic-then-oindex, oindex-then-basic-on-another-axis, basic-then-basic, basic-then-vindex). Crossing the same cases over chunked and chunk-free wrappers is what verifies the two resolution strategies agree.EdgeDimensionGrid: all fourDimensionGridLikemethods against hand-computed values, including a size-1 axis, a single-chunk axis, and an irregular grid; plusdimension_grids_from_chunksnormalization and its threeValueErrors.shape/ndim/size/dtype/len/repr, view shape coming from the transform rather than the source, chunk discovery agreeing across all four flavours, view-aligned.chunksfor basic-slice views, and.chunks is Nonefor fancy views.importorskip):da.from_array(LazyArray(zarr_array))computes correctly with no translation ceremony, and the wrapper's.chunkshands straight back to dask. Verified locally with dask installed; it skips in CI since dask is not in the repo's test group.np.array(..., copy=False).379 passed, 1 skipped (tensorstore) for the package suite.
just lint,just typecheck(pyright, 0 errors), andjust docs-check(strict) are all clean, as are the prek hooks.Logistics
Stacked on #249 — base branch is
feat/zarr-indexing-package. Targets the 0.2.0 milestone; it does not block the 0.1.0 publish.