From ede1eaa5b9fd4539ce2c95e62d77169c37ddb693 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 11:53:21 +0200 Subject: [PATCH 01/32] =?UTF-8?q?feat(zarr-indexing):=20LazyArray=20?= =?UTF-8?q?=E2=80=94=20generic=20lazy=20indexing=20over=20array-API=20arra?= =?UTF-8?q?ys?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- packages/zarr-indexing/changes/267.feature.md | 11 + packages/zarr-indexing/docs/api/boundary.md | 5 + packages/zarr-indexing/docs/api/index.md | 13 +- packages/zarr-indexing/docs/api/lazy_array.md | 5 + packages/zarr-indexing/docs/design-notes.md | 165 +++ packages/zarr-indexing/docs/index.md | 147 +++ packages/zarr-indexing/mkdocs.yml | 3 + .../src/zarr_indexing/__init__.py | 18 +- .../src/zarr_indexing/boundary.py | 322 +++++ .../zarr-indexing/src/zarr_indexing/grid.py | 232 +++- .../src/zarr_indexing/lazy_array.py | 1145 +++++++++++++++++ .../src/zarr_indexing/transform.py | 174 ++- .../zarr-indexing/tests/test_lazy_array.py | 1090 ++++++++++++++++ 13 files changed, 3278 insertions(+), 52 deletions(-) create mode 100644 packages/zarr-indexing/changes/267.feature.md create mode 100644 packages/zarr-indexing/docs/api/boundary.md create mode 100644 packages/zarr-indexing/docs/api/lazy_array.md create mode 100644 packages/zarr-indexing/docs/design-notes.md create mode 100644 packages/zarr-indexing/src/zarr_indexing/boundary.py create mode 100644 packages/zarr-indexing/src/zarr_indexing/lazy_array.py create mode 100644 packages/zarr-indexing/tests/test_lazy_array.py diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md new file mode 100644 index 0000000000..40744cbfd9 --- /dev/null +++ b/packages/zarr-indexing/changes/267.feature.md @@ -0,0 +1,11 @@ +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`; every other NumPy operation materializes the view through `__array__`. `__dask_tokenize__`, `__len__`, `__iter__`, `__bool__`, `__int__`, `__float__`, and `__index__` round out 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` (the source's vocabulary; the wrapper's own surface speaks only of 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 and nothing else, 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 either way — dense only when every stride is 1, a lattice within the hull otherwise, a mere superset for a query, and `None` for an empty selection — and `LazyArray.strides()` gives the per-dimension step that completes the description of a box. `Partition` gained `box`, the part's box in the wrapped array's global coordinates, since `Partition.array.bounding_box()` is part-local and cannot tell two parts apart. 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) reports "no axis" as `None` 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`. diff --git a/packages/zarr-indexing/docs/api/boundary.md b/packages/zarr-indexing/docs/api/boundary.md new file mode 100644 index 0000000000..4f9fa99531 --- /dev/null +++ b/packages/zarr-indexing/docs/api/boundary.md @@ -0,0 +1,5 @@ +--- +title: boundary +--- + +::: zarr_indexing.boundary diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index b9a58b70fa..09b6bb6b32 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/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, plus `Partition` and the `parts()` / + `with_parts()` pair that decide the boxes a read is broken into +- [`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)) diff --git a/packages/zarr-indexing/docs/api/lazy_array.md b/packages/zarr-indexing/docs/api/lazy_array.md new file mode 100644 index 0000000000..1ad928d84d --- /dev/null +++ b/packages/zarr-indexing/docs/api/lazy_array.md @@ -0,0 +1,5 @@ +--- +title: lazy_array +--- + +::: zarr_indexing.lazy_array diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md new file mode 100644 index 0000000000..6cca05b163 --- /dev/null +++ b/packages/zarr-indexing/docs/design-notes.md @@ -0,0 +1,165 @@ +--- +title: Design notes +--- + +# Design notes + +Three things that are easier to explain once than to infer from the API: where +this library sits relative to TensorStore, why rectangular selections are a +category rather than a fast path, and what is deliberately not implemented yet. + +## Relationship to TensorStore + +The core is [TensorStore's](https://google.github.io/tensorstore/index_space.html) +index-transform model, reimplemented in Python against NumPy. The parts that +are the same are the same on purpose: + +- **The model.** An `IndexTransform` pairs an input `IndexDomain` — a + rectangular region with an explicit, possibly non-zero origin — with one + output index map per storage dimension, in the three flavours TensorStore + defines: constant, single-input-dimension (affine), and index array. + Composition is the single operation that stacks views. +- **Slice semantics.** Slice bounds are literal domain coordinates: no + clamping, no negative wrapping, non-empty intervals must be contained in the + domain, and a strided slice's domain origin is `trunc(start/step)` rounded + toward zero. Every one of those rules was executed against tensorstore 0.1.84 + and is pinned in `tests/test_tensorstore_parity.py`. +- **The wire format.** A canonical [ndsel](ndsel.md) transform body is, + field-for-field, a TensorStore `IndexTransform` minus the `kind` + discriminator, and `tests/test_ndsel_tensorstore.py` loads our bodies into + `tensorstore.IndexTransform(json=...)` and round-trips them back through our + engine layer. + +Index arrays are the one place the representations differ in practice. Ours are +normalized to the transform's full input rank — full-sized on the axis a map +varies over, singleton elsewhere — so the orthogonal/vectorized distinction is +derivable from the shape. TensorStore permits a lower-rank array that broadcasts +against the input domain; we accept those on load and normalize them. +`ArrayMap` also records the input dimension an *orthogonal* (`oindex`) array +varies over, a field TensorStore's format has no slot for, so +[the serializer](api/json.md) collapses it on the way out and reconstructs it on +the way in. + +Four deliberate inversions: + +| | TensorStore | `zarr-indexing` | +| --- | --- | --- | +| Dialect | One strict dialect everywhere: literal coordinates, no negative wrapping | The algebra keeps that dialect; each public boundary picks its own. [`LazyArray`](api/lazy_array.md) speaks positional NumPy, `zarr.Array.lazy` speaks literal. [`zarr_indexing.boundary`](api/boundary.md) is the translation | +| Scheduling | An internal C++ scheduler owns concurrency and chunk ordering | [`parts()`](api/lazy_array.md) hands the partition structure to whatever scheduler the caller already has — dask, a thread pool, a task queue | +| Wire format | Implementation-defined JSON, specified by what the implementation accepts | [ndsel](ndsel.md) is spec-first, with a vendored language-agnostic conformance corpus every implementation runs | +| Backends | A driver ecosystem (zarr, N5, neuroglancer, GCS, …) built into the library | No drivers. The wrapper takes anything with `shape`, `dtype`, and `__getitem__`, and prefers the array's own `__array_namespace__` | + +The asymmetries run the other way too, and are worth stating plainly. +TensorStore is a mature, heavily optimized C++ system with a performance +ceiling this cannot approach: our resolution is Python-level bookkeeping over +NumPy, and the per-part overhead is real. What this library has instead is +small size and no dependency beyond NumPy, which is what makes the algebra +adoptable by a Python project that wants the model without the runtime. + +## Bounding-box selections vs query selections + +Every selection this library can express falls into exactly one of two +categories, and the boundary between them is structural, not a heuristic: + +**A box** is a transform whose output maps are all `ConstantMap` or +`DimensionMap` — no `ArrayMap`. Such a map is affine and monotone: storage +coordinate `offset + stride * i` for `i` running over an interval. The whole +selection is therefore described by `O(ndim)` integers — an interval and a +stride per dimension — composition and intersection are interval arithmetic, +and the coordinates it touches form a regular lattice. Basic indexing produces +one, and composing basic indexing with basic indexing keeps one. + +**A query** is a transform with at least one `ArrayMap` — an explicit lookup +table of coordinates. It costs `O(n)` to store, it has no locality (the +coordinates may repeat, reverse, or scatter arbitrarily), and intersecting it +with a region means scanning it. `oindex`, `vindex`, and boolean masks all +produce one, and once an axis is a query no amount of subsequent basic +indexing makes it a box again. + +[ndsel](ndsel.md) encodes the same split in its message kinds: `point`, `box`, +and `slice` desugar to constant and affine output maps and are always boxes; +`points` desugars to `index_array` maps, and a `transform` body is a box +exactly when none of its output maps carries an `index_array`. A consumer can +therefore classify a selection off the wire without materializing anything: + +```python +from zarr_indexing import IndexTransform, transform_to_canonical + +transform_to_canonical(IndexTransform.from_shape((100, 80))[10:50, ::4])["output"] +# [{'offset': 0, 'stride': 1, 'input_dimension': 0}, +# {'offset': 0, 'stride': 4, 'input_dimension': 1}] + +import numpy as np +gather = IndexTransform.from_shape((100, 80)).oindex[np.array([90, 3, 3]), slice(None)] +transform_to_canonical(gather)["output"][0] +# {'offset': 0, 'stride': 1, 'index_array': [[90], [3], [3]], +# 'index_array_bounds': ['-inf', '+inf']} +``` + +The distinction matters to anyone downstream of a selection. A box can be +tiled into rectangular dask chunks or handed to a viewer or tile server that +only understands rectangles; a query cannot, and has to be resolved into a +gather. A box is also servable as a single strided slab read — but only a +strided one: reading its bounding box and discarding the rest is a +proportionally larger transfer as soon as any stride exceeds 1. The two also behave differently under +partitioning: a box touches a contiguous run of parts, while a query can touch +any subset of them, in any order, more than once. + +[`LazyArray`](api/lazy_array.md) exposes the category directly: + +```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) + +slab = lazy.lazy[10:50, ::4] +slab.is_box # True +slab.bounding_box() # ((10, 50), (0, 77)) +slab.strides() # (1, 4) +slab.shape # (40, 20) + +gather = lazy.lazy.oindex[[90, 3, 3], :] +gather.is_box # False +gather.bounding_box() # ((3, 91), (0, 80)) +gather.strides() # None +gather.shape # (3, 80) +``` + +`bounding_box()` is defined for both — it is the hull, the smallest interval per +storage dimension containing everything the selection reaches. `strides()` is +defined only for a box, and is the other half of the description: the hull says +*where*, the strides say *how densely*. + +Both halves matter, because a box is dense in its hull only when every stride is +1. The slab above spans a 40x77 hull over the 40x20 cells it actually selects — +a consumer that issued one rectangular read of the hull and discarded the rest +would move 3.85x the data. A query's hull is looser still and carries no stride +at all: 88 rows of hull over three selected rows. An empty selection returns +`None` from both, having no coordinate to report an interval around. + +There is deliberately no separate `BoxView` type today. A statically-typed +rectangular-only view is the obvious next step, but it should be introduced by +a consumer that actually needs the guarantee in its signatures rather than +speculatively; `is_box` is the runtime answer until then. + +## Current scope + +Three limits are known, intentional, and expected to lift: + +- **No negative steps.** `a[::-1]` raises rather than reversing. The transform + algebra can represent a negative stride and the lowering engine resolves one, + but the slice boundary rejects it while the ndsel spec change pinning the + domain-origin rule for negative steps is in flight. *Planned.* +- **Finite explicit bounds only.** `IndexDomain` has no implicit or unbounded + dimensions; the message layer will normalize a body with `"-inf"`/`"+inf"` + bounds, but the engine layer refuses to lower one into a transform. + TensorStore supports both. *Planned.* +- **Labels are carried, not propagated.** `IndexDomain` holds optional + dimension labels and the wire format round-trips them, but indexing + operations build new domains without them, so a label does not survive a + slice. *Planned.* diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index d123cbd32b..af297286c1 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -66,6 +66,10 @@ format stays TensorStore-loadable. Chunk resolution and the `oindex`/`vindex` helpers exist to serve NumPy-shaped selection semantics, which TensorStore does not have to model. +The [design notes](design-notes.md) set out the comparison in full, including +the four places the two libraries deliberately diverge and the ones where +TensorStore is simply ahead. + ## Quickstart Indexing a transform produces a new transform. No I/O happens, and no @@ -142,9 +146,152 @@ 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) + +view = lazy.lazy[10:50, ::4].lazy.oindex[[3, 0, 0], :] +view.shape # (3, 20) +repr(view) # '' + +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. + +Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a `LazyArray` a +drop-in source for `dask.array.from_array`. Everything else NumPy might do to it +— arithmetic, `numpy.sum`, `numpy.stack` — materializes the whole view through +`__array__` and hands back a plain NumPy array: laziness here is about indexing, +not about deferred compute. + +### Parts + +A read is broken up along a **partitioning**: a grid of boxes the wrapper walks +through [chunk resolution](api/chunk_resolution.md), so each box is read once +with plain basic slicing and the selection is applied to the block in memory. +`parts()` exposes that walk, projected through the view: + +```python +part = next(iter(lazy.lazy[0:35, 0:10].parts())) + +part.base_coords # (0, 0) — which box of the base grid +part.box # ((0, 30), (0, 40)) — that box in the array's own coordinates +part.array.shape # (30, 10) — a LazyArray for this box's cells +part.out_selection # (slice(0, 30, None), slice(0, 10, None)) +part.is_complete # False — the view does not cover the whole box +``` + +Each `part.array` is an ordinary `LazyArray`, so parts resolve independently and +concurrently, and `result()` is nothing but the assembly of that walk. Note that +`part.array.bounding_box()` is *part-local*; `part.box` is the global one, and is +what tells two parts apart. + +The partitioning is discovered from the wrapped array — `read_chunk_sizes`, then +`chunks`. Those are the *source's* words; the wrapper's own surface speaks only +of parts, and `with_parts` chooses a different one without touching the data or +the view: + +```python +view = lazy.lazy[10:50, ::4] + +[part.base_coords for part in view.with_parts((50, 20)).parts()] +# [(0, 0), (0, 1), (0, 2), (0, 3)] — uniform boxes, tail clipped + +[part.base_coords for part in view.with_parts(((50, 50), (20, 20, 20, 20))).parts()] +# [(0, 0), (0, 1), (0, 2), (0, 3)] — explicit per-axis sizes + +[part.base_coords for part in view.with_parts(None).parts()] +# [(0, 0)] — one whole-array part; resolve in one shot +``` + +`result()` is identical under any of them — repartitioning changes how the read +is cut up, never what it returns. Boxes that deliberately straddle the source's +own (to bound peak memory, say) are legal; they cost I/O, not correctness. + +### Boxes and queries + +A selection is either **rectangular** — describable by an interval and a stride +per dimension — or a **query**, an explicit list of coordinates. `is_box` says +which; `bounding_box()` gives the storage region touched either way, and +`strides()` how densely a box covers it: + +```python +slab = lazy.lazy[10:50, ::4] +slab.is_box # True +slab.bounding_box() # ((10, 50), (0, 77)) +slab.strides() # (1, 4) +slab.shape # (40, 20) + +gather = lazy.lazy.oindex[[90, 3, 3], :] +gather.is_box # False +gather.bounding_box() # ((3, 91), (0, 80)) — 88 rows of hull over 3 selected +gather.strides() # None +``` + +The bounding box is *dense* — every cell in it selected — only when every stride +is 1. The slab above spans a 40x77 hull over the 40x20 it selects, so slabbing +the hull and discarding the rest would read 3.85x what is needed; a consumer +that wants one rectangular read has to honour `strides()` too. + +It is a structural category, not a heuristic: basic indexing composes to a box +forever, and one `oindex`, `vindex`, or mask anywhere in the chain leaves it +forever. The [design notes](design-notes.md) explain why that matters to +anything downstream of a selection. + ## Reference - [The ndsel wire format](ndsel.md) +- [Design notes](design-notes.md) - [API reference](api/index.md) - [Changelog](https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/CHANGELOG.md) - [License (MIT)](https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/LICENSE.txt) diff --git a/packages/zarr-indexing/mkdocs.yml b/packages/zarr-indexing/mkdocs.yml index d7261f32e1..d341c127a8 100644 --- a/packages/zarr-indexing/mkdocs.yml +++ b/packages/zarr-indexing/mkdocs.yml @@ -14,6 +14,7 @@ use_directory_urls: true nav: - index.md - ndsel.md + - design-notes.md - API Reference: - api/index.md - ' zarr_indexing.transform': api/transform.md @@ -22,6 +23,8 @@ nav: - ' zarr_indexing.composition': api/composition.md - ' zarr_indexing.chunk_resolution': api/chunk_resolution.md - ' zarr_indexing.grid': api/grid.md + - ' zarr_indexing.lazy_array': api/lazy_array.md + - ' zarr_indexing.boundary': api/boundary.md - ' zarr_indexing.json': api/json.md - ' zarr_indexing.messages': api/messages.md - ' zarr_indexing.errors': api/errors.md diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index effe38ca88..c152d7d6c6 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/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, Partition 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,20 @@ "ConstantMap", "DimensionGridLike", "DimensionMap", + "EdgeDimensionGrid", "IndexDomain", "IndexDomainJSON", "IndexTransform", "IndexTransformJSON", + "LazyArray", "NdselError", "OutputIndexMap", "OutputIndexMapJSON", + "Partition", "__version__", + "array_map_dependent_axis", "compose", + "dimension_grids_from_chunks", "index_domain_from_json", "index_domain_to_json", "index_transform_from_json", diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py new file mode 100644 index 0000000000..a2c35bf595 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -0,0 +1,322 @@ +"""The positional (NumPy) selection dialect, lowered onto the transform algebra. + +The transform algebra speaks **literal domain coordinates**: an index is a +point in the view's own coordinate system, which after `view = arr[10:50]` runs +from 10 to 49, and a negative index is genuinely negative rather than counted +from the end (TensorStore's convention — see `zarr_indexing.transform`). + +NumPy speaks **positions**: index 0 always means the first element of the thing +you are indexing, and `-1` means the last. This module is the translation layer +between the two. It validates a selection against the view's shape with NumPy +semantics, then shifts every coordinate by the domain's origin so the transform +layer sees literal coordinates. + +Note +---- +`zarr.Array` currently carries its own copy of this normalization, tuned to a +different boundary contract (`Array.lazy[...]` deliberately exposes the literal +dialect, so a view's coordinates keep their meaning across composition). This +module is the generic, zarr-free version used by `LazyArray`; consolidating +zarr's copy onto it is left to a follow-up. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np + +if TYPE_CHECKING: + from zarr_indexing.domain import IndexDomain + +SelectionMode = Literal["basic", "orthogonal", "vectorized"] + + +def _is_bool_scalar(sel: Any) -> bool: + """Return True for a Python or NumPy boolean scalar. + + `isinstance(True, int)` holds, so booleans have to be rejected before any + integer handling — `arr[True]` would otherwise silently mean `arr[1]`. + """ + return isinstance(sel, (bool, np.bool_)) + + +def _as_index_array(sel: Any) -> np.ndarray[Any, np.dtype[Any]] | None: + """Return `sel` as an ndarray if it is array-like, else None.""" + if isinstance(sel, np.ndarray): + return sel + if isinstance(sel, (list, tuple)): + arr = np.asarray(sel) + if arr.dtype.kind in "biu": + return arr + raise IndexError( + f"arrays used as indices must be of integer or boolean type; got dtype {arr.dtype}" + ) + return None + + +def _axes_consumed(sel: Any, mode: SelectionMode) -> int: + """How many axes of the view a single selection entry consumes.""" + if sel is None: + return 0 + arr = _as_index_array(sel) + # A multidimensional boolean mask consumes one axis per mask dimension. + # Orthogonal indexing is per-axis by construction, so a mask there is 1-D + # and consumes exactly one axis. + if mode == "vectorized" and arr is not None and arr.dtype == np.bool_: + return arr.ndim + return 1 + + +def _normalize_int(value: int, size: int, axis: int) -> int: + """Bounds-check a positional integer index, wrapping negatives NumPy-style.""" + idx = value + if idx < 0: + idx += size + if idx < 0 or idx >= size: + raise IndexError(f"index {value} is out of bounds for axis {axis} with size {size}") + return idx + + +def _normalize_slice(sel: slice, size: int, axis: int) -> tuple[int, int, int]: + """Resolve a positional slice to `(start, stop, step)` within `[0, size]`.""" + if sel.step is not None and not isinstance(sel.step, (int, np.integer)): + raise IndexError(f"slice step must be an integer; got {sel.step!r}") + step = 1 if sel.step is None else int(sel.step) + if step < 1: + raise IndexError( + f"slice step must be positive; got {step} for axis {axis}. Reversed " + "and zero-step slices are not supported." + ) + start, stop, step = sel.indices(size) + # `slice.indices` can hand back stop < start for an empty forward slice + # (e.g. `5:2`); the transform layer wants a canonical empty interval. + stop = max(stop, start) + return start, stop, step + + +def _normalize_int_array( + arr: np.ndarray[Any, np.dtype[Any]], size: int, axis: int +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Bounds-check a positional integer index array, wrapping negatives.""" + if arr.dtype.kind not in "iu": + raise IndexError( + f"arrays used as indices must be of integer or boolean type; got dtype {arr.dtype}" + ) + # Cast before wrapping: an unsigned array cannot represent the intermediate + # negative values, and `intp` covers every index NumPy can address. + out = arr.astype(np.intp, copy=True) + if out.size > 0: + negative = out < 0 + if bool(negative.any()): + out = np.where(negative, out + size, out) + lo, hi = int(out.min()), int(out.max()) + if lo < 0 or hi >= size: + bad = lo if lo < 0 else hi + raise IndexError(f"index {bad} is out of bounds for axis {axis} with size {size}") + return out + + +def _expanded_axis_walk(entries: tuple[Any, ...], ndim: int, mode: SelectionMode) -> list[int]: + """The starting axis each entry addresses, with an ellipsis expanded. + + The returned list has one entry per element of `entries`; the value for an + `Ellipsis` (or a `newaxis`) is the axis it starts at, which is also the axis + the following entry resumes from once the skipped axes are accounted for. + """ + for sel in entries: + if _is_bool_scalar(sel): + raise IndexError( + "boolean scalars are not valid indices; use a boolean array " + "matching the shape of the axes it selects" + ) + if sum(1 for sel in entries if sel is Ellipsis) > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + consumed = sum(_axes_consumed(sel, mode) for sel in entries if sel is not Ellipsis) + if consumed > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, but {consumed} were indexed" + ) + + axes: list[int] = [] + axis = 0 + for sel in entries: + axes.append(axis) + axis += (ndim - consumed) if sel is Ellipsis else _axes_consumed(sel, mode) + return axes + + +def split_scalar_axes( + selection: Any, + domain: IndexDomain, + mode: SelectionMode, +) -> tuple[tuple[Any, ...] | None, Any]: + """Peel scalar integer indices out of a fancy selection. + + In NumPy, a scalar integer is a *basic* index wherever it appears: it drops + its axis, and is applied before the advanced indices rather than broadcast + against them. `a[0, [1, 2], :]` is `a[0][[1, 2], :]`. Neither the orthogonal + nor the vectorized path of the transform algebra models that — both widen a + scalar into a length-1 index array, which keeps the axis — so the scalars + are split off here and applied as a separate basic step first. + + Parameters + ---------- + selection + A positional orthogonal or vectorized selection. + domain + The domain of the view being indexed. + mode + `"orthogonal"` or `"vectorized"`; controls how many axes each entry + covers, which decides where the scalars sit. + + Returns + ------- + tuple + `(basic_selection, remaining_selection)`. `basic_selection` is a + full-rank basic selection in **literal** domain coordinates that drops + the scalar axes, or `None` when the selection has no scalar entries (in + which case `remaining_selection` is `selection` unchanged). + + Raises + ------ + IndexError + If a boolean scalar is used as an index, an index is out of bounds, or + too many indices are supplied. + """ + entries = selection if isinstance(selection, tuple) else (selection,) + axes = _expanded_axis_walk(entries, domain.ndim, mode) + + scalar_axes: dict[int, int] = {} + remaining: list[Any] = [] + for sel, axis in zip(entries, axes, strict=True): + if isinstance(sel, (int, np.integer)) and not _is_bool_scalar(sel): + scalar_axes[axis] = _normalize_int(int(sel), domain.shape[axis], axis) + else: + remaining.append(sel) + + if len(scalar_axes) == 0: + return None, selection + + basic: list[Any] = [] + for axis in range(domain.ndim): + lo = domain.inclusive_min[axis] + if axis in scalar_axes: + basic.append(lo + scalar_axes[axis]) + else: + basic.append(slice(lo, domain.exclusive_max[axis])) + return tuple(basic), tuple(remaining) + + +def normalize_positional_selection( + selection: Any, + domain: IndexDomain, + mode: SelectionMode, +) -> Any: + """Translate a positional (NumPy-dialect) selection into literal coordinates. + + Positions are zero-based offsets into the *current view*; negatives wrap + from the end. The returned selection addresses the same cells in the + literal coordinate system `domain` uses, ready for + `zarr_indexing.transform.selection_to_transform`. + + Parameters + ---------- + selection + A NumPy-style selection: integers, slices, `Ellipsis`, integer arrays or + lists, or boolean arrays. + domain + The domain of the view being indexed. Its shape defines the positional + bounds and its origin the coordinate shift. + mode + Which selection dialect the entries follow: `"basic"` (integers and + slices), `"orthogonal"` (per-axis arrays, outer product), or + `"vectorized"` (correlated coordinate arrays or a single mask). + + Returns + ------- + tuple + The selection with every coordinate expressed in literal domain + coordinates. + + Raises + ------ + IndexError + If a boolean scalar is used as an index, a boolean mask does not match + the shape of the axes it covers, an index is out of bounds, or too many + indices are supplied. + """ + entries = selection if isinstance(selection, tuple) else (selection,) + shape = domain.shape + origin = domain.inclusive_min + ndim = domain.ndim + + for sel in entries: + if _is_bool_scalar(sel): + raise IndexError( + "boolean scalars are not valid indices; use a boolean array " + "matching the shape of the axes it selects" + ) + + n_ellipsis = sum(1 for sel in entries if sel is Ellipsis) + if n_ellipsis > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + consumed = sum(_axes_consumed(sel, mode) for sel in entries if sel is not Ellipsis) + if consumed > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, but {consumed} were indexed" + ) + + result: list[Any] = [] + axis = 0 + for sel in entries: + if sel is Ellipsis: + # Passed through rather than expanded: every mode of + # `selection_to_transform` expands an ellipsis (and pads short + # selections) to whole-axis slices itself, and `vectorized` mode + # rejects an explicit slice, so expanding here would turn a legal + # partial coordinate selection into an error. + result.append(Ellipsis) + axis += ndim - consumed + continue + if sel is None: + # newaxis: no axis of the view is consumed, and there is no + # coordinate to shift. The transform layer decides whether the mode + # accepts it. + result.append(None) + continue + + arr = _as_index_array(sel) + if arr is not None and arr.dtype == np.bool_: + n_axes = _axes_consumed(sel, mode) + expected = shape[axis : axis + n_axes] + if arr.shape != tuple(expected): + raise IndexError( + f"boolean index has shape {arr.shape} but the axes it " + f"covers have shape {tuple(expected)}" + ) + for offset, positions in enumerate(np.nonzero(arr)): + result.append(positions.astype(np.intp) + origin[axis + offset]) + axis += n_axes + continue + + if axis >= ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, " + f"but {consumed} were indexed" + ) + size = shape[axis] + if arr is not None: + result.append(_normalize_int_array(arr, size, axis) + origin[axis]) + elif isinstance(sel, slice): + start, stop, step = _normalize_slice(sel, size, axis) + result.append(slice(start + origin[axis], stop + origin[axis], step)) + elif isinstance(sel, (int, np.integer)): + result.append(_normalize_int(int(sel), size, axis) + origin[axis]) + else: + raise IndexError(f"unsupported selection type: {type(sel)!r}") + axis += 1 + + # Axes the selection did not mention are left to `selection_to_transform`, + # which pads them with whole-axis slices in every mode. + return tuple(result) diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py index de1dae2dfc..5bb01dea0a 100644 --- a/packages/zarr-indexing/src/zarr_indexing/grid.py +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -1,18 +1,27 @@ -"""Structural typing for the chunk-grid surface used by chunk resolution. +"""Per-dimension chunk grids: the Protocol and a concrete implementation. `chunk_resolution` needs only a narrow slice of a chunk grid: the per-dimension mapping between storage indices and chunk coordinates, passed as one `DimensionGridLike` per storage dimension. Rather than import zarr's concrete grid types, we type against this Protocol; zarr's per-dimension grids satisfy it structurally, so no zarr import is needed here. + +`EdgeDimensionGrid` is a concrete implementation for callers that do not have a +zarr grid to hand — it is built from an explicit tuple of per-chunk sizes along +one axis, so a clipped ("edge") boundary chunk is expressed directly rather than +derived from a uniform chunk shape. `dimension_grids_from_chunks` builds one per +axis from either chunk convention. """ from __future__ import annotations -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, Any, Protocol, cast + +import numpy as np if TYPE_CHECKING: - import numpy as np + from collections.abc import Iterable, Sequence + import numpy.typing as npt @@ -23,3 +32,220 @@ def index_to_chunk(self, idx: int) -> int: ... def chunk_offset(self, chunk_ix: int) -> int: ... def chunk_size(self, chunk_ix: int) -> int: ... def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: ... + + +class EdgeDimensionGrid: + """A one-dimensional chunk grid described by an explicit tuple of chunk sizes. + + The sizes follow the dask `Array.chunks` convention for a single axis: one + entry per chunk giving that chunk's extent, with a possibly-smaller final + ("edge") chunk. Chunk offsets are the exclusive prefix sums, so lookups are + a binary search rather than a division — the grid does not have to be + regular. + + Parameters + ---------- + sizes + Per-chunk extents along this axis. Every entry must be positive. An + empty sequence describes a zero-length axis. + + Examples + -------- + >>> grid = EdgeDimensionGrid((3, 3, 1)) + >>> grid.index_to_chunk(4) + 1 + >>> grid.chunk_offset(2), grid.chunk_size(2) + (6, 1) + """ + + __slots__ = ("_offsets", "sizes") + + sizes: tuple[int, ...] + + def __init__(self, sizes: Sequence[int]) -> None: + normalized = tuple(int(s) for s in sizes) + for i, s in enumerate(normalized): + if s <= 0: + raise ValueError( + f"chunk sizes must be positive; got {s} at position {i} of {normalized}" + ) + self.sizes = normalized + # Exclusive prefix sums, length len(sizes) + 1: `_offsets[c]` is the + # first index of chunk `c`, and `_offsets[-1]` the total extent. + offsets: np.ndarray[Any, np.dtype[np.intp]] = np.zeros(len(normalized) + 1, dtype=np.intp) + if len(normalized) > 0: + np.cumsum(np.asarray(normalized, dtype=np.intp), out=offsets[1:]) + self._offsets = offsets + + @property + def num_chunks(self) -> int: + """The number of chunks along this axis.""" + return len(self.sizes) + + @property + def extent(self) -> int: + """The total number of indices this axis spans.""" + return int(self._offsets[-1]) + + def __repr__(self) -> str: + return f"EdgeDimensionGrid(sizes={self.sizes})" + + def __eq__(self, other: object) -> bool: + if not isinstance(other, EdgeDimensionGrid): + return NotImplemented + return self.sizes == other.sizes + + def __hash__(self) -> int: + return hash((type(self).__name__, self.sizes)) + + def index_to_chunk(self, idx: int) -> int: + """Return the index of the chunk holding index `idx`.""" + if idx < 0 or idx >= self.extent: + raise IndexError(f"index {idx} is out of bounds for an axis of extent {self.extent}") + return int(np.searchsorted(self._offsets, idx, side="right")) - 1 + + def chunk_offset(self, chunk_ix: int) -> int: + """Return the first index belonging to chunk `chunk_ix`.""" + if chunk_ix < 0 or chunk_ix >= len(self.sizes): + raise IndexError( + f"chunk index {chunk_ix} is out of bounds for {len(self.sizes)} chunks" + ) + return int(self._offsets[chunk_ix]) + + def chunk_size(self, chunk_ix: int) -> int: + """Return the extent of chunk `chunk_ix`.""" + if chunk_ix < 0 or chunk_ix >= len(self.sizes): + raise IndexError( + f"chunk index {chunk_ix} is out of bounds for {len(self.sizes)} chunks" + ) + return self.sizes[chunk_ix] + + def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: + """Return the chunk holding each index in `indices`, elementwise.""" + arr = np.asarray(indices, dtype=np.intp) + if arr.size > 0 and (int(arr.min()) < 0 or int(arr.max()) >= self.extent): + raise IndexError( + f"indices must lie in [0, {self.extent}); got [{int(arr.min())}, {int(arr.max())}]" + ) + return (np.searchsorted(self._offsets, arr, side="right") - 1).astype(np.intp) + + +def _uniform_chunk_sizes(extent: int, chunk: int) -> tuple[int, ...]: + """Expand a single uniform chunk length into per-chunk sizes, tail clipped.""" + if chunk <= 0: + raise ValueError(f"chunk shape entries must be positive; got {chunk}") + n_full, remainder = divmod(extent, chunk) + sizes = (chunk,) * n_full + if remainder > 0: + sizes = (*sizes, remainder) + return sizes + + +def _entry_kind(entry: Any) -> str: + """Classify one `chunks` entry as `"int"`, `"sequence"`, or `"neither"`.""" + if isinstance(entry, (int, np.integer)) and not isinstance(entry, bool): + return "int" + if isinstance(entry, (str, bytes)): + return "neither" + try: + iter(cast("Iterable[Any]", entry)) + except TypeError: + return "neither" + return "sequence" + + +def dimension_grids_from_chunks( + chunks: Sequence[int] | Sequence[Sequence[int]], + shape: Sequence[int], +) -> tuple[EdgeDimensionGrid, ...]: + """Build one `EdgeDimensionGrid` per axis from either chunk convention. + + Two conventions are accepted, disambiguated by element type: + + - a **uniform chunk shape** — one integer per axis, the chunk length along + that axis. The trailing chunk is clipped to the array extent, matching how + a zarr regular chunk grid covers a shape that is not a multiple of the + chunk shape. + - **dask-convention per-axis sizes** — one sequence of per-chunk extents per + axis, e.g. `((3, 3, 1), (4,))`. Each sequence must sum to the + corresponding entry of `shape`. + + Parameters + ---------- + chunks + A uniform chunk shape or dask-convention per-axis chunk sizes. + shape + The array shape the grids cover. + + Returns + ------- + tuple of EdgeDimensionGrid + One grid per axis, in axis order. + + Raises + ------ + ValueError + If `chunks` has a different length than `shape`, mixes the two + conventions, or declares per-axis sizes that do not sum to `shape`. + + Examples + -------- + >>> grids = dimension_grids_from_chunks((3, 4), (7, 4)) + >>> (grids[0].sizes, grids[1].sizes) + ((3, 3, 1), (4,)) + """ + shape_t = tuple(int(s) for s in shape) + entries: tuple[Any, ...] = tuple(chunks) + if len(entries) != len(shape_t): + raise ValueError( + f"chunks must have one entry per dimension; got {len(entries)} " + f"entries for shape {shape_t}" + ) + + _CONVENTIONS = ( + "chunks must be either a uniform chunk shape (one integer per " + "dimension) or per-axis chunk sizes (one sequence of integers per " + "dimension)" + ) + # Classify every entry first, so the diagnosis names what is actually wrong. + # An entry that belongs to *neither* convention (a float, a string, None) is + # a different mistake from a well-formed mixture of the two, and saying + # "not a mixture" about `(2.5, 2, 2)` sends the reader looking for the wrong + # thing. `sum`/explicit loops rather than `all`/`any` over a comprehension: + # the latter narrows `entry` to Never on the fall-through branch. + kinds = [_entry_kind(entry) for entry in entries] + neither = [(axis, entries[axis]) for axis, kind in enumerate(kinds) if kind == "neither"] + if len(neither) > 0: + described = ", ".join(f"{entry!r} at dimension {axis}" for axis, entry in neither) + verb = "is" if len(neither) == 1 else "are" + raise ValueError(f"{_CONVENTIONS}; {described} {verb} neither") + + n_int = sum(1 for kind in kinds if kind == "int") + if len(entries) > 0 and n_int == len(entries): + return tuple( + EdgeDimensionGrid(_uniform_chunk_sizes(extent, int(entry))) + for entry, extent in zip(entries, shape_t, strict=True) + ) + if n_int > 0: + raise ValueError(f"{_CONVENTIONS}, not a mixture; got {entries!r}") + + grids: list[EdgeDimensionGrid] = [] + for axis, (entry, extent) in enumerate(zip(entries, shape_t, strict=True)): + # Every remaining entry is a sequence; its *elements* may still not be + # integers, which is a third distinct mistake. + # `int()` would silently truncate a float, so check the element type + # rather than relying on the conversion to fail. + elements: tuple[Any, ...] = tuple(cast("Iterable[Any]", entry)) + if any(not isinstance(e, (int, np.integer)) or isinstance(e, bool) for e in elements): + raise ValueError( + f"per-axis chunk sizes must be integers; dimension {axis} has {entry!r}" + ) + sizes = tuple(int(e) for e in elements) + total = sum(sizes) + if total != extent: + raise ValueError( + f"per-axis chunk sizes for dimension {axis} sum to {total}, " + f"but the array extent is {extent}" + ) + grids.append(EdgeDimensionGrid(sizes)) + return tuple(grids) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py new file mode 100644 index 0000000000..eb5afc2c86 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -0,0 +1,1145 @@ +"""`LazyArray` — TensorStore-style lazy indexing over any array-API-like array. + +`LazyArray` wraps an existing array — NumPy, zarr, CuPy, anything with `shape`, +`dtype`, and `__getitem__` — and adds a `.lazy` accessor whose indexing +operations build up an [`IndexTransform`](transform.md) instead of reading data: + +```python +view = LazyArray(source).lazy[10:50, ::2].lazy.oindex[[3, 1, 1], :] +view.shape # known without touching the data +values = view.result() +``` + +Nothing is read until `result()` (or `__array__`, or an eager `__getitem__`). +Composition is free: a view of a view is still a single transform. + +Parts +----- +A `LazyArray` carries a **partitioning** of the array it wraps: a grid of boxes +that a read is broken into. `parts()` walks those boxes as they fall through the +view, yielding a [`Partition`](#zarr_indexing.lazy_array.Partition) per box — +the base coordinates of the box, a `LazyArray` covering exactly the cells of the +view that live in it, and where those cells belong in the result. `result()` is +built on that walk and nothing else: + +```python +for part in view.parts(): + out[part.out_selection] = part.array.result() +``` + +so each box is read once, with plain basic slicing, and the selection is applied +to the block in memory. + +The partitioning is discovered from the wrapped array at construction — first +`read_chunk_sizes` (zarr's clipped per-axis sizes, sharding-aware), then +`chunks`, read as per-axis sizes if its entries are sequences and as a uniform +box shape if they are integers. Those attribute names are the *source's* +vocabulary; the wrapper's own surface speaks only of parts. An array that +advertises neither is a single whole-array part, and resolving it lowers the +whole view to one pass of array operations (`take`, basic slicing, a flat +gather) rather than materializing a block first. + +`with_parts` replaces the partitioning without touching the data or the view: + +```python +view.with_parts((64, 64)) # uniform boxes, tail clipped +view.with_parts(((3, 3, 1),)) # explicit per-axis sizes +view.with_parts(None) # one whole-array part; resolve in one shot +``` + +Repartitioning never changes what `result()` returns — only how the read is cut +up. Deliberately misaligning the parts with the source's own boxes is legal and +useful (to bound peak memory, or to batch small reads); it costs I/O, not +correctness. + +Boxes and queries +----------------- +A selection is either **rectangular** — an interval and a stride per dimension, +which is what basic indexing composes to, however deeply — or a **query**, an +explicit list of coordinates, which is what `oindex`, `vindex`, and masks +produce and which no amount of subsequent basic indexing undoes. `is_box` +reports the category and `bounding_box()` the storage region touched, exact for +a box and a hull for a query. The distinction is structural rather than an +optimization; [the design notes](../design-notes.md) say why it matters to +anything downstream of a selection. + +The positional dialect +---------------------- +Selections on `LazyArray` are **positional, NumPy-style**: index 0 is the first +element of the current view, `-1` is the last, boolean masks must match the +view's shape, and every index is bounds-checked against the view. + +This is deliberately different from `zarr.Array.lazy[...]`, which exposes the +**literal** TensorStore dialect: a zarr view keeps the coordinate system of the +array it came from, so after `v = arr.lazy[10:50]` the first element of `v` is +`v[10]` and a negative index is out of bounds rather than counted from the end. +That dialect is the right one for zarr, where a view's coordinates are meant to +stay comparable with the parent array's. `LazyArray` is a duck array — it has to +behave like the thing it wraps to be usable as a drop-in for NumPy or as a dask +source — so it re-zeroes its coordinates on every view and speaks positions. +`zarr_indexing.boundary` performs the translation between the two. + +Two more NumPy rules the dialect keeps, 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. So + `lazy.oindex[0]` has the shape of `x[0]`, `lazy.oindex[0, [1, 2], :]` means + `x[0][numpy.ix_([1, 2], ...)]`, and `lazy.oindex[0, 1, 2]` and + `lazy.vindex[0, 1, 2]` are both zero-rank. Use a length-1 list to keep an + axis. +- **Advanced indices land where NumPy puts them.** For a `vindex` selection that + leaves some axes unindexed, the gathered dimensions sit where the coordinate + arrays sat when those arrays are adjacent, and lead when a slice separates + them — so `lazy.vindex[..., i, j]` has shape `(x.shape[0], *broadcast)`, + matching `x[..., i, j]`. + +Materializing on fallback +------------------------- +`LazyArray` implements `__array__` but deliberately implements neither +`__array_ufunc__`/`__array_function__` nor `__array_namespace__`. **Every NumPy +operation other than indexing therefore materializes the whole view**: +`numpy.sum(view)`, `view + 1`, `numpy.stack([view, view])` all convert through +`__array__` first and give you a plain NumPy array back. That is intentional — +laziness here is about *indexing*, not about building a deferred compute graph — +but it means a `LazyArray` is not a drop-in for arithmetic on a large array. Use +`.lazy[...]` to narrow first, or hand the wrapper to `dask.array.from_array` and +let dask own the compute graph. + +Device caveat +------------- +Resolving a partitioned view builds its output buffer with NumPy, because the +resolver's scatter indices are host-side integer arrays. Wrapping a device array +that advertises parts therefore returns host memory. A single whole-array part +lowers directly and stays in the wrapped array's own namespace. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import operator +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +import numpy as np + +from zarr_indexing.boundary import ( + SelectionMode, + normalize_positional_selection, + split_scalar_axes, +) +from zarr_indexing.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) +from zarr_indexing.grid import EdgeDimensionGrid, dimension_grids_from_chunks +from zarr_indexing.json import transform_to_canonical +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import ( + IndexTransform, + array_map_dependent_axis, + selection_to_transform, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Sequence + + SelectFn = Callable[[Any, SelectionMode], "LazyArray"] + +__all__ = ["LazyArray", "Partition"] + +# Above this many bytes, the no-dask token fallback describes an array +# structurally instead of digesting its contents. See `_wrapped_token`. +_TOKEN_DIGEST_LIMIT = 1 << 20 + + +class ArrayLike(Protocol): + """The surface `LazyArray` needs from the array it wraps.""" + + @property + def shape(self) -> tuple[int, ...]: ... + @property + def dtype(self) -> Any: ... + def __getitem__(self, key: Any) -> Any: ... + + +# --------------------------------------------------------------------------- # +# Array-namespace helpers +# --------------------------------------------------------------------------- # + + +def _namespace(x: Any) -> Any: + """Return `x`'s array-API namespace, or None if it does not declare one.""" + getter = getattr(x, "__array_namespace__", None) + if getter is None: + return None + try: + return getter() + except Exception: # pragma: no cover - a namespace that refuses to load + return None + + +def _take(x: Any, indices: np.ndarray[Any, np.dtype[np.intp]], axis: int) -> Any: + """Gather along one axis, preferring the wrapped array's own namespace.""" + xp = _namespace(x) + take = getattr(xp, "take", None) + if take is not None: + return take(x, xp.asarray(indices), axis=axis) + return x[(slice(None),) * axis + (indices,)] + + +def _reshape(x: Any, shape: tuple[int, ...]) -> Any: + xp = _namespace(x) + reshape = getattr(xp, "reshape", None) + if reshape is not None: + return reshape(x, shape) + return x.reshape(shape) + + +def _transpose(x: Any, perm: tuple[int, ...]) -> Any: + xp = _namespace(x) + for name in ("permute_dims", "transpose"): + fn = getattr(xp, name, None) + if fn is not None: + return fn(x, perm) + return np.transpose(x, perm) + + +def _expand_dims(x: Any, axis: int) -> Any: + xp = _namespace(x) + fn = getattr(xp, "expand_dims", None) + if fn is not None: + return fn(x, axis=axis) + return np.expand_dims(x, axis=axis) + + +# --------------------------------------------------------------------------- # +# Partition discovery +# --------------------------------------------------------------------------- # + + +def _read_source_attribute(array: Any, name: str) -> Any: + """Read a partition-describing attribute, treating any failure as "absent". + + Discovery inspects an object we did not write. A missing attribute is the + common case, but zarr raises an `AttributeError` subclass from + `read_chunk_sizes` on a lazy view, and other backends compute the attribute + lazily and may fail for their own reasons. Any failure here means "this + array does not advertise a partitioning", never a hard error. + """ + try: + return getattr(array, name, None) + except Exception: + return None + + +def _discover_parts(array: Any, shape: tuple[int, ...]) -> tuple[EdgeDimensionGrid, ...] | None: + """Resolve the partitioning advertised by `array`, or None for one whole part. + + Discovery parses external input: an attribute that does not describe a + partitioning of `shape` means "this object does not advertise one I + understand", and the array is treated as unpartitioned rather than rejected. + A partitioning is an I/O strategy, so reading the whole array is always a + correct fallback. `with_parts` is *our* API, and validates strictly. + """ + declared = _read_source_attribute(array, "read_chunk_sizes") + if declared is None: + declared = _read_source_attribute(array, "chunks") + if declared is None: + return None + try: + return dimension_grids_from_chunks(declared, shape) + except (ValueError, TypeError): + return None + + +def _whole_array_grids(shape: tuple[int, ...]) -> tuple[EdgeDimensionGrid, ...]: + """A partitioning with a single part covering the whole array.""" + return tuple(EdgeDimensionGrid((extent,) if extent > 0 else ()) for extent in shape) + + +# --------------------------------------------------------------------------- # +# The lowering engine +# --------------------------------------------------------------------------- # + + +def _is_identity_transform(transform: IndexTransform, shape: tuple[int, ...]) -> bool: + """True when `transform` maps every coordinate of `shape` to itself. + + Structural rather than an `==` against `IndexTransform.from_shape`: an + `ArrayMap` holds an ndarray, so equality on two transforms that both carry + one would try to take the truth value of an array. + """ + domain = transform.domain + if domain.inclusive_min != (0,) * len(shape) or domain.exclusive_max != shape: + return False + if len(transform.output) != len(shape): + return False + return all( + isinstance(m, DimensionMap) and m.input_dimension == i and m.offset == 0 and m.stride == 1 + for i, m in enumerate(transform.output) + ) + + +def _is_correlated(transform: IndexTransform) -> bool: + """True when the transform gathers a list of points rather than an outer product.""" + return any(isinstance(m, ArrayMap) and m.input_dimension is None for m in transform.output) + + +def _dimension_map_coords( + m: DimensionMap, transform: IndexTransform +) -> np.ndarray[Any, np.dtype[np.intp]]: + """The storage coordinates a DimensionMap enumerates, in view order.""" + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = transform.domain.exclusive_max[d] + return (m.offset + m.stride * np.arange(lo, hi, dtype=np.intp)).astype(np.intp) + + +def _array_map_coords(m: ArrayMap) -> np.ndarray[Any, np.dtype[np.intp]]: + """The storage coordinates an ArrayMap enumerates, flattened.""" + return (m.offset + m.stride * m.index_array).astype(np.intp).reshape(-1) + + +def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: int) -> Any: + """Permute `result`'s axes into input-domain order, restoring dropped axes. + + `axis_input_dims[k]` is the input (domain) dimension that axis `k` of + `result` corresponds to. Axes are permuted so that they appear in increasing + domain-dimension order, and any domain dimension no output map depends on + (only reachable via `newaxis`) is reinserted as a singleton. + + This is where NumPy's advanced-index placement rules are absorbed: whatever + order the gather produced, the lowered result always comes back in the + view's own axis order. + """ + if len(set(axis_input_dims)) != len(axis_input_dims): + raise NotImplementedError( + "resolving a transform whose output maps share an input dimension " + "(a diagonal view) is not supported" + ) + order = sorted(range(len(axis_input_dims)), key=lambda k: axis_input_dims[k]) + if order != list(range(len(order))): + result = _transpose(result, tuple(order)) + covered = sorted(axis_input_dims) + for dim in range(rank): + if dim not in covered: + result = _expand_dims(result, dim) + return result + + +def _lower(array: Any, transform: IndexTransform) -> Any: + """Lower a transform to one pass of array operations over `array`. + + The single lowering engine: every read, partitioned or not, ends here. The + result is always in the transform's own domain axis order and of exactly its + domain shape. + """ + if _is_correlated(transform): + result = _lower_correlated(array, transform) + else: + result = _lower_orthogonal(array, transform) + if isinstance(result, np.generic): + # Basic indexing every axis of a NumPy array yields a scalar; `result()` + # documents a zero-dimensional array. + return np.asarray(result) + return result + + +def _lower_orthogonal(array: Any, transform: IndexTransform) -> Any: + """Basic slicing plus one `take` per fancy-indexed axis (an outer product). + + Orthogonal `ArrayMap`s vary over distinct input axes, so gathering them one + axis at a time is exact — a `take` along one storage axis leaves every other + axis's coordinates untouched. + """ + outputs = transform.output + gathered: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + for out_dim, m in enumerate(outputs): + if isinstance(m, ArrayMap): + gathered[out_dim] = _array_map_coords(m) + elif isinstance(m, DimensionMap) and m.stride < 0: + # A reversing map has no positive-step slice; gather it explicitly. + gathered[out_dim] = _dimension_map_coords(m, transform) + + result = array + for out_dim, coords in gathered.items(): + result = _take(result, coords, axis=out_dim) + + selection: list[Any] = [] + axis_input_dims: list[int] = [] + for out_dim, m in enumerate(outputs): + if isinstance(m, ConstantMap): + selection.append(m.offset) + continue + if out_dim in gathered: + selection.append(slice(None)) + else: + assert isinstance(m, DimensionMap) + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = transform.domain.exclusive_max[d] + selection.append(slice(m.offset + m.stride * lo, m.offset + m.stride * hi, m.stride)) + if isinstance(m, ArrayMap): + axis = array_map_dependent_axis(m) + if axis is None: + raise NotImplementedError( + "resolving an orthogonal ArrayMap that varies over no input " + "dimension is not supported; such a map should have been " + "collapsed to a ConstantMap" + ) + axis_input_dims.append(axis) + else: + axis_input_dims.append(m.input_dimension) + result = result[tuple(selection)] + return _restore_domain_axis_order(result, axis_input_dims, transform.input_rank) + + +def _lower_correlated(array: Any, transform: IndexTransform) -> Any: + """Flatten the correlated axes, gather the points once, reshape back. + + Correlated (`vindex`) `ArrayMap`s address a list of *points*, not an outer + product, so they cannot be gathered one axis at a time. The correlated + storage axes are moved to the front and flattened, the per-point coordinates + are converted to offsets into that flat axis with row-major strides, and a + single `take` collects them. + """ + outputs = transform.output + correlated_dims = [ + d for d, m in enumerate(outputs) if isinstance(m, ArrayMap) and m.input_dimension is None + ] + if any(isinstance(m, ArrayMap) and m.input_dimension is not None for m in outputs): + raise NotImplementedError( + "resolving a transform with both correlated and orthogonal ArrayMaps is not supported" + ) + + slice_input_dims = {m.input_dimension for m in outputs if isinstance(m, DimensionMap)} + broadcast_axes = [d for d in range(transform.input_rank) if d not in slice_input_dims] + broadcast_shape = tuple(transform.domain.shape[d] for d in broadcast_axes) + + # Gather any reversing slice axis first, then take the basic-slice cut. The + # correlated axes keep their full extent: their coordinates are absolute. + gathered: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = { + d: _dimension_map_coords(m, transform) + for d, m in enumerate(outputs) + if isinstance(m, DimensionMap) and m.stride < 0 + } + result = array + for out_dim, coords in gathered.items(): + result = _take(result, coords, axis=out_dim) + + selection: list[Any] = [] + residual_axis_dims: list[int] = [] + correlated_positions: list[int] = [] + residual_positions: list[int] = [] + axis = 0 + for out_dim, m in enumerate(outputs): + if isinstance(m, ConstantMap): + selection.append(m.offset) + continue + if out_dim in correlated_dims: + selection.append(slice(None)) + correlated_positions.append(axis) + elif out_dim in gathered: + selection.append(slice(None)) + residual_positions.append(axis) + assert isinstance(m, DimensionMap) + residual_axis_dims.append(m.input_dimension) + else: + assert isinstance(m, DimensionMap) + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = transform.domain.exclusive_max[d] + selection.append(slice(m.offset + m.stride * lo, m.offset + m.stride * hi, m.stride)) + residual_positions.append(axis) + residual_axis_dims.append(d) + axis += 1 + result = result[tuple(selection)] + + # Correlated axes to the front, in output order, so the flattening strides + # below match the order the coordinates are combined in. + perm = tuple(correlated_positions) + tuple(residual_positions) + if perm != tuple(range(len(perm))): + result = _transpose(result, perm) + + n_corr = len(correlated_dims) + corr_sizes = tuple(int(s) for s in result.shape[:n_corr]) + tail_shape = tuple(int(s) for s in result.shape[n_corr:]) + result = _reshape(result, (math.prod(corr_sizes), *tail_shape)) + + flat_index = np.zeros(math.prod(broadcast_shape), dtype=np.intp) + stride = 1 + for position in range(n_corr - 1, -1, -1): + m = outputs[correlated_dims[position]] + assert isinstance(m, ArrayMap) + flat_index = flat_index + _array_map_coords(m) * stride + stride *= corr_sizes[position] + + result = _take(result, flat_index, axis=0) + result = _reshape(result, broadcast_shape + tail_shape) + return _restore_domain_axis_order( + result, list(broadcast_axes) + residual_axis_dims, transform.input_rank + ) + + +# --------------------------------------------------------------------------- # +# Partitions +# --------------------------------------------------------------------------- # + + +def _partition_out_selection( + sub_transform: IndexTransform, + out_indices: Any, + out_shape: tuple[int, ...], +) -> tuple[Any, ...]: + """Where a partition's values belong in an array of shape `out_shape`. + + A correlated sub-transform scatters through flat offsets into the row-major + result; unravelling them turns that into an index tuple for the result's own + shape, so callers never need a flat working buffer. + """ + _chunk_selection, out_selection, _drop_axes = sub_transform_to_selections( + sub_transform, out_indices + ) + if not _is_correlated(sub_transform): + return out_selection + if len(out_shape) == 0: + return () + flat = out_selection[0] + if isinstance(flat, slice): + flat = np.arange(flat.start, flat.stop, dtype=np.intp) + return tuple(np.unravel_index(np.asarray(flat, dtype=np.intp), out_shape)) + + +def _covers_whole_part(local_transform: IndexTransform, part_shape: tuple[int, ...]) -> bool: + """Whether a part-local transform addresses every cell of its part. + + Conservative: an `ArrayMap` could in principle enumerate a whole part, but + proving it costs more than the answer is worth, so a fancy-indexed axis + always reports incomplete. + """ + domain = local_transform.domain + for out_dim, m in enumerate(local_transform.output): + extent = part_shape[out_dim] + if isinstance(m, ConstantMap): + if extent != 1 or m.offset != 0: + return False + elif isinstance(m, DimensionMap): + if m.stride != 1: + return False + d = m.input_dimension + if m.offset + domain.inclusive_min[d] != 0: + return False + if m.offset + domain.exclusive_max[d] != extent: + return False + else: + return False + return True + + +@dataclass(frozen=True) +class Partition: + """One box of a `LazyArray`'s partitioning, as it falls through the view. + + Yielded by [`LazyArray.parts`][zarr_indexing.lazy_array.LazyArray.parts]. + The parts of a view tile it exactly and disjointly: assembling every + `array.result()` at its `out_selection` reproduces the view's `result()`, + and each part can be resolved independently and concurrently. + + Attributes + ---------- + base_coords + Which box of the base partitioning this is, one coordinate per dimension + of the wrapped array. + box + The box itself, in the **global** storage coordinates of the wrapped + array: one `[inclusive_min, exclusive_max)` interval per dimension. + `array.bounding_box()` is part-*local* and so cannot tell two parts + apart; this can, and it is what pairs with `is_complete` to decide a + read-modify-write. + array + A `LazyArray` covering exactly the cells of the view that live in this + box. Resolving it reads the box once with basic slicing and applies the + selection to the block in memory. + out_selection + Where `array.result()` belongs in an array of the view's shape — a + NumPy index tuple, usable directly as `out[part.out_selection] = ...`. + is_complete + Whether the view covers the whole box. Useful to a writer deciding + between a blind overwrite and a read-modify-write. + """ + + base_coords: tuple[int, ...] + box: tuple[tuple[int, int], ...] + array: LazyArray + out_selection: tuple[Any, ...] + is_complete: bool + + +# --------------------------------------------------------------------------- # +# Tokenization +# --------------------------------------------------------------------------- # + + +def _wrapped_token(array: Any) -> Any: + """A deterministic token for the wrapped array. + + Determinism scope, in order of preference: the array's own + `__dask_tokenize__`; `dask.base.tokenize` when dask is importable (imported + lazily — this package never requires it); otherwise a local fallback that + digests the contents of a small array and, above `_TOKEN_DIGEST_LIMIT`, + falls back to a structural description that is stable across processes but + **not** sensitive to the array's contents. + """ + hook = getattr(array, "__dask_tokenize__", None) + if hook is not None: + try: + return hook() + except Exception: # pragma: no cover - a hook that refuses to run + pass + try: + # dask is an optional peer, never a dependency of this package, so it is + # imported here and its absence is ordinary. + from dask.base import tokenize # pyright: ignore[reportMissingImports] + except ImportError: + pass + else: + return tokenize(array) + + shape = tuple(int(s) for s in getattr(array, "shape", ())) + dtype = getattr(array, "dtype", None) + structural = (type(array).__qualname__, shape, str(dtype)) + + # Decide whether to digest the contents from the *declared* size. Measuring + # it by converting first would read the whole array — a multi-gigabyte store + # pulled into memory by a token call, which is the opposite of the point. + itemsize = getattr(dtype, "itemsize", None) + if not isinstance(itemsize, int) or itemsize * math.prod(shape) > _TOKEN_DIGEST_LIMIT: + return structural + try: + contents = np.ascontiguousarray(array) + except Exception: + return structural + return (*structural, hashlib.sha256(contents.tobytes()).hexdigest()) + + +# --------------------------------------------------------------------------- # +# The wrapper +# --------------------------------------------------------------------------- # + + +class LazyArray: + """A lazily-indexable view over an array-API-like array. + + Wrapping is cheap and non-destructive: the wrapped array is never copied and + never read at construction time. Indexing through `.lazy` composes an + `IndexTransform` and returns another `LazyArray`; `result()` materializes. + + Selections use the **positional NumPy dialect**, and reads are broken up + along a **partitioning** discovered from the wrapped array — see the module + docstring, which also covers how the dialect differs from `zarr.Array.lazy` + and why every non-indexing NumPy operation materializes the view. + + Parameters + ---------- + array + The array to wrap. It must expose `shape`, `dtype`, and `__getitem__` + with basic (integer/slice) indexing. Its partitioning, if it advertises + one, is discovered here; use `with_parts` to choose a different one. + + Examples + -------- + >>> import numpy as np + >>> source = np.arange(12).reshape(3, 4) + >>> view = LazyArray(source).with_parts((2, 2)).lazy[1:, ::2] + >>> view.shape + (2, 2) + >>> view.result() + array([[ 4, 6], + [ 8, 10]]) + """ + + __slots__ = ("_array", "_parts", "_transform", "_window") + + def __init__(self, array: ArrayLike) -> None: + shape = tuple(int(s) for s in array.shape) + self._array = array + self._window: tuple[slice, ...] | None = None + self._transform = IndexTransform.from_shape(shape) + self._parts = _discover_parts(array, shape) + + @classmethod + def _derive( + cls, + array: ArrayLike, + transform: IndexTransform, + parts: tuple[EdgeDimensionGrid, ...] | None, + window: tuple[slice, ...] | None, + ) -> LazyArray: + """Build a wrapper sharing `array` but carrying a new transform or partitioning.""" + view = cls.__new__(cls) + view._array = array + # Views re-zero their coordinate system: the positional dialect means a + # view's first element is at position 0 whatever it was sliced from. + view._transform = transform.translate_domain_to((0,) * transform.input_rank) + view._parts = parts + view._window = window + return view + + @property + def _base_shape(self) -> tuple[int, ...]: + """The shape of what this wrapper treats as its base array.""" + if self._window is None: + return tuple(int(s) for s in self._array.shape) + return tuple(s.stop - s.start for s in self._window) + + def _read_base(self) -> Any: + """The base array, materializing the window if this wrapper has one.""" + if self._window is None: + return self._array + return np.asarray(self._array[self._window]) + + # -- array-API surface -------------------------------------------------- + + @property + def array(self) -> ArrayLike: + """The wrapped array.""" + return self._array + + @property + def transform(self) -> IndexTransform: + """The composed transform from this view's coordinates to storage.""" + return self._transform + + @property + def shape(self) -> tuple[int, ...]: + """The shape of this view — the transform's input domain, not the source's.""" + return self._transform.domain.shape + + @property + def ndim(self) -> int: + return self._transform.input_rank + + @property + def size(self) -> int: + return math.prod(self.shape) + + @property + def dtype(self) -> Any: + """The wrapped array's dtype; views never change it.""" + return self._array.dtype + + # -- shape of the selection --------------------------------------------- + + @property + def is_box(self) -> bool: + """Whether this view selects a rectangular region rather than a point list. + + True exactly when the composed transform's output maps are all + `ConstantMap` or `DimensionMap` — no `ArrayMap`. Such a selection is + affine and monotone along every axis, so it is described completely by + an interval **and a stride** per dimension: + [`bounding_box`][zarr_indexing.lazy_array.LazyArray.bounding_box] + together with + [`strides`][zarr_indexing.lazy_array.LazyArray.strides]. Basic indexing, + however deeply composed, stays a box; one `oindex`, `vindex`, or mask + anywhere in the chain leaves it forever. + + A box is *dense* — every cell of its bounding box selected — only when + every stride is 1. A strided box covers its hull sparsely: + `lazy[10:50, ::4]` selects 40x20 cells out of a 40x77 hull, so a + consumer that slabs the hull and discards the rest reads 3.85x what it + needs. Check `strides` before treating a box as a single slab read. + + The distinction is what lets a consumer decide between a slab read and a + gather; see [the design notes](../design-notes.md) for why it is a + category rather than an optimization. + + Examples + -------- + >>> import numpy as np + >>> array = LazyArray(np.arange(12).reshape(3, 4)) + >>> (array.lazy[1:, ::2].is_box, array.lazy.oindex[[2, 0], :].is_box) + (True, False) + """ + return not any(isinstance(m, ArrayMap) for m in self._transform.output) + + def bounding_box(self) -> tuple[tuple[int, int], ...] | None: + """The storage region this view touches, one interval per storage dimension. + + Defined for any selection, box or not, as the hull: the smallest + `[inclusive_min, exclusive_max)` interval per dimension of the array + this view reads from that contains every coordinate the selection + reaches. + + The hull is only *dense* — every cell in it selected — for a box whose + every stride is 1. A strided box selects a sublattice of its hull (pair + this with [`strides`][zarr_indexing.lazy_array.LazyArray.strides] to + describe it fully), and a query's hull is merely a superset that can be + arbitrarily loose: `oindex[[0, 999]]` has a 1000-wide hull over two + rows. + + Returns + ------- + tuple of (int, int), or None + One interval per storage dimension, or `None` when the view is + empty (`size == 0`) and so touches no coordinate at all — an empty + selection has no meaningful origin to report an interval around. + + Notes + ----- + The coordinates are those of the array this view reads from, which for + the `array` of a [`Partition`][zarr_indexing.lazy_array.Partition] is + **part-local** — relative to that part's own box, not the wrapped array. + Two parts of the same view can therefore report identical bounding + boxes; [`Partition.box`][zarr_indexing.lazy_array.Partition] gives the + global box they sit in. + + Examples + -------- + >>> import numpy as np + >>> array = LazyArray(np.arange(12).reshape(3, 4)) + >>> array.lazy[1:, ::2].bounding_box() + ((1, 3), (0, 3)) + >>> array.lazy.oindex[[2, 0], :].bounding_box() + ((0, 3), (0, 4)) + >>> array.lazy[1:1].bounding_box() is None + True + """ + if self.size == 0: + return None + domain = self._transform.domain + bounds: list[tuple[int, int]] = [] + for m in self._transform.output: + if isinstance(m, ConstantMap): + bounds.append((m.offset, m.offset + 1)) + elif isinstance(m, DimensionMap): + d = m.input_dimension + first = m.offset + m.stride * domain.inclusive_min[d] + last = m.offset + m.stride * (domain.exclusive_max[d] - 1) + bounds.append((min(first, last), max(first, last) + 1)) + else: + coords = m.offset + m.stride * m.index_array + bounds.append((int(coords.min()), int(coords.max()) + 1)) + return tuple(bounds) + + def strides(self) -> tuple[int, ...] | None: + """The step between selected coordinates, one per storage dimension. + + Completes the description a box selection needs: `bounding_box()` says + *where*, `strides()` says *how densely*. A stride of 1 means every cell + of the hull along that dimension is selected; `k` means every `k`-th. + Dimensions fixed by an integer index report 1 — they span a single + coordinate. + + Returns + ------- + tuple of int, or None + One positive stride per storage dimension, or `None` when + [`is_box`][zarr_indexing.lazy_array.LazyArray.is_box] is false: a + query's coordinates are a lookup table and have no step. + + Notes + ----- + The magnitude only. A reversing view (`lazy[::-1]`) selects the same + *set* of coordinates as its forward twin, so it reports the same + bounding box and the same strides; the traversal direction lives in the + transform, not in this description of the region touched. + + Examples + -------- + >>> import numpy as np + >>> array = LazyArray(np.arange(24).reshape(4, 6)) + >>> (array.lazy[1:, ::2].bounding_box(), array.lazy[1:, ::2].strides()) + (((1, 4), (0, 5)), (1, 2)) + >>> array.lazy[2, ::3].strides() + (1, 3) + >>> array.lazy.oindex[[2, 0], :].strides() is None + True + """ + if not self.is_box: + return None + return tuple( + 1 if isinstance(m, ConstantMap) else abs(m.stride) for m in self._transform.output + ) + + # -- partitioning ------------------------------------------------------- + + def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> LazyArray: + """Return the same view, read in different parts. + + The transform, the wrapped array, and therefore `result()` are all + unchanged; only the boxes the read is broken into differ. Cheap: nothing + is copied and nothing is read. + + Parameters + ---------- + parts + Either a uniform box shape (one integer per dimension of the base + array, with the trailing box clipped to the extent), dask-convention + per-axis sizes (one sequence of box extents per dimension, each + summing to the base extent), or `None` for a single part covering + the whole array — which makes `result()` lower the view in one shot + instead of assembling it block by block. + + Returns + ------- + LazyArray + The same view with a new partitioning. + + Raises + ------ + ValueError + If `parts` has the wrong length, mixes the two conventions, contains + a non-positive extent, or declares per-axis sizes that do not sum to + the base shape. Unlike discovery, this is our own API and validates + strictly. + + Examples + -------- + >>> import numpy as np + >>> view = LazyArray(np.arange(12).reshape(3, 4)) + >>> [part.base_coords for part in view.with_parts((2, 3)).parts()] + [(0, 0), (0, 1), (1, 0), (1, 1)] + """ + grids = None if parts is None else dimension_grids_from_chunks(parts, self._base_shape) + return LazyArray._derive(self._array, self._transform, grids, self._window) + + def parts(self) -> Iterator[Partition]: + """Iterate the base partitioning, projected through this view. + + Yields one [`Partition`][zarr_indexing.lazy_array.Partition] per box the + view actually touches. The parts tile the view exactly and disjointly, + and each carries a `LazyArray` that can be resolved on its own — in + another thread, in another order, or not at all. + + A wrapper with no partitioning (see `with_parts`) yields a single part + covering the whole array. + + Yields + ------ + Partition + One per touched box, in the resolver's own order. + + Examples + -------- + >>> import numpy as np + >>> view = LazyArray(np.arange(12).reshape(3, 4)).with_parts((2, 2)) + >>> part = next(view.lazy[:, 1:].parts()) + >>> (part.base_coords, part.array.shape, part.is_complete) + ((0, 0), (2, 1), False) + """ + base_shape = self._base_shape + grids = self._parts if self._parts is not None else _whole_array_grids(base_shape) + out_shape = self.shape + rank = len(base_shape) + + for base_coords, local, out_indices in iter_chunk_transforms(self._transform, grids): + origin = tuple(grid.chunk_offset(c) for grid, c in zip(grids, base_coords, strict=True)) + extent = tuple(grid.chunk_size(c) for grid, c in zip(grids, base_coords, strict=True)) + if origin == (0,) * rank and extent == base_shape: + # The part is the whole base: lowering directly against the + # source beats materializing a block that is the source. + window = self._window + elif self._window is None: + window = tuple(slice(o, o + e) for o, e in zip(origin, extent, strict=True)) + else: + window = tuple( + slice(w.start + o, w.start + o + e) + for w, o, e in zip(self._window, origin, extent, strict=True) + ) + # The global box, computed from the origin directly rather than from + # `window`: a part covering the whole base carries no window (so + # nothing is pre-materialized) but still sits somewhere concrete. + if self._window is None: + global_origin = origin + else: + global_origin = tuple( + w.start + o for w, o in zip(self._window, origin, strict=True) + ) + yield Partition( + base_coords=base_coords, + box=tuple((o, o + e) for o, e in zip(global_origin, extent, strict=True)), + array=LazyArray._derive(self._array, local, None, window), + out_selection=_partition_out_selection(local, out_indices, out_shape), + is_complete=_covers_whole_part(local, extent), + ) + + # -- indexing ----------------------------------------------------------- + + @property + def lazy(self) -> _LazyIndexer: + """Lazy indexing: `lazy[...]`, `lazy.oindex[...]`, `lazy.vindex[...]`. + + Each returns a new `LazyArray` view; no data is read. + """ + return _LazyIndexer(self._select) + + def _select(self, selection: Any, mode: SelectionMode) -> LazyArray: + transform = self._transform + if mode != "basic": + # NumPy applies scalar integers as basic indices before the advanced + # ones, dropping their axes. Split them into their own step. + scalar_selection, selection = split_scalar_axes(selection, transform.domain, mode) + if scalar_selection is not None: + transform = selection_to_transform(scalar_selection, transform, "basic") + transform = transform.translate_domain_to((0,) * transform.input_rank) + literal = normalize_positional_selection(selection, transform.domain, mode) + composed = selection_to_transform(literal, transform, mode) + return LazyArray._derive(self._array, composed, self._parts, self._window) + + def __getitem__(self, selection: Any) -> Any: + """Read a basic selection **eagerly**, like `numpy.ndarray.__getitem__`. + + Eager, not lazy, so that a `LazyArray` is a drop-in duck array for + consumers (dask's `from_array`, `numpy.asarray`) that expect indexing to + produce data. Use `.lazy[...]` for the lazy form. + """ + return self._select(selection, "basic").result() + + def result(self) -> Any: + """Materialize this view. + + Assembles the view from its `parts()`, reading each box once. A wrapper + with a single whole-array part skips the buffer entirely and lowers + straight to array operations. + + Returns + ------- + array + An array of shape `self.shape`, identical whatever partitioning is + in force. A partitioned view produces a NumPy array (see the module + docstring's device caveat); a single whole-array part produces + whatever the wrapped array's namespace produces. A view with a + zero-rank domain returns a zero-dimensional array, not a scalar. + """ + if self._parts is None: + return _lower(self._read_base(), self._transform) + + out_shape = self.shape + out = np.empty(out_shape, dtype=np.dtype(self.dtype)) + if math.prod(out_shape) > 0: + for part in self.parts(): + value = np.asarray(part.array.result()) + if len(out_shape) == 0: + value = value.reshape(()) + out[part.out_selection] = value + return out + + # -- protocols ---------------------------------------------------------- + + def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: + if copy is False: + raise ValueError( + "a LazyArray cannot be converted to a NumPy array without a " + "copy: materializing a view always produces new data" + ) + return np.asarray(self.result(), dtype=dtype) + + def __dask_tokenize__(self) -> Any: + """A deterministic token: the wrapped array, the view, and the parts. + + Two wrappers token-equal when they wrap the same data, address the same + cells, and read them in the same boxes. The view contributes its + canonical ndsel body, so transforms that differ only in representation + token alike. See `_wrapped_token` for the determinism scope of the + wrapped array's contribution; dask is imported lazily and is never a + requirement of this package. + """ + return ( + type(self).__qualname__, + _wrapped_token(self._array), + None if self._window is None else tuple((s.start, s.stop) for s in self._window), + json.dumps(transform_to_canonical(self._transform), sort_keys=True), + None if self._parts is None else tuple(grid.sizes for grid in self._parts), + ) + + def __len__(self) -> int: + if self.ndim == 0: + raise TypeError("len() of unsized object") + return self.shape[0] + + def __iter__(self) -> Iterator[Any]: + """Iterate eagerly over the first axis, like a NumPy array. + + The rank check happens in `__iter__` itself rather than in the + generator, so `iter(view)` on a zero-rank view raises immediately as + NumPy's does, instead of waiting for the first `next`. + """ + if self.ndim == 0: + raise TypeError("iteration over a 0-d array") + return (self[position] for position in range(self.shape[0])) + + # NumPy's own conversions decide what a size-1 (or wrong-sized) view means, + # including which exception it raises, so these delegate rather than + # reimplement. Each materializes the view first. + def __bool__(self) -> bool: + return bool(self.result()) + + def __int__(self) -> int: + return int(self.result()) + + def __float__(self) -> float: + return float(self.result()) + + def __index__(self) -> int: + return operator.index(self.result()) + + def __repr__(self) -> str: + wrapped = type(self._array).__name__ + described = [f"{wrapped} shape={self.shape} dtype={self.dtype}"] + if not _is_identity_transform(self._transform, self._base_shape): + described.append(f"view={self._transform.selection_repr}") + return f"" + + +class _LazyIndexer: + """The `.lazy` accessor: builds views instead of reading data. + + Holds the owning view's bound `_select` rather than the view itself, so the + accessor classes never reach into another object's internals. + """ + + __slots__ = ("_select",) + + def __init__(self, select: SelectFn) -> None: + self._select = select + + def __getitem__(self, selection: Any) -> LazyArray: + """Basic (integer / slice / ellipsis) indexing, lazily.""" + return self._select(selection, "basic") + + @property + def oindex(self) -> _LazyOIndex: + """Orthogonal (outer-product) indexing, lazily.""" + return _LazyOIndex(self._select) + + @property + def vindex(self) -> _LazyVIndex: + """Vectorized (coordinate / mask) indexing, lazily.""" + return _LazyVIndex(self._select) + + +class _LazyOIndex: + """`lazy.oindex[...]` — one selection per axis, combined as an outer product.""" + + __slots__ = ("_select",) + + def __init__(self, select: SelectFn) -> None: + self._select = select + + def __getitem__(self, selection: Any) -> LazyArray: + return self._select(selection, "orthogonal") + + +class _LazyVIndex: + """`lazy.vindex[...]` — correlated coordinate arrays, or a single mask.""" + + __slots__ = ("_select",) + + def __init__(self, select: SelectFn) -> None: + self._select = select + + def __getitem__(self, selection: Any) -> LazyArray: + return self._select(selection, "vectorized") diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index e1a3898b1d..1adbbf74f9 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -30,7 +30,7 @@ import math from dataclasses import dataclass -from typing import Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np @@ -38,6 +38,9 @@ from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +if TYPE_CHECKING: + from collections.abc import Sequence + @dataclass(frozen=True, slots=True) class IndexTransform: @@ -363,7 +366,14 @@ def _intersect_orthogonal( # axis, or `input_dimension` for a degenerate length-1 array). Filter # along that axis and keep the array at full input rank so the # singleton axes it broadcasts over are preserved. - d = _array_map_dependent_axis(m) + axis = array_map_dependent_axis(m) + if axis is None: + raise ValueError( + f"output[{out_dim}] is an ArrayMap that varies over no input " + "dimension; a map with no dependency axis should have been " + "collapsed to a ConstantMap" + ) + d = axis storage = m.offset + m.stride * m.index_array mask = (storage >= lo) & (storage < hi) # The array is singleton on every axis but `d`, so its mask reduces @@ -421,8 +431,6 @@ def _intersect_correlated( `out_indices` is the flat scatter index into the (row-major flattened) output buffer, of shape `(surviving_points,) + (residual slice sizes)`. """ - corr_maps = [cast("ArrayMap", transform.output[i]) for i in correlated_dims] - # Mixing correlated and orthogonal ArrayMaps in one transform is not produced # by any single selection and is not supported here. orthogonal_array_dims = [ @@ -436,10 +444,15 @@ def _intersect_correlated( "ArrayMaps is not supported" ) - # The broadcast (dependency) axes are shared by every correlated map; they are - # the leading axes of the domain, followed by the residual slice axes. - broadcast_axes = _array_map_dependency_axes(corr_maps[0].index_array) - broadcast_shape = tuple(corr_maps[0].index_array.shape[a] for a in broadcast_axes) + # The broadcast axes are exactly the input axes no `DimensionMap` binds: a + # correlated transform's input domain is its residual slice axes plus the + # collapsed broadcast block. Deriving them by complement rather than from the + # index array's non-singleton axes keeps this correct when a broadcast axis + # is itself size 1, and when NumPy's placement rule puts the broadcast block + # somewhere other than the front (see `_broadcast_insertion_point`). + bound_axes = {m.input_dimension for m in transform.output if isinstance(m, DimensionMap)} + broadcast_axes = tuple(a for a in range(transform.input_rank) if a not in bound_axes) + broadcast_shape = tuple(transform.domain.shape[a] for a in broadcast_axes) # Joint bounds mask over the broadcast block. combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None @@ -523,23 +536,35 @@ def _intersect_correlated( ) result = IndexTransform(domain=new_domain, output=tuple(new_output)) - # Flat scatter index into the row-major output buffer of shape - # (broadcast points, residual slice sizes...): flat = point * prod(slice) + - # (row-major offset within the slice block). - prod_slice = 1 - for _d, _lo, _hi, full, _m in slice_dims: - prod_slice *= full - out_indices: np.ndarray[Any, np.dtype[np.intp]] = (surviving * prod_slice).reshape( + # Flat scatter index into the caller's row-major output buffer, whose shape + # is the *input* domain's shape. The buffer is addressed positionally, so + # this assumes a zero-origin domain — the resolvers normalize with + # `translate_domain_to` before resolving. + # + # Each surviving point is a flat index into the broadcast block; unravel it + # to per-axis coordinates so the buffer stride of each broadcast axis is + # applied at its real position, wherever NumPy's placement rule put it. + domain_shape = transform.domain.shape + buffer_strides = [1] * len(domain_shape) + for axis in range(len(domain_shape) - 2, -1, -1): + buffer_strides[axis] = buffer_strides[axis + 1] * domain_shape[axis + 1] + + point_offsets = np.zeros(n_points, dtype=np.intp) + if len(broadcast_shape) > 0: + for axis, coords_along_axis in zip( + broadcast_axes, np.unravel_index(surviving, broadcast_shape), strict=True + ): + point_offsets = point_offsets + coords_along_axis.astype(np.intp) * buffer_strides[axis] + + out_indices: np.ndarray[Any, np.dtype[np.intp]] = point_offsets.reshape( (n_points,) + (1,) * n_slice ) - running = 1 - for j in range(n_slice - 1, -1, -1): - _d, nlo, nhi, full, _m = slice_dims[j] - coords = np.arange(nlo, nhi, dtype=np.intp) * running + for j in range(n_slice): + d, nlo, nhi, _full, _m = slice_dims[j] + coords = np.arange(nlo, nhi, dtype=np.intp) * buffer_strides[d] shape = [1] * (1 + n_slice) shape[1 + j] = coords.size out_indices = out_indices + coords.reshape(shape) - running *= full return (result, out_indices.astype(np.intp)) @@ -792,18 +817,31 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra raise RuntimeError(f"unexpected: dimension {d} not handled") else: # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) - new_arr = _reindex_array(m, normalized, transform.domain) - array_input_dim: int | None = None + old_axes = set(_array_map_dependency_axes(m.index_array)) if m.input_dimension is not None: - array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) - new_output.append( - ArrayMap( - index_array=new_arr, - offset=m.offset, - stride=m.stride, - input_dimension=array_input_dim, + old_axes.add(m.input_dimension) + new_arr = _reindex_array(m, normalized, transform.domain) + if len(old_axes) > 0 and old_axes <= dropped_dims and new_arr.size == 1: + # Degenerate collapse: every input axis this map varied over was + # consumed by an integer index, so it now designates exactly one + # storage coordinate. Keeping it as an ArrayMap would leave a + # dangling `input_dimension` pointing at an axis that no longer + # exists — and, after renumbering, aliasing a *different* axis. + new_output.append( + ConstantMap(offset=m.offset + m.stride * int(new_arr.reshape(-1)[0])) + ) + else: + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) ) - ) return IndexTransform(domain=new_domain, output=tuple(new_output)) @@ -821,18 +859,35 @@ def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) -def _array_map_dependent_axis(m: ArrayMap) -> int: +def array_map_dependent_axis(m: ArrayMap) -> int | None: """Return the single input axis an orthogonal `ArrayMap` varies over. Normally this is the array's one non-singleton axis. A degenerate length-1 orthogonal selection normalizes to an all-singleton shape (its dependency axes are empty and indistinguishable by shape from a scalar), so - `input_dimension` breaks the tie — it records the axis the map binds. + `input_dimension` breaks the tie — it records the *live* axis the map binds. + + Returns + ------- + int or None + The axis the map varies over, or `None` when it varies over no input + axis at all — a correlated (`vindex`) map, or a map that has become + constant. `None` is a real answer, not an error: callers that need an + axis must handle it rather than reading a stale `input_dimension`. + Indexing operations collapse a map whose every dependency axis is + consumed by an integer index into a `ConstantMap`, so an + `input_dimension` surviving here always names a live axis. + + Raises + ------ + ValueError + If the map varies over more than one axis, which makes it correlated + rather than orthogonal. """ dep = _array_map_dependency_axes(m.index_array) if len(dep) == 1: return dep[0] - if m.input_dimension is not None: + if len(dep) == 0: return m.input_dimension raise ValueError( f"orthogonal ArrayMap must vary over exactly one axis; got dependency " @@ -1014,6 +1069,27 @@ def __getitem__(self, selection: Any) -> IndexTransform: return _apply_vindex(self._transform, selection) +def _broadcast_insertion_point(array_dims: Sequence[int], slice_dims: Sequence[int]) -> int: + """Where the broadcast dimensions land, as a count of leading slice dimensions. + + NumPy's advanced-indexing placement rule: when the advanced indices are all + next to each other in the index tuple, the broadcast dimensions are inserted + at the spot they occupied; when a slice separates them, they lead. So + `a[:, i, j]` has shape `(len(a), *broadcast)` while `a[i, :, j]` has shape + `(*broadcast, a.shape[1])`. + + Returns the number of slice dimensions that precede the broadcast block; `0` + means the broadcast dimensions lead. + """ + if len(array_dims) == 0: + return 0 + first, last = array_dims[0], array_dims[-1] + separated = any(first < d < last for d in slice_dims) + if separated: + return 0 + return sum(1 for d in slice_dims if d < first) + + def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: """Apply vectorized indexing to an IndexTransform. @@ -1092,27 +1168,30 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: broadcast_arrays = [] broadcast_shape = () - # Build new domain: broadcast dims first, then slice dims - new_inclusive_min: list[int] = [] - new_exclusive_max: list[int] = [] - - # Broadcast dimensions - for s in broadcast_shape: - new_inclusive_min.append(0) - new_exclusive_max.append(s) - # Slice dimensions (preserved-domain literal semantics, like basic indexing) slice_dim_params: dict[int, tuple[int, int, int]] = {} + slice_bounds: list[tuple[int, int]] = [] for old_dim in slice_dims: sel = processed[old_dim] assert isinstance(sel, slice) lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) - new_inclusive_min.append(origin) - new_exclusive_max.append(origin + size) + slice_bounds.append((origin, origin + size)) slice_dim_params[old_dim] = (start, step, origin) + n_before = _broadcast_insertion_point(array_dims, slice_dims) + + # Build the new domain with NumPy's placement rule: the broadcast + # (correlated) dimensions sit where the advanced indices sat when those are + # adjacent, and lead when a slice separates them. + new_inclusive_min = [lo for lo, _ in slice_bounds[:n_before]] + new_exclusive_max = [hi for _, hi in slice_bounds[:n_before]] + new_inclusive_min.extend([0] * len(broadcast_shape)) + new_exclusive_max.extend(broadcast_shape) + new_inclusive_min.extend(lo for lo, _ in slice_bounds[n_before:]) + new_exclusive_max.extend(hi for _, hi in slice_bounds[n_before:]) + new_domain = IndexDomain( inclusive_min=tuple(new_inclusive_min), exclusive_max=tuple(new_exclusive_max), @@ -1139,7 +1218,9 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: # dependency axes derived from the shape coincide — the signature # of a pointwise scatter rather than an outer product. broadcast_arr = array_dim_to_broadcast[d] - full_arr = broadcast_arr.reshape(broadcast_shape + (1,) * len(slice_dims)) + full_arr = broadcast_arr.reshape( + (1,) * n_before + broadcast_shape + (1,) * (len(slice_dims) - n_before) + ) new_output.append( ArrayMap( index_array=full_arr, @@ -1152,7 +1233,8 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: start, step, origin = slice_dim_params[d] new_offset = m.offset + m.stride * (start - step * origin) new_stride = m.stride * step - new_input_dim = n_broadcast_dims + slice_dims.index(d) + position = slice_dims.index(d) + new_input_dim = position if position < n_before else position + n_broadcast_dims new_output.append( DimensionMap( input_dimension=new_input_dim, offset=new_offset, stride=new_stride diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py new file mode 100644 index 0000000000..0a42fa1aec --- /dev/null +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -0,0 +1,1090 @@ +"""Tests for `zarr_indexing.grid` grids and the `LazyArray` wrapper. + +The happy-path suite is a single oracle test: every selection case is applied +both to a `LazyArray` and to the NumPy array it wraps, and the results must +match. The case list is crossed with four source flavours — an unchunked NumPy +array, NumPy with each of the two declared chunk conventions, and a zarr array +whose chunking is auto-discovered — so the chunked and unchunked resolution +strategies are held to the same answers. +""" + +from __future__ import annotations + +import operator +import pickle +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Any + +import numpy as np +import pytest + +from zarr_indexing import ( + ArrayMap, + ConstantMap, + DimensionMap, + EdgeDimensionGrid, + IndexTransform, + LazyArray, + array_map_dependent_axis, + dimension_grids_from_chunks, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + +SHAPE = (7, 5, 4) +PART_SHAPE = (3, 2, 3) +EXPLICIT_PARTS = ((3, 3, 1), (2, 2, 1), (3, 1)) + + +def reference() -> np.ndarray[Any, np.dtype[np.int64]]: + """The array every source flavour holds, and the oracle for every case.""" + return np.arange(int(np.prod(SHAPE)), dtype=np.int64).reshape(SHAPE) + + +def outer(ref: np.ndarray[Any, Any], selections: Sequence[Any]) -> np.ndarray[Any, Any]: + """NumPy oracle for orthogonal indexing: the outer product of per-axis selections. + + Scalar integers are basic indices — NumPy applies them first and drops the + axis — so they are peeled off before the outer product is formed. + """ + scalars = tuple( + sel if isinstance(sel, (int, np.integer)) and not isinstance(sel, bool) else slice(None) + for sel in selections + ) + reduced = ref[scalars] + axes = [ + np.arange(size)[sel] + for size, sel in zip( + reduced.shape, + [ + s + for s in selections + if not (isinstance(s, (int, np.integer)) and not isinstance(s, bool)) + ], + strict=True, + ) + ] + if len(axes) == 0: + return reduced + return reduced[np.ix_(*axes)] + + +# --------------------------------------------------------------------------- +# EdgeDimensionGrid +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("sizes", "expected_index_to_chunk", "expected_offsets", "expected_sizes"), + [ + # A clipped edge chunk. + ((3, 3, 1), [0, 0, 0, 1, 1, 1, 2], [0, 3, 6], [3, 3, 1]), + # A single chunk covering the whole axis. + ((4,), [0, 0, 0, 0], [0], [4]), + # A size-1 axis. + ((1,), [0], [0], [1]), + # Irregular sizes: the grid does not have to be regular. + ((1, 2, 1), [0, 1, 1, 2], [0, 1, 3], [1, 2, 1]), + ], +) +def test_edge_dimension_grid( + sizes: tuple[int, ...], + expected_index_to_chunk: list[int], + expected_offsets: list[int], + expected_sizes: list[int], +) -> None: + """All four `DimensionGridLike` methods agree with hand-computed values.""" + grid = EdgeDimensionGrid(sizes) + extent = sum(sizes) + + assert grid.num_chunks == len(sizes) + assert grid.extent == extent + assert [grid.index_to_chunk(i) for i in range(extent)] == expected_index_to_chunk + assert [grid.chunk_offset(c) for c in range(len(sizes))] == expected_offsets + assert [grid.chunk_size(c) for c in range(len(sizes))] == expected_sizes + np.testing.assert_array_equal( + grid.indices_to_chunks(np.arange(extent, dtype=np.intp)), + np.asarray(expected_index_to_chunk, dtype=np.intp), + ) + + +def test_edge_dimension_grid_rejects_nonpositive_size() -> None: + with pytest.raises(ValueError, match="chunk sizes must be positive"): + EdgeDimensionGrid((3, 0, 2)) + + +def test_edge_dimension_grid_rejects_out_of_bounds_index() -> None: + with pytest.raises(IndexError, match="out of bounds for an axis of extent 4"): + EdgeDimensionGrid((3, 1)).index_to_chunk(4) + + +def test_edge_dimension_grid_rejects_out_of_bounds_chunk() -> None: + with pytest.raises(IndexError, match="chunk index 2 is out of bounds"): + EdgeDimensionGrid((3, 1)).chunk_offset(2) + + +@pytest.mark.parametrize( + ("chunks", "shape", "expected"), + [ + # Uniform chunk shape, tail clipped. + ((3, 4), (7, 4), ((3, 3, 1), (4,))), + # Dask-convention per-axis sizes, passed through. + (((3, 3, 1), (2, 2)), (7, 4), ((3, 3, 1), (2, 2))), + # A chunk longer than the axis collapses to one clipped chunk. + ((10,), (4,), ((4,),)), + # A zero-length axis has no chunks at all. + ((3,), (0,), ((),)), + ], +) +def test_dimension_grids_from_chunks( + chunks: Any, shape: tuple[int, ...], expected: tuple[tuple[int, ...], ...] +) -> None: + grids = dimension_grids_from_chunks(chunks, shape) + assert tuple(grid.sizes for grid in grids) == expected + + +def test_dimension_grids_from_chunks_rejects_wrong_length() -> None: + with pytest.raises(ValueError, match="one entry per dimension"): + dimension_grids_from_chunks((3,), (7, 4)) + + +def test_dimension_grids_from_chunks_rejects_mixed_conventions() -> None: + with pytest.raises(ValueError, match="not a mixture"): + dimension_grids_from_chunks((3, (2, 2)), (7, 4)) + + +def test_dimension_grids_from_chunks_rejects_wrong_total() -> None: + with pytest.raises(ValueError, match="sum to 5, but the array extent is 7"): + dimension_grids_from_chunks(((3, 2), (4,)), (7, 4)) + + +def test_dimension_grids_from_chunks_rejects_non_sequence_entry() -> None: + """A float entry is neither convention, and says so instead of raising TypeError.""" + with pytest.raises(ValueError, match=r"3\.5 at dimension 0 is neither"): + dimension_grids_from_chunks((3.5, (4,)), (7, 4)) + + +def test_array_map_dependent_axis_reports_no_axis() -> None: + """A map varying over nothing answers None rather than a stale binding.""" + correlated = ArrayMap(index_array=np.array([[1], [2]], dtype=np.intp)) + assert array_map_dependent_axis(correlated) == 0 + assert array_map_dependent_axis(ArrayMap(index_array=np.array([[3]], dtype=np.intp))) is None + + +def test_scalar_on_a_fancy_axis_collapses_to_a_constant() -> None: + """The degenerate-collapse rule: an all-singleton ArrayMap becomes a ConstantMap. + + Without it the map keeps an `input_dimension` naming an axis the integer + index just removed, which after renumbering aliases a different axis. + """ + view = IndexTransform.from_shape((7, 5)).oindex[np.array([3, 1]), slice(None)][0] + assert view.output[0] == ConstantMap(offset=3) + assert isinstance(view.output[1], DimensionMap) + assert view.domain.shape == (5,) + + +# --------------------------------------------------------------------------- +# Sources +# --------------------------------------------------------------------------- + + +def make_zarr_source() -> Any: + """A zarr array holding `reference()`, whose parts are auto-discovered.""" + zarr = pytest.importorskip("zarr") + array = zarr.create_array({}, shape=SHAPE, chunks=PART_SHAPE, dtype="int64") + array[:] = reference() + return array + + +def make_source(flavour: str) -> LazyArray: + """Build a `LazyArray` over `reference()` with the requested partitioning.""" + data = reference() + if flavour == "numpy-whole": + return LazyArray(data) + if flavour == "numpy-explicit-parts": + return LazyArray(data).with_parts(EXPLICIT_PARTS) + if flavour == "numpy-uniform-parts": + return LazyArray(data).with_parts(PART_SHAPE) + if flavour == "zarr": + return LazyArray(make_zarr_source()) + if flavour == "zarr-misaligned": + # Parts that deliberately straddle the zarr array's own chunks. + return LazyArray(make_zarr_source()).with_parts((4, 3, 3)) + raise AssertionError(f"unknown source flavour {flavour!r}") + + +FLAVOURS = [ + "numpy-whole", + "numpy-explicit-parts", + "numpy-uniform-parts", + "zarr", + "zarr-misaligned", +] + + +@pytest.fixture(params=FLAVOURS) +def source(request: pytest.FixtureRequest) -> LazyArray: + return make_source(request.param) + + +# --------------------------------------------------------------------------- +# The oracle +# --------------------------------------------------------------------------- + +MASK = (reference() % 11) == 0 +# A mask over the two trailing axes, for the `vindex[..., mask]` idiom. +TRAILING_MASK = (reference()[0] % 3) == 0 + +# (id, lazy view builder, NumPy oracle). Integer scalars are deliberately absent +# from the orthogonal cases: the transform algebra keeps an int-selected axis as +# a length-1 axis under `oindex`, which `np.ix_` cannot express. +CASES: list[tuple[str, Callable[[LazyArray], LazyArray], Callable[[Any], Any]]] = [ + ( + "basic-strided-and-int-drop", + lambda a: a.lazy[1:6:2, :, -1], + lambda r: r[1:6:2, :, -1], + ), + ("basic-ellipsis", lambda a: a.lazy[..., -2], lambda r: r[..., -2]), + ("basic-negative-scalar", lambda a: a.lazy[-3], lambda r: r[-3]), + ("basic-empty", lambda a: a.lazy[:, 2:2, :], lambda r: r[:, 2:2, :]), + ("basic-all-scalars", lambda a: a.lazy[-1, 0, 2], lambda r: r[-1, 0, 2]), + ( + "oindex-unsorted-duplicates-multi-axis", + lambda a: a.lazy.oindex[[4, 0, 0, 2], :, [3, 1]], + lambda r: outer(r, ([4, 0, 0, 2], slice(None), [3, 1])), + ), + ( + "oindex-negative-and-slice", + lambda a: a.lazy.oindex[:, [-1, 0], 1:4], + lambda r: outer(r, (slice(None), [-1, 0], slice(1, 4))), + ), + ( + "oindex-boolean-axis", + lambda a: a.lazy.oindex[np.array([True, False, True, False, False, False, True]), :, :], + lambda r: outer( + r, + (np.array([True, False, True, False, False, False, True]), slice(None), slice(None)), + ), + ), + ( + "vindex-coordinates", + lambda a: a.lazy.vindex[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])], + lambda r: r[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])], + ), + ( + "vindex-broadcast-pair", + lambda a: a.lazy.vindex[np.array([[0], [6]]), np.array([1, 4]), np.array([2, 0])], + lambda r: r[np.array([[0], [6]]), np.array([1, 4]), np.array([2, 0])], + ), + ( + "vindex-negative-coordinates", + lambda a: a.lazy.vindex[np.array([-1, -7]), np.array([-2, 0]), np.array([0, -1])], + lambda r: r[np.array([-1, -7]), np.array([-2, 0]), np.array([0, -1])], + ), + ("vindex-mask", lambda a: a.lazy.vindex[MASK], lambda r: r[MASK]), + ( + "compose-basic-then-oindex", + lambda a: a.lazy[1:6].lazy.oindex[[3, 0, 0], [4, 1], :], + lambda r: outer(r[1:6], ([3, 0, 0], [4, 1], slice(None))), + ), + ( + "compose-oindex-then-basic-other-axis", + lambda a: a.lazy.oindex[[4, 0, 2], :, :].lazy[:, 1:4, ::2], + lambda r: outer(r, ([4, 0, 2], slice(None), slice(None)))[:, 1:4, ::2], + ), + ( + "compose-basic-then-basic", + lambda a: a.lazy[2:, 1:].lazy[::2, -1], + lambda r: r[2:, 1:][::2, -1], + ), + ( + "compose-basic-then-vindex", + lambda a: a.lazy[1:6, :, 1:].lazy.vindex[ + np.array([0, 4]), np.array([2, 0]), np.array([1, 2]) + ], + lambda r: r[1:6, :, 1:][np.array([0, 4]), np.array([2, 0]), np.array([1, 2])], + ), + # Scalar integers are basic indices in the positional dialect: they drop the + # axis, in every mode, exactly as NumPy does. + ("oindex-scalar-drops-axis", lambda a: a.lazy.oindex[0], lambda r: r[0]), + ( + "oindex-scalar-with-arrays", + lambda a: a.lazy.oindex[0, [1, 2], :], + lambda r: outer(r, (0, [1, 2], slice(None))), + ), + ( + "oindex-scalar-middle-axis", + lambda a: a.lazy.oindex[[3, 1], -1, :], + lambda r: outer(r, ([3, 1], -1, slice(None))), + ), + ("oindex-all-scalars", lambda a: a.lazy.oindex[0, 1, 2], lambda r: r[0, 1, 2]), + ("vindex-all-scalars", lambda a: a.lazy.vindex[0, 1, 2], lambda r: r[0, 1, 2]), + ( + "vindex-scalar-with-arrays", + lambda a: a.lazy.vindex[0, [1, 2], [3, 0]], + lambda r: r[0, [1, 2], [3, 0]], + ), + ( + "vindex-scalar-on-middle-axis", + lambda a: a.lazy.vindex[[1, 2], 0, [3, 0]], + lambda r: r[[1, 2], 0, [3, 0]], + ), + # Scalar applied to a previously fancy-indexed axis, both orders. + ( + "compose-oindex-then-scalar", + lambda a: a.lazy.oindex[[3, 1], :, :].lazy[0], + lambda r: outer(r, ([3, 1], slice(None), slice(None)))[0], + ), + ( + "compose-oindex-then-scalar-negative", + lambda a: a.lazy.oindex[[3, 1, 1], :, :].lazy[-1, 2], + lambda r: outer(r, ([3, 1, 1], slice(None), slice(None)))[-1, 2], + ), + ( + "compose-scalar-then-oindex", + lambda a: a.lazy[0].lazy.oindex[[3, 1], :], + lambda r: outer(r[0], ([3, 1], slice(None))), + ), + ( + "compose-vindex-then-scalar", + lambda a: a.lazy.vindex[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])].lazy[ + 1 + ], + lambda r: r[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])][1], + ), + # Partial vindex whose coordinate arrays are NOT on the leading axes: NumPy + # inserts the gathered axis where the (adjacent) advanced indices sat. + ( + "vindex-trailing-arrays", + lambda a: a.lazy.vindex[..., np.array([1, 4, 0]), np.array([2, 0, 1])], + lambda r: r[..., np.array([1, 4, 0]), np.array([2, 0, 1])], + ), + ( + "vindex-single-trailing-array", + lambda a: a.lazy.vindex[..., np.array([3, 0, 1])], + lambda r: r[..., np.array([3, 0, 1])], + ), + ( + "vindex-trailing-mask", + lambda a: a.lazy.vindex[..., TRAILING_MASK], + lambda r: r[..., TRAILING_MASK], + ), + ( + "vindex-leading-partial", + lambda a: a.lazy.vindex[np.array([1, 2, 2])], + lambda r: r[np.array([1, 2, 2])], + ), + ( + "compose-basic-then-vindex-trailing", + lambda a: a.lazy[2:, 1:].lazy.vindex[..., np.array([1, 3, 0])], + lambda r: r[2:, 1:][..., np.array([1, 3, 0])], + ), + # A ConstantMap sitting between a slice and the coordinate arrays: NumPy + # counts the integer as an advanced index, so the gathered axis moves to the + # front of the chunk block even though the arrays are trailing. + ( + "compose-vindex-trailing-then-scalar", + lambda a: a.lazy.vindex[..., np.array([3, 0])].lazy[2], + lambda r: r[..., np.array([3, 0])][2], + ), + # Two fancy axes, then a scalar on the first: the surviving ArrayMap ends up + # behind a ConstantMap and a slice, which is where NumPy's integer-counts-as- + # advanced rule bites. + ( + "compose-oindex-two-axes-then-scalar", + lambda a: a.lazy.oindex[[3, 1, 0], 3:5, [2, 0, 2]].lazy[0], + lambda r: outer(r, ([3, 1, 0], slice(3, 5), [2, 0, 2]))[0], + ), + ( + "compose-vindex-trailing-then-scalar-and-slice", + lambda a: a.lazy.vindex[..., np.array([3, 0, 1])].lazy[-1, 1:4], + lambda r: r[..., np.array([3, 0, 1])][-1, 1:4], + ), +] + + +@pytest.mark.parametrize(("build", "oracle"), [c[1:] for c in CASES], ids=[c[0] for c in CASES]) +def test_selection_matches_numpy( + source: LazyArray, + build: Callable[[LazyArray], LazyArray], + oracle: Callable[[Any], Any], +) -> None: + """Every selection resolves to what NumPy computes positionally on the same data.""" + view = build(source) + expected = np.asarray(oracle(reference())) + + assert view.shape == expected.shape + assert view.ndim == expected.ndim + np.testing.assert_array_equal(np.asarray(view.result()), expected) + # `__getitem__` is eager, and `__array__` routes through `result()`. + np.testing.assert_array_equal(np.asarray(view), expected) + + +# --------------------------------------------------------------------------- +# Randomized chain sweep +# --------------------------------------------------------------------------- + + +def _random_basic(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: + selection: list[Any] = [] + for size in shape: + roll = rng.random() + if roll < 0.3: + selection.append(int(rng.integers(-size, size))) + elif roll < 0.7: + start = int(rng.integers(0, size)) + stop = int(rng.integers(start, size + 1)) + selection.append(slice(start, stop, int(rng.integers(1, 4)))) + else: + selection.append(slice(None)) + return tuple(selection) + + +def _random_oindex(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: + selection: list[Any] = [] + for size in shape: + roll = rng.random() + if roll < 0.25: + selection.append(int(rng.integers(-size, size))) + elif roll < 0.65: + count = int(rng.integers(1, 5)) + selection.append(rng.integers(-size, size, size=count).tolist()) + elif roll < 0.8: + mask = rng.random(size) < 0.5 + mask[int(rng.integers(0, size))] = True + selection.append(mask) + else: + start = int(rng.integers(0, size)) + selection.append(slice(start, size)) + return tuple(selection) + + +def _random_vindex(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: + ndim = len(shape) + count = int(rng.integers(1, ndim + 1)) + trailing = bool(rng.random() < 0.5) + axes = range(ndim - count, ndim) if trailing else range(count) + sizes = [shape[axis] for axis in axes] + + entries: list[Any] + if rng.random() < 0.2: + # A single boolean mask spanning the whole covered block. + mask = rng.random(tuple(sizes)) < 0.5 + mask.flat[int(rng.integers(0, mask.size))] = True + entries = [mask] + else: + length = int(rng.integers(1, 5)) + entries = [ + int(rng.integers(-size, size)) + if rng.random() < 0.25 + else rng.integers(-size, size, size=length) + for size in sizes + ] + return (Ellipsis, *entries) if trailing else tuple(entries) + + +def _apply_oracle( + ref: np.ndarray[Any, Any], mode: str, selection: tuple[Any, ...] +) -> np.ndarray[Any, Any]: + if mode == "orthogonal": + return outer(ref, selection) + # NumPy's own semantics *are* basic and vectorized indexing. + return ref[selection] + + +def _apply_view(view: LazyArray, mode: str, selection: tuple[Any, ...]) -> LazyArray: + if mode == "basic": + return view.lazy[selection] + if mode == "orthogonal": + return view.lazy.oindex[selection] + return view.lazy.vindex[selection] + + +def _random_chain(rng: np.random.Generator) -> list[tuple[str, tuple[Any, ...]]]: + """A chain of 2-4 steps with at most one fancy step, in either order.""" + n_steps = int(rng.integers(2, 5)) + fancy_at = int(rng.integers(0, n_steps)) + fancy_mode = "orthogonal" if rng.random() < 0.5 else "vectorized" + + chain: list[tuple[str, tuple[Any, ...]]] = [] + running = reference() + for step in range(n_steps): + if running.ndim == 0 or running.size == 0: + break + mode = fancy_mode if step == fancy_at else "basic" + if mode == "basic": + selection = _random_basic(rng, running.shape) + elif mode == "orthogonal": + selection = _random_oindex(rng, running.shape) + else: + selection = _random_vindex(rng, running.shape) + chain.append((mode, selection)) + running = _apply_oracle(running, mode, selection) + return chain + + +@pytest.mark.parametrize("flavour", FLAVOURS) +def test_random_chains_match_numpy(flavour: str) -> None: + """A seeded sweep of composed chains, under every partitioning. + + The partitioning invariant as a property: `result()` is identical whatever + boxes the read is broken into, including boxes deliberately misaligned with + the source's own. + """ + rng = np.random.default_rng(20260730) + source = make_source(flavour) + partitionings: list[Any] = [None, (2, 2, 2), (7, 5, 4), ((4, 3), (1, 3, 1), (3, 1))] + # Reads against a real store cost more per chain; the NumPy flavours carry + # the bulk of the sweep and exercise the identical code path. + n_chains = 120 if flavour.startswith("zarr") else 400 + + for _ in range(n_chains): + chain = _random_chain(rng) + expected = reference() + for mode, selection in chain: + expected = _apply_oracle(expected, mode, selection) + + view = source + for mode, selection in chain: + view = _apply_view(view, mode, selection) + + assert view.shape == expected.shape, f"{flavour}: {chain}" + np.testing.assert_array_equal( + np.asarray(view.result()), np.asarray(expected), err_msg=f"{flavour}: {chain}" + ) + for parts in partitionings: + np.testing.assert_array_equal( + np.asarray(view.with_parts(parts).result()), + np.asarray(expected), + err_msg=f"{flavour} parts={parts}: {chain}", + ) + + +def test_eager_getitem_returns_data(source: LazyArray) -> None: + """`arr[...]` reads immediately, like `numpy.ndarray.__getitem__`.""" + np.testing.assert_array_equal(np.asarray(source[1:3, ::2, -1]), reference()[1:3, ::2, -1]) + + +# --------------------------------------------------------------------------- +# Attribute forwarding +# --------------------------------------------------------------------------- + + +def test_forwards_array_attributes(source: LazyArray) -> None: + assert source.shape == SHAPE + assert source.ndim == len(SHAPE) + assert source.size == int(np.prod(SHAPE)) + assert source.dtype == np.dtype("int64") + assert len(source) == SHAPE[0] + assert "LazyArray" in repr(source) + + +def test_view_shape_comes_from_the_transform(source: LazyArray) -> None: + """A view reports its own shape, not the wrapped array's.""" + view = source.lazy[1:6:2, :, -1] + assert view.shape == (3, 5) + assert view.ndim == 2 + assert view.size == 15 + assert view.dtype == source.dtype + assert "view=" in repr(view) + + +def test_the_wrapper_has_no_chunks_vocabulary() -> None: + """`chunks` is the *source's* word, never the wrapper's.""" + assert not hasattr(make_source("numpy-whole"), "chunks") + assert not hasattr(make_source("zarr"), "chunks") + with pytest.raises(TypeError): + LazyArray(reference(), chunks=PART_SHAPE) # type: ignore[call-arg] + + +def test_scalar_is_basic_even_when_a_slice_separates_it_from_the_arrays() -> None: + """A documented, deliberate departure from one NumPy corner. + + NumPy counts an integer as an *advanced* index for its placement rule, so + `x[0, :, i]` has shape `(len(i), x.shape[1])` — the arrays lead because the + slice separates the integer from them. The positional dialect instead treats + every scalar as a basic index applied first, which is the rule the rest of + the surface follows and the only one `oindex` can express, so the same + selection reads as `x[0][:, i]`. + """ + data = reference() + view = make_source("numpy-uniform-parts").lazy.vindex[0, ..., np.array([3, 0])] + np.testing.assert_array_equal(np.asarray(view.result()), data[0][..., np.array([3, 0])]) + assert view.shape == data[0][..., np.array([3, 0])].shape + assert view.shape != data[0, ..., np.array([3, 0])].shape + + +def test_zero_dimensional_result_is_an_array(source: LazyArray) -> None: + """Both resolvers agree on the kind of a zero-rank result.""" + result = source.lazy[0, 1, 2].result() + assert isinstance(result, np.ndarray) + assert result.ndim == 0 + assert result[()] == reference()[0, 1, 2] + + +class ForeignArray: + """A minimal array-like whose advertised `chunks` we do not control.""" + + def __init__(self, data: np.ndarray[Any, Any], chunks: Any) -> None: + self._data = data + self.chunks = chunks + + @property + def shape(self) -> tuple[int, ...]: + return self._data.shape + + @property + def dtype(self) -> Any: + return self._data.dtype + + def __getitem__(self, key: Any) -> Any: + return self._data[key] + + +def test_malformed_discovered_parts_are_ignored() -> None: + """A foreign object's unusable `chunks` falls back to one whole-array part. + + Discovery parses external input, so an attribute that does not describe a + partitioning of the shape means "none I understand", not an error. + """ + data = reference() + for bogus in (((3, 3), (5,), (4,)), (3, 2), "nope", (3.5, 2, 2)): + wrapped = LazyArray(ForeignArray(data, bogus)) + assert len(list(wrapped.parts())) == 1 + np.testing.assert_array_equal(np.asarray(wrapped.lazy[1:3].result()), data[1:3]) + + +def test_discovered_parts_come_from_the_source_vocabulary() -> None: + """`read_chunk_sizes` wins over `chunks`, and both are read as parts.""" + data = reference() + wrapped = LazyArray(ForeignArray(data, PART_SHAPE)) + assert [part.base_coords for part in wrapped.parts()][:3] == [(0, 0, 0), (0, 0, 1), (0, 1, 0)] + assert len(list(wrapped.parts())) == 3 * 3 * 2 + + +# --------------------------------------------------------------------------- +# Boxes +# --------------------------------------------------------------------------- + +# (id, build, expected is_box, expected bounding_box). The bounds are storage +# intervals of the wrapped (7, 5, 4) array, computed by hand. +BOX_CASES: list[tuple[str, Callable[[LazyArray], LazyArray], bool, Any]] = [ + ("identity", lambda a: a, True, ((0, 7), (0, 5), (0, 4))), + ("basic-slice", lambda a: a.lazy[1:6, :, 1:3], True, ((1, 6), (0, 5), (1, 3))), + # Stride 2 over axis 1 touches 0, 2, 4; the hull is the closed span. + ("strided", lambda a: a.lazy[::3, ::2, :], True, ((0, 7), (0, 5), (0, 4))), + ("int-drop", lambda a: a.lazy[2, :, -1], True, ((2, 3), (0, 5), (3, 4))), + ("all-scalars", lambda a: a.lazy[0, 1, 2], True, ((0, 1), (1, 2), (2, 3))), + ("negative-and-open", lambda a: a.lazy[-2:], True, ((5, 7), (0, 5), (0, 4))), + ( + "oindex", + lambda a: a.lazy.oindex[[4, 0, 0], :, [3, 1]], + False, + ((0, 5), (0, 5), (1, 4)), + ), + ( + "vindex", + lambda a: a.lazy.vindex[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])], + False, + ((0, 7), (0, 5), (0, 3)), + ), + ("mask", lambda a: a.lazy.vindex[MASK], False, ((0, 7), (0, 5), (0, 4))), + # Composition preserves the category in both directions. + ("box-of-box", lambda a: a.lazy[1:6].lazy[:, 1:3], True, ((1, 6), (1, 3), (0, 4))), + ( + "box-after-fancy", + lambda a: a.lazy.oindex[[4, 0, 2], :, :].lazy[0:2], + False, + ((0, 5), (0, 5), (0, 4)), + ), + ("empty", lambda a: a.lazy[2:2], True, None), + ( + "empty-fancy", + lambda a: a.lazy.oindex[np.array([], dtype=np.intp), :, :], + False, + None, + ), +] + + +@pytest.mark.parametrize( + ("build", "expected_is_box", "expected_bounds"), + [case[1:] for case in BOX_CASES], + ids=[case[0] for case in BOX_CASES], +) +def test_is_box_and_bounding_box( + source: LazyArray, + build: Callable[[LazyArray], LazyArray], + expected_is_box: bool, + expected_bounds: Any, +) -> None: + """The box/query taxonomy, and the hull every selection has either way.""" + view = build(source) + assert view.is_box is expected_is_box + assert view.bounding_box() == expected_bounds + + +def test_a_unit_stride_box_is_dense_in_its_bounding_box() -> None: + """Stride 1 everywhere: the hull is exactly what the view selects.""" + data = reference() + view = make_source("numpy-uniform-parts").lazy[1:6, :, 1:3] + assert view.is_box + assert view.strides() == (1, 1, 1) + bounds = view.bounding_box() + assert bounds is not None + np.testing.assert_array_equal( + np.asarray(view.result()), + data[tuple(slice(lo, hi) for lo, hi in bounds)], + ) + + +def test_a_strided_box_is_sparse_in_its_bounding_box() -> None: + """Stride > 1: the hull is a superset, and `strides()` is what says by how much.""" + data = reference() + view = make_source("numpy-uniform-parts").lazy[1:6, ::2, :] + assert view.is_box + assert view.strides() == (1, 2, 1) + bounds = view.bounding_box() + assert bounds is not None + assert bounds == ((1, 6), (0, 5), (0, 4)) + + hull = data[tuple(slice(lo, hi) for lo, hi in bounds)] + assert hull.size == 5 * 5 * 4 + assert view.size == 5 * 3 * 4 + # A consumer that slabbed the hull and threw the rest away would over-read. + assert hull.size > view.size + # Applying the strides to the hull recovers the selection exactly. + np.testing.assert_array_equal( + np.asarray(view.result()), + hull[tuple(slice(None, None, step) for step in view.strides() or ())], + ) + + +def test_strides_are_none_for_a_query() -> None: + view = make_source("numpy-uniform-parts").lazy.oindex[[5, 1], :, :] + assert not view.is_box + assert view.strides() is None + + +def test_an_integer_indexed_dimension_has_stride_one() -> None: + view = make_source("numpy-uniform-parts").lazy[2, ::3, :] + assert view.strides() == (1, 3, 1) + assert view.bounding_box() == ((2, 3), (0, 4), (0, 4)) + + +def test_a_query_bounding_box_is_only_a_hull() -> None: + """For a fancy selection the box is a superset, and `is_box` says so.""" + view = make_source("numpy-uniform-parts").lazy.oindex[[5, 1], :, :] + assert not view.is_box + assert view.bounding_box() == ((1, 6), (0, 5), (0, 4)) + # The hull spans 5 rows; the selection touches 2 of them. + assert view.shape[0] == 2 + + +# --------------------------------------------------------------------------- +# Parts +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flavour", FLAVOURS) +@pytest.mark.parametrize( + "build", + [ + lambda a: a, + lambda a: a.lazy[1:6, :, 1:], + lambda a: a.lazy.oindex[[4, 0, 0], :, [3, 1]], + lambda a: a.lazy.vindex[..., np.array([1, 4, 0]), np.array([2, 0, 1])], + ], + ids=["identity", "basic", "oindex", "vindex"], +) +def test_parts_tile_the_view_exactly_and_disjointly( + flavour: str, build: Callable[[LazyArray], LazyArray] +) -> None: + """Assembling the parts reproduces `result()`; the placements cover with no overlap.""" + view = build(make_source(flavour)) + expected = np.asarray(view.result()) + + assembled = np.zeros(view.shape, dtype=view.dtype) + hits = np.zeros(view.shape, dtype=np.int64) + for part in view.parts(): + assembled[part.out_selection] = np.asarray(part.array.result()) + # `np.add.at` accumulates per element; `+= 1` on a fancy index would + # count a repeated position once and hide a real overlap. + np.add.at(hits, part.out_selection, 1) + + np.testing.assert_array_equal(assembled, expected) + np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64)) + + +def test_parts_resolve_independently_and_concurrently() -> None: + """Each part's `array` is a standalone `LazyArray` with no shared mutable state.""" + view = make_source("zarr").lazy[1:7, :, 1:].with_parts((2, 2, 2)) + parts = list(view.parts()) + assert len(parts) > 1 + + with ThreadPoolExecutor(max_workers=4) as pool: + values = list(pool.map(lambda part: np.asarray(part.array.result()), parts)) + + assembled = np.zeros(view.shape, dtype=view.dtype) + for part, value in zip(parts, values, strict=True): + assembled[part.out_selection] = value + np.testing.assert_array_equal(assembled, np.asarray(view.result())) + + +def test_parts_report_completeness() -> None: + """`is_complete` distinguishes a fully-covered box from a partial one.""" + array = make_source("numpy-uniform-parts") + assert all(part.is_complete for part in array.parts()) + # Boxes along axis 2 are [0, 3) and [3, 4); dropping column 0 leaves the + # first partially covered and the second whole. + trimmed = {part.base_coords[2]: part.is_complete for part in array.lazy[:, :, 1:].parts()} + assert trimmed == {0: False, 1: True} + # A fancy axis is always reported incomplete. + assert not any(part.is_complete for part in array.lazy.oindex[[4, 0, 0], :, :].parts()) + + +def test_partition_boxes_are_global_and_tile_the_base() -> None: + """`box` is in the wrapped array's coordinates, so parts are distinguishable. + + `array.bounding_box()` is part-local and can repeat across parts; `box` is + what says which part of the base a `Partition` actually sits in. + """ + view = make_source("numpy-uniform-parts").lazy[1:6, :, 1:] + parts = {part.base_coords: part for part in view.parts()} + + # The reviewer's repro: two parts of the same view whose part-local + # bounding boxes coincide but which cover different regions of the base. + first, second = parts[0, 0, 0], parts[0, 1, 0] + assert first.array.bounding_box() == second.array.bounding_box() + assert first.box != second.box + assert first.box == ((0, 3), (0, 2), (0, 3)) + assert second.box == ((0, 3), (2, 4), (0, 3)) + + # Every box is the base partitioning's own box for those coordinates, and + # the touched boxes tile the region the view reads without overlapping. + grids = dimension_grids_from_chunks(PART_SHAPE, SHAPE) + for coords, part in parts.items(): + expected = tuple( + (grid.chunk_offset(c), grid.chunk_offset(c) + grid.chunk_size(c)) + for grid, c in zip(grids, coords, strict=True) + ) + assert part.box == expected + covered = np.zeros(SHAPE, dtype=np.int64) + for part in parts.values(): + covered[tuple(slice(lo, hi) for lo, hi in part.box)] += 1 + assert covered.max() == 1 + # The view reads rows 1..5 and columns 1.., so every box it touches + # intersects that region and no box outside it is visited. + assert covered[1:6, :, 1:].sum() > 0 + assert covered[6:, :, :].sum() == 0 + + +def test_partition_box_of_a_whole_array_part() -> None: + part = next(iter(make_source("numpy-whole").parts())) + assert part.box == tuple((0, extent) for extent in SHAPE) + + +def test_an_unpartitioned_wrapper_has_a_single_whole_array_part() -> None: + view = make_source("numpy-whole") + parts = list(view.parts()) + assert len(parts) == 1 + assert parts[0].base_coords == (0, 0, 0) + assert parts[0].array.shape == SHAPE + assert parts[0].is_complete + np.testing.assert_array_equal(np.asarray(parts[0].array.result()), reference()) + + +def test_with_parts_keeps_the_view_and_the_base() -> None: + view = make_source("zarr").lazy[1:6, ::2] + repartitioned = view.with_parts((2, 1, 4)) + assert repartitioned.shape == view.shape + assert repartitioned.array is view.array + assert repartitioned.transform == view.transform + # A different partitioning really is a different set of boxes. + assert [part.base_coords for part in repartitioned.parts()] != [ + part.base_coords for part in view.parts() + ] + np.testing.assert_array_equal(np.asarray(repartitioned.result()), np.asarray(view.result())) + + +def test_with_parts_none_forces_one_shot_resolution() -> None: + view = make_source("zarr").lazy.oindex[[4, 0, 0], :, :] + whole = view.with_parts(None) + assert len(list(whole.parts())) == 1 + np.testing.assert_array_equal(np.asarray(whole.result()), np.asarray(view.result())) + + +@pytest.mark.parametrize( + ("parts", "match"), + [ + ((3,), "one entry per dimension"), + ((3, (2, 2), 4), "not a mixture"), + (((3, 3), (2, 2, 1), (3, 1)), "sum to 6, but the array extent is 7"), + ((3, 0, 3), "chunk shape entries must be positive"), + # A float or a None belongs to neither convention; saying "not a + # mixture" would send the reader looking for the wrong mistake. + ((3.5, 2, 2), r"3\.5 at dimension 0 is neither"), + ((None, 5, 4), "None at dimension 0 is neither"), + ((3.5, 2.5, 2.5), r"3\.5 at dimension 0, 2\.5 at dimension 1, .* are neither"), + (((1.5, 5.5), (5,), (4,)), "per-axis chunk sizes must be integers; dimension 0"), + ], +) +def test_with_parts_validates_strictly(parts: Any, match: str) -> None: + """`with_parts` is our own API, so a malformed partitioning raises.""" + with pytest.raises(ValueError, match=match): + make_source("numpy-whole").with_parts(parts) + + +# --------------------------------------------------------------------------- +# Protocols +# --------------------------------------------------------------------------- + + +def test_dask_token_is_deterministic_and_discriminating() -> None: + """Same data, same view, same parts token alike; anything else differs.""" + data = reference() + base = LazyArray(data) + assert base.__dask_tokenize__() == LazyArray(reference()).__dask_tokenize__() + + tokens = { + "base": base.__dask_tokenize__(), + "view": base.lazy[1:3].__dask_tokenize__(), + "other view": base.lazy[2:4].__dask_tokenize__(), + "parts": base.with_parts((2, 2, 2)).__dask_tokenize__(), + "other data": LazyArray(data + 1).__dask_tokenize__(), + } + assert len({repr(token) for token in tokens.values()}) == len(tokens) + # Equivalent transforms reached different ways still token alike. + assert base.lazy[1:5].lazy[0:2].__dask_tokenize__() == base.lazy[1:3].__dask_tokenize__() + + +def test_iteration_yields_eager_slices(source: LazyArray) -> None: + rows = list(source.lazy[2:5]) + assert len(rows) == 3 + for row, expected in zip(rows, reference()[2:5], strict=True): + np.testing.assert_array_equal(np.asarray(row), expected) + + +def test_iteration_over_a_zero_dimensional_view_is_rejected() -> None: + with pytest.raises(TypeError, match="iteration over a 0-d array"): + iter(make_source("numpy-uniform-parts").lazy[0, 0, 0]) + + +def test_len_of_a_zero_dimensional_view_is_rejected() -> None: + with pytest.raises(TypeError, match="len\\(\\) of unsized object"): + len(make_source("numpy-uniform-parts").lazy[0, 0, 0]) + + +@pytest.mark.parametrize( + ("convert", "selection"), + [ + (bool, (1, 1, 1)), + (int, (1, 1, 1)), + (float, (1, 1, 1)), + (operator.index, (1, 1, 1)), + (bool, (1, 1, slice(0, 1))), + ], + ids=["bool-0d", "int-0d", "float-0d", "index-0d", "bool-size-1"], +) +def test_scalar_conversions_match_numpy(convert: Any, selection: Any) -> None: + """Size-1 conversions delegate to NumPy, values and all.""" + data = reference() + view = make_source("numpy-uniform-parts").lazy[selection] + assert convert(view) == convert(data[selection]) + + +@pytest.mark.parametrize( + ("convert", "selection", "error"), + [ + (bool, (slice(0, 2), 0, 0), ValueError), + (bool, (slice(0, 0), 0, 0), ValueError), + (int, (slice(0, 1), 0, 0), TypeError), + (float, (slice(0, 2), 0, 0), TypeError), + (operator.index, (slice(0, 1), 0, 0), TypeError), + ], + ids=["bool-many", "bool-empty", "int-1d", "float-many", "index-1d"], +) +def test_scalar_conversions_raise_what_numpy_raises( + convert: Any, selection: Any, error: type[Exception] +) -> None: + data = reference() + view = make_source("numpy-uniform-parts").lazy[selection] + with pytest.raises(error): + convert(view) + with pytest.raises(error): + convert(data[selection]) + + +def test_pickle_round_trip() -> None: + """A wrapper over a picklable base survives a round trip, view and parts intact.""" + view = LazyArray(reference()).with_parts((2, 2, 2)).lazy[1:6, ::2].lazy.oindex[[3, 0, 0], :, :] + restored = pickle.loads(pickle.dumps(view)) + assert restored.shape == view.shape + assert restored.__dask_tokenize__() == view.__dask_tokenize__() + np.testing.assert_array_equal(np.asarray(restored.result()), np.asarray(view.result())) + + +# --------------------------------------------------------------------------- +# dask interop +# --------------------------------------------------------------------------- + + +def test_dask_from_array_roundtrip() -> None: + """A `LazyArray` is a drop-in dask source — no translation ceremony.""" + da = pytest.importorskip("dask.array") + source = make_source("zarr") + + lazy = da.from_array(source) + np.testing.assert_array_equal(lazy.compute(), reference()) + + # dask chooses its own blocks; the wrapper reads each of them through its + # own parts, so the two partitionings need not agree. + blocked = da.from_array(source, chunks=(4, 3, 3)) + assert blocked.chunks == ((4, 3), (3, 2), (3, 1)) + np.testing.assert_array_equal(blocked[2:, ::2].compute(), reference()[2:, ::2]) + + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +def test_boolean_scalar_is_rejected() -> None: + with pytest.raises(IndexError, match="boolean scalars are not valid indices"): + make_source("numpy-uniform-parts").lazy[True] + + +def test_mask_shape_must_match_the_view() -> None: + array = make_source("numpy-uniform-parts") + with pytest.raises(IndexError, match="boolean index has shape"): + array.lazy.vindex[np.ones((2, 2, 2), dtype=bool)] + + +def test_scalar_index_out_of_bounds() -> None: + with pytest.raises(IndexError, match="index 7 is out of bounds for axis 0 with size 7"): + make_source("numpy-uniform-parts").lazy[7] + + +def test_scalar_index_out_of_bounds_in_a_view() -> None: + """Bounds are the *view's*, not the wrapped array's.""" + array = make_source("numpy-uniform-parts").lazy[1:4] + with pytest.raises(IndexError, match="index 3 is out of bounds for axis 0 with size 3"): + array.lazy[3] + + +def test_index_array_out_of_bounds() -> None: + array = make_source("numpy-uniform-parts") + with pytest.raises(IndexError, match="index 99 is out of bounds for axis 0 with size 7"): + array.lazy.oindex[[0, 99], :, :] + + +def test_too_many_indices() -> None: + with pytest.raises(IndexError, match="too many indices"): + make_source("numpy-uniform-parts").lazy[0, 0, 0, 0] + + +def test_copy_false_conversion_is_rejected() -> None: + array = make_source("numpy-uniform-parts") + with pytest.raises(ValueError, match="cannot be converted to a NumPy array without a copy"): + np.array(array, copy=False) From 7266a45381c777d7fd7b95d7f1a8a69c620505be Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 10:57:08 +0200 Subject: [PATCH 02/32] feat(zarr-indexing): negative-step slices, per merged ndsel 1.0-draft.2 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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 --- packages/zarr-indexing/changes/267.feature.md | 3 + packages/zarr-indexing/docs/design-notes.md | 17 +- .../src/zarr_indexing/boundary.py | 22 ++- .../src/zarr_indexing/lazy_array.py | 4 +- .../src/zarr_indexing/messages.py | 23 ++- .../src/zarr_indexing/transform.py | 77 +++++--- .../tests/conformance/PROVENANCE.md | 19 +- .../tests/conformance/errors.json | 4 +- .../tests/conformance/slice.json | 95 ++++++++++ .../zarr-indexing/tests/test_lazy_array.py | 111 ++++++++++- packages/zarr-indexing/tests/test_messages.py | 42 +++++ .../tests/test_tensorstore_parity.py | 172 ++++++++++++++++++ 12 files changed, 540 insertions(+), 49 deletions(-) diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 40744cbfd9..6506332453 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -9,3 +9,6 @@ A read is broken up along a **partitioning** — a grid of boxes, discovered fro 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 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. Note that a reversing slice normally produces a negative domain origin (`[::-1]` on a length-20 axis gives `[-19, 1)`), which is the result staying anchored to the source coordinate frame; `LazyArray` re-bases every view to origin 0, so its positional dialect is unaffected. diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index 6cca05b163..48c7bbd6fb 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -149,12 +149,19 @@ speculatively; `is_box` is the runtime answer until then. ## Current scope -Three limits are known, intentional, and expected to lift: +Negative steps are supported as of ndsel 1.0-draft.2: `a[::-1]` reverses, one +desugaring rule covers both signs, and a reversed interval is an error rather +than a silently empty selection. One consequence is worth stating loudly: +**a negative step normally produces a negative domain origin.** Reversing a +length-20 zero-origin axis gives the domain `[-19, 1)`, because the result stays +anchored to the source coordinate frame and a reversing map runs that frame +backwards. `LazyArray` re-bases every view to origin 0, so the positional +dialect never shows it; a caller working with `IndexTransform` directly will, +and re-bases explicitly with `translate_domain_to` if it wants NumPy-shaped +coordinates. + +Two limits remain, both intentional and expected to lift: -- **No negative steps.** `a[::-1]` raises rather than reversing. The transform - algebra can represent a negative stride and the lowering engine resolves one, - but the slice boundary rejects it while the ndsel spec change pinning the - domain-origin rule for negative steps is in flight. *Planned.* - **Finite explicit bounds only.** `IndexDomain` has no implicit or unbounded dimensions; the message layer will normalize a body with `"-inf"`/`"+inf"` bounds, but the engine layer refuses to lower one into a transform. diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py index a2c35bf595..eb41c5a925 100644 --- a/packages/zarr-indexing/src/zarr_indexing/boundary.py +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -79,19 +79,23 @@ def _normalize_int(value: int, size: int, axis: int) -> int: def _normalize_slice(sel: slice, size: int, axis: int) -> tuple[int, int, int]: - """Resolve a positional slice to `(start, stop, step)` within `[0, size]`.""" + """Resolve a positional slice to `(start, stop, step)`, either direction. + + `slice.indices` already applies NumPy's rules — negative bounds count from + the end, out-of-range bounds clamp, and a reversed slice runs downward with + `stop` one *below* the last selected position. The one thing it does not do + is canonicalize an empty result: it can hand back a stop on the far side of + the start (`5:2` going up, `2:5` going down), which the transform layer + reads as a direction error rather than an empty selection. Collapsing it to + `stop == start` keeps NumPy's "empty, not an error" answer. + """ if sel.step is not None and not isinstance(sel.step, (int, np.integer)): raise IndexError(f"slice step must be an integer; got {sel.step!r}") step = 1 if sel.step is None else int(sel.step) - if step < 1: - raise IndexError( - f"slice step must be positive; got {step} for axis {axis}. Reversed " - "and zero-step slices are not supported." - ) + if step == 0: + raise IndexError(f"slice step cannot be zero (axis {axis})") start, stop, step = sel.indices(size) - # `slice.indices` can hand back stop < start for an empty forward slice - # (e.g. `5:2`); the transform layer wants a canonical empty interval. - stop = max(stop, start) + stop = max(stop, start) if step > 0 else min(stop, start) return start, stop, step diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index eb5afc2c86..fd57c16890 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -845,7 +845,9 @@ def strides(self) -> tuple[int, ...] | None: The magnitude only. A reversing view (`lazy[::-1]`) selects the same *set* of coordinates as its forward twin, so it reports the same bounding box and the same strides; the traversal direction lives in the - transform, not in this description of the region touched. + transform, not in this description of the region touched. A consumer + that cares about order reads the transform, or simply reverses the block + it gets back. Examples -------- diff --git a/packages/zarr-indexing/src/zarr_indexing/messages.py b/packages/zarr-indexing/src/zarr_indexing/messages.py index d5761d1a38..a255f9a2a6 100644 --- a/packages/zarr-indexing/src/zarr_indexing/messages.py +++ b/packages/zarr-indexing/src/zarr_indexing/messages.py @@ -54,6 +54,9 @@ "output_map_conflict", "rank_mismatch", "step_zero", + # Retired in 1.0-draft.2, when negative `step` became specified. Kept in + # the set so a message carrying the code is still recognized, but no + # condition in this implementation emits it. "negative_step_unsupported", } ) @@ -408,17 +411,27 @@ def _normalize_slice(obj: dict[str, Any]) -> dict[str, Any]: for k, s in enumerate(step): if s == 0: raise NdselError("step_zero", f"step[{k}] is zero") - if s < 0: - raise NdselError("negative_step_unsupported", f"step[{k}] is negative ({s})") inclusive_min: list[Any] = [] exclusive_max: list[Any] = [] output: list[dict[str, Any]] = [] for k in range(n): a, b, s = start[k], stop[k], step[k] - m = max(0, -(-(b - a) // s)) # ceil((b - a) / s) - o = _trunc_div(a, s) # trunc(a / s), toward zero - offset = a - s * o # lattice phase, in (-s, s) + # One rule for both signs (spec 5.3): the traversal runs from `a` + # toward `b`, so the source interval's length is `b - a` going up and + # `a - b` going down. + length = (b - a) if s > 0 else (a - b) + if length < 0: + # A reversed interval is a mistake about the direction of travel, + # not an empty selection. `b == a` is the way to select nothing. + raise NdselError( + "bounds_out_of_order", + f"start[{k}]={a} and stop[{k}]={b} with step {s} run the wrong " + "way; an empty selection is spelled stop == start", + ) + m = -(-length // abs(s)) # ceil(length / |s|) + o = _trunc_div(a, s) # trunc(a / s), toward zero, both signs + offset = a - s * o # lattice phase, |offset| < |s| inclusive_min.append(o) exclusive_max.append(o + m) output.append({"offset": offset, "stride": s, "input_dimension": k}) diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 1adbbf74f9..4bf96b5fe0 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -608,6 +608,20 @@ def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | return tuple(result) +def _positional_slice(pos: int, size: int, step: int) -> slice: + """A NumPy slice selecting `size` elements from `pos`, walking by `step`. + + The stop is `pos + size*step`, except that a downward walk reaching the + start of the array must stop at `None`: NumPy would read a negative stop as + counting from the end, so `slice(6, -1, -1)` selects **nothing** where + `slice(6, None, -1)` selects the first seven elements in reverse. + """ + stop = pos + size * step + if step < 0 and stop < 0: + return slice(pos, None, step) + return slice(pos, stop, step) + + def _reindex_array( m: ArrayMap, normalized: tuple[int | slice | None, ...], @@ -658,7 +672,7 @@ def _reindex_array( # indexed positionally, so shift by the domain origin. start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) pos = start - lo - idx.append(slice(pos, pos + size * step, step)) + idx.append(_positional_slice(pos, size, step)) else: # Broadcast axis: preserve the singleton (it still broadcasts # over the narrowed domain), regardless of the slice bounds. @@ -724,7 +738,7 @@ def _reindex_array_oindex( hi = domain.exclusive_max[old_dim] start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) pos = start - lo - idx.append(slice(pos, pos + size * step, step)) + idx.append(_positional_slice(pos, size, step)) else: idx.append(slice(None)) @@ -1276,39 +1290,56 @@ def _resolve_slice_ts(sel: slice, dim: int, lo: int, hi: int) -> tuple[int, int, """Resolve a slice against domain `[lo, hi)` with TensorStore semantics. Slice bounds are **literal domain coordinates** — never from-the-end, never - clamped. Rules (each verified against tensorstore 0.1.84): - - - defaults: `start = lo`, `stop = hi`; + clamped. One rule covers both signs of the step (each part verified against + tensorstore 0.1.84, and matching ndsel 1.0-draft.2 section 5.3): + + - defaults follow the direction of travel: `start = lo`, `stop = hi` going + up; `start = hi - 1`, `stop = lo - 1` going down; + - the traversal runs from `start` toward `stop`, which is excluded, so the + source interval is `[start, stop)` going up and `[stop + 1, start + 1)` + going down; - a non-empty interval must be contained in the domain (no clamping — a NumPy-style out-of-range or negative bound is an error, not a shorter or wrapped result); - - an **empty** interval (`start == stop`) is valid anywhere; - - reversed bounds (`start > stop` with positive step) are an error, not - an empty result; - - the result's domain origin is `trunc(start/step)` (rounded toward - zero) and coordinate `origin + k` maps to input `start + k*step`. + - an **empty** interval is valid anywhere, for either sign; + - an interval running the wrong way (`stop` on the far side of `start` from + the direction of travel) is an error, not an empty result; + - the result's domain origin is `trunc(start/step)` — toward zero, for both + signs — and coordinate `origin + k` maps to input `start + k*step`. + + A negative step normally produces a **negative origin**: reversing a + zero-origin axis of length 20 gives the domain `[-19, 1)`. That is the + coordinate frame staying anchored to the source; a caller that needs + non-negative coordinates re-bases explicitly (`translate_domain_to`). Returns `(start, step, origin, size)` in domain coordinates. """ step = 1 if sel.step is None else sel.step - if step <= 0: - # Negative steps are valid in TensorStore but not yet supported here; - # step 0 is invalid everywhere. - raise IndexError("slice step must be positive") - start = lo if sel.start is None else sel.start - stop = hi if sel.stop is None else sel.stop - if stop < start: + if step == 0: + raise IndexError("slice step must not be zero") + if step > 0: + start = lo if sel.start is None else sel.start + stop = hi if sel.stop is None else sel.stop + interval_lo, interval_hi = start, stop + else: + start = hi - 1 if sel.start is None else sel.start + stop = lo - 1 if sel.stop is None else sel.stop + interval_lo, interval_hi = stop + 1, start + 1 + length = interval_hi - interval_lo + if length < 0: raise IndexError( - f"slice interval [{start}, {stop}) with step {step} does not specify " - f"a valid interval for dimension {dim} (start > stop)" + f"slice from {start} to {stop} with step {step} does not specify a " + f"valid interval for dimension {dim}: the derived interval " + f"[{interval_lo}, {interval_hi}) runs the wrong way. An empty " + "selection is spelled stop == start." ) - size = -(-(stop - start) // step) # ceil((stop - start) / step) - if size > 0 and (start < lo or stop > hi): + if length > 0 and (interval_lo < lo or interval_hi > hi): hint = _LITERAL_HINT if (start < 0 or stop < 0) and lo >= 0 else "" raise BoundsCheckError( - f"slice interval [{start}, {stop}) is not contained within domain " - f"[{lo}, {hi}) for dimension {dim}{hint}" + f"slice interval [{interval_lo}, {interval_hi}) is not contained " + f"within domain [{lo}, {hi}) for dimension {dim}{hint}" ) + size = -(-length // abs(step)) # ceil(length / |step|) origin = _trunc_div(start, step) return start, step, origin, size diff --git a/packages/zarr-indexing/tests/conformance/PROVENANCE.md b/packages/zarr-indexing/tests/conformance/PROVENANCE.md index 0a6faef60b..787bcb1787 100644 --- a/packages/zarr-indexing/tests/conformance/PROVENANCE.md +++ b/packages/zarr-indexing/tests/conformance/PROVENANCE.md @@ -5,9 +5,9 @@ The JSON fixtures in this directory (`point.json`, `box.json`, `slice.json`, unmodified**, from the ndsel reference repository. - **Source:** -- **Branch:** `main` (merge of d-v-b/ndsel#1, `fix/slice-origin-trunc`) -- **Commit:** `c59bc556c` (fixtures byte-identical to the previously vendored - `c132b4c1caa3205830ce35a42502363171f650a7`) +- **Branch:** `main` (merge of d-v-b/ndsel#2, negative `step`) +- **Commit:** `92d6a32df0cd1ac47d548f14f42909a95997cf19` (previously vendored: + `c59bc556c`, itself byte-identical to `c132b4c1caa3205830ce35a42502363171f650a7`) - **Path in source:** `conformance/` **Do not edit these files.** They are vendored as-is so that @@ -18,3 +18,16 @@ corpus, re-vendor from a newer ndsel commit and update the commit SHA above. ndsel PR #1 (merged) corrected the `slice` desugaring origin from `floor(a/s)` to `trunc(a/s)` (rounding toward zero), which matches `zarr_indexing`' existing `_trunc_div` semantics. + +ndsel PR #2 (merged) specified negative `step`, which changed two fixtures: + +- `slice.json` gained the negative-step cases (full reverse, `|s| > 1` + non-divisible, negative-coordinate intervals, empty-at-any-coordinate). +- `errors.json` retired `error/negative-step` — the reason code + `negative_step_unsupported` is retired with it — and replaced it with three + `bounds_out_of_order` fixtures pinning that a reversed interval is an error + for either sign of the step, rather than being clamped to empty. + +Re-vendoring those two files and teaching `zarr_indexing.messages` the new +desugaring are one change: the corpus is the definition of correct here, so it +lands in the same commit as the code that satisfies it. diff --git a/packages/zarr-indexing/tests/conformance/errors.json b/packages/zarr-indexing/tests/conformance/errors.json index e5e08bab3c..072acd0af5 100644 --- a/packages/zarr-indexing/tests/conformance/errors.json +++ b/packages/zarr-indexing/tests/conformance/errors.json @@ -1,6 +1,8 @@ [ { "name": "error/step-zero", "input": { "kind": "slice", "start": [0], "stop": [4], "step": [0] }, "error": "step_zero" }, - { "name": "error/negative-step", "input": { "kind": "slice", "start": [9], "stop": [0], "step": [-2] }, "error": "negative_step_unsupported" }, + { "name": "error/slice-reversed-interval-unit-step", "input": { "kind": "slice", "start": [9], "stop": [0] }, "error": "bounds_out_of_order" }, + { "name": "error/slice-reversed-interval-positive-step", "input": { "kind": "slice", "start": [9], "stop": [0], "step": [2] }, "error": "bounds_out_of_order" }, + { "name": "error/slice-reversed-interval-negative-step", "input": { "kind": "slice", "start": [5], "stop": [6], "step": [-1] }, "error": "bounds_out_of_order" }, { "name": "error/multiple-upper-bounds", "input": { "kind": "box", "shape": [3], "exclusive_max": [3] }, "error": "multiple_upper_bounds" }, { "name": "error/rank-mismatch", "input": { "kind": "slice", "start": [0, 0], "stop": [4] }, "error": "rank_mismatch" }, { "name": "error/unknown-kind", "input": { "kind": "bogus" }, "error": "unknown_kind" }, diff --git a/packages/zarr-indexing/tests/conformance/slice.json b/packages/zarr-indexing/tests/conformance/slice.json index 2f1a0694ce..959ebc6247 100644 --- a/packages/zarr-indexing/tests/conformance/slice.json +++ b/packages/zarr-indexing/tests/conformance/slice.json @@ -57,5 +57,100 @@ "input_labels": [""], "output": [ { "offset": -2, "stride": 3, "input_dimension": 0 } ] } + }, + { + "name": "slice/negative-step-full-reverse", + "input": { "kind": "slice", "start": [19], "stop": [-1], "step": [-1] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-19], "input_exclusive_max": [1], + "input_labels": [""], + "output": [ { "offset": 0, "stride": -1, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-divisible-span", + "input": { "kind": "slice", "start": [15], "stop": [5], "step": [-2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-7], "input_exclusive_max": [-2], + "input_labels": [""], + "output": [ { "offset": 1, "stride": -2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-nondivisible-span", + "input": { "kind": "slice", "start": [15], "stop": [5], "step": [-4] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-3], "input_exclusive_max": [0], + "input_labels": [""], + "output": [ { "offset": 3, "stride": -4, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-down-to-zero", + "input": { "kind": "slice", "start": [9], "stop": [0], "step": [-2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-4], "input_exclusive_max": [1], + "input_labels": [""], + "output": [ { "offset": 1, "stride": -2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-negative-interval", + "input": { "kind": "slice", "start": [-1], "stop": [-6], "step": [-2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ { "offset": -1, "stride": -2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-negative-interval-step3", + "input": { "kind": "slice", "start": [-2], "stop": [-9], "step": [-3] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ { "offset": -2, "stride": -3, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-single-point", + "input": { "kind": "slice", "start": [5], "stop": [4], "step": [-3] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-1], "input_exclusive_max": [0], + "input_labels": [""], + "output": [ { "offset": 2, "stride": -3, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-empty", + "input": { "kind": "slice", "start": [5], "stop": [5], "step": [-1] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-5], "input_exclusive_max": [-5], + "input_labels": [""], + "output": [ { "offset": 0, "stride": -1, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-step-keeps-labels", + "input": { "kind": "slice", "start": [19], "stop": [-1], "step": [-1], "labels": ["x"] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-19], "input_exclusive_max": [1], + "input_labels": ["x"], + "output": [ { "offset": 0, "stride": -1, "input_dimension": 0 } ] + } + }, + { + "name": "slice/2d-mixed-sign-step", + "input": { "kind": "slice", "start": [19, 0], "stop": [-1, 10], "step": [-1, 2] }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [-19, 0], + "input_exclusive_max": [1, 5], + "input_labels": ["", ""], + "output": [ + { "offset": 0, "stride": -1, "input_dimension": 0 }, + { "offset": 0, "stride": 2, "input_dimension": 1 } + ] + } } ] diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 0a42fa1aec..9ac5daf3c5 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -396,6 +396,51 @@ def source(request: pytest.FixtureRequest) -> LazyArray: lambda a: a.lazy.oindex[[3, 1, 0], 3:5, [2, 0, 2]].lazy[0], lambda r: outer(r, ([3, 1, 0], slice(3, 5), [2, 0, 2]))[0], ), + # Negative steps: the positional dialect is NumPy's, including the empty + # cases NumPy allows where the transform algebra alone would object. + ("reverse", lambda a: a.lazy[::-1], lambda r: r[::-1]), + ("reverse-every-axis", lambda a: a.lazy[::-1, ::-1, ::-1], lambda r: r[::-1, ::-1, ::-1]), + ("reverse-strided", lambda a: a.lazy[::-2], lambda r: r[::-2]), + ("reverse-nondivisible", lambda a: a.lazy[::-3], lambda r: r[::-3]), + ("reverse-bounded", lambda a: a.lazy[5:1:-1], lambda r: r[5:1:-1]), + ("reverse-negative-start", lambda a: a.lazy[-1:None:-1], lambda r: r[-1:None:-1]), + ("reverse-past-the-start", lambda a: a.lazy[:-8:-1], lambda r: r[:-8:-1]), + ("reverse-empty", lambda a: a.lazy[2:2:-1], lambda r: r[2:2:-1]), + # NumPy reads a reversed *positional* interval as empty; only the literal + # layer calls it a direction error. + ("reverse-inverted-is-empty", lambda a: a.lazy[2:5:-1], lambda r: r[2:5:-1]), + ("reverse-with-int-drop", lambda a: a.lazy[::-2, 2, ::-1], lambda r: r[::-2, 2, ::-1]), + ("reverse-trailing-axis", lambda a: a.lazy[..., ::-1], lambda r: r[..., ::-1]), + ( + "compose-reverse-then-reverse", + lambda a: a.lazy[::-1].lazy[::-1], + lambda r: r[::-1][::-1], + ), + ( + "compose-strided-then-reverse", + lambda a: a.lazy[::2].lazy[::-1], + lambda r: r[::2][::-1], + ), + ( + "compose-reverse-then-strided", + lambda a: a.lazy[::-1].lazy[::2], + lambda r: r[::-1][::2], + ), + ( + "compose-reverse-then-oindex", + lambda a: a.lazy[::-1].lazy.oindex[[3, 0, 0], :, :], + lambda r: outer(r[::-1], ([3, 0, 0], slice(None), slice(None))), + ), + ( + "compose-oindex-then-reverse", + lambda a: a.lazy.oindex[[3, 1, 2], :, :].lazy[::-1], + lambda r: outer(r, ([3, 1, 2], slice(None), slice(None)))[::-1], + ), + ( + "compose-reverse-then-vindex", + lambda a: a.lazy[::-1].lazy.vindex[..., np.array([1, 3, 0])], + lambda r: r[::-1][..., np.array([1, 3, 0])], + ), ( "compose-vindex-trailing-then-scalar-and-slice", lambda a: a.lazy.vindex[..., np.array([3, 0, 1])].lazy[-1, 1:4], @@ -432,10 +477,17 @@ def _random_basic(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any roll = rng.random() if roll < 0.3: selection.append(int(rng.integers(-size, size))) - elif roll < 0.7: + elif roll < 0.55: start = int(rng.integers(0, size)) stop = int(rng.integers(start, size + 1)) selection.append(slice(start, stop, int(rng.integers(1, 4)))) + elif roll < 0.8: + # Downward: `start >= stop` and the stop may fall off the front, + # which is spelled `None`. + start = int(rng.integers(0, size)) + stop_choice = int(rng.integers(-1, start + 1)) + stop = None if stop_choice < 0 else stop_choice + selection.append(slice(start, stop, -int(rng.integers(1, 4)))) else: selection.append(slice(None)) return tuple(selection) @@ -795,8 +847,13 @@ def test_a_query_bounding_box_is_only_a_hull() -> None: lambda a: a.lazy[1:6, :, 1:], lambda a: a.lazy.oindex[[4, 0, 0], :, [3, 1]], lambda a: a.lazy.vindex[..., np.array([1, 4, 0]), np.array([2, 0, 1])], + # A reversing view drives the negative-stride branches of + # `_intersect_dimension_map` and `iter_chunk_transforms`, which were + # written defensively long before anything could reach them. + lambda a: a.lazy[::-1, ::-2, :], + lambda a: a.lazy[5:1:-1, :, ::-1], ], - ids=["identity", "basic", "oindex", "vindex"], + ids=["identity", "basic", "oindex", "vindex", "reversed", "reversed-bounded"], ) def test_parts_tile_the_view_exactly_and_disjointly( flavour: str, build: Callable[[LazyArray], LazyArray] @@ -1088,3 +1145,53 @@ def test_copy_false_conversion_is_rejected() -> None: array = make_source("numpy-uniform-parts") with pytest.raises(ValueError, match="cannot be converted to a NumPy array without a copy"): np.array(array, copy=False) + + +# --------------------------------------------------------------------------- +# Negative steps +# --------------------------------------------------------------------------- + + +def test_reversed_box_reports_a_positive_stride(source: LazyArray) -> None: + """A reversal is still a box; `strides()` is magnitudes, so it matches the forward twin.""" + reversed_view = source.lazy[::-2] + forward = source.lazy[::2] + assert reversed_view.is_box + assert reversed_view.strides() == forward.strides() == (2, 1, 1) + assert reversed_view.bounding_box() == ((0, 7), (0, 5), (0, 4)) + + +def test_a_reversed_view_is_re_based_to_origin_zero() -> None: + """The literal domain of a reversal is negative; the positional dialect hides it.""" + view = make_source("numpy-whole").lazy[::-1] + # The algebra's own answer keeps the source frame. + assert IndexTransform.from_shape(SHAPE)[::-1].domain.inclusive_min[0] == -6 + # The wrapper re-bases, so positions start at 0 as NumPy expects. + assert view.transform.domain.inclusive_min == (0, 0, 0) + assert view.shape == SHAPE + np.testing.assert_array_equal(np.asarray(view.result()), reference()[::-1]) + + +def test_zero_step_is_rejected() -> None: + with pytest.raises(IndexError, match="step cannot be zero"): + make_source("numpy-whole").lazy[::0] + + +def test_reversed_positional_interval_is_empty_not_an_error() -> None: + """NumPy's rule at the boundary; the literal layer keeps TensorStore's.""" + view = make_source("numpy-uniform-parts").lazy[2:5:-1] + assert view.shape == (0, 5, 4) + np.testing.assert_array_equal(np.asarray(view.result()), reference()[2:5:-1]) + # Literal coordinates, on the other hand, call it a direction error. + with pytest.raises(IndexError, match="valid interval"): + IndexTransform.from_shape(SHAPE)[2:5:-1] + + +def test_negative_step_over_a_fancy_axis_reverses_the_coordinates(source: LazyArray) -> None: + """Reversing a gathered axis materializes, rather than attaching a stride.""" + view = source.lazy.oindex[[3, 1, 2], :, :].lazy[::-1] + expected = outer(reference(), ([3, 1, 2], slice(None), slice(None)))[::-1] + np.testing.assert_array_equal(np.asarray(view.result()), expected) + m = view.transform.output[0] + assert isinstance(m, ArrayMap) + np.testing.assert_array_equal(m.index_array.reshape(-1), np.array([2, 1, 3])) diff --git a/packages/zarr-indexing/tests/test_messages.py b/packages/zarr-indexing/tests/test_messages.py index 14bed66448..2ebd41a263 100644 --- a/packages/zarr-indexing/tests/test_messages.py +++ b/packages/zarr-indexing/tests/test_messages.py @@ -89,3 +89,45 @@ def test_empty_string_kind_is_unknown_kind() -> None: with pytest.raises(NdselError) as excinfo: normalize_ndsel({"kind": ""}) assert excinfo.value.reason == "unknown_kind" + + +class TestNegativeStep: + """ndsel 1.0-draft.2 section 5.3: one desugaring rule, both signs.""" + + def test_full_reverse(self) -> None: + """The spec's own worked example: reversing a length-20 axis.""" + body = normalize_ndsel({"kind": "slice", "start": [19], "stop": [-1], "step": [-1]}) + assert body["input_inclusive_min"] == [-19] + assert body["input_exclusive_max"] == [1] + assert body["output"] == [{"offset": 0, "stride": -1, "input_dimension": 0}] + + def test_trunc_origin_for_a_negative_step(self) -> None: + """`trunc(15 / -2) == -7`; `floor` would give -8.""" + body = normalize_ndsel({"kind": "slice", "start": [15], "stop": [5], "step": [-2]}) + assert body["input_inclusive_min"] == [-7] + assert body["input_exclusive_max"] == [-2] + assert body["output"] == [{"offset": 1, "stride": -2, "input_dimension": 0}] + + def test_empty_is_legal_at_any_coordinate(self) -> None: + body = normalize_ndsel({"kind": "slice", "start": [5], "stop": [5], "step": [-1]}) + assert body["input_inclusive_min"] == body["input_exclusive_max"] == [-5] + + @pytest.mark.parametrize( + "message", + [ + {"kind": "slice", "start": [9], "stop": [0]}, + {"kind": "slice", "start": [9], "stop": [0], "step": [2]}, + {"kind": "slice", "start": [5], "stop": [6], "step": [-1]}, + ], + ids=["unit-step", "positive-step", "negative-step"], + ) + def test_a_reversed_interval_is_an_error(self, message: dict[str, object]) -> None: + """Not clamped to empty: travelling the wrong way is a mistake, either sign.""" + with pytest.raises(NdselError) as excinfo: + normalize_ndsel(message) + assert excinfo.value.reason == "bounds_out_of_order" + + def test_zero_step_still_errors(self) -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": "slice", "start": [0], "stop": [4], "step": [0]}) + assert excinfo.value.reason == "step_zero" diff --git a/packages/zarr-indexing/tests/test_tensorstore_parity.py b/packages/zarr-indexing/tests/test_tensorstore_parity.py index 1ed99046e9..9f4123fc02 100644 --- a/packages/zarr-indexing/tests/test_tensorstore_parity.py +++ b/packages/zarr-indexing/tests/test_tensorstore_parity.py @@ -208,6 +208,178 @@ def test_reversed_bounds_raise(self, sel: slice) -> None: _a()[sel] +class TestNegativeStep: + """Section 1 of the negative-step study: one desugaring rule, both signs. + + Every expectation below is TensorStore 0.1.84's recorded output (study + sections 1.3-1.5), which an exhaustive 32,980-case sweep found zero + disagreements with. + """ + + # (domain, slice, expected domain, expected offset, expected stride, cells) + RECORDED: ClassVar[list[tuple[tuple[int, int], slice, tuple[int, int], int, int]]] = [ + ((0, 20), slice(15, 5, -1), (-15, -5), 0, -1), + ((0, 20), slice(15, 5, -2), (-7, -2), 1, -2), + ((0, 20), slice(15, 4, -2), (-7, -1), 1, -2), + ((0, 20), slice(None, None, -1), (-19, 1), 0, -1), + ((0, 20), slice(None, None, -2), (-9, 1), 1, -2), + ((0, 20), slice(5, None, -1), (-5, 1), 0, -1), + ((0, 20), slice(None, 5, -1), (-19, -5), 0, -1), + ((0, 20), slice(5, 5, -1), (-5, -5), 0, -1), + ((0, 20), slice(5, 4, -1), (-5, -4), 0, -1), + ((0, 20), slice(5, 4, -3), (-1, 0), 2, -3), + ((0, 20), slice(15, 5, -4), (-3, 0), 3, -4), + ((0, 20), slice(15, 5, -7), (-2, 0), 1, -7), + ((5, 25), slice(None, None, -2), (-12, -2), 0, -2), + ((-10, 10), slice(-1, -6, -2), (0, 3), -1, -2), + ((-10, 10), slice(-2, -9, -3), (0, 3), -2, -3), + ] + + @pytest.mark.parametrize( + ("domain", "sel", "expected_domain", "offset", "stride"), + RECORDED, + ids=[f"{d}{s}" for d, s, _, _, _ in RECORDED], + ) + def test_recorded_desugaring( + self, + domain: tuple[int, int], + sel: slice, + expected_domain: tuple[int, int], + offset: int, + stride: int, + ) -> None: + t = _identity(*domain)[sel] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == expected_domain + m = _dim(t) + assert (m.offset, m.stride) == (offset, stride) + + @pytest.mark.parametrize( + ("domain", "sel", "cells"), + [ + ((0, 20), slice(15, 5, -1), list(range(15, 5, -1))), + ((0, 20), slice(15, 5, -2), [15, 13, 11, 9, 7]), + ((0, 20), slice(None, None, -1), list(range(19, -1, -1))), + ((0, 20), slice(15, 5, -7), [15, 8]), + ((5, 25), slice(None, None, -2), list(range(24, 4, -2))), + ((-10, 10), slice(-1, -6, -2), [-1, -3, -5]), + ((-10, 10), slice(-2, -9, -3), [-2, -5, -8]), + ], + ) + def test_recorded_cells(self, domain: tuple[int, int], sel: slice, cells: list[int]) -> None: + assert _base_cells(_identity(*domain)[sel]) == cells + + def test_trunc_not_floor_or_ceil(self) -> None: + """The three rows of study section 1.2 that discriminate the rounding.""" + # floor would give -8 here, trunc gives -7. + assert _identity(0, 20)[15:5:-2].domain.inclusive_min[0] == -7 + # ceil would give 1 for both of these; trunc gives 0. + assert _identity(-10, 10)[-1:-6:-2].domain.inclusive_min[0] == 0 + assert _identity(-10, 10)[-2:-9:-3].domain.inclusive_min[0] == 0 + + @pytest.mark.parametrize("sel", [slice(5, 15, -1), slice(5, 6, -1)]) + def test_inverted_interval_raises(self, sel: slice) -> None: + with pytest.raises(IndexError, match="valid.*interval"): + _identity(0, 20)[sel] + + @pytest.mark.parametrize("sel", [slice(20, 0, -1), slice(20, 19, -1), slice(15, -5, -1)]) + def test_uncontained_interval_raises(self, sel: slice) -> None: + with pytest.raises(BoundsCheckError, match="not contained"): + _identity(0, 20)[sel] + + def test_zero_step_raises(self) -> None: + with pytest.raises(IndexError, match="step must not be zero"): + _identity(0, 20)[15:5:0] + + def test_empty_interval_legal_outside_the_domain(self) -> None: + t = _identity(0, 20)[25:25:-1] + assert t.domain.shape == (0,) + assert t.domain.inclusive_min[0] == -25 + + @pytest.mark.parametrize( + ("domain", "first", "second", "expected_domain", "offset", "stride", "cells"), + [ + # Study section 1.5, recorded verbatim. Each row applies `second` + # to the view `first` produced — strides multiply, and a double + # reverse recovers the identity. + ( + (0, 20), + slice(0, 20, 2), + slice(None, None, -1), + (-9, 1), + 0, + -2, + list(range(18, -2, -2)), + ), + ((0, 20), slice(0, 20, 2), slice(None, None, -2), (-4, 1), 2, -4, [18, 14, 10, 6, 2]), + ( + (0, 20), + slice(None, None, -1), + slice(None, None, -1), + (0, 20), + 0, + 1, + list(range(20)), + ), + ( + (0, 20), + slice(None, None, -1), + slice(None, None, 2), + (-9, 1), + 1, + -2, + list(range(19, -1, -2)), + ), + ( + (-10, 10), + slice(None, None, -1), + slice(None, None, -3), + (-3, 4), + -1, + 3, + [-10, -7, -4, -1, 2, 5, 8], + ), + ], + ids=[ + "strided-then-reverse", + "strided-then-reverse-by-two", + "double-reverse-is-identity", + "reverse-then-strided", + "reverse-then-strided-on-negative-origin", + ], + ) + def test_recorded_composition( + self, + domain: tuple[int, int], + first: slice, + second: slice, + expected_domain: tuple[int, int], + offset: int, + stride: int, + cells: list[int], + ) -> None: + t = _identity(*domain)[first][second] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == expected_domain + m = _dim(t) + assert (m.offset, m.stride) == (offset, stride) + assert _base_cells(t) == cells + + def test_negative_step_over_an_index_array_reverses_it(self) -> None: + """Study section 1.5: a reversing step materializes, it does not stride. + + Recorded: `oindex[[3, 1, 2]]` then `[2:-1:-1]` gives index array + `[[2], [1], [3]]` over domain `[-2, 1)`. + """ + base = IndexTransform.identity(IndexDomain(inclusive_min=(0, 0), exclusive_max=(4, 5))) + gathered = base.oindex[np.array([3, 1, 2]), slice(None)] + reversed_view = gathered[2:-1:-1, :] + + assert reversed_view.domain.inclusive_min[0] == -2 + assert reversed_view.domain.exclusive_max[0] == 1 + m = reversed_view.output[0] + assert isinstance(m, ArrayMap) + np.testing.assert_array_equal(m.index_array.reshape(-1), np.array([2, 1, 3])) + + class TestTranslate: """Oracle sections 4 and 12: translate_by / translate_to preserve the cell mapping.""" From 02e4c6fed969f3bfdebe5f3e615d2f620f86873e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 11:12:01 +0200 Subject: [PATCH 03/32] =?UTF-8?q?polish(zarr-indexing):=20re-review=20mino?= =?UTF-8?q?rs=20=E2=80=94=20step-zero=20ValueError,=20kw-only=20Partition,?= =?UTF-8?q?=20strides=20docs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- packages/zarr-indexing/src/zarr_indexing/boundary.py | 2 +- packages/zarr-indexing/src/zarr_indexing/lazy_array.py | 8 ++++++-- packages/zarr-indexing/tests/test_lazy_array.py | 2 +- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py index eb41c5a925..594d66b5ba 100644 --- a/packages/zarr-indexing/src/zarr_indexing/boundary.py +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -93,7 +93,7 @@ def _normalize_slice(sel: slice, size: int, axis: int) -> tuple[int, int, int]: raise IndexError(f"slice step must be an integer; got {sel.step!r}") step = 1 if sel.step is None else int(sel.step) if step == 0: - raise IndexError(f"slice step cannot be zero (axis {axis})") + raise ValueError(f"slice step cannot be zero (axis {axis})") # ValueError: NumPy parity start, stop, step = sel.indices(size) stop = max(stop, start) if step > 0 else min(stop, start) return start, stop, step diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index fd57c16890..c33b13b41c 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -538,7 +538,7 @@ def _covers_whole_part(local_transform: IndexTransform, part_shape: tuple[int, . return True -@dataclass(frozen=True) +@dataclass(frozen=True, kw_only=True) class Partition: """One box of a `LazyArray`'s partitioning, as it falls through the view. @@ -838,7 +838,11 @@ def strides(self) -> tuple[int, ...] | None: tuple of int, or None One positive stride per storage dimension, or `None` when [`is_box`][zarr_indexing.lazy_array.LazyArray.is_box] is false: a - query's coordinates are a lookup table and have no step. + query's coordinates are a lookup table and have no step. An + **empty** box still reports its strides even though + [`bounding_box`][zarr_indexing.lazy_array.LazyArray.bounding_box] + returns `None` — the step is a property of the selection's shape, + not of the (empty) region it touches. Notes ----- diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 9ac5daf3c5..cda0b14513 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -1173,7 +1173,7 @@ def test_a_reversed_view_is_re_based_to_origin_zero() -> None: def test_zero_step_is_rejected() -> None: - with pytest.raises(IndexError, match="step cannot be zero"): + with pytest.raises(ValueError, match="step cannot be zero"): make_source("numpy-whole").lazy[::0] From 7bc3ad110098a3b3703a863cc1e663dc6cbee896 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 12:14:51 +0200 Subject: [PATCH 04/32] docs(zarr-indexing): plain technical language throughout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- packages/zarr-indexing/changes/267.feature.md | 10 +- packages/zarr-indexing/docs/api/index.md | 4 +- packages/zarr-indexing/docs/design-notes.md | 98 ++++----- packages/zarr-indexing/docs/index.md | 144 ++++++------- packages/zarr-indexing/docs/ndsel.md | 36 ++-- .../src/zarr_indexing/boundary.py | 20 +- .../src/zarr_indexing/lazy_array.py | 190 +++++++++--------- .../src/zarr_indexing/transform.py | 20 +- 8 files changed, 262 insertions(+), 260 deletions(-) diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 6506332453..99a2a3f098 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -1,14 +1,14 @@ -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`; every other NumPy operation materializes the view through `__array__`. `__dask_tokenize__`, `__len__`, `__iter__`, `__bool__`, `__int__`, `__float__`, and `__index__` round out the surface, and a wrapper pickles as long as its base does. +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`; every other NumPy operation materializes the view through `__array__`. `__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` (the source's vocabulary; the wrapper's own surface speaks only of 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 and nothing else, 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. +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 either way — dense only when every stride is 1, a lattice within the hull otherwise, a mere superset for a query, and `None` for an empty selection — and `LazyArray.strides()` gives the per-dimension step that completes the description of a box. `Partition` gained `box`, the part's box in the wrapped array's global coordinates, since `Partition.array.bounding_box()` is part-local and cannot tell two parts apart. 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. +`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.array.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) reports "no axis" as `None` instead of falling back to that stale binding. +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. Note that a reversing slice normally produces a negative domain origin (`[::-1]` on a length-20 axis gives `[-19, 1)`), which is the result staying anchored to the source coordinate frame; `LazyArray` re-bases every view to origin 0, so its positional dialect is unaffected. +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. diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index 09b6bb6b32..17cb401863 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -36,8 +36,8 @@ and the wire format built on top of it. - [`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, plus `Partition` and the `parts()` / - `with_parts()` pair that decide the boxes a read is broken into + TensorStore-style deferred indexing, plus `Partition` and `parts()` / + `with_parts()`, which determine the boxes a read is broken into - [`zarr_indexing.boundary`](boundary.md) — the translation between NumPy's positional dialect and the transform algebra's literal coordinates diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index 48c7bbd6fb..6977aca76b 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -4,15 +4,15 @@ title: Design notes # Design notes -Three things that are easier to explain once than to infer from the API: where -this library sits relative to TensorStore, why rectangular selections are a -category rather than a fast path, and what is deliberately not implemented yet. +Three topics the API does not state directly: how this library relates to +TensorStore, why rectangular selections are a category rather than a fast path, +and what is deliberately not implemented yet. ## Relationship to TensorStore The core is [TensorStore's](https://google.github.io/tensorstore/index_space.html) -index-transform model, reimplemented in Python against NumPy. The parts that -are the same are the same on purpose: +index-transform model, reimplemented in Python against NumPy. The following +parts are intentionally identical: - **The model.** An `IndexTransform` pairs an input `IndexDomain` — a rectangular region with an explicit, possibly non-zero origin — with one @@ -30,36 +30,36 @@ are the same are the same on purpose: `tensorstore.IndexTransform(json=...)` and round-trips them back through our engine layer. -Index arrays are the one place the representations differ in practice. Ours are -normalized to the transform's full input rank — full-sized on the axis a map -varies over, singleton elsewhere — so the orthogonal/vectorized distinction is -derivable from the shape. TensorStore permits a lower-rank array that broadcasts +The representations differ in one place: index arrays. Ours are normalized to +the transform's full input rank — full-sized on the axis a map varies over, +singleton elsewhere — so the orthogonal/vectorized distinction is derivable +from the shape. TensorStore permits a lower-rank array that broadcasts against the input domain; we accept those on load and normalize them. `ArrayMap` also records the input dimension an *orthogonal* (`oindex`) array varies over, a field TensorStore's format has no slot for, so [the serializer](api/json.md) collapses it on the way out and reconstructs it on the way in. -Four deliberate inversions: +Four deliberate differences: | | TensorStore | `zarr-indexing` | | --- | --- | --- | | Dialect | One strict dialect everywhere: literal coordinates, no negative wrapping | The algebra keeps that dialect; each public boundary picks its own. [`LazyArray`](api/lazy_array.md) speaks positional NumPy, `zarr.Array.lazy` speaks literal. [`zarr_indexing.boundary`](api/boundary.md) is the translation | -| Scheduling | An internal C++ scheduler owns concurrency and chunk ordering | [`parts()`](api/lazy_array.md) hands the partition structure to whatever scheduler the caller already has — dask, a thread pool, a task queue | +| Scheduling | An internal C++ scheduler owns concurrency and chunk ordering | [`parts()`](api/lazy_array.md) exposes the partition structure so the caller's own scheduler — dask, a thread pool, a task queue — drives it | | Wire format | Implementation-defined JSON, specified by what the implementation accepts | [ndsel](ndsel.md) is spec-first, with a vendored language-agnostic conformance corpus every implementation runs | | Backends | A driver ecosystem (zarr, N5, neuroglancer, GCS, …) built into the library | No drivers. The wrapper takes anything with `shape`, `dtype`, and `__getitem__`, and prefers the array's own `__array_namespace__` | -The asymmetries run the other way too, and are worth stating plainly. -TensorStore is a mature, heavily optimized C++ system with a performance -ceiling this cannot approach: our resolution is Python-level bookkeeping over -NumPy, and the per-part overhead is real. What this library has instead is -small size and no dependency beyond NumPy, which is what makes the algebra -adoptable by a Python project that wants the model without the runtime. +The comparison also runs the other way. TensorStore is a mature, heavily +optimized C++ system whose performance this library cannot approach: resolution +here is Python-level bookkeeping over NumPy, and the per-part overhead is +significant. This library is small and depends on nothing beyond NumPy, so the +algebra can be adopted by a Python project that wants the model without the C++ +runtime. ## Bounding-box selections vs query selections Every selection this library can express falls into exactly one of two -categories, and the boundary between them is structural, not a heuristic: +categories. The boundary between them is structural, not a heuristic: **A box** is a transform whose output maps are all `ConstantMap` or `DimensionMap` — no `ArrayMap`. Such a map is affine and monotone: storage @@ -73,8 +73,8 @@ one, and composing basic indexing with basic indexing keeps one. table of coordinates. It costs `O(n)` to store, it has no locality (the coordinates may repeat, reverse, or scatter arbitrarily), and intersecting it with a region means scanning it. `oindex`, `vindex`, and boolean masks all -produce one, and once an axis is a query no amount of subsequent basic -indexing makes it a box again. +produce one, and once an axis is a query, subsequent basic indexing cannot make +it a box again. [ndsel](ndsel.md) encodes the same split in its message kinds: `point`, `box`, and `slice` desugar to constant and affine output maps and are always boxes; @@ -96,12 +96,12 @@ transform_to_canonical(gather)["output"][0] # 'index_array_bounds': ['-inf', '+inf']} ``` -The distinction matters to anyone downstream of a selection. A box can be -tiled into rectangular dask chunks or handed to a viewer or tile server that -only understands rectangles; a query cannot, and has to be resolved into a -gather. A box is also servable as a single strided slab read — but only a -strided one: reading its bounding box and discarding the rest is a -proportionally larger transfer as soon as any stride exceeds 1. The two also behave differently under +The distinction matters to consumers of a selection. A box can be tiled into +rectangular dask chunks or passed to a viewer or tile server that only accepts +rectangles; a query cannot, and has to be resolved into a gather. A box can also +be served as a single strided slab read, but the read has to be strided: reading +its bounding box and discarding the rest transfers proportionally more data as +soon as any stride exceeds 1. The two also behave differently under partitioning: a box touches a contiguous run of parts, while a query can touch any subset of them, in any order, more than once. @@ -130,37 +130,37 @@ gather.strides() # None gather.shape # (3, 80) ``` -`bounding_box()` is defined for both — it is the hull, the smallest interval per -storage dimension containing everything the selection reaches. `strides()` is -defined only for a box, and is the other half of the description: the hull says -*where*, the strides say *how densely*. +`bounding_box()` is defined for both: it is the hull, the smallest interval per +storage dimension containing every coordinate the selection reaches. +`strides()` is defined only for a box and gives the step per dimension. +Together the two describe a box selection completely. -Both halves matter, because a box is dense in its hull only when every stride is -1. The slab above spans a 40x77 hull over the 40x20 cells it actually selects — -a consumer that issued one rectangular read of the hull and discarded the rest -would move 3.85x the data. A query's hull is looser still and carries no stride -at all: 88 rows of hull over three selected rows. An empty selection returns -`None` from both, having no coordinate to report an interval around. +Both are needed, because a box is dense in its hull only when every stride is +1. The slab above spans a 40x77 hull over the 40x20 cells it selects, so a +consumer that issued one rectangular read of the hull and discarded the rest +would transfer 3.85x the data. A query's hull is looser still and carries no +stride at all: 88 rows of hull over three selected rows. An empty selection +returns `None` from both, because it touches no coordinate to report an interval +around. There is deliberately no separate `BoxView` type today. A statically-typed -rectangular-only view is the obvious next step, but it should be introduced by -a consumer that actually needs the guarantee in its signatures rather than -speculatively; `is_box` is the runtime answer until then. +rectangular-only view is a plausible next step, but it should be introduced by +a consumer that needs the guarantee in its signatures rather than +speculatively; `is_box` is the runtime check until then. ## Current scope Negative steps are supported as of ndsel 1.0-draft.2: `a[::-1]` reverses, one desugaring rule covers both signs, and a reversed interval is an error rather -than a silently empty selection. One consequence is worth stating loudly: -**a negative step normally produces a negative domain origin.** Reversing a -length-20 zero-origin axis gives the domain `[-19, 1)`, because the result stays -anchored to the source coordinate frame and a reversing map runs that frame -backwards. `LazyArray` re-bases every view to origin 0, so the positional -dialect never shows it; a caller working with `IndexTransform` directly will, -and re-bases explicitly with `translate_domain_to` if it wants NumPy-shaped -coordinates. - -Two limits remain, both intentional and expected to lift: +than a silently empty selection. One consequence: a negative step normally +produces a negative domain origin. Reversing a length-20 zero-origin axis gives +the domain `[-19, 1)`, because the result stays anchored to the source +coordinate frame and a reversing map traverses that frame backwards. `LazyArray` +re-bases every view to origin 0, so the positional dialect never exposes it; a +caller working with `IndexTransform` directly will see it, and re-bases +explicitly with `translate_domain_to` for NumPy-shaped coordinates. + +Two limits remain, both intentional and expected to be lifted: - **Finite explicit bounds only.** `IndexDomain` has no implicit or unbounded dimensions; the message layer will normalize a body with `"-inf"`/`"+inf"` diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index af297286c1..c8f5bcba42 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -12,14 +12,13 @@ pip install zarr-indexing ## What this is -An indexing operation — a slice, an integer, a fancy index array — is a -*mapping* from the coordinates a user asks for to the coordinates that live in -storage. This library makes that mapping a first-class value: an -[`IndexTransform`](api/transform.md). Transforms compose, so a view of a view -of an array is still a single transform, and nothing is read until someone -asks for data. +An indexing operation — a slice, an integer, a fancy index array — is a mapping +from the coordinates a user asks for to the coordinates in storage. This library +represents that mapping as a value: an [`IndexTransform`](api/transform.md). +Transforms compose, so a view of a view of an array is still a single transform, +and no data is read until a caller asks for it. -Three pieces do the work: +The package has three parts: - **The transform algebra** ([`zarr_indexing.transform`](api/transform.md), [`zarr_indexing.domain`](api/domain.md), @@ -31,9 +30,9 @@ Three pieces do the work: `ArrayMap` are three representations of the same thing, a set of integer coordinates, traded off against each other for efficiency. - **Chunk resolution** ([`zarr_indexing.chunk_resolution`](api/chunk_resolution.md)): - given a transform and a chunk grid, which chunks does this selection touch, - which coordinates does it touch *inside* each chunk, and where do the values - land in the output buffer? The resolver is dependency-aware: correlated + given a transform and a chunk grid, compute which chunks the selection + touches, which coordinates it touches inside each chunk, and where the values + land in the output buffer. The resolver is dependency-aware: correlated (`vindex`) array maps are enumerated jointly rather than as a cartesian product, and orthogonal (`oindex`) array maps contribute only the chunks their index arrays actually land in, so resolution scales with the number of @@ -54,11 +53,11 @@ The model is [TensorStore's](https://google.github.io/tensorstore/index_space.ht index transform, reimplemented in Python against NumPy: index domains with explicit origins, output index maps of constant / single-input-dimension / index-array flavour, and composition as the single operation that stacks -views. Names and semantics follow TensorStore where they overlap — notably, -negative indices are literal coordinates, not Python-style offsets from the -end, and it is the caller's job to normalize them. +views. Names and semantics follow TensorStore where they overlap: in +particular, negative indices are literal coordinates, not Python-style offsets +from the end, and the caller is responsible for normalizing them. -The differences are the ones NumPy compatibility forces. `ArrayMap` records +The differences are those NumPy compatibility requires. `ArrayMap` records the input dimension an *orthogonal* (`oindex`) index array varies over, which TensorStore's format has no field for; the [serializer collapses or reconstructs that field](api/json.md) so the wire @@ -66,9 +65,9 @@ format stays TensorStore-loadable. Chunk resolution and the `oindex`/`vindex` helpers exist to serve NumPy-shaped selection semantics, which TensorStore does not have to model. -The [design notes](design-notes.md) set out the comparison in full, including -the four places the two libraries deliberately diverge and the ones where -TensorStore is simply ahead. +The [design notes](design-notes.md) give the comparison in full, including the +four places the two libraries deliberately differ and the areas where +TensorStore is more capable. ## Quickstart @@ -85,12 +84,12 @@ view.domain # IndexDomain(inclusive_min=(10,), exclusive_max=(50,)) view.selection_repr # '{ [10, 50), 5 }' ``` -The domain describes what the *user* sees (here a single dimension, 40 long, -with origin 10); the output maps describe what *storage* sees (a stride-1 -`DimensionMap` and the `ConstantMap` for the dropped dimension). +The domain describes the coordinates the user sees (here a single dimension of +length 40 with origin 10); the output maps describe the storage coordinates (a +stride-1 `DimensionMap` and a `ConstantMap` for the dropped dimension). -Fancy indexing works the same way, in both flavours, and still materializes -nothing but the index arrays themselves: +Fancy indexing works the same way in both flavours, and materializes nothing but +the index arrays themselves: ```python import numpy as np @@ -99,9 +98,9 @@ transform.oindex[np.array([3, 1, 90]), 0:4] # '{ {3, 1, 90}, [0, 4) }', shape transform.vindex[np.array([0, 40, 99]), np.array([1, 2, 3])] # shape (3,) ``` -Transforms built independently stack with +Transforms built independently are combined with [`compose`](api/composition.md), which is what `transform[...]` uses -internally when you index an already-indexed view: +internally when indexing an already-indexed view: ```python from zarr_indexing import compose @@ -112,10 +111,10 @@ outer = IndexTransform.from_shape((50,))[10:20] compose(outer, inner).selection_repr # '{ [20, 40) step 2 }' ``` -Resolution against a chunk grid is where a transform finally meets storage. -Chunk resolution asks the grid only for the per-dimension index-to-chunk -mapping described by [`DimensionGridLike`](api/grid.md), so any object with -those four methods will do: +Chunk resolution applies a transform to a chunk grid. It asks the grid only for +the per-dimension index-to-chunk mapping described by +[`DimensionGridLike`](api/grid.md), so any object with those four methods +works: ```python from dataclasses import dataclass @@ -142,16 +141,16 @@ for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(view, grid ``` Each yielded `sub_transform` is the original transform restricted to one chunk -and translated into chunk-local coordinates — exactly what a codec pipeline +and translated into chunk-local coordinates, which is 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: +[`LazyArray`](api/lazy_array.md) applies all of the above to an existing array — +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 @@ -174,22 +173,22 @@ view.result()[:, :4] # [ 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: +In the transform algebra underneath, indices are literal coordinates in a view's +own (possibly shifted) domain. `LazyArray` instead uses **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: +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. Advanced indices are placed as NumPy places 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) @@ -206,20 +205,20 @@ LazyArray(x).lazy.vindex[..., [3, 0]].result() # [11, 8]]) ``` -Use a length-1 list (`oindex[[1]]`) where you want to keep an axis. +Use a length-1 list (`oindex[[1]]`) to keep an axis. -Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a `LazyArray` a -drop-in source for `dask.array.from_array`. Everything else NumPy might do to it -— arithmetic, `numpy.sum`, `numpy.stack` — materializes the whole view through -`__array__` and hands back a plain NumPy array: laziness here is about indexing, -not about deferred compute. +Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a `LazyArray` usable +as a source for `dask.array.from_array`. Every other NumPy operation — +arithmetic, `numpy.sum`, `numpy.stack` — materializes the whole view through +`__array__` and returns a plain NumPy array. Laziness here applies to indexing, +not to deferred computation. ### Parts A read is broken up along a **partitioning**: a grid of boxes the wrapper walks -through [chunk resolution](api/chunk_resolution.md), so each box is read once -with plain basic slicing and the selection is applied to the block in memory. -`parts()` exposes that walk, projected through the view: +through [chunk resolution](api/chunk_resolution.md). Each box is read once with +basic slicing, and the selection is applied to the block in memory. `parts()` +exposes that walk, projected through the view: ```python part = next(iter(lazy.lazy[0:35, 0:10].parts())) @@ -231,14 +230,14 @@ part.out_selection # (slice(0, 30, None), slice(0, 10, None)) part.is_complete # False — the view does not cover the whole box ``` -Each `part.array` is an ordinary `LazyArray`, so parts resolve independently and -concurrently, and `result()` is nothing but the assembly of that walk. Note that -`part.array.bounding_box()` is *part-local*; `part.box` is the global one, and is -what tells two parts apart. +Each `part.array` is an ordinary `LazyArray`, so parts can be resolved +independently and concurrently, and `result()` assembles that walk. +`part.array.bounding_box()` is part-local; `part.box` gives the box in the +wrapped array's coordinates, which is what distinguishes two parts. The partitioning is discovered from the wrapped array — `read_chunk_sizes`, then -`chunks`. Those are the *source's* words; the wrapper's own surface speaks only -of parts, and `with_parts` chooses a different one without touching the data or +`chunks`. Those attribute names belong to the source; this API refers only to +parts. `with_parts` selects a different partitioning without touching the data or the view: ```python @@ -254,16 +253,17 @@ view = lazy.lazy[10:50, ::4] # [(0, 0)] — one whole-array part; resolve in one shot ``` -`result()` is identical under any of them — repartitioning changes how the read -is cut up, never what it returns. Boxes that deliberately straddle the source's -own (to bound peak memory, say) are legal; they cost I/O, not correctness. +`result()` returns the same values under any of them: repartitioning changes how +the read is divided, not what it returns. Boxes that straddle the source's own +(to bound peak memory, for example) are permitted; they cost extra I/O but do +not affect correctness. ### Boxes and queries A selection is either **rectangular** — describable by an interval and a stride -per dimension — or a **query**, an explicit list of coordinates. `is_box` says -which; `bounding_box()` gives the storage region touched either way, and -`strides()` how densely a box covers it: +per dimension — or a **query**, an explicit list of coordinates. `is_box` +reports which. `bounding_box()` gives the storage region touched in either case, +and `strides()` gives the step per dimension for a box: ```python slab = lazy.lazy[10:50, ::4] @@ -278,15 +278,15 @@ gather.bounding_box() # ((3, 91), (0, 80)) — 88 rows of hull over 3 selected gather.strides() # None ``` -The bounding box is *dense* — every cell in it selected — only when every stride -is 1. The slab above spans a 40x77 hull over the 40x20 it selects, so slabbing -the hull and discarding the rest would read 3.85x what is needed; a consumer -that wants one rectangular read has to honour `strides()` too. +The bounding box is dense — every cell in it selected — only when every stride +is 1. The slab above spans a 40x77 hull over the 40x20 cells it selects, so +reading the hull and discarding the rest would transfer 3.85x the required data; +a consumer that wants one rectangular read has to apply `strides()` too. -It is a structural category, not a heuristic: basic indexing composes to a box -forever, and one `oindex`, `vindex`, or mask anywhere in the chain leaves it -forever. The [design notes](design-notes.md) explain why that matters to -anything downstream of a selection. +The category is structural, not a heuristic: basic indexing always composes to a +box, and one `oindex`, `vindex`, or mask anywhere in the chain makes the +selection a query permanently. The [design notes](design-notes.md) explain why +that matters to consumers of a selection. ## Reference diff --git a/packages/zarr-indexing/docs/ndsel.md b/packages/zarr-indexing/docs/ndsel.md index 74f394b758..2b03685d87 100644 --- a/packages/zarr-indexing/docs/ndsel.md +++ b/packages/zarr-indexing/docs/ndsel.md @@ -7,7 +7,7 @@ title: The ndsel wire format [ndsel](https://github.com/zarr-developers/ndsel) is a draft JSON representation of NumPy-style n-dimensional selections, adapted from TensorStore's `IndexTransform` model. `zarr-indexing` implements it in two -layers, and the split between them is the thing worth understanding: +layers: | Layer | Module | Depends on | Job | | --- | --- | --- | --- | @@ -15,16 +15,16 @@ layers, and the split between them is the thing worth understanding: | Engine | [`zarr_indexing.json`](api/json.md) | NumPy | Lowers a *canonical* body into an in-memory [`IndexTransform`](api/transform.md), and back. | Constraints that only make sense for a real array — finite bounds, index -arrays as `ndarray`s — live in the engine layer and nowhere else. That is why -`messages` can happily normalize a message with `"-inf"` bounds that -`json.transform_from_canonical` will refuse to lower. +arrays as `ndarray`s — live in the engine layer and nowhere else. As a result, +`messages` normalizes a message with `"-inf"` bounds that +`json.transform_from_canonical` refuses to lower. ## Two entry points [`parse_ndsel`](api/messages.md#zarr_indexing.messages.parse_ndsel) -structurally validates a message of any kind and returns it **unchanged** — -use it when you want to keep a message in its compact shorthand form but -confirm it is well formed. +structurally validates a message of any kind and returns it unchanged. Use it +to confirm that a message is well formed while keeping it in its compact +shorthand form. [`normalize_ndsel`](api/messages.md#zarr_indexing.messages.normalize_ndsel) desugars a message into the single deterministic **canonical transform body** @@ -112,10 +112,10 @@ varies over. The serializer bridges that gap in both directions: inherently ambiguous between the two flavours and reconstructs as orthogonal, which is behaviorally identical in that case. -There is one deliberate exception, worth calling out because it is the one -place a round trip changes representation rather than preserving it. An -all-singleton `index_array` — size 1 — selects the same coordinate no matter -what the input is, so it is **collapsed to a `constant` map** on serialize: +There is one deliberate exception, and it is the only place a round trip changes +representation rather than preserving it. An all-singleton `index_array` — size +1 — selects the same coordinate regardless of the input, so it is collapsed to +a `constant` map on serialize: ```python from zarr_indexing import IndexTransform, transform_to_canonical @@ -129,10 +129,10 @@ transform_to_canonical(IndexTransform.from_shape((100, 100)).oindex[[5], 0:2]) # {'offset': 0, 'stride': 1, 'input_dimension': 1}]} ``` -The size-1 input dimension stays in the domain, unconsumed by any output map — -still a valid transform, and still the right output shape. A length-1 `oindex` -selection therefore round-trips *behaviorally* (an `ArrayMap` comes back as a -`ConstantMap`) rather than by object identity. +The size-1 input dimension stays in the domain, unconsumed by any output map. +The transform is still valid and the output shape is unchanged. A length-1 +`oindex` selection therefore round-trips behaviorally (an `ArrayMap` comes back +as a `ConstantMap`) rather than by object identity. ## Conformance @@ -151,6 +151,6 @@ Do not edit the vendored files; to pick up spec changes, re-vendor from a newer ndsel commit and update the recorded SHA. A second, optional test (`tests/test_ndsel_tensorstore.py`, skipped unless -`tensorstore` is installed) closes the loop against a real TensorStore by -loading canonical bodies into `tensorstore.IndexTransform` and re-loading -TensorStore's own `to_json()` output back through the engine layer. +`tensorstore` is installed) checks against TensorStore itself by loading +canonical bodies into `tensorstore.IndexTransform` and re-loading TensorStore's +own `to_json()` output back through the engine layer. diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py index 594d66b5ba..093f1295b7 100644 --- a/packages/zarr-indexing/src/zarr_indexing/boundary.py +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -1,15 +1,15 @@ """The positional (NumPy) selection dialect, lowered onto the transform algebra. -The transform algebra speaks **literal domain coordinates**: an index is a -point in the view's own coordinate system, which after `view = arr[10:50]` runs -from 10 to 49, and a negative index is genuinely negative rather than counted +The transform algebra uses **literal domain coordinates**: an index is a point +in the view's own coordinate system, which after `view = arr[10:50]` runs from +10 to 49, and a negative index is a negative coordinate rather than an offset from the end (TensorStore's convention — see `zarr_indexing.transform`). -NumPy speaks **positions**: index 0 always means the first element of the thing -you are indexing, and `-1` means the last. This module is the translation layer -between the two. It validates a selection against the view's shape with NumPy -semantics, then shifts every coordinate by the domain's origin so the transform -layer sees literal coordinates. +NumPy uses **positions**: index 0 always means the first element of the object +being indexed, and `-1` means the last. This module translates between the two. +It validates a selection against the view's shape with NumPy semantics, then +shifts every coordinate by the domain's origin so the transform layer sees +literal coordinates. Note ---- @@ -157,7 +157,7 @@ def split_scalar_axes( ) -> tuple[tuple[Any, ...] | None, Any]: """Peel scalar integer indices out of a fancy selection. - In NumPy, a scalar integer is a *basic* index wherever it appears: it drops + In NumPy, a scalar integer is a basic index wherever it appears: it drops its axis, and is applied before the advanced indices rather than broadcast against them. `a[0, [1, 2], :]` is `a[0][[1, 2], :]`. Neither the orthogonal nor the vectorized path of the transform algebra models that — both widen a @@ -219,7 +219,7 @@ def normalize_positional_selection( ) -> Any: """Translate a positional (NumPy-dialect) selection into literal coordinates. - Positions are zero-based offsets into the *current view*; negatives wrap + Positions are zero-based offsets into the current view; negatives wrap from the end. The returned selection addresses the same cells in the literal coordinate system `domain` uses, ready for `zarr_indexing.transform.selection_to_transform`. diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index c33b13b41c..19439727fd 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -11,7 +11,8 @@ ``` Nothing is read until `result()` (or `__array__`, or an eager `__getitem__`). -Composition is free: a view of a view is still a single transform. +Composition does not accumulate layers: a view of a view is still a single +transform. Parts ----- @@ -20,24 +21,24 @@ view, yielding a [`Partition`](#zarr_indexing.lazy_array.Partition) per box — the base coordinates of the box, a `LazyArray` covering exactly the cells of the view that live in it, and where those cells belong in the result. `result()` is -built on that walk and nothing else: +built on that walk: ```python for part in view.parts(): out[part.out_selection] = part.array.result() ``` -so each box is read once, with plain basic slicing, and the selection is applied -to the block in memory. +Each box is therefore read once, with basic slicing, and the selection is +applied to the block in memory. The partitioning is discovered from the wrapped array at construction — first `read_chunk_sizes` (zarr's clipped per-axis sizes, sharding-aware), then `chunks`, read as per-axis sizes if its entries are sequences and as a uniform -box shape if they are integers. Those attribute names are the *source's* -vocabulary; the wrapper's own surface speaks only of parts. An array that -advertises neither is a single whole-array part, and resolving it lowers the -whole view to one pass of array operations (`take`, basic slicing, a flat -gather) rather than materializing a block first. +box shape if they are integers. Those attribute names belong to the wrapped +array; this API refers only to parts. An array that advertises neither gets a +single whole-array part, and resolving it lowers the whole view to one pass of +array operations (`take`, basic slicing, a flat gather) rather than +materializing a block first. `with_parts` replaces the partitioning without touching the data or the view: @@ -47,21 +48,21 @@ view.with_parts(None) # one whole-array part; resolve in one shot ``` -Repartitioning never changes what `result()` returns — only how the read is cut -up. Deliberately misaligning the parts with the source's own boxes is legal and -useful (to bound peak memory, or to batch small reads); it costs I/O, not -correctness. +Repartitioning changes how the read is divided, not what `result()` returns. +Parts that do not align with the source's own boxes are permitted and can be +useful (to bound peak memory, or to batch small reads); they cost extra I/O but +do not affect correctness. Boxes and queries ----------------- A selection is either **rectangular** — an interval and a stride per dimension, -which is what basic indexing composes to, however deeply — or a **query**, an +which is what basic indexing composes to at any depth — or a **query**, an explicit list of coordinates, which is what `oindex`, `vindex`, and masks -produce and which no amount of subsequent basic indexing undoes. `is_box` -reports the category and `bounding_box()` the storage region touched, exact for -a box and a hull for a query. The distinction is structural rather than an -optimization; [the design notes](../design-notes.md) say why it matters to -anything downstream of a selection. +produce and which subsequent basic indexing cannot undo. `is_box` reports the +category and `bounding_box()` reports the storage region touched: exact for a +box, a hull for a query. The distinction is structural rather than an +optimization; [the design notes](../design-notes.md) describe why it matters to +consumers of a selection. The positional dialect ---------------------- @@ -69,41 +70,41 @@ element of the current view, `-1` is the last, boolean masks must match the view's shape, and every index is bounds-checked against the view. -This is deliberately different from `zarr.Array.lazy[...]`, which exposes the +This differs deliberately from `zarr.Array.lazy[...]`, which exposes the **literal** TensorStore dialect: a zarr view keeps the coordinate system of the array it came from, so after `v = arr.lazy[10:50]` the first element of `v` is `v[10]` and a negative index is out of bounds rather than counted from the end. -That dialect is the right one for zarr, where a view's coordinates are meant to -stay comparable with the parent array's. `LazyArray` is a duck array — it has to -behave like the thing it wraps to be usable as a drop-in for NumPy or as a dask -source — so it re-zeroes its coordinates on every view and speaks positions. -`zarr_indexing.boundary` performs the translation between the two. +That dialect suits zarr, where a view's coordinates stay comparable with the +parent array's. `LazyArray` is a duck array and has to behave like the array it +wraps to be usable as a NumPy drop-in or as a dask source, so it re-zeroes its +coordinates on every view and uses positions. `zarr_indexing.boundary` performs +the translation between the two. Two more NumPy rules the dialect keeps, in every mode: -- **A scalar integer drops its axis.** It is a basic index wherever it appears, +- A scalar integer drops its axis. It is a basic index wherever it appears, applied before any advanced index rather than broadcast against one. So `lazy.oindex[0]` has the shape of `x[0]`, `lazy.oindex[0, [1, 2], :]` means `x[0][numpy.ix_([1, 2], ...)]`, and `lazy.oindex[0, 1, 2]` and `lazy.vindex[0, 1, 2]` are both zero-rank. Use a length-1 list to keep an axis. -- **Advanced indices land where NumPy puts them.** For a `vindex` selection that - leaves some axes unindexed, the gathered dimensions sit where the coordinate - arrays sat when those arrays are adjacent, and lead when a slice separates - them — so `lazy.vindex[..., i, j]` has shape `(x.shape[0], *broadcast)`, - matching `x[..., i, j]`. +- Advanced indices are placed as NumPy places them. For a `vindex` selection + that leaves some axes unindexed, the gathered dimensions sit where the + coordinate arrays sat when those arrays are adjacent, and lead when a slice + separates them — so `lazy.vindex[..., i, j]` has shape + `(x.shape[0], *broadcast)`, matching `x[..., i, j]`. Materializing on fallback ------------------------- `LazyArray` implements `__array__` but deliberately implements neither -`__array_ufunc__`/`__array_function__` nor `__array_namespace__`. **Every NumPy -operation other than indexing therefore materializes the whole view**: +`__array_ufunc__`/`__array_function__` nor `__array_namespace__`. Every NumPy +operation other than indexing therefore materializes the whole view: `numpy.sum(view)`, `view + 1`, `numpy.stack([view, view])` all convert through -`__array__` first and give you a plain NumPy array back. That is intentional — -laziness here is about *indexing*, not about building a deferred compute graph — -but it means a `LazyArray` is not a drop-in for arithmetic on a large array. Use -`.lazy[...]` to narrow first, or hand the wrapper to `dask.array.from_array` and -let dask own the compute graph. +`__array__` first and return a plain NumPy array. That is intentional: laziness +here applies to indexing, not to building a deferred compute graph. It does mean +a `LazyArray` is not a drop-in for arithmetic on a large array. Use `.lazy[...]` +to narrow the view first, or pass the wrapper to `dask.array.from_array` so that +dask owns the compute graph. Device caveat ------------- @@ -241,7 +242,7 @@ def _discover_parts(array: Any, shape: tuple[int, ...]) -> tuple[EdgeDimensionGr partitioning of `shape` means "this object does not advertise one I understand", and the array is treated as unpartitioned rather than rejected. A partitioning is an I/O strategy, so reading the whole array is always a - correct fallback. `with_parts` is *our* API, and validates strictly. + correct fallback. `with_parts` is a public API and validates strictly. """ declared = _read_source_attribute(array, "read_chunk_sizes") if declared is None: @@ -332,9 +333,9 @@ def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: in def _lower(array: Any, transform: IndexTransform) -> Any: """Lower a transform to one pass of array operations over `array`. - The single lowering engine: every read, partitioned or not, ends here. The - result is always in the transform's own domain axis order and of exactly its - domain shape. + Every read, partitioned or not, goes through this function. The result is + always in the transform's own domain axis order and of exactly its domain + shape. """ if _is_correlated(transform): result = _lower_correlated(array, transform) @@ -516,7 +517,7 @@ def _covers_whole_part(local_transform: IndexTransform, part_shape: tuple[int, . """Whether a part-local transform addresses every cell of its part. Conservative: an `ArrayMap` could in principle enumerate a whole part, but - proving it costs more than the answer is worth, so a fancy-indexed axis + checking that costs more than the answer is worth, so a fancy-indexed axis always reports incomplete. """ domain = local_transform.domain @@ -553,10 +554,10 @@ class Partition: Which box of the base partitioning this is, one coordinate per dimension of the wrapped array. box - The box itself, in the **global** storage coordinates of the wrapped + The box itself, in the global storage coordinates of the wrapped array: one `[inclusive_min, exclusive_max)` interval per dimension. - `array.bounding_box()` is part-*local* and so cannot tell two parts - apart; this can, and it is what pairs with `is_complete` to decide a + `array.bounding_box()` is part-local and so cannot distinguish two + parts; this attribute can, and together with `is_complete` it decides a read-modify-write. array A `LazyArray` covering exactly the cells of the view that live in this @@ -590,7 +591,7 @@ def _wrapped_token(array: Any) -> Any: lazily — this package never requires it); otherwise a local fallback that digests the contents of a small array and, above `_TOKEN_DIGEST_LIMIT`, falls back to a structural description that is stable across processes but - **not** sensitive to the array's contents. + not sensitive to the array's contents. """ hook = getattr(array, "__dask_tokenize__", None) if hook is not None: @@ -632,9 +633,9 @@ def _wrapped_token(array: Any) -> Any: class LazyArray: """A lazily-indexable view over an array-API-like array. - Wrapping is cheap and non-destructive: the wrapped array is never copied and - never read at construction time. Indexing through `.lazy` composes an - `IndexTransform` and returns another `LazyArray`; `result()` materializes. + Wrapping neither copies nor reads the wrapped array at construction time. + Indexing through `.lazy` composes an `IndexTransform` and returns another + `LazyArray`; `result()` materializes. Selections use the **positional NumPy dialect**, and reads are broken up along a **partitioning** discovered from the wrapped array — see the module @@ -739,20 +740,21 @@ def is_box(self) -> bool: True exactly when the composed transform's output maps are all `ConstantMap` or `DimensionMap` — no `ArrayMap`. Such a selection is affine and monotone along every axis, so it is described completely by - an interval **and a stride** per dimension: + an interval and a stride per dimension: [`bounding_box`][zarr_indexing.lazy_array.LazyArray.bounding_box] together with [`strides`][zarr_indexing.lazy_array.LazyArray.strides]. Basic indexing, - however deeply composed, stays a box; one `oindex`, `vindex`, or mask - anywhere in the chain leaves it forever. + at any depth of composition, stays a box; one `oindex`, `vindex`, or + mask anywhere in the chain makes the selection a query permanently. - A box is *dense* — every cell of its bounding box selected — only when + A box is dense — every cell of its bounding box selected — only when every stride is 1. A strided box covers its hull sparsely: `lazy[10:50, ::4]` selects 40x20 cells out of a 40x77 hull, so a - consumer that slabs the hull and discards the rest reads 3.85x what it - needs. Check `strides` before treating a box as a single slab read. + consumer that reads the whole hull and discards the rest transfers 3.85x + the data it needs. Check `strides` before treating a box as a single + slab read. - The distinction is what lets a consumer decide between a slab read and a + The distinction lets a consumer decide between a slab read and a gather; see [the design notes](../design-notes.md) for why it is a category rather than an optimization. @@ -773,10 +775,10 @@ def bounding_box(self) -> tuple[tuple[int, int], ...] | None: this view reads from that contains every coordinate the selection reaches. - The hull is only *dense* — every cell in it selected — for a box whose + The hull is dense — every cell in it selected — only for a box whose every stride is 1. A strided box selects a sublattice of its hull (pair this with [`strides`][zarr_indexing.lazy_array.LazyArray.strides] to - describe it fully), and a query's hull is merely a superset that can be + describe it fully), and a query's hull is a superset that can be arbitrarily loose: `oindex[[0, 999]]` has a 1000-wide hull over two rows. @@ -784,14 +786,14 @@ def bounding_box(self) -> tuple[tuple[int, int], ...] | None: ------- tuple of (int, int), or None One interval per storage dimension, or `None` when the view is - empty (`size == 0`) and so touches no coordinate at all — an empty - selection has no meaningful origin to report an interval around. + empty (`size == 0`) and so touches no coordinate at all, leaving no + interval to report. Notes ----- The coordinates are those of the array this view reads from, which for the `array` of a [`Partition`][zarr_indexing.lazy_array.Partition] is - **part-local** — relative to that part's own box, not the wrapped array. + part-local — relative to that part's own box, not the wrapped array. Two parts of the same view can therefore report identical bounding boxes; [`Partition.box`][zarr_indexing.lazy_array.Partition] gives the global box they sit in. @@ -827,30 +829,30 @@ def bounding_box(self) -> tuple[tuple[int, int], ...] | None: def strides(self) -> tuple[int, ...] | None: """The step between selected coordinates, one per storage dimension. - Completes the description a box selection needs: `bounding_box()` says - *where*, `strides()` says *how densely*. A stride of 1 means every cell - of the hull along that dimension is selected; `k` means every `k`-th. - Dimensions fixed by an integer index report 1 — they span a single - coordinate. + Together with `bounding_box()`, this fully describes a box selection: + `bounding_box()` gives the interval per dimension, `strides()` gives the + step per dimension. A stride of 1 means every cell of the hull along + that dimension is selected; `k` means every `k`-th. Dimensions fixed by + an integer index report 1 — they span a single coordinate. Returns ------- tuple of int, or None One positive stride per storage dimension, or `None` when [`is_box`][zarr_indexing.lazy_array.LazyArray.is_box] is false: a - query's coordinates are a lookup table and have no step. An - **empty** box still reports its strides even though + query's coordinates are a lookup table and have no step. An empty + box still reports its strides even though [`bounding_box`][zarr_indexing.lazy_array.LazyArray.bounding_box] - returns `None` — the step is a property of the selection's shape, - not of the (empty) region it touches. + returns `None`, because the step is a property of the selection's + shape, not of the (empty) region it touches. Notes ----- - The magnitude only. A reversing view (`lazy[::-1]`) selects the same - *set* of coordinates as its forward twin, so it reports the same - bounding box and the same strides; the traversal direction lives in the - transform, not in this description of the region touched. A consumer - that cares about order reads the transform, or simply reverses the block + Magnitudes only. A reversing view (`lazy[::-1]`) selects the same set of + coordinates as the equivalent forward view, so it reports the same + bounding box and the same strides. The traversal direction is recorded + in the transform, not in this description of the region touched. A + consumer that needs the order reads the transform, or reverses the block it gets back. Examples @@ -876,8 +878,8 @@ def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> L """Return the same view, read in different parts. The transform, the wrapped array, and therefore `result()` are all - unchanged; only the boxes the read is broken into differ. Cheap: nothing - is copied and nothing is read. + unchanged; only the boxes the read is broken into differ. Nothing is + copied and nothing is read. Parameters ---------- @@ -886,7 +888,7 @@ def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> L array, with the trailing box clipped to the extent), dask-convention per-axis sizes (one sequence of box extents per dimension, each summing to the base extent), or `None` for a single part covering - the whole array — which makes `result()` lower the view in one shot + the whole array, which makes `result()` lower the view in one pass instead of assembling it block by block. Returns @@ -899,8 +901,8 @@ def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> L ValueError If `parts` has the wrong length, mixes the two conventions, contains a non-positive extent, or declares per-axis sizes that do not sum to - the base shape. Unlike discovery, this is our own API and validates - strictly. + the base shape. Unlike partition discovery, this is a public API and + validates strictly. Examples -------- @@ -917,7 +919,7 @@ def parts(self) -> Iterator[Partition]: Yields one [`Partition`][zarr_indexing.lazy_array.Partition] per box the view actually touches. The parts tile the view exactly and disjointly, - and each carries a `LazyArray` that can be resolved on its own — in + and each carries a `LazyArray` that can be resolved on its own: in another thread, in another order, or not at all. A wrapper with no partitioning (see `with_parts`) yields a single part @@ -996,11 +998,11 @@ def _select(self, selection: Any, mode: SelectionMode) -> LazyArray: return LazyArray._derive(self._array, composed, self._parts, self._window) def __getitem__(self, selection: Any) -> Any: - """Read a basic selection **eagerly**, like `numpy.ndarray.__getitem__`. + """Read a basic selection eagerly, like `numpy.ndarray.__getitem__`. - Eager, not lazy, so that a `LazyArray` is a drop-in duck array for - consumers (dask's `from_array`, `numpy.asarray`) that expect indexing to - produce data. Use `.lazy[...]` for the lazy form. + Reads here are eager, not lazy, so that a `LazyArray` works as a duck + array for consumers (dask's `from_array`, `numpy.asarray`) that expect + indexing to produce data. Use `.lazy[...]` for the lazy form. """ return self._select(selection, "basic").result() @@ -1008,8 +1010,8 @@ def result(self) -> Any: """Materialize this view. Assembles the view from its `parts()`, reading each box once. A wrapper - with a single whole-array part skips the buffer entirely and lowers - straight to array operations. + with a single whole-array part skips the output buffer and lowers + directly to array operations. Returns ------- @@ -1046,12 +1048,12 @@ def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: def __dask_tokenize__(self) -> Any: """A deterministic token: the wrapped array, the view, and the parts. - Two wrappers token-equal when they wrap the same data, address the same - cells, and read them in the same boxes. The view contributes its - canonical ndsel body, so transforms that differ only in representation - token alike. See `_wrapped_token` for the determinism scope of the - wrapped array's contribution; dask is imported lazily and is never a - requirement of this package. + Two wrappers produce equal tokens when they wrap the same data, address + the same cells, and read them in the same boxes. The view contributes + its canonical ndsel body, so transforms that differ only in + representation produce the same token. See `_wrapped_token` for the + determinism scope of the wrapped array's contribution; dask is imported + lazily and is never a requirement of this package. """ return ( type(self).__qualname__, diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 4bf96b5fe0..1834787064 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -612,8 +612,8 @@ def _positional_slice(pos: int, size: int, step: int) -> slice: """A NumPy slice selecting `size` elements from `pos`, walking by `step`. The stop is `pos + size*step`, except that a downward walk reaching the - start of the array must stop at `None`: NumPy would read a negative stop as - counting from the end, so `slice(6, -1, -1)` selects **nothing** where + start of the array must stop at `None`: NumPy reads a negative stop as + counting from the end, so `slice(6, -1, -1)` selects nothing where `slice(6, None, -1)` selects the first seven elements in reverse. """ stop = pos + size * step @@ -631,7 +631,7 @@ def _reindex_array( The array's axes correspond to the transform's input dimensions (0-indexed over the domain shape). Each axis is either a **dependency axis** — the array - genuinely varies with that input dimension — or a **singleton** axis it + varies with that input dimension — or a **singleton** axis it broadcasts over. Integer indexing, slicing, or newaxis is applied to the array only along its dependency axes; a selection on a singleton axis does not touch the array's values (it just narrows or drops that broadcast axis). @@ -879,14 +879,14 @@ def array_map_dependent_axis(m: ArrayMap) -> int | None: Normally this is the array's one non-singleton axis. A degenerate length-1 orthogonal selection normalizes to an all-singleton shape (its dependency axes are empty and indistinguishable by shape from a scalar), so - `input_dimension` breaks the tie — it records the *live* axis the map binds. + `input_dimension` breaks the tie — it records the live axis the map binds. Returns ------- int or None The axis the map varies over, or `None` when it varies over no input axis at all — a correlated (`vindex`) map, or a map that has become - constant. `None` is a real answer, not an error: callers that need an + constant. `None` is a valid result, not an error: callers that need an axis must handle it rather than reading a stale `input_dimension`. Indexing operations collapse a map whose every dependency axis is consumed by an integer index into a `ConstantMap`, so an @@ -1301,16 +1301,16 @@ def _resolve_slice_ts(sel: slice, dim: int, lo: int, hi: int) -> tuple[int, int, - a non-empty interval must be contained in the domain (no clamping — a NumPy-style out-of-range or negative bound is an error, not a shorter or wrapped result); - - an **empty** interval is valid anywhere, for either sign; + - an empty interval is valid anywhere, for either sign; - an interval running the wrong way (`stop` on the far side of `start` from the direction of travel) is an error, not an empty result; - the result's domain origin is `trunc(start/step)` — toward zero, for both signs — and coordinate `origin + k` maps to input `start + k*step`. - A negative step normally produces a **negative origin**: reversing a - zero-origin axis of length 20 gives the domain `[-19, 1)`. That is the - coordinate frame staying anchored to the source; a caller that needs - non-negative coordinates re-bases explicitly (`translate_domain_to`). + A negative step normally produces a negative origin: reversing a + zero-origin axis of length 20 gives the domain `[-19, 1)`. The coordinate + frame stays anchored to the source; a caller that needs non-negative + coordinates re-bases explicitly with `translate_domain_to`. Returns `(start, step, origin, size)` in domain coordinates. """ From 252ec55620278e1b98d594f5d5ba5069143c1dc4 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 12:22:22 +0200 Subject: [PATCH 05/32] docs(zarr-indexing): American spelling (flavour -> flavor) Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-indexing/docs/design-notes.md | 2 +- packages/zarr-indexing/docs/index.md | 6 +-- packages/zarr-indexing/docs/ndsel.md | 2 +- .../zarr-indexing/src/zarr_indexing/json.py | 2 +- .../src/zarr_indexing/output_map.py | 2 +- .../src/zarr_indexing/transform.py | 2 +- .../tests/test_chunk_resolution.py | 2 +- .../zarr-indexing/tests/test_lazy_array.py | 44 +++++++++---------- 8 files changed, 31 insertions(+), 31 deletions(-) diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index 6977aca76b..bdfc7adf40 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -16,7 +16,7 @@ parts are intentionally identical: - **The model.** An `IndexTransform` pairs an input `IndexDomain` — a rectangular region with an explicit, possibly non-zero origin — with one - output index map per storage dimension, in the three flavours TensorStore + output index map per storage dimension, in the three flavors TensorStore defines: constant, single-input-dimension (affine), and index array. Composition is the single operation that stacks views. - **Slice semantics.** Slice bounds are literal domain coordinates: no diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index c8f5bcba42..3f6f7afe19 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -52,7 +52,7 @@ per-dimension grids satisfy structurally. The model is [TensorStore's](https://google.github.io/tensorstore/index_space.html) index transform, reimplemented in Python against NumPy: index domains with explicit origins, output index maps of constant / single-input-dimension / -index-array flavour, and composition as the single operation that stacks +index-array flavor, and composition as the single operation that stacks views. Names and semantics follow TensorStore where they overlap: in particular, negative indices are literal coordinates, not Python-style offsets from the end, and the caller is responsible for normalizing them. @@ -88,7 +88,7 @@ The domain describes the coordinates the user sees (here a single dimension of length 40 with origin 10); the output maps describe the storage coordinates (a stride-1 `DimensionMap` and a `ConstantMap` for the dropped dimension). -Fancy indexing works the same way in both flavours, and materializes nothing but +Fancy indexing works the same way in both flavors, and materializes nothing but the index arrays themselves: ```python @@ -177,7 +177,7 @@ In the transform algebra underneath, indices are literal coordinates in a view's own (possibly shifted) domain. `LazyArray` instead uses **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: +orthogonal and vectorized flavors: ```python LazyArray(np.arange(12).reshape(3, 4)).lazy.vindex[[0, 2], [1, 3]].result() diff --git a/packages/zarr-indexing/docs/ndsel.md b/packages/zarr-indexing/docs/ndsel.md index 2b03685d87..e02b850e95 100644 --- a/packages/zarr-indexing/docs/ndsel.md +++ b/packages/zarr-indexing/docs/ndsel.md @@ -109,7 +109,7 @@ varies over. The serializer bridges that gap in both directions: solely owns a single non-singleton axis is orthogonal; arrays that share non-singleton axes, or vary over several, are correlated (`vindex`), and get `input_dimension = None`. A single 1-D array over a rank-1 domain is - inherently ambiguous between the two flavours and reconstructs as + inherently ambiguous between the two flavors and reconstructs as orthogonal, which is behaviorally identical in that case. There is one deliberate exception, and it is the only place a round trip changes diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py index c95696f309..89c8a04c42 100644 --- a/packages/zarr-indexing/src/zarr_indexing/json.py +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -34,7 +34,7 @@ single non-singleton axis is orthogonal (`input_dimension = that axis`); arrays that share non-singleton axes, or vary over several, are correlated (`vindex`, `input_dimension = None`). A single 1-D array over a rank-1 domain is - inherently ambiguous between the two flavours; it reconstructs as orthogonal, + inherently ambiguous between the two flavors; it reconstructs as orthogonal, which is behaviorally identical for the single-array case. `index_transform_to_json` / `index_transform_from_json` (and the `*_domain_*` diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py index 581229bd22..ca828a8ff9 100644 --- a/packages/zarr-indexing/src/zarr_indexing/output_map.py +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -78,7 +78,7 @@ class ArrayMap: enclosing transform: `index_array` has the enclosing domain's rank, sized fully on the axes it varies over and singleton (size 1) elsewhere. The dependency axes are therefore derivable from the shape (see - `transform._array_map_dependency_axes`), which distinguishes the two flavours + `transform._array_map_dependency_axes`), which distinguishes the two flavors of multi-array fancy indexing: - **orthogonal** (`oindex`): each array varies along a single, *distinct* diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 1834787064..f056ca119e 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -263,7 +263,7 @@ def _intersect( For each output dimension, restrict to storage coordinates within `[output_domain.inclusive_min[d], output_domain.exclusive_max[d])`. - Two flavours of fancy indexing require different treatment, distinguished by + Two flavors of fancy indexing require different treatment, distinguished by the ArrayMaps' dependency axes (see `_array_map_dependency_axes`): - **orthogonal** (`oindex`): each ArrayMap varies over a single, distinct diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index 0738384e2b..f6a37c7012 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -461,7 +461,7 @@ def test_mixed_maps_2d(self) -> None: assert drop_axes == () -class TestChunkResolutionArrayMapFlavours: +class TestChunkResolutionArrayMapFlavors: """Chunk resolution must yield outer-product (np.ix_) selectors for orthogonal ArrayMaps and shared flat-scatter selectors for correlated ones, and must return early for empty fancy selections.""" diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index cda0b14513..e85604bf4d 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2,7 +2,7 @@ The happy-path suite is a single oracle test: every selection case is applied both to a `LazyArray` and to the NumPy array it wraps, and the results must -match. The case list is crossed with four source flavours — an unchunked NumPy +match. The case list is crossed with four source flavors — an unchunked NumPy array, NumPy with each of the two declared chunk conventions, and a zarr array whose chunking is auto-discovered — so the chunked and unchunked resolution strategies are held to the same answers. @@ -38,7 +38,7 @@ def reference() -> np.ndarray[Any, np.dtype[np.int64]]: - """The array every source flavour holds, and the oracle for every case.""" + """The array every source flavor holds, and the oracle for every case.""" return np.arange(int(np.prod(SHAPE)), dtype=np.int64).reshape(SHAPE) @@ -197,24 +197,24 @@ def make_zarr_source() -> Any: return array -def make_source(flavour: str) -> LazyArray: +def make_source(flavor: str) -> LazyArray: """Build a `LazyArray` over `reference()` with the requested partitioning.""" data = reference() - if flavour == "numpy-whole": + if flavor == "numpy-whole": return LazyArray(data) - if flavour == "numpy-explicit-parts": + if flavor == "numpy-explicit-parts": return LazyArray(data).with_parts(EXPLICIT_PARTS) - if flavour == "numpy-uniform-parts": + if flavor == "numpy-uniform-parts": return LazyArray(data).with_parts(PART_SHAPE) - if flavour == "zarr": + if flavor == "zarr": return LazyArray(make_zarr_source()) - if flavour == "zarr-misaligned": + if flavor == "zarr-misaligned": # Parts that deliberately straddle the zarr array's own chunks. return LazyArray(make_zarr_source()).with_parts((4, 3, 3)) - raise AssertionError(f"unknown source flavour {flavour!r}") + raise AssertionError(f"unknown source flavor {flavor!r}") -FLAVOURS = [ +FLAVORS = [ "numpy-whole", "numpy-explicit-parts", "numpy-uniform-parts", @@ -223,7 +223,7 @@ def make_source(flavour: str) -> LazyArray: ] -@pytest.fixture(params=FLAVOURS) +@pytest.fixture(params=FLAVORS) def source(request: pytest.FixtureRequest) -> LazyArray: return make_source(request.param) @@ -576,8 +576,8 @@ def _random_chain(rng: np.random.Generator) -> list[tuple[str, tuple[Any, ...]]] return chain -@pytest.mark.parametrize("flavour", FLAVOURS) -def test_random_chains_match_numpy(flavour: str) -> None: +@pytest.mark.parametrize("flavor", FLAVORS) +def test_random_chains_match_numpy(flavor: str) -> None: """A seeded sweep of composed chains, under every partitioning. The partitioning invariant as a property: `result()` is identical whatever @@ -585,11 +585,11 @@ def test_random_chains_match_numpy(flavour: str) -> None: the source's own. """ rng = np.random.default_rng(20260730) - source = make_source(flavour) + source = make_source(flavor) partitionings: list[Any] = [None, (2, 2, 2), (7, 5, 4), ((4, 3), (1, 3, 1), (3, 1))] - # Reads against a real store cost more per chain; the NumPy flavours carry + # Reads against a real store cost more per chain; the NumPy flavors carry # the bulk of the sweep and exercise the identical code path. - n_chains = 120 if flavour.startswith("zarr") else 400 + n_chains = 120 if flavor.startswith("zarr") else 400 for _ in range(n_chains): chain = _random_chain(rng) @@ -601,15 +601,15 @@ def test_random_chains_match_numpy(flavour: str) -> None: for mode, selection in chain: view = _apply_view(view, mode, selection) - assert view.shape == expected.shape, f"{flavour}: {chain}" + assert view.shape == expected.shape, f"{flavor}: {chain}" np.testing.assert_array_equal( - np.asarray(view.result()), np.asarray(expected), err_msg=f"{flavour}: {chain}" + np.asarray(view.result()), np.asarray(expected), err_msg=f"{flavor}: {chain}" ) for parts in partitionings: np.testing.assert_array_equal( np.asarray(view.with_parts(parts).result()), np.asarray(expected), - err_msg=f"{flavour} parts={parts}: {chain}", + err_msg=f"{flavor} parts={parts}: {chain}", ) @@ -839,7 +839,7 @@ def test_a_query_bounding_box_is_only_a_hull() -> None: # --------------------------------------------------------------------------- -@pytest.mark.parametrize("flavour", FLAVOURS) +@pytest.mark.parametrize("flavor", FLAVORS) @pytest.mark.parametrize( "build", [ @@ -856,10 +856,10 @@ def test_a_query_bounding_box_is_only_a_hull() -> None: ids=["identity", "basic", "oindex", "vindex", "reversed", "reversed-bounded"], ) def test_parts_tile_the_view_exactly_and_disjointly( - flavour: str, build: Callable[[LazyArray], LazyArray] + flavor: str, build: Callable[[LazyArray], LazyArray] ) -> None: """Assembling the parts reproduces `result()`; the placements cover with no overlap.""" - view = build(make_source(flavour)) + view = build(make_source(flavor)) expected = np.asarray(view.result()) assembled = np.zeros(view.shape, dtype=view.dtype) From 934ea2980eefc1bb392adc66597e175b11165e07 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 12:31:41 +0200 Subject: [PATCH 06/32] ci(zarr-indexing): scoped lint ignores for the deliberate blind excepts 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 --- packages/zarr-indexing/pyproject.toml | 8 ++++++++ packages/zarr-indexing/src/zarr_indexing/lazy_array.py | 6 ++++++ 2 files changed, 14 insertions(+) diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index 21ba4ef7e7..b996833d9c 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -75,6 +75,14 @@ packages = ["src/zarr_indexing"] extend = "../../pyproject.toml" target-version = "py312" +[tool.ruff.lint.per-file-ignores] +# Chunk discovery and __dask_tokenize__ deliberately catch Exception: a +# foreign source's attributes may fail arbitrarily and discovery must degrade +# to "no information"; a token call must never raise. Configured here (not as +# noqa comments) because the pinned pre-commit ruff and the floating CI ruff +# disagree on whether these rules fire, and RUF100 strips the comments. +"src/zarr_indexing/lazy_array.py" = ["BLE001", "S110"] + [tool.pytest.ini_options] minversion = "7" testpaths = ["tests"] diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 19439727fd..3ece61c40a 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -177,6 +177,8 @@ def _namespace(x: Any) -> Any: return None try: return getter() + # A foreign namespace may fail in arbitrary ways; discovery must degrade + # to "no information", never raise. except Exception: # pragma: no cover - a namespace that refuses to load return None @@ -231,6 +233,8 @@ def _read_source_attribute(array: Any, name: str) -> Any: """ try: return getattr(array, name, None) + # Guarded properties (e.g. zarr's LazyViewError) and broken foreign + # attributes may raise anything; discovery must degrade to None. except Exception: return None @@ -597,6 +601,7 @@ def _wrapped_token(array: Any) -> Any: if hook is not None: try: return hook() + # A token must never raise; fall through to the structural fallback. except Exception: # pragma: no cover - a hook that refuses to run pass try: @@ -620,6 +625,7 @@ def _wrapped_token(array: Any) -> Any: return structural try: contents = np.ascontiguousarray(array) + # A token must never raise; an unreadable source keeps the structural token. except Exception: return structural return (*structural, hashlib.sha256(contents.tobytes()).hexdigest()) From bffec52c50180baffa9a23f167419570da34cecb Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 12:35:45 +0200 Subject: [PATCH 07/32] ci(zarr-indexing): pin ruff in the lint job and justfile Mirrors the main-branch pin (d-v-b#271) so this PR's workflow runs the same ruff version; bump together with the pyproject pin. Assisted-by: ClaudeCode:claude-fable-5 --- .github/workflows/zarr-indexing.yml | 3 ++- packages/zarr-indexing/justfile | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.github/workflows/zarr-indexing.yml b/.github/workflows/zarr-indexing.yml index 2106b10916..8685521a0b 100644 --- a/.github/workflows/zarr-indexing.yml +++ b/.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 diff --git a/packages/zarr-indexing/justfile b/packages/zarr-indexing/justfile index 20e1b2b837..51cf8f0cad 100644 --- a/packages/zarr-indexing/justfile +++ b/packages/zarr-indexing/justfile @@ -15,9 +15,10 @@ default: test *args: uv run --project ../.. --group test --with-editable . python -m pytest tests {{ args }} -# Lint with the same invocation CI uses +# Lint with the same invocation CI uses. Ruff is pinned to the repo-wide +# version (see pyproject.toml [dependency-groups] docs); bump together. lint: - uvx ruff check . + uvx ruff@0.15.22 check . # Type-check the package sources typecheck: From cfd84e77c1681dd116b23165027c380b89ca08a8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 13:57:45 +0200 Subject: [PATCH 08/32] docs: add lazy-indexing examples for NumPy and Dask 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 --- .../user-guide/examples/lazy_indexing_dask.md | 7 + .../examples/lazy_indexing_numpy.md | 7 + examples/lazy_indexing_dask/README.md | 35 +++++ .../lazy_indexing_dask/lazy_indexing_dask.py | 117 ++++++++++++++++ examples/lazy_indexing_numpy/README.md | 36 +++++ .../lazy_indexing_numpy.py | 130 ++++++++++++++++++ mkdocs.yml | 2 + 7 files changed, 334 insertions(+) create mode 100644 docs/user-guide/examples/lazy_indexing_dask.md create mode 100644 docs/user-guide/examples/lazy_indexing_numpy.md create mode 100644 examples/lazy_indexing_dask/README.md create mode 100644 examples/lazy_indexing_dask/lazy_indexing_dask.py create mode 100644 examples/lazy_indexing_numpy/README.md create mode 100644 examples/lazy_indexing_numpy/lazy_indexing_numpy.py diff --git a/docs/user-guide/examples/lazy_indexing_dask.md b/docs/user-guide/examples/lazy_indexing_dask.md new file mode 100644 index 0000000000..d4b06eb316 --- /dev/null +++ b/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" +``` diff --git a/docs/user-guide/examples/lazy_indexing_numpy.md b/docs/user-guide/examples/lazy_indexing_numpy.md new file mode 100644 index 0000000000..f835d55c4a --- /dev/null +++ b/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" +``` diff --git a/examples/lazy_indexing_dask/README.md b/examples/lazy_indexing_dask/README.md new file mode 100644 index 0000000000..ae1e1e08d6 --- /dev/null +++ b/examples/lazy_indexing_dask/README.md @@ -0,0 +1,35 @@ +# 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. + +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 + +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. + +## 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 +``` diff --git a/examples/lazy_indexing_dask/lazy_indexing_dask.py b/examples/lazy_indexing_dask/lazy_indexing_dask.py new file mode 100644 index 0000000000..939dc2004e --- /dev/null +++ b/examples/lazy_indexing_dask/lazy_indexing_dask.py @@ -0,0 +1,117 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = [ +# "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main", +# "zarr-indexing @ git+https://github.com/zarr-developers/zarr-python.git@main#subdirectory=packages/zarr-indexing", +# "dask[array]==2025.3.0", +# "numpy==2.4.3", +# "pytest==9.0.2" +# ] +# /// +# + +""" +Demonstrate using zarr_indexing.LazyArray with Dask +""" + +import sys + +import dask +import dask.array as da +import numpy as np +import pytest +from dask.base import tokenize +from zarr_indexing import LazyArray + +import zarr + + +@pytest.fixture +def source() -> zarr.Array: + """A chunked Zarr array to wrap.""" + array = zarr.create_array(store={}, shape=(40, 30), chunks=(10, 10), dtype="i4") + array[:] = np.arange(40 * 30).reshape(40, 30) + return array + + +def test_from_array(source: zarr.Array) -> None: + """Hand a LazyArray to `dask.array.from_array`.""" + lazy = LazyArray(source) + + # `from_array` needs `shape`, `dtype`, and `__getitem__`, which the wrapper + # provides. Each Dask block reads its own region through the wrapper. + array = da.from_array(lazy, chunks=(10, 10)) + print(array) + assert np.array_equal(array.compute(scheduler="threads"), source[:]) + + # A view works the same way, and its shape is the shape of the selection. + view = LazyArray(source).lazy[5:35, 3:27] + array = da.from_array(view, chunks=(10, 10)) + assert array.shape == (30, 24) + assert np.array_equal(array.compute(scheduler="threads"), source[5:35, 3:27]) + + +def test_parts_as_tasks(source: zarr.Array) -> None: + """Build one task per partition and compute them in parallel.""" + view = LazyArray(source).lazy[5:35, 3:27] + + # The partitioning is discovered from the wrapped array's chunks, so each + # partition of the view lies within one stored chunk. + parts = list(view.parts()) + print(f"{len(parts)} parts for a {view.shape} view of a {source.shape} array") + + # A partition carries a sub-view to resolve and where its result belongs, so + # the reads are independent and the placement needs no coordination. + @dask.delayed + def read(part: object) -> np.ndarray: + return part.array.result() + + blocks = dask.compute(*[read(part) for part in parts], scheduler="threads") + + result = np.empty(view.shape, dtype=view.dtype) + for part, block in zip(parts, blocks, strict=True): + result[part.out_selection] = block + assert np.array_equal(result, source[5:35, 3:27]) + + # `is_complete` reports whether a partition covers its whole partition of + # the base array, which a writer uses to choose between overwriting a chunk + # and reading it first. + complete = [part.box for part in parts if part.is_complete] + print(f"{len(complete)} of {len(parts)} parts cover their chunk completely") + + +def test_tokenize(source: zarr.Array) -> None: + """Deterministic tokens let Dask cache and deduplicate work.""" + lazy = LazyArray(source) + + # Two wrappers over the same array and the same selection are the same task + # to Dask, whether or not they are the same Python object. + assert tokenize(lazy) == tokenize(LazyArray(source)) + assert tokenize(lazy.lazy[0:10]) == tokenize(LazyArray(source).lazy[0:10]) + + # Different selections are different tasks. + assert tokenize(lazy.lazy[0:10]) != tokenize(lazy.lazy[10:20]) + + # Selections that describe the same region are the same task, however they + # were composed. + assert tokenize(lazy.lazy[0:20].lazy[5:10]) == tokenize(lazy.lazy[5:10]) + + +if __name__ == "__main__": + # Run the example with printed output, and a dummy pytest configuration file specified. + # Without the dummy configuration file, at test time pytest will attempt to use the + # configuration file in the project root, which will error because Zarr is using some + # plugins that are not installed in this example. + sys.exit( + pytest.main( + [ + "-s", + __file__, + f"-c {__file__}", + # Suppress: "PytestAssertRewriteWarning: Module already imported so + # cannot be rewritten; zarr" + "-W", + "ignore::pytest.PytestAssertRewriteWarning", + ] + ) + ) diff --git a/examples/lazy_indexing_numpy/README.md b/examples/lazy_indexing_numpy/README.md new file mode 100644 index 0000000000..f14933fe22 --- /dev/null +++ b/examples/lazy_indexing_numpy/README.md @@ -0,0 +1,36 @@ +# Lazy Indexing a NumPy Array + +This example demonstrates how to wrap an array in `zarr_indexing.LazyArray` and +index it without reading data. + +The example shows how to: + +- Wrap a NumPy array and read the forwarded `shape`, `dtype`, and `ndim` +- Compose selections through `.lazy[...]`, `.lazy.oindex[...]`, and + `.lazy.vindex[...]`, and materialize the composed view once with `result()` +- Tell a box selection (slices and integers, described by an interval and a step + per dimension) from a query selection (points gathered through an index array) + using `is_box`, `bounding_box()`, and `strides()` +- Declare a partitioning with `with_parts()`, iterate it with `parts()`, and + assemble a result from the partitions + +`LazyArray` wraps any object exposing `shape`, `dtype`, and `__getitem__`, so the +same API applies to a Zarr array, and the partitioning is then discovered from +the array's chunks. The Dask example covers that case. + +## 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_numpy/lazy_indexing_numpy.py +``` + +Alternatively, run it with plain Python, in which case you must first install +`zarr-indexing`, `numpy`, and `pytest` yourself: + +```bash +python examples/lazy_indexing_numpy/lazy_indexing_numpy.py +``` diff --git a/examples/lazy_indexing_numpy/lazy_indexing_numpy.py b/examples/lazy_indexing_numpy/lazy_indexing_numpy.py new file mode 100644 index 0000000000..3ae02e9182 --- /dev/null +++ b/examples/lazy_indexing_numpy/lazy_indexing_numpy.py @@ -0,0 +1,130 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = [ +# "zarr-indexing @ git+https://github.com/zarr-developers/zarr-python.git@main#subdirectory=packages/zarr-indexing", +# "numpy==2.4.3", +# "pytest==9.0.2" +# ] +# /// +# + +""" +Demonstrate lazy indexing over a plain NumPy array with zarr_indexing.LazyArray +""" + +import sys + +import numpy as np +import pytest +from zarr_indexing import LazyArray + + +def test_wrap_and_compose() -> None: + """Wrap an array, compose selections without reading, then materialize once.""" + data = np.arange(12 * 8).reshape(12, 8) + lazy = LazyArray(data) + + # The wrapper forwards the attributes an array consumer expects. + assert lazy.shape == (12, 8) + assert lazy.dtype == data.dtype + assert lazy.ndim == 2 + + # `.lazy[...]` returns another LazyArray. No element of `data` is read. + view = lazy.lazy[2:10, ::2] + print(view) + assert view.shape == (8, 4) + + # Selections compose. Each step narrows the view; still nothing is read. + smaller = view.lazy[1:5, 1:3] + + # `result()` performs the read. NumPy is the reference for the whole chain. + assert np.array_equal(smaller.result(), data[2:10, ::2][1:5, 1:3]) + + # Selections use positional NumPy semantics: indices count from zero within + # the current view, and negative indices count from the end. + assert np.array_equal(lazy.lazy[-1].result(), data[-1]) + assert np.array_equal(lazy.lazy[::-1].result(), data[::-1]) + + # Orthogonal and vectorized indexing are available under the same accessor. + rows = np.array([9, 1, 4]) + assert np.array_equal(lazy.lazy.oindex[rows, :].result(), data[rows, :]) + cols = np.array([0, 3, 7]) + assert np.array_equal(lazy.lazy.vindex[rows, cols].result(), data[rows, cols]) + + # A LazyArray is also an ordinary duck array: __getitem__ reads immediately, + # and np.asarray materializes the view. + assert np.array_equal(lazy[2:4, 0], data[2:4, 0]) + assert np.array_equal(np.asarray(view), data[2:10, ::2]) + + +def test_box_and_query_selections() -> None: + """Distinguish selections that describe a region from selections that gather points.""" + data = np.arange(12 * 8).reshape(12, 8) + lazy = LazyArray(data) + + # A box selection is built from slices and integers alone. It is described + # completely by an interval and a step per dimension, so a consumer can + # serve it as one strided read. + box = lazy.lazy[2:10, ::2] + print(f"box: is_box={box.is_box} bounding_box={box.bounding_box()} strides={box.strides()}") + assert box.is_box + assert box.bounding_box() == ((2, 10), (0, 7)) + assert box.strides() == (1, 2) + + # A query selection gathers points through an index array. Its coordinates + # are a lookup table, so `strides()` is undefined and `bounding_box()` is + # the hull of the points rather than an exact description. + query = lazy.lazy.oindex[np.array([9, 1, 4]), :] + print(f"query: is_box={query.is_box} bounding_box={query.bounding_box()}") + assert not query.is_box + assert query.strides() is None + assert query.bounding_box() == ((1, 10), (0, 8)) + + # Composing a box onto a query keeps it a query. + assert not query.lazy[0:2, 0:2].is_box + + +def test_parts() -> None: + """Iterate the partitions a view covers, and assemble the result from them.""" + data = np.arange(12 * 8).reshape(12, 8) + + # A plain NumPy array declares no partitioning, so `with_parts` states one. + # Partitioning changes the granularity of reads, never the result. + lazy = LazyArray(data).with_parts((4, 4)) + view = lazy.lazy[2:10, ::2] + + parts = list(view.parts()) + print(f"{len(parts)} parts") + for part in parts[:2]: + print(f" base_coords={part.base_coords} box={part.box} complete={part.is_complete}") + + # Each part carries a sub-view of its own, where that sub-view lands in the + # result, and whether it covers its partition completely. Resolving the + # parts and placing them is what `result()` does. + assembled = np.empty(view.shape, dtype=view.dtype) + for part in parts: + assembled[part.out_selection] = part.array.result() + assert np.array_equal(assembled, view.result()) + + # The partitioning is a read strategy, so a different one gives the same data. + assert np.array_equal(LazyArray(data).with_parts((5, 3)).lazy[2:10, ::2].result(), assembled) + + +if __name__ == "__main__": + # Run the example with printed output, and a dummy pytest configuration file specified. + # Without the dummy configuration file, at test time pytest will attempt to use the + # configuration file in the project root, which will error because Zarr is using some + # plugins that are not installed in this example. + sys.exit( + pytest.main( + [ + "-s", + __file__, + f"-c {__file__}", + # Suppress: "PytestAssertRewriteWarning: Module already imported so + # cannot be rewritten; zarr" + "-W", + "ignore::pytest.PytestAssertRewriteWarning", + ] + ) + ) diff --git a/mkdocs.yml b/mkdocs.yml index 6a0d94052e..d388d86231 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -30,6 +30,8 @@ nav: - user-guide/glossary.md - Examples: - user-guide/examples/custom_dtype.md + - user-guide/examples/lazy_indexing_dask.md + - user-guide/examples/lazy_indexing_numpy.md - user-guide/examples/rectilinear_chunks.md - user-guide/examples/codec_pipeline_performance.md - user-guide/examples/sharding_coalescing.md From 3a4b1c792a3e5143716b19e7a688d81828b776cc Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 14:22:52 +0200 Subject: [PATCH 09/32] docs: compare dask task graphs with fused transforms in the dask example 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 --- examples/lazy_indexing_dask/README.md | 20 +++++- .../lazy_indexing_dask/lazy_indexing_dask.py | 64 +++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/examples/lazy_indexing_dask/README.md b/examples/lazy_indexing_dask/README.md index ae1e1e08d6..f2294d6327 100644 --- a/examples/lazy_indexing_dask/README.md +++ b/examples/lazy_indexing_dask/README.md @@ -1,7 +1,8 @@ # 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. +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: @@ -12,11 +13,28 @@ The example shows how to: - 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 diff --git a/examples/lazy_indexing_dask/lazy_indexing_dask.py b/examples/lazy_indexing_dask/lazy_indexing_dask.py index 939dc2004e..5f8d93d6c2 100644 --- a/examples/lazy_indexing_dask/lazy_indexing_dask.py +++ b/examples/lazy_indexing_dask/lazy_indexing_dask.py @@ -15,6 +15,7 @@ """ import sys +import time import dask import dask.array as da @@ -97,6 +98,69 @@ def test_tokenize(source: zarr.Array) -> None: assert tokenize(lazy.lazy[0:20].lazy[5:10]) == tokenize(lazy.lazy[5:10]) +def test_indexing_only_workload() -> None: + """Compare an accumulating task graph with a fused transform. + + Dask records each indexing operation as another graph layer, and slices the + chunk grid to build it, so composing selections costs time proportional to + the number of selections and the number of chunks. `LazyArray` composes each + selection into the single transform it already holds, so the cost of + composing does not grow with the depth of the chain, and reading resolves + that one transform rather than walking a graph. + + Timings are printed rather than asserted, since they depend on the machine. + """ + data = np.zeros((2000, 4), dtype="i4") # 2000 chunks, one row each + + def dask_chain(depth: int) -> da.Array: + array = da.from_array(data, chunks=(1, 4)) + for _ in range(depth): + array = array[1:] + return array + + def lazy_chain(depth: int) -> LazyArray: + view = LazyArray(data) + for _ in range(depth): + view = view.lazy[1:] + return view + + # Read once through each path first, so the timings below exclude the cost + # of importing and initializing the machinery. + dask_chain(1)[:2].compute(scheduler="synchronous") + lazy_chain(1).lazy[:2].result() + + header = ( + f"{'selections':>10} {'dask compose':>13} {'dask read':>10} {'layers':>7}" + f" {'LazyArray compose':>18} {'LazyArray read':>15}" + ) + print(header) + for depth in (1, 5, 20): + start = time.perf_counter() + chained = dask_chain(depth) + dask_compose = time.perf_counter() - start + + start = time.perf_counter() + from_dask = chained[:2].compute(scheduler="synchronous") + dask_read = time.perf_counter() - start + + start = time.perf_counter() + view = lazy_chain(depth) + lazy_compose = time.perf_counter() - start + + start = time.perf_counter() + from_lazy = view.lazy[:2].result() + lazy_read = time.perf_counter() - start + + # Both paths describe the same selection, so they read the same data. + assert np.array_equal(from_dask, from_lazy) + + layers = len(chained.__dask_graph__().layers) + print( + f"{depth:>10} {dask_compose * 1e3:>12.2f}ms {dask_read * 1e3:>9.2f}ms {layers:>7}" + f" {lazy_compose * 1e3:>17.3f}ms {lazy_read * 1e3:>14.3f}ms" + ) + + if __name__ == "__main__": # Run the example with printed output, and a dummy pytest configuration file specified. # Without the dummy configuration file, at test time pytest will attempt to use the From 681d94744c591147b618c63e96b47e4b211c7459 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 14:47:34 +0200 Subject: [PATCH 10/32] test: run examples against this repository's local packages MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- tests/test_examples.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 9f8085e8c2..607207b453 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -18,6 +18,15 @@ # This is the absolute path to the local Zarr installation. Moving this test to a different directory will break it. ZARR_PROJECT_PATH = Path(".").absolute() +# Packages that live in this repository. An example declares them by their +# published name, so that the example works for a reader who just runs it, and +# the test substitutes the local checkout, so that CI exercises the code in the +# working tree rather than a released or `main` version. +LOCAL_PACKAGES: Final = { + "zarr": ZARR_PROJECT_PATH, + "zarr-indexing": ZARR_PROJECT_PATH / "packages" / "zarr-indexing", +} + def set_dep(script: str, dependency: str) -> str: """ @@ -54,11 +63,15 @@ def set_dep(script: str, dependency: str) -> str: def resave_script(source_path: Path, dest_path: Path) -> None: """ - Read a script from source_path and save it to dest_path after inserting the absolute path to the - local Zarr project directory in the PEP-723 header. + Read a script from source_path and save it to dest_path after inserting the absolute paths to + the local copies of this repository's packages in the PEP-723 header. + + A dependency the script does not declare is left alone, so each example only gets the local + packages it actually uses. """ - source_text = source_path.read_text() - dest_text = set_dep(source_text, f"zarr @ file:///{ZARR_PROJECT_PATH}") + dest_text = source_path.read_text() + for name, path in LOCAL_PACKAGES.items(): + dest_text = set_dep(dest_text, f"{name} @ file:///{path}") dest_path.write_text(dest_text) From fb29f79bf1763a51ab5acbba2d594485eaae3d83 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 15:30:14 +0200 Subject: [PATCH 11/32] feat(zarr-indexing): negotiate what indexing a source supports MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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 --- packages/zarr-indexing/changes/267.feature.md | 4 + packages/zarr-indexing/docs/api/index.md | 4 + packages/zarr-indexing/docs/api/support.md | 5 + packages/zarr-indexing/docs/design-notes.md | 19 + packages/zarr-indexing/mkdocs.yml | 1 + .../src/zarr_indexing/__init__.py | 8 + .../src/zarr_indexing/lazy_array.py | 363 +++++++++++++++++- .../src/zarr_indexing/support.py | 164 ++++++++ .../zarr-indexing/tests/test_lazy_array.py | 290 +++++++++++++- 9 files changed, 836 insertions(+), 22 deletions(-) create mode 100644 packages/zarr-indexing/docs/api/support.md create mode 100644 packages/zarr-indexing/src/zarr_indexing/support.py diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 99a2a3f098..86ee1ce61f 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -12,3 +12,7 @@ Fixed the domain layout and partitioned resolution of a vectorized selection who 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: a NumPy array or an object carrying zarr's `oindex`/`vindex` pair reads as `VECTORIZED`, a source's own `__zarr_indexing_support__` attribute is honored when it holds an `IndexingSupport`, and everything else is `BASIC`, 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. diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index 17cb401863..fe112adb6b 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -38,6 +38,10 @@ and the wire format built on top of it. gives any array-API-like array (NumPy, zarr, CuPy, …) a `.lazy` accessor for TensorStore-style deferred indexing, plus `Partition` and `parts()` / `with_parts()`, which determine the boxes a read is broken into +- [`zarr_indexing.support`](support.md) — `IndexingSupport`, the four levels of + indexing a wrapped array may accept, and the detection that picks one, which + together decide how much of a selection `LazyArray` hands to its source and + how much it finishes in memory - [`zarr_indexing.boundary`](boundary.md) — the translation between NumPy's positional dialect and the transform algebra's literal coordinates diff --git a/packages/zarr-indexing/docs/api/support.md b/packages/zarr-indexing/docs/api/support.md new file mode 100644 index 0000000000..c19527d8a4 --- /dev/null +++ b/packages/zarr-indexing/docs/api/support.md @@ -0,0 +1,5 @@ +--- +title: support +--- + +::: zarr_indexing.support diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index bdfc7adf40..f73272ed81 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -148,6 +148,25 @@ rectangular-only view is a plausible next step, but it should be introduced by a consumer that needs the guarantee in its signatures rather than speculatively; `is_box` is the runtime check until then. +## Related work + +The closest neighbour to this package in the Python ecosystem is xarray's +[`xarray/core/indexing.py`](https://github.com/pydata/xarray/blob/main/xarray/core/indexing.py). +It solves the same problem — carrying a selection around as a value, and +deciding how much of it a given backend can be asked to perform — and the +`IndexingSupport` taxonomy here (`BASIC`, `OUTER`, `OUTER_1VECTOR`, +`VECTORIZED`) is taken from it, member names and meanings included, so that a +reader who knows one knows the other. The split of a selection into a part +pushed to the source and a part finished with NumPy follows xarray's +`decompose_indexer`, and the heuristic for choosing which axis keeps its +coordinate array when only one can is xarray's. + +This package is a standalone implementation rather than a port: no code is +shared, the selection is carried as an `IndexTransform` rather than as xarray's +explicit indexer classes, and the decomposition is applied per chunk partition +as well as per array, so a source is asked for one key per box rather than one +key per read. + ## Current scope Negative steps are supported as of ndsel 1.0-draft.2: `a[::-1]` reverses, one diff --git a/packages/zarr-indexing/mkdocs.yml b/packages/zarr-indexing/mkdocs.yml index d341c127a8..9a95be87a2 100644 --- a/packages/zarr-indexing/mkdocs.yml +++ b/packages/zarr-indexing/mkdocs.yml @@ -24,6 +24,7 @@ nav: - ' zarr_indexing.chunk_resolution': api/chunk_resolution.md - ' zarr_indexing.grid': api/grid.md - ' zarr_indexing.lazy_array': api/lazy_array.md + - ' zarr_indexing.support': api/support.md - ' zarr_indexing.boundary': api/boundary.md - ' zarr_indexing.json': api/json.md - ' zarr_indexing.messages': api/messages.md diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index c152d7d6c6..c1ef695d7a 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -47,6 +47,11 @@ from zarr_indexing.lazy_array import LazyArray, Partition from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.support import ( + IndexingSupport, + infer_indexing_support, + resolve_indexing_support, +) from zarr_indexing.transform import ( IndexTransform, array_map_dependent_axis, @@ -65,6 +70,7 @@ "IndexDomainJSON", "IndexTransform", "IndexTransformJSON", + "IndexingSupport", "LazyArray", "NdselError", "OutputIndexMap", @@ -78,9 +84,11 @@ "index_domain_to_json", "index_transform_from_json", "index_transform_to_json", + "infer_indexing_support", "iter_chunk_transforms", "normalize_ndsel", "parse_ndsel", + "resolve_indexing_support", "selection_to_transform", "sub_transform_to_selections", "transform_from_canonical", diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 3ece61c40a..7672478cbe 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -28,7 +28,7 @@ out[part.out_selection] = part.array.result() ``` -Each box is therefore read once, with basic slicing, and the selection is +Each box is therefore read once, and whatever the source could not do itself is applied to the block in memory. The partitioning is discovered from the wrapped array at construction — first @@ -36,9 +36,8 @@ `chunks`, read as per-axis sizes if its entries are sequences and as a uniform box shape if they are integers. Those attribute names belong to the wrapped array; this API refers only to parts. An array that advertises neither gets a -single whole-array part, and resolving it lowers the whole view to one pass of -array operations (`take`, basic slicing, a flat gather) rather than -materializing a block first. +single whole-array part, and resolving it lowers the whole view in one pass +rather than materializing a block first. `with_parts` replaces the partitioning without touching the data or the view: @@ -53,6 +52,29 @@ useful (to bound peak memory, or to batch small reads); they cost extra I/O but do not affect correctness. +What the source is asked for +--------------------------- +Sources differ in what their indexing surface accepts, so a read is negotiated +rather than assumed. Each wrapper carries an +[`IndexingSupport`](support.md) level, and every request is split in two: the +largest part of the selection that level can express is asked of the source in a +single call — through `oindex` or `vindex` when the source has them — and the +remainder is applied to the block that comes back, with NumPy. + +The level is detected at construction (a NumPy array or zarr's +`oindex`/`vindex` pair reads as `VECTORIZED`, a source's own +`__zarr_indexing_support__` declaration is honored, and anything else is +`BASIC`, which is the only assumption that is always safe). +`with_indexing_support` overrides the detection: + +```python +view.with_indexing_support(IndexingSupport.OUTER) +``` + +The level changes 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. + Boxes and queries ----------------- A selection is either **rectangular** — an interval and a stride per dimension, @@ -111,7 +133,7 @@ Resolving a partitioned view builds its output buffer with NumPy, because the resolver's scatter indices are host-side integer arrays. Wrapping a device array that advertises parts therefore returns host memory. A single whole-array part -lowers directly and stays in the wrapped array's own namespace. +skips that buffer and stays in the wrapped array's own namespace. """ from __future__ import annotations @@ -136,7 +158,12 @@ ) from zarr_indexing.grid import EdgeDimensionGrid, dimension_grids_from_chunks from zarr_indexing.json import transform_to_canonical -from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.support import ( + IndexingSupport, + has_accessor, + resolve_indexing_support, +) from zarr_indexing.transform import ( IndexTransform, array_map_dependent_axis, @@ -488,6 +515,217 @@ def _lower_correlated(array: Any, transform: IndexTransform) -> Any: ) +# --------------------------------------------------------------------------- # +# Source-capability negotiation +# --------------------------------------------------------------------------- # +# +# `_lower` is exact but assumes the array under it accepts anything NumPy +# accepts. Most sources do not, and the ones that do would rather be asked for +# what is wanted in one request than be walked axis by axis. So a read is split +# in two: a *pushed key*, the largest part of the selection the source can +# express, and a *residual transform*, whatever is left, which `_lower` applies +# to the block that comes back. The split is the same idea as xarray's +# `decompose_indexer`, expressed over `IndexTransform` rather than over +# xarray's explicit indexer classes. +# +# The pushed key holds exactly one selector per storage dimension, always a +# slice or a one-dimensional integer array and never a bare integer, so the +# block has one axis per storage dimension whatever the level. That invariant is +# what lets the residual be built by rewriting output maps in place: only the +# coordinates change, from storage coordinates to positions within the block. + + +def _push_slice_for_dimension_map( + m: DimensionMap, transform: IndexTransform +) -> tuple[slice, DimensionMap]: + """The positive-step slice covering a `DimensionMap`, and its block-local map. + + A negative step is read forwards and reversed by the residual: a source is + only ever asked for a slice that walks upwards, which is the one form every + array-like agrees on. + """ + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = max(transform.domain.exclusive_max[d], lo) + if hi == lo: + return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) + first = m.offset + m.stride * lo + last = m.offset + m.stride * (hi - 1) + if m.stride > 0: + return ( + slice(first, last + 1, m.stride), + DimensionMap(input_dimension=d, offset=-lo, stride=1), + ) + # Descending: the block holds the same coordinates in ascending order, so + # the residual walks it backwards from the last block position. + return ( + slice(last, first + 1, -m.stride), + DimensionMap(input_dimension=d, offset=hi - 1, stride=-1), + ) + + +def _push_gain(coords: np.ndarray[Any, np.dtype[np.intp]]) -> float: + """How much a slab read of `coords` would over-read, as a ratio. + + xarray's heuristic for choosing which axis keeps its array when only one + can: the span of the coordinates divided by how many distinct ones there + are. An axis whose coordinates are dense and contiguous scores 1 and loses + nothing to a slab; a sparse axis scores high and is the one worth spending + the single array slot on. + """ + if coords.size == 0: + return 0.0 + span = int(coords.max()) - int(coords.min()) + 1 + return span / int(np.unique(coords).size) + + +def _array_push_capacity(support: IndexingSupport, has_oindex: bool, wanted: int) -> int: + """How many axes may carry an integer array in one request to the source.""" + if support is IndexingSupport.BASIC: + return 0 + if support is IndexingSupport.OUTER_1VECTOR: + return 1 + # OUTER and VECTORIZED both accept several array axes, but only through an + # `oindex` accessor: a bare `__getitem__` key with two arrays means an outer + # product to HDF5 and a correlated gather to NumPy, and nothing about the + # source says which reading applies. One array is the key both readings + # agree on, so that is the fallback. + return wanted if has_oindex else 1 + + +def _decompose( + transform: IndexTransform, support: IndexingSupport, has_oindex: bool +) -> tuple[tuple[Any, ...], IndexTransform, int]: + """Split `transform` into a key for the source and a transform for the block. + + Returns `(key, residual, n_arrays)`. `key` has one selector per storage + dimension; `n_arrays` is how many of them are integer arrays, which decides + whether the request goes to `oindex` or to `__getitem__`. Applying `residual` + to the block the key returns is equivalent to applying `transform` to the + source. + """ + outputs = transform.output + coords: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = { + d: (m.offset + m.stride * m.index_array).astype(np.intp) + for d, m in enumerate(outputs) + if isinstance(m, ArrayMap) + } + # An empty axis carries no coordinates to push and needs no array slot. + array_axes = [d for d, c in coords.items() if c.size > 0] + capacity = _array_push_capacity(support, has_oindex, len(array_axes)) + if capacity >= len(array_axes): + pushed_axes = set(array_axes) + else: + ranked = sorted(array_axes, key=lambda d: _push_gain(coords[d]), reverse=True) + pushed_axes = set(ranked[:capacity]) + + key: list[Any] = [] + residual: list[OutputIndexMap] = [] + for d, m in enumerate(outputs): + if isinstance(m, ConstantMap): + # A slice rather than the integer: every axis survives into the + # block, so residual output `d` addresses block axis `d`. + key.append(slice(m.offset, m.offset + 1, 1)) + residual.append(ConstantMap(offset=0)) + elif isinstance(m, DimensionMap): + pushed, local = _push_slice_for_dimension_map(m, transform) + key.append(pushed) + residual.append(local) + else: + full = coords[d] + if full.size == 0: + key.append(slice(0, 0, 1)) + local_index = full + elif d in pushed_axes: + # Sorted and deduplicated: a source is never asked for the same + # coordinate twice, and the residual looks positions back up. + unique = np.unique(full) + key.append(unique) + local_index = np.searchsorted(unique, full).astype(np.intp) + else: + # No array slot for this axis: read the slab it spans and + # gather from that instead. + origin = int(full.min()) + key.append(slice(origin, int(full.max()) + 1, 1)) + local_index = (full - origin).astype(np.intp) + residual.append(ArrayMap(index_array=local_index, input_dimension=m.input_dimension)) + return ( + tuple(key), + IndexTransform(domain=transform.domain, output=tuple(residual)), + len(pushed_axes), + ) + + +def _point_key(transform: IndexTransform) -> tuple[Any, ...] | None: + """The whole selection as one flat coordinate array per axis, or None. + + Defined only when every storage dimension is addressed by a correlated + `ArrayMap` or fixed by a `ConstantMap`, which is what a plain `vindex[i, j, + k]` composes to. The points are then exactly the cells the view wants, in + the view's own order, so a source that can gather points has nothing left to + do — whereas the outer superset those same points span can be very much + larger than the points themselves. + + A selection with a sliced axis is not expressible this way without + enumerating the slice, so it goes through `_decompose` instead. + """ + shape = transform.domain.shape + if math.prod(shape) == 0: + return None + key: list[Any] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + key.append(np.full(math.prod(shape), m.offset, dtype=np.intp)) + elif isinstance(m, ArrayMap) and m.input_dimension is None: + coords = (m.offset + m.stride * m.index_array).astype(np.intp) + key.append(np.broadcast_to(coords, shape).reshape(-1)) + else: + return None + return tuple(key) + + +def _shift_key(key: tuple[Any, ...], window: tuple[slice, ...] | None) -> tuple[Any, ...]: + """Rebase a key built in part-local coordinates onto the wrapped array.""" + if window is None: + return key + shifted: list[Any] = [] + for selector, w in zip(key, window, strict=True): + if isinstance(selector, slice): + shifted.append(slice(w.start + selector.start, w.start + selector.stop, selector.step)) + else: + shifted.append(selector + w.start) + return tuple(shifted) + + +def _read( + source: Any, + window: tuple[slice, ...] | None, + transform: IndexTransform, + support: IndexingSupport, +) -> Any: + """Materialize `transform` over `source`, asking for no more than `support` allows. + + `window` restricts the read to one part's box, in the wrapped array's + coordinates; it is folded into the key rather than materialized first, so a + part costs one request whatever its selection looks like. + """ + has_vindex = has_accessor(source, "vindex") + has_oindex = has_accessor(source, "oindex") + + if support is IndexingSupport.VECTORIZED and _is_correlated(transform): + points = _point_key(transform) + if points is not None: + gathered = source.vindex if has_vindex else source + # The key is flat so that a `vindex` accepting only one-dimensional + # coordinate lists is served too; the shape is restored here. + return _reshape(gathered[_shift_key(points, window)], transform.domain.shape) + + key, residual, n_arrays = _decompose(transform, support, has_oindex) + key = _shift_key(key, window) + block = source.oindex[key] if n_arrays > 0 and has_oindex else source[key] + return _lower(block, residual) + + # --------------------------------------------------------------------------- # # Partitions # --------------------------------------------------------------------------- # @@ -643,8 +881,10 @@ class LazyArray: Indexing through `.lazy` composes an `IndexTransform` and returns another `LazyArray`; `result()` materializes. - Selections use the **positional NumPy dialect**, and reads are broken up - along a **partitioning** discovered from the wrapped array — see the module + Selections use the **positional NumPy dialect**; reads are broken up along a + **partitioning** discovered from the wrapped array; and how much of a + selection each read hands to that array is **negotiated** from an + [`IndexingSupport`](support.md) level, also detected here. See the module docstring, which also covers how the dialect differs from `zarr.Array.lazy` and why every non-indexing NumPy operation materializes the view. @@ -652,8 +892,10 @@ class LazyArray: ---------- array The array to wrap. It must expose `shape`, `dtype`, and `__getitem__` - with basic (integer/slice) indexing. Its partitioning, if it advertises - one, is discovered here; use `with_parts` to choose a different one. + with basic (integer/slice) indexing; anything beyond that is used only + if detected or declared. Its partitioning, if it advertises one, is + discovered here; use `with_parts` to choose a different one, and + `with_indexing_support` to choose a different level. Examples -------- @@ -667,7 +909,7 @@ class LazyArray: [ 8, 10]]) """ - __slots__ = ("_array", "_parts", "_transform", "_window") + __slots__ = ("_array", "_parts", "_support", "_transform", "_window") def __init__(self, array: ArrayLike) -> None: shape = tuple(int(s) for s in array.shape) @@ -675,6 +917,7 @@ def __init__(self, array: ArrayLike) -> None: self._window: tuple[slice, ...] | None = None self._transform = IndexTransform.from_shape(shape) self._parts = _discover_parts(array, shape) + self._support = resolve_indexing_support(array) @classmethod def _derive( @@ -683,6 +926,7 @@ def _derive( transform: IndexTransform, parts: tuple[EdgeDimensionGrid, ...] | None, window: tuple[slice, ...] | None, + support: IndexingSupport, ) -> LazyArray: """Build a wrapper sharing `array` but carrying a new transform or partitioning.""" view = cls.__new__(cls) @@ -692,6 +936,7 @@ def _derive( view._transform = transform.translate_domain_to((0,) * transform.input_rank) view._parts = parts view._window = window + view._support = support return view @property @@ -701,12 +946,6 @@ def _base_shape(self) -> tuple[int, ...]: return tuple(int(s) for s in self._array.shape) return tuple(s.stop - s.start for s in self._window) - def _read_base(self) -> Any: - """The base array, materializing the window if this wrapper has one.""" - if self._window is None: - return self._array - return np.asarray(self._array[self._window]) - # -- array-API surface -------------------------------------------------- @property @@ -737,6 +976,83 @@ def dtype(self) -> Any: """The wrapped array's dtype; views never change it.""" return self._array.dtype + # -- source capability -------------------------------------------------- + + @property + def indexing_support(self) -> IndexingSupport: + """How much of a selection this wrapper hands to the array it wraps. + + Read-only; use + [`with_indexing_support`][zarr_indexing.lazy_array.LazyArray.with_indexing_support] + to change it. Resolved once, at construction, in this order: + + 1. an explicit `with_indexing_support` on this wrapper or the one it was + derived from; + 2. the wrapped array's own `__zarr_indexing_support__` declaration; + 3. inference from the wrapped array — `VECTORIZED` for a NumPy array or + for anything carrying zarr's `oindex`/`vindex` pair, + [`BASIC`][zarr_indexing.support.IndexingSupport] otherwise. + + See [`zarr_indexing.support`](support.md) for what each level means and + why the default is the most restrictive one. + + Examples + -------- + >>> import numpy as np + >>> LazyArray(np.arange(6)).indexing_support + + """ + return self._support + + def with_indexing_support(self, support: IndexingSupport) -> LazyArray: + """Return the same view, negotiating differently with the wrapped array. + + The transform, the wrapped array, the partitioning, and therefore + `result()` are all unchanged; only how much of a selection is asked of + the source, and how much is finished in memory, differs. Nothing is + copied and nothing is read. + + Use this when the wrapped array supports more than can be detected from + the outside, or when it supports less than it appears to and a request + it cannot answer must not be made. + + Parameters + ---------- + support + The level to use, overriding both the source's own declaration and + inference. + + Returns + ------- + LazyArray + The same view, read through a differently-negotiated source. + + Raises + ------ + TypeError + If `support` is not an `IndexingSupport`. Unlike detection, which + treats a malformed declaration as absent, this is a public API and + validates strictly. + + Examples + -------- + >>> import numpy as np + >>> from zarr_indexing.support import IndexingSupport + >>> view = LazyArray(np.arange(6)).with_indexing_support(IndexingSupport.BASIC) + >>> view.indexing_support + + >>> view.lazy.oindex[[4, 1, 1]].result() + array([4, 1, 1]) + """ + # The annotation already says this, but the check is for callers who are + # not type-checked: a bare string is the obvious mistake to make here, + # and it would otherwise be read as "no arrays" and silently over-read. + if not isinstance(support, IndexingSupport): # pyright: ignore[reportUnnecessaryIsInstance] + raise TypeError( + f"indexing support must be an IndexingSupport, got {type(support).__name__}" + ) + return LazyArray._derive(self._array, self._transform, self._parts, self._window, support) + # -- shape of the selection --------------------------------------------- @property @@ -918,7 +1234,7 @@ def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> L [(0, 0), (0, 1), (1, 0), (1, 1)] """ grids = None if parts is None else dimension_grids_from_chunks(parts, self._base_shape) - return LazyArray._derive(self._array, self._transform, grids, self._window) + return LazyArray._derive(self._array, self._transform, grids, self._window, self._support) def parts(self) -> Iterator[Partition]: """Iterate the base partitioning, projected through this view. @@ -975,7 +1291,7 @@ def parts(self) -> Iterator[Partition]: yield Partition( base_coords=base_coords, box=tuple((o, o + e) for o, e in zip(global_origin, extent, strict=True)), - array=LazyArray._derive(self._array, local, None, window), + array=LazyArray._derive(self._array, local, None, window, self._support), out_selection=_partition_out_selection(local, out_indices, out_shape), is_complete=_covers_whole_part(local, extent), ) @@ -1001,7 +1317,7 @@ def _select(self, selection: Any, mode: SelectionMode) -> LazyArray: transform = transform.translate_domain_to((0,) * transform.input_rank) literal = normalize_positional_selection(selection, transform.domain, mode) composed = selection_to_transform(literal, transform, mode) - return LazyArray._derive(self._array, composed, self._parts, self._window) + return LazyArray._derive(self._array, composed, self._parts, self._window, self._support) def __getitem__(self, selection: Any) -> Any: """Read a basic selection eagerly, like `numpy.ndarray.__getitem__`. @@ -1029,7 +1345,7 @@ def result(self) -> Any: zero-rank domain returns a zero-dimensional array, not a scalar. """ if self._parts is None: - return _lower(self._read_base(), self._transform) + return _read(self._array, self._window, self._transform, self._support) out_shape = self.shape out = np.empty(out_shape, dtype=np.dtype(self.dtype)) @@ -1060,6 +1376,11 @@ def __dask_tokenize__(self) -> Any: representation produce the same token. See `_wrapped_token` for the determinism scope of the wrapped array's contribution; dask is imported lazily and is never a requirement of this package. + + The `IndexingSupport` level is deliberately absent. It decides which + requests are made of the source, not which values come back, so two + wrappers that differ only in their level describe the same data and + should token alike. """ return ( type(self).__qualname__, diff --git a/packages/zarr-indexing/src/zarr_indexing/support.py b/packages/zarr-indexing/src/zarr_indexing/support.py new file mode 100644 index 0000000000..1bd8371041 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/support.py @@ -0,0 +1,164 @@ +"""How much indexing a wrapped array can do for itself. + +[`LazyArray`](lazy_array.md) reads through the object it wraps, and the objects +it can wrap differ widely in what they accept. A NumPy array takes any key NumPy +takes. A zarr array takes basic slicing through `__getitem__` and offers +`oindex` and `vindex` accessors for the two fancy dialects. An HDF5 dataset +takes integer arrays but only one axis at a time. A minimal duck array may take +nothing but integers and slices. + +`IndexingSupport` names those four capabilities, using the taxonomy and the +member names of xarray's [`IndexingSupport`][1], so a reader who knows one knows +the other. `LazyArray` uses the level to decide how much of a selection to hand +to the source and how much to finish in memory: the part the source cannot +express is applied to the block that comes back, with NumPy. Every level +therefore produces the same answer — the level chooses how much data crosses the +boundary, not what the read means. + +[1]: https://github.com/pydata/xarray/blob/main/xarray/core/indexing.py + +Declaring a level +----------------- +A source can declare its own level by exposing a `__zarr_indexing_support__` +attribute holding an `IndexingSupport`. That attribute is read defensively — a +value of the wrong type, or an attribute access that raises, means "no +declaration" rather than an error. + +Otherwise the level is inferred, conservatively: + +- a `numpy.ndarray` (or a subclass) supports everything NumPy supports: + `VECTORIZED` +- an object exposing both an `oindex` and a `vindex` accessor — zarr's surface — + is taken at its word: `VECTORIZED` +- anything else: `BASIC` + +`BASIC` is the default for an unrecognized source because it is the only +assumption that is always correct: an object that answers `shape`, `dtype`, and +`__getitem__` is guaranteed to understand integers and slices and nothing more. +Inference deliberately ignores `__array_namespace__`: the array API standard +requires neither integer-array nor boolean-mask indexing, so declaring one is no +evidence of fancy indexing. + +`LazyArray.with_indexing_support` overrides both, for a source that supports more +(or less) than can be detected from the outside. +""" + +from __future__ import annotations + +import enum +from typing import Any + +import numpy as np + +__all__ = [ + "SUPPORT_ATTRIBUTE", + "IndexingSupport", + "infer_indexing_support", + "resolve_indexing_support", +] + +SUPPORT_ATTRIBUTE = "__zarr_indexing_support__" +"""The attribute a source sets to declare its own `IndexingSupport`.""" + + +class IndexingSupport(enum.Enum): + """What a source's indexing surface accepts. + + Each member describes the widest key the source can be given in one request. + A level is a promise about `__getitem__` unless the source also exposes an + `oindex` or `vindex` accessor, in which case the orthogonal and correlated + forms are sent there instead — see the notes on each member. + + The four levels are not a numeric scale, and this enum is deliberately not + ordered: `OUTER_1VECTOR` sits *between* `BASIC` and `OUTER` in capability, + which no ordering of xarray's member names would make obvious. Compare + members by identity. + """ + + BASIC = "basic" + """Integers and slices only. + + `source[key]` accepts a tuple of integers and slices with a positive step. + This is the floor every array-like meets, and the default for a source whose + capability cannot be detected. A fancy selection against a `BASIC` source is + read as the smallest enclosing slab and gathered from that block in memory. + """ + + OUTER = "outer" + """Adds a one-dimensional integer array per axis, combined as an outer product. + + Several array axes can be requested at once, and each varies independently: + the result is the outer product of the per-axis selections, which is what + `zarr.Array.oindex` and `h5py`'s multi-list keys mean. + + A multi-array outer request is only ever sent to an `oindex` accessor. + `__getitem__` alone is ambiguous — the same key means an outer product to + HDF5 and a correlated gather to NumPy — so a source that declares `OUTER` + without exposing `oindex` is given at most one array axis per request, the + key both readings agree on. + """ + + OUTER_1VECTOR = "outer_1vector" + """Outer indexing, but with at most one array axis per request. + + The limit some backends impose (HDF5's is the classic one). The remaining + array axes are read as slabs and gathered in memory. The axis that keeps its + array is the one where a slab would over-read the most. + """ + + VECTORIZED = "vectorized" + """Full NumPy-style fancy indexing: several arrays broadcast against each other. + + `source[key]` follows NumPy's rules, so a tuple of coordinate arrays gathers + a list of points rather than an outer product. A source that exposes + `vindex` is given point gathers there instead, and one that exposes `oindex` + is given outer-product requests there, since both accessors say what they + mean without relying on NumPy's placement rules. + """ + + +def _read_attribute(array: Any, name: str) -> Any: + """Read an attribute off a foreign object, treating any failure as "absent".""" + try: + return getattr(array, name, None) + # Detection inspects an object we did not write; a guarded property may + # raise anything, and that can only ever mean "no information". + except Exception: + return None + + +def has_accessor(array: Any, name: str) -> bool: + """Whether `array` exposes a usable `oindex`/`vindex`-style accessor.""" + return _read_attribute(array, name) is not None + + +def declared_indexing_support(array: Any) -> IndexingSupport | None: + """The level `array` declares for itself, or None if it declares none. + + Reads `__zarr_indexing_support__` defensively: a missing attribute, an + attribute of the wrong type, and an attribute access that raises are all + "no declaration". + """ + declared = _read_attribute(array, SUPPORT_ATTRIBUTE) + return declared if isinstance(declared, IndexingSupport) else None + + +def infer_indexing_support(array: Any) -> IndexingSupport: + """Guess what `array` accepts from its type and its accessors. + + Conservative by construction: only a NumPy array and an object carrying + zarr's full `oindex`/`vindex` surface are credited with more than `BASIC`. + See the module docstring for why an under-estimate is the safe direction and + why `__array_namespace__` is not evidence. + """ + if isinstance(array, np.ndarray): + return IndexingSupport.VECTORIZED + if has_accessor(array, "oindex") and has_accessor(array, "vindex"): + return IndexingSupport.VECTORIZED + return IndexingSupport.BASIC + + +def resolve_indexing_support(array: Any) -> IndexingSupport: + """The level to use for `array`: what it declares, else what can be inferred.""" + declared = declared_indexing_support(array) + return infer_indexing_support(array) if declared is None else declared diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index e85604bf4d..817e9cd084 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -28,6 +28,7 @@ array_map_dependent_axis, dimension_grids_from_chunks, ) +from zarr_indexing.support import IndexingSupport if TYPE_CHECKING: from collections.abc import Callable, Sequence @@ -211,7 +212,7 @@ def make_source(flavor: str) -> LazyArray: if flavor == "zarr-misaligned": # Parts that deliberately straddle the zarr array's own chunks. return LazyArray(make_zarr_source()).with_parts((4, 3, 3)) - raise AssertionError(f"unknown source flavor {flavor!r}") + raise TypeError(f"unknown source flavor {flavor!r}") FLAVORS = [ @@ -1195,3 +1196,290 @@ def test_negative_step_over_a_fancy_axis_reverses_the_coordinates(source: LazyAr m = view.transform.output[0] assert isinstance(m, ArrayMap) np.testing.assert_array_equal(m.index_array.reshape(-1), np.array([2, 1, 3])) + + +# --------------------------------------------------------------------------- +# Source capability +# --------------------------------------------------------------------------- + +LEVELS = list(IndexingSupport) + + +class MinimalSource: + """The floor of the wrapped-array protocol: `shape`, `dtype`, `__getitem__`. + + Delegates every key straight to NumPy, so it is capable of far more than it + advertises. That is the point: nothing about it says so, and detection has + to take the cautious reading. + """ + + def __init__(self, data: np.ndarray[Any, Any]) -> None: + self._data = data + + @property + def shape(self) -> tuple[int, ...]: + return self._data.shape + + @property + def dtype(self) -> Any: + return self._data.dtype + + def __getitem__(self, key: Any) -> Any: + return self._data[key] + + +def _require_basic(key: tuple[Any, ...], channel: str) -> None: + """Raise unless every selector is a slice walking upwards.""" + for selector in key: + if not isinstance(selector, slice): + raise TypeError(f"{channel} was given a non-slice selector {selector!r}") + if selector.step is not None and selector.step < 0: + raise TypeError(f"{channel} was given a negative step {selector!r}") + + +class _OuterAccessor: + """An `oindex` that accepts slices and up to `limit` one-dimensional arrays.""" + + def __init__(self, data: np.ndarray[Any, Any], limit: int | None) -> None: + self._data = data + self._limit = limit + + def __getitem__(self, key: tuple[Any, ...]) -> Any: + arrays = 0 + for selector in key: + if isinstance(selector, slice): + if selector.step is not None and selector.step < 0: + raise TypeError(f"oindex was given a negative step {selector!r}") + elif isinstance(selector, np.ndarray) and selector.ndim == 1: + arrays += 1 + else: + raise TypeError(f"oindex was given {selector!r}") + if self._limit is not None and arrays > self._limit: + raise TypeError(f"oindex takes at most {self._limit} array axis, got {arrays}") + axes = [ + np.arange(size)[selector] if isinstance(selector, slice) else selector + for selector, size in zip(key, self._data.shape, strict=True) + ] + return self._data[np.ix_(*axes)] + + +class _VIndexAccessor: + """A `vindex` that, like zarr's, gathers points and takes nothing but coordinates.""" + + def __init__(self, data: np.ndarray[Any, Any]) -> None: + self._data = data + + def __getitem__(self, key: tuple[Any, ...]) -> Any: + if not all(isinstance(selector, np.ndarray) for selector in key): + raise TypeError(f"vindex takes coordinate arrays only, got {key!r}") + return self._data[key] + + +class StrictSource(MinimalSource): + """A source that raises when handed a key its declared level forbids. + + The three channels are separate on purpose: `__getitem__` stays basic + whatever the level, arrays arrive through `oindex`, points through `vindex`, + and each accessor exists only at a level that has one. A wrapper that + exceeds the declaration therefore fails loudly instead of quietly working + because the data underneath happens to be a NumPy array. + """ + + __zarr_indexing_support__ = IndexingSupport.BASIC + + def __getitem__(self, key: Any) -> Any: + _require_basic(key, "__getitem__") + return super().__getitem__(key) + + +class Outer1VectorSource(StrictSource): + __zarr_indexing_support__ = IndexingSupport.OUTER_1VECTOR + + @property + def oindex(self) -> _OuterAccessor: + return _OuterAccessor(self._data, limit=1) + + +class OuterSource(StrictSource): + __zarr_indexing_support__ = IndexingSupport.OUTER + + @property + def oindex(self) -> _OuterAccessor: + return _OuterAccessor(self._data, limit=None) + + +class VectorizedSource(OuterSource): + __zarr_indexing_support__ = IndexingSupport.VECTORIZED + + @property + def vindex(self) -> _VIndexAccessor: + return _VIndexAccessor(self._data) + + +STRICT_SOURCES: dict[IndexingSupport, Callable[[np.ndarray[Any, Any]], MinimalSource]] = { + IndexingSupport.BASIC: StrictSource, + IndexingSupport.OUTER_1VECTOR: Outer1VectorSource, + IndexingSupport.OUTER: OuterSource, + IndexingSupport.VECTORIZED: VectorizedSource, +} + + +@pytest.mark.parametrize("level", LEVELS, ids=[level.name for level in LEVELS]) +@pytest.mark.parametrize(("build", "oracle"), [c[1:] for c in CASES], ids=[c[0] for c in CASES]) +def test_every_level_gives_the_same_answer( + level: IndexingSupport, + build: Callable[[LazyArray], LazyArray], + oracle: Callable[[Any], Any], +) -> None: + """The invariant: the level chooses how the data is fetched, never what it is. + + Declaring `BASIC` on a NumPy array must answer exactly what declaring + `VECTORIZED` answers, having gathered more of it in memory to get there. + """ + expected = np.asarray(oracle(reference())) + for flavor in ("numpy-whole", "numpy-uniform-parts", "zarr"): + view = build(make_source(flavor).with_indexing_support(level)) + assert view.shape == expected.shape, f"{flavor} at {level}" + np.testing.assert_array_equal( + np.asarray(view.result()), expected, err_msg=f"{flavor} at {level}" + ) + + +@pytest.mark.parametrize("level", LEVELS, ids=[level.name for level in LEVELS]) +@pytest.mark.parametrize(("build", "oracle"), [c[1:] for c in CASES], ids=[c[0] for c in CASES]) +def test_a_source_is_never_asked_for_more_than_it_declares( + level: IndexingSupport, + build: Callable[[LazyArray], LazyArray], + oracle: Callable[[Any], Any], +) -> None: + """A source that refuses what its level forbids still gets the right answer. + + This is the test the negotiation exists for. The equivalence test above runs + against sources that would have tolerated an over-broad key; this one runs + against sources that will not. + """ + expected = np.asarray(oracle(reference())) + source = STRICT_SOURCES[level](reference()) + for parts in (None, PART_SHAPE): + view = build(LazyArray(source).with_parts(parts)) + np.testing.assert_array_equal( + np.asarray(view.result()), expected, err_msg=f"{level} parts={parts}" + ) + + +def test_random_chains_agree_across_levels() -> None: + """The seeded sweep, crossed with the four levels rather than the sources.""" + rng = np.random.default_rng(20260731) + data = reference() + for _ in range(50): + chain = _random_chain(rng) + expected = data + for mode, selection in chain: + expected = _apply_oracle(expected, mode, selection) + + for level in LEVELS: + for parts in (None, (2, 2, 2)): + view = LazyArray(data).with_indexing_support(level).with_parts(parts) + for mode, selection in chain: + view = _apply_view(view, mode, selection) + np.testing.assert_array_equal( + np.asarray(view.result()), + np.asarray(expected), + err_msg=f"{level} parts={parts}: {chain}", + ) + + +# -- detection -------------------------------------------------------------- + + +def test_a_numpy_array_is_detected_as_vectorized() -> None: + assert LazyArray(reference()).indexing_support is IndexingSupport.VECTORIZED + + +def test_a_zarr_array_is_detected_as_vectorized() -> None: + """zarr carries both accessors, which is the surface inference recognizes.""" + assert LazyArray(make_zarr_source()).indexing_support is IndexingSupport.VECTORIZED + + +def test_a_bare_array_like_is_detected_as_basic() -> None: + """Nothing but `shape`, `dtype`, and `__getitem__` says nothing about fancy keys.""" + assert LazyArray(MinimalSource(reference())).indexing_support is IndexingSupport.BASIC + + +@pytest.mark.parametrize("level", LEVELS, ids=[level.name for level in LEVELS]) +def test_a_declared_level_is_honored(level: IndexingSupport) -> None: + source = MinimalSource(reference()) + source.__zarr_indexing_support__ = level # type: ignore[attr-defined] + assert LazyArray(source).indexing_support is level + + +def test_a_malformed_declaration_falls_back_to_inference() -> None: + """Detection parses external input: a wrong-typed declaration is no declaration.""" + source = MinimalSource(reference()) + source.__zarr_indexing_support__ = "vectorized" # type: ignore[attr-defined] + assert LazyArray(source).indexing_support is IndexingSupport.BASIC + + +def test_a_declaration_that_raises_falls_back_to_inference() -> None: + """A guarded property that fails means "no information", never an error.""" + + class Hostile(MinimalSource): + @property + def __zarr_indexing_support__(self) -> IndexingSupport: + raise RuntimeError("no") + + assert LazyArray(Hostile(reference())).indexing_support is IndexingSupport.BASIC + + +def test_with_indexing_support_overrides_declaration_and_inference() -> None: + """The explicit setting wins over both other sources of the level.""" + declared = MinimalSource(reference()) + declared.__zarr_indexing_support__ = IndexingSupport.OUTER # type: ignore[attr-defined] + assert ( + LazyArray(declared).with_indexing_support(IndexingSupport.BASIC).indexing_support + is IndexingSupport.BASIC + ) + # And over inference, in the direction that matters least for correctness + # and most for I/O: a NumPy array told to behave like a minimal backend. + assert ( + LazyArray(reference()).with_indexing_support(IndexingSupport.BASIC).indexing_support + is IndexingSupport.BASIC + ) + + +def test_with_indexing_support_keeps_the_view_and_the_parts() -> None: + view = make_source("zarr").lazy[1:6, ::2] + renegotiated = view.with_indexing_support(IndexingSupport.OUTER) + assert renegotiated.array is view.array + assert renegotiated.transform == view.transform + assert [part.base_coords for part in renegotiated.parts()] == [ + part.base_coords for part in view.parts() + ] + np.testing.assert_array_equal(np.asarray(renegotiated.result()), np.asarray(view.result())) + # Parts inherit the level, so a part never exceeds it either. + assert all( + part.array.indexing_support is IndexingSupport.OUTER for part in renegotiated.parts() + ) + + +@pytest.mark.parametrize("bogus", ["vectorized", 3, None]) +def test_with_indexing_support_validates_strictly(bogus: Any) -> None: + """Our own API, so a level we do not recognize raises rather than degrading.""" + with pytest.raises(TypeError, match="must be an IndexingSupport"): + LazyArray(reference()).with_indexing_support(bogus) + + +def test_indexing_support_is_read_only() -> None: + with pytest.raises(AttributeError): + LazyArray(reference()).indexing_support = IndexingSupport.BASIC # type: ignore[misc] + + +def test_the_strict_doubles_notice_an_over_broad_key() -> None: + """A negative control: the respect test only means something if the double bites.""" + data = reference() + with pytest.raises(TypeError, match="non-slice selector"): + StrictSource(data)[np.array([0, 1]), slice(None), slice(None)] + with pytest.raises(TypeError, match="at most 1 array axis"): + Outer1VectorSource(data).oindex[np.array([0, 1]), np.array([0, 1]), slice(None)] + with pytest.raises(TypeError, match="coordinate arrays only"): + VectorizedSource(data).vindex[np.array([0, 1]), slice(None), slice(None)] From fc2d7d69ea592933946cfc208230b917cbe2de89 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 15:34:14 +0200 Subject: [PATCH 12/32] fix(zarr-indexing): token the data, not how it is read 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 --- .../src/zarr_indexing/lazy_array.py | 29 ++++++++++--------- .../zarr-indexing/tests/test_lazy_array.py | 21 ++++++++++++-- 2 files changed, 34 insertions(+), 16 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 7672478cbe..28607654ee 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -1368,26 +1368,27 @@ def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: return np.asarray(self.result(), dtype=dtype) def __dask_tokenize__(self) -> Any: - """A deterministic token: the wrapped array, the view, and the parts. - - Two wrappers produce equal tokens when they wrap the same data, address - the same cells, and read them in the same boxes. The view contributes - its canonical ndsel body, so transforms that differ only in - representation produce the same token. See `_wrapped_token` for the - determinism scope of the wrapped array's contribution; dask is imported - lazily and is never a requirement of this package. - - The `IndexingSupport` level is deliberately absent. It decides which - requests are made of the source, not which values come back, so two - wrappers that differ only in their level describe the same data and - should token alike. + """A deterministic token: the wrapped array and the view. + + Two wrappers produce equal tokens when they wrap the same data and + address the same cells. The view contributes its canonical ndsel body, + so transforms that differ only in representation produce the same + token. See `_wrapped_token` for the determinism scope of the wrapped + array's contribution; dask is imported lazily and is never a + requirement of this package. + + The partitioning and the `IndexingSupport` level are deliberately + absent. Both decide how the data is read — in which boxes, and through + which requests to the source — and neither changes the values that come + back, so two wrappers differing only in those describe the same data. + A token identifies data, so they token alike and a consumer that caches + on tokens reuses one result for both. """ return ( type(self).__qualname__, _wrapped_token(self._array), None if self._window is None else tuple((s.start, s.stop) for s in self._window), json.dumps(transform_to_canonical(self._transform), sort_keys=True), - None if self._parts is None else tuple(grid.sizes for grid in self._parts), ) def __len__(self) -> int: diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 817e9cd084..3c71a9fec2 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -1000,7 +1000,7 @@ def test_with_parts_validates_strictly(parts: Any, match: str) -> None: def test_dask_token_is_deterministic_and_discriminating() -> None: - """Same data, same view, same parts token alike; anything else differs.""" + """Same data and same view token alike; a different selection differs.""" data = reference() base = LazyArray(data) assert base.__dask_tokenize__() == LazyArray(reference()).__dask_tokenize__() @@ -1009,7 +1009,6 @@ def test_dask_token_is_deterministic_and_discriminating() -> None: "base": base.__dask_tokenize__(), "view": base.lazy[1:3].__dask_tokenize__(), "other view": base.lazy[2:4].__dask_tokenize__(), - "parts": base.with_parts((2, 2, 2)).__dask_tokenize__(), "other data": LazyArray(data + 1).__dask_tokenize__(), } assert len({repr(token) for token in tokens.values()}) == len(tokens) @@ -1017,6 +1016,24 @@ def test_dask_token_is_deterministic_and_discriminating() -> None: assert base.lazy[1:5].lazy[0:2].__dask_tokenize__() == base.lazy[1:3].__dask_tokenize__() +def test_dask_token_ignores_how_the_data_is_read() -> None: + """Partitioning and indexing support are read strategies, not data.""" + base = LazyArray(reference()) + token = base.__dask_tokenize__() + + assert base.with_parts((2, 2, 2)).__dask_tokenize__() == token + assert base.with_parts((3, 5, 4)).__dask_tokenize__() == token + assert base.with_indexing_support(IndexingSupport.BASIC).__dask_tokenize__() == token + assert base.with_indexing_support(IndexingSupport.OUTER).__dask_tokenize__() == token + + # The same holds for a view, where both strategies also apply. + view = base.lazy[1:6, 2:7] + assert ( + view.with_parts((2, 2, 2)).with_indexing_support(IndexingSupport.BASIC).__dask_tokenize__() + == view.__dask_tokenize__() + ) + + def test_iteration_yields_eager_slices(source: LazyArray) -> None: rows = list(source.lazy[2:5]) assert len(rows) == 3 From b700ad3abbea99fcfa5702235d07612ec7824225 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 15:51:20 +0200 Subject: [PATCH 13/32] fix(zarr-indexing): an empty downward walk selects nothing 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 --- .../src/zarr_indexing/transform.py | 12 ++++++--- .../zarr-indexing/tests/test_lazy_array.py | 26 ++++++++++++++++--- 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index f056ca119e..c17063b9f0 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -611,11 +611,15 @@ def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | def _positional_slice(pos: int, size: int, step: int) -> slice: """A NumPy slice selecting `size` elements from `pos`, walking by `step`. - The stop is `pos + size*step`, except that a downward walk reaching the - start of the array must stop at `None`: NumPy reads a negative stop as - counting from the end, so `slice(6, -1, -1)` selects nothing where - `slice(6, None, -1)` selects the first seven elements in reverse. + The stop is `pos + size*step`, except in two cases. An empty selection is + written out explicitly, because the arithmetic form can be a negative stop + that NumPy would read as counting from the end. And a downward walk + reaching the start of the array must stop at `None`, for the same reason: + `slice(6, -1, -1)` selects nothing where `slice(6, None, -1)` selects the + first seven elements in reverse. """ + if size <= 0: + return slice(0, 0, 1) stop = pos + size * step if step < 0 and stop < 0: return slice(pos, None, step) diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 3c71a9fec2..9c949b607b 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -447,6 +447,24 @@ def source(request: pytest.FixtureRequest) -> LazyArray: lambda a: a.lazy.vindex[..., np.array([3, 0, 1])].lazy[-1, 1:4], lambda r: r[..., np.array([3, 0, 1])][-1, 1:4], ), + # A downward walk that begins off the front of the axis selects nothing. + # Written against an oindex axis and a vindex axis, where the selection is + # carried by an index array rather than by the domain. + ( + "compose-oindex-then-empty-downward-walk", + lambda a: a.lazy.oindex[np.array([3, 1, 4]), :, :].lazy[-8::-1], + lambda r: r[np.array([3, 1, 4])][-8::-1], + ), + ( + "compose-vindex-then-empty-downward-walk", + lambda a: a.lazy.vindex[np.array([3, 1]), np.array([2, 0])].lazy[-9::-2], + lambda r: r[np.array([3, 1]), np.array([2, 0])][-9::-2], + ), + ( + "empty-downward-walk-on-a-plain-axis", + lambda a: a.lazy[-11::-1], + lambda r: r[-11::-1], + ), ] @@ -484,9 +502,11 @@ def _random_basic(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any selection.append(slice(start, stop, int(rng.integers(1, 4)))) elif roll < 0.8: # Downward: `start >= stop` and the stop may fall off the front, - # which is spelled `None`. - start = int(rng.integers(0, size)) - stop_choice = int(rng.integers(-1, start + 1)) + # which is spelled `None`. The start is drawn from below `-size` as + # well, where the walk begins off the front and selects nothing — + # a case that reads as an ordinary negative index but is empty. + start = int(rng.integers(-2 * size - 1, size)) if size else 0 + stop_choice = int(rng.integers(-1, max(start, 0) + 1)) stop = None if stop_choice < 0 else stop_choice selection.append(slice(start, stop, -int(rng.integers(1, 4)))) else: From 4b8917745b36f5cde0925a05abf6262eee49fafc Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 16:33:49 +0200 Subject: [PATCH 14/32] fix(zarr-indexing): a selection of slices is not a fancy selection MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- packages/zarr-indexing/changes/267.feature.md | 8 + .../src/zarr_indexing/chunk_resolution.py | 42 +++- .../src/zarr_indexing/composition.py | 38 +++- .../zarr-indexing/src/zarr_indexing/json.py | 43 +++- .../src/zarr_indexing/transform.py | 52 ++++- .../tests/test_chunk_resolution.py | 54 +++++ .../zarr-indexing/tests/test_composition.py | 87 ++++++++ packages/zarr-indexing/tests/test_json.py | 64 ++++++ .../zarr-indexing/tests/test_lazy_array.py | 189 +++++++++++++++++- .../zarr-indexing/tests/test_transform.py | 20 +- 10 files changed, 558 insertions(+), 39 deletions(-) diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 86ee1ce61f..9bf55a1ae6 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -16,3 +16,11 @@ Negative-step slices are supported, following the merged ndsel 1.0-draft.2 (PR # `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: a NumPy array or an object carrying zarr's `oindex`/`vindex` pair reads as `VECTORIZED`, a source's own `__zarr_indexing_support__` attribute is honored when it holds an `IndexingSupport`, and everything else is `BASIC`, 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 the chunk selection `sub_transform_to_selections` builds for a negative-stride dimension: 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. + +`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. diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 7aea86ad02..9372aff028 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -39,7 +39,16 @@ from zarr_indexing.domain import IndexDomain from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap -from zarr_indexing.transform import IndexTransform +from zarr_indexing.transform import ( + IndexTransform, + _positional_slice, # pyright: ignore[reportPrivateUsage] +) + +# `_positional_slice` is a leading-underscore helper in `transform.py`, shared +# with this module on purpose: the slice an `ArrayMap` axis is reindexed by and +# the slice a `DimensionMap` is lowered to are the same walk, and the two must +# agree on where a downward one stops. See `json.py` for the same suppression +# and the same open question about promoting these helpers out of "private". if TYPE_CHECKING: from collections.abc import Iterator, Sequence @@ -128,6 +137,13 @@ def iter_chunk_transforms( indices (integer array). `None` for basic/slice indexing. """ + if any(size == 0 for size in transform.domain.shape): + # An empty view touches no chunk. Checked on the domain rather than on + # the index arrays: an axis of genuine extent 1 is stored as a broadcast + # singleton, so a slice that empties the domain does not shrink the + # array, and the emptiness shows only here. + return + array_map_1d = _one_dimensional_correlated_array_map(transform) if array_map_1d is not None: sorted_map, storage = array_map_1d @@ -261,6 +277,18 @@ def iter_chunk_transforms( yield (chunk_coords, local, surviving) +def _dimension_map_slice(m: DimensionMap, dim_lo: int, dim_hi: int) -> slice: + """The NumPy slice a `DimensionMap` reads over input range `[dim_lo, dim_hi)`. + + The walk starts at the storage coordinate of `dim_lo` and takes `dim_hi - + dim_lo` steps of `m.stride`. A downward walk that reaches the start of the + axis must stop at `None`: an integer stop below zero counts from the end + instead, and swapping the endpoints while keeping the negative step (which + this did) produces a slice that selects nothing at all. + """ + return _positional_slice(m.offset + m.stride * dim_lo, dim_hi - dim_lo, m.stride) + + def sub_transform_to_selections( sub_transform: IndexTransform, out_indices: OutIndices = None, @@ -327,11 +355,7 @@ def sub_transform_to_selections( chunk_sel.append(m.offset) elif isinstance(m, DimensionMap): d = m.input_dimension - start = m.offset + m.stride * inclusive_min[d] - stop = m.offset + m.stride * exclusive_max[d] - if m.stride < 0: - start, stop = stop + 1, start + 1 - chunk_sel.append(slice(start, stop, m.stride)) + chunk_sel.append(_dimension_map_slice(m, inclusive_min[d], exclusive_max[d])) else: # ArrayMap idx = m.index_array.reshape(-1) chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) @@ -362,11 +386,7 @@ def sub_transform_to_selections( d = m.input_dimension dim_lo = inclusive_min[d] dim_hi = exclusive_max[d] - start = m.offset + m.stride * dim_lo - stop = m.offset + m.stride * dim_hi - if m.stride < 0: - start, stop = stop + 1, start + 1 - chunk_sel.append(slice(start, stop, m.stride)) + chunk_sel.append(_dimension_map_slice(m, dim_lo, dim_hi)) out_sel.append(slice(dim_lo, dim_hi)) else: # ArrayMap (orthogonal: full-rank, raveled to its 1-D fancy coords) idx = m.index_array.reshape(-1) diff --git a/packages/zarr-indexing/src/zarr_indexing/composition.py b/packages/zarr-indexing/src/zarr_indexing/composition.py index f5cc82599c..cb5e661a45 100644 --- a/packages/zarr-indexing/src/zarr_indexing/composition.py +++ b/packages/zarr-indexing/src/zarr_indexing/composition.py @@ -41,12 +41,16 @@ def compose(outer: IndexTransform, inner: IndexTransform) -> IndexTransform: f"({inner.domain.ndim})" ) - result_output = [_compose_single(outer, inner_map) for inner_map in inner.output] + result_output = [ + _compose_single(outer, inner_map, inner.domain.inclusive_min) for inner_map in inner.output + ] return IndexTransform(domain=outer.domain, output=tuple(result_output)) -def _compose_single(outer: IndexTransform, inner_map: OutputIndexMap) -> OutputIndexMap: +def _compose_single( + outer: IndexTransform, inner_map: OutputIndexMap, inner_origin: tuple[int, ...] +) -> OutputIndexMap: """Compose a single inner output map with the full outer transform.""" if isinstance(inner_map, ConstantMap): return ConstantMap(offset=inner_map.offset) @@ -55,7 +59,7 @@ def _compose_single(outer: IndexTransform, inner_map: OutputIndexMap) -> OutputI return _compose_dimension(outer, inner_map) # inner_map: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) - return _compose_array(outer, inner_map) + return _compose_array(outer, inner_map, inner_origin) def _compose_dimension(outer: IndexTransform, inner_map: DimensionMap) -> OutputIndexMap: @@ -91,11 +95,20 @@ def _compose_dimension(outer: IndexTransform, inner_map: DimensionMap) -> Output ) -def _compose_array(outer: IndexTransform, inner_map: ArrayMap) -> OutputIndexMap: +def _compose_array( + outer: IndexTransform, inner_map: ArrayMap, inner_origin: tuple[int, ...] +) -> OutputIndexMap: """Compose when inner is an ArrayMap. storage = offset_i + stride_i * arr_i[intermediate] We need to evaluate arr_i at the intermediate coordinates produced by outer. + + Both domains carry their own origin, and neither is necessarily 0 — a step-1 + slice keeps its literal bounds and a negative step produces a negative + origin, so non-zero origins are the ordinary case here rather than the exotic + one. The intermediate coordinates are read over the *outer* domain's own + range, and the inner array is addressed positionally from the *inner* + domain's origin. """ arr_i = inner_map.index_array offset_i = inner_map.offset @@ -106,7 +119,12 @@ def _compose_array(outer: IndexTransform, inner_map: ArrayMap) -> OutputIndexMap if all_constant: # Evaluate arr_i at the single constant point - idx = tuple(m.offset for m in outer.output if isinstance(m, ConstantMap)) + idx = tuple( + m.offset - lo + for m, lo in zip( + (m for m in outer.output if isinstance(m, ConstantMap)), inner_origin, strict=True + ) + ) value = int(arr_i[idx]) return ConstantMap(offset=offset_i + stride_i * value) @@ -115,15 +133,17 @@ def _compose_array(outer: IndexTransform, inner_map: ArrayMap) -> OutputIndexMap outer_map = outer.output[0] if isinstance(outer_map, DimensionMap): - dim_size = outer.domain.shape[outer_map.input_dimension] - user_indices = np.arange(dim_size, dtype=np.intp) + dim = outer_map.input_dimension + user_indices = np.arange( + outer.domain.inclusive_min[dim], outer.domain.exclusive_max[dim], dtype=np.intp + ) intermediate_vals = outer_map.offset + outer_map.stride * user_indices - new_arr = arr_i[intermediate_vals] + new_arr = arr_i[intermediate_vals - inner_origin[0]] return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) if isinstance(outer_map, ArrayMap): intermediate_vals = outer_map.offset + outer_map.stride * outer_map.index_array - new_arr = arr_i[intermediate_vals] + new_arr = arr_i[intermediate_vals - inner_origin[0]] return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) # General multi-dim case: not yet implemented diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py index 89c8a04c42..7be579c8bc 100644 --- a/packages/zarr-indexing/src/zarr_indexing/json.py +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -5,12 +5,16 @@ (spec section 4.3, as produced by `zarr_indexing.messages.normalize_ndsel`) into the numpy-backed `IndexTransform` the chunk engine runs on, and back. -Two engine constraints live **here and only here**: +Three engine constraints live **here and only here**: - **Finite bounds.** An `IndexDomain` addresses a finite array, so a canonical body carrying a `"-inf"`/`"+inf"` bound cannot be lowered; `from_json` raises. - **Implicit bounds lower by value.** The `[n]`-bracket implicit/explicit flag is a message-layer concern; the engine keeps only the integer value. +- **Integer `index_array` content.** The message layer carries `index_array` + verbatim (the spec defers its shape and type), so lowering is where a float, + boolean or string array is rejected — as an `NdselError`, rather than + truncating `[0.9, 1.9]` to cells 0 and 1 or leaking a raw NumPy error. ## The `index_array` wire format (and the degenerate-collapse it documents) @@ -49,7 +53,7 @@ import numpy as np from zarr_indexing.domain import IndexDomain -from zarr_indexing.messages import normalize_ndsel +from zarr_indexing.messages import NdselError, normalize_ndsel from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_indexing.transform import ( IndexTransform, @@ -132,6 +136,33 @@ def _lower_bound(bound: BoundJSON, where: str) -> int: return int(value) +def _lower_index_array(raw: Any, where: str) -> np.ndarray[Any, np.dtype[np.intp]]: + """Lower a canonical `index_array` to `intp`, rejecting non-integer content. + + The message layer carries `index_array` verbatim — the spec defers its shape + and type to the engine — so this is where the content is checked. An index + array names storage cells, and nothing but an integer names one: converting + `[0.9, 1.9]` would silently read cells 0 and 1, and `[true, false]` cells 1 + and 0. Strings raise here rather than leaking NumPy's own conversion error. + """ + try: + arr = np.asarray(raw) + except (TypeError, ValueError) as exc: + raise NdselError("invalid_json", f"{where} is not an array: {exc}") from exc + if arr.size == 0 and arr.dtype.kind == "f": + # An empty JSON list carries no element type and NumPy defaults it to + # float64. An empty selection is legal, so take it as an empty index array. + return np.zeros(arr.shape, dtype=np.intp) + if arr.dtype.kind not in "iu": + raise NdselError( + "invalid_json", + f"{where} must hold integers, got an array of {arr.dtype.name}; an " + f"index array names storage cells, which floats, booleans and " + f"strings do not", + ) + return np.asarray(arr, dtype=np.intp) + + def _lower_labels(labels: list[str]) -> tuple[str, ...] | None: """All-empty labels collapse to `None` so a label-free domain round-trips.""" return None if all(label == "" for label in labels) else tuple(labels) @@ -209,7 +240,7 @@ def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: share axes. """ if "index_array" in data: - arr = np.asarray(data["index_array"], dtype=np.intp) + arr = _lower_index_array(data["index_array"], "index_array") return ArrayMap( index_array=arr, offset=data.get("offset", 0), @@ -284,11 +315,13 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: # Classify index_array maps globally: an axis owned by exactly one array map # (and the map's sole non-singleton axis) marks that map orthogonal; shared # or multiple non-singleton axes mark the maps correlated (vindex). + arrays: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} array_axes: dict[int, tuple[int, ...]] = {} axis_owners: Counter[int] = Counter() for i, om in enumerate(output_raw): if "index_array" in om: - arr = np.asarray(om["index_array"], dtype=np.intp) + arr = _lower_index_array(om["index_array"], f"output[{i}].index_array") + arrays[i] = arr dep = _array_map_dependency_axes(arr) array_axes[i] = dep axis_owners.update(dep) @@ -300,7 +333,7 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: input_dim = dep[0] if len(dep) == 1 and axis_owners[dep[0]] == 1 else None output.append( ArrayMap( - index_array=np.asarray(om["index_array"], dtype=np.intp), + index_array=arrays[i], offset=om.get("offset", 0), stride=om.get("stride", 1), input_dimension=input_dim, diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index c17063b9f0..ec05e62aa5 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -285,6 +285,14 @@ def _intersect( f"transform output rank ({transform.output_rank})" ) + if any(size == 0 for size in transform.domain.shape): + # An empty input domain addresses no coordinates at all, so it meets no + # output domain. Deciding it here keeps the per-flavor intersections from + # having to reconcile an empty domain with an index array that is *not* + # empty: a genuine extent-1 axis is stored as a broadcast singleton, so + # emptying its domain leaves the array at size 1. + return None + correlated_dims = [ i for i, m in enumerate(transform.output) @@ -720,19 +728,37 @@ def _guard_fancy_after_fancy(m: ArrayMap, fancy_dims: set[int] | list[int]) -> N def _reindex_array_oindex( - arr: np.ndarray[Any, np.dtype[np.intp]], + m: ArrayMap, normalized: tuple[Any, ...] | list[Any], domain: IndexDomain, ) -> np.ndarray[Any, np.dtype[np.intp]]: - """Apply oindex/vindex selection to an existing ArrayMap's index_array. - - Each old input dimension gets either an array (fancy index that axis) - or a slice applied to the corresponding array axis. + """Apply an oindex/vindex selection to an existing ArrayMap's index_array. + + Dependency-aware in exactly the way `_reindex_array` is: an entry applies to + the array only along the axes the map genuinely varies over (its dependency + axes, plus the `input_dimension` recorded for a degenerate length-1 + orthogonal selection). An entry landing on a **singleton** axis the map + merely broadcasts over does not touch its values — the singleton still + broadcasts over the narrowed domain, so it is preserved. Applying such a + slice positionally instead would index a size-1 axis out of range and + truncate the whole array to size 0. + + Only slices ever take the broadcast path: `_guard_fancy_after_fancy` rejects + coordinates aimed at a broadcast axis before this is called. """ + arr = m.index_array + dependent = set(_array_map_dependency_axes(arr)) + if m.input_dimension is not None: + dependent.add(m.input_dimension) + idx: list[Any] = [] for old_dim, sel in enumerate(normalized): if old_dim >= arr.ndim: break + if old_dim not in dependent: + # Broadcast axis: keep the singleton whatever the entry says. + idx.append(slice(None)) + continue lo = domain.inclusive_min[old_dim] if isinstance(sel, np.ndarray): # Values are literal domain coordinates; the stored array is @@ -870,11 +896,15 @@ def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, Normalized `ArrayMap` index arrays carry the full input rank of their enclosing transform: an axis the array varies over has its full size, while an axis the array is independent of is a singleton (size 1). The dependency - axes are therefore exactly the non-singleton axes. An orthogonal (`oindex`) - array depends on a single axis; a vectorized (`vindex`) array depends on all - of the (shared) broadcast axes. + axes are therefore exactly the axes of size 2 or more. An orthogonal + (`oindex`) array depends on a single axis; a vectorized (`vindex`) array + depends on all of the (shared) broadcast axes. + + A size-**0** axis carries no dependency either: the array has no values to + vary, so an empty selection stays the flavor it was made as rather than + reading as correlated with every other axis. """ - return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) + return tuple(axis for axis, size in enumerate(index_array.shape) if size > 1) def array_map_dependent_axis(m: ArrayMap) -> int | None: @@ -1061,7 +1091,7 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: else: # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) _guard_fancy_after_fancy(m, list(dim_array.keys())) - new_arr = _reindex_array_oindex(m.index_array, normalized, transform.domain) + new_arr = _reindex_array_oindex(m, normalized, transform.domain) array_input_dim: int | None = None if m.input_dimension is not None: array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) @@ -1261,7 +1291,7 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: else: # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) _guard_fancy_after_fancy(m, array_dims) - new_arr = _reindex_array_oindex(m.index_array, processed, transform.domain) + new_arr = _reindex_array_oindex(m, processed, transform.domain) new_output.append( ArrayMap( index_array=new_arr, diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index f6a37c7012..8ff3b19b6c 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -445,6 +445,60 @@ def test_array_map_with_offset_stride(self) -> None: np.testing.assert_array_equal(chunk_sel[0], np.array([10, 15, 20])) assert drop_axes == () + def test_dimension_map_reversed(self) -> None: + """A negative stride walks the axis backwards, stopping off the front. + + `slice(0, 8, -1)` — the endpoints swapped while the step stays negative — + selects nothing; the axis reversed is `slice(7, None, -1)`. + """ + t = IndexTransform.identity(IndexDomain.from_shape((8,)))[::-1].translate_domain_to((0,)) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (slice(7, None, -1),) + assert out_sel == (slice(0, 8),) + assert drop_axes == () + np.testing.assert_array_equal(np.arange(8)[chunk_sel[0]], np.arange(8)[::-1]) + + def test_dimension_map_reversed_bounded(self) -> None: + """A reversed walk that stops short keeps an ordinary integer stop.""" + t = IndexTransform.identity(IndexDomain.from_shape((8,)))[6:2:-1].translate_domain_to((0,)) + chunk_sel, _out_sel, _drop = sub_transform_to_selections(t) + np.testing.assert_array_equal(np.arange(8)[chunk_sel[0]], np.arange(8)[6:2:-1]) + + def test_dimension_map_reversed_strided(self) -> None: + """The same, with a stride whose walk does not land on the origin.""" + t = IndexTransform.identity(IndexDomain.from_shape((9,)))[::-3].translate_domain_to((0,)) + chunk_sel, _out_sel, _drop = sub_transform_to_selections(t) + np.testing.assert_array_equal(np.arange(9)[chunk_sel[0]], np.arange(9)[::-3]) + + def test_dimension_map_slice_reads_the_coordinates_the_map_names(self) -> None: + """Over every offset/stride/extent in range, the emitted slice picks out + exactly `offset + stride * i` for `i` in the input range.""" + extent = 12 + cells = np.arange(extent) + for stride in (-4, -3, -2, -1, 1, 2, 3, 4): + for size in range(extent + 1): + span = 0 if size == 0 else abs(stride) * (size - 1) + offsets = range(extent - span) if span < extent else () + for offset in offsets: + start = offset if stride > 0 else offset + span + t = IndexTransform( + domain=IndexDomain.from_shape((size,)), + output=(DimensionMap(input_dimension=0, offset=start, stride=stride),), + ) + chunk_sel, _out_sel, _drop = sub_transform_to_selections(t) + expected = np.array([start + stride * i for i in range(size)]) + np.testing.assert_array_equal( + cells[chunk_sel[0]], expected, err_msg=f"{stride=} {size=} {start=}" + ) + + def test_correlated_residual_dimension_map_reversed(self) -> None: + """The correlated branch builds the same reversed slice for its residual dim.""" + t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] + reversed_slice = t[:, ::-1].translate_domain_to((0, 0)) + chunk_sel, _out_sel, _drop = sub_transform_to_selections(reversed_slice) + assert chunk_sel[2] == slice(4, None, -1) + np.testing.assert_array_equal(np.arange(5)[chunk_sel[2]], np.arange(5)[::-1]) + def test_mixed_maps_2d(self) -> None: """Mix of ConstantMap and DimensionMap.""" t = IndexTransform( diff --git a/packages/zarr-indexing/tests/test_composition.py b/packages/zarr-indexing/tests/test_composition.py index dd92f59b80..3c2aed351b 100644 --- a/packages/zarr-indexing/tests/test_composition.py +++ b/packages/zarr-indexing/tests/test_composition.py @@ -105,6 +105,93 @@ def test_array_inner_array_outer(self) -> None: np.testing.assert_array_equal(result.output[0].index_array, expected) +def _storage_of(transform: IndexTransform, point: tuple[int, ...]) -> tuple[int, ...]: + """The storage coordinates a transform assigns to one input point. + + The point is given in the transform's own domain coordinates; an index array + carries the full input rank, singleton on the axes it does not vary over, and + is addressed positionally from the domain origin. + """ + coords: list[int] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + coords.append(m.offset) + elif isinstance(m, DimensionMap): + coords.append(m.offset + m.stride * point[m.input_dimension]) + else: + origin = transform.domain.inclusive_min + idx = tuple( + 0 if m.index_array.shape[axis] == 1 else point[axis] - origin[axis] + for axis in range(m.index_array.ndim) + ) + coords.append(m.offset + m.stride * int(m.index_array[idx])) + return tuple(coords) + + +def _assert_composes_pointwise(outer: IndexTransform, inner: IndexTransform) -> None: + """`compose(outer, inner)` must agree with running the two in sequence.""" + composed = compose(outer, inner) + assert composed.domain == outer.domain + lo = outer.domain.inclusive_min + hi = outer.domain.exclusive_max + for coord in np.ndindex(*outer.domain.shape): + point = tuple(int(c) + int(o) for c, o in zip(coord, lo, strict=True)) + assert all(point[d] < hi[d] for d in range(len(hi))) + intermediate = _storage_of(outer, point) + assert _storage_of(composed, point) == _storage_of(inner, intermediate) + + +class TestComposeOverANonZeroOriginDomain: + """The outer domain need not start at 0 — a step-1 slice preserves its + literal bounds and a negative step produces a negative origin — so the inner + map has to be evaluated over the outer domain's real range.""" + + def test_array_inner_sliced_outer(self) -> None: + inner = IndexTransform.from_shape((10,)).oindex[np.array([3, 1, 4, 1, 5])] + outer = IndexTransform.identity(IndexDomain.from_shape((5,)))[1:4] + assert outer.domain.inclusive_min == (1,) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal( + result.output[0].index_array, np.array([1, 4, 1], dtype=np.intp) + ) + _assert_composes_pointwise(outer, inner) + + def test_array_inner_reversed_outer(self) -> None: + inner = IndexTransform.from_shape((10,)).oindex[np.array([3, 1, 4, 1, 5])] + outer = IndexTransform.identity(IndexDomain.from_shape((5,)))[::-1] + assert outer.domain.inclusive_min == (-4,) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal( + result.output[0].index_array, np.array([5, 1, 4, 1, 3], dtype=np.intp) + ) + _assert_composes_pointwise(outer, inner) + + def test_array_inner_strided_outer(self) -> None: + inner = IndexTransform.from_shape((10,)).oindex[np.array([3, 1, 4, 1, 5, 9])] + outer = IndexTransform.identity(IndexDomain.from_shape((6,)))[1::2] + _assert_composes_pointwise(outer, inner) + + def test_array_inner_translated_outer(self) -> None: + inner = IndexTransform.from_shape((10,)).oindex[np.array([3, 1, 4, 1, 5])] + outer = IndexTransform.identity(IndexDomain.from_shape((5,))).translate_domain_to((-2,)) + _assert_composes_pointwise(outer, inner) + + def test_array_inner_array_outer_over_a_shifted_domain(self) -> None: + inner = IndexTransform.from_shape((10,)).oindex[np.array([3, 1, 4, 1, 5])] + outer = IndexTransform.from_shape((5,)).oindex[np.array([4, 0, 2])] + _assert_composes_pointwise(outer, inner) + + def test_array_inner_constant_outer_over_a_shifted_domain(self) -> None: + inner = IndexTransform.from_shape((10,)).oindex[np.array([3, 1, 4, 1, 5])] + outer = IndexTransform( + domain=IndexDomain.from_shape((4,)), + output=(ConstantMap(offset=3),), + ) + _assert_composes_pointwise(outer, inner) + + class TestComposeMultiDim: def test_2d_identity_compose(self) -> None: a = IndexTransform.from_shape((10, 20)) diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py index 42b59b2c30..a348ef0fb8 100644 --- a/packages/zarr-indexing/tests/test_json.py +++ b/packages/zarr-indexing/tests/test_json.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Any + import numpy as np import pytest @@ -13,6 +15,7 @@ output_index_map_from_json, output_index_map_to_json, ) +from zarr_indexing.messages import NdselError from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import IndexTransform @@ -324,6 +327,67 @@ def test_slices_and_constants(self) -> None: assert _transforms_equal(rt, t) +def _index_array_body(index_array: Any, rank: int = 1) -> IndexTransformJSON: + return { + "input_rank": rank, + "input_inclusive_min": [0] * rank, + "input_exclusive_max": [2] * rank, + "input_labels": [""] * rank, + "output": [{"offset": 0, "stride": 1, "index_array": index_array}], + } + + +@pytest.mark.parametrize( + ("index_array", "detail"), + [ + ([0.9, 1.9], "float64"), + ([0, 1.5], "float64"), + ([True, False], "bool"), + (["a", "b"], "str"), + ("abc", "str"), + ([None, None], "object"), + ], + ids=["floats", "mixed", "bools", "strings", "string", "nulls"], +) +def test_a_non_integer_index_array_is_rejected(index_array: Any, detail: str) -> None: + """An `index_array` addresses storage cells, so it must be integral. + + Lowering a float array silently truncated it (`[0.9, 1.9]` selected cells 0 + and 1), a bool array coerced to 0/1, and a string array leaked a raw NumPy + `ValueError` from the middle of the conversion. + """ + with pytest.raises(NdselError) as excinfo: + index_transform_from_json(_index_array_body(index_array)) + assert excinfo.value.reason == "invalid_json" + assert "index_array" in str(excinfo.value) + assert detail in str(excinfo.value) + + +def test_a_ragged_index_array_is_rejected() -> None: + """A nested list that is not rectangular is not an array at all.""" + with pytest.raises(NdselError) as excinfo: + index_transform_from_json(_index_array_body([[0, 1], [2]])) + assert excinfo.value.reason == "invalid_json" + + +@pytest.mark.parametrize( + ("index_array", "rank"), [([0, 1], 1), ([[0], [1]], 2), ([], 1)], ids=["1d", "2d", "empty"] +) +def test_an_integer_index_array_is_accepted(index_array: Any, rank: int) -> None: + """Integers of any nesting still lower, including an empty selection.""" + t = index_transform_from_json(_index_array_body(index_array, rank)) + m = t.output[0] + assert isinstance(m, ArrayMap) + assert m.index_array.dtype == np.intp + + +def test_a_non_integer_index_array_is_rejected_by_the_map_loader() -> None: + """The single-map loader enforces the same constraint as the transform one.""" + with pytest.raises(NdselError) as excinfo: + output_index_map_from_json({"index_array": [0.5, 1.5]}) + assert excinfo.value.reason == "invalid_json" + + def test_infinite_bound_rejected_on_lowering() -> None: body: IndexTransformJSON = { "input_rank": 1, diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 9c949b607b..5a9a12e12b 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -465,6 +465,33 @@ def source(request: pytest.FixtureRequest) -> LazyArray: lambda a: a.lazy[-11::-1], lambda r: r[-11::-1], ), + # A fancy *spelling* whose entries are all slices is not a fancy selection: + # it narrows the view's own axes and must compose exactly like basic + # indexing. The slices start past 0, so a step that applied them to the + # broadcast (singleton) axes of the existing index array would truncate it. + ( + "compose-oindex-then-oindex-slices-only", + lambda a: a.lazy.oindex[[4, 0, 0], :, :].lazy.oindex[:, 2:5, 1:], + lambda r: outer(r, ([4, 0, 0], slice(None), slice(None)))[:, 2:5, 1:], + ), + ( + "compose-oindex-then-oindex-slices-only-strided", + lambda a: a.lazy.oindex[:, [3, 1, 1], :].lazy.oindex[1::2, :, ::-1], + lambda r: outer(r, (slice(None), [3, 1, 1], slice(None)))[1::2, :, ::-1], + ), + ( + "compose-vindex-then-oindex-slices-only", + lambda a: a.lazy.vindex[np.array([4, 0, 2]), np.array([1, 3, 0])].lazy.oindex[1:, 2:], + lambda r: r[np.array([4, 0, 2]), np.array([1, 3, 0])][1:, 2:], + ), + ( + "compose-oindex-then-oindex-array-on-its-own-axis", + lambda a: a.lazy.oindex[[4, 0, 0], :, :].lazy.oindex[[2, 0], 3:, :], + lambda r: outer( + outer(r, ([4, 0, 0], slice(None), slice(None))), + ([2, 0], slice(3, None), slice(None)), + ), + ), ] @@ -574,11 +601,44 @@ def _apply_view(view: LazyArray, mode: str, selection: tuple[Any, ...]) -> LazyA return view.lazy.vindex[selection] +def _random_slices_only(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: + """A selection of slices alone, spelled through `oindex`. + + `oindex` entries that are all slices are not a fancy selection — they narrow + the view's own axes and must compose like basic indexing. A start past 0 is + what distinguishes a step that walks the index array's dependency axes from + one that walks its broadcast singletons, so slices are drawn to reach past + the origin. (`vindex` is coordinate-only and rejects a slice outright, so + this spelling has no vectorized counterpart.) + """ + selection: list[Any] = [] + for size in shape: + roll = rng.random() + if roll < 0.4: + start = int(rng.integers(0, size)) if size else 0 + selection.append(slice(start, size)) + elif roll < 0.7: + start = int(rng.integers(0, size)) if size else 0 + selection.append(slice(start, size, int(rng.integers(1, 3)))) + elif roll < 0.85: + selection.append(slice(None, None, -1)) + else: + selection.append(slice(None)) + return tuple(selection) + + def _random_chain(rng: np.random.Generator) -> list[tuple[str, tuple[Any, ...]]]: - """A chain of 2-4 steps with at most one fancy step, in either order.""" + """A chain of 2-4 steps with at most one fancy step, in either order. + + A step spelled through `oindex` but carrying only slices is drawn separately + (`slices-only`): it is legal after a fancy step — genuine fancy-after-fancy + is not, and raises — and it exercises the reindexing of an existing index + array by a *slice* rather than by coordinates. + """ n_steps = int(rng.integers(2, 5)) fancy_at = int(rng.integers(0, n_steps)) fancy_mode = "orthogonal" if rng.random() < 0.5 else "vectorized" + slices_only_at = int(rng.integers(0, n_steps)) if rng.random() < 0.4 else -1 chain: list[tuple[str, tuple[Any, ...]]] = [] running = reference() @@ -586,7 +646,10 @@ def _random_chain(rng: np.random.Generator) -> list[tuple[str, tuple[Any, ...]]] if running.ndim == 0 or running.size == 0: break mode = fancy_mode if step == fancy_at else "basic" - if mode == "basic": + if step == slices_only_at and step != fancy_at: + mode = "orthogonal" + selection = _random_slices_only(rng, running.shape) + elif mode == "basic": selection = _random_basic(rng, running.shape) elif mode == "orthogonal": selection = _random_oindex(rng, running.shape) @@ -634,6 +697,128 @@ def test_random_chains_match_numpy(flavor: str) -> None: ) +@pytest.mark.parametrize("flavor", ["numpy-whole", "numpy-uniform-parts"]) +def test_random_chains_have_parts_that_tile_the_view(flavor: str) -> None: + """The seeded sweep, read through `parts()` rather than through `result()`. + + `result()` can absorb a defect that `parts()` cannot — an empty view assembles + correctly from no parts at all — so the iteration contract needs its own + sweep: every part places its own values, and together they cover the view + exactly once. + """ + rng = np.random.default_rng(20260731) + source = make_source(flavor) + + for _ in range(300): + chain = _random_chain(rng) + expected = reference() + for mode, selection in chain: + expected = _apply_oracle(expected, mode, selection) + + view = source + for mode, selection in chain: + view = _apply_view(view, mode, selection) + + assembled = np.zeros(view.shape, dtype=view.dtype) + hits = np.zeros(view.shape, dtype=np.int64) + for part in view.parts(): + value = np.asarray(part.array.result()) + if len(view.shape) == 0: + # A zero-rank view is assembled the way `result()` assembles it: + # its single part still carries the collapsed correlated axis, so + # the value arrives with rank 1. + value = value.reshape(()) + assembled[part.out_selection] = value + np.add.at(hits, part.out_selection, 1) + + np.testing.assert_array_equal(assembled, np.asarray(expected), err_msg=f"{chain}") + np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64), err_msg=f"{chain}") + + +PARTITIONINGS_1D: list[Any] = [None, (1,), (2,), (5,)] + + +def test_a_slice_only_fancy_step_after_a_fancy_step_reads_real_data() -> None: + """`oindex[:, 2:8]` after an `oindex` narrows the view, it does not re-index it. + + The second step carries no coordinates, so it is not fancy-after-fancy: it + must compose like basic indexing. Applying its slices to the *broadcast* + axes of the first step's index array instead truncates that array to size 0, + which leaves the resolver with no parts to read and `result()` handing back + an unwritten buffer. + """ + base = np.arange(24).reshape(3, 8) + expected = base[np.ix_([0, 2], range(8))][:, 2:8] + + for parts in (None, (2, 4), (1, 8), (3, 3)): + view = ( + LazyArray(base).with_parts(parts).lazy.oindex[np.array([0, 2]), :].lazy.oindex[:, 2:8] + ) + assert view.shape == expected.shape, f"parts={parts}" + np.testing.assert_array_equal(np.asarray(view.result()), expected, err_msg=f"{parts}") + + +def test_a_genuine_fancy_step_after_a_fancy_step_is_still_rejected() -> None: + """Coordinates landing on a broadcast axis of an existing selection raise. + + The documented limit: a second fancy step may re-index the axes the first one + varies over, but not the axes it merely broadcasts along. + """ + base = np.arange(24).reshape(3, 8) + view = LazyArray(base).lazy.oindex[np.array([0, 2]), :] + with pytest.raises(NotImplementedError, match="fancy-after-fancy"): + view.lazy.oindex[:, np.array([1, 3])] + with pytest.raises(NotImplementedError, match="fancy-after-fancy"): + view.lazy.vindex[np.array([0, 1]), np.array([1, 3])] + + +def test_a_zero_length_axis_resolves_the_same_way_under_every_partitioning() -> None: + """A size-0 axis carries no dependency, so it cannot make a map correlated.""" + base = np.zeros((3, 0)) + expected = base[np.ix_([1, 0, 0], np.arange(0, dtype=int))] + + for parts in (None, ((1, 1, 1), ()), ((3,), ())): + view = ( + LazyArray(base) + .with_parts(parts) + .lazy.oindex[np.array([1, -3, -3]), :] + .lazy.oindex[:, :] + ) + assert view.shape == expected.shape, f"parts={parts}" + np.testing.assert_array_equal(np.asarray(view.result()), expected, err_msg=f"{parts}") + assert list(view.parts()) == [] + + +def test_an_empty_slice_of_a_length_one_correlated_axis_has_no_parts() -> None: + """`parts()` agrees with `result()` that an emptied view selects nothing. + + A correlated selection of exactly one point normalizes to an all-singleton + index array, indistinguishable by shape from an axis the map broadcasts + over — so a later slice that empties the domain leaves the array at size 1. + """ + base = np.arange(5) + mask = np.array([False, True, False, False, False]) + + for parts in PARTITIONINGS_1D: + view = LazyArray(base).with_parts(parts).lazy.vindex[mask].lazy[1:-2] + assert view.shape == (0,), f"parts={parts}" + assert list(view.parts()) == [], f"parts={parts}" + np.testing.assert_array_equal(np.asarray(view.result()), base[mask][1:-2]) + + +def test_an_empty_slice_of_a_length_one_vindex_pair_has_no_parts() -> None: + """The same emptied view, spelled with explicit coordinates over two axes.""" + base = np.arange(35).reshape(7, 5) + + for parts in (None, (2, 2), (7, 5)): + view = LazyArray(base).with_parts(parts).lazy.vindex[np.array([6]), np.array([0])].lazy[0:0] + assert view.shape == (0,), f"parts={parts}" + assert list(view.parts()) == [], f"parts={parts}" + np.testing.assert_array_equal( + np.asarray(view.result()), base[np.array([6]), np.array([0])][0:0] + ) + + def test_eager_getitem_returns_data(source: LazyArray) -> None: """`arr[...]` reads immediately, like `numpy.ndarray.__getitem__`.""" np.testing.assert_array_equal(np.asarray(source[1:3, ::2, -1]), reference()[1:3, ::2, -1]) diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index baecd9ada2..237e4271a0 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -5,7 +5,11 @@ from zarr_indexing.domain import IndexDomain from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap -from zarr_indexing.transform import IndexTransform, selection_to_transform +from zarr_indexing.transform import ( + IndexTransform, + array_map_dependent_axis, + selection_to_transform, +) class TestIndexTransformConstruction: @@ -577,6 +581,20 @@ def test_scalar_array_has_no_dependency(self) -> None: assert _array_map_dependency_axes(np.ones((1, 1), dtype=np.intp)) == () + def test_zero_length_axis_has_no_dependency(self) -> None: + """An axis of size 0 carries no dependency either: it selects nothing, so + the array does not vary along it any more than along a singleton.""" + from zarr_indexing.transform import _array_map_dependency_axes + + assert _array_map_dependency_axes(np.zeros((0, 4), dtype=np.intp)) == (1,) + assert _array_map_dependency_axes(np.zeros((3, 0), dtype=np.intp)) == (0,) + assert _array_map_dependency_axes(np.zeros((0, 1), dtype=np.intp)) == () + + def test_zero_length_axis_does_not_make_a_map_correlated(self) -> None: + """An empty orthogonal selection is legal, so it must classify as one.""" + m = ArrayMap(index_array=np.zeros((0, 4), dtype=np.intp), input_dimension=1) + assert array_map_dependent_axis(m) == 1 + class TestIntersectArrayMapClassification: """`_intersect` must distinguish orthogonal (outer-product) ArrayMaps from From b0bc23a800534d9c40f40c73c5a3939814306ede Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 17:10:35 +0200 Subject: [PATCH 15/32] fix(zarr-indexing): count a domain axis no output map depends on MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../src/zarr_indexing/chunk_resolution.py | 45 ++++++- .../src/zarr_indexing/lazy_array.py | 66 ++++++++-- .../tests/test_chunk_resolution.py | 11 +- .../zarr-indexing/tests/test_lazy_array.py | 120 ++++++++++++++++++ 4 files changed, 219 insertions(+), 23 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 9372aff028..edc03228ab 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -42,6 +42,7 @@ from zarr_indexing.transform import ( IndexTransform, _positional_slice, # pyright: ignore[reportPrivateUsage] + array_map_dependent_axis, ) # `_positional_slice` is a leading-underscore helper in `transform.py`, shared @@ -289,6 +290,17 @@ def _dimension_map_slice(m: DimensionMap, dim_lo: int, dim_hi: int) -> slice: return _positional_slice(m.offset + m.stride * dim_lo, dim_hi - dim_lo, m.stride) +def _array_input_dim(m: ArrayMap, out_dim: int) -> int: + """The input dimension an orthogonal `ArrayMap` places its values along.""" + axis = array_map_dependent_axis(m) + if axis is None: + raise ValueError( + f"output[{out_dim}] is an ArrayMap that varies over no input dimension; " + "a map with no dependency axis should have been collapsed to a ConstantMap" + ) + return axis + + def sub_transform_to_selections( sub_transform: IndexTransform, out_indices: OutIndices = None, @@ -312,9 +324,19 @@ def sub_transform_to_selections( ------- tuple `(chunk_selection, out_selection, drop_axes)` + + Notes + ----- + `out_selection` carries one entry per **input (domain)** dimension, not one + per output map. The two counts usually agree, but a domain dimension that no + output map depends on — a `vindex` broadcast axis whose partner a later basic + index consumed — has no map to be derived from and is filled in from the + domain. Emitting fewer entries than the buffer has axes would place the + values against the leading axes instead. """ inclusive_min = sub_transform.domain.inclusive_min exclusive_max = sub_transform.domain.exclusive_max + input_rank = sub_transform.input_rank # Orthogonal outer product: >= 2 ArrayMaps each bound to a distinct input # dimension. out_indices is a per-output-dim dict of surviving positions. The @@ -324,7 +346,7 @@ def sub_transform_to_selections( # size-1 in chunk space and squeezed out via drop_axes. if isinstance(out_indices, dict): chunk_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] - out_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + by_input_dim: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} drop_axes: list[int] = [] for out_dim, m in enumerate(sub_transform.output): if isinstance(m, ConstantMap): @@ -333,11 +355,15 @@ def sub_transform_to_selections( elif isinstance(m, DimensionMap): rng = np.arange(inclusive_min[m.input_dimension], exclusive_max[m.input_dimension]) chunk_arrays.append((m.offset + m.stride * rng).astype(np.intp)) - out_arrays.append(rng.astype(np.intp)) + by_input_dim[m.input_dimension] = rng.astype(np.intp) else: # ArrayMap idx = m.index_array.ravel() chunk_arrays.append((m.offset + m.stride * idx).astype(np.intp)) - out_arrays.append(out_indices[out_dim]) + by_input_dim[_array_input_dim(m, out_dim)] = out_indices[out_dim] + out_arrays = [ + by_input_dim.get(d, np.arange(inclusive_min[d], exclusive_max[d], dtype=np.intp)) + for d in range(input_rank) + ] return np.ix_(*chunk_arrays), np.ix_(*out_arrays), tuple(drop_axes) # Correlated (vindex) sub-transforms carry ArrayMaps with `input_dimension` @@ -375,11 +401,11 @@ def sub_transform_to_selections( return tuple(chunk_sel), (out_scatter,), () chunk_sel = [] # annotated in the correlated branch above (same function scope) - out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + out_by_dim: dict[int, slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = {} # Single-pass build for the basic / single-orthogonal-array cases. # ConstantMap dims are dropped (no out_sel entry). - for m in sub_transform.output: + for out_dim, m in enumerate(sub_transform.output): if isinstance(m, ConstantMap): chunk_sel.append(m.offset) elif isinstance(m, DimensionMap): @@ -387,7 +413,7 @@ def sub_transform_to_selections( dim_lo = inclusive_min[d] dim_hi = exclusive_max[d] chunk_sel.append(_dimension_map_slice(m, dim_lo, dim_hi)) - out_sel.append(slice(dim_lo, dim_hi)) + out_by_dim[d] = slice(dim_lo, dim_hi) else: # ArrayMap (orthogonal: full-rank, raveled to its 1-D fancy coords) idx = m.index_array.reshape(-1) if m.offset == 0 and m.stride == 1: @@ -395,6 +421,11 @@ def sub_transform_to_selections( else: chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) # Orthogonal ArrayMap: out_indices holds the surviving positions. - out_sel.append(out_indices if out_indices is not None else slice(0, idx.size)) + out_by_dim[_array_input_dim(m, out_dim)] = ( + out_indices if out_indices is not None else slice(0, idx.size) + ) + out_sel = [ + out_by_dim.get(d, slice(inclusive_min[d], exclusive_max[d])) for d in range(input_rank) + ] return tuple(chunk_sel), tuple(out_sel), () diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 28607654ee..9fbfd099ea 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -334,17 +334,49 @@ def _array_map_coords(m: ArrayMap) -> np.ndarray[Any, np.dtype[np.intp]]: return (m.offset + m.stride * m.index_array).astype(np.intp).reshape(-1) -def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: int) -> Any: +def _correlated_map_coords( + m: ArrayMap, broadcast_axes: list[int], broadcast_shape: tuple[int, ...], input_rank: int +) -> np.ndarray[Any, np.dtype[np.intp]]: + """One storage coordinate per point of the correlated block, flattened. + + A correlated `ArrayMap`'s index array carries the transform's full input + rank, with a singleton on every axis it does not vary over — including + broadcast axes it shares with the *other* correlated maps but is itself + constant along. Flattening it directly would then yield fewer coordinates + than there are points, so it is reduced to the broadcast block and + broadcast up to it explicitly. + """ + coords = (m.offset + m.stride * m.index_array).astype(np.intp) + if coords.ndim == input_rank: + # Drop the axes bound by a slice, which the map is singleton along. + # Removing size-1 axes by reshape preserves element order wherever they + # sit, so no transpose is needed. + coords = coords.reshape(tuple(coords.shape[axis] for axis in broadcast_axes)) + return np.ascontiguousarray(np.broadcast_to(coords, broadcast_shape)).reshape(-1) + + +def _restore_domain_axis_order( + result: Any, axis_input_dims: list[int], domain_shape: tuple[int, ...] +) -> Any: """Permute `result`'s axes into input-domain order, restoring dropped axes. `axis_input_dims[k]` is the input (domain) dimension that axis `k` of `result` corresponds to. Axes are permuted so that they appear in increasing - domain-dimension order, and any domain dimension no output map depends on - (only reachable via `newaxis`) is reinserted as a singleton. - - This is where NumPy's advanced-index placement rules are absorbed: whatever - order the gather produced, the lowered result always comes back in the - view's own axis order. + domain-dimension order, and any domain dimension no output map depends on is + reinserted at **its own extent**. + + An unreferenced dimension is not always a singleton. A `vindex` coordinate + array with a broadcast axis it does not vary over leaves that axis in the + domain; a later basic index that consumes the axis the array *does* vary + over collapses the map to a `ConstantMap` and leaves the broadcast axis + behind, with whatever extent the basic index gave it — including 0. Every + position along such an axis holds the same values, so it is restored by + repeating the block, and an extent of 0 restores an empty result rather than + fabricating a row. + + This is also where NumPy's advanced-index placement rules are absorbed: + whatever order the gather produced, the lowered result always comes back in + the view's own axis order. """ if len(set(axis_input_dims)) != len(axis_input_dims): raise NotImplementedError( @@ -354,10 +386,13 @@ def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: in order = sorted(range(len(axis_input_dims)), key=lambda k: axis_input_dims[k]) if order != list(range(len(order))): result = _transpose(result, tuple(order)) - covered = sorted(axis_input_dims) - for dim in range(rank): - if dim not in covered: - result = _expand_dims(result, dim) + covered = set(axis_input_dims) + for dim, extent in enumerate(domain_shape): + if dim in covered: + continue + result = _expand_dims(result, dim) + if extent != 1: + result = _take(result, np.zeros(extent, dtype=np.intp), axis=dim) return result @@ -425,7 +460,7 @@ def _lower_orthogonal(array: Any, transform: IndexTransform) -> Any: else: axis_input_dims.append(m.input_dimension) result = result[tuple(selection)] - return _restore_domain_axis_order(result, axis_input_dims, transform.input_rank) + return _restore_domain_axis_order(result, axis_input_dims, transform.domain.shape) def _lower_correlated(array: Any, transform: IndexTransform) -> Any: @@ -505,13 +540,16 @@ def _lower_correlated(array: Any, transform: IndexTransform) -> Any: for position in range(n_corr - 1, -1, -1): m = outputs[correlated_dims[position]] assert isinstance(m, ArrayMap) - flat_index = flat_index + _array_map_coords(m) * stride + flat_index = flat_index + ( + _correlated_map_coords(m, broadcast_axes, broadcast_shape, transform.input_rank) + * stride + ) stride *= corr_sizes[position] result = _take(result, flat_index, axis=0) result = _reshape(result, broadcast_shape + tail_shape) return _restore_domain_axis_order( - result, list(broadcast_axes) + residual_axis_dims, transform.input_rank + result, list(broadcast_axes) + residual_axis_dims, transform.domain.shape ) diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index 8ff3b19b6c..dc8846661b 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -387,14 +387,21 @@ def test_2d_correlated_vindex_diagonal_is_linear_in_points( class TestSubTransformToSelections: def test_constant_map(self) -> None: - """ConstantMap produces int selection + drop axis.""" + """ConstantMap produces an int chunk selection and no drop axis. + + The domain dimension here is referenced by no output map, so its + out-selection cannot come from a map: it is the whole axis, taken from + the domain. One entry per domain dimension is what `out[out_sel] = value` + needs — an empty tuple would leave the buffer's rank unaddressed and + place a lower-rank part against the leading axes. + """ t = IndexTransform( domain=IndexDomain.from_shape((10,)), output=(ConstantMap(offset=5),), ) chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) assert chunk_sel == (5,) - assert out_sel == () + assert out_sel == (slice(0, 10),) assert drop_axes == () def test_dimension_map_stride_1(self) -> None: diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 5a9a12e12b..2a8b7dc3f4 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -560,6 +560,30 @@ def _random_oindex(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[An return tuple(selection) +def _broadcast_singleton_axes( + rng: np.random.Generator, entries: list[Any], length: int +) -> list[Any]: + """Reshape 1-D coordinate arrays so the selection carries singleton axes. + + A coordinate array of shape `(1, n)` or `(n, 1)` contributes a broadcast axis + it does not vary over. That axis stays in the view's domain, and a later + basic index that consumes its partner leaves it referenced by no output map + at all — the shape that makes a broadcast axis and a genuine extent-1 axis + indistinguishable from the index array alone. + """ + rank = int(rng.integers(2, 4)) + reshaped: list[Any] = [] + for entry in entries: + if not isinstance(entry, np.ndarray): + reshaped.append(entry) + continue + varying = int(rng.integers(0, rank)) + reshaped.append( + entry.reshape(tuple(length if axis == varying else 1 for axis in range(rank))) + ) + return reshaped + + def _random_vindex(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: ndim = len(shape) count = int(rng.integers(1, ndim + 1)) @@ -581,6 +605,8 @@ def _random_vindex(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[An else rng.integers(-size, size, size=length) for size in sizes ] + if rng.random() < 0.35: + entries = _broadcast_singleton_axes(rng, entries, length) return (Ellipsis, *entries) if trailing else tuple(entries) @@ -772,6 +798,100 @@ def test_a_genuine_fancy_step_after_a_fancy_step_is_still_rejected() -> None: view.lazy.vindex[np.array([0, 1]), np.array([1, 3])] +# --------------------------------------------------------------------------- +# Domain dimensions no output map depends on +# --------------------------------------------------------------------------- +# +# A `vindex` coordinate array with a *singleton* broadcast axis leaves that axis +# in the view's domain while the map varies only over its partner. A later basic +# index that consumes the partner collapses the map to a `ConstantMap`, and the +# singleton axis survives with nothing referencing it. Every stage of resolution +# has to keep counting it: the lowered block needs the axis back at its true +# extent, and a part has to say where its values belong along it. + +UNREFERENCED_AXIS_PARTITIONINGS: list[Any] = [None, (1, 1, 1), (2, 2, 2), (3, 4, 5), (3, 1, 2)] + +# (id, view builder, NumPy oracle) over `np.arange(60).reshape(3, 4, 5)`. +UNREFERENCED_AXIS_CASES: list[ + tuple[str, Callable[[LazyArray], LazyArray], Callable[[Any], Any]] +] = [ + ( + "leading-singleton-row", + lambda a: a.lazy.vindex[np.array([[2, 0]])].lazy[:, 0], + lambda r: r[np.array([[2, 0]])][:, 0], + ), + ( + "trailing-singleton-column", + lambda a: a.lazy.vindex[np.array([[2], [0]])].lazy[0], + lambda r: r[np.array([[2], [0]])][0], + ), + ( + "repeated-coordinates-over-a-singleton", + lambda a: a.lazy.vindex[np.array([[1, 1]]), np.array([[3, 3]])].lazy[:, 0], + lambda r: r[np.array([[1, 1]]), np.array([[3, 3]])][:, 0], + ), + ( + "unreferenced-axis-emptied", + lambda a: a.lazy.vindex[np.array([[2, 0]])].lazy[0:0, 0], + lambda r: r[np.array([[2, 0]])][0:0, 0], + ), + ( + "partial-vindex-with-a-residual-slice", + lambda a: a.lazy.vindex[np.array([[2, 0]]), np.array([[1, 3]])].lazy[:, 0, 1:4], + lambda r: r[np.array([[2, 0]]), np.array([[1, 3]])][:, 0, 1:4], + ), +] + + +def unreferenced_axis_reference() -> np.ndarray[Any, np.dtype[np.int64]]: + return np.arange(60, dtype=np.int64).reshape(3, 4, 5) + + +@pytest.mark.parametrize( + ("build", "oracle"), + [case[1:] for case in UNREFERENCED_AXIS_CASES], + ids=[case[0] for case in UNREFERENCED_AXIS_CASES], +) +@pytest.mark.parametrize("parts", UNREFERENCED_AXIS_PARTITIONINGS) +def test_a_domain_axis_no_output_map_depends_on_still_resolves( + build: Callable[[LazyArray], LazyArray], + oracle: Callable[[Any], Any], + parts: Any, +) -> None: + """The value, the shape and the tiling all hold when an axis is unreferenced.""" + data = unreferenced_axis_reference() + expected = np.asarray(oracle(data)) + view = build(LazyArray(data).with_parts(parts)) + + assert view.shape == expected.shape + np.testing.assert_array_equal(np.asarray(view.result()), expected) + + hits = np.zeros(view.shape, dtype=np.int64) + assembled = np.zeros(view.shape, dtype=view.dtype) + for part in view.parts(): + assembled[part.out_selection] = np.asarray(part.array.result()) + np.add.at(hits, part.out_selection, 1) + np.testing.assert_array_equal(assembled, expected) + np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64)) + + +def test_an_unreferenced_domain_axis_of_extent_zero_stays_empty() -> None: + """An emptied broadcast axis must not be restored as a fabricated row. + + The lowered block has no axis for a dimension nothing depends on, so the + resolver puts one back. Putting it back at extent 1 invents a row of data + for a selection whose own `shape` says it is empty. + """ + data = np.arange(140, dtype=np.int64).reshape(7, 5, 4) + coords = np.array([[6], [3], [0]]) + for parts in (None, (1, 1, 1), (3, 2, 3), (7, 5, 4)): + view = LazyArray(data).with_parts(parts).lazy.vindex[coords, -4].lazy[0, 0:0] + expected = data[coords, -4][0, 0:0] + assert view.shape == expected.shape == (0, 4), f"parts={parts}" + np.testing.assert_array_equal(np.asarray(view.result()), expected, err_msg=f"{parts}") + assert list(view.parts()) == [] + + def test_a_zero_length_axis_resolves_the_same_way_under_every_partitioning() -> None: """A size-0 axis carries no dependency, so it cannot make a map correlated.""" base = np.zeros((3, 0)) From 22c784213280c82e61c4d663f59932a90985e3e8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 17:14:16 +0200 Subject: [PATCH 16/32] fix(zarr-indexing): a materialized view never hands back the source MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../zarr-indexing/src/zarr_indexing/grid.py | 10 +- .../src/zarr_indexing/lazy_array.py | 244 ++++++++++++++--- .../src/zarr_indexing/support.py | 28 +- .../zarr-indexing/tests/test_lazy_array.py | 256 ++++++++++++++++++ 4 files changed, 503 insertions(+), 35 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py index 5bb01dea0a..f8f349012f 100644 --- a/packages/zarr-indexing/src/zarr_indexing/grid.py +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -67,7 +67,9 @@ def __init__(self, sizes: Sequence[int]) -> None: for i, s in enumerate(normalized): if s <= 0: raise ValueError( - f"chunk sizes must be positive; got {s} at position {i} of {normalized}" + f"chunk sizes must be positive; got {s} at position {i} of " + f"{normalized}. A zero-length axis is spelled as no chunks " + f"at all: EdgeDimensionGrid(())" ) self.sizes = normalized # Exclusive prefix sums, length len(sizes) + 1: `_offsets[c]` is the @@ -247,5 +249,11 @@ def dimension_grids_from_chunks( f"per-axis chunk sizes for dimension {axis} sum to {total}, " f"but the array extent is {extent}" ) + if extent == 0 and all(size == 0 for size in sizes): + # A zero-length axis has no chunks. `(0,)` and `(0, 0)` say the same + # thing as `()`, and the uniform convention spells it `(n,)` for any + # `n`, so all three spellings are accepted rather than one of them + # being a positivity error. + sizes = () grids.append(EdgeDimensionGrid(sizes)) return tuple(grids) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 9fbfd099ea..036aa392e2 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -142,6 +142,7 @@ import json import math import operator +import uuid from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol @@ -735,6 +736,25 @@ def _shift_key(key: tuple[Any, ...], window: tuple[slice, ...] | None) -> tuple[ return tuple(shifted) +def _gatherable(block: Any) -> Any: + """The block the residual is applied to, as something that can be gathered from. + + The residual finishes the read with `take`, `reshape` and `transpose`, which + is more than the `BASIC` floor asks of a *source*: that floor is `shape`, + `dtype` and a `__getitem__` understanding integers and slices, and a source + meeting exactly it may hand back a block that understands no more. Such a + block is converted, so that the floor is a promise about the source alone. + + A block that is already a NumPy array, or that answers with an array-API + namespace of its own, is left where it is: the namespace is how a device + array keeps a read on its own device (see the module docstring's device + caveat), and converting it would defeat that. + """ + if isinstance(block, np.ndarray) or _namespace(block) is not None: + return block + return np.asarray(block) + + def _read( source: Any, window: tuple[slice, ...] | None, @@ -756,12 +776,13 @@ def _read( gathered = source.vindex if has_vindex else source # The key is flat so that a `vindex` accepting only one-dimensional # coordinate lists is served too; the shape is restored here. - return _reshape(gathered[_shift_key(points, window)], transform.domain.shape) + block = _gatherable(gathered[_shift_key(points, window)]) + return _reshape(block, transform.domain.shape) key, residual, n_arrays = _decompose(transform, support, has_oindex) key = _shift_key(key, window) block = source.oindex[key] if n_arrays > 0 and has_oindex else source[key] - return _lower(block, residual) + return _lower(_gatherable(block), residual) # --------------------------------------------------------------------------- # @@ -769,6 +790,37 @@ def _read( # --------------------------------------------------------------------------- # +def _wrapped_base(array: Any) -> Any: + """The concrete array at the bottom of a (possibly nested) wrapper.""" + while isinstance(array, LazyArray): + array = array.array + return array + + +def _detach(value: Any, source: Any) -> Any: + """Return `value` sharing no memory with the array `source` reads from. + + An unpartitioned read of a selection that lowers to basic slicing hands back + a *view* of the wrapped array, so writing to it would reach through and + change the source — while the same read under any partitioning returns a + fresh buffer. `result()` is answering "what does this view hold", not "lend + me the storage", and its answer must not depend on how the read was divided. + + NumPy is the namespace whose sharing can be settled here, from the arrays' + memory bounds; a foreign namespace whose basic indexing returns views is out + of reach, and its blocks are handed back as they come (see the module + docstring's device caveat). + """ + base = _wrapped_base(source) + if ( + isinstance(value, np.ndarray) + and isinstance(base, np.ndarray) + and np.may_share_memory(value, base) + ): + return value.copy() + return value + + def _partition_out_selection( sub_transform: IndexTransform, out_indices: Any, @@ -793,12 +845,52 @@ def _partition_out_selection( return tuple(np.unravel_index(np.asarray(flat, dtype=np.intp), out_shape)) +def _out_selection_cell_count(selection: tuple[Any, ...], out_shape: tuple[int, ...]) -> int: + """How many cells of an array of shape `out_shape` a `Partition.out_selection` writes. + + Counted from the selectors' own shapes, so nothing is read and no index + array is materialized. The selectors are slices and integer arrays: the + slices contribute their lengths, and the arrays broadcast against each other + exactly as NumPy's advanced indexing broadcasts them, whether they arrive as + an open mesh (`numpy.ix_`) or as parallel coordinates (`unravel_index`). + """ + total = 1 + array_shapes: list[tuple[int, ...]] = [] + if len(selection) != len(out_shape): + raise AssertionError( + f"a partition addressed {len(selection)} of the view's {len(out_shape)} " + "dimensions; this is a bug in zarr-indexing's partition walk" + ) + for selector, extent in zip(selection, out_shape, strict=True): + if isinstance(selector, slice): + start: Any = selector.start + stop: Any = selector.stop + step: Any = selector.step + if step is None and start is not None and stop is not None and 0 <= start <= stop: + # The shape a partition walk actually produces: a concrete, + # forward, in-bounds interval. Sized directly, so the common + # path allocates neither a tuple nor a range. + total *= int(stop) - int(start) + else: + total *= len(range(*selector.indices(extent))) + else: + array_shapes.append(tuple(int(s) for s in np.shape(selector))) + if len(array_shapes) > 0: + total *= math.prod(np.broadcast_shapes(*array_shapes)) + return total + + def _covers_whole_part(local_transform: IndexTransform, part_shape: tuple[int, ...]) -> bool: """Whether a part-local transform addresses every cell of its part. - Conservative: an `ArrayMap` could in principle enumerate a whole part, but - checking that costs more than the answer is worth, so a fancy-indexed axis - always reports incomplete. + Direction does not matter: a reversing view (`lazy[::-1]`) reads every cell + of the box, just back to front, so a stride of -1 covers a part exactly as a + stride of 1 does. Only the set of coordinates is asked about here, which is + the same question `bounding_box()` and `strides()` answer. + + Conservative in one place: an `ArrayMap` could in principle enumerate a whole + part, but checking that costs more than the answer is worth, so a + fancy-indexed axis always reports incomplete. """ domain = local_transform.domain for out_dim, m in enumerate(local_transform.output): @@ -807,12 +899,19 @@ def _covers_whole_part(local_transform: IndexTransform, part_shape: tuple[int, . if extent != 1 or m.offset != 0: return False elif isinstance(m, DimensionMap): - if m.stride != 1: + if abs(m.stride) != 1: return False d = m.input_dimension - if m.offset + domain.inclusive_min[d] != 0: - return False - if m.offset + domain.exclusive_max[d] != extent: + lo = domain.inclusive_min[d] + hi = domain.exclusive_max[d] + if hi <= lo: + # An empty walk covers only an empty part. + if extent != 0: + return False + continue + first = m.offset + m.stride * lo + last = m.offset + m.stride * (hi - 1) + if min(first, last) != 0 or max(first, last) != extent - 1: return False else: return False @@ -845,10 +944,13 @@ class Partition: selection to the block in memory. out_selection Where `array.result()` belongs in an array of the view's shape — a - NumPy index tuple, usable directly as `out[part.out_selection] = ...`. + NumPy index tuple with one entry per dimension of the view, usable + directly as `out[part.out_selection] = ...`. is_complete Whether the view covers the whole box. Useful to a writer deciding - between a blind overwrite and a read-modify-write. + between a blind overwrite and a read-modify-write. True for a reversing + view, which reads every cell of the box back to front; false for every + fancy-indexed axis, which is conservative rather than exact. """ base_coords: tuple[int, ...] @@ -864,14 +966,24 @@ class Partition: def _wrapped_token(array: Any) -> Any: - """A deterministic token for the wrapped array. - - Determinism scope, in order of preference: the array's own - `__dask_tokenize__`; `dask.base.tokenize` when dask is importable (imported - lazily — this package never requires it); otherwise a local fallback that - digests the contents of a small array and, above `_TOKEN_DIGEST_LIMIT`, - falls back to a structural description that is stable across processes but - not sensitive to the array's contents. + """A token for the wrapped array. + + In order of preference: the array's own `__dask_tokenize__`; + `dask.base.tokenize` when dask is importable (imported lazily — this package + never requires it); otherwise a local fallback that digests the contents of + a small array. + + The two environments do not agree, and neither is a translation of the + other: a token taken with dask installed is meaningless to a process without + it, and the reverse. A token is an identifier within one process, not a + portable name. + + Above `_TOKEN_DIGEST_LIMIT` the local fallback has nothing left to identify + the contents with — reading them is exactly what a token call must not do — + so it declines to claim equality at all and returns a value that matches + nothing, including itself. A cache keyed on it misses; the alternative, a + structural description, is a cache that hands one array's result to a + different array of the same shape and dtype. """ hook = getattr(array, "__dask_tokenize__", None) if hook is not None: @@ -892,18 +1004,22 @@ def _wrapped_token(array: Any) -> Any: shape = tuple(int(s) for s in getattr(array, "shape", ())) dtype = getattr(array, "dtype", None) structural = (type(array).__qualname__, shape, str(dtype)) + # A token nothing can equal, for when the contents cannot be identified. It + # is the shape and dtype that would otherwise be mistaken for an identity, + # so they are kept alongside it for a reader looking at a graph. + unidentified = (*structural, "unidentified", uuid.uuid4().hex) # Decide whether to digest the contents from the *declared* size. Measuring # it by converting first would read the whole array — a multi-gigabyte store # pulled into memory by a token call, which is the opposite of the point. itemsize = getattr(dtype, "itemsize", None) if not isinstance(itemsize, int) or itemsize * math.prod(shape) > _TOKEN_DIGEST_LIMIT: - return structural + return unidentified try: contents = np.ascontiguousarray(array) - # A token must never raise; an unreadable source keeps the structural token. + # A token must never raise; an unreadable source is simply unidentified. except Exception: - return structural + return unidentified return (*structural, hashlib.sha256(contents.tobytes()).hexdigest()) @@ -950,6 +1066,17 @@ class LazyArray: __slots__ = ("_array", "_parts", "_support", "_transform", "_window") def __init__(self, array: ArrayLike) -> None: + if isinstance(array, np.matrix): + # `np.matrix` keeps every result two-dimensional, so `m[1]` has shape + # `(1, n)` where every other array-like gives `(n,)`. A view's shape + # comes from the transform, which follows NumPy's rule, so the two + # disagree on every rank-reducing selection. Refused at the door + # rather than resolved into a shape the view did not promise. + raise TypeError( + "numpy.matrix cannot be wrapped: it never reduces rank, so a " + "view's shape and its result would disagree. Convert it first, " + "with numpy.asarray(m)." + ) shape = tuple(int(s) for s in array.shape) self._array = array self._window: tuple[slice, ...] | None = None @@ -1370,38 +1497,91 @@ def result(self) -> Any: """Materialize this view. Assembles the view from its `parts()`, reading each box once. A wrapper - with a single whole-array part skips the output buffer and lowers - directly to array operations. + with **no partitioning** — `with_parts(None)`, and the default for a + source that advertises none — skips the output buffer and lowers the + whole view directly to array operations. A partitioning with a single + whole-array box still goes through the buffer: the branch is on whether + a partitioning is in force, not on how many boxes it has. + + The result never shares memory with the wrapped array, so writing to it + is safe whichever of those two routes it came by. Returns ------- array An array of shape `self.shape`, identical whatever partitioning is in force. A partitioned view produces a NumPy array (see the module - docstring's device caveat); a single whole-array part produces - whatever the wrapped array's namespace produces. A view with a - zero-rank domain returns a zero-dimensional array, not a scalar. + docstring's device caveat); an unpartitioned one produces whatever + the wrapped array's namespace produces. A view with a zero-rank + domain returns a zero-dimensional array, not a scalar. + + Raises + ------ + AssertionError + If the partition walk does not cover the view. The output buffer is + uninitialized where nothing was written, so an incomplete walk is + reported rather than returned. """ if self._parts is None: - return _read(self._array, self._window, self._transform, self._support) + block = _read(self._array, self._window, self._transform, self._support) + return _detach(block, self._array) out_shape = self.shape - out = np.empty(out_shape, dtype=np.dtype(self.dtype)) - if math.prod(out_shape) > 0: + out = self._output_buffer(out_shape) + size = math.prod(out_shape) + if size > 0: + written = 0 for part in self.parts(): - value = np.asarray(part.array.result()) + # Counted before the scatter, so a part addressing the wrong + # number of axes is named rather than reported as a broadcast + # failure against the buffer. + written += _out_selection_cell_count(part.out_selection, out_shape) + value = np.asanyarray(part.array.result()) if len(out_shape) == 0: value = value.reshape(()) out[part.out_selection] = value + if written != size: + # The buffer is uninitialized where no part wrote, so a partition + # walk that does not tile the view exactly would otherwise hand + # back process memory dressed as data. The parts are disjoint by + # contract, so counting the cells each addresses is enough: + # a gap undercounts and an overlap overcounts. + raise AssertionError( + f"the partition walk addressed {written} of the view's {size} " + "cells; this is a bug in zarr-indexing's partition walk" + ) return out + def _output_buffer(self, out_shape: tuple[int, ...]) -> Any: + """The buffer `result()` scatters parts into. + + Deliberately uninitialized: every cell is written by exactly one part, + and `result()` verifies that before returning. A masked source gets a + masked buffer so that scattering preserves the mask; nothing else about + the wrapped array's own type survives a partitioned read (see the module + docstring's device caveat). + """ + dtype = np.dtype(self.dtype) + if isinstance(self._array, np.ma.MaskedArray): + return np.ma.masked_all(out_shape, dtype=dtype) + return np.empty(out_shape, dtype=dtype) + # -- protocols ---------------------------------------------------------- def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: + """Materialize the view as a NumPy array. + + The result never shares memory with the wrapped array, whatever `copy` + asks for: `result()` already detaches, so `copy=True` gets an array the + caller owns and `copy=None` gets the same one rather than a second + allocation. `copy=False` is refused, because materializing means reading + — the values do not exist as a NumPy array until this call makes them. + """ if copy is False: raise ValueError( "a LazyArray cannot be converted to a NumPy array without a " - "copy: materializing a view always produces new data" + "copy: a view is a description of a read, and the values only " + "exist once the read is made" ) return np.asarray(self.result(), dtype=dtype) diff --git a/packages/zarr-indexing/src/zarr_indexing/support.py b/packages/zarr-indexing/src/zarr_indexing/support.py index 1bd8371041..4165ab643e 100644 --- a/packages/zarr-indexing/src/zarr_indexing/support.py +++ b/packages/zarr-indexing/src/zarr_indexing/support.py @@ -22,7 +22,11 @@ A source can declare its own level by exposing a `__zarr_indexing_support__` attribute holding an `IndexingSupport`. That attribute is read defensively — a value of the wrong type, or an attribute access that raises, means "no -declaration" rather than an error. +declaration" rather than an error. A member of a *different* enum naming one of +these four levels (xarray's own `IndexingSupport`, whose members these are +borrowed from) is honored: falling through to inference would answer with a +level the source never asked for, and inference is the more permissive of the +two. Otherwise the level is inferred, conservatively: @@ -138,9 +142,29 @@ def declared_indexing_support(array: Any) -> IndexingSupport | None: Reads `__zarr_indexing_support__` defensively: a missing attribute, an attribute of the wrong type, and an attribute access that raises are all "no declaration". + + One wrong type is read rather than discarded: a member of a *different* + enum that names one of these four levels. This module borrows xarray's + taxonomy and its member names, and xarray's own `IndexingSupport` is the + obvious thing for a source to reach for. Discarding it would fall through to + inference, which for a source carrying `oindex`/`vindex` answers + `VECTORIZED` — a *more* permissive level than the source asked for, which is + the one direction detection must never take. """ declared = _read_attribute(array, SUPPORT_ATTRIBUTE) - return declared if isinstance(declared, IndexingSupport) else None + if isinstance(declared, IndexingSupport): + return declared + return _foreign_indexing_support(declared) + + +def _foreign_indexing_support(declared: Any) -> IndexingSupport | None: + """The `IndexingSupport` a member of a look-alike enum names, if any.""" + if not isinstance(declared, enum.Enum): + return None + try: + return IndexingSupport[declared.name] + except KeyError: + return None def infer_indexing_support(array: Any) -> IndexingSupport: diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 2a8b7dc3f4..7fa476b4c2 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -1825,3 +1825,259 @@ def test_the_strict_doubles_notice_an_over_broad_key() -> None: Outer1VectorSource(data).oindex[np.array([0, 1]), np.array([0, 1]), slice(None)] with pytest.raises(TypeError, match="coordinate arrays only"): VectorizedSource(data).vindex[np.array([0, 1]), slice(None), slice(None)] + + +# --------------------------------------------------------------------------- +# Materializing +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("parts", [None, (2, 2, 2), SHAPE]) +@pytest.mark.parametrize( + ("build", "description"), + [ + (lambda a: a.lazy[1:3, :, :], "a basic slice"), + (lambda a: a.lazy[:, :, :], "the whole array"), + (lambda a: a.lazy[::-1, :, :], "a reversal"), + (lambda a: a.lazy.oindex[[2, 0], :, :], "a gather"), + ], +) +def test_materializing_never_hands_back_the_wrapped_array( + parts: Any, build: Callable[[LazyArray], LazyArray], description: str +) -> None: + """Writing to a materialized result must never reach the source. + + An unpartitioned read of a basic selection can be answered with a *view* of + the wrapped array, and NumPy 2 hands whatever `__array__` returns straight to + the caller. Every route out of the wrapper detaches, so the answer does not + depend on how the read happened to be divided. + """ + data = reference() + view = build(LazyArray(data).with_parts(parts)) + + for materialize in ( + lambda v: v.result(), + lambda v: np.array(v, copy=True), + lambda v: np.asarray(v), + lambda v: np.array(v), + ): + before = data.copy() + materialized = np.asarray(materialize(view)) + assert not np.shares_memory(materialized, data), description + materialized[...] = -1 + np.testing.assert_array_equal(data, before, err_msg=description) + + +def test_an_eager_getitem_never_hands_back_the_wrapped_array() -> None: + data = reference() + block = LazyArray(data)[1:3] + block[...] = -1 + np.testing.assert_array_equal(data, reference()) + + +def test_result_refuses_to_return_a_partly_written_buffer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A partition walk that leaves a gap must raise, not return process memory. + + `result()` scatters into an uninitialized buffer, which is only safe because + the parts tile the view. This is the guard that turns any future break of + that contract into a failure instead of into plausible-looking numbers. + """ + view = LazyArray(reference()).with_parts(PART_SHAPE).lazy[:, 1:, :] + complete = LazyArray.parts + + def drop_one(self: LazyArray) -> Any: + return list(complete(self))[:-1] + + monkeypatch.setattr(LazyArray, "parts", drop_one) + with pytest.raises(AssertionError, match="partition walk addressed"): + view.result() + + +def test_result_refuses_a_partition_of_the_wrong_rank(monkeypatch: pytest.MonkeyPatch) -> None: + """A part addressing fewer axes than the view has is caught by name.""" + view = LazyArray(reference()).with_parts(PART_SHAPE).lazy[:, 1:, :] + complete = LazyArray.parts + + def truncate(self: LazyArray) -> Any: + from dataclasses import replace + + return [replace(part, out_selection=part.out_selection[:-1]) for part in complete(self)] + + monkeypatch.setattr(LazyArray, "parts", truncate) + with pytest.raises(AssertionError, match="of the view's 3 dimensions"): + view.result() + + +class DuckBlock: + """An array-like meeting exactly the documented `BASIC` floor and no more. + + `shape`, `dtype`, and a `__getitem__` that understands integers and slices. + Anything else — an integer array, `take`, `reshape` — raises, and indexing it + yields another one of itself, so a block that comes back from it is as + limited as the source was. + """ + + __zarr_indexing_support__ = IndexingSupport.BASIC + + def __init__(self, data: np.ndarray[Any, Any]) -> None: + self._data = data + + @property + def shape(self) -> tuple[int, ...]: + return self._data.shape + + @property + def dtype(self) -> Any: + return self._data.dtype + + def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: + return np.array(self._data, dtype=dtype, copy=True if copy is None else copy) + + def __getitem__(self, key: Any) -> DuckBlock: + selectors = key if isinstance(key, tuple) else (key,) + for selector in selectors: + if not isinstance(selector, (int, np.integer, slice)): + raise TypeError(f"basic indexing only, got {selector!r}") + return DuckBlock(self._data[key]) + + +@pytest.mark.parametrize( + ("build", "oracle"), + [ + (lambda a: a.lazy[1:5, ::2, :], lambda r: r[1:5, ::2, :]), + ( + lambda a: a.lazy.oindex[[4, 0, 0], :, :], + lambda r: r[np.ix_([4, 0, 0], range(5), range(4))], + ), + (lambda a: a.lazy.vindex[[4, 0], [1, 1]], lambda r: r[[4, 0], [1, 1]]), + (lambda a: a.lazy[::-1, :, :], lambda r: r[::-1, :, :]), + ], + ids=["basic", "oindex", "vindex", "reversal"], +) +@pytest.mark.parametrize("parts", [None, PART_SHAPE]) +def test_a_source_meeting_only_the_basic_floor_resolves_any_selection( + build: Callable[[LazyArray], LazyArray], oracle: Callable[[Any], Any], parts: Any +) -> None: + """The floor is a promise about the source; its blocks are coerced, not trusted.""" + expected = np.asarray(oracle(reference())) + view = build(LazyArray(DuckBlock(reference())).with_parts(parts)) + assert view.shape == expected.shape + np.testing.assert_array_equal(np.asarray(view.result()), expected) + + +def test_the_duck_block_double_refuses_a_fancy_key() -> None: + """A negative control: the floor test only means something if the double bites.""" + with pytest.raises(TypeError, match="basic indexing only"): + DuckBlock(reference())[np.array([1, 0])] + + +# --------------------------------------------------------------------------- +# Sources with their own opinions +# --------------------------------------------------------------------------- + + +@pytest.mark.filterwarnings("ignore::PendingDeprecationWarning") +def test_numpy_matrix_is_refused() -> None: + """`np.matrix` never reduces rank, so a view's shape could not be honored.""" + with pytest.raises(TypeError, match="numpy.matrix cannot be wrapped"): + LazyArray(np.matrix(np.arange(12).reshape(3, 4))) + + +@pytest.mark.parametrize("parts", [None, (2, 2), (1, 4), (3, 4)]) +def test_a_masked_source_keeps_its_mask_under_every_partitioning(parts: Any) -> None: + data = np.ma.masked_greater(np.arange(12).reshape(3, 4), 7) + got = LazyArray(data).with_parts(parts).lazy[:, 1:].result() + expected = data[:, 1:] + assert isinstance(got, np.ma.MaskedArray), parts + np.testing.assert_array_equal(np.ma.getmaskarray(got), np.ma.getmaskarray(expected)) + np.testing.assert_array_equal(np.ma.filled(got, 0), np.ma.filled(expected, 0)) + + +def test_a_look_alike_enum_declaration_is_honored_not_upgraded() -> None: + """xarray's `IndexingSupport` has these member names; it must not read as absent. + + Falling through to inference would answer `VECTORIZED` for this source, which + carries the `oindex`/`vindex` pair — a *more* permissive level than the one + it declared. + """ + import enum as _enum + + foreign = _enum.Enum("IndexingSupport", ["BASIC", "OUTER", "OUTER_1VECTOR", "VECTORIZED"]) + + class Declaring(VectorizedSource): + __zarr_indexing_support__ = foreign.BASIC # type: ignore[assignment] + + source = Declaring(reference()) + assert LazyArray(source).indexing_support is IndexingSupport.BASIC + + +def test_an_unrecognized_foreign_enum_still_falls_back_to_inference() -> None: + import enum as _enum + + unrelated = _enum.Enum("Colour", ["RED", "GREEN"]) + + class Declaring(MinimalSource): + __zarr_indexing_support__ = unrelated.RED + + assert LazyArray(Declaring(reference())).indexing_support is IndexingSupport.BASIC + + +def test_a_large_array_without_dask_refuses_to_claim_equality( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Above the digest limit the fallback must miss a cache rather than lie. + + Two arrays differing in one element used to token identically, because the + fallback described the shape and dtype and gave up on the contents. + """ + import sys + + monkeypatch.setitem(sys.modules, "dask.base", None) + big = np.zeros(1 << 19, dtype=np.int64) + other = big.copy() + other[0] = 1 + + assert LazyArray(big).__dask_tokenize__() != LazyArray(other).__dask_tokenize__() + assert LazyArray(big).__dask_tokenize__() != LazyArray(big).__dask_tokenize__() + + # Below the limit the contents are digested, so equal data still tokens alike. + small = np.zeros(8, dtype=np.int64) + assert LazyArray(small).__dask_tokenize__() == LazyArray(small.copy()).__dask_tokenize__() + + +# --------------------------------------------------------------------------- +# Completeness and partition spellings +# --------------------------------------------------------------------------- + + +def test_a_reversing_view_covers_its_parts() -> None: + """A reversal reads every cell of every box, back to front.""" + data = np.arange(48).reshape(8, 6) + forward = LazyArray(data).with_parts((2, 2)).lazy[:, :] + reversed_view = LazyArray(data).with_parts((2, 2)).lazy[::-1, ::-1] + assert [part.is_complete for part in reversed_view.parts()] == [ + part.is_complete for part in forward.parts() + ] + assert all(part.is_complete for part in reversed_view.parts()) + + +def test_a_strided_reversal_is_still_incomplete() -> None: + data = np.arange(48).reshape(8, 6) + view = LazyArray(data).with_parts((2, 2)).lazy[::-2, :] + assert not any(part.is_complete for part in view.parts()) + + +@pytest.mark.parametrize("parts", [((0,), (3,)), ((), (3,)), (1, 1), ((0, 0), (3,))]) +def test_a_zero_length_axis_accepts_every_spelling_of_no_chunks(parts: Any) -> None: + """`(0,)`, `(0, 0)`, `()` and a uniform shape all describe an axis with no cells.""" + data = np.zeros((0, 3)) + view = LazyArray(data).with_parts(parts) + assert view.result().shape == (0, 3) + assert list(view.parts()) == [] + + +def test_a_zero_chunk_on_a_nonempty_axis_is_still_rejected() -> None: + with pytest.raises(ValueError, match="chunk sizes must be positive"): + LazyArray(np.zeros((4, 3))).with_parts(((0, 4), (3,))) From 807847381a1a87f365a832f5b7eb318861204a51 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 17:24:39 +0200 Subject: [PATCH 17/32] docs(zarr-indexing): correct claims a reviewer found false MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- packages/zarr-indexing/README.md | 4 ++ packages/zarr-indexing/changes/267.feature.md | 10 ++- packages/zarr-indexing/docs/design-notes.md | 63 ++++++++++++++----- packages/zarr-indexing/docs/index.md | 20 +++--- .../src/zarr_indexing/lazy_array.py | 46 +++++++++----- 5 files changed, 101 insertions(+), 42 deletions(-) diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index ccdfe595a5..2e7249c919 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -11,6 +11,10 @@ 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 - `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single output diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 9bf55a1ae6..14b20175e2 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -1,4 +1,4 @@ -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`; every other NumPy operation materializes the view through `__array__`. `__dask_tokenize__`, `__len__`, `__iter__`, `__bool__`, `__int__`, `__float__`, and `__index__` complete the surface, and a wrapper pickles as long as its base does. +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. @@ -13,7 +13,7 @@ Fixed the domain layout and partitioned resolution of a vectorized selection who 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: a NumPy array or an object carrying zarr's `oindex`/`vindex` pair reads as `VECTORIZED`, a source's own `__zarr_indexing_support__` attribute is honored when it holds an `IndexingSupport`, and everything else is `BASIC`, 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. +`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. @@ -23,4 +23,10 @@ Fixed `parts()` raising on a view emptied by a slice over an axis of extent 1: a Fixed the chunk selection `sub_transform_to_selections` builds for a negative-stride dimension: 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 `()`. + `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. diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index f73272ed81..cc4a60c69c 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -30,15 +30,24 @@ parts are intentionally identical: `tensorstore.IndexTransform(json=...)` and round-trips them back through our engine layer. -The representations differ in one place: index arrays. Ours are normalized to -the transform's full input rank — full-sized on the axis a map varies over, -singleton elsewhere — so the orthogonal/vectorized distinction is derivable -from the shape. TensorStore permits a lower-rank array that broadcasts -against the input domain; we accept those on load and normalize them. -`ArrayMap` also records the input dimension an *orthogonal* (`oindex`) array -varies over, a field TensorStore's format has no slot for, so -[the serializer](api/json.md) collapses it on the way out and reconstructs it on -the way in. +The representations differ in one place: index arrays. Both models want an index +array at the transform's full input rank, with singleton axes for the dimensions +a map does not vary over. TensorStore enforces it — its JSON parser rejects a +rank-1 array over a rank-2 domain outright, with `Index array for output +dimension 0 has rank 1 but must have rank 2` (checked against tensorstore +0.1.84) — while our loader is the more permissive of the two and also accepts a +lower-rank array that broadcasts against the input domain. That is a +compatibility affordance, not a difference in the model: ndsel leaves index-array +rank to [the engine layer](ndsel.md#lowering-to-a-transform), and everything the +algebra builds itself is at full rank. + +The reason full rank matters here is that we *derive* meaning from those +singletons rather than merely tolerating them: an array full-sized on one axis +and singleton elsewhere is orthogonal, and one varying over several shared axes +is vectorized, so the distinction is readable off the shape. `ArrayMap` also +records the input dimension an *orthogonal* (`oindex`) array varies over, a field +TensorStore's format has no slot for, so [the serializer](api/json.md) collapses +it on the way out and reconstructs it on the way in. Four deliberate differences: @@ -74,7 +83,8 @@ table of coordinates. It costs `O(n)` to store, it has no locality (the coordinates may repeat, reverse, or scatter arbitrarily), and intersecting it with a region means scanning it. `oindex`, `vindex`, and boolean masks all produce one, and once an axis is a query, subsequent basic indexing cannot make -it a box again. +it a box again. Nor can a second query be composed onto the axes an existing one +broadcasts along — see [Current scope](#current-scope). [ndsel](ndsel.md) encodes the same split in its message kinds: `point`, `box`, and `slice` desugar to constant and affine output maps and are always boxes; @@ -102,8 +112,10 @@ rectangles; a query cannot, and has to be resolved into a gather. A box can also be served as a single strided slab read, but the read has to be strided: reading its bounding box and discarding the rest transfers proportionally more data as soon as any stride exceeds 1. The two also behave differently under -partitioning: a box touches a contiguous run of parts, while a query can touch -any subset of them, in any order, more than once. +partitioning: a box touches a regularly-spaced run of parts, in increasing +order, each at most once — a stride larger than a part's extent skips parts +outright, so the run is not contiguous — while a query can touch any subset of +them, in any order, more than once. [`LazyArray`](api/lazy_array.md) exposes the category directly: @@ -139,9 +151,10 @@ Both are needed, because a box is dense in its hull only when every stride is 1. The slab above spans a 40x77 hull over the 40x20 cells it selects, so a consumer that issued one rectangular read of the hull and discarded the rest would transfer 3.85x the data. A query's hull is looser still and carries no -stride at all: 88 rows of hull over three selected rows. An empty selection -returns `None` from both, because it touches no coordinate to report an interval -around. +stride at all: 88 rows of hull over three selected rows. An empty *box* touches +no coordinate to report an interval around, so `bounding_box()` is `None` while +`strides()` still answers — the step is a property of the selection's shape, not +of the region it reaches. Only a query returns `None` from both. There is deliberately no separate `BoxView` type today. A statically-typed rectangular-only view is a plausible next step, but it should be introduced by @@ -179,8 +192,24 @@ re-bases every view to origin 0, so the positional dialect never exposes it; a caller working with `IndexTransform` directly will see it, and re-bases explicitly with `translate_domain_to` for NumPy-shaped coordinates. -Two limits remain, both intentional and expected to be lifted: - +Five limits remain, all intentional and all expected to be lifted. The first +three raise `NotImplementedError` at the point of use rather than returning +something approximate: + +- **Fancy after fancy.** A second `oindex`/`vindex`/mask step may re-index the + axes an existing index array *varies over*, but not the axes it merely + broadcasts along, so `lazy.oindex[[2, 0], :].lazy.oindex[:, [1, 3]]` raises. + Composing the two index arrays means either an outer product of lookup tables + or a joint gather, and which one is meant depends on how the axes line up. + A step spelled through `oindex` but carrying only slices is *not* a fancy + step: it narrows the view's own axes and composes like basic indexing. + *Planned.* +- **Diagonal views.** A transform whose output maps share an input dimension — + reachable by building one directly, not through `LazyArray`'s selection + surface — cannot be resolved. *Planned.* +- **Mixed correlated and orthogonal index arrays.** No single selection produces + a transform holding both a `vindex`-style and an `oindex`-style `ArrayMap`, so + resolving one is not implemented. *Planned.* - **Finite explicit bounds only.** `IndexDomain` has no implicit or unbounded dimensions; the message layer will normalize a body with `"-inf"`/`"+inf"` bounds, but the engine layer refuses to lower one into a transform. diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index 3f6f7afe19..e701840999 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -208,10 +208,12 @@ LazyArray(x).lazy.vindex[..., [3, 0]].result() Use a length-1 list (`oindex[[1]]`) to keep an axis. Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a `LazyArray` usable -as a source for `dask.array.from_array`. Every other NumPy operation — -arithmetic, `numpy.sum`, `numpy.stack` — materializes the whole view through -`__array__` and returns a plain NumPy array. Laziness here applies to indexing, -not to deferred computation. +as a source for `dask.array.from_array`. Any NumPy *function* given a view — +`numpy.sum(view)`, `numpy.add(view, 1)`, `numpy.stack([view, view])` — +materializes the whole thing through `__array__` and works on the result. +Python's arithmetic operators do not: `view + 1` raises `TypeError`, because the +wrapper defines no arithmetic dunders. Laziness here applies to indexing, not to +deferred computation. ### Parts @@ -254,7 +256,8 @@ view = lazy.lazy[10:50, ::4] ``` `result()` returns the same values under any of them: repartitioning changes how -the read is divided, not what it returns. Boxes that straddle the source's own +the read is divided, not what it returns, and the result never shares memory +with the wrapped array. Parts that do not line up with the source's own boxes (to bound peak memory, for example) are permitted; they cost extra I/O but do not affect correctness. @@ -285,8 +288,11 @@ a consumer that wants one rectangular read has to apply `strides()` too. The category is structural, not a heuristic: basic indexing always composes to a box, and one `oindex`, `vindex`, or mask anywhere in the chain makes the -selection a query permanently. The [design notes](design-notes.md) explain why -that matters to consumers of a selection. +selection a query permanently. A *second* fancy step is limited — it may +re-index the axes the first one varies over, but not the axes it broadcasts +along, and raises `NotImplementedError` if it tries. The +[design notes](design-notes.md) list that limit with the others and explain why +the category matters to consumers of a selection. ## Reference diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 036aa392e2..011b3eda79 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -81,10 +81,11 @@ which is what basic indexing composes to at any depth — or a **query**, an explicit list of coordinates, which is what `oindex`, `vindex`, and masks produce and which subsequent basic indexing cannot undo. `is_box` reports the -category and `bounding_box()` reports the storage region touched: exact for a -box, a hull for a query. The distinction is structural rather than an -optimization; [the design notes](../design-notes.md) describe why it matters to -consumers of a selection. +category and `bounding_box()` reports the storage region touched: the exact +interval per dimension for a box, a hull for a query. A box is only *dense* in +that interval when every entry of `strides()` is 1. The distinction is +structural rather than an optimization; [the design +notes](../design-notes.md) describe why it matters to consumers of a selection. The positional dialect ---------------------- @@ -119,21 +120,34 @@ Materializing on fallback ------------------------- `LazyArray` implements `__array__` but deliberately implements neither -`__array_ufunc__`/`__array_function__` nor `__array_namespace__`. Every NumPy -operation other than indexing therefore materializes the whole view: -`numpy.sum(view)`, `view + 1`, `numpy.stack([view, view])` all convert through -`__array__` first and return a plain NumPy array. That is intentional: laziness -here applies to indexing, not to building a deferred compute graph. It does mean -a `LazyArray` is not a drop-in for arithmetic on a large array. Use `.lazy[...]` -to narrow the view first, or pass the wrapper to `dask.array.from_array` so that -dask owns the compute graph. +`__array_ufunc__`/`__array_function__` nor `__array_namespace__`. A NumPy +*function* given a view therefore materializes the whole thing through +`__array__` and works on the resulting array: `numpy.sum(view)`, +`numpy.add(view, 1)` and `numpy.stack([view, view])` all do, and so does +`numpy.ones(view.shape) + view`, where the ndarray on the left dispatches. + +Python's arithmetic *operators* do not: `view + 1` raises `TypeError`, because +the wrapper defines no arithmetic dunders and an `int` has nothing to dispatch +to. Both facts follow from the same intent — laziness here applies to indexing, +not to building a deferred compute graph — and a `LazyArray` is not a drop-in +for arithmetic on a large array either way. Use `.lazy[...]` to narrow the view +first, or pass the wrapper to `dask.array.from_array` so that dask owns the +compute graph. Device caveat ------------- -Resolving a partitioned view builds its output buffer with NumPy, because the -resolver's scatter indices are host-side integer arrays. Wrapping a device array -that advertises parts therefore returns host memory. A single whole-array part -skips that buffer and stays in the wrapped array's own namespace. +Resolving a **partitioned** view builds its output buffer with NumPy, because +the resolver's scatter indices are host-side integer arrays. Wrapping a device +array that advertises parts therefore returns host memory, and any information +carried by the wrapped array's own type is lost with it — a `numpy.ma` mask is +the exception, kept by allocating a masked buffer. A wrapper with **no** +partitioning (`with_parts(None)`) skips that buffer and stays in the wrapped +array's own namespace. + +`result()` never returns memory shared with the wrapped array. That is settled +from the arrays' memory bounds, which only NumPy exposes: a foreign namespace +whose basic indexing returns views is beyond reach here, and an unpartitioned +read of such a source can hand back one of its views. """ from __future__ import annotations From bfeba4f77c231ade38e83a54088a8076e5c04100 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 17:40:10 +0200 Subject: [PATCH 18/32] fix(zarr-indexing): hold the full-rank invariant inside the engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The index-array rank was checked only from above, so a lower-rank array could exist inside the engine and be read for dependency axes it did not have. Nothing produced one: the tolerance was there for a test asserting compatibility with a body TensorStore itself rejects (verified against 0.1.84 — a rank-1 array over a rank-3 domain is an error in its JSON parser). Require the full input rank in the type, widen a lower-rank array at the JSON boundary where external input arrives, and give the test the shape TensorStore accepts. Assisted-by: ClaudeCode:claude-fable-5 --- .../zarr-indexing/src/zarr_indexing/json.py | 6 ++++ .../src/zarr_indexing/transform.py | 17 +++++----- packages/zarr-indexing/tests/test_json.py | 34 +++++++++++++++++-- 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py index 7be579c8bc..b0b29e1a9b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/json.py +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -321,6 +321,12 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: for i, om in enumerate(output_raw): if "index_array" in om: arr = _lower_index_array(om["index_array"], f"output[{i}].index_array") + # ndsel leaves index-array rank unvalidated, so an external producer + # may send an array of lower rank that broadcasts against the domain. + # Widen it here, on the way in, so every transform that exists holds + # the full-rank invariant the engine reads dependency axes from. + if arr.ndim < domain.ndim: + arr = arr.reshape((1,) * (domain.ndim - arr.ndim) + arr.shape) arrays[i] = arr dep = _array_map_dependency_axes(arr) array_axes[i] = dep diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index ec05e62aa5..fd00cb80c6 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -69,15 +69,14 @@ def __post_init__(self) -> None: f"output[{i}].input_dimension = {m.input_dimension} " f"is out of range for input rank {self.domain.ndim}" ) - elif isinstance(m, ArrayMap) and m.index_array.ndim > self.domain.ndim: - # ArrayMap index arrays produced by indexing and chunk resolution - # are normalized to the full input rank (an axis the array varies - # over is full-sized, every other axis a singleton). A rank - # *exceeding* the domain is always a bug. A rank *below* it is - # tolerated: TensorStore-format JSON (external input) may supply a - # lower-rank index array that broadcasts against the input domain, - # and `_array_map_dependency_axes` treats any missing leading axes - # as singleton dependencies. + elif isinstance(m, ArrayMap) and m.index_array.ndim != self.domain.ndim: + # An index array carries the transform's full input rank: the axis + # a map varies over is full-sized, every other axis a singleton. + # The rank is what makes the dependency axes readable from the + # shape, so a mismatch is a bug rather than a spelling. External + # JSON may use a lower-rank array that broadcasts against the + # domain; `transform_from_canonical` widens those on the way in, + # so the invariant holds for every transform that exists. raise ValueError( f"output[{i}].index_array has {m.index_array.ndim} dims " f"but input domain has {self.domain.ndim} dims" diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py index a348ef0fb8..40be46cf2e 100644 --- a/packages/zarr-indexing/tests/test_json.py +++ b/packages/zarr-indexing/tests/test_json.py @@ -260,7 +260,9 @@ def test_tensorstore_compatible_format(self) -> None: "output": [ {"offset": 5}, {"offset": 10, "stride": 2, "input_dimension": 1}, - {"offset": 0, "stride": 1, "index_array": [1, 2, 0]}, + # Full input rank, which is what TensorStore itself requires: + # it rejects a rank-1 array over a rank-3 domain outright. + {"offset": 0, "stride": 1, "index_array": [[[1, 2, 0]]]}, ], } t = index_transform_from_json(json) @@ -273,7 +275,7 @@ def test_tensorstore_compatible_format(self) -> None: assert t.output[1].stride == 2 assert t.output[1].input_dimension == 1 assert isinstance(t.output[2], ArrayMap) - np.testing.assert_array_equal(t.output[2].index_array, [1, 2, 0]) + np.testing.assert_array_equal(t.output[2].index_array, [[[1, 2, 0]]]) # Roundtrip json_rt = index_transform_to_json(t) @@ -398,3 +400,31 @@ def test_infinite_bound_rejected_on_lowering() -> None: } with pytest.raises(ValueError, match="infinite"): index_transform_from_json(body) + + +def test_a_lower_rank_index_array_is_widened_on_the_way_in() -> None: + """External JSON may broadcast a lower-rank array; the engine never holds one. + + ndsel leaves index-array rank unvalidated, so a conformant producer may send + an array of lower rank. It is widened at the boundary, which keeps the + full-rank invariant true of every transform the engine builds. + """ + body: IndexTransformJSON = { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [3, 4], + "output": [{"index_array": [1, 2, 0]}, {"input_dimension": 1}], + } + transform = index_transform_from_json(body) + array_map = transform.output[0] + assert isinstance(array_map, ArrayMap) + assert array_map.index_array.shape == (1, 3) + assert array_map.index_array.ndim == transform.domain.ndim + + +def test_an_index_array_of_the_wrong_rank_is_rejected() -> None: + """Inside the engine, a rank that does not match the domain is a bug.""" + with pytest.raises(ValueError, match="index_array has 1 dims"): + IndexTransform( + domain=IndexDomain.from_shape((3, 4)), + output=(ArrayMap(index_array=np.array([1, 2, 0], dtype=np.intp)),), + ) From 74c3c1ad1a7c8e90976046814d8de409c187edf0 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 18:08:14 +0200 Subject: [PATCH 19/32] fix(zarr-indexing): an index array spans the domain it is read over MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An index array axis must be the domain's extent or a singleton it broadcasts over. Any other size leaves input coordinates with no entry, which read as a smaller selection rather than as the error it is: the truncated array behind one of this review's silent-corruption bugs was a (3, 0) array over a (3, 2) domain, which this rejects at construction. Two fixtures carried the inconsistency they were meant to exercise — an empty array over a domain with room for two coordinates, and a widening case whose array covered three of four positions — and now describe domains their arrays span. Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/transform.py | 30 +++++++++++--- packages/zarr-indexing/tests/test_json.py | 41 +++++++++++++++---- 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index fd00cb80c6..0d8b2d8214 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -69,7 +69,7 @@ def __post_init__(self) -> None: f"output[{i}].input_dimension = {m.input_dimension} " f"is out of range for input rank {self.domain.ndim}" ) - elif isinstance(m, ArrayMap) and m.index_array.ndim != self.domain.ndim: + elif isinstance(m, ArrayMap): # An index array carries the transform's full input rank: the axis # a map varies over is full-sized, every other axis a singleton. # The rank is what makes the dependency axes readable from the @@ -77,10 +77,30 @@ def __post_init__(self) -> None: # JSON may use a lower-rank array that broadcasts against the # domain; `transform_from_canonical` widens those on the way in, # so the invariant holds for every transform that exists. - raise ValueError( - f"output[{i}].index_array has {m.index_array.ndim} dims " - f"but input domain has {self.domain.ndim} dims" - ) + if m.index_array.ndim != self.domain.ndim: + raise ValueError( + f"output[{i}].index_array has {m.index_array.ndim} dims " + f"but input domain has {self.domain.ndim} dims" + ) + # Every axis is either the domain's extent or a singleton it + # broadcasts over. Any other size addresses input coordinates the + # array has no entry for, which reads as a smaller selection + # rather than as the error it is. + bad = [ + (axis, size, extent) + for axis, (size, extent) in enumerate( + zip(m.index_array.shape, self.domain.shape, strict=True) + ) + if size not in (1, extent) + ] + if len(bad) > 0: + axis, size, extent = bad[0] + raise ValueError( + f"output[{i}].index_array has {size} entries on axis {axis}, " + f"which is neither 1 nor the domain's extent of {extent} " + f"(index_array shape {m.index_array.shape}, " + f"domain shape {self.domain.shape})" + ) @property def input_rank(self) -> int: diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py index 40be46cf2e..5ca9f30234 100644 --- a/packages/zarr-indexing/tests/test_json.py +++ b/packages/zarr-indexing/tests/test_json.py @@ -329,11 +329,11 @@ def test_slices_and_constants(self) -> None: assert _transforms_equal(rt, t) -def _index_array_body(index_array: Any, rank: int = 1) -> IndexTransformJSON: +def _index_array_body(index_array: Any, rank: int = 1, extent: int = 2) -> IndexTransformJSON: return { "input_rank": rank, "input_inclusive_min": [0] * rank, - "input_exclusive_max": [2] * rank, + "input_exclusive_max": [extent] * rank, "input_labels": [""] * rank, "output": [{"offset": 0, "stride": 1, "index_array": index_array}], } @@ -373,11 +373,18 @@ def test_a_ragged_index_array_is_rejected() -> None: @pytest.mark.parametrize( - ("index_array", "rank"), [([0, 1], 1), ([[0], [1]], 2), ([], 1)], ids=["1d", "2d", "empty"] + ("index_array", "rank", "extent"), + [([0, 1], 1, 2), ([[0], [1]], 2, 2), ([], 1, 0)], + ids=["1d", "2d", "empty"], ) -def test_an_integer_index_array_is_accepted(index_array: Any, rank: int) -> None: - """Integers of any nesting still lower, including an empty selection.""" - t = index_transform_from_json(_index_array_body(index_array, rank)) +def test_an_integer_index_array_is_accepted(index_array: Any, rank: int, extent: int) -> None: + """Integers of any nesting still lower, including an empty selection. + + An empty array selects nothing, so the domain it is read over is empty too; + a domain with room for coordinates the array does not supply is rejected + (see `test_an_index_array_that_does_not_span_its_domain_is_rejected`). + """ + t = index_transform_from_json(_index_array_body(index_array, rank, extent)) m = t.output[0] assert isinstance(m, ArrayMap) assert m.index_array.dtype == np.intp @@ -409,15 +416,17 @@ def test_a_lower_rank_index_array_is_widened_on_the_way_in() -> None: an array of lower rank. It is widened at the boundary, which keeps the full-rank invariant true of every transform the engine builds. """ + # A rank-1 array widens into the trailing axis, so it spans that axis's + # extent of four. body: IndexTransformJSON = { "input_inclusive_min": [0, 0], "input_exclusive_max": [3, 4], - "output": [{"index_array": [1, 2, 0]}, {"input_dimension": 1}], + "output": [{"index_array": [1, 2, 0, 2]}, {"input_dimension": 1}], } transform = index_transform_from_json(body) array_map = transform.output[0] assert isinstance(array_map, ArrayMap) - assert array_map.index_array.shape == (1, 3) + assert array_map.index_array.shape == (1, 4) assert array_map.index_array.ndim == transform.domain.ndim @@ -428,3 +437,19 @@ def test_an_index_array_of_the_wrong_rank_is_rejected() -> None: domain=IndexDomain.from_shape((3, 4)), output=(ArrayMap(index_array=np.array([1, 2, 0], dtype=np.intp)),), ) + + +def test_an_index_array_that_does_not_span_its_domain_is_rejected() -> None: + """An array with entries for only part of an axis is not a smaller selection. + + Reading it that way is how a truncated index array turned into a partially + written result rather than an error. + """ + with pytest.raises(ValueError, match="neither 1 nor the domain's extent"): + IndexTransform( + domain=IndexDomain.from_shape((3, 2)), + output=( + ArrayMap(index_array=np.zeros((3, 0), dtype=np.intp)), + DimensionMap(input_dimension=1), + ), + ) From 0ad17411ef3cd25ebe21d9c39e8706af95b3d8fa Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 18:27:53 +0200 Subject: [PATCH 20/32] test(zarr-indexing): a state machine for chained indexing, and the rank-0 part it found MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `zarr_indexing.testing`, behind a `testing` extra: a Hypothesis state machine that composes indexing steps onto a LazyArray and checks each step's shape, `result()`, and `parts()` assembly against NumPy, plus the selection strategies on their own. A project can point it at its own array by overriding one method. The machine asserts the documented assembly literally — a part's values must arrive at the shape its out_selection addresses — which is how it found the defect it also fixes: intersecting a correlated transform with a part's bounds collapsed the surviving broadcast block into one axis even when the block was already rank 0, so a view narrowed to a single point produced parts of rank 1. A rank-0 block now stays rank 0, and `result()` drops the reshape that was absorbing the mismatch. Merged from the branch that produced it, which predates the coverage guard in `result()`; the guard stays and the reshape it compensated with goes. Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-indexing/changes/267.feature.md | 5 + packages/zarr-indexing/docs/api/index.md | 11 + .../docs/api/testing_stateful.md | 5 + .../docs/api/testing_strategies.md | 5 + packages/zarr-indexing/docs/index.md | 31 + packages/zarr-indexing/mkdocs.yml | 2 + packages/zarr-indexing/pyproject.toml | 11 +- .../src/zarr_indexing/chunk_resolution.py | 9 +- .../src/zarr_indexing/lazy_array.py | 5 +- .../src/zarr_indexing/testing/__init__.py | 55 ++ .../src/zarr_indexing/testing/stateful.py | 362 +++++++++ .../src/zarr_indexing/testing/strategies.py | 165 ++++ .../src/zarr_indexing/transform.py | 29 +- .../zarr-indexing/tests/test_lazy_array.py | 75 +- .../tests/test_lazy_array_stateful.py | 62 ++ packages/zarr-indexing/uv.lock | 769 ++++++++++++++++++ 16 files changed, 1551 insertions(+), 50 deletions(-) create mode 100644 packages/zarr-indexing/docs/api/testing_stateful.md create mode 100644 packages/zarr-indexing/docs/api/testing_strategies.md create mode 100644 packages/zarr-indexing/src/zarr_indexing/testing/__init__.py create mode 100644 packages/zarr-indexing/src/zarr_indexing/testing/stateful.py create mode 100644 packages/zarr-indexing/src/zarr_indexing/testing/strategies.py create mode 100644 packages/zarr-indexing/tests/test_lazy_array_stateful.py create mode 100644 packages/zarr-indexing/uv.lock diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 14b20175e2..3e19c22a42 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -28,5 +28,10 @@ Fixed a domain dimension no output map depends on, which a `vindex` coordinate a `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 `array.result()` came back with rank 1 where the view had rank 0 and the documented `out[part.out_selection] = part.array.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. diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index fe112adb6b..0eb5882f73 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -57,6 +57,17 @@ and the wire format built on top of it. - [`zarr_indexing.errors`](errors.md) — the canonical index-error types, which `zarr.errors` re-exports by identity +**Test support** (needs the `testing` extra) + +- [`zarr_indexing.testing.stateful`](testing_stateful.md) — + `ChainedIndexingStateMachine`, a Hypothesis state machine that composes + indexing steps onto a `LazyArray` wrapping your array and checks every step + against NumPy, plus `apply_selection`, the NumPy model it checks against +- [`zarr_indexing.testing.strategies`](testing_strategies.md) — the selection + strategies the machine draws from, for a project that has its own harness + Every name listed in `zarr_indexing.__all__` is re-exported at the top level, so `from zarr_indexing import IndexTransform` and `from zarr_indexing.transform import IndexTransform` are equivalent. +`zarr_indexing.testing` is deliberately not among them: it imports +`hypothesis`, which the rest of the package does not. diff --git a/packages/zarr-indexing/docs/api/testing_stateful.md b/packages/zarr-indexing/docs/api/testing_stateful.md new file mode 100644 index 0000000000..0aeb57df29 --- /dev/null +++ b/packages/zarr-indexing/docs/api/testing_stateful.md @@ -0,0 +1,5 @@ +--- +title: testing.stateful +--- + +::: zarr_indexing.testing.stateful diff --git a/packages/zarr-indexing/docs/api/testing_strategies.md b/packages/zarr-indexing/docs/api/testing_strategies.md new file mode 100644 index 0000000000..1dd7d59457 --- /dev/null +++ b/packages/zarr-indexing/docs/api/testing_strategies.md @@ -0,0 +1,5 @@ +--- +title: testing.strategies +--- + +::: zarr_indexing.testing.strategies diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index e701840999..c7d3afc559 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -294,6 +294,37 @@ along, and raises `NotImplementedError` if it tries. The [design notes](design-notes.md) list that limit with the others and explain why the category matters to consumers of a selection. +## Checking your own array + +If your array is one `LazyArray` will be asked to read, you can check that it +holds up under composed selections without writing the generators yourself. +`zarr_indexing.testing` ships the Hypothesis state machine this package tests +itself with: it composes indexing steps onto a view of your array — basic, +orthogonal, vectorized, 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. + +```python +from zarr_indexing.testing import ChainedIndexingStateMachine, state_machine_test + +class MyArrayIndexing(ChainedIndexingStateMachine): + def make_source(self, data): + array = my_format.create(shape=data.shape, dtype=data.dtype) + array[:] = data + return array + +TestMyArrayIndexing = state_machine_test(MyArrayIndexing) +``` + +That is the whole subclass. `data`, `partitionings`, and `supports` are class +attributes if you want to widen or narrow what is drawn — `supports` in +particular names the [`IndexingSupport`](api/support.md) levels your source +really serves, and a chain read at one level must answer what the same chain +answers at every other. A failure shrinks to the shortest chain that causes it. + +Install with `pip install zarr-indexing[testing]`; the rest of the package does +not depend on Hypothesis. + ## Reference - [The ndsel wire format](ndsel.md) diff --git a/packages/zarr-indexing/mkdocs.yml b/packages/zarr-indexing/mkdocs.yml index 9a95be87a2..9541fa4d37 100644 --- a/packages/zarr-indexing/mkdocs.yml +++ b/packages/zarr-indexing/mkdocs.yml @@ -29,6 +29,8 @@ nav: - ' zarr_indexing.json': api/json.md - ' zarr_indexing.messages': api/messages.md - ' zarr_indexing.errors': api/errors.md + - ' zarr_indexing.testing.stateful': api/testing_stateful.md + - ' zarr_indexing.testing.strategies': api/testing_strategies.md - Changelog: https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/CHANGELOG.md watch: diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index b996833d9c..5bb615fdd8 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -34,6 +34,12 @@ dependencies = [ "numpy>=2", ] +[project.optional-dependencies] +# `zarr_indexing.testing` — a Hypothesis state machine and selection strategies +# for projects checking their own array against this package. Nothing else in +# the package imports hypothesis. +testing = ["hypothesis>=6.160.0"] + [project.urls] Homepage = "https://github.com/zarr-developers/zarr-python" Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-indexing" @@ -47,7 +53,10 @@ Documentation = "https://zarr-indexing.readthedocs.io/" # test suite, which already has zarr installed. `zarr` is intentionally NOT # listed here to avoid a workspace dependency cycle; run these tests from the # repo root (`uv run pytest packages/zarr-indexing/tests`), not in isolation. -test = ["pytest"] +# `hypothesis` arrives via the `testing` extra, which is what +# `zarr_indexing.testing` needs; the repo-root `test` group pins the exact +# version CI runs against. Bump the two together. +test = ["pytest", "hypothesis>=6.160.0"] docs = [ # Pins match the zarr-python docs environment in the repo-root # pyproject.toml so the two sites render with the same toolchain. diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index edc03228ab..c6068cf733 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -375,6 +375,13 @@ def sub_transform_to_selections( isinstance(m, ArrayMap) and m.input_dimension is None for m in sub_transform.output ) if correlated: + # The broadcast block is the input axis no `DimensionMap` binds. It is + # absent when the block is rank 0 (every coordinate a scalar), and then + # the coordinate arrays index as scalars: the block carries the residual + # slice axes alone, matching a domain that has no points axis either. + bound = {m.input_dimension for m in sub_transform.output if isinstance(m, DimensionMap)} + has_points_axis = any(d not in bound for d in range(sub_transform.input_rank)) + points_shape: tuple[int, ...] = (-1,) if has_points_axis else () chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] for m in sub_transform.output: if isinstance(m, ConstantMap): @@ -383,7 +390,7 @@ def sub_transform_to_selections( d = m.input_dimension chunk_sel.append(_dimension_map_slice(m, inclusive_min[d], exclusive_max[d])) else: # ArrayMap - idx = m.index_array.reshape(-1) + idx = m.index_array.reshape(points_shape) chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) # Chunk resolution always supplies the flat scatter index for a # correlated transform. Absent one (a bare sub-transform), fall back to an diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 011b3eda79..694b59cfc6 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -1550,10 +1550,7 @@ def result(self) -> Any: # number of axes is named rather than reported as a broadcast # failure against the buffer. written += _out_selection_cell_count(part.out_selection, out_shape) - value = np.asanyarray(part.array.result()) - if len(out_shape) == 0: - value = value.reshape(()) - out[part.out_selection] = value + out[part.out_selection] = np.asanyarray(part.array.result()) if written != size: # The buffer is uninitialized where no part wrote, so a partition # walk that does not tile the view exactly would otherwise hand diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py new file mode 100644 index 0000000000..babbe31623 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py @@ -0,0 +1,55 @@ +"""Test support for projects whose arrays are read through `LazyArray`. + +A Hypothesis state machine that composes indexing steps onto a `LazyArray` +wrapping your array and checks every step against NumPy +([`stateful`][zarr_indexing.testing.stateful]), and the selection strategies it +draws from, exported on their own for a project that has its own harness +([`strategies`][zarr_indexing.testing.strategies]). + +```python +from zarr_indexing.testing import ChainedIndexingStateMachine, state_machine_test + +class MyArrayIndexing(ChainedIndexingStateMachine): + def make_source(self, data): + array = my_format.create(shape=data.shape, dtype=data.dtype) + array[:] = data + return array + +TestMyArrayIndexing = state_machine_test(MyArrayIndexing) +``` + +This subpackage needs `hypothesis`, which the rest of `zarr_indexing` does not: +install it with the `testing` extra (`pip install zarr-indexing[testing]`). +""" + +from zarr_indexing.testing.stateful import ( + DEFAULT_DATA, + DEFAULT_PARTITIONINGS, + DEFAULT_SETTINGS, + ChainedIndexingStateMachine, + apply_selection, + outer_selection, + state_machine_test, +) +from zarr_indexing.testing.strategies import ( + basic_selections, + masks, + orthogonal_selections, + slice_selections, + vectorized_selections, +) + +__all__ = [ + "DEFAULT_DATA", + "DEFAULT_PARTITIONINGS", + "DEFAULT_SETTINGS", + "ChainedIndexingStateMachine", + "apply_selection", + "basic_selections", + "masks", + "orthogonal_selections", + "outer_selection", + "slice_selections", + "state_machine_test", + "vectorized_selections", +] diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py new file mode 100644 index 0000000000..60d9a9183d --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -0,0 +1,362 @@ +"""A stateful property test for indexing an array through `LazyArray`. + +`ChainedIndexingStateMachine` composes indexing steps onto a `LazyArray` +wrapping *your* array — `lazy[...]`, `lazy.oindex[...]`, `lazy.vindex[...]`, +each step applied to the view the last one produced — while applying the same +steps to a NumPy array holding the same values. After every step the view must +still agree with that model three ways: its shape, its `result()`, and the +assembly of its `parts()`. + +Point it at an array by subclassing and overriding `make_source`: + +```python +from zarr_indexing.testing import ChainedIndexingStateMachine, state_machine_test + +class MyArrayIndexing(ChainedIndexingStateMachine): + def make_source(self, data): + array = my_format.create(shape=data.shape, dtype=data.dtype) + array[:] = data + return array + +TestMyArrayIndexing = state_machine_test(MyArrayIndexing) +``` + +`data`, `partitionings`, and `supports` are class attributes; override any of +them to widen or narrow what is drawn. The base class needs no `make_source` at +all — left alone it wraps the NumPy array itself, which is a useful smoke test +of this package but says nothing about yours. + +What it is checking +------------------- +The parts invariant is the one with teeth. +[`Partition`][zarr_indexing.lazy_array.Partition] documents +`out[part.out_selection] = part.array.result()` as the assembly procedure, so +this checks that literally: a part's values must arrive at exactly the shape its +`out_selection` addresses — not merely a shape that broadcasts into it — land +there, and cover the view once. Checking through `result()` alone would prove +only that `result()` is self-consistent. + +The `declare_support` rule draws an +[`IndexingSupport`][zarr_indexing.support.IndexingSupport] level and applies it +to the view, so the level a source is read at becomes part of the chain. This is +the property a source author most needs pinned: the level chooses how much data +crosses the boundary, never what the read means, so a chain read at `BASIC` must +answer exactly what the same chain read at `VECTORIZED` answers. `BASIC` is +always honest — at that level `LazyArray` sends nothing but integers and slices +through `__getitem__` — so it is drawn for every source; the levels beyond it are +only as trustworthy as the accessors they name, which is why `supports` defaults +to the declared or inferred level alone and widening it is your call. + +Requires the `testing` extra (`pip install zarr-indexing[testing]`). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar + +import numpy as np +from hypothesis import HealthCheck, settings +from hypothesis import strategies as st +from hypothesis.stateful import RuleBasedStateMachine, initialize, invariant, precondition, rule + +from zarr_indexing.lazy_array import LazyArray +from zarr_indexing.support import IndexingSupport +from zarr_indexing.testing.strategies import ( + basic_selections, + orthogonal_selections, + slice_selections, + vectorized_selections, +) + +if TYPE_CHECKING: + from collections.abc import Sequence + + from zarr_indexing.boundary import SelectionMode + +__all__ = [ + "DEFAULT_DATA", + "DEFAULT_PARTITIONINGS", + "DEFAULT_SETTINGS", + "ChainedIndexingStateMachine", + "apply_selection", + "outer_selection", + "state_machine_test", +] + +DEFAULT_DATA = np.arange(7 * 5 * 4, dtype=np.int64).reshape(7, 5, 4) +"""The values the source holds by default: distinct, so a misplaced cell shows.""" + +DEFAULT_PARTITIONINGS: tuple[Any, ...] = ( + None, + (2, 2, 2), + (7, 5, 4), + (3, 2, 3), + ((3, 3, 1), (2, 2, 1), (3, 1)), + (4, 3, 3), +) +"""Partitionings to read under: a single whole-array part, uniform boxes of +several shapes (some of which do not divide the extent), and explicit per-axis +sizes. Boxes that straddle whatever the source declares are deliberate — they +cost extra I/O but must not change an answer.""" + +DEFAULT_SETTINGS = settings( + max_examples=250, + stateful_step_count=10, + deadline=None, + suppress_health_check=[ + HealthCheck.data_too_large, + HealthCheck.filter_too_much, + HealthCheck.too_slow, + ], +) +"""Enough examples to find a defect reachable only through a narrow chain, at a +few seconds per run when there is nothing to find. Every step is followed by +checks that each materialize the whole view, so the budget buys examples rather +than long chains — a chain runs out of axes to index within a few steps anyway. + +`filter_too_much` is suppressed because a chain that reaches a rank-0 view +leaves only `repartition` enabled, so a run that opens there is discarded.""" + + +# --------------------------------------------------------------------------- # +# The NumPy model +# --------------------------------------------------------------------------- # + + +def outer_selection(array: Any, selection: Sequence[Any]) -> Any: + """Apply an orthogonal selection to a NumPy array: the outer product of its axes. + + NumPy has no operator for this, so the model is built from `numpy.ix_`. + Scalar integers are basic indices — NumPy applies them first and drops the + axis — so they are peeled off before the outer product is formed. + """ + + def is_scalar(sel: Any) -> bool: + return isinstance(sel, (int, np.integer)) and not isinstance(sel, bool) + + scalars = tuple(sel if is_scalar(sel) else slice(None) for sel in selection) + reduced = array[scalars] + axes = [ + np.arange(size)[sel] + for size, sel in zip(reduced.shape, [s for s in selection if not is_scalar(s)], strict=True) + ] + if len(axes) == 0: + return reduced + return reduced[np.ix_(*axes)] + + +def apply_selection(array: Any, selection: tuple[Any, ...], mode: SelectionMode) -> Any: + """Apply a selection to a NumPy array in the given mode — the model a view is checked against. + + NumPy's own semantics *are* basic and vectorized indexing, so only the + orthogonal mode needs building (see `outer_selection`). + """ + if mode == "orthogonal": + return outer_selection(array, selection) + return array[selection] + + +def state_machine_test( + machine: type[RuleBasedStateMachine], *, config: settings = DEFAULT_SETTINGS +) -> Any: + """The pytest-collectable `TestCase` for a machine, with settings applied. + + Hypothesis builds a fresh `TestCase` per state-machine class, so settings + set on a base class do not reach a subclass's; this applies them where they + land. Assign the result to a module-level name beginning with `Test`. + """ + case = machine.TestCase + case.settings = config + return case + + +# --------------------------------------------------------------------------- # +# The machine +# --------------------------------------------------------------------------- # + +_SOURCE = "_zarr_indexing_cached_source" + + +class ChainedIndexingStateMachine(RuleBasedStateMachine): + """Indexing steps composed onto one `LazyArray`, against NumPy as the model. + + Subclass and override `make_source` to point it at your own array. See the + module docstring for the shape of that subclass and for what the invariants + check. + + Attributes + ---------- + data + The values the source holds, and the model every step is checked + against. Any shape and dtype NumPy supports; every axis must be + non-empty. + partitionings + Drawn once per run, before any indexing: `with_parts` is a pure setter + that carries through composition untouched and is read only when a view + resolves, so choosing it up front reaches the same states choosing it + mid-chain does, and spends the whole step budget on indexing. + supports + `IndexingSupport` levels `declare_support` may draw, beyond `BASIC`, + which is always drawn. `None` means the level `LazyArray` detects for + the source. Name a level here only if the source really serves it — a + level that overstates the source is a bug in the test, not in the code. + """ + + data: ClassVar[Any] = DEFAULT_DATA + partitionings: ClassVar[Sequence[Any]] = DEFAULT_PARTITIONINGS + supports: ClassVar[Sequence[IndexingSupport] | None] = None + + def make_source(self, data: Any) -> Any: + """Build the array under test, holding `data`. + + Called once per machine class and cached, not once per example: an + example is cheap and a source may not be. The machine only ever reads, + so the same object serves every run — but it must therefore not be + mutated by anything else while the test runs. + + The default returns `data` itself, so an unsubclassed machine exercises + this package against NumPy. + """ + return data + + def __init__(self) -> None: + super().__init__() + cls = type(self) + self.model: Any = np.asarray(cls.data) + source = cls.__dict__.get(_SOURCE) + if source is None: + source = self.make_source(self.model) + setattr(cls, _SOURCE, source) + self.view = LazyArray(source) + self.levels: tuple[IndexingSupport, ...] = _support_levels(self.view, cls.supports) + # Genuine fancy-after-fancy raises `NotImplementedError` by design, so a + # chain carries at most one fancy step — in either order relative to the + # basic ones. + self.fancy_used = False + self.chain: list[tuple[str, Any]] = [] + + def _indexable(self) -> bool: + """Whether there is anything left to index. + + A rank-0 or empty view takes no further step — NumPy would reject one + too — so the chain ends there, and the invariants keep checking. + """ + return self.model.ndim > 0 and self.model.size > 0 + + def _step(self, mode: SelectionMode, selection: tuple[Any, ...], *, fancy: bool) -> None: + self.chain.append((mode, selection)) + self.model = apply_selection(self.model, selection, mode) + if mode == "basic": + self.view = self.view.lazy[selection] + elif mode == "orthogonal": + self.view = self.view.lazy.oindex[selection] + else: + self.view = self.view.lazy.vindex[selection] + self.fancy_used = self.fancy_used or fancy + + # -- rules -------------------------------------------------------------- + + @initialize(data=st.data()) + def choose_partitioning(self, data: st.DataObject) -> None: + """Fix how the read is broken up, before any indexing.""" + parts = data.draw(st.sampled_from(list(type(self).partitionings))) + self.view = self.view.with_parts(parts) + self.chain.append(("parts", parts)) + + @precondition(lambda self: self._indexable()) + @rule(data=st.data()) + def basic(self, data: st.DataObject) -> None: + self._step("basic", data.draw(basic_selections(self.model.shape)), fancy=False) + + @precondition(lambda self: self._indexable() and not self.fancy_used) + @rule(data=st.data()) + def orthogonal(self, data: st.DataObject) -> None: + self._step("orthogonal", data.draw(orthogonal_selections(self.model.shape)), fancy=True) + + @precondition(lambda self: self._indexable() and not self.fancy_used) + @rule(data=st.data()) + def vectorized(self, data: st.DataObject) -> None: + self._step("vectorized", data.draw(vectorized_selections(self.model.shape)), fancy=True) + + @precondition(lambda self: self._indexable()) + @rule(data=st.data()) + def slices_only(self, data: st.DataObject) -> None: + """An `oindex` step carrying only slices is not a fancy selection. + + It narrows the view's own axes and composes like basic indexing, so it + is legal after a fancy step, where genuine coordinates are not. + """ + self._step("orthogonal", data.draw(slice_selections(self.model.shape)), fancy=False) + + @rule(data=st.data()) + def declare_support(self, data: st.DataObject) -> None: + """Read the rest of the chain at a different level. This must change no answer.""" + level = data.draw(st.sampled_from(list(self.levels))) + self.view = self.view.with_indexing_support(level) + self.chain.append(("support", level)) + + @precondition(lambda self: not self._indexable()) + @rule(data=st.data()) + def repartition(self, data: st.DataObject) -> None: + """Re-box a chain that has run out of axes to index. + + Something must stay enabled once the view is rank-0 or empty, or + Hypothesis has no move to make and abandons the run. Re-boxing is the + useful thing to do there: it changes nothing the invariants may see, and + a rank-0 view read through every partitioning is exactly the state a + collapsed correlated selection reaches. + """ + parts = data.draw(st.sampled_from(list(type(self).partitionings))) + self.view = self.view.with_parts(parts) + self.chain.append(("parts", parts)) + + # -- invariants --------------------------------------------------------- + + @invariant() + def the_view_has_the_models_shape(self) -> None: + assert self.view.shape == self.model.shape, self.chain + + @invariant() + def result_matches_the_model(self) -> None: + np.testing.assert_array_equal( + np.asarray(self.view.result()), self.model, err_msg=str(self.chain) + ) + + @invariant() + def parts_tile_the_view(self) -> None: + """The documented assembly, run literally. + + Every part's values arrive at exactly the shape its `out_selection` + addresses — not merely a shape that broadcasts into it — and together + the parts cover the view once. + """ + assembled = np.zeros(self.view.shape, dtype=self.view.dtype) + hits = np.zeros(self.view.shape, dtype=np.int64) + for part in self.view.parts(): + value = np.asarray(part.array.result()) + assert value.shape == assembled[part.out_selection].shape, ( + f"part {part.base_coords} carries {value.shape} for an out_selection " + f"addressing {assembled[part.out_selection].shape}: {self.chain}" + ) + assembled[part.out_selection] = value + np.add.at(hits, part.out_selection, 1) + + np.testing.assert_array_equal(assembled, self.model, err_msg=str(self.chain)) + np.testing.assert_array_equal( + hits, np.ones(self.view.shape, dtype=np.int64), err_msg=str(self.chain) + ) + + +def _support_levels( + view: LazyArray, declared: Sequence[IndexingSupport] | None +) -> tuple[IndexingSupport, ...]: + """The levels `declare_support` draws from, `BASIC` always among them. + + `IndexingSupport` is deliberately unordered — `OUTER_1VECTOR` sits between + `BASIC` and `OUTER` in capability — so this is a set of levels the source is + asserted to serve, not a ceiling to count down from. + """ + levels = list(declared) if declared is not None else [view.indexing_support] + if IndexingSupport.BASIC not in levels: + levels.insert(0, IndexingSupport.BASIC) + return tuple(dict.fromkeys(levels)) diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py new file mode 100644 index 0000000000..f49d3a360d --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py @@ -0,0 +1,165 @@ +"""Hypothesis strategies for the selections `LazyArray` accepts. + +Each strategy takes the shape of the array being indexed and generates one +selection for it — an index tuple with one entry per axis, in the spelling its +mode expects. They are the generators behind +[`ChainedIndexingStateMachine`][zarr_indexing.testing.stateful.ChainedIndexingStateMachine] +and are exported on their own for a project that has its own test harness and +wants only the hard part. + +```python +from hypothesis import given, strategies as st +from zarr_indexing.testing.strategies import basic_selections + +@given(selection=basic_selections((7, 5, 4))) +def test_my_array_slices_like_numpy(selection): + assert_array_equal(my_array[selection], reference[selection]) +``` + +Every axis of `shape` must be non-empty: a selection over an axis of extent 0 +has no coordinates to draw. Filter or narrow the shape before calling. + +Requires the `testing` extra (`pip install zarr-indexing[testing]`). +""" + +from __future__ import annotations + +import operator +from typing import TYPE_CHECKING, Any + +import numpy as np +from hypothesis import strategies as st + +if TYPE_CHECKING: + from collections.abc import Callable + +__all__ = [ + "basic_selections", + "masks", + "orthogonal_selections", + "slice_selections", + "vectorized_selections", +] + + +def _entries(shape: tuple[int, ...], entry: Callable[[int], st.SearchStrategy[Any]]) -> Any: + """One `entry` strategy per axis, as an index tuple.""" + return st.tuples(*[entry(size) for size in shape]) + + +def _basic_entry(size: int) -> st.SearchStrategy[Any]: + steps = st.integers(1, 3) + return st.one_of( + # A scalar integer drops its axis, in every mode, exactly as NumPy does. + st.integers(-size, size - 1), + st.builds(slice, st.integers(0, size), st.integers(0, size), steps), + # Downward. The start is drawn from below `-size` as well, where the walk + # begins off the front and selects nothing — a case that reads as an + # ordinary negative index but is empty — and a stop that falls off the + # front is spelled `None`. + st.builds( + slice, + st.integers(-2 * size - 1, size - 1), + st.none() | st.integers(0, size), + steps.map(operator.neg), + ), + st.just(slice(None)), + ) + + +def _orthogonal_entry(size: int) -> st.SearchStrategy[Any]: + coordinate = st.integers(-size, size - 1) + return st.one_of( + coordinate, + st.lists(coordinate, min_size=1, max_size=4), + masks((size,)), + st.builds(slice, st.integers(0, size - 1), st.just(size)), + ) + + +def _slice_entry(size: int) -> st.SearchStrategy[slice]: + return st.one_of( + st.builds(slice, st.integers(0, size - 1), st.just(size), st.integers(1, 2)), + st.just(slice(None, None, -1)), + st.just(slice(None)), + ) + + +@st.composite +def masks(draw: st.DrawFn, shape: tuple[int, ...]) -> np.ndarray[Any, np.dtype[np.bool_]]: + """Boolean masks over `shape`, each selecting at least one cell. + + An all-False mask is legal but is a separate concern — it empties the view, + and a chain of selections is more interesting when every step leaves + something to index — so one cell is always forced True. + """ + size = int(np.prod(shape)) + flags = np.array(draw(st.lists(st.booleans(), min_size=size, max_size=size))) + flags[draw(st.integers(0, size - 1))] = True + return flags.reshape(shape) + + +def basic_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]: + """Basic selections: one scalar integer or slice per axis. + + Slices run in both directions, including the two empty spellings — a + forward slice whose stop precedes its start, and a backward one whose start + is off the front of the axis. + """ + return _entries(shape, _basic_entry) + + +def orthogonal_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]: + """Orthogonal (`oindex`) selections: an outer product of per-axis choices. + + Each axis draws a scalar, a coordinate list (unsorted, with duplicates), a + boolean mask, or a slice. + """ + return _entries(shape, _orthogonal_entry) + + +def slice_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]: + """Selections of slices alone, for the `oindex` spelling that carries no coordinates. + + Such a step is not a fancy selection — it narrows the view's own axes and + composes like basic indexing — so it is legal after a fancy step, where + genuine coordinates are not. The starts reach past the origin, which is what + distinguishes a step that walks an existing index array's dependency axes + from one that walks its broadcast singletons. + """ + return _entries(shape, _slice_entry) + + +@st.composite +def vectorized_selections(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[Any, ...]: + """Vectorized (`vindex`) selections over a leading or trailing block of axes. + + `vindex` is coordinate-only — it rejects a slice outright — so a partial + selection names its axes by position: a leading block, or a trailing one + reached through an ellipsis. Either a single boolean mask spanning the whole + covered block, or one entry per axis, each a coordinate array or a scalar + (a scalar being a basic index NumPy applies before the coordinates). + """ + ndim = len(shape) + trailing = draw(st.booleans()) + count = draw(st.integers(1, ndim)) + axes = range(ndim - count, ndim) if trailing else range(count) + sizes = [shape[axis] for axis in axes] + + entries: list[Any] + if draw(st.booleans()): + entries = [draw(masks(tuple(sizes)))] + else: + length = draw(st.integers(1, 4)) + entries = [ + draw( + st.one_of( + st.integers(-size, size - 1), + st.lists(st.integers(-size, size - 1), min_size=length, max_size=length).map( + lambda values: np.array(values, dtype=np.intp) + ), + ) + ) + for size in sizes + ] + return (Ellipsis, *entries) if trailing else tuple(entries) diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 0d8b2d8214..1d6cc3d6c2 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -457,6 +457,12 @@ def _intersect_correlated( The surviving broadcast axes collapse to a single axis; the returned `out_indices` is the flat scatter index into the (row-major flattened) output buffer, of shape `(surviving_points,) + (residual slice sizes)`. + + A rank-0 broadcast block — every coordinate array a scalar, as after + `vindex[...]` narrowed to a single point — has no axis to collapse and stays + rank 0: the block either survives whole or the intersection is empty. The + result keeps only the residual slice axes and `out_indices` loses its leading + points axis, so the sub-transform's rank still matches the view's. """ # Mixing correlated and orthogonal ArrayMaps in one transform is not produced # by any single selection and is not supported here. @@ -528,17 +534,23 @@ def _intersect_correlated( for out_dim in correlated_dims } - # New domain: the collapsed broadcast axis, then one axis per residual slice. - new_min = [0] - new_max = [n_points] + # A rank-0 broadcast block contributes no axis: the leading `(n_points,)` of + # the domain, of every index array, and of `out_indices` is present only when + # there was a block to collapse. + points_shape = (n_points,) if len(broadcast_shape) > 0 else () + + # New domain: the collapsed broadcast axis if there is one, then one axis per + # residual slice. + new_min = [0] * len(points_shape) + new_max = list(points_shape) new_input_dim_of = {} - for new_axis, (d, nlo, nhi, _full, _m) in enumerate(slice_dims, start=1): + for new_axis, (d, nlo, nhi, _full, _m) in enumerate(slice_dims, start=len(points_shape)): new_min.append(nlo) new_max.append(nhi) new_input_dim_of[d] = new_axis new_domain = IndexDomain(inclusive_min=tuple(new_min), exclusive_max=tuple(new_max)) - corr_shape = (n_points,) + (1,) * n_slice + corr_shape = points_shape + (1,) * n_slice new_output: list[OutputIndexMap] = [] for out_dim, m in enumerate(transform.output): if out_dim in correlated_dims: @@ -583,14 +595,15 @@ def _intersect_correlated( ): point_offsets = point_offsets + coords_along_axis.astype(np.intp) * buffer_strides[axis] + n_lead = len(points_shape) out_indices: np.ndarray[Any, np.dtype[np.intp]] = point_offsets.reshape( - (n_points,) + (1,) * n_slice + points_shape + (1,) * n_slice ) for j in range(n_slice): d, nlo, nhi, _full, _m = slice_dims[j] coords = np.arange(nlo, nhi, dtype=np.intp) * buffer_strides[d] - shape = [1] * (1 + n_slice) - shape[1 + j] = coords.size + shape = [1] * (n_lead + n_slice) + shape[n_lead + j] = coords.size out_indices = out_indices + coords.reshape(shape) return (result, out_indices.astype(np.intp)) diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 7fa476b4c2..26b3a9d362 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -723,42 +723,12 @@ def test_random_chains_match_numpy(flavor: str) -> None: ) -@pytest.mark.parametrize("flavor", ["numpy-whole", "numpy-uniform-parts"]) -def test_random_chains_have_parts_that_tile_the_view(flavor: str) -> None: - """The seeded sweep, read through `parts()` rather than through `result()`. - - `result()` can absorb a defect that `parts()` cannot — an empty view assembles - correctly from no parts at all — so the iteration contract needs its own - sweep: every part places its own values, and together they cover the view - exactly once. - """ - rng = np.random.default_rng(20260731) - source = make_source(flavor) - - for _ in range(300): - chain = _random_chain(rng) - expected = reference() - for mode, selection in chain: - expected = _apply_oracle(expected, mode, selection) - - view = source - for mode, selection in chain: - view = _apply_view(view, mode, selection) - - assembled = np.zeros(view.shape, dtype=view.dtype) - hits = np.zeros(view.shape, dtype=np.int64) - for part in view.parts(): - value = np.asarray(part.array.result()) - if len(view.shape) == 0: - # A zero-rank view is assembled the way `result()` assembles it: - # its single part still carries the collapsed correlated axis, so - # the value arrives with rank 1. - value = value.reshape(()) - assembled[part.out_selection] = value - np.add.at(hits, part.out_selection, 1) - - np.testing.assert_array_equal(assembled, np.asarray(expected), err_msg=f"{chain}") - np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64), err_msg=f"{chain}") +# `result()` can absorb a defect that `parts()` cannot — an empty view assembles +# correctly from no parts at all, and a rank-0 one can be reshaped into place — +# so the iteration contract needs its own sweep, asserting the documented +# assembly literally rather than through `result()`. That sweep is the +# `ChainedIndexing` state machine in `test_lazy_array_stateful.py`, which drives +# the same operations and shrinks a failure to the chain that caused it. PARTITIONINGS_1D: list[Any] = [None, (1,), (2,), (5,)] @@ -939,6 +909,39 @@ def test_an_empty_slice_of_a_length_one_vindex_pair_has_no_parts() -> None: ) +def test_a_correlated_view_narrowed_to_one_point_has_parts_of_the_views_rank() -> None: + """A rank-0 broadcast block survives intersection without gaining an axis. + + `Partition` documents `out[part.out_selection] = part.array.result()` as the + assembly, so a part's values must arrive at the rank the view has — including + the rank 0 a correlated selection reaches once every coordinate is a scalar. + """ + base = np.arange(140).reshape(7, 5, 4) + cases = [ + # No residual slice: the whole view is the one gathered point. + ((np.array([-1]), np.array([-1]), np.array([2])), 0), + # One residual slice dimension survives alongside the collapsed block. + ((np.array([6, 1]), np.array([4, 0])), 1), + ] + + for parts in (None, (2, 2, 2), (3, 3, 3), (7, 5, 4)): + for selection, tail in cases: + expected = base[selection][tail] + view = LazyArray(base).with_parts(parts).lazy.vindex[selection].lazy[tail] + assert view.shape == expected.shape, f"parts={parts}, {selection}" + + assembled = np.zeros(view.shape, dtype=view.dtype) + hits = np.zeros(view.shape, dtype=np.int64) + for part in view.parts(): + assembled[part.out_selection] = np.asarray(part.array.result()) + np.add.at(hits, part.out_selection, 1) + + err = f"parts={parts}, {selection}" + np.testing.assert_array_equal(assembled, expected, err_msg=err) + np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64), err_msg=err) + np.testing.assert_array_equal(np.asarray(view.result()), expected, err_msg=err) + + def test_eager_getitem_returns_data(source: LazyArray) -> None: """`arr[...]` reads immediately, like `numpy.ndarray.__getitem__`.""" np.testing.assert_array_equal(np.asarray(source[1:3, ::2, -1]), reference()[1:3, ::2, -1]) diff --git a/packages/zarr-indexing/tests/test_lazy_array_stateful.py b/packages/zarr-indexing/tests/test_lazy_array_stateful.py new file mode 100644 index 0000000000..a50cd8f59c --- /dev/null +++ b/packages/zarr-indexing/tests/test_lazy_array_stateful.py @@ -0,0 +1,62 @@ +"""This package's own use of the state machine it exports. + +`ChainedIndexingStateMachine` composes indexing steps onto a `LazyArray` and +checks each step against NumPy — see +`zarr_indexing.testing.stateful` for what the invariants assert and why. + +Two sources: a NumPy array, which is the machine's default and reads at +`VECTORIZED`, and a real zarr array, whose partitioning is discovered from the +store rather than declared. The zarr case runs a smaller budget: it reads +through a store, and it exercises the same code paths. + +This replaces a seeded `_random_chain` sweep in `test_lazy_array` that read +chained selections through `parts()`. That sweep did reach the states it was +meant to, but a rank-0 correlated view was absorbed by a reshape in `result()` +and mirrored into the sweep rather than read as a failure; asserting the +documented assembly literally makes that impossible to paper over. +`test_lazy_array` keeps its `result()`-based sweep, which is the deterministic +cross-flavor coverage this does not attempt. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +from hypothesis import settings + +from zarr_indexing.support import IndexingSupport +from zarr_indexing.testing import ( + DEFAULT_SETTINGS, + ChainedIndexingStateMachine, + state_machine_test, +) + +TestNumpyIndexing = state_machine_test(ChainedIndexingStateMachine) + + +class ZarrIndexing(ChainedIndexingStateMachine): + """The same chains against a zarr array, read at every level it serves. + + A zarr array carries `oindex` and `vindex`, so `VECTORIZED` is what + `LazyArray` detects; naming the levels explicitly also draws the two outer + ones, which route an orthogonal request to `oindex` one array axis at a time. + """ + + supports = ( + IndexingSupport.BASIC, + IndexingSupport.OUTER_1VECTOR, + IndexingSupport.OUTER, + IndexingSupport.VECTORIZED, + ) + + def make_source(self, data: Any) -> Any: + zarr = pytest.importorskip("zarr") + array = zarr.create_array({}, shape=data.shape, chunks=(3, 2, 3), dtype=data.dtype) + array[:] = data + return array + + +TestZarrIndexing = state_machine_test( + ZarrIndexing, config=settings(DEFAULT_SETTINGS, max_examples=50) +) diff --git a/packages/zarr-indexing/uv.lock b/packages/zarr-indexing/uv.lock new file mode 100644 index 0000000000..ab554e375b --- /dev/null +++ b/packages/zarr-indexing/uv.lock @@ -0,0 +1,769 @@ +version = 1 +revision = 3 +requires-python = ">=3.12" + +[[package]] +name = "babel" +version = "2.18.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/b2/51899539b6ceeeb420d40ed3cd4b7a40519404f9baf3d4ac99dc413a834b/babel-2.18.0.tar.gz", hash = "sha256:b80b99a14bd085fcacfa15c9165f651fbb3406e66cc603abf11c5750937c992d", size = 9959554, upload-time = "2026-02-01T12:30:56.078Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/f5/21d2de20e8b8b0408f0681956ca2c69f1320a3848ac50e6e7f39c6159675/babel-2.18.0-py3-none-any.whl", hash = "sha256:e2b422b277c2b9a9630c1d7903c2a00d0830c409c59ac8cae9081c92f1aeba35", size = 10196845, upload-time = "2026-02-01T12:30:53.445Z" }, +] + +[[package]] +name = "backrefs" +version = "8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/56/4744bcd0c82184e80c52b0ac4076c261a8ffa1f1b343ff2f6e89ce0e1cef/backrefs-8.0.tar.gz", hash = "sha256:b556cd7d36c3a3a2f256b89590b176b8eddfb73bcfaee3a3ddd84ea66d21ce50", size = 7013081, upload-time = "2026-07-26T19:54:24.638Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/fd/9bf53b6a6f6f519ffaac765df2f2a25e5c2fc6d32cfd2b2747099e72c911/backrefs-8.0-py310-none-any.whl", hash = "sha256:4a627b817fd2dce43b79ab48da63613340509381cd8ce0897078a0bce79a2ab8", size = 380377, upload-time = "2026-07-26T19:54:17.457Z" }, + { url = "https://files.pythonhosted.org/packages/e1/29/4bd7ae72a2634da00379c2b3bcc5439e7c94620235c6afea8af15229a973/backrefs-8.0-py311-none-any.whl", hash = "sha256:f0c35cf0102ba6b6070c12a492be3c1c1d3f5839529784b9a9565d6d04569a01", size = 392169, upload-time = "2026-07-26T19:54:18.782Z" }, + { url = "https://files.pythonhosted.org/packages/29/13/232505664e8e2a0c7a2eb0c505cfade9d715538f89a5d62bc4c272968f62/backrefs-8.0-py312-none-any.whl", hash = "sha256:87f0fae8c5f207fe9f4b2887efc71d42f4900ac78faa1af08d675ef303692dc5", size = 398084, upload-time = "2026-07-26T19:54:19.954Z" }, + { url = "https://files.pythonhosted.org/packages/8a/69/47a3dc20abc4fa5486655fde681bd55e63211b46c886d8c02223d6468431/backrefs-8.0-py313-none-any.whl", hash = "sha256:601ce68ca12385dbda06ce264406b4c4210cf5b79fd0fd627592365c92f29a88", size = 400040, upload-time = "2026-07-26T19:54:21.194Z" }, + { url = "https://files.pythonhosted.org/packages/1c/cf/e5f9b68a5b0e939a2fb933a66c20180d0c9241bf8927f7a47fa48c1675e9/backrefs-8.0-py314-none-any.whl", hash = "sha256:9ec96efa080938be92323e8e730e57718c9c88eb15ad70bbef4e1766df591408", size = 411903, upload-time = "2026-07-26T19:54:23.221Z" }, +] + +[[package]] +name = "certifi" +version = "2026.7.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bd/2a/23f34ec9d04624958e137efdc394888716353190e75f25dd22c7a2c7a8aa/charset_normalizer-3.4.9.tar.gz", hash = "sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b", size = 152439, upload-time = "2026-07-07T14:34:58.454Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/70/4a/ecbd131485c07fcdfad54e28946d513e3da22ef3b4bd854dcafae54ec739/charset_normalizer-3.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:45b0cc4e3556cd875e09102988d1ab8356c998b596c9fced84547c8138b487a0", size = 319300, upload-time = "2026-07-07T14:33:15.666Z" }, + { url = "https://files.pythonhosted.org/packages/ec/96/5d9364e3342d69f3a045e1777bc47c85c383e6e9466d561b33fdb419d1f9/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b2aff1c7b3884512b9512c3eaadd9bab39fb45042ffaaa1dd08ff2b9f8109d9", size = 215802, upload-time = "2026-07-07T14:33:17.031Z" }, + { url = "https://files.pythonhosted.org/packages/4b/4c/5361f9aa7f2cb58d94f2ab831b3d493f69efb1d239654b4744e3c09527cb/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9104ed0bd76a429d46f9ec0dbc9b08ad1d2dcdf2b00a5a0daa1c145329b35b44", size = 237171, upload-time = "2026-07-07T14:33:18.576Z" }, + { url = "https://files.pythonhosted.org/packages/50/78/ce342ca4ff30b2eb49fe6d9578df85974f90c67d294113e94efdd9664cbd/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7b86a2b16095d250c6f58b3d9b2eee6f4147754344f3dab0922f7c9bf7d226c9", size = 233075, upload-time = "2026-07-07T14:33:20.084Z" }, + { url = "https://files.pythonhosted.org/packages/01/c4/4fa4c8b3097a11f3c5f09a35b72ed6855fb1d332469504962ab7bafcc702/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5e226f6218febc71f6c1fc2fafb91c226f75bdc1d8fb12d66823716e891608fd", size = 224256, upload-time = "2026-07-07T14:33:21.747Z" }, + { url = "https://files.pythonhosted.org/packages/87/3a/ad914516df7e358a81aae018caa5e0470ba827fa6d763b1d2e87d920a5f6/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:90c44bc373b7687f6948b693cceaea1348ae0975d7474746559494468e3c1d84", size = 208784, upload-time = "2026-07-07T14:33:23.313Z" }, + { url = "https://files.pythonhosted.org/packages/d7/74/3c12f9755717dfe5c5c87da63f35d765fa0c00382ec26bf23f7fae34f2ba/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9cdef90ae47919cae358d8ab15797a800ed41da7aba5d72419fb510729e2ed4b", size = 219928, upload-time = "2026-07-07T14:33:24.814Z" }, + { url = "https://files.pythonhosted.org/packages/33/9a/895095b83e7907abd6d3d99aad3a38ad0d9686cc186cb0c94c24320fe63e/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:60f44ade2cf573dad7a277e6f8ca9a51a21dda572b13bd7d8539bb3cd5dbedde", size = 218489, upload-time = "2026-07-07T14:33:26.42Z" }, + { url = "https://files.pythonhosted.org/packages/a1/34/ef5c05f412f42520d7709b7d3784d19640839eb7366ded1755511585429f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a1786910334ed46ab1dd73222f2cd1e05c2c3bb39f6dddb4f8b36fc382058a39", size = 210267, upload-time = "2026-07-07T14:33:27.952Z" }, + { url = "https://files.pythonhosted.org/packages/83/dc/9b29fa4412b318bf3bfea985c35d67eb55e04b59a7c3f2237168b0e0be6f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:03d07803992c6c7bbc976327f34b18b6160327fc81cb82c9d504720ac0be3b62", size = 226030, upload-time = "2026-07-07T14:33:29.397Z" }, + { url = "https://files.pythonhosted.org/packages/0e/42/6dbc00b8cd16011691203e33570fa42ed5746599a2e878112d16eab403a3/charset_normalizer-3.4.9-cp312-cp312-win32.whl", hash = "sha256:78841cccf1af7b40f6f716338d50c0902dbe88d9f800b3c973b7a9a0a693a642", size = 151185, upload-time = "2026-07-07T14:33:30.781Z" }, + { url = "https://files.pythonhosted.org/packages/80/cc/f920afd1a23c58ccd53c1d36085a71893a4737ff5e66e0371efab6809850/charset_normalizer-3.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:4b3dac63058cc36820b0dd072f89898604e2d39686fe05321729d00d8ac185a0", size = 162557, upload-time = "2026-07-07T14:33:32.176Z" }, + { url = "https://files.pythonhosted.org/packages/f0/e6/0386d43a261ff4e4b30c5857af7df877254b46bec7b9d1b74b6bf969a90b/charset_normalizer-3.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:78fa18e436a1a0e58dbd7e02fc4473f3f32cceb12df9dfca542d075961c307d2", size = 152665, upload-time = "2026-07-07T14:33:33.711Z" }, + { url = "https://files.pythonhosted.org/packages/b2/06/97ec2aeae780b31d742b6352218b43841a6871e2564578ca522dce4a45c3/charset_normalizer-3.4.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:440eede837960000d74978f0eba527be106b5b9aee0daf779d395276ed0b0614", size = 317688, upload-time = "2026-07-07T14:33:35.408Z" }, + { url = "https://files.pythonhosted.org/packages/d0/39/8ff066c672434225f8d25f8b739f992af250944392173dcc88362681c9bf/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21e764fd1e70b6a3e205a0e46f3051701f98a8cb3fad66eeb80e48bb502f8698", size = 214982, upload-time = "2026-07-07T14:33:36.996Z" }, + { url = "https://files.pythonhosted.org/packages/92/8f/3a47a3667c83c2df9483d91644c6c107de3bf8874aa1793da9d3012eb986/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e4fd89cc178bced6ad29cb3e6dd4aa63fa5017c3524dbd0b25998fb64a87cc8b", size = 236460, upload-time = "2026-07-07T14:33:38.536Z" }, + { url = "https://files.pythonhosted.org/packages/f1/60/b22cdbee7e4013dab8b0d7647fc6181120fbbbc8f7025c226d15bd5a47fc/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bd47ba7fc3ca94896759ea0109775132d3e7ab921fbf54038e1bab2e46c313c9", size = 232003, upload-time = "2026-07-07T14:33:40.059Z" }, + { url = "https://files.pythonhosted.org/packages/ea/f8/72eb13dcabe7257035cea8aefd922caad2f110d252bf9f67c4c2ca763aee/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:84fd18bcc17526fc2b3c1af7d2b9217d32c9c04448c16ec693b9b4f1985c3d33", size = 223149, upload-time = "2026-07-07T14:33:41.631Z" }, + { url = "https://files.pythonhosted.org/packages/b0/3e/faee8f9de92b14ee1198e9163252bb15efee7301b31256a3b6d9ebfdd0dd/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:5b10cd92fc5c498b35a8635df6d5a100207f88b63a4dc1de7ef9a548e1e2cd63", size = 207901, upload-time = "2026-07-07T14:33:43.209Z" }, + { url = "https://files.pythonhosted.org/packages/3a/25/45f30093ae27dd7b92a793b61882a38685f993700113ca36e0c9c14965e1/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a4fbdde9dd4a9ce5fd52c2b3a347bb50cc89483ef783f1cb00d408c13f7a96c0", size = 219176, upload-time = "2026-07-07T14:33:44.725Z" }, + { url = "https://files.pythonhosted.org/packages/48/18/c8f397329c35e32f6a837e488986f4ae03bd2abebc453b48714991630c2f/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:416c229f77e5ea25b3dfd4b582f8d73d7e43c22320302b9ab128a2d3a0b38efe", size = 217356, upload-time = "2026-07-07T14:33:46.192Z" }, + { url = "https://files.pythonhosted.org/packages/86/7e/5ce0bba863470fd1902d5e5843968951bddf38abe4742fc97116ef4598b3/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:75286256590a6320cf106a0d28970d3560aad9ee09aa7b34fb40524792436d35", size = 209614, upload-time = "2026-07-07T14:33:47.705Z" }, + { url = "https://files.pythonhosted.org/packages/6c/ef/2473d3c4d869155be4af1191111d59c4d5c4e0173026f7e85b176e23bf65/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:69b157c5d3292bcd443faca052f3096f637f1e074b98212a933c074ae23dc3b8", size = 224991, upload-time = "2026-07-07T14:33:49.238Z" }, + { url = "https://files.pythonhosted.org/packages/d0/a3/53ddae3db108a088156aa8ddfafd411ebbc1340f48c5573f697b27f69a39/charset_normalizer-3.4.9-cp313-cp313-win32.whl", hash = "sha256:51307f5c71007673a2bf8232ad973483d281e74cb99c8c5a990af1eefa6277d9", size = 150622, upload-time = "2026-07-07T14:33:50.711Z" }, + { url = "https://files.pythonhosted.org/packages/e8/ef/6953a77c7cf2c2ff9998e6f575ab3e380119f100223381565a4f94c1f836/charset_normalizer-3.4.9-cp313-cp313-win_amd64.whl", hash = "sha256:fe2c7201c642b7c308f1675355ad7ff7b66acfe3541625efe5a3ad38f29d6115", size = 161947, upload-time = "2026-07-07T14:33:52.197Z" }, + { url = "https://files.pythonhosted.org/packages/6e/fb/d560d1d1555debbfe7849d9cac6145c1b537709d79576bf22557ed803b82/charset_normalizer-3.4.9-cp313-cp313-win_arm64.whl", hash = "sha256:611057cc5d5c0afc743ba8be6bd828c17e0aaa8643f9d0a9b9bb7dea80eb8012", size = 152594, upload-time = "2026-07-07T14:33:53.486Z" }, + { url = "https://files.pythonhosted.org/packages/7e/8d/496817fa0944239ecae662dd57ea765cfeaec6a735f9f025d4b7b72e7143/charset_normalizer-3.4.9-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:0327fcd59a935777d83410750c50600ee9571af2846f71ce40f25b13da1ef380", size = 317253, upload-time = "2026-07-07T14:33:54.994Z" }, + { url = "https://files.pythonhosted.org/packages/2b/f9/ef4a69ea338ad3c0deceea0f5f7d2380ae8b52132b06d652cb0d2cd86706/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a79d9f4d8001473a30c163556b3c3bfebec837495a412dde78b51672f6134f9", size = 215898, upload-time = "2026-07-07T14:33:56.334Z" }, + { url = "https://files.pythonhosted.org/packages/8c/e7/5ddfd76fc061eb52de219658a4aa431cbacadf0a0219c8854f00da50d289/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:33bdcc2a32c0a0e861f60841a512c8acc658c87c2ac59d89e3a46dacf7d866e4", size = 236718, upload-time = "2026-07-07T14:33:57.9Z" }, + { url = "https://files.pythonhosted.org/packages/49/ba/768fa3f36048d81c477a0ce61f813bc1454d80917ccfe550abd9f44f5e24/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f840ed6d8ecba8255df8c42b87fadeda98ddfc6eeec05e2dc66e26d46dd6f58a", size = 232519, upload-time = "2026-07-07T14:33:59.811Z" }, + { url = "https://files.pythonhosted.org/packages/f4/c4/b3e049d2aa3766180c78507110543d9d50894cc97f57de543f1be521dcdc/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c25fe15c70c59eb7c5ce8c06a1f3fa1da0ecc5ea1e7a5922c40fd2fa9b0d5046", size = 223143, upload-time = "2026-07-07T14:34:01.517Z" }, + { url = "https://files.pythonhosted.org/packages/19/79/55c32d06d76ae4feafe053f061f3e3ab70bcf19f4007797ce8c3efda7830/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_armv7l.whl", hash = "sha256:f7fb7d750cfa0a070d2c24e831fd3481019a60dd317ea2b39acbcebc08b6ed81", size = 206742, upload-time = "2026-07-07T14:34:03.04Z" }, + { url = "https://files.pythonhosted.org/packages/10/e0/47c079dd82d217c807479cd59ffd30af56307ea31c108b75758970459ad3/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4d1c96a7a18b9690a4d46df09e3e3382406ae3213727cd1019ebade1c4a81917", size = 219191, upload-time = "2026-07-07T14:34:04.657Z" }, + { url = "https://files.pythonhosted.org/packages/42/ab/b9bc2e77d6b44a7e46ef62ec5cac1c9a6ba7b9135a5d560f002696ec9995/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a4cfde78a9f2880208d16a93b795726a3017d5977e08d1e162a7a31322479c41", size = 218328, upload-time = "2026-07-07T14:34:06.115Z" }, + { url = "https://files.pythonhosted.org/packages/f1/78/c9c71d599f5aa2d42bcdd35cbbd46d7f535351a57e40ff7d8e5a7e219401/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:d4d6fcde76f94f5cb9e43e9e9a61f16dacefd228cbbf6f1a09bd9b219a92f1a1", size = 207406, upload-time = "2026-07-07T14:34:07.554Z" }, + { url = "https://files.pythonhosted.org/packages/f6/39/c914445c321a845097ce4f6ac7de9a18228a77b766272125a1ce00d851eb/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:898f0e9068ca27d37f8e83a5b962821df851532e6c4a7d615c1c033f9da6eedf", size = 225157, upload-time = "2026-07-07T14:34:09.061Z" }, + { url = "https://files.pythonhosted.org/packages/9b/f2/c0d4b8508565a36bc5c624e88ed297f5b0b1095011034d7f5b83a69908b5/charset_normalizer-3.4.9-cp314-cp314-win32.whl", hash = "sha256:c1c948747b03be832dceed96ca815cef7360de9aa19d37c730f8e3f6101aca48", size = 151095, upload-time = "2026-07-07T14:34:10.901Z" }, + { url = "https://files.pythonhosted.org/packages/49/fd/a1d26144398c67486422a72bf5812cda22cb4ccfcd95a290fb41ceb4b8e2/charset_normalizer-3.4.9-cp314-cp314-win_amd64.whl", hash = "sha256:16b65ea0f2465b6fb52aa22de5eca612aa964ddfec00a912e26f4656cbef890b", size = 162796, upload-time = "2026-07-07T14:34:12.47Z" }, + { url = "https://files.pythonhosted.org/packages/20/95/d75e82f8ce9fd323ebf059c16c9aadefb22a1ecde13b7840b35835e4886c/charset_normalizer-3.4.9-cp314-cp314-win_arm64.whl", hash = "sha256:40a126142a56b2dfc0aacbad1de8310cbf60da7656db0e6b16eebd48e3e93519", size = 153334, upload-time = "2026-07-07T14:34:14.044Z" }, + { url = "https://files.pythonhosted.org/packages/00/5e/17398df3a139985ba9d11ed072531986f408c8fca952835ef1ab1820c02b/charset_normalizer-3.4.9-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:609b3ba8fcc0fb5ab7af00719d0fb6ad0cb518e48e7712d12fd68f1327951198", size = 338848, upload-time = "2026-07-07T14:34:15.688Z" }, + { url = "https://files.pythonhosted.org/packages/cd/91/7253a32e86b7e1d1239b1b36ba6dd0f021a21107ab33054b53119cc083b9/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:51447e9aa2684679af07ca5021c3db526e0284347ebf4ffcec1154c3350cfe32", size = 223022, upload-time = "2026-07-07T14:34:17.248Z" }, + { url = "https://files.pythonhosted.org/packages/cb/32/2e64bd2be10e89c61e57ebe6a93fd98ae88eb7ebe414b5121f22c96c69eb/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cc1b0fff8ead343dae06305f954eb8468ba0ec1a97881f42489d198e4ce3c632", size = 241590, upload-time = "2026-07-07T14:34:18.813Z" }, + { url = "https://files.pythonhosted.org/packages/3d/ef/d96ec496cfea0c21db43b0ad03891308b02388d054cc902cf0e5a1ad6a88/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fa36ec09ef71d158186bc79e359ff5fdd6e7996fe8ab638f00d6b93139ba4fcf", size = 239584, upload-time = "2026-07-07T14:34:20.52Z" }, + { url = "https://files.pythonhosted.org/packages/d4/ce/9af95f7876194bd7a14e3dfe4a4de2e0bff02666a3910d72beafd06cc297/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:df115d4d83168fdf2cae48ef1ff6d1cb4c466364e30861b37121de0f3bf1b990", size = 230224, upload-time = "2026-07-07T14:34:22.189Z" }, + { url = "https://files.pythonhosted.org/packages/52/94/af74dde74a3996bd959c350709bfe50e297823d70a8c1cbd54b838880863/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:f86c6358749bd4fda175388691e3ba8c46e24c5347d0afd20f9b7edfc9faf07d", size = 212667, upload-time = "2026-07-07T14:34:23.857Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f0/f1c4fe746c395922961b5916ed1d7d6e7d4c84851d19ed43cc89980ec953/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:32286a2c8d167e897177b673176c1e3e00d4057caf5d2b64eef9a3666b03018e", size = 227179, upload-time = "2026-07-07T14:34:25.586Z" }, + { url = "https://files.pythonhosted.org/packages/e4/56/6c745619ac397e8871e2bcd3cea1eec86b877488f33888b3aef5c3ed506e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:83aed2c10721ddd90f68140685391b50811a880af20654c59af6b6c66c40513c", size = 225372, upload-time = "2026-07-07T14:34:27.212Z" }, + { url = "https://files.pythonhosted.org/packages/78/ad/98aae8630ac71f16711968e38a5acfecce41b778bf2f0312851020f565a8/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:cd6c3d4b783c556fa00bf540854e42f135e2f256abd29669fcd0da0f2dec79c2", size = 215222, upload-time = "2026-07-07T14:34:28.774Z" }, + { url = "https://files.pythonhosted.org/packages/f7/40/9593d54209765207a7f11073c06494c1721e4ca4a0a426c597679bf7f91e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ee2f2a527e3c1a6e6411eb4209642e138b544a2d72fe5d0d76daf77b24063534", size = 231958, upload-time = "2026-07-07T14:34:30.345Z" }, + { url = "https://files.pythonhosted.org/packages/b1/27/693ee5e8a18191eb38647360c51cd505013e2bd3b366aa43fd5344c21e3c/charset_normalizer-3.4.9-cp314-cp314t-win32.whl", hash = "sha256:0d861473f743244d349b50f850d10eb87aeb22bbdcc8e64f79273c94af5a8226", size = 155580, upload-time = "2026-07-07T14:34:31.884Z" }, + { url = "https://files.pythonhosted.org/packages/80/3f/bd97d3d9c613013d07cb7733d299385b41df37f0471310f5a73dc359f0b8/charset_normalizer-3.4.9-cp314-cp314t-win_amd64.whl", hash = "sha256:9b8e0f3107e2200b76f6054de99016eac3ee6762713587b36baaa7e4bd2ae177", size = 167620, upload-time = "2026-07-07T14:34:33.438Z" }, + { url = "https://files.pythonhosted.org/packages/3d/c6/eee9dca4439b1061f76373f06ea855678cc4a64c1c3c90b50e479edbb8eb/charset_normalizer-3.4.9-cp314-cp314t-win_arm64.whl", hash = "sha256:19ac87f93086ce37b86e098888555c4b4bc48102279bae3350098c0ed664b501", size = 158037, upload-time = "2026-07-07T14:34:35.018Z" }, + { url = "https://files.pythonhosted.org/packages/98/2b/f97f1c193fb855c345d678f5077d6926034db0722df74c8f057020e05a25/charset_normalizer-3.4.9-py3-none-any.whl", hash = "sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5", size = 64538, upload-time = "2026-07-07T14:34:56.993Z" }, +] + +[[package]] +name = "click" +version = "8.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/d4/81420972a676e8ffea40450d8c8c92943e7218a78fe9b64359836cc9876b/click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6", size = 338000, upload-time = "2026-06-24T17:45:15.148Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "ghp-import" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d9/29/d40217cbe2f6b1359e00c6c307bb3fc876ba74068cbab3dde77f03ca0dc4/ghp-import-2.1.0.tar.gz", hash = "sha256:9c535c4c61193c2df8871222567d7fd7e5014d835f97dc7b7439069e2413d343", size = 10943, upload-time = "2022-05-02T15:47:16.11Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f7/ec/67fbef5d497f86283db54c22eec6f6140243aae73265799baaaa19cd17fb/ghp_import-2.1.0-py3-none-any.whl", hash = "sha256:8337dd7b50877f163d4c0289bc1f1c7f127550241988d568c1db512c4324a619", size = 11034, upload-time = "2022-05-02T15:47:14.552Z" }, +] + +[[package]] +name = "griffe-inherited-docstrings" +version = "1.1.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "griffelib" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/cb/da/fd002dc5f215cd896bfccaebe8b4aa1cdeed8ea1d9d60633685bd61ff933/griffe_inherited_docstrings-1.1.3.tar.gz", hash = "sha256:cd1f937ec9336a790e5425e7f9b92f5a5ab17f292ba86917f1c681c0704cb64e", size = 26738, upload-time = "2026-02-21T09:38:44.312Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/20/4bc15f242181daad1c104e0a7d33be49e712461ea89e548152be0365b9ea/griffe_inherited_docstrings-1.1.3-py3-none-any.whl", hash = "sha256:aa7f6e624515c50d9325a5cfdf4b2acac547f1889aca89092d5da7278f739695", size = 6710, upload-time = "2026-02-20T11:06:38.75Z" }, +] + +[[package]] +name = "griffelib" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/33/e4/8d187ea29c2e30b3a09505c567513077d6117861bde1fbd997a167f262ec/griffelib-2.1.0.tar.gz", hash = "sha256:762a186d2c6fd6794d4ea20d428d597ffb857cb56b66421651cbba15bdd5e813", size = 216234, upload-time = "2026-06-19T12:05:42.278Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e4/d3/5268aeabf2ad82658c4e2ff3a060648d0f02f3926cb53247c0e4d0dab49e/griffelib-2.1.0-py3-none-any.whl", hash = "sha256:cc7b3d2d2865ad0b909fcc38086e3f554b5ea7acbaa7bbb7ecaa3f5dfb7d9f00", size = 142560, upload-time = "2026-06-19T12:05:38.742Z" }, +] + +[[package]] +name = "hypothesis" +version = "6.164.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "sortedcontainers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7a/ac/7b76103bd74d8457e4de0c6a6c3a26ac6327016438bde125e0a3de83a5b8/hypothesis-6.164.0.tar.gz", hash = "sha256:5d63d263d8c71b571638c18d9591f6e34b836c60a12469e9d9105c1c785f00f1", size = 492022, upload-time = "2026-07-30T12:39:49.085Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/fe/d5b75a55892b33e72945f82efc71f645d29c0bfdb9f00727f7535a52edcc/hypothesis-6.164.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:14b861ac3353f8643b82a3ba76b8a0a54d2a06160c32b9a1f64a8ab41b179089", size = 771561, upload-time = "2026-07-30T12:39:00.404Z" }, + { url = "https://files.pythonhosted.org/packages/1c/b0/2f01e9efc7267446bad0e2a68f7472daa174a72553d213b16aefe44b2bda/hypothesis-6.164.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3d8c8bb00a4b86ae90b9ad41f3e1c99d016ec3e64c0ff9d676a4bb7be4f56948", size = 767079, upload-time = "2026-07-30T12:39:23.123Z" }, + { url = "https://files.pythonhosted.org/packages/c3/26/d7bcd26b58e1df2bd39116b924b2a72676215d9650e68cbff9a629c3ce30/hypothesis-6.164.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e80e3ba8eaf37664eaa0f2625cef120b330b128a7df570210cf8be4f5ae65aaa", size = 1096364, upload-time = "2026-07-30T12:38:49.972Z" }, + { url = "https://files.pythonhosted.org/packages/9d/17/99fe7ea866935da83444c3ef7885a14fc7349d96ff61c6faebd37ef4edf2/hypothesis-6.164.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8cdf70f821e2d2f3a0bccaab29830aea8aefb63a77806e7e91246fb65a10c8d3", size = 1124963, upload-time = "2026-07-30T12:39:13.1Z" }, + { url = "https://files.pythonhosted.org/packages/38/e8/df08be6296cbc1271d44e81f8ff9dcd6267a07552fb768e0fdc166e93d40/hypothesis-6.164.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bcc3743e22b3cffa7267b4bc74d03628606e4a115495728e986a7be220987315", size = 1145886, upload-time = "2026-07-30T12:39:45.612Z" }, + { url = "https://files.pythonhosted.org/packages/4e/72/d5cf6fbfac40891d4281f630e16a6eb217ff56f97e350a06e0fd9322aa6a/hypothesis-6.164.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:730f09d4afcd8a918b3d589bfb6421e3b41c057aa57652a773ef4f512cc60836", size = 1101181, upload-time = "2026-07-30T12:39:05.194Z" }, + { url = "https://files.pythonhosted.org/packages/fb/ff/7ceb002329febffb678b65835ca6e9479a916325d088aadb0210d07f8252/hypothesis-6.164.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:9651cb48cb5a995295b442138d15d381547b935dcb0066fca7148a7955347400", size = 1137970, upload-time = "2026-07-30T12:39:16.076Z" }, + { url = "https://files.pythonhosted.org/packages/7c/8f/c12c697b73ca9ca24d8a913879e3e0a9db86479754c7221554247c701565/hypothesis-6.164.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:51d161d2655dd86143b370c577267b5b7b4c2e8fcb8a3f22c1a787572aad707c", size = 1270184, upload-time = "2026-07-30T12:38:54.436Z" }, + { url = "https://files.pythonhosted.org/packages/0e/2f/93f1c850c794fc9c80f5e61b3b20652126b865e6f57b348ae530446aadc7/hypothesis-6.164.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:e8a250552390128b57e3afe55035ce2c2cb1f6f0919817657854244f071bc5be", size = 1397987, upload-time = "2026-07-30T12:38:21.113Z" }, + { url = "https://files.pythonhosted.org/packages/c5/b8/bab2546325e15e87c8518dfbca263c81dbc35d566c516d66c9da98a38b77/hypothesis-6.164.0-cp310-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:570cd51944e1cc3443847d8afa3d17fcf8aac475a1f744c9e7318a5ad7ef5c9f", size = 1270755, upload-time = "2026-07-30T12:38:51.571Z" }, + { url = "https://files.pythonhosted.org/packages/6a/4e/ea97dd39678a42dc5a24e3e2a64d3b950fad9fb1dcce8d7be5afb52a0335/hypothesis-6.164.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3a423e543055b3de5af7a7624c4285422541658367211fa293a3a57dd0ad01ba", size = 1312888, upload-time = "2026-07-30T12:38:30.847Z" }, + { url = "https://files.pythonhosted.org/packages/44/84/a6f2d5b12b23d65f16eb398750e430065f9d1f40f4418569e3b87ef58d23/hypothesis-6.164.0-cp310-abi3-win32.whl", hash = "sha256:f5e51490b2ce64c66138f24477d83c71b6224ab0ef65700da10187c464b54e94", size = 657401, upload-time = "2026-07-30T12:39:11.581Z" }, + { url = "https://files.pythonhosted.org/packages/f5/d3/c5ee410daa594cac2d3fe1fbe5473f2390e35f4369e168a817e43341ce2f/hypothesis-6.164.0-cp310-abi3-win_amd64.whl", hash = "sha256:c9059dfbb039342b6590bbce207f90e0f9a80fdf45a404c68c2d3e598be78ab3", size = 663566, upload-time = "2026-07-30T12:39:30.27Z" }, + { url = "https://files.pythonhosted.org/packages/90/91/4942fe3f2f08b920368ed5a2937346259e843e382205513b4a0e70d2de9d/hypothesis-6.164.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:6bc3373fe550cf4d7cadb94ceaeb91e431e1418a96b7baa330487366eaa67d3c", size = 773152, upload-time = "2026-07-30T12:38:33.328Z" }, + { url = "https://files.pythonhosted.org/packages/eb/df/e66d052386a2b6c3e2f3eab32a02d7de3c9c59cd21d5dd58c08ecfa715f0/hypothesis-6.164.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2780297ca68929b153eff7effb2ebe67e9487d2fd9f49fa961007f8f2d236c9e", size = 764713, upload-time = "2026-07-30T12:38:48.59Z" }, + { url = "https://files.pythonhosted.org/packages/f5/d5/5a50d14b8f04809e973c4dea884b367fef3663ff253c1205fa9e96229ef9/hypothesis-6.164.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6b400bb4eb5a4a1e19cd5af3cc63817909e6b54b4603e04022bdba46860913d7", size = 1095160, upload-time = "2026-07-30T12:38:58.925Z" }, + { url = "https://files.pythonhosted.org/packages/58/01/781b19ce4382ec239c4dc6ec3bd9f195e69e5570f2814bbf04b5781ecb18/hypothesis-6.164.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7fca6632933fc506dd96926d9383483e4c0066c7ff62c748d059a3276da761e7", size = 1145199, upload-time = "2026-07-30T12:39:09.904Z" }, + { url = "https://files.pythonhosted.org/packages/e9/64/30e016863515ca01c1c738b05dd50491353d3ccae6432362e56e0c15d0da/hypothesis-6.164.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:b9e1f6e89e5ec34735b727f3ce41d12e7f3b8efc162c91c8a225e10b54b504b4", size = 1267980, upload-time = "2026-07-30T12:38:18.733Z" }, + { url = "https://files.pythonhosted.org/packages/84/23/17eb8d67d59ecd3a820c905fbdf514e371dd7d01631e62a304cdd5793abe/hypothesis-6.164.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:51b0f967f608707b24ed37a298174ae6eec7899bfe3f271d1c3062c39ad66c06", size = 1312181, upload-time = "2026-07-30T12:38:36.056Z" }, + { url = "https://files.pythonhosted.org/packages/42/69/cff9f3cd9524252adda7c8e0e129dfc176e72f64fdf0bf1552d1ea43d78d/hypothesis-6.164.0-cp312-cp312-win_amd64.whl", hash = "sha256:5770df7d518bf867a9379e9081abd9e44db1d15473430e26a0946438c08c5926", size = 660690, upload-time = "2026-07-30T12:38:28.107Z" }, + { url = "https://files.pythonhosted.org/packages/ba/b4/729697380a22dc2ce8feae3c64b08bf3bd3c27e99c3706cb9bdac40c6fc8/hypothesis-6.164.0-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:29e7cb48974cb9fd87602e20625c890385793c6b56c18a957085a9c291f56ef8", size = 773046, upload-time = "2026-07-30T12:39:40.473Z" }, + { url = "https://files.pythonhosted.org/packages/38/35/72374f02d90dfda198afd8aac6b1e7d1184506f97e62ebcf3d2c1e5bf761/hypothesis-6.164.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1ff8c3819345be8dd15ee6588ee9383869a54c9a3d2232cce5e26b456424135d", size = 764659, upload-time = "2026-07-30T12:38:55.896Z" }, + { url = "https://files.pythonhosted.org/packages/6e/75/fb26388915d71e5949b98ccd0c9d95edcbe6b45d0370f177d43633d81ae2/hypothesis-6.164.0-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:33e88be13fac3ff7cb789a0b4cc43d99fb297db085f529fbb363188141c7d5bf", size = 1095078, upload-time = "2026-07-30T12:38:34.677Z" }, + { url = "https://files.pythonhosted.org/packages/be/63/f6da6e39667d39a1e44c5df82fbe6cff070c29aaffa9beb62a5322e7d8ae/hypothesis-6.164.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2e296d03a77355ce2e1c32e85a636b555edf0ddaaef277f98f1b84fe38a4595", size = 1145015, upload-time = "2026-07-30T12:39:26.487Z" }, + { url = "https://files.pythonhosted.org/packages/88/c7/55ba09727da3d9a60628c50e31e6083a36f403cb230f5e1a7bd1749a5c39/hypothesis-6.164.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:53698a1b246714539dd0ecc2d556cde613d74e9f7385ec4109e0651ab2d382d6", size = 1268027, upload-time = "2026-07-30T12:38:25.676Z" }, + { url = "https://files.pythonhosted.org/packages/ff/35/4789cade332f799b0e8f2f7ea0fe2aae6157a85e60f74497e316dd17a7e3/hypothesis-6.164.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:004c92c4b869f8e258f0641101b7743cae8420436f4465383f681c086ef95c9d", size = 1311895, upload-time = "2026-07-30T12:39:14.621Z" }, + { url = "https://files.pythonhosted.org/packages/12/8a/18d85e624f8631aec42daa8a2f07c6edcedb7385b2c0f375ba8a30cbd065/hypothesis-6.164.0-cp313-cp313-win_amd64.whl", hash = "sha256:4878f81fa92a580d3e16b53e64e01a9d9fe1dca5973783558493a003138dbd36", size = 660656, upload-time = "2026-07-30T12:38:37.696Z" }, + { url = "https://files.pythonhosted.org/packages/c7/06/3c144d427799c7c72befb0bb3b199d419a89b96e1002fd8f0cc94c84ffb7/hypothesis-6.164.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:9110010bdf6deb3ba9134f8ce8b683e8bb0fba108a351045c96d60c410eb6963", size = 773254, upload-time = "2026-07-30T12:38:38.919Z" }, + { url = "https://files.pythonhosted.org/packages/74/2d/b61a10d9e70df04aa7e8f34efef8e4afe364e8995c59f894e1c35b428214/hypothesis-6.164.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4df103e5d32b47d574c6e857d45361e2cba5a198d6dae4e4ee1bd248b3a2cbfa", size = 764786, upload-time = "2026-07-30T12:38:24.464Z" }, + { url = "https://files.pythonhosted.org/packages/99/68/7f80ac7bdffe78686135311c919534be411d4565c2a5ba38fd389880c553/hypothesis-6.164.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4abec95020960c0ed08e5be318d2bcdde79f2c6fc7785e368a9389d31d3e802a", size = 1095578, upload-time = "2026-07-30T12:38:57.422Z" }, + { url = "https://files.pythonhosted.org/packages/45/f9/97dcbac776bcf33cb4241b52111527821f707b60a84d03d0ea670b09a134/hypothesis-6.164.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9b106756cc9abd50ab1632541ea7b7223792d877a084726aa0304237d758181e", size = 1145207, upload-time = "2026-07-30T12:38:23.387Z" }, + { url = "https://files.pythonhosted.org/packages/a4/df/68184b6f71540435c895cf35ad1d67a3634a887c597ab38d3372c0d20186/hypothesis-6.164.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:11c4aab2ae6757fc4bc3bbf009487e24fd3490365817bbf40b9ec85a7e02fabb", size = 1268357, upload-time = "2026-07-30T12:38:52.946Z" }, + { url = "https://files.pythonhosted.org/packages/a0/76/6a6851dc8af89a5c0418937d38456417b2a1fc9db15c992b9cb43d53a7a3/hypothesis-6.164.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:4713edecbc0969557ca135769a36d1e524c8e3b7a2b271de48d98fa29f681bf6", size = 1312183, upload-time = "2026-07-30T12:39:28.604Z" }, + { url = "https://files.pythonhosted.org/packages/2f/19/83adeb1f8f045bd8a1ab9822d0c3db28b337d37fff01d809fcd6e3ea70f8/hypothesis-6.164.0-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:e6882d316c390d33c55ec8f1675f35ab238d7c0473ccf8d235c69eaef6c621b9", size = 604771, upload-time = "2026-07-30T12:39:33.579Z" }, + { url = "https://files.pythonhosted.org/packages/0b/62/fcb48ebfbccdc5b695de175b9d1d344b3688782150f0603124bb70c0891b/hypothesis-6.164.0-cp314-cp314-win_amd64.whl", hash = "sha256:7c3357633b38bca8c927fd90d02b39a0a3f35f24cdbcfb2fb1dcf69a3f63bd85", size = 660570, upload-time = "2026-07-30T12:39:43.898Z" }, + { url = "https://files.pythonhosted.org/packages/42/61/5857da7db0435fa69df658a9eafba62eb8a1319454005ce2a0d97f6f9e4d/hypothesis-6.164.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:53152cb549f52d661c47768d0d12a192ef26a7a9758a7f13b8ec41e8e63d6325", size = 771839, upload-time = "2026-07-30T12:38:26.842Z" }, + { url = "https://files.pythonhosted.org/packages/6c/9c/22292a9dab1c544362d1759244132c7d71a9d9d5eda5d454ec735fba6bd3/hypothesis-6.164.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cee7898ad84b63da6506ae48483bb36f319a25ea4c2b1d2df47d021cc4080c24", size = 763363, upload-time = "2026-07-30T12:38:20.042Z" }, + { url = "https://files.pythonhosted.org/packages/e8/29/cc0c6e9a065a32f93fe52dde746232f007d2cabf619d4e7b1b37bd34c424/hypothesis-6.164.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0def33f0d236e54144a5218997e4492925144d4615f25fdbb4ac8e47b7b709e6", size = 1094171, upload-time = "2026-07-30T12:39:35.158Z" }, + { url = "https://files.pythonhosted.org/packages/a7/59/37040d0776a29d4bc6d0ca9a50ca2755200007e4a8ddc27b010115b69c85/hypothesis-6.164.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:471fd80d70f2df606b1320276168bc2c6007a586124a1d81628264ccb9266f68", size = 1144089, upload-time = "2026-07-30T12:38:43.024Z" }, + { url = "https://files.pythonhosted.org/packages/fa/10/5235ed3c090a2f12fa15cc1d08e5a36cfa31bc0607c45199b0806e930ab4/hypothesis-6.164.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:2eb285756aee62890fd08d6e97cf77651dfe7c093ceac094df52120a7a8dbe68", size = 1266595, upload-time = "2026-07-30T12:39:36.979Z" }, + { url = "https://files.pythonhosted.org/packages/7f/97/ffc4cee4dfdffe658e839d5f4df72ae3fa7bfea9401550b475d9700e0ee2/hypothesis-6.164.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:7c5215b5568968c35c6e124e5a4a8068f80419d6171414ddf735b49e1df1ab59", size = 1310998, upload-time = "2026-07-30T12:38:45.788Z" }, + { url = "https://files.pythonhosted.org/packages/dd/08/681d4a272cd2812151581c3328e41a80a34e420d676e419a25b4b9dc2291/hypothesis-6.164.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a845e59fae87bb47a6fb84e0d5adb5679b3b55042fc3f8791da91486103cfbf0", size = 660724, upload-time = "2026-07-30T12:38:40.341Z" }, +] + +[[package]] +name = "idna" +version = "3.18" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "markdown" +version = "3.10.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/29/6f/da4c6aea59b3001f2e8c0ec7497475aadaf3b021c10cab5b2858f0f32b26/markdown-3.10.3.tar.gz", hash = "sha256:3589362618f743188b4d955b874402bc814f4f83f544dc207719f4baa7d9c45f", size = 372596, upload-time = "2026-07-30T19:05:29.005Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/69/4a5af2bc115a9a33fefe51709749de8262be3f9ba063d1753a837cdbc49c/markdown-3.10.3-py3-none-any.whl", hash = "sha256:fa6c92a00a4a3c98b22728c64a935ae1928250ae65058a6ded814d2cc29a4cea", size = 110757, upload-time = "2026-07-30T19:05:27.883Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, + { url = "https://files.pythonhosted.org/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, + { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, + { url = "https://files.pythonhosted.org/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, + { url = "https://files.pythonhosted.org/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, + { url = "https://files.pythonhosted.org/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, + { url = "https://files.pythonhosted.org/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, + { url = "https://files.pythonhosted.org/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, + { url = "https://files.pythonhosted.org/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, + { url = "https://files.pythonhosted.org/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, + { url = "https://files.pythonhosted.org/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, + { url = "https://files.pythonhosted.org/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, + { url = "https://files.pythonhosted.org/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, + { url = "https://files.pythonhosted.org/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, + { url = "https://files.pythonhosted.org/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, + { url = "https://files.pythonhosted.org/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, + { url = "https://files.pythonhosted.org/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, + { url = "https://files.pythonhosted.org/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, + { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, +] + +[[package]] +name = "mergedeep" +version = "1.3.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/41/580bb4006e3ed0361b8151a01d324fb03f420815446c7def45d02f74c270/mergedeep-1.3.4.tar.gz", hash = "sha256:0096d52e9dad9939c3d975a774666af186eda617e6ca84df4c94dec30004f2a8", size = 4661, upload-time = "2021-02-05T18:55:30.623Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/19/04f9b178c2d8a15b076c8b5140708fa6ffc5601fb6f1e975537072df5b2a/mergedeep-1.3.4-py3-none-any.whl", hash = "sha256:70775750742b25c0d8f36c55aed03d24c3384d17c951b3175d898bd778ef0307", size = 6354, upload-time = "2021-02-05T18:55:29.583Z" }, +] + +[[package]] +name = "mkdocs" +version = "1.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "ghp-import" }, + { name = "jinja2" }, + { name = "markdown" }, + { name = "markupsafe" }, + { name = "mergedeep" }, + { name = "mkdocs-get-deps" }, + { name = "packaging" }, + { name = "pathspec" }, + { name = "pyyaml" }, + { name = "pyyaml-env-tag" }, + { name = "watchdog" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bc/c6/bbd4f061bd16b378247f12953ffcb04786a618ce5e904b8c5a01a0309061/mkdocs-1.6.1.tar.gz", hash = "sha256:7b432f01d928c084353ab39c57282f29f92136665bdd6abf7c1ec8d822ef86f2", size = 3889159, upload-time = "2024-08-30T12:24:06.899Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/5b/dbc6a8cddc9cfa9c4971d59fb12bb8d42e161b7e7f8cc89e49137c5b279c/mkdocs-1.6.1-py3-none-any.whl", hash = "sha256:db91759624d1647f3f34aa0c3f327dd2601beae39a366d6e064c03468d35c20e", size = 3864451, upload-time = "2024-08-30T12:24:05.054Z" }, +] + +[[package]] +name = "mkdocs-autorefs" +version = "1.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown" }, + { name = "markupsafe" }, + { name = "mkdocs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/52/c0/f641843de3f612a6b48253f39244165acff36657a91cc903633d456ae1ac/mkdocs_autorefs-1.4.4.tar.gz", hash = "sha256:d54a284f27a7346b9c38f1f852177940c222da508e66edc816a0fa55fc6da197", size = 56588, upload-time = "2026-02-10T15:23:55.105Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/de/a3e710469772c6a89595fc52816da05c1e164b4c866a89e3cb82fb1b67c5/mkdocs_autorefs-1.4.4-py3-none-any.whl", hash = "sha256:834ef5408d827071ad1bc69e0f39704fa34c7fc05bc8e1c72b227dfdc5c76089", size = 25530, upload-time = "2026-02-10T15:23:53.817Z" }, +] + +[[package]] +name = "mkdocs-get-deps" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mergedeep" }, + { name = "platformdirs" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/25/b3cccb187655b9393572bde9b09261d267c3bf2f2cdabe347673be5976a6/mkdocs_get_deps-0.2.2.tar.gz", hash = "sha256:8ee8d5f316cdbbb2834bc1df6e69c08fe769a83e040060de26d3c19fad3599a1", size = 11047, upload-time = "2026-03-10T02:46:33.632Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/29/744136411e785c4b0b744d5413e56555265939ab3a104c6a4b719dad33fd/mkdocs_get_deps-0.2.2-py3-none-any.whl", hash = "sha256:e7878cbeac04860b8b5e0ca31d3abad3df9411a75a32cde82f8e44b6c16ff650", size = 9555, upload-time = "2026-03-10T02:46:32.256Z" }, +] + +[[package]] +name = "mkdocs-material" +version = "9.7.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "babel" }, + { name = "backrefs" }, + { name = "colorama" }, + { name = "jinja2" }, + { name = "markdown" }, + { name = "mkdocs" }, + { name = "mkdocs-material-extensions" }, + { name = "paginate" }, + { name = "pygments" }, + { name = "pymdown-extensions" }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/45/29/6d2bcf41ae40802c4beda2432396fff97b8456fb496371d1bc7aad6512ec/mkdocs_material-9.7.6.tar.gz", hash = "sha256:00bdde50574f776d328b1862fe65daeaf581ec309bd150f7bff345a098c64a69", size = 4097959, upload-time = "2026-03-19T15:41:58.161Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2c/01/bc663630c510822c95c47a66af9fa7a443c295b47d5f041e5e6ae62ef659/mkdocs_material-9.7.6-py3-none-any.whl", hash = "sha256:71b84353921b8ea1ba84fe11c50912cc512da8fe0881038fcc9a0761c0e635ba", size = 9305470, upload-time = "2026-03-19T15:41:55.217Z" }, +] + +[[package]] +name = "mkdocs-material-extensions" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/79/9b/9b4c96d6593b2a541e1cb8b34899a6d021d208bb357042823d4d2cabdbe7/mkdocs_material_extensions-1.3.1.tar.gz", hash = "sha256:10c9511cea88f568257f960358a467d12b970e1f7b2c0e5fb2bb48cab1928443", size = 11847, upload-time = "2023-11-22T19:09:45.208Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/54/662a4743aa81d9582ee9339d4ffa3c8fd40a4965e033d77b9da9774d3960/mkdocs_material_extensions-1.3.1-py3-none-any.whl", hash = "sha256:adff8b62700b25cb77b53358dad940f3ef973dd6db797907c49e3c2ef3ab4e31", size = 8728, upload-time = "2023-11-22T19:09:43.465Z" }, +] + +[[package]] +name = "mkdocstrings" +version = "1.0.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jinja2" }, + { name = "markdown" }, + { name = "markupsafe" }, + { name = "mkdocs" }, + { name = "mkdocs-autorefs" }, + { name = "pymdown-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1d/5d/f888d4d3eb31359b327bc9b17a212d6ef03fe0b0682fbb3fc2cb849fb12b/mkdocstrings-1.0.4.tar.gz", hash = "sha256:3969a6515b77db65fd097b53c1b7aa4ae840bd71a2ee62a6a3e89503446d7172", size = 100088, upload-time = "2026-04-15T09:16:53.376Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/94/be70f8ee9c45f2f62b39a1f0e9303bc20e138a8f3b8e50ffd89498e177e1/mkdocstrings-1.0.4-py3-none-any.whl", hash = "sha256:63464b4b29053514f32a1dbbf604e52876d5e638111b0c295ab7ed3cac73ca9b", size = 35560, upload-time = "2026-04-15T09:16:51.436Z" }, +] + +[[package]] +name = "mkdocstrings-python" +version = "2.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "griffelib" }, + { name = "mkdocs-autorefs" }, + { name = "mkdocstrings" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/b6/e858701499d57eee8b3fd8e78168083956c6683ddbe727b46758b19e1119/mkdocstrings_python-2.0.5.tar.gz", hash = "sha256:3a4d92556ad39637e88af94a5374213af9a8e3040c3824ceaed04b486c017594", size = 199578, upload-time = "2026-06-19T10:41:08.868Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/fc/10ab7e80650a9c9e8f4f1105f8c8e73567f88ed0c06ada589ab81d38687c/mkdocstrings_python-2.0.5-py3-none-any.whl", hash = "sha256:30c837bbff016549f659fcba6539ac351303f0fd7e713c89a040611072236e9d", size = 104951, upload-time = "2026-06-19T10:41:07.378Z" }, +] + +[[package]] +name = "numpy" +version = "2.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/22/fd/89965aa4ac08c74998539fcbf24fa3540f3e15237fbeb6bcf9c908f4aade/numpy-2.5.1.tar.gz", hash = "sha256:a48a113e6afea91f5608793bafa7ef2ad481fefbda87ec5069f483de61cb9fa3", size = 20755553, upload-time = "2026-07-04T17:08:00.933Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/7b/14687aa674250e5e546f616f486b0d56d3631cd5b2415739141ce40bdcea/numpy-2.5.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2c889b56fe48b1018f764b0eec8df59ab654e9148aa91faa12596043500de277", size = 16801574, upload-time = "2026-07-04T17:06:12.423Z" }, + { url = "https://files.pythonhosted.org/packages/e1/19/cc5bb2a3f2913d27d6dbb2c78d25921fabaedc6741d4a5a615a11f3c5bf3/numpy-2.5.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ab451b59c5643c570974c43aef780703ef1d3b4965d2be07afd530615a9358d1", size = 11772250, upload-time = "2026-07-04T17:06:15.726Z" }, + { url = "https://files.pythonhosted.org/packages/42/77/fdf34a71dd30f54979b18603bee915e0aaf825b07afe79acd60b04b691e2/numpy-2.5.1-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:78798bd5b9ad744056af8efa90e3b9ddaa53272a0848a483084a1cc0a13b2dc0", size = 5331516, upload-time = "2026-07-04T17:06:17.913Z" }, + { url = "https://files.pythonhosted.org/packages/ce/e2/eb7efa015b4cce41e2517bf182a7fce0d7d5b9d9ed76a29bfa0f4fe4505c/numpy-2.5.1-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:2ae0ca40bcb22d6ba59c1dfd5446f49940b0f2d821fde133f10dda11f816b84e", size = 6664863, upload-time = "2026-07-04T17:06:20.02Z" }, + { url = "https://files.pythonhosted.org/packages/a9/4b/a2b32dd94ee9ffbeecb28152240042a3949db33b1c834d44090b80e1b3b8/numpy-2.5.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:61ac47e772e6b8ea489e1d2f441a34c5c3ac17327e7ce294cbdf535795ad4e75", size = 15167977, upload-time = "2026-07-04T17:06:21.621Z" }, + { url = "https://files.pythonhosted.org/packages/b8/a9/6e73d68500f80773f65f0654ea932019d6694329a0eb0ed0533de38df376/numpy-2.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:59fda5e192b570217ec2580c96f00e9a7e12ef6866a900eb089b62c1a32545ca", size = 16672469, upload-time = "2026-07-04T17:06:24.064Z" }, + { url = "https://files.pythonhosted.org/packages/24/7d/ad3e59015135f5261c95fd4cafeff159c955febd83a99a1d9250c4233815/numpy-2.5.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:f7119ebff1a9829e9f431a4f9d28e703023bb6b9fe7c8f724467dbfc27c94ab3", size = 16527531, upload-time = "2026-07-04T17:06:26.69Z" }, + { url = "https://files.pythonhosted.org/packages/83/d0/a39b2fbcde9cb17a1dac678f254b33a6336298af9df338824c685425d5e8/numpy-2.5.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e824c2acf8862052246be5a44c15da1777940c60d010dd2aab897824d9c430f9", size = 18431940, upload-time = "2026-07-04T17:06:29.521Z" }, + { url = "https://files.pythonhosted.org/packages/04/12/cff070947791c1ed425ff76413189adbdc2fbe215eba7ce7fa454a03c7f8/numpy-2.5.1-cp312-cp312-win32.whl", hash = "sha256:08d60c810432eb83360958dea0999ac4cfb94531ea8efcbf0b7f277c2068aeb2", size = 6066764, upload-time = "2026-07-04T17:06:32.571Z" }, + { url = "https://files.pythonhosted.org/packages/65/66/53f31807a48a750f9d748da273bc3fcedd12b27ff1f3e373bfec55ef2dc0/numpy-2.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:f7d60026c0bdb1380e83bfa7a0419c4577ee4b9a08880afcb6dadeb74c649fa2", size = 12430966, upload-time = "2026-07-04T17:06:34.926Z" }, + { url = "https://files.pythonhosted.org/packages/2b/2a/d1a88066b1c14186f5d3c0d18c94f17b064511982bab0578d49ee9d43c29/numpy-2.5.1-cp312-cp312-win_arm64.whl", hash = "sha256:17a25e09640602e10bc8de0e6fa2b3fd68eedd84ba6d7842dc8f32f9ab87bd0b", size = 10350488, upload-time = "2026-07-04T17:06:37.785Z" }, + { url = "https://files.pythonhosted.org/packages/eb/07/ec2a3f0c91761581d4b7104a740791800025983f9a4dc4e73f91a99aeac4/numpy-2.5.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0bfebd8695f9863592fe744be833a258120b14a9f39da255e8aa8fade2c0ddd1", size = 16796419, upload-time = "2026-07-04T17:06:40.37Z" }, + { url = "https://files.pythonhosted.org/packages/ab/ab/ddb499fc4f8780354395face5b65c7fd107bcd6e1d667a5f07d046956f6f/numpy-2.5.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:30b44a6b53a7ae63c54c089a8726e5563ed302716c5b7ccc85afade40b0e7ff6", size = 11765832, upload-time = "2026-07-04T17:06:42.768Z" }, + { url = "https://files.pythonhosted.org/packages/88/b3/3c28c558a09fc72100c646dac6d2fce8e834c471b0edca01a29996706117/numpy-2.5.1-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:6165343f81b56ef8f514f396989e529b61d9dc709b99421b07e9f3e698e2287d", size = 5325143, upload-time = "2026-07-04T17:06:45.466Z" }, + { url = "https://files.pythonhosted.org/packages/5e/0e/ce19b985bb15c596f4f05954e76cccc77c845083b3b8f938a6c68e523128/numpy-2.5.1-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:4939237038ada79308dda3204ac6462df056b5672b2e25db1149cf873668b3e1", size = 6659749, upload-time = "2026-07-04T17:06:47.288Z" }, + { url = "https://files.pythonhosted.org/packages/2e/20/1ee6614d64332a1bba6411f38e68cb79eec1b2459e20a623777c5c5492a2/numpy-2.5.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1c6759f538fb912fc46de0a6b1758ccf7b57bc7c7ebebc23974fdac3de8db0cd", size = 15164716, upload-time = "2026-07-04T17:06:49.494Z" }, + { url = "https://files.pythonhosted.org/packages/ed/a7/2bcd3fdbb87804755c35b729bf8709d62025c5f4cfd7d5b2415997097515/numpy-2.5.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9726558e8db4a5bf7929a70ae50f63abda4daf0efe810e3bfbab95976f75fc1a", size = 16661440, upload-time = "2026-07-04T17:06:52.061Z" }, + { url = "https://files.pythonhosted.org/packages/fc/d7/a41e3310c886fe457d36e670bbf24fae411aca8a7b6ad92a32afd924077c/numpy-2.5.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:3935f3b419b244a02732676fa5317a9193cc596a4c0646db07e5b421229ac9f7", size = 16526305, upload-time = "2026-07-04T17:06:54.605Z" }, + { url = "https://files.pythonhosted.org/packages/53/75/4333a9a707c1edd3a4e1a0c58eca52c0f31e55089fa80db02b5565b24df7/numpy-2.5.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dc932a65ded7ce9013d120845a2514dcccb1a67bfc8deb8d37633762951904a6", size = 18423008, upload-time = "2026-07-04T17:06:57.54Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/e314a32b1c11a2ffe818ddad3a57b50b4b6e1b6c487192eb50cdef0415d0/numpy-2.5.1-cp313-cp313-win32.whl", hash = "sha256:4b4ff1608417eb7a59da7b967bbb798cacfe071d2caf526a24281cd562072ed9", size = 6063885, upload-time = "2026-07-04T17:07:00.14Z" }, + { url = "https://files.pythonhosted.org/packages/10/70/800b3fca480af32df9e8ea9f3d4a0c8feb4b32d7f195d174eabbda4829ad/numpy-2.5.1-cp313-cp313-win_amd64.whl", hash = "sha256:6c3fe51bc6a16453d452997053454f309e8e0ed7b42d6b361ce4ac8c32913d74", size = 12425674, upload-time = "2026-07-04T17:07:02.387Z" }, + { url = "https://files.pythonhosted.org/packages/8b/0b/196350c122f50f6ca56846f2d71efd5e0d24b7b2e07355e019b2e2c7a11e/numpy-2.5.1-cp313-cp313-win_arm64.whl", hash = "sha256:f7feb014281029e628ba2d5a007407443b06e418b6fe451d1e2adcbc8eba0107", size = 10350256, upload-time = "2026-07-04T17:07:04.878Z" }, + { url = "https://files.pythonhosted.org/packages/db/f4/731b6085a83faf6ca843394cbd5e217280c214399f7e8b21b9f552af0ae2/numpy-2.5.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:7c786fe9a5bbe360022e584c5a34cf6b54265c71bd7ec8ac3d8fec38968071f8", size = 16795063, upload-time = "2026-07-04T17:07:07.374Z" }, + { url = "https://files.pythonhosted.org/packages/bf/64/0e215f2048dd11a55bb989ed41b3585ef57452404e638d703a211a3e4157/numpy-2.5.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:32985c896d897419ef8da6917872d80b78ad0ea26d85b23245c7366ffde76d75", size = 11776652, upload-time = "2026-07-04T17:07:09.907Z" }, + { url = "https://files.pythonhosted.org/packages/b5/59/2b844c7a6e9deff69b404a66221e1542937734f65d5e6e39411876053862/numpy-2.5.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:efd736408cc97c79b9e6917338dfc8f06013b2274f992e96b1d9a81a71e2a2c2", size = 5335944, upload-time = "2026-07-04T17:07:12.227Z" }, + { url = "https://files.pythonhosted.org/packages/86/51/9bf7cb2cabcebc9e017e4ec7e6322b378317a542c08b4cb68479c1efc716/numpy-2.5.1-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:ab84dc6b074fa881cae55bea94cc4f68e285181ba7f32497bf7dee6b1496165b", size = 6656266, upload-time = "2026-07-04T17:07:14.368Z" }, + { url = "https://files.pythonhosted.org/packages/83/3e/fb7615b211b82a32f44d5180a6d421b61f84d4fadd578b48ba4ac34e189f/numpy-2.5.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:caf3e317d33d60c37986b452613f4ab51246d0691350c03d0cb4a898627f4a95", size = 15179720, upload-time = "2026-07-04T17:07:16.272Z" }, + { url = "https://files.pythonhosted.org/packages/41/5f/0f992cb24560673496c5d68de61913b57166ce530ffda07c1f280e0cc464/numpy-2.5.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:54ad769f17bc2d833b620851989f62054fb9ab93c969d9e1dc3c8e3d56beea21", size = 16664835, upload-time = "2026-07-04T17:07:19.021Z" }, + { url = "https://files.pythonhosted.org/packages/a2/2f/97d6475ee91afe2587797d09446f9d3e475ad4cb681662d824809327b75a/numpy-2.5.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c12afb53450fa976d4c681c50a7423729a4c51c0465ed9f32b8a9cabbc472373", size = 16539135, upload-time = "2026-07-04T17:07:22.015Z" }, + { url = "https://files.pythonhosted.org/packages/c4/5b/4db81e4ba0be7e2776b1de68c82aa862c7f8ec27e1b4927d4ae075e20678/numpy-2.5.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e8c11c405efc5ff6816d5983c96cdfa215bab3428961243af3ff59b228490438", size = 18426684, upload-time = "2026-07-04T17:07:24.941Z" }, + { url = "https://files.pythonhosted.org/packages/1f/64/c0ba2d90724d450279a7df8f32057241070250a26a7e2b5337d77347f481/numpy-2.5.1-cp314-cp314-win32.whl", hash = "sha256:f2479a47f8d5932d1718168a681ad6e536a9df484c83cfcf9de365e164537ace", size = 6116103, upload-time = "2026-07-04T17:07:27.622Z" }, + { url = "https://files.pythonhosted.org/packages/c1/1a/837f9ed7405adcd7a40538792eb169eddd8fa5630c16a1ef49dae71a30f4/numpy-2.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:24d0eb82c0541d3415a33425db64ae439dffccd7b4dbcb30e7c35120205c506a", size = 12562177, upload-time = "2026-07-04T17:07:29.887Z" }, + { url = "https://files.pythonhosted.org/packages/22/ed/49707938b6dd0a78a9178dd93227dc89e4c11af47f5c798d70366e8d0483/numpy-2.5.1-cp314-cp314-win_arm64.whl", hash = "sha256:5a4c988b38d261deeeaad9954e3deb091ad905c94e8bb6708654ef1d97f286b0", size = 10627739, upload-time = "2026-07-04T17:07:32.568Z" }, + { url = "https://files.pythonhosted.org/packages/a6/c7/bb4b882cfe7f299cbc8b66e42e7dd78cf9d14e40f9469fc5e3db7e15b3bd/numpy-2.5.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:a33276be12fa045805f477f22482088b66bb758ffbe89a9d21457de863a32e22", size = 11894709, upload-time = "2026-07-04T17:07:34.941Z" }, + { url = "https://files.pythonhosted.org/packages/40/3f/5af7f4a7f6224aef48017aa82bb6174c7a659d724be0c75017b7e64a55b4/numpy-2.5.1-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f089d7b00756190aacf1f5d34bdf38c3c430ac82b4f868f8cede73380460fce7", size = 5453810, upload-time = "2026-07-04T17:07:37.495Z" }, + { url = "https://files.pythonhosted.org/packages/20/c9/3474309bc94d634d3f9c3eddf03250ecb8c22cd948ef16fef69a77cc5d7b/numpy-2.5.1-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:09e9bfd8d2cf479c7d174804fb3811c53a8e9f20a37444008606b57d6b7a826d", size = 6761189, upload-time = "2026-07-04T17:07:39.563Z" }, + { url = "https://files.pythonhosted.org/packages/90/8a/558ae39fdd55d7e7f7fef9a84a6e964ac6b23edbd2a07e52bb084500507d/numpy-2.5.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e68d8dd1e7eba712948f2053a29ec86917bc70ba1358df869d9f06649ef9cf09", size = 15225039, upload-time = "2026-07-04T17:07:41.682Z" }, + { url = "https://files.pythonhosted.org/packages/63/27/ca7392b2d030277bdf0273e7d23255b3ee57d57a7c170a6f4fb3981e1e5d/numpy-2.5.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:99d5095fa265a0c4152e7bb12759e14381ef5496152f1ce58f44bdf55c44beb4", size = 16701306, upload-time = "2026-07-04T17:07:44.611Z" }, + { url = "https://files.pythonhosted.org/packages/02/42/03d53ae7996c44d4374a8262e9dc41671fd56cbb98f7d47ef85cf5da4c6b/numpy-2.5.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ab87a91b3cc3382b8956095bd8f95e00cf679bb81554339be1a2ba404a1473c1", size = 16589955, upload-time = "2026-07-04T17:07:47.694Z" }, + { url = "https://files.pythonhosted.org/packages/7b/15/6c1784ae469640e65db111e9a34b3d0f14d91e8a38b9ce34810ced370dbb/numpy-2.5.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:224ca51130ef7da85bea2191625181cb4f337f9cb64b471f10c1a12aa8b60077", size = 18464252, upload-time = "2026-07-04T17:07:50.684Z" }, + { url = "https://files.pythonhosted.org/packages/94/a8/f98e50356cf167df656c526c2dfeec2d7dde182f2a3da4b458a5938e2776/numpy-2.5.1-cp314-cp314t-win32.whl", hash = "sha256:6eab239876581b2b3c5a242281b6007bbdbcd1c7085d7709bb57c5929b11e6bf", size = 6263298, upload-time = "2026-07-04T17:07:53.445Z" }, + { url = "https://files.pythonhosted.org/packages/72/ac/96ae880cdecad0b3275d9359fcec72667b49a4863c9f12942e43679dda02/numpy-2.5.1-cp314-cp314t-win_amd64.whl", hash = "sha256:83ce9c80d5b521b0d77ddcbe5447c218d247929b6cc056ca5351342accfff0af", size = 12748623, upload-time = "2026-07-04T17:07:55.384Z" }, + { url = "https://files.pythonhosted.org/packages/a1/5a/4d2b1601df3602dba7a14f3348ba9bfe94a18adb428e693df6154c293831/numpy-2.5.1-cp314-cp314t-win_arm64.whl", hash = "sha256:5a6db61f9aaa57e369905c67d852045d3c4f7126405b29d09b19dec118e9c9cb", size = 10697674, upload-time = "2026-07-04T17:07:58.506Z" }, +] + +[[package]] +name = "packaging" +version = "26.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d7/f1/e7a6dd94a8d4a5626c03e4e99c87f241ba9e350cd9e6d75123f992427270/packaging-26.2.tar.gz", hash = "sha256:ff452ff5a3e828ce110190feff1178bb1f2ea2281fa2075aadb987c2fb221661", size = 228134, upload-time = "2026-04-24T20:15:23.917Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, +] + +[[package]] +name = "paginate" +version = "0.5.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/46/68dde5b6bc00c1296ec6466ab27dddede6aec9af1b99090e1107091b3b84/paginate-0.5.7.tar.gz", hash = "sha256:22bd083ab41e1a8b4f3690544afb2c60c25e5c9a63a30fa2f483f6c60c8e5945", size = 19252, upload-time = "2024-08-25T14:17:24.139Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/96/04b8e52da071d28f5e21a805b19cb9390aa17a47462ac87f5e2696b9566d/paginate-0.5.7-py2.py3-none-any.whl", hash = "sha256:b885e2af73abcf01d9559fd5216b57ef722f8c42affbb63942377668e35c7591", size = 13746, upload-time = "2024-08-25T14:17:22.55Z" }, +] + +[[package]] +name = "pathspec" +version = "1.1.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/82/42f767fc1c1143d6fd36efb827202a2d997a375e160a71eb2888a925aac1/pathspec-1.1.1.tar.gz", hash = "sha256:17db5ecd524104a120e173814c90367a96a98d07c45b2e10c2f3919fff91bf5a", size = 135180, upload-time = "2026-04-27T01:46:08.907Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, +] + +[[package]] +name = "platformdirs" +version = "4.11.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/78/9b/560e4be8e26f6fd133a03630a8df0c663b9e8d61b4ade152b72005aec83b/platformdirs-4.11.0.tar.gz", hash = "sha256:0555d18370482847566ffabcaa53ad7c6c1c29f195989ae1ed634a05f76ea1e0", size = 31953, upload-time = "2026-07-21T13:09:36.565Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/68/d8d58938dfb1370b266a1a729e6d77a985be23689a0496498ee17b2cbf90/platformdirs-4.11.0-py3-none-any.whl", hash = "sha256:360ccded2b7fce0af0ff80cc8f5942a1c5d99b0e856033acb030bfc634709e74", size = 23247, upload-time = "2026-07-21T13:09:35.422Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "pygments" +version = "2.20.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" }, +] + +[[package]] +name = "pymdown-extensions" +version = "11.0.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/21/a9/5f0c535ba3b08fe09270c16808e053a968868242ecbd5676d4e3a488bf28/pymdown_extensions-11.0.1.tar.gz", hash = "sha256:dd2905ae6fc5b75582fafb139a1266ffc754705efa902aa50067fa7ff4f94ec0", size = 857113, upload-time = "2026-07-02T17:59:22.955Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d6/54/da572c98c0b77626a91b5d3b89f0231d8bff5125c225420908632f8b342d/pymdown_extensions-11.0.1-py3-none-any.whl", hash = "sha256:db3943a62bab7e03af1364f0c4083e64b91fb097675a4b6cceccfbe9a77e5eb2", size = 269455, upload-time = "2026-07-02T17:59:21.271Z" }, +] + +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, + { url = "https://files.pythonhosted.org/packages/d1/11/0fd08f8192109f7169db964b5707a2f1e8b745d4e239b784a5a1dd80d1db/pyyaml-6.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8", size = 181669, upload-time = "2025-09-25T21:32:23.673Z" }, + { url = "https://files.pythonhosted.org/packages/b1/16/95309993f1d3748cd644e02e38b75d50cbc0d9561d21f390a76242ce073f/pyyaml-6.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1", size = 173252, upload-time = "2025-09-25T21:32:25.149Z" }, + { url = "https://files.pythonhosted.org/packages/50/31/b20f376d3f810b9b2371e72ef5adb33879b25edb7a6d072cb7ca0c486398/pyyaml-6.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c", size = 767081, upload-time = "2025-09-25T21:32:26.575Z" }, + { url = "https://files.pythonhosted.org/packages/49/1e/a55ca81e949270d5d4432fbbd19dfea5321eda7c41a849d443dc92fd1ff7/pyyaml-6.0.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5", size = 841159, upload-time = "2025-09-25T21:32:27.727Z" }, + { url = "https://files.pythonhosted.org/packages/74/27/e5b8f34d02d9995b80abcef563ea1f8b56d20134d8f4e5e81733b1feceb2/pyyaml-6.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6", size = 801626, upload-time = "2025-09-25T21:32:28.878Z" }, + { url = "https://files.pythonhosted.org/packages/f9/11/ba845c23988798f40e52ba45f34849aa8a1f2d4af4b798588010792ebad6/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6", size = 753613, upload-time = "2025-09-25T21:32:30.178Z" }, + { url = "https://files.pythonhosted.org/packages/3d/e0/7966e1a7bfc0a45bf0a7fb6b98ea03fc9b8d84fa7f2229e9659680b69ee3/pyyaml-6.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be", size = 794115, upload-time = "2025-09-25T21:32:31.353Z" }, + { url = "https://files.pythonhosted.org/packages/de/94/980b50a6531b3019e45ddeada0626d45fa85cbe22300844a7983285bed3b/pyyaml-6.0.3-cp313-cp313-win32.whl", hash = "sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26", size = 137427, upload-time = "2025-09-25T21:32:32.58Z" }, + { url = "https://files.pythonhosted.org/packages/97/c9/39d5b874e8b28845e4ec2202b5da735d0199dbe5b8fb85f91398814a9a46/pyyaml-6.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c", size = 154090, upload-time = "2025-09-25T21:32:33.659Z" }, + { url = "https://files.pythonhosted.org/packages/73/e8/2bdf3ca2090f68bb3d75b44da7bbc71843b19c9f2b9cb9b0f4ab7a5a4329/pyyaml-6.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb", size = 140246, upload-time = "2025-09-25T21:32:34.663Z" }, + { url = "https://files.pythonhosted.org/packages/9d/8c/f4bd7f6465179953d3ac9bc44ac1a8a3e6122cf8ada906b4f96c60172d43/pyyaml-6.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac", size = 181814, upload-time = "2025-09-25T21:32:35.712Z" }, + { url = "https://files.pythonhosted.org/packages/bd/9c/4d95bb87eb2063d20db7b60faa3840c1b18025517ae857371c4dd55a6b3a/pyyaml-6.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310", size = 173809, upload-time = "2025-09-25T21:32:36.789Z" }, + { url = "https://files.pythonhosted.org/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, + { url = "https://files.pythonhosted.org/packages/02/9e/e5e9b168be58564121efb3de6859c452fccde0ab093d8438905899a3a483/pyyaml-6.0.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788", size = 836355, upload-time = "2025-09-25T21:32:39.178Z" }, + { url = "https://files.pythonhosted.org/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, + { url = "https://files.pythonhosted.org/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, + { url = "https://files.pythonhosted.org/packages/23/20/bb6982b26a40bb43951265ba29d4c246ef0ff59c9fdcdf0ed04e0687de4d/pyyaml-6.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac", size = 156429, upload-time = "2025-09-25T21:32:57.844Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f4/a4541072bb9422c8a883ab55255f918fa378ecf083f5b85e87fc2b4eda1b/pyyaml-6.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3", size = 143912, upload-time = "2025-09-25T21:32:59.247Z" }, + { url = "https://files.pythonhosted.org/packages/7c/f9/07dd09ae774e4616edf6cda684ee78f97777bdd15847253637a6f052a62f/pyyaml-6.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3", size = 189108, upload-time = "2025-09-25T21:32:44.377Z" }, + { url = "https://files.pythonhosted.org/packages/4e/78/8d08c9fb7ce09ad8c38ad533c1191cf27f7ae1effe5bb9400a46d9437fcf/pyyaml-6.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba", size = 183641, upload-time = "2025-09-25T21:32:45.407Z" }, + { url = "https://files.pythonhosted.org/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, + { url = "https://files.pythonhosted.org/packages/8b/cc/dff0684d8dc44da4d22a13f35f073d558c268780ce3c6ba1b87055bb0b87/pyyaml-6.0.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702", size = 861132, upload-time = "2025-09-25T21:32:50.149Z" }, + { url = "https://files.pythonhosted.org/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, + { url = "https://files.pythonhosted.org/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, + { url = "https://files.pythonhosted.org/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/7a/1c7270340330e575b92f397352af856a8c06f230aa3e76f86b39d01b416a/pyyaml-6.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9", size = 174062, upload-time = "2025-09-25T21:32:55.767Z" }, + { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, +] + +[[package]] +name = "pyyaml-env-tag" +version = "1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/2e/79c822141bfd05a853236b504869ebc6b70159afc570e1d5a20641782eaa/pyyaml_env_tag-1.1.tar.gz", hash = "sha256:2eb38b75a2d21ee0475d6d97ec19c63287a7e140231e4214969d0eac923cd7ff", size = 5737, upload-time = "2025-05-13T15:24:01.64Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/04/11/432f32f8097b03e3cd5fe57e88efb685d964e2e5178a48ed61e841f7fdce/pyyaml_env_tag-1.1-py3-none-any.whl", hash = "sha256:17109e1a528561e32f026364712fee1264bc2ea6715120891174ed1b980d2e04", size = 4722, upload-time = "2025-05-13T15:23:59.629Z" }, +] + +[[package]] +name = "requests" +version = "2.34.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, +] + +[[package]] +name = "ruff" +version = "0.15.20" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/dc/35b341fc554ba02f217fc10da57d1a75168cfbcf75b0ef2202176d4c4f2d/ruff-0.15.20.tar.gz", hash = "sha256:1416eb04349192646b54de98f146c4f59afe37d0decfc02c3cbbf396f3a28566", size = 4755489, upload-time = "2026-06-25T17:20:37.578Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/d9/2d5014f0253ba541d2061d9fa7193f48e941c8b21bb88a7ff9bbe0bd0596/ruff-0.15.20-py3-none-linux_armv6l.whl", hash = "sha256:00e188c53e499c3c1637f73c91dcf2fb56d576cab76ce1be50a27c4e80e37078", size = 10839665, upload-time = "2026-06-25T17:19:44.702Z" }, + { url = "https://files.pythonhosted.org/packages/c6/d3/ac1798ba64f670698867fcfc591d50e7e421bef137db564858f619a30fcf/ruff-0.15.20-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9ebd1fd9b9c95fc0bd7b2761aebec1f030013d2e193a2901b224af68fe47251b", size = 11208649, upload-time = "2026-06-25T17:19:48.787Z" }, + { url = "https://files.pythonhosted.org/packages/47/47/d3ac899991202095dfcf3d5176be4272642be3cf981a2f1a30f72a2afb95/ruff-0.15.20-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c5b16cdd67ca108185cd36dce98c576350c03b1660a751de725fb049193a0632", size = 10622638, upload-time = "2026-06-25T17:19:51.354Z" }, + { url = "https://files.pythonhosted.org/packages/33/13/4e043fe30aa94d4ff5213a9881fc296d12960f5971b234a5263fdc225312/ruff-0.15.20-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3413bb3c3d2ca6a8208f1f4809cd2dca3c6de6d0b491c0e70847672bde6e6efd", size = 10984227, upload-time = "2026-06-25T17:19:54.044Z" }, + { url = "https://files.pythonhosted.org/packages/76/e6/92e7bf40388bc5800073b96564f56264f7e48bfd1a498f5ced6ae6d5a769/ruff-0.15.20-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bd7ec42b3bb3da066488db093308a69c4ac5ee6d2af333a86ba6e2eb2e7dd44b", size = 10622882, upload-time = "2026-06-25T17:19:57.037Z" }, + { url = "https://files.pythonhosted.org/packages/13/7a/43460be3f24495a3aa46d4b16873e2c4941b3b5f0b00cf88c03b7b94b339/ruff-0.15.20-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e1a36ad0eb77fba9aabfb69ede54de6f376d04ac18ebea022847046d340a8267", size = 11474808, upload-time = "2026-06-25T17:20:00.357Z" }, + { url = "https://files.pythonhosted.org/packages/27/a0/f37077884873221c6b33b4ab49eb18f9f88e54a16a25a5bca59bef46dd66/ruff-0.15.20-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b6df3b1e4610432f0386dba04d853b5f08cbbc903410c6fcc02f620f05aff53c", size = 12293094, upload-time = "2026-06-25T17:20:03.446Z" }, + { url = "https://files.pythonhosted.org/packages/a6/74/165545b60256a9704c21ac0ec4a0d07933b320812f9584836c9f4aca4292/ruff-0.15.20-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e89f198a1ea6ef0d727c1cf16088bc91a6cb0ab947dedc966715691647186eae", size = 11526176, upload-time = "2026-06-25T17:20:06.301Z" }, + { url = "https://files.pythonhosted.org/packages/86/b1/a976a136d40ade83ce743578399865f57001003a409acadc0ecbb3051082/ruff-0.15.20-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:309809086c2acb67624950a3c8133e80f32d0d3e27106c0cd60ff26657c9f24b", size = 11520767, upload-time = "2026-06-25T17:20:09.191Z" }, + { url = "https://files.pythonhosted.org/packages/19/0f/f032696cb01c9b54c0263fa393474d7758f1cdc021a01b04e3cbc2500999/ruff-0.15.20-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:2d2374caa2f2c2f9e2b7da0a50802cfb8b79f55a9b5e49379f564544fbf56487", size = 11500132, upload-time = "2026-06-25T17:20:13.602Z" }, + { url = "https://files.pythonhosted.org/packages/4b/f4/51b1a14bc69e8c224b15dab9cce8e99b425e0455d462caa2b3c9be2b6a8e/ruff-0.15.20-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:a1ed17b65293e0c2f22fc387bc13198a5de94bf4429589b0ff6946b0feaf21a3", size = 10943828, upload-time = "2026-06-25T17:20:16.635Z" }, + { url = "https://files.pythonhosted.org/packages/71/4b/fe267640783cd02bf6c5cc290b1df1051be2ec294c678b5c15fe19e52343/ruff-0.15.20-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:f701305e66b38ea6c91882490eb73459796808e4c6362a1b765255e0cdcd4053", size = 10645418, upload-time = "2026-06-25T17:20:19.4Z" }, + { url = "https://files.pythonhosted.org/packages/b0/c0/a65aa4ec2f5e87a1df32dc3ec1fede434fe3dfd5cbcf3b503cafc676ab54/ruff-0.15.20-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5b9c0c367ad8e5d0d5b5b8537864c469a0a0e55417aadfbeca41fa61333be9f4", size = 11211770, upload-time = "2026-06-25T17:20:22.033Z" }, + { url = "https://files.pythonhosted.org/packages/5a/a4/0caa331d954ae2723d729d351c989cb4ca8b6077d5c6c2cb6de75e98c041/ruff-0.15.20-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:01cc00dd58f0df339d0e902219dd53990ea99996a0344e5d9cc8d45d5307e460", size = 11618698, upload-time = "2026-06-25T17:20:25.259Z" }, + { url = "https://files.pythonhosted.org/packages/10/9b/5f14927848d2fd4aa891fd88d883788c5a7baba561c7874732364045708c/ruff-0.15.20-py3-none-win32.whl", hash = "sha256:ed65ef510e43a137207e0f01cfcf998aeddb1aeeda5c9d35023e910284d7cf21", size = 10857322, upload-time = "2026-06-25T17:20:28.612Z" }, + { url = "https://files.pythonhosted.org/packages/fa/f0/fe47c501f9dea92a26d788ff98bb5d92ed4cb4c88792c5c88af6b697dc8e/ruff-0.15.20-py3-none-win_amd64.whl", hash = "sha256:a525c81c70fb0380344dd1d8745d8cc1c890b7fc94a58d5a07bd8eb9557b8415", size = 11993274, upload-time = "2026-06-25T17:20:31.871Z" }, + { url = "https://files.pythonhosted.org/packages/d7/2b/9555445e1201d92b3195f45cdb153a0b68f24e0a4273f6e3d5ab46e212bb/ruff-0.15.20-py3-none-win_arm64.whl", hash = "sha256:2f5b2a6d614e8700388806a14996c40fab2c47b819ef57d790a34878858ed9ca", size = 11343498, upload-time = "2026-06-25T17:20:35.03Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "sortedcontainers" +version = "2.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e8/c4/ba2f8066cceb6f23394729afe52f3bf7adec04bf9ed2c820b39e19299111/sortedcontainers-2.4.0.tar.gz", hash = "sha256:25caa5a06cc30b6b83d11423433f65d1f9d76c4c6a0c90e3379eaa43b9bfdb88", size = 30594, upload-time = "2021-05-16T22:03:42.897Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/46/9cb0e58b2deb7f82b84065f37f3bffeb12413f947f9388e4cac22c4621ce/sortedcontainers-2.4.0-py2.py3-none-any.whl", hash = "sha256:a163dcaede0f1c021485e957a39245190e74249897e2ae4b2aa38595db237ee0", size = 29575, upload-time = "2021-05-16T22:03:41.177Z" }, +] + +[[package]] +name = "urllib3" +version = "2.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, +] + +[[package]] +name = "watchdog" +version = "6.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/db/7d/7f3d619e951c88ed75c6037b246ddcf2d322812ee8ea189be89511721d54/watchdog-6.0.0.tar.gz", hash = "sha256:9ddf7c82fda3ae8e24decda1338ede66e1c99883db93711d8fb941eaa2d8c282", size = 131220, upload-time = "2024-11-01T14:07:13.037Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/ea/3930d07dafc9e286ed356a679aa02d777c06e9bfd1164fa7c19c288a5483/watchdog-6.0.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:bdd4e6f14b8b18c334febb9c4425a878a2ac20efd1e0b231978e7b150f92a948", size = 96471, upload-time = "2024-11-01T14:06:37.745Z" }, + { url = "https://files.pythonhosted.org/packages/12/87/48361531f70b1f87928b045df868a9fd4e253d9ae087fa4cf3f7113be363/watchdog-6.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c7c15dda13c4eb00d6fb6fc508b3c0ed88b9d5d374056b239c4ad1611125c860", size = 88449, upload-time = "2024-11-01T14:06:39.748Z" }, + { url = "https://files.pythonhosted.org/packages/5b/7e/8f322f5e600812e6f9a31b75d242631068ca8f4ef0582dd3ae6e72daecc8/watchdog-6.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6f10cb2d5902447c7d0da897e2c6768bca89174d0c6e1e30abec5421af97a5b0", size = 89054, upload-time = "2024-11-01T14:06:41.009Z" }, + { url = "https://files.pythonhosted.org/packages/68/98/b0345cabdce2041a01293ba483333582891a3bd5769b08eceb0d406056ef/watchdog-6.0.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:490ab2ef84f11129844c23fb14ecf30ef3d8a6abafd3754a6f75ca1e6654136c", size = 96480, upload-time = "2024-11-01T14:06:42.952Z" }, + { url = "https://files.pythonhosted.org/packages/85/83/cdf13902c626b28eedef7ec4f10745c52aad8a8fe7eb04ed7b1f111ca20e/watchdog-6.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:76aae96b00ae814b181bb25b1b98076d5fc84e8a53cd8885a318b42b6d3a5134", size = 88451, upload-time = "2024-11-01T14:06:45.084Z" }, + { url = "https://files.pythonhosted.org/packages/fe/c4/225c87bae08c8b9ec99030cd48ae9c4eca050a59bf5c2255853e18c87b50/watchdog-6.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a175f755fc2279e0b7312c0035d52e27211a5bc39719dd529625b1930917345b", size = 89057, upload-time = "2024-11-01T14:06:47.324Z" }, + { url = "https://files.pythonhosted.org/packages/a9/c7/ca4bf3e518cb57a686b2feb4f55a1892fd9a3dd13f470fca14e00f80ea36/watchdog-6.0.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7607498efa04a3542ae3e05e64da8202e58159aa1fa4acddf7678d34a35d4f13", size = 79079, upload-time = "2024-11-01T14:06:59.472Z" }, + { url = "https://files.pythonhosted.org/packages/5c/51/d46dc9332f9a647593c947b4b88e2381c8dfc0942d15b8edc0310fa4abb1/watchdog-6.0.0-py3-none-manylinux2014_armv7l.whl", hash = "sha256:9041567ee8953024c83343288ccc458fd0a2d811d6a0fd68c4c22609e3490379", size = 79078, upload-time = "2024-11-01T14:07:01.431Z" }, + { url = "https://files.pythonhosted.org/packages/d4/57/04edbf5e169cd318d5f07b4766fee38e825d64b6913ca157ca32d1a42267/watchdog-6.0.0-py3-none-manylinux2014_i686.whl", hash = "sha256:82dc3e3143c7e38ec49d61af98d6558288c415eac98486a5c581726e0737c00e", size = 79076, upload-time = "2024-11-01T14:07:02.568Z" }, + { url = "https://files.pythonhosted.org/packages/ab/cc/da8422b300e13cb187d2203f20b9253e91058aaf7db65b74142013478e66/watchdog-6.0.0-py3-none-manylinux2014_ppc64.whl", hash = "sha256:212ac9b8bf1161dc91bd09c048048a95ca3a4c4f5e5d4a7d1b1a7d5752a7f96f", size = 79077, upload-time = "2024-11-01T14:07:03.893Z" }, + { url = "https://files.pythonhosted.org/packages/2c/3b/b8964e04ae1a025c44ba8e4291f86e97fac443bca31de8bd98d3263d2fcf/watchdog-6.0.0-py3-none-manylinux2014_ppc64le.whl", hash = "sha256:e3df4cbb9a450c6d49318f6d14f4bbc80d763fa587ba46ec86f99f9e6876bb26", size = 79078, upload-time = "2024-11-01T14:07:05.189Z" }, + { url = "https://files.pythonhosted.org/packages/62/ae/a696eb424bedff7407801c257d4b1afda455fe40821a2be430e173660e81/watchdog-6.0.0-py3-none-manylinux2014_s390x.whl", hash = "sha256:2cce7cfc2008eb51feb6aab51251fd79b85d9894e98ba847408f662b3395ca3c", size = 79077, upload-time = "2024-11-01T14:07:06.376Z" }, + { url = "https://files.pythonhosted.org/packages/b5/e8/dbf020b4d98251a9860752a094d09a65e1b436ad181faf929983f697048f/watchdog-6.0.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:20ffe5b202af80ab4266dcd3e91aae72bf2da48c0d33bdb15c66658e685e94e2", size = 79078, upload-time = "2024-11-01T14:07:07.547Z" }, + { url = "https://files.pythonhosted.org/packages/07/f6/d0e5b343768e8bcb4cda79f0f2f55051bf26177ecd5651f84c07567461cf/watchdog-6.0.0-py3-none-win32.whl", hash = "sha256:07df1fdd701c5d4c8e55ef6cf55b8f0120fe1aef7ef39a1c6fc6bc2e606d517a", size = 79065, upload-time = "2024-11-01T14:07:09.525Z" }, + { url = "https://files.pythonhosted.org/packages/db/d9/c495884c6e548fce18a8f40568ff120bc3a4b7b99813081c8ac0c936fa64/watchdog-6.0.0-py3-none-win_amd64.whl", hash = "sha256:cbafb470cf848d93b5d013e2ecb245d4aa1c8fd0504e863ccefa32445359d680", size = 79070, upload-time = "2024-11-01T14:07:10.686Z" }, + { url = "https://files.pythonhosted.org/packages/33/e8/e40370e6d74ddba47f002a32919d91310d6074130fe4e17dabcafc15cbf1/watchdog-6.0.0-py3-none-win_ia64.whl", hash = "sha256:a1914259fa9e1454315171103c6a30961236f508b9b623eae470268bbcc6a22f", size = 79067, upload-time = "2024-11-01T14:07:11.845Z" }, +] + +[[package]] +name = "zarr-indexing" +source = { editable = "." } +dependencies = [ + { name = "numpy" }, +] + +[package.optional-dependencies] +testing = [ + { name = "hypothesis" }, +] + +[package.dev-dependencies] +docs = [ + { name = "griffe-inherited-docstrings" }, + { name = "mkdocs" }, + { name = "mkdocs-material" }, + { name = "mkdocstrings" }, + { name = "mkdocstrings-python" }, + { name = "ruff" }, +] +test = [ + { name = "hypothesis" }, + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "hypothesis", marker = "extra == 'testing'", specifier = ">=6.160.0" }, + { name = "numpy", specifier = ">=2" }, +] +provides-extras = ["testing"] + +[package.metadata.requires-dev] +docs = [ + { name = "griffe-inherited-docstrings", specifier = "==1.1.3" }, + { name = "mkdocs", specifier = "==1.6.1" }, + { name = "mkdocs-material", specifier = "==9.7.6" }, + { name = "mkdocstrings", specifier = "==1.0.4" }, + { name = "mkdocstrings-python", specifier = "==2.0.5" }, + { name = "ruff", specifier = "==0.15.20" }, +] +test = [ + { name = "hypothesis", specifier = ">=6.160.0" }, + { name = "pytest" }, +] From a9f77feed486dda1d1497f0b9dc4c6f27149f294 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 21:19:33 +0200 Subject: [PATCH 21/32] fix(zarr-indexing): the defects an adversarial review found at the boundaries MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six reviewers went at the package, each proving findings by execution. The algebra held: ~85k chained selections against NumPy, ~24k part assemblies with poisoned buffers, 3.8k transform round-trips evaluated as coordinate maps, all clean. Everything below was at a boundary. `result()` and `__array__(copy=True)` could hand back a live view of the source. `_detach` asked whether the *source* was an ndarray, but a duck array that merely stores its data in NumPy returns NumPy views, and those went straight to the caller — while three docstrings promised the opposite unconditionally. It now asks whether the *result* owns its buffer, so memory is released only when sharing is disproved rather than when it cannot be established. The wire format could not reload its own output. `tolist()` renders every empty array as `[]` once the leading axis is the zero-length one, so an ordinary empty selection lost the axis it varied over, and the loader put it back on a different one by prepending singletons. Nested lists cannot express the shape either, so the body carries it. `index_domain_from_json` was a second undefended way into the same objects: a bare `int()` that truncated 3.9, coerced "3" and True, and let a non-string label into a tuple[str, ...]. It goes through the message layer now, as a transform body always did. With it: a rank ceiling, i64 checks on desugared bounds so normalization stays idempotent, ordered index_array_bounds, and typed errors where raw ones leaked. `sub_transform_to_selections` transposed its blocks two ways. A ConstantMap was emitted as a bare integer, which NumPy counts among the *advanced* indices whenever an index array is present, moving the broadcast axis to the front when a slice separates them — while `out_selection` is built positionally. And the correlated branch documented a points-major block but assembled the chunk selection in output order, so a residual slice before the coordinates arrived slice-major. Constants are length-one slices now, named in drop_axes, and the correlated scatter is permuted to the block NumPy actually returns. Nothing caught either one because `LazyArray` resolves a part through its own lowering and never reads `chunk_selection`; one of the two was even asserted as correct in a passing test's comment. Three further failures were one stale field: `_apply_vindex` carried `input_dimension` onto a map the vindex had just made correlated. The value outlived the shape that justified it and was believed later by a scatter that filed positions under the wrong axis — which is why one view's answer depended on how it was partitioned. The dependency is read back off the array now, and `__post_init__` checks the field it was wrong about: ArrayMap was the one map whose `input_dimension` nothing validated. Also: `oindex` over a correlated view applied its index tuple positionally, NumPy's vectorized rule, collapsing two arrays into one axis; `compose()` indexed an inner array's broadcast singletons by the raw coordinate and sized its one-dimensional shortcut by the input rank while gating on the output rank; and neither checked that the outer transform's output lands inside the inner domain, where a negative would have wrapped. Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/chunk_resolution.py | 85 ++++++++++- .../src/zarr_indexing/composition.py | 75 ++++++++- .../zarr-indexing/src/zarr_indexing/json.py | 128 ++++++++++++++-- .../src/zarr_indexing/lazy_array.py | 59 ++++--- .../src/zarr_indexing/messages.py | 120 +++++++++++++-- .../src/zarr_indexing/transform.py | 68 ++++++++- .../tests/test_chunk_resolution.py | 97 ++++++++++-- .../zarr-indexing/tests/test_composition.py | 54 +++++++ packages/zarr-indexing/tests/test_json.py | 144 +++++++++++++++++- .../zarr-indexing/tests/test_lazy_array.py | 35 ++++- .../zarr-indexing/tests/test_transform.py | 68 +++++++++ 11 files changed, 855 insertions(+), 78 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index c6068cf733..61c4f0d36b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -33,6 +33,7 @@ from __future__ import annotations +import math from typing import TYPE_CHECKING, Any import numpy as np @@ -290,6 +291,46 @@ def _dimension_map_slice(m: DimensionMap, dim_lo: int, dim_hi: int) -> slice: return _positional_slice(m.offset + m.stride * dim_lo, dim_hi - dim_lo, m.stride) +def _scatter_in_block_order( + scatter: np.ndarray[Any, np.dtype[np.intp]], + *, + residual_extents: list[int], + basic_before_arrays: int, + array_positions: list[int], + has_points_axis: bool, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Reorder a points-major scatter to match the block the chunk read returns. + + The scatter enumerates the correlated selection points-major: point varying + slowest, then the residual sliced axes in their own order. NumPy does not + always hand the block back that way. With the coordinate arrays adjacent it + places the points axis where they sit, so a residual slice *before* them puts + a residual axis first; only when a slice separates the arrays does the points + axis lead. Assigning one order into the other transposed the values silently + when the extents happened to agree, and raised a broadcast error when they + did not. + + Nothing is read here: the scatter is an index array, so matching the block is + a permutation of it. + """ + if not has_points_axis or len(array_positions) == 0: + return scatter + separated = array_positions[-1] - array_positions[0] + 1 != len(array_positions) + points_axis = 0 if separated else basic_before_arrays + if points_axis == 0: + return scatter + + residual_cells = math.prod(residual_extents) if len(residual_extents) > 0 else 0 + if residual_cells == 0: + return scatter + n_points = scatter.size // residual_cells + # Points-major, then move the points axis to where the block carries it. The + # scatter keeps the block's shape rather than being flattened: the caller + # assigns the block into it directly, so the two must agree axis for axis. + grid = scatter.reshape((n_points, *residual_extents)) + return np.moveaxis(grid, 0, points_axis) + + def _array_input_dim(m: ArrayMap, out_dim: int) -> int: """The input dimension an orthogonal `ArrayMap` places its values along.""" axis = array_map_dependent_axis(m) @@ -383,14 +424,30 @@ def sub_transform_to_selections( has_points_axis = any(d not in bound for d in range(sub_transform.input_rank)) points_shape: tuple[int, ...] = (-1,) if has_points_axis else () chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + # Where each entry lands in the block, so the scatter below can be put in + # the same order: the residual slices in the order they appear, and the + # positions the coordinate arrays occupy. + residual_extents: list[int] = [] + array_positions: list[int] = [] + basic_before_arrays = 0 for m in sub_transform.output: if isinstance(m, ConstantMap): - chunk_sel.append(m.offset) + # See the basic branch below on why this is a slice: an integer + # would join the advanced indices and move the points axis. + chunk_sel.append(slice(m.offset, m.offset + 1)) + if len(array_positions) == 0: + basic_before_arrays += 1 + residual_extents.append(1) elif isinstance(m, DimensionMap): d = m.input_dimension - chunk_sel.append(_dimension_map_slice(m, inclusive_min[d], exclusive_max[d])) + sl = _dimension_map_slice(m, inclusive_min[d], exclusive_max[d]) + chunk_sel.append(sl) + if len(array_positions) == 0: + basic_before_arrays += 1 + residual_extents.append(exclusive_max[d] - inclusive_min[d]) else: # ArrayMap idx = m.index_array.reshape(points_shape) + array_positions.append(len(chunk_sel)) chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) # Chunk resolution always supplies the flat scatter index for a # correlated transform. Absent one (a bare sub-transform), fall back to an @@ -404,17 +461,33 @@ def sub_transform_to_selections( n *= s out_scatter = slice(0, n) else: - out_scatter = out_indices + out_scatter = _scatter_in_block_order( + out_indices, + residual_extents=residual_extents, + basic_before_arrays=basic_before_arrays, + array_positions=array_positions, + has_points_axis=has_points_axis, + ) return tuple(chunk_sel), (out_scatter,), () chunk_sel = [] # annotated in the correlated branch above (same function scope) out_by_dim: dict[int, slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = {} + basic_drop_axes: list[int] = [] # Single-pass build for the basic / single-orthogonal-array cases. - # ConstantMap dims are dropped (no out_sel entry). + # ConstantMap dims read one cell and are reported in drop_axes. for out_dim, m in enumerate(sub_transform.output): if isinstance(m, ConstantMap): - chunk_sel.append(m.offset) + # A length-one slice, not the bare integer this used to emit. NumPy + # counts an integer among the *advanced* indices when the tuple also + # holds an index array, and moves the broadcast axis to the front + # whenever a slice separates the two — while `out_selection` below is + # built positionally and keeps the original order. The two sides then + # described transposed blocks. A slice is a basic index wherever it + # sits, so the order is the one it looks like, and the axis it leaves + # behind is squeezed by the caller through drop_axes. + chunk_sel.append(slice(m.offset, m.offset + 1)) + basic_drop_axes.append(out_dim) elif isinstance(m, DimensionMap): d = m.input_dimension dim_lo = inclusive_min[d] @@ -435,4 +508,4 @@ def sub_transform_to_selections( out_sel = [ out_by_dim.get(d, slice(inclusive_min[d], exclusive_max[d])) for d in range(input_rank) ] - return tuple(chunk_sel), tuple(out_sel), () + return tuple(chunk_sel), tuple(out_sel), tuple(basic_drop_axes) diff --git a/packages/zarr-indexing/src/zarr_indexing/composition.py b/packages/zarr-indexing/src/zarr_indexing/composition.py index cb5e661a45..ff4865397d 100644 --- a/packages/zarr-indexing/src/zarr_indexing/composition.py +++ b/packages/zarr-indexing/src/zarr_indexing/composition.py @@ -20,8 +20,11 @@ from __future__ import annotations +from typing import Any + import numpy as np +from zarr_indexing.errors import BoundsCheckError from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_indexing.transform import IndexTransform @@ -95,6 +98,37 @@ def _compose_dimension(outer: IndexTransform, inner_map: DimensionMap) -> Output ) +def _checked_position(position: int, size: int, axis: int) -> int: + """A single positional index into an inner index array, bounds-checked. + + The intermediate coordinate comes from the outer transform's outputs, which + nothing here has established lie inside the inner domain. Unchecked, a + negative one wraps NumPy-style and reads a cell at the far end of the axis — + silently, and with no way for a caller to tell. + """ + if position < 0 or position >= size: + raise BoundsCheckError( + f"composing reads position {position} on axis {axis} of an index " + f"array of extent {size}; the outer transform's output lies outside " + f"the inner transform's domain" + ) + return position + + +def _checked_positions( + positions: np.ndarray[Any, np.dtype[np.intp]], size: int +) -> np.ndarray[Any, np.dtype[np.intp]]: + """The same check, for a whole array of positional indices.""" + if positions.size > 0 and (int(positions.min()) < 0 or int(positions.max()) >= size): + raise BoundsCheckError( + f"composing reads positions in " + f"[{int(positions.min())}, {int(positions.max())}] of an index array " + f"of extent {size}; the outer transform's output lies outside the " + f"inner transform's domain" + ) + return positions + + def _compose_array( outer: IndexTransform, inner_map: ArrayMap, inner_origin: tuple[int, ...] ) -> OutputIndexMap: @@ -118,11 +152,19 @@ def _compose_array( all_constant = all(isinstance(m, ConstantMap) for m in outer.output) if all_constant: - # Evaluate arr_i at the single constant point + # Evaluate arr_i at the single constant point. A non-dependency axis of + # the inner array is a singleton that broadcasts over the whole domain, + # so the coordinate there is always 0 whatever the intermediate names — + # indexing such an axis by the raw coordinate walked off the end of it. idx = tuple( - m.offset - lo - for m, lo in zip( - (m for m in outer.output if isinstance(m, ConstantMap)), inner_origin, strict=True + 0 if size == 1 else _checked_position(m.offset - lo, size, axis) + for axis, (m, lo, size) in enumerate( + zip( + (m for m in outer.output if isinstance(m, ConstantMap)), + inner_origin, + arr_i.shape, + strict=True, + ) ) ) value = int(arr_i[idx]) @@ -131,6 +173,7 @@ def _compose_array( # For 1D inner array with a single outer output (simple case) if arr_i.ndim == 1 and len(outer.output) == 1: outer_map = outer.output[0] + rank = outer.input_rank if isinstance(outer_map, DimensionMap): dim = outer_map.input_dimension @@ -138,13 +181,29 @@ def _compose_array( outer.domain.inclusive_min[dim], outer.domain.exclusive_max[dim], dtype=np.intp ) intermediate_vals = outer_map.offset + outer_map.stride * user_indices - new_arr = arr_i[intermediate_vals - inner_origin[0]] - return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + positions = _checked_positions(intermediate_vals - inner_origin[0], arr_i.shape[0]) + # Shaped to the *outer* input rank, not left 1-D. The gate above is on + # the output rank, so a rank-2 outer reached here and built an array + # of the wrong rank for the domain it belongs to. + shape = (1,) * dim + (len(positions),) + (1,) * (rank - dim - 1) + return ArrayMap( + index_array=arr_i[positions].reshape(shape), + offset=offset_i, + stride=stride_i, + input_dimension=dim, + ) if isinstance(outer_map, ArrayMap): intermediate_vals = outer_map.offset + outer_map.stride * outer_map.index_array - new_arr = arr_i[intermediate_vals - inner_origin[0]] - return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + positions = _checked_positions(intermediate_vals - inner_origin[0], arr_i.shape[0]) + # The outer array already carries the outer input rank, so indexing + # with it keeps that rank; the dependency travels with it. + return ArrayMap( + index_array=arr_i[positions], + offset=offset_i, + stride=stride_i, + input_dimension=outer_map.input_dimension, + ) # General multi-dim case: not yet implemented raise NotImplementedError( diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py index b0b29e1a9b..bd446306b2 100644 --- a/packages/zarr-indexing/src/zarr_indexing/json.py +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -47,6 +47,7 @@ from __future__ import annotations +import math from collections import Counter from typing import Any, Required, TypedDict @@ -108,6 +109,7 @@ class OutputIndexMapJSON(TypedDict, total=False): input_dimension: int index_array: NestedIntList index_array_bounds: list[IndexValueJSON] + index_array_shape: list[int] class IndexTransformJSON(TypedDict, total=False): @@ -126,12 +128,18 @@ class IndexTransformJSON(TypedDict, total=False): def _lower_bound(bound: BoundJSON, where: str) -> int: - """Lower a canonical bound to a finite integer, rejecting infinities.""" + """Lower a canonical bound to a finite integer, rejecting infinities. + + Reached only with a bound the message layer has already validated as an + `index-value`, so the one thing left to rule out is a sentinel: an + `IndexDomain` addresses a finite array. + """ value = bound[0] if isinstance(bound, list) else bound if value == "-inf" or value == "+inf": - raise ValueError( + raise NdselError( + "invalid_json", f"{where} is infinite ({value!r}); an IndexDomain addresses a finite " - f"array and cannot lower an infinite bound" + f"array and cannot lower an infinite bound", ) return int(value) @@ -145,6 +153,13 @@ def _lower_index_array(raw: Any, where: str) -> np.ndarray[Any, np.dtype[np.intp `[0.9, 1.9]` would silently read cells 0 and 1, and `[true, false]` cells 1 and 0. Strings raise here rather than leaking NumPy's own conversion error. """ + if not isinstance(raw, list): + # A bare integer would become a rank-0 array and then be widened into a + # length-1 map, so a document that names no cells would select one. + raise NdselError( + "invalid_json", + f"{where} must be an array of integers, got {raw!r}", + ) try: arr = np.asarray(raw) except (TypeError, ValueError) as exc: @@ -188,17 +203,29 @@ def index_domain_to_json(domain: IndexDomain) -> IndexDomainJSON: def index_domain_from_json(data: IndexDomainJSON) -> IndexDomain: - """Construct an IndexDomain from its canonical JSON representation.""" + """Construct an IndexDomain from its canonical JSON representation. + + The document is validated by the message layer first, exactly as a transform + body is. Reading the keys directly would be a second, undefended way into + the same objects: `int(value)` alone accepts `3.9`, `"3"` and `True`, and + each of those builds a domain that is not the document's. + """ + if not isinstance(data, dict): + raise NdselError("invalid_json", f"an index domain must be a JSON object, got {data!r}") + body = normalize_ndsel({**data, "kind": "transform"}) inclusive_min = tuple( _lower_bound(b, f"input_inclusive_min[{i}]") - for i, b in enumerate(data["input_inclusive_min"]) + for i, b in enumerate(body["input_inclusive_min"]) ) exclusive_max = tuple( _lower_bound(b, f"input_exclusive_max[{i}]") - for i, b in enumerate(data["input_exclusive_max"]) + for i, b in enumerate(body["input_exclusive_max"]) + ) + return IndexDomain( + inclusive_min=inclusive_min, + exclusive_max=exclusive_max, + labels=_lower_labels(body["input_labels"]), ) - labels = _lower_labels(list(data["input_labels"])) - return IndexDomain(inclusive_min=inclusive_min, exclusive_max=exclusive_max, labels=labels) # --------------------------------------------------------------------------- @@ -223,12 +250,20 @@ def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON: if m.index_array.size == 1: value = int(m.index_array.reshape(-1)[0]) return {"offset": m.offset + m.stride * value} - return { + body: OutputIndexMapJSON = { "offset": m.offset, "stride": m.stride, "index_array": m.index_array.tolist(), "index_array_bounds": ["-inf", "+inf"], } + if m.index_array.size == 0: + # `tolist()` renders every empty array as `[]` once the leading axis is + # the zero-length one, so rank — and with it the axis the map varies + # over — does not survive the trip. Nested lists cannot express it + # either: `[[]]` is shape (1, 0) and nothing spells (0, 1). Carry the + # shape alongside so the map that comes back is the map that went out. + body["index_array_shape"] = list(m.index_array.shape) + return body def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: @@ -264,6 +299,54 @@ def _solo_dependency_axis(arr: np.ndarray[Any, Any]) -> int | None: return dep[0] if len(dep) == 1 else None +def _full_rank_index_array( + arr: np.ndarray[Any, np.dtype[np.intp]], + declared_shape: list[int] | None, + domain: IndexDomain, + where: str, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Give an incoming `index_array` the input rank the engine requires. + + Three cases, in order of how much the document tells us: + + - An explicit `index_array_shape` settles it. This is what our own emitter + writes for an empty array, whose shape nested lists cannot carry. + - An empty array with no declared shape has to be reconstructed from the + domain: it can only be empty because some input dimension is, so that + dimension takes the zero and every other takes a broadcast singleton. + Ambiguous when the domain is empty on more than one axis, and impossible + when it is empty on none — both are rejected rather than guessed at. + - A lower-rank non-empty array is aligned to the *trailing* input + dimensions, which is how NumPy broadcasts and how a producer that omits + leading singletons means it to be read. + """ + if declared_shape is not None: + if math.prod(declared_shape) != arr.size: + raise NdselError( + "invalid_json", + f"{where}.index_array_shape is {declared_shape}, which holds " + f"{math.prod(declared_shape)} entries, but index_array holds {arr.size}", + ) + return arr.reshape(tuple(declared_shape)) + + if arr.size == 0 and arr.ndim != domain.ndim: + empty_axes = [k for k, extent in enumerate(domain.shape) if extent == 0] + if len(empty_axes) != 1: + raise NdselError( + "invalid_json", + f"{where}.index_array is empty, but the input domain has " + f"{len(empty_axes)} zero-length dimensions, so the axis it varies " + f"over cannot be recovered; send 'index_array_shape' to say which", + ) + shape = [1] * domain.ndim + shape[empty_axes[0]] = 0 + return arr.reshape(tuple(shape)) + + if arr.ndim < domain.ndim: + return arr.reshape((1,) * (domain.ndim - arr.ndim) + arr.shape) + return arr + + # --------------------------------------------------------------------------- # IndexTransform serialization # --------------------------------------------------------------------------- @@ -294,7 +377,17 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: `input_dimension` values are reconstructed by global dependency-axis ownership (see the module docstring). """ - body = normalize_ndsel({"kind": "transform", **data}) + if not isinstance(data, dict): + raise NdselError("invalid_json", f"a transform body must be a JSON object, got {data!r}") + if data.get("kind", "transform") != "transform": + # Spelled last below, so a body carrying its own `kind` cannot reinterpret + # the document as some other message and return a selection this function + # never promised. + raise NdselError( + "invalid_json", + f"a transform body cannot carry kind {data['kind']!r}", + ) + body = normalize_ndsel({**data, "kind": "transform"}) inclusive_min = tuple( _lower_bound(b, f"input_inclusive_min[{i}]") @@ -320,13 +413,13 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: axis_owners: Counter[int] = Counter() for i, om in enumerate(output_raw): if "index_array" in om: - arr = _lower_index_array(om["index_array"], f"output[{i}].index_array") + where = f"output[{i}]" + arr = _lower_index_array(om["index_array"], f"{where}.index_array") # ndsel leaves index-array rank unvalidated, so an external producer # may send an array of lower rank that broadcasts against the domain. # Widen it here, on the way in, so every transform that exists holds # the full-rank invariant the engine reads dependency axes from. - if arr.ndim < domain.ndim: - arr = arr.reshape((1,) * (domain.ndim - arr.ndim) + arr.shape) + arr = _full_rank_index_array(arr, om.get("index_array_shape"), domain, where) arrays[i] = arr dep = _array_map_dependency_axes(arr) array_axes[i] = dep @@ -356,7 +449,14 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: else: output.append(ConstantMap(offset=om.get("offset", 0))) - return IndexTransform(domain=domain, output=tuple(output)) + try: + return IndexTransform(domain=domain, output=tuple(output)) + except ValueError as exc: + # The engine's invariants are the last gate a document passes, and they + # speak in the engine's vocabulary. A document that fails them is invalid + # input, so it leaves here as one — with the engine's account of what was + # wrong kept, since it names the offending output map and axis. + raise NdselError("rank_mismatch", str(exc)) from exc # Historical names, now pointing at the canonical converters. diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 694b59cfc6..50f350fbfb 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -820,19 +820,28 @@ def _detach(value: Any, source: Any) -> Any: fresh buffer. `result()` is answering "what does this view hold", not "lend me the storage", and its answer must not depend on how the read was divided. - NumPy is the namespace whose sharing can be settled here, from the arrays' - memory bounds; a foreign namespace whose basic indexing returns views is out - of reach, and its blocks are handed back as they come (see the module - docstring's device caveat). + Sharing is settled two ways, because the source is not always something + `may_share_memory` can be pointed at. When the wrapped array is itself a + NumPy array the two are compared directly, and the copy is made only when + they really overlap. When it is not — a duck array that stores its data in + NumPy and returns views from `__getitem__`, which is exactly the `BASIC` + source this package invites people to wrap — there is nothing to compare + against, so the question becomes whether the result owns its buffer. A NumPy + array that does not own its buffer is a view of something, and the only + thing it can be a view of here is storage the source lent us. + + Deciding it that way means the answer never depends on knowing what the + source is: memory is released only when sharing is *disproved*, never merely + because it could not be established. An array whose namespace is foreign + enough that it is not a NumPy array at all is still handed back as it comes + (see the module docstring's device caveat). """ + if not isinstance(value, np.ndarray): + return value base = _wrapped_base(source) - if ( - isinstance(value, np.ndarray) - and isinstance(base, np.ndarray) - and np.may_share_memory(value, base) - ): - return value.copy() - return value + if isinstance(base, np.ndarray): + return value.copy() if np.may_share_memory(value, base) else value + return value if value.flags.owndata else value.copy() def _partition_out_selection( @@ -845,18 +854,24 @@ def _partition_out_selection( A correlated sub-transform scatters through flat offsets into the row-major result; unravelling them turns that into an index tuple for the result's own shape, so callers never need a flat working buffer. + + The correlated scatter is taken directly rather than through + `sub_transform_to_selections`, whose business is to match a scatter to the + block a *NumPy chunk read* returns — which is not the block a part produces. + A part is resolved by lowering its own transform, and that lowering keeps the + points axis first; the bridge has to answer for NumPy's placement rules + instead, and permutes its scatter to suit. Reading the bridge's answer here + would apply a correction for a block this path never sees. """ - _chunk_selection, out_selection, _drop_axes = sub_transform_to_selections( - sub_transform, out_indices - ) if not _is_correlated(sub_transform): - return out_selection + return sub_transform_to_selections(sub_transform, out_indices)[1] if len(out_shape) == 0: return () - flat = out_selection[0] - if isinstance(flat, slice): - flat = np.arange(flat.start, flat.stop, dtype=np.intp) - return tuple(np.unravel_index(np.asarray(flat, dtype=np.intp), out_shape)) + if out_indices is None: + flat = np.arange(math.prod(sub_transform.domain.shape), dtype=np.intp) + else: + flat = np.asarray(out_indices, dtype=np.intp) + return tuple(np.unravel_index(flat, out_shape)) def _out_selection_cell_count(selection: tuple[Any, ...], out_shape: tuple[int, ...]) -> int: @@ -1518,7 +1533,11 @@ def result(self) -> Any: a partitioning is in force, not on how many boxes it has. The result never shares memory with the wrapped array, so writing to it - is safe whichever of those two routes it came by. + is safe whichever of those two routes it came by. That holds for a + source that stores its data in NumPy whether or not the source is itself + a NumPy array; a source whose namespace is foreign enough that its + blocks are not NumPy arrays is out of reach (see the module docstring's + device caveat). Returns ------- diff --git a/packages/zarr-indexing/src/zarr_indexing/messages.py b/packages/zarr-indexing/src/zarr_indexing/messages.py index a255f9a2a6..0cad32684f 100644 --- a/packages/zarr-indexing/src/zarr_indexing/messages.py +++ b/packages/zarr-indexing/src/zarr_indexing/messages.py @@ -91,9 +91,22 @@ def __init__(self, reason: str, detail: str = "") -> None: _TRANSFORM_UPPER = ("input_exclusive_max", "input_inclusive_max", "input_shape") _OUTPUT_MAP_FIELDS = frozenset( - {"offset", "stride", "input_dimension", "index_array", "index_array_bounds"} + { + "offset", + "stride", + "input_dimension", + "index_array", + "index_array_bounds", + "index_array_shape", + } ) +# An upper bound on `input_rank`, because normalization allocates proportionally +# to it — an identity `output`, a bound per dimension, a label per dimension — +# from a document that carries no data behind the number. Matches the rank +# TensorStore accepts, which is well above any real array. +_MAX_RANK = 32 + # --------------------------------------------------------------------------- # Leaf value validators @@ -262,28 +275,50 @@ def _resolve_upper_bound( implicit = _bound_is_implicit(raw) value = _bound_value(raw) if kind_of == "inclusive": - new = _inclusive_to_exclusive(value) + new = _inclusive_to_exclusive(value, f"{upper_field}[{k}]") else: # shape - new = _shape_to_exclusive(_bound_value(inclusive_min[k]), value) + new = _shape_to_exclusive(_bound_value(inclusive_min[k]), value, f"{upper_field}[{k}]") result.append(_rewrap(new, implicit=implicit)) return result -def _inclusive_to_exclusive(value: int | str) -> int | str: +def _checked_i64(value: int, where: str) -> int: + """An arithmetic result that must still be a 64-bit signed integer. + + Normalization is idempotent (spec section 4.3): whatever it emits must pass + the same validation on the way back in. Desugaring adds — `inclusive_max + 1`, + `inclusive_min + shape` — so a bound at the top of the range would otherwise + be emitted one past it and rejected by the next call on our own output. + """ + if value < _I64_MIN or value > _I64_MAX: + raise NdselError( + "invalid_json", + f"{where} is {value}, which is outside the 64-bit signed range; the " + f"normalized form cannot represent it", + ) + return value + + +def _inclusive_to_exclusive(value: int | str, where: str) -> int | str: if value == "+inf" or value == "-inf": return value assert isinstance(value, int) - return value + 1 + return _checked_i64(value + 1, f"{where} converted to an exclusive bound") -def _shape_to_exclusive(min_value: int | str, shape_value: int | str) -> int | str: +def _shape_to_exclusive(min_value: int | str, shape_value: int | str, where: str) -> int | str: + if shape_value == "-inf": + raise NdselError( + "invalid_json", + f"{where} is '-inf'; a shape counts cells and cannot be negatively infinite", + ) if shape_value == "+inf" or min_value == "+inf": return "+inf" if min_value == "-inf": return "-inf" assert isinstance(min_value, int) assert isinstance(shape_value, int) - return min_value + shape_value + return _checked_i64(min_value + shape_value, f"{where} added to its inclusive_min") def _validate_domain(inclusive_min: list[Any], exclusive_max: list[Any], *, prefix: str) -> None: @@ -512,12 +547,17 @@ def _normalize_output_map(raw: Any, where: str) -> dict[str, Any]: else ["-inf", "+inf"] ) # index_array is carried verbatim (spec section 7 defers shape validation). - return { + normalized: dict[str, Any] = { "offset": offset, "stride": stride, "index_array": raw["index_array"], "index_array_bounds": bounds, } + if "index_array_shape" in raw: + normalized["index_array_shape"] = _check_index_array_shape( + raw["index_array_shape"], where + ) + return normalized if has_input_dim: input_dim = _check_int(raw["input_dimension"], f"{where}.input_dimension") @@ -541,10 +581,41 @@ def _check_index_array_bounds(value: Any, where: str) -> list[int | str]: "invalid_json", f"{where}.index_array_bounds must be a two-element array, got {value!r}", ) - return [ - _check_index_value(value[0], f"{where}.index_array_bounds[0]"), - _check_index_value(value[1], f"{where}.index_array_bounds[1]"), - ] + lo = _check_index_value(value[0], f"{where}.index_array_bounds[0]") + hi = _check_index_value(value[1], f"{where}.index_array_bounds[1]") + if _ext_key(lo) > _ext_key(hi): + raise NdselError( + "bounds_out_of_order", + f"{where}.index_array_bounds: lower bound {lo!r} > upper bound {hi!r}", + ) + return [lo, hi] + + +def _check_index_array_shape(value: Any, where: str) -> list[int]: + """Validate an `index_array_shape`: non-negative extents, one per input dimension. + + Carried because JSON nested lists cannot express the shape of an array with + no elements — `[]` is the only spelling of every empty shape, so a leading + zero axis loses both its rank and which dimension it varies over. See the + `json` module docstring. + """ + if not isinstance(value, list): + raise NdselError( + "invalid_json", f"{where}.index_array_shape must be an array, got {value!r}" + ) + if len(value) > _MAX_RANK: + raise NdselError( + "invalid_json", + f"{where}.index_array_shape has rank {len(value)}, above the maximum {_MAX_RANK}", + ) + extents = [_check_int(v, f"{where}.index_array_shape[{i}]") for i, v in enumerate(value)] + for i, extent in enumerate(extents): + if extent < 0: + raise NdselError( + "invalid_json", + f"{where}.index_array_shape[{i}] must be >= 0, got {extent}", + ) + return extents def _normalize_transform(obj: dict[str, Any]) -> dict[str, Any]: @@ -567,6 +638,14 @@ def _normalize_transform(obj: dict[str, Any]) -> dict[str, Any]: declared_rank = _check_int(obj["input_rank"], "input_rank") if declared_rank < 0: raise NdselError("invalid_json", f"input_rank must be >= 0, got {declared_rank}") + if declared_rank > _MAX_RANK: + # Normalization fills a bound, a label and an identity output map per + # dimension, so an unbacked rank is a request to allocate from a + # document that carries nothing. + raise NdselError( + "invalid_json", + f"input_rank must be <= {_MAX_RANK}, got {declared_rank}", + ) inclusive_min_raw = ( _check_bound_list(obj["input_inclusive_min"], "input_inclusive_min") @@ -603,6 +682,23 @@ def _normalize_transform(obj: dict[str, Any]) -> dict[str, Any]: if not isinstance(obj["output"], list): raise NdselError("invalid_json", f"output must be an array, got {obj['output']!r}") output = [_normalize_output_map(m, f"output[{i}]") for i, m in enumerate(obj["output"])] + for i, m in enumerate(output): + # An `input_dimension` names one of *this* transform's input + # dimensions, so the rank is what bounds it. Checked here rather than + # in `_normalize_output_map`, which sees one map and not the rank. + if "input_dimension" in m and m["input_dimension"] >= rank: + raise NdselError( + "rank_mismatch", + f"output[{i}].input_dimension is {m['input_dimension']}, " + f"outside the valid range [0, {rank}) for input_rank {rank}", + ) + declared = m.get("index_array_shape") + if declared is not None and len(declared) != rank: + raise NdselError( + "rank_mismatch", + f"output[{i}].index_array_shape has rank {len(declared)} but " + f"input_rank is {rank}", + ) else: output = _identity_output(rank) diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 1d6cc3d6c2..e8f3f9d074 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -101,6 +101,29 @@ def __post_init__(self) -> None: f"(index_array shape {m.index_array.shape}, " f"domain shape {self.domain.shape})" ) + # `input_dimension` claims which axis the map varies over, and + # readers fall back to it when the shape alone cannot say. A + # value that names no axis, or names one the array does not + # actually vary over, is a claim the array contradicts — and it + # survives to be believed much later, by a scatter that files + # positions under the wrong axis. `DimensionMap` has always been + # range-checked here; this is the same check for the same field. + if m.input_dimension is not None: + if m.input_dimension < 0 or m.input_dimension >= self.domain.ndim: + raise ValueError( + f"output[{i}].input_dimension = {m.input_dimension} " + f"is out of range for input rank {self.domain.ndim}" + ) + dependency = _array_map_dependency_axes(m.index_array) + if len(dependency) > 1 or ( + len(dependency) == 1 and dependency[0] != m.input_dimension + ): + raise ValueError( + f"output[{i}] claims input_dimension=" + f"{m.input_dimension} but its index array varies over " + f"{dependency}; an orthogonal map varies over the one " + f"axis it names, and a correlated one names none" + ) @property def input_rank(self) -> int: @@ -763,9 +786,18 @@ def _reindex_array_oindex( m: ArrayMap, normalized: tuple[Any, ...] | list[Any], domain: IndexDomain, + *, + outer: bool, ) -> np.ndarray[Any, np.dtype[np.intp]]: """Apply an oindex/vindex selection to an existing ArrayMap's index_array. + `outer` says which dialect the caller is in, and the two disagree whenever + more than one entry is an index array. Orthogonal indexing means the outer + product of the per-axis selections; vectorized indexing means the arrays + broadcast against each other into one axis. Applying the tuple positionally + gives NumPy's vectorized answer, so an orthogonal caller was getting a + selection one rank smaller than it asked for. + Dependency-aware in exactly the way `_reindex_array` is: an entry applies to the array only along the axes the map genuinely varies over (its dependency axes, plus the `input_dimension` recorded for a degenerate length-1 @@ -804,7 +836,22 @@ def _reindex_array_oindex( else: idx.append(slice(None)) - result = arr[tuple(idx)] if idx else arr + if len(idx) == 0: + return np.asarray(arr, dtype=np.intp) + if outer and sum(isinstance(entry, np.ndarray) for entry in idx) > 1: + # An open mesh, so each axis's selection applies independently. Only + # built when it would differ from the positional form, to leave the + # single-array path — every existing orthogonal selection — untouched. + selectors = [ + entry + if isinstance(entry, np.ndarray) + else np.arange(arr.shape[axis])[entry] # a slice, resolved on its own axis + for axis, entry in enumerate(idx) + ] + selectors.extend(np.arange(arr.shape[axis]) for axis in range(len(idx), arr.ndim)) + result = arr[np.ix_(*selectors)] + else: + result = arr[tuple(idx)] return np.asarray(result, dtype=np.intp) @@ -1123,7 +1170,7 @@ def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: else: # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) _guard_fancy_after_fancy(m, list(dim_array.keys())) - new_arr = _reindex_array_oindex(m, normalized, transform.domain) + new_arr = _reindex_array_oindex(m, normalized, transform.domain, outer=True) array_input_dim: int | None = None if m.input_dimension is not None: array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) @@ -1323,13 +1370,26 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: else: # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) _guard_fancy_after_fancy(m, array_dims) - new_arr = _reindex_array_oindex(m, processed, transform.domain) + new_arr = _reindex_array_oindex(m, processed, transform.domain, outer=False) + # Read the dependency back off the array that was just built, rather + # than carrying the one the old array had. The selection can change + # which axes a map varies over — a vectorized index applied to an + # orthogonal map makes it correlated — and the stale value outlived + # the shape that justified it, to be believed later by readers that + # trust it when the shape alone cannot say. + dependency = _array_map_dependency_axes(new_arr) + broadcast_axes = range(n_before, n_before + n_broadcast_dims) + new_input_dim = ( + dependency[0] + if len(dependency) == 1 and dependency[0] not in broadcast_axes + else None + ) new_output.append( ArrayMap( index_array=new_arr, offset=m.offset, stride=m.stride, - input_dimension=m.input_dimension, + input_dimension=new_input_dim, ) ) diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index dc8846661b..f7b4815a96 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -1,19 +1,18 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import Any import numpy as np +import pytest from zarr.core.chunk_grids import ChunkGrid, FixedDimension, VaryingDimension from zarr_indexing import chunk_resolution from zarr_indexing.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections from zarr_indexing.domain import IndexDomain +from zarr_indexing.grid import dimension_grids_from_chunks from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import IndexTransform -if TYPE_CHECKING: - import pytest - class TestChunkResolutionIdentity: def test_single_chunk(self) -> None: @@ -387,7 +386,14 @@ def test_2d_correlated_vindex_diagonal_is_linear_in_points( class TestSubTransformToSelections: def test_constant_map(self) -> None: - """ConstantMap produces an int chunk selection and no drop axis. + """ConstantMap reads one cell through a length-one slice, and says so. + + Not the bare integer this used to emit. NumPy counts an integer among + the *advanced* indices whenever the tuple also holds an index array, and + moves the broadcast axis to the front when a slice separates them, while + `out_selection` is built positionally — so the two sides described + transposed blocks. A slice is basic wherever it sits, and the size-one + axis it leaves behind is named in `drop_axes` for the caller to squeeze. The domain dimension here is referenced by no output map, so its out-selection cannot come from a map: it is the whole axis, taken from @@ -400,9 +406,9 @@ def test_constant_map(self) -> None: output=(ConstantMap(offset=5),), ) chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) - assert chunk_sel == (5,) + assert chunk_sel == (slice(5, 6),) assert out_sel == (slice(0, 10),) - assert drop_axes == () + assert drop_axes == (0,) def test_dimension_map_stride_1(self) -> None: """DimensionMap with stride=1 produces contiguous slice.""" @@ -516,10 +522,12 @@ def test_mixed_maps_2d(self) -> None: ), ) chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) - assert chunk_sel[0] == 5 + assert chunk_sel[0] == slice(5, 6) assert chunk_sel[1] == slice(0, 10, 1) - # drop_axes is empty — integer in chunk_sel naturally drops the dim via numpy - assert drop_axes == () + # The constant's axis survives the read at size one and is squeezed by + # the caller; letting NumPy drop it through an integer index put the + # constant in the advanced-index group, where it reordered the block. + assert drop_axes == (0,) class TestChunkResolutionArrayMapFlavors: @@ -580,3 +588,72 @@ def test_correlated_scatter_with_residual_slice(self) -> None: # Output side: a single flat scatter index of shape (points, slice) = (2, 5). assert len(out_sel) == 1 assert np.asarray(out_sel[0]).shape == (2, 5) + + +class TestBridgeBlockOrder: + """`out[out_sel] = chunk[chunk_sel]` must address the same block on both sides. + + Nothing exercised this pairing before: `LazyArray` resolves a part by + lowering its own transform and only ever reads `out_selection`, so a + `chunk_selection` describing a differently-ordered block went unnoticed by + the whole suite while corrupting any consumer that followed the contract. + """ + + @staticmethod + def _assemble( + data: np.ndarray[Any, Any], chunks: tuple[int, ...], transform: IndexTransform + ) -> np.ndarray[Any, Any]: + grids = dimension_grids_from_chunks(chunks, data.shape) + out = np.zeros(transform.domain.shape, dtype=data.dtype) + for coords, local, out_indices in iter_chunk_transforms(transform, grids): + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(local, out_indices) + block_slice = tuple( + slice(c * size, c * size + grid.chunk_size(c)) + for c, size, grid in zip(coords, chunks, grids, strict=True) + ) + block = data[block_slice][chunk_sel] + if len(drop_axes) > 0: + block = np.squeeze(block, axis=drop_axes) + if len(out_sel) == 1 and isinstance(out_sel[0], np.ndarray): + out.reshape(-1)[out_sel[0]] = block + else: + out[out_sel] = block + return out + + @pytest.mark.parametrize("chunks", [(2, 3, 3), (1, 2, 2)], ids=["one-chunk", "many-chunks"]) + def test_a_constant_before_an_array_keeps_the_blocks_axis_order( + self, chunks: tuple[int, int, int] + ) -> None: + """`oindex[int, :, array]` — a constant, a slice, then an index array. + + NumPy folds an integer into the advanced-index group, and moves the + broadcast axis to the front when a slice separates the two. Emitting the + constant as an integer therefore transposed the chunk block against an + `out_selection` built positionally. + """ + data = np.arange(18).reshape(2, 3, 3) + transform = ( + IndexTransform.from_shape(data.shape)[1] + .translate_domain_to((0, 0)) + .oindex[:, np.array([0, 1, 2])] + ) + np.testing.assert_array_equal( + self._assemble(data, chunks, transform), data[1][:, [0, 1, 2]] + ) + + def test_a_residual_slice_before_the_points_keeps_the_scatter_in_step(self) -> None: + """A correlated read whose residual dimension precedes its coordinates. + + The scatter enumerates points-major, but NumPy leaves the points axis + where the coordinate arrays sit, so the block arrived slice-major. Equal + extents transposed the values silently; unequal ones raised. + """ + data = np.arange(12).reshape(3, 4) + transform = IndexTransform.from_shape(data.shape).vindex[..., np.array([0, 1, 2])] + np.testing.assert_array_equal(self._assemble(data, (3, 4), transform), data[..., [0, 1, 2]]) + + def test_a_residual_slice_before_the_points_of_a_different_length(self) -> None: + """The same shape, with extents that cannot coincide — this one raised.""" + data = np.arange(8).reshape(2, 4) + transform = IndexTransform.from_shape(data.shape).vindex[..., np.array([0, 1, 2])] + np.testing.assert_array_equal(self._assemble(data, (2, 4), transform), data[..., [0, 1, 2]]) diff --git a/packages/zarr-indexing/tests/test_composition.py b/packages/zarr-indexing/tests/test_composition.py index 3c2aed351b..263cb168ae 100644 --- a/packages/zarr-indexing/tests/test_composition.py +++ b/packages/zarr-indexing/tests/test_composition.py @@ -5,6 +5,7 @@ from zarr_indexing.composition import compose from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import IndexTransform @@ -251,3 +252,56 @@ def test_three_transforms(self) -> None: assert isinstance(abc.output[0], DimensionMap) assert abc.output[0].offset == 25 assert abc.output[0].stride == 2 + + +def test_composing_an_inner_array_with_a_broadcast_axis_wider_than_one_cell() -> None: + """A non-dependency axis is a singleton that broadcasts, so its coordinate is 0. + + Indexing it by the raw intermediate coordinate walked off the end of an axis + the array only has one entry for. + """ + outer = IndexTransform( + domain=IndexDomain(inclusive_min=(-2,), exclusive_max=(0,)), + output=(ConstantMap(offset=1), ConstantMap(offset=0)), + ) + inner = IndexTransform( + domain=IndexDomain(inclusive_min=(0, 0), exclusive_max=(2, 1)), + output=(ArrayMap(index_array=np.array([[13]], dtype=np.intp), offset=-2, stride=2),), + ) + composed = compose(outer, inner) + assert composed.output[0] == ConstantMap(offset=24) + + +def test_composing_a_one_dimensional_inner_array_under_a_higher_rank_outer() -> None: + """The shortcut gated on the output rank but sized by the input rank. + + A rank-2 outer therefore built a rank-1 array for a rank-2 domain, which the + engine's own invariant then rejected. + """ + outer = IndexTransform( + domain=IndexDomain.from_shape((2, 3)), + output=(DimensionMap(input_dimension=0, offset=1, stride=1),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((4,)), + output=(ArrayMap(index_array=np.array([7, 3, 5, 1], dtype=np.intp), input_dimension=0),), + ) + composed = compose(outer, inner) + array_map = composed.output[0] + assert isinstance(array_map, ArrayMap) + assert array_map.index_array.shape == (2, 1) + np.testing.assert_array_equal(array_map.index_array, np.array([[3], [5]])) + + +def test_composing_out_of_the_inner_domain_is_refused() -> None: + """An intermediate outside the inner domain wrapped NumPy-style and read a cell.""" + outer = IndexTransform( + domain=IndexDomain.from_shape((1,)), + output=(ConstantMap(offset=-3),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((2,)), + output=(ArrayMap(index_array=np.array([7, 3], dtype=np.intp), input_dimension=0),), + ) + with pytest.raises(BoundsCheckError, match="outside the inner transform's domain"): + compose(outer, inner) diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py index 5ca9f30234..bcc68e4640 100644 --- a/packages/zarr-indexing/tests/test_json.py +++ b/packages/zarr-indexing/tests/test_json.py @@ -346,10 +346,15 @@ def _index_array_body(index_array: Any, rank: int = 1, extent: int = 2) -> Index ([0, 1.5], "float64"), ([True, False], "bool"), (["a", "b"], "str"), - ("abc", "str"), + # Not lists at all, so they are turned away before their content is + # looked at: a bare string would be iterated into characters, and a bare + # integer would become a rank-0 array and then a length-1 map, so a + # document naming no cells would select one. + ("abc", "must be an array of integers"), + (5, "must be an array of integers"), ([None, None], "object"), ], - ids=["floats", "mixed", "bools", "strings", "string", "nulls"], + ids=["floats", "mixed", "bools", "strings", "string", "scalar", "nulls"], ) def test_a_non_integer_index_array_is_rejected(index_array: Any, detail: str) -> None: """An `index_array` addresses storage cells, so it must be integral. @@ -453,3 +458,138 @@ def test_an_index_array_that_does_not_span_its_domain_is_rejected() -> None: DimensionMap(input_dimension=1), ), ) + + +def test_an_empty_index_array_round_trips_with_its_axis_intact() -> None: + """Selecting nothing must survive a trip through JSON. + + `tolist()` renders every empty array as `[]` once the leading axis is the + zero-length one, so both the rank and the axis the map varies over were lost + on the way out — and the loader, widening by prepending singletons, put the + dependency back on a different axis. The body carries the shape for exactly + this case, so the map that comes back is the map that went out. + """ + for shape, selection in ( + ((5, 3), (np.array([], dtype=np.intp), slice(None))), + ((5, 5), (np.array([], dtype=np.intp), np.array([], dtype=np.intp))), + ): + transform = IndexTransform.from_shape(shape).oindex[selection] + body = index_transform_to_json(transform) + reloaded = index_transform_from_json(body) + + before = [m.index_array.shape for m in transform.output if isinstance(m, ArrayMap)] + after = [m.index_array.shape for m in reloaded.output if isinstance(m, ArrayMap)] + assert after == before + assert index_transform_to_json(reloaded) == body + + +def test_an_empty_index_array_without_a_declared_shape_is_recovered_from_the_domain() -> None: + """A producer that omits the shape is still readable when the domain settles it. + + An index array can only be empty because an input dimension is, so a domain + with exactly one zero-length dimension names the axis unambiguously. + """ + body: IndexTransformJSON = { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [0, 4], + "output": [{"index_array": []}, {"input_dimension": 1}], + } + array_map = index_transform_from_json(body).output[0] + assert isinstance(array_map, ArrayMap) + assert array_map.index_array.shape == (0, 1) + + +def test_an_ambiguous_empty_index_array_is_rejected() -> None: + """Two zero-length dimensions leave nothing to recover the axis from.""" + body: IndexTransformJSON = { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [0, 0], + "output": [{"index_array": []}, {"input_dimension": 1}], + } + with pytest.raises(NdselError) as excinfo: + index_transform_from_json(body) + assert excinfo.value.reason == "invalid_json" + assert "index_array_shape" in str(excinfo.value) + + +def test_a_declared_index_array_shape_that_does_not_fit_is_rejected() -> None: + body: IndexTransformJSON = { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [3, 4], + "output": [{"index_array": [1, 2, 0, 2], "index_array_shape": [1, 3]}], + } + with pytest.raises(NdselError) as excinfo: + index_transform_from_json(body) + assert excinfo.value.reason == "invalid_json" + assert "index_array_shape" in str(excinfo.value) + + +@pytest.mark.parametrize( + ("document", "reason", "detail"), + [ + ( + {"input_inclusive_min": [0.0], "input_exclusive_max": [3], "input_labels": [""]}, + "invalid_json", + "must be an integer", + ), + ( + {"input_inclusive_min": [0], "input_exclusive_max": ["3"], "input_labels": [""]}, + "invalid_json", + "must be an integer", + ), + ( + {"input_inclusive_min": [False], "input_exclusive_max": [True], "input_labels": [""]}, + "invalid_json", + "must be an integer", + ), + ( + {"input_inclusive_min": [0], "input_exclusive_max": [3], "input_labels": [5]}, + "invalid_json", + "must be a string", + ), + ( + {"input_inclusive_min": [0], "input_exclusive_max": [2**200], "input_labels": [""]}, + "invalid_json", + "64-bit signed range", + ), + ], + ids=["float", "string", "bool", "non-string-label", "out-of-range"], +) +def test_a_malformed_domain_document_is_rejected(document: Any, reason: str, detail: str) -> None: + """The domain loader validates what the message layer validates. + + Reading the keys directly was a second, undefended way into the same + objects: a bare `int()` truncated `3.9` to 3, coerced `"3"` and `True`, and + let a non-string label into a `tuple[str, ...]` — each building a domain + that was not the document's, and re-dumping as a different document. + """ + with pytest.raises(NdselError) as excinfo: + index_domain_from_json(document) + assert excinfo.value.reason == reason + assert detail in str(excinfo.value) + + +def test_a_transform_body_cannot_reinterpret_itself_as_another_message() -> None: + """A `kind` inside the body must not change which message is being read.""" + with pytest.raises(NdselError) as excinfo: + index_transform_from_json({"kind": "points", "coords": [[1, 2], [3, 4]]}) + assert excinfo.value.reason == "invalid_json" + assert "kind" in str(excinfo.value) + + +def test_an_engine_invariant_failure_leaves_the_loader_as_a_typed_error() -> None: + """A document is invalid input however deep the check that catches it lives. + + The engine's rank and span invariants are the last gate a document passes, + and they raised a bare `ValueError` written in the engine's vocabulary. + """ + body: IndexTransformJSON = { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [5], + "input_labels": [""], + "output": [{"index_array": [[1, 2], [3, 4]]}], + } + with pytest.raises(NdselError) as excinfo: + index_transform_from_json(body) + assert excinfo.value.reason == "rank_mismatch" diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 26b3a9d362..d2b3e2e428 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -1835,7 +1835,31 @@ def test_the_strict_doubles_notice_an_over_broad_key() -> None: # --------------------------------------------------------------------------- +class NumpyBackedSource: + """A duck array that stores its data in NumPy and returns views from reads. + + Not an `np.ndarray`, so nothing about the source can be compared against the + result's memory — but its blocks are NumPy views of storage the caller must + not be handed. The `BASIC` source this package invites people to wrap. + """ + + def __init__(self, data: np.ndarray[Any, Any]) -> None: + self.data = data + + @property + def shape(self) -> tuple[int, ...]: + return self.data.shape + + @property + def dtype(self) -> Any: + return self.data.dtype + + def __getitem__(self, selection: Any) -> Any: + return self.data[selection] + + @pytest.mark.parametrize("parts", [None, (2, 2, 2), SHAPE]) +@pytest.mark.parametrize("wrap", [lambda d: d, NumpyBackedSource], ids=["ndarray", "duck"]) @pytest.mark.parametrize( ("build", "description"), [ @@ -1846,7 +1870,10 @@ def test_the_strict_doubles_notice_an_over_broad_key() -> None: ], ) def test_materializing_never_hands_back_the_wrapped_array( - parts: Any, build: Callable[[LazyArray], LazyArray], description: str + parts: Any, + wrap: Callable[[np.ndarray[Any, Any]], Any], + build: Callable[[LazyArray], LazyArray], + description: str, ) -> None: """Writing to a materialized result must never reach the source. @@ -1854,9 +1881,13 @@ def test_materializing_never_hands_back_the_wrapped_array( the wrapped array, and NumPy 2 hands whatever `__array__` returns straight to the caller. Every route out of the wrapper detaches, so the answer does not depend on how the read happened to be divided. + + Run over both a raw `ndarray` and a duck array that merely stores its data in + NumPy: the second is the case where the source cannot be compared against the + result, so detaching has to decide from the result's own buffer instead. """ data = reference() - view = build(LazyArray(data).with_parts(parts)) + view = build(LazyArray(wrap(data)).with_parts(parts)) for materialize in ( lambda v: v.result(), diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index 237e4271a0..42643b7cc5 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -4,6 +4,7 @@ import pytest from zarr_indexing.domain import IndexDomain +from zarr_indexing.lazy_array import LazyArray from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import ( IndexTransform, @@ -644,3 +645,70 @@ def test_length1_orthogonal_not_treated_as_correlated(self) -> None: assert result is not None _restricted, out_indices = result assert isinstance(out_indices, dict) + + +class TestDerivedMapDependency: + """A map's `input_dimension` must describe the array it is built with. + + Three separate failures came from one stale value: a vectorized index applied + to an orthogonal map makes it correlated, but the old dependency was carried + onto the new array anyway. Readers fall back to that field when the shape + alone cannot say, so the wrong axis was believed much later — by a scatter + that filed positions under it, which is why the answer depended on how the + read was partitioned. + """ + + def test_a_vindex_over_a_fancy_view_is_marked_correlated(self) -> None: + base = np.arange(6) + view = ( + LazyArray(base) + .lazy.oindex[np.array([0, 1])] + .lazy.vindex[np.array([[0, 1, 0], [1, 0, 1]])] + ) + np.testing.assert_array_equal( + np.asarray(view.result()), base[[0, 1]][[[0, 1, 0], [1, 0, 1]]] + ) + + def test_the_same_view_resolves_alike_however_it_is_partitioned(self) -> None: + base = np.arange(36).reshape(6, 6) + + def build(array: LazyArray) -> LazyArray: + return array.lazy.oindex[np.array([-3, -6, -4]), -4].lazy.vindex[np.array([[-2, -3]])] + + unpartitioned = np.asarray(build(LazyArray(base)).result()) + partitioned = np.asarray(build(LazyArray(base).with_parts((3, 3))).result()) + np.testing.assert_array_equal(partitioned, unpartitioned) + np.testing.assert_array_equal(unpartitioned, np.array([[2, 20]])) + + def test_an_array_map_claiming_an_axis_it_does_not_vary_over_is_rejected(self) -> None: + """The validation that would have caught the two above at their source.""" + with pytest.raises(ValueError, match="varies over"): + IndexTransform( + domain=IndexDomain.from_shape((2, 3)), + output=( + ArrayMap(index_array=np.array([[0, 1, 2]], dtype=np.intp), input_dimension=0), + ), + ) + + def test_an_array_map_input_dimension_out_of_range_is_rejected(self) -> None: + with pytest.raises(ValueError, match="out of range"): + IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=np.array([0], dtype=np.intp), input_dimension=99),), + ) + + +def test_an_orthogonal_step_over_a_correlated_view_is_an_outer_product() -> None: + """`oindex` after `vindex` means the outer product, not a joint gather. + + The reindexing applied its index tuple positionally, which is NumPy's + *vectorized* rule, so two arrays collapsed into one axis and the result came + back a rank short of what was asked for. + """ + base = np.arange(14).reshape(7, 2) + view = LazyArray(base).lazy.vindex[ + np.array([[5, 5], [1, 2], [0, 4]]), np.array([[1, 1], [1, 0], [1, 0]]) + ] + result = np.asarray(view.lazy.oindex[np.array([1, 1, 0]), np.array([1, 1, 0, 1])].result()) + assert result.shape == (3, 4) + np.testing.assert_array_equal(result, np.array([[4, 4, 3, 4], [4, 4, 3, 4], [11, 11, 11, 11]])) From c7c279b2534bd73fd09b5e549c6dfe3c8a2ad84e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 21:41:47 +0200 Subject: [PATCH 22/32] test(zarr-indexing): generate the selections that were never generated MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A mutation audit ran 31 mutations and 10 survived. They were not scattered: the invariants check thoroughly what a view *returns* and never what it *claims about itself*, and the strategies could not draw whole classes of selection. The generators now draw them. Orthogonal slices carry a step and may stop early, so a strided or reversed slice reaches `oindex` at all. Coordinate lists and masks can be empty, so a fancy selection that selects nothing exists — the shape that lost its axis on the way through JSON. Vectorized coordinate arrays can be multi-dimensional, so a rank-raising `vindex` is generated. Measured over 4000 draws each, every one of those counts was previously 0. The first run of the widened generators found a live defect: an empty Python list carries no element type and NumPy defaults it to float64, so `oindex[[]]` was refused as a non-integer index array. `json.py` already had that case; the boundary did not. NumPy takes `a[np.ix_([])]`, and so does this now. `Partition.is_complete` gets an invariant. It is what a consumer reads to decide it may take a whole-box read, so a wrongly-`True` one is silent corruption — and three separate mutations to `_covers_whole_part` survived the entire suite, including ones reporting `is_complete` for a part carrying two of its three cells. Asserted one way only: the flag is documented as conservative and only the claim to cover everything has to be earned. Two more state machines. The sorted one-dimensional fancy path needs both ranks to be 1, and the output rank is the *source's*, so the rank-3 default source walled it off entirely — 0 hits from the machine against 105 from the unit tests, in the path where reordering and duplicate coordinates are partitioned. A rank-1 source now reaches it 322 times per run. A source with an extent-1 axis covers the other side of a distinction the code draws from the domain rather than the array. Finally the two remaining mutants, both checked by mutating and confirming the failure: the index-array bound checks are probed one past the boundary rather than comfortably outside it, and `_out_selection_cell_count`'s guard is tested directly, since a partition walk only ever produces the forward in-bounds intervals that never reach it. Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/boundary.py | 6 ++ .../src/zarr_indexing/testing/stateful.py | 15 +++++ .../src/zarr_indexing/testing/strategies.py | 52 ++++++++++++++-- .../zarr-indexing/tests/test_lazy_array.py | 60 +++++++++++++++++++ .../tests/test_lazy_array_stateful.py | 37 +++++++++++- .../zarr-indexing/tests/test_transform.py | 27 +++++++++ 6 files changed, 191 insertions(+), 6 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py index 093f1295b7..16f9506d62 100644 --- a/packages/zarr-indexing/src/zarr_indexing/boundary.py +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -47,6 +47,12 @@ def _as_index_array(sel: Any) -> np.ndarray[Any, np.dtype[Any]] | None: return sel if isinstance(sel, (list, tuple)): arr = np.asarray(sel) + if arr.size == 0 and arr.dtype.kind == "f": + # An empty Python list carries no element type and NumPy defaults it + # to float64. Selecting nothing is legal — NumPy takes `a[np.ix_([])]` + # — so read it as the empty integer selection it spells, rather than + # rejecting it for a dtype it never had a chance to have. + return np.zeros(arr.shape, dtype=np.intp) if arr.dtype.kind in "biu": return arr raise IndexError( diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 60d9a9183d..901480149a 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -52,6 +52,7 @@ def make_source(self, data): from __future__ import annotations +import math from typing import TYPE_CHECKING, Any, ClassVar import numpy as np @@ -338,6 +339,20 @@ def parts_tile_the_view(self) -> None: f"part {part.base_coords} carries {value.shape} for an out_selection " f"addressing {assembled[part.out_selection].shape}: {self.chain}" ) + # `is_complete` is what a consumer reads to decide it may take a + # whole-box read and skip assembling anything, so a wrongly-`True` + # one is the silent-corruption case. Asserted one way only: the flag + # is documented as conservative, free to say `False` about a part it + # does cover (a strided walk over a one-cell box, a fancy axis that + # happens to enumerate everything), and only the claim to cover + # everything has to be earned. Both quantities are already in hand + # here, and nothing else in the suite compares them. + if part.is_complete: + box_cells = math.prod(stop - start for start, stop in part.box) + assert value.size == box_cells, ( + f"part {part.base_coords} reports is_complete but carries " + f"{value.size} of its box's {box_cells} cells: {self.chain}" + ) assembled[part.out_selection] = value np.add.at(hits, part.out_selection, 1) diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py index f49d3a360d..63a3d351a9 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py @@ -35,6 +35,7 @@ def test_my_array_slices_like_numpy(selection): __all__ = [ "basic_selections", + "empty_masks", "masks", "orthogonal_selections", "slice_selections", @@ -68,12 +69,29 @@ def _basic_entry(size: int) -> st.SearchStrategy[Any]: def _orthogonal_entry(size: int) -> st.SearchStrategy[Any]: + """One axis of an `oindex` selection. + + The slices carry a step and are free to stop early. Drawing them as + `slice(start, size)` alone meant no strided or reversed slice ever reached + `oindex`, and no orthogonal selection ever stopped short of the axis end. + + An empty coordinate list is drawn too. It selects nothing, which is legal + and is exactly the shape that lost its axis on the way through JSON — but + with `min_size=1` no fancy selection was ever empty. + """ coordinate = st.integers(-size, size - 1) return st.one_of( coordinate, st.lists(coordinate, min_size=1, max_size=4), + st.just([]), masks((size,)), - st.builds(slice, st.integers(0, size - 1), st.just(size)), + empty_masks((size,)), + st.builds( + slice, + st.integers(0, size - 1), + st.integers(0, size) | st.none(), + st.integers(1, 3) | st.integers(-3, -1), + ), ) @@ -99,6 +117,17 @@ def masks(draw: st.DrawFn, shape: tuple[int, ...]) -> np.ndarray[Any, np.dtype[n return flags.reshape(shape) +def empty_masks(shape: tuple[int, ...]) -> st.SearchStrategy[np.ndarray[Any, np.dtype[np.bool_]]]: + """The all-False mask over `shape` — a fancy selection that empties the view. + + Split out from `masks`, which forces a cell True so a chain has something + left to index at the next step. Drawn on its own because an empty fancy + selection is a shape the code paths treat separately, and nothing generated + one. + """ + return st.just(np.zeros(shape, dtype=np.bool_)) + + def basic_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]: """Basic selections: one scalar integer or slice per axis. @@ -150,14 +179,27 @@ def vectorized_selections(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[Any, if draw(st.booleans()): entries = [draw(masks(tuple(sizes)))] else: - length = draw(st.integers(1, 4)) + # The coordinate arrays share one shape, which is what makes the + # selection correlated. That shape is not always one-dimensional: a + # vectorized read of a (2, 3) block of points is an ordinary thing to + # ask for and produces a result of that rank. Drawing only 1-D arrays + # meant no rank-raising vindex was ever generated — and a length of 0 + # covers the empty case the same way `_orthogonal_entry` does. + coordinate_shape = draw( + st.one_of( + st.integers(0, 4).map(lambda length: (length,)), + st.tuples(st.integers(1, 2), st.integers(1, 3)), + ) + ) entries = [ draw( st.one_of( st.integers(-size, size - 1), - st.lists(st.integers(-size, size - 1), min_size=length, max_size=length).map( - lambda values: np.array(values, dtype=np.intp) - ), + st.lists( + st.integers(-size, size - 1), + min_size=int(np.prod(coordinate_shape)), + max_size=int(np.prod(coordinate_shape)), + ).map(lambda values: np.array(values, dtype=np.intp).reshape(coordinate_shape)), ) ) for size in sizes diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index d2b3e2e428..f0ee8e2c06 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -28,6 +28,7 @@ array_map_dependent_axis, dimension_grids_from_chunks, ) +from zarr_indexing.lazy_array import _out_selection_cell_count from zarr_indexing.support import IndexingSupport if TYPE_CHECKING: @@ -2115,3 +2116,62 @@ def test_a_zero_length_axis_accepts_every_spelling_of_no_chunks(parts: Any) -> N def test_a_zero_chunk_on_a_nonempty_axis_is_still_rejected() -> None: with pytest.raises(ValueError, match="chunk sizes must be positive"): LazyArray(np.zeros((4, 3))).with_parts(((0, 4), (3,))) + + +@pytest.mark.parametrize( + "selection", + [ + (slice(None, None, -1), slice(None), slice(None)), + (slice(3, 1, -1), slice(None), slice(None)), + (slice(None, None, -2), slice(None, None, -1), slice(None)), + ], + ids=["reversed", "reversed-partial", "reversed-strided"], +) +def test_the_coverage_count_agrees_with_numpy_for_reversed_selections( + selection: tuple[Any, ...], +) -> None: + """The safety net behind `result()`'s coverage assertion, checked on its own. + + `_out_selection_cell_count` sizes a partition's `out_selection` without + materializing it, and `result()` trusts that count to decide whether the + walk covered the view. Nothing pinned it for a reversed slice, so dropping + its `start <= stop` guard — or wrapping the subtraction in `abs()` — left + the suite green. A net nobody tests only matters once something else breaks, + which is exactly when it needs to be right. + """ + data = reference() + view = LazyArray(data).with_parts((2, 2, 2)).lazy[selection] + out_shape = view.shape + for part in view.parts(): + counted = _out_selection_cell_count(part.out_selection, out_shape) + assert counted == np.empty(out_shape)[part.out_selection].size + + +@pytest.mark.parametrize( + ("selection", "out_shape", "expected"), + [ + (((slice(2, 5)),), (10,), 3), + # A backwards interval selects nothing. The fast path subtracts, which + # would make this negative and let an incomplete walk sum to the view's + # own size — so the count falls back to `range` whenever the interval is + # not a forward, in-bounds one. + (((slice(5, 2)),), (10,), 0), + # Counts from the end, to index 8 — past the stop, so nothing. + (((slice(-2, 5)),), (10,), 0), + ((slice(5, 2), slice(0, 3)), (10, 10), 0), + ], + ids=["forward", "backwards", "negative-start", "backwards-in-a-pair"], +) +def test_the_coverage_count_matches_numpy_for_intervals_the_fast_path_declines( + selection: tuple[Any, ...], out_shape: tuple[int, ...], expected: int +) -> None: + """The guard on `result()`'s safety net, exercised where the walk cannot reach it. + + A partition walk only ever produces concrete forward in-bounds intervals, so + the guard that keeps everything else off the subtraction fast path is not + reachable through `parts()` at all — which is why removing it left the whole + suite green. It is the net's own contract, so it is checked directly. + """ + counted = _out_selection_cell_count(selection, out_shape) + assert counted == np.empty(out_shape)[selection].size + assert counted == expected diff --git a/packages/zarr-indexing/tests/test_lazy_array_stateful.py b/packages/zarr-indexing/tests/test_lazy_array_stateful.py index a50cd8f59c..475dd53515 100644 --- a/packages/zarr-indexing/tests/test_lazy_array_stateful.py +++ b/packages/zarr-indexing/tests/test_lazy_array_stateful.py @@ -20,8 +20,9 @@ from __future__ import annotations -from typing import Any +from typing import Any, ClassVar +import numpy as np import pytest from hypothesis import settings @@ -35,6 +36,40 @@ TestNumpyIndexing = state_machine_test(ChainedIndexingStateMachine) +class OneDimensionalIndexing(ChainedIndexingStateMachine): + """The same chains over a rank-1 source. + + The sorted one-dimensional fancy path in `chunk_resolution` is entered only + when both the input and output ranks are 1, and the output rank is the + *source's* — so the rank-3 default walls that path off from the machine + entirely, and from every downstream project told to subclass it. That path + is where reordering and duplicate coordinates are partitioned, which is the + corruption class this whole harness exists to catch. + """ + + data = np.arange(30, dtype=np.int64) + partitionings: ClassVar[tuple[Any, ...]] = (None, (4,), (30,), ((7, 8, 15),)) + + +TestOneDimensionalIndexing = state_machine_test(OneDimensionalIndexing) + + +class SingletonAxisIndexing(ChainedIndexingStateMachine): + """A source with an extent-1 axis, which the code must not read as a broadcast one. + + An index array's axis is a singleton either because the map broadcasts over + it or because the domain is genuinely one cell wide there, and the two are + told apart by the domain rather than the array. Nothing generated the second + kind. + """ + + data = np.arange(2 * 1 * 3, dtype=np.int64).reshape(2, 1, 3) + partitionings: ClassVar[tuple[Any, ...]] = (None, (1, 1, 1), (2, 1, 2)) + + +TestSingletonAxisIndexing = state_machine_test(SingletonAxisIndexing) + + class ZarrIndexing(ChainedIndexingStateMachine): """The same chains against a zarr array, read at every level it serves. diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index 42643b7cc5..3e2540f51b 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -4,6 +4,7 @@ import pytest from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError from zarr_indexing.lazy_array import LazyArray from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import ( @@ -712,3 +713,29 @@ def test_an_orthogonal_step_over_a_correlated_view_is_an_outer_product() -> None result = np.asarray(view.lazy.oindex[np.array([1, 1, 0]), np.array([1, 1, 0, 1])].result()) assert result.shape == (3, 4) np.testing.assert_array_equal(result, np.array([[4, 4, 3, 4], [4, 4, 3, 4], [11, 11, 11, 11]])) + + +@pytest.mark.parametrize( + ("value", "description"), + [(1, "one below the lower bound"), (10, "the exclusive upper bound itself")], +) +def test_an_index_array_value_just_outside_the_domain_is_refused( + value: int, description: str +) -> None: + """The bound checks are probed at the boundary, not comfortably past it. + + Both were only ever exercised from well outside the domain, so relaxing + either by one — `lo - 1` instead of `lo` — went unnoticed while letting a + view read a cell it does not address. + """ + transform = IndexTransform.from_shape((12,))[2:10] + with pytest.raises(BoundsCheckError, match="out of bounds"): + transform.oindex[np.array([value, 3])] + + +def test_an_index_array_value_at_each_end_of_the_domain_is_accepted() -> None: + """The other side of the same boundary: the extremes themselves are in range.""" + transform = IndexTransform.from_shape((12,))[2:10] + array_map = transform.oindex[np.array([2, 9])].output[0] + assert isinstance(array_map, ArrayMap) + np.testing.assert_array_equal(array_map.index_array, np.array([2, 9])) From 5f986b27be57a6b48c0d94f2c2389a246eaa2f0e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 31 Jul 2026 22:14:30 +0200 Subject: [PATCH 23/32] refactor(zarr-indexing)!: settle the API decisions that get dearer after 1.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `with_parts` decided what it had been given by inspecting the type of it: a sequence of integers meant uniform boxes, a sequence of sequences meant per-axis sizes, and `None` meant no partitioning at all — which also sent `result()` down an entirely different code path. Three semantics behind one parameter, and no way to ask for one of them and be told when you had spelled it wrong. They are now `with_parts`, `with_parts_per_axis` and `unpartitioned`. A harness that draws from a list of mixed partitionings still needs the dispatch, so it exists once, as `zarr_indexing.testing.repartition`. The sizes those methods take are relative to the array being read, not to the view reading it, and a narrowed view partitions the base extents — which could only be discovered from an error message. `base_shape` says it. `ArrayMap`, `IndexTransform` and `Partition` are `frozen=True`, which reads as a promise that a value can be compared and hashed. Both raised: `==` on the index arrays returned an array and then `ValueError: the truth value of an array is ambiguous`, and `hash()` refused an ndarray outright. So no transform could enter a set or key a cache, and the package had already grown an internal `_is_identity_transform` because of it. An `ArrayMap` was frozen but the array inside it was not, so reaching through a view's transform to `index_array[0] = 9` silently changed what the view returned. It is held through a read-only view now — a view rather than a flag on the caller's array, since constructing a map should not take away the right to write to an array you still own. `IndexDomain.narrow` clamped a slice bound to the domain, in a package whose stated invariant is no clamping and no negative wrapping. `narrow(slice(-3, None))` on `[0, 10)` therefore returned the whole axis — reading as the NumPy spelling of "the last three" and answering with something else — and a stop past the end returned a domain its own parent did not contain. Both raise `BoundsCheckError` now, and a stride raises `ValueError` like every other unimplemented request rather than `IndexError`. `Partition.array` was a view of the array while `LazyArray.array` was the raw source: adjacent types, one name, inverted meanings. The part's is `view`. Smaller: `BoundsCheckError` and `VindexInvalidSelectionError` are exported at the top level, being what exported functions raise; `errors.py` no longer claims `zarr.errors` re-exports them by identity, which is false and would have been believed; `parts()` says it is single-use; and `sub_transform_to_selections` says it is provisional rather than implying the rest of the API's stability. Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-indexing/docs/api/index.md | 5 +- packages/zarr-indexing/docs/index.md | 10 +- .../src/zarr_indexing/__init__.py | 3 + .../src/zarr_indexing/chunk_resolution.py | 12 +- .../zarr-indexing/src/zarr_indexing/domain.py | 39 ++++- .../zarr-indexing/src/zarr_indexing/errors.py | 12 +- .../src/zarr_indexing/lazy_array.py | 147 ++++++++++++++---- .../src/zarr_indexing/output_map.py | 54 ++++++- .../src/zarr_indexing/testing/__init__.py | 2 + .../src/zarr_indexing/testing/stateful.py | 27 +++- .../src/zarr_indexing/transform.py | 12 ++ packages/zarr-indexing/tests/test_domain.py | 30 +++- .../zarr-indexing/tests/test_lazy_array.py | 60 +++---- 13 files changed, 313 insertions(+), 100 deletions(-) diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index 0eb5882f73..581f0cdb50 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -54,8 +54,9 @@ and the wire format built on top of it. **Errors** -- [`zarr_indexing.errors`](errors.md) — the canonical index-error types, which - `zarr.errors` re-exports by identity +- [`zarr_indexing.errors`](errors.md) — the index-error types this package + raises, also exported at the top level. `zarr.errors` defines classes of the + same names, which are different objects; both subclass `IndexError` **Test support** (needs the `testing` extra) diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index c7d3afc559..d840fc632c 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -227,14 +227,14 @@ part = next(iter(lazy.lazy[0:35, 0:10].parts())) part.base_coords # (0, 0) — which box of the base grid part.box # ((0, 30), (0, 40)) — that box in the array's own coordinates -part.array.shape # (30, 10) — a LazyArray for this box's cells +part.view.shape # (30, 10) — a LazyArray for this box's cells part.out_selection # (slice(0, 30, None), slice(0, 10, None)) part.is_complete # False — the view does not cover the whole box ``` -Each `part.array` is an ordinary `LazyArray`, so parts can be resolved +Each `part.view` is an ordinary `LazyArray`, so parts can be resolved independently and concurrently, and `result()` assembles that walk. -`part.array.bounding_box()` is part-local; `part.box` gives the box in the +`part.view.bounding_box()` is part-local; `part.box` gives the box in the wrapped array's coordinates, which is what distinguishes two parts. The partitioning is discovered from the wrapped array — `read_chunk_sizes`, then @@ -248,10 +248,10 @@ view = lazy.lazy[10:50, ::4] [part.base_coords for part in view.with_parts((50, 20)).parts()] # [(0, 0), (0, 1), (0, 2), (0, 3)] — uniform boxes, tail clipped -[part.base_coords for part in view.with_parts(((50, 50), (20, 20, 20, 20))).parts()] +[part.base_coords for part in view.with_parts_per_axis(((50, 50), (20, 20, 20, 20))).parts()] # [(0, 0), (0, 1), (0, 2), (0, 3)] — explicit per-axis sizes -[part.base_coords for part in view.with_parts(None).parts()] +[part.base_coords for part in view.unpartitioned().parts()] # [(0, 0)] — one whole-array part; resolve in one shot ``` diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index c1ef695d7a..41c11c3955 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -28,6 +28,7 @@ ) from zarr_indexing.composition import compose from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError from zarr_indexing.grid import ( DimensionGridLike, EdgeDimensionGrid, @@ -62,6 +63,7 @@ __all__ = [ "ArrayMap", + "BoundsCheckError", "ConstantMap", "DimensionGridLike", "DimensionMap", @@ -76,6 +78,7 @@ "OutputIndexMap", "OutputIndexMapJSON", "Partition", + "VindexInvalidSelectionError", "__version__", "array_map_dependent_axis", "compose", diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 61c4f0d36b..610ebd4db1 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -27,8 +27,16 @@ `sub_transform_to_selections` bridges from the transform representation back to the raw `(chunk_selection, out_selection, drop_axes)` tuples that -the current codec pipeline expects. This bridge will go away when the codec -pipeline accepts transforms natively. +the current codec pipeline expects. + +That bridge is **provisional**. Its return shape is a codec pipeline's +internal vocabulary rather than anything this package wants to keep saying, and +it is expected to go away once a pipeline accepts transforms natively. It is +exported because zarr's read path consumes it today; it is not covered by +whatever compatibility promise the rest of this API carries, and a consumer +outside zarr should expect to follow it through a change. The contract while it +exists is exactly `out[out_selection] = chunk[chunk_selection]`, with the axes +named in `drop_axes` squeezed out of the block first. """ from __future__ import annotations diff --git a/packages/zarr-indexing/src/zarr_indexing/domain.py b/packages/zarr-indexing/src/zarr_indexing/domain.py index f20d5bf7bd..4a9c244677 100644 --- a/packages/zarr-indexing/src/zarr_indexing/domain.py +++ b/packages/zarr-indexing/src/zarr_indexing/domain.py @@ -16,6 +16,8 @@ from dataclasses import dataclass, field from typing import Any +from zarr_indexing.errors import BoundsCheckError + @dataclass(frozen=True, slots=True) class IndexDomain: @@ -127,8 +129,22 @@ def translate(self, offset: tuple[int, ...]) -> IndexDomain: def narrow(self, selection: Any) -> IndexDomain: """Apply a basic selection and return a narrowed domain. - Indices are absolute coordinates. Integer indices produce length-1 extent. - Strided slices are not supported — use IndexTransform for strides. + + Indices are absolute coordinates, not NumPy-style offsets: `-3` names + the coordinate `-3`, and is out of bounds unless the domain contains it. + Integer indices produce a length-1 extent. Strided slices are not + supported — use `IndexTransform` for strides. + + Raises + ------ + BoundsCheckError + If a bound lies outside this domain. A slice bound used to be + clamped instead, so `narrow(slice(-3, None))` on `[0, 10)` quietly + returned the whole axis — reading as the NumPy spelling of "the last + three" and answering with something else — and `narrow(slice(20, + 30))` returned a domain its own parent did not contain. The rest of + the algebra states no clamping and no negative wrapping as an + invariant and enforces it; this is the one place that did not. """ normalized = _normalize_selection(selection, self.ndim) new_inclusive_min: list[int] = [] @@ -138,7 +154,7 @@ def narrow(self, selection: Any) -> IndexDomain: ): if isinstance(sel, int): if sel < dim_lo or sel >= dim_hi: - raise IndexError( + raise BoundsCheckError( f"index {sel} is out of bounds for dimension {dim_idx} " f"with domain [{dim_lo}, {dim_hi})" ) @@ -147,17 +163,24 @@ def narrow(self, selection: Any) -> IndexDomain: else: start, stop, step = sel.start, sel.stop, sel.step if step is not None and step != 1: - raise IndexError( + raise ValueError( "IndexDomain.narrow only supports step=1 slices. " f"Got step={step}. Use IndexTransform for strided access." ) abs_start = dim_lo if start is None else start abs_stop = dim_hi if stop is None else stop - abs_start = max(abs_start, dim_lo) - abs_stop = min(abs_stop, dim_hi) - abs_stop = max(abs_stop, abs_start) + for bound, name in ((abs_start, "start"), (abs_stop, "stop")): + if bound < dim_lo or bound > dim_hi: + raise BoundsCheckError( + f"slice {name} {bound} is out of bounds for dimension " + f"{dim_idx} with domain [{dim_lo}, {dim_hi}); indices " + f"here are absolute coordinates, so they are neither " + f"clamped to the domain nor counted from its end" + ) + # An empty interval is legal; a reversed one is the same request + # spelled backwards, and reads as empty rather than as an error. new_inclusive_min.append(abs_start) - new_exclusive_max.append(abs_stop) + new_exclusive_max.append(max(abs_stop, abs_start)) return IndexDomain( inclusive_min=tuple(new_inclusive_min), exclusive_max=tuple(new_exclusive_max), diff --git a/packages/zarr-indexing/src/zarr_indexing/errors.py b/packages/zarr-indexing/src/zarr_indexing/errors.py index fa2f6fc5d3..42f610af47 100644 --- a/packages/zarr-indexing/src/zarr_indexing/errors.py +++ b/packages/zarr-indexing/src/zarr_indexing/errors.py @@ -1,10 +1,12 @@ """Canonical index-error types raised by the transform algebra. -These are the authoritative class definitions. `zarr.errors` re-exports the -same objects (`from zarr_indexing.errors import ...`) so that, e.g., -`zarr.errors.BoundsCheckError is zarr_indexing.errors.BoundsCheckError`. -Both subclass the built-in `IndexError`, so existing `except IndexError` (or -`except zarr.errors.BoundsCheckError`) catch sites keep working unchanged. +Both subclass the built-in `IndexError`, so an `except IndexError` catch site +keeps working unchanged whichever library raised. + +`zarr.errors` defines classes of the same names, and they are *not* these +objects: `zarr.errors.BoundsCheckError is BoundsCheckError` is false. Catching +zarr's around a call into this package therefore catches nothing but their +shared `IndexError` base. Import these from here. """ from __future__ import annotations diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 50f350fbfb..aa456ddaf4 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -25,7 +25,7 @@ ```python for part in view.parts(): - out[part.out_selection] = part.array.result() + out[part.out_selection] = part.view.result() ``` Each box is therefore read once, and whatever the source could not do itself is @@ -43,8 +43,8 @@ ```python view.with_parts((64, 64)) # uniform boxes, tail clipped -view.with_parts(((3, 3, 1),)) # explicit per-axis sizes -view.with_parts(None) # one whole-array part; resolve in one shot +view.with_parts_per_axis(((3, 3, 1),)) # explicit per-axis sizes +view.unpartitioned() # one whole-array part; resolve in one shot ``` Repartitioning changes how the read is divided, not what `result()` returns. @@ -141,7 +141,7 @@ array that advertises parts therefore returns host memory, and any information carried by the wrapped array's own type is lost with it — a `numpy.ma` mask is the exception, kept by allocating a masked buffer. A wrapper with **no** -partitioning (`with_parts(None)`) skips that buffer and stays in the wrapped +partitioning (`unpartitioned()`) skips that buffer and stays in the wrapped array's own namespace. `result()` never returns memory shared with the wrapped array. That is settled @@ -157,6 +157,7 @@ import math import operator import uuid +from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol @@ -171,7 +172,7 @@ iter_chunk_transforms, sub_transform_to_selections, ) -from zarr_indexing.grid import EdgeDimensionGrid, dimension_grids_from_chunks +from zarr_indexing.grid import DimensionGridLike, EdgeDimensionGrid, dimension_grids_from_chunks from zarr_indexing.json import transform_to_canonical from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_indexing.support import ( @@ -186,7 +187,7 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Iterator, Sequence + from collections.abc import Callable, Iterator SelectFn = Callable[[Any, SelectionMode], "LazyArray"] @@ -953,8 +954,8 @@ class Partition: Yielded by [`LazyArray.parts`][zarr_indexing.lazy_array.LazyArray.parts]. The parts of a view tile it exactly and disjointly: assembling every - `array.result()` at its `out_selection` reproduces the view's `result()`, - and each part can be resolved independently and concurrently. + `view.result()` at its `out_selection` reproduces the whole view's + `result()`, and each part can be resolved independently and concurrently. Attributes ---------- @@ -964,27 +965,31 @@ class Partition: box The box itself, in the global storage coordinates of the wrapped array: one `[inclusive_min, exclusive_max)` interval per dimension. - `array.bounding_box()` is part-local and so cannot distinguish two + `view.bounding_box()` is part-local and so cannot distinguish two parts; this attribute can, and together with `is_complete` it decides a read-modify-write. - array + view A `LazyArray` covering exactly the cells of the view that live in this box. Resolving it reads the box once with basic slicing and applies the - selection to the block in memory. + selection to the block in memory. Named `view` rather than `array` + because `LazyArray.array` is the opposite thing — the raw wrapped source + — and the two sat next to each other meaning inverses. out_selection - Where `array.result()` belongs in an array of the view's shape — a + Where `view.result()` belongs in an array of the whole view's shape — a NumPy index tuple with one entry per dimension of the view, usable directly as `out[part.out_selection] = ...`. is_complete Whether the view covers the whole box. Useful to a writer deciding between a blind overwrite and a read-modify-write. True for a reversing view, which reads every cell of the box back to front; false for every - fancy-indexed axis, which is conservative rather than exact. + fancy-indexed axis, which is conservative rather than exact — so it may + be `False` for a part it does in fact cover, but never `True` for one it + does not. """ base_coords: tuple[int, ...] box: tuple[tuple[int, int], ...] - array: LazyArray + view: LazyArray out_selection: tuple[Any, ...] is_complete: bool @@ -1147,6 +1152,19 @@ def array(self) -> ArrayLike: """The wrapped array.""" return self._array + @property + def base_shape(self) -> tuple[int, ...]: + """The shape the partitioning is expressed in — not this view's shape. + + `with_parts` and `with_parts_per_axis` describe boxes of the array being + read, not of the view reading it, so a narrowed view still partitions + the extents named here. For a part's own `array`, this is the part's + box, which is why the same call means different sizes there. Without + somewhere to read it, the frame in force could only be inferred from an + error message. + """ + return self._base_shape + @property def transform(self) -> IndexTransform: """The composed transform from this view's coordinates to storage.""" @@ -1390,22 +1408,26 @@ def strides(self) -> tuple[int, ...] | None: # -- partitioning ------------------------------------------------------- - def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> LazyArray: - """Return the same view, read in different parts. + def with_parts(self, parts: Sequence[int]) -> LazyArray: + """Return the same view, read in uniform boxes of shape `parts`. - The transform, the wrapped array, and therefore `result()` are all - unchanged; only the boxes the read is broken into differ. Nothing is - copied and nothing is read. + One integer per dimension of `base_shape`, with the trailing box in each + dimension clipped to the extent. The transform, the wrapped array, and + therefore `result()` are all unchanged; only the boxes the read is + broken into differ. Nothing is copied and nothing is read. + + For per-axis sizes see + [`with_parts_per_axis`][zarr_indexing.lazy_array.LazyArray.with_parts_per_axis], + and to read in one pass see + [`unpartitioned`][zarr_indexing.lazy_array.LazyArray.unpartitioned]. + The three were one parameter whose meaning was decided by inspecting the + type of what it was given, which left no way to ask for one of them and + be told when you had spelled it wrong. Parameters ---------- parts - Either a uniform box shape (one integer per dimension of the base - array, with the trailing box clipped to the extent), dask-convention - per-axis sizes (one sequence of box extents per dimension, each - summing to the base extent), or `None` for a single part covering - the whole array, which makes `result()` lower the view in one pass - instead of assembling it block by block. + The box shape, one integer per dimension of `base_shape`. Returns ------- @@ -1415,10 +1437,9 @@ def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> L Raises ------ ValueError - If `parts` has the wrong length, mixes the two conventions, contains - a non-positive extent, or declares per-axis sizes that do not sum to - the base shape. Unlike partition discovery, this is a public API and - validates strictly. + If `parts` has the wrong length or contains a non-positive extent. + Unlike partition discovery, this is a public API and validates + strictly. Examples -------- @@ -1427,12 +1448,70 @@ def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> L >>> [part.base_coords for part in view.with_parts((2, 3)).parts()] [(0, 0), (0, 1), (1, 0), (1, 1)] """ - grids = None if parts is None else dimension_grids_from_chunks(parts, self._base_shape) + if any(isinstance(entry, Sequence) for entry in parts): + raise ValueError( + "with_parts takes one integer per dimension; for per-axis box " + "sizes use with_parts_per_axis" + ) + return self._with_grids(dimension_grids_from_chunks(parts, self._base_shape)) + + def with_parts_per_axis(self, sizes: Sequence[Sequence[int]]) -> LazyArray: + """Return the same view, read in boxes of explicitly listed sizes. + + The dask convention: one sequence of box extents per dimension of + `base_shape`, each summing to that dimension's extent. Use it when the + boxes are not uniform — a partitioning discovered from a store, or one + whose last box differs by more than clipping. + + Parameters + ---------- + sizes + One sequence of box extents per dimension of `base_shape`. + + Returns + ------- + LazyArray + The same view with a new partitioning. + + Raises + ------ + ValueError + If `sizes` has the wrong length, contains a non-positive extent, or + declares sizes that do not sum to `base_shape`. + + Examples + -------- + >>> import numpy as np + >>> view = LazyArray(np.arange(12).reshape(3, 4)) + >>> [part.box for part in view.with_parts_per_axis(((1, 2), (4,))).parts()] + [((0, 1), (0, 4)), ((1, 3), (0, 4))] + """ + return self._with_grids(dimension_grids_from_chunks(sizes, self._base_shape)) + + def unpartitioned(self) -> LazyArray: + """Return the same view, read in one pass. + + `result()` then lowers the whole view to array operations directly + rather than allocating an output buffer and assembling it box by box. + `parts()` still yields a single part covering everything. + + Returns + ------- + LazyArray + The same view with no partitioning. + """ + return self._with_grids(None) + + def _with_grids(self, grids: tuple[DimensionGridLike, ...] | None) -> LazyArray: return LazyArray._derive(self._array, self._transform, grids, self._window, self._support) def parts(self) -> Iterator[Partition]: """Iterate the base partitioning, projected through this view. + Single-use: this is a generator, so it is consumed by the first walk and + a second `for` over the same object yields nothing. Call `parts()` again + for a fresh walk, or keep a `list` of it if you need to revisit. + Yields one [`Partition`][zarr_indexing.lazy_array.Partition] per box the view actually touches. The parts tile the view exactly and disjointly, and each carries a `LazyArray` that can be resolved on its own: in @@ -1451,7 +1530,7 @@ def parts(self) -> Iterator[Partition]: >>> import numpy as np >>> view = LazyArray(np.arange(12).reshape(3, 4)).with_parts((2, 2)) >>> part = next(view.lazy[:, 1:].parts()) - >>> (part.base_coords, part.array.shape, part.is_complete) + >>> (part.base_coords, part.view.shape, part.is_complete) ((0, 0), (2, 1), False) """ base_shape = self._base_shape @@ -1485,7 +1564,7 @@ def parts(self) -> Iterator[Partition]: yield Partition( base_coords=base_coords, box=tuple((o, o + e) for o, e in zip(global_origin, extent, strict=True)), - array=LazyArray._derive(self._array, local, None, window, self._support), + view=LazyArray._derive(self._array, local, None, window, self._support), out_selection=_partition_out_selection(local, out_indices, out_shape), is_complete=_covers_whole_part(local, extent), ) @@ -1526,7 +1605,7 @@ def result(self) -> Any: """Materialize this view. Assembles the view from its `parts()`, reading each box once. A wrapper - with **no partitioning** — `with_parts(None)`, and the default for a + with **no partitioning** — `unpartitioned()`, and the default for a source that advertises none — skips the output buffer and lowers the whole view directly to array operations. A partitioning with a single whole-array box still goes through the buffer: the branch is on whether @@ -1569,7 +1648,7 @@ def result(self) -> Any: # number of axes is named rather than reported as a broadcast # failure against the buffer. written += _out_selection_cell_count(part.out_selection, out_shape) - out[part.out_selection] = np.asanyarray(part.array.result()) + out[part.out_selection] = np.asanyarray(part.view.result()) if written != size: # The buffer is uninitialized where no part wrote, so a partition # walk that does not tile the view exactly would otherwise hand diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py index ca828a8ff9..871d778620 100644 --- a/packages/zarr-indexing/src/zarr_indexing/output_map.py +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -37,8 +37,9 @@ from dataclasses import dataclass from typing import TYPE_CHECKING +import numpy as np + if TYPE_CHECKING: - import numpy as np import numpy.typing as npt @@ -101,5 +102,56 @@ class ArrayMap: stride: int = 1 input_dimension: int | None = None + def __post_init__(self) -> None: + """Hold the index array through a read-only view. + + A map is frozen, but the array inside it was not: reaching through a + view's transform to `index_array[0] = 9` silently changed what the view + returned, in a package whose whole contract is that a view is a + description of a read and resolving it twice answers alike. The view is + taken rather than the caller's array being marked, so constructing a map + does not take away the ability to write to an array you still own. + """ + # `asarray` first: a NumPy scalar reaches here from paths that index an + # array down to one element, and a scalar has no flags to set. + frozen = np.asarray(self.index_array).view() + frozen.flags.writeable = False + object.__setattr__(self, "index_array", frozen) + + def __eq__(self, other: object) -> bool: + """Value equality, comparing index arrays element-wise. + + The generated `__eq__` compares them with `==`, whose result for two + arrays is an array — so asking whether two maps are equal raised + `ValueError: the truth value of an array ... is ambiguous`. `frozen=True` + reads as a promise that a value can be compared and hashed, and this is + what makes good on it. + """ + if not isinstance(other, ArrayMap): + return NotImplemented + return ( + self.offset == other.offset + and self.stride == other.stride + and self.input_dimension == other.input_dimension + and self.index_array.shape == other.index_array.shape + and bool(np.array_equal(self.index_array, other.index_array)) + ) + + def __hash__(self) -> int: + """Hashed by the array's contents, so equal maps hash alike. + + The generated `__hash__` hashed the ndarray itself, which is unhashable; + a map could therefore not go in a set, or key a cache. + """ + return hash( + ( + self.offset, + self.stride, + self.input_dimension, + self.index_array.shape, + self.index_array.tobytes(), + ) + ) + OutputIndexMap = ConstantMap | DimensionMap | ArrayMap diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py index babbe31623..e98d051900 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py @@ -29,6 +29,7 @@ def make_source(self, data): ChainedIndexingStateMachine, apply_selection, outer_selection, + repartition, state_machine_test, ) from zarr_indexing.testing.strategies import ( @@ -49,6 +50,7 @@ def make_source(self, data): "masks", "orthogonal_selections", "outer_selection", + "repartition", "slice_selections", "state_machine_test", "vectorized_selections", diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 901480149a..88d22e5383 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -30,7 +30,7 @@ def make_source(self, data): ------------------- The parts invariant is the one with teeth. [`Partition`][zarr_indexing.lazy_array.Partition] documents -`out[part.out_selection] = part.array.result()` as the assembly procedure, so +`out[part.out_selection] = part.view.result()` as the assembly procedure, so this checks that literally: a part's values must arrive at exactly the shape its `out_selection` addresses — not merely a shape that broadcasts into it — land there, and cover the view once. Checking through `result()` alone would prove @@ -53,6 +53,7 @@ def make_source(self, data): from __future__ import annotations import math +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, ClassVar import numpy as np @@ -70,8 +71,6 @@ def make_source(self, data): ) if TYPE_CHECKING: - from collections.abc import Sequence - from zarr_indexing.boundary import SelectionMode __all__ = [ @@ -81,6 +80,7 @@ def make_source(self, data): "ChainedIndexingStateMachine", "apply_selection", "outer_selection", + "repartition", "state_machine_test", ] @@ -171,6 +171,21 @@ def state_machine_test( return case +def repartition(view: LazyArray, parts: Any) -> LazyArray: + """Apply one of `partitionings` to a view. + + The three partitioning spellings are three named methods, so a list holding + a mix of them needs a dispatch somewhere. Choosing among them is what a test + harness drawing from that list is doing, so it lives here rather than being + pushed back into the public API as a type-inspecting parameter. + """ + if parts is None: + return view.unpartitioned() + if any(isinstance(entry, Sequence) for entry in parts): + return view.with_parts_per_axis(parts) + return view.with_parts(parts) + + # --------------------------------------------------------------------------- # # The machine # --------------------------------------------------------------------------- # @@ -261,7 +276,7 @@ def _step(self, mode: SelectionMode, selection: tuple[Any, ...], *, fancy: bool) def choose_partitioning(self, data: st.DataObject) -> None: """Fix how the read is broken up, before any indexing.""" parts = data.draw(st.sampled_from(list(type(self).partitionings))) - self.view = self.view.with_parts(parts) + self.view = repartition(self.view, parts) self.chain.append(("parts", parts)) @precondition(lambda self: self._indexable()) @@ -308,7 +323,7 @@ def repartition(self, data: st.DataObject) -> None: collapsed correlated selection reaches. """ parts = data.draw(st.sampled_from(list(type(self).partitionings))) - self.view = self.view.with_parts(parts) + self.view = repartition(self.view, parts) self.chain.append(("parts", parts)) # -- invariants --------------------------------------------------------- @@ -334,7 +349,7 @@ def parts_tile_the_view(self) -> None: assembled = np.zeros(self.view.shape, dtype=self.view.dtype) hits = np.zeros(self.view.shape, dtype=np.int64) for part in self.view.parts(): - value = np.asarray(part.array.result()) + value = np.asarray(part.view.result()) assert value.shape == assembled[part.out_selection].shape, ( f"part {part.base_coords} carries {value.shape} for an out_selection " f"addressing {assembled[part.out_selection].shape}: {self.chain}" diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index e8f3f9d074..1975ee687e 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -125,6 +125,18 @@ def __post_init__(self) -> None: f"axis it names, and a correlated one names none" ) + def __eq__(self, other: object) -> bool: + """Value equality. `ArrayMap` compares its index array element-wise, so + a transform holding one can be compared at all — the generated `__eq__` + raised `ValueError: the truth value of an array ... is ambiguous`.""" + if not isinstance(other, IndexTransform): + return NotImplemented + return self.domain == other.domain and self.output == other.output + + def __hash__(self) -> int: + """Hashed by value, so a transform can key a cache or enter a set.""" + return hash((self.domain, self.output)) + @property def input_rank(self) -> int: return self.domain.ndim diff --git a/packages/zarr-indexing/tests/test_domain.py b/packages/zarr-indexing/tests/test_domain.py index 9664a0b08a..0278ecd714 100644 --- a/packages/zarr-indexing/tests/test_domain.py +++ b/packages/zarr-indexing/tests/test_domain.py @@ -3,6 +3,7 @@ import pytest from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError class TestIndexDomainConstruction: @@ -171,19 +172,32 @@ def test_narrow_non_zero_origin(self) -> None: def test_narrow_int_out_of_bounds(self) -> None: d = IndexDomain.from_shape((10,)) - with pytest.raises(IndexError, match="out of bounds"): + with pytest.raises(BoundsCheckError, match="out of bounds"): d.narrow((10,)) def test_narrow_int_below_origin(self) -> None: d = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) - with pytest.raises(IndexError, match="out of bounds"): + with pytest.raises(BoundsCheckError, match="out of bounds"): d.narrow((4,)) - def test_narrow_clamps_to_domain(self) -> None: + def test_narrow_refuses_a_bound_outside_the_domain(self) -> None: + """Clamping made two different requests answer alike, and neither well. + + Indices here are absolute coordinates, so `-5` is a coordinate this + domain does not contain rather than NumPy's "five from the end" — and + clamping returned the whole axis for it, which is what a caller writing + the NumPy spelling would least expect. A stop past the end produced a + domain the parent did not contain. + """ d = IndexDomain.from_shape((10,)) - result = d.narrow((slice(-5, 100),)) - assert result.inclusive_min == (0,) - assert result.exclusive_max == (10,) + with pytest.raises(BoundsCheckError, match="absolute coordinates"): + d.narrow((slice(-5, 100),)) + with pytest.raises(BoundsCheckError, match="out of bounds"): + d.narrow((slice(20, 30),)) + + def test_narrow_accepts_the_bounds_of_the_domain_itself(self) -> None: + d = IndexDomain.from_shape((10,)) + assert d.narrow((slice(0, 10),)) == d def test_narrow_bare_slice(self) -> None: d = IndexDomain.from_shape((10,)) @@ -197,6 +211,8 @@ def test_narrow_too_many_indices(self) -> None: d.narrow((1, 2)) def test_narrow_step_not_one(self) -> None: + """A stride is not a bounds failure — the rest of the algebra raises + `ValueError` for a request it does not implement, and so does this.""" d = IndexDomain.from_shape((10,)) - with pytest.raises(IndexError, match="step=1"): + with pytest.raises(ValueError, match="step=1"): d.narrow((slice(0, 10, 2),)) diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index f0ee8e2c06..a51ad4df11 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -30,6 +30,7 @@ ) from zarr_indexing.lazy_array import _out_selection_cell_count from zarr_indexing.support import IndexingSupport +from zarr_indexing.testing import repartition if TYPE_CHECKING: from collections.abc import Callable, Sequence @@ -205,7 +206,7 @@ def make_source(flavor: str) -> LazyArray: if flavor == "numpy-whole": return LazyArray(data) if flavor == "numpy-explicit-parts": - return LazyArray(data).with_parts(EXPLICIT_PARTS) + return LazyArray(data).with_parts_per_axis(EXPLICIT_PARTS) if flavor == "numpy-uniform-parts": return LazyArray(data).with_parts(PART_SHAPE) if flavor == "zarr": @@ -718,7 +719,7 @@ def test_random_chains_match_numpy(flavor: str) -> None: ) for parts in partitionings: np.testing.assert_array_equal( - np.asarray(view.with_parts(parts).result()), + np.asarray(repartition(view, parts).result()), np.asarray(expected), err_msg=f"{flavor} parts={parts}: {chain}", ) @@ -749,7 +750,7 @@ def test_a_slice_only_fancy_step_after_a_fancy_step_reads_real_data() -> None: for parts in (None, (2, 4), (1, 8), (3, 3)): view = ( - LazyArray(base).with_parts(parts).lazy.oindex[np.array([0, 2]), :].lazy.oindex[:, 2:8] + repartition(LazyArray(base), parts).lazy.oindex[np.array([0, 2]), :].lazy.oindex[:, 2:8] ) assert view.shape == expected.shape, f"parts={parts}" np.testing.assert_array_equal(np.asarray(view.result()), expected, err_msg=f"{parts}") @@ -832,7 +833,7 @@ def test_a_domain_axis_no_output_map_depends_on_still_resolves( """The value, the shape and the tiling all hold when an axis is unreferenced.""" data = unreferenced_axis_reference() expected = np.asarray(oracle(data)) - view = build(LazyArray(data).with_parts(parts)) + view = build(repartition(LazyArray(data), parts)) assert view.shape == expected.shape np.testing.assert_array_equal(np.asarray(view.result()), expected) @@ -840,7 +841,7 @@ def test_a_domain_axis_no_output_map_depends_on_still_resolves( hits = np.zeros(view.shape, dtype=np.int64) assembled = np.zeros(view.shape, dtype=view.dtype) for part in view.parts(): - assembled[part.out_selection] = np.asarray(part.array.result()) + assembled[part.out_selection] = np.asarray(part.view.result()) np.add.at(hits, part.out_selection, 1) np.testing.assert_array_equal(assembled, expected) np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64)) @@ -856,7 +857,7 @@ def test_an_unreferenced_domain_axis_of_extent_zero_stays_empty() -> None: data = np.arange(140, dtype=np.int64).reshape(7, 5, 4) coords = np.array([[6], [3], [0]]) for parts in (None, (1, 1, 1), (3, 2, 3), (7, 5, 4)): - view = LazyArray(data).with_parts(parts).lazy.vindex[coords, -4].lazy[0, 0:0] + view = repartition(LazyArray(data), parts).lazy.vindex[coords, -4].lazy[0, 0:0] expected = data[coords, -4][0, 0:0] assert view.shape == expected.shape == (0, 4), f"parts={parts}" np.testing.assert_array_equal(np.asarray(view.result()), expected, err_msg=f"{parts}") @@ -870,8 +871,7 @@ def test_a_zero_length_axis_resolves_the_same_way_under_every_partitioning() -> for parts in (None, ((1, 1, 1), ()), ((3,), ())): view = ( - LazyArray(base) - .with_parts(parts) + repartition(LazyArray(base), parts) .lazy.oindex[np.array([1, -3, -3]), :] .lazy.oindex[:, :] ) @@ -891,7 +891,7 @@ def test_an_empty_slice_of_a_length_one_correlated_axis_has_no_parts() -> None: mask = np.array([False, True, False, False, False]) for parts in PARTITIONINGS_1D: - view = LazyArray(base).with_parts(parts).lazy.vindex[mask].lazy[1:-2] + view = repartition(LazyArray(base), parts).lazy.vindex[mask].lazy[1:-2] assert view.shape == (0,), f"parts={parts}" assert list(view.parts()) == [], f"parts={parts}" np.testing.assert_array_equal(np.asarray(view.result()), base[mask][1:-2]) @@ -902,7 +902,9 @@ def test_an_empty_slice_of_a_length_one_vindex_pair_has_no_parts() -> None: base = np.arange(35).reshape(7, 5) for parts in (None, (2, 2), (7, 5)): - view = LazyArray(base).with_parts(parts).lazy.vindex[np.array([6]), np.array([0])].lazy[0:0] + view = ( + repartition(LazyArray(base), parts).lazy.vindex[np.array([6]), np.array([0])].lazy[0:0] + ) assert view.shape == (0,), f"parts={parts}" assert list(view.parts()) == [], f"parts={parts}" np.testing.assert_array_equal( @@ -913,7 +915,7 @@ def test_an_empty_slice_of_a_length_one_vindex_pair_has_no_parts() -> None: def test_a_correlated_view_narrowed_to_one_point_has_parts_of_the_views_rank() -> None: """A rank-0 broadcast block survives intersection without gaining an axis. - `Partition` documents `out[part.out_selection] = part.array.result()` as the + `Partition` documents `out[part.out_selection] = part.view.result()` as the assembly, so a part's values must arrive at the rank the view has — including the rank 0 a correlated selection reaches once every coordinate is a scalar. """ @@ -928,13 +930,13 @@ def test_a_correlated_view_narrowed_to_one_point_has_parts_of_the_views_rank() - for parts in (None, (2, 2, 2), (3, 3, 3), (7, 5, 4)): for selection, tail in cases: expected = base[selection][tail] - view = LazyArray(base).with_parts(parts).lazy.vindex[selection].lazy[tail] + view = repartition(LazyArray(base), parts).lazy.vindex[selection].lazy[tail] assert view.shape == expected.shape, f"parts={parts}, {selection}" assembled = np.zeros(view.shape, dtype=view.dtype) hits = np.zeros(view.shape, dtype=np.int64) for part in view.parts(): - assembled[part.out_selection] = np.asarray(part.array.result()) + assembled[part.out_selection] = np.asarray(part.view.result()) np.add.at(hits, part.out_selection, 1) err = f"parts={parts}, {selection}" @@ -1195,7 +1197,7 @@ def test_parts_tile_the_view_exactly_and_disjointly( assembled = np.zeros(view.shape, dtype=view.dtype) hits = np.zeros(view.shape, dtype=np.int64) for part in view.parts(): - assembled[part.out_selection] = np.asarray(part.array.result()) + assembled[part.out_selection] = np.asarray(part.view.result()) # `np.add.at` accumulates per element; `+= 1` on a fancy index would # count a repeated position once and hide a real overlap. np.add.at(hits, part.out_selection, 1) @@ -1211,7 +1213,7 @@ def test_parts_resolve_independently_and_concurrently() -> None: assert len(parts) > 1 with ThreadPoolExecutor(max_workers=4) as pool: - values = list(pool.map(lambda part: np.asarray(part.array.result()), parts)) + values = list(pool.map(lambda part: np.asarray(part.view.result()), parts)) assembled = np.zeros(view.shape, dtype=view.dtype) for part, value in zip(parts, values, strict=True): @@ -1243,7 +1245,7 @@ def test_partition_boxes_are_global_and_tile_the_base() -> None: # The reviewer's repro: two parts of the same view whose part-local # bounding boxes coincide but which cover different regions of the base. first, second = parts[0, 0, 0], parts[0, 1, 0] - assert first.array.bounding_box() == second.array.bounding_box() + assert first.view.bounding_box() == second.view.bounding_box() assert first.box != second.box assert first.box == ((0, 3), (0, 2), (0, 3)) assert second.box == ((0, 3), (2, 4), (0, 3)) @@ -1277,9 +1279,9 @@ def test_an_unpartitioned_wrapper_has_a_single_whole_array_part() -> None: parts = list(view.parts()) assert len(parts) == 1 assert parts[0].base_coords == (0, 0, 0) - assert parts[0].array.shape == SHAPE + assert parts[0].view.shape == SHAPE assert parts[0].is_complete - np.testing.assert_array_equal(np.asarray(parts[0].array.result()), reference()) + np.testing.assert_array_equal(np.asarray(parts[0].view.result()), reference()) def test_with_parts_keeps_the_view_and_the_base() -> None: @@ -1297,7 +1299,7 @@ def test_with_parts_keeps_the_view_and_the_base() -> None: def test_with_parts_none_forces_one_shot_resolution() -> None: view = make_source("zarr").lazy.oindex[[4, 0, 0], :, :] - whole = view.with_parts(None) + whole = view.unpartitioned() assert len(list(whole.parts())) == 1 np.testing.assert_array_equal(np.asarray(whole.result()), np.asarray(view.result())) @@ -1320,7 +1322,7 @@ def test_with_parts_none_forces_one_shot_resolution() -> None: def test_with_parts_validates_strictly(parts: Any, match: str) -> None: """`with_parts` is our own API, so a malformed partitioning raises.""" with pytest.raises(ValueError, match=match): - make_source("numpy-whole").with_parts(parts) + repartition(make_source("numpy-whole"), parts) # --------------------------------------------------------------------------- @@ -1707,7 +1709,7 @@ def test_a_source_is_never_asked_for_more_than_it_declares( expected = np.asarray(oracle(reference())) source = STRICT_SOURCES[level](reference()) for parts in (None, PART_SHAPE): - view = build(LazyArray(source).with_parts(parts)) + view = build(repartition(LazyArray(source), parts)) np.testing.assert_array_equal( np.asarray(view.result()), expected, err_msg=f"{level} parts={parts}" ) @@ -1725,7 +1727,7 @@ def test_random_chains_agree_across_levels() -> None: for level in LEVELS: for parts in (None, (2, 2, 2)): - view = LazyArray(data).with_indexing_support(level).with_parts(parts) + view = repartition(LazyArray(data).with_indexing_support(level), parts) for mode, selection in chain: view = _apply_view(view, mode, selection) np.testing.assert_array_equal( @@ -1803,9 +1805,7 @@ def test_with_indexing_support_keeps_the_view_and_the_parts() -> None: ] np.testing.assert_array_equal(np.asarray(renegotiated.result()), np.asarray(view.result())) # Parts inherit the level, so a part never exceeds it either. - assert all( - part.array.indexing_support is IndexingSupport.OUTER for part in renegotiated.parts() - ) + assert all(part.view.indexing_support is IndexingSupport.OUTER for part in renegotiated.parts()) @pytest.mark.parametrize("bogus", ["vectorized", 3, None]) @@ -1888,7 +1888,7 @@ def test_materializing_never_hands_back_the_wrapped_array( result, so detaching has to decide from the result's own buffer instead. """ data = reference() - view = build(LazyArray(wrap(data)).with_parts(parts)) + view = build(repartition(LazyArray(wrap(data)), parts)) for materialize in ( lambda v: v.result(), @@ -1997,7 +1997,7 @@ def test_a_source_meeting_only_the_basic_floor_resolves_any_selection( ) -> None: """The floor is a promise about the source; its blocks are coerced, not trusted.""" expected = np.asarray(oracle(reference())) - view = build(LazyArray(DuckBlock(reference())).with_parts(parts)) + view = build(repartition(LazyArray(DuckBlock(reference())), parts)) assert view.shape == expected.shape np.testing.assert_array_equal(np.asarray(view.result()), expected) @@ -2023,7 +2023,7 @@ def test_numpy_matrix_is_refused() -> None: @pytest.mark.parametrize("parts", [None, (2, 2), (1, 4), (3, 4)]) def test_a_masked_source_keeps_its_mask_under_every_partitioning(parts: Any) -> None: data = np.ma.masked_greater(np.arange(12).reshape(3, 4), 7) - got = LazyArray(data).with_parts(parts).lazy[:, 1:].result() + got = repartition(LazyArray(data), parts).lazy[:, 1:].result() expected = data[:, 1:] assert isinstance(got, np.ma.MaskedArray), parts np.testing.assert_array_equal(np.ma.getmaskarray(got), np.ma.getmaskarray(expected)) @@ -2108,14 +2108,14 @@ def test_a_strided_reversal_is_still_incomplete() -> None: def test_a_zero_length_axis_accepts_every_spelling_of_no_chunks(parts: Any) -> None: """`(0,)`, `(0, 0)`, `()` and a uniform shape all describe an axis with no cells.""" data = np.zeros((0, 3)) - view = LazyArray(data).with_parts(parts) + view = repartition(LazyArray(data), parts) assert view.result().shape == (0, 3) assert list(view.parts()) == [] def test_a_zero_chunk_on_a_nonempty_axis_is_still_rejected() -> None: with pytest.raises(ValueError, match="chunk sizes must be positive"): - LazyArray(np.zeros((4, 3))).with_parts(((0, 4), (3,))) + LazyArray(np.zeros((4, 3))).with_parts_per_axis(((0, 4), (3,))) @pytest.mark.parametrize( From 89edc6c4c392d8a0076d39e8e26a8d768e3577f8 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 07:48:02 +0200 Subject: [PATCH 24/32] docs(zarr-indexing): correct the claims a reviewer could check, and two dialect gaps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `boundary.py` justified applying scalars before the advanced indices by asserting NumPy does the same, citing `a[0, [1, 2], :]`. NumPy groups a scalar *with* the advanced indices for placement, so the two disagree the moment a slice separates them: `a[0, ..., [1, 2]]` is `(2, 3)` where `a[0][..., [1, 2]]` is `(3, 2)`. Scalar-first is this package's documented dialect and stays; the reasoning was wrong, and a wrong reason invites someone to "fix" the correct end later. Two places where the dialect really did diverge, both now matching NumPy. A zero-dimensional integer array is a scalar — `a[np.array(2), :]` drops its axis — but only Python and NumPy integers counted, so a 0-d array was widened into a length-1 index array and kept an axis; that was a third answer, agreeing with neither NumPy nor eager zarr. And a multi-dimensional array in an orthogonal selection is refused where the rule lives, instead of surfacing two layers down as a rank complaint about an `index_array` the caller never wrote. `iter_chunk_transforms` documented two shapes for `out_indices` and returns three: the `dict[int, ndarray]` an orthogonal selection with several index arrays produces was missing, from the function downstream integrators use most. Packaging: the README is the PyPI long description, so the monorepo-only development section moved to CONTRIBUTING.md and the dead relative link to the justfile went with it. The examples pinned `zarr-indexing` to a moving `main` rather than to a release they document. Three of the six docs pins claimed to match the repo root and did not, and the justfile's ruff comment contradicted the package's own pin. Assisted-by: ClaudeCode:claude-fable-5 --- .../lazy_indexing_dask/lazy_indexing_dask.py | 4 +- .../lazy_indexing_numpy.py | 4 +- packages/zarr-indexing/CONTRIBUTING.md | 27 +++++++++ packages/zarr-indexing/README.md | 28 ++------- packages/zarr-indexing/changes/267.feature.md | 4 +- packages/zarr-indexing/changes/272.bugfix.md | 9 +++ packages/zarr-indexing/changes/273.misc.md | 6 ++ packages/zarr-indexing/docs/index.md | 8 ++- packages/zarr-indexing/pyproject.toml | 6 +- .../src/zarr_indexing/__init__.py | 6 +- .../src/zarr_indexing/boundary.py | 38 +++++++++--- .../src/zarr_indexing/chunk_resolution.py | 8 ++- .../zarr-indexing/src/zarr_indexing/json.py | 12 ++-- .../src/zarr_indexing/lazy_array.py | 4 +- .../src/zarr_indexing/transform.py | 17 +++++- .../zarr-indexing/tests/test_lazy_array.py | 33 ++++++++++ packages/zarr-indexing/uv.lock | 60 +++++++++---------- 17 files changed, 190 insertions(+), 84 deletions(-) create mode 100644 packages/zarr-indexing/CONTRIBUTING.md create mode 100644 packages/zarr-indexing/changes/272.bugfix.md create mode 100644 packages/zarr-indexing/changes/273.misc.md diff --git a/examples/lazy_indexing_dask/lazy_indexing_dask.py b/examples/lazy_indexing_dask/lazy_indexing_dask.py index 5f8d93d6c2..5658a45130 100644 --- a/examples/lazy_indexing_dask/lazy_indexing_dask.py +++ b/examples/lazy_indexing_dask/lazy_indexing_dask.py @@ -2,7 +2,7 @@ # requires-python = ">=3.12" # dependencies = [ # "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main", -# "zarr-indexing @ git+https://github.com/zarr-developers/zarr-python.git@main#subdirectory=packages/zarr-indexing", +# "zarr-indexing>=0.1", # "dask[array]==2025.3.0", # "numpy==2.4.3", # "pytest==9.0.2" @@ -65,7 +65,7 @@ def test_parts_as_tasks(source: zarr.Array) -> None: # the reads are independent and the placement needs no coordination. @dask.delayed def read(part: object) -> np.ndarray: - return part.array.result() + return part.view.result() blocks = dask.compute(*[read(part) for part in parts], scheduler="threads") diff --git a/examples/lazy_indexing_numpy/lazy_indexing_numpy.py b/examples/lazy_indexing_numpy/lazy_indexing_numpy.py index 3ae02e9182..d7e1a91adf 100644 --- a/examples/lazy_indexing_numpy/lazy_indexing_numpy.py +++ b/examples/lazy_indexing_numpy/lazy_indexing_numpy.py @@ -1,7 +1,7 @@ # /// script # requires-python = ">=3.12" # dependencies = [ -# "zarr-indexing @ git+https://github.com/zarr-developers/zarr-python.git@main#subdirectory=packages/zarr-indexing", +# "zarr-indexing>=0.1", # "numpy==2.4.3", # "pytest==9.0.2" # ] @@ -103,7 +103,7 @@ def test_parts() -> None: # parts and placing them is what `result()` does. assembled = np.empty(view.shape, dtype=view.dtype) for part in parts: - assembled[part.out_selection] = part.array.result() + assembled[part.out_selection] = part.view.result() assert np.array_equal(assembled, view.result()) # The partitioning is a read strategy, so a different one gives the same data. diff --git a/packages/zarr-indexing/CONTRIBUTING.md b/packages/zarr-indexing/CONTRIBUTING.md new file mode 100644 index 0000000000..632e24929f --- /dev/null +++ b/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/`. + +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. diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index 2e7249c919..c3ab8d085a 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -31,27 +31,9 @@ repository and consumed by `zarr` to resolve array indexing operations. pip install zarr-indexing ``` -## Developing +## Contributing -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/`. - -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 +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). diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 3e19c22a42..fcf9bcf48d 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -2,7 +2,7 @@ Added `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, C 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.array.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. +`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. @@ -28,7 +28,7 @@ Fixed a domain dimension no output map depends on, which a `vindex` coordinate a `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 `array.result()` came back with rank 1 where the view had rank 0 and the documented `out[part.out_selection] = part.array.result()` assembly raised `ValueError`. A rank-0 block now stays rank 0: it survives whole or the intersection is empty. +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. diff --git a/packages/zarr-indexing/changes/272.bugfix.md b/packages/zarr-indexing/changes/272.bugfix.md new file mode 100644 index 0000000000..8ef3ac9ae4 --- /dev/null +++ b/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 +could not reload its own output for a selection that selects nothing, and its +domain loader validated nothing; `sub_transform_to_selections` 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. diff --git a/packages/zarr-indexing/changes/273.misc.md b/packages/zarr-indexing/changes/273.misc.md new file mode 100644 index 0000000000..d5210f4c8d --- /dev/null +++ b/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. diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index d840fc632c..9bb0ab1b74 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -80,7 +80,7 @@ from zarr_indexing import IndexTransform transform = IndexTransform.from_shape((100, 100)) view = transform[10:50, 5] -view.domain # IndexDomain(inclusive_min=(10,), exclusive_max=(50,)) +view.domain # IndexDomain(inclusive_min=(10,), exclusive_max=(50,), labels=None) view.selection_repr # '{ [10, 50), 5 }' ``` @@ -239,8 +239,10 @@ wrapped array's coordinates, which is what distinguishes two parts. The partitioning is discovered from the wrapped array — `read_chunk_sizes`, then `chunks`. Those attribute names belong to the source; this API refers only to -parts. `with_parts` selects a different partitioning without touching the data or -the view: +parts. `with_parts`, `with_parts_per_axis` and `unpartitioned` select a +different partitioning without touching the data or the view. The sizes are +expressed in `base_shape`, the shape of the array being read, not of the view +reading it: ```python view = lazy.lazy[10:50, ::4] diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index 5bb615fdd8..5e99141d07 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -60,13 +60,13 @@ test = ["pytest", "hypothesis>=6.160.0"] docs = [ # Pins match the zarr-python docs environment in the repo-root # pyproject.toml so the two sites render with the same toolchain. - "mkdocs-material==9.7.6", + "mkdocs-material==9.7.7", "mkdocs==1.6.1", - "mkdocstrings==1.0.4", + "mkdocstrings==1.0.6", "mkdocstrings-python==2.0.5", "griffe-inherited-docstrings==1.1.3", # mkdocstrings uses ruff to format rendered signatures - "ruff==0.15.20", + "ruff==0.15.22", ] [tool.hatch.version] diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 41c11c3955..1ef089d4b2 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -13,10 +13,14 @@ output dimension can depend on the input (see `output_map.py`) - `compose` — chain two transforms into one +`LazyArray` wraps an array-API-like array and gives it deferred indexing +through `.lazy[...]`, yielding its reads as `Partition`s. + The chunk-resolution helpers (`iter_chunk_transforms`, `sub_transform_to_selections`) and `selection_to_transform` are also exported here: they form the surface the zarr integration layer (array indexing) depends -on. The `*Like` grid Protocols describe the chunk-grid surface chunk resolution +on — see `chunk_resolution` on why the second of those is provisional. The +`DimensionGridLike` Protocol describes the chunk-grid surface chunk resolution consumes without importing zarr. """ diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py index 16f9506d62..f00cf54270 100644 --- a/packages/zarr-indexing/src/zarr_indexing/boundary.py +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -61,6 +61,22 @@ def _as_index_array(sel: Any) -> np.ndarray[Any, np.dtype[Any]] | None: return None +def _is_scalar_index(sel: Any) -> bool: + """Whether an entry is a scalar integer index. + + A zero-dimensional integer array counts. NumPy treats `a[np.array(2), :]` + exactly as `a[2, :]` — the axis is dropped — but this took only Python and + NumPy integers, so a 0-d array fell through to the fancy path and was widened + into a length-1 index array, keeping an axis NumPy drops. That was a third + answer, agreeing with neither NumPy nor zarr. + """ + if _is_bool_scalar(sel): + return False + if isinstance(sel, (int, np.integer)): + return True + return isinstance(sel, np.ndarray) and sel.ndim == 0 and sel.dtype.kind in "iu" + + def _axes_consumed(sel: Any, mode: SelectionMode) -> int: """How many axes of the view a single selection entry consumes.""" if sel is None: @@ -163,12 +179,20 @@ def split_scalar_axes( ) -> tuple[tuple[Any, ...] | None, Any]: """Peel scalar integer indices out of a fancy selection. - In NumPy, a scalar integer is a basic index wherever it appears: it drops - its axis, and is applied before the advanced indices rather than broadcast - against them. `a[0, [1, 2], :]` is `a[0][[1, 2], :]`. Neither the orthogonal - nor the vectorized path of the transform algebra models that — both widen a - scalar into a length-1 index array, which keeps the axis — so the scalars - are split off here and applied as a separate basic step first. + A scalar integer drops its axis, and neither the orthogonal nor the + vectorized path of the transform algebra models that — both widen a scalar + into a length-1 index array, which keeps the axis — so the scalars are split + off here and applied as a separate basic step first. + + Applying them *first* is this package's rule, not NumPy's. NumPy groups a + scalar with the advanced indices for the purpose of placing the broadcast + result, so the two disagree when a scalar and an index array are separated: + `a[0, ..., [1, 2]]` has shape `(2, 3)` for a `(2, 3, 4)` array, where + `a[0][..., [1, 2]]` has shape `(3, 2)`. The earlier claim here that they + always agree rested on `a[0, [1, 2], :]`, where the indices are adjacent and + they happen to. Scalar-first is the documented dialect (see the `lazy_array` + module docstring) — the divergence is deliberate, and this note exists so + that the correct end is not "fixed" later. Parameters ---------- @@ -200,7 +224,7 @@ def split_scalar_axes( scalar_axes: dict[int, int] = {} remaining: list[Any] = [] for sel, axis in zip(entries, axes, strict=True): - if isinstance(sel, (int, np.integer)) and not _is_bool_scalar(sel): + if _is_scalar_index(sel): scalar_axes[axis] = _normalize_int(int(sel), domain.shape[axis], axis) else: remaining.append(sel) diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 610ebd4db1..def3cccf18 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -143,8 +143,12 @@ def iter_chunk_transforms( - `chunk_coords`: which chunk to access. - `sub_transform`: maps output buffer coords to chunk-local coords. - - `out_indices`: for vectorized/array indexing, the output scatter - indices (integer array). `None` for basic/slice indexing. + - `out_indices`: where this chunk's values belong in the output. `None` + for basic/slice indexing, where the sub-transform says it on its own; an + integer array of scatter indices for a correlated (vectorized) selection; + and a `dict[int, ndarray]` of surviving positions per output dimension + for an orthogonal selection carrying more than one index array, which + `sub_transform_to_selections` turns into an open mesh. """ if any(size == 0 for size in transform.domain.shape): diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py index bd446306b2..ab705d92bc 100644 --- a/packages/zarr-indexing/src/zarr_indexing/json.py +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -210,7 +210,10 @@ def index_domain_from_json(data: IndexDomainJSON) -> IndexDomain: the same objects: `int(value)` alone accepts `3.9`, `"3"` and `True`, and each of those builds a domain that is not the document's. """ - if not isinstance(data, dict): + # The annotation says what a well-formed caller passes; this is a parser of + # documents that arrive from elsewhere, so the shape is checked rather than + # assumed. + if not isinstance(data, dict): # pyright: ignore[reportUnnecessaryIsInstance] raise NdselError("invalid_json", f"an index domain must be a JSON object, got {data!r}") body = normalize_ndsel({**data, "kind": "transform"}) inclusive_min = tuple( @@ -377,15 +380,16 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: `input_dimension` values are reconstructed by global dependency-axis ownership (see the module docstring). """ - if not isinstance(data, dict): + if not isinstance(data, dict): # pyright: ignore[reportUnnecessaryIsInstance] raise NdselError("invalid_json", f"a transform body must be a JSON object, got {data!r}") - if data.get("kind", "transform") != "transform": + kind = data.get("kind", "transform") + if kind != "transform": # Spelled last below, so a body carrying its own `kind` cannot reinterpret # the document as some other message and return a selection this function # never promised. raise NdselError( "invalid_json", - f"a transform body cannot carry kind {data['kind']!r}", + f"a transform body cannot carry kind {kind!r}", ) body = normalize_ndsel({**data, "kind": "transform"}) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index aa456ddaf4..69fd0b9241 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -172,7 +172,7 @@ iter_chunk_transforms, sub_transform_to_selections, ) -from zarr_indexing.grid import DimensionGridLike, EdgeDimensionGrid, dimension_grids_from_chunks +from zarr_indexing.grid import EdgeDimensionGrid, dimension_grids_from_chunks from zarr_indexing.json import transform_to_canonical from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_indexing.support import ( @@ -1502,7 +1502,7 @@ def unpartitioned(self) -> LazyArray: """ return self._with_grids(None) - def _with_grids(self, grids: tuple[DimensionGridLike, ...] | None) -> LazyArray: + def _with_grids(self, grids: tuple[EdgeDimensionGrid, ...] | None) -> LazyArray: return LazyArray._derive(self._array, self._transform, grids, self._window, self._support) def parts(self) -> Iterator[Partition]: diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 1975ee687e..af8895ea02 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -21,9 +21,9 @@ - **compose(outer, inner)** — chain two transforms. See `composition.py`. The transform is the atomic unit that connects user-facing indexing to -chunk-level I/O. Every `Array` holds a transform (identity by default). -`Array.lazy[...]` composes a new transform lazily. Reading resolves the -transform against the chunk grid via intersect + translate. +chunk-level I/O. A wrapper holds one — `LazyArray` starts from the identity — +and `.lazy[...]` composes a new transform lazily rather than reading. Reading +resolves the transform against the chunk grid via intersect + translate. """ from __future__ import annotations @@ -1522,6 +1522,17 @@ def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) if sel is Ellipsis or isinstance(sel, (int, np.integer)): continue if isinstance(sel, (list, np.ndarray)): + if mode == "orthogonal": + array = np.asarray(sel) + # An orthogonal selection is per-axis, so an integer array names + # coordinates along one axis and can only be one-dimensional. + # Left to the engine, this surfaced much later as a rank + # complaint about an `index_array` the caller never wrote. + if array.dtype.kind in "iu" and array.ndim > 1: + raise IndexError( + f"integer arrays in an orthogonal selection must be " + f"1-dimensional only; got one with {array.ndim} dimensions" + ) continue raise IndexError(f"unsupported selection type for {mode} indexing: {type(sel)!r}") diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index a51ad4df11..5481e93b70 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2175,3 +2175,36 @@ def test_the_coverage_count_matches_numpy_for_intervals_the_fast_path_declines( counted = _out_selection_cell_count(selection, out_shape) assert counted == np.empty(out_shape)[selection].size assert counted == expected + + +def test_a_zero_dimensional_index_array_drops_its_axis_like_a_scalar() -> None: + """`a[np.array(2), :]` is `a[2, :]` in NumPy, and now here too. + + Only Python and NumPy integers counted as scalars, so a 0-d array fell + through to the fancy path and was widened into a length-1 index array — + keeping an axis NumPy drops. That was a third answer, agreeing with neither + NumPy nor eager zarr, which rejects it. + """ + data = np.arange(20).reshape(4, 5) + for mode, expected in ( + ("oindex", data[np.array(2), :]), + ("vindex", data[np.array(2), np.array(3)]), + ): + view = ( + LazyArray(data).lazy.oindex[np.array(2), slice(None)] + if mode == "oindex" + else LazyArray(data).lazy.vindex[np.array(2), np.array(3)] + ) + assert view.shape == expected.shape, mode + np.testing.assert_array_equal(np.asarray(view.result()), expected, err_msg=mode) + + +def test_a_multidimensional_array_in_an_orthogonal_selection_is_refused() -> None: + """The rule belongs to the selection, so the message speaks its vocabulary. + + Left to the engine, this surfaced as a rank complaint about an `index_array` + the caller never wrote — the transform layer's words for a mistake made two + layers above it. + """ + with pytest.raises(IndexError, match="must be 1-dimensional"): + LazyArray(np.arange(20).reshape(4, 5)).lazy.oindex[[[0, 1], [2, 3]], slice(None)] diff --git a/packages/zarr-indexing/uv.lock b/packages/zarr-indexing/uv.lock index ab554e375b..2687ba71d8 100644 --- a/packages/zarr-indexing/uv.lock +++ b/packages/zarr-indexing/uv.lock @@ -366,7 +366,7 @@ wheels = [ [[package]] name = "mkdocs-material" -version = "9.7.6" +version = "9.7.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "babel" }, @@ -381,9 +381,9 @@ dependencies = [ { name = "pymdown-extensions" }, { name = "requests" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/45/29/6d2bcf41ae40802c4beda2432396fff97b8456fb496371d1bc7aad6512ec/mkdocs_material-9.7.6.tar.gz", hash = "sha256:00bdde50574f776d328b1862fe65daeaf581ec309bd150f7bff345a098c64a69", size = 4097959, upload-time = "2026-03-19T15:41:58.161Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f1/cd/c05d3a530ba7934f144fb45f7203cd236adc25c7bdcc34673d202f4b0278/mkdocs_material-9.7.7.tar.gz", hash = "sha256:c0649c065b1b0512d60aad8c10f947f8e455284475239b364b610f2deb4d0855", size = 4097923, upload-time = "2026-07-17T16:21:33.156Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/01/bc663630c510822c95c47a66af9fa7a443c295b47d5f041e5e6ae62ef659/mkdocs_material-9.7.6-py3-none-any.whl", hash = "sha256:71b84353921b8ea1ba84fe11c50912cc512da8fe0881038fcc9a0761c0e635ba", size = 9305470, upload-time = "2026-03-19T15:41:55.217Z" }, + { url = "https://files.pythonhosted.org/packages/ad/21/17c1bc9e6f47c972ad66fb2ac2568f99f90f1207eeb6fc3b34d094dba7b5/mkdocs_material-9.7.7-py3-none-any.whl", hash = "sha256:8ea9bb1737a5b524a5f9dcf2e1b4ebda8274ae3008aa7845720a97083bef708f", size = 9305438, upload-time = "2026-07-17T16:21:30.017Z" }, ] [[package]] @@ -397,7 +397,7 @@ wheels = [ [[package]] name = "mkdocstrings" -version = "1.0.4" +version = "1.0.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jinja2" }, @@ -407,9 +407,9 @@ dependencies = [ { name = "mkdocs-autorefs" }, { name = "pymdown-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1d/5d/f888d4d3eb31359b327bc9b17a212d6ef03fe0b0682fbb3fc2cb849fb12b/mkdocstrings-1.0.4.tar.gz", hash = "sha256:3969a6515b77db65fd097b53c1b7aa4ae840bd71a2ee62a6a3e89503446d7172", size = 100088, upload-time = "2026-04-15T09:16:53.376Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/71/f85bdf13355073ae15a7375f09879375a830553552e58c1c4b7e0bbc5c8b/mkdocstrings-1.0.6.tar.gz", hash = "sha256:a0b8c2bdd29a6416c80d717aa369bbf7831946bd9f23c2a66db1b1dbe7693dbd", size = 100649, upload-time = "2026-07-11T19:38:05.732Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6e/94/be70f8ee9c45f2f62b39a1f0e9303bc20e138a8f3b8e50ffd89498e177e1/mkdocstrings-1.0.4-py3-none-any.whl", hash = "sha256:63464b4b29053514f32a1dbbf604e52876d5e638111b0c295ab7ed3cac73ca9b", size = 35560, upload-time = "2026-04-15T09:16:51.436Z" }, + { url = "https://files.pythonhosted.org/packages/5d/5b/4c1902e8bdd5c4db63284e9d101dece4038d4025d6d88850ffe0a1578980/mkdocstrings-1.0.6-py3-none-any.whl", hash = "sha256:2703708697487d1b6d6d7b412e176fa436edf120c1bf81dc9e126b12d00893c7", size = 35787, upload-time = "2026-07-11T19:38:04.417Z" }, ] [[package]] @@ -647,27 +647,27 @@ wheels = [ [[package]] name = "ruff" -version = "0.15.20" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/43/dc/35b341fc554ba02f217fc10da57d1a75168cfbcf75b0ef2202176d4c4f2d/ruff-0.15.20.tar.gz", hash = "sha256:1416eb04349192646b54de98f146c4f59afe37d0decfc02c3cbbf396f3a28566", size = 4755489, upload-time = "2026-06-25T17:20:37.578Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/d9/2d5014f0253ba541d2061d9fa7193f48e941c8b21bb88a7ff9bbe0bd0596/ruff-0.15.20-py3-none-linux_armv6l.whl", hash = "sha256:00e188c53e499c3c1637f73c91dcf2fb56d576cab76ce1be50a27c4e80e37078", size = 10839665, upload-time = "2026-06-25T17:19:44.702Z" }, - { url = "https://files.pythonhosted.org/packages/c6/d3/ac1798ba64f670698867fcfc591d50e7e421bef137db564858f619a30fcf/ruff-0.15.20-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9ebd1fd9b9c95fc0bd7b2761aebec1f030013d2e193a2901b224af68fe47251b", size = 11208649, upload-time = "2026-06-25T17:19:48.787Z" }, - { url = "https://files.pythonhosted.org/packages/47/47/d3ac899991202095dfcf3d5176be4272642be3cf981a2f1a30f72a2afb95/ruff-0.15.20-py3-none-macosx_11_0_arm64.whl", hash = "sha256:c5b16cdd67ca108185cd36dce98c576350c03b1660a751de725fb049193a0632", size = 10622638, upload-time = "2026-06-25T17:19:51.354Z" }, - { url = "https://files.pythonhosted.org/packages/33/13/4e043fe30aa94d4ff5213a9881fc296d12960f5971b234a5263fdc225312/ruff-0.15.20-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3413bb3c3d2ca6a8208f1f4809cd2dca3c6de6d0b491c0e70847672bde6e6efd", size = 10984227, upload-time = "2026-06-25T17:19:54.044Z" }, - { url = "https://files.pythonhosted.org/packages/76/e6/92e7bf40388bc5800073b96564f56264f7e48bfd1a498f5ced6ae6d5a769/ruff-0.15.20-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bd7ec42b3bb3da066488db093308a69c4ac5ee6d2af333a86ba6e2eb2e7dd44b", size = 10622882, upload-time = "2026-06-25T17:19:57.037Z" }, - { url = "https://files.pythonhosted.org/packages/13/7a/43460be3f24495a3aa46d4b16873e2c4941b3b5f0b00cf88c03b7b94b339/ruff-0.15.20-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e1a36ad0eb77fba9aabfb69ede54de6f376d04ac18ebea022847046d340a8267", size = 11474808, upload-time = "2026-06-25T17:20:00.357Z" }, - { url = "https://files.pythonhosted.org/packages/27/a0/f37077884873221c6b33b4ab49eb18f9f88e54a16a25a5bca59bef46dd66/ruff-0.15.20-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b6df3b1e4610432f0386dba04d853b5f08cbbc903410c6fcc02f620f05aff53c", size = 12293094, upload-time = "2026-06-25T17:20:03.446Z" }, - { url = "https://files.pythonhosted.org/packages/a6/74/165545b60256a9704c21ac0ec4a0d07933b320812f9584836c9f4aca4292/ruff-0.15.20-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e89f198a1ea6ef0d727c1cf16088bc91a6cb0ab947dedc966715691647186eae", size = 11526176, upload-time = "2026-06-25T17:20:06.301Z" }, - { url = "https://files.pythonhosted.org/packages/86/b1/a976a136d40ade83ce743578399865f57001003a409acadc0ecbb3051082/ruff-0.15.20-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:309809086c2acb67624950a3c8133e80f32d0d3e27106c0cd60ff26657c9f24b", size = 11520767, upload-time = "2026-06-25T17:20:09.191Z" }, - { url = "https://files.pythonhosted.org/packages/19/0f/f032696cb01c9b54c0263fa393474d7758f1cdc021a01b04e3cbc2500999/ruff-0.15.20-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:2d2374caa2f2c2f9e2b7da0a50802cfb8b79f55a9b5e49379f564544fbf56487", size = 11500132, upload-time = "2026-06-25T17:20:13.602Z" }, - { url = "https://files.pythonhosted.org/packages/4b/f4/51b1a14bc69e8c224b15dab9cce8e99b425e0455d462caa2b3c9be2b6a8e/ruff-0.15.20-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:a1ed17b65293e0c2f22fc387bc13198a5de94bf4429589b0ff6946b0feaf21a3", size = 10943828, upload-time = "2026-06-25T17:20:16.635Z" }, - { url = "https://files.pythonhosted.org/packages/71/4b/fe267640783cd02bf6c5cc290b1df1051be2ec294c678b5c15fe19e52343/ruff-0.15.20-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:f701305e66b38ea6c91882490eb73459796808e4c6362a1b765255e0cdcd4053", size = 10645418, upload-time = "2026-06-25T17:20:19.4Z" }, - { url = "https://files.pythonhosted.org/packages/b0/c0/a65aa4ec2f5e87a1df32dc3ec1fede434fe3dfd5cbcf3b503cafc676ab54/ruff-0.15.20-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5b9c0c367ad8e5d0d5b5b8537864c469a0a0e55417aadfbeca41fa61333be9f4", size = 11211770, upload-time = "2026-06-25T17:20:22.033Z" }, - { url = "https://files.pythonhosted.org/packages/5a/a4/0caa331d954ae2723d729d351c989cb4ca8b6077d5c6c2cb6de75e98c041/ruff-0.15.20-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:01cc00dd58f0df339d0e902219dd53990ea99996a0344e5d9cc8d45d5307e460", size = 11618698, upload-time = "2026-06-25T17:20:25.259Z" }, - { url = "https://files.pythonhosted.org/packages/10/9b/5f14927848d2fd4aa891fd88d883788c5a7baba561c7874732364045708c/ruff-0.15.20-py3-none-win32.whl", hash = "sha256:ed65ef510e43a137207e0f01cfcf998aeddb1aeeda5c9d35023e910284d7cf21", size = 10857322, upload-time = "2026-06-25T17:20:28.612Z" }, - { url = "https://files.pythonhosted.org/packages/fa/f0/fe47c501f9dea92a26d788ff98bb5d92ed4cb4c88792c5c88af6b697dc8e/ruff-0.15.20-py3-none-win_amd64.whl", hash = "sha256:a525c81c70fb0380344dd1d8745d8cc1c890b7fc94a58d5a07bd8eb9557b8415", size = 11993274, upload-time = "2026-06-25T17:20:31.871Z" }, - { url = "https://files.pythonhosted.org/packages/d7/2b/9555445e1201d92b3195f45cdb153a0b68f24e0a4273f6e3d5ab46e212bb/ruff-0.15.20-py3-none-win_arm64.whl", hash = "sha256:2f5b2a6d614e8700388806a14996c40fab2c47b819ef57d790a34878858ed9ca", size = 11343498, upload-time = "2026-06-25T17:20:35.03Z" }, +version = "0.15.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3a/06/ae069393fc66e8ff33036d4b368003833bf6e88ccf182e17e7a2f1c754fd/ruff-0.15.22.tar.gz", hash = "sha256:3f15175b1fb580126f58285a5dae6b2ea89000136d980c64499211f116b54809", size = 4785063, upload-time = "2026-07-16T15:14:13.244Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/18/ee54b7ae1e121be7a28ea6da4b67564ebb0530e183a54415ab7e3bcd2c4e/ruff-0.15.22-py3-none-linux_armv6l.whl", hash = "sha256:44423e73493737f5e7c5b41d475483898ff37afcdae38bc3da5085e29af1c2d8", size = 10781258, upload-time = "2026-07-16T15:13:19.452Z" }, + { url = "https://files.pythonhosted.org/packages/2f/d2/2520cb14761ddbeaf57642a76942fc36adcbdbe53b4532241995f6fc485c/ruff-0.15.22-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:b82c6482946e9eda7ff2e091d25b8bad3f718684e1916d41bd56873cee05b697", size = 10999477, upload-time = "2026-07-16T15:13:23.318Z" }, + { url = "https://files.pythonhosted.org/packages/c9/10/74e53572aa758dfaa678c2a2646b5c5515d884b7ca56be4d2ce03ca4b560/ruff-0.15.22-py3-none-macosx_11_0_arm64.whl", hash = "sha256:11c1c715af53a09f714e011106bffc419751ec8232fcb5da42173284ea3fec6f", size = 10466716, upload-time = "2026-07-16T15:13:26.162Z" }, + { url = "https://files.pythonhosted.org/packages/1e/cc/44eaaf0844e028182f2d0a8f2190d0f359159aed0a9e5ab861d892f1ae2a/ruff-0.15.22-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:742a29cf29bddb7c8327895d6a10e0e6c5b38a96dd407af9b5d0857f809c0576", size = 10892644, upload-time = "2026-07-16T15:13:29.229Z" }, + { url = "https://files.pythonhosted.org/packages/9f/21/8edf559014d2b0f82beea19cfb713993ad802ccda16868769979c6090a84/ruff-0.15.22-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:72af58b951b0ae395935ae79763dc349bc0eb706319d28f7a33ad2cfb3cfc178", size = 10576719, upload-time = "2026-07-16T15:13:32.35Z" }, + { url = "https://files.pythonhosted.org/packages/bf/1e/3a13abd392a3b50b62e5938a831f9ab6e588358cacad5c18545b716d2182/ruff-0.15.22-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:62d425005c1835eb24e2ee4161cb90e8db263415f4a71c8c72c33abaa6c0c224", size = 11376494, upload-time = "2026-07-16T15:13:35.958Z" }, + { url = "https://files.pythonhosted.org/packages/bf/3e/422d3d95bcf04dd78e1aeac22184d4f9a8fb2c01865d39d44618484a0317/ruff-0.15.22-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e8b9b3f8779a4f08c969defc3c8c35abffaa757e601ed5ae66d6d1db6519969a", size = 12208370, upload-time = "2026-07-16T15:13:39.185Z" }, + { url = "https://files.pythonhosted.org/packages/1e/91/5d065a0e0a02bf4813f5119ad278462eed081d2b832eb7c021ade0ec9e65/ruff-0.15.22-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:1e0dd1b2e4d3d585f897a0d137cbf4eaf6223bef4e8ce34d6bb12556c5f9249e", size = 11581098, upload-time = "2026-07-16T15:13:42.132Z" }, + { url = "https://files.pythonhosted.org/packages/f6/f9/a0d4871d12fae702eb1f41b686caf05f1f8b124dc6db6f784f53d74918fa/ruff-0.15.22-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:365523eb91d9224e1bcb03b022fbf0facb8f9e23792a2c53d9d4b3924bdbdebb", size = 11399422, upload-time = "2026-07-16T15:13:45.2Z" }, + { url = "https://files.pythonhosted.org/packages/18/80/c843a5176cddbceb0b7e8dd41cf9993490796c1c469348d384f5a5c13c56/ruff-0.15.22-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:fabfd168afdf29fee5be98b831efa9683c94d7c5a3b58b9ce5a2e38444589a74", size = 11381683, upload-time = "2026-07-16T15:13:48.46Z" }, + { url = "https://files.pythonhosted.org/packages/d4/00/8485de0ae92239438a36cfc51350db9b9e85c9ebdfaea91b18e422706662/ruff-0.15.22-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:225dbf095a87f1d9f90f5fd7924d2613ee452a75a4308c63a8f50f761787aa7c", size = 10850295, upload-time = "2026-07-16T15:13:51.655Z" }, + { url = "https://files.pythonhosted.org/packages/fa/91/24977ec2ec72eaf15e4394ace2959fdff2dd1e14f03e005e838023407169/ruff-0.15.22-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:1877d63b9d24ed278744f1523fd11b85540566d54641f97c566d7d9dc5ca5296", size = 10579640, upload-time = "2026-07-16T15:13:54.79Z" }, + { url = "https://files.pythonhosted.org/packages/9c/47/9b51216951974df1f263ac19da550d34252e0ed7218c25f10c5ef9ed7517/ruff-0.15.22-py3-none-musllinux_1_2_i686.whl", hash = "sha256:a1606c510bd7215680d32efab38965f7cdec3ef69f5170a3f4791404ffdd5262", size = 11105077, upload-time = "2026-07-16T15:13:57.915Z" }, + { url = "https://files.pythonhosted.org/packages/c2/47/20e9d4a3b8016778acea5fc32bb50d35d207500a17ddb529ffa6996feef8/ruff-0.15.22-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:630479b18625f5ffc373f77603a22a9f8ac0acd7ff0501178b5db28ec71e9c64", size = 11490980, upload-time = "2026-07-16T15:14:01.032Z" }, + { url = "https://files.pythonhosted.org/packages/4d/76/3f72d8fc38c1cb77b38c56a70da9d0c17700cc1cc50f9649c9d3c8f5ba71/ruff-0.15.22-py3-none-win32.whl", hash = "sha256:e5ba0e4a13fd14abbed2a77b517a3911290c6c6c59ef67784328d1668fab76cf", size = 10789165, upload-time = "2026-07-16T15:14:04.16Z" }, + { url = "https://files.pythonhosted.org/packages/cb/46/4965251734c2b6fcdca1b1b187d20bcac3af0ee5b083b89c910bb961ce3a/ruff-0.15.22-py3-none-win_amd64.whl", hash = "sha256:9be63ba1eb936acd2d1342fb8337c356353706fce233b2a15a09a97037e6acde", size = 11938297, upload-time = "2026-07-16T15:14:07.316Z" }, + { url = "https://files.pythonhosted.org/packages/57/c9/e69b1ff4c8b69093ef08b8919ab767af0569666865b39c30a8795d88d3c6/ruff-0.15.22-py3-none-win_arm64.whl", hash = "sha256:e1168075b72158510839f250027659cdd78476f40507dd517892304c41318661", size = 11298172, upload-time = "2026-07-16T15:14:10.51Z" }, ] [[package]] @@ -758,10 +758,10 @@ provides-extras = ["testing"] docs = [ { name = "griffe-inherited-docstrings", specifier = "==1.1.3" }, { name = "mkdocs", specifier = "==1.6.1" }, - { name = "mkdocs-material", specifier = "==9.7.6" }, - { name = "mkdocstrings", specifier = "==1.0.4" }, + { name = "mkdocs-material", specifier = "==9.7.7" }, + { name = "mkdocstrings", specifier = "==1.0.6" }, { name = "mkdocstrings-python", specifier = "==2.0.5" }, - { name = "ruff", specifier = "==0.15.20" }, + { name = "ruff", specifier = "==0.15.22" }, ] test = [ { name = "hypothesis", specifier = ">=6.160.0" }, From 5934c1ee00f8f61566cdb4b13308f1108e23be4d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 09:17:35 +0200 Subject: [PATCH 25/32] fix(zarr-indexing): collapse an empty index array instead of extending the format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit fixed the empty-selection round trip by carrying an `index_array_shape` field, on the reasoning that JSON nested lists cannot express the shape of an array with a leading zero axis — `[]` is the only spelling of every empty shape, and `[[]]` is (1, 0) with nothing for (0, 1). That much is true, but the conclusion was wrong: it invented a field ndsel does not define, so the documents this package wrote stopped being ndsel documents. The reference implementation does not have the problem, because it never emits an empty index array. `t[ts.d[0][[]]]` in TensorStore is `out[0] = 0`, emitted as `{}` — a constant map. An empty index array names no cell, and it can only be empty because an input dimension is, since the full-rank invariant makes every axis either 1 or the domain's extent. Nothing is ever read through it, the emptiness is carried by the domain, and the map is degenerate in exactly the way a size-1 array is — which this format already collapses. So it collapses the same way, and the extension is gone. The loader still recovers the axis from the domain for an empty array arriving from a producer that does emit one, since ndsel does not forbid it; that path just no longer has a first-party caller. Checked against tensorstore 0.1.84: every canonical body from 4000 randomized transforms loads into `ts.IndexTransform`, and every one re-emits itself unchanged. No spec change needed. Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-indexing/changes/272.bugfix.md | 4 +- .../zarr-indexing/src/zarr_indexing/json.py | 78 ++++++++++--------- .../src/zarr_indexing/messages.py | 39 ---------- packages/zarr-indexing/tests/test_json.py | 43 ++++------ 4 files changed, 61 insertions(+), 103 deletions(-) diff --git a/packages/zarr-indexing/changes/272.bugfix.md b/packages/zarr-indexing/changes/272.bugfix.md index 8ef3ac9ae4..807447ee38 100644 --- a/packages/zarr-indexing/changes/272.bugfix.md +++ b/packages/zarr-indexing/changes/272.bugfix.md @@ -1,8 +1,8 @@ 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 -could not reload its own output for a selection that selects nothing, and its -domain loader validated nothing; `sub_transform_to_selections` described a +emitted a document nothing could load for a selection that selects nothing, and +its domain loader validated nothing; `sub_transform_to_selections` 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 diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py index ab705d92bc..595354fce4 100644 --- a/packages/zarr-indexing/src/zarr_indexing/json.py +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -30,7 +30,17 @@ domain, unconsumed — a valid transform. This makes a length-1 `oindex` selection round-trip *behaviorally* (an `ArrayMap` becomes a `ConstantMap`) rather than by object identity. - 2. Non-degenerate `index_array` maps are emitted **without** `input_dimension`. + 2. An **empty** `index_array` (size 0) collapses the same way, to + `{offset: 0}`. It names no cell, and it can only be empty because an input + dimension is — the full-rank invariant makes every axis either 1 or the + domain's extent — so nothing is ever read through it and the emptiness is + carried by the domain, which is emitted separately. TensorStore does the + same: `t[ts.d[0][[]]]` is `out[0] = 0`, emitted as `{}`. Emitting the array + instead would produce a document neither implementation could load, because + `ndarray.tolist()` renders every empty array as `[]` once the leading axis + is the zero-length one, and nested lists cannot spell the shape back — + `[[]]` is `(1, 0)` and nothing spells `(0, 1)`. + 3. Non-degenerate `index_array` maps are emitted **without** `input_dimension`. - **On load** (`transform_from_canonical`): the in-memory `input_dimension` is reconstructed from the full-rank array's dependency axes (its non-singleton @@ -47,7 +57,6 @@ from __future__ import annotations -import math from collections import Counter from typing import Any, Required, TypedDict @@ -109,7 +118,6 @@ class OutputIndexMapJSON(TypedDict, total=False): input_dimension: int index_array: NestedIntList index_array_bounds: list[IndexValueJSON] - index_array_shape: list[int] class IndexTransformJSON(TypedDict, total=False): @@ -253,20 +261,28 @@ def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON: if m.index_array.size == 1: value = int(m.index_array.reshape(-1)[0]) return {"offset": m.offset + m.stride * value} - body: OutputIndexMapJSON = { + if m.index_array.size == 0: + # An empty index array names no cell, and it can only be empty because + # an input dimension is — the full-rank invariant makes every axis + # either 1 or the domain's extent, so a 0 there means the domain has a 0 + # there too. Nothing is ever read through this map, which makes it + # degenerate in exactly the way a size-1 array is, and it collapses the + # same way. The empty dimension stays in the domain, so the transform + # that comes back describes the same (empty) selection. + # + # This is also what the reference implementation does: TensorStore + # renders `t[ts.d[0][[]]]` as `out[0] = 0` and emits `{}`. Emitting the + # array instead would produce a document neither it nor this package + # could load, because `tolist()` renders every empty array as `[]` once + # the leading axis is the zero-length one — and nested lists cannot + # spell the shape back, `[[]]` being (1, 0) with nothing for (0, 1). + return {"offset": 0} + return { "offset": m.offset, "stride": m.stride, "index_array": m.index_array.tolist(), "index_array_bounds": ["-inf", "+inf"], } - if m.index_array.size == 0: - # `tolist()` renders every empty array as `[]` once the leading axis is - # the zero-length one, so rank — and with it the axis the map varies - # over — does not survive the trip. Nested lists cannot express it - # either: `[[]]` is shape (1, 0) and nothing spells (0, 1). Carry the - # shape alongside so the map that comes back is the map that went out. - body["index_array_shape"] = list(m.index_array.shape) - return body def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: @@ -304,34 +320,24 @@ def _solo_dependency_axis(arr: np.ndarray[Any, Any]) -> int | None: def _full_rank_index_array( arr: np.ndarray[Any, np.dtype[np.intp]], - declared_shape: list[int] | None, domain: IndexDomain, where: str, ) -> np.ndarray[Any, np.dtype[np.intp]]: """Give an incoming `index_array` the input rank the engine requires. - Three cases, in order of how much the document tells us: - - - An explicit `index_array_shape` settles it. This is what our own emitter - writes for an empty array, whose shape nested lists cannot carry. - - An empty array with no declared shape has to be reconstructed from the - domain: it can only be empty because some input dimension is, so that - dimension takes the zero and every other takes a broadcast singleton. - Ambiguous when the domain is empty on more than one axis, and impossible - when it is empty on none — both are rejected rather than guessed at. - - A lower-rank non-empty array is aligned to the *trailing* input - dimensions, which is how NumPy broadcasts and how a producer that omits - leading singletons means it to be read. + ndsel leaves index-array rank unvalidated, so a conformant producer may send + an array of lower rank that broadcasts against the domain. A non-empty one + is aligned to the *trailing* input dimensions, which is how NumPy broadcasts + and how a producer omitting leading singletons means it to be read. + + An empty array is a different matter: `[]` is the only spelling of every + empty shape once the leading axis is the zero-length one, so the axis it + varies over cannot be read off it. It is recovered from the domain, which + can only be empty on the axis in question — and rejected when the domain + leaves that ambiguous. This package never emits such a document (an empty + map is degenerate and collapses to a constant, as TensorStore's does), so + this path exists for external producers alone. """ - if declared_shape is not None: - if math.prod(declared_shape) != arr.size: - raise NdselError( - "invalid_json", - f"{where}.index_array_shape is {declared_shape}, which holds " - f"{math.prod(declared_shape)} entries, but index_array holds {arr.size}", - ) - return arr.reshape(tuple(declared_shape)) - if arr.size == 0 and arr.ndim != domain.ndim: empty_axes = [k for k, extent in enumerate(domain.shape) if extent == 0] if len(empty_axes) != 1: @@ -339,7 +345,7 @@ def _full_rank_index_array( "invalid_json", f"{where}.index_array is empty, but the input domain has " f"{len(empty_axes)} zero-length dimensions, so the axis it varies " - f"over cannot be recovered; send 'index_array_shape' to say which", + f"over cannot be recovered", ) shape = [1] * domain.ndim shape[empty_axes[0]] = 0 @@ -423,7 +429,7 @@ def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: # may send an array of lower rank that broadcasts against the domain. # Widen it here, on the way in, so every transform that exists holds # the full-rank invariant the engine reads dependency axes from. - arr = _full_rank_index_array(arr, om.get("index_array_shape"), domain, where) + arr = _full_rank_index_array(arr, domain, where) arrays[i] = arr dep = _array_map_dependency_axes(arr) array_axes[i] = dep diff --git a/packages/zarr-indexing/src/zarr_indexing/messages.py b/packages/zarr-indexing/src/zarr_indexing/messages.py index 0cad32684f..71ac4c5e1a 100644 --- a/packages/zarr-indexing/src/zarr_indexing/messages.py +++ b/packages/zarr-indexing/src/zarr_indexing/messages.py @@ -97,7 +97,6 @@ def __init__(self, reason: str, detail: str = "") -> None: "input_dimension", "index_array", "index_array_bounds", - "index_array_shape", } ) @@ -553,10 +552,6 @@ def _normalize_output_map(raw: Any, where: str) -> dict[str, Any]: "index_array": raw["index_array"], "index_array_bounds": bounds, } - if "index_array_shape" in raw: - normalized["index_array_shape"] = _check_index_array_shape( - raw["index_array_shape"], where - ) return normalized if has_input_dim: @@ -591,33 +586,6 @@ def _check_index_array_bounds(value: Any, where: str) -> list[int | str]: return [lo, hi] -def _check_index_array_shape(value: Any, where: str) -> list[int]: - """Validate an `index_array_shape`: non-negative extents, one per input dimension. - - Carried because JSON nested lists cannot express the shape of an array with - no elements — `[]` is the only spelling of every empty shape, so a leading - zero axis loses both its rank and which dimension it varies over. See the - `json` module docstring. - """ - if not isinstance(value, list): - raise NdselError( - "invalid_json", f"{where}.index_array_shape must be an array, got {value!r}" - ) - if len(value) > _MAX_RANK: - raise NdselError( - "invalid_json", - f"{where}.index_array_shape has rank {len(value)}, above the maximum {_MAX_RANK}", - ) - extents = [_check_int(v, f"{where}.index_array_shape[{i}]") for i, v in enumerate(value)] - for i, extent in enumerate(extents): - if extent < 0: - raise NdselError( - "invalid_json", - f"{where}.index_array_shape[{i}] must be >= 0, got {extent}", - ) - return extents - - def _normalize_transform(obj: dict[str, Any]) -> dict[str, Any]: allowed = frozenset( { @@ -692,13 +660,6 @@ def _normalize_transform(obj: dict[str, Any]) -> dict[str, Any]: f"output[{i}].input_dimension is {m['input_dimension']}, " f"outside the valid range [0, {rank}) for input_rank {rank}", ) - declared = m.get("index_array_shape") - if declared is not None and len(declared) != rank: - raise NdselError( - "rank_mismatch", - f"output[{i}].index_array_shape has rank {len(declared)} but " - f"input_rank is {rank}", - ) else: output = _identity_output(rank) diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py index bcc68e4640..83ea55d04f 100644 --- a/packages/zarr-indexing/tests/test_json.py +++ b/packages/zarr-indexing/tests/test_json.py @@ -460,14 +460,18 @@ def test_an_index_array_that_does_not_span_its_domain_is_rejected() -> None: ) -def test_an_empty_index_array_round_trips_with_its_axis_intact() -> None: +def test_an_empty_index_array_collapses_to_a_constant() -> None: """Selecting nothing must survive a trip through JSON. + An empty index array names no cell, and can only be empty because an input + dimension is, so nothing is ever read through it. It is degenerate in + exactly the way a size-1 array is, and collapses the same way — which is + also what TensorStore emits for `t[ts.d[0][[]]]`. + + Emitting the array instead produced a document nothing could load: `tolist()` renders every empty array as `[]` once the leading axis is the - zero-length one, so both the rank and the axis the map varies over were lost - on the way out — and the loader, widening by prepending singletons, put the - dependency back on a different axis. The body carries the shape for exactly - this case, so the map that comes back is the map that went out. + zero-length one, so the rank went with it, and the loader put the dependency + back on a different axis by prepending singletons. """ for shape, selection in ( ((5, 3), (np.array([], dtype=np.intp), slice(None))), @@ -475,19 +479,18 @@ def test_an_empty_index_array_round_trips_with_its_axis_intact() -> None: ): transform = IndexTransform.from_shape(shape).oindex[selection] body = index_transform_to_json(transform) - reloaded = index_transform_from_json(body) - before = [m.index_array.shape for m in transform.output if isinstance(m, ArrayMap)] - after = [m.index_array.shape for m in reloaded.output if isinstance(m, ArrayMap)] - assert after == before + assert all("index_array" not in m for m in body["output"]) + reloaded = index_transform_from_json(body) + assert reloaded.domain == transform.domain assert index_transform_to_json(reloaded) == body -def test_an_empty_index_array_without_a_declared_shape_is_recovered_from_the_domain() -> None: - """A producer that omits the shape is still readable when the domain settles it. +def test_an_empty_index_array_from_elsewhere_is_recovered_from_the_domain() -> None: + """A producer that does emit one is still readable when the domain settles it. - An index array can only be empty because an input dimension is, so a domain - with exactly one zero-length dimension names the axis unambiguously. + This package never writes such a document, but ndsel does not forbid it, and + the domain names the axis unambiguously when exactly one dimension is empty. """ body: IndexTransformJSON = { "input_inclusive_min": [0, 0], @@ -509,19 +512,7 @@ def test_an_ambiguous_empty_index_array_is_rejected() -> None: with pytest.raises(NdselError) as excinfo: index_transform_from_json(body) assert excinfo.value.reason == "invalid_json" - assert "index_array_shape" in str(excinfo.value) - - -def test_a_declared_index_array_shape_that_does_not_fit_is_rejected() -> None: - body: IndexTransformJSON = { - "input_inclusive_min": [0, 0], - "input_exclusive_max": [3, 4], - "output": [{"index_array": [1, 2, 0, 2], "index_array_shape": [1, 3]}], - } - with pytest.raises(NdselError) as excinfo: - index_transform_from_json(body) - assert excinfo.value.reason == "invalid_json" - assert "index_array_shape" in str(excinfo.value) + assert "zero-length" in str(excinfo.value) @pytest.mark.parametrize( From d9464ffab218d2ce62bbd59ad1cfbdf9efacdb66 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 11:15:04 +0200 Subject: [PATCH 26/32] fix(indexing): address lazy array review findings Assisted-by: Codex:gpt-5 --- .../src/zarr_indexing/lazy_array.py | 59 ++++++++++++------- .../src/zarr_indexing/output_map.py | 25 +++++--- .../zarr-indexing/tests/test_lazy_array.py | 31 ++++++++++ .../zarr-indexing/tests/test_output_map.py | 26 ++++++++ 4 files changed, 112 insertions(+), 29 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 69fd0b9241..7caeb51fc1 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -1589,7 +1589,13 @@ def _select(self, selection: Any, mode: SelectionMode) -> LazyArray: transform = selection_to_transform(scalar_selection, transform, "basic") transform = transform.translate_domain_to((0,) * transform.input_rank) literal = normalize_positional_selection(selection, transform.domain, mode) - composed = selection_to_transform(literal, transform, mode) + if mode == "basic": + # IndexTransform's basic path includes NumPy's `None`/newaxis. + # `selection_to_transform` intentionally exposes a narrower basic + # selection contract and rejects it. + composed = transform[literal] + else: + composed = selection_to_transform(literal, transform, mode) return LazyArray._derive(self._array, composed, self._parts, self._window, self._support) def __getitem__(self, selection: Any) -> Any: @@ -1634,33 +1640,44 @@ def result(self) -> Any: uninitialized where nothing was written, so an incomplete walk is reported rather than returned. """ + out_shape = self.shape + size = math.prod(out_shape) + if size == 0: + return self._empty_result(out_shape) + if self._parts is None: block = _read(self._array, self._window, self._transform, self._support) return _detach(block, self._array) - out_shape = self.shape out = self._output_buffer(out_shape) - size = math.prod(out_shape) - if size > 0: - written = 0 - for part in self.parts(): - # Counted before the scatter, so a part addressing the wrong - # number of axes is named rather than reported as a broadcast - # failure against the buffer. - written += _out_selection_cell_count(part.out_selection, out_shape) - out[part.out_selection] = np.asanyarray(part.view.result()) - if written != size: - # The buffer is uninitialized where no part wrote, so a partition - # walk that does not tile the view exactly would otherwise hand - # back process memory dressed as data. The parts are disjoint by - # contract, so counting the cells each addresses is enough: - # a gap undercounts and an overlap overcounts. - raise AssertionError( - f"the partition walk addressed {written} of the view's {size} " - "cells; this is a bug in zarr-indexing's partition walk" - ) + written = 0 + for part in self.parts(): + # Counted before the scatter, so a part addressing the wrong + # number of axes is named rather than reported as a broadcast + # failure against the buffer. + written += _out_selection_cell_count(part.out_selection, out_shape) + out[part.out_selection] = np.asanyarray(part.view.result()) + if written != size: + # The buffer is uninitialized where no part wrote, so a partition + # walk that does not tile the view exactly would otherwise hand + # back process memory dressed as data. The parts are disjoint by + # contract, so counting the cells each addresses is enough: + # a gap undercounts and an overlap overcounts. + raise AssertionError( + f"the partition walk addressed {written} of the view's {size} " + "cells; this is a bug in zarr-indexing's partition walk" + ) return out + def _empty_result(self, out_shape: tuple[int, ...]) -> Any: + """Allocate an empty result without asking the source for any data.""" + if self._parts is None: + xp = _namespace(self._array) + empty = getattr(xp, "empty", None) + if empty is not None: + return empty(out_shape, dtype=self.dtype) + return self._output_buffer(out_shape) + def _output_buffer(self, out_shape: tuple[int, ...]) -> Any: """The buffer `result()` scatters parts into. diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py index 871d778620..7f7285e4ff 100644 --- a/packages/zarr-indexing/src/zarr_indexing/output_map.py +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -103,21 +103,30 @@ class ArrayMap: input_dimension: int | None = None def __post_init__(self) -> None: - """Hold the index array through a read-only view. + """Own the index array and expose it read-only. A map is frozen, but the array inside it was not: reaching through a view's transform to `index_array[0] = 9` silently changed what the view returned, in a package whose whole contract is that a view is a - description of a read and resolving it twice answers alike. The view is - taken rather than the caller's array being marked, so constructing a map - does not take away the ability to write to an array you still own. + description of a read and resolving it twice answers alike. Owning the + array also prevents the caller from changing the contents behind the + read-only view, which would invalidate this value object's hash. """ - # `asarray` first: a NumPy scalar reaches here from paths that index an - # array down to one element, and a scalar has no flags to set. - frozen = np.asarray(self.index_array).view() - frozen.flags.writeable = False + # Immutable bytes are the ultimate owner so callers cannot re-enable + # the WRITEABLE flag, as they can on a read-only array that owns its + # allocation. `asarray` also accepts the NumPy scalars that reach here + # after indexing an array down to one element. + array = np.asarray(self.index_array) + frozen = np.frombuffer(array.tobytes(), dtype=array.dtype).reshape(array.shape) object.__setattr__(self, "index_array", frozen) + def __reduce__(self) -> tuple[object, tuple[object, int, int, int | None]]: + """Reconstruct through `__init__`, preserving the ownership invariant.""" + return ( + type(self), + (self.index_array, self.offset, self.stride, self.input_dimension), + ) + def __eq__(self, other: object) -> bool: """Value equality, comparing index arrays element-wise. diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 5481e93b70..17f1297835 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -250,6 +250,7 @@ def source(request: pytest.FixtureRequest) -> LazyArray: ), ("basic-ellipsis", lambda a: a.lazy[..., -2], lambda r: r[..., -2]), ("basic-negative-scalar", lambda a: a.lazy[-3], lambda r: r[-3]), + ("basic-newaxis", lambda a: a.lazy[None, :, :, None], lambda r: r[None, :, :, None]), ("basic-empty", lambda a: a.lazy[:, 2:2, :], lambda r: r[:, 2:2, :]), ("basic-all-scalars", lambda a: a.lazy[-1, 0, 2], lambda r: r[-1, 0, 2]), ( @@ -1576,6 +1577,36 @@ def __getitem__(self, key: Any) -> Any: return self._data[key] +class RecordingSource(MinimalSource): + """A minimal source that records every attempted data read.""" + + def __init__(self, data: np.ndarray[Any, Any]) -> None: + super().__init__(data) + self.reads: list[Any] = [] + + def __getitem__(self, key: Any) -> Any: + self.reads.append(key) + return super().__getitem__(key) + + +@pytest.mark.parametrize( + "build", + [ + lambda a: a.lazy[:, 2:2, :], + lambda a: a.lazy.vindex[np.array([[6], [3], [0]]), -4].lazy[0, 0:0], + ], + ids=["ordinary", "unreferenced-axis"], +) +def test_an_empty_view_does_not_read_its_source( + build: Callable[[LazyArray], LazyArray], +) -> None: + source = RecordingSource(reference()) + result = build(LazyArray(source)).result() + + assert result.size == 0 + assert source.reads == [] + + def _require_basic(key: tuple[Any, ...], channel: str) -> None: """Raise unless every selector is a slice walking upwards.""" for selector in key: diff --git a/packages/zarr-indexing/tests/test_output_map.py b/packages/zarr-indexing/tests/test_output_map.py index 498101444e..2efe26e713 100644 --- a/packages/zarr-indexing/tests/test_output_map.py +++ b/packages/zarr-indexing/tests/test_output_map.py @@ -1,6 +1,9 @@ from __future__ import annotations +import pickle + import numpy as np +import pytest from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap @@ -54,3 +57,26 @@ def test_frozen(self) -> None: arr = np.array([0], dtype=np.intp) m = ArrayMap(index_array=arr) assert isinstance(m, ArrayMap) + + def test_owns_index_array_and_keeps_hash_stable(self) -> None: + arr = np.array([1, 3, 5], dtype=np.intp) + m = ArrayMap(index_array=arr) + lookup = {m: "value"} + + arr[:] = 9 + + np.testing.assert_array_equal(m.index_array, [1, 3, 5]) + assert lookup[m] == "value" + + def test_pickle_round_trip_keeps_index_array_read_only(self) -> None: + restored = pickle.loads(pickle.dumps(ArrayMap(index_array=np.array([1, 3, 5])))) + + assert not restored.index_array.flags.writeable + with pytest.raises(ValueError, match="read-only"): + restored.index_array[0] = 9 + + def test_index_array_cannot_be_made_writeable(self) -> None: + m = ArrayMap(index_array=np.array([1, 3, 5])) + + with pytest.raises(ValueError): + m.index_array.flags.writeable = True From 2018fd4a718f6b1763567d0901941e96d3dad889 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 11:31:28 +0200 Subject: [PATCH 27/32] fix(zarr-indexing): keep an empty masked result masked whichever parts are in force MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An empty view is now answered without reading the source, and the shortcut reached for the array namespace's own `empty`, which knows nothing about masks. An unpartitioned empty view over a masked source therefore came back a plain array while the same view partitioned came back masked — no cells either way, so no value changed, but the caller's type depended on how the read had been divided, which `result()` promises it never does. A masked source goes through `_output_buffer`, which is the branch that knows. Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/lazy_array.py | 13 +++++++++++-- packages/zarr-indexing/tests/test_lazy_array.py | 17 +++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 7caeb51fc1..6833e878d2 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -1670,8 +1670,17 @@ def result(self) -> Any: return out def _empty_result(self, out_shape: tuple[int, ...]) -> Any: - """Allocate an empty result without asking the source for any data.""" - if self._parts is None: + """Allocate an empty result without asking the source for any data. + + A masked source is answered from `_output_buffer`, which knows to build + a masked buffer, whatever partitioning is in force. Reaching for the + namespace's own `empty` first would hand back a plain array on the + unpartitioned path and a masked one on the partitioned path — the same + answer differing by how the read was divided, which `result()` promises + it never does. There are no cells either way, so this is about the type + the caller gets back rather than about any value. + """ + if self._parts is None and not isinstance(self._array, np.ma.MaskedArray): xp = _namespace(self._array) empty = getattr(xp, "empty", None) if empty is not None: diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 17f1297835..432d4c1187 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2061,6 +2061,23 @@ def test_a_masked_source_keeps_its_mask_under_every_partitioning(parts: Any) -> np.testing.assert_array_equal(np.ma.filled(got, 0), np.ma.filled(expected, 0)) +@pytest.mark.parametrize("parts", [None, (2, 2), (3, 4)]) +def test_a_masked_source_keeps_its_mask_when_the_view_is_empty(parts: Any) -> None: + """An empty result is still a result, and its type must not depend on the parts. + + An empty view is answered without reading the source at all, and that + shortcut reached for the array namespace's own `empty` — which knows nothing + about masks — so an unpartitioned empty view came back a plain array while + the same view partitioned came back masked. No cells either way, so nothing + about the values changed; the caller just got a different type depending on + how the read had been divided. + """ + data = np.ma.masked_greater(np.arange(12).reshape(3, 4), 7) + got = repartition(LazyArray(data), parts).lazy[:, 2:2].result() + assert isinstance(got, np.ma.MaskedArray), parts + assert np.asarray(got).shape == (3, 0), parts + + def test_a_look_alike_enum_declaration_is_honored_not_upgraded() -> None: """xarray's `IndexingSupport` has these member names; it must not read as absent. From a8552695a034bdab94d3a7e2e25a8c1c81abdbf4 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 13:14:42 +0200 Subject: [PATCH 28/32] docs(indexing): design chunk projection API Assisted-by: Codex:gpt-5.6 --- .../2026-08-01-chunk-projection-design.md | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 docs/superpowers/specs/2026-08-01-chunk-projection-design.md diff --git a/docs/superpowers/specs/2026-08-01-chunk-projection-design.md b/docs/superpowers/specs/2026-08-01-chunk-projection-design.md new file mode 100644 index 0000000000..452bad55ee --- /dev/null +++ b/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. From 8cc1b4bc3eb537afe5e321b89581187e2cbb4ba6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 13:21:44 +0200 Subject: [PATCH 29/32] feat(indexing): add reusable chunk plans Assisted-by: Codex:gpt-5.6 --- .../src/zarr_indexing/__init__.py | 8 ++ .../src/zarr_indexing/chunk_resolution.py | 130 +++++++++++++++++- .../tests/test_chunk_resolution.py | 34 ++++- 3 files changed, 170 insertions(+), 2 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 1ef089d4b2..e097f10a22 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -27,7 +27,11 @@ from importlib.metadata import version from zarr_indexing.chunk_resolution import ( + ChunkCoverage, + ChunkPlan, + ChunkProjection, iter_chunk_transforms, + plan_chunks, sub_transform_to_selections, ) from zarr_indexing.composition import compose @@ -68,6 +72,9 @@ __all__ = [ "ArrayMap", "BoundsCheckError", + "ChunkCoverage", + "ChunkPlan", + "ChunkProjection", "ConstantMap", "DimensionGridLike", "DimensionMap", @@ -95,6 +102,7 @@ "iter_chunk_transforms", "normalize_ndsel", "parse_ndsel", + "plan_chunks", "resolve_indexing_support", "selection_to_transform", "sub_transform_to_selections", diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index def3cccf18..3ebf58701b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -42,7 +42,8 @@ from __future__ import annotations import math -from typing import TYPE_CHECKING, Any +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Literal import numpy as np @@ -75,6 +76,73 @@ OutIndices, ] +type ChunkCoverage = Literal["full", "partial", "unknown"] + + +@dataclass(frozen=True, slots=True) +class ChunkProjection: + """One source-independent projection of a request through a chunk. + + Both transforms share a synthetic input domain. ``chunk_transform`` maps + that domain to chunk-local storage coordinates; ``cell_transform`` maps it + to the original request domain. + """ + + chunk_coords: tuple[int, ...] + chunk_domain: IndexDomain + chunk_transform: IndexTransform + cell_transform: IndexTransform + coverage: ChunkCoverage + + def __post_init__(self) -> None: + if self.chunk_transform.domain != self.cell_transform.domain: + raise ValueError( + "chunk_transform and cell_transform must share an input domain; " + f"got {self.chunk_transform.domain!r} and {self.cell_transform.domain!r}" + ) + + +@dataclass(frozen=True, slots=True) +class ChunkPlan: + """A reusable, lazy partition of an index transform over a chunk grid.""" + + transform: IndexTransform + dimension_grids: tuple[DimensionGridLike, ...] + + def projections(self) -> Iterator[ChunkProjection]: + """Return a fresh iterator over the chunks touched by this plan.""" + return _iter_chunk_projections(self.transform, self.dimension_grids) + + def __iter__(self) -> Iterator[ChunkProjection]: + return self.projections() + + +def plan_chunks( + transform: IndexTransform, + dimension_grids: Sequence[DimensionGridLike], +) -> ChunkPlan: + """Plan a transform against a caller-selected chunk grid. + + Parameters + ---------- + transform + Mapping from the request domain to storage coordinates. + dimension_grids + One storage grid per transform output dimension. + + Returns + ------- + ChunkPlan + A reusable plan whose projections are computed lazily. + """ + grids = tuple(dimension_grids) + if len(grids) != transform.output_rank: + raise ValueError( + "dimension_grids must have one entry per transform output dimension; " + f"got {len(grids)} grids for output rank {transform.output_rank}" + ) + return ChunkPlan(transform=transform, dimension_grids=grids) + def _one_dimensional_correlated_array_map( transform: IndexTransform, @@ -291,6 +359,66 @@ def iter_chunk_transforms( yield (chunk_coords, local, surviving) +def _covers_whole_chunk(transform: IndexTransform, chunk_shape: tuple[int, ...]) -> bool: + """Whether an affine chunk-local transform addresses every chunk cell.""" + domain = transform.domain + for out_dim, m in enumerate(transform.output): + extent = chunk_shape[out_dim] + if isinstance(m, ConstantMap): + if extent != 1 or m.offset != 0: + return False + elif isinstance(m, DimensionMap): + if abs(m.stride) != 1: + return False + lo = domain.inclusive_min[m.input_dimension] + hi = domain.exclusive_max[m.input_dimension] + if hi <= lo: + if extent != 0: + return False + continue + first = m.offset + m.stride * lo + last = m.offset + m.stride * (hi - 1) + if min(first, last) != 0 or max(first, last) != extent - 1: + return False + else: + return False + return True + + +def _iter_chunk_projections( + transform: IndexTransform, + dim_grids: Sequence[DimensionGridLike], +) -> Iterator[ChunkProjection]: + """Convert private intersection results into public paired projections.""" + for chunk_coords, chunk_transform, survivors in iter_chunk_transforms(transform, dim_grids): + chunk_min = tuple( + grid.chunk_offset(coord) for grid, coord in zip(dim_grids, chunk_coords, strict=True) + ) + chunk_shape = tuple( + grid.chunk_size(coord) for grid, coord in zip(dim_grids, chunk_coords, strict=True) + ) + chunk_domain = IndexDomain( + inclusive_min=chunk_min, + exclusive_max=tuple( + origin + extent for origin, extent in zip(chunk_min, chunk_shape, strict=True) + ), + ) + cell_transform = IndexTransform.identity(chunk_transform.domain) + if survivors is not None or any(isinstance(m, ArrayMap) for m in chunk_transform.output): + coverage: ChunkCoverage = "unknown" + elif _covers_whole_chunk(chunk_transform, chunk_shape): + coverage = "full" + else: + coverage = "partial" + yield ChunkProjection( + chunk_coords=chunk_coords, + chunk_domain=chunk_domain, + chunk_transform=chunk_transform, + cell_transform=cell_transform, + coverage=coverage, + ) + + def _dimension_map_slice(m: DimensionMap, dim_lo: int, dim_hi: int) -> slice: """The NumPy slice a `DimensionMap` reads over input range `[dim_lo, dim_hi)`. diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index f7b4815a96..32f73d8a80 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -6,7 +6,7 @@ import pytest from zarr.core.chunk_grids import ChunkGrid, FixedDimension, VaryingDimension -from zarr_indexing import chunk_resolution +from zarr_indexing import ChunkPlan, ChunkProjection, chunk_resolution, plan_chunks from zarr_indexing.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections from zarr_indexing.domain import IndexDomain from zarr_indexing.grid import dimension_grids_from_chunks @@ -14,6 +14,38 @@ from zarr_indexing.transform import IndexTransform +def test_basic_plan_is_reiterable_and_projects_both_spaces() -> None: + """A plan can be revisited without losing either side of each projection.""" + transform = IndexTransform.from_shape((6,))[1:6] + grids = dimension_grids_from_chunks((3,), (6,)) + + plan = plan_chunks(transform, grids) + first = list(plan) + second = list(plan.projections()) + + assert isinstance(plan, ChunkPlan) + assert all(isinstance(projection, ChunkProjection) for projection in first) + assert [projection.chunk_coords for projection in first] == [(0,), (1,)] + assert [projection.chunk_domain for projection in first] == [ + IndexDomain((0,), (3,)), + IndexDomain((3,), (6,)), + ] + assert [projection.coverage for projection in first] == ["partial", "full"] + assert first == second + assert all( + projection.chunk_transform.domain == projection.cell_transform.domain + for projection in first + ) + + +def test_plan_rejects_grid_rank_different_from_transform_output_rank() -> None: + """A missing storage grid dimension is rejected before iteration.""" + transform = IndexTransform.from_shape((2, 3)) + + with pytest.raises(ValueError, match="1 grids for output rank 2"): + plan_chunks(transform, dimension_grids_from_chunks((2,), (2,))) + + class TestChunkResolutionIdentity: def test_single_chunk(self) -> None: """Array fits in one chunk.""" From 8cdb1e58e28aaa6987454a9847fe91452ca07457 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 13:24:59 +0200 Subject: [PATCH 30/32] feat(indexing): project chunks through paired transforms Assisted-by: Codex:gpt-5.6 --- .../src/zarr_indexing/chunk_resolution.py | 91 ++++++++++++++++++- .../tests/test_chunk_resolution.py | 80 ++++++++++++++++ 2 files changed, 170 insertions(+), 1 deletion(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 3ebf58701b..8b1a9479a8 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -385,6 +385,95 @@ def _covers_whole_chunk(transform: IndexTransform, chunk_shape: tuple[int, ...]) return True +def _orthogonal_cell_transform( + original: IndexTransform, + restricted: IndexTransform, + survivors: dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]], +) -> IndexTransform: + """Map a compacted orthogonal intersection back to request coordinates.""" + by_input_dimension: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + if isinstance(survivors, dict): + survivor_items = survivors.items() + else: + array_output_dimensions = [ + output_dimension + for output_dimension, output_map in enumerate(original.output) + if isinstance(output_map, ArrayMap) + ] + if len(array_output_dimensions) != 1: + raise ValueError( + "one survivor array requires exactly one orthogonal ArrayMap; " + f"found output dimensions {array_output_dimensions}" + ) + survivor_items = ((array_output_dimensions[0], survivors),) + + for output_dimension, positions in survivor_items: + output_map = original.output[output_dimension] + if not isinstance(output_map, ArrayMap): + raise TypeError( + f"survivors for output dimension {output_dimension} do not describe an ArrayMap" + ) + input_dimension = array_map_dependent_axis(output_map) + if input_dimension is None: + raise ValueError( + f"output dimension {output_dimension} has no orthogonal input dimension" + ) + by_input_dimension[input_dimension] = np.asarray(positions, dtype=np.intp) + + output: list[ConstantMap | DimensionMap | ArrayMap] = [] + rank = original.input_rank + for input_dimension in range(rank): + positions = by_input_dimension.get(input_dimension) + if positions is None: + output.append(DimensionMap(input_dimension=input_dimension)) + continue + shape = (1,) * input_dimension + (positions.size,) + (1,) * (rank - input_dimension - 1) + output.append( + ArrayMap( + index_array=positions.reshape(shape), + offset=original.domain.inclusive_min[input_dimension], + input_dimension=input_dimension, + ) + ) + return IndexTransform(domain=restricted.domain, output=tuple(output)) + + +def _correlated_cell_transform( + original: IndexTransform, + restricted: IndexTransform, + survivors: np.ndarray[Any, np.dtype[np.intp]], +) -> IndexTransform: + """Map compacted correlated points back through the request's row-major domain.""" + positions = np.asarray(survivors, dtype=np.intp) + coordinates = np.unravel_index(positions, original.domain.shape) + output = tuple( + ArrayMap( + index_array=np.asarray(coordinate, dtype=np.intp), + offset=origin, + ) + for coordinate, origin in zip(coordinates, original.domain.inclusive_min, strict=True) + ) + return IndexTransform(domain=restricted.domain, output=output) + + +def _cell_transform( + original: IndexTransform, + restricted: IndexTransform, + survivors: OutIndices, +) -> IndexTransform: + """Convert private survivor bookkeeping into a direction-neutral transform.""" + if survivors is None: + return IndexTransform.identity(restricted.domain) + if any( + isinstance(output_map, ArrayMap) and output_map.input_dimension is None + for output_map in original.output + ): + if isinstance(survivors, dict): + raise ValueError("correlated intersections require one shared survivor array") + return _correlated_cell_transform(original, restricted, survivors) + return _orthogonal_cell_transform(original, restricted, survivors) + + def _iter_chunk_projections( transform: IndexTransform, dim_grids: Sequence[DimensionGridLike], @@ -403,7 +492,7 @@ def _iter_chunk_projections( origin + extent for origin, extent in zip(chunk_min, chunk_shape, strict=True) ), ) - cell_transform = IndexTransform.identity(chunk_transform.domain) + cell_transform = _cell_transform(transform, chunk_transform, survivors) if survivors is not None or any(isinstance(m, ArrayMap) for m in chunk_transform.output): coverage: ChunkCoverage = "unknown" elif _covers_whole_chunk(chunk_transform, chunk_shape): diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index 32f73d8a80..d71e0bf96c 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -46,6 +46,86 @@ def test_plan_rejects_grid_rank_different_from_transform_output_rank() -> None: plan_chunks(transform, dimension_grids_from_chunks((2,), (2,))) +def _storage_of(transform: IndexTransform, point: tuple[int, ...]) -> tuple[int, ...]: + """Evaluate the three map forms at one point, independently of planning.""" + result: list[int] = [] + for output_map in transform.output: + if isinstance(output_map, ConstantMap): + result.append(output_map.offset) + elif isinstance(output_map, DimensionMap): + result.append(output_map.offset + output_map.stride * point[output_map.input_dimension]) + else: + index = tuple( + 0 + if output_map.index_array.shape[axis] == 1 + else point[axis] - transform.domain.inclusive_min[axis] + for axis in range(output_map.index_array.ndim) + ) + result.append( + output_map.offset + output_map.stride * int(output_map.index_array[index]) + ) + return tuple(result) + + +@pytest.mark.parametrize( + ("transform", "grids"), + [ + ( + IndexTransform.from_shape((6,)).oindex[np.array([4, 0, 4, 2])], + dimension_grids_from_chunks((3,), (6,)), + ), + ( + IndexTransform.from_shape((4, 5)).oindex[np.array([3, 0]), np.array([4, 1, 1])], + dimension_grids_from_chunks(((1, 3), (2, 3)), (4, 5)), + ), + ( + IndexTransform.from_shape((2, 4, 5)).vindex[ + ..., np.array([3, 0, 3]), np.array([4, 1, 1]) + ], + dimension_grids_from_chunks((1, 2, 3), (2, 4, 5)), + ), + ], + ids=["repeated-oindex", "irregular-oindex", "vindex-with-residual"], +) +def test_projection_invariants_for_fancy_selections( + transform: IndexTransform, grids: tuple[Any, ...] +) -> None: + """Both transforms agree pointwise and cell ranges tile request space once.""" + plan = plan_chunks(transform, grids) + request_points: list[tuple[int, ...]] = [] + + for projection in plan: + assert projection.coverage == "unknown" + cell_domain = projection.cell_transform.domain + for positional_point in np.ndindex(*cell_domain.shape): + cell_point = tuple( + coordinate + origin + for coordinate, origin in zip( + positional_point, cell_domain.inclusive_min, strict=True + ) + ) + request_point = _storage_of(projection.cell_transform, cell_point) + chunk_point = _storage_of(projection.chunk_transform, cell_point) + storage_point = _storage_of(plan.transform, request_point) + chunk_origin = projection.chunk_domain.inclusive_min + + assert chunk_point == tuple( + value - origin for value, origin in zip(storage_point, chunk_origin, strict=True) + ) + request_points.append(request_point) + + expected_request_points = [ + tuple( + coordinate + origin + for coordinate, origin in zip( + positional_point, transform.domain.inclusive_min, strict=True + ) + ) + for positional_point in np.ndindex(*transform.domain.shape) + ] + assert sorted(request_points) == sorted(expected_request_points) + + class TestChunkResolutionIdentity: def test_single_chunk(self) -> None: """Array fits in one chunk.""" From 86db502e83ccde15402119cd05fc856a78df5bd3 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 13:29:20 +0200 Subject: [PATCH 31/32] refactor(indexing): build lazy parts from projections Assisted-by: Codex:gpt-5.6 --- .../src/zarr_indexing/chunk_resolution.py | 3 + .../src/zarr_indexing/lazy_array.py | 155 +++++++++--------- .../zarr-indexing/tests/test_lazy_array.py | 69 ++++++++ 3 files changed, 154 insertions(+), 73 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 8b1a9479a8..82ef05bf33 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -493,6 +493,9 @@ def _iter_chunk_projections( ), ) cell_transform = _cell_transform(transform, chunk_transform, survivors) + synthetic_origin = (0,) * chunk_transform.input_rank + chunk_transform = chunk_transform.translate_domain_to(synthetic_origin) + cell_transform = cell_transform.translate_domain_to(synthetic_origin) if survivors is not None or any(isinstance(m, ArrayMap) for m in chunk_transform.output): coverage: ChunkCoverage = "unknown" elif _covers_whole_chunk(chunk_transform, chunk_shape): diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 6833e878d2..7e95eb73d2 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -169,8 +169,8 @@ split_scalar_axes, ) from zarr_indexing.chunk_resolution import ( - iter_chunk_transforms, - sub_transform_to_selections, + ChunkProjection, + plan_chunks, ) from zarr_indexing.grid import EdgeDimensionGrid, dimension_grids_from_chunks from zarr_indexing.json import transform_to_canonical @@ -846,33 +846,71 @@ def _detach(value: Any, source: Any) -> Any: def _partition_out_selection( - sub_transform: IndexTransform, - out_indices: Any, - out_shape: tuple[int, ...], + cell_transform: IndexTransform, ) -> tuple[Any, ...]: - """Where a partition's values belong in an array of shape `out_shape`. - - A correlated sub-transform scatters through flat offsets into the row-major - result; unravelling them turns that into an index tuple for the result's own - shape, so callers never need a flat working buffer. - - The correlated scatter is taken directly rather than through - `sub_transform_to_selections`, whose business is to match a scatter to the - block a *NumPy chunk read* returns — which is not the block a part produces. - A part is resolved by lowering its own transform, and that lowering keeps the - points axis first; the bridge has to answer for NumPy's placement rules - instead, and permutes its scatter to suit. Reading the bridge's answer here - would apply a correction for a block this path never sees. - """ - if not _is_correlated(sub_transform): - return sub_transform_to_selections(sub_transform, out_indices)[1] - if len(out_shape) == 0: - return () - if out_indices is None: - flat = np.arange(math.prod(sub_transform.domain.shape), dtype=np.intp) - else: - flat = np.asarray(out_indices, dtype=np.intp) - return tuple(np.unravel_index(flat, out_shape)) + """Lower ``cell_transform`` to NumPy selectors on the request buffer.""" + domain = cell_transform.domain + if _is_correlated(cell_transform): + selectors: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + for output_map in cell_transform.output: + if isinstance(output_map, ConstantMap): + coordinates = np.full(domain.shape, output_map.offset, dtype=np.intp) + elif isinstance(output_map, DimensionMap): + input_dimension = output_map.input_dimension + axis = np.arange( + domain.inclusive_min[input_dimension], + domain.exclusive_max[input_dimension], + dtype=np.intp, + ) + shape = ( + (1,) * input_dimension + + (axis.size,) + + ((1,) * (domain.ndim - input_dimension - 1)) + ) + coordinates = np.broadcast_to(axis.reshape(shape), domain.shape) + coordinates = output_map.offset + output_map.stride * coordinates + else: + coordinates = output_map.offset + output_map.stride * np.broadcast_to( + output_map.index_array, domain.shape + ) + selectors.append(np.asarray(coordinates, dtype=np.intp)) + return tuple(selectors) + + selectors: list[int | slice | np.ndarray[Any, np.dtype[np.intp]]] = [] + n_array_maps = sum(isinstance(output_map, ArrayMap) for output_map in cell_transform.output) + for output_map in cell_transform.output: + if isinstance(output_map, ConstantMap): + selectors.append(output_map.offset) + elif isinstance(output_map, DimensionMap): + input_dimension = output_map.input_dimension + lo = domain.inclusive_min[input_dimension] + hi = domain.exclusive_max[input_dimension] + selectors.append( + slice( + output_map.offset + output_map.stride * lo, + output_map.offset + output_map.stride * hi, + output_map.stride, + ) + ) + else: + selectors.append( + (output_map.offset + output_map.stride * output_map.index_array.ravel()).astype( + np.intp + ) + ) + if n_array_maps > 1: + axes = [ + np.asarray([selector], dtype=np.intp) + if isinstance(selector, int) + else ( + np.arange(selector.start, selector.stop, selector.step, dtype=np.intp) + if isinstance(selector, slice) + else selector + ) + for selector in selectors + ] + return np.ix_(*axes) + return tuple(selectors) def _out_selection_cell_count(selection: tuple[Any, ...], out_shape: tuple[int, ...]) -> int: @@ -910,44 +948,6 @@ def _out_selection_cell_count(selection: tuple[Any, ...], out_shape: tuple[int, return total -def _covers_whole_part(local_transform: IndexTransform, part_shape: tuple[int, ...]) -> bool: - """Whether a part-local transform addresses every cell of its part. - - Direction does not matter: a reversing view (`lazy[::-1]`) reads every cell - of the box, just back to front, so a stride of -1 covers a part exactly as a - stride of 1 does. Only the set of coordinates is asked about here, which is - the same question `bounding_box()` and `strides()` answer. - - Conservative in one place: an `ArrayMap` could in principle enumerate a whole - part, but checking that costs more than the answer is worth, so a - fancy-indexed axis always reports incomplete. - """ - domain = local_transform.domain - for out_dim, m in enumerate(local_transform.output): - extent = part_shape[out_dim] - if isinstance(m, ConstantMap): - if extent != 1 or m.offset != 0: - return False - elif isinstance(m, DimensionMap): - if abs(m.stride) != 1: - return False - d = m.input_dimension - lo = domain.inclusive_min[d] - hi = domain.exclusive_max[d] - if hi <= lo: - # An empty walk covers only an empty part. - if extent != 0: - return False - continue - first = m.offset + m.stride * lo - last = m.offset + m.stride * (hi - 1) - if min(first, last) != 0 or max(first, last) != extent - 1: - return False - else: - return False - return True - - @dataclass(frozen=True, kw_only=True) class Partition: """One box of a `LazyArray`'s partitioning, as it falls through the view. @@ -987,11 +987,20 @@ class Partition: does not. """ - base_coords: tuple[int, ...] + projection: ChunkProjection box: tuple[tuple[int, int], ...] view: LazyArray out_selection: tuple[Any, ...] - is_complete: bool + + @property + def base_coords(self) -> tuple[int, ...]: + """Coordinates of this partition in the selected base grid.""" + return self.projection.chunk_coords + + @property + def is_complete(self) -> bool: + """Whether the projection proves it covers the entire selected cell.""" + return self.projection.coverage == "full" # --------------------------------------------------------------------------- # @@ -1535,10 +1544,11 @@ def parts(self) -> Iterator[Partition]: """ base_shape = self._base_shape grids = self._parts if self._parts is not None else _whole_array_grids(base_shape) - out_shape = self.shape rank = len(base_shape) - for base_coords, local, out_indices in iter_chunk_transforms(self._transform, grids): + for projection in plan_chunks(self._transform, grids): + base_coords = projection.chunk_coords + local = projection.chunk_transform origin = tuple(grid.chunk_offset(c) for grid, c in zip(grids, base_coords, strict=True)) extent = tuple(grid.chunk_size(c) for grid, c in zip(grids, base_coords, strict=True)) if origin == (0,) * rank and extent == base_shape: @@ -1562,11 +1572,10 @@ def parts(self) -> Iterator[Partition]: w.start + o for w, o in zip(self._window, origin, strict=True) ) yield Partition( - base_coords=base_coords, + projection=projection, box=tuple((o, o + e) for o, e in zip(global_origin, extent, strict=True)), view=LazyArray._derive(self._array, local, None, window, self._support), - out_selection=_partition_out_selection(local, out_indices, out_shape), - is_complete=_covers_whole_part(local, extent), + out_selection=_partition_out_selection(projection.cell_transform), ) # -- indexing ----------------------------------------------------------- diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 432d4c1187..79334f7be6 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -20,6 +20,7 @@ from zarr_indexing import ( ArrayMap, + ChunkProjection, ConstantMap, DimensionMap, EdgeDimensionGrid, @@ -1207,6 +1208,74 @@ def test_parts_tile_the_view_exactly_and_disjointly( np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64)) +def _transform_point(transform: IndexTransform, point: tuple[int, ...]) -> tuple[int, ...]: + """Evaluate a transform pointwise for projection placement assertions.""" + result: list[int] = [] + for output_map in transform.output: + if isinstance(output_map, ConstantMap): + result.append(output_map.offset) + elif isinstance(output_map, DimensionMap): + result.append(output_map.offset + output_map.stride * point[output_map.input_dimension]) + else: + index = tuple( + 0 + if output_map.index_array.shape[axis] == 1 + else point[axis] - transform.domain.inclusive_min[axis] + for axis in range(output_map.index_array.ndim) + ) + result.append( + output_map.offset + output_map.stride * int(output_map.index_array[index]) + ) + return tuple(result) + + +@pytest.mark.parametrize( + "build", + [ + lambda array: array.lazy.oindex[[6, 0, 2], :, [3, 1]], + lambda array: array.lazy.vindex[..., np.array([4, 0, 4]), np.array([3, 1, 1])], + ], + ids=["orthogonal", "vectorized"], +) +def test_partition_exposes_projection_as_its_placement_authority( + build: Callable[[LazyArray], LazyArray], +) -> None: + """A part's NumPy placement addresses exactly its cell-transform range.""" + view = build(LazyArray(reference()).with_parts(PART_SHAPE)) + assembled = np.zeros(view.shape, dtype=view.dtype) + + for part in view.parts(): + projection = part.projection + assert isinstance(projection, ChunkProjection) + assert part.view.transform == projection.chunk_transform + assert part.base_coords == projection.chunk_coords + assert part.box == tuple( + zip( + projection.chunk_domain.inclusive_min, + projection.chunk_domain.exclusive_max, + strict=True, + ) + ) + + expected_hits = np.zeros(view.shape, dtype=np.int8) + cell_domain = projection.cell_transform.domain + for positional_point in np.ndindex(*cell_domain.shape): + cell_point = tuple( + coordinate + origin + for coordinate, origin in zip( + positional_point, cell_domain.inclusive_min, strict=True + ) + ) + expected_hits[_transform_point(projection.cell_transform, cell_point)] = 1 + actual_hits = np.zeros(view.shape, dtype=np.int8) + actual_hits[part.out_selection] = 1 + np.testing.assert_array_equal(actual_hits, expected_hits) + + assembled[part.out_selection] = np.asarray(part.view.result()) + + np.testing.assert_array_equal(assembled, np.asarray(view.result())) + + def test_parts_resolve_independently_and_concurrently() -> None: """Each part's `array` is a standalone `LazyArray` with no shared mutable state.""" view = make_source("zarr").lazy[1:7, :, 1:].with_parts((2, 2, 2)) From 3041e33206b33005e2ef58bad59aadef924e211d Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sat, 1 Aug 2026 13:42:05 +0200 Subject: [PATCH 32/32] refactor(indexing): expose projection-only chunk planning Assisted-by: Codex:gpt-5.6 --- packages/zarr-indexing/README.md | 3 + packages/zarr-indexing/changes/267.feature.md | 4 +- packages/zarr-indexing/changes/272.bugfix.md | 2 +- packages/zarr-indexing/docs/api/index.md | 5 +- packages/zarr-indexing/docs/design-notes.md | 18 + packages/zarr-indexing/docs/index.md | 17 +- .../src/zarr_indexing/__init__.py | 13 +- .../src/zarr_indexing/chunk_resolution.py | 321 +------ .../src/zarr_indexing/lazy_array.py | 32 +- .../tests/test_chunk_resolution.py | 838 +++++------------- .../zarr-indexing/tests/test_lazy_array.py | 2 +- 11 files changed, 311 insertions(+), 944 deletions(-) diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index c3ab8d085a..f1285cd166 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -17,6 +17,9 @@ Key types: `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 diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index fcf9bcf48d..6376f0db6f 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -21,7 +21,7 @@ Fixed an `oindex`/`vindex` step whose entries are all slices applied to a view t 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 the chunk selection `sub_transform_to_selections` builds for a negative-stride dimension: 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 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. @@ -35,3 +35,5 @@ A new `zarr_indexing.testing` subpackage, behind a `testing` extra (`pip install 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. diff --git a/packages/zarr-indexing/changes/272.bugfix.md b/packages/zarr-indexing/changes/272.bugfix.md index 807447ee38..0717e1e384 100644 --- a/packages/zarr-indexing/changes/272.bugfix.md +++ b/packages/zarr-indexing/changes/272.bugfix.md @@ -2,7 +2,7 @@ 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; `sub_transform_to_selections` described a +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 diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index 581f0cdb50..b206c9c808 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -23,9 +23,8 @@ and the wire format built on top of it. **Chunk resolution** - [`zarr_indexing.chunk_resolution`](chunk_resolution.md) — - `iter_chunk_transforms` (transform + chunk grid → per-chunk transforms) and - `sub_transform_to_selections` (the bridge back to the selection tuples the - current codec pipeline expects) + `plan_chunks`, which lazily projects a request through a caller-selected grid, + plus the reusable `ChunkPlan` and paired-transform `ChunkProjection` values - [`zarr_indexing.grid`](grid.md) — `DimensionGridLike`, the Protocol describing the narrow chunk-grid surface chunk resolution consumes, so that nothing here imports `zarr`, plus `EdgeDimensionGrid` and diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index cc4a60c69c..c52afcac94 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -58,6 +58,24 @@ Four deliberate differences: | Wire format | Implementation-defined JSON, specified by what the implementation accepts | [ndsel](ndsel.md) is spec-first, with a vendored language-agnostic conformance corpus every implementation runs | | Backends | A driver ecosystem (zarr, N5, neuroglancer, GCS, …) built into the library | No drivers. The wrapper takes anything with `shape`, `dtype`, and `__getitem__`, and prefers the array's own `__array_namespace__` | +Chunk planning deliberately follows TensorStore's paired-transform boundary. +TensorStore exposes a transform from a shared cell domain into a decoded chunk +and another into the request; `ChunkProjection.chunk_transform` and +`ChunkProjection.cell_transform` preserve the same directionality. This is more +general than returning a NumPy read key plus scatter indices: slices, outer +products, and correlated gathers remain ordinary transforms, and a consumer can +lower them to its own execution vocabulary. + +The ownership boundary differs. `plan_chunks` retains only the logical request +and a caller-supplied grid, and lazily yields source-independent projections. It +does not own reads, writes, buffers, locks, or scheduling. Zarr can therefore +plan reads against an inner codec-chunk grid and writes against an atomic shard +grid; napari or dask can turn the same projections into tasks without putting a +dask dependency in this package. `coverage` is relative to that selected grid: +`full` proves a blind replacement safe, `partial` proves it is not, and +`unknown` conservatively covers fancy selections whose duplicates would require +additional work to classify. + The comparison also runs the other way. TensorStore is a mature, heavily optimized C++ system whose performance this library cannot approach: resolution here is Python-level bookkeeping over NumPy, and the per-part overhead is diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index 9bb0ab1b74..8b08dcd0fd 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -119,7 +119,7 @@ works: ```python from dataclasses import dataclass -from zarr_indexing import iter_chunk_transforms +from zarr_indexing import plan_chunks @dataclass(frozen=True) @@ -134,16 +134,19 @@ class RegularDimensionGrid: grids = [RegularDimensionGrid(32), RegularDimensionGrid(32)] -for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(view, grids): - print(chunk_coords, sub_transform.selection_repr) +for projection in plan_chunks(view, grids): + print(projection.chunk_coords, projection.chunk_transform.selection_repr) # (0, 0) { [10, 32), 5 } # (1, 0) { [0, 18), 5 } ``` -Each yielded `sub_transform` is the original transform restricted to one chunk -and translated into chunk-local coordinates, which is 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. +`plan_chunks` returns a lazy, reusable plan. Each projection has two transforms +over the same compact domain: `chunk_transform` locates cells in chunk-local +storage, while `cell_transform` places those cells in the original request. +This paired representation works unchanged for slices, orthogonal indexing, and +vectorized indexing. `chunk_domain` describes the selected grid cell in global +storage coordinates, and `coverage` tells a writer whether full replacement is +proven safe. ## Lazy views over any array diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index e097f10a22..4dacbbb7c1 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -16,11 +16,10 @@ `LazyArray` wraps an array-API-like array and gives it deferred indexing through `.lazy[...]`, yielding its reads as `Partition`s. -The chunk-resolution helpers (`iter_chunk_transforms`, -`sub_transform_to_selections`) and `selection_to_transform` are also exported -here: they form the surface the zarr integration layer (array indexing) depends -on — see `chunk_resolution` on why the second of those is provisional. The -`DimensionGridLike` Protocol describes the chunk-grid surface chunk resolution +`plan_chunks` projects a transform through a caller-selected chunk grid without +coupling the result to a storage backend or scheduler. `selection_to_transform` +is also exported for consumers starting with a NumPy-style selection. The +`DimensionGridLike` Protocol describes the narrow grid surface chunk resolution consumes without importing zarr. """ @@ -30,9 +29,7 @@ ChunkCoverage, ChunkPlan, ChunkProjection, - iter_chunk_transforms, plan_chunks, - sub_transform_to_selections, ) from zarr_indexing.composition import compose from zarr_indexing.domain import IndexDomain @@ -99,13 +96,11 @@ "index_transform_from_json", "index_transform_to_json", "infer_indexing_support", - "iter_chunk_transforms", "normalize_ndsel", "parse_ndsel", "plan_chunks", "resolve_indexing_support", "selection_to_transform", - "sub_transform_to_selections", "transform_from_canonical", "transform_to_canonical", ] diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 82ef05bf33..669e1a5291 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -18,30 +18,21 @@ 3. **Translate** — shift the restricted transform to chunk-local coordinates via `transform.translate(-chunk_origin)`. -4. **Yield** — produce `(chunk_coords, local_transform, surviving_indices)` - triples that the codec pipeline consumes. +4. **Project** — pair the chunk-local storage transform with a transform back + to the request's cells. Both use the same compact, zero-origin domain. Sorted one-dimensional correlated array maps can be partitioned directly because every touched chunk owns a contiguous slice of the index array. That case bypasses candidate enumeration and repeated intersection. -`sub_transform_to_selections` bridges from the transform representation -back to the raw `(chunk_selection, out_selection, drop_axes)` tuples that -the current codec pipeline expects. - -That bridge is **provisional**. Its return shape is a codec pipeline's -internal vocabulary rather than anything this package wants to keep saying, and -it is expected to go away once a pipeline accepts transforms natively. It is -exported because zarr's read path consumes it today; it is not covered by -whatever compatibility promise the rest of this API carries, and a consumer -outside zarr should expect to follow it through a change. The contract while it -exists is exactly `out[out_selection] = chunk[chunk_selection]`, with the axes -named in `drop_axes` squeezed out of the block first. +The public result is a lazy, reusable `ChunkPlan`. Each `ChunkProjection` is +source-independent: it identifies the chunk and expresses both sides of the +gather without assuming NumPy selectors, a codec pipeline, or an execution +scheduler. """ from __future__ import annotations -import math from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal @@ -49,31 +40,21 @@ from zarr_indexing.domain import IndexDomain from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap -from zarr_indexing.transform import ( - IndexTransform, - _positional_slice, # pyright: ignore[reportPrivateUsage] - array_map_dependent_axis, -) - -# `_positional_slice` is a leading-underscore helper in `transform.py`, shared -# with this module on purpose: the slice an `ArrayMap` axis is reindexed by and -# the slice a `DimensionMap` is lowered to are the same walk, and the two must -# agree on where a downward one stops. See `json.py` for the same suppression -# and the same open question about promoting these helpers out of "private". +from zarr_indexing.transform import IndexTransform, array_map_dependent_axis if TYPE_CHECKING: from collections.abc import Iterator, Sequence from zarr_indexing.grid import DimensionGridLike -OutIndices = ( +_OutIndices = ( dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None ) -ChunkTransformResult = tuple[ +_ChunkTransformResult = tuple[ tuple[int, ...], IndexTransform, - OutIndices, + _OutIndices, ] type ChunkCoverage = Literal["full", "partial", "unknown"] @@ -86,6 +67,20 @@ class ChunkProjection: Both transforms share a synthetic input domain. ``chunk_transform`` maps that domain to chunk-local storage coordinates; ``cell_transform`` maps it to the original request domain. + + Attributes + ---------- + chunk_coords + Coordinates of the selected cell in the caller's grid. + chunk_domain + Bounds of that grid cell in global storage coordinates. + chunk_transform + Mapping from the shared synthetic domain to chunk-local storage. + cell_transform + Mapping from the shared synthetic domain to request coordinates. + coverage + Whether the request is proven to cover the whole grid cell exactly + once. Fancy selections are conservatively ``"unknown"``. """ chunk_coords: tuple[int, ...] @@ -104,7 +99,11 @@ def __post_init__(self) -> None: @dataclass(frozen=True, slots=True) class ChunkPlan: - """A reusable, lazy partition of an index transform over a chunk grid.""" + """A reusable, lazy partition of an index transform over a chunk grid. + + Construct plans with `plan_chunks`; iterating either the plan or + `projections()` performs a fresh chunk walk. + """ transform: IndexTransform dimension_grids: tuple[DimensionGridLike, ...] @@ -172,7 +171,7 @@ def _iter_sorted_1d_array_map( m: ArrayMap, storage: np.ndarray[Any, np.dtype[np.intp]], dim_grid: DimensionGridLike, -) -> Iterator[ChunkTransformResult]: +) -> Iterator[_ChunkTransformResult]: """Resolve a sorted 1-D ArrayMap one touched chunk at a time.""" start = 0 while start < storage.size: @@ -199,24 +198,14 @@ def _iter_sorted_1d_array_map( start = stop -def iter_chunk_transforms( +def _iter_chunk_transform_results( transform: IndexTransform, dim_grids: Sequence[DimensionGridLike], -) -> Iterator[ChunkTransformResult]: - """Resolve a composed IndexTransform against per-dimension chunk grids. - - `dim_grids` holds one `DimensionGridLike` per output (storage) dimension — - for zarr this is the chunk grid's per-dimension sequence. Yields - `(chunk_coords, sub_transform, out_indices)` triples: - - - `chunk_coords`: which chunk to access. - - `sub_transform`: maps output buffer coords to chunk-local coords. - - `out_indices`: where this chunk's values belong in the output. `None` - for basic/slice indexing, where the sub-transform says it on its own; an - integer array of scatter indices for a correlated (vectorized) selection; - and a `dict[int, ndarray]` of surviving positions per output dimension - for an orthogonal selection carrying more than one index array, which - `sub_transform_to_selections` turns into an open mesh. +) -> Iterator[_ChunkTransformResult]: + """Resolve a transform into private intersection bookkeeping. + + The survivor arrays are an implementation detail immediately converted to + a public `cell_transform` by `_iter_chunk_projections`. """ if any(size == 0 for size in transform.domain.shape): @@ -459,7 +448,7 @@ def _correlated_cell_transform( def _cell_transform( original: IndexTransform, restricted: IndexTransform, - survivors: OutIndices, + survivors: _OutIndices, ) -> IndexTransform: """Convert private survivor bookkeeping into a direction-neutral transform.""" if survivors is None: @@ -479,7 +468,9 @@ def _iter_chunk_projections( dim_grids: Sequence[DimensionGridLike], ) -> Iterator[ChunkProjection]: """Convert private intersection results into public paired projections.""" - for chunk_coords, chunk_transform, survivors in iter_chunk_transforms(transform, dim_grids): + for chunk_coords, chunk_transform, survivors in _iter_chunk_transform_results( + transform, dim_grids + ): chunk_min = tuple( grid.chunk_offset(coord) for grid, coord in zip(dim_grids, chunk_coords, strict=True) ) @@ -509,235 +500,3 @@ def _iter_chunk_projections( cell_transform=cell_transform, coverage=coverage, ) - - -def _dimension_map_slice(m: DimensionMap, dim_lo: int, dim_hi: int) -> slice: - """The NumPy slice a `DimensionMap` reads over input range `[dim_lo, dim_hi)`. - - The walk starts at the storage coordinate of `dim_lo` and takes `dim_hi - - dim_lo` steps of `m.stride`. A downward walk that reaches the start of the - axis must stop at `None`: an integer stop below zero counts from the end - instead, and swapping the endpoints while keeping the negative step (which - this did) produces a slice that selects nothing at all. - """ - return _positional_slice(m.offset + m.stride * dim_lo, dim_hi - dim_lo, m.stride) - - -def _scatter_in_block_order( - scatter: np.ndarray[Any, np.dtype[np.intp]], - *, - residual_extents: list[int], - basic_before_arrays: int, - array_positions: list[int], - has_points_axis: bool, -) -> np.ndarray[Any, np.dtype[np.intp]]: - """Reorder a points-major scatter to match the block the chunk read returns. - - The scatter enumerates the correlated selection points-major: point varying - slowest, then the residual sliced axes in their own order. NumPy does not - always hand the block back that way. With the coordinate arrays adjacent it - places the points axis where they sit, so a residual slice *before* them puts - a residual axis first; only when a slice separates the arrays does the points - axis lead. Assigning one order into the other transposed the values silently - when the extents happened to agree, and raised a broadcast error when they - did not. - - Nothing is read here: the scatter is an index array, so matching the block is - a permutation of it. - """ - if not has_points_axis or len(array_positions) == 0: - return scatter - separated = array_positions[-1] - array_positions[0] + 1 != len(array_positions) - points_axis = 0 if separated else basic_before_arrays - if points_axis == 0: - return scatter - - residual_cells = math.prod(residual_extents) if len(residual_extents) > 0 else 0 - if residual_cells == 0: - return scatter - n_points = scatter.size // residual_cells - # Points-major, then move the points axis to where the block carries it. The - # scatter keeps the block's shape rather than being flattened: the caller - # assigns the block into it directly, so the two must agree axis for axis. - grid = scatter.reshape((n_points, *residual_extents)) - return np.moveaxis(grid, 0, points_axis) - - -def _array_input_dim(m: ArrayMap, out_dim: int) -> int: - """The input dimension an orthogonal `ArrayMap` places its values along.""" - axis = array_map_dependent_axis(m) - if axis is None: - raise ValueError( - f"output[{out_dim}] is an ArrayMap that varies over no input dimension; " - "a map with no dependency axis should have been collapsed to a ConstantMap" - ) - return axis - - -def sub_transform_to_selections( - sub_transform: IndexTransform, - out_indices: OutIndices = None, -) -> tuple[ - tuple[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], - tuple[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], - tuple[int, ...], -]: - """Convert a chunk-local sub-transform to raw selections for the codec pipeline. - - Parameters - ---------- - sub_transform - A chunk-local IndexTransform (output maps already translated to - chunk-local coordinates). - out_indices - For vectorized indexing: the output scatter indices for this chunk. - None for orthogonal/basic indexing. - - Returns - ------- - tuple - `(chunk_selection, out_selection, drop_axes)` - - Notes - ----- - `out_selection` carries one entry per **input (domain)** dimension, not one - per output map. The two counts usually agree, but a domain dimension that no - output map depends on — a `vindex` broadcast axis whose partner a later basic - index consumed — has no map to be derived from and is filled in from the - domain. Emitting fewer entries than the buffer has axes would place the - values against the leading axes instead. - """ - inclusive_min = sub_transform.domain.inclusive_min - exclusive_max = sub_transform.domain.exclusive_max - input_rank = sub_transform.input_rank - - # Orthogonal outer product: >= 2 ArrayMaps each bound to a distinct input - # dimension. out_indices is a per-output-dim dict of surviving positions. The - # codec applies chunk_array[chunk_sel] / out[out_sel] with NumPy semantics, so - # build np.ix_-style selections (mirroring the legacy OrthogonalIndexer): one - # 1-D selector per dimension, expanded to an open mesh. ConstantMap dims are - # size-1 in chunk space and squeezed out via drop_axes. - if isinstance(out_indices, dict): - chunk_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] - by_input_dim: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} - drop_axes: list[int] = [] - for out_dim, m in enumerate(sub_transform.output): - if isinstance(m, ConstantMap): - chunk_arrays.append(np.array([m.offset], dtype=np.intp)) - drop_axes.append(out_dim) - elif isinstance(m, DimensionMap): - rng = np.arange(inclusive_min[m.input_dimension], exclusive_max[m.input_dimension]) - chunk_arrays.append((m.offset + m.stride * rng).astype(np.intp)) - by_input_dim[m.input_dimension] = rng.astype(np.intp) - else: # ArrayMap - idx = m.index_array.ravel() - chunk_arrays.append((m.offset + m.stride * idx).astype(np.intp)) - by_input_dim[_array_input_dim(m, out_dim)] = out_indices[out_dim] - out_arrays = [ - by_input_dim.get(d, np.arange(inclusive_min[d], exclusive_max[d], dtype=np.intp)) - for d in range(input_rank) - ] - return np.ix_(*chunk_arrays), np.ix_(*out_arrays), tuple(drop_axes) - - # Correlated (vindex) sub-transforms carry ArrayMaps with `input_dimension` - # None. They scatter through a single flat index (`out_indices`) into the - # row-major-flattened output buffer; the chunk selection reads a - # (points, residual-slice) block via the raveled coordinate arrays and any - # residual DimensionMap slices. - correlated = any( - isinstance(m, ArrayMap) and m.input_dimension is None for m in sub_transform.output - ) - if correlated: - # The broadcast block is the input axis no `DimensionMap` binds. It is - # absent when the block is rank 0 (every coordinate a scalar), and then - # the coordinate arrays index as scalars: the block carries the residual - # slice axes alone, matching a domain that has no points axis either. - bound = {m.input_dimension for m in sub_transform.output if isinstance(m, DimensionMap)} - has_points_axis = any(d not in bound for d in range(sub_transform.input_rank)) - points_shape: tuple[int, ...] = (-1,) if has_points_axis else () - chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] - # Where each entry lands in the block, so the scatter below can be put in - # the same order: the residual slices in the order they appear, and the - # positions the coordinate arrays occupy. - residual_extents: list[int] = [] - array_positions: list[int] = [] - basic_before_arrays = 0 - for m in sub_transform.output: - if isinstance(m, ConstantMap): - # See the basic branch below on why this is a slice: an integer - # would join the advanced indices and move the points axis. - chunk_sel.append(slice(m.offset, m.offset + 1)) - if len(array_positions) == 0: - basic_before_arrays += 1 - residual_extents.append(1) - elif isinstance(m, DimensionMap): - d = m.input_dimension - sl = _dimension_map_slice(m, inclusive_min[d], exclusive_max[d]) - chunk_sel.append(sl) - if len(array_positions) == 0: - basic_before_arrays += 1 - residual_extents.append(exclusive_max[d] - inclusive_min[d]) - else: # ArrayMap - idx = m.index_array.reshape(points_shape) - array_positions.append(len(chunk_sel)) - chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) - # Chunk resolution always supplies the flat scatter index for a - # correlated transform. Absent one (a bare sub-transform), fall back to an - # identity scatter over the whole flattened output buffer. - # `out_indices` is narrowed to a flat scatter array or None here (the - # per-dimension dict is an orthogonal outer product, handled above). - out_scatter: slice | np.ndarray[Any, np.dtype[np.intp]] - if out_indices is None: - n = 1 - for s in sub_transform.domain.shape: - n *= s - out_scatter = slice(0, n) - else: - out_scatter = _scatter_in_block_order( - out_indices, - residual_extents=residual_extents, - basic_before_arrays=basic_before_arrays, - array_positions=array_positions, - has_points_axis=has_points_axis, - ) - return tuple(chunk_sel), (out_scatter,), () - - chunk_sel = [] # annotated in the correlated branch above (same function scope) - out_by_dim: dict[int, slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = {} - basic_drop_axes: list[int] = [] - - # Single-pass build for the basic / single-orthogonal-array cases. - # ConstantMap dims read one cell and are reported in drop_axes. - for out_dim, m in enumerate(sub_transform.output): - if isinstance(m, ConstantMap): - # A length-one slice, not the bare integer this used to emit. NumPy - # counts an integer among the *advanced* indices when the tuple also - # holds an index array, and moves the broadcast axis to the front - # whenever a slice separates the two — while `out_selection` below is - # built positionally and keeps the original order. The two sides then - # described transposed blocks. A slice is a basic index wherever it - # sits, so the order is the one it looks like, and the axis it leaves - # behind is squeezed by the caller through drop_axes. - chunk_sel.append(slice(m.offset, m.offset + 1)) - basic_drop_axes.append(out_dim) - elif isinstance(m, DimensionMap): - d = m.input_dimension - dim_lo = inclusive_min[d] - dim_hi = exclusive_max[d] - chunk_sel.append(_dimension_map_slice(m, dim_lo, dim_hi)) - out_by_dim[d] = slice(dim_lo, dim_hi) - else: # ArrayMap (orthogonal: full-rank, raveled to its 1-D fancy coords) - idx = m.index_array.reshape(-1) - if m.offset == 0 and m.stride == 1: - chunk_sel.append(idx) - else: - chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) - # Orthogonal ArrayMap: out_indices holds the surviving positions. - out_by_dim[_array_input_dim(m, out_dim)] = ( - out_indices if out_indices is not None else slice(0, idx.size) - ) - - out_sel = [ - out_by_dim.get(d, slice(inclusive_min[d], exclusive_max[d])) for d in range(input_rank) - ] - return tuple(chunk_sel), tuple(out_sel), tuple(basic_drop_axes) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 7e95eb73d2..4fcbb3ed5b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -18,10 +18,10 @@ ----- A `LazyArray` carries a **partitioning** of the array it wraps: a grid of boxes that a read is broken into. `parts()` walks those boxes as they fall through the -view, yielding a [`Partition`](#zarr_indexing.lazy_array.Partition) per box — -the base coordinates of the box, a `LazyArray` covering exactly the cells of the -view that live in it, and where those cells belong in the result. `result()` is -built on that walk: +view, yielding a [`Partition`](#zarr_indexing.lazy_array.Partition) per box. Its +paired projection describes the chunk-local read and where its cells land in the +request; `view` is the directly resolvable form. `result()` is built on that +walk: ```python for part in view.parts(): @@ -851,7 +851,7 @@ def _partition_out_selection( """Lower ``cell_transform`` to NumPy selectors on the request buffer.""" domain = cell_transform.domain if _is_correlated(cell_transform): - selectors: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + correlated_selectors: list[np.ndarray[Any, np.dtype[np.intp]]] = [] for output_map in cell_transform.output: if isinstance(output_map, ConstantMap): coordinates = np.full(domain.shape, output_map.offset, dtype=np.intp) @@ -873,8 +873,8 @@ def _partition_out_selection( coordinates = output_map.offset + output_map.stride * np.broadcast_to( output_map.index_array, domain.shape ) - selectors.append(np.asarray(coordinates, dtype=np.intp)) - return tuple(selectors) + correlated_selectors.append(np.asarray(coordinates, dtype=np.intp)) + return tuple(correlated_selectors) selectors: list[int | slice | np.ndarray[Any, np.dtype[np.intp]]] = [] n_array_maps = sum(isinstance(output_map, ArrayMap) for output_map in cell_transform.output) @@ -959,6 +959,12 @@ class Partition: Attributes ---------- + projection + The source-independent description of this part. Its paired + `chunk_transform` and `cell_transform` share one compact synthetic + domain, mapping each selected cell to chunk-local storage and request + coordinates respectively. This is the authoritative placement model; + `base_coords` and `is_complete` are conveniences derived from it. base_coords Which box of the base partitioning this is, one coordinate per dimension of the wrapped array. @@ -966,8 +972,8 @@ class Partition: The box itself, in the global storage coordinates of the wrapped array: one `[inclusive_min, exclusive_max)` interval per dimension. `view.bounding_box()` is part-local and so cannot distinguish two - parts; this attribute can, and together with `is_complete` it decides a - read-modify-write. + parts; this attribute can. For a nested or repartitioned view this box + may be narrower than `projection.chunk_domain`. view A `LazyArray` covering exactly the cells of the view that live in this box. Resolving it reads the box once with basic slicing and applies the @@ -980,11 +986,9 @@ class Partition: directly as `out[part.out_selection] = ...`. is_complete Whether the view covers the whole box. Useful to a writer deciding - between a blind overwrite and a read-modify-write. True for a reversing - view, which reads every cell of the box back to front; false for every - fancy-indexed axis, which is conservative rather than exact — so it may - be `False` for a part it does in fact cover, but never `True` for one it - does not. + between a blind overwrite and a read-modify-write. Fancy projections + report `False` because their coverage is deliberately `unknown` until + duplicate-aware proof is added. """ projection: ChunkProjection diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index d71e0bf96c..9b36970f57 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -6,14 +6,59 @@ import pytest from zarr.core.chunk_grids import ChunkGrid, FixedDimension, VaryingDimension +import zarr_indexing from zarr_indexing import ChunkPlan, ChunkProjection, chunk_resolution, plan_chunks -from zarr_indexing.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections from zarr_indexing.domain import IndexDomain from zarr_indexing.grid import dimension_grids_from_chunks -from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.output_map import ConstantMap, DimensionMap from zarr_indexing.transform import IndexTransform +def _storage_of(transform: IndexTransform, point: tuple[int, ...]) -> tuple[int, ...]: + """Evaluate the three map forms at one point, independently of planning.""" + result: list[int] = [] + for output_map in transform.output: + if isinstance(output_map, ConstantMap): + result.append(output_map.offset) + elif isinstance(output_map, DimensionMap): + result.append(output_map.offset + output_map.stride * point[output_map.input_dimension]) + else: + index = tuple( + 0 + if output_map.index_array.shape[axis] == 1 + else point[axis] - transform.domain.inclusive_min[axis] + for axis in range(output_map.index_array.ndim) + ) + result.append( + output_map.offset + output_map.stride * int(output_map.index_array[index]) + ) + return tuple(result) + + +def _points(domain: IndexDomain) -> list[tuple[int, ...]]: + """Enumerate a small finite domain in its own coordinates.""" + return [ + tuple( + coordinate + origin + for coordinate, origin in zip(position, domain.inclusive_min, strict=True) + ) + for position in np.ndindex(*domain.shape) + ] + + +def _count_intersect_calls(monkeypatch: pytest.MonkeyPatch) -> dict[str, int]: + """Count real intersections to protect touched-only candidate enumeration.""" + calls = {"n": 0} + original = IndexTransform.intersect + + def counting(self: IndexTransform, output_domain: IndexDomain) -> object: + calls["n"] += 1 + return original(self, output_domain) + + monkeypatch.setattr(IndexTransform, "intersect", counting) + return calls + + def test_basic_plan_is_reiterable_and_projects_both_spaces() -> None: """A plan can be revisited without losing either side of each projection.""" transform = IndexTransform.from_shape((6,))[1:6] @@ -36,6 +81,28 @@ def test_basic_plan_is_reiterable_and_projects_both_spaces() -> None: projection.chunk_transform.domain == projection.cell_transform.domain for projection in first ) + assert all(projection.chunk_transform.domain.origin == (0,) for projection in first) + + +def test_projection_requires_one_shared_synthetic_domain() -> None: + """Paired transforms with different cell domains are rejected as incoherent.""" + with pytest.raises(ValueError, match="must share an input domain"): + ChunkProjection( + chunk_coords=(0,), + chunk_domain=IndexDomain.from_shape((3,)), + chunk_transform=IndexTransform.identity(IndexDomain.from_shape((2,))), + cell_transform=IndexTransform.identity(IndexDomain.from_shape((1,))), + coverage="partial", + ) + + +def test_projection_plan_is_the_only_public_chunk_resolution_surface() -> None: + """The greenfield API does not retain tuple or NumPy-selector bridges.""" + assert {"ChunkCoverage", "ChunkPlan", "ChunkProjection", "plan_chunks"} <= set( + zarr_indexing.__all__ + ) + assert "iter_chunk_transforms" not in zarr_indexing.__all__ + assert "sub_transform_to_selections" not in zarr_indexing.__all__ def test_plan_rejects_grid_rank_different_from_transform_output_rank() -> None: @@ -46,25 +113,62 @@ def test_plan_rejects_grid_rank_different_from_transform_output_rank() -> None: plan_chunks(transform, dimension_grids_from_chunks((2,), (2,))) -def _storage_of(transform: IndexTransform, point: tuple[int, ...]) -> tuple[int, ...]: - """Evaluate the three map forms at one point, independently of planning.""" - result: list[int] = [] - for output_map in transform.output: - if isinstance(output_map, ConstantMap): - result.append(output_map.offset) - elif isinstance(output_map, DimensionMap): - result.append(output_map.offset + output_map.stride * point[output_map.input_dimension]) - else: - index = tuple( - 0 - if output_map.index_array.shape[axis] == 1 - else point[axis] - transform.domain.inclusive_min[axis] - for axis in range(output_map.index_array.ndim) - ) - result.append( - output_map.offset + output_map.stride * int(output_map.index_array[index]) - ) - return tuple(result) +@pytest.mark.parametrize( + ("transform", "expected"), + [ + (IndexTransform.from_shape((5,)), ["full", "full"]), + (IndexTransform.from_shape((5,))[::-1], ["full", "full"]), + (IndexTransform.from_shape((5,))[::2], ["partial", "partial"]), + (IndexTransform.from_shape((5,))[2], ["partial"]), + ( + IndexTransform.from_shape((5,)).oindex[np.array([0, 1, 2, 3, 4])], + ["unknown", "unknown"], + ), + ], + ids=["clipped-edge", "reverse", "strided", "scalar", "fancy-is-conservative"], +) +def test_coverage_classification(transform: IndexTransform, expected: list[str]) -> None: + """Coverage is exact for affine requests and conservative for gathers.""" + grids = dimension_grids_from_chunks((3,), (5,)) + + assert [projection.coverage for projection in plan_chunks(transform, grids)] == expected + + +@pytest.mark.parametrize( + ("transform", "grids", "expected_coords"), + [ + ( + IndexTransform.from_shape((30,)), + dimension_grids_from_chunks((10,), (30,)), + [(0,), (1,), (2,)], + ), + ( + IndexTransform.from_shape((20, 30)), + dimension_grids_from_chunks((10, 10), (20, 30)), + [(i, j) for i in range(2) for j in range(3)], + ), + ( + IndexTransform.from_shape((100, 100))[25, :], + dimension_grids_from_chunks((10, 10), (100, 100)), + [(2, j) for j in range(10)], + ), + ( + IndexTransform.from_shape((100,))[8:15], + dimension_grids_from_chunks((10,), (100,)), + [(0,), (1,)], + ), + ], + ids=["one-dimensional", "two-dimensional", "constant-map", "slice"], +) +def test_affine_plans_touch_the_expected_chunks( + transform: IndexTransform, + grids: tuple[Any, ...], + expected_coords: list[tuple[int, ...]], +) -> None: + """Identity, constant, and sliced transforms enumerate literal grid cells.""" + assert [projection.chunk_coords for projection in plan_chunks(transform, grids)] == ( + expected_coords + ) @pytest.mark.parametrize( @@ -96,676 +200,156 @@ def test_projection_invariants_for_fancy_selections( for projection in plan: assert projection.coverage == "unknown" - cell_domain = projection.cell_transform.domain - for positional_point in np.ndindex(*cell_domain.shape): - cell_point = tuple( - coordinate + origin - for coordinate, origin in zip( - positional_point, cell_domain.inclusive_min, strict=True - ) - ) + assert projection.chunk_transform.domain == projection.cell_transform.domain + for cell_point in _points(projection.cell_transform.domain): request_point = _storage_of(projection.cell_transform, cell_point) chunk_point = _storage_of(projection.chunk_transform, cell_point) storage_point = _storage_of(plan.transform, request_point) chunk_origin = projection.chunk_domain.inclusive_min - assert chunk_point == tuple( value - origin for value, origin in zip(storage_point, chunk_origin, strict=True) ) + assert all( + 0 <= value < extent + for value, extent in zip(chunk_point, projection.chunk_domain.shape, strict=True) + ) request_points.append(request_point) - expected_request_points = [ - tuple( - coordinate + origin - for coordinate, origin in zip( - positional_point, transform.domain.inclusive_min, strict=True - ) - ) - for positional_point in np.ndindex(*transform.domain.shape) - ] - assert sorted(request_points) == sorted(expected_request_points) - - -class TestChunkResolutionIdentity: - def test_single_chunk(self) -> None: - """Array fits in one chunk.""" - t = IndexTransform.from_shape((10,)) - grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=10),)) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 1 - coords, sub_t, _ = results[0] - assert coords == (0,) - assert sub_t.domain.shape == (10,) - - def test_multiple_chunks_1d(self) -> None: - """1D array spanning 3 chunks.""" - t = IndexTransform.from_shape((30,)) - grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 3 - coords_list = [r[0] for r in results] - assert (0,) in coords_list - assert (1,) in coords_list - assert (2,) in coords_list - - def test_multiple_chunks_2d(self) -> None: - """2D array spanning 2x3 chunks.""" - t = IndexTransform.from_shape((20, 30)) - grid = ChunkGrid( - dimensions=( - FixedDimension(size=10, extent=20), - FixedDimension(size=10, extent=30), - ) - ) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 6 - coords_list = [r[0] for r in results] - assert (0, 0) in coords_list - assert (1, 2) in coords_list - - -class TestChunkResolutionSliced: - def test_slice_within_chunk(self) -> None: - """Slice that falls within a single chunk.""" - # Chunk resolution consumes zero-origin transforms: the I/O layer - # normalizes preserved (user-facing) domains via translate_domain_to - # before resolving, so mirror that contract here. - t = IndexTransform.from_shape((100,))[5:8].translate_domain_to((0,)) - grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 1 - coords, sub_t, _ = results[0] - assert coords == (0,) - assert isinstance(sub_t.output[0], DimensionMap) - assert sub_t.output[0].offset == 5 - - def test_slice_across_chunks(self) -> None: - """Slice that spans two chunks.""" - t = IndexTransform.from_shape((100,))[8:15] - grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 2 - coords_list = [r[0] for r in results] - assert (0,) in coords_list - assert (1,) in coords_list - - -class TestChunkResolutionConstant: - def test_integer_index(self) -> None: - """Integer index produces constant map — single chunk per constant dim.""" - t = IndexTransform.from_shape((100, 100))[25, :] - grid = ChunkGrid( - dimensions=( - FixedDimension(size=10, extent=100), - FixedDimension(size=10, extent=100), - ) - ) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 10 - for coords, _, _ in results: - assert coords[0] == 2 - - -class TestChunkResolutionArray: - def test_array_index(self) -> None: - """Array index map — chunks determined by array values.""" - idx = np.array([5, 15, 25], dtype=np.intp) - t = IndexTransform( - domain=IndexDomain.from_shape((3,)), - output=(ArrayMap(index_array=idx),), - ) - grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) - results = list(iter_chunk_transforms(t, grid._dimensions)) - coords_list = [r[0] for r in results] - assert (0,) in coords_list - assert (1,) in coords_list - assert (2,) in coords_list + assert sorted(request_points) == sorted(_points(transform.domain)) + + +def test_empty_request_has_no_projections() -> None: + """An empty fancy selection does not fabricate a touched chunk.""" + transform = IndexTransform.from_shape((10,)).oindex[np.array([], dtype=np.intp)] + grids = dimension_grids_from_chunks((3,), (10,)) + assert list(plan_chunks(transform, grids)) == [] -class TestChunkResolutionSorted1D: - def test_matches_general_resolution_for_randomized_sorted_selections( + +class TestSortedOneDimensionalPlan: + def test_matches_general_resolution_for_randomized_selections( self, monkeypatch: pytest.MonkeyPatch ) -> None: - """Direct partitioning matches the original resolver across varied inputs.""" + """The direct sorted path has the same paired transforms as intersection.""" rng = np.random.default_rng(0) grids = ( ChunkGrid(dimensions=(FixedDimension(size=7, extent=30),)), ChunkGrid(dimensions=(VaryingDimension(edges=(3, 4, 8, 5, 10), extent=30),)), ) - for grid in grids: for _ in range(50): - idx = np.sort(rng.integers(0, 30, size=int(rng.integers(1, 80)))).astype(np.intp) - transform = IndexTransform.from_shape((30,)).vindex[idx] - direct = list(iter_chunk_transforms(transform, grid._dimensions)) - + indices = np.sort(rng.integers(0, 30, size=int(rng.integers(1, 80)))).astype( + np.intp + ) + transform = IndexTransform.from_shape((30,)).vindex[indices] + direct = list(plan_chunks(transform, grid._dimensions)) with monkeypatch.context() as context: context.setattr( chunk_resolution, "_one_dimensional_correlated_array_map", lambda _transform: None, ) - general = list(iter_chunk_transforms(transform, grid._dimensions)) - - assert [result[0] for result in direct] == [result[0] for result in general] - for direct_result, general_result in zip(direct, general, strict=True): - _, direct_t, direct_out = direct_result - _, general_t, general_out = general_result - assert direct_t.domain == general_t.domain - - direct_chunk_sel, direct_out_sel, direct_drop = sub_transform_to_selections( - direct_t, direct_out - ) - general_chunk_sel, general_out_sel, general_drop = sub_transform_to_selections( - general_t, general_out - ) - assert direct_drop == general_drop - np.testing.assert_array_equal(direct_chunk_sel[0], general_chunk_sel[0]) - np.testing.assert_array_equal(direct_out_sel[0], general_out_sel[0]) + general = list(plan_chunks(transform, grid._dimensions)) + assert direct == general - def test_sorted_vindex_partitions_chunks_without_intersection( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Sorted vectorized coordinates are sliced directly per touched chunk.""" - idx = np.array([0, 3, 4, 4, 9, 11], dtype=np.intp) - t = IndexTransform.from_shape((12,)).vindex[idx] + def test_sorted_coordinates_bypass_intersection(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Sorted coordinates partition directly at touched chunk boundaries.""" + transform = IndexTransform.from_shape((12,)).vindex[ + np.array([0, 3, 4, 4, 9, 11], dtype=np.intp) + ] grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) - - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) - - assert [result[0] for result in results] == [(0,), (1,), (2,)] - assert calls["n"] == 0 - - expected_chunk_indices = ([0, 3], [0, 0], [1, 3]) - expected_out_indices = ([0, 1], [2, 3], [4, 5]) - for result, expected_chunk, expected_out in zip( - results, expected_chunk_indices, expected_out_indices, strict=True - ): - _, sub_t, out_indices = result - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices) - np.testing.assert_array_equal(chunk_sel[0], expected_chunk) - np.testing.assert_array_equal(out_sel[0], expected_out) - assert drop_axes == () - - def test_sorted_array_map_preserves_offset_and_stride(self) -> None: - """Storage partitioning retains the ArrayMap's offset and stride.""" - t = IndexTransform( - domain=IndexDomain.from_shape((3,)), - output=( - ArrayMap( - index_array=np.array([0, 1, 2], dtype=np.intp), - offset=1, - stride=3, - ), - ), - ) - grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=8),)) - - results = list(iter_chunk_transforms(t, grid._dimensions)) - - assert [result[0] for result in results] == [(0,), (1,)] - expected_chunk_indices = ([1], [0, 3]) - expected_out_indices = ([0], [1, 2]) - for result, expected_chunk, expected_out in zip( - results, expected_chunk_indices, expected_out_indices, strict=True - ): - _, sub_t, out_indices = result - chunk_sel, out_sel, _ = sub_transform_to_selections(sub_t, out_indices) - np.testing.assert_array_equal(chunk_sel[0], expected_chunk) - np.testing.assert_array_equal(out_sel[0], expected_out) - - def test_sorted_vindex_with_varying_chunks(self) -> None: - """Touched-boundary searches also support a non-uniform 1-D grid.""" - idx = np.array([0, 1, 2, 3, 5, 9], dtype=np.intp) - t = IndexTransform.from_shape((10,)).vindex[idx] - grid = ChunkGrid(dimensions=(VaryingDimension(edges=(2, 3, 5), extent=10),)) - - results = list(iter_chunk_transforms(t, grid._dimensions)) - - assert [result[0] for result in results] == [(0,), (1,), (2,)] - expected_chunk_indices = ([0, 1], [0, 1], [0, 4]) - for result, expected_chunk in zip(results, expected_chunk_indices, strict=True): - _, sub_t, out_indices = result - chunk_sel, _, _ = sub_transform_to_selections(sub_t, out_indices) - np.testing.assert_array_equal(chunk_sel[0], expected_chunk) - - def test_sorted_vindex_with_zero_sized_dimension_uses_general_resolution( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - """A zero-sized grid cannot be partitioned by touched boundaries.""" - t = IndexTransform.from_shape((10,)).vindex[np.array([1], dtype=np.intp)] - grid = ChunkGrid(dimensions=(FixedDimension(size=0, extent=10),)) - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert results == [] - assert calls["n"] == 1 + projections = list(plan_chunks(transform, grid._dimensions)) - def test_unsorted_vindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Unsorted coordinates continue through the general intersection logic.""" - t = IndexTransform.from_shape((12,)).vindex[np.array([9, 0, 4], dtype=np.intp)] + assert [projection.chunk_coords for projection in projections] == [(0,), (1,), (2,)] + assert calls["n"] == 0 + assert [ + [ + _storage_of(projection.cell_transform, point)[0] + for point in _points(projection.cell_transform.domain) + ] + for projection in projections + ] == [[0, 1], [2, 3], [4, 5]] + + def test_unsorted_coordinates_use_intersection(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Unsorted coordinates retain the general intersection path.""" + transform = IndexTransform.from_shape((12,)).vindex[np.array([9, 0, 4], dtype=np.intp)] grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) - - assert [result[0] for result in results] == [(0,), (1,), (2,)] - assert calls["n"] == 3 - def test_sorted_oindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Orthogonal ArrayMaps retain their existing domain-aware resolution.""" - t = IndexTransform.from_shape((12,)).oindex[np.array([0, 4, 9], dtype=np.intp)] - grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) - - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) + projections = list(plan_chunks(transform, grid._dimensions)) - assert [result[0] for result in results] == [(0,), (1,), (2,)] + assert [projection.chunk_coords for projection in projections] == [(0,), (1,), (2,)] assert calls["n"] == 3 -def _count_intersect_calls(monkeypatch: pytest.MonkeyPatch) -> dict[str, int]: - """Wrap `IndexTransform.intersect` with a call counter. - - Returns a mutable dict whose `"n"` entry is the number of times - `intersect` is invoked. Used to assert that candidate-chunk enumeration is - proportional to the *touched* chunks, not the dense bounding box between the - min and max touched chunk. - """ - calls = {"n": 0} - original = IndexTransform.intersect - - def counting(self: IndexTransform, output_domain: IndexDomain) -> object: - calls["n"] += 1 - return original(self, output_domain) - - monkeypatch.setattr(IndexTransform, "intersect", counting) - return calls - - -class TestChunkResolutionTouchedOnly: - """`iter_chunk_transforms` must enumerate only the chunks a fancy selection - actually touches — never the dense `range(min_chunk, max_chunk + 1)` bounding - box. These guard against a regression to bounding-box enumeration, whose cost - scales with grid size rather than with the number of selected coordinates. - """ - - def test_1d_sparse_vindex_enumerates_only_touched_chunks( +class TestTouchedOnlyCandidateEnumeration: + def test_sparse_one_dimensional_selection_skips_the_dense_span( self, monkeypatch: pytest.MonkeyPatch ) -> None: - """Two far-apart coordinates on a 1000-chunk grid touch exactly 2 chunks. - - A dense bounding-box enumeration would intersect ~1000 candidate chunks; - touched-only enumeration intersects exactly 2. - """ - # 4000 elements, chunk size 4 -> 1000 chunks. coords 1 and 3997 land in - # chunk 0 and chunk 999 respectively (998 empty chunks between them). + """Two sorted points on a 1000-cell grid require no intersections.""" + transform = IndexTransform.from_shape((4000,)).vindex[np.array([1, 3997], dtype=np.intp)] grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=4000),)) - t = IndexTransform.from_shape((4000,)).vindex[np.array([1, 3997], dtype=np.intp)] - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) - coords = sorted(r[0] for r in results) - assert coords == [(0,), (999,)] - # Sorted 1-D coordinates are partitioned directly, without intersecting - # either the touched chunks or the 998 empty chunks between them. + projections = list(plan_chunks(transform, grid._dimensions)) + + assert [projection.chunk_coords for projection in projections] == [(0,), (999,)] assert calls["n"] == 0 - def test_2d_orthogonal_enumerates_only_touched_chunks( - self, monkeypatch: pytest.MonkeyPatch + @pytest.mark.parametrize( + ("mode", "expected_coords", "expected_calls"), + [ + ("orthogonal", [(0, 0), (0, 999), (999, 0), (999, 999)], 4), + ("correlated", [(0, 0), (999, 999)], 2), + ], + ) + def test_sparse_two_dimensional_selection_uses_only_touched_combinations( + self, + monkeypatch: pytest.MonkeyPatch, + mode: str, + expected_coords: list[tuple[int, int]], + expected_calls: int, ) -> None: - """Orthogonal outer product of two 2-coordinate arrays touches 2x2 chunks. - - Per-dimension distinct touched chunks: {0, 999} on each axis. The outer - product is 2*2 = 4 candidate chunks (all survive), versus ~1e6 for a - dense 1000x1000 bounding box. - """ - grid = ChunkGrid( - dimensions=( - FixedDimension(size=4, extent=4000), - FixedDimension(size=4, extent=4000), - ) + """Orthogonal points use their outer product; correlated points remain paired.""" + base = IndexTransform.from_shape((4000, 4000)) + first = np.array([1, 3997], dtype=np.intp) + second = np.array([2, 3998], dtype=np.intp) + transform = ( + base.oindex[first, second] if mode == "orthogonal" else base.vindex[first, second] ) - t = IndexTransform.from_shape((4000, 4000)).oindex[ - np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) - ] - - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) - - coords = sorted(r[0] for r in results) - assert coords == [(0, 0), (0, 999), (999, 0), (999, 999)] - assert calls["n"] == 4 - - def test_2d_correlated_vindex_enumerates_joint_touched_chunks( - self, monkeypatch: pytest.MonkeyPatch - ) -> None: - """Two correlated (vindex) coordinate arrays scatter to 2 diagonal chunks. - - The two points (1, 2) and (3997, 3998) touch chunks (0, 0) and - (999, 999). Correlated coordinate arrays are grouped *jointly*, so - enumeration intersects exactly the 2 touched chunks — never the 2x2 - cartesian product of per-dimension distinct chunks, and never the dense - 1e6 grid. - """ grid = ChunkGrid( dimensions=( FixedDimension(size=4, extent=4000), FixedDimension(size=4, extent=4000), ) ) - t = IndexTransform.from_shape((4000, 4000)).vindex[ - np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) - ] - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) - coords = sorted(r[0] for r in results) - assert coords == [(0, 0), (999, 999)] - assert calls["n"] == 2 + projections = list(plan_chunks(transform, grid._dimensions)) + + assert sorted(projection.chunk_coords for projection in projections) == expected_coords + assert calls["n"] == expected_calls - def test_2d_correlated_vindex_diagonal_is_linear_in_points( + def test_correlated_diagonal_scales_with_points_not_their_product( self, monkeypatch: pytest.MonkeyPatch ) -> None: - """A diagonal of P correlated points touches P chunks with O(P) intersections. - - Enumerating the cartesian product of per-dimension distinct chunk sets - would cost P**2 intersections (2500 here) — quadratic in the number of - selected points for the scattered selections of zarr-python gh-4174. - Joint grouping keeps resolution work proportional to the touched chunks. - """ - p = 50 + """Fifty diagonal points require fifty, rather than 2500, intersections.""" + n_points = 50 + coordinates = np.arange(n_points, dtype=np.intp) * 8 + transform = IndexTransform.from_shape((4000, 4000)).vindex[coordinates, coordinates] grid = ChunkGrid( dimensions=( FixedDimension(size=4, extent=4000), FixedDimension(size=4, extent=4000), ) ) - # point i lands in chunk (2i, 2i): all per-dimension chunks distinct - coords_1d = np.arange(p, dtype=np.intp) * 8 - t = IndexTransform.from_shape((4000, 4000)).vindex[coords_1d, coords_1d] - calls = _count_intersect_calls(monkeypatch) - results = list(iter_chunk_transforms(t, grid._dimensions)) - - assert sorted(r[0] for r in results) == [(2 * i, 2 * i) for i in range(p)] - assert calls["n"] == p - - -class TestSubTransformToSelections: - def test_constant_map(self) -> None: - """ConstantMap reads one cell through a length-one slice, and says so. - - Not the bare integer this used to emit. NumPy counts an integer among - the *advanced* indices whenever the tuple also holds an index array, and - moves the broadcast axis to the front when a slice separates them, while - `out_selection` is built positionally — so the two sides described - transposed blocks. A slice is basic wherever it sits, and the size-one - axis it leaves behind is named in `drop_axes` for the caller to squeeze. - - The domain dimension here is referenced by no output map, so its - out-selection cannot come from a map: it is the whole axis, taken from - the domain. One entry per domain dimension is what `out[out_sel] = value` - needs — an empty tuple would leave the buffer's rank unaddressed and - place a lower-rank part against the leading axes. - """ - t = IndexTransform( - domain=IndexDomain.from_shape((10,)), - output=(ConstantMap(offset=5),), - ) - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) - assert chunk_sel == (slice(5, 6),) - assert out_sel == (slice(0, 10),) - assert drop_axes == (0,) - - def test_dimension_map_stride_1(self) -> None: - """DimensionMap with stride=1 produces contiguous slice.""" - t = IndexTransform( - domain=IndexDomain.from_shape((10,)), - output=(DimensionMap(input_dimension=0, offset=3, stride=1),), - ) - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) - assert chunk_sel == (slice(3, 13, 1),) - assert out_sel == (slice(0, 10),) - assert drop_axes == () - - def test_dimension_map_strided(self) -> None: - """DimensionMap with stride>1 produces strided slice.""" - t = IndexTransform( - domain=IndexDomain.from_shape((5,)), - output=(DimensionMap(input_dimension=0, offset=2, stride=3),), - ) - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) - assert chunk_sel == (slice(2, 17, 3),) - assert out_sel == (slice(0, 5),) - assert drop_axes == () - - def test_array_map(self) -> None: - """ArrayMap produces integer array selection.""" - arr = np.array([1, 5, 9], dtype=np.intp) - t = IndexTransform( - domain=IndexDomain.from_shape((3,)), - output=(ArrayMap(index_array=arr, offset=0, stride=1),), - ) - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) - assert isinstance(chunk_sel[0], np.ndarray) - np.testing.assert_array_equal(chunk_sel[0], arr) - # Without chunk_mask, out_sel falls back to domain-based slices - assert out_sel == (slice(0, 3),) - assert drop_axes == () - - def test_array_map_with_offset_stride(self) -> None: - """ArrayMap with offset and stride computes storage coords.""" - arr = np.array([0, 1, 2], dtype=np.intp) - t = IndexTransform( - domain=IndexDomain.from_shape((3,)), - output=(ArrayMap(index_array=arr, offset=10, stride=5),), - ) - chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) - assert isinstance(chunk_sel[0], np.ndarray) - np.testing.assert_array_equal(chunk_sel[0], np.array([10, 15, 20])) - assert drop_axes == () - - def test_dimension_map_reversed(self) -> None: - """A negative stride walks the axis backwards, stopping off the front. - - `slice(0, 8, -1)` — the endpoints swapped while the step stays negative — - selects nothing; the axis reversed is `slice(7, None, -1)`. - """ - t = IndexTransform.identity(IndexDomain.from_shape((8,)))[::-1].translate_domain_to((0,)) - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) - assert chunk_sel == (slice(7, None, -1),) - assert out_sel == (slice(0, 8),) - assert drop_axes == () - np.testing.assert_array_equal(np.arange(8)[chunk_sel[0]], np.arange(8)[::-1]) - - def test_dimension_map_reversed_bounded(self) -> None: - """A reversed walk that stops short keeps an ordinary integer stop.""" - t = IndexTransform.identity(IndexDomain.from_shape((8,)))[6:2:-1].translate_domain_to((0,)) - chunk_sel, _out_sel, _drop = sub_transform_to_selections(t) - np.testing.assert_array_equal(np.arange(8)[chunk_sel[0]], np.arange(8)[6:2:-1]) - - def test_dimension_map_reversed_strided(self) -> None: - """The same, with a stride whose walk does not land on the origin.""" - t = IndexTransform.identity(IndexDomain.from_shape((9,)))[::-3].translate_domain_to((0,)) - chunk_sel, _out_sel, _drop = sub_transform_to_selections(t) - np.testing.assert_array_equal(np.arange(9)[chunk_sel[0]], np.arange(9)[::-3]) - - def test_dimension_map_slice_reads_the_coordinates_the_map_names(self) -> None: - """Over every offset/stride/extent in range, the emitted slice picks out - exactly `offset + stride * i` for `i` in the input range.""" - extent = 12 - cells = np.arange(extent) - for stride in (-4, -3, -2, -1, 1, 2, 3, 4): - for size in range(extent + 1): - span = 0 if size == 0 else abs(stride) * (size - 1) - offsets = range(extent - span) if span < extent else () - for offset in offsets: - start = offset if stride > 0 else offset + span - t = IndexTransform( - domain=IndexDomain.from_shape((size,)), - output=(DimensionMap(input_dimension=0, offset=start, stride=stride),), - ) - chunk_sel, _out_sel, _drop = sub_transform_to_selections(t) - expected = np.array([start + stride * i for i in range(size)]) - np.testing.assert_array_equal( - cells[chunk_sel[0]], expected, err_msg=f"{stride=} {size=} {start=}" - ) - def test_correlated_residual_dimension_map_reversed(self) -> None: - """The correlated branch builds the same reversed slice for its residual dim.""" - t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] - reversed_slice = t[:, ::-1].translate_domain_to((0, 0)) - chunk_sel, _out_sel, _drop = sub_transform_to_selections(reversed_slice) - assert chunk_sel[2] == slice(4, None, -1) - np.testing.assert_array_equal(np.arange(5)[chunk_sel[2]], np.arange(5)[::-1]) - - def test_mixed_maps_2d(self) -> None: - """Mix of ConstantMap and DimensionMap.""" - t = IndexTransform( - domain=IndexDomain.from_shape((10,)), - output=( - ConstantMap(offset=5), - DimensionMap(input_dimension=0, offset=0, stride=1), - ), - ) - chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) - assert chunk_sel[0] == slice(5, 6) - assert chunk_sel[1] == slice(0, 10, 1) - # The constant's axis survives the read at size one and is squeezed by - # the caller; letting NumPy drop it through an integer index put the - # constant in the advanced-index group, where it reordered the block. - assert drop_axes == (0,) - - -class TestChunkResolutionArrayMapFlavors: - """Chunk resolution must yield outer-product (np.ix_) selectors for - orthogonal ArrayMaps and shared flat-scatter selectors for correlated ones, - and must return early for empty fancy selections.""" - - def test_empty_array_selection_yields_nothing(self) -> None: - """An empty ArrayMap selection produces no chunk transforms (no crash).""" - t = IndexTransform( - domain=IndexDomain.from_shape((0,)), - output=(ArrayMap(index_array=np.array([], dtype=np.intp)),), - ) - grid = ChunkGrid(dimensions=(FixedDimension(size=3, extent=10),)) - assert list(iter_chunk_transforms(t, grid._dimensions)) == [] + projections = list(plan_chunks(transform, grid._dimensions)) - def test_orthogonal_outer_product_selectors(self) -> None: - """Two independent arrays produce np.ix_-style (mesh) chunk/out selectors.""" - t = IndexTransform.from_shape((10, 10)).oindex[np.array([1, 3]), np.array([2, 4, 6])] - grid = ChunkGrid( - dimensions=(FixedDimension(size=10, extent=10), FixedDimension(size=10, extent=10)) - ) - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 1 - _coords, sub_t, out_indices = results[0] - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices) - # np.ix_ produces one 2-D open-mesh selector per axis, for both sides. - assert len(chunk_sel) == 2 - assert len(out_sel) == 2 - assert isinstance(chunk_sel[0], np.ndarray) - assert isinstance(chunk_sel[1], np.ndarray) - assert chunk_sel[0].shape == (2, 1) - assert chunk_sel[1].shape == (1, 3) - assert drop_axes == () - - def test_correlated_scatter_with_residual_slice(self) -> None: - """Correlated arrays + a residual slice dim scatter through a single flat - index whose shape matches the (points, slice) block read from the chunk.""" - t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] - grid = ChunkGrid( - dimensions=( - FixedDimension(size=4, extent=4), - FixedDimension(size=3, extent=3), - FixedDimension(size=5, extent=5), - ) - ) - # One chunk holds everything: both points survive, slice dim spans [0,5). - results = list(iter_chunk_transforms(t, grid._dimensions)) - assert len(results) == 1 - _coords, sub_t, out_indices = results[0] - chunk_sel, out_sel, _drop = sub_transform_to_selections(sub_t, out_indices) - # Chunk side: flat coordinate arrays for the two correlated dims plus a - # slice for the residual dim. - assert len(chunk_sel) == 3 - np.testing.assert_array_equal(np.asarray(chunk_sel[0]), [1, 3]) - np.testing.assert_array_equal(np.asarray(chunk_sel[1]), [2, 0]) - assert chunk_sel[2] == slice(0, 5, 1) - # Output side: a single flat scatter index of shape (points, slice) = (2, 5). - assert len(out_sel) == 1 - assert np.asarray(out_sel[0]).shape == (2, 5) - - -class TestBridgeBlockOrder: - """`out[out_sel] = chunk[chunk_sel]` must address the same block on both sides. - - Nothing exercised this pairing before: `LazyArray` resolves a part by - lowering its own transform and only ever reads `out_selection`, so a - `chunk_selection` describing a differently-ordered block went unnoticed by - the whole suite while corrupting any consumer that followed the contract. - """ - - @staticmethod - def _assemble( - data: np.ndarray[Any, Any], chunks: tuple[int, ...], transform: IndexTransform - ) -> np.ndarray[Any, Any]: - grids = dimension_grids_from_chunks(chunks, data.shape) - out = np.zeros(transform.domain.shape, dtype=data.dtype) - for coords, local, out_indices in iter_chunk_transforms(transform, grids): - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(local, out_indices) - block_slice = tuple( - slice(c * size, c * size + grid.chunk_size(c)) - for c, size, grid in zip(coords, chunks, grids, strict=True) - ) - block = data[block_slice][chunk_sel] - if len(drop_axes) > 0: - block = np.squeeze(block, axis=drop_axes) - if len(out_sel) == 1 and isinstance(out_sel[0], np.ndarray): - out.reshape(-1)[out_sel[0]] = block - else: - out[out_sel] = block - return out - - @pytest.mark.parametrize("chunks", [(2, 3, 3), (1, 2, 2)], ids=["one-chunk", "many-chunks"]) - def test_a_constant_before_an_array_keeps_the_blocks_axis_order( - self, chunks: tuple[int, int, int] - ) -> None: - """`oindex[int, :, array]` — a constant, a slice, then an index array. - - NumPy folds an integer into the advanced-index group, and moves the - broadcast axis to the front when a slice separates the two. Emitting the - constant as an integer therefore transposed the chunk block against an - `out_selection` built positionally. - """ - data = np.arange(18).reshape(2, 3, 3) - transform = ( - IndexTransform.from_shape(data.shape)[1] - .translate_domain_to((0, 0)) - .oindex[:, np.array([0, 1, 2])] - ) - np.testing.assert_array_equal( - self._assemble(data, chunks, transform), data[1][:, [0, 1, 2]] - ) - - def test_a_residual_slice_before_the_points_keeps_the_scatter_in_step(self) -> None: - """A correlated read whose residual dimension precedes its coordinates. - - The scatter enumerates points-major, but NumPy leaves the points axis - where the coordinate arrays sit, so the block arrived slice-major. Equal - extents transposed the values silently; unequal ones raised. - """ - data = np.arange(12).reshape(3, 4) - transform = IndexTransform.from_shape(data.shape).vindex[..., np.array([0, 1, 2])] - np.testing.assert_array_equal(self._assemble(data, (3, 4), transform), data[..., [0, 1, 2]]) - - def test_a_residual_slice_before_the_points_of_a_different_length(self) -> None: - """The same shape, with extents that cannot coincide — this one raised.""" - data = np.arange(8).reshape(2, 4) - transform = IndexTransform.from_shape(data.shape).vindex[..., np.array([0, 1, 2])] - np.testing.assert_array_equal(self._assemble(data, (2, 4), transform), data[..., [0, 1, 2]]) + assert sorted(projection.chunk_coords for projection in projections) == [ + (2 * index, 2 * index) for index in range(n_points) + ] + assert calls["n"] == n_points diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 79334f7be6..79b43851fc 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -1182,7 +1182,7 @@ def test_a_query_bounding_box_is_only_a_hull() -> None: lambda a: a.lazy.oindex[[4, 0, 0], :, [3, 1]], lambda a: a.lazy.vindex[..., np.array([1, 4, 0]), np.array([2, 0, 1])], # A reversing view drives the negative-stride branches of - # `_intersect_dimension_map` and `iter_chunk_transforms`, which were + # `_intersect_dimension_map` and chunk projection, which were # written defensively long before anything could reach them. lambda a: a.lazy[::-1, ::-2, :], lambda a: a.lazy[5:1:-1, :, ::-1],