feat(zarr-indexing): LazyArray — generic lazy indexing over array-API arrays

.github/workflows/zarr-indexing.yml

@@ -63,7 +63,8 @@ jobs:
       - name: Install uv
         uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2
       - name: Run ruff
-        run: uvx ruff check .
+        # Pinned to the repo-wide ruff version (see pyproject.toml [dependency-groups] docs); bump together.
+        run: uvx ruff@0.15.22 check .
 
   pyright:
     name: pyright
@@ -85,7 +86,7 @@ jobs:
       - name: Sync test dependency group
         run: uv sync --group test --python 3.12
       - name: Run pyright
-        run: uv run --group test --with pyright pyright src
+        run: uv run --group test --with pyright pyright
 
   docs:
     name: docs

packages/zarr-indexing/CONTRIBUTING.md

@@ -0,0 +1,27 @@
+# Contributing to zarr-indexing
+
+Package-scoped development commands live in the [`justfile`](./justfile)
+(requires [just](https://github.com/casey/just)):
+
+```
+just test        # run the test suite (extra args go to pytest)
+just lint        # ruff, same invocation as CI
+just typecheck   # pyright, same invocation as CI
+just docs-check  # strict build of the docs site
+just check       # all of the above
+just docs-serve  # serve the docs site locally
+```
+
+Run them from this directory, or from anywhere in the repository as
+`just packages/zarr-indexing/<recipe>`.
+
+The test recipe runs against the workspace-root environment, because the
+chunk-resolution tests exercise this package against `zarr`'s chunk grids and
+`zarr` is deliberately not a dependency of this package.
+
+## License
+
+MIT
+
+The package lives at `packages/zarr-indexing` inside the
+[zarr-python](https://github.com/zarr-developers/zarr-python) repository.

packages/zarr-indexing/README.md

@@ -11,8 +11,15 @@ I/O until you explicitly read or write.
 
 Key types:
 
+- `LazyArray` — wraps any array-API-like array (NumPy, zarr, CuPy, ...) and adds
+  a `.lazy` accessor: `LazyArray(x).lazy[10:50, ::2].lazy.oindex[[3, 1, 1], :]`
+  composes a transform and returns a new view without reading data, and
+  `result()` materializes it
 - `IndexDomain` — a rectangular region of integer coordinates
 - `IndexTransform` — maps input coordinates to storage coordinates
+- `ChunkPlan` and `ChunkProjection` — lazily partition a selection over a
+  caller-selected grid and pair each chunk-local transform with its placement in
+  the request, without binding a storage backend or scheduler
 - `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single output
   dimension can depend on the input
 - `compose` — chain two transforms into one
@@ -27,27 +34,14 @@ repository and consumed by `zarr` to resolve array indexing operations.
 pip install zarr-indexing
 ```
 
-## Developing
+## Examples
 
-Package-scoped development commands live in the [`justfile`](./justfile)
-(requires [just](https://github.com/casey/just)):
+- [Lazy indexing a NumPy array](examples/lazy_indexing_numpy/README.md)
+- [Lazy indexing with Dask](examples/lazy_indexing_dask/README.md)
 
-```
-just test        # run the test suite (extra args go to pytest)
-just lint        # ruff, same invocation as CI
-just typecheck   # pyright, same invocation as CI
-just docs-check  # strict build of the docs site
-just check       # all of the above
-just docs-serve  # serve the docs site locally
-```
-
-Run them from this directory, or from anywhere in the repository as
-`just packages/zarr-indexing/<recipe>`.
-
-The test recipe runs against the workspace-root environment, because the
-chunk-resolution tests exercise this package against `zarr`'s chunk grids and
-`zarr` is deliberately not a dependency of this package.
-
-## License
+## Contributing
 
-MIT
+Development commands, the test suite and the docs build are described in
+[CONTRIBUTING.md](https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/CONTRIBUTING.md)
+in the repository. Issues and pull requests go to
+[zarr-developers/zarr-python](https://github.com/zarr-developers/zarr-python).

packages/zarr-indexing/changes/267.feature.md

@@ -0,0 +1,41 @@
+Added `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, CuPy, ...) TensorStore-style lazy indexing. `LazyArray(x).lazy[...]`, `.lazy.oindex[...]`, and `.lazy.vindex[...]` compose an `IndexTransform` and return a new view without reading data; `result()` materializes it. Selections use positional NumPy semantics (zero-based within the current view, negatives wrap, masks must match the view shape, scalar integers drop their axis, and a partial `vindex` places its gathered dimensions where NumPy would), which the new `zarr_indexing.boundary` module translates into the transform algebra's literal coordinates. `__getitem__` reads eagerly, so a `LazyArray` can be used as a source for `dask.array.from_array`; any NumPy function given a view materializes it through `__array__` first (Python's arithmetic operators do not — the wrapper defines no arithmetic dunders, so `view + 1` raises `TypeError`). `__dask_tokenize__`, `__len__`, `__iter__`, `__bool__`, `__int__`, `__float__`, and `__index__` complete the surface, and a wrapper pickles as long as its base does.
+
+A read is broken up along a **partitioning** — a grid of boxes, discovered from the wrapped array's `read_chunk_sizes` or `chunks` (those attribute names belong to the wrapped array; this API refers only to parts). `LazyArray.parts()` walks those boxes as they fall through the view, yielding a `Partition` with the box's base coordinates, a resolvable `LazyArray` for the cells of the view inside it, where those cells belong in the result, and whether the box is fully covered. `result()` is the assembly of that walk, so one partitioner and one lowering engine serve every read. `LazyArray.with_parts()` chooses a different partitioning — uniform box shape, explicit per-axis sizes, or `None` for a single whole-array part — without touching the data or the view, and never changes what `result()` returns.
+
+`LazyArray.is_box` reports whether a selection is rectangular — true exactly when the composed transform carries no `index_array` output map, which basic indexing preserves under composition and any `oindex`/`vindex`/mask ends permanently. `LazyArray.bounding_box()` gives the storage region a selection touches in either case — dense only when every stride is 1, a lattice within the hull otherwise, a superset for a query, and `None` for an empty selection — and `LazyArray.strides()` gives the per-dimension step that, with the bounding box, describes a box completely. `Partition` gained `box`, the part's box in the wrapped array's global coordinates, since `Partition.view.bounding_box()` is part-local and cannot distinguish two parts. A new [design notes](https://zarr-indexing.readthedocs.io/en/latest/design-notes/) page covers the relationship to TensorStore, the box/query category, and the current scope limits.
+
+`zarr_indexing.grid` gained `EdgeDimensionGrid`, a concrete `DimensionGridLike` built from explicit per-part sizes, and `dimension_grids_from_chunks`, which normalizes either convention into one grid per axis.
+
+Fixed an integer index applied to an axis a previous orthogonal or vectorized selection had already indexed. The composed `ArrayMap` became all-singleton but kept an `input_dimension` naming the axis the integer had just removed, which after renumbering aliased a different axis. Such a map now collapses to a `ConstantMap` at composition time, and `array_map_dependent_axis` (promoted from a private helper) returns `None` for "no axis" instead of falling back to that stale binding.
+
+Fixed the domain layout and partitioned resolution of a vectorized selection whose coordinate arrays are not on the leading axes — `vindex[..., i, j]`, `vindex[..., mask]`. The gathered dimensions now follow NumPy's placement rule (in the spot the advanced indices occupied when they are adjacent, leading when a slice separates them), and the per-part gather is realigned to match the scatter indices instead of raising a shape-mismatch `ValueError`.
+
+
+Negative-step slices are supported, following the merged ndsel 1.0-draft.2 (PR #2) and TensorStore 0.1.84: `arr[::-1]`, `arr[5:1:-2]`, and reversal composed over an already-strided or already-gathered view. One desugaring rule covers both signs — omitted bounds resolve on the side the traversal starts and stops, the source interval is `[start, stop)` going up and `[stop + 1, start + 1)` going down, the origin is `trunc(start / step)` for either sign — and a reversed interval is an error rather than a silently empty selection. The reason code `negative_step_unsupported` is retired with the spec change, and the vendored conformance corpus is re-vendored from ndsel `92d6a32` to match. A reversing slice normally produces a negative domain origin (`[::-1]` on a length-20 axis gives `[-19, 1)`), because the result stays anchored to the source coordinate frame; `LazyArray` re-bases every view to origin 0, so its positional dialect is unaffected.
+
+`LazyArray` now negotiates with the array it wraps instead of assuming what it accepts. A wrapper carries an `IndexingSupport` level — `BASIC`, `OUTER`, `OUTER_1VECTOR`, or `VECTORIZED`, the taxonomy and member names of [xarray's `IndexingSupport`](https://github.com/pydata/xarray/blob/main/xarray/core/indexing.py) — and every read is split into the largest part of the selection that level can express, asked of the source in a single call, and a remainder applied to the block that comes back with NumPy. The level is detected at construction, declaration first: a source's own `__zarr_indexing_support__` attribute wins when it names one of the four levels, and otherwise the level is inferred — `VECTORIZED` for a NumPy array or an object carrying zarr's `oindex`/`vindex` pair, `BASIC` for everything else, the only assumption that is always correct. `LazyArray.with_indexing_support()` overrides both, and `LazyArray.indexing_support` reports the level in force. The level decides how much data crosses the boundary, never what `result()` returns: declaring `BASIC` on a NumPy array answers exactly what `VECTORIZED` answers, having read the enclosing slab and gathered from it instead.
+
+The negotiation applies within a partition as well as to a whole-array read, so a part costs one request rather than a whole-box read followed by an in-memory selection: an orthogonal selection goes to `oindex` in one call instead of one `take` per fancy axis, a coordinate gather goes to `vindex` instead of reading the outer product the points span, and a source with neither accessor is asked only for slices. A multi-array outer request is only ever sent to an `oindex` accessor, because a bare `__getitem__` key with two arrays means an outer product to HDF5 and a correlated gather to NumPy and nothing about a source says which reading applies; a source declaring `OUTER` without an `oindex` accessor is given one array axis per request.
+
+Fixed an `oindex`/`vindex` step whose entries are all slices applied to a view that already has a fancy-indexed axis. Such a step carries no coordinates — it narrows the view's own axes and composes like basic indexing — but its slices were applied positionally to every axis of the existing index array, including the singleton axes the array only broadcasts over. A slice starting past 0 therefore truncated the index array to size 0, and a view with no coordinates left resolves to no parts at all, so `result()` returned an unwritten buffer. Index-array reindexing is now dependency-aware on both the basic and the fancy path; a genuine fancy-after-fancy step (coordinates aimed at a broadcast axis) still raises `NotImplementedError`. Relatedly, an axis of length 0 no longer counts as an axis an index array varies over, so an empty orthogonal selection is no longer misread as a correlated one.
+
+Fixed `parts()` raising on a view emptied by a slice over an axis of extent 1: a correlated selection of a single point normalizes to an all-singleton index array, so emptying the domain leaves the array at size 1 and the resolver looked for a chunk to read. An empty domain now yields no parts, matching what `result()` already returned.
+
+Fixed negative-stride chunk projection: the endpoints were swapped while the step stayed negative, producing a slice that selects nothing where the reversed axis was meant. Fixed `compose()` evaluating an inner index array over `range(size)` rather than over the outer domain's own range, which resolved every coordinate to the wrong cell whenever the outer domain did not start at 0 — which a step-1 slice and a negative-step slice both produce routinely.
+
+Fixed a domain dimension no output map depends on, which a `vindex` coordinate array with a singleton broadcast axis leaves behind once a later basic index consumes the axis the array varies over. The partition walk emitted an out-selection of lower rank than the output buffer (so a partitioned read placed each part against the leading axes and broadcast the rest), the lowering engine restored the axis as a singleton whatever the domain said (so an emptied one fabricated a row), and the correlated gather combined a broadcast domain shape with an un-broadcast index array. All three now count such an axis.
+
+`result()` never returns memory shared with the wrapped array. An unpartitioned read of a basic selection lowers to plain slicing and so used to hand back a view of the source, which a partitioned read never did; `numpy.array(view, copy=True)` inherited the alias. `result()` also verifies that the parts covered the output before returning it, so a partition walk that leaves a gap raises instead of returning the uninitialized buffer underneath.
+
+A source meeting exactly the documented `BASIC` floor — `shape`, `dtype`, and a `__getitem__` taking integers and slices — is now enough: the block it returns is coerced before the residual selection is applied to it, rather than being handed `take` and `reshape` itself. A `numpy.ma` source keeps its mask through a partitioned read. `numpy.matrix` is refused at construction, since it never reduces rank and so cannot honor a view's shape. A `__zarr_indexing_support__` declaration holding a member of a look-alike enum (xarray's own `IndexingSupport`) is honored rather than discarded in favor of a more permissive inference. `Partition.is_complete` is true for a reversing view, which reads every cell of its box. `with_parts` accepts `(0,)` for a zero-length axis, as it already accepted `()`.
+Fixed the rank of the parts of a correlated view narrowed to a single point — `vindex[..., i, j, k][0]` and the like. Intersecting a correlated transform with a part's bounds collapses the surviving broadcast block into one input axis, which it did even when the block was already rank 0, so the part's `view.result()` came back with rank 1 where the view had rank 0 and the documented `out[part.out_selection] = part.view.result()` assembly raised `ValueError`. A rank-0 block now stays rank 0: it survives whole or the intersection is empty.
+
+A new `zarr_indexing.testing` subpackage, behind a `testing` extra (`pip install zarr-indexing[testing]`), carries the Hypothesis machinery this package tests itself with. `ChainedIndexingStateMachine` composes indexing steps onto a `LazyArray` wrapping an array you supply — basic, orthogonal, and vectorized selections, each applied to the view the last one produced — and after every step checks the view's shape, its `result()`, and the assembly of its `parts()` against a NumPy array holding the same values. Point it at your own array by overriding one method; `data`, `partitionings`, and `supports` are class attributes. A `declare_support` rule draws an `IndexingSupport` level and applies it mid-chain, which pins the property the level is documented to have and nothing asserted before: the level decides how much data crosses the boundary, never what the read means. `zarr_indexing.testing.strategies` exports the selection strategies on their own, for a project that has its own harness. Nothing outside this subpackage imports Hypothesis.
+
+That machine replaces the seeded sweep that read chained selections through `parts()`, and asserts the documented assembly literally: a part's values must arrive at exactly the shape its `out_selection` addresses, not merely one that broadcasts into it. The defect above was reachable by the sweep it replaces, but was absorbed by a reshape in `result()` and mirrored into the sweep rather than read as a failure; the invariant makes that impossible to paper over, and a failure now shrinks to the shortest chain that causes it.
+
+`transform_from_canonical` now rejects a non-integer `index_array` with an `NdselError` carrying the `invalid_json` reason code, rather than silently truncating a float array to its integer parts, coercing booleans to 0/1, or leaking NumPy's own conversion error for strings.
+
+Added source-independent chunk planning. `plan_chunks(transform, grids)` returns a lazy, reusable `ChunkPlan`; each `ChunkProjection` pairs a chunk-local storage transform with a transform back to the request over one shared synthetic domain. The same representation covers basic, orthogonal, and vectorized indexing, carries global chunk bounds and conservative full/partial/unknown coverage, and leaves I/O, buffering, and scheduling to zarr, dask, napari, or another consumer. `LazyArray.parts()` now exposes this projection directly and derives its NumPy placement from the paired transform. The provisional tuple resolver and selector bridge have been removed.
+
+Added a visual indexing guide with executable examples that follows lazy views from a NumPy selection through chunk projections, explains arbitrary-origin coordinates and prepended chunks, and includes diagrams tested in both light and dark themes.

packages/zarr-indexing/changes/272.bugfix.md

@@ -0,0 +1,9 @@
+An adversarial review of the whole package found, and this fixes, several
+defects at its boundaries: `result()` and `__array__(copy=True)` could hand back
+a live view of a source that merely stored its data in NumPy; the wire format
+emitted a document nothing could load for a selection that selects nothing, and
+its domain loader validated nothing; chunk-selection lowering described a
+transposed block in two separate cases; a map derived from a vectorized
+selection carried a stale `input_dimension`, which made one view's answer depend
+on how it was partitioned; and `oindex` over a correlated view applied NumPy's
+vectorized rule instead of the outer product.

packages/zarr-indexing/changes/273.misc.md

@@ -0,0 +1,6 @@
+`with_parts` is now three named methods — `with_parts`, `with_parts_per_axis`
+and `unpartitioned` — instead of one parameter whose meaning was decided by the
+type of what it was given. `Partition.array` is `Partition.view`, no longer the
+inverse of `LazyArray.array`. `ArrayMap`, `IndexTransform` and `Partition` can
+be compared and hashed, which `frozen=True` had implied and neither could do.
+`LazyArray.base_shape` says which shape a partitioning is expressed in.

packages/zarr-indexing/docs/api/boundary.md

@@ -0,0 +1,5 @@
+---
+title: boundary
+---
+
+::: zarr_indexing.boundary