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

docs/superpowers/specs/2026-08-01-chunk-projection-design.md

@@ -0,0 +1,154 @@
+# Chunk Projection Design
+
+## Goal
+
+Replace the provisional chunk-resolution tuple with a stable, source-independent
+projection model based on TensorStore's paired-transform design. The result must be useful
+to zarr-python's codec pipeline and to external schedulers such as Dask or napari without
+embedding source objects, NumPy-specific scatter indices, caching, or execution policy.
+
+This package is greenfield. There are no dependent projects and the existing public API
+does not require compatibility shims.
+
+## Decision
+
+`IndexTransform` remains the logical selection: it maps the caller-visible request domain
+to storage coordinates. Chunk planning partitions that transform over a caller-supplied
+grid and yields one `ChunkProjection` per touched grid cell.
+
+Each projection uses a synthetic cell domain and carries two transforms with exactly that
+input domain:
+
+- `chunk_transform`: synthetic cell coordinates to chunk-local storage coordinates.
+- `cell_transform`: synthetic cell coordinates to the original request domain.
+
+This replaces the current `(chunk_coords, sub_transform, out_indices)` result. In
+particular, `out_indices` is not exposed: it is an incomplete second representation whose
+shape changes between basic, orthogonal, and vectorized indexing. A transform represents
+all three uniformly and composes with the rest of the package.
+
+## Public API
+
+```python
+from dataclasses import dataclass
+from typing import Literal
+
+ChunkCoverage = Literal["full", "partial", "unknown"]
+
+
+@dataclass(frozen=True, slots=True)
+class ChunkProjection:
+    chunk_coords: tuple[int, ...]
+    chunk_domain: IndexDomain
+    chunk_transform: IndexTransform
+    cell_transform: IndexTransform
+    coverage: ChunkCoverage
+
+
+@dataclass(frozen=True, slots=True)
+class ChunkPlan:
+    transform: IndexTransform
+    dimension_grids: tuple[DimensionGridLike, ...]
+
+    def projections(self) -> Iterator[ChunkProjection]: ...
+    def __iter__(self) -> Iterator[ChunkProjection]: ...
+
+
+def plan_chunks(
+    transform: IndexTransform,
+    dimension_grids: Sequence[DimensionGridLike],
+) -> ChunkPlan: ...
+```
+
+`ChunkPlan` is re-iterable and cheap to construct. It retains only the logical request and
+the selected grid; projections remain lazy. The caller chooses the grid appropriate for
+the operation. A zarr integration should pass the no-read-amplification grid for reads and
+the atomic/no-write-amplification grid for writes. The indexing package does not invent a
+read/write layout abstraction before it has one to consume.
+
+`iter_chunk_transforms` and its tuple result cease to be public. The existing intersection
+machinery may continue to use temporary survivor arrays internally, but they must be
+converted into `cell_transform` before crossing the chunk-resolution boundary.
+
+## Projection Invariants
+
+For every projection:
+
+1. `chunk_transform.domain == cell_transform.domain`.
+2. `chunk_transform` maps every cell into `[0, chunk_shape)`.
+3. `cell_transform` is one-to-one and maps every cell into the plan transform's input
+   domain.
+4. Composing the plan transform with `cell_transform`, then subtracting the chunk origin,
+   is equivalent to `chunk_transform`.
+5. Across one plan, `cell_transform` ranges are disjoint and their union is the request
+   domain. Storage coordinates may repeat because fancy indexing may request the same
+   element more than once.
+
+Basic and orthogonal intersections retain their natural rank. Correlated intersections may
+collapse their broadcast block to a synthetic points dimension, just as the existing
+resolver does; `cell_transform` maps those synthetic points back to the original broadcast
+axes.
+
+## Coverage
+
+Coverage is computed against the grid supplied to `plan_chunks`:
+
+- `"full"`: the projection is proven to address every cell of the grid cell exactly once.
+- `"partial"`: it is proven not to cover the entire grid cell.
+- `"unknown"`: exact proof would require materializing or sorting fancy coordinates.
+
+Affine unit-stride projections can be classified exactly. Fancy projections are
+`"unknown"` unless the resolver already has a cheap exact proof. Consumers may skip a
+read-modify-write only for `"full"`; both other values are conservative.
+
+Coverage is an optimization, not part of selection correctness. A caller planning against
+a shard/write grid therefore receives shard coverage, while one planning against an inner
+read grid receives inner-chunk coverage.
+
+## `LazyArray` Integration
+
+`LazyArray.parts()` constructs a `ChunkPlan` from the view transform and discovered read
+partitioning. Each `Partition` contains its `ChunkProjection` and a source-bound `LazyArray`
+view for resolving that projection. Source binding remains above the pure plan.
+
+NumPy placement used by `LazyArray.result()` is derived internally from
+`projection.cell_transform`. It is not stored in `ChunkProjection` and is not the public
+chunk-planning contract. This keeps execution conveniences available without forcing
+NumPy indexing tuples on Dask, napari, codec, or non-NumPy consumers.
+
+Caching, scheduling, cancellation, locks, and I/O are deliberately absent from both
+`ChunkPlan` and `ChunkProjection`. Those belong to source/execution adapters. A later async
+executor can stream `plan.projections()` and stop between projections without changing the
+planning model.
+
+## Error Handling
+
+`plan_chunks` validates that the number of dimension grids equals the transform output
+rank. Construction otherwise performs no chunk walk. Errors that depend on coordinates or
+grid boundaries arise while iterating projections and identify the offending dimension or
+chunk coordinate.
+
+`ChunkProjection` validates its shared-domain invariant at construction. Violations are
+internal errors and raise `ValueError` with both domains in the message.
+
+Empty request domains yield an empty projection iterator.
+
+## Testing
+
+Tests exercise observable transform semantics rather than dataclass field presence alone.
+
+- Basic, reversed, orthogonal, correlated, repeated, scalar, empty, and irregular-grid
+  selections verify all projection invariants.
+- For each selection, applying `chunk_transform` to the chunk and placing the result through
+  `cell_transform` reconstructs the direct selection exactly.
+- Tests verify that cell ranges tile the request domain once even when storage coordinates
+  repeat.
+- Coverage tests exercise full, partial, and unknown outcomes, including clipped edge cells.
+- A plan is iterated twice to prove it is reusable and lazy.
+- `LazyArray.parts()` and `result()` tests prove source binding and assembly use the new
+  projection contract.
+- Public export and documentation tests ensure the tuple resolver is no longer advertised.
+
+The implementation follows test-driven development: each new public behavior is first
+introduced by a focused failing test, then implemented minimally, then covered by the full
+`zarr-indexing` test suite and lint/type checks.

docs/user-guide/examples/lazy_indexing_dask.md

@@ -0,0 +1,7 @@
+--8<-- "examples/lazy_indexing_dask/README.md"
+
+## Source Code
+
+```python exec="false" reason="pymdownx snippet include directive, not python source"
+--8<-- "examples/lazy_indexing_dask/lazy_indexing_dask.py"
+```

docs/user-guide/examples/lazy_indexing_numpy.md

@@ -0,0 +1,7 @@
+--8<-- "examples/lazy_indexing_numpy/README.md"
+
+## Source Code
+
+```python exec="false" reason="pymdownx snippet include directive, not python source"
+--8<-- "examples/lazy_indexing_numpy/lazy_indexing_numpy.py"
+```

examples/lazy_indexing_dask/README.md

@@ -0,0 +1,53 @@
+# Lazy Indexing with Dask
+
+This example demonstrates how to use `zarr_indexing.LazyArray` with Dask, both as
+an array Dask can wrap and as a source of independent tasks, and compares the two
+ways of deferring an indexing operation.
+
+The example shows how to:
+
+- Pass a `LazyArray` — over a Zarr array or over a view of one — to
+  `dask.array.from_array`
+- Build one Dask task per partition from `parts()`, compute them in parallel, and
+  place each result with the partition's `out_selection`
+- Read `is_complete` to tell which partitions cover a stored chunk completely
+- Rely on `__dask_tokenize__`, so that equal selections produce equal tokens and
+  Dask can cache and deduplicate the work
+- Measure what a task graph costs for indexing-only work, against composing the
+  same selections into one transform
+
+A `LazyArray` exposes no `chunks` attribute, so `dask.array.from_array` chooses
+its own block size unless one is given. The partitioning that `parts()` reports
+is discovered from the wrapped array and is independent of Dask's blocks.
+
+## Choosing Between Them
+
+If Dask is doing arithmetic across chunks, reductions, rechunking, or distributed
+execution, it is the right tool, and its task graph is what makes that work.
+
+If Dask is used *only* to defer indexing — take a view now, read it later, with
+no computation in between — then the graph is overhead. Dask slices the chunk
+grid on every indexing operation and records another layer, so composing
+selections costs time proportional to both the depth of the chain and the number
+of chunks in the array, and reading walks what was accumulated. `LazyArray`
+composes each selection into the single transform it already holds, so composing
+is independent of the depth of the chain, and reading enumerates only the
+partitions the selection touches. The last test in this example prints both, and
+the gap widens with the number of chunks and the number of selections.
+
+## Running the Example
+
+The script declares its dependencies inline
+([PEP 723](https://peps.python.org/pep-0723/)), so the easiest way to run it is
+with [uv](https://docs.astral.sh/uv/), which installs them automatically:
+
+```bash
+uv run examples/lazy_indexing_dask/lazy_indexing_dask.py
+```
+
+Alternatively, run it with plain Python, in which case you must first install
+`zarr`, `zarr-indexing`, `dask[array]`, `numpy`, and `pytest` yourself:
+
+```bash
+python examples/lazy_indexing_dask/lazy_indexing_dask.py
+```