From bf6b01da0c5944a79bd97ef54c495bb350b906e6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 13 Sep 2026 18:50:30 +0200 Subject: [PATCH 01/12] docs(indexing): ground design and integration claims in current behavior Assisted-by: Codex:GPT-6 --- packages/zarr-indexing/CONTRIBUTING.md | 4 +- packages/zarr-indexing/README.md | 14 ++- packages/zarr-indexing/docs/design-notes.md | 109 +++++++++--------- packages/zarr-indexing/docs/guide/index.md | 11 +- .../zarr-indexing/docs/guide/integrations.md | 43 +++---- packages/zarr-indexing/docs/guide/patterns.md | 9 +- packages/zarr-indexing/docs/index.md | 8 +- packages/zarr-indexing/justfile | 5 +- packages/zarr-indexing/pyproject.toml | 26 ++--- 9 files changed, 118 insertions(+), 111 deletions(-) diff --git a/packages/zarr-indexing/CONTRIBUTING.md b/packages/zarr-indexing/CONTRIBUTING.md index 632e24929f..a591d31bd6 100644 --- a/packages/zarr-indexing/CONTRIBUTING.md +++ b/packages/zarr-indexing/CONTRIBUTING.md @@ -8,11 +8,11 @@ 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 check # the checks above plus TensorStore parity just docs-serve # serve the docs site locally ``` -Run them from this directory, or from anywhere in the repository as +Run them from this directory, or from the repository root as `just packages/zarr-indexing/`. The test recipe runs against the workspace-root environment, because the diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index 4328de1cc1..fb631d37e6 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -5,9 +5,9 @@ Composable, lazy coordinate transforms for Zarr array indexing. Documentation: This package implements TensorStore-inspired index transforms. The core idea: -every indexing operation (slicing, fancy indexing, etc.) produces a coordinate -mapping from user space to storage space. These mappings compose lazily — no -I/O until you explicitly read or write. +each supported indexing operation (slicing, fancy indexing, etc.) produces a coordinate +mapping from user space to storage space. These mappings compose without reading selected source values. `LazyArray` +materializes them on request; the transform algebra itself performs no source I/O. Key types: @@ -32,9 +32,11 @@ Key types: dimension can depend on the input - `compose` — chain two transforms into one -The package depends only on NumPy and the standard library; it does not import -`zarr`. It is developed in the [zarr-python](https://github.com/zarr-developers/zarr-python) -repository and consumed by `zarr` to resolve array indexing operations. +The base package depends on NumPy and the standard library; its optional testing +module also requires Hypothesis. The package does not import `zarr`. It is developed +in the [zarr-python](https://github.com/zarr-developers/zarr-python) repository, +and its integration tests exercise Zarr chunk grids. Installing it does not +replace Zarr's indexing implementation. ## Installation diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index 59eef4dacd..f5db95500c 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -28,13 +28,14 @@ about the deliberately matching semantics: - **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. + toward zero. `tests/test_tensorstore_parity.py` compares the enumerated cases with + TensorStore when that optional dependency is installed. +- **The wire format.** [ndsel](ndsel.md) uses TensorStore's domain and + output-map field names for transform bodies, with an additional `kind` + discriminator. `tests/test_ndsel_tensorstore.py` checks interoperability for + the tested cases. Their validation rules differ, and loading and + re-emitting a message through the engine can normalize or discard metadata; + see [lowering to a transform](ndsel.md#lowering-to-a-transform). - **Chunk partitioning.** Both factor a transform over a grid before visiting any cell, rather than intersecting the whole transform with each chunk. TensorStore's `IndexTransformGridPartition` holds strided sets and index @@ -55,9 +56,9 @@ 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 | +| Dialect | Coordinate indices are literal; negative coordinates do not wrap | The algebra keeps that dialect; each public boundary picks its own. [`LazyArray`](api/lazy_array.md) speaks positional NumPy, `IndexTransform` 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) 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 | +| Wire format | [Documented JSON schema](https://google.github.io/tensorstore/index_space.html#index-transform) | [ndsel](ndsel.md) is spec-first, with a vendored conformance corpus exercised by this implementation | | Backends | A driver ecosystem (zarr, N5, neuroglancer, GCS, …) built into the library | No drivers. The default reader needs `shape`, `dtype`, basic integer/slice indexing, and selected slabs convertible to NumPy system memory; other backends use explicit custom readers. A device reader owns transfer into the supplied system-memory output | The mechanics of a @@ -74,19 +75,18 @@ and caller-supplied grid; 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. Independent strided planning -here is per axis, but each materialized `ChunkProjection` is -Python-level bookkeeping over NumPy — two domains, two transforms and the -projection itself — so the per-part overhead of the object view is -significant; a consumer that reads the partition's tables directly pays no -per-chunk object construction. 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. +`coverage` describes selection coverage relative to that grid. A `full` +classification can help a writer avoid reading old values, but does not by +itself prove that a write is safe: encoding requirements, conflicts, duplicate +semantics, and concurrency remain consumer responsibilities. `unknown` means +the planner has not established complete or partial coverage. + +This implementation performs Python-level bookkeeping over NumPy. This page +provides no benchmark establishing a general performance ordering against +TensorStore; costs depend on the selection and execution backend. + +The partition tables can be consumed without constructing a `ChunkProjection` +for each chunk. Materializing projections adds Python object construction. ## Bounding-box selections vs query selections @@ -102,11 +102,11 @@ 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, subsequent basic indexing cannot make -it a box again. A second query composes onto any axis of an existing one — +table of coordinates. Its stored coordinate arrays cost space proportional to their stored size. +Coordinates may repeat or scatter, but can also be contiguous and local. +Current query-resolution paths inspect these arrays. Fancy indexing can produce +a query, but singleton or constant selections can collapse to `ConstantMap`; +subsequent indexing can therefore make a query affine again. A second query composes onto any axis of an existing one — including the axes it merely broadcasts along — by evaluating the existing lookup tables at the new coordinates. @@ -116,9 +116,10 @@ planning and materialization. [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: +`points` desugars to `index_array` maps. A transform without index-array maps +is a box in the engine’s structural classification. Loading can further +simplify degenerate index arrays to constants, so an arbitrary incoming body +with `index_array` fields need not remain a query. For example: ```python from zarr_indexing import IndexTransform @@ -134,16 +135,13 @@ gather.to_json()["output"][0] # 'index_array_bounds': ['-inf', '+inf']} ``` -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 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. +The representation helps a consumer choose a lowering strategy. Independent +affine axes can often be read with slices plus reversal, permutation, or +broadcasting. Arbitrary affine maps can also express diagonals, so the absence +of `ArrayMap` alone is not proof of a rectangular slab. Queries may be lowered +through gathers or covers, and can sometimes simplify to slices. Chunk plans +group selected coordinates by chunk while preserving result placement; repeated +coordinates do not imply repeated visits to the same chunk. [`LazyArray`](api/lazy_array.md) exposes the category directly: @@ -173,16 +171,20 @@ gather.shape # (3, 80) `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 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 +These summaries omit traversal direction, input-axis correspondence, and +result layout. For example, forward and reversed views have identical bounds +and stride magnitudes but different ordered results. Use the transform for the +complete selection. + +For independent axes with multiple selected coordinates, a stride magnitude +greater than one leaves gaps in the hull. Singleton axes are an exception, +and a query can also cover every cell of its hull. 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 *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. +of the region it reaches. An empty 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 @@ -248,8 +250,8 @@ could accept and finishing the rest elsewhere. A reader lowers the complete transform and can compose through delegation instead. This resembles [zarrita.js store extensions](https://zarrita.dev/packages/zarrita.html), where storage-specific behavior is an explicit extension point rather than an -inferred array capability. The implementation remains independently authored: -no code is shared with TensorStore, xarray, or zarrita.js. +inferred array capability. This is an architectural analogy, not a claim of API or implementation +compatibility. ## Current scope @@ -263,7 +265,7 @@ 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. -Fancy selections compose without restriction: a second `oindex`/`vindex`/mask +Supported fancy selections compose across already-fancy views: a second `oindex`/`vindex`/mask step may land on any axis of an already-fancy view, including axes an existing index array merely broadcasts along, so `lazy.oindex[[2, 0], :].lazy.oindex[:, [1, 3]]` selects the outer product it @@ -273,13 +275,12 @@ which evaluates the existing lookup tables at the new coordinates — rather tha rewritten in place. Resolution classifies the result by structure (`index_array_structure`): pure per-axis outer products keep the orthogonal resolvers, and everything else — correlated maps, mixtures, index arrays -sharing an input axis (a diagonal gather, reachable only by hand-building a -transform) — takes the general reader/intersection path. Chunk planning +sharing an input axis (as in paired vectorized coordinates) — takes the general reader/intersection path. Chunk planning factors index arrays into connected dependency components before flattening, so independent groups do not expand one another. Vectorized selection preserves broadcast singletons to retain those dependencies. -Three limits remain, all intentional and all expected to be lifted: +Some current limits are: - **Affine diagonals.** A hand-built transform in which two output maps read one input dimension — two slice maps, or a slice map and an orthogonal index @@ -287,15 +288,15 @@ Three limits remain, all intentional and all expected to be lifted: varying over a dimension a slice map also reads is rejected with `NotImplementedError`. No selection dialect produces either. Supporting them needs a strided set per *input* dimension spanning all dependent storage - axes, TensorStore's connected-component representation. *Planned.* + axes, TensorStore's connected-component representation. - **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.* + TensorStore supports both. - **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.* + slice. ## Selection to chunk operations diff --git a/packages/zarr-indexing/docs/guide/index.md b/packages/zarr-indexing/docs/guide/index.md index 68192f58cb..b7e748e262 100644 --- a/packages/zarr-indexing/docs/guide/index.md +++ b/packages/zarr-indexing/docs/guide/index.md @@ -216,12 +216,17 @@ description; the assertion's call to `result()` is the first operation in the example that materializes the selected data. !!! warning "Stop here: the materialization boundary" - Indexing through `.lazy[...]` never reads. These do: + Indexing through `.lazy[...]` composes a selection without reading source values. + These operations request values: - `result()` - eager indexing of the wrapper: `view[...]` - - `numpy.asarray(view)`, or passing the view to any NumPy function - (`numpy.add(view, 1)` converts, and therefore materializes, the view) + - `numpy.asarray(view)` and NumPy operations that convert the view + (`numpy.add(view, 1)` does so; `numpy.shape(view)` and `numpy.ndim(view)` + can use metadata without reading values) + + Dask tokenization may also inspect values, depending on the wrapped source + and tokenization path. Python arithmetic such as `view + 1` raises `TypeError` instead: this wrapper defers indexing, not a general compute graph. diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index f787d11f6c..9a79e5e4ae 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -3,8 +3,9 @@ For the complete path from indexing syntax to chunk coordinates, local selectors, and result positions, start with [From a selection to chunk operations](selection-flow.md). -This package supplies indexing plans. It does **not** supply scheduling, -caching, codecs, or async orchestration. A consumer decides when projections +The core package supplies indexing plans and synchronous readers. It does not +provide a general scheduler, codec pipeline, or async execution engine. A +synchronous cache is included as an example. A consumer decides when projections run, how decoded chunks are obtained, and where completed values are retained. An `IndexTransform` says which source values belong in a result; a `Reader` lowers that complete transform for one backend. The reader does not choose @@ -61,21 +62,21 @@ afterward, so accessing a row does not recalculate every point. ## One slab read or many part reads -A backend with its own native subset read — a Rust or C zarr implementation, -a database, an HTTP range endpoint — resolves a **dense box** (`is_box` with -every stride 1) best as a single read: hand it the whole selection and let it -dispatch to chunks, decode in parallel, and partial-decode shards on its own -side of the boundary. Splitting that read along this library's partitioning -only adds round-trips. Every **other** selection — a strided box, an `oindex` -or `vindex` gather — is where the partitioning earns its keep. The **cover** -of a read is the smallest step-1 slab enclosing every coordinate it needs; -partitioned, each part's cover is bounded by that part's box, so a sparse -selection can never force one read of its whole bounding hull (the smallest -rectangle containing every selected coordinate — a thousand rows for the two -of `oindex[[0, 999]]`). +A backend with an efficient native subset operation may benefit from receiving +one complete dense selection so it can choose its own chunk dispatch. Other +backends may benefit from partitioning, including for strided or fancy +selections. The tradeoff depends on the backend, chunk layout, latency, memory, +and selection; a single read is not universally fastest. + +A read's **cover** is the smallest unit-step slab enclosing its coordinates. +With this package's basic readers, partitioning limits each source read to the +part's selected cover. This may reduce over-reading, but a partition that spans +the source can still require the entire hull. Custom readers choose their own +source operations under the reader contract. The composed view carries enough to make that call at materialization time, -and re-partitioning is a pure setter, so the policy is three lines: +and `with_parts()` returns a view with a new partitioning. This example uses +unit strides as a sufficient condition for its independently mapped selections: ```python --8<-- "snippets/integrations.py:dense-box-repartition" @@ -87,10 +88,10 @@ the dense box becomes exactly one backend call. Both regimes go through ### Sources that accept only unit-step slices -The default `basic_reader` pushes strided and descending selections down as -positive-step slices, which reads the minimum but assumes the source accepts -any step. Many backends do not: FFI bindings and range requests often -support nothing but `slice(start, stop, 1)`. Select +For affine selections, `basic_reader` uses positive-step slices and applies +reversal or layout changes in memory. Fancy selections can require reading a +cover containing unselected values. The source must accept the emitted steps. +For a source that accepts only unit steps, use `slice(start, stop, 1)` with [`unit_step_reader`][zarr_indexing.reader.UnitStepReader] for such a source and every key it receives is an ascending unit-step slice per axis, with strides, reversals, and gathers applied to the in-memory block instead: @@ -99,8 +100,8 @@ strides, reversals, and gathers applied to the in-memory block instead: view = LazyArray(source).with_reader(unit_step_reader) ``` -A strided selection then over-reads its cover by the stride factor, which the -partitioning above bounds by one part. +For strided selections, the ratio of cover cells to selected cells depends on +stride, length, and endpoint alignment. Partitioning can reduce that cover. ## napari-like consumer diff --git a/packages/zarr-indexing/docs/guide/patterns.md b/packages/zarr-indexing/docs/guide/patterns.md index 9d22c4eb4e..f2f2ddbe2a 100644 --- a/packages/zarr-indexing/docs/guide/patterns.md +++ b/packages/zarr-indexing/docs/guide/patterns.md @@ -302,10 +302,11 @@ so they equal the zero-origin models after `translate_domain_to`: --8<-- "snippets/indexing_patterns.py:indexing-patterns" ``` -`LazyArray` adds nothing to these semantics: it is a regular array-like API -whose `.lazy`, `.lazy.oindex`, and `.lazy.vindex` accessors compile the same -dialects to the same transforms — the only difference is the return type, a -view instead of an array. The test suite holds the wrapper to this matrix. +`LazyArray` exposes the transform machinery through a positional array-like +API. Its `.lazy`, `.lazy.oindex`, and `.lazy.vindex` accessors return views and +normalize positions before composition. This boundary differs from the literal +coordinate semantics of `IndexTransform`, as the following table shows. The +executable matrix checks the documented cases, not every possible NumPy expression. ## Positions vs literal coordinates diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index 87e5a4b7a1..598282b698 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -5,8 +5,7 @@ the *declaration* of an array indexing expression from the result of that expres Developed for use in [`zarr`](https://zarr.readthedocs.io). -Inspired by [TensorStore](https://google.github.io/tensorstore/), which pioneered -the approach used here. +Inspired by [TensorStore's index-transform model](https://google.github.io/tensorstore/index_space.html). ## Install @@ -28,8 +27,9 @@ you want its values: --8<-- "snippets/canonical_slice.py:landing-quickstart" ``` -Nothing is read until the `result()` call, however many selections are -composed. [Lazy views compose](guide/index.md#lazy-views-compose) shows how +Composing these selections does not read source values; the example reads them +at `result()`. Construction inspects source metadata, and Dask tokenization can +inspect source values. [Lazy views compose](guide/index.md#lazy-views-compose) shows how the chain stays one description, and where the materialization boundary is. ## Learn more diff --git a/packages/zarr-indexing/justfile b/packages/zarr-indexing/justfile index 44874bc0be..24834a960d 100644 --- a/packages/zarr-indexing/justfile +++ b/packages/zarr-indexing/justfile @@ -22,7 +22,7 @@ test-tensorstore *args: uv run --project ../.. --group test --with-editable . --with 'tensorstore>=0.1.84' python -m pytest tests/test_ndsel_tensorstore.py tests/test_tensorstore_parity.py {{ args }} # Lint with the same invocation CI uses. Ruff is pinned to the repo-wide -# version (see pyproject.toml [dependency-groups] docs); bump together. +# version in the root .pre-commit-config.yaml; bump together. lint: uvx ruff@0.16.0 check . @@ -30,7 +30,8 @@ lint: typecheck: uv run --group test --with pyright pyright -# Run everything CI runs for this package +# Run these checks with the locally selected interpreter and dependencies; +# CI also exercises its configured Python-version matrix check: lint typecheck test test-tensorstore docs-check # Preview the changelog that the next release would generate diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index 953a432f07..5a2c9e8fd4 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -48,19 +48,14 @@ Changelog = "https://github.com/zarr-developers/zarr-python/blob/main/packages/z Documentation = "https://zarr-indexing.readthedocs.io/" [dependency-groups] -# The package and its tests import nothing from `zarr`: chunk resolution -# consumes the `DimensionGridLike` protocol and the tests use this package's -# own grids (`zarr_indexing.grid`). The suite is nevertheless run from the -# repo-root environment (`just test`, the invocation CI uses) so that it sees -# the same pinned toolchain — hypothesis, ruff, pyright — as the parent -# project; the repo is not a uv workspace, so this package is layered in as an -# editable overlay. `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. +# Package tests use the local grid protocol and grid implementations. The +# repository recipes layer this package into the root test environment. +# This test group and the optional testing extra both declare Hypothesis; +# keep their supported minimum versions aligned. 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. + # Documentation tooling for this package. These pins are maintained here + # and can differ from the parent project's tooling versions. "mkdocs-material==9.7.7", "mkdocs==1.6.1", "mkdocstrings==1.0.6", @@ -83,7 +78,7 @@ packages = ["src/zarr_indexing"] # An allowlist, so nothing that merely happens to sit in the package directory # — a scratch script, a stray notebook — can ride along in a release. The list -# keeps an sdist self-testing: `tests/` carries the vendored ndsel conformance +# includes test assets in the sdist: `tests/` carries the vendored ndsel conformance # corpus, and `tests/test_doc_examples.py` executes `docs/snippets/*.py` and # `examples/*/*.py`, so those are part of the suite rather than decoration. # `pyproject.toml`, `README.md` and `LICENSE.txt` are added by hatchling itself. @@ -107,14 +102,15 @@ target-version = "py312" # 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. +# noqa comments) because different ruff versions +# have differed on whether these rules fire; RUF100 can remove unused noqa comments. "src/zarr_indexing/lazy_array.py" = ["BLE001", "S110"] [tool.pytest.ini_options] minversion = "7" # src is collected for its doctests: every public object's Examples section -# executes under --doctest-modules, so the documented examples cannot rot. +# is eligible for --doctest-modules. Passing examples do not verify all prose +# contracts or optional-dependency paths. testpaths = ["tests", "src/zarr_indexing"] pythonpath = ["."] xfail_strict = true From eebbe3eab55f832b6b73ae555d895eceb3e30e26 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 13 Sep 2026 18:48:26 +0200 Subject: [PATCH 02/12] docs(indexing): correct reader lazy-array and cache contracts Assisted-by: Codex:GPT-6 --- .../examples/lazy_indexing_dask/README.md | 25 +- .../lazy_indexing_dask/lazy_indexing_dask.py | 26 +- .../examples/lazy_indexing_numpy/README.md | 5 +- .../lazy_indexing_numpy.py | 5 +- .../system_memory_chunk_cache/README.md | 8 +- .../system_memory_chunk_cache.py | 5 +- .../src/zarr_indexing/lazy_array.py | 231 ++++++++---------- .../zarr-indexing/src/zarr_indexing/reader.py | 26 +- .../zarr-indexing/tests/test_doc_examples.py | 4 +- packages/zarr-indexing/tests/test_reader.py | 2 +- 10 files changed, 164 insertions(+), 173 deletions(-) diff --git a/packages/zarr-indexing/examples/lazy_indexing_dask/README.md b/packages/zarr-indexing/examples/lazy_indexing_dask/README.md index c9edb903ef..bdf5220ae0 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_dask/README.md +++ b/packages/zarr-indexing/examples/lazy_indexing_dask/README.md @@ -10,9 +10,9 @@ The example shows how 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 +- Read `is_complete` to inspect coverage of a partition cell +- Inspect `__dask_tokenize__` for the example's equal source/selection pairs; + token equality can support task deduplication but does not promise persistent caching - Measure what a task graph costs for indexing-only work, against composing the same selections into one transform @@ -25,15 +25,16 @@ is discovered from the wrapped array and is independent of Dask's blocks. 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. +For indexing-only workloads, graph construction and scheduling can be an +additional cost. The example measures repeated leading slices and reports graph +layers and timings for the selected Dask version. It does not establish general +complexity bounds or a guaranteed speedup: Dask can optimize graphs, and costs +depend on the selection, chunk layout, and scheduler. + +`LazyArray` stores one composed transform rather than retaining a wrapper for +each prior selection. Applying a chain still costs work for every operation; +index-array composition may process arrays whose size depends on earlier +selections. Reading also incurs partition planning and source I/O. ## Running the Example diff --git a/packages/zarr-indexing/examples/lazy_indexing_dask/lazy_indexing_dask.py b/packages/zarr-indexing/examples/lazy_indexing_dask/lazy_indexing_dask.py index d6ed43f322..e814bc3c76 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_dask/lazy_indexing_dask.py +++ b/packages/zarr-indexing/examples/lazy_indexing_dask/lazy_indexing_dask.py @@ -62,7 +62,8 @@ def test_parts_as_tasks(source: zarr.Array) -> None: 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. + # reads can run concurrently with this source and reader; assembly below + # places the returned blocks sequentially. @dask.delayed def read(part: object) -> np.ndarray: return part.view.result() @@ -75,38 +76,35 @@ def read(part: object) -> np.ndarray: 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. + # the base array. A writer also needs storage-unit alignment, value-order, + # and concurrency checks before using coverage to skip a read. 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.""" + """Check token equality for these unchanged source and selection pairs.""" 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. + # These wrappers have equal tokens despite being different Python objects. + # Source mutation and token hooks affect whether cached results remain valid. 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. + # These two slice chains serialize to the same transform and token. 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. + This measures repeated leading slices at several depths. LazyArray retains + one composed transform, but the loop still performs each selection. Dask's + graph construction and execution costs depend on graph optimizations and + chunk layout; these measurements do not establish general complexity bounds. Timings are printed rather than asserted, since they depend on the machine. """ diff --git a/packages/zarr-indexing/examples/lazy_indexing_numpy/README.md b/packages/zarr-indexing/examples/lazy_indexing_numpy/README.md index e76e065047..26540f4d4f 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_numpy/README.md +++ b/packages/zarr-indexing/examples/lazy_indexing_numpy/README.md @@ -14,8 +14,9 @@ The example shows how to: - 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 +`LazyArray` wraps compatible sources exposing `shape`, `dtype`, and basic +slicing whose results can be converted to NumPy system memory. This includes +a Zarr array, and the partitioning is then discovered from the array's chunks. The Dask example covers that case. ## Running the Example diff --git a/packages/zarr-indexing/examples/lazy_indexing_numpy/lazy_indexing_numpy.py b/packages/zarr-indexing/examples/lazy_indexing_numpy/lazy_indexing_numpy.py index ae29a7fc51..a545be1227 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_numpy/lazy_indexing_numpy.py +++ b/packages/zarr-indexing/examples/lazy_indexing_numpy/lazy_indexing_numpy.py @@ -63,7 +63,7 @@ def test_box_and_query_selections() -> None: data = np.arange(12 * 8).reshape(12, 8) lazy = LazyArray.from_numpy(data) - # A box selection is built from slices and integers alone. It is described + # This box selection is built from slices and integers. 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] @@ -81,7 +81,8 @@ def test_box_and_query_selections() -> None: assert query.strides() is None assert query.bounding_box() == ((1, 10), (0, 8)) - # Composing a box onto a query keeps it a query. + # This slice retains an index-array map, so this particular view stays a query. + # Singleton gathers and later scalar indexing can instead collapse to a box. assert not query.lazy[0:2, 0:2].is_box diff --git a/packages/zarr-indexing/examples/system_memory_chunk_cache/README.md b/packages/zarr-indexing/examples/system_memory_chunk_cache/README.md index d6271f9401..0369f290fb 100644 --- a/packages/zarr-indexing/examples/system_memory_chunk_cache/README.md +++ b/packages/zarr-indexing/examples/system_memory_chunk_cache/README.md @@ -23,7 +23,7 @@ explicit retry. The reader owns cache state and source reads, but not result shape or assembly. Each request calls `view.parts()` once and keeps the resulting tuple. The cache -pins the tuple's chunk coordinates, then materializes with +queues the tuple's chunk coordinates and defers eviction, then materializes with `view.result(parts=parts)`, so scheduling and assembly reuse one plan. Every reader call consumes the exact projection attached to its `ReadContext`; the reader does not invoke the chunk planner again. @@ -34,6 +34,12 @@ a retained failure that does not retry implicitly, and an explicit retry after the source is repaired. The integration guide contains the detailed request table. +Capacity counts resident chunks, not bytes, and is enforced after a successful +request. A request can temporarily exceed it; a failed request skips that eviction +step. Records, event history, and output/coordinate buffers are outside this count. +The example assumes an unchanged source and one source/grid per reader; it has no +cache invalidation for source mutation, and its mutable state is not thread-safe. + This is synchronous system-memory reference architecture, not a production-ready cache, scheduler, renderer, or complete napari integration. Its types are intentionally not exported by `zarr_indexing`. diff --git a/packages/zarr-indexing/examples/system_memory_chunk_cache/system_memory_chunk_cache.py b/packages/zarr-indexing/examples/system_memory_chunk_cache/system_memory_chunk_cache.py index 1645d89b2c..e3aad6669a 100644 --- a/packages/zarr-indexing/examples/system_memory_chunk_cache/system_memory_chunk_cache.py +++ b/packages/zarr-indexing/examples/system_memory_chunk_cache/system_memory_chunk_cache.py @@ -182,7 +182,7 @@ def retry(self, chunk_coords: ChunkCoords) -> None: @contextmanager def request(self, required: tuple[ChunkCoords, ...]) -> Iterator[None]: - """Prepare every part and defer eviction until one request completes.""" + """Prepare parts; evict after the outermost request succeeds, not on failure.""" self._prepare(required) self._requests += 1 try: @@ -331,7 +331,8 @@ def _read(self, key: Any, *, orthogonal: bool) -> np.ndarray[Any, Any]: self.reader.projection_uses.clear() lazy = self._lazy.lazy view = lazy.oindex[key] if orthogonal else lazy[key] - # One prepared tuple is the request plan: pin from it, then hand the + # One prepared tuple is the request plan: queue its chunks and defer + # eviction during the request, then hand the # same owned parts back to LazyArray for assembly without replanning. parts = tuple(view.parts()) required = tuple(dict.fromkeys(part.base_coords for part in parts)) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index cae10ba65c..c5c2967b94 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -11,9 +11,12 @@ values = view.result() ``` -Nothing is read until `result()` (or `__array__`, or an eager `__getitem__`). -Every `.lazy` operation is metadata-only. Composition does not accumulate -layers: a view of a view is still a single transform and retains its reader. +Selection construction does not read source values: `result()`, `__array__`, +and eager `__getitem__` perform reads. Tokenization can also read or hash source +data, depending on the source and tokenization path. `.lazy` operations inspect +selection metadata and may copy or process supplied index arrays. Composition +does not accumulate wrapper layers: a view of a view is still a single transform +and retains its reader. Parts ----- @@ -26,8 +29,10 @@ into the final buffer or an owned temporary for fancy placement. The part view's transform directly addresses its raw wrapped array. The paired -projection deliberately retains the chunk-local frame; both travel together in -the `ReadContext` passed to the reader. +projection deliberately retains the chunk-local frame. Parent materialization +passes both in `ReadContext`; calling `part.view.result()` directly uses an +unpartitioned context with `projection=None`. Readers that require the projection +should be used through the parent `result(parts=parts)` path. The partitioning is discovered from the wrapped array at construction — first `read_chunk_sizes` (zarr's clipped per-axis sizes, sharding-aware), then @@ -47,8 +52,9 @@ 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. +useful for controlling per-read sizes or batching small reads. They can change +I/O costs; the intended selected values remain the same for an unchanged source +and a conforming reader. The full result buffer is still allocated. Readers ------- @@ -64,63 +70,52 @@ parts. Consumers may materialize part views concurrently; `LazyArray` does not serialize calls, so a stateful reader must synchronize its own mutable state. -Both built-in readers lower through NumPy system memory. They do not implicitly -transfer device arrays. A device source requires an explicit custom reader that -performs any needed transfer into the supplied system-memory output buffer. +The built-in readers lower through NumPy system memory using NumPy conversion +of source slices. Device arrays that refuse implicit conversion need a custom +reader that explicitly transfers values into the output buffer. Boxes and queries ----------------- -A selection is either **rectangular** — an interval and a stride per dimension, -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: 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. +`is_box` reports whether the current transform contains only constant and affine +output maps. An index-array gather commonly produces a query, but singleton +indices and later selections can remove its index-array maps and make it a box. +`bounding_box()` gives a coordinate hull, and `strides()` gives stride magnitudes +for a box. These describe storage coordinates, not the full result layout or +traversal order. A query can fill its hull, and a singleton box can fill its hull +even when its recorded stride exceeds one. 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 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 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. Any non-boolean object implementing Python's - `SupportsIndex` protocol is accepted as one, including in slice bounds and - steps; an `__int__` method alone is deliberately not enough. A scalar 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 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]`. +view's shape, integer coordinates are bounds-checked, and slices are clipped +to the view's extent. + +This differs from the low-level `IndexTransform` literal-coordinate dialect: +a transform can retain a nonzero domain origin, while `LazyArray` re-zeroes +positions on each derived view. The current main Zarr `Array` does not expose +this wrapper as an `Array.lazy` attribute; use `LazyArray(array)` explicitly. + +Scalar integers drop axes. Non-boolean objects implementing `SupportsIndex` +are accepted as scalar indices and in slice bounds; `__int__` alone is not enough. +For orthogonal and vectorized modes this wrapper applies scalar indices first, +then the remaining advanced selection. Orthogonal indices form an outer product. +Vectorized indexing accepts coordinate arrays or a shape-matching boolean mask. +An ellipsis can retain unindexed axes (`[..., i, j]`), but explicit slice entries +in vectorized selections are currently rejected. Scalar-first processing can +also differ from NumPy's advanced-axis placement: for shape `(2, 3, 4)`, +`lazy.vindex[0, ..., [1, 2]]` has shape `(3, 2)`, whereas NumPy's same selection +has shape `(2, 3)`. These modes do not implement every NumPy indexing form. Materializing on fallback ------------------------- `LazyArray` implements `__array__` but deliberately implements neither -`__array_ufunc__` nor `__array_function__`. A NumPy *function* given a view -therefore materializes the whole thing through -`__array__` and works on the resulting array: `numpy.sum(view)`, +`__array_ufunc__` nor `__array_function__`. Many NumPy operations therefore materialize through `__array__` and work 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. +Metadata queries such as `numpy.shape(view)` and `numpy.ndim(view)` can use the +exposed attributes without materializing; other unsupported operations may fail. Python's arithmetic *operators* do not: `view + 1` raises `TypeError`, because the wrapper defines no arithmetic dunders and an `int` has nothing to dispatch @@ -178,8 +173,8 @@ __all__ = ["LazyArray", "Partition"] -# Above this many bytes, the no-dask token fallback describes an array -# structurally instead of digesting its contents. See `_wrapped_token`. +# Above this declared byte count, the no-Dask token fallback adds a fresh UUID +# instead of digesting contents. See `_wrapped_token`. _TOKEN_DIGEST_LIMIT = 1 << 20 @@ -244,8 +239,9 @@ def _discover_parts(array: Any, shape: tuple[int, ...]) -> tuple[DimensionGrid, 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 a public API and validates strictly. + With a compatible reader, resolving the view without partitioning preserves + its values, though it may require larger reads. `with_parts` is a public API + and validates strictly. """ declared = _read_source_attribute(array, "read_chunk_sizes") if declared is None: @@ -415,12 +411,14 @@ class Partition: Yielded by [`LazyArray.parts`][zarr_indexing.lazy_array.LazyArray.parts]. The parts of a view tile it exactly and disjointly: assembling every `view.result()` at its `out_selection` reproduces the whole view's - `result()`, and each part can be resolved independently and concurrently. + `result()` for readers supporting contexts without projections. Parts can be + resolved concurrently when the source and reader permit it. Derived parts retain the same reader object; a shared stateful reader owns synchronization for concurrent calls. A consumer that needs the plan before materialization can prepare it once - and reuse the same immutable parts for both scheduling and assembly: + and reuse the same part records for both scheduling and assembly. The frozen + records do not snapshot the mutable source or reader: ```python parts = tuple(view.parts()) @@ -464,7 +462,8 @@ 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. Fancy projections + whether coverage is complete; it does not by itself establish buffer order, + storage-unit alignment, or concurrent-write safety. Fancy projections report `False` because their coverage is deliberately `unknown` until duplicate-aware proof is added. @@ -539,23 +538,20 @@ def _wrapped_token(array: Any) -> Any: 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. + Tokens can differ depending on whether Dask is available and on the source + hook. This fallback does not provide a portable content identifier. + + Above `_TOKEN_DIGEST_LIMIT`, or when conversion is unavailable, the local + fallback adds a fresh UUID on each call, so repeated calls normally differ. + The returned tuple is still equal to itself. Below the limit, conversion + can read the source, and the hash is of its NumPy buffer bytes. Object-array + buffer bytes contain object references, not a recursive content snapshot. """ hook = getattr(array, "__dask_tokenize__", None) if hook is not None: try: return hook() - # A token must never raise; fall through to the structural fallback. + # A failing source hook falls through to the remaining tokenization paths. except Exception: # pragma: no cover - a hook that refuses to run pass try: @@ -570,7 +566,7 @@ 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 + # A fresh identifier per call when 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) @@ -583,7 +579,7 @@ def _wrapped_token(array: Any) -> Any: return unidentified try: contents = np.ascontiguousarray(array) - # A token must never raise; an unreadable source is simply unidentified. + # Failed NumPy conversion leaves this source unidentified. except Exception: return unidentified return (*structural, hashlib.sha256(contents.tobytes()).hexdigest()) @@ -605,8 +601,8 @@ class LazyArray: along a **partitioning** discovered from the wrapped array. Every derived view retains its reader; that reader receives the complete projected transform once per part. 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. + dialect differs from low-level literal transforms and which NumPy operations + materialize the view. This wrapper describes **reads**. It defines no `__setitem__`, so assigning into a view raises `TypeError`. Writing belongs to the @@ -643,7 +639,8 @@ class LazyArray: def __init__(self, array: _WrappedArray) -> None: """Wrap `array` without reading it; parameters are documented on the class. - The only validation here is the `numpy.matrix` rejection (`TypeError`). + Reject `numpy.matrix`, normalize the shape, and inspect partition metadata. + Shape conversion and transform construction can also reject invalid input. """ if isinstance(array, np.matrix): # `np.matrix` keeps every result two-dimensional, so `m[1]` has shape @@ -666,7 +663,7 @@ def __init__(self, array: _WrappedArray) -> None: @classmethod def from_numpy(cls, array: np.ndarray[Any, Any]) -> LazyArray: - """Wrap a NumPy array with its explicitly selected optimized reader.""" + """Wrap a NumPy array with its explicitly selected NumPy reader.""" if not isinstance(cast(object, array), np.ndarray): raise TypeError( f"LazyArray.from_numpy requires a numpy.ndarray, got {type(array).__name__}" @@ -769,26 +766,13 @@ def with_reader(self, reader: Reader) -> LazyArray: 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, - 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 - 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 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 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. + True exactly when the composed transform has no `ArrayMap` output maps. + Basic indexing of a box stays a box. Advanced indexing can introduce + index arrays, but singleton gathers and later selections may remove them. + + `bounding_box()` and `strides()` describe the touched coordinate region + and stride magnitudes; use the transform for order, result axes, and + repetitions. Neither this flag nor the hull alone proves dense coverage. Examples -------- @@ -807,12 +791,10 @@ def bounding_box(self) -> tuple[tuple[int, int], ...] | None: this view reads from that contains every coordinate the selection reaches. - 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 a superset that can be - arbitrarily loose: `oindex[[0, 999]]` has a 1000-wide hull over two - rows. + A strided box can leave gaps, and a query hull can be arbitrarily loose: + `oindex[[0, 999]]` spans a 1000-wide hull over two selected rows. Conversely, + a query may cover every coordinate in its hull. Inspect the transform + when coverage or traversal order matters. Returns ------- @@ -859,7 +841,7 @@ def bounding_box(self) -> tuple[tuple[int, int], ...] | None: def strides(self) -> tuple[int, ...] | None: """The step between selected coordinates, one per storage dimension. - Together with `bounding_box()`, this fully describes a box selection: + Together with `bounding_box()`, this describes a box's per-axis coordinates: `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 @@ -1111,8 +1093,8 @@ def lazy(self) -> _LazyIndexer: 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. + # This frontend applies scalar integers first, dropping their axes, + # then normalizes its supported advanced-indexing forms. scalar_selection, selection = split_scalar_axes(selection, transform.domain, mode) if scalar_selection is not None: transform = transform.select(scalar_selection, "basic") @@ -1214,7 +1196,9 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any: # 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. + # a gap alone undercounts and an overlap alone overcounts. This + # count cannot detect a compensating gap and overlap; disjointness + # still relies on the planner contract. if prepared_parts is not None: raise ValueError( "prepared parts do not tile the view exactly: " @@ -1230,7 +1214,8 @@ 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 + as required by the planner contract. `result()` checks the total count; + caller-supplied parts additionally undergo coverage validation. A masked source gets a masked buffer so that reader writes preserve the mask; other source- specific array types do not survive materializing. """ @@ -1246,9 +1231,9 @@ def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: The result never shares memory with the wrapped array, whatever `copy` asks for: `result()` already allocates, 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. + caller owns. `copy=None` need not copy again when the dtype already + matches; requesting another dtype can allocate a conversion buffer. `copy=False` is refused, because materializing means reading + — this API always materializes into its own result buffer. """ if copy is False: raise ValueError( @@ -1259,22 +1244,18 @@ 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 and the view. - - Two wrappers produce equal tokens when they wrap the same data and - address the same cells. The view contributes a digest of its canonical - ndsel body, so transforms that differ only in representation produce - the same token, and a fancy selection with a large index array does not - embed that array's JSON in the 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 reader are deliberately absent. Both decide how - the data is read — in which boxes, and through which request strategy — - 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. + """Tokenize the source and serialized transform. + + The transform contributes a digest of JSON produced by `to_json()`, so + its JSON is not embedded in the returned token, but is allocated while + computing the digest. Equal serialized transforms and equal source + tokens produce equal tokens; arbitrary semantically equivalent mappings + are not guaranteed to serialize identically. + + Source tokenization can read or hash data and need not be deterministic + on every fallback path; see `_wrapped_token`. The reader and partitioning + are omitted under the contract that they preserve values. Cache users + must also account for source mutation and the source's token semantics. """ canonical = json.dumps(self._transform.to_json(), sort_keys=True) return ( diff --git a/packages/zarr-indexing/src/zarr_indexing/reader.py b/packages/zarr-indexing/src/zarr_indexing/reader.py index 8d47d51d21..2527972e02 100644 --- a/packages/zarr-indexing/src/zarr_indexing/reader.py +++ b/packages/zarr-indexing/src/zarr_indexing/reader.py @@ -103,7 +103,8 @@ class BasicReader: Each transform is decomposed into the smallest enclosing positive-slice slab and a residual transform. The slab is read once with basic indexing, - so fancy or negative-step selections may over-read, and the residual is + so fancy selections may over-read. Reversals read their selected coordinates + in ascending order without requiring a negative source slice. The residual is then lowered through NumPy system-memory operations into the supplied buffer. @@ -132,12 +133,13 @@ def read_into(self, source: Any, context: ReadContext, out: Any, /) -> None: class NumPyReader: - """Reader optimized for NumPy system-memory arrays. + """Explicit reader for NumPy system-memory arrays. This is the reader selected by [`LazyArray.from_numpy`][zarr_indexing.lazy_array.LazyArray.from_numpy]. It applies the complete transform with NumPy operations and is applicable to - `numpy.ndarray` sources, including `numpy.ma.MaskedArray`. + `numpy.ndarray` sources, including `numpy.ma.MaskedArray`. Its current + implementation uses the same slab and residual path as `BasicReader`. Examples -------- @@ -166,12 +168,13 @@ class UnitStepReader: unit-step slab and a residual transform, so the source only ever receives `slice(start, stop, 1)` on every axis — the one form an API without general strided reads (an FFI binding, an HTTP range endpoint) supports. - `BasicReader` instead pushes strided and reversed slices down, which - reads less but asks more of the source. + `BasicReader` instead requests positive-stride slices for both forward and + reversed affine selections, then reverses in memory when needed. The residual lowering applies strides, reversals, and gathers to the - in-memory block, so a strided selection over-reads its cover by the - stride factor. Partitioning the wrapping + in-memory block. For n selected points at positive spacing k on one axis, + the cover contains (n - 1) * k + 1 elements when n is nonzero; its over-read + ratio approaches k for long selections. Partitioning the wrapping [`LazyArray`][zarr_indexing.lazy_array.LazyArray] (`with_parts`) bounds each cover by a part. @@ -489,8 +492,8 @@ def _push_slice_for_dimension_map( """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. + only asked for positive-stride slices. Supporting those slices is part of + this reader's source contract. """ d = m.input_dimension lo = transform.domain.inclusive_min[d] @@ -524,9 +527,8 @@ def _push_unit_slice_for_dimension_map( Strides and reversals stay in the residual: the source is only ever asked for a contiguous ascending slice, and the original stride is replayed - against the in-memory block. The cover therefore over-reads a strided - selection by its stride factor, which is the price of a source that - accepts nothing but `slice(start, stop, 1)`. + against the in-memory block. The cover may include unselected elements; + for n distinct points at spacing k it has (n - 1) * k + 1 elements. """ d = m.input_dimension lo = transform.domain.inclusive_min[d] diff --git a/packages/zarr-indexing/tests/test_doc_examples.py b/packages/zarr-indexing/tests/test_doc_examples.py index a4833296da..802e39e7f8 100644 --- a/packages/zarr-indexing/tests/test_doc_examples.py +++ b/packages/zarr-indexing/tests/test_doc_examples.py @@ -72,7 +72,7 @@ def _markdown_includes() -> tuple[tuple[str, str, str | None], ...]: PATTERN_CASES: tuple[dict[str, Any], ...] = PATTERN_NAMESPACE["PATTERN_CASES"] CACHE_NAMESPACE: dict[str, Any] = runpy.run_path(str(CACHE_EXAMPLE)) -# The documented pattern matrix must keep covering every selection family. +# The documented pattern matrix must keep its listed selection families. REQUIRED_PATTERNS = { "basic-slice", "integer-axis-removal", @@ -133,7 +133,7 @@ def test_documentation_example_executes(example: Path) -> None: @pytest.mark.parametrize("script", CLI_EXAMPLES, ids=lambda path: path.stem) def test_cli_example_runs_as_a_subprocess(script: Path) -> None: - """The CLI examples exit 0 when run the way their READMEs instruct.""" + """CLI examples exit 0 with this test environment's installed dependencies.""" if "dask" in script.stem: pytest.importorskip("dask.array") completed = subprocess.run( diff --git a/packages/zarr-indexing/tests/test_reader.py b/packages/zarr-indexing/tests/test_reader.py index 11241a316c..115cebc52e 100644 --- a/packages/zarr-indexing/tests/test_reader.py +++ b/packages/zarr-indexing/tests/test_reader.py @@ -370,7 +370,7 @@ def test_empty_domain_composed_fancy_transform_reads_as_empty() -> None: def test_unit_step_reader_reads_through_lazy_array() -> None: - """The full dialect resolves through a source that only accepts unit-step slices. + """The parametrized selections resolve through a unit-step-only source. `UnitStepOnlySource` asserts the shape of every key it receives, so each selection here also proves no strided, descending, or non-slice key From ba5e0808cf8871976b4f9d7169b20ad92a07ae95 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 13 Sep 2026 18:49:20 +0200 Subject: [PATCH 03/12] fix(indexing): validate wire boundaries and clarify format contracts Assisted-by: Codex:GPT-6 --- packages/zarr-indexing/docs/ndsel.md | 63 +++++++++++-------- .../zarr-indexing/src/zarr_indexing/_wire.py | 35 ++++++++--- .../zarr-indexing/src/zarr_indexing/errors.py | 4 +- .../zarr-indexing/src/zarr_indexing/json.py | 22 ++++--- .../src/zarr_indexing/messages.py | 43 +++++++------ .../src/zarr_indexing/testing/stateful.py | 26 ++++---- .../src/zarr_indexing/testing/strategies.py | 12 ++-- packages/zarr-indexing/tests/test_json.py | 5 +- .../tests/test_ndsel_tensorstore.py | 8 +-- .../tests/test_wire_boundaries.py | 63 +++++++++++++++++++ 10 files changed, 191 insertions(+), 90 deletions(-) create mode 100644 packages/zarr-indexing/tests/test_wire_boundaries.py diff --git a/packages/zarr-indexing/docs/ndsel.md b/packages/zarr-indexing/docs/ndsel.md index 94971d49bf..090dce64c4 100644 --- a/packages/zarr-indexing/docs/ndsel.md +++ b/packages/zarr-indexing/docs/ndsel.md @@ -14,11 +14,13 @@ before a transform is serialized. `zarr-indexing` implements ndsel in two layers | Layer | Module | Depends on | Job | | --- | --- | --- | --- | -| Message | [`zarr_indexing.messages`](api/messages.md) | stdlib only | JSON in, canonical JSON out. Validates and desugars. Never rounds, clamps, or drops information. | +| Message | [`zarr_indexing.messages`](api/messages.md) | stdlib only | Validates and desugars JSON, removing redundant constant-map fields. | | 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. As a result, +arrays as `ndarray`s — are checked during engine lowering. The message layer +also limits input rank to 32 and checks affine input-dimension references. +These checks are stricter than the draft's required message validation. As a result, `messages` normalizes a message with `"-inf"` bounds that `IndexTransform.from_json` refuses to lower. @@ -47,9 +49,13 @@ normalize_ndsel({"kind": "box", "inclusive_min": [10, 5], "shape": [40, 1]}) ``` Normalization is idempotent: re-tag the output with `kind: "transform"` and -normalizing it again returns the same body. Because the canonical body is -field-for-field a TensorStore `IndexTransform` minus `kind`, a normalized -message loads directly into `tensorstore.IndexTransform(json=...)`. +normalizing it again returns the same body. The canonical body uses +TensorStore's `IndexTransform` field vocabulary, but normalization does not +guarantee acceptance by TensorStore. For example, the message layer leaves +index-array content unchecked, whereas TensorStore validates it; TensorStore +also requires unique nonempty labels and restricts finite index values to +`[-(2**62 - 2), 2**62 - 2]`. See +[TensorStore's index-space constraints](https://google.github.io/tensorstore/index_space.html). Both entry points raise [`NdselError`](api/messages.md#zarr_indexing.messages.NdselError), which @@ -64,12 +70,13 @@ Four are shorthands; the fifth is the canonical form itself. | `kind` | Fields | Selects | | --- | --- | --- | | `point` | `coords` | A single element. Normalizes to rank 0 with one `constant` output map per dimension. | -| `box` | `inclusive_min`, one of `exclusive_max` / `inclusive_max` / `shape`, `labels` | A rectangular region. Exactly one upper-bound spelling may appear. | +| `box` | `inclusive_min`, one of `exclusive_max` / `inclusive_max` / `shape`, `labels` | A rectangular region. At most one upper-bound spelling may appear; omission gives implicit positive infinity. | | `slice` | `start`, `stop`, `step`, `labels` | A strided region, one Python-style slice per dimension. | | `points` | `coords` (a list of coordinate rows) | An explicit list of points — the `vindex` case. Normalizes to one `index_array` output map per dimension over a shared rank-1 input domain. | | `transform` | `input_rank`, `input_inclusive_min`, one of the three `input_*` upper bounds, `input_labels`, `output` | The full canonical form. | -Value rules the message layer enforces throughout: every integer is a 64-bit +Value rules for validated fields (excluding verbatim `index_array` payloads +and discarded constant-map fields): every integer is a 64-bit signed value; JSON booleans are **not** integers (Python's `isinstance(True, int)` is guarded against explicitly); the `"-inf"` / `"+inf"` sentinels are legal only in bound positions; and an implicit bound is the @@ -81,17 +88,21 @@ normalization intact. The engine layer converts between canonical bodies and `IndexTransform`s: ```python -from zarr_indexing import IndexTransform +from zarr_indexing import IndexTransform, normalize_ndsel +canonical = normalize_ndsel({"kind": "box", "shape": [2, 3]}) t = IndexTransform.from_json(canonical) -t.to_json() == canonical +assert t.to_json() == canonical ``` `IndexDomain` carries the same pair for a bare domain body, and each output map kind has a `to_json`; `output_index_map_from_json` dispatches the wire's -tagged union back to the right kind. +structurally discriminated union back to the right kind. Exact JSON equality +in this example is not a general round-trip guarantee: implicit flags are +removed, finite `index_array_bounds` are not retained or enforced by the +engine, and degenerate array maps are collapsed. -Two engine constraints apply here and only here. A canonical body carrying a +A canonical body carrying a `"-inf"` or `"+inf"` bound cannot be lowered — an `IndexDomain` addresses a finite array — so `IndexTransform.from_json` raises. And implicit bounds lower *by value*: the `[n]`-bracket flag is a message-layer concern, and the engine @@ -102,21 +113,17 @@ keeps only the integer. ndsel and TensorStore both **reject** an output map that carries both `input_dimension` and `index_array`. The in-memory [`ArrayMap`](api/output_map.md#zarr_indexing.output_map.ArrayMap), though, -records an `input_dimension` to pin the axis an orthogonal (`oindex`) array -varies over. The serializer bridges that gap in both directions: +records its dependency axes in its full-rank array shape, with no +`input_dimension` field: - **On serialize**, a non-degenerate `index_array` map is emitted *without* `input_dimension`. -- **On load**, the in-memory `input_dimension` is reconstructed from the - full-rank array's dependency axes — its non-singleton axes. An array that - 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 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 -representation rather than preserving it. An all-singleton `index_array` — size +- **On load**, dependency axes are the full-rank array's non-singleton axes. + Maps sharing these axes describe correlated coordinates. The engine also + accepts lower-rank nonempty arrays by prepending singleton axes; that + convenience is not a guarantee of compatibility with other ndsel consumers. + +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: @@ -137,6 +144,11 @@ 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. +Empty index arrays also serialize as constant maps with offset zero: the +empty input domain carries the fact that no coordinates are selected. This +avoids losing trailing shape information in JSON when an array has a leading +zero-length axis. + ## Conformance The package is checked against the language-agnostic ndsel conformance corpus, @@ -145,8 +157,9 @@ vendored unmodified under — one JSON file per message kind plus `errors.json`, with the source commit recorded in `PROVENANCE.md`. Each fixture is either a *success* case (`input` + expected `normalized` body) or an *error* case (`input` + expected -reason code), and an implementation is conformant iff `normalize` reproduces -every one. `tests/test_conformance.py` runs the whole corpus as one +reason code). Matching every fixture establishes corpus conformance; the +fixtures do not prove correctness for every possible input or universal +TensorStore compatibility. `tests/test_conformance.py` runs the whole corpus as one parametrized test per fixture, so a corpus update reports failures fixture by fixture rather than as a single opaque assertion. diff --git a/packages/zarr-indexing/src/zarr_indexing/_wire.py b/packages/zarr-indexing/src/zarr_indexing/_wire.py index c3e6fd82b6..e2fc28d73c 100644 --- a/packages/zarr-indexing/src/zarr_indexing/_wire.py +++ b/packages/zarr-indexing/src/zarr_indexing/_wire.py @@ -47,12 +47,28 @@ def lower_index_array(raw: Any, where: str) -> np.ndarray[Any, np.dtype[np.intp] 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. + # The wire representation requires a nested array, not a scalar. raise NdselError( "invalid_json", f"{where} must be an array of integers, got {raw!r}", ) + # Validate before NumPy inference can coerce mixed booleans to integers or + # conversion to intp can wrap unsigned coordinates. + limits = np.iinfo(np.intp) + pending = [raw] + while pending: + for value in pending.pop(): + if isinstance(value, list): + pending.append(value) + elif isinstance(value, (bool, np.bool_)) or not isinstance(value, (int, np.integer)): + dtype = np.asarray(value).dtype.name + raise NdselError( + "invalid_json", f"{where} must hold integers, got {dtype}: {value!r}" + ) + elif not limits.min <= value <= limits.max: + raise NdselError( + "invalid_json", f"{where} coordinate {value} is outside intp range" + ) try: arr = np.asarray(raw) except (TypeError, ValueError) as exc: @@ -88,16 +104,17 @@ def full_rank_index_array( ) -> np.ndarray[Any, np.dtype[np.intp]]: """Give an incoming `index_array` the input rank the engine requires. - 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. + ndsel intends index arrays to have the input rank, but defers validating + that constraint. This engine also accepts lower-rank arrays, aligning them + to the *trailing* input dimensions as in NumPy broadcasting. This extension + does not imply other ndsel consumers accept the same document. 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 + varies over cannot be read off it. When ranks differ, this engine recovers + a shape with the domain's single empty axis and singleton axes elsewhere; + it rejects recovery if the domain has zero or multiple empty axes. + 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. """ diff --git a/packages/zarr-indexing/src/zarr_indexing/errors.py b/packages/zarr-indexing/src/zarr_indexing/errors.py index efd3b1ecd6..74780aa4fc 100644 --- a/packages/zarr-indexing/src/zarr_indexing/errors.py +++ b/packages/zarr-indexing/src/zarr_indexing/errors.py @@ -5,8 +5,8 @@ `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. +zarr's class does not catch this package's errors. Import these from here, +or catch their shared built-in `IndexError` base to handle both libraries. """ from __future__ import annotations diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py index 8bf42c74a5..70d29e1585 100644 --- a/packages/zarr-indexing/src/zarr_indexing/json.py +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -1,7 +1,7 @@ """The canonical ndsel wire vocabulary, and the rules for lowering it. -This is the **engine layer**. Where `messages.py` is pure JSON→JSON and imposes -no array constraints, this module holds the JSON shapes a canonical ndsel body +This is the **engine layer**. Where `messages.py` is pure JSON→JSON, +this module holds the JSON shapes a canonical ndsel body takes (spec section 4.3, as produced by `zarr_indexing.messages.normalize_ndsel`) together with the lowering rules that turn one into the numpy-backed engine representation. @@ -10,9 +10,9 @@ its one serialization: `IndexTransform.to_json` / `from_json`, `IndexDomain.to_json` / `from_json`, and `to_json` on each output map kind, with `output_index_map_from_json` in `zarr_indexing.output_map` dispatching the -wire's tagged union back to the right kind. This module is what they share. +wire's structurally discriminated union back to the right kind. This module is what they share. -Three engine constraints live **here and only here**: +The engine lowering rules include: - **Finite bounds.** An `IndexDomain` addresses a finite array, so a canonical body carrying a `"-inf"`/`"+inf"` bound cannot be lowered; `from_json` raises. @@ -42,10 +42,10 @@ 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 — + same: `t[ts.d[0][[]]]` is `out[0] = 0`, emitted as `{}`. This avoids losing + trailing shape information: `ndarray.tolist()` renders every empty array as + `[]` once the leading axis is zero-length, and nested lists cannot spell + the shape back — `[[]]` is `(1, 0)` and nothing spells `(0, 1)`. 3. Non-degenerate `index_array` maps are emitted with their array and bounds only. @@ -130,7 +130,11 @@ class OutputIndexMapJSON(TypedDict, total=False): """Nested lists of output coordinates, one nesting level per input dimension.""" index_array_bounds: list[IndexValueJSON] - """Bounds the `index_array` values are promised to lie in; `["-inf", "+inf"]` if unconstrained.""" + """Wire bounds on index-array values; `["-inf", "+inf"]` if unconstrained. + + The engine currently discards this field on load and does not enforce + finite bounds against the array values. Serialization emits unconstrained bounds. + """ class IndexTransformJSON(TypedDict, total=False): diff --git a/packages/zarr-indexing/src/zarr_indexing/messages.py b/packages/zarr-indexing/src/zarr_indexing/messages.py index df0e0c87eb..6cd5447cc1 100644 --- a/packages/zarr-indexing/src/zarr_indexing/messages.py +++ b/packages/zarr-indexing/src/zarr_indexing/messages.py @@ -3,10 +3,10 @@ This module implements the [ndsel](https://github.com/zarr-developers/ndsel) draft wire format: a JSON-serializable representation of NumPy-style n-dimensional selections that adapts TensorStore's `IndexTransform` model. It is a **pure -JSON→JSON** layer: it depends on nothing but the standard library, imposes no -engine (numpy/array) constraints, and never rounds, clamps, or drops -information. Engine constraints (finite bounds, in-memory `IndexTransform` -construction) live one layer up, in `json.py`. +JSON→JSON** layer depending only on the standard library. It validates and +desugars messages, removing redundant fields on constant maps. It limits input +rank to 32 and checks affine references against that rank. Finite bounds and +in-memory array construction are enforced by the engine lowering layer. Two entry points: @@ -18,11 +18,12 @@ `IndexTransform` JSON body, without the `kind` discriminator. `normalize` is idempotent when its output is re-tagged with `kind: "transform"`. -The canonical body is, field-for-field, a TensorStore `IndexTransform` (minus -`kind`), so a normalized `transform` loads directly into TensorStore once -`kind` is stripped. +The canonical body uses TensorStore's `IndexTransform` field vocabulary. +Normalization alone does not guarantee TensorStore acceptance: index-array +content is deferred, and TensorStore has additional coordinate and label limits. -Value rules enforced here: every integer is a 64-bit signed value; JSON +Value rules for validated fields (excluding the verbatim `index_array` payload +and discarded constant-map fields): integers are 64-bit signed values; JSON booleans are **not** integers (Python's `isinstance(True, int)` is guarded against explicitly); the `"-inf"`/`"+inf"` sentinels are legal only in bound positions; an implicit bound is the one-element `[n]`-bracket form, and its @@ -43,7 +44,7 @@ # Error taxonomy # --------------------------------------------------------------------------- -#: The complete set of ndsel reason codes (spec section 6). +#: Current ndsel reason codes plus the recognized retired negative-step code. REASON_CODES = frozenset( { "invalid_json", @@ -117,7 +118,7 @@ def __init__(self, reason: str, detail: str = "") -> None: # 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. +# TensorStore accepts. This is an implementation limit, not an ndsel limit. _MAX_RANK = 32 @@ -480,8 +481,8 @@ def _normalize_slice(obj: dict[str, Any]) -> dict[str, Any]: 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) + inclusive_min.append(_checked_i64(o, f"input_inclusive_min[{k}]")) + exclusive_max.append(_checked_i64(o + m, f"input_exclusive_max[{k}]")) output.append({"offset": offset, "stride": s, "input_dimension": k}) labels = labels_raw if labels_raw is not None else [""] * n @@ -730,17 +731,23 @@ def normalize_ndsel(obj: Any) -> dict[str, Any]: """ message = _require_object(obj) kind = _message_kind(message) - return _NORMALIZERS[kind](message) + canonical = _NORMALIZERS[kind](message) + # Apply the same limit to inferred and shorthand ranks as to explicit + # transform ranks, so every result can be normalized again. + if canonical["input_rank"] > _MAX_RANK: + raise NdselError( + "invalid_json", f"input_rank must be <= {_MAX_RANK}, got {canonical['input_rank']}" + ) + return canonical def parse_ndsel(obj: Any) -> dict[str, Any]: """Structurally validate an ndsel message, returning it unchanged. - A lighter gate than `normalize_ndsel`: it confirms the message is a - well-formed ndsel message of a recognized kind (correct field membership, - JSON types, upper-bound exclusivity, domain ordering, step signs) and - raises `NdselError` otherwise, but does not desugar it. Useful for - validating a message you intend to keep in its compact shorthand form. + Runs the same validation and desugaring as `normalize_ndsel`, discards + the canonical body, and returns the original object. Useful for validating + a message you intend to keep in its compact shorthand form. Index-array + payload validation remains the engine's responsibility. Examples -------- diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 4458d7922a..b407f121e9 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -39,7 +39,7 @@ def make_source(self, data): The `choose_reader` rule draws a reader and applies it to the view, so the execution strategy becomes part of the chain. Every reader listed by a subclass must preserve the NumPy model for its source. The universal `basic_reader` is -always exercised, even when a subclass lists only specialized readers; with no +always eligible to be drawn, even when a subclass lists only specialized readers; with no declared readers it is the sole strategy drawn. Requires the `testing` extra (`pip install zarr-indexing[testing]`). @@ -105,13 +105,12 @@ def make_source(self, data): 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. +"""A budget of 250 examples with at most 10 state-machine steps per example. -`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.""" +Invariants materialize the whole view after steps. Deadlines and selected +health checks are disabled to accommodate those potentially expensive reads; +the budget does not guarantee discovery of every defect. +""" # --------------------------------------------------------------------------- # @@ -244,8 +243,9 @@ def __init__(self) -> None: 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. + This harness stops drawing selections for rank-0 or empty views. + Reader changes, repartitioning and invariant checks remain possible. + NumPy itself permits some indexing of such arrays, such as `[()]`. """ return self.model.ndim > 0 and self.model.size > 0 @@ -307,11 +307,9 @@ def choose_reader(self, data: st.DataObject) -> None: 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. + Once the harness stops drawing selections, this rule tests another + partitioning of the same view while the reader-change rule also remains + enabled. The invariants must continue to match the NumPy model. """ parts = data.draw(st.sampled_from(list(type(self).partitionings))) self.view = repartition(self.view, parts) diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py index 63a3d351a9..1c62b925db 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py @@ -1,8 +1,8 @@ """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 +selection for it, in the spelling its mode expects. Some vectorized selections +use a partial coordinate tuple or a mask with an ellipsis. 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. @@ -16,8 +16,9 @@ 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. +Coordinate-drawing strategies require non-empty axes; vectorized selections +also require positive rank. The `empty_masks` strategy can generate a mask +for an empty shape because it does not draw element coordinates. Requires the `testing` extra (`pip install zarr-indexing[testing]`). """ @@ -151,8 +152,7 @@ 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 + composes like basic indexing. 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. """ diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py index 9848c6801f..a15ea60573 100644 --- a/packages/zarr-indexing/tests/test_json.py +++ b/packages/zarr-indexing/tests/test_json.py @@ -347,9 +347,8 @@ def _index_array_body(index_array: Any, rank: int = 1, extent: int = 2) -> Index ([True, False], "bool"), (["a", "b"], "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. + # looked at: the wire representation requires a nested array rather + # than a scalar string or integer. ("abc", "must be an array of integers"), (5, "must be an array of integers"), ([None, None], "object"), diff --git a/packages/zarr-indexing/tests/test_ndsel_tensorstore.py b/packages/zarr-indexing/tests/test_ndsel_tensorstore.py index ab9465c8d1..29b5b8a2e4 100644 --- a/packages/zarr-indexing/tests/test_ndsel_tensorstore.py +++ b/packages/zarr-indexing/tests/test_ndsel_tensorstore.py @@ -1,14 +1,14 @@ """Cross-check canonical ndsel bodies against a real TensorStore. -A normalized ndsel `transform` body is, field-for-field, a TensorStore -`IndexTransform` (minus the `kind` discriminator, which the canonical body never -carries). This test loads a handful of finite-bound canonical bodies into +Canonical ndsel bodies use TensorStore's `IndexTransform` field vocabulary, +but the consumers have different validation constraints. This test loads a +handful of finite-bound canonical bodies supported by both implementations into `tensorstore.IndexTransform(json=...)` and confirms that TensorStore's own `to_json()` re-loads, through our engine layer, into an equivalent transform. Skipped when tensorstore is not installed. Run it explicitly with: - uv run --with tensorstore pytest \ + hatch run test.py3.12-optional:pytest \ packages/zarr-indexing/tests/test_ndsel_tensorstore.py -q """ diff --git a/packages/zarr-indexing/tests/test_wire_boundaries.py b/packages/zarr-indexing/tests/test_wire_boundaries.py new file mode 100644 index 0000000000..81074f1c40 --- /dev/null +++ b/packages/zarr-indexing/tests/test_wire_boundaries.py @@ -0,0 +1,63 @@ +"""Boundary validation must preserve coordinates and canonical messages.""" + +from typing import Any + +import numpy as np +import pytest + +from zarr_indexing._wire import lower_index_array +from zarr_indexing.messages import NdselError, normalize_ndsel + + +@pytest.mark.parametrize("raw", [[True, 2], [[0, False]], [True, False]]) +def test_index_array_rejects_boolean_elements(raw: Any) -> None: + with pytest.raises(NdselError, match="integers"): + lower_index_array(raw, "index_array") + + +@pytest.mark.parametrize("value", [int(np.iinfo(np.intp).max) + 1, int(np.iinfo(np.intp).min) - 1]) +def test_index_array_rejects_out_of_range_coordinates(value: int) -> None: + with pytest.raises(NdselError): + lower_index_array([value], "index_array") + + +@pytest.mark.parametrize( + "raw", [[], [[], []], [0, -1, 2], [int(np.iinfo(np.intp).min), int(np.iinfo(np.intp).max)]] +) +def test_index_array_preserves_integer_coordinates(raw: Any) -> None: + result = lower_index_array(raw, "index_array") + assert result.dtype == np.dtype(np.intp) + assert result.tolist() == raw + + +@pytest.mark.parametrize(("start", "stop"), [(-(2**63), -(2**63)), (-(2**63) + 1, -(2**63))]) +def test_slice_rejects_unrepresentable_canonical_bounds(start: int, stop: int) -> None: + with pytest.raises(NdselError): + normalize_ndsel({"kind": "slice", "start": [start], "stop": [stop], "step": [-1]}) + + +@pytest.mark.parametrize( + "message", + [ + {"kind": "box", "shape": [1] * 33}, + {"kind": "slice", "start": [0] * 33, "stop": [1] * 33}, + {"kind": "transform", "input_shape": [1] * 33}, + ], +) +def test_normalize_rejects_excessive_inferred_rank(message: dict[str, Any]) -> None: + with pytest.raises(NdselError, match="rank"): + normalize_ndsel(message) + + +@pytest.mark.parametrize( + "message", + [ + {"kind": "box", "shape": [1] * 32}, + {"kind": "transform", "input_rank": 32}, + {"kind": "slice", "start": [-(2**63) + 1], "stop": [-(2**63) + 1], "step": [-1]}, + {"kind": "slice", "start": [2**63 - 1], "stop": [-(2**63) + 1], "step": [-(2**63)]}, + ], +) +def test_normalize_boundary_results_are_idempotent(message: dict[str, Any]) -> None: + canonical = normalize_ndsel(message) + assert normalize_ndsel({"kind": "transform", **canonical}) == canonical From 88a2ed4827d6faebfe43c1a9389ef89a5f7f0e15 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 13 Sep 2026 18:50:11 +0200 Subject: [PATCH 04/12] fix(indexing): validate selector bounds and shared dependencies Correct mathematical API documentation to match supported coordinate, grid, and chunk projection contracts. Assisted-by: Codex:GPT-6 --- packages/zarr-indexing/docs/api/grid.md | 5 +- packages/zarr-indexing/docs/api/index.md | 7 +-- .../docs/snippets/output_maps.py | 8 +-- .../src/zarr_indexing/_composition.py | 8 ++- .../src/zarr_indexing/boundary.py | 4 +- .../src/zarr_indexing/chunk_resolution.py | 4 +- .../zarr-indexing/src/zarr_indexing/domain.py | 7 ++- .../zarr-indexing/src/zarr_indexing/grid.py | 17 +++-- .../src/zarr_indexing/output_map.py | 39 ++++++------ .../src/zarr_indexing/transform.py | 62 ++++++++++++------- packages/zarr-indexing/tests/test_boundary.py | 47 ++++++++++++++ .../zarr-indexing/tests/test_transform.py | 12 ++++ 12 files changed, 154 insertions(+), 66 deletions(-) create mode 100644 packages/zarr-indexing/tests/test_boundary.py diff --git a/packages/zarr-indexing/docs/api/grid.md b/packages/zarr-indexing/docs/api/grid.md index b7c376eb85..012c9df6dc 100644 --- a/packages/zarr-indexing/docs/api/grid.md +++ b/packages/zarr-indexing/docs/api/grid.md @@ -11,8 +11,9 @@ objects whose `shape` is the valid data size and whose `codec_shape` preserves the full codec-buffer size at a regular-grid boundary. `dimension_grids_from_chunks` returns these compact dimensions: integer chunk -shapes become `FixedDimension` instances and explicit per-axis edge sequences -become `VaryingDimension` instances. `DimensionGridLike` remains the narrow +shapes become `FixedDimension` instances and positive per-axis edge sequences +become `VaryingDimension` instances. An empty-axis sequence of zeros (including +an empty sequence) becomes `FixedDimension(size=0, extent=0)`. `DimensionGridLike` remains the narrow protocol used by the chunk planner, while `EdgeDimensionGrid` is kept for explicit edge-based and coordinate-origin examples. diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index a5b7163bbd..8b7a5d6515 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -25,12 +25,11 @@ and the wire format built on top of it. - [`zarr_indexing.domain`](domain.md) — `IndexDomain`, a rectangular region of integer coordinates with an explicit (possibly non-zero) origin - [`zarr_indexing.output_map`](output_map.md) — `ConstantMap`, `DimensionMap`, - and `ArrayMap`: three representations of a set of integer coordinates, one - per storage dimension + and `ArrayMap`: three coordinate mappings that preserve order and duplicates, + one per storage dimension - [`zarr_indexing.transform`](transform.md) — `IndexTransform`, which pairs a domain with output maps, plus the indexing (`[...]`, `.oindex`, `.vindex`), - `intersect`, and `translate` operations, and `selection_to_transform` - transforms into one + `intersect`, `translate`, and `compose` operations, and `selection_to_transform` **Chunk resolution** diff --git a/packages/zarr-indexing/docs/snippets/output_maps.py b/packages/zarr-indexing/docs/snippets/output_maps.py index a78fc73b50..7daed4b63e 100644 --- a/packages/zarr-indexing/docs/snippets/output_maps.py +++ b/packages/zarr-indexing/docs/snippets/output_maps.py @@ -54,11 +54,9 @@ def resolve(transform: IndexTransform, values: np.ndarray[Any, Any]) -> np.ndarr # --8<-- [end:array-map] # --8<-- [start:constant-map] -# ConstantMap reads one source coordinate for every request cell. No NumPy -# selection spells this operation: source[0] drops the axis, and a repeated -# fancy index source[[0, 0, 0, 0]] matches the values but degrades the -# description to a coordinate list. The value-faithful counterpart is a -# broadcast. +# ConstantMap reads one source coordinate for every request cell. NumPy can +# express the same values with source[[0, 0, 0, 0]] or with a broadcast. +# The constant map represents the repeated coordinate without storing a list. repeat = IndexTransform( domain=IndexDomain.from_shape((4,)), output=(ConstantMap(offset=0),), diff --git a/packages/zarr-indexing/src/zarr_indexing/_composition.py b/packages/zarr-indexing/src/zarr_indexing/_composition.py index f90ea26fa7..ab92ce45bf 100644 --- a/packages/zarr-indexing/src/zarr_indexing/_composition.py +++ b/packages/zarr-indexing/src/zarr_indexing/_composition.py @@ -15,7 +15,8 @@ composing it with an outer `ArrayMap` leaves the index array alone and rescales around it. - An `ArrayMap` inner map must be *evaluated* at the coordinates the outer - transform produces, which is the only case that touches array data. + transform produces. This case gathers from the inner lookup table; other + cases can still inspect or copy index arrays during validation/construction. """ from __future__ import annotations @@ -43,7 +44,10 @@ def compose(outer: IndexTransform, inner: IndexTransform) -> IndexTransform: `inner` maps intermediate coords (rank n) to output coords (rank p). The result maps user coords (rank m) to output coords (rank p). - Precondition: `outer.output_rank == inner.domain.ndim`. + The output rank of `outer` must equal the input rank of `inner`, and every + coordinate produced by a nonempty `outer` must lie in `inner.domain`. + A rank mismatch raises `ValueError`; out-of-domain coordinates raise + `BoundsCheckError`. Examples -------- diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py index c4831715c4..1b3e96ad1c 100644 --- a/packages/zarr-indexing/src/zarr_indexing/boundary.py +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -132,8 +132,8 @@ def _expanded_axis_walk(entries: tuple[Any, ...], ndim: int, mode: SelectionMode """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. + `Ellipsis` (or a `newaxis`) is the axis it starts at. The following entry + advances past the ellipsis's skipped axes; a `newaxis` consumes no axis. """ for sel in entries: if is_bool_scalar(sel): diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index 6d14f0d2f6..f914e83c2f 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -89,7 +89,9 @@ class ChunkProjection: 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"``. + once. Projections retaining ArrayMaps or survivor arrays are + conservatively ``"unknown"``. A singleton fancy selection can simplify + to a ConstantMap and receive a proven ``"full"`` or ``"partial"`` result. Examples -------- diff --git a/packages/zarr-indexing/src/zarr_indexing/domain.py b/packages/zarr-indexing/src/zarr_indexing/domain.py index b42bd04ccc..dc9531fb25 100644 --- a/packages/zarr-indexing/src/zarr_indexing/domain.py +++ b/packages/zarr-indexing/src/zarr_indexing/domain.py @@ -139,10 +139,11 @@ def contains(self, index: tuple[int, ...]) -> bool: ) def contains_domain(self, other: IndexDomain) -> bool: - """Whether every coordinate of `other` lies inside this domain. + """Whether `other` has the same rank and bounds enclosed by this domain. - An empty `other` within this domain's bounds is contained. A rank - mismatch returns `False` rather than raising. + Empty domains are still checked by their bounds: an empty `other` + located outside this domain returns `False`, unlike empty-set + containment. A rank mismatch returns `False` rather than raising. """ if other.ndim != self.ndim: return False diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py index 20ad3f95c2..bd350eeebe 100644 --- a/packages/zarr-indexing/src/zarr_indexing/grid.py +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -41,7 +41,9 @@ class DimensionGridLike(Protocol): def index_to_chunk(self, idx: int) -> int: """Map a global source index to the index of the chunk that contains it. - Implementers must raise `IndexError` when `idx` lies outside `[0, extent)`. + Raise `IndexError` outside the coordinates supported by the grid. + The built-in bounded grids use `[0, extent)`; custom grids may support + negative coordinates or an unbounded region. """ ... @@ -56,7 +58,8 @@ def chunk_size(self, chunk_ix: int) -> int: def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: """Vectorized `index_to_chunk`: map global source indices to chunk indices. - Implementers must raise `IndexError` if any index lies outside `[0, extent)`. + Raise `IndexError` if any index is outside the grid's supported region, + consistently with the scalar method. """ ... @@ -163,7 +166,10 @@ def with_extent(self, new_extent: int) -> FixedDimension: return FixedDimension(size=self.size, extent=new_extent) def resize(self, new_extent: int) -> FixedDimension: - """Return a copy resized to `new_extent`; the fixed chunk size covers any new extent.""" + """Return a copy resized to `new_extent`, retaining the chunk size. + + A zero chunk size is supported only for an empty extent and cannot grow. + """ return FixedDimension(size=self.size, extent=new_extent) @property @@ -351,8 +357,9 @@ def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.int def with_extent(self, new_extent: int) -> DimensionGrid: """Return a grid with the existing chunk layout re-clipped to `new_extent`. - Implementers must not invent new grid cells: raise `ValueError` when the declared - layout cannot cover `new_extent`. + Preserve the existing layout rule: a fixed positive chunk size can + cover a larger extent without changing its size; explicit edge lists + raise `ValueError` if their sum cannot cover `new_extent`. """ ... diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py index 16ad7aeb9b..515ddb803d 100644 --- a/packages/zarr-indexing/src/zarr_indexing/output_map.py +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -56,13 +56,10 @@ 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 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. + axes are therefore exactly the axes of size 2 or more. This is structural: + values are not inspected for constant or repeated coordinates. An array + with shape `(0, 2)` is empty but still reports axis 1; the zero-size axis + itself is not reported. """ return tuple(axis for axis, size in enumerate(index_array.shape) if size > 1) @@ -76,8 +73,9 @@ class ConstantMap: Examples -------- - Every input cell maps to the same output coordinate — the NumPy analogy - is a broadcast (`np.broadcast_to(5, (3,))`), not an index: + Every input cell maps to the same output coordinate, like broadcasting + coordinate 5 with `np.broadcast_to(5, (3,))`. Repeated fancy indices can + also describe these coordinates, using an explicit list: >>> from zarr_indexing.domain import IndexDomain >>> from zarr_indexing.transform import IndexTransform @@ -154,11 +152,11 @@ class ArrayMap: the result. Arises from fancy indexing (e.g., `arr[[5, 1, 1]]` or boolean masks). - Freshly constructed maps are normalized to the **full input rank** of their - enclosing transform: `index_array` has the enclosing domain's rank, sized + A map used in a transform must have its **full input rank**: + `index_array` has the enclosing domain's rank, sized fully on the axes it varies over and singleton (size 1) elsewhere. The shape is the single source of truth for what the map depends on — its - **dependency axes** are exactly its non-singleton axes (see + **dependency axes** are exactly its axes of size greater than one (see `_array_map_dependency_axes`) — and it distinguishes the two flavors of multi-array fancy indexing: @@ -273,11 +271,12 @@ def __hash__(self) -> int: @property def dependency_axes(self) -> tuple[int, ...]: - """Every input axis this map varies over: its non-singleton axes. + """Structural dependency axes: axes of size greater than one. - One axis means orthogonal, several mean correlated, and none means - the map is degenerate — the shape is the single source of truth for - all three. + Axes of size greater than one are reported, regardless of coordinate + values or a zero-size axis elsewhere. Whether the whole transform is + orthogonal also depends on how other maps use these axes; a single + map's shape does not establish independence. Examples -------- @@ -302,7 +301,7 @@ def dependent_axis(self) -> int | None: ------- int or None The axis the map varies over, or `None` when it varies over no input - axis at all — an empty map, or a hand-built all-singleton one. `None` + axis of size greater than one, such as an all-singleton map. `None` is a valid result, not an error; such maps resolve through the pointwise (general) path. @@ -366,9 +365,9 @@ def to_json(self) -> OutputIndexMapJSON: def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: """Construct the output map a canonical wire form names. - The wire form is a tagged union — `index_array`, then `input_dimension`, - else constant — so loading it dispatches to the right kind here rather - than on any one of them. + The wire form is structurally discriminated: the presence of `index_array` + selects an array map, `input_dimension` selects a dimension map, and + neither selects a constant map. Examples -------- diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 975494490c..fca9cdd5f3 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -8,8 +8,9 @@ Key operations: - **Indexing** (`transform[2:8]`, `.oindex[idx]`, `.vindex[idx]`) — - produces a new transform with a narrower input domain and adjusted output - maps. No I/O occurs. This is how lazy slicing works. + produces a new transform with a new input domain and adjusted output maps. + Slices may narrow a domain; fancy indices can repeat cells and enlarge it, + and integer/newaxis indexing changes its rank. No I/O occurs. - **intersect(output_domain)** — restrict to output coordinates within a region. This is chunk resolution: "which of my coordinates fall in this @@ -526,10 +527,20 @@ def intersect( Returns `(restricted_transform, out_indices)` or None if empty. - `out_indices` carries the surviving output positions: `None` when all - positions survive (ConstantMap/DimensionMap only), a single integer array - for one ArrayMap (or correlated/vectorized ArrayMaps), or a dict keyed by - output dimension for >= 2 orthogonal ArrayMaps (an outer product). + `out_indices` is bookkeeping for chunk resolution. `None` means there + are no ArrayMaps; the restricted domain itself records surviving input + coordinates, which may be fewer than the original domain's. One + orthogonal ArrayMap returns its surviving positional indices; multiple + orthogonal ArrayMaps return a dict keyed by output dimension. General + array maps return a shared scatter array. For paired public coordinate + mappings, use `plan_chunks` and its `cell_transform`. + + Raises + ------ + NotImplementedError + For a nonempty transform whose ArrayMap has a non-singleton axis + also referenced by a DimensionMap. These dependencies cannot be + filtered independently by the current intersection implementation. """ return _intersect(self, output_domain) @@ -567,8 +578,10 @@ def __getitem__(self, selection: Any) -> IndexTransform: No I/O occurs. Integers and slice bounds are literal domain coordinates (TensorStore convention): negative values are not counted from the end, - and out-of-domain values raise `BoundsCheckError`. Integer indices drop - their input dimension; `None` inserts a size-1 dimension. + and nonempty out-of-domain slice intervals raise `BoundsCheckError`. + Empty slices (`start == stop`) are accepted even outside the domain. + Integer indices must lie inside the domain and drop their input + dimension; `None` inserts a size-1 dimension. """ return _apply_basic_indexing(self, selection) @@ -637,14 +650,17 @@ def index_array_structure(self) -> Literal["none", "orthogonal", "general"]: Returns ------- `"none"` when no output map is an `ArrayMap`; `"orthogonal"` when every - `ArrayMap` varies over exactly one input axis, each its own (an outer - product, one independent gather per axis); `"general"` otherwise — + `ArrayMap` has exactly one non-singleton input axis, distinct from all + other ArrayMaps' axes and all DimensionMap input dimensions (one + independent gather per axis); `"general"` otherwise — correlated (`vindex`) maps sharing their non-singleton axes, maps produced by composing fancy steps, maps sharing an input axis (a diagonal gather), - and empty or hand-built all-singleton maps whose shape names no axis. The + and maps whose shape names no dependency axis. A map can have a + non-singleton axis even when another axis has size zero. The orthogonal resolvers narrow one axis at a time and are only sound for `"orthogonal"`; everything else takes the pointwise path that collapses - the joint block. Everything is read off the index arrays' shapes. + the joint block or rejects unsupported shared affine/array dependencies. + Classification uses index-array shapes and DimensionMap input dimensions. Examples -------- @@ -662,7 +678,7 @@ def index_array_structure(self) -> Literal["none", "orthogonal", "general"]: >>> t.vindex[np.array([0, 2]), np.array([1, 3])].index_array_structure 'general' """ - seen: set[int] = set() + seen = {m.input_dimension for m in self.output if isinstance(m, DimensionMap)} has_array = False for m in self.output: if not isinstance(m, ArrayMap): @@ -1113,7 +1129,7 @@ def _intersect_general( transform: IndexTransform, output_domain: IndexDomain, ) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]]] | None: - """Intersect a transform with any index-array structure, pointwise. + """Intersect array maps jointly when they are independent of residual affine axes. Every `ArrayMap` — correlated, orthogonal, or several sharing an axis — is treated as a lookup table over the joint block of non-slice axes: a block @@ -1126,8 +1142,10 @@ def _intersect_general( keeps its own resolver. 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)`. + `out_indices` combines positional broadcast offsets with literal residual + slice coordinates, of shape `(surviving_points,) + (residual slice sizes)`. + Chunk resolution removes residual origins before treating it as a flat + offset into the request buffer. 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 @@ -1221,10 +1239,9 @@ def _intersect_general( ) result = IndexTransform(domain=new_domain, output=tuple(new_output)) - # 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. + # Scatter bookkeeping with row-major input-domain strides. Broadcast + # coordinates below are positional; residual coordinates remain literal. + # `_correlated_cell_transform` removes their origins before unraveling. # # 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 @@ -1706,8 +1723,9 @@ def _selection_axis_count(selection: Any) -> int: def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: """Apply vectorized indexing to an IndexTransform. - All array indices are broadcast together. Broadcast dimensions are prepended, - followed by non-array (slice) dimensions. + All array indices are broadcast together. Their broadcast dimensions occupy + the advanced indices' position when adjacent, and lead when a slice axis + separates the advanced indices (see `_broadcast_insertion_point`). A transform that already carries index arrays takes the composition path (`_compose_selection`) instead of being rewritten in place; see diff --git a/packages/zarr-indexing/tests/test_boundary.py b/packages/zarr-indexing/tests/test_boundary.py new file mode 100644 index 0000000000..92a2e99d4f --- /dev/null +++ b/packages/zarr-indexing/tests/test_boundary.py @@ -0,0 +1,47 @@ +"""Positional array selectors retain their values through bounds validation.""" + +from typing import Literal + +import numpy as np +import pytest + +from zarr_indexing.boundary import normalize_positional_selection +from zarr_indexing.domain import IndexDomain + + +@pytest.mark.parametrize("dtype", ["i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8"]) +@pytest.mark.parametrize("mode", ["orthogonal", "vectorized"]) +@pytest.mark.parametrize("origin", [-3, 0, 10]) +def test_positional_integer_arrays( + dtype: str, mode: Literal["orthogonal", "vectorized"], origin: int +) -> None: + domain = IndexDomain((origin,), (origin + 4,)) + values = [-4, -1, 0, 3] if np.dtype(dtype).kind == "i" else [0, 3, 0, 3] + selector = np.array(values, dtype=dtype) + snapshot = selector.copy() + normalized = normalize_positional_selection(selector, domain, mode) + expected = np.array([value % 4 + origin for value in values], dtype=np.intp) + np.testing.assert_array_equal(normalized[0], expected) + np.testing.assert_array_equal(selector, snapshot) + assert normalized[0].dtype == np.dtype(np.intp) + + +@pytest.mark.parametrize( + ("dtype", "value"), + [ + (dtype, value) + for dtype in ("i1", "i2", "i4", "i8", "u1", "u2", "u4", "u8") + for value in ( + (int(np.iinfo(dtype).min), -5, 4, int(np.iinfo(dtype).max)) + if np.dtype(dtype).kind == "i" + else (4, int(np.iinfo(dtype).max)) + ) + ], +) +@pytest.mark.parametrize("mode", ["orthogonal", "vectorized"]) +def test_positional_integer_array_out_of_bounds( + dtype: str, value: int, mode: Literal["orthogonal", "vectorized"] +) -> None: + selector = np.array([value], dtype=dtype) + with pytest.raises(IndexError, match="out of bounds"): + normalize_positional_selection(selector, IndexDomain.from_shape((4,)), mode) diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index 9e8b24bf4a..ef2fad0a81 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -1274,3 +1274,15 @@ def test_index_array_structure_classifies_the_three_shapes() -> None: ), ) assert diagonal.index_array_structure == "general" + + +@pytest.mark.parametrize("array_first", [False, True]) +def test_intersect_rejects_array_and_affine_shared_input_axis(array_first: bool) -> None: + """Joint filtering of affine and lookup coordinates must not return extra cells.""" + affine = DimensionMap(0) + lookup = ArrayMap(np.array([0, 3, 1, 2])) + output = (lookup, affine) if array_first else (affine, lookup) + transform = IndexTransform(IndexDomain.from_shape((4,)), output) + output_domain = IndexDomain((0, 1), (2, 4)) if array_first else IndexDomain((1, 0), (4, 2)) + with pytest.raises(NotImplementedError, match="also bound by a slice map"): + transform.intersect(output_domain) From 7de59755b0b5c5cc9e84381138d182618b413bbd Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 13 Sep 2026 18:54:06 +0200 Subject: [PATCH 05/12] docs(indexing): reconcile reader contracts and record audit fixes Assisted-by: Codex:GPT-6 --- packages/zarr-indexing/README.md | 6 +++--- .../changes/+factual-audit.bugfix.md | 1 + .../changes/+factual-audit.doc.md | 1 + packages/zarr-indexing/docs/api/lazy_array.md | 16 ++++++++-------- packages/zarr-indexing/docs/api/reader.md | 5 +++-- packages/zarr-indexing/docs/design-notes.md | 2 +- packages/zarr-indexing/docs/guide/index.md | 5 +++-- .../zarr-indexing/docs/guide/integrations.md | 18 ++++++++---------- packages/zarr-indexing/pyproject.toml | 13 ++++--------- .../src/zarr_indexing/lazy_array.py | 2 +- 10 files changed, 33 insertions(+), 36 deletions(-) create mode 100644 packages/zarr-indexing/changes/+factual-audit.bugfix.md create mode 100644 packages/zarr-indexing/changes/+factual-audit.doc.md diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index fb631d37e6..033e94b9e8 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -15,9 +15,9 @@ Key types: accessor: `LazyArray.from_numpy(numpy_array).lazy[10:50, ::2].lazy.oindex[[3, 1, 1], :]` composes a transform and returns a new view without reading data, and `result()` materializes it into owned system memory. `LazyArray(source)` uses - the conservative basic reader; `from_numpy` explicitly selects NumPy's - optimized reader. Device arrays require an explicit custom reader responsible - for transferring values into the supplied system-memory output buffer. + the basic reader; `from_numpy` selects `numpy_reader`, which currently uses + the same slab-and-gather implementation. Device sources that refuse NumPy + conversion need a custom reader to transfer values into the output buffer. - `Reader` — the explicit backend execution boundary: transforms say which values belong in the result, while readers say how a backend obtains them - `IndexDomain` — a rectangular region of integer coordinates diff --git a/packages/zarr-indexing/changes/+factual-audit.bugfix.md b/packages/zarr-indexing/changes/+factual-audit.bugfix.md new file mode 100644 index 0000000000..20451c142c --- /dev/null +++ b/packages/zarr-indexing/changes/+factual-audit.bugfix.md @@ -0,0 +1 @@ +Validate signed and unsigned positional indices before conversion, reject invalid wire index-array values and unrepresentable normalized bounds/ranks, and raise an explicit error for unsupported intersections sharing an affine and lookup input axis instead of returning incorrect coordinates. diff --git a/packages/zarr-indexing/changes/+factual-audit.doc.md b/packages/zarr-indexing/changes/+factual-audit.doc.md new file mode 100644 index 0000000000..86f035319b --- /dev/null +++ b/packages/zarr-indexing/changes/+factual-audit.doc.md @@ -0,0 +1 @@ +Correct indexing, reader, serialization, cache, and integration descriptions to match supported behavior; qualify NumPy/TensorStore compatibility and performance claims. diff --git a/packages/zarr-indexing/docs/api/lazy_array.md b/packages/zarr-indexing/docs/api/lazy_array.md index f855511d46..13740981cd 100644 --- a/packages/zarr-indexing/docs/api/lazy_array.md +++ b/packages/zarr-indexing/docs/api/lazy_array.md @@ -7,21 +7,21 @@ reader and composes its transform without reading data. `result()` allocates owned system memory, then calls that reader once for each projected part. Rectangular parts write directly into their final slices; advanced placement may first use an owned dense temporary. `LazyArray(source)` assumes only basic -indexing, while `LazyArray.from_numpy(array)` explicitly selects NumPy's -optimized reader. +indexing, while `LazyArray.from_numpy(array)` selects `numpy_reader`. Both +currently use the same slab-and-gather implementation. The built-in readers lower through NumPy system memory and support sources -whose basic reads can be converted there. They do not implicitly transfer -device arrays; a device source needs an explicit custom reader that transfers -into the supplied system-memory output. Derived views and parts share their +whose basic reads can be converted there. A device source that refuses NumPy +conversion needs a custom reader that transfers into the output buffer. Derived views and parts share their reader and part views may be materialized concurrently, so stateful readers must synchronize their own mutable state. Every public `Partition.view.transform` directly maps that view's zero-origin coordinates into its raw `Partition.view.array`, including for non-first partitions. `Partition.projection.chunk_transform` intentionally stays local to -the selected chunk. During materialization the reader receives both frames in -one `ReadContext`: the public global transform in `context.transform` and the -same local plan in `context.projection`. +the selected chunk. During parent materialization (`view.result(parts=parts)`) the reader receives +both frames in one `ReadContext`: the public global transform in `context.transform` and the +same local plan in `context.projection`. Direct `part.view.result()` calls +provide no projection; readers requiring it must use parent assembly. ::: zarr_indexing.lazy_array diff --git a/packages/zarr-indexing/docs/api/reader.md b/packages/zarr-indexing/docs/api/reader.md index 39129c4d47..3944336745 100644 --- a/packages/zarr-indexing/docs/api/reader.md +++ b/packages/zarr-indexing/docs/api/reader.md @@ -34,7 +34,7 @@ class RecordingReader: self.calls = [] def read_into(self, source, context, out, /): - self.calls.append((source, context, out)) + self.calls.append((source, context, out.shape, out.dtype)) self.inner.read_into(source, context, out) @@ -44,7 +44,8 @@ view = LazyArray.from_numpy(array).with_reader(outer) values = view.result() ``` -Both wrappers observe the same three objects, in outer-to-inner order. This +Both wrappers observe the same arguments, in outer-to-inner order, and log +output metadata without retaining the output buffer. This delegation pattern supports policies such as logging and caching without library-defined wrapper primitives. diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index f5db95500c..dfa92d4c95 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -16,7 +16,7 @@ visual guide owns the mechanics of ## 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 visual guide +index-transform model, implemented here in Python against NumPy. The visual guide introduces the shared model in [Coordinates are addresses](guide/index.md#coordinates-are-addresses) and [Lazy views compose](guide/index.md#lazy-views-compose); the comparison here is diff --git a/packages/zarr-indexing/docs/guide/index.md b/packages/zarr-indexing/docs/guide/index.md index b7e748e262..f55794c4c0 100644 --- a/packages/zarr-indexing/docs/guide/index.md +++ b/packages/zarr-indexing/docs/guide/index.md @@ -389,8 +389,9 @@ each bundles a sub-view of the request (`.view`), that chunk's projection Within one `Partition`, the frames divide: `Partition.view.transform` is a different, global transform — it maps the part view directly into the raw wrapped source — while only `Partition.projection.chunk_transform` uses -zero-origin chunk-local coordinates. Readers receive both so the global -source address and the local planning frame cannot be confused. +zero-origin chunk-local coordinates. Parent assembly passes both frames to +the reader. Direct `part.view.result()` calls supply the global transform +with `projection=None`. | Projection field | What its output coordinates mean | | --- | --- | diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index 9a79e5e4ae..a0bc721175 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -91,9 +91,8 @@ the dense box becomes exactly one backend call. Both regimes go through For affine selections, `basic_reader` uses positive-step slices and applies reversal or layout changes in memory. Fancy selections can require reading a cover containing unselected values. The source must accept the emitted steps. -For a source that accepts only unit steps, use `slice(start, stop, 1)` with -[`unit_step_reader`][zarr_indexing.reader.UnitStepReader] for such a source -and every key it receives is an ascending unit-step slice per axis, with +Use [`unit_step_reader`][zarr_indexing.reader.UnitStepReader] for a source that +accepts only unit steps. Every key it receives is an ascending unit-step slice per axis, with strides, reversals, and gathers applied to the in-memory block instead: ```python @@ -206,13 +205,12 @@ decoded chunks. Every read delta follows directly from the viewport request: | 4 | `image[1:5, 2]` | `(0, 0)` | `(0, 0)`, `(1, 0)` | The evicted chunk is reloaded while the required ready chunk is retained. | | 5 | `image[3:5, 4:6]` | `(1, 1)` fails; no repeated read; `(1, 1)` succeeds after retry | `(0, 0)`, `(1, 1)` | Failure is retained until explicit retry; the repaired source then returns `[[28, 29], [36, 37]]`. | -Chunks required by an active request are pinned through assembly, so a request -may temporarily span more chunks than the steady-state capacity. Capacity is -counted in decoded chunks—not records or bytes—and eviction occurs only after -all requested values have been placed. Because pinning and materialization use -the same prepared tuple, those lifecycle decisions cannot drift from the parts -that are actually read, and the cache never has to infer or reconstruct a -projection. +The example defers eviction until a successful outermost request finishes, so +it can temporarily exceed capacity during assembly. Capacity counts decoded +chunks, not bytes, event records, or temporary arrays. Failed requests skip +that eviction step. The example assumes an unchanged source and one source/grid +per reader; it does not implement invalidation or synchronization for concurrent +requests. The prepared tuple supplies the projections used for each read. The event log makes the failure boundary equally explicit: diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index 5a2c9e8fd4..f2e286ab42 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -100,8 +100,8 @@ 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 +# foreign source's attributes or token hooks may fail, so these paths provide +# fallback metadata or tokens when ordinary exceptions occur. Configured here (not as # noqa comments) because different ruff versions # have differed on whether these rules fire; RUF100 can remove unused noqa comments. "src/zarr_indexing/lazy_array.py" = ["BLE001", "S110"] @@ -133,13 +133,8 @@ include = [ enableExperimentalFeatures = true typeCheckingMode = "strict" pythonVersion = "3.12" -# This strict config was written for zarr-metadata's JSON/dataclass-shaped -# code. zarr-indexing is numpy-heavy, and numpy's stubs return partially -# unknown types (e.g. `ndarray[Unknown, Unknown]`, `dtype[Unknown]`) even for -# fully-typed call sites, so the reportUnknown* family below cannot reasonably -# be satisfied here. Downgraded to warnings (not silenced) rather than -# disabled outright, and CI (which only fails the pyright job on errors, not -# warnings) still surfaces them for visibility. +# Partially unknown types from NumPy stubs are reported as warnings. CI's +# pyright invocation fails on errors, while retaining these warnings in its output. reportUnknownVariableType = "warning" reportUnknownArgumentType = "warning" reportUnknownMemberType = "warning" diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index c5c2967b94..b23115f1b5 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -1197,7 +1197,7 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any: # back process memory dressed as data. The parts are disjoint by # contract, so counting the cells each addresses is enough: # a gap alone undercounts and an overlap alone overcounts. This - # count cannot detect a compensating gap and overlap; disjointness + # count cannot detect a compensating gap and overlap; absence of overlap # still relies on the planner contract. if prepared_parts is not None: raise ValueError( From 560ccd205f57a8d7ef041bcba0212a9edcfa4cb2 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 13 Sep 2026 18:57:53 +0200 Subject: [PATCH 06/12] docs(indexing): reconcile audit with current partition implementation Retain the existing unsigned selector fix and update the unsupported mixed-dependency error assertion for general intersection routing. Assisted-by: Codex:GPT-6 --- packages/zarr-indexing/CONTRIBUTING.md | 6 +++--- packages/zarr-indexing/README.md | 2 +- .../changes/+factual-audit.bugfix.md | 2 +- packages/zarr-indexing/docs/design-notes.md | 20 +++++++++---------- packages/zarr-indexing/justfile | 7 +++---- .../src/zarr_indexing/transform.py | 2 +- .../tests/test_chunk_resolution.py | 2 +- 7 files changed, 20 insertions(+), 21 deletions(-) diff --git a/packages/zarr-indexing/CONTRIBUTING.md b/packages/zarr-indexing/CONTRIBUTING.md index a591d31bd6..bcbe51c079 100644 --- a/packages/zarr-indexing/CONTRIBUTING.md +++ b/packages/zarr-indexing/CONTRIBUTING.md @@ -15,9 +15,9 @@ just docs-serve # serve the docs site locally Run them from this directory, or from the repository root 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. +The test recipe layers this package into the repository-root environment. +Chunk-resolution tests use this package’s own grids; the Dask example also +uses `zarr`, which is not a dependency of the base indexing package. ## License diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index 033e94b9e8..2992dcf280 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -35,7 +35,7 @@ Key types: The base package depends on NumPy and the standard library; its optional testing module also requires Hypothesis. The package does not import `zarr`. It is developed in the [zarr-python](https://github.com/zarr-developers/zarr-python) repository, -and its integration tests exercise Zarr chunk grids. Installing it does not +and its examples include reading Zarr arrays through Dask. Installing it does not replace Zarr's indexing implementation. ## Installation diff --git a/packages/zarr-indexing/changes/+factual-audit.bugfix.md b/packages/zarr-indexing/changes/+factual-audit.bugfix.md index 20451c142c..810bd1a74e 100644 --- a/packages/zarr-indexing/changes/+factual-audit.bugfix.md +++ b/packages/zarr-indexing/changes/+factual-audit.bugfix.md @@ -1 +1 @@ -Validate signed and unsigned positional indices before conversion, reject invalid wire index-array values and unrepresentable normalized bounds/ranks, and raise an explicit error for unsupported intersections sharing an affine and lookup input axis instead of returning incorrect coordinates. +Reject invalid wire index-array values and unrepresentable normalized bounds/ranks, and raise an explicit error for unsupported intersections sharing an affine and lookup input axis instead of returning incorrect coordinates. diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index dfa92d4c95..cb80e5ba19 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -48,9 +48,11 @@ about the deliberately matching semantics: components within a mixed request. Both derive the per-chunk transforms from the partition ([the guide](guide/index.md#a-plan-is-a-product-of-per-axis-tables) shows the tables). TensorStore keeps strided sets implicit, while this - library materializes their per-axis rows for vectorized consumers. Diagonals are rejected here; supporting them needs a strided set per - *input* dimension spanning every storage axis that reads it, TensorStore's - representation. + library materializes their per-axis rows for vectorized consumers. Pure + affine diagonals need grouping by input dimension; mixed affine/index-array + dependencies need joint partitioning. TensorStore classifies a connected + component containing index-array edges as an index-array set + ([source](https://github.com/google/tensorstore/blob/66b2ce5290fa2ec5c8019682391421062ce767a2/tensorstore/internal/grid_partition.h#L58-L67)). Four deliberate differences: @@ -282,13 +284,11 @@ broadcast singletons to retain those dependencies. Some current limits are: -- **Affine diagonals.** A hand-built transform in which two output maps read - one input dimension — two slice maps, or a slice map and an orthogonal index - array — is rejected at planning with `ValueError`; a correlated index array - varying over a dimension a slice map also reads is rejected with - `NotImplementedError`. No selection dialect produces either. Supporting them - needs a strided set per *input* dimension spanning all dependent storage - axes, TensorStore's connected-component representation. +- **Shared affine dependencies.** Planning rejects two affine output maps + sharing an input axis with `ValueError`. An index array sharing a varying + input axis with an affine map takes the general classification and raises + `NotImplementedError`. Pure affine diagonals would need grouping dependent + storage axes by input dimension; mixed components need joint partitioning. - **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/justfile b/packages/zarr-indexing/justfile index 24834a960d..e32f177ff0 100644 --- a/packages/zarr-indexing/justfile +++ b/packages/zarr-indexing/justfile @@ -5,10 +5,9 @@ default: @just --list -# Nothing here imports `zarr`; the suite runs against the repo-root -# environment so it shares the parent project's pinned test toolchain -# (hypothesis, ruff, pyright). The repo is not a uv workspace, so this package -# is layered in as an editable overlay — the same invocation CI uses. +# The suite runs against the root test environment, which also supplies +# Zarr for the Dask example. This package is layered in as an editable +# overlay, using the same test invocation as CI. # Run the test suite; extra args are passed to pytest test *args: uv run --project ../.. --group test --with-editable . python -m pytest tests src/zarr_indexing {{ args }} diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index fca9cdd5f3..5c01c9e85b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -1241,7 +1241,7 @@ def _intersect_general( # Scatter bookkeeping with row-major input-domain strides. Broadcast # coordinates below are positional; residual coordinates remain literal. - # `_correlated_cell_transform` removes their origins before unraveling. + # These are not uniformly zero-origin flat indices for nonzero domains. # # 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 diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index ef67c2cedb..c5b2c803ae 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -782,7 +782,7 @@ def test_mixed_affine_array_dependency_is_rejected() -> None: IndexDomain.from_shape((4,)), (DimensionMap(0), ArrayMap(np.array([3, 2, 1, 0]))) ) grids = dimension_grids_from_chunks((2, 2), shape=(4, 4)) - with pytest.raises(ValueError, match="read input axis 0"): + with pytest.raises(NotImplementedError, match="also bound by a slice map"): list(plan_chunks(transform, grids)) From 1f7f020439704a36e26cfe6f53a86816e919dc07 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Sun, 13 Sep 2026 18:57:17 +0200 Subject: [PATCH 07/12] docs(indexing): clarify planning coverage and benchmark measurement boundaries Assisted-by: Codex:GPT-6 --- packages/zarr-indexing/benchmarks/README.md | 13 +++++++++++++ packages/zarr-indexing/benchmarks/chunk_planning.py | 7 +++++++ packages/zarr-indexing/docs/guide/selection-flow.md | 5 ++++- packages/zarr-indexing/docs/release-notes.md | 3 ++- .../zarr-indexing/docs/snippets/grid_partition.py | 6 +++--- packages/zarr-indexing/mkdocs.yml | 2 +- 6 files changed, 30 insertions(+), 6 deletions(-) diff --git a/packages/zarr-indexing/benchmarks/README.md b/packages/zarr-indexing/benchmarks/README.md index 28ef3e6a16..e5d39b1a89 100644 --- a/packages/zarr-indexing/benchmarks/README.md +++ b/packages/zarr-indexing/benchmarks/README.md @@ -22,6 +22,19 @@ whether input arrays and transform construction are included. Repeated local column access measures cache reuse, which trades allocation against retained memory. Bounded coordinate batches avoid constructing the full coordinate array. +The current script reports `peak_mib` as the incremental peak tracked by +`tracemalloc` during a separate invocation after the timed calls. It is not +process RSS or total retained memory. Inputs and most transforms are constructed +before measurement. The case walks construct plans and consume projections +without retaining them; `all_coordinates` materializes the complete coordinate +array. There is no bounded-batch coordinate workload in this script. + +`local_rows` reuses one prepared table: its first timed invocation populates the +local-coordinate cache, while later timed invocations and the allocation probe +reuse that cache. Its median mixes a cold first call with warm calls, and its +reported peak excludes the already-retained cache. A separate fresh-table +measurement would be needed to quantify cold cache construction. + These scripts measure planning rather than codec or storage throughput. Repeat measurements with alternating operation order before interpreting small timing differences. Preserve raw benchmark output as an experiment artifact rather than diff --git a/packages/zarr-indexing/benchmarks/chunk_planning.py b/packages/zarr-indexing/benchmarks/chunk_planning.py index 934d68dcbc..d7eadf248d 100644 --- a/packages/zarr-indexing/benchmarks/chunk_planning.py +++ b/packages/zarr-indexing/benchmarks/chunk_planning.py @@ -25,6 +25,11 @@ def measure(operation: Callable[[], Any], repeats: int) -> dict[str, float]: + """Time repeated calls, then trace one additional call's incremental peak. + + State retained by earlier calls can affect the extra call; this is not a + measurement of process RSS or all memory owned by the operation's inputs. + """ samples = [] for _ in range(repeats): start = time.perf_counter() @@ -84,6 +89,8 @@ def read_local_rows() -> None: for row in range(len(table)): table.local[table.run(row)] + # The first timed call fills table.local; later calls, including the traced + # allocation probe, reuse it. The table and retained cache are not re-created. results["local_rows"] = measure(read_local_rows, args.repeats) part = plan_chunks( IndexTransform.from_shape((100, 100, 100)), diff --git a/packages/zarr-indexing/docs/guide/selection-flow.md b/packages/zarr-indexing/docs/guide/selection-flow.md index 571baa3869..749a9c36ae 100644 --- a/packages/zarr-indexing/docs/guide/selection-flow.md +++ b/packages/zarr-indexing/docs/guide/selection-flow.md @@ -171,7 +171,7 @@ pairs: | `(0,)` | `[1]` | `[1]` | | `(1,)` | `[2, 2, 0]` | `[0, 2, 3]` | -Chunk `(1,)` is fetched once, but its local value `2` contributes to two result +Chunk `(1,)` is fetched once, but the value at local coordinate `2` contributes to two result positions. Chunk visitation order need not be result order; `out_selection` restores the requested arrangement. @@ -188,6 +188,9 @@ than one. Repeated coverage is not full; fancy coverage remains conservative. Zarr's merge operation can skip a read for a full data-extent write and allocate a fill-valued codec buffer when the selected data is smaller than that buffer. +That shortcut also requires the consumer's expected value layout and order; +complete coverage by a reversal or reordered gather alone does not establish +that the supplied buffer can be copied directly into codec order. Touching every chunk alone does not prove full coverage of each chunk. Planning does not fetch bytes, decode buffers, choose concurrency, or define diff --git a/packages/zarr-indexing/docs/release-notes.md b/packages/zarr-indexing/docs/release-notes.md index 767cd56a84..d0687d8758 100644 --- a/packages/zarr-indexing/docs/release-notes.md +++ b/packages/zarr-indexing/docs/release-notes.md @@ -1,5 +1,6 @@