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

[d-v-b]

🤖 AI text below 🤖

Adds a generic LazyArray wrapper to packages/zarr-indexing. Nothing in src/zarr is touched.

What it is

LazyArray wraps any array-API-like array — NumPy, zarr, CuPy, anything with shape, dtype, and __getitem__ — forwards the array-API attributes, and adds a .lazy accessor for TensorStore-style deferred indexing:

view = LazyArray(zarr_array).lazy[10:50, ::4].lazy.oindex[[3, 0, 0], :]
view.shape      # known without touching data
view.result()   # one resolution pass

.lazy[...], .lazy.oindex[...], and .lazy.vindex[...] compose an IndexTransform and return a new LazyArray. __getitem__ (no .lazy) is eager — mirroring zarr.Array's design, and what makes the wrapper a drop-in duck array for dask.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.lazy deliberately exposes that dialect, because a zarr view's coordinates are meant to stay comparable with the parent array's.

LazyArray cannot 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 (minding isinstance(True, int)), masks must match the view shape exactly, and integer arrays are bounds-checked positionally. The new zarr_indexing.boundary module is the generic (zarr-free) translation from positions to literal coordinates before selection_to_transform. It duplicates what zarr/core/array.py does 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 guarded getattr since zarr raises an AttributeError-flavoured LazyViewError on 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.

• Chunked source: iterate iter_chunk_transforms(transform, grids); read each touched chunk once with plain basic slicing, convert the local transform via sub_transform_to_selections, and scatter into the output buffer (flat scatter for correlated maps). This mirrors zarr.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.
• Unchunked source: one-shot lowering of the composed transform to array operations — pure-slice transforms become a single basic __getitem__, orthogonal ArrayMaps become successive per-axis take (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 concrete DimensionGridLike built from an explicit tuple of per-chunk sizes (offsets by cumsum, lookups by searchsorted, so the grid need not be regular), plus dimension_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.

• Oracle matrix: 16 selection cases x 4 source flavours (NumPy unchunked, NumPy with declared dask-convention chunks, NumPy with a uniform chunk-shape argument, and an auto-discovered zarr array), every result compared against NumPy computed positionally on the same data. Cases cover basic (strided, negative, ellipsis, int-drop, empty, all-scalar), 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 four DimensionGridLike methods against hand-computed values, including a size-1 axis, a single-chunk axis, and an irregular grid; plus dimension_grids_from_chunks normalization and its three ValueErrors.
• Forwarding: 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 .chunks for basic-slice views, and .chunks is None for fancy views.
• dask interop (importorskip): da.from_array(LazyArray(zarr_array)) computes correctly with no translation ceremony, and the wrapper's .chunks hands straight back to dask. Verified locally with dask installed; it skips in CI since dask is not in the repo's test group.
• Errors: boolean scalar, wrong-shape mask, out-of-bounds scalar (including inside a view, where the bounds are the view's), out-of-bounds index array, too many indices, chunks that do not sum to the shape, and np.array(..., copy=False).

379 passed, 1 skipped (tensorstore) for the package suite. just lint, just typecheck (pyright, 0 errors), and just 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.