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

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

@@ -0,0 +1,5 @@
+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` is a drop-in source for `dask.array.from_array`. Chunking is taken from an explicit `chunks=` argument, then `read_chunk_sizes`, then `chunks`; a chunked source resolves through `iter_chunk_transforms` so each touched chunk is read exactly once, and an unchunked source lowers to a single pass of array operations. `zarr_indexing.grid` gained `EdgeDimensionGrid`, a concrete `DimensionGridLike` built from explicit per-chunk sizes, and `dimension_grids_from_chunks`, which normalizes either chunk 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) reports "no axis" as `None` instead of falling back to that stale binding.
+
+Fixed the domain layout and chunked 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-chunk gather is realigned to match the scatter indices instead of raising a shape-mismatch `ValueError`.

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

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

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

@@ -28,7 +28,18 @@ and the wire format built on top of it.
   current codec pipeline expects)
 - [`zarr_indexing.grid`](grid.md) — `DimensionGridLike`, the Protocol
   describing the narrow chunk-grid surface chunk resolution consumes, so that
-  nothing here imports `zarr`
+  nothing here imports `zarr`, plus `EdgeDimensionGrid` and
+  `dimension_grids_from_chunks`, a concrete per-axis grid for callers with no
+  zarr grid to hand
+
+**Lazy arrays**
+
+- [`zarr_indexing.lazy_array`](lazy_array.md) — `LazyArray`, a wrapper that
+  gives any array-API-like array (NumPy, zarr, CuPy, …) a `.lazy` accessor for
+  TensorStore-style deferred indexing, resolving chunk-aware when the wrapped
+  array declares chunks
+- [`zarr_indexing.boundary`](boundary.md) — the translation between NumPy's
+  positional dialect and the transform algebra's literal coordinates
 
 **The ndsel wire format** (see [the guide](../ndsel.md))
 

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

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

packages/zarr-indexing/docs/index.md

@@ -142,6 +142,76 @@ and translated into chunk-local coordinates — exactly what a codec pipeline
 needs to decode that chunk and scatter the result. `out_indices` carries the
 output scatter indices for array selections, and is `None` for basic indexing.
 
+## Lazy views over any array
+
+[`LazyArray`](api/lazy_array.md) packages all of the above into a wrapper you
+can put around an array you already have — NumPy, zarr, CuPy, anything with
+`shape`, `dtype`, and `__getitem__`. Its `.lazy` accessor composes transforms
+instead of reading, and `result()` resolves the composed view in one pass:
+
+```python
+import numpy as np
+import zarr
+
+from zarr_indexing import LazyArray
+
+arr = zarr.create_array({}, shape=(100, 80), chunks=(30, 40), dtype="int32")
+arr[:] = np.arange(8000).reshape(100, 80)
+
+lazy = LazyArray(arr)
+lazy.chunks  # ((30, 30, 30, 10), (40, 40))
+
+view = lazy.lazy[10:50, ::4].lazy.oindex[[3, 0, 0], :]
+view.shape   # (3, 20)
+repr(view)   # '<LazyArray Array shape=(3, 20) dtype=int32 view={ {13, 10, 10}, [0, 80) step 4 }>'
+
+view.result()[:, :4]
+# array([[1040, 1044, 1048, 1052],
+#        [ 800,  804,  808,  812],
+#        [ 800,  804,  808,  812]], dtype=int32)
+```
+
+Unlike the transform algebra underneath it — where indices are literal
+coordinates in a view's own (possibly shifted) domain — `LazyArray` speaks
+**positional NumPy semantics**: a view is re-zeroed, `-1` is the last element,
+and boolean masks must match the view's shape. `.lazy.oindex` and `.lazy.vindex`
+give the orthogonal and vectorized flavours:
+
+```python
+LazyArray(np.arange(12).reshape(3, 4)).lazy.vindex[[0, 2], [1, 3]].result()
+# array([ 1, 11])
+```
+
+Two further NumPy rules hold in every mode. A **scalar integer drops its axis**
+— it is a basic index wherever it appears, applied before any advanced index
+rather than broadcast against one — and **advanced indices land where NumPy puts
+them**, next to the axes they replaced when they are adjacent and leading when a
+slice separates them:
+
+```python
+x = np.arange(12).reshape(3, 4)
+
+LazyArray(x).lazy.oindex[1].result()
+# array([4, 5, 6, 7])
+
+LazyArray(x).lazy.oindex[[2, 0], 1].result()
+# array([9, 1])
+
+LazyArray(x).lazy.vindex[..., [3, 0]].result()
+# array([[ 3,  0],
+#        [ 7,  4],
+#        [11,  8]])
+```
+
+Use a length-1 list (`oindex[[1]]`) where you want to keep an axis.
+
+Chunking is discovered from the wrapped array (`read_chunk_sizes`, then
+`chunks`) or declared with `chunks=`. When it is known, a view resolves through
+[chunk resolution](api/chunk_resolution.md) so every touched chunk is read
+exactly once; when it is not, the view lowers to a single pass of array
+operations. Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a
+`LazyArray` a drop-in source for `dask.array.from_array`.
+
 ## Reference
 
 - [The ndsel wire format](ndsel.md)

packages/zarr-indexing/mkdocs.yml

@@ -22,6 +22,8 @@ nav:
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.composition</code>': api/composition.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.chunk_resolution</code>': api/chunk_resolution.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.grid</code>': api/grid.md
+      - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.lazy_array</code>': api/lazy_array.md
+      - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.boundary</code>': api/boundary.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.json</code>': api/json.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.messages</code>': api/messages.md
       - '<code class="doc-symbol doc-symbol-toc doc-symbol-module"></code> <code>zarr_indexing.errors</code>': api/errors.md

packages/zarr-indexing/src/zarr_indexing/__init__.py

@@ -28,7 +28,11 @@
 )
 from zarr_indexing.composition import compose
 from zarr_indexing.domain import IndexDomain
-from zarr_indexing.grid import DimensionGridLike
+from zarr_indexing.grid import (
+    DimensionGridLike,
+    EdgeDimensionGrid,
+    dimension_grids_from_chunks,
+)
 from zarr_indexing.json import (
     IndexDomainJSON,
     IndexTransformJSON,
@@ -40,9 +44,14 @@
     transform_from_canonical,
     transform_to_canonical,
 )
+from zarr_indexing.lazy_array import LazyArray
 from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel
 from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap
-from zarr_indexing.transform import IndexTransform, selection_to_transform
+from zarr_indexing.transform import (
+    IndexTransform,
+    array_map_dependent_axis,
+    selection_to_transform,
+)
 
 __version__ = version("zarr-indexing")
 
@@ -51,15 +60,19 @@
     "ConstantMap",
     "DimensionGridLike",
     "DimensionMap",
+    "EdgeDimensionGrid",
     "IndexDomain",
     "IndexDomainJSON",
     "IndexTransform",
     "IndexTransformJSON",
+    "LazyArray",
     "NdselError",
     "OutputIndexMap",
     "OutputIndexMapJSON",
     "__version__",
+    "array_map_dependent_axis",
     "compose",
+    "dimension_grids_from_chunks",
     "index_domain_from_json",
     "index_domain_to_json",
     "index_transform_from_json",