diff --git a/.github/workflows/check_changelogs.yml b/.github/workflows/check_changelogs.yml index 0033b43db2..d7a54fc2c4 100644 --- a/.github/workflows/check_changelogs.yml +++ b/.github/workflows/check_changelogs.yml @@ -29,3 +29,6 @@ jobs: - name: Check zarr-metadata changelog entries run: uv run --no-sync python ci/check_changelog_entries.py packages/zarr-metadata/changes + + - name: Check zarr-indexing changelog entries + run: uv run --no-sync python ci/check_changelog_entries.py packages/zarr-indexing/changes diff --git a/.github/workflows/zarr-indexing-release.yml b/.github/workflows/zarr-indexing-release.yml new file mode 100644 index 0000000000..bdba2c0db0 --- /dev/null +++ b/.github/workflows/zarr-indexing-release.yml @@ -0,0 +1,117 @@ +name: zarr-indexing release + +on: + workflow_dispatch: + push: + tags: + - 'zarr_indexing-v*' + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: false + +jobs: + build: + name: Build wheel and sdist + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-indexing + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + fetch-depth: 0 # hatch-vcs needs full history + tags + + - name: Install Hatch + uses: pypa/hatch@257e27e51a6a5616ed08a39a408a21c35c9931bc + with: + version: '1.16.5' + + - name: Build + run: hatch build + + - uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: zarr-indexing-dist + path: packages/zarr-indexing/dist + + test_artifacts: + name: Test built artifacts + needs: [build] + runs-on: ubuntu-latest + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-indexing-dist + path: dist + + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + with: + enable-cache: false + + - name: Set up Python + run: uv python install 3.12 + + - name: Install built wheel and run import smoke test + run: | + wheel=$(ls dist/*.whl) + uv run --with "${wheel}" --python 3.12 --no-project \ + python -c "import zarr_indexing; print('zarr_indexing', zarr_indexing.__version__)" + + upload_pypi: + name: Upload to PyPI + needs: [build, test_artifacts] + if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/zarr_indexing-v') + runs-on: ubuntu-latest + environment: + name: zarr-indexing-releases + url: https://pypi.org/p/zarr-indexing + permissions: + id-token: write # required for OIDC trusted publishing + attestations: write # required for artifact attestations + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-indexing-dist + path: dist + + - name: Generate artifact attestation + uses: actions/attest@59d89421af93a897026c735860bf21b6eb4f7b26 # v4.1.0 + with: + subject-path: dist/* + + - name: Publish package to PyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 + + upload_testpypi: + name: Upload to TestPyPI + needs: [build, test_artifacts] + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + environment: + name: zarr-indexing-releases-test + url: https://test.pypi.org/p/zarr-indexing + permissions: + id-token: write + attestations: write + steps: + - uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: zarr-indexing-dist + path: dist + + - name: Generate artifact attestation + uses: actions/attest@59d89421af93a897026c735860bf21b6eb4f7b26 # v4.1.0 + with: + subject-path: dist/* + + - name: Publish package to TestPyPI + uses: pypa/gh-action-pypi-publish@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0 + with: + repository-url: https://test.pypi.org/legacy/ diff --git a/.github/workflows/zarr-indexing.yml b/.github/workflows/zarr-indexing.yml new file mode 100644 index 0000000000..7e1e456a7a --- /dev/null +++ b/.github/workflows/zarr-indexing.yml @@ -0,0 +1,102 @@ +name: zarr-indexing + +on: + push: + branches: [main] + paths: + - 'packages/zarr-indexing/**' + - '.github/workflows/zarr-indexing.yml' + pull_request: + paths: + - 'packages/zarr-indexing/**' + - '.github/workflows/zarr-indexing.yml' + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + name: pytest py=${{ matrix.python-version }} + runs-on: ubuntu-latest + defaults: + run: + shell: bash + strategy: + fail-fast: false + matrix: + python-version: ['3.12', '3.13', '3.14'] + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + with: + enable-cache: true + - name: Set up Python ${{ matrix.python-version }} + run: uv python install ${{ matrix.python-version }} + # The transform tests exercise chunk resolution against zarr's ChunkGrid, + # so they run from the repo root against the full workspace (which installs + # both `zarr` and the `zarr-indexing` workspace member) rather than in + # package isolation. + - name: Sync test dependency group + run: uv sync --group test --python ${{ matrix.python-version }} + - name: Run pytest + run: uv run --group test pytest packages/zarr-indexing/tests + + ruff: + name: ruff + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-indexing + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + - name: Run ruff + run: uvx ruff check . + + pyright: + name: pyright + runs-on: ubuntu-latest + defaults: + run: + shell: bash + working-directory: packages/zarr-indexing + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + persist-credentials: false + - name: Install uv + uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0 + with: + enable-cache: true + - name: Set up Python + run: uv python install 3.12 + - name: Sync test dependency group + run: uv sync --group test --python 3.12 + - name: Run pyright + run: uv run --group test --with pyright pyright src + + zarr-indexing-complete: + name: zarr-indexing complete + needs: [test, ruff, pyright] + if: always() + runs-on: ubuntu-latest + steps: + - name: Check failure + if: | + contains(needs.*.result, 'failure') || + contains(needs.*.result, 'cancelled') + run: exit 1 + - name: Success + run: echo Success! diff --git a/changes/3906.feature.md b/changes/3906.feature.md new file mode 100644 index 0000000000..d2d6e80d95 --- /dev/null +++ b/changes/3906.feature.md @@ -0,0 +1 @@ +Added a `.lazy` accessor to `Array` for lazy (deferred) indexing. `array.lazy[...]`, `array.lazy.oindex[...]`, and `array.lazy.vindex[...]` return a new `Array` view composed on a TensorStore-style index-transform layer; slicing and indexing a view never triggers I/O, and views can be composed further before reading. Call `.result()` on a view to resolve it and return the underlying data. Regular (eager) indexing semantics on `Array` and `AsyncArray` are unchanged, with two deliberate exceptions. First, the `repr` of a lazy view now carries a `domain={...}` suffix describing the view's coordinate box; the `repr` of an ordinary (identity) `Array` or `AsyncArray` is unchanged from 3.x and never shows this suffix. Second, iterating a 0-d `Array` now raises `TypeError("iteration over a 0-d array")` to match NumPy (previously it silently yielded nothing). Arrays also gain `Array.chunk_projections`, which enumerates the stored chunks a view (or a whole array) projects onto, for view-aware chunk partitioning. The index-transform machinery that powers this lives in a new NumPy-only `zarr-indexing` workspace package (import name `zarr_indexing`), which `zarr` now depends on at runtime. diff --git a/docs/api/zarr/array.md b/docs/api/zarr/array.md index ff61cb1fe2..2445ee58d3 100644 --- a/docs/api/zarr/array.md +++ b/docs/api/zarr/array.md @@ -1,2 +1,3 @@ ::: zarr.Array ::: zarr.AsyncArray +::: zarr.ChunkProjection diff --git a/docs/user-guide/lazy_indexing.md b/docs/user-guide/lazy_indexing.md new file mode 100644 index 0000000000..25b06658bf --- /dev/null +++ b/docs/user-guide/lazy_indexing.md @@ -0,0 +1,306 @@ +# Lazy indexing + +Zarr arrays support *lazy* indexing through the `Array.lazy` accessor. Where +ordinary indexing reads or writes data immediately, `a.lazy[selection]` returns a +lightweight **view** — itself a `zarr.Array` — without touching storage. Views +compose, support orthogonal and coordinate selection, write through to the +backing array, and materialize on demand. + +Zarr's lazy indexing follows [TensorStore's indexing +model](https://google.github.io/tensorstore/python/indexing.html): a view has a +**domain** — a box of coordinates — and a positive index is a **literal +coordinate** in that domain (a negative index counts from the domain's end, as +in NumPy). + +```python exec="true" session="lazy" +import numpy as np +import zarr +``` + +```python exec="true" session="lazy" source="above" result="ansi" +a = zarr.create_array(store="memory://lazy-demo", shape=(12,), chunks=(3,), dtype="int32") +a[...] = np.arange(12) + +view = a.lazy[2:10] # no I/O happens here +print(view) +print(view.result()) # I/O happens here +``` + +## Theory: indexing as a declaration + +Eager indexing is an *action*: `a[2:10]` performs I/O now and hands back the +bytes. Lazy indexing is a *declaration*: `a.lazy[2:10]` records **which cells +you mean** — an index transform mapping the view's coordinates to storage +coordinates — and defers the I/O. Because the selection is data rather than an +action, zarr can: + +- **compose** it with further selections without reading anything, +- **write through** it (the same transform routes values back to storage), +- **plan** with it (`chunk_projections` enumerates exactly the stored chunks + the declaration touches). + +### Domains are preserved: an index is a name, not a position + +A view keeps the coordinates of the cells it selects. Slicing `[2:10]` does not +renumber anything — the view's domain *is* `[2, 10)`, and coordinate 3 still +means what it meant on the parent: + +```python exec="true" session="lazy" source="above" result="ansi" +v = a.lazy[2:10] +print(v.shape) # (8,) — eight cells... +print(v.lazy[3].result()) # ...and coordinate 3 is still base cell 3 +print(v.lazy[3:7].result()) # coordinates [3, 7) — literal, stable +``` + +This is what makes composition safe: a coordinate means the same cell no matter +how many views deep you are. `a.lazy[2:10].lazy[3:7]` and `a.lazy[3:7]` are the +same selection. + +The price of stable names: **positions are not valid indices.** The first +element of `v` is coordinate 2, not 0 — and `v[0]` is an error, not the first +element: + +```python exec="true" session="lazy" source="above" result="ansi" +try: + v[0] +except IndexError as e: + print(e) +``` + +To renumber a view explicitly, move its domain with `translate_to` (or shift it +with `translate_by`) — the data does not move, only the labels: + +```python exec="true" session="lazy" source="above" result="ansi" +z = v.translate_to((0,)) # same cells, coordinates now [0, 8) +print(z.lazy[0].result()) # coordinate 0 -> base cell 2 +``` + +### Negative indices count from the end + +*Positive* indices are literal coordinates (so positions still are not indices — +see above), but a **negative** index counts from the end of the current view's +domain, exactly like NumPy: `k` maps to `exclusive_max + k`. This holds in every +syntactic form — integers, slice bounds, and index arrays — at the public +boundary (`a.lazy[...]`, `v.lazy[...]`, `v[...]`, `.oindex`, `.vindex`, and the +`get_/set_*_selection` methods). For a fresh, zero-origin array both rules +coincide with NumPy in full: + +```python exec="true" session="lazy" source="above" result="ansi" +print(a.lazy[-1].result()) # last element, like NumPy +print(a.lazy[-3:].result()) # the last three +print(a.lazy.oindex[[-1, -2]].result()) # index arrays wrap too +``` + +On a view the wrap is relative to the *view's* end: on the domain-`[2, 10)` +view `v`, `-1` names coordinate `9`, and `-8` names the first cell (coordinate +`2`). Only an out-of-range index raises — a positive past the end, or a +negative below `-size` (one that would wrap past the domain origin): + +```python exec="true" session="lazy" source="above" result="ansi" +n = a.shape[0] +for make in (lambda: a.lazy[n], lambda: a.lazy[-n - 1], lambda: a.lazy.oindex[[-n - 1]]): + try: + make() + except IndexError as e: + print(type(e).__name__, "-", str(e).split(";")[0]) +``` + +### No clamping: intervals must fit the domain + +A slice interval must be contained in the domain — an out-of-range bound is an +error, not a silently shorter result. Empty intervals are the one exception: +they are valid anywhere. Reversed bounds are an error, not an empty result. + +```python exec="true" session="lazy" source="above" result="ansi" +for sel in (slice(5, 100), slice(5, 2)): + try: + a.lazy[sel] + except IndexError as e: + print(type(e).__name__, "-", str(e).split(";")[0]) +print(a.lazy[5:5].shape) # empty is fine, anywhere +``` + +### Strided views renumber by division + +A strided slice produces a domain in *strided units*: for step `k`, the new +origin is `start / k` rounded toward zero, and coordinate `origin + i` maps to +base cell `start + i*k` (TensorStore's rule): + +```python exec="true" session="lazy" source="above" result="ansi" +s = a.lazy[1:10:3] # base cells 1, 4, 7 +print(s) # domain [0, 3) +print(s.lazy[1].result()) # coordinate 1 -> base cell 4 +``` + +### One coordinate system per view + +Every way of indexing a view — `v[...]`, `v.lazy[...]`, `v.oindex`, `v.vindex` +— uses the same rule: positive indices are literal domain coordinates, and +negative indices wrap from the domain's end. (The base array's ordinary `a[...]` +API is unchanged: it keeps full NumPy semantics.) NumPy-style zero-based access +to a view's *positive* range is spelled explicitly: materialize with `result()` +/ `np.asarray`, or renumber with `translate_to`. + +```python exec="true" session="lazy" source="above" result="ansi" +print(v[3], v.lazy[3].result()) # same coordinate, same cell +print(np.asarray(v)[0]) # materialized: NumPy rules apply +print(a[-1]) # base arrays keep NumPy semantics +``` + +## Common patterns + +### Crop, analyze, crop again + +```python exec="true" session="lazy" source="above" result="ansi" +img = zarr.create_array(store="memory://lazy-img", shape=(100, 100), chunks=(10, 10), dtype="float64") +img[...] = np.arange(100 * 100).reshape(100, 100) + +crop = img.lazy[25:75, 25:75] # no I/O; domain [25,75) x [25,75) +inner = crop.lazy[35:65, 35:65] # coordinates are literal: this is img[35:65, 35:65] +print(crop.shape, inner.shape) +print(float(np.mean(inner))) # I/O happens here, for the inner crop only +``` + +### Write through a view + +Assignment through the accessor, or through a view, routes values back to +storage — including strided and composed selections: + +```python exec="true" session="lazy" source="above" result="ansi" +img.lazy[30:50, 40:60] = 0.0 # region write +tile = img.lazy[30:50, 40:60] +tile[30:35, 40:45] = 7.0 # write through the view, same coordinates +img.lazy[::2, ::2] = -1.0 # strided write, NumPy-equivalent cells +print(img[29:33, 39:43]) +``` + +### Orthogonal and coordinate selection + +`lazy.oindex` selects an outer product per axis; `lazy.vindex` selects points. +Index-array *values* are domain coordinates; the dimension a fancy selection +*creates* gets a fresh `[0, n)` domain (there is no meaningful coordinate to +preserve for "the i-th pick"): + +```python exec="true" session="lazy" source="above" result="ansi" +rows = img.lazy.oindex[[3, 17, 42], :] # picked dim: domain [0, 3); row dim preserved +sub = rows.lazy[:, 10:20] # column window, literal coords +print(rows.shape, sub.shape) + +pts = img.lazy.vindex[[3, 17], [5, 9]] # two points -> fresh domain [0, 2) +print(pts.result()) +``` + +Boolean masks are array selections, so they go through `oindex`/`vindex`; the +positions of `True` values become **coordinates, counted from 0** — not offsets +from the view's origin. On a view whose domain starts at 2, a mask `True` at +position 3 addresses coordinate 3, and a `True` at position 0 or 1 is out of +the domain (matching TensorStore, where a mask is sugar for the coordinate +array of its `True` positions): + +```python exec="true" session="lazy" source="above" result="ansi" +mask = np.zeros(12, dtype=bool) +mask[[1, 4, 6]] = True +print(a.lazy.oindex[mask].result()) +``` + +### Materializing + +`view.result()`, `view[...]`, and `np.asarray(view)` are equivalent whole-view +reads; views also work directly with NumPy reductions. Views are **not** +iterable (iterate the materialized result instead): + +```python exec="true" session="lazy" source="above" result="ansi" +w2 = a.lazy[3:9] +print(w2.result(), float(np.mean(w2))) +try: + iter(w2) +except TypeError as e: + print(e) +``` + +### Chunk-aware processing + +`chunk_projections` enumerates the stored chunks a view touches: which store +object (`coord`, `key`), its stored `shape`, the region *within the chunk* the +view covers (`chunk_selection`), the region *of the view* it maps to +(`array_selection`, positional — 0-based into the view's extent), and whether +the chunk is only partially covered (`is_partial` — a partial write requires a +read-modify-write): + +```python exec="true" session="lazy" source="above" result="ansi" +for p in a.lazy[2:10].chunk_projections(): + print(p.key, p.chunk_selection, p.array_selection, p.is_partial) +print(a.lazy[2:10].is_chunk_aligned()) +print(a.lazy[3:9].is_chunk_aligned()) # starts and ends on chunk boundaries +``` + +This is the supported way to partition any selection for parallel or +chunk-at-a-time work — compose the selection through `.lazy`, then project. +Since `array_selection` is positional, re-zero the view (or materialize) to use +it: + +```python exec="true" session="lazy" source="above" result="ansi" +crop0 = img.lazy[25:75, 25:75].translate_to((0, 0)) +total = 0.0 +for p in crop0.chunk_projections(): + total += float(np.sum(crop0[tuple(slice(s.start, s.stop) for s in p.array_selection)])) +print(total == float(np.sum(crop0))) +``` + +For sharded arrays, pass `unit="write"` to enumerate at shard (write-unit) +granularity; read-unit projections for sharded arrays are not yet implemented. + +### What a view will not tell you + +Members that describe the chunk grid assume the array *fills* its grid, which a +view generally does not. On a view they raise `zarr.errors.LazyViewError` +instead of silently describing the backing array: + +```python exec="true" session="lazy" source="above" result="ansi" +try: + v.chunks +except zarr.errors.LazyViewError as e: + print(e) +``` + +Logical members (`shape`, `size`, `nbytes`, `dtype`, `attrs`, ...) reflect the +view; `metadata` and `chunk_grid` remain available and describe the *backing* +array. + +## Coming from NumPy + +- **A view's positive indices are coordinates, not positions.** `a.lazy[2:10]` + is indexed with 2..9, not 0..7. Renumber explicitly with + `view.translate_to((0, ...))` if you want positions. +- **Negative indices count from the end**, exactly like NumPy — in every form + (integer, slice bound, index array). `k` maps to `exclusive_max + k` in the + current view's domain; only a negative below `-size` (or a positive past the + end) is out of bounds. +- **No clamping**: out-of-range slice bounds raise (including a negative that + would wrap past the domain origin); reversed bounds raise; only empty + intervals are allowed anywhere. +- **No negative steps**: `a.lazy[::-1]` raises; reversal is not yet supported. +- **No `newaxis`**: `a.lazy[None]` raises; insert axes on the materialized + result instead. +- **The basic accessor takes basic selections only** (integers, slices, + ellipsis). Lists, arrays, and boolean masks go through `lazy.oindex` / + `lazy.vindex`. +- **Views are not iterable**; iterate `view.result()`. +- **Base arrays are unchanged**: `a[-1]`, `a[5:100]`, and friends keep full + NumPy semantics on non-view arrays. + +## Current limitations + +- Negative slice steps (reversal) are not yet supported. +- Integer indexing a dimension *created by* an `oindex`/`vindex` selection + (e.g. `rows.lazy[0]` after `rows = a.lazy.oindex[[3, 17, 42], :]`) is not yet + supported reliably; slice the view instead (`rows.lazy[0:1]`). +- Composing a fancy (`oindex`/`vindex`) selection onto a view that already has a + fancy-indexed axis (fancy-after-fancy — e.g. `a.lazy.oindex[[0, 2, 5], :]` + then `.oindex[:, [5]]`) is not supported and raises `NotImplementedError`. + Materialize the view first with `.result()` and index the array, or reorder + the selections so the fancy step is applied last. +- `chunk_projections(unit="read")` on sharded arrays (inner-chunk granularity) + is not yet implemented; use `unit="write"`. +- Views cannot be resized or appended to, and block selection is not defined + for views. diff --git a/mkdocs.yml b/mkdocs.yml index 46bfc1764c..6d19c239cf 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -15,6 +15,7 @@ nav: - user-guide/index.md - user-guide/installation.md - user-guide/arrays.md + - user-guide/lazy_indexing.md - user-guide/groups.md - user-guide/attributes.md - user-guide/data_types.md diff --git a/packages/zarr-indexing/CHANGELOG.md b/packages/zarr-indexing/CHANGELOG.md new file mode 100644 index 0000000000..7c4bc92cad --- /dev/null +++ b/packages/zarr-indexing/CHANGELOG.md @@ -0,0 +1,3 @@ +# Release notes + + diff --git a/packages/zarr-indexing/LICENSE.txt b/packages/zarr-indexing/LICENSE.txt new file mode 100644 index 0000000000..1e8da4d242 --- /dev/null +++ b/packages/zarr-indexing/LICENSE.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2015-2025 Zarr Developers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md new file mode 100644 index 0000000000..86029614e9 --- /dev/null +++ b/packages/zarr-indexing/README.md @@ -0,0 +1,30 @@ +# zarr-indexing + +Composable, lazy coordinate transforms for Zarr array indexing. + +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. + +Key types: + +- `IndexDomain` — a rectangular region of integer coordinates +- `IndexTransform` — maps input coordinates to storage coordinates +- `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single output + 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. + +## Installation + +```bash +pip install zarr-indexing +``` + +## License + +MIT diff --git a/packages/zarr-indexing/changes/3906.feature.md b/packages/zarr-indexing/changes/3906.feature.md new file mode 100644 index 0000000000..ca39164018 --- /dev/null +++ b/packages/zarr-indexing/changes/3906.feature.md @@ -0,0 +1 @@ +Reworked the JSON layer to conform to the [ndsel](https://github.com/d-v-b/ndsel) draft wire format, which adapts TensorStore's `IndexTransform`. A new `zarr_indexing.messages` module (`parse_ndsel`, `normalize_ndsel`, `NdselError`) is a pure JSON-to-JSON layer that accepts all five message kinds (`point`/`box`/`slice`/`points`/`transform`) and normalizes them to the canonical transform body, enforcing the full ndsel error taxonomy. The package is checked against the vendored, language-agnostic ndsel conformance corpus. `index_transform_to_json`/`index_transform_from_json` (and the domain variants) now produce and consume the canonical body. On serialization, orthogonal (`oindex`) `index_array` maps no longer emit `input_dimension` alongside `index_array` (a combination both ndsel and TensorStore reject), and degenerate all-singleton index arrays collapse to constant maps; the in-memory `input_dimension` is reconstructed from the array's dependency axes on load. diff --git a/packages/zarr-indexing/changes/README.md b/packages/zarr-indexing/changes/README.md new file mode 100644 index 0000000000..feb3f8674e --- /dev/null +++ b/packages/zarr-indexing/changes/README.md @@ -0,0 +1,25 @@ +Writing a changelog entry for `zarr-indexing` +----------------------------------------------- + +Fragments in **this** directory are release notes for the `zarr-indexing` +package only — kept separate from the parent zarr-python `changes/` +directory so a PR touching only `packages/zarr-indexing/` produces a +release note for this package only. + +Please put a new file in this directory named `xxxx..md`, where + +- `xxxx` is the pull request number associated with this entry +- `` is one of: + - feature + - bugfix + - doc + - removal + - misc + +Inside the file, please write a short description of what you have +changed, and how it impacts users of `zarr-indexing`. + +A `zarr-indexing` release runs `towncrier build` in `packages/zarr-indexing/`, +which consumes the fragments here and updates `CHANGELOG.md`. Fragments +that describe parent zarr-python changes (not the transforms package) +belong in the top-level `changes/` directory, not here. diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml new file mode 100644 index 0000000000..f29afe2f4a --- /dev/null +++ b/packages/zarr-indexing/pyproject.toml @@ -0,0 +1,113 @@ +[build-system] +requires = ["hatchling>=1.29.0", "hatch-vcs"] +build-backend = "hatchling.build" + +[project] +name = "zarr-indexing" +dynamic = ["version"] +description = "Composable, lazy coordinate transforms for Zarr array indexing." +readme = "README.md" +requires-python = ">=3.12" +license = "MIT" +license-files = ["LICENSE.txt"] +authors = [ + { name = "Davis Bennett", email = "davis.v.bennett@gmail.com" }, +] +classifiers = [ + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Information Technology", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Operating System :: OS Independent", + "Programming Language :: Python", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", + "Topic :: Scientific/Engineering", + "Topic :: Software Development :: Libraries :: Python Modules", + "Typing :: Typed", +] +keywords = ["zarr"] +dependencies = [ + "numpy>=2", +] + +[project.urls] +Homepage = "https://github.com/zarr-developers/zarr-python" +Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-indexing" +Issues = "https://github.com/zarr-developers/zarr-python/issues" +Changelog = "https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/CHANGELOG.md" +Documentation = "https://github.com/zarr-developers/zarr-python/blob/main/packages/zarr-indexing/README.md" + +[dependency-groups] +# The transform tests exercise chunk resolution against zarr's ChunkGrid +# (tests/test_chunk_resolution.py) and are collected by the parent zarr-python +# test suite, which already has zarr installed. `zarr` is intentionally NOT +# listed here to avoid a workspace dependency cycle; run these tests from the +# repo root (`uv run pytest packages/zarr-indexing/tests`), not in isolation. +test = ["pytest"] + +[tool.hatch.version] +source = "vcs" +tag-pattern = '^zarr_indexing-v(?P.+)$' +# `git_describe_command` ensures we get the zarr_indexing tags instead of latest. +# `local_scheme` strips the git commit info so the appending info is just a counter from latest tag. +# test-pypi doesn't accept git commit info in tags, and the count should be enough to distinguish unique runs. +raw-options = { root = "../..", git_describe_command = "git describe --dirty --tags --long --match zarr_indexing-v*", local_scheme = "no-local-version" } + +[tool.hatch.build.targets.wheel] +packages = ["src/zarr_indexing"] + +[tool.ruff] +extend = "../../pyproject.toml" +target-version = "py312" + +[tool.pytest.ini_options] +minversion = "7" +testpaths = ["tests"] +xfail_strict = true +addopts = ["-ra", "--strict-config", "--strict-markers"] +filterwarnings = [ + "error", +] + +[tool.pyright] +include = ["src"] +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. +reportUnknownVariableType = "warning" +reportUnknownArgumentType = "warning" +reportUnknownMemberType = "warning" +reportUnknownParameterType = "warning" + +[tool.numpydoc_validation] +checks = [ + "GL10", + "SS04", + "PR02", + "PR03", + "PR05", + "PR06", +] + +[tool.towncrier] +# Fragments for this package live alongside the package source, separate +# from the parent zarr-python `changes/` directory, so a PR touching only +# `packages/zarr-indexing/` produces a release note for this package only. +directory = "changes" +filename = "CHANGELOG.md" +package = "zarr_indexing" +underlines = ["", "", ""] +title_format = "## {version} ({project_date})" +issue_format = "[#{issue}](https://github.com/zarr-developers/zarr-python/issues/{issue})" +start_string = "\n" diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py new file mode 100644 index 0000000000..6e999417d5 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -0,0 +1,75 @@ +"""Composable, lazy coordinate transforms for zarr array indexing. + +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. + +Key types: + +- `IndexDomain` — a rectangular region of integer coordinates +- `IndexTransform` — maps input coordinates to storage coordinates +- `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single + output dimension can depend on the input (see `output_map.py`) +- `compose` — chain two transforms into one + +The chunk-resolution helpers (`iter_chunk_transforms`, +`sub_transform_to_selections`) and `selection_to_transform` are also exported +here: they form the surface the zarr integration layer (array indexing) depends +on. The `*Like` grid Protocols describe the chunk-grid surface chunk resolution +consumes without importing zarr. +""" + +from importlib.metadata import version + +from zarr_indexing.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) +from zarr_indexing.composition import compose +from zarr_indexing.domain import IndexDomain +from zarr_indexing.grid import ChunkGridLike, DimensionGridLike +from zarr_indexing.json import ( + IndexDomainJSON, + IndexTransformJSON, + OutputIndexMapJSON, + index_domain_from_json, + index_domain_to_json, + index_transform_from_json, + index_transform_to_json, + transform_from_canonical, + transform_to_canonical, +) +from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.transform import IndexTransform, selection_to_transform + +__version__ = version("zarr-indexing") + +__all__ = [ + "ArrayMap", + "ChunkGridLike", + "ConstantMap", + "DimensionGridLike", + "DimensionMap", + "IndexDomain", + "IndexDomainJSON", + "IndexTransform", + "IndexTransformJSON", + "NdselError", + "OutputIndexMap", + "OutputIndexMapJSON", + "__version__", + "compose", + "index_domain_from_json", + "index_domain_to_json", + "index_transform_from_json", + "index_transform_to_json", + "iter_chunk_transforms", + "normalize_ndsel", + "parse_ndsel", + "selection_to_transform", + "sub_transform_to_selections", + "transform_from_canonical", + "transform_to_canonical", +] diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py new file mode 100644 index 0000000000..0dec19346c --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -0,0 +1,355 @@ +"""Chunk resolution — mapping transforms to chunk-level I/O. + +Given an `IndexTransform` (which coordinates a user wants to access) and a +`ChunkGrid` (how storage is divided into chunks), chunk resolution answers: + + For each chunk, which storage coordinates does this transform touch, + and where do those values land in the output buffer? + +The algorithm is: + +1. **Enumerate candidate chunks** — determine which chunks could possibly + be touched by the transform's output coordinate ranges. + +2. **Intersect** — for each candidate chunk, call + `transform.intersect(chunk_domain)` to restrict the transform to + coordinates within that chunk. If the intersection is empty, skip it. + +3. **Translate** — shift the restricted transform to chunk-local coordinates + via `transform.translate(-chunk_origin)`. + +4. **Yield** — produce `(chunk_coords, local_transform, surviving_indices)` + triples that the codec pipeline consumes. + +Sorted one-dimensional correlated array maps can be partitioned directly +because every touched chunk owns a contiguous slice of the index array. That +case bypasses candidate enumeration and repeated intersection. + +`sub_transform_to_selections` bridges from the transform representation +back to the raw `(chunk_selection, out_selection, drop_axes)` tuples that +the current codec pipeline expects. This bridge will go away when the codec +pipeline accepts transforms natively. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + + from zarr_indexing.grid import ChunkGridLike, DimensionGridLike + +OutIndices = ( + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None +) + +ChunkTransformResult = tuple[ + tuple[int, ...], + IndexTransform, + OutIndices, +] + + +def _one_dimensional_correlated_array_map( + transform: IndexTransform, +) -> tuple[ArrayMap, np.ndarray[Any, np.dtype[np.intp]]] | None: + """Return a nonempty correlated 1-D ArrayMap and its storage coordinates. + + A one-dimensional array selection has no cross-dimensional correlation to + preserve. The computed storage coordinates are also reused by general + resolution when they are unsorted. + """ + if transform.input_rank != 1 or transform.output_rank != 1: + return None + + m = transform.output[0] + if ( + not isinstance(m, ArrayMap) + or m.input_dimension is not None + or m.index_array.ndim != 1 + or m.index_array.size == 0 + ): + return None + + return m, m.offset + m.stride * m.index_array + + +def _iter_sorted_1d_array_map( + m: ArrayMap, + storage: np.ndarray[Any, np.dtype[np.intp]], + dim_grid: DimensionGridLike, +) -> Iterator[ChunkTransformResult]: + """Resolve a sorted 1-D ArrayMap one touched chunk at a time.""" + start = 0 + while start < storage.size: + chunk = dim_grid.index_to_chunk(int(storage[start])) + chunk_start = dim_grid.chunk_offset(chunk) + chunk_stop = chunk_start + dim_grid.chunk_size(chunk) + stop = int(np.searchsorted(storage, chunk_stop, side="left")) + + restricted = IndexTransform( + domain=IndexDomain(inclusive_min=(0,), exclusive_max=(stop - start,)), + output=( + ArrayMap( + index_array=m.index_array[start:stop], + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ), + ), + ) + local = restricted.translate((-chunk_start,)) + surviving = np.arange(start, stop, dtype=np.intp) + + yield (chunk,), local, surviving + start = stop + + +def iter_chunk_transforms( + transform: IndexTransform, + chunk_grid: ChunkGridLike, +) -> Iterator[ChunkTransformResult]: + """Resolve a composed IndexTransform against a ChunkGrid. + + Yields `(chunk_coords, sub_transform, out_indices)` triples: + + - `chunk_coords`: which chunk to access. + - `sub_transform`: maps output buffer coords to chunk-local coords. + - `out_indices`: for vectorized/array indexing, the output scatter + indices (integer array). `None` for basic/slice indexing. + """ + # `_dimensions` is declared on the `ChunkGridLike` Protocol in `grid.py` so + # that zarr's concrete `ChunkGrid` (whose own `_dimensions` is a private + # attribute) satisfies it structurally without a hard zarr import. Pyright's + # `reportPrivateUsage` still fires on the underscore name; renaming the + # Protocol member (or zarr's attribute) is an open pre-publish API + # decision, not made here. + dim_grids = chunk_grid._dimensions # pyright: ignore[reportPrivateUsage] + + array_map_1d = _one_dimensional_correlated_array_map(transform) + if array_map_1d is not None: + sorted_map, storage = array_map_1d + if storage[0] <= storage[-1] and bool(np.all(storage[1:] >= storage[:-1])): + dim_grid = dim_grids[0] + first_chunk = dim_grid.index_to_chunk(int(storage[0])) + if dim_grid.chunk_size(first_chunk) > 0: + yield from _iter_sorted_1d_array_map(sorted_map, storage, dim_grid) + return + + # Enumerate candidate chunks via the cartesian product of per-dimension + # candidate chunk ids, then for each candidate intersect the transform with + # the chunk domain (`transform.intersect` handles orthogonal and vectorized + # cases alike, filtering out combinations it does not actually touch). + # + # Each dimension contributes exactly the chunk ids it can touch: + # + # - `ConstantMap`/`DimensionMap` dims contribute a contiguous `range` — a + # single chunk for a constant, and the span between the first and last + # chunk for a slice. These are already tight (or nearly so). + # - `ArrayMap` (fancy) dims contribute only the *distinct* chunk ids the + # index array actually lands in (`np.unique`), never the dense + # `range(min_chunk, max_chunk + 1)` between them. A sparse fancy selection + # (e.g. two far-apart coordinates) would otherwise enumerate every chunk + # in the bounding box, making resolution scale with grid size instead of + # with the number of selected coordinates. + # + # For >= 2 correlated (vindex) ArrayMaps the per-dimension distinct sets + # over-approximate the *joint* touched set (their cartesian product includes + # combinations no single point lands in), but `intersect` filters those out, + # so the yielded chunks are identical either way — and the work stays bounded + # by the per-dimension distinct chunk counts, not the grid size. + chunk_candidates: list[Sequence[int]] = [] + for out_dim, m in enumerate(transform.output): + dg = dim_grids[out_dim] + if isinstance(m, ConstantMap): + # Single chunk + c = dg.index_to_chunk(m.offset) + chunk_candidates.append((c,)) + elif isinstance(m, DimensionMap): + d = m.input_dimension + dim_lo = transform.domain.inclusive_min[d] + dim_hi = transform.domain.exclusive_max[d] + if dim_lo >= dim_hi: + return # empty domain + if m.stride > 0: + s_min = m.offset + m.stride * dim_lo + s_max = m.offset + m.stride * (dim_hi - 1) + else: + s_min = m.offset + m.stride * (dim_hi - 1) + s_max = m.offset + m.stride * dim_lo + first = dg.index_to_chunk(s_min) + last = dg.index_to_chunk(s_max) + chunk_candidates.append(range(first, last + 1)) + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap). + # Storage coordinates were already computed for a correlated 1-D map. + storage = ( + array_map_1d[1] if array_map_1d is not None else m.offset + m.stride * m.index_array + ) + flat = storage.ravel().astype(np.intp) + if flat.size == 0: + # Empty fancy selection: no coordinates, so no chunks are touched. + return + chunk_ids = dg.indices_to_chunks(flat) + # Enumerate only the distinct chunks the coordinates land in. + chunk_candidates.append([int(c) for c in np.unique(chunk_ids)]) + + import itertools + + for chunk_coords_tuple in itertools.product(*chunk_candidates): + chunk_coords = tuple(int(c) for c in chunk_coords_tuple) + + # Build the chunk domain in storage space + chunk_min: list[int] = [] + chunk_max: list[int] = [] + chunk_shift: list[int] = [] + for out_dim, c in enumerate(chunk_coords): + dg = dim_grids[out_dim] + c_start = dg.chunk_offset(c) + c_size = dg.chunk_size(c) + chunk_min.append(c_start) + chunk_max.append(c_start + c_size) + chunk_shift.append(-c_start) + + chunk_domain = IndexDomain( + inclusive_min=tuple(chunk_min), + exclusive_max=tuple(chunk_max), + ) + + # Intersect transform with chunk domain + result = transform.intersect(chunk_domain) + if result is None: + continue + + restricted, surviving = result + + # Translate to chunk-local coordinates + local = restricted.translate(tuple(chunk_shift)) + + yield (chunk_coords, local, surviving) + + +def sub_transform_to_selections( + sub_transform: IndexTransform, + out_indices: OutIndices = None, +) -> tuple[ + tuple[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], + tuple[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]], ...], + tuple[int, ...], +]: + """Convert a chunk-local sub-transform to raw selections for the codec pipeline. + + Parameters + ---------- + sub_transform + A chunk-local IndexTransform (output maps already translated to + chunk-local coordinates). + out_indices + For vectorized indexing: the output scatter indices for this chunk. + None for orthogonal/basic indexing. + + Returns + ------- + tuple + `(chunk_selection, out_selection, drop_axes)` + """ + inclusive_min = sub_transform.domain.inclusive_min + exclusive_max = sub_transform.domain.exclusive_max + + # Orthogonal outer product: >= 2 ArrayMaps each bound to a distinct input + # dimension. out_indices is a per-output-dim dict of surviving positions. The + # codec applies chunk_array[chunk_sel] / out[out_sel] with NumPy semantics, so + # build np.ix_-style selections (mirroring the legacy OrthogonalIndexer): one + # 1-D selector per dimension, expanded to an open mesh. ConstantMap dims are + # size-1 in chunk space and squeezed out via drop_axes. + if isinstance(out_indices, dict): + chunk_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + out_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + drop_axes: list[int] = [] + for out_dim, m in enumerate(sub_transform.output): + if isinstance(m, ConstantMap): + chunk_arrays.append(np.array([m.offset], dtype=np.intp)) + drop_axes.append(out_dim) + elif isinstance(m, DimensionMap): + rng = np.arange(inclusive_min[m.input_dimension], exclusive_max[m.input_dimension]) + chunk_arrays.append((m.offset + m.stride * rng).astype(np.intp)) + out_arrays.append(rng.astype(np.intp)) + else: # ArrayMap + idx = m.index_array.ravel() + chunk_arrays.append((m.offset + m.stride * idx).astype(np.intp)) + out_arrays.append(out_indices[out_dim]) + return np.ix_(*chunk_arrays), np.ix_(*out_arrays), tuple(drop_axes) + + # Correlated (vindex) sub-transforms carry ArrayMaps with `input_dimension` + # None. They scatter through a single flat index (`out_indices`) into the + # row-major-flattened output buffer; the chunk selection reads a + # (points, residual-slice) block via the raveled coordinate arrays and any + # residual DimensionMap slices. + correlated = any( + isinstance(m, ArrayMap) and m.input_dimension is None for m in sub_transform.output + ) + if correlated: + chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + for m in sub_transform.output: + if isinstance(m, ConstantMap): + chunk_sel.append(m.offset) + elif isinstance(m, DimensionMap): + d = m.input_dimension + start = m.offset + m.stride * inclusive_min[d] + stop = m.offset + m.stride * exclusive_max[d] + if m.stride < 0: + start, stop = stop + 1, start + 1 + chunk_sel.append(slice(start, stop, m.stride)) + else: # ArrayMap + idx = m.index_array.reshape(-1) + chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) + # Chunk resolution always supplies the flat scatter index for a + # correlated transform. Absent one (a bare sub-transform), fall back to an + # identity scatter over the whole flattened output buffer. + # `out_indices` is narrowed to a flat scatter array or None here (the + # per-dimension dict is an orthogonal outer product, handled above). + out_scatter: slice | np.ndarray[Any, np.dtype[np.intp]] + if out_indices is None: + n = 1 + for s in sub_transform.domain.shape: + n *= s + out_scatter = slice(0, n) + else: + out_scatter = out_indices + return tuple(chunk_sel), (out_scatter,), () + + chunk_sel = [] # annotated in the correlated branch above (same function scope) + out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + + # Single-pass build for the basic / single-orthogonal-array cases. + # ConstantMap dims are dropped (no out_sel entry). + for m in sub_transform.output: + if isinstance(m, ConstantMap): + chunk_sel.append(m.offset) + elif isinstance(m, DimensionMap): + d = m.input_dimension + dim_lo = inclusive_min[d] + dim_hi = exclusive_max[d] + start = m.offset + m.stride * dim_lo + stop = m.offset + m.stride * dim_hi + if m.stride < 0: + start, stop = stop + 1, start + 1 + chunk_sel.append(slice(start, stop, m.stride)) + out_sel.append(slice(dim_lo, dim_hi)) + else: # ArrayMap (orthogonal: full-rank, raveled to its 1-D fancy coords) + idx = m.index_array.reshape(-1) + if m.offset == 0 and m.stride == 1: + chunk_sel.append(idx) + else: + chunk_sel.append((m.offset + m.stride * idx).astype(np.intp)) + # Orthogonal ArrayMap: out_indices holds the surviving positions. + out_sel.append(out_indices if out_indices is not None else slice(0, idx.size)) + + return tuple(chunk_sel), tuple(out_sel), () diff --git a/packages/zarr-indexing/src/zarr_indexing/composition.py b/packages/zarr-indexing/src/zarr_indexing/composition.py new file mode 100644 index 0000000000..8fa71f3e5b --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/composition.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import numpy as np + +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.transform import IndexTransform + + +def compose(outer: IndexTransform, inner: IndexTransform) -> IndexTransform: + """Compose two IndexTransforms. + + `outer` maps user coords (rank m) to intermediate coords (rank n). + `inner` maps intermediate coords (rank n) to storage coords (rank p). + The result maps user coords (rank m) to storage coords (rank p). + + Precondition: `outer.output_rank == inner.domain.ndim`. + """ + if outer.output_rank != inner.domain.ndim: + raise ValueError( + f"outer output rank ({outer.output_rank}) must match inner input rank " + f"({inner.domain.ndim})" + ) + + result_output = [_compose_single(outer, inner_map) for inner_map in inner.output] + + return IndexTransform(domain=outer.domain, output=tuple(result_output)) + + +def _compose_single(outer: IndexTransform, inner_map: OutputIndexMap) -> OutputIndexMap: + """Compose a single inner output map with the full outer transform.""" + if isinstance(inner_map, ConstantMap): + return ConstantMap(offset=inner_map.offset) + + if isinstance(inner_map, DimensionMap): + return _compose_dimension(outer, inner_map) + + # inner_map: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + return _compose_array(outer, inner_map) + + +def _compose_dimension(outer: IndexTransform, inner_map: DimensionMap) -> OutputIndexMap: + """Compose when inner is a DimensionMap. + + storage = offset_i + stride_i * intermediate[dim_i] + where intermediate[dim_i] = outer.output[dim_i](user_input) + """ + dim_i = inner_map.input_dimension + offset_i = inner_map.offset + stride_i = inner_map.stride + outer_map = outer.output[dim_i] + + if isinstance(outer_map, ConstantMap): + return ConstantMap(offset=offset_i + stride_i * outer_map.offset) + + if isinstance(outer_map, DimensionMap): + return DimensionMap( + input_dimension=outer_map.input_dimension, + offset=offset_i + stride_i * outer_map.offset, + stride=stride_i * outer_map.stride, + ) + + # outer_map: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + # Affine post-composition leaves the index array (and hence its full + # input rank and dependency axes) untouched; carry the orthogonal + # binding through unchanged. + return ArrayMap( + index_array=outer_map.index_array, + offset=offset_i + stride_i * outer_map.offset, + stride=stride_i * outer_map.stride, + input_dimension=outer_map.input_dimension, + ) + + +def _compose_array(outer: IndexTransform, inner_map: ArrayMap) -> OutputIndexMap: + """Compose when inner is an ArrayMap. + + storage = offset_i + stride_i * arr_i[intermediate] + We need to evaluate arr_i at the intermediate coordinates produced by outer. + """ + arr_i = inner_map.index_array + offset_i = inner_map.offset + stride_i = inner_map.stride + + # Check if all outer outputs are constant + all_constant = all(isinstance(m, ConstantMap) for m in outer.output) + + if all_constant: + # Evaluate arr_i at the single constant point + idx = tuple(m.offset for m in outer.output if isinstance(m, ConstantMap)) + value = int(arr_i[idx]) + return ConstantMap(offset=offset_i + stride_i * value) + + # For 1D inner array with a single outer output (simple case) + if arr_i.ndim == 1 and len(outer.output) == 1: + outer_map = outer.output[0] + + if isinstance(outer_map, DimensionMap): + dim_size = outer.domain.shape[outer_map.input_dimension] + user_indices = np.arange(dim_size, dtype=np.intp) + intermediate_vals = outer_map.offset + outer_map.stride * user_indices + new_arr = arr_i[intermediate_vals] + return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + + if isinstance(outer_map, ArrayMap): + intermediate_vals = outer_map.offset + outer_map.stride * outer_map.index_array + new_arr = arr_i[intermediate_vals] + return ArrayMap(index_array=new_arr, offset=offset_i, stride=stride_i) + + # General multi-dim case: not yet implemented + raise NotImplementedError( + "Composing a multi-dimensional inner array map with non-constant outer maps " + "is not yet supported." + ) diff --git a/packages/zarr-indexing/src/zarr_indexing/domain.py b/packages/zarr-indexing/src/zarr_indexing/domain.py new file mode 100644 index 0000000000..f20d5bf7bd --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/domain.py @@ -0,0 +1,189 @@ +"""Index domains — rectangular regions in N-dimensional integer space. + +An `IndexDomain` represents the set of valid coordinates for an array or +array view. It is the cartesian product of per-dimension integer ranges:: + + IndexDomain(inclusive_min=(2, 5), exclusive_max=(10, 20)) + # represents {(i, j) : 2 <= i < 10, 5 <= j < 20} + +Unlike NumPy, domains can have **non-zero origins**. After slicing +`arr[5:10]`, the result has origin 5 and shape 5 — coordinates 5 through +9 are valid. This follows the TensorStore convention. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True, slots=True) +class IndexDomain: + """A rectangular region in N-dimensional index space. + + The valid coordinates are the integers in + `[inclusive_min[d], exclusive_max[d])` for each dimension `d`. + """ + + inclusive_min: tuple[int, ...] + exclusive_max: tuple[int, ...] + labels: tuple[str, ...] | None = None + # Lazily-memoized shape. Excluded from init/repr/eq/hash: it is derived + # state, not part of the domain's identity. The domain is frozen, so the + # value is computed at most once (see `shape`). `None` is the unset + # sentinel; an empty shape caches as `()`. + _shape: tuple[int, ...] | None = field(default=None, init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + if len(self.inclusive_min) != len(self.exclusive_max): + raise ValueError( + f"inclusive_min and exclusive_max must have the same length. " + f"Got {len(self.inclusive_min)} and {len(self.exclusive_max)}." + ) + for i, (lo, hi) in enumerate(zip(self.inclusive_min, self.exclusive_max, strict=True)): + if lo > hi: + raise ValueError( + f"inclusive_min must be <= exclusive_max for all dimensions. " + f"Dimension {i}: {lo} > {hi}" + ) + if self.labels is not None and len(self.labels) != len(self.inclusive_min): + raise ValueError( + f"labels must have the same length as dimensions. " + f"Got {len(self.labels)} labels for {len(self.inclusive_min)} dimensions." + ) + + @classmethod + def from_shape(cls, shape: tuple[int, ...]) -> IndexDomain: + """Create a domain with origin at zero.""" + return cls( + inclusive_min=(0,) * len(shape), + exclusive_max=shape, + ) + + @property + def ndim(self) -> int: + return len(self.inclusive_min) + + @property + def origin(self) -> tuple[int, ...]: + return self.inclusive_min + + @property + def shape(self) -> tuple[int, ...]: + cached = self._shape + if cached is None: + cached = tuple( + hi - lo for lo, hi in zip(self.inclusive_min, self.exclusive_max, strict=True) + ) + object.__setattr__(self, "_shape", cached) + return cached + + def contains(self, index: tuple[int, ...]) -> bool: + if len(index) != self.ndim: + return False + return all( + lo <= idx < hi + for lo, hi, idx in zip(self.inclusive_min, self.exclusive_max, index, strict=True) + ) + + def contains_domain(self, other: IndexDomain) -> bool: + if other.ndim != self.ndim: + return False + return all( + self_lo <= other_lo and other_hi <= self_hi + for self_lo, self_hi, other_lo, other_hi in zip( + self.inclusive_min, + self.exclusive_max, + other.inclusive_min, + other.exclusive_max, + strict=True, + ) + ) + + def intersect(self, other: IndexDomain) -> IndexDomain | None: + if other.ndim != self.ndim: + raise ValueError( + f"Cannot intersect domains with different ranks: {self.ndim} vs {other.ndim}" + ) + new_min = tuple( + max(a, b) for a, b in zip(self.inclusive_min, other.inclusive_min, strict=True) + ) + new_max = tuple( + min(a, b) for a, b in zip(self.exclusive_max, other.exclusive_max, strict=True) + ) + if any(lo >= hi for lo, hi in zip(new_min, new_max, strict=True)): + return None + return IndexDomain(inclusive_min=new_min, exclusive_max=new_max) + + def translate(self, offset: tuple[int, ...]) -> IndexDomain: + if len(offset) != self.ndim: + raise ValueError( + f"Offset must have same length as domain dimensions. " + f"Domain has {self.ndim} dimensions, offset has {len(offset)}." + ) + new_min = tuple(lo + off for lo, off in zip(self.inclusive_min, offset, strict=True)) + new_max = tuple(hi + off for hi, off in zip(self.exclusive_max, offset, strict=True)) + return IndexDomain(inclusive_min=new_min, exclusive_max=new_max) + + def narrow(self, selection: Any) -> IndexDomain: + """Apply a basic selection and return a narrowed domain. + Indices are absolute coordinates. Integer indices produce length-1 extent. + Strided slices are not supported — use IndexTransform for strides. + """ + normalized = _normalize_selection(selection, self.ndim) + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + for dim_idx, (sel, dim_lo, dim_hi) in enumerate( + zip(normalized, self.inclusive_min, self.exclusive_max, strict=True) + ): + if isinstance(sel, int): + if sel < dim_lo or sel >= dim_hi: + raise IndexError( + f"index {sel} is out of bounds for dimension {dim_idx} " + f"with domain [{dim_lo}, {dim_hi})" + ) + new_inclusive_min.append(sel) + new_exclusive_max.append(sel + 1) + else: + start, stop, step = sel.start, sel.stop, sel.step + if step is not None and step != 1: + raise IndexError( + "IndexDomain.narrow only supports step=1 slices. " + f"Got step={step}. Use IndexTransform for strided access." + ) + abs_start = dim_lo if start is None else start + abs_stop = dim_hi if stop is None else stop + abs_start = max(abs_start, dim_lo) + abs_stop = min(abs_stop, dim_hi) + abs_stop = max(abs_stop, abs_start) + new_inclusive_min.append(abs_start) + new_exclusive_max.append(abs_stop) + return IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + +def _normalize_selection(selection: Any, ndim: int) -> tuple[int | slice, ...]: + """Normalize a basic selection to a tuple of ints/slices with length ndim.""" + if not isinstance(selection, tuple): + selection = (selection,) + result: list[int | slice] = [] + ellipsis_seen = False + for sel in selection: + if sel is Ellipsis: + if ellipsis_seen: + raise IndexError("an index can only have a single ellipsis ('...')") + ellipsis_seen = True + num_missing = ndim - (len(selection) - 1) + result.extend([slice(None)] * num_missing) + else: + result.append(sel) + while len(result) < ndim: + result.append(slice(None)) + if len(result) > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, " + f"but {len(result)} were indexed" + ) + return tuple(result) diff --git a/packages/zarr-indexing/src/zarr_indexing/errors.py b/packages/zarr-indexing/src/zarr_indexing/errors.py new file mode 100644 index 0000000000..fa2f6fc5d3 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/errors.py @@ -0,0 +1,21 @@ +"""Canonical index-error types raised by the transform algebra. + +These are the authoritative class definitions. `zarr.errors` re-exports the +same objects (`from zarr_indexing.errors import ...`) so that, e.g., +`zarr.errors.BoundsCheckError is zarr_indexing.errors.BoundsCheckError`. +Both subclass the built-in `IndexError`, so existing `except IndexError` (or +`except zarr.errors.BoundsCheckError`) catch sites keep working unchanged. +""" + +from __future__ import annotations + +__all__ = [ + "BoundsCheckError", + "VindexInvalidSelectionError", +] + + +class VindexInvalidSelectionError(IndexError): ... + + +class BoundsCheckError(IndexError): ... diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py new file mode 100644 index 0000000000..6db1025720 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -0,0 +1,39 @@ +"""Structural typing for the chunk-grid surface used by chunk resolution. + +`chunk_resolution` needs only a narrow slice of a chunk grid: the per-dimension +mapping between storage indices and chunk coordinates. Rather than import +zarr's concrete `ChunkGrid`, we type against these Protocols. zarr's +`ChunkGrid` / `DimensionGrid` satisfy them structurally, so no zarr import is +needed here. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Protocol + +if TYPE_CHECKING: + from collections.abc import Sequence + + import numpy as np + import numpy.typing as npt + + +class DimensionGridLike(Protocol): + """The per-dimension chunk-mapping surface consumed by chunk resolution.""" + + def index_to_chunk(self, idx: int) -> int: ... + def chunk_offset(self, chunk_ix: int) -> int: ... + def chunk_size(self, chunk_ix: int) -> int: ... + def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: ... + + +class ChunkGridLike(Protocol): + """A chunk grid exposing its per-dimension grids via `_dimensions`. + + Typed as a read-only, covariant `Sequence` so a concrete grid whose + `_dimensions` is a `tuple` of a more specific dimension type structurally + satisfies this Protocol. + """ + + @property + def _dimensions(self) -> Sequence[DimensionGridLike]: ... diff --git a/packages/zarr-indexing/src/zarr_indexing/json.py b/packages/zarr-indexing/src/zarr_indexing/json.py new file mode 100644 index 0000000000..c95696f309 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/json.py @@ -0,0 +1,325 @@ +"""Lowering between canonical ndsel bodies and in-memory `IndexTransform`s. + +This is the **engine layer**. Where `messages.py` is pure JSON→JSON and imposes +no array constraints, this module converts a *canonical* ndsel transform body +(spec section 4.3, as produced by `zarr_indexing.messages.normalize_ndsel`) +into the numpy-backed `IndexTransform` the chunk engine runs on, and back. + +Two engine constraints live **here and only here**: + +- **Finite bounds.** An `IndexDomain` addresses a finite array, so a canonical + body carrying a `"-inf"`/`"+inf"` bound cannot be lowered; `from_json` raises. +- **Implicit bounds lower by value.** The `[n]`-bracket implicit/explicit flag + is a message-layer concern; the engine keeps only the integer value. + +## The `index_array` wire format (and the degenerate-collapse it documents) + +ndsel and TensorStore both **reject** an output map that carries *both* +`input_dimension` and `index_array`. The in-memory `ArrayMap`, however, records +an `input_dimension` to pin the axis an orthogonal (`oindex`) array varies over. +This module bridges the gap: + +- **On serialize** (`transform_to_canonical`): + 1. An all-singleton `index_array` (size 1) selects a single coordinate + regardless of input, so it is **collapsed to a `constant` map** + `{offset: offset + stride*value}`. The size-1 input dimension stays in the + domain, unconsumed — a valid transform. This makes a length-1 `oindex` + selection round-trip *behaviorally* (an `ArrayMap` becomes a `ConstantMap`) + rather than by object identity. + 2. Non-degenerate `index_array` maps are emitted **without** `input_dimension`. + +- **On load** (`transform_from_canonical`): the in-memory `input_dimension` is + reconstructed from the full-rank array's dependency axes (its non-singleton + axes, see `transform._array_map_dependency_axes`). An array that solely owns a + single non-singleton axis is orthogonal (`input_dimension = that axis`); arrays + that share non-singleton axes, or vary over several, are correlated (`vindex`, + `input_dimension = None`). A single 1-D array over a rank-1 domain is + inherently ambiguous between the two flavours; it reconstructs as orthogonal, + which is behaviorally identical for the single-array case. + +`index_transform_to_json` / `index_transform_from_json` (and the `*_domain_*` +variants) are these canonical converters under their historical names. +""" + +from __future__ import annotations + +from collections import Counter +from typing import Any, Required, TypedDict + +import numpy as np + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.messages import normalize_ndsel +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr_indexing.transform import ( + IndexTransform, + _array_map_dependency_axes, # pyright: ignore[reportPrivateUsage] +) + +# `_array_map_dependency_axes` is a leading-underscore helper in `transform.py`, +# but it is deliberately shared with this module (the engine-level JSON <-> +# `IndexTransform` lowering below needs the same dependency-axis logic that +# `transform.py`'s own array-reindexing helpers use). It is not part of the +# package's public API; pyright's `reportPrivateUsage` flags the cross-module +# import anyway. See `chunk_resolution.py`'s `_dimensions` suppression for the +# analogous rationale — whether to promote either symbol out of "private" is +# an open pre-publish API decision, not resolved here. + +# --------------------------------------------------------------------------- +# TypedDict definitions (canonical JSON shapes) +# --------------------------------------------------------------------------- + +# An `index_array` serializes via `ndarray.tolist()`, so it is a nested list of +# ints whose nesting depth equals the array rank. +NestedIntList = list[Any] + +# A canonical *lowered* body carries only finite integer bounds, but the JSON +# shape admits the full ndsel `bound` grammar: an explicit int / sentinel, or a +# one-element implicit `[value]` array. +IndexValueJSON = int | str +BoundJSON = int | str | list[IndexValueJSON] + + +class IndexDomainJSON(TypedDict, total=False): + """Canonical JSON representation of an IndexDomain.""" + + input_inclusive_min: Required[list[BoundJSON]] + input_exclusive_max: Required[list[BoundJSON]] + input_labels: Required[list[str]] + + +class OutputIndexMapJSON(TypedDict, total=False): + """Canonical JSON representation of a single output index map. + + Exactly one of three forms (distinguished by which fields are present): + + - `{"offset": 5}` — constant + - `{"offset": 0, "stride": 1, "input_dimension": 0}` — single_input_dimension + - `{"offset": 0, "stride": 1, "index_array": [...], + "index_array_bounds": ["-inf", "+inf"]}` — index_array + """ + + offset: int + stride: int + input_dimension: int + index_array: NestedIntList + index_array_bounds: list[IndexValueJSON] + + +class IndexTransformJSON(TypedDict, total=False): + """Canonical JSON representation of an IndexTransform (spec section 4.3).""" + + input_rank: Required[int] + input_inclusive_min: Required[list[BoundJSON]] + input_exclusive_max: Required[list[BoundJSON]] + input_labels: Required[list[str]] + output: Required[list[OutputIndexMapJSON]] + + +# --------------------------------------------------------------------------- +# Bound / label lowering (engine constraints) +# --------------------------------------------------------------------------- + + +def _lower_bound(bound: BoundJSON, where: str) -> int: + """Lower a canonical bound to a finite integer, rejecting infinities.""" + value = bound[0] if isinstance(bound, list) else bound + if value == "-inf" or value == "+inf": + raise ValueError( + f"{where} is infinite ({value!r}); an IndexDomain addresses a finite " + f"array and cannot lower an infinite bound" + ) + return int(value) + + +def _lower_labels(labels: list[str]) -> tuple[str, ...] | None: + """All-empty labels collapse to `None` so a label-free domain round-trips.""" + return None if all(label == "" for label in labels) else tuple(labels) + + +def _emit_labels(labels: tuple[str, ...] | None, rank: int) -> list[str]: + """Emit canonical labels: `[""]*rank` when the domain is unlabeled.""" + return [""] * rank if labels is None else list(labels) + + +# --------------------------------------------------------------------------- +# IndexDomain serialization +# --------------------------------------------------------------------------- + + +def index_domain_to_json(domain: IndexDomain) -> IndexDomainJSON: + """Convert an IndexDomain to its canonical JSON representation.""" + return { + "input_inclusive_min": list(domain.inclusive_min), + "input_exclusive_max": list(domain.exclusive_max), + "input_labels": _emit_labels(domain.labels, domain.ndim), + } + + +def index_domain_from_json(data: IndexDomainJSON) -> IndexDomain: + """Construct an IndexDomain from its canonical JSON representation.""" + inclusive_min = tuple( + _lower_bound(b, f"input_inclusive_min[{i}]") + for i, b in enumerate(data["input_inclusive_min"]) + ) + exclusive_max = tuple( + _lower_bound(b, f"input_exclusive_max[{i}]") + for i, b in enumerate(data["input_exclusive_max"]) + ) + labels = _lower_labels(list(data["input_labels"])) + return IndexDomain(inclusive_min=inclusive_min, exclusive_max=exclusive_max, labels=labels) + + +# --------------------------------------------------------------------------- +# OutputIndexMap serialization +# --------------------------------------------------------------------------- + + +def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON: + """Convert an output index map to its canonical JSON representation. + + A degenerate all-singleton `ArrayMap` collapses to a `constant` map; a + non-degenerate one is emitted without `input_dimension` (see the module + docstring on the wire format). + """ + if isinstance(m, ConstantMap): + return {"offset": m.offset} + + if isinstance(m, DimensionMap): + return {"offset": m.offset, "stride": m.stride, "input_dimension": m.input_dimension} + + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + if m.index_array.size == 1: + value = int(m.index_array.reshape(-1)[0]) + return {"offset": m.offset + m.stride * value} + return { + "offset": m.offset, + "stride": m.stride, + "index_array": m.index_array.tolist(), + "index_array_bounds": ["-inf", "+inf"], + } + + +def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: + """Construct an output index map from its canonical JSON representation. + + An `index_array` map's `input_dimension` is reconstructed from the array's + dependency axes in isolation (single non-singleton axis → orthogonal). The + transform-level loader classifies globally; use it when several maps may + share axes. + """ + if "index_array" in data: + arr = np.asarray(data["index_array"], dtype=np.intp) + return ArrayMap( + index_array=arr, + offset=data.get("offset", 0), + stride=data.get("stride", 1), + input_dimension=_solo_dependency_axis(arr), + ) + + if "input_dimension" in data: + return DimensionMap( + input_dimension=data["input_dimension"], + offset=data.get("offset", 0), + stride=data.get("stride", 1), + ) + + return ConstantMap(offset=data.get("offset", 0)) + + +def _solo_dependency_axis(arr: np.ndarray[Any, Any]) -> int | None: + """The single axis a lone `index_array` varies over, or `None` if not exactly one.""" + dep = _array_map_dependency_axes(arr) + return dep[0] if len(dep) == 1 else None + + +# --------------------------------------------------------------------------- +# IndexTransform serialization +# --------------------------------------------------------------------------- + + +def transform_to_canonical(transform: IndexTransform) -> IndexTransformJSON: + """Convert an IndexTransform to its canonical ndsel transform body. + + The result is fully explicit (spec section 4.3): `input_rank`, fully written + bounds and labels, and an explicit `output` with `offset`/`stride` present + on every affine and array map. + """ + return { + "input_rank": transform.domain.ndim, + "input_inclusive_min": list(transform.domain.inclusive_min), + "input_exclusive_max": list(transform.domain.exclusive_max), + "input_labels": _emit_labels(transform.domain.labels, transform.domain.ndim), + "output": [output_index_map_to_json(m) for m in transform.output], + } + + +def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform: + """Construct an IndexTransform from a canonical (or canonicalizable) body. + + The body is first run through the message layer (`normalize_ndsel`) so that + omitted fields — identity `output`, default bounds/labels — are filled and + validated, then lowered to the engine representation. `index_array` maps' + `input_dimension` values are reconstructed by global dependency-axis + ownership (see the module docstring). + """ + body = normalize_ndsel({"kind": "transform", **data}) + + inclusive_min = tuple( + _lower_bound(b, f"input_inclusive_min[{i}]") + for i, b in enumerate(body["input_inclusive_min"]) + ) + exclusive_max = tuple( + _lower_bound(b, f"input_exclusive_max[{i}]") + for i, b in enumerate(body["input_exclusive_max"]) + ) + domain = IndexDomain( + inclusive_min=inclusive_min, + exclusive_max=exclusive_max, + labels=_lower_labels(body["input_labels"]), + ) + + output_raw: list[dict[str, Any]] = body["output"] + + # Classify index_array maps globally: an axis owned by exactly one array map + # (and the map's sole non-singleton axis) marks that map orthogonal; shared + # or multiple non-singleton axes mark the maps correlated (vindex). + array_axes: dict[int, tuple[int, ...]] = {} + axis_owners: Counter[int] = Counter() + for i, om in enumerate(output_raw): + if "index_array" in om: + arr = np.asarray(om["index_array"], dtype=np.intp) + dep = _array_map_dependency_axes(arr) + array_axes[i] = dep + axis_owners.update(dep) + + output: list[OutputIndexMap] = [] + for i, om in enumerate(output_raw): + if "index_array" in om: + dep = array_axes[i] + input_dim = dep[0] if len(dep) == 1 and axis_owners[dep[0]] == 1 else None + output.append( + ArrayMap( + index_array=np.asarray(om["index_array"], dtype=np.intp), + offset=om.get("offset", 0), + stride=om.get("stride", 1), + input_dimension=input_dim, + ) + ) + elif "input_dimension" in om: + output.append( + DimensionMap( + input_dimension=om["input_dimension"], + offset=om.get("offset", 0), + stride=om.get("stride", 1), + ) + ) + else: + output.append(ConstantMap(offset=om.get("offset", 0))) + + return IndexTransform(domain=domain, output=tuple(output)) + + +# Historical names, now pointing at the canonical converters. +index_transform_to_json = transform_to_canonical +index_transform_from_json = transform_from_canonical diff --git a/packages/zarr-indexing/src/zarr_indexing/messages.py b/packages/zarr-indexing/src/zarr_indexing/messages.py new file mode 100644 index 0000000000..0ee6cb97ab --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/messages.py @@ -0,0 +1,657 @@ +"""The ndsel message layer — pure JSON in, canonical JSON out. + +This module implements the [ndsel](https://github.com/d-v-b/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`. + +Two entry points: + +- `parse_ndsel(obj)` — structurally validate an ndsel message of any of the + five kinds (`point`/`box`/`slice`/`points`/`transform`), returning it + unchanged. Raises `NdselError` (carrying a spec reason code) on any defect. +- `normalize_ndsel(obj)` — desugar and canonicalize a message to the single + deterministic **canonical transform body** of the spec (section 4.3): a bare + `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. + +Value rules enforced here: 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; an implicit bound is the one-element `[n]`-bracket form, and its +implicit/explicit flag is preserved through normalization. +""" + +from __future__ import annotations + +from typing import Any + +__all__ = [ + "NdselError", + "normalize_ndsel", + "parse_ndsel", +] + +# --------------------------------------------------------------------------- +# Error taxonomy +# --------------------------------------------------------------------------- + +#: The complete set of ndsel reason codes (spec section 6). +REASON_CODES = frozenset( + { + "invalid_json", + "unknown_kind", + "unknown_field", + "multiple_upper_bounds", + "bounds_out_of_order", + "output_map_conflict", + "rank_mismatch", + "step_zero", + "negative_step_unsupported", + } +) + + +class NdselError(ValueError): + """An ndsel message failed validation. + + Carries the spec `reason` code (one of `REASON_CODES`) so callers and the + conformance harness can assert on it directly, plus a human-readable + `detail`. + """ + + def __init__(self, reason: str, detail: str = "") -> None: + self.reason = reason + self.detail = detail + super().__init__(f"{reason}: {detail}" if detail else reason) + + +# --------------------------------------------------------------------------- +# 64-bit signed integer range (spec section 3.5) +# --------------------------------------------------------------------------- + +_I64_MIN = -(2**63) +_I64_MAX = 2**63 - 1 + +_KNOWN_KINDS = frozenset({"point", "box", "slice", "points", "transform"}) + +# The two upper-bound spellings, keyed by message prefix. Only one of the three +# per group may appear (spec section 4.1 / 5.2). +_BOX_UPPER = ("exclusive_max", "inclusive_max", "shape") +_TRANSFORM_UPPER = ("input_exclusive_max", "input_inclusive_max", "input_shape") + +_OUTPUT_MAP_FIELDS = frozenset( + {"offset", "stride", "input_dimension", "index_array", "index_array_bounds"} +) + + +# --------------------------------------------------------------------------- +# Leaf value validators +# --------------------------------------------------------------------------- + + +def _is_int(value: Any) -> bool: + """True iff `value` is a JSON integer — an `int` that is not a `bool`. + + JSON has no boolean-as-integer: `True`/`False` are rejected even though + Python makes `bool` a subclass of `int` (spec section 3.6). + """ + return isinstance(value, int) and not isinstance(value, bool) + + +def _check_int(value: Any, where: str) -> int: + """Validate a plain-integer position: an in-range i64, never a sentinel.""" + if not _is_int(value): + raise NdselError("invalid_json", f"{where} must be an integer, got {value!r}") + if value < _I64_MIN or value > _I64_MAX: + raise NdselError("invalid_json", f"{where} is outside the 64-bit signed range: {value}") + return int(value) + + +def _is_sentinel(value: Any) -> bool: + return value in ("-inf", "+inf") + + +def _check_index_value(value: Any, where: str) -> int | str: + """Validate an `index-value`: an in-range i64 or a `"-inf"`/`"+inf"` sentinel.""" + if _is_sentinel(value): + return str(value) + return _check_int(value, where) + + +def _check_bound(value: Any, where: str) -> int | str | list[int | str]: + """Validate a `bound`: an explicit `index-value`, or a one-element implicit `[index-value]`.""" + if isinstance(value, list): + if len(value) != 1: + raise NdselError( + "invalid_json", + f"{where} implicit bound must be a one-element array, got {value!r}", + ) + return [_check_index_value(value[0], where)] + return _check_index_value(value, where) + + +def _check_int_list(value: Any, where: str) -> list[int]: + if not isinstance(value, list): + raise NdselError("invalid_json", f"{where} must be an array, got {value!r}") + return [_check_int(v, f"{where}[{i}]") for i, v in enumerate(value)] + + +def _check_bound_list(value: Any, where: str) -> list[Any]: + if not isinstance(value, list): + raise NdselError("invalid_json", f"{where} must be an array, got {value!r}") + return [_check_bound(v, f"{where}[{i}]") for i, v in enumerate(value)] + + +def _check_label_list(value: Any, where: str) -> list[str]: + if not isinstance(value, list): + raise NdselError("invalid_json", f"{where} must be an array, got {value!r}") + for i, v in enumerate(value): + if not isinstance(v, str): + raise NdselError("invalid_json", f"{where}[{i}] must be a string, got {v!r}") + return list(value) + + +# --------------------------------------------------------------------------- +# Extended-integer order for bounds (spec section 4.1) +# --------------------------------------------------------------------------- + + +def _bound_value(bound: int | str | list[int | str]) -> int | str: + """The underlying `index-value` of a bound, dropping the implicit bracket.""" + return bound[0] if isinstance(bound, list) else bound + + +def _bound_is_implicit(bound: int | str | list[int | str]) -> bool: + return isinstance(bound, list) + + +def _ext_key(value: int | str) -> tuple[int, int]: + """A sort key giving the extended-integer order `-inf < n < +inf` exactly. + + Uses an integer tier plus the value, so no float rounding of near-`2**63` + integers can misorder the `inclusive_min <= exclusive_max` check. + """ + if value == "-inf": + return (0, 0) + if value == "+inf": + return (2, 0) + assert isinstance(value, int) + return (1, value) + + +def _rewrap(value: int | str, *, implicit: bool) -> int | str | list[int | str]: + return [value] if implicit else value + + +# --------------------------------------------------------------------------- +# Message-level helpers +# --------------------------------------------------------------------------- + + +def _require_object(obj: Any) -> dict[str, Any]: + if not isinstance(obj, dict): + raise NdselError("invalid_json", f"message must be a JSON object, got {type(obj).__name__}") + return obj + + +def _message_kind(obj: dict[str, Any]) -> str: + kind = obj.get("kind") + if not isinstance(kind, str): + raise NdselError("invalid_json", "message must have a string 'kind' field") + if kind not in _KNOWN_KINDS: + raise NdselError("unknown_kind", f"unknown kind {kind!r}") + return kind + + +def _check_membership(obj: dict[str, Any], allowed: frozenset[str], what: str) -> None: + """Strict membership (spec section 3.7): reject any undefined member.""" + for key in obj: + if key not in allowed: + raise NdselError("unknown_field", f"{what} has undefined member {key!r}") + + +def _single_upper_bound(obj: dict[str, Any], fields: tuple[str, str, str]) -> str | None: + present = [f for f in fields if f in obj] + if len(present) > 1: + raise NdselError( + "multiple_upper_bounds", + f"at most one of {fields} may be present; got {present}", + ) + return present[0] if present else None + + +def _resolve_upper_bound( + upper_field: str | None, + upper_raw: list[Any] | None, + inclusive_min: list[Any], + rank: int, + *, + kind_of: str, +) -> list[int | str | list[int | str]]: + """Produce `exclusive_max` from whichever upper-bound spelling was supplied. + + - `exclusive_max`/`input_exclusive_max` → used directly. + - `inclusive_max`/`input_inclusive_max` → each element `+1`. + - `shape`/`input_shape` → `inclusive_min + shape` per element. + - none → an **implicit `+inf`** in every dimension. + + The implicit/explicit bracket travels with the extent-bearing field (the + upper bound, or `shape`), matching the spec's `[n]`-bracket convention. + """ + if upper_field is None: + return [["+inf"] for _ in range(rank)] + + assert upper_raw is not None + if kind_of == "exclusive": + return list(upper_raw) + + result: list[int | str | list[int | str]] = [] + for k in range(rank): + raw = upper_raw[k] + implicit = _bound_is_implicit(raw) + value = _bound_value(raw) + if kind_of == "inclusive": + new = _inclusive_to_exclusive(value) + else: # shape + new = _shape_to_exclusive(_bound_value(inclusive_min[k]), value) + result.append(_rewrap(new, implicit=implicit)) + return result + + +def _inclusive_to_exclusive(value: int | str) -> int | str: + if value == "+inf" or value == "-inf": + return value + assert isinstance(value, int) + return value + 1 + + +def _shape_to_exclusive(min_value: int | str, shape_value: int | str) -> int | str: + if shape_value == "+inf" or min_value == "+inf": + return "+inf" + if min_value == "-inf": + return "-inf" + assert isinstance(min_value, int) + assert isinstance(shape_value, int) + return min_value + shape_value + + +def _validate_domain(inclusive_min: list[Any], exclusive_max: list[Any], *, prefix: str) -> None: + """Every dimension must satisfy `inclusive_min <= exclusive_max` (empty is valid).""" + for k, (lo, hi) in enumerate(zip(inclusive_min, exclusive_max, strict=True)): + if _ext_key(_bound_value(lo)) > _ext_key(_bound_value(hi)): + raise NdselError( + "bounds_out_of_order", + f"{prefix}[{k}]: inclusive_min {_bound_value(lo)!r} > " + f"exclusive_max {_bound_value(hi)!r}", + ) + + +def _identity_output(rank: int) -> list[dict[str, Any]]: + return [{"offset": 0, "stride": 1, "input_dimension": k} for k in range(rank)] + + +# --------------------------------------------------------------------------- +# Per-kind desugaring +# --------------------------------------------------------------------------- + + +def _normalize_point(obj: dict[str, Any]) -> dict[str, Any]: + _check_membership(obj, frozenset({"kind", "coords"}), "point") + if "coords" not in obj: + raise NdselError("invalid_json", "point requires 'coords'") + coords = _check_int_list(obj["coords"], "coords") + return { + "input_rank": 0, + "input_inclusive_min": [], + "input_exclusive_max": [], + "input_labels": [], + "output": [{"offset": c} for c in coords], + } + + +def _infer_rank( + obj: dict[str, Any], + named_lengths: list[tuple[str, int]], + *, + declared: int | None, +) -> int: + """Reconcile a declared rank (if any) with every present array's length.""" + rank = declared + for name, length in named_lengths: + if rank is None: + rank = length + elif rank != length: + raise NdselError( + "rank_mismatch", + f"{name} has length {length}, inconsistent with rank {rank}", + ) + return rank if rank is not None else 0 + + +def _normalize_box(obj: dict[str, Any]) -> dict[str, Any]: + allowed = frozenset( + {"kind", "inclusive_min", "exclusive_max", "inclusive_max", "shape", "labels"} + ) + _check_membership(obj, allowed, "box") + + inclusive_min_raw = ( + _check_bound_list(obj["inclusive_min"], "inclusive_min") if "inclusive_min" in obj else None + ) + upper_field = _single_upper_bound(obj, _BOX_UPPER) + upper_raw = _check_bound_list(obj[upper_field], upper_field) if upper_field else None + labels_raw = _check_label_list(obj["labels"], "labels") if "labels" in obj else None + + named_lengths: list[tuple[str, int]] = [] + if inclusive_min_raw is not None: + named_lengths.append(("inclusive_min", len(inclusive_min_raw))) + if upper_raw is not None: + named_lengths.append((upper_field or "", len(upper_raw))) + if labels_raw is not None: + named_lengths.append(("labels", len(labels_raw))) + rank = _infer_rank(obj, named_lengths, declared=None) + + inclusive_min = inclusive_min_raw if inclusive_min_raw is not None else [0] * rank + exclusive_max = _resolve_upper_bound( + upper_field, upper_raw, inclusive_min, rank, kind_of=_upper_kind(upper_field, _BOX_UPPER) + ) + labels = labels_raw if labels_raw is not None else [""] * rank + _validate_domain(inclusive_min, exclusive_max, prefix="box") + + return { + "input_rank": rank, + "input_inclusive_min": inclusive_min, + "input_exclusive_max": exclusive_max, + "input_labels": labels, + "output": _identity_output(rank), + } + + +def _upper_kind(upper_field: str | None, fields: tuple[str, str, str]) -> str: + if upper_field is None or upper_field == fields[0]: + return "exclusive" + if upper_field == fields[1]: + return "inclusive" + return "shape" + + +def _normalize_slice(obj: dict[str, Any]) -> dict[str, Any]: + allowed = frozenset({"kind", "start", "stop", "step", "labels"}) + _check_membership(obj, allowed, "slice") + if "start" not in obj: + raise NdselError("invalid_json", "slice requires 'start'") + if "stop" not in obj: + raise NdselError("invalid_json", "slice requires 'stop'") + start = _check_int_list(obj["start"], "start") + stop = _check_int_list(obj["stop"], "stop") + step = _check_int_list(obj["step"], "step") if "step" in obj else [1] * len(start) + labels_raw = _check_label_list(obj["labels"], "labels") if "labels" in obj else None + + n = len(start) + for name, arr in (("stop", stop), ("step", step)): + if len(arr) != n: + raise NdselError( + "rank_mismatch", f"{name} has length {len(arr)}, expected {n} (from start)" + ) + if labels_raw is not None and len(labels_raw) != n: + raise NdselError( + "rank_mismatch", f"labels has length {len(labels_raw)}, expected {n} (from start)" + ) + + for k, s in enumerate(step): + if s == 0: + raise NdselError("step_zero", f"step[{k}] is zero") + if s < 0: + raise NdselError("negative_step_unsupported", f"step[{k}] is negative ({s})") + + inclusive_min: list[Any] = [] + exclusive_max: list[Any] = [] + output: list[dict[str, Any]] = [] + for k in range(n): + a, b, s = start[k], stop[k], step[k] + m = max(0, -(-(b - a) // s)) # ceil((b - a) / s) + o = _trunc_div(a, s) # trunc(a / s), toward zero + offset = a - s * o # lattice phase, in (-s, s) + inclusive_min.append(o) + exclusive_max.append(o + m) + output.append({"offset": offset, "stride": s, "input_dimension": k}) + + labels = labels_raw if labels_raw is not None else [""] * n + return { + "input_rank": n, + "input_inclusive_min": inclusive_min, + "input_exclusive_max": exclusive_max, + "input_labels": labels, + "output": output, + } + + +def _normalize_points(obj: dict[str, Any]) -> dict[str, Any]: + _check_membership(obj, frozenset({"kind", "coords"}), "points") + if "coords" not in obj: + raise NdselError("invalid_json", "points requires 'coords'") + coords = obj["coords"] + if not isinstance(coords, list): + raise NdselError("invalid_json", f"points coords must be an array, got {coords!r}") + + rows: list[list[int]] = [] + n: int | None = None + for i, row in enumerate(coords): + if not isinstance(row, list): + raise NdselError("invalid_json", f"points coords[{i}] must be an array, got {row!r}") + row_ints = [_check_int(v, f"coords[{i}][{j}]") for j, v in enumerate(row)] + if n is None: + n = len(row_ints) + elif len(row_ints) != n: + raise NdselError( + "rank_mismatch", + f"points coords[{i}] has length {len(row_ints)}, expected {n} (ragged)", + ) + rows.append(row_ints) + + m = len(rows) + n = n if n is not None else 0 + output = [ + { + "offset": 0, + "stride": 1, + "index_array": [rows[i][k] for i in range(m)], + "index_array_bounds": ["-inf", "+inf"], + } + for k in range(n) + ] + return { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [m], + "input_labels": [""], + "output": output, + } + + +def _normalize_output_map(raw: Any, where: str) -> dict[str, Any]: + if not isinstance(raw, dict): + raise NdselError("invalid_json", f"{where} must be a JSON object, got {raw!r}") + _check_membership(raw, _OUTPUT_MAP_FIELDS, where) + + has_index_array = "index_array" in raw + has_input_dim = "input_dimension" in raw + if has_index_array and has_input_dim: + raise NdselError( + "output_map_conflict", + f"{where} carries both 'input_dimension' and 'index_array'", + ) + + offset = _check_int(raw["offset"], f"{where}.offset") if "offset" in raw else 0 + + if has_index_array: + stride = _check_int(raw["stride"], f"{where}.stride") if "stride" in raw else 1 + bounds = ( + _check_index_array_bounds(raw["index_array_bounds"], where) + if "index_array_bounds" in raw + else ["-inf", "+inf"] + ) + # index_array is carried verbatim (spec section 7 defers shape validation). + return { + "offset": offset, + "stride": stride, + "index_array": raw["index_array"], + "index_array_bounds": bounds, + } + + if has_input_dim: + input_dim = _check_int(raw["input_dimension"], f"{where}.input_dimension") + if input_dim < 0: + raise NdselError( + "invalid_json", f"{where}.input_dimension must be >= 0, got {input_dim}" + ) + stride = _check_int(raw["stride"], f"{where}.stride") if "stride" in raw else 1 + return {"offset": offset, "stride": stride, "input_dimension": input_dim} + + # Constant map: only offset survives. A stray `stride`/`index_array_bounds` + # is schema-valid (the output-map schema permits those members on any map), + # so it is silently dropped rather than rejected — a constant carries only + # `offset` in canonical form (spec section 4.3). + return {"offset": offset} + + +def _check_index_array_bounds(value: Any, where: str) -> list[int | str]: + if not isinstance(value, list) or len(value) != 2: + raise NdselError( + "invalid_json", + f"{where}.index_array_bounds must be a two-element array, got {value!r}", + ) + return [ + _check_index_value(value[0], f"{where}.index_array_bounds[0]"), + _check_index_value(value[1], f"{where}.index_array_bounds[1]"), + ] + + +def _normalize_transform(obj: dict[str, Any]) -> dict[str, Any]: + allowed = frozenset( + { + "kind", + "input_rank", + "input_inclusive_min", + "input_exclusive_max", + "input_inclusive_max", + "input_shape", + "input_labels", + "output", + } + ) + _check_membership(obj, allowed, "transform") + + declared_rank: int | None = None + if "input_rank" in obj: + declared_rank = _check_int(obj["input_rank"], "input_rank") + if declared_rank < 0: + raise NdselError("invalid_json", f"input_rank must be >= 0, got {declared_rank}") + + inclusive_min_raw = ( + _check_bound_list(obj["input_inclusive_min"], "input_inclusive_min") + if "input_inclusive_min" in obj + else None + ) + upper_field = _single_upper_bound(obj, _TRANSFORM_UPPER) + upper_raw = _check_bound_list(obj[upper_field], upper_field) if upper_field else None + labels_raw = ( + _check_label_list(obj["input_labels"], "input_labels") if "input_labels" in obj else None + ) + + named_lengths: list[tuple[str, int]] = [] + if inclusive_min_raw is not None: + named_lengths.append(("input_inclusive_min", len(inclusive_min_raw))) + if upper_raw is not None: + named_lengths.append((upper_field or "", len(upper_raw))) + if labels_raw is not None: + named_lengths.append(("input_labels", len(labels_raw))) + rank = _infer_rank(obj, named_lengths, declared=declared_rank) + + inclusive_min = inclusive_min_raw if inclusive_min_raw is not None else [0] * rank + exclusive_max = _resolve_upper_bound( + upper_field, + upper_raw, + inclusive_min, + rank, + kind_of=_upper_kind(upper_field, _TRANSFORM_UPPER), + ) + labels = labels_raw if labels_raw is not None else [""] * rank + _validate_domain(inclusive_min, exclusive_max, prefix="input") + + if "output" in obj: + if not isinstance(obj["output"], list): + raise NdselError("invalid_json", f"output must be an array, got {obj['output']!r}") + output = [_normalize_output_map(m, f"output[{i}]") for i, m in enumerate(obj["output"])] + else: + output = _identity_output(rank) + + return { + "input_rank": rank, + "input_inclusive_min": inclusive_min, + "input_exclusive_max": exclusive_max, + "input_labels": labels, + "output": output, + } + + +_NORMALIZERS = { + "point": _normalize_point, + "box": _normalize_box, + "slice": _normalize_slice, + "points": _normalize_points, + "transform": _normalize_transform, +} + + +# --------------------------------------------------------------------------- +# trunc division (spec section 5.3 correction, matches _trunc_div in transform.py) +# --------------------------------------------------------------------------- + + +def _trunc_div(a: int, b: int) -> int: + """Integer division rounded toward zero (C semantics).""" + q = a // b + if q < 0 and q * b != a: + q += 1 + return q + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def normalize_ndsel(obj: Any) -> dict[str, Any]: + """Desugar and canonicalize an ndsel message to its canonical transform body. + + Accepts any of the five message kinds and returns the bare canonical + `IndexTransform` body of spec section 4.3 — no `kind` field. Raises + `NdselError` (carrying a reason code) for any invalid input. + """ + message = _require_object(obj) + kind = _message_kind(message) + return _NORMALIZERS[kind](message) + + +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. + """ + message = _require_object(obj) + _message_kind(message) + # Validation and desugaring share one pass; run it and discard the body. + normalize_ndsel(message) + return message diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py new file mode 100644 index 0000000000..581229bd22 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -0,0 +1,105 @@ +"""Output index maps — three representations of a set of integer coordinates. + +An output index map describes, for one dimension of storage, which coordinates +an array access will touch. Conceptually it is a **set of integers**. Three +representations cover the cases that arise in practice: + +- `ConstantMap(offset=5)` — a singleton set: `{5}` +- `DimensionMap(input_dimension=0, offset=3, stride=2)` over input `[0, 5)` + — an arithmetic progression: `{3, 5, 7, 9, 11}` +- `ArrayMap(index_array=[1, 5, 9])` — an explicit enumeration: `{1, 5, 9}` + +Every output map supports two set-theoretic operations (defined on +`IndexTransform`, which provides the input domain context these maps lack): + +- **intersect** — restrict to coordinates within a range (e.g., a chunk). + `{3, 5, 7, 9, 11} ∩ [4, 8) = {5, 7}` +- **translate** — shift every coordinate by a constant (e.g., make chunk-local). + `{5, 7} - 4 = {1, 3}` + +These two operations are the foundation of chunk resolution: for each chunk, +intersect the map with the chunk's range, then translate to chunk-local +coordinates. + +The three types exist because they trade off generality for efficiency: + +- `ConstantMap`: O(1) storage, O(1) intersection +- `DimensionMap`: O(1) storage, O(1) intersection (analytical) +- `ArrayMap`: O(n) storage, O(n) intersection (must scan the array) + +Collapsing everything to `ArrayMap` would be correct but wasteful — a +billion-element slice would materialize a billion coordinates just to group +them by chunk, when `DimensionMap` does it with three integers. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import numpy as np + import numpy.typing as npt + + +@dataclass(frozen=True, slots=True) +class ConstantMap: + """A singleton set: one storage coordinate. + + Represents `{offset}`. Arises from integer indexing (e.g., `arr[5]` + fixes one dimension to coordinate 5). + """ + + offset: int = 0 + + +@dataclass(frozen=True, slots=True) +class DimensionMap: + """An arithmetic progression of storage coordinates. + + Represents `{offset + stride * i : i in input_range}`, where the input + range comes from the enclosing `IndexTransform`'s domain. Arises from + slice indexing (e.g., `arr[2:10:3]` gives offset=2, stride=3). + """ + + input_dimension: int + offset: int = 0 + stride: int = 1 + + +@dataclass(frozen=True, slots=True) +class ArrayMap: + """An explicit enumeration of storage coordinates. + + Represents `{offset + stride * index_array[i] : i in input_range}`. + Arises from fancy indexing (e.g., `arr[[1, 5, 9]]` 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 + fully on the axes it varies over and singleton (size 1) elsewhere. The + dependency axes are therefore derivable from the shape (see + `transform._array_map_dependency_axes`), which distinguishes the two flavours + of multi-array fancy indexing: + + - **orthogonal** (`oindex`): each array varies along a single, *distinct* + axis (all others singleton); the result is their outer product. + - **vectorized** (`vindex`): the arrays are correlated and share the same + non-singleton (broadcast) axes; the result is a pointwise scatter. + + `input_dimension` records the single axis an orthogonal array varies over + (`None` for vectorized), binding it the way `DimensionMap` is bound. It is + usually redundant with the shape-derived classifier, but stays authoritative + for the shapes the classifier cannot distinguish: a length-1 orthogonal + selection normalizes to an all-singleton array (no non-singleton axis), and + length-1 vectorized arrays are equally degenerate. `None` therefore marks a + map as correlated, and an integer pins the dependency axis of a degenerate + orthogonal map (see `transform._array_map_dependent_axis`). + """ + + index_array: npt.NDArray[np.intp] + offset: int = 0 + stride: int = 1 + input_dimension: int | None = None + + +OutputIndexMap = ConstantMap | DimensionMap | ArrayMap diff --git a/packages/zarr-indexing/src/zarr_indexing/py.typed b/packages/zarr-indexing/src/zarr_indexing/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py new file mode 100644 index 0000000000..e1a3898b1d --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -0,0 +1,1311 @@ +"""Index transforms — composable, lazy coordinate mappings. + +An `IndexTransform` pairs an **input domain** (the coordinates a user sees) +with a tuple of **output maps** (the storage coordinates those inputs map to). +One output map per storage dimension. See `output_map.py` for the three +output map types. + +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. + +- **intersect(output_domain)** — restrict to storage coordinates within a + region. This is chunk resolution: "which of my coordinates fall in this + chunk?" + +- **translate(shift)** — shift all output coordinates. This makes coordinates + chunk-local: "express my coordinates relative to the chunk origin." + +- **compose(outer, inner)** — chain two transforms. See `composition.py`. + +The transform is the atomic unit that connects user-facing indexing to +chunk-level I/O. Every `Array` holds a transform (identity by default). +`Array.lazy[...]` composes a new transform lazily. Reading resolves the +transform against the chunk grid via intersect + translate. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Literal, cast + +import numpy as np + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap + + +@dataclass(frozen=True, slots=True) +class IndexTransform: + """A composable mapping from input coordinates to storage coordinates. + + An `IndexTransform` has: + + - `domain`: an `IndexDomain` describing the valid input coordinates + (the user-facing shape, possibly with non-zero origin). + - `output`: a tuple of output maps (one per storage dimension), each + describing which storage coordinates the inputs touch. + + For a freshly opened array, the transform is the identity: input + coordinate `i` maps to storage coordinate `i`. Indexing operations + compose new transforms without I/O. + """ + + domain: IndexDomain + output: tuple[OutputIndexMap, ...] + + def __post_init__(self) -> None: + for i, m in enumerate(self.output): + if isinstance(m, DimensionMap): + if m.input_dimension < 0 or m.input_dimension >= self.domain.ndim: + raise ValueError( + f"output[{i}].input_dimension = {m.input_dimension} " + f"is out of range for input rank {self.domain.ndim}" + ) + elif isinstance(m, ArrayMap) and m.index_array.ndim > self.domain.ndim: + # ArrayMap index arrays produced by indexing and chunk resolution + # are normalized to the full input rank (an axis the array varies + # over is full-sized, every other axis a singleton). A rank + # *exceeding* the domain is always a bug. A rank *below* it is + # tolerated: TensorStore-format JSON (external input) may supply a + # lower-rank index array that broadcasts against the input domain, + # and `_array_map_dependency_axes` treats any missing leading axes + # as singleton dependencies. + raise ValueError( + f"output[{i}].index_array has {m.index_array.ndim} dims " + f"but input domain has {self.domain.ndim} dims" + ) + + @property + def input_rank(self) -> int: + return self.domain.ndim + + @property + def output_rank(self) -> int: + return len(self.output) + + @classmethod + def identity(cls, domain: IndexDomain) -> IndexTransform: + output = tuple(DimensionMap(input_dimension=i) for i in range(domain.ndim)) + return cls(domain=domain, output=output) + + @classmethod + def from_shape(cls, shape: tuple[int, ...]) -> IndexTransform: + return cls.identity(IndexDomain.from_shape(shape)) + + @property + def selection_repr(self) -> str: + """Compact domain string, e.g. `'{ [2, 8), [0, 10) }'`. + + Follows TensorStore's IndexDomain notation: each dimension shown + as `[inclusive_min, exclusive_max)` with stride annotation if not 1. + Constant (integer-indexed) dimensions show as a single value. + Array-indexed dimensions show the set of selected coordinates. + """ + parts: list[str] = [] + for m in self.output: + if isinstance(m, ConstantMap): + parts.append(str(m.offset)) + elif isinstance(m, DimensionMap): + d = m.input_dimension + lo = self.domain.inclusive_min[d] + hi = self.domain.exclusive_max[d] + start = m.offset + m.stride * lo + stop = m.offset + m.stride * hi + if m.stride == 1: + parts.append(f"[{start}, {stop})") + else: + parts.append(f"[{start}, {stop}) step {m.stride}") + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + storage = m.offset + m.stride * m.index_array + n = int(storage.size) # .size, not len(): index_array may be 0-d + if n <= 5: + vals = ", ".join(str(int(v)) for v in storage.ravel()) + parts.append("{" + vals + "}") + else: + parts.append("{" + f"array({n})" + "}") + return "{ " + ", ".join(parts) + " }" + + def __repr__(self) -> str: + maps: list[str] = [] + for i, m in enumerate(self.output): + if isinstance(m, ConstantMap): + maps.append(f"out[{i}] = {m.offset}") + elif isinstance(m, DimensionMap): + maps.append(f"out[{i}] = {m.offset} + {m.stride} * in[{m.input_dimension}]") + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + maps.append(f"out[{i}] = {m.offset} + {m.stride} * arr{m.index_array.shape}[in]") + maps_str = ", ".join(maps) + return f"IndexTransform(domain={self.domain}, {maps_str})" + + def intersect( + self, output_domain: IndexDomain + ) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] + | np.ndarray[Any, np.dtype[np.intp]] + | None, + ] + | None + ): + """Restrict this transform to storage coordinates within output_domain. + + 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). + """ + return _intersect(self, output_domain) + + def translate(self, shift: tuple[int, ...]) -> IndexTransform: + """Shift all output coordinates by `shift`.""" + if len(shift) != self.output_rank: + raise ValueError(f"shift must have length {self.output_rank}, got {len(shift)}") + new_output: list[OutputIndexMap] = [] + for m, s in zip(self.output, shift, strict=True): + if isinstance(m, ConstantMap): + new_output.append(ConstantMap(offset=m.offset + s)) + elif isinstance(m, DimensionMap): + new_output.append( + DimensionMap( + input_dimension=m.input_dimension, + offset=m.offset + s, + stride=m.stride, + ) + ) + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + new_output.append( + ArrayMap( + index_array=m.index_array, + offset=m.offset + s, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) + return IndexTransform(domain=self.domain, output=tuple(new_output)) + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_basic_indexing(self, selection) + + def translate_domain_by(self, shift: tuple[int, ...]) -> IndexTransform: + """Shift the *input* domain by `shift`, preserving which cells are addressed. + + TensorStore's `translate_by`: the domain moves, and every output map is + re-offset so that new coordinate `c` addresses the cell that `c - shift` + addressed before. ArrayMaps are indexed positionally over the domain, so + their index arrays are unchanged. + """ + if len(shift) != self.input_rank: + raise ValueError(f"shift must have length {self.input_rank}, got {len(shift)}") + new_domain = self.domain.translate(shift) + new_output: list[OutputIndexMap] = [] + for m in self.output: + if isinstance(m, DimensionMap): + s = shift[m.input_dimension] + new_output.append( + DimensionMap( + input_dimension=m.input_dimension, + offset=m.offset - m.stride * s, + stride=m.stride, + ) + ) + else: + # ConstantMap: no input dependence. ArrayMap: positional over + # the domain, invariant under domain translation. + new_output.append(m) + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + def translate_domain_to(self, origins: tuple[int, ...]) -> IndexTransform: + """Move the input domain so its per-dimension origins equal `origins`. + + TensorStore's `translate_to`; `translate_domain_to((0,) * rank)` + re-zeros a view's coordinate system without changing which cells it + addresses. + """ + if len(origins) != self.input_rank: + raise ValueError(f"origins must have length {self.input_rank}, got {len(origins)}") + shift = tuple(o - m for o, m in zip(origins, self.domain.inclusive_min, strict=True)) + return self.translate_domain_by(shift) + + @property + def oindex(self) -> _OIndexHelper: + return _OIndexHelper(self) + + @property + def vindex(self) -> _VIndexHelper: + return _VIndexHelper(self) + + +def _intersect( + transform: IndexTransform, output_domain: IndexDomain +) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None, + ] + | None +): + """Intersect a transform with an output domain (e.g., a chunk's bounds). + + For each output dimension, restrict to storage coordinates within + `[output_domain.inclusive_min[d], output_domain.exclusive_max[d])`. + + Two flavours of fancy indexing require different treatment, distinguished by + the ArrayMaps' dependency axes (see `_array_map_dependency_axes`): + + - **orthogonal** (`oindex`): each ArrayMap varies over a single, distinct + input axis, forming an outer product. Every output dimension is intersected + independently and the input domain narrowed per axis. + - **correlated** (`vindex`): the ArrayMaps share their (broadcast) dependency + axes and scatter through a single flat index. A point survives only if ALL + its storage coordinates fall within the output domain; residual slice + dimensions are intersected independently, as in the orthogonal case. + + A `None` `input_dimension` marks a correlated map, so any such map routes the + whole transform through the correlated intersection. + + Returns `None` if the intersection is empty. + """ + if output_domain.ndim != transform.output_rank: + raise ValueError( + f"output_domain rank ({output_domain.ndim}) != " + f"transform output rank ({transform.output_rank})" + ) + + correlated_dims = [ + i + for i, m in enumerate(transform.output) + if isinstance(m, ArrayMap) and m.input_dimension is None + ] + if len(correlated_dims) > 0: + return _intersect_correlated(transform, output_domain, correlated_dims) + return _intersect_orthogonal(transform, output_domain) + + +def _intersect_dimension_map( + m: DimensionMap, input_lo: int, input_hi: int, lo: int, hi: int +) -> tuple[int, int] | None: + """Narrow a DimensionMap's input range to storage coordinates in `[lo, hi)`. + + `input_lo`/`input_hi` are the current (possibly already narrowed) input + range for the map's axis. Returns the new `(input_lo, input_hi)` or `None` + if no input produces an in-bounds storage coordinate. + """ + if input_lo >= input_hi: + return None + if m.stride > 0: + new_input_lo = max(input_lo, math.ceil((lo - m.offset) / m.stride)) + new_input_hi = min(input_hi, math.ceil((hi - m.offset) / m.stride)) + elif m.stride < 0: + new_input_lo = max(input_lo, math.ceil((hi - 1 - m.offset) / m.stride)) + new_input_hi = min(input_hi, math.ceil((lo - 1 - m.offset) / m.stride)) + else: + if lo <= m.offset < hi: + new_input_lo, new_input_hi = input_lo, input_hi + else: + return None + if new_input_lo >= new_input_hi: + return None + return new_input_lo, new_input_hi + + +def _intersect_orthogonal( + transform: IndexTransform, output_domain: IndexDomain +) -> ( + tuple[ + IndexTransform, + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None, + ] + | None +): + """Intersect a transform with no correlated ArrayMaps. + + Every output dimension is intersected independently. Multiple ArrayMaps bound + to distinct input dimensions form an outer product, so each array's surviving + *output* positions are tracked separately. + """ + new_min = list(transform.domain.inclusive_min) + new_max = list(transform.domain.exclusive_max) + new_output: list[OutputIndexMap] = [] + out_positions: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + + for out_dim, m in enumerate(transform.output): + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + + if isinstance(m, ConstantMap): + if lo <= m.offset < hi: + new_output.append(m) + else: + return None + + elif isinstance(m, DimensionMap): + d = m.input_dimension + narrowed = _intersect_dimension_map(m, new_min[d], new_max[d], lo, hi) + if narrowed is None: + return None + new_min[d], new_max[d] = narrowed + new_output.append(m) + + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + # Orthogonal: the array varies over a single axis (its dependency + # axis, or `input_dimension` for a degenerate length-1 array). Filter + # along that axis and keep the array at full input rank so the + # singleton axes it broadcasts over are preserved. + d = _array_map_dependent_axis(m) + storage = m.offset + m.stride * m.index_array + mask = (storage >= lo) & (storage < hi) + # The array is singleton on every axis but `d`, so its mask reduces + # to a 1-D vector along `d`. + survivors = np.nonzero(mask.reshape(-1))[0].astype(np.intp) + if survivors.size == 0: + return None + filtered = np.take(m.index_array, survivors, axis=d) + new_output.append( + ArrayMap( + index_array=np.asarray(filtered, dtype=np.intp), + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) + new_max[d] = new_min[d] + int(survivors.size) + out_positions[out_dim] = survivors + + new_domain = IndexDomain( + inclusive_min=tuple(new_min), + exclusive_max=tuple(new_max), + ) + result = IndexTransform(domain=new_domain, output=tuple(new_output)) + + # Hand back the surviving output positions in the shape the bridge expects: + # None (no arrays), a single vector (one array), or a per-output-dim dict + # (>= 2 orthogonal arrays → outer product). + out_indices: ( + dict[int, np.ndarray[Any, np.dtype[np.intp]]] | np.ndarray[Any, np.dtype[np.intp]] | None + ) + if len(out_positions) == 0: + out_indices = None + elif len(out_positions) == 1: + out_indices = next(iter(out_positions.values())) + else: + out_indices = out_positions + return (result, out_indices) + + +def _intersect_correlated( + transform: IndexTransform, + output_domain: IndexDomain, + correlated_dims: list[int], +) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]]] | None: + """Intersect a correlated (vindex) transform with an output domain. + + The correlated ArrayMaps share their broadcast (dependency) axes; a broadcast + point survives only if ALL its storage coordinates fall within the output + domain. Residual DimensionMap dimensions are intersected independently (as in + the orthogonal case) and preserved, so a partial vindex — e.g. two coordinate + arrays over a 3-D array, leaving one slice dimension — resolves correctly. + + 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)`. + """ + corr_maps = [cast("ArrayMap", transform.output[i]) for i in correlated_dims] + + # Mixing correlated and orthogonal ArrayMaps in one transform is not produced + # by any single selection and is not supported here. + orthogonal_array_dims = [ + i + for i, m in enumerate(transform.output) + if isinstance(m, ArrayMap) and m.input_dimension is not None + ] + if len(orthogonal_array_dims) > 0: + raise NotImplementedError( + "intersecting a transform with both correlated and orthogonal " + "ArrayMaps is not supported" + ) + + # The broadcast (dependency) axes are shared by every correlated map; they are + # the leading axes of the domain, followed by the residual slice axes. + broadcast_axes = _array_map_dependency_axes(corr_maps[0].index_array) + broadcast_shape = tuple(corr_maps[0].index_array.shape[a] for a in broadcast_axes) + + # Joint bounds mask over the broadcast block. + combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None + for out_dim in correlated_dims: + cm = cast("ArrayMap", transform.output[out_dim]) + storage = cm.offset + cm.stride * cm.index_array + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + mask = (storage >= lo) & (storage < hi) + combined = mask if combined is None else (combined & mask) + assert combined is not None + # The correlated maps are singleton on every non-broadcast axis, so the mask + # collapses (C-order) to the broadcast block. + combined_bcast = combined.reshape(broadcast_shape) + surviving = np.nonzero(combined_bcast.reshape(-1))[0].astype(np.intp) + if surviving.size == 0: + return None + + # Intersect residual (slice / constant) dimensions independently. Slice dims + # are ordered by input dimension so their flat-buffer strides are row-major. + slice_dims: list[tuple[int, int, int, int, DimensionMap]] = [] # (in_dim, lo, hi, full, m) + for out_dim, m in enumerate(transform.output): + if out_dim in correlated_dims: + continue + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + if isinstance(m, ConstantMap): + if not (lo <= m.offset < hi): + return None + elif isinstance(m, DimensionMap): + d = m.input_dimension + input_lo = transform.domain.inclusive_min[d] + input_hi = transform.domain.exclusive_max[d] + narrowed = _intersect_dimension_map(m, input_lo, input_hi, lo, hi) + if narrowed is None: + return None + slice_dims.append((d, narrowed[0], narrowed[1], input_hi - input_lo, m)) + slice_dims.sort(key=lambda item: item[0]) + + n_points = int(surviving.size) + n_slice = len(slice_dims) + corr_values = { + out_dim: cast("ArrayMap", transform.output[out_dim]) + .index_array.reshape(broadcast_shape) + .reshape(-1)[surviving] + for out_dim in correlated_dims + } + + # New domain: the collapsed broadcast axis, then one axis per residual slice. + new_min = [0] + new_max = [n_points] + new_input_dim_of = {} + for new_axis, (d, nlo, nhi, _full, _m) in enumerate(slice_dims, start=1): + new_min.append(nlo) + new_max.append(nhi) + new_input_dim_of[d] = new_axis + new_domain = IndexDomain(inclusive_min=tuple(new_min), exclusive_max=tuple(new_max)) + + corr_shape = (n_points,) + (1,) * n_slice + new_output: list[OutputIndexMap] = [] + for out_dim, m in enumerate(transform.output): + if out_dim in correlated_dims: + corr = cast("ArrayMap", m) + new_output.append( + ArrayMap( + index_array=corr_values[out_dim].reshape(corr_shape).astype(np.intp), + offset=corr.offset, + stride=corr.stride, + ) + ) + elif isinstance(m, ConstantMap): + new_output.append(m) + else: + assert isinstance(m, DimensionMap) + new_output.append( + DimensionMap( + input_dimension=new_input_dim_of[m.input_dimension], + offset=m.offset, + stride=m.stride, + ) + ) + result = IndexTransform(domain=new_domain, output=tuple(new_output)) + + # Flat scatter index into the row-major output buffer of shape + # (broadcast points, residual slice sizes...): flat = point * prod(slice) + + # (row-major offset within the slice block). + prod_slice = 1 + for _d, _lo, _hi, full, _m in slice_dims: + prod_slice *= full + out_indices: np.ndarray[Any, np.dtype[np.intp]] = (surviving * prod_slice).reshape( + (n_points,) + (1,) * n_slice + ) + running = 1 + for j in range(n_slice - 1, -1, -1): + _d, nlo, nhi, full, _m = slice_dims[j] + coords = np.arange(nlo, nhi, dtype=np.intp) * running + shape = [1] * (1 + n_slice) + shape[1 + j] = coords.size + out_indices = out_indices + coords.reshape(shape) + running *= full + return (result, out_indices.astype(np.intp)) + + +def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | None, ...]: + """Normalize a selection to a tuple of int, slice, or None (newaxis), + expanding ellipsis and padding with slice(None) as needed. + """ + if not isinstance(selection, tuple): + selection = (selection,) + + # Count non-newaxis, non-ellipsis entries to determine how many real dims are addressed + n_newaxis = sum(1 for s in selection if s is None) + has_ellipsis = any(s is Ellipsis for s in selection) + n_real = len(selection) - n_newaxis - (1 if has_ellipsis else 0) + + if n_real > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, but {n_real} were indexed" + ) + + result: list[int | slice | None] = [] + ellipsis_seen = False + for sel in selection: + if sel is Ellipsis: + if ellipsis_seen: + raise IndexError("an index can only have a single ellipsis ('...')") + ellipsis_seen = True + num_missing = ndim - n_real + result.extend([slice(None)] * num_missing) + elif isinstance(sel, (int, np.integer)): + result.append(int(sel)) + elif isinstance(sel, slice) or sel is None: + result.append(sel) + else: + raise IndexError(f"unsupported selection type for basic indexing: {type(sel)!r}") + + # Pad remaining dimensions with slice(None) + while sum(1 for s in result if s is not None) < ndim: + result.append(slice(None)) + + return tuple(result) + + +def _reindex_array( + m: ArrayMap, + normalized: tuple[int | slice | None, ...], + domain: IndexDomain, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Apply basic indexing operations to an ArrayMap's index_array. + + The array's axes correspond to the transform's input dimensions (0-indexed + over the domain shape). Each axis is either a **dependency axis** — the array + genuinely varies with that input dimension — or a **singleton** axis it + broadcasts over. Integer indexing, slicing, or newaxis is applied to the + array only along its dependency axes; a selection on a singleton axis does not + touch the array's values (it just narrows or drops that broadcast axis). + """ + dependent = set(_array_map_dependency_axes(m.index_array)) + if m.input_dimension is not None: + # Degenerate length-1 orthogonal selection: the recorded axis is a + # dependency even though its size (1) makes it look singleton. + dependent.add(m.input_dimension) + arr = m.index_array + + # Build a numpy indexing tuple: one entry per old input dimension + idx: list[Any] = [] + old_dim = 0 + newaxis_positions: list[int] = [] + result_axis = 0 + + for sel in normalized: + if sel is None: + newaxis_positions.append(result_axis) + result_axis += 1 + elif isinstance(sel, int): + if old_dim < arr.ndim: + if old_dim in dependent: + # Convert absolute domain coordinate to 0-based array index + idx.append(sel - domain.inclusive_min[old_dim]) + else: + # Broadcast axis: keep the single element and drop the axis. + idx.append(0) + old_dim += 1 + else: + # sel: slice (normalized: tuple[int | slice | None, ...]) + if old_dim < arr.ndim: + if old_dim in dependent: + lo = domain.inclusive_min[old_dim] + hi = domain.exclusive_max[old_dim] + # Bounds are literal domain coordinates; the stored array is + # indexed positionally, so shift by the domain origin. + start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + pos = start - lo + idx.append(slice(pos, pos + size * step, step)) + else: + # Broadcast axis: preserve the singleton (it still broadcasts + # over the narrowed domain), regardless of the slice bounds. + idx.append(slice(None)) + old_dim += 1 + result_axis += 1 + + result = arr[tuple(idx)] if idx else arr + + for pos in newaxis_positions: + result = np.expand_dims(result, axis=pos) + + return np.asarray(result, dtype=np.intp) + + +_FANCY_AFTER_FANCY_MSG = ( + "applying a fancy (orthogonal/vectorized) selection to a view that already " + "has a fancy-indexed axis is not supported (fancy-after-fancy composition): " + "the new coordinates would index a broadcast axis of the existing selection. " + "Materialize the view first with `.result()` and index the array, or reorder " + "the selections so the fancy step is applied last." +) + + +def _guard_fancy_after_fancy(m: ArrayMap, fancy_dims: set[int] | list[int]) -> None: + """Reject a fancy step that lands on a broadcast axis of an existing ArrayMap. + + A new orthogonal/vectorized selection can only be absorbed into an existing + ArrayMap along the axes that map genuinely varies over (its dependency axes, + plus the recorded `input_dimension` for a degenerate length-1 orthogonal + selection). A fancy index targeting any other axis — a singleton axis the map + merely broadcasts over — cannot be reindexed and used to leak a raw NumPy + `IndexError` at resolve time. Raise a clear `NotImplementedError` instead. + """ + dependent = set(_array_map_dependency_axes(m.index_array)) + if m.input_dimension is not None: + dependent.add(m.input_dimension) + for d in fancy_dims: + if d < m.index_array.ndim and d not in dependent: + raise NotImplementedError(_FANCY_AFTER_FANCY_MSG) + + +def _reindex_array_oindex( + arr: np.ndarray[Any, np.dtype[np.intp]], + normalized: tuple[Any, ...] | list[Any], + domain: IndexDomain, +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Apply oindex/vindex selection to an existing ArrayMap's index_array. + + Each old input dimension gets either an array (fancy index that axis) + or a slice applied to the corresponding array axis. + """ + idx: list[Any] = [] + for old_dim, sel in enumerate(normalized): + if old_dim >= arr.ndim: + break + lo = domain.inclusive_min[old_dim] + if isinstance(sel, np.ndarray): + # Values are literal domain coordinates; the stored array is + # indexed positionally, so shift by the domain origin. + idx.append(sel - lo) + elif isinstance(sel, slice): + hi = domain.exclusive_max[old_dim] + start, step, _origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + pos = start - lo + idx.append(slice(pos, pos + size * step, step)) + else: + idx.append(slice(None)) + + result = arr[tuple(idx)] if idx else arr + return np.asarray(result, dtype=np.intp) + + +def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply basic indexing (int, slice, ellipsis, newaxis) to an IndexTransform.""" + normalized = _normalize_basic_selection(selection, transform.domain.ndim) + + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + old_dim = 0 + new_dim_idx = 0 + old_to_new_dim: dict[int, int] = {} + dropped_dims: set[int] = set() + + # Per old-dim: the slice parameters (for computing new output maps) + dim_slice_params: dict[int, tuple[int, int, int]] = {} # old_dim -> (start, stop, step) + dim_int_val: dict[int, int] = {} # old_dim -> integer index value + + for sel in normalized: + if sel is None: + # newaxis: add a size-1 dimension + new_inclusive_min.append(0) + new_exclusive_max.append(1) + new_dim_idx += 1 + elif isinstance(sel, int): + # Integer index: drop this input dimension. + # Negative indices are literal coordinates (TensorStore convention), + # NOT "from the end" like NumPy. The Array layer handles conversion. + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + idx = sel + if idx < lo or idx >= hi: + hint = _LITERAL_HINT if sel < 0 else "" + raise BoundsCheckError( + f"index {sel} is out of bounds for dimension {old_dim} " + f"(valid indices [{lo}, {hi})){hint}" + ) + dropped_dims.add(old_dim) + dim_int_val[old_dim] = idx + old_dim += 1 + else: + # sel: slice (normalized: tuple[int | slice | None, ...]) + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + + # TensorStore semantics: bounds are literal coordinates; a step-1 + # slice keeps them as the new domain, a strided slice's domain is + # [trunc(start/step), trunc(start/step) + size). + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + dim_slice_params[old_dim] = (start, step, origin) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + old_dim += 1 + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + # Now update output maps + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in dropped_dims: + # Integer index: this output becomes constant + new_offset = m.offset + m.stride * dim_int_val[d] + new_output.append(ConstantMap(offset=new_offset)) + elif d in old_to_new_dim: + # Slice: new coordinate `origin + k` maps to old coordinate + # `start + k*step`, i.e. old = start - step*origin + step*new. + start, step, origin = dim_slice_params[d] + new_offset = m.offset + m.stride * (start - step * origin) + new_stride = m.stride * step + new_input_dim = old_to_new_dim[d] + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + raise RuntimeError(f"unexpected: dimension {d} not handled") + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + new_arr = _reindex_array(m, normalized, transform.domain) + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) + ) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, ...]: + """Return the input axes on which a normalized index array varies. + + Normalized `ArrayMap` index arrays carry the full input rank of their + enclosing transform: an axis the array varies over has its full size, while + an axis the array is independent of is a singleton (size 1). The dependency + axes are therefore exactly the non-singleton axes. An orthogonal (`oindex`) + array depends on a single axis; a vectorized (`vindex`) array depends on all + of the (shared) broadcast axes. + """ + return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) + + +def _array_map_dependent_axis(m: ArrayMap) -> int: + """Return the single input axis an orthogonal `ArrayMap` varies over. + + Normally this is the array's one non-singleton axis. A degenerate length-1 + orthogonal selection normalizes to an all-singleton shape (its dependency + axes are empty and indistinguishable by shape from a scalar), so + `input_dimension` breaks the tie — it records the axis the map binds. + """ + dep = _array_map_dependency_axes(m.index_array) + if len(dep) == 1: + return dep[0] + if m.input_dimension is not None: + return m.input_dimension + raise ValueError( + f"orthogonal ArrayMap must vary over exactly one axis; got dependency " + f"axes {dep} with input_dimension={m.input_dimension}" + ) + + +def _reshape_to_axis( + values: np.ndarray[Any, np.dtype[np.intp]], axis: int, ndim: int +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Reshape a 1-D selection to full rank `ndim` varying only along `axis`. + + The result has `values` laid out along `axis` and singleton (size-1) axes + everywhere else, so its dependency axis is derivable from its shape. + """ + flat = np.asarray(values, dtype=np.intp).ravel() + shape = [1] * ndim + shape[axis] = flat.shape[0] + return flat.reshape(shape) + + +class _OIndexHelper: + """Helper that provides orthogonal (outer) indexing via `transform.oindex[...]`.""" + + def __init__(self, transform: IndexTransform) -> None: + self._transform = transform + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_oindex(self._transform, selection) + + +def _normalize_oindex_selection( + selection: Any, ndim: int +) -> tuple[np.ndarray[Any, np.dtype[np.intp]] | slice, ...]: + """Normalize an oindex selection: arrays, slices, booleans, integers.""" + if not isinstance(selection, tuple): + selection = (selection,) + + # Expand ellipsis + has_ellipsis = any(s is Ellipsis for s in selection) + n_ellipsis = 1 if has_ellipsis else 0 + n_real = len(selection) - n_ellipsis + + result: list[np.ndarray[Any, np.dtype[np.intp]] | slice] = [] + for sel in selection: + if sel is Ellipsis: + num_missing = ndim - n_real + result.extend([slice(None)] * num_missing) + elif isinstance(sel, np.ndarray) and sel.dtype == np.bool_: + # Boolean array -> integer indices + (indices,) = np.nonzero(sel) + result.append(indices.astype(np.intp)) + elif isinstance(sel, np.ndarray): + result.append(sel.astype(np.intp)) + elif isinstance(sel, slice): + result.append(sel) + elif isinstance(sel, (int, np.integer)): + # Convert integer scalars to 1-element arrays for orthogonal indexing + result.append(np.array([int(sel)], dtype=np.intp)) + elif isinstance(sel, (list, tuple)): + result.append(np.asarray(sel, dtype=np.intp)) + else: + result.append(sel) + + # Pad with slice(None) + while len(result) < ndim: + result.append(slice(None)) + + return tuple(result) + + +def _apply_oindex(transform: IndexTransform, selection: Any) -> IndexTransform: + """Apply orthogonal indexing to an IndexTransform. + + Each index array is applied independently per dimension (outer product). + """ + normalized = _normalize_oindex_selection(selection, transform.domain.ndim) + + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + new_dim_idx = 0 + old_to_new_dim: dict[int, int] = {} + + # Info per old dim + dim_array: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + dim_slice_params: dict[int, tuple[int, int, int]] = {} + + for old_dim, sel in enumerate(normalized): + if isinstance(sel, np.ndarray): + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + # Index-array values are literal domain coordinates; the fancy dim + # they create gets a fresh zero-origin [0, n) domain (TensorStore). + _check_array_in_bounds(sel, lo, hi) + dim_array[old_dim] = sel + new_inclusive_min.append(0) + new_exclusive_max.append(len(sel)) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + else: + # sel: slice (_normalize_oindex_selection returns + # tuple[np.ndarray | slice, ...]) + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + dim_slice_params[old_dim] = (start, step, origin) + old_to_new_dim[old_dim] = new_dim_idx + new_dim_idx += 1 + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in dim_array: + new_axis = old_to_new_dim[d] + # Normalize to full input rank: the selection varies along its own + # new axis and is singleton on every other axis. The dependency + # axis is then derivable from the shape (a single non-singleton + # axis marks the selection orthogonal / outer-product rather than + # vectorized). `input_dimension` is kept populated as a + # compatibility shim for consumers not yet migrated to the + # shape-derived classifier. + full_arr = _reshape_to_axis(dim_array[d], new_axis, new_dim_idx) + new_output.append( + ArrayMap( + index_array=full_arr, + offset=m.offset, + stride=m.stride, + input_dimension=new_axis, + ) + ) + elif d in dim_slice_params: + start, step, origin = dim_slice_params[d] + new_offset = m.offset + m.stride * (start - step * origin) + new_stride = m.stride * step + new_input_dim = old_to_new_dim[d] + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + raise RuntimeError(f"unexpected: dimension {d} not handled") + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + _guard_fancy_after_fancy(m, list(dim_array.keys())) + new_arr = _reindex_array_oindex(m.index_array, normalized, transform.domain) + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) + ) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +class _VIndexHelper: + """Helper that provides vectorized (fancy) indexing via `transform.vindex[...]`.""" + + def __init__(self, transform: IndexTransform) -> None: + self._transform = transform + + def __getitem__(self, selection: Any) -> IndexTransform: + return _apply_vindex(self._transform, selection) + + +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. + """ + if not isinstance(selection, tuple): + selection = (selection,) + + # Expand ellipsis and count consumed dimensions + # Boolean arrays with ndim > 1 consume ndim dims + n_consumed = 0 + for s in selection: + if s is Ellipsis: + continue + if isinstance(s, np.ndarray) and s.dtype == np.bool_ and s.ndim > 1: + n_consumed += s.ndim + else: + n_consumed += 1 + ndim = transform.domain.ndim + + expanded: list[Any] = [] + for sel in selection: + if sel is Ellipsis: + num_missing = ndim - n_consumed + expanded.extend([slice(None)] * num_missing) + else: + expanded.append(sel) + # Count dimensions already consumed by expanded entries + n_expanded_dims = 0 + for sel in expanded: + if isinstance(sel, np.ndarray) and sel.dtype == np.bool_ and sel.ndim > 1: + n_expanded_dims += sel.ndim + else: + n_expanded_dims += 1 + while n_expanded_dims < ndim: + expanded.append(slice(None)) + n_expanded_dims += 1 + + # Convert booleans, lists, ints to integer arrays + processed: list[np.ndarray[Any, np.dtype[np.intp]] | slice] = [] + for sel in expanded: + if isinstance(sel, np.ndarray) and sel.dtype == np.bool_: + indices_tuple = np.nonzero(sel) + processed.extend(indices.astype(np.intp) for indices in indices_tuple) + elif isinstance(sel, np.ndarray): + processed.append(sel.astype(np.intp)) + elif isinstance(sel, (list, tuple)): + processed.append(np.asarray(sel, dtype=np.intp)) + elif isinstance(sel, (int, np.integer)): + processed.append(np.array([int(sel)], dtype=np.intp)) + else: + processed.append(sel) + + # Separate array dims and slice dims + array_dims: list[int] = [] + slice_dims: list[int] = [] + arrays: list[np.ndarray[Any, np.dtype[np.intp]]] = [] + + for i, sel in enumerate(processed): + if isinstance(sel, np.ndarray): + lo = transform.domain.inclusive_min[i] + hi = transform.domain.exclusive_max[i] + _check_array_in_bounds(sel, lo, hi) + array_dims.append(i) + arrays.append(sel) + else: + slice_dims.append(i) + + # Broadcast all arrays together + broadcast_arrays: list[np.ndarray[Any, np.dtype[np.intp]]] + if len(arrays) > 0: + broadcast_arrays = list(np.broadcast_arrays(*arrays)) + broadcast_shape = broadcast_arrays[0].shape + else: + broadcast_arrays = [] + broadcast_shape = () + + # Build new domain: broadcast dims first, then slice dims + new_inclusive_min: list[int] = [] + new_exclusive_max: list[int] = [] + + # Broadcast dimensions + for s in broadcast_shape: + new_inclusive_min.append(0) + new_exclusive_max.append(s) + + # Slice dimensions (preserved-domain literal semantics, like basic indexing) + slice_dim_params: dict[int, tuple[int, int, int]] = {} + for old_dim in slice_dims: + sel = processed[old_dim] + assert isinstance(sel, slice) + lo = transform.domain.inclusive_min[old_dim] + hi = transform.domain.exclusive_max[old_dim] + start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) + new_inclusive_min.append(origin) + new_exclusive_max.append(origin + size) + slice_dim_params[old_dim] = (start, step, origin) + + new_domain = IndexDomain( + inclusive_min=tuple(new_inclusive_min), + exclusive_max=tuple(new_exclusive_max), + ) + + # Build output maps + array_dim_to_broadcast: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + for i, d in enumerate(array_dims): + array_dim_to_broadcast[d] = broadcast_arrays[i] + + # New dim index for slice dims starts after broadcast dims + n_broadcast_dims = len(broadcast_shape) + + new_output: list[OutputIndexMap] = [] + for m in transform.output: + if isinstance(m, ConstantMap): + new_output.append(m) + elif isinstance(m, DimensionMap): + d = m.input_dimension + if d in array_dim_to_broadcast: + # Normalize to full input rank: the broadcast (correlated) axes + # come first, followed by a singleton axis per slice dimension. + # Every vectorized array shares the same broadcast axes, so the + # dependency axes derived from the shape coincide — the signature + # of a pointwise scatter rather than an outer product. + broadcast_arr = array_dim_to_broadcast[d] + full_arr = broadcast_arr.reshape(broadcast_shape + (1,) * len(slice_dims)) + new_output.append( + ArrayMap( + index_array=full_arr, + offset=m.offset, + stride=m.stride, + ) + ) + else: + # Slice dim: new coord `origin + k` maps to old `start + k*step` + start, step, origin = slice_dim_params[d] + new_offset = m.offset + m.stride * (start - step * origin) + new_stride = m.stride * step + new_input_dim = n_broadcast_dims + slice_dims.index(d) + new_output.append( + DimensionMap( + input_dimension=new_input_dim, offset=new_offset, stride=new_stride + ) + ) + else: + # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) + _guard_fancy_after_fancy(m, array_dims) + new_arr = _reindex_array_oindex(m.index_array, processed, transform.domain) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) + + return IndexTransform(domain=new_domain, output=tuple(new_output)) + + +_LITERAL_HINT = ( + "; within this transform layer, indices are literal domain coordinates (the " + "public Array boundary wraps NumPy-style negatives before they reach here)" +) + + +def _trunc_div(a: int, b: int) -> int: + """Integer division rounded toward zero (C semantics), as TensorStore uses + for strided-slice domain origins — distinct from Python's floor division + for negative operands (`trunc(-9/2) == -4` where `-9 // 2 == -5`).""" + q = a // b + if q < 0 and q * b != a: + q += 1 + return q + + +def _resolve_slice_ts(sel: slice, dim: int, lo: int, hi: int) -> tuple[int, int, int, int]: + """Resolve a slice against domain `[lo, hi)` with TensorStore semantics. + + Slice bounds are **literal domain coordinates** — never from-the-end, never + clamped. Rules (each verified against tensorstore 0.1.84): + + - defaults: `start = lo`, `stop = hi`; + - a non-empty interval must be contained in the domain (no clamping — a + NumPy-style out-of-range or negative bound is an error, not a shorter or + wrapped result); + - an **empty** interval (`start == stop`) is valid anywhere; + - reversed bounds (`start > stop` with positive step) are an error, not + an empty result; + - the result's domain origin is `trunc(start/step)` (rounded toward + zero) and coordinate `origin + k` maps to input `start + k*step`. + + Returns `(start, step, origin, size)` in domain coordinates. + """ + step = 1 if sel.step is None else sel.step + if step <= 0: + # Negative steps are valid in TensorStore but not yet supported here; + # step 0 is invalid everywhere. + raise IndexError("slice step must be positive") + start = lo if sel.start is None else sel.start + stop = hi if sel.stop is None else sel.stop + if stop < start: + raise IndexError( + f"slice interval [{start}, {stop}) with step {step} does not specify " + f"a valid interval for dimension {dim} (start > stop)" + ) + size = -(-(stop - start) // step) # ceil((stop - start) / step) + if size > 0 and (start < lo or stop > hi): + hint = _LITERAL_HINT if (start < 0 or stop < 0) and lo >= 0 else "" + raise BoundsCheckError( + f"slice interval [{start}, {stop}) is not contained within domain " + f"[{lo}, {hi}) for dimension {dim}{hint}" + ) + origin = _trunc_div(start, step) + return start, step, origin, size + + +def _check_array_in_bounds(arr: np.ndarray[Any, np.dtype[np.intp]], lo: int, hi: int) -> None: + """Reject index-array values outside the domain `[lo, hi)`. + + Index-array values are literal domain coordinates (TensorStore semantics): + a value below `inclusive_min` is out of bounds rather than counting from + the end. Out-of-range values raise instead of silently wrapping. + """ + if arr.size == 0: + return + lo_val, hi_val = int(arr.min()), int(arr.max()) + if lo_val < lo: + hint = _LITERAL_HINT if lo_val < 0 and lo >= 0 else "" + raise BoundsCheckError( + f"index {lo_val} is out of bounds (valid indices [{lo}, {hi})){hint}" + ) + if hi_val >= hi: + raise BoundsCheckError(f"index {hi_val} is out of bounds (valid indices [{lo}, {hi}))") + + +def _validate_array_selection(selection: Any, shape: tuple[int, ...], mode: str) -> None: + """Validate array-based selections (orthogonal, vectorized). + + Rejects types that are not valid for coordinate/vectorized indexing. + Does not check bounds — the transform operations handle that. + """ + items = selection if isinstance(selection, tuple) else (selection,) + for sel in items: + if isinstance(sel, slice): + # vindex is coordinate-only (matches eager zarr): every axis needs an + # integer/boolean array, never a slice. Orthogonal (oindex) allows slices. + if mode == "vectorized": + raise VindexInvalidSelectionError( + "unsupported selection type for vectorized indexing; only " + "coordinate selection (tuple of integer arrays) and mask selection " + f"(single Boolean array) are supported; got {selection!r}" + ) + continue + if sel is Ellipsis or isinstance(sel, (int, np.integer)): + continue + if isinstance(sel, (list, np.ndarray)): + continue + raise IndexError(f"unsupported selection type for {mode} indexing: {type(sel)!r}") + + +def _validate_basic_selection(selection: Any) -> None: + """Validate that a selection only contains basic indexing types (int, slice, Ellipsis). + + Rejects None (newaxis), arrays, lists, floats, strings, etc. + """ + items = selection if isinstance(selection, tuple) else (selection,) + for s in items: + if s is Ellipsis or isinstance(s, (int, np.integer, slice)): + continue + raise IndexError(f"unsupported selection type for basic indexing: {type(s)!r}") + + +def selection_to_transform( + selection: Any, + transform: IndexTransform, + mode: Literal["basic", "orthogonal", "vectorized"], +) -> IndexTransform: + """Convert a user selection into a composed IndexTransform. + + Negative indices are treated as literal coordinates (TensorStore convention). + The caller (Array layer) is responsible for converting numpy-style negative + indices before calling this function. + """ + if mode == "basic": + _validate_basic_selection(selection) + return transform[selection] + elif mode == "orthogonal": + _validate_array_selection(selection, transform.domain.shape, mode) + return transform.oindex[selection] + elif mode == "vectorized": + _validate_array_selection(selection, transform.domain.shape, mode) + return transform.vindex[selection] + else: + raise ValueError(f"Unknown mode: {mode!r}") diff --git a/packages/zarr-indexing/tests/conformance/PROVENANCE.md b/packages/zarr-indexing/tests/conformance/PROVENANCE.md new file mode 100644 index 0000000000..6f75de56c3 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/PROVENANCE.md @@ -0,0 +1,20 @@ +# Provenance of the ndsel conformance corpus + +The JSON fixtures in this directory (`point.json`, `box.json`, `slice.json`, +`points.json`, `transform.json`, `errors.json`) and `README.md` are **vendored, +unmodified**, from the ndsel reference repository. + +- **Source:** +- **Branch:** `main` (merge of PR #1, `fix/slice-origin-trunc`) +- **Commit:** `c59bc556c` (fixtures byte-identical to the previously vendored + `c132b4c1caa3205830ce35a42502363171f650a7`) +- **Path in source:** `conformance/` + +**Do not edit these files.** They are vendored as-is so that +`zarr_indexing`' ndsel message layer can be checked against the same +language-agnostic corpus every other ndsel implementation runs. To update the +corpus, re-vendor from a newer ndsel commit and update the commit SHA above. + +ndsel PR #1 (merged) corrected the `slice` desugaring origin from +`floor(a/s)` to `trunc(a/s)` (rounding toward zero), which matches +`zarr_indexing`' existing `_trunc_div` semantics. diff --git a/packages/zarr-indexing/tests/conformance/README.md b/packages/zarr-indexing/tests/conformance/README.md new file mode 100644 index 0000000000..ecb0c57ca2 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/README.md @@ -0,0 +1,16 @@ +# ndsel conformance corpus + +Language-agnostic fixtures. Each file is a JSON array of cases. + +A **success** case: + { "name": "...", "input": , "normalized": } + +An **error** case: + { "name": "...", "input": , "error": "" } + +An implementation is conformant iff, for every success case, +`normalize(input)` equals `normalized` by structural JSON equality, and for +every error case, `normalize(input)` is rejected with the given reason code. + +The `normalized` value is a canonical `transform` body (the `kind` field is +omitted; implementations compare the transform structure). diff --git a/packages/zarr-indexing/tests/conformance/box.json b/packages/zarr-indexing/tests/conformance/box.json new file mode 100644 index 0000000000..e847872f86 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/box.json @@ -0,0 +1,50 @@ +[ + { + "name": "box/2d-min-max", + "input": { "kind": "box", "inclusive_min": [0, 0], "exclusive_max": [3, 4] }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [0, 0], + "input_exclusive_max": [3, 4], + "input_labels": ["", ""], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "box/shape-only-origin-zero", + "input": { "kind": "box", "shape": [5] }, + "normalized": { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [5], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 1, "input_dimension": 0 } ] + } + }, + { + "name": "box/inclusive-max", + "input": { "kind": "box", "inclusive_min": [2], "inclusive_max": [9] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [2], "input_exclusive_max": [10], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 1, "input_dimension": 0 } ] + } + }, + { + "name": "box/implicit-and-infinite-bounds", + "input": { "kind": "box", "inclusive_min": [["-inf"], 0], "exclusive_max": [["+inf"], 4], "labels": ["t", ""] }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [["-inf"], 0], + "input_exclusive_max": [["+inf"], 4], + "input_labels": ["t", ""], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/errors.json b/packages/zarr-indexing/tests/conformance/errors.json new file mode 100644 index 0000000000..e5e08bab3c --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/errors.json @@ -0,0 +1,23 @@ +[ + { "name": "error/step-zero", "input": { "kind": "slice", "start": [0], "stop": [4], "step": [0] }, "error": "step_zero" }, + { "name": "error/negative-step", "input": { "kind": "slice", "start": [9], "stop": [0], "step": [-2] }, "error": "negative_step_unsupported" }, + { "name": "error/multiple-upper-bounds", "input": { "kind": "box", "shape": [3], "exclusive_max": [3] }, "error": "multiple_upper_bounds" }, + { "name": "error/rank-mismatch", "input": { "kind": "slice", "start": [0, 0], "stop": [4] }, "error": "rank_mismatch" }, + { "name": "error/unknown-kind", "input": { "kind": "bogus" }, "error": "unknown_kind" }, + { "name": "error/transform-multiple-upper-bounds", "input": { "kind": "transform", "input_shape": [3], "input_exclusive_max": [3] }, "error": "multiple_upper_bounds" }, + { "name": "error/transform-rank-mismatch", "input": { "kind": "transform", "input_rank": 2, "input_inclusive_min": [0] }, "error": "rank_mismatch" }, + { "name": "error/missing-kind", "input": { "coords": [1, 2] }, "error": "invalid_json" }, + { "name": "error/point-missing-coords", "input": { "kind": "point" }, "error": "invalid_json" }, + { "name": "error/point-bool-coord", "input": { "kind": "point", "coords": [true] }, "error": "invalid_json" }, + { "name": "error/slice-missing-stop", "input": { "kind": "slice", "start": [0] }, "error": "invalid_json" }, + { "name": "error/box-non-list-bound", "input": { "kind": "box", "inclusive_min": 5 }, "error": "invalid_json" }, + { "name": "error/points-bool-coord", "input": { "kind": "points", "coords": [[true]] }, "error": "invalid_json" }, + { "name": "error/integer-out-of-i64-range", "input": { "kind": "point", "coords": [99999999999999999999] }, "error": "invalid_json" }, + { "name": "error/box-inverted-bounds", "input": { "kind": "box", "inclusive_min": [5], "exclusive_max": [3] }, "error": "bounds_out_of_order" }, + { "name": "error/box-negative-shape", "input": { "kind": "box", "shape": [-3] }, "error": "bounds_out_of_order" }, + { "name": "error/transform-inverted-bounds", "input": { "kind": "transform", "input_inclusive_min": [0], "input_exclusive_max": [-1] }, "error": "bounds_out_of_order" }, + { "name": "error/output-map-conflict", "input": { "kind": "transform", "output": [{ "input_dimension": 0, "index_array": [1, 2] }] }, "error": "output_map_conflict" }, + { "name": "error/box-unknown-field", "input": { "kind": "box", "shapee": [3] }, "error": "unknown_field" }, + { "name": "error/point-unknown-field", "input": { "kind": "point", "coords": [1], "extra": true }, "error": "unknown_field" }, + { "name": "error/output-map-unknown-field", "input": { "kind": "transform", "output": [{ "offset": 0, "bogus": 1 }] }, "error": "unknown_field" } +] diff --git a/packages/zarr-indexing/tests/conformance/point.json b/packages/zarr-indexing/tests/conformance/point.json new file mode 100644 index 0000000000..99a5ea16d8 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/point.json @@ -0,0 +1,30 @@ +[ + { + "name": "point/2d", + "input": { "kind": "point", "coords": [4, 7] }, + "normalized": { + "input_rank": 0, + "input_inclusive_min": [], + "input_exclusive_max": [], + "input_labels": [], + "output": [ { "offset": 4 }, { "offset": 7 } ] + } + }, + { + "name": "point/scalar-0d", + "input": { "kind": "point", "coords": [] }, + "normalized": { + "input_rank": 0, "input_inclusive_min": [], "input_exclusive_max": [], + "input_labels": [], "output": [] + } + }, + { + "name": "point/large-i64", + "input": { "kind": "point", "coords": [1152921504606846976] }, + "normalized": { + "input_rank": 0, "input_inclusive_min": [], "input_exclusive_max": [], + "input_labels": [], + "output": [ { "offset": 1152921504606846976 } ] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/points.json b/packages/zarr-indexing/tests/conformance/points.json new file mode 100644 index 0000000000..1ad92e12b1 --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/points.json @@ -0,0 +1,34 @@ +[ + { + "name": "points/three-2d", + "input": { "kind": "points", "coords": [[1, 10], [2, 20], [3, 30]] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ + { "offset": 0, "stride": 1, "index_array": [1, 2, 3], "index_array_bounds": ["-inf", "+inf"] }, + { "offset": 0, "stride": 1, "index_array": [10, 20, 30], "index_array_bounds": ["-inf", "+inf"] } + ] + } + }, + { + "name": "points/1d", + "input": { "kind": "points", "coords": [[5], [9], [2]] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ + { "offset": 0, "stride": 1, "index_array": [5, 9, 2], "index_array_bounds": ["-inf", "+inf"] } + ] + } + }, + { + "name": "points/empty", + "input": { "kind": "points", "coords": [] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [0], "input_exclusive_max": [0], + "input_labels": [""], + "output": [] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/slice.json b/packages/zarr-indexing/tests/conformance/slice.json new file mode 100644 index 0000000000..2f1a0694ce --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/slice.json @@ -0,0 +1,61 @@ +[ + { + "name": "slice/unit-step-preserves-frame", + "input": { "kind": "slice", "start": [5], "stop": [10] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [5], "input_exclusive_max": [10], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 1, "input_dimension": 0 } ] + } + }, + { + "name": "slice/divisible-stride", + "input": { "kind": "slice", "start": [4], "stop": [10], "step": [2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [2], "input_exclusive_max": [5], + "input_labels": [""], + "output": [ { "offset": 0, "stride": 2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/nondivisible-stride-phase-offset", + "input": { "kind": "slice", "start": [5], "stop": [10], "step": [2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [2], "input_exclusive_max": [5], + "input_labels": [""], + "output": [ { "offset": 1, "stride": 2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/2d-mixed-step", + "input": { "kind": "slice", "start": [0, 5], "stop": [10, 10], "step": [2, 1] }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [0, 5], + "input_exclusive_max": [5, 10], + "input_labels": ["", ""], + "output": [ + { "offset": 0, "stride": 2, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "slice/negative-start-trunc-origin", + "input": { "kind": "slice", "start": [-9], "stop": [5], "step": [2] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-4], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ { "offset": -1, "stride": 2, "input_dimension": 0 } ] + } + }, + { + "name": "slice/negative-start-trunc-origin-step3", + "input": { "kind": "slice", "start": [-8], "stop": [6], "step": [3] }, + "normalized": { + "input_rank": 1, "input_inclusive_min": [-2], "input_exclusive_max": [3], + "input_labels": [""], + "output": [ { "offset": -2, "stride": 3, "input_dimension": 0 } ] + } + } +] diff --git a/packages/zarr-indexing/tests/conformance/transform.json b/packages/zarr-indexing/tests/conformance/transform.json new file mode 100644 index 0000000000..f26157aaab --- /dev/null +++ b/packages/zarr-indexing/tests/conformance/transform.json @@ -0,0 +1,57 @@ +[ + { + "name": "transform/omitted-output-identity", + "input": { "kind": "transform", "input_inclusive_min": [0, 0], "input_exclusive_max": [3, 4] }, + "normalized": { + "input_rank": 2, "input_inclusive_min": [0, 0], "input_exclusive_max": [3, 4], + "input_labels": ["", ""], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "transform/implicit-bounds-and-labels", + "input": { + "kind": "transform", + "input_inclusive_min": [["-inf"], 7], + "input_exclusive_max": [["+inf"], 11], + "input_labels": ["x", "y"] + }, + "normalized": { + "input_rank": 2, + "input_inclusive_min": [["-inf"], 7], + "input_exclusive_max": [["+inf"], 11], + "input_labels": ["x", "y"], + "output": [ + { "offset": 0, "stride": 1, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "input_dimension": 1 } + ] + } + }, + { + "name": "transform/explicit-output-all-three-map-kinds", + "input": { + "kind": "transform", + "input_inclusive_min": [0], + "input_exclusive_max": [3], + "output": [ + { "offset": 7 }, + { "input_dimension": 0, "stride": 2 }, + { "index_array": [1, 2, 3] } + ] + }, + "normalized": { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [3], + "input_labels": [""], + "output": [ + { "offset": 7 }, + { "offset": 0, "stride": 2, "input_dimension": 0 }, + { "offset": 0, "stride": 1, "index_array": [1, 2, 3], "index_array_bounds": ["-inf", "+inf"] } + ] + } + } +] diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py new file mode 100644 index 0000000000..f2d4cab012 --- /dev/null +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -0,0 +1,496 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from zarr.core.chunk_grids import ChunkGrid, FixedDimension, VaryingDimension + +import zarr_indexing.chunk_resolution as chunk_resolution +from zarr_indexing.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + +if TYPE_CHECKING: + import pytest + + +class TestChunkResolutionIdentity: + def test_single_chunk(self) -> None: + """Array fits in one chunk.""" + t = IndexTransform.from_shape((10,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=10),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + coords, sub_t, _ = results[0] + assert coords == (0,) + assert sub_t.domain.shape == (10,) + + def test_multiple_chunks_1d(self) -> None: + """1D array spanning 3 chunks.""" + t = IndexTransform.from_shape((30,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 3 + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + assert (2,) in coords_list + + def test_multiple_chunks_2d(self) -> None: + """2D array spanning 2x3 chunks.""" + t = IndexTransform.from_shape((20, 30)) + grid = ChunkGrid( + dimensions=( + FixedDimension(size=10, extent=20), + FixedDimension(size=10, extent=30), + ) + ) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 6 + coords_list = [r[0] for r in results] + assert (0, 0) in coords_list + assert (1, 2) in coords_list + + +class TestChunkResolutionSliced: + def test_slice_within_chunk(self) -> None: + """Slice that falls within a single chunk.""" + # Chunk resolution consumes zero-origin transforms: the I/O layer + # normalizes preserved (user-facing) domains via translate_domain_to + # before resolving, so mirror that contract here. + t = IndexTransform.from_shape((100,))[5:8].translate_domain_to((0,)) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + coords, sub_t, _ = results[0] + assert coords == (0,) + assert isinstance(sub_t.output[0], DimensionMap) + assert sub_t.output[0].offset == 5 + + def test_slice_across_chunks(self) -> None: + """Slice that spans two chunks.""" + t = IndexTransform.from_shape((100,))[8:15] + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=100),)) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 2 + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + + +class TestChunkResolutionConstant: + def test_integer_index(self) -> None: + """Integer index produces constant map — single chunk per constant dim.""" + t = IndexTransform.from_shape((100, 100))[25, :] + grid = ChunkGrid( + dimensions=( + FixedDimension(size=10, extent=100), + FixedDimension(size=10, extent=100), + ) + ) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 10 + for coords, _, _ in results: + assert coords[0] == 2 + + +class TestChunkResolutionArray: + def test_array_index(self) -> None: + """Array index map — chunks determined by array values.""" + idx = np.array([5, 15, 25], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=idx),), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=10, extent=30),)) + results = list(iter_chunk_transforms(t, grid)) + coords_list = [r[0] for r in results] + assert (0,) in coords_list + assert (1,) in coords_list + assert (2,) in coords_list + + +class TestChunkResolutionSorted1D: + def test_matches_general_resolution_for_randomized_sorted_selections( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Direct partitioning matches the original resolver across varied inputs.""" + rng = np.random.default_rng(0) + grids = ( + ChunkGrid(dimensions=(FixedDimension(size=7, extent=30),)), + ChunkGrid(dimensions=(VaryingDimension(edges=(3, 4, 8, 5, 10), extent=30),)), + ) + + for grid in grids: + for _ in range(50): + idx = np.sort(rng.integers(0, 30, size=int(rng.integers(1, 80)))).astype(np.intp) + transform = IndexTransform.from_shape((30,)).vindex[idx] + direct = list(iter_chunk_transforms(transform, grid)) + + with monkeypatch.context() as context: + context.setattr( + chunk_resolution, + "_one_dimensional_correlated_array_map", + lambda _transform: None, + ) + general = list(iter_chunk_transforms(transform, grid)) + + assert [result[0] for result in direct] == [result[0] for result in general] + for direct_result, general_result in zip(direct, general, strict=True): + _, direct_t, direct_out = direct_result + _, general_t, general_out = general_result + assert direct_t.domain == general_t.domain + + direct_chunk_sel, direct_out_sel, direct_drop = sub_transform_to_selections( + direct_t, direct_out + ) + general_chunk_sel, general_out_sel, general_drop = sub_transform_to_selections( + general_t, general_out + ) + assert direct_drop == general_drop + np.testing.assert_array_equal(direct_chunk_sel[0], general_chunk_sel[0]) + np.testing.assert_array_equal(direct_out_sel[0], general_out_sel[0]) + + def test_sorted_vindex_partitions_chunks_without_intersection( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Sorted vectorized coordinates are sliced directly per touched chunk.""" + idx = np.array([0, 3, 4, 4, 9, 11], dtype=np.intp) + t = IndexTransform.from_shape((12,)).vindex[idx] + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + assert calls["n"] == 0 + + expected_chunk_indices = ([0, 3], [0, 0], [1, 3]) + expected_out_indices = ([0, 1], [2, 3], [4, 5]) + for result, expected_chunk, expected_out in zip( + results, expected_chunk_indices, expected_out_indices, strict=True + ): + _, sub_t, out_indices = result + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices) + np.testing.assert_array_equal(chunk_sel[0], expected_chunk) + np.testing.assert_array_equal(out_sel[0], expected_out) + assert drop_axes == () + + def test_sorted_array_map_preserves_offset_and_stride(self) -> None: + """Storage partitioning retains the ArrayMap's offset and stride.""" + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=( + ArrayMap( + index_array=np.array([0, 1, 2], dtype=np.intp), + offset=1, + stride=3, + ), + ), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=8),)) + + results = list(iter_chunk_transforms(t, grid)) + + assert [result[0] for result in results] == [(0,), (1,)] + expected_chunk_indices = ([1], [0, 3]) + expected_out_indices = ([0], [1, 2]) + for result, expected_chunk, expected_out in zip( + results, expected_chunk_indices, expected_out_indices, strict=True + ): + _, sub_t, out_indices = result + chunk_sel, out_sel, _ = sub_transform_to_selections(sub_t, out_indices) + np.testing.assert_array_equal(chunk_sel[0], expected_chunk) + np.testing.assert_array_equal(out_sel[0], expected_out) + + def test_sorted_vindex_with_varying_chunks(self) -> None: + """Touched-boundary searches also support a non-uniform 1-D grid.""" + idx = np.array([0, 1, 2, 3, 5, 9], dtype=np.intp) + t = IndexTransform.from_shape((10,)).vindex[idx] + grid = ChunkGrid(dimensions=(VaryingDimension(edges=(2, 3, 5), extent=10),)) + + results = list(iter_chunk_transforms(t, grid)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + expected_chunk_indices = ([0, 1], [0, 1], [0, 4]) + for result, expected_chunk in zip(results, expected_chunk_indices, strict=True): + _, sub_t, out_indices = result + chunk_sel, _, _ = sub_transform_to_selections(sub_t, out_indices) + np.testing.assert_array_equal(chunk_sel[0], expected_chunk) + + def test_sorted_vindex_with_zero_sized_dimension_uses_general_resolution( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A zero-sized grid cannot be partitioned by touched boundaries.""" + t = IndexTransform.from_shape((10,)).vindex[np.array([1], dtype=np.intp)] + grid = ChunkGrid(dimensions=(FixedDimension(size=0, extent=10),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + assert results == [] + assert calls["n"] == 1 + + def test_unsorted_vindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Unsorted coordinates continue through the general intersection logic.""" + t = IndexTransform.from_shape((12,)).vindex[np.array([9, 0, 4], dtype=np.intp)] + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + assert calls["n"] == 3 + + def test_sorted_oindex_uses_general_resolution(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Orthogonal ArrayMaps retain their existing domain-aware resolution.""" + t = IndexTransform.from_shape((12,)).oindex[np.array([0, 4, 9], dtype=np.intp)] + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + assert [result[0] for result in results] == [(0,), (1,), (2,)] + assert calls["n"] == 3 + + +def _count_intersect_calls(monkeypatch: pytest.MonkeyPatch) -> dict[str, int]: + """Wrap `IndexTransform.intersect` with a call counter. + + Returns a mutable dict whose `"n"` entry is the number of times + `intersect` is invoked. Used to assert that candidate-chunk enumeration is + proportional to the *touched* chunks, not the dense bounding box between the + min and max touched chunk. + """ + calls = {"n": 0} + original = IndexTransform.intersect + + def counting(self: IndexTransform, output_domain: IndexDomain) -> object: + calls["n"] += 1 + return original(self, output_domain) + + monkeypatch.setattr(IndexTransform, "intersect", counting) + return calls + + +class TestChunkResolutionTouchedOnly: + """`iter_chunk_transforms` must enumerate only the chunks a fancy selection + actually touches — never the dense `range(min_chunk, max_chunk + 1)` bounding + box. These guard against a regression to bounding-box enumeration, whose cost + scales with grid size rather than with the number of selected coordinates. + """ + + def test_1d_sparse_vindex_enumerates_only_touched_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two far-apart coordinates on a 1000-chunk grid touch exactly 2 chunks. + + A dense bounding-box enumeration would intersect ~1000 candidate chunks; + touched-only enumeration intersects exactly 2. + """ + # 4000 elements, chunk size 4 -> 1000 chunks. coords 1 and 3997 land in + # chunk 0 and chunk 999 respectively (998 empty chunks between them). + grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=4000),)) + t = IndexTransform.from_shape((4000,)).vindex[np.array([1, 3997], dtype=np.intp)] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + coords = sorted(r[0] for r in results) + assert coords == [(0,), (999,)] + # Sorted 1-D coordinates are partitioned directly, without intersecting + # either the touched chunks or the 998 empty chunks between them. + assert calls["n"] == 0 + + def test_2d_orthogonal_enumerates_only_touched_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Orthogonal outer product of two 2-coordinate arrays touches 2x2 chunks. + + Per-dimension distinct touched chunks: {0, 999} on each axis. The outer + product is 2*2 = 4 candidate chunks (all survive), versus ~1e6 for a + dense 1000x1000 bounding box. + """ + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4000), + FixedDimension(size=4, extent=4000), + ) + ) + t = IndexTransform.from_shape((4000, 4000)).oindex[ + np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) + ] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + coords = sorted(r[0] for r in results) + assert coords == [(0, 0), (0, 999), (999, 0), (999, 999)] + assert calls["n"] == 4 + + def test_2d_correlated_vindex_enumerates_per_dim_distinct_chunks( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Two correlated (vindex) coordinate arrays scatter to 2 diagonal chunks. + + The two points (1, 2) and (3997, 3998) touch chunks (0, 0) and + (999, 999). Per-dimension distinct touched chunks are {0, 999} on each + axis, so enumeration intersects the 2x2 = 4 combinations; the two + off-diagonal combinations are filtered out by `intersect`, leaving 2 + surviving chunks. The key guarantee is that the work is bounded by the + per-dimension distinct touched chunks (4), not the dense 1e6 grid. + """ + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4000), + FixedDimension(size=4, extent=4000), + ) + ) + t = IndexTransform.from_shape((4000, 4000)).vindex[ + np.array([1, 3997], dtype=np.intp), np.array([2, 3998], dtype=np.intp) + ] + + calls = _count_intersect_calls(monkeypatch) + results = list(iter_chunk_transforms(t, grid)) + + coords = sorted(r[0] for r in results) + assert coords == [(0, 0), (999, 999)] + # 2x2 per-dim-distinct combinations enumerated; 2 survive intersection. + assert calls["n"] == 4 + + +class TestSubTransformToSelections: + def test_constant_map(self) -> None: + """ConstantMap produces int selection + drop axis.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (5,) + assert out_sel == () + assert drop_axes == () + + def test_dimension_map_stride_1(self) -> None: + """DimensionMap with stride=1 produces contiguous slice.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=3, stride=1),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (slice(3, 13, 1),) + assert out_sel == (slice(0, 10),) + assert drop_axes == () + + def test_dimension_map_strided(self) -> None: + """DimensionMap with stride>1 produces strided slice.""" + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(DimensionMap(input_dimension=0, offset=2, stride=3),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel == (slice(2, 17, 3),) + assert out_sel == (slice(0, 5),) + assert drop_axes == () + + def test_array_map(self) -> None: + """ArrayMap produces integer array selection.""" + arr = np.array([1, 5, 9], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=0, stride=1),), + ) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(t) + assert isinstance(chunk_sel[0], np.ndarray) + np.testing.assert_array_equal(chunk_sel[0], arr) + # Without chunk_mask, out_sel falls back to domain-based slices + assert out_sel == (slice(0, 3),) + assert drop_axes == () + + def test_array_map_with_offset_stride(self) -> None: + """ArrayMap with offset and stride computes storage coords.""" + arr = np.array([0, 1, 2], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=10, stride=5),), + ) + chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) + assert isinstance(chunk_sel[0], np.ndarray) + np.testing.assert_array_equal(chunk_sel[0], np.array([10, 15, 20])) + assert drop_axes == () + + def test_mixed_maps_2d(self) -> None: + """Mix of ConstantMap and DimensionMap.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + chunk_sel, _out_sel, drop_axes = sub_transform_to_selections(t) + assert chunk_sel[0] == 5 + assert chunk_sel[1] == slice(0, 10, 1) + # drop_axes is empty — integer in chunk_sel naturally drops the dim via numpy + assert drop_axes == () + + +class TestChunkResolutionArrayMapFlavours: + """Chunk resolution must yield outer-product (np.ix_) selectors for + orthogonal ArrayMaps and shared flat-scatter selectors for correlated ones, + and must return early for empty fancy selections.""" + + def test_empty_array_selection_yields_nothing(self) -> None: + """An empty ArrayMap selection produces no chunk transforms (no crash).""" + t = IndexTransform( + domain=IndexDomain.from_shape((0,)), + output=(ArrayMap(index_array=np.array([], dtype=np.intp)),), + ) + grid = ChunkGrid(dimensions=(FixedDimension(size=3, extent=10),)) + assert list(iter_chunk_transforms(t, grid)) == [] + + def test_orthogonal_outer_product_selectors(self) -> None: + """Two independent arrays produce np.ix_-style (mesh) chunk/out selectors.""" + t = IndexTransform.from_shape((10, 10)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + grid = ChunkGrid( + dimensions=(FixedDimension(size=10, extent=10), FixedDimension(size=10, extent=10)) + ) + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + _coords, sub_t, out_indices = results[0] + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_t, out_indices) + # np.ix_ produces one 2-D open-mesh selector per axis, for both sides. + assert len(chunk_sel) == 2 + assert len(out_sel) == 2 + assert isinstance(chunk_sel[0], np.ndarray) + assert isinstance(chunk_sel[1], np.ndarray) + assert chunk_sel[0].shape == (2, 1) + assert chunk_sel[1].shape == (1, 3) + assert drop_axes == () + + def test_correlated_scatter_with_residual_slice(self) -> None: + """Correlated arrays + a residual slice dim scatter through a single flat + index whose shape matches the (points, slice) block read from the chunk.""" + t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] + grid = ChunkGrid( + dimensions=( + FixedDimension(size=4, extent=4), + FixedDimension(size=3, extent=3), + FixedDimension(size=5, extent=5), + ) + ) + # One chunk holds everything: both points survive, slice dim spans [0,5). + results = list(iter_chunk_transforms(t, grid)) + assert len(results) == 1 + _coords, sub_t, out_indices = results[0] + chunk_sel, out_sel, _drop = sub_transform_to_selections(sub_t, out_indices) + # Chunk side: flat coordinate arrays for the two correlated dims plus a + # slice for the residual dim. + assert len(chunk_sel) == 3 + np.testing.assert_array_equal(np.asarray(chunk_sel[0]), [1, 3]) + np.testing.assert_array_equal(np.asarray(chunk_sel[1]), [2, 0]) + assert chunk_sel[2] == slice(0, 5, 1) + # Output side: a single flat scatter index of shape (points, slice) = (2, 5). + assert len(out_sel) == 1 + assert np.asarray(out_sel[0]).shape == (2, 5) diff --git a/packages/zarr-indexing/tests/test_composition.py b/packages/zarr-indexing/tests/test_composition.py new file mode 100644 index 0000000000..dd92f59b80 --- /dev/null +++ b/packages/zarr-indexing/tests/test_composition.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.composition import compose +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + + +class TestComposeConstantInner: + """Inner = constant. Result is always constant.""" + + def test_constant_inner_any_outer(self) -> None: + outer = IndexTransform.from_shape((5,)) + inner = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ConstantMap(offset=42),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 42 + + +class TestComposeDimensionInner: + """Inner = DimensionMap.""" + + def test_dimension_inner_constant_outer(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 25 + + def test_dimension_inner_dimension_outer(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=5, stride=2),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 25 + assert result.output[0].stride == 6 + assert result.output[0].input_dimension == 0 + + def test_dimension_inner_array_outer(self) -> None: + arr = np.array([0, 2, 4], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr, offset=5, stride=2),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=3),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].offset == 25 + assert result.output[0].stride == 6 + np.testing.assert_array_equal(result.output[0].index_array, arr) + + +class TestComposeArrayInner: + """Inner = ArrayMap.""" + + def test_array_inner_constant_outer(self) -> None: + inner_arr = np.array([10, 20, 30], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ConstantMap(offset=1),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=inner_arr, offset=0, stride=1),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 20 + + def test_array_inner_array_outer(self) -> None: + outer_arr = np.array([0, 2, 1], dtype=np.intp) + inner_arr = np.array([10, 20, 30], dtype=np.intp) + outer = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=outer_arr, offset=0, stride=1),), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=inner_arr, offset=0, stride=1),), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ArrayMap) + expected = np.array([10, 30, 20], dtype=np.intp) + np.testing.assert_array_equal(result.output[0].index_array, expected) + + +class TestComposeMultiDim: + def test_2d_identity_compose(self) -> None: + a = IndexTransform.from_shape((10, 20)) + b = IndexTransform.from_shape((10, 20)) + result = compose(a, b) + assert result.domain.shape == (10, 20) + for i in range(2): + m = result.output[i] + assert isinstance(m, DimensionMap) + assert m.input_dimension == i + assert m.offset == 0 + assert m.stride == 1 + + def test_mixed_map_types(self) -> None: + outer = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + inner = IndexTransform( + domain=IndexDomain.from_shape((10, 10)), + output=( + DimensionMap(input_dimension=0, offset=2, stride=3), + DimensionMap(input_dimension=1, offset=0, stride=1), + ), + ) + result = compose(outer, inner) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 17 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 0 + assert result.output[1].offset == 0 + assert result.output[1].stride == 1 + + def test_rank_mismatch_raises(self) -> None: + outer = IndexTransform.from_shape((10,)) + inner = IndexTransform.from_shape((10, 20)) + with pytest.raises(ValueError, match="rank"): + compose(outer, inner) + + +class TestComposeChain: + def test_three_transforms(self) -> None: + a = IndexTransform.from_shape((100,)) + b = IndexTransform( + domain=IndexDomain.from_shape((100,)), + output=(DimensionMap(input_dimension=0, offset=10, stride=1),), + ) + c = IndexTransform( + domain=IndexDomain.from_shape((100,)), + output=(DimensionMap(input_dimension=0, offset=5, stride=2),), + ) + bc = compose(b, c) + abc = compose(a, bc) + assert isinstance(abc.output[0], DimensionMap) + assert abc.output[0].offset == 25 + assert abc.output[0].stride == 2 diff --git a/packages/zarr-indexing/tests/test_conformance.py b/packages/zarr-indexing/tests/test_conformance.py new file mode 100644 index 0000000000..207a9d8236 --- /dev/null +++ b/packages/zarr-indexing/tests/test_conformance.py @@ -0,0 +1,55 @@ +"""ndsel conformance corpus harness. + +Runs the vendored, language-agnostic ndsel fixtures (see +`tests/conformance/PROVENANCE.md`) against this package's message layer +(`zarr_indexing.messages`). An implementation is conformant iff: + +- for every *success* fixture, `normalize_ndsel(input)` equals the fixture's + `normalized` value by structural JSON equality; +- for every *error* fixture, `normalize_ndsel(input)` is rejected with an + `NdselError` carrying the fixture's `error` reason code. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +from zarr_indexing.messages import NdselError, normalize_ndsel + +_CONFORMANCE_DIR = Path(__file__).parent / "conformance" + + +def _load_cases() -> list[tuple[str, dict[str, Any]]]: + cases: list[tuple[str, dict[str, Any]]] = [] + for path in sorted(_CONFORMANCE_DIR.glob("*.json")): + data = json.loads(path.read_text()) + cases.extend((f"{path.stem}::{case['name']}", case) for case in data) + return cases + + +_CASES = _load_cases() +_SUCCESS = [(name, c) for name, c in _CASES if "normalized" in c] +_ERROR = [(name, c) for name, c in _CASES if "error" in c] + + +def test_corpus_is_present() -> None: + # Guard against an empty/missing vendored corpus silently passing. + assert len(_SUCCESS) > 0 + assert len(_ERROR) > 0 + + +@pytest.mark.parametrize(("name", "case"), _SUCCESS, ids=[name for name, _ in _SUCCESS]) +def test_success_fixture(name: str, case: dict[str, Any]) -> None: + result = normalize_ndsel(case["input"]) + assert result == case["normalized"] + + +@pytest.mark.parametrize(("name", "case"), _ERROR, ids=[name for name, _ in _ERROR]) +def test_error_fixture(name: str, case: dict[str, Any]) -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel(case["input"]) + assert excinfo.value.reason == case["error"] diff --git a/packages/zarr-indexing/tests/test_domain.py b/packages/zarr-indexing/tests/test_domain.py new file mode 100644 index 0000000000..9664a0b08a --- /dev/null +++ b/packages/zarr-indexing/tests/test_domain.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import pytest + +from zarr_indexing.domain import IndexDomain + + +class TestIndexDomainConstruction: + def test_from_shape(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.inclusive_min == (0, 0) + assert d.exclusive_max == (10, 20) + assert d.ndim == 2 + assert d.origin == (0, 0) + assert d.shape == (10, 20) + + def test_from_shape_0d(self) -> None: + d = IndexDomain.from_shape(()) + assert d.ndim == 0 + assert d.shape == () + + def test_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(5, 10), exclusive_max=(15, 30)) + assert d.origin == (5, 10) + assert d.shape == (10, 20) + assert d.ndim == 2 + + def test_validation_mismatched_lengths(self) -> None: + with pytest.raises(ValueError, match="same length"): + IndexDomain(inclusive_min=(0,), exclusive_max=(10, 20)) + + def test_validation_min_greater_than_max(self) -> None: + with pytest.raises(ValueError, match="inclusive_min must be <="): + IndexDomain(inclusive_min=(10,), exclusive_max=(5,)) + + def test_empty_domain(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(5,)) + assert d.shape == (0,) + + def test_labels(self) -> None: + d = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + assert d.labels == ("x", "y") + + def test_labels_none(self) -> None: + d = IndexDomain.from_shape((10,)) + assert d.labels is None + + +class TestIndexDomainContains: + def test_contains_inside(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((0, 0)) is True + assert d.contains((9, 19)) is True + assert d.contains((5, 10)) is True + + def test_contains_outside(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((10, 0)) is False + assert d.contains((-1, 0)) is False + assert d.contains((0, 20)) is False + + def test_contains_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + assert d.contains((5,)) is True + assert d.contains((9,)) is True + assert d.contains((4,)) is False + assert d.contains((10,)) is False + + def test_contains_wrong_ndim(self) -> None: + d = IndexDomain.from_shape((10, 20)) + assert d.contains((5,)) is False + + def test_contains_domain_inside(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain(inclusive_min=(2, 3), exclusive_max=(8, 15)) + assert outer.contains_domain(inner) is True + + def test_contains_domain_outside(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain(inclusive_min=(2, 3), exclusive_max=(11, 15)) + assert outer.contains_domain(inner) is False + + def test_contains_domain_wrong_ndim(self) -> None: + outer = IndexDomain.from_shape((10, 20)) + inner = IndexDomain.from_shape((5,)) + assert outer.contains_domain(inner) is False + + +class TestIndexDomainIntersect: + def test_overlapping(self) -> None: + a = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 10)) + b = IndexDomain(inclusive_min=(5, 5), exclusive_max=(15, 15)) + result = a.intersect(b) + assert result is not None + assert result.inclusive_min == (5, 5) + assert result.exclusive_max == (10, 10) + + def test_disjoint(self) -> None: + a = IndexDomain(inclusive_min=(0,), exclusive_max=(5,)) + b = IndexDomain(inclusive_min=(10,), exclusive_max=(15,)) + assert a.intersect(b) is None + + def test_touching_boundary(self) -> None: + a = IndexDomain(inclusive_min=(0,), exclusive_max=(5,)) + b = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + assert a.intersect(b) is None + + def test_contained(self) -> None: + a = IndexDomain.from_shape((20,)) + b = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + result = a.intersect(b) + assert result is not None + assert result.inclusive_min == (5,) + assert result.exclusive_max == (10,) + + def test_wrong_ndim(self) -> None: + a = IndexDomain.from_shape((10,)) + b = IndexDomain.from_shape((10, 20)) + with pytest.raises(ValueError, match="different ranks"): + a.intersect(b) + + +class TestIndexDomainTranslate: + def test_translate_positive(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.translate((5, 10)) + assert result.inclusive_min == (5, 10) + assert result.exclusive_max == (15, 30) + + def test_translate_negative(self) -> None: + d = IndexDomain(inclusive_min=(10, 20), exclusive_max=(30, 40)) + result = d.translate((-10, -20)) + assert result.inclusive_min == (0, 0) + assert result.exclusive_max == (20, 20) + + def test_translate_wrong_length(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(ValueError, match="same length"): + d.translate((1, 2)) + + +class TestIndexDomainNarrow: + def test_narrow_slice(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.narrow((slice(2, 8), slice(5, 15))) + assert result.inclusive_min == (2, 5) + assert result.exclusive_max == (8, 15) + + def test_narrow_int(self) -> None: + d = IndexDomain.from_shape((10, 20)) + result = d.narrow((3, slice(None))) + assert result.inclusive_min == (3, 0) + assert result.exclusive_max == (4, 20) + + def test_narrow_ellipsis(self) -> None: + d = IndexDomain.from_shape((10, 20, 30)) + result = d.narrow((slice(1, 5), ...)) + assert result.inclusive_min == (1, 0, 0) + assert result.exclusive_max == (5, 20, 30) + + def test_narrow_slice_none(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow((slice(None),)) + assert result == d + + def test_narrow_non_zero_origin(self) -> None: + d = IndexDomain(inclusive_min=(10,), exclusive_max=(20,)) + result = d.narrow((slice(12, 18),)) + assert result.inclusive_min == (12,) + assert result.exclusive_max == (18,) + + def test_narrow_int_out_of_bounds(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="out of bounds"): + d.narrow((10,)) + + def test_narrow_int_below_origin(self) -> None: + d = IndexDomain(inclusive_min=(5,), exclusive_max=(10,)) + with pytest.raises(IndexError, match="out of bounds"): + d.narrow((4,)) + + def test_narrow_clamps_to_domain(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow((slice(-5, 100),)) + assert result.inclusive_min == (0,) + assert result.exclusive_max == (10,) + + def test_narrow_bare_slice(self) -> None: + d = IndexDomain.from_shape((10,)) + result = d.narrow(slice(2, 8)) + assert result.inclusive_min == (2,) + assert result.exclusive_max == (8,) + + def test_narrow_too_many_indices(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="too many indices"): + d.narrow((1, 2)) + + def test_narrow_step_not_one(self) -> None: + d = IndexDomain.from_shape((10,)) + with pytest.raises(IndexError, match="step=1"): + d.narrow((slice(0, 10, 2),)) diff --git a/packages/zarr-indexing/tests/test_json.py b/packages/zarr-indexing/tests/test_json.py new file mode 100644 index 0000000000..42b59b2c30 --- /dev/null +++ b/packages/zarr-indexing/tests/test_json.py @@ -0,0 +1,336 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.json import ( + IndexTransformJSON, + index_domain_from_json, + index_domain_to_json, + index_transform_from_json, + index_transform_to_json, + output_index_map_from_json, + output_index_map_to_json, +) +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + + +def _maps_equal(a: object, b: object) -> bool: + if type(a) is not type(b): + return False + if isinstance(a, ConstantMap): + assert isinstance(b, ConstantMap) + return a.offset == b.offset + if isinstance(a, DimensionMap): + assert isinstance(b, DimensionMap) + return (a.input_dimension, a.offset, a.stride) == (b.input_dimension, b.offset, b.stride) + assert isinstance(a, ArrayMap) + assert isinstance(b, ArrayMap) + return ( + a.offset == b.offset + and a.stride == b.stride + and a.input_dimension == b.input_dimension + and np.array_equal(a.index_array, b.index_array) + ) + + +def _transforms_equal(a: IndexTransform, b: IndexTransform) -> bool: + """Structural equality that compares `ArrayMap` index arrays element-wise + (`IndexTransform`'s dataclass `__eq__` cannot, as numpy `==` is ambiguous).""" + return ( + a.domain == b.domain + and len(a.output) == len(b.output) + and all(_maps_equal(x, y) for x, y in zip(a.output, b.output, strict=True)) + ) + + +class TestIndexDomainJSON: + def test_roundtrip(self) -> None: + domain = IndexDomain(inclusive_min=(2, 5), exclusive_max=(10, 20)) + json = index_domain_to_json(domain) + assert json == { + "input_inclusive_min": [2, 5], + "input_exclusive_max": [10, 20], + "input_labels": ["", ""], + } + restored = index_domain_from_json(json) + assert restored == domain + + def test_with_labels(self) -> None: + domain = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + json = index_domain_to_json(domain) + assert json["input_labels"] == ["x", "y"] + restored = index_domain_from_json(json) + assert restored.labels == ("x", "y") + + def test_without_labels_emits_empty_and_round_trips_to_none(self) -> None: + domain = IndexDomain.from_shape((5,)) + json = index_domain_to_json(domain) + # Canonical form always writes labels; an unlabeled domain gets [""]*rank. + assert json["input_labels"] == [""] + restored = index_domain_from_json(json) + assert restored.labels is None + + def test_zero_origin(self) -> None: + domain = IndexDomain.from_shape((10, 20, 30)) + json = index_domain_to_json(domain) + assert json == { + "input_inclusive_min": [0, 0, 0], + "input_exclusive_max": [10, 20, 30], + "input_labels": ["", "", ""], + } + assert index_domain_from_json(json) == domain + + +class TestOutputIndexMapJSON: + def test_constant(self) -> None: + m = ConstantMap(offset=42) + json = output_index_map_to_json(m) + assert json == {"offset": 42} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 42 + + def test_constant_zero(self) -> None: + m = ConstantMap(offset=0) + json = output_index_map_to_json(m) + assert json == {"offset": 0} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 0 + + def test_dimension(self) -> None: + m = DimensionMap(input_dimension=1, offset=10, stride=3) + json = output_index_map_to_json(m) + assert json == {"offset": 10, "stride": 3, "input_dimension": 1} + restored = output_index_map_from_json(json) + assert isinstance(restored, DimensionMap) + assert restored.input_dimension == 1 + assert restored.offset == 10 + assert restored.stride == 3 + + def test_dimension_stride_1_written(self) -> None: + """Canonical form writes stride even at its default of 1.""" + m = DimensionMap(input_dimension=0) + json = output_index_map_to_json(m) + assert json == {"offset": 0, "stride": 1, "input_dimension": 0} + restored = output_index_map_from_json(json) + assert isinstance(restored, DimensionMap) + assert restored.stride == 1 + + def test_array(self) -> None: + arr = np.array([1, 5, 9], dtype=np.intp) + m = ArrayMap(index_array=arr, offset=2, stride=3) + json = output_index_map_to_json(m) + # Canonical: stride/offset present, index_array_bounds present, and + # no input_dimension (ndsel/TensorStore reject it beside index_array). + assert json == { + "offset": 2, + "stride": 3, + "index_array": [1, 5, 9], + "index_array_bounds": ["-inf", "+inf"], + } + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + np.testing.assert_array_equal(restored.index_array, arr) + assert restored.offset == 2 + assert restored.stride == 3 + + def test_array_stride_1_written(self) -> None: + arr = np.array([0, 1, 2], dtype=np.intp) + m = ArrayMap(index_array=arr) + json = output_index_map_to_json(m) + assert json["stride"] == 1 + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + assert restored.stride == 1 + + def test_array_2d(self) -> None: + arr = np.array([[1, 2], [3, 4]], dtype=np.intp) + m = ArrayMap(index_array=arr) + json = output_index_map_to_json(m) + assert json["index_array"] == [[1, 2], [3, 4]] + restored = output_index_map_from_json(json) + assert isinstance(restored, ArrayMap) + np.testing.assert_array_equal(restored.index_array, arr) + + def test_degenerate_singleton_array_collapses_to_constant(self) -> None: + """An all-singleton index_array selects one coordinate -> constant map.""" + m = ArrayMap(index_array=np.array([[4]], dtype=np.intp), offset=1, stride=2) + json = output_index_map_to_json(m) + assert json == {"offset": 1 + 2 * 4} + restored = output_index_map_from_json(json) + assert isinstance(restored, ConstantMap) + assert restored.offset == 9 + + +class TestIndexTransformJSON: + def test_identity(self) -> None: + t = IndexTransform.from_shape((10, 20)) + json = index_transform_to_json(t) + assert json == { + "input_rank": 2, + "input_inclusive_min": [0, 0], + "input_exclusive_max": [10, 20], + "input_labels": ["", ""], + "output": [ + {"offset": 0, "stride": 1, "input_dimension": 0}, + {"offset": 0, "stride": 1, "input_dimension": 1}, + ], + } + restored = index_transform_from_json(json) + assert restored.domain == t.domain + assert len(restored.output) == 2 + for orig, rest in zip(t.output, restored.output, strict=True): + assert type(orig) is type(rest) + + def test_sliced(self) -> None: + t = IndexTransform.from_shape((100,))[10:50:2] + json = index_transform_to_json(t) + restored = index_transform_from_json(json) + assert restored.domain.shape == t.domain.shape + assert isinstance(restored.output[0], DimensionMap) + orig = t.output[0] + assert isinstance(orig, DimensionMap) + assert restored.output[0].offset == orig.offset + assert restored.output[0].stride == orig.stride + + def test_with_constant(self) -> None: + t = IndexTransform.from_shape((10, 20))[3] + json = index_transform_to_json(t) + restored = index_transform_from_json(json) + assert isinstance(restored.output[0], ConstantMap) + assert restored.output[0].offset == 3 + assert isinstance(restored.output[1], DimensionMap) + + def test_with_array(self) -> None: + idx = np.array([1, 5, 9], dtype=np.intp) + t = IndexTransform.from_shape((10, 20)).oindex[idx, :] + json = index_transform_to_json(t) + # The oindex array must not carry input_dimension on the wire. + assert "input_dimension" not in json["output"][0] + restored = index_transform_from_json(json) + assert isinstance(restored.output[0], ArrayMap) + # Orthogonal arrays are normalized to full input rank with a singleton + # axis on the dimension they do not vary over. + assert restored.output[0].index_array.shape == (3, 1) + np.testing.assert_array_equal(restored.output[0].index_array, idx.reshape(3, 1)) + # input_dimension is reconstructed from the sole non-singleton axis. + assert restored.output[0].input_dimension == 0 + assert isinstance(restored.output[1], DimensionMap) + + def test_roundtrip_preserves_singleton_axes(self) -> None: + """Full-rank orthogonal arrays keep their singleton axes across JSON.""" + t = IndexTransform.from_shape((10, 20)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + restored = index_transform_from_json(index_transform_to_json(t)) + orig0, orig1 = t.output[0], t.output[1] + rest0, rest1 = restored.output[0], restored.output[1] + assert isinstance(orig0, ArrayMap) + assert isinstance(orig1, ArrayMap) + assert isinstance(rest0, ArrayMap) + assert isinstance(rest1, ArrayMap) + assert rest0.index_array.shape == (2, 1) + assert rest1.index_array.shape == (1, 3) + np.testing.assert_array_equal(rest0.index_array, orig0.index_array) + np.testing.assert_array_equal(rest1.index_array, orig1.index_array) + # Distinct, exclusively-owned axes -> reconstructed as orthogonal. + assert rest0.input_dimension == 0 + assert rest1.input_dimension == 1 + + def test_with_labels(self) -> None: + domain = IndexDomain(inclusive_min=(0, 0), exclusive_max=(10, 20), labels=("x", "y")) + t = IndexTransform.identity(domain) + json = index_transform_to_json(t) + assert json["input_labels"] == ["x", "y"] + restored = index_transform_from_json(json) + assert restored.domain.labels == ("x", "y") + + def test_tensorstore_compatible_format(self) -> None: + """A canonical body loads and round-trips through the engine layer.""" + json: IndexTransformJSON = { + "input_rank": 3, + "input_inclusive_min": [0, 0, 0], + "input_exclusive_max": [100, 200, 3], + "input_labels": ["x", "y", "channel"], + "output": [ + {"offset": 5}, + {"offset": 10, "stride": 2, "input_dimension": 1}, + {"offset": 0, "stride": 1, "index_array": [1, 2, 0]}, + ], + } + t = index_transform_from_json(json) + assert t.domain.shape == (100, 200, 3) + assert t.domain.labels == ("x", "y", "channel") + assert isinstance(t.output[0], ConstantMap) + assert t.output[0].offset == 5 + assert isinstance(t.output[1], DimensionMap) + assert t.output[1].offset == 10 + assert t.output[1].stride == 2 + assert t.output[1].input_dimension == 1 + assert isinstance(t.output[2], ArrayMap) + np.testing.assert_array_equal(t.output[2].index_array, [1, 2, 0]) + + # Roundtrip + json_rt = index_transform_to_json(t) + t_rt = index_transform_from_json(json_rt) + assert t_rt.domain == t.domain + + +class TestCanonicalRoundTrips: + """Round-trip `transform == from(to(transform))`, up to the documented + degenerate-collapse (all-singleton ArrayMap -> ConstantMap).""" + + def test_oindex_multi_axis(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)).oindex[np.array([1, 3]), :, np.array([2, 4, 6])] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_oindex_with_slice(self) -> None: + t = IndexTransform.from_shape((10, 20))[2:8].oindex[np.array([3, 5, 7]), :] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_vindex(self) -> None: + t = IndexTransform.from_shape((10, 20)).vindex[np.array([1, 3, 5]), np.array([2, 4, 6])] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_vindex_with_residual_slice(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)).vindex[np.array([1, 3]), np.array([2, 4]), :] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + def test_length1_degenerate_oindex_collapses(self) -> None: + """A length-1 oindex array becomes an all-singleton ArrayMap; the JSON + round-trip collapses it to a ConstantMap (behaviorally identical).""" + t = IndexTransform.from_shape((10, 20)).oindex[np.array([7]), :] + m = t.output[0] + assert isinstance(m, ArrayMap) + assert m.index_array.size == 1 + + rt = index_transform_from_json(index_transform_to_json(t)) + # The degenerate array collapsed to a constant selecting the same cell. + rm = rt.output[0] + assert isinstance(rm, ConstantMap) + assert rm.offset == 7 + # The size-1 input dimension survives, unconsumed, in the domain. + assert rt.domain == t.domain + + def test_slices_and_constants(self) -> None: + t = IndexTransform.from_shape((10, 20, 30))[2:8:2, 5, :] + rt = index_transform_from_json(index_transform_to_json(t)) + assert _transforms_equal(rt, t) + + +def test_infinite_bound_rejected_on_lowering() -> None: + body: IndexTransformJSON = { + "input_rank": 1, + "input_inclusive_min": [0], + "input_exclusive_max": [["+inf"]], + "input_labels": [""], + "output": [{"offset": 0, "stride": 1, "input_dimension": 0}], + } + with pytest.raises(ValueError, match="infinite"): + index_transform_from_json(body) diff --git a/packages/zarr-indexing/tests/test_messages.py b/packages/zarr-indexing/tests/test_messages.py new file mode 100644 index 0000000000..14bed66448 --- /dev/null +++ b/packages/zarr-indexing/tests/test_messages.py @@ -0,0 +1,91 @@ +"""Message-layer tests beyond the vendored conformance corpus. + +The corpus (see `test_conformance.py`) covers the desugaring matrix and error +codes. These tests pin behaviors the corpus does not: `normalize` idempotence, +`parse_ndsel`, 64-bit boundary handling, and schema-valid-but-redundant maps. +""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel + +_MESSAGES = [ + {"kind": "point", "coords": [4, 7]}, + {"kind": "box", "inclusive_min": [0, 0], "exclusive_max": [3, 4]}, + {"kind": "box", "inclusive_min": [["-inf"], 0], "exclusive_max": [["+inf"], 4]}, + {"kind": "slice", "start": [5], "stop": [10], "step": [2]}, + {"kind": "points", "coords": [[1, 10], [2, 20]]}, + { + "kind": "transform", + "input_inclusive_min": [0], + "input_exclusive_max": [3], + "output": [{"offset": 7}, {"input_dimension": 0, "stride": 2}, {"index_array": [1, 2, 3]}], + }, +] + + +@pytest.mark.parametrize("message", _MESSAGES) +def test_normalize_is_idempotent(message: dict[str, Any]) -> None: + once = normalize_ndsel(message) + twice = normalize_ndsel({"kind": "transform", **once}) + assert twice == once + + +@pytest.mark.parametrize("message", _MESSAGES) +def test_parse_returns_message_unchanged(message: dict[str, Any]) -> None: + assert parse_ndsel(message) == message + + +def test_parse_rejects_invalid() -> None: + with pytest.raises(NdselError) as excinfo: + parse_ndsel({"kind": "slice", "start": [0]}) + assert excinfo.value.reason == "invalid_json" + + +def test_constant_map_drops_redundant_stride() -> None: + # A constant map (no input_dimension, no index_array) is schema-valid even + # with a stray stride; it canonicalizes to offset-only. + result = normalize_ndsel( + {"kind": "transform", "input_rank": 0, "output": [{"offset": 5, "stride": 9}]} + ) + assert result["output"] == [{"offset": 5}] + + +def test_i64_min_and_max_round_trip() -> None: + i64_min, i64_max = -(2**63), 2**63 - 1 + result = normalize_ndsel({"kind": "point", "coords": [i64_min, i64_max]}) + assert result["output"] == [{"offset": i64_min}, {"offset": i64_max}] + + +def test_i64_overflow_rejected() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": "point", "coords": [2**63]}) + assert excinfo.value.reason == "invalid_json" + + +def test_bool_in_output_offset_rejected() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": "transform", "input_rank": 0, "output": [{"offset": True}]}) + assert excinfo.value.reason == "invalid_json" + + +def test_sentinel_not_allowed_in_plain_integer_position() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": "point", "coords": ["+inf"]}) + assert excinfo.value.reason == "invalid_json" + + +def test_not_an_object_rejected() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel([1, 2, 3]) + assert excinfo.value.reason == "invalid_json" + + +def test_empty_string_kind_is_unknown_kind() -> None: + with pytest.raises(NdselError) as excinfo: + normalize_ndsel({"kind": ""}) + assert excinfo.value.reason == "unknown_kind" diff --git a/packages/zarr-indexing/tests/test_ndsel_tensorstore.py b/packages/zarr-indexing/tests/test_ndsel_tensorstore.py new file mode 100644 index 0000000000..794ac5f3d5 --- /dev/null +++ b/packages/zarr-indexing/tests/test_ndsel_tensorstore.py @@ -0,0 +1,52 @@ +"""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 +`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 \ + packages/zarr-indexing/tests/test_ndsel_tensorstore.py -q +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.json import transform_from_canonical, transform_to_canonical +from zarr_indexing.transform import IndexTransform + +ts = pytest.importorskip("tensorstore") + + +def _canonical_transforms() -> list[IndexTransform]: + base = IndexTransform.from_shape((10, 20)) + return [ + base, # identity + base[2:8:2, :], # strided DimensionMap + identity + base[3, :], # integer index -> ConstantMap + DimensionMap + base.oindex[np.array([1, 5, 9]), :], # orthogonal index_array + IndexTransform.from_shape((10, 20, 30)).vindex[ + np.array([1, 3]), np.array([2, 4]), : + ], # correlated index_arrays + residual slice + ] + + +@pytest.mark.parametrize("transform", _canonical_transforms()) +def test_body_loads_in_tensorstore_and_round_trips(transform: IndexTransform) -> None: + body = transform_to_canonical(transform) + + # (1) The canonical body loads directly as a TensorStore IndexTransform. + ts_transform = ts.IndexTransform(json=body) + + # (2) TensorStore's own JSON re-loads, through our engine, to an equivalent + # transform. Comparing via our canonical form normalizes away + # representational choices (index_array_bounds, default omissions) that + # both sides make differently but that denote the same selection. + ts_json = ts_transform.to_json() + reloaded = transform_from_canonical(ts_json) + assert transform_to_canonical(reloaded) == transform_to_canonical(transform) diff --git a/packages/zarr-indexing/tests/test_output_map.py b/packages/zarr-indexing/tests/test_output_map.py new file mode 100644 index 0000000000..498101444e --- /dev/null +++ b/packages/zarr-indexing/tests/test_output_map.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import numpy as np + +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap + + +class TestConstantMap: + def test_construction(self) -> None: + m = ConstantMap(offset=42) + assert m.offset == 42 + + def test_default_offset(self) -> None: + m = ConstantMap() + assert m.offset == 0 + + def test_frozen(self) -> None: + m = ConstantMap(offset=5) + assert isinstance(m, ConstantMap) + + +class TestDimensionMap: + def test_construction(self) -> None: + m = DimensionMap(input_dimension=3, offset=5, stride=2) + assert m.input_dimension == 3 + assert m.offset == 5 + assert m.stride == 2 + + def test_defaults(self) -> None: + m = DimensionMap(input_dimension=0) + assert m.offset == 0 + assert m.stride == 1 + + def test_frozen(self) -> None: + m = DimensionMap(input_dimension=0) + assert isinstance(m, DimensionMap) + + +class TestArrayMap: + def test_construction(self) -> None: + arr = np.array([1, 3, 5], dtype=np.intp) + m = ArrayMap(index_array=arr, offset=10, stride=2) + assert m.offset == 10 + assert m.stride == 2 + np.testing.assert_array_equal(m.index_array, arr) + + def test_defaults(self) -> None: + arr = np.array([0, 1], dtype=np.intp) + m = ArrayMap(index_array=arr) + assert m.offset == 0 + assert m.stride == 1 + + def test_frozen(self) -> None: + arr = np.array([0], dtype=np.intp) + m = ArrayMap(index_array=arr) + assert isinstance(m, ArrayMap) diff --git a/packages/zarr-indexing/tests/test_tensorstore_parity.py b/packages/zarr-indexing/tests/test_tensorstore_parity.py new file mode 100644 index 0000000000..1ed99046e9 --- /dev/null +++ b/packages/zarr-indexing/tests/test_tensorstore_parity.py @@ -0,0 +1,263 @@ +"""TensorStore-parity oracle tests for IndexTransform semantics. + +Every case in this module was executed against tensorstore 0.1.84 (see the +lazy-indexing design notes): the expected domains, values, and error conditions +are TensorStore's observed behavior, which zarr's lazy indexing matches by +design. Core rules pinned here: + +- **Domain preservation**: a step-1 slice keeps the literal coordinates of the + selected interval (`a[2:10]` has domain `[2, 10)`); nothing re-zeros + implicitly. Re-zeroing is explicit via `translate_to`. +- **Strided-domain rule**: for step ``k``, ``origin = trunc(start/k)`` (rounded + toward zero), ``shape = ceil((stop - start)/k)``, and coordinate + ``origin + i`` maps to base cell ``start + i*k``. +- **Strict containment**: non-empty slice intervals must lie within the domain + — no clamping, no negative-wrapping; empty intervals are valid anywhere; + reversed non-empty bounds are an error, not an empty result. +- **Fancy-dim rule**: index-array dims get fresh explicit ``[0, n)`` domains; + index-array values are absolute domain coordinates. +- **Translate rules**: ``translate_by``/``translate_to`` shift the input domain + while preserving which cells are addressed. +""" + +from __future__ import annotations + +from typing import ClassVar + +import numpy as np +import pytest + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.errors import BoundsCheckError +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform + + +def _identity(lo: int, hi: int) -> IndexTransform: + """Identity transform over the 1-D domain [lo, hi).""" + return IndexTransform.identity(IndexDomain(inclusive_min=(lo,), exclusive_max=(hi,))) + + +def _a() -> IndexTransform: + """The oracle's base fixture: identity over [0, 12).""" + return _identity(0, 12) + + +def _w() -> IndexTransform: + """The oracle's translated fixture: identity over [-10, 2), cell c -> base c + 10.""" + return _a().translate_domain_by((-10,)) + + +def _dim(t: IndexTransform) -> DimensionMap: + m = t.output[0] + assert isinstance(m, DimensionMap) + return m + + +def _base_cells(t: IndexTransform) -> list[int]: + """The base cells a 1-D single-DimensionMap transform addresses, in order.""" + m = _dim(t) + lo, hi = t.domain.inclusive_min[0], t.domain.exclusive_max[0] + return [m.offset + m.stride * c for c in range(lo, hi)] + + +class TestDomainPreservation: + """Oracle section 1-2: step-1 slices keep literal coordinates.""" + + def test_slice_preserves_domain(self) -> None: + t = _a()[2:10] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((2,), (10,)) + assert _base_cells(t) == list(range(2, 10)) + + def test_integer_on_preserved_domain_is_a_coordinate(self) -> None: + v = _a()[2:10] + assert isinstance(v[3].output[0], ConstantMap) + assert v[3].output[0].offset == 3 # coordinate 3 = base cell 3 + assert v[2].output[0].offset == 2 + assert v[9].output[0].offset == 9 + + @pytest.mark.parametrize("bad", [0, -1, 10]) + def test_out_of_domain_integer_raises(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[2, 10\)"): + _a()[2:10][bad] + + def test_slice_of_slice_is_literal(self) -> None: + v = _a()[2:10] + t = v[3:7] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((3,), (7,)) + assert _base_cells(t) == [3, 4, 5, 6] + + def test_ellipsis_preserves_domain(self) -> None: + v = _a()[2:10] + t = v[...] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((2,), (10,)) + + +class TestNegativeOriginDomain: + """Oracle section 3: on domain [-10, 2), -1 is just another index.""" + + def test_translated_domain(self) -> None: + w = _w() + assert (w.domain.inclusive_min, w.domain.exclusive_max) == ((-10,), (2,)) + assert _base_cells(w) == list(range(12)) + + @pytest.mark.parametrize(("coord", "base"), [(-5, 5), (-10, 0), (-1, 9), (1, 11)]) + def test_negative_coordinates_address_cells(self, coord: int, base: int) -> None: + t = _w()[coord] + assert isinstance(t.output[0], ConstantMap) + assert t.output[0].offset == base + + @pytest.mark.parametrize("bad", [-11, 2]) + def test_out_of_domain_raises(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[-10, 2\)"): + _w()[bad] + + def test_negative_slice_bounds_are_coordinates(self) -> None: + t = _w()[-5:] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((-5,), (2,)) + assert _base_cells(t) == [5, 6, 7, 8, 9, 10, 11] + t2 = _w()[-5:-2] + assert (t2.domain.inclusive_min, t2.domain.exclusive_max) == ((-5,), (-2,)) + assert _base_cells(t2) == [5, 6, 7] + + +class TestStridedDomains: + """Oracle section 5: origin = trunc(start/step), coord origin+i -> start + i*step.""" + + # (slice, expected (lo, hi), expected base cells) — verbatim oracle rows. + CASES: ClassVar[list[tuple[slice, tuple[int, int], list[int]]]] = [ + (slice(1, 10, 3), (0, 3), [1, 4, 7]), + (slice(None, None, 2), (0, 6), [0, 2, 4, 6, 8, 10]), + (slice(2, 11, 3), (0, 3), [2, 5, 8]), + (slice(0, 12, 4), (0, 3), [0, 4, 8]), + (slice(5, 12, 2), (2, 6), [5, 7, 9, 11]), + (slice(6, 12, 2), (3, 6), [6, 8, 10]), + (slice(7, 12, 3), (2, 4), [7, 10]), + ] + + @pytest.mark.parametrize(("sel", "dom", "cells"), CASES) + def test_strided_domain_and_cells( + self, sel: slice, dom: tuple[int, int], cells: list[int] + ) -> None: + t = _a()[sel] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == dom + assert _base_cells(t) == cells + + def test_strided_on_negative_origin(self) -> None: + # w[-9:2:2] -> domain [-4, 2), base cells 1,3,5,7,9,11 + t = _w()[-9:2:2] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (-4, 2) + assert _base_cells(t) == [1, 3, 5, 7, 9, 11] + # w[::2] -> domain [-5, 1), base cells 0,2,4,6,8,10 + t2 = _w()[::2] + assert (t2.domain.inclusive_min[0], t2.domain.exclusive_max[0]) == (-5, 1) + assert _base_cells(t2) == [0, 2, 4, 6, 8, 10] + + def test_strided_composition(self) -> None: + s = _a()[1:10:3] # domain [0, 3), cells 1,4,7 + assert [s[k].output[0].offset for k in range(3)] == [1, 4, 7] + t = s[1:3] + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (1, 3) + assert _base_cells(t) == [4, 7] + t2 = _a()[::2][1:4] + assert (t2.domain.inclusive_min[0], t2.domain.exclusive_max[0]) == (1, 4) + assert _base_cells(t2) == [2, 4, 6] + t3 = _a()[::2][::2] + assert (t3.domain.inclusive_min[0], t3.domain.exclusive_max[0]) == (0, 3) + assert _base_cells(t3) == [0, 4, 8] + + @pytest.mark.parametrize("bad", [-2, -1, 3, 4]) + def test_strided_bounds(self, bad: int) -> None: + with pytest.raises(BoundsCheckError, match=r"valid indices \[0, 3\)"): + _a()[1:10:3][bad] + + +class TestStrictContainment: + """Oracle section 11: no clamping, no wrapping; empty intervals valid anywhere.""" + + @pytest.mark.parametrize( + "sel", + [ + slice(5, 100), + slice(-3, None), + slice(-3, -1), + slice(0, 13), + slice(12, 14), + slice(100, 200), + ], + ) + def test_uncontained_interval_raises(self, sel: slice) -> None: + with pytest.raises(BoundsCheckError, match="not contained"): + _a()[sel] + + def test_uncontained_on_negative_origin(self) -> None: + with pytest.raises(BoundsCheckError, match="not contained"): + _w()[-20:] + + @pytest.mark.parametrize( + ("sel", "pos"), [(slice(5, 5), 5), (slice(0, 0), 0), (slice(13, 13), 13)] + ) + def test_empty_interval_valid_anywhere(self, sel: slice, pos: int) -> None: + t = _a()[sel] + assert t.domain.shape == (0,) + assert t.domain.inclusive_min[0] == pos + + @pytest.mark.parametrize("sel", [slice(5, 2), slice(100, 50)]) + def test_reversed_bounds_raise(self, sel: slice) -> None: + with pytest.raises(IndexError, match="valid.*interval|interval"): + _a()[sel] + + +class TestTranslate: + """Oracle sections 4 and 12: translate_by / translate_to preserve the cell mapping.""" + + def test_translate_to_zero(self) -> None: + t = _a()[2:10].translate_domain_to((0,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((0,), (8,)) + assert _base_cells(t) == list(range(2, 10)) + + def test_translate_to_offset(self) -> None: + t = _a().translate_domain_to((5,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((5,), (17,)) + assert _base_cells(t) == list(range(12)) + + def test_translate_by_composes_with_stride(self) -> None: + # a[::2].translate_by[5] -> domain [5, 11), base = 2*(coord-5) + t = _a()[::2].translate_domain_by((5,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((5,), (11,)) + assert _base_cells(t) == [0, 2, 4, 6, 8, 10] + assert t[5].output[0].offset == 0 + assert t[10].output[0].offset == 10 + with pytest.raises(BoundsCheckError, match=r"valid indices \[5, 11\)"): + t[0] + + def test_translate_strided_to(self) -> None: + t = _a()[1:10:3].translate_domain_to((100,)) + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((100,), (103,)) + assert _base_cells(t) == [1, 4, 7] + + +class TestFancyDims: + """Oracle section 7: fancy dims get fresh [0, n); values are absolute coordinates.""" + + def test_index_array_values_are_coordinates(self) -> None: + v = _a()[2:10] + t = v.oindex[(np.array([3, 5], dtype=np.intp),)] + assert (t.domain.inclusive_min, t.domain.exclusive_max) == ((0,), (2,)) + m = t.output[0] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + np.testing.assert_array_equal(np.asarray(storage).ravel(), [3, 5]) + + def test_index_array_on_negative_origin(self) -> None: + t = _w().oindex[(np.array([-10, -1], dtype=np.intp),)] + m = t.output[0] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + np.testing.assert_array_equal(np.asarray(storage).ravel(), [0, 9]) + + def test_index_array_out_of_domain_raises(self) -> None: + v = _a()[2:10] + for bad in ([0, 3], [-1, 3], [3, 10]): + with pytest.raises(BoundsCheckError): + v.oindex[(np.array(bad, dtype=np.intp),)] diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py new file mode 100644 index 0000000000..baecd9ada2 --- /dev/null +++ b/packages/zarr-indexing/tests/test_transform.py @@ -0,0 +1,628 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform, selection_to_transform + + +class TestIndexTransformConstruction: + def test_from_shape(self) -> None: + t = IndexTransform.from_shape((10, 20)) + assert t.input_rank == 2 + assert t.output_rank == 2 + assert t.domain.shape == (10, 20) + assert t.domain.origin == (0, 0) + for i, m in enumerate(t.output): + assert isinstance(m, DimensionMap) + assert m.input_dimension == i + assert m.offset == 0 + assert m.stride == 1 + + def test_identity(self) -> None: + domain = IndexDomain(inclusive_min=(5,), exclusive_max=(15,)) + t = IndexTransform.identity(domain) + assert t.input_rank == 1 + assert t.output_rank == 1 + assert t.domain == domain + assert isinstance(t.output[0], DimensionMap) + assert t.output[0].input_dimension == 0 + + def test_from_shape_0d(self) -> None: + t = IndexTransform.from_shape(()) + assert t.input_rank == 0 + assert t.output_rank == 0 + assert t.domain.shape == () + + def test_custom_output_maps(self) -> None: + domain = IndexDomain.from_shape((10,)) + maps = (ConstantMap(offset=42), DimensionMap(input_dimension=0, offset=5, stride=2)) + t = IndexTransform(domain=domain, output=maps) + assert t.input_rank == 1 + assert t.output_rank == 2 + + def test_validation_input_dimension_out_of_range(self) -> None: + domain = IndexDomain.from_shape((10,)) + maps = (DimensionMap(input_dimension=5),) + with pytest.raises(ValueError, match="input_dimension"): + IndexTransform(domain=domain, output=maps) + + +class TestIndexTransformBasicIndexing: + def test_slice_identity(self) -> None: + """slice(None) on identity transform is a no-op.""" + t = IndexTransform.from_shape((10, 20)) + result = t[slice(None), slice(None)] + assert result.domain.shape == (10, 20) + assert result.input_rank == 2 + assert result.output_rank == 2 + + def test_slice_narrows(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[2:8, 5:15] + # Domains are preserved (TensorStore): the slice keeps its literal + # coordinates, so the map stays the identity (out = in). + assert result.domain.shape == (6, 10) + assert result.domain.origin == (2, 5) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + assert result.output[0].input_dimension == 0 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == 0 + assert result.output[1].input_dimension == 1 + + def test_strided_slice(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t[::2] + assert result.domain.shape == (5,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 2 + + def test_strided_slice_with_start(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t[1:9:3] + # indices: 1, 4, 7 -> 3 elements + assert result.domain.shape == (3,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 1 + assert result.output[0].stride == 3 + + def test_int_drops_dimension(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[3] + assert result.input_rank == 1 + assert result.output_rank == 2 + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 3 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 0 + + def test_int_middle_dimension(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + result = t[:, 5, :] + assert result.input_rank == 2 + assert result.output_rank == 3 + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].input_dimension == 0 + assert isinstance(result.output[1], ConstantMap) + assert result.output[1].offset == 5 + assert isinstance(result.output[2], DimensionMap) + assert result.output[2].input_dimension == 1 + + def test_ellipsis(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + result = t[2:8, ...] + assert result.input_rank == 3 + assert result.domain.shape == (6, 20, 30) + + def test_newaxis(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[np.newaxis, :, :] + assert result.input_rank == 3 + assert result.domain.shape == (1, 10, 20) + assert result.output_rank == 2 + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].input_dimension == 1 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 2 + + def test_int_out_of_bounds(self) -> None: + t = IndexTransform.from_shape((10,)) + with pytest.raises(IndexError): + t[10] + + def test_negative_int_is_literal(self) -> None: + """Negative indices are literal coordinates (TensorStore convention), + not 'from the end' like NumPy.""" + t = IndexTransform.from_shape((10,)) + with pytest.raises(IndexError): + t[-1] # -1 is out of bounds for domain [0, 10) + + def test_negative_int_valid_with_negative_origin(self) -> None: + """Negative index is valid if the domain includes negative coordinates.""" + domain = IndexDomain(inclusive_min=(-5,), exclusive_max=(5,)) + t = IndexTransform.identity(domain) + result = t[-3] + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == -3 + + def test_composition_of_slices(self) -> None: + """Slicing a sliced transform re-selects in literal domain coordinates.""" + t = IndexTransform.from_shape((100,)) + result = t[10:50][15:30] + assert result.domain.shape == (15,) + assert result.domain.origin == (15,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + + def test_composition_of_strides(self) -> None: + t = IndexTransform.from_shape((100,)) + result = t[::2][::3] + # t[::2] -> shape (50,), offset=0, stride=2 + # [::3] -> shape ceil(50/3)=17, offset=0, stride=2*3=6 + assert result.domain.shape == (17,) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].stride == 6 + + def test_bare_int(self) -> None: + """Non-tuple selection.""" + t = IndexTransform.from_shape((10, 20)) + result = t[3] + assert result.input_rank == 1 + + def test_bare_slice(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t[2:8] + assert result.domain.shape == (6, 20) + + +class TestBasicIndexingOnArrayMaps: + """When a transform already has ArrayMap outputs, basic indexing must + apply the corresponding operation to the index_array's axes.""" + + def test_int_on_array_map_drops_axis(self) -> None: + """Integer index on a dimension referenced by an ArrayMap should + index into the array on that axis.""" + arr = np.array([[10, 20], [30, 40], [50, 60]], dtype=np.intp) + # 2D input domain (3, 2), one ArrayMap output + t = IndexTransform( + domain=IndexDomain.from_shape((3, 2)), + output=(ArrayMap(index_array=arr),), + ) + # Index with int on dim 0 -> pick row 1 -> arr[1, :] = [30, 40] + result = t[1] + assert result.input_rank == 1 + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([30, 40])) + + def test_slice_on_array_map(self) -> None: + """Slice on a dimension referenced by an ArrayMap should slice the array.""" + arr = np.array([10, 20, 30, 40, 50], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[1:4] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([20, 30, 40])) + + def test_strided_slice_on_array_map(self) -> None: + """Strided slice on ArrayMap should stride the array.""" + arr = np.array([10, 20, 30, 40, 50], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[::2] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, np.array([10, 30, 50])) + + def test_newaxis_on_array_map(self) -> None: + """Newaxis should insert an axis in the index_array.""" + arr = np.array([10, 20, 30], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr),), + ) + result = t[np.newaxis, :] + assert result.input_rank == 2 + assert result.domain.shape == (1, 3) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].index_array.shape == (1, 3) + np.testing.assert_array_equal(result.output[0].index_array, np.array([[10, 20, 30]])) + + def test_int_drops_one_of_two_array_dims(self) -> None: + """2D array map, int on dim 0, slice on dim 1.""" + arr = np.array([[10, 20, 30], [40, 50, 60]], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((2, 3)), + output=(ArrayMap(index_array=arr),), + ) + result = t[0, 1:3] + assert result.input_rank == 1 + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + # arr[0, 1:3] = [20, 30] + np.testing.assert_array_equal(result.output[0].index_array, np.array([20, 30])) + + +class TestIndexTransformOindex: + def test_oindex_int_array(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.oindex[idx, :] + assert result.input_rank == 2 + assert result.domain.shape == (3, 20) + assert isinstance(result.output[0], ArrayMap) + # Full input rank: the array varies along its own axis (0), singleton on 1. + assert result.output[0].index_array.shape == (3, 1) + np.testing.assert_array_equal(result.output[0].index_array, idx.reshape(3, 1)) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].input_dimension == 1 + + def test_oindex_bool_array(self) -> None: + t = IndexTransform.from_shape((5,)) + mask = np.array([True, False, True, False, True]) + result = t.oindex[mask] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal( + result.output[0].index_array, np.array([0, 2, 4], dtype=np.intp) + ) + + def test_oindex_mixed(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([2, 4], dtype=np.intp) + result = t.oindex[idx, 5:15] + assert result.input_rank == 2 + assert result.domain.shape == (2, 10) + # fancy dim: fresh zero-origin; slice dim: preserved literal coords + assert result.domain.origin == (0, 5) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == 0 + + def test_oindex_multiple_arrays(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + idx0 = np.array([1, 3], dtype=np.intp) + idx1 = np.array([5, 10, 15], dtype=np.intp) + result = t.oindex[idx0, :, idx1] + assert result.input_rank == 3 + assert result.domain.shape == (2, 20, 3) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], DimensionMap) + assert isinstance(result.output[2], ArrayMap) + + def test_oindex_multiple_arrays_preserves_independent_axes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.oindex[np.array([1, 3]), np.array([2, 4, 6])] + assert result.domain.shape == (2, 3) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + assert result.output[0].index_array.shape == (2, 1) + assert result.output[1].index_array.shape == (1, 3) + + +class TestIndexTransformVindex: + def test_vindex_single_array(self) -> None: + t = IndexTransform.from_shape((10,)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.vindex[idx] + assert result.input_rank == 1 + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, idx) + + def test_vindex_broadcast(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([[1, 2], [3, 4]], dtype=np.intp) + idx1 = np.array([[10, 11], [12, 13]], dtype=np.intp) + result = t.vindex[idx0, idx1] + assert result.input_rank == 2 + assert result.domain.shape == (2, 2) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + np.testing.assert_array_equal(result.output[0].index_array, idx0) + np.testing.assert_array_equal(result.output[1].index_array, idx1) + + def test_vindex_with_slice(self) -> None: + t = IndexTransform.from_shape((10, 20, 30)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = t.vindex[idx, :, :] + assert result.input_rank == 3 + assert result.domain.shape == (3, 20, 30) + assert isinstance(result.output[0], ArrayMap) + + def test_vindex_bool_mask(self) -> None: + t = IndexTransform.from_shape((5,)) + mask = np.array([True, False, True, False, True]) + result = t.vindex[mask] + assert result.domain.shape == (3,) + assert isinstance(result.output[0], ArrayMap) + + def test_vindex_broadcast_different_shapes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([1, 2, 3], dtype=np.intp) + idx1 = np.array([[10], [11]], dtype=np.intp) + result = t.vindex[idx0, idx1] + assert result.input_rank == 2 + assert result.domain.shape == (2, 3) + + def test_vindex_multiple_arrays_preserves_shared_axes(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.vindex[np.array([1, 3]), np.array([2, 4])] + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + assert result.output[0].index_array.shape == (2,) + assert result.output[1].index_array.shape == (2,) + + +class TestSelectionToTransform: + def test_basic_slice(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform((slice(2, 8), slice(5, 15)), t, "basic") + assert result.domain.shape == (6, 10) + assert result.domain.origin == (2, 5) # preserved literal coordinates + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + + def test_basic_int(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform((3, slice(None)), t, "basic") + assert result.input_rank == 1 + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 3 + + def test_basic_ellipsis(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = selection_to_transform(Ellipsis, t, "basic") + assert result.domain.shape == (10, 20) + + def test_orthogonal(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx = np.array([1, 3, 5], dtype=np.intp) + result = selection_to_transform((idx, slice(None)), t, "orthogonal") + assert result.domain.shape == (3, 20) + assert isinstance(result.output[0], ArrayMap) + + def test_vectorized(self) -> None: + t = IndexTransform.from_shape((10, 20)) + idx0 = np.array([1, 3], dtype=np.intp) + idx1 = np.array([5, 7], dtype=np.intp) + result = selection_to_transform((idx0, idx1), t, "vectorized") + assert result.domain.shape == (2,) + assert isinstance(result.output[0], ArrayMap) + assert isinstance(result.output[1], ArrayMap) + + def test_composition_with_non_identity(self) -> None: + """Indexing a sliced transform uses literal domain coordinates. + + The slice [10:50] preserves its domain, so a follow-up [15:30] + re-selects coordinates 15..29 of the base (TensorStore semantics), and + the composed map stays the identity (out = in). + """ + t = IndexTransform.from_shape((100,))[10:50] + result = selection_to_transform(slice(15, 30), t, "basic") + assert (result.domain.inclusive_min, result.domain.exclusive_max) == ((15,), (30,)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == 0 + assert result.output[0].stride == 1 + + +class TestIndexTransformIntersect: + def test_constant_inside(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.intersect(IndexDomain(inclusive_min=(0,), exclusive_max=(10,))) + assert result is not None + restricted, surviving = result + assert isinstance(restricted.output[0], ConstantMap) + assert restricted.output[0].offset == 5 + assert surviving is None + + def test_constant_outside(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.intersect(IndexDomain(inclusive_min=(10,), exclusive_max=(20,))) + assert result is None + + def test_dimension_partial(self) -> None: + """DimensionMap over [0,10) intersected with [5,15) narrows input to [5,10).""" + t = IndexTransform.from_shape((10,)) + result = t.intersect(IndexDomain(inclusive_min=(5,), exclusive_max=(15,))) + assert result is not None + restricted, surviving = result + assert restricted.domain.inclusive_min == (5,) + assert restricted.domain.exclusive_max == (10,) + assert surviving is None + + def test_dimension_no_overlap(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t.intersect(IndexDomain(inclusive_min=(20,), exclusive_max=(30,))) + assert result is None + + def test_dimension_strided(self) -> None: + """stride=2, offset=1 over [0,5): storage 1,3,5,7,9. Chunk [4,8).""" + t = IndexTransform( + domain=IndexDomain.from_shape((5,)), + output=(DimensionMap(input_dimension=0, offset=1, stride=2),), + ) + result = t.intersect(IndexDomain(inclusive_min=(4,), exclusive_max=(8,))) + assert result is not None + restricted, _surviving = result + # input 2->5, input 3->7. Both in [4,8). + assert restricted.domain.inclusive_min == (2,) + assert restricted.domain.exclusive_max == (4,) + + def test_array_partial(self) -> None: + arr = np.array([3, 8, 15, 22], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((4,)), + output=(ArrayMap(index_array=arr),), + ) + result = t.intersect(IndexDomain(inclusive_min=(5,), exclusive_max=(20,))) + assert result is not None + restricted, surviving = result + assert isinstance(restricted.output[0], ArrayMap) + np.testing.assert_array_equal(restricted.output[0].index_array, np.array([8, 15])) + assert surviving is not None + np.testing.assert_array_equal(surviving, np.array([1, 2])) + + def test_array_none_inside(self) -> None: + arr = np.array([1, 2, 3], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((3,)), + output=(ArrayMap(index_array=arr),), + ) + assert t.intersect(IndexDomain(inclusive_min=(10,), exclusive_max=(20,))) is None + + def test_2d_mixed(self) -> None: + """2D: ConstantMap on dim 0, DimensionMap on dim 1.""" + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=( + ConstantMap(offset=5), + DimensionMap(input_dimension=0, offset=0, stride=1), + ), + ) + chunk = IndexDomain(inclusive_min=(0, 5), exclusive_max=(10, 15)) + result = t.intersect(chunk) + assert result is not None + restricted, _ = result + assert isinstance(restricted.output[0], ConstantMap) + assert restricted.output[0].offset == 5 + assert isinstance(restricted.output[1], DimensionMap) + assert restricted.domain.inclusive_min == (5,) + assert restricted.domain.exclusive_max == (10,) + + +class TestIndexTransformTranslate: + def test_translate_constant(self) -> None: + t = IndexTransform( + domain=IndexDomain.from_shape((10,)), + output=(ConstantMap(offset=5),), + ) + result = t.translate((-5,)) + assert isinstance(result.output[0], ConstantMap) + assert result.output[0].offset == 0 + + def test_translate_dimension(self) -> None: + t = IndexTransform.from_shape((10,)) + result = t.translate((-3,)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == -3 + assert result.output[0].stride == 1 + + def test_translate_array(self) -> None: + arr = np.array([5, 10], dtype=np.intp) + t = IndexTransform( + domain=IndexDomain.from_shape((2,)), + output=(ArrayMap(index_array=arr, offset=3),), + ) + result = t.translate((-3,)) + assert isinstance(result.output[0], ArrayMap) + assert result.output[0].offset == 0 + np.testing.assert_array_equal(result.output[0].index_array, arr) + + def test_translate_2d(self) -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.translate((-5, -10)) + assert isinstance(result.output[0], DimensionMap) + assert result.output[0].offset == -5 + assert isinstance(result.output[1], DimensionMap) + assert result.output[1].offset == -10 + + +class TestArrayMapDependencyAxes: + """`_array_map_dependency_axes` derives the input axes an array varies on + from its (full-rank) shape: non-singleton axes vary, singleton axes do not.""" + + def test_orthogonal_single_axis(self) -> None: + from zarr_indexing.transform import _array_map_dependency_axes + + t = IndexTransform.from_shape((10, 20)).oindex[np.array([1, 3]), np.array([2, 4, 6])] + m0, m1 = t.output[0], t.output[1] + assert isinstance(m0, ArrayMap) + assert isinstance(m1, ArrayMap) + assert _array_map_dependency_axes(m0.index_array) == (0,) + assert _array_map_dependency_axes(m1.index_array) == (1,) + + def test_vectorized_shares_axes(self) -> None: + from zarr_indexing.transform import _array_map_dependency_axes + + t = IndexTransform.from_shape((10, 20)).vindex[np.array([1, 3]), np.array([2, 4])] + m0, m1 = t.output[0], t.output[1] + assert isinstance(m0, ArrayMap) + assert isinstance(m1, ArrayMap) + assert _array_map_dependency_axes(m0.index_array) == (0,) + assert _array_map_dependency_axes(m1.index_array) == (0,) + + def test_scalar_array_has_no_dependency(self) -> None: + from zarr_indexing.transform import _array_map_dependency_axes + + assert _array_map_dependency_axes(np.ones((1, 1), dtype=np.intp)) == () + + +class TestIntersectArrayMapClassification: + """`_intersect` must distinguish orthogonal (outer-product) ArrayMaps from + correlated (vectorized) ones by their dependency axes, keep surviving arrays + at full input rank, and preserve residual (slice) dimensions.""" + + def test_orthogonal_outer_product_keeps_full_rank(self) -> None: + """Two arrays on distinct axes narrow independently and stay full rank; + out_indices is a per-output-dim dict of surviving positions.""" + t = IndexTransform.from_shape((10, 10)).oindex[np.array([1, 3, 8]), np.array([2, 6, 9])] + # Chunk covering storage [0,5) x [0,5): rows 1,3 survive (out pos 0,1), + # cols 2 survives (out pos 0). + chunk = IndexDomain(inclusive_min=(0, 0), exclusive_max=(5, 5)) + result = t.intersect(chunk) + assert result is not None + restricted, out_indices = result + assert isinstance(restricted.output[0], ArrayMap) + assert isinstance(restricted.output[1], ArrayMap) + # Full input rank preserved (not raveled to 1-D). + assert restricted.output[0].index_array.ndim == 2 + assert restricted.output[1].index_array.ndim == 2 + assert restricted.domain.ndim == 2 + assert isinstance(out_indices, dict) + np.testing.assert_array_equal(out_indices[0], np.array([0, 1])) + np.testing.assert_array_equal(out_indices[1], np.array([0])) + + def test_correlated_with_residual_slice_preserves_slice_dim(self) -> None: + """A vindex transform with two correlated arrays plus a residual slice + dim intersects without a rank error and keeps the DimensionMap.""" + t = IndexTransform.from_shape((4, 3, 5)).vindex[np.array([1, 3]), np.array([2, 0])] + # Chunk covering storage [0,2) x [2,3) x [0,5): only point (1,2,*) is in + # bounds on both array dims -> one surviving broadcast point. + chunk = IndexDomain(inclusive_min=(0, 2, 0), exclusive_max=(2, 3, 5)) + result = t.intersect(chunk) + assert result is not None + restricted, out_indices = result + # A DimensionMap for the residual slice dim survives (no post-init error). + assert any(isinstance(m, DimensionMap) for m in restricted.output) + assert out_indices is not None + + def test_length1_orthogonal_not_treated_as_correlated(self) -> None: + """A length-1 orthogonal array (all-singleton shape) is still an outer + product with the length-3 axis: out_indices is a dict, not a flat array.""" + t = IndexTransform.from_shape((6, 6)).oindex[np.array([2]), np.array([1, 3, 5])] + chunk = IndexDomain(inclusive_min=(0, 0), exclusive_max=(6, 6)) + result = t.intersect(chunk) + assert result is not None + _restricted, out_indices = result + assert isinstance(out_indices, dict) diff --git a/pyproject.toml b/pyproject.toml index 727071b5a3..7ad2d3839b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ 'google-crc32c>=1.5', 'typing_extensions>=4.14', 'donfig>=0.8', + 'zarr-indexing>=0.1.0', ] dynamic = [ @@ -116,6 +117,11 @@ remote-tests = [ "moto[s3,server]==5.2.2", "requests==2.34.2", ] +interop-tests = [ + # Third-party array consumers whose duck-typing contract we test against; + # left floating so the `optional` matrix exercises their latest releases. + "dask[array]", +] release = [ "towncrier==25.8.0", ] @@ -146,12 +152,31 @@ dev = [ "mypy==2.1.0", ] +# The repo is a uv workspace so that `zarr` can depend on the in-tree +# `packages/zarr-indexing` during development and CI. Released wheels depend on +# the published zarr-indexing distribution; the workspace source below only +# affects resolution from a checkout (uv.lock, `uv run`, hatch envs using uv). +[tool.uv.workspace] +members = ["packages/zarr-indexing"] + +[tool.uv.sources] +zarr-indexing = { workspace = true } + [tool.coverage.report] exclude_also = [ 'if TYPE_CHECKING:', ] [tool.coverage.run] +# Measured source lives in two trees: zarr's own `src` and the in-tree +# `zarr-indexing` workspace package. Configured here (not via `--source=src` +# on the CLI) so every `coverage run` invocation — the hatch run-coverage / +# run-coverage-html / run-hypothesis scripts and ad-hoc local runs — measures +# both consistently. +source = [ + "src", + "packages/zarr-indexing/src", +] omit = [ "bench/compress_normal.py", ] @@ -160,9 +185,10 @@ omit = [ version.source = "vcs" # Only consider zarr-python's own `v*` tags when deriving the version. Without # this filter `git describe` matches the most recent tag of any shape, -# including the `zarr_metadata-v*` tags used to release the zarr-metadata -# subpackage — which would make a from-source build report a `0.2.x` version -# instead of `3.x`. +# including the `zarr_metadata-v*` and `zarr_indexing-v*` tags used to release +# the in-tree subpackages — which would make a from-source build report a +# `0.x` subpackage version instead of `3.x`. (`--match v*` excludes both, since +# neither tag starts with a bare `v`.) version.raw-options = { git_describe_command = "git describe --dirty --tags --long --match v*" } [tool.hatch.build] @@ -189,21 +215,22 @@ matrix.deps.features = [ ] matrix.deps.dependency-groups = [ {value = "remote-tests", if = ["optional"]}, + {value = "interop-tests", if = ["optional"]}, ] [tool.hatch.envs.test.scripts] run-coverage = [ - "coverage run --source=src -m pytest --ignore tests/benchmarks --junitxml=junit.xml -o junit_family=legacy {args:}", + "coverage run -m pytest --ignore tests/benchmarks --junitxml=junit.xml -o junit_family=legacy {args:}", "coverage xml", ] run-coverage-html = [ - "coverage run --source=src -m pytest --ignore tests/benchmarks {args:}", + "coverage run -m pytest --ignore tests/benchmarks {args:}", "coverage html", ] run = "pytest --ignore tests/benchmarks" run-verbose = "run-coverage --verbose" run-hypothesis = [ - "coverage run --source=src -m pytest -nauto --run-slow-hypothesis tests/test_properties.py tests/test_store/test_stateful* {args:}", + "coverage run -m pytest -nauto --run-slow-hypothesis tests/test_properties.py tests/test_store/test_stateful* {args:}", "coverage xml", ] run-benchmark = "pytest --benchmark-enable tests/benchmarks" @@ -226,7 +253,7 @@ python = ["3.12", "3.13"] [tool.hatch.envs.gputest.scripts] run-coverage = [ - "coverage run --source=src -m pytest -m gpu --junitxml=junit.xml -o junit_family=legacy --ignore tests/benchmarks {args:}", + "coverage run -m pytest -m gpu --junitxml=junit.xml -o junit_family=legacy --ignore tests/benchmarks {args:}", "coverage xml", ] run = "pytest -m gpu --ignore tests/benchmarks" @@ -427,7 +454,12 @@ ignore_errors = true [tool.pytest.ini_options] minversion = "9" -testpaths = ["src", "tests", "docs/user-guide"] +# `packages/zarr-indexing/tests` is collected here rather than run in an +# isolated package env: those tests exercise chunk resolution against zarr's +# concrete `ChunkGrid`, so they need `zarr` importable (which the root env +# provides). Keeping them in testpaths means the main test suite / `run-coverage` +# continues to exercise the transform algebra after the move to the subpackage. +testpaths = ["src", "tests", "docs/user-guide", "packages/zarr-indexing/tests"] log_cli_level = "INFO" log_level = "INFO" # Enables strict_config, strict_markers, strict_xfail, strict_parametrization_ids, and diff --git a/src/zarr/__init__.py b/src/zarr/__init__.py index cdf3840c3b..cbda7ad0ae 100644 --- a/src/zarr/__init__.py +++ b/src/zarr/__init__.py @@ -35,6 +35,7 @@ zeros_like, ) from zarr.core.array import Array, AsyncArray +from zarr.core.chunk_partition import ChunkProjection from zarr.core.config import config from zarr.core.group import AsyncGroup, Group @@ -146,6 +147,7 @@ def set_format(log_format: str) -> None: "Array", "AsyncArray", "AsyncGroup", + "ChunkProjection", "Group", "__version__", "array", diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index f75ef72415..69971c190f 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -19,6 +19,12 @@ import numpy as np from typing_extensions import Sentinel, deprecated +from zarr_indexing.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr_indexing.output_map import ArrayMap +from zarr_indexing.transform import ( + IndexTransform, + selection_to_transform, +) import zarr from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec @@ -42,6 +48,7 @@ from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, + ChunkSpec, _is_rectilinear_chunks, as_regular_shape, guess_chunks, @@ -55,6 +62,7 @@ V2ChunkKeyEncoding, parse_chunk_key_encoding, ) +from zarr.core.chunk_partition import ChunkProjection, iter_chunk_projections from zarr.core.common import ( JSON, ZARR_JSON, @@ -104,10 +112,14 @@ _iter_regions, check_fields, check_no_multi_fields, + ensure_tuple, + is_basic_selection, + is_coordinate_selection, is_pure_fancy_indexing, is_pure_orthogonal_indexing, is_scalar, pop_fields, + replace_lists, ) from zarr.core.metadata import ( ArrayMetadata, @@ -133,6 +145,7 @@ from zarr.errors import ( ArrayNotFoundError, ChunkNotFoundError, + LazyViewError, MetadataValidationError, ZarrDeprecationWarning, ZarrUserWarning, @@ -370,6 +383,7 @@ class AsyncArray[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: store_path: StorePath codec_pipeline: CodecPipeline = field(init=False) _chunk_grid: ChunkGrid = field(init=False) + _transform: IndexTransform = field(init=False) config: ArrayConfig @overload @@ -406,6 +420,7 @@ def __init__( "codec_pipeline", create_codec_pipeline(metadata=metadata_parsed, store=store_path.store), ) + object.__setattr__(self, "_transform", IndexTransform.from_shape(metadata_parsed.shape)) @classmethod async def _create( @@ -808,6 +823,35 @@ async def example(): _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict) return cls(store_path=store_path, metadata=_metadata_dict) + def _with_transform(self, transform: IndexTransform) -> AsyncArray[T_ArrayMetadata]: + """Return a new AsyncArray sharing storage but with a different transform.""" + new = object.__new__(type(self)) + object.__setattr__(new, "metadata", self.metadata) + object.__setattr__(new, "store_path", self.store_path) + object.__setattr__(new, "config", self.config) + object.__setattr__(new, "_chunk_grid", self._chunk_grid) + object.__setattr__(new, "codec_pipeline", self.codec_pipeline) + object.__setattr__(new, "_transform", transform) + return new + + def translate_by(self, shift: tuple[int, ...]) -> AsyncArray[T_ArrayMetadata]: + """Shift this array's domain by `shift`, preserving which cells it addresses. + + TensorStore's `translate_by`: the view's coordinate labels move; the data + does not. `a.translate_by((-10,))` gives a view whose domain starts at + -10, where coordinate -10 addresses the cell that 0 addressed before. + """ + return self._with_transform(self._transform.translate_domain_by(tuple(shift))) + + def translate_to(self, origins: tuple[int, ...]) -> AsyncArray[T_ArrayMetadata]: + """Move this array's domain so its per-dimension origins equal `origins`. + + TensorStore's `translate_to`; `view.translate_to((0,) * view.ndim)` + re-zeros a view's coordinate system without changing which cells it + addresses. + """ + return self._with_transform(self._transform.translate_domain_to(tuple(origins))) + @property def store(self) -> Store: return self.store_path.store @@ -826,7 +870,7 @@ def ndim(self) -> int: int The number of dimensions in the Array. """ - return len(self.metadata.shape) + return len(self.shape) @property def shape(self) -> tuple[int, ...]: @@ -837,6 +881,114 @@ def shape(self) -> tuple[int, ...]: tuple The shape of the Array. """ + return self._transform.domain.shape + + @property + def _is_identity(self) -> bool: + """Whether this array's transform is the identity over the full storage domain. + + A freshly-opened or resized array has the identity transform: input coord + `i` maps to storage coord `i` over the whole storage shape. Eager + indexing on such an array produces the same coordinates the legacy + indexers compute, so it can take the legacy fast path and skip + transform resolution. Lazy views (created via `_with_transform`) + carry a non-identity transform and must go through the transform path. + Cheap (O(ndim)); the domain's shape lookup it relies on is memoized. + """ + return _transform_is_identity(self._transform, self.metadata.shape) + + def _require_identity(self, name: str) -> None: + """Raise `LazyViewError` if this array is a non-identity lazy view. + + Grid-describing/mutating members assume the array fills its chunk grid, + which a view (a sliced/indexed array) generally does not. See + `LazyViewError`; `name` is the member being guarded, for the message. + """ + if not self._is_identity: + raise LazyViewError( + f"`{name}` is not defined for a lazy view (a sliced or indexed array that " + f"does not fill its chunk grid). Use `chunk_projections` for this view's " + f"chunk granularity; the backing array's stored structure is available via " + f"`metadata` / `chunk_grid`." + ) + + def _require_identity_for_async(self, name: str) -> None: + """Raise `LazyViewError` if this async entry point is used on a lazy view. + + The async surface does not yet route selections through a view's index + transform: it builds a legacy indexer from `metadata.shape`, which for a + non-identity view would silently read or write the *wrong* region. Until + async transform routing lands, guard these entry points and steer users + to the synchronous surface. `name` is the member being guarded. + """ + if not self._is_identity: + raise LazyViewError( + f"`{name}` on the async surface does not yet support lazy views " + f"(sliced or indexed arrays with a non-identity transform); it would " + f"silently read or write the wrong region. Use the synchronous `Array` " + f"surface, or materialize the view first with `.result()` (for example, " + f"index `view.result()` instead of `await view.async_array...`)." + ) + + @property + def _is_sharded(self) -> bool: + """Whether the array stores inner chunks inside shards (a sharding codec). + + A sharding codec is the array-bytes codec in the chain: array-array + filters may precede it and bytes-bytes codecs (e.g. an outer compressor) + may follow it, so its presence — not the chain's length — is what makes + an array sharded. `validate_codecs` permits such trailing codecs, and + they round-trip data correctly; keying off `len(codecs) == 1` would + misclassify them as unsharded. + """ + from zarr.codecs.sharding import ShardingCodec + + codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) + return any(isinstance(codec, ShardingCodec) for codec in codecs) + + def chunk_projections( + self, *, unit: Literal["read", "write"] = "read" + ) -> Iterator[ChunkProjection]: + """Enumerate the stored chunks this array (or lazy view) projects onto. + + Yields a `ChunkProjection` per stored chunk: its coordinate, store key, and + (extent-clipped) shape; the region of the chunk this array covers; the region + of this array it maps to; and whether the coverage is partial (a partial + write is a read-modify-write). For an identity array every chunk is fully + covered and the projections tile the whole domain; for a view only the touched + chunks appear. + + `unit` selects the granularity: `"write"` is the store-object grid (the shard + when sharded, else the chunk); `"read"` is the chunk grid. They coincide + unless the array is sharded. Read-unit (inner-chunk) partitioning of a sharded + array is not yet implemented; use `unit="write"` there. + + To partition an arbitrary selection, compose through the lazy accessor: + `array.lazy[sel].chunk_projections()`. + """ + if unit not in ("read", "write"): + raise ValueError(f"unit must be 'read' or 'write', got {unit!r}") + if unit == "read" and self._is_sharded: + raise NotImplementedError( + "read-unit (inner-chunk) `chunk_projections` for sharded arrays is not yet " + "implemented; use `unit='write'` for shard-granularity projections." + ) + return iter_chunk_projections( + self._transform, self._chunk_grid, self.metadata.encode_chunk_key + ) + + def is_chunk_aligned(self) -> bool: + """Whether this array/view aligns to write-unit (store-object) boundaries. + + True iff no stored write unit is only partially covered — i.e. every unit can + be written without a read-modify-write. A cheap wrapper over + `chunk_projections`. + """ + return all(not p.is_partial for p in self.chunk_projections(unit="write")) + + @property + def storage_shape(self) -> tuple[int, ...]: + """The shape of the underlying storage array (ignoring any view transform).""" return self.metadata.shape @property @@ -852,6 +1004,7 @@ def chunks(self) -> tuple[int, ...]: tuple[int, ...]: The chunk shape of the Array. """ + self._require_identity("chunks") # TODO: move sharding awareness out of metadata return self.metadata.chunks @@ -878,6 +1031,7 @@ def read_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: >>> arr.read_chunk_sizes ((30, 30, 30, 10), (40, 40)) """ + self._require_identity("read_chunk_sizes") from zarr.codecs.sharding import ShardingCodec @@ -909,7 +1063,7 @@ def write_chunk_sizes(self) -> tuple[tuple[int, ...], ...]: >>> arr.write_chunk_sizes ((30, 30, 30, 10), (40, 40)) """ - + self._require_identity("write_chunk_sizes") return self._chunk_grid.chunk_sizes @property @@ -925,6 +1079,7 @@ def shards(self) -> tuple[int, ...] | None: tuple[int, ...]: The shard shape of the Array. """ + self._require_identity("shards") return self.metadata.shards @property @@ -936,7 +1091,9 @@ def size(self) -> int: int Total number of elements in the array """ - return math.prod(self.metadata.shape) + # `self.shape` is the view's logical shape (== metadata.shape for a + # non-view); using it keeps `size`/`nbytes` correct for lazy views. + return math.prod(self.shape) @property def filters(self) -> tuple[Numcodec, ...] | tuple[ArrayArrayCodec, ...]: @@ -1101,6 +1258,7 @@ def cdata_shape(self) -> tuple[int, ...]: tuple[int, ...] The number of chunks along each dimension. """ + self._require_identity("cdata_shape") return self._chunk_grid_shape @property @@ -1157,6 +1315,7 @@ def nchunks(self) -> int: int The total number of chunks in the array. """ + self._require_identity("nchunks") return product(self._chunk_grid_shape) @property @@ -1241,6 +1400,7 @@ async def example(): result = asyncio.run(example()) ``` """ + self._require_identity("nchunks_initialized") return await _nchunks_initialized(self) async def _nshards_initialized(self) -> int: @@ -1282,6 +1442,7 @@ async def example(): return await _nshards_initialized(self) async def nbytes_stored(self) -> int: + self._require_identity("nbytes_stored") return await _nbytes_stored(self.store_path) def _iter_chunk_coords( @@ -1500,7 +1661,7 @@ async def getitem( >>> asyncio.run(example()) np.int32(0) """ - + self._require_identity_for_async("getitem") return await _getitem( self.store_path, self.metadata, @@ -1519,6 +1680,7 @@ async def get_orthogonal_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: + self._require_identity_for_async("get_orthogonal_selection") if prototype is None: prototype = default_buffer_prototype() indexer = OrthogonalIndexer(selection, self.metadata.shape, self._chunk_grid) @@ -1542,6 +1704,7 @@ async def get_mask_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: + self._require_identity_for_async("get_mask_selection") if prototype is None: prototype = default_buffer_prototype() indexer = MaskIndexer(mask, self.metadata.shape, self._chunk_grid) @@ -1565,6 +1728,7 @@ async def get_coordinate_selection( fields: Fields | None = None, prototype: BufferPrototype | None = None, ) -> NDArrayLikeOrScalar: + self._require_identity_for_async("get_coordinate_selection") if prototype is None: prototype = default_buffer_prototype() indexer = CoordinateIndexer(selection, self.metadata.shape, self._chunk_grid) @@ -1610,6 +1774,42 @@ async def _set_selection( fields=fields, ) + async def _get_selection_t( + self, + transform: IndexTransform, + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, + ) -> NDArrayLikeOrScalar: + return await _get_selection_via_transform( + self.store_path, + self.metadata, + self.config, + transform, + self.codec_pipeline, + prototype=prototype, + out=out, + chunk_grid=self._chunk_grid, + ) + + async def _set_selection_t( + self, + transform: IndexTransform, + value: npt.ArrayLike, + *, + prototype: BufferPrototype, + ) -> None: + return await _set_selection_via_transform( + self.store_path, + self.metadata, + self.config, + transform, + value, + self.codec_pipeline, + prototype=prototype, + chunk_grid=self._chunk_grid, + ) + async def setitem( self, selection: BasicSelection, @@ -1649,6 +1849,7 @@ async def setitem( - This method is asynchronous and should be awaited. - Supports basic indexing, where the selection is contiguous and does not involve advanced indexing. """ + self._require_identity_for_async("setitem") return await _setitem( self.store_path, self.metadata, @@ -1700,6 +1901,7 @@ async def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True) ----- - This method is asynchronous and should be awaited. """ + self._require_identity("resize") return await _resize(self, new_shape, delete_outside_chunks) async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: @@ -1721,6 +1923,7 @@ async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: The size of all dimensions other than `axis` must match between this array and `data`. """ + self._require_identity("append") return await _append(self, data, axis) async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: @@ -1753,7 +1956,12 @@ async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: return self def __repr__(self) -> str: - return f"" + base = f"" + return f"{base} domain={self._transform.selection_repr}>" @property def info(self) -> Any: @@ -1792,6 +2000,7 @@ def info(self) -> Any: Compressors : (ZstdCodec(level=0, checksum=False),) No. bytes : 480 """ + self._require_identity("info") return self._info() async def info_complete(self) -> Any: @@ -1811,6 +2020,7 @@ async def info_complete(self) -> Any: ------- [zarr.AsyncArray.info][] - A property giving just the statically known information about an array. """ + self._require_identity("info_complete") return await _info_complete(self) def _info( @@ -1873,6 +2083,49 @@ def _chunk_grid(self) -> ChunkGrid: """The chunk grid for this array, bound to the array's shape.""" return self.async_array._chunk_grid + def _with_transform(self, transform: IndexTransform) -> Array[T_ArrayMetadata]: + """Return a new Array sharing storage but with a different transform.""" + new_async = self._async_array._with_transform(transform) + return type(self)(new_async) + + def translate_by(self, shift: tuple[int, ...]) -> Array[T_ArrayMetadata]: + """Shift this array's domain by `shift`, preserving which cells it addresses. + + TensorStore's `translate_by`: the view's coordinate labels move; the data + does not. `a.translate_by((-10,))` gives a view whose domain starts at + -10, where coordinate -10 addresses the cell that 0 addressed before. + """ + return self._with_transform(self._async_array._transform.translate_domain_by(tuple(shift))) + + def translate_to(self, origins: tuple[int, ...]) -> Array[T_ArrayMetadata]: + """Move this array's domain so its per-dimension origins equal `origins`. + + TensorStore's `translate_to`; `view.translate_to((0,) * view.ndim)` + re-zeros a view's coordinate system without changing which cells it + addresses. + """ + return self._with_transform( + self._async_array._transform.translate_domain_to(tuple(origins)) + ) + + def __iter__(self) -> Any: + """Iterate over the first axis (identity arrays only). + + Lazy views are not iterable, matching TensorStore: the Python iteration + protocol counts positions from 0, which are not valid coordinates in a + preserved (possibly non-zero-origin) domain — silently yielding nothing + or the wrong cells. Read the view first: `iter(view.result())`. + """ + if not self._async_array._is_identity: + raise TypeError( + "lazy views are not iterable; materialize first, e.g. iterate `view.result()`" + ) + if self.ndim == 0: + raise TypeError("iteration over a 0-d array") + # A plain generator function would defer these raises to the first + # next(); returning an inner generator keeps them eager at iter(). + return (self[i] for i in range(self.shape[0])) + @classmethod def _create( cls, @@ -2208,7 +2461,11 @@ def cdata_shape(self) -> tuple[int, ...]: When sharding is used, this counts inner chunks (not shards) per dimension. """ - return self._chunk_grid_shape + # Deliberately NOT main's `self._chunk_grid_shape` delegation: the public + # property must keep routing through the async `cdata_shape`, which + # carries the lazy-view guard (`_require_identity`); the private + # `_chunk_grid_shape` is the unguarded internal member. + return self.async_array.cdata_shape @property def _chunk_grid_shape(self) -> tuple[int, ...]: @@ -2286,6 +2543,27 @@ def nbytes(self) -> int: """ return self.async_array.nbytes + def chunk_projections( + self, *, unit: Literal["read", "write"] = "read" + ) -> Iterator[ChunkProjection]: + """Enumerate the stored chunks this array (or lazy view) projects onto. + + See [zarr.AsyncArray.chunk_projections][] for the full description. Each + `ChunkProjection` reports a stored chunk's coordinate/key/shape, the region of + it this array covers, the region of this array it maps to, and whether the + coverage is partial. Compose through `lazy` to partition an arbitrary + selection: `array.lazy[sel].chunk_projections()`. + """ + return self.async_array.chunk_projections(unit=unit) + + def is_chunk_aligned(self) -> bool: + """Whether this array/view aligns to write-unit (store-object) boundaries. + + True iff no stored write unit is only partially covered. See + [zarr.AsyncArray.is_chunk_aligned][]. + """ + return self.async_array.is_chunk_aligned() + @property def nchunks_initialized(self) -> int: """ @@ -2747,6 +3025,104 @@ def __setitem__(self, selection: Selection, value: npt.ArrayLike) -> None: else: self.set_basic_selection(cast("BasicSelection", pure_selection), value, fields=fields) + def _reject_bool_scalar(self, sel: Any) -> None: + # `isinstance(True, int)` is True, so a bare Python/NumPy bool scalar would + # otherwise pass the integer validators and silently read element 0/1 where + # eager NumPy indexing raises. Boolean *arrays* (masks) remain valid. + if isinstance(sel, (bool, np.bool_)): + raise IndexError( + "boolean scalar indices are not supported; use a boolean array " + "(mask) for boolean indexing" + ) + + def _normalize_public_selection(self, selection: Any) -> Any: + """Normalize a user selection at the public Array boundary. + + Rejects boolean *scalar* indices, validates boolean-mask shapes, and wraps + negative integers, negative index-array elements, and negative slice bounds + to literal coordinates in the current view's domain: an index `k` with + `-size <= k < 0` on axis `d` becomes `exclusive_max[d] + k` (from the end of + the view's domain — exactly NumPy for a zero-origin domain). Positive indices + are left as literal domain coordinates. A boolean mask consumes `mask.ndim` + dimensions and must exactly match the view domain's shape on them (NumPy + parity), else IndexError — never a silent truncation. The transform layer + stays literal; it performs the final bounds check, so an out-of-range value + (a positive past the end, or `k < -size`) is rejected there before any chunk + access. + """ + domain = self._async_array._transform.domain + exclusive_max = domain.exclusive_max + ndim = len(exclusive_max) + + is_tuple = isinstance(selection, tuple) + items = list(selection) if is_tuple else [selection] + + for item in items: + self._reject_bool_scalar(item) + + # A boolean mask consumes one dimension per mask axis; everything else + # consumes exactly one (newaxis/Ellipsis consume none). + n_real = 0 + has_mask = False + for item in items: + if item is Ellipsis or item is None: + continue + mask = _as_bool_mask(item) + if mask is not None: + has_mask = True + n_real += mask.ndim + else: + n_real += 1 + if n_real > ndim: + if has_mask: + # The transform layer would raise an unhelpful ValueError for a + # too-high-rank mask; raise the canonical IndexError here. + raise IndexError( + f"too many indices for array: array has {ndim} dimension(s), " + f"but {n_real} were indexed" + ) + # Otherwise leave the selection untouched and let the transform layer + # raise its canonical error with the correct message. + return selection + + normalized: list[Any] = [] + axis = 0 + for item in items: + if item is Ellipsis: + axis += ndim - n_real # ellipsis fills the unaddressed axes + normalized.append(item) + elif item is None: + normalized.append(item) # newaxis consumes no domain axis + else: + mask = _as_bool_mask(item) + if mask is not None: + expected = domain.shape[axis : axis + mask.ndim] + for offset, (want, got) in enumerate(zip(expected, mask.shape, strict=True)): + if want != got: + raise IndexError( + "boolean index did not match indexed array along " + f"dimension {axis + offset}; dimension is {want} but " + f"corresponding boolean dimension is {got}" + ) + normalized.append(item) + axis += mask.ndim + else: + normalized.append(_wrap_negative_index(item, exclusive_max[axis])) + axis += 1 + return tuple(normalized) if is_tuple else normalized[0] + + def _reject_fields_on_view(self, fields: Fields | None) -> None: + # Routing `fields=` on a non-identity view through the legacy indexer + # ignores the transform and reads/writes the wrong storage region. The + # transform layer cannot preserve field-aware semantics, so reject before + # any storage access rather than corrupt data. + if fields is not None and not self._async_array._is_identity: + raise NotImplementedError( + "field selection (`fields=`) is not supported on a lazy view " + "(non-identity transform); materialize the view first " + "(e.g. zarr.array(view[...]))" + ) + def get_basic_selection( self, selection: BasicSelection = Ellipsis, @@ -2869,14 +3245,21 @@ def get_basic_selection( if prototype is None: prototype = default_buffer_prototype() - return sync( - self.async_array._get_selection( - BasicIndexer(selection, self.shape, self._chunk_grid), - out=out, - fields=fields, - prototype=prototype, + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + # Eager (identity-transform) arrays and structured-dtype field + # selection use the original indexer path directly. + return sync( + self.async_array._get_selection( + BasicIndexer(selection, self.shape, self._chunk_grid), + out=out, + fields=fields, + prototype=prototype, + ) ) - ) + selection = self._normalize_public_selection(selection) + transform = selection_to_transform(selection, self._async_array._transform, "basic") + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_basic_selection( self, @@ -2978,8 +3361,18 @@ def set_basic_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = BasicIndexer(selection, self.shape, self._chunk_grid) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + # Eager (identity-transform) arrays and structured-dtype field + # selection use the original indexer path directly. + indexer = BasicIndexer(selection, self.shape, self._chunk_grid) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) + ) + return + selection = self._normalize_public_selection(selection) + transform = selection_to_transform(selection, self._async_array._transform, "basic") + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_orthogonal_selection( self, @@ -3106,12 +3499,25 @@ def get_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) - return sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + # Eager (identity) arrays and structured-dtype field selection use the + # original indexer path directly. + indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) + return sync( + self.async_array._get_selection( + indexer=indexer, out=out, fields=fields, prototype=prototype + ) ) + # Lazy view (non-identity transform): route through the transform so the + # view's offset/stride are honored. Use basic mode for plain int/slice + # selections (so integer axes drop), orthogonal mode for fancy selections. + selection = self._normalize_public_selection(selection) + mode: Literal["basic", "orthogonal"] = ( + "basic" if is_basic_selection(selection) else "orthogonal" ) + transform = selection_to_transform(selection, self._async_array._transform, mode) + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_orthogonal_selection( self, @@ -3224,10 +3630,24 @@ def set_orthogonal_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) - return sync( - self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + # Eager (identity) arrays and structured-dtype field selection use the + # original indexer path directly. + indexer = OrthogonalIndexer(selection, self.shape, self._chunk_grid) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) + ) + return + # Lazy view (non-identity transform): route through the transform so the + # view's offset/stride are honored. Use basic mode for plain int/slice + # selections (so integer axes drop), orthogonal mode for fancy selections. + selection = self._normalize_public_selection(selection) + mode: Literal["basic", "orthogonal"] = ( + "basic" if is_basic_selection(selection) else "orthogonal" ) + transform = selection_to_transform(selection, self._async_array._transform, mode) + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_mask_selection( self, @@ -3312,12 +3732,29 @@ def get_mask_selection( if prototype is None: prototype = default_buffer_prototype() - indexer = MaskIndexer(mask, self.shape, self._chunk_grid) - return sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + indexer = MaskIndexer(mask, self.shape, self._chunk_grid) + return sync( + self.async_array._get_selection( + indexer=indexer, out=out, fields=fields, prototype=prototype + ) ) - ) + # Unwrap if VIndex passed a tuple + if isinstance(mask, tuple) and len(mask) == 1: + mask = mask[0] + # Validate mask + mask_arr = np.asarray(mask) + if mask_arr.dtype != np.bool_: + raise IndexError("invalid mask selection; expected Boolean array") + if mask_arr.shape != self.shape: + raise IndexError( + f"invalid mask selection; expected Boolean array with shape {self.shape}, " + f"got {mask_arr.shape}" + ) + selection = (mask_arr,) + transform = selection_to_transform(selection, self._async_array._transform, "vectorized") + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_mask_selection( self, @@ -3401,8 +3838,26 @@ def set_mask_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = MaskIndexer(mask, self.shape, self._chunk_grid) - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + indexer = MaskIndexer(mask, self.shape, self._chunk_grid) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) + ) + return + if isinstance(mask, tuple) and len(mask) == 1: + mask = mask[0] + mask_arr = np.asarray(mask) + if mask_arr.dtype != np.bool_: + raise IndexError("invalid mask selection; expected Boolean array") + if mask_arr.shape != self.shape: + raise IndexError( + f"invalid mask selection; expected Boolean array with shape {self.shape}, " + f"got {mask_arr.shape}" + ) + selection = (mask_arr,) + transform = selection_to_transform(selection, self._async_array._transform, "vectorized") + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_coordinate_selection( self, @@ -3420,7 +3875,14 @@ def get_coordinate_selection( selection : tuple An integer (coordinate) array for each dimension of the array. out : NDBuffer, optional - If given, load the selected data directly into this buffer. + If given, load the selected data directly into this buffer. The expected + shape of `out` currently differs between eager arrays and lazy views: + on an ordinary (identity) array `out` must be a flat buffer of shape + `(n,)`, where `n` is the number of selected points; on a lazy view `out` + must have the broadcast selection shape (the shape the coordinate arrays + broadcast to). The broadcast-shape contract is the long-term direction; + the flat contract on eager arrays is retained for backwards + compatibility and is expected to be aligned to it in a future release. fields : str or sequence of str, optional For arrays with a structured dtype, one or more fields can be specified to extract data for. @@ -3489,17 +3951,40 @@ def get_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) - out_array = sync( - self.async_array._get_selection( - indexer=indexer, out=out, fields=fields, prototype=prototype + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) + out_array = sync( + self.async_array._get_selection( + indexer=indexer, out=out, fields=fields, prototype=prototype + ) ) + if hasattr(out_array, "shape"): + out_array = np.array(out_array).reshape(indexer.sel_shape) + return out_array + # Validate and normalize as coordinate selection + sel_normalized = ensure_tuple(selection) + sel_normalized = tuple( + np.asarray([s], dtype=np.intp) if isinstance(s, (int, np.integer)) else s + for s in sel_normalized ) - - if hasattr(out_array, "shape"): - # restore shape - out_array = np.array(out_array).reshape(indexer.sel_shape) - return out_array + sel_normalized = replace_lists(sel_normalized) + if not is_coordinate_selection(sel_normalized, self.shape): + raise IndexError( + "invalid coordinate selection; expected one integer " + "(coordinate) array per dimension of the target array, " + f"got {selection!r}" + ) + # Wrap negative coordinates against the view's domain (see + # _normalize_public_selection); positive coordinates stay literal. + sel_normalized = self._normalize_public_selection(sel_normalized) + transform = selection_to_transform( + sel_normalized, self._async_array._transform, "vectorized" + ) + # `_get_selection_via_transform` already returns the broadcast selection + # shape (transform.domain.shape == the broadcast shape of the coordinate + # arrays), so no further reshape/copy is needed here. + return sync(self._async_array._get_selection_t(transform, out=out, prototype=prototype)) def set_coordinate_selection( self, @@ -3580,30 +4065,57 @@ def set_coordinate_selection( """ if prototype is None: prototype = default_buffer_prototype() - # setup indexer - indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) - - # handle value - need ndarray-like flatten value - if not is_scalar(value, self.dtype): - try: - from numcodecs.compat import ensure_ndarray_like - - value = ensure_ndarray_like(value) # TODO replace with agnostic - except TypeError: - # Handle types like `list` or `tuple` - value = np.array(value) # TODO replace with agnostic - if hasattr(value, "shape") and len(value.shape) > 1: - value = np.array(value).reshape(-1) - - if not is_scalar(value, self.dtype) and ( - isinstance(value, NDArrayLike) and indexer.shape != value.shape - ): - raise ValueError( - f"Attempting to set a selection of {indexer.sel_shape[0]} " - f"elements with an array of {value.shape[0]} elements." + # Normalize a caller-supplied empty fields list/tuple to None, so that + # `fields is not None` uniformly means "a field was requested" (matching + # pop_fields, which yields None for no fields); keep the exact-emptiness + # check rather than a falsy one. + if fields == [] or fields == (): + fields = None + self._reject_fields_on_view(fields) + if fields is not None or self._async_array._is_identity: + indexer = CoordinateIndexer(selection, self.shape, self._chunk_grid) + if not is_scalar(value, self.dtype): + try: + from numcodecs.compat import ensure_ndarray_like + + value = ensure_ndarray_like(value) + except TypeError: + value = np.array(value) + if hasattr(value, "shape") and len(value.shape) > 1: + value = np.array(value).reshape(-1) + if not is_scalar(value, self.dtype) and ( + isinstance(value, NDArrayLike) and indexer.shape != value.shape + ): + raise ValueError( + f"Attempting to set a selection of {indexer.sel_shape[0]} " + f"elements with an array of {value.shape[0]} elements." + ) + sync( + self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype) ) - - sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) + return + sel_normalized = ensure_tuple(selection) + sel_normalized = tuple( + np.asarray([s], dtype=np.intp) if isinstance(s, (int, np.integer)) else s + for s in sel_normalized + ) + sel_normalized = replace_lists(sel_normalized) + if not is_coordinate_selection(sel_normalized, self.shape): + raise IndexError( + "invalid coordinate selection; expected one integer " + "(coordinate) array per dimension of the target array, " + f"got {selection!r}" + ) + # Wrap negative coordinates against the view's domain (see + # _normalize_public_selection); positive coordinates stay literal. + sel_normalized = self._normalize_public_selection(sel_normalized) + transform = selection_to_transform( + sel_normalized, self._async_array._transform, "vectorized" + ) + # Flatten value for coordinate selection + if not is_scalar(value, self.dtype) and hasattr(value, "shape") and len(value.shape) > 1: + value = np.asarray(value).reshape(-1) + sync(self._async_array._set_selection_t(transform, value, prototype=prototype)) def get_block_selection( self, @@ -3702,6 +4214,14 @@ def get_block_selection( """ if prototype is None: prototype = default_buffer_prototype() + if not self._async_array._is_identity: + # Block selection addresses the storage chunk grid; it is ill-defined + # on a lazy view (offset/stride) and the legacy indexer would read raw + # storage blocks, ignoring the transform. Reject rather than corrupt. + raise NotImplementedError( + "block selection is not supported on a lazy view; materialize the " + "view first (e.g. zarr.array(view[...]))" + ) indexer = BlockIndexer(selection, self.shape, self._chunk_grid) return sync( self.async_array._get_selection( @@ -3802,6 +4322,13 @@ def set_block_selection( """ if prototype is None: prototype = default_buffer_prototype() + if not self._async_array._is_identity: + # See get_block_selection: block selection is ill-defined on a lazy + # view, so reject rather than write to the wrong storage region. + raise NotImplementedError( + "block selection is not supported on a lazy view; materialize the " + "view first (e.g. zarr.array(view[...]))" + ) indexer = BlockIndexer(selection, self.shape, self._chunk_grid) sync(self.async_array._set_selection(indexer, value, fields=fields, prototype=prototype)) @@ -3831,6 +4358,19 @@ def blocks(self) -> BlockIndex: examples.""" return BlockIndex(self) + @property + def lazy(self) -> _LazyIndexAccessor: + """Lazy indexing accessor. Returns a new Array with composed transform, no I/O.""" + return _LazyIndexAccessor(self) + + def result(self, prototype: BufferPrototype | None = None) -> NDArrayLikeOrScalar: + """Read and return the data for this array view. + + Equivalent to `self[...]` but more explicit for lazy views, and forwards + `prototype`. Named after `tensorstore.Future.result`. + """ + return self.get_basic_selection(Ellipsis, prototype=prototype) + def resize(self, new_shape: ShapeLike) -> None: """ Change the shape of the array by growing or shrinking one or more @@ -3934,7 +4474,12 @@ def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: return type(self)(new_array) def __repr__(self) -> str: - return f"" + base = f"" + return f"{base} domain={self._async_array._transform.selection_repr}>" @property def info(self) -> Any: @@ -4051,6 +4596,116 @@ async def _shards_initialized( type SerializerLike = dict[str, JSON] | ArrayBytesCodec | Literal["auto"] +def _as_bool_mask(item: Any) -> np.ndarray[Any, np.dtype[np.bool_]] | None: + """Return `item` as a boolean ndarray if it is a boolean mask (ndarray or + list), else None. Boolean *scalars* are not masks (rejected separately).""" + if isinstance(item, (list, np.ndarray)): + arr = np.asarray(item) + if arr.dtype == np.bool_: + return arr + return None + + +def _wrap_negative_index(item: Any, exclusive_max: int) -> Any: + """Wrap negative values in a single per-axis selection element to literal + coordinates against `exclusive_max` (from-the-end within the axis's domain). + + A negative integer `k` becomes `exclusive_max + k`; negative slice bounds and + negative integer-array elements are wrapped the same way. Positive values, + `None` slice bounds, boolean masks, and non-integer items pass through + unchanged. This performs no bounds checking — the transform layer rejects an + out-of-range value (a positive past the end, or one that wraps below the + domain origin) before any chunk access. + """ + if isinstance(item, slice): + start = item.start + stop = item.stop + if start is not None and start < 0: + start += exclusive_max + if stop is not None and stop < 0: + stop += exclusive_max + return slice(start, stop, item.step) + if isinstance(item, (int, np.integer)): + k = int(item) + return k + exclusive_max if k < 0 else k + if isinstance(item, (list, np.ndarray)): + arr = np.asarray(item) + if arr.dtype == np.bool_ or not np.issubdtype(arr.dtype, np.integer): + return item # boolean mask or non-integer array: no wrapping + return np.where(arr < 0, arr + exclusive_max, arr) + return item + + +class _LazyOIndex: + """Lazy orthogonal indexing via `array.lazy.oindex[...]`.""" + + __slots__ = ("_array",) + + def __init__(self, array: Array[Any]) -> None: + self._array = array + + def __getitem__(self, selection: Any) -> Array[Any]: + selection = self._array._normalize_public_selection(selection) + new_t = selection_to_transform(selection, self._array._async_array._transform, "orthogonal") + return self._array._with_transform(new_t) + + def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: + selection = self._array._normalize_public_selection(selection) + new_t = selection_to_transform(selection, self._array._async_array._transform, "orthogonal") + self._array._with_transform(new_t)[...] = value + + +class _LazyVIndex: + """Lazy vectorized indexing via `array.lazy.vindex[...]`.""" + + __slots__ = ("_array",) + + def __init__(self, array: Array[Any]) -> None: + self._array = array + + def __getitem__(self, selection: Any) -> Array[Any]: + selection = self._array._normalize_public_selection(selection) + new_t = selection_to_transform(selection, self._array._async_array._transform, "vectorized") + return self._array._with_transform(new_t) + + def __setitem__(self, selection: Any, value: npt.ArrayLike) -> None: + selection = self._array._normalize_public_selection(selection) + new_t = selection_to_transform(selection, self._array._async_array._transform, "vectorized") + self._array._with_transform(new_t)[...] = value + + +class _LazyIndexAccessor: + """Provides lazy indexing via `array.lazy[...]`.""" + + __slots__ = ("_array",) + + def __init__(self, array: Array[Any]) -> None: + self._array = array + + def __getitem__(self, selection: Selection) -> Array[Any]: + selection = self._array._normalize_public_selection(selection) + new_t = selection_to_transform(selection, self._array._async_array._transform, "basic") + return self._array._with_transform(new_t) + + def __setitem__(self, selection: Selection, value: npt.ArrayLike) -> None: + selection = self._array._normalize_public_selection(selection) + new_t = selection_to_transform(selection, self._array._async_array._transform, "basic") + self._array._with_transform(new_t)[...] = value + + @property + def shape(self) -> tuple[int, ...]: + """Shape of the array (or view) this accessor indexes into.""" + return self._array.shape + + @property + def oindex(self) -> _LazyOIndex: + return _LazyOIndex(self._array) + + @property + def vindex(self) -> _LazyVIndex: + return _LazyVIndex(self._array) + + class ShardsConfigParam(TypedDict): shape: tuple[int, ...] index_location: IndexLocation | None @@ -4256,6 +4911,23 @@ async def from_array( [0, 0]]) """ mode: Literal["a"] = "a" + # Reject lazy-view sources up front, before touching the store or creating the + # target array. The copy loop reads via `data.async_array.getitem`, which does + # not yet route through a view's index transform (it would copy the wrong + # region); guard here so the user gets one clear error and no half-created + # target is left behind. + _source_async: AsyncArray[Any] | None = None + if isinstance(data, Array): + _source_async = data._async_array + elif isinstance(data, AsyncArray): + _source_async = data + if _source_async is not None and not _source_async._is_identity: + raise LazyViewError( + "Creating an array from a lazy view (a sliced or indexed array with a " + "non-identity transform) via `from_array` is not yet supported. " + "Materialize the view first, e.g. `zarr.from_array(store, data=view.result())`, " + "or use the synchronous copy idiom." + ) config_parsed = parse_array_config(config) store_path = await make_store_path(store, path=name, mode=mode, storage_options=storage_options) @@ -5364,17 +6036,18 @@ async def _nbytes_stored( return await store_path.store.getsize_prefix(store_path.path) -def _get_chunk_spec( +def _array_spec_from_chunk_spec( metadata: ArrayMetadata, - chunk_grid: ChunkGrid, - chunk_coords: tuple[int, ...], + spec: ChunkSpec, array_config: ArrayConfig, prototype: BufferPrototype, ) -> ArraySpec: - """Build an ArraySpec for a single chunk using the ChunkGrid.""" - spec = chunk_grid[chunk_coords] - if spec is None: - raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") + """Build an ArraySpec from an already-resolved ChunkSpec. + + Split out from `_get_chunk_spec` so the transform read/write path can + resolve `chunk_grid[chunk_coords]` once per chunk and feed the same + `spec` to both this and `_is_complete_chunk`. + """ return ArraySpec( shape=spec.codec_shape, dtype=metadata.dtype, @@ -5384,6 +6057,338 @@ def _get_chunk_spec( ) +def _get_chunk_spec( + metadata: ArrayMetadata, + chunk_grid: ChunkGrid, + chunk_coords: tuple[int, ...], + array_config: ArrayConfig, + prototype: BufferPrototype, +) -> ArraySpec: + """Build an ArraySpec for a single chunk using the ChunkGrid.""" + spec = chunk_grid[chunk_coords] + if spec is None: + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") + return _array_spec_from_chunk_spec(metadata, spec, array_config, prototype) + + +def _is_vectorized_transform(transform: IndexTransform) -> bool: + """Return True for a vectorized (vindex/coordinate) transform. + + Coordinate ArrayMaps (`input_dimension is None`) are jointly indexed over + the broadcast input domain and scatter through a single flat index, so the + output buffer is flattened during read/write and reshaped afterwards. This + holds whenever *any* such map is present — including a single coordinate map + with a multi-dimensional index array (`vindex[idx_2d]` on a 1-D array), + whose flat result still needs a flat buffer. Orthogonal ArrayMaps (oindex), + each bound to a distinct input dimension, form an outer product and keep a + multi-dimensional buffer — even when every output is an ArrayMap + (e.g. `oindex[i0, i1]`). + """ + return any(isinstance(m, ArrayMap) and m.input_dimension is None for m in transform.output) + + +def _transform_is_identity(transform: IndexTransform, storage_shape: tuple[int, ...]) -> bool: + """Return True if `transform` is the identity over the full storage domain. + + An identity transform maps input coordinate `i` to storage coordinate `i` + across the array's whole storage shape (origin 0, unit stride, dimensions in + order). Such an array is an ordinary eager array — indexing it produces the + same coordinates the legacy indexers compute, so the legacy fast path is + safe. Any narrowing, striding, reordering, or fancy selection (i.e. a lazy + view) yields a non-identity transform that must go through the transform + resolver. Cheap: O(ndim), no array work. + """ + from zarr_indexing.output_map import DimensionMap + + domain = transform.domain + ndim = len(storage_shape) + if domain.ndim != ndim or len(transform.output) != ndim: + return False + if domain.inclusive_min != (0,) * ndim or domain.exclusive_max != storage_shape: + return False + for i, m in enumerate(transform.output): + if not ( + type(m) is DimensionMap and m.input_dimension == i and m.offset == 0 and m.stride == 1 + ): + return False + return True + + +def _is_complete_chunk(sub_transform: IndexTransform, spec: ChunkSpec) -> bool: + """Check if a sub-transform covers an entire chunk. + + `spec` is the chunk's already-resolved `ChunkSpec` (the caller looks + it up once and shares it with `_array_spec_from_chunk_spec`). + """ + from zarr_indexing.output_map import ConstantMap, DimensionMap + + shape = spec.shape + for out_dim, m in enumerate(sub_transform.output): + if isinstance(m, ConstantMap): + # A ConstantMap means a single element is selected along this output dimension, + # so the write does not cover the full chunk along this dimension. + chunk_dim_size = shape[out_dim] + if chunk_dim_size > 1: + return False + continue # chunk dim size is 1, so selecting the single element is complete + if isinstance(m, DimensionMap): + chunk_dim_size = shape[out_dim] + # Compute actual storage range: storage = offset + stride * input_coord + dim_lo = sub_transform.domain.inclusive_min[m.input_dimension] + dim_hi = sub_transform.domain.exclusive_max[m.input_dimension] + if m.stride == 1: + storage_start = m.offset + dim_lo + storage_stop = m.offset + dim_hi + if storage_start != 0 or storage_stop != chunk_dim_size: + return False + else: + return False # strided access is never a complete chunk + else: + return False # ArrayMap is never a "complete chunk" + return True + + +async def _get_selection_via_transform( + store_path: StorePath, + metadata: ArrayMetadata, + config: ArrayConfig, + transform: IndexTransform, + codec_pipeline: CodecPipeline, + *, + prototype: BufferPrototype, + out: NDBuffer | None = None, + chunk_grid: ChunkGrid | None = None, +) -> NDArrayLikeOrScalar: + """Read data using an IndexTransform.""" + # The user-facing transform may carry a preserved (possibly non-zero-origin) + # domain; the I/O layer addresses output buffers positionally, so normalize + # to a zero-origin domain here. Translation preserves the cell mapping. + transform = transform.translate_domain_to((0,) * transform.input_rank) + if chunk_grid is None: + chunk_grid = ChunkGrid.from_metadata(metadata) + + # Get dtype (same logic as existing _get_selection) + if metadata.zarr_format == 2: + zdtype = metadata.dtype + order = metadata.order + else: + zdtype = metadata.data_type + order = config.order + dtype = zdtype.to_native_dtype() + + out_shape = transform.domain.shape + + # Vectorized indexing (any correlated ArrayMap, input_dimension None — even a + # single one, e.g. vindex[idx_2d]) scatters through a single flat index, so + # the working buffer must be 1D during the read and reshaped to out_shape + # afterwards. Orthogonal indexing — including the case where every output is + # an ArrayMap bound to a distinct input dimension (oindex[i0, i1]) — keeps a + # multi-dimensional buffer, one out_sel entry per dim. A zero-rank domain + # still needs a flat (1,) working buffer for the scatter, but its single cell + # maps to a scalar on return (see the out_shape == () branch below). + needs_flat_buffer = _is_vectorized_transform(transform) + buffer_shape = (product(out_shape),) if needs_flat_buffer else out_shape + + # Setup output buffer. The caller-visible `out` contract is the broadcast + # selection shape (out_shape). Vectorized reads still scatter through a flat + # index internally, so when `out` is supplied for a vectorized selection we + # read into a flat temporary and copy the reshaped result into `out` below. + copy_into_out: NDBuffer | None = None + if out is not None: + if not isinstance(out, NDBuffer): + raise TypeError(f"out argument needs to be an NDBuffer. Got {type(out)!r}") + if out.shape != out_shape: + raise ValueError( + f"shape of out argument doesn't match. Expected {out_shape}, got {out.shape}" + ) + if needs_flat_buffer: + copy_into_out = out + out_buffer = prototype.nd_buffer.empty(shape=buffer_shape, dtype=dtype, order=order) + else: + out_buffer = out + else: + out_buffer = prototype.nd_buffer.empty(shape=buffer_shape, dtype=dtype, order=order) + + if product(out_shape) > 0: + _config = config + if metadata.zarr_format == 2: + _config = replace(_config, order=order) + + # Build batch_info using transforms + batch_info = [] + drop_axes: tuple[int, ...] = () + for chunk_coords, sub_transform, out_indices in iter_chunk_transforms( + transform, chunk_grid + ): + chunk_spec = chunk_grid[chunk_coords] + if chunk_spec is None: + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") + chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) + drop_axes = da # same for all chunks + batch_info.append( + ( + store_path / metadata.encode_chunk_key(chunk_coords), + _array_spec_from_chunk_spec(metadata, chunk_spec, _config, prototype), + chunk_sel, + out_sel, + _is_complete_chunk(sub_transform, chunk_spec), + ) + ) + + results = await codec_pipeline.read(batch_info, out_buffer, drop_axes=drop_axes) + + # Handle read_missing_chunks + if _config.read_missing_chunks is False: + missing_info = [] + for i, result in enumerate(results): + if result["status"] == "missing": + coords_path = batch_info[i][0] + missing_info.append(f" chunk at '{coords_path}'") + if len(missing_info) > 0: + chunks_str = "\n".join(missing_info) + raise ChunkNotFoundError( + f"{len(missing_info)} chunk(s) not found in store '{store_path}'.\n" + f"Set the 'array.read_missing_chunks' config to True to fill " + f"missing chunks with the fill value.\n" + f"Missing chunks:\n{chunks_str}" + ) + + # Return scalar for 0-d results. A correlated read used a flat (1,) working + # buffer for the scatter, so collapse it to a 0-d value before extracting. + if out_shape == (): + if copy_into_out is not None: + copy_into_out[...] = out_buffer.as_ndarray_like().reshape(()) + return copy_into_out.as_scalar() + # Mirror NDBuffer.as_scalar (GPU-safe conversion), but collapse the flat + # (1,) correlated working buffer to 0-d first. + return cast("NDArrayLikeOrScalar", out_buffer.as_numpy_array().reshape(())[()]) + out_result = out_buffer.as_ndarray_like() + # Reshape the flat working buffer back to the broadcast selection shape. The + # buffer is freshly allocated and contiguous, so reshape returns a view — no + # copy — and preserves the backend array type. + if needs_flat_buffer and hasattr(out_result, "reshape"): + out_result = out_result.reshape(out_shape) + # A vectorized read filled a flat temporary; copy the reshaped result into + # the caller-provided out buffer, whose shape is the broadcast selection shape. + if copy_into_out is not None: + copy_into_out[...] = out_result + return copy_into_out.as_ndarray_like() + return out_result + + +async def _set_selection_via_transform( + store_path: StorePath, + metadata: ArrayMetadata, + config: ArrayConfig, + transform: IndexTransform, + value: npt.ArrayLike, + codec_pipeline: CodecPipeline, + *, + prototype: BufferPrototype, + chunk_grid: ChunkGrid | None = None, +) -> None: + """Write data using an IndexTransform.""" + # See _get_selection_via_transform: normalize to a zero-origin domain + # so value/buffer placement is positional. + transform = transform.translate_domain_to((0,) * transform.input_rank) + if chunk_grid is None: + chunk_grid = ChunkGrid.from_metadata(metadata) + + # Get dtype from metadata + if metadata.zarr_format == 2: + zdtype = metadata.dtype + else: + zdtype = metadata.data_type + dtype = zdtype.to_native_dtype() + + # check value shape + if np.isscalar(value): + array_like = prototype.buffer.create_zero_length().as_array_like() + if isinstance(array_like, np._typing._SupportsArrayFunc): + array_like_ = cast("np._typing._SupportsArrayFunc", array_like) + value = np.asanyarray(value, dtype=dtype, like=array_like_) + else: + if not hasattr(value, "shape"): + value = np.asarray(value, dtype) + if not hasattr(value, "dtype") or value.dtype.name != dtype.name: + if hasattr(value, "astype"): + value = value.astype(dtype=dtype, order="A") + else: + value = np.array(value, dtype=dtype, order="A") + value = cast("NDArrayLike", value) + + # Validate value shape against selection shape + sel_shape = transform.domain.shape + needs_flat_buffer = _is_vectorized_transform(transform) + if hasattr(value, "shape") and value.shape != () and value.shape != sel_shape: + if needs_flat_buffer: + # For ArrayMap (coordinate/vindex), values are flattened so check total size + sel_size = product(sel_shape) + val_size = product(value.shape) + if val_size != sel_size and val_size != 1: + raise ValueError( + f"Attempting to set a selection with a value of incompatible shape. " + f"The selection has shape {sel_shape}, but the value has shape {value.shape}." + ) + else: + # Check if value is broadcastable to sel_shape + try: + np.broadcast_shapes(value.shape, sel_shape) + except ValueError: + raise ValueError( + f"Attempting to set a selection with a value of incompatible shape. " + f"The selection has shape {sel_shape}, but the value has shape {value.shape}." + ) from None + + if product(sel_shape) == 0: + return + + # When the transform has ArrayMap outputs, chunk resolution produces + # flat scatter indices (out_indices). The value buffer must be 1D + # during the write, matching the flat index layout. + if ( + needs_flat_buffer + and hasattr(value, "reshape") + and not np.isscalar(value) + and np.ndim(value) > 0 + ): + value = np.asarray(value).reshape(-1) + + # Convert to NDBuffer + value_buffer = prototype.nd_buffer.from_ndarray_like(value) + + # Determine memory order + if metadata.zarr_format == 2: + order = metadata.order + else: + order = config.order + + _config = config + if metadata.zarr_format == 2: + _config = replace(_config, order=order) + + # Build batch_info using transforms + batch_info = [] + drop_axes: tuple[int, ...] = () + for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, chunk_grid): + chunk_spec = chunk_grid[chunk_coords] + if chunk_spec is None: + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") + chunk_sel, out_sel, da = sub_transform_to_selections(sub_transform, out_indices) + drop_axes = da # same for all chunks + batch_info.append( + ( + store_path / metadata.encode_chunk_key(chunk_coords), + _array_spec_from_chunk_spec(metadata, chunk_spec, _config, prototype), + chunk_sel, + out_sel, + _is_complete_chunk(sub_transform, chunk_spec), + ) + ) + + await codec_pipeline.write(batch_info, value_buffer, drop_axes=drop_axes) + + async def _get_selection( store_path: StorePath, metadata: ArrayMetadata, @@ -5494,7 +6499,7 @@ async def _get_selection( coords = indexed_chunks[i][0] key = metadata.encode_chunk_key(coords) missing_info.append(f" chunk '{key}' (grid position {coords})") - if missing_info: + if len(missing_info) > 0: chunks_str = "\n".join(missing_info) raise ChunkNotFoundError( f"{len(missing_info)} chunk(s) not found in store '{store_path}'.\n" @@ -5765,9 +6770,10 @@ async def _delete_key(key: str) -> None: # Write new metadata await save_metadata(array.store_path, new_metadata) - # Update metadata and chunk_grid (in place) + # Update metadata, chunk_grid, and transform (in place) object.__setattr__(array, "metadata", new_metadata) object.__setattr__(array, "_chunk_grid", new_chunk_grid) + object.__setattr__(array, "_transform", IndexTransform.from_shape(new_shape)) async def _append( diff --git a/src/zarr/core/chunk_partition.py b/src/zarr/core/chunk_partition.py new file mode 100644 index 0000000000..5e054a2141 --- /dev/null +++ b/src/zarr/core/chunk_partition.py @@ -0,0 +1,117 @@ +"""View-aware chunk partitioning (Layer B of the lazy-view chunk-layout design). + +A lazy view is an `IndexTransform` applied to a backing array — i.e. a stored +selection already resolved. `chunk_projections` enumerates the stored chunks that +selection (or a whole identity array) projects onto, reusing the same resolution +machinery the read/write I/O path uses (`iter_chunk_transforms` + +`sub_transform_to_selections`). Each projection reports the stored chunk, the +region of it this array covers, the region of this array it maps to, and whether the +coverage is partial (a partial write is a read-modify-write). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from zarr_indexing.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from zarr_indexing.transform import IndexTransform + + from zarr.core.chunk_grids import ChunkGrid + +__all__ = ["ChunkProjection", "iter_chunk_projections"] + + +@dataclass(frozen=True) +class ChunkProjection: + """How one stored chunk contributes to an array (or lazy view). + + Attributes + ---------- + coord + Coordinate of the stored chunk in the chunk grid. + key + Store key of the stored chunk (which file). + shape + Stored size of the chunk, clipped to the array extent. + chunk_selection + The region *within* the stored chunk that this array covers. + array_selection + The region of this array (or view) that the chunk maps to. + is_partial + Whether the chunk is only partially covered — a write to a partial chunk + is a read-modify-write. + """ + + coord: tuple[int, ...] + key: str + shape: tuple[int, ...] + chunk_selection: tuple[Any, ...] + array_selection: tuple[Any, ...] + is_partial: bool + + +def _covers_full_chunk(chunk_selection: tuple[Any, ...], shape: tuple[int, ...]) -> bool: + """Whether `chunk_selection` selects every element of a chunk of `shape`. + + A per-dimension full `0:size:1` slice is a full cover. An integer entry is + also a full cover *iff* the chunk's extent along that dimension is 1 — fixing + the only element of a length-1 dimension covers it exactly, the same as an + equivalent length-1 slice would. This mirrors the write path's + `_is_complete_chunk` (`ConstantMap` over a dimension of size 1 is complete); + keeping the two in sync matters because a divergence would make + `chunk_projections` plan needless read-modify-writes for coverage the write + path already treats as complete. An array selection always touches a subset + (or requires per-element comparison this cheap check does not attempt), so it + is never a full cover. + """ + if len(chunk_selection) != len(shape): + return False + for sel, size in zip(chunk_selection, shape, strict=True): + if isinstance(sel, slice): + start, stop, step = sel.indices(size) + if not (start == 0 and stop == size and step == 1): + return False + elif isinstance(sel, int): + if size > 1: + return False + else: + return False + return True + + +def iter_chunk_projections( + transform: IndexTransform, + chunk_grid: ChunkGrid, + encode_key: Callable[[tuple[int, ...]], str], +) -> Iterator[ChunkProjection]: + """Yield a `ChunkProjection` for each stored chunk `transform` projects onto. + + `array_selection` is positional (0-based into the view's extent), so the + transform is normalized to a zero-origin domain first — translation + preserves which cells are addressed. + """ + transform = transform.translate_domain_to((0,) * transform.input_rank) + chunk_sizes = chunk_grid.chunk_sizes # per-dimension, extent-clipped + for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, chunk_grid): + if chunk_grid[chunk_coords] is None: + raise IndexError(f"Chunk coordinates {chunk_coords} are out of bounds.") + chunk_selection, array_selection, _drop_axes = sub_transform_to_selections( + sub_transform, out_indices + ) + shape = tuple(chunk_sizes[dim][c] for dim, c in enumerate(chunk_coords)) + yield ChunkProjection( + coord=chunk_coords, + key=encode_key(chunk_coords), + shape=shape, + chunk_selection=chunk_selection, + array_selection=array_selection, + is_partial=not _covers_full_chunk(chunk_selection, shape), + ) diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index f6eb495cd9..006dec36fc 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1029,6 +1029,7 @@ class AsyncOIndex[T_ArrayMetadata: (ArrayV2Metadata, ArrayV3Metadata)]: async def getitem(self, selection: OrthogonalSelection | AnyArray) -> NDArrayLikeOrScalar: from zarr.core.array import Array + self.array._require_identity_for_async("oindex.getitem") # if input is a Zarr array, we materialize it now. if isinstance(selection, Array): selection = _zarr_array_to_int_or_bool_array(selection) @@ -1377,6 +1378,7 @@ async def getitem( # TODO requires solving this circular sync issue: https://github.com/zarr-developers/zarr-python/pull/3083#discussion_r2230737448 from zarr.core.array import Array + self.array._require_identity_for_async("vindex.getitem") # if input is a Zarr array, we materialize it now. if isinstance(selection, Array): selection = _zarr_array_to_int_or_bool_array(selection) @@ -1442,8 +1444,16 @@ def pop_fields(selection: SelectionWithFields) -> tuple[Fields | None, Selection return None, cast("Selection", selection) else: # multiple items, split fields from selection items - fields: Fields = [f for f in selection if isinstance(f, str)] - fields = fields[0] if len(fields) == 1 else fields + field_names = [f for f in selection if isinstance(f, str)] + # No fields -> None (consistent with the non-tuple branch above), so callers + # can use `fields is not None` to mean "a field was requested". + fields: Fields | None + if len(field_names) == 0: + fields = None + elif len(field_names) == 1: + fields = field_names[0] + else: + fields = field_names selection_tuple = tuple(s for s in selection if not isinstance(s, str)) selection = cast( "Selection", selection_tuple[0] if len(selection_tuple) == 1 else selection_tuple diff --git a/src/zarr/errors.py b/src/zarr/errors.py index 781bebe534..e7f72a827e 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -1,3 +1,8 @@ +# `BoundsCheckError` and `VindexInvalidSelectionError` are defined canonically in +# `zarr_indexing.errors` (the transform algebra raises them) and re-exported here +# as aliases so `zarr.errors.BoundsCheckError` stays the identical class object. +from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError + __all__ = [ "ArrayIndexError", "ArrayNotFoundError", @@ -9,6 +14,7 @@ "ContainsGroupError", "DataTypeValidationError", "GroupNotFoundError", + "LazyViewError", "MetadataValidationError", "NegativeStepError", "NodeTypeValidationError", @@ -139,16 +145,26 @@ class ZarrRuntimeWarning(RuntimeWarning): """ -class VindexInvalidSelectionError(IndexError): ... - - class NegativeStepError(IndexError): ... -class BoundsCheckError(IndexError): ... +class ArrayIndexError(IndexError): ... -class ArrayIndexError(IndexError): ... +class LazyViewError(NotImplementedError, AttributeError): + """Raised when an operation that assumes an array fills its chunk grid is used + on a non-identity lazy view (created via `Array.lazy[...]`). + + Grid-describing members (`chunks`, `shards`, `nchunks`, `read_chunk_sizes`, + ...) and grid-mutating ones (`resize`, `append`) have no well-defined answer + for a view onto a subset of the backing grid. Use `chunk_projections` for the + view's granularity; the backing array's stored structure is available via + `metadata` / `chunk_grid`. Subclasses + `NotImplementedError` so existing consumers that catch it keep working, and + `AttributeError` so duck-typing probes (`hasattr(view, "chunks")`, + `getattr(x, "chunks", None)` — e.g. `dask.array.from_array`) treat guarded + members as absent instead of crashing. + """ class ChunkNotFoundError(BaseZarrError): diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index 99e81b0389..1de454cfbc 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -2,6 +2,7 @@ import math import sys from collections.abc import Callable, Mapping +from dataclasses import dataclass from typing import Any, Literal import hypothesis.extra.numpy as npst @@ -27,6 +28,7 @@ from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding from zarr.core.common import JSON, AccessModeLiteral, ZarrFormat from zarr.core.dtype import get_data_type_from_native_dtype +from zarr.core.indexing import Selection from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata from zarr.core.sync import sync @@ -612,6 +614,210 @@ def orthogonal_indices( return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer)) +IndexMode = Literal["basic", "oindex", "vindex", "mask"] + + +@st.composite +def windows(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[slice, ...]: + """A non-negative, full-rank tuple of slice windows — one per axis. + + A rank-preserving sub-region selection: each axis gets `start:stop` with + `0 <= start < stop <= size` (an empty `0:0` slice for a zero-length axis). + Bounds stay non-negative so the window is valid for any consumer, including + those that treat negative indices as literal coordinates rather than + from-the-end (e.g. building a sub-array view). + """ + out: list[slice] = [] + for size in shape: + if size == 0: + out.append(slice(0, 0)) + continue + start = draw(st.integers(min_value=0, max_value=size - 1)) + stop = draw(st.integers(min_value=start + 1, max_value=size)) + out.append(slice(start, stop)) + return tuple(out) + + +@st.composite +def numpy_array_indexers( + draw: st.DrawFn, *, mode: IndexMode, shape: tuple[int, ...] +) -> tuple[Selection, Selection]: + """A `(zarr_selection, numpy_selection)` pair for `mode` on `shape`. + + Scope: the *element-space* indexing modes that have a direct NumPy-array + equivalent, so a NumPy array of `shape` can serve as the correctness oracle. + One strategy covers them all, so a test can be written once and parametrized + over mode instead of re-deriving selection setup per test: + + - `"basic"` — slices / ints / ellipsis (no newaxis, no negative slices) + - `"oindex"` — per-axis integer arrays or slices (orthogonal / outer product) + - `"vindex"` — broadcast integer coordinate arrays (vectorized) + - `"mask"` — a boolean array of `shape` + + The two returned selections differ only for `oindex` (zarr's per-axis + spelling vs the `np.ix_`-style spelling numpy needs); for the other modes + the same object indexes both a zarr array and its numpy reference. The + array-based modes (`oindex`/`vindex`/`mask`) need `shape` to have no + zero-length axis; `basic` has no such requirement. + + Deliberately excluded is **block** indexing (`Array.blocks` / + `get_block_selection`): it addresses the *chunk grid*, not array elements, + so it is parametrized by the chunk grid rather than `shape` and has no + NumPy-array equivalent to compare against — its oracle is a coordinate + translation built by the separate `block_indices` strategy. + """ + if mode == "basic": + sel = draw(basic_indices(shape=shape)) + return sel, sel + if mode == "oindex": + return draw(orthogonal_indices(shape=shape)) + if mode == "vindex": + idx = draw( + npst.integer_array_indices( + shape=shape, result_shape=npst.array_shapes(min_side=1, max_dims=None) + ) + ) + return idx, idx + if mode == "mask": + m = draw(npst.arrays(dtype=np.bool_, shape=st.just(shape))) + return m, m + raise ValueError(f"unknown indexing mode: {mode!r}") + + +# --- Composed indexing programs ------------------------------------------------- +# +# A *program* is a short sequence of indexing operations composed on a single +# array, plus an execution recipe describing how to realize the composed +# selection (read whole, read a sub-selection eagerly, read into an ``out=`` +# buffer, or write). It exists to differentially test the index-transform layer: +# every operation composes a real transform, and the runner (in +# ``tests/test_properties.py``) checks the composed result against a NumPy oracle +# that applies the equivalent selections in lockstep. +# +# The mode vocabulary here is the *public dialect* ("orthogonal"/"vectorized"), +# not the internal accessor names ("oindex"/"vindex"); the runner maps between +# them. +ProgramMode = Literal["basic", "orthogonal", "vectorized"] +ProgramExecution = Literal["materialize", "eager_on_lazy", "out", "set_scalar", "set_array"] + + +@dataclass(frozen=True) +class IndexOperation: + """One indexing step in an `IndexingProgram`. + + ``selection`` carries exactly what the runner needs to reproduce the step on + *both* a zarr array and its NumPy reference: + + - ``"basic"`` — a single object used verbatim for both (slice / int / + ellipsis / tuple thereof). + - ``"orthogonal"`` — a ``(zarr_selection, numpy_selection)`` pair, where the + NumPy side is the ``np.ix_`` outer-product spelling of the per-axis + selection. + - ``"vectorized"`` — a single coordinate selection used verbatim for both + (plain NumPy fancy-indexing semantics). + """ + + mode: ProgramMode + selection: Any + + +@dataclass(frozen=True) +class IndexingProgram: + """A composed sequence of `IndexOperation` plus an execution recipe.""" + + operations: tuple[IndexOperation, ...] + execution: ProgramExecution + + +# Small arrays keep 300 Hypothesis examples cheap while still crossing chunk +# boundaries (the runner chunks each axis into ~2 pieces). ``min_side=0`` lets a +# program start from — or a basic step produce — a zero-extent axis, which is the +# required "empty result" coverage class. +program_shapes = npst.array_shapes(min_dims=1, max_dims=3, min_side=0, max_side=4) + + +def _basic_result_shape(shape: tuple[int, ...], selection: Any) -> tuple[int, ...]: + """The NumPy result shape of applying a *basic* ``selection`` to ``shape``. + + Used to carry the visible shape forward while generating a program, so each + step's selection is drawn against the shape the previous steps leave behind. + Uses a tiny placeholder array (program shapes are capped small). + """ + result_shape: tuple[int, ...] = np.empty(shape, dtype=np.int8)[selection].shape + return result_shape + + +@st.composite +def indexing_programs(draw: st.DrawFn, *, shape: tuple[int, ...]) -> IndexingProgram: + """Generate a composable indexing program over an array of ``shape``. + + Structure (one to three operations): a prefix of zero-to-two **basic** steps, + optionally followed by a single **fancy** (orthogonal/vectorized) step as the + *last* operation. The visible shape is carried forward so every step is valid + against what its predecessors produce. + + **Why fancy is confined to the last position** — the exclusions here are + deliberate and documented so they can shrink as the underlying bugs are + fixed: + + - *Fancy-after-fancy* composition (applying oindex/vindex to a view that + already carries an orthogonal ArrayMap axis) is unsupported: it raises a + clear ``NotImplementedError`` at composition time rather than resolving, so + at most one fancy step is generated. + - *Basic-after-fancy* is excluded wholesale. Integer basic indexing on an + oindex-picked axis is a strict xfail + (``TestKnownFancyIntBugs::test_int_read_on_oindex_view``); rather than track + which composed axes became "fancy-picked" and admit only the working + basic-on-non-fancy-axis subset, we keep fancy last and let the basic prefix + cover multi-step basic composition. (When those bugs are fixed this + restriction can be relaxed to interleave basic and fancy freely.) + + A fancy step is only emitted when the current shape has rank ``>= 1`` and no + zero-length axis (integer/orthogonal selections cannot address an empty axis); + otherwise the program stays basic-only. If nothing else was generated, a + single basic step is appended so a program always has at least one operation. + """ + n_basic = draw(st.integers(min_value=0, max_value=2)) + want_fancy = draw(st.booleans()) + + operations: list[IndexOperation] = [] + cur = shape + for _ in range(n_basic): + if len(cur) == 0: + break # a rank-0 (scalar) view has no further indexing to compose + sel = draw(basic_indices(shape=cur)) + operations.append(IndexOperation("basic", sel)) + cur = _basic_result_shape(cur, sel) + + fancy_ok = len(cur) > 0 and all(s > 0 for s in cur) + if want_fancy and fancy_ok: + mode = draw(st.sampled_from(("orthogonal", "vectorized"))) + if mode == "orthogonal": + zsel, npsel = draw(orthogonal_indices(shape=cur)) + operations.append(IndexOperation("orthogonal", (zsel, npsel))) + else: + idx = draw( + npst.integer_array_indices( + shape=cur, + result_shape=npst.array_shapes(min_side=1, max_side=4, max_dims=2), + ) + ) + operations.append(IndexOperation("vectorized", idx)) + + if len(operations) == 0: + operations.append(IndexOperation("basic", draw(basic_indices(shape=cur)))) + + executions: tuple[ProgramExecution, ...] = ( + "materialize", + "eager_on_lazy", + "out", + "set_scalar", + "set_array", + ) + execution = draw(st.sampled_from(executions)) + return IndexingProgram(tuple(operations), execution) + + @st.composite def block_indices( draw: st.DrawFn, *, chunk_sizes: tuple[tuple[int, ...], ...] @@ -620,17 +826,17 @@ def block_indices( Strategy for block-selection indexers over a chunk grid. Block indexing is basic indexing applied to the block grid (the grid of - chunks), so each axis is drawn with ``basic_indices`` over that axis's chunk - count, mirroring how ``orthogonal_indices`` reuses ``basic_indices`` per - axis. ``chunk_sizes`` gives the per-chunk data sizes of the array's *outer* - (block) grid for every axis — i.e. ``Array.write_chunk_sizes``, the grid that - ``Array.blocks`` addresses (the shard grid when sharding is used). For - example ``(3, 3, 3, 1)`` for a length-10 axis with a regular chunk size of 3, - or the explicit edges of a rectilinear axis; ``nchunks`` for an axis is - ``len(chunk_sizes[axis])``. + chunks), so each axis is drawn with `basic_indices` over that axis's chunk + count, mirroring how `orthogonal_indices` reuses `basic_indices` per + axis. `chunk_sizes` gives the per-chunk data sizes of the array's *outer* + (block) grid for every axis — i.e. `Array.write_chunk_sizes`, the grid that + `Array.blocks` addresses (the shard grid when sharding is used). For + example `(3, 3, 3, 1)` for a length-10 axis with a regular chunk size of 3, + or the explicit edges of a rectilinear axis; `nchunks` for an axis is + `len(chunk_sizes[axis])`. The array-space translation uses the cumulative sum of those sizes, matching - ``BlockIndexer``'s use of ``dim_grid.chunk_offset``. Because the sizes are + `BlockIndexer`'s use of `dim_grid.chunk_offset`. Because the sizes are clipped to the array extent, the final offset equals the extent and the translation is exact for regular (uniform), rectilinear, and sharded grids alike. @@ -643,7 +849,7 @@ def block_indices( ------- block_indexer A per-axis tuple of ints / step-1 slices addressing whole chunks, - suitable for ``Array.blocks`` / ``get_block_selection`` / ``set_block_selection``. + suitable for `Array.blocks` / `get_block_selection` / `set_block_selection`. array_indexer The equivalent array-space selection (a tuple of slices) for indexing the corresponding numpy array, used as the comparison oracle. @@ -679,7 +885,7 @@ def predicate(value: tuple[Any, ...]) -> bool: # basic_indices draws slices far more often than bare integers, so the # integer (single-block) branch below would only be hit on rare draws. # Union in an explicit integer so it is reliably exercised — keeping - # coverage deterministic under the derandomized ``ci`` Hypothesis profile. + # coverage deterministic under the derandomized `ci` Hypothesis profile. (dim_sel,) = draw( dim_strategy | st.integers(min_value=0, max_value=nchunks - 1).map(lambda i: (i,)) ) @@ -704,9 +910,9 @@ def block_test_arrays( - **regular**: a regular chunk grid, optionally wrapped in sharding. - **rectilinear**: a variable (rectilinear) chunk grid, always unsharded. - Returns ``(zarray, nparray)``. The per-axis block sizes the oracle needs are - ``zarray.write_chunk_sizes`` — the array's *outer* (block / shard) grid, which - is exactly the grid ``Array.blocks`` addresses; the caller reads it directly. + Returns `(zarray, nparray)`. The per-axis block sizes the oracle needs are + `zarray.write_chunk_sizes` — the array's *outer* (block / shard) grid, which + is exactly the grid `Array.blocks` addresses; the caller reads it directly. """ chunks: tuple[int, ...] | list[list[int]] if draw(st.booleans()): @@ -752,10 +958,10 @@ def key_ranges( [(key, byte_request), (key, byte_request),...] - where ``byte_request`` is ``None`` or any of the concrete ``ByteRequest`` + where `byte_request` is `None` or any of the concrete `ByteRequest` subtypes. The bounds are drawn independently of each value's length, so the offsets/suffixes routinely exceed the data and exercise the clamping logic - in ``_normalize_byte_range_index``. + in `_normalize_byte_range_index`. """ def make_range(start: int, length: int) -> RangeByteRequest: @@ -784,7 +990,7 @@ def complex_rectilinear_arrays( """Generate a rectilinear array with many small chunks. The shape is derived from the chunk edges (5-10 chunks per dim, - sizes 1-5), exercising higher chunk counts than ``rectilinear_arrays``. + sizes 1-5), exercising higher chunk counts than `rectilinear_arrays`. """ ndim = draw(st.integers(min_value=1, max_value=3)) nchunks = draw(st.integers(min_value=5, max_value=10)) diff --git a/tests/test_array.py b/tests/test_array.py index b1a7a3c0f2..cb26659de3 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -1707,6 +1707,19 @@ def test_scalar_array(value: Any, zarr_format: ZarrFormat) -> None: assert isinstance(arr[()], NDArrayLikeOrScalar) +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_iter_0d_array_raises(zarr_format: ZarrFormat) -> None: + """Iterating a 0-d array raises TypeError, matching NumPy. + + This pins an intentional behavior change: previously `Array.__iter__` fell + through to the generic (length-based) sequence protocol and silently produced + an empty iterator for a 0-d array; it now raises eagerly, like `numpy.ndarray`. + """ + arr = zarr.array(np.array(1), zarr_format=zarr_format) + with pytest.raises(TypeError, match="iteration over a 0-d array"): + iter(arr) + + @pytest.mark.parametrize("store", ["local"], indirect=True) @pytest.mark.parametrize("store2", ["local"], indirect=["store2"]) @pytest.mark.parametrize("src_format", [2, 3]) @@ -1989,7 +2002,29 @@ def test_array_repr(store: Store) -> None: shape = (2, 3, 4) dtype = "uint8" arr = zarr.create_array(store, shape=shape, dtype=dtype) + # Identity (non-view) arrays use the legacy repr with no `domain=` suffix, + # byte-identical across `Array` and `AsyncArray`. assert str(arr) == f"" + assert str(arr._async_array) == f"" + + +@pytest.mark.parametrize("store", ["memory"], indirect=["store"]) +def test_array_repr_lazy_view(store: Store) -> None: + """Non-identity lazy views carry a `domain={...}` suffix on both the sync + `Array` and the `AsyncArray` classes; identity arrays never do.""" + shape = (2, 3, 4) + arr = zarr.create_array(store, shape=shape, dtype="uint8") + # Sync view. + view = arr.lazy[1:2] + assert "domain={ [1, 2)" in repr(view) + # Async view (reachable via the async translate_by surface). A pure domain + # translation leaves the addressed storage box unchanged, so the suffix + # shows the storage coordinates; what matters here is that the suffix appears. + async_view = arr._async_array.translate_by((1, 0, 0)) + assert not async_view._is_identity + assert "domain={ [0, 2), [0, 3), [0, 4) }" in repr(async_view) + # And the identity async array still has no suffix. + assert "domain=" not in repr(arr._async_array) class UnknownObjectCodecDtype(VariableLengthUTF8): diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py new file mode 100644 index 0000000000..7d9eea8e4e --- /dev/null +++ b/tests/test_lazy_indexing.py @@ -0,0 +1,1504 @@ +"""End-to-end tests for the public lazy-indexing API (``Array.lazy``). + +These exercise lazy views — ``arr.lazy[sel]``, ``arr.lazy.oindex[sel]``, +``arr.lazy.vindex[sel]`` — and their write-through and composition behaviour, +in every case comparing against an equivalent NumPy reference. + +Coverage is parametrized over a matrix of array geometries so the same dense +set of selections runs against 0-D arrays, and 1/2/3-D arrays in both +**unsharded** and **sharded** form. Sharded arrays route through inner-chunk +geometry (partial-shard read-modify-write), so running the full matrix against +them guards the transform read/write path where it is most likely to break. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any +from unittest import mock + +import numpy as np +import numpy.typing as npt +import pytest +from zarr_indexing.domain import IndexDomain +from zarr_indexing.output_map import ArrayMap +from zarr_indexing.transform import IndexTransform + +import zarr +from zarr.codecs.bytes import BytesCodec +from zarr.codecs.gzip import GzipCodec +from zarr.codecs.sharding import ShardingCodec +from zarr.core.buffer import default_buffer_prototype +from zarr.core.sync import sync +from zarr.errors import BoundsCheckError, LazyViewError, ZarrUserWarning +from zarr.storage import MemoryStore + +if TYPE_CHECKING: + from collections.abc import Callable + + +@dataclass(frozen=True) +class Config: + """A single array geometry: shape, inner chunk shape, optional shard shape.""" + + id: str + shape: tuple[int, ...] + chunks: tuple[int, ...] + shards: tuple[int, ...] | None + + +CONFIGS = [ + Config("0d", (), (), None), + Config("1d-unsharded", (24,), (4,), None), + Config("1d-sharded", (24,), (4,), (12,)), + Config("2d-unsharded", (20, 30), (5, 10), None), + Config("2d-sharded", (20, 30), (5, 10), (10, 30)), + Config("3d-unsharded", (8, 6, 10), (2, 3, 5), None), + Config("3d-sharded", (8, 6, 10), (2, 3, 5), (4, 6, 10)), +] + +# Basic (slice/int/ellipsis) selections valid for each rank. Includes +# across-boundary and positive-strided cases. Negative steps are unsupported +# by zarr (see TestLazyErrors.test_negative_step_raises), so they are excluded. +BASIC_SELECTIONS: dict[int, list[Any]] = { + 0: [Ellipsis], + 1: [slice(3, 20), 7, Ellipsis, slice(None, None, 2), slice(5, 22)], + 2: [ + (slice(2, 8), slice(5, 15)), + 3, + (3, 5), + Ellipsis, + (slice(None, None, 2), slice(None, None, 3)), + (slice(3, 17), slice(8, 22)), + (slice(None), 5), + ], + 3: [ + (slice(1, 7), slice(0, 4), slice(2, 9)), + 2, + (1, 2, 3), + Ellipsis, + (slice(None, None, 2), slice(None, None, 2), slice(None, None, 2)), + (slice(2, 7), slice(1, 5), slice(3, 9)), + ], +} + +# Per-rank integer index arrays for orthogonal / vectorized fancy indexing. +FANCY_INDICES: dict[int, tuple[npt.NDArray[np.intp], ...]] = { + 1: (np.array([1, 3, 7, 20], dtype=np.intp),), + 2: (np.array([1, 5, 10, 18], dtype=np.intp), np.array([0, 3, 9, 29], dtype=np.intp)), + 3: ( + np.array([1, 5, 7], dtype=np.intp), + np.array([0, 3, 5], dtype=np.intp), + np.array([2, 7, 9], dtype=np.intp), + ), +} + +ND_CONFIGS = [c for c in CONFIGS if len(c.shape) >= 1] + + +def _make(cfg: Config) -> tuple[zarr.Array[Any], npt.NDArray[Any]]: + """Build a zarr array for ``cfg`` and an identical NumPy reference.""" + a = zarr.create_array( + MemoryStore(), shape=cfg.shape, chunks=cfg.chunks, shards=cfg.shards, dtype="i4" + ) + n = int(np.prod(cfg.shape, dtype=int)) # prod(()) == 1 -> a single 0-D element + ref = np.arange(n, dtype="i4").reshape(cfg.shape) + a[...] = ref + return a, ref + + +def _value_like(expected: npt.NDArray[Any]) -> Any: + """A distinctly-valued array (or scalar) shaped like ``expected`` for writes.""" + shape = np.shape(expected) + val = (np.arange(int(np.prod(shape, dtype=int)), dtype="i4").reshape(shape) + 1) * 7 + 1 + return val[()] if shape == () else val # 0-D -> python/np scalar + + +def _fmt(sel: Any) -> str: + if sel is Ellipsis: + return "..." + if isinstance(sel, tuple): + return ",".join(_fmt(s) for s in sel) + if isinstance(sel, slice): + start = "" if sel.start is None else str(sel.start) + stop = "" if sel.stop is None else str(sel.stop) + step = f":{sel.step}" if sel.step is not None else "" + return f"{start}:{stop}{step}" + if isinstance(sel, np.ndarray): + return "x".join(map(str, sel.shape)) + return str(sel) + + +BASIC_CASES = [ + pytest.param(cfg, sel, id=f"{cfg.id}:{_fmt(sel)}") + for cfg in CONFIGS + for sel in BASIC_SELECTIONS[len(cfg.shape)] +] +ND_CASES = [pytest.param(cfg, id=cfg.id) for cfg in ND_CONFIGS] +# Configs with >= 2 dimensions, i.e. enough to exercise orthogonal indexing +# across multiple array axes (the outer-product case). +MULTI_AXIS = [cfg for cfg in ND_CONFIGS if len(cfg.shape) >= 2] +MULTI_AXIS_CASES = [pytest.param(cfg, id=cfg.id) for cfg in MULTI_AXIS] +MULTI_AXIS_UNSHARDED_CASES = [ + pytest.param(cfg, id=cfg.id) for cfg in MULTI_AXIS if cfg.shards is None +] +MULTI_AXIS_SHARDED_CASES = [ + pytest.param(cfg, id=cfg.id) for cfg in MULTI_AXIS if cfg.shards is not None +] + + +def _oindex_one_axis(cfg: Config) -> tuple[Any, ...]: + """An orthogonal selection with a single fancy axis (axis 0) and slices elsewhere.""" + idx = FANCY_INDICES[len(cfg.shape)][0] + return (idx, *([slice(None)] * (len(cfg.shape) - 1))) + + +def _rand_slice(rng: np.random.Generator, size: int) -> slice: + """A random non-empty positive-step slice within ``[0, size)``.""" + start = int(rng.integers(0, size)) + stop = int(rng.integers(start + 1, size + 1)) + step = int(rng.integers(1, 4)) + return slice(start, stop, step) + + +def _unique_coords( + rng: np.random.Generator, shape: tuple[int, ...], n: int +) -> tuple[npt.NDArray[np.intp], ...]: + """``n`` distinct coordinate tuples (no duplicate points → write order irrelevant).""" + total = int(np.prod(shape, dtype=int)) + flat = rng.choice(total, size=min(n, total), replace=False) + return tuple(c.astype(np.intp) for c in np.unravel_index(flat, shape)) + + +# Multi-dim geometries used for the randomized model-based round-trips (the 0-D / +# 1-D configs have too small a selection space to be interesting here). +RANDOM_CASES = [pytest.param(cfg, id=cfg.id) for cfg in CONFIGS if len(cfg.shape) >= 2] + + +class TestLazyBasicRead: + @pytest.mark.parametrize(("cfg", "sel"), BASIC_CASES) + def test_matches_numpy(self, cfg: Config, sel: Any) -> None: + """A lazy basic view reads identically to NumPy (and to eager indexing).""" + a, ref = _make(cfg) + expected = ref[sel] + view = a.lazy[sel] + assert tuple(view.shape) == np.shape(expected) + np.testing.assert_array_equal(view[...], expected) + np.testing.assert_array_equal(a[sel], expected) # eager parity, all geometries + + @pytest.mark.parametrize(("cfg", "sel"), BASIC_CASES) + def test_result_and_asarray(self, cfg: Config, sel: Any) -> None: + """``view.result()`` and ``np.asarray(view)`` agree with direct resolution.""" + a, ref = _make(cfg) + expected = ref[sel] + view = a.lazy[sel] + np.testing.assert_array_equal(view.result(), expected) + np.testing.assert_array_equal(np.asarray(view), np.asarray(expected)) + + +class TestLazyBasicWrite: + @pytest.mark.parametrize(("cfg", "sel"), BASIC_CASES) + def test_write_through(self, cfg: Config, sel: Any) -> None: + """Assigning through a lazy basic view mutates exactly the selected region.""" + a, ref = _make(cfg) + expected = ref.copy() + val = _value_like(ref[sel]) + expected[sel] = val + a.lazy[sel] = val + np.testing.assert_array_equal(a[...], expected) + + +class TestLazyOIndex: + @pytest.mark.parametrize("cfg", ND_CASES) + def test_read_single_axis(self, cfg: Config) -> None: + """Lazy orthogonal indexing on a single fancy axis matches NumPy.""" + a, ref = _make(cfg) + sel = _oindex_one_axis(cfg) + expected = ref[sel] + view = a.lazy.oindex[sel] + assert tuple(view.shape) == expected.shape + np.testing.assert_array_equal(view[...], expected) + + @pytest.mark.parametrize("cfg", ND_CASES) + def test_write_single_axis(self, cfg: Config) -> None: + """Write-through lazy orthogonal indexing on a single fancy axis updates the region.""" + a, ref = _make(cfg) + sel = _oindex_one_axis(cfg) + expected = ref.copy() + val = _value_like(ref[sel]) + expected[sel] = val + a.lazy.oindex[sel] = val + np.testing.assert_array_equal(a[...], expected) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_multi_axis_read(self, cfg: Config) -> None: + """Lazy orthogonal indexing across >=2 array axes is the outer product (np.ix_).""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref[np.ix_(*idx)] + view = a.lazy.oindex[idx] + assert tuple(view.shape) == expected.shape + np.testing.assert_array_equal(view[...], expected) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_UNSHARDED_CASES) + def test_multi_axis_write(self, cfg: Config) -> None: + """Write-through lazy orthogonal indexing across >=2 array axes (unsharded).""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref.copy() + val = _value_like(ref[np.ix_(*idx)]) + expected[np.ix_(*idx)] = val + a.lazy.oindex[idx] = val + np.testing.assert_array_equal(a[...], expected) + + @pytest.mark.xfail( + strict=True, + reason="orthogonal multi-array writes are unsupported by the sharding " + "partial-write codec (eager oindex raises the same way); a strict xfail " + "flags if sharded support is added", + ) + @pytest.mark.parametrize("cfg", MULTI_AXIS_SHARDED_CASES) + def test_multi_axis_write_sharded_unsupported(self, cfg: Config) -> None: + """Sharded orthogonal multi-array write — a pre-existing codec limitation.""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + a.lazy.oindex[idx] = _value_like(ref[np.ix_(*idx)]) + + +class TestLazyVIndex: + def test_empty_writes_are_noops(self) -> None: + a, ref = _make(CONFIGS[1]) + view = a.lazy[2:10] + + view.lazy.vindex[(np.array([], dtype=np.intp),)] = np.array([], dtype="i4") + view.set_mask_selection(np.zeros(view.shape, dtype=bool), np.array([], dtype="i4")) + + np.testing.assert_array_equal(a[...], ref) + + @pytest.mark.parametrize("cfg", ND_CASES) + def test_read(self, cfg: Config) -> None: + """Lazy vectorized indexing matches NumPy's point (coordinate) selection.""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref[idx] + view = a.lazy.vindex[idx] + assert tuple(view.shape) == expected.shape + np.testing.assert_array_equal(view[...], expected) + + @pytest.mark.parametrize("cfg", ND_CASES) + def test_write(self, cfg: Config) -> None: + """Write-through lazy vectorized indexing updates the selected points.""" + a, ref = _make(cfg) + idx = FANCY_INDICES[len(cfg.shape)] + expected = ref.copy() + val = _value_like(ref[idx]) + expected[idx] = val + a.lazy.vindex[idx] = val + np.testing.assert_array_equal(a[...], expected) + + +class TestLazyComposition: + @pytest.mark.parametrize("cfg", ND_CASES) + def test_chained_views_compose(self, cfg: Config) -> None: + """Composing two lazy slices equals applying them in sequence on NumPy. + + Positive bounds are literal domain coordinates, so the inner slice + re-selects in the outer view's preserved coordinate system. + """ + a, ref = _make(cfg) + n = cfg.shape[0] + view = a.lazy[1 : n - 1].lazy[2 : n - 2] + expected = ref[2 : n - 2] # literal coordinates: the inner slice re-selects + assert tuple(view.shape) == expected.shape + np.testing.assert_array_equal(view[...], expected) + # the composed view's domain is the literal interval it covers + t = view._async_array._transform + assert (t.domain.inclusive_min[0], t.domain.exclusive_max[0]) == (2, n - 2) + + +class TestLazyViewMethods: + """Indexing methods called on a lazy *view* object (``v = arr.lazy[...]``) must + honor the view's transform, not silently fall back to the storage grid. + + The accessor (``arr.lazy.oindex[...]``) was covered elsewhere, but the methods + on the *returned* non-identity Array (``v.oindex[...]``, ``v[..., -1]``) route + through ``Array``'s own dispatch, which had correctness gaps. + """ + + def test_zero_rank_result_is_scalar(self) -> None: + """A zero-dimensional array reads back as a scalar (shape ``()``) eagerly and lazily.""" + a = zarr.create_array({}, shape=(), chunks=(), dtype="i4") + a[...] = 5 + assert np.shape(a[...]) == () + assert np.shape(a[()]) == () + assert np.shape(a.lazy[...].result()) == () + assert a.lazy[...].result() == 5 + + def test_coordinate_methods_wrap_negative_indices(self) -> None: + a, ref = _make(CONFIGS[1]) + view = a.lazy[2:10] + + np.testing.assert_array_equal( + view.get_coordinate_selection((np.array([-1], dtype=np.intp),)), ref[[9]] + ) + view.set_coordinate_selection((np.array([-1], dtype=np.intp),), 999) + + expected = ref.copy() + expected[9] = 999 + np.testing.assert_array_equal(a[...], expected) + + @pytest.mark.parametrize("cfg", ND_CASES) + def test_view_oindex_respects_transform(self, cfg: Config) -> None: + """``v.oindex[idx]`` on a sub-view reads from the view, not the base array.""" + a, ref = _make(cfg) + n0 = cfg.shape[0] + cut = n0 // 2 + vslice = (slice(cut, n0), *([slice(None)] * (len(cfg.shape) - 1))) + v = a.lazy[vslice] + idx = np.array([cut, cfg.shape[0] - 1], dtype=np.intp) # domain coordinates + osel: Any = (idx, *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v.oindex[osel], ref[osel]) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_view_trailing_index_after_ellipsis(self, cfg: Config) -> None: + """``v[..., k]``: positive ``k`` is a literal domain coordinate; a negative + ``k`` wraps from the end of the view's domain (NumPy parity at the boundary).""" + a, ref = _make(cfg) + v = a.lazy[1 : cfg.shape[0]] + last = cfg.shape[-1] - 1 + np.testing.assert_array_equal(v[..., last], ref[1:, ..., last]) + np.testing.assert_array_equal(v[..., -1], ref[1:, ..., -1]) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_view_basic_tuple_and_int(self, cfg: Config) -> None: + """Basic tuple selections on a view (incl. integer axes that drop) honor the view.""" + a, ref = _make(cfg) + cut = cfg.shape[0] // 2 + v = a.lazy[cut:] + full: Any = (slice(cut, cut + 2), *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v[full], ref[full]) + intsel: Any = (cut + 1, *([slice(None)] * (len(cfg.shape) - 1))) + np.testing.assert_array_equal(v[intsel], ref[intsel]) # int axis drops + + @pytest.mark.parametrize("cfg", MULTI_AXIS_CASES) + def test_view_vindex(self, cfg: Config) -> None: + """``v.vindex[...]`` (coordinate selection) on a view honors the view.""" + a, ref = _make(cfg) + cut = cfg.shape[0] // 2 + v = a.lazy[cut:] + # coordinates: axis 0 lives in [cut, n); the other axes are untouched + idx = ( + np.array([cut, cut + 1, cut + 2], dtype=np.intp), + *(np.array([0, 1, 2], dtype=np.intp) for _ in cfg.shape[1:]), + ) + np.testing.assert_array_equal(v.vindex[idx], ref[idx]) + + def test_view_vindex_with_broadcast_out_buffer(self) -> None: + """vindex with a multi-dim result and out= on a view takes a broadcast-shaped buffer. + + Vectorized indexing scatters through a single flat index internally, but + the caller-visible `out` contract is the broadcast selection shape: the + flat temporary is an implementation detail. The out buffer must therefore + match `expected.shape`, and the results must land in it. + """ + a, ref = _make(CONFIGS[3]) # 2d-unsharded + v = a.lazy[2:18] + i0 = np.array([[2, 3], [4, 5]], dtype=np.intp) # coordinates within [2, 18) + i1 = np.array([[0, 5], [10, 15]], dtype=np.intp) + expected = ref[i0, i1] + assert expected.shape == (2, 2) + buf = default_buffer_prototype().nd_buffer.empty(shape=expected.shape, dtype=np.dtype("i4")) + v.get_coordinate_selection((i0, i1), out=buf) + np.testing.assert_array_equal(np.asarray(buf.as_ndarray_like()), expected) + + @pytest.mark.parametrize("cfg", MULTI_AXIS_UNSHARDED_CASES) + def test_view_write_through_tuple(self, cfg: Config) -> None: + """Writing through a view with a basic tuple selection lands in view coords.""" + a, ref = _make(cfg) + cut = cfg.shape[0] // 2 + v = a.lazy[cut:] + expected = ref.copy() + wsel: Any = (slice(cut, cut + 2), *([slice(None)] * (len(cfg.shape) - 1))) + val = _value_like(ref[wsel]) + expected[wsel] = val + v[wsel] = val + np.testing.assert_array_equal(a[...], expected) + + +class TestLazyAccessorSurface: + def test_shape(self) -> None: + """`.lazy.shape` mirrors the wrapped array's shape, including on views.""" + a, _ = _make(CONFIGS[3]) # 2d-unsharded (20, 30) + assert a.lazy.shape == (20, 30) + v = a.lazy[2:8, 5:15] + assert v.lazy.shape == (6, 10) + + +class TestDaskInterop: + def test_from_array_lazy_view(self) -> None: + """A zero-origin lazy view works as a `dask.array.from_array` source. + + This is the dask-free-wrapper use case (napari): the view exposes + `shape`/`dtype`/`ndim` and eager `__getitem__`, which is all dask needs. + """ + da = pytest.importorskip("dask.array") + a, ref = _make(CONFIGS[3]) # 2d-unsharded (20, 30) + v = a.lazy[2:18, 5:25].translate_to((0, 0)) + d = da.from_array(v, chunks=(8, 10)) + assert d.shape == (16, 20) + np.testing.assert_array_equal(d.compute(scheduler="synchronous"), ref[2:18, 5:25]) + + +class TestLazyErrors: + def test_negative_step_raises(self) -> None: + """A negative slice step is unsupported and raises on the lazy path.""" + a, _ = _make(CONFIGS[1]) # 1d-unsharded + with pytest.raises(IndexError, match="step must be positive"): + a.lazy[::-1] + + def test_accessor_negative_index_wraps(self) -> None: + """The public boundary wraps negative indices NumPy-style against the + current view's ``exclusive_max`` (positive indices stay literal domain + coordinates). Only out-of-range negatives (``k < -size``) or positive + indices past the end raise — before any chunk access. + """ + a, ref = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + np.testing.assert_array_equal(a.lazy[-1].result(), ref[-1]) + np.testing.assert_array_equal( + a.lazy.oindex[(np.array([-1], dtype=np.intp),)].result(), ref[[-1]] + ) + np.testing.assert_array_equal( + a.lazy.vindex[(np.array([-1], dtype=np.intp),)].result(), ref[[-1]] + ) + # too-negative (< -size) and positive-past-the-end still raise + for sel in (np.array([-25], dtype=np.intp), np.array([24], dtype=np.intp)): + with pytest.raises(IndexError, match="out of bounds"): + _ = a.lazy.oindex[(sel,)].result() + with pytest.raises(IndexError, match="out of bounds"): + _ = a.lazy.vindex[(sel,)].result() + + def test_negative_slice_bounds_wrap(self) -> None: + """Negative slice bounds wrap from the end of the current view's domain + (NumPy parity at the boundary). Bounds that wrap out of range still raise, + and bounds that resolve reversed are invalid intervals. + """ + a, ref = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + for sel in (slice(-3, None), slice(-24, None), slice(None, -1), slice(1, -1)): + np.testing.assert_array_equal(a.lazy[sel].result(), ref[sel]) + np.testing.assert_array_equal(a.lazy.oindex[(sel,)].result(), ref[sel]) + # on a view: wraps against the view's exclusive_max, not the base's + v = a.lazy[2:10] # domain [2, 10) + np.testing.assert_array_equal(v.lazy[-2:].result(), ref[8:10]) + # writes wrap too + a.lazy[-3:] = 0 + expected = ref.copy() + expected[-3:] = 0 + np.testing.assert_array_equal(a[...], expected) + # out-of-range wrap (start < -size) still raises + with pytest.raises(BoundsCheckError, match="not contained"): + a.lazy[-25:] + # bounds that RESOLVE reversed (stop < start after wrapping) are invalid + with pytest.raises(IndexError, match="interval"): + a.lazy[-1:1] + + def test_slice_bounds_strict_containment(self) -> None: + """Non-empty slice intervals must be contained in the domain — no + clamping (TensorStore semantics); empty intervals are valid anywhere; + reversed bounds are an error, not an empty result.""" + a, _ = _make(CONFIGS[1]) # shape (24,) + for sel in (slice(5, 100), slice(100, 200), slice(0, 25)): + with pytest.raises(BoundsCheckError, match="not contained"): + a.lazy[sel] + assert a.lazy[5:5].shape == (0,) + assert a.lazy[30:30].shape == (0,) # empty is valid even outside the domain + with pytest.raises(IndexError, match="interval"): + a.lazy[5:2] + + def test_view_positive_literal_negative_wraps(self) -> None: + """A view's positive indices are literal domain coordinates; negative + indices wrap from the view's ``exclusive_max`` (NumPy parity at the + boundary). Positive coordinates below the domain origin remain out of + bounds — the view keeps its preserved coordinate system for positives. + """ + a, ref = _make(CONFIGS[1]) + v = a.lazy[2:10] # domain [2, 10) + np.testing.assert_array_equal(v[2:5], ref[2:5]) # positive literal + np.testing.assert_array_equal(v[3], ref[3]) # positive literal + np.testing.assert_array_equal(v[-1], ref[9]) # wrap: exclusive_max 10 - 1 + np.testing.assert_array_equal(v[-3:], ref[7:10]) # wrap + # positive coordinates below the origin are still out of bounds + below_origin: list[Callable[[], Any]] = [lambda: v[1], lambda: v[0:3]] + for bad in below_origin: + with pytest.raises(BoundsCheckError): + bad() + np.testing.assert_array_equal(np.asarray(v), ref[2:10]) + z = v.translate_to((0,)) + np.testing.assert_array_equal(z[0:3], ref[2:5]) + + def test_lazy_bounds_errors_share_one_type(self) -> None: + """Out-of-range indices — positive past the end, or negative past + ``-size`` — all raise BoundsCheckError (an IndexError), with one message + shape naming the valid range.""" + a, _ = _make(CONFIGS[1]) # shape (24,) + triggers: list[Callable[[], Any]] = [ + lambda: a.lazy[24], + lambda: a.lazy[-25], + lambda: a.lazy[20:30], + lambda: a.lazy.oindex[(np.array([-25], dtype=np.intp),)], + ] + for trigger in triggers: + with pytest.raises(BoundsCheckError, match="out of bounds|not contained"): + trigger() + + def test_guard_message_names_real_apis(self) -> None: + """LazyViewError points at APIs that exist (chunk_projections/metadata), + not the not-yet-added chunk_layout.""" + a, _ = _make(CONFIGS[1]) + with pytest.raises(LazyViewError) as ei: + _ = a.lazy[2:10].chunks + assert "chunk_projections" in str(ei.value) + assert "chunk_layout" not in str(ei.value) + + def test_mask_positions_are_absolute_coordinates(self) -> None: + """Boolean-mask True-positions are absolute coordinates counted from 0, + NOT offsets from the view's origin (TensorStore semantics, verified + against tensorstore 0.1.84: on a domain-[2,10) view, a mask with True + at {3,5,7} addresses cells 3,5,7 — and a True below the origin is out + of the domain, rejected eagerly here where TensorStore defers to read). + """ + a, ref = _make(CONFIGS[1]) # shape (24,) + v = a.lazy[2:10] # domain [2, 10) + mask = np.zeros(8, dtype=bool) + mask[[3, 5, 7]] = True + np.testing.assert_array_equal(v.lazy.oindex[mask].result(), ref[[3, 5, 7]]) + v.lazy.oindex[mask] = 99 # write path: same coordinates + assert list(np.flatnonzero(a[...] == 99)) == [3, 5, 7] + below_origin = np.zeros(8, dtype=bool) + below_origin[0] = True # coordinate 0 is not in [2, 10) + with pytest.raises(BoundsCheckError, match="out of bounds"): + v.lazy.oindex[below_origin] + + def test_views_are_not_iterable(self) -> None: + """iter() on a view raises eagerly (TensorStore parity): the getitem + protocol counts positions from 0, which are not domain coordinates.""" + a, ref = _make(CONFIGS[1]) + with pytest.raises(TypeError, match="not iterable"): + iter(a.lazy[2:10]) + assert [int(x) for x in a][:3] == [int(v) for v in ref[:3]] # base unchanged + + def test_fancy_view_repr_does_not_crash(self) -> None: + """repr of an integer-indexed fancy view must not raise (0-d index array + in selection_repr).""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + ov = b.lazy.oindex[np.array([1, 3]), slice(None)] + assert isinstance(repr(ov.lazy[0]), str) + + +class TestLazyComposedAdvanced: + """Basic-slice-then-fancy composition routes through the view's transform + (the view's domain origin is preserved), with negatives wrapped at the + boundary against the view's ``exclusive_max``.""" + + def test_slice_then_oindex_read(self) -> None: + """``v.oindex[idx]`` on a sliced view reads in the view's domain + coordinates; negatives wrap from the view's end.""" + a, ref = _make(CONFIGS[1]) # shape (24,) + v = a.lazy[10:20] # domain [10, 20) + np.testing.assert_array_equal(v.oindex[np.array([10, 15, 19])], ref[[10, 15, 19]]) + np.testing.assert_array_equal(v.oindex[np.array([-1, -10])], ref[[19, 10]]) + + def test_slice_then_lazy_oindex_read(self) -> None: + """The lazy accessor chain ``v.lazy.oindex[idx]`` composes the same way.""" + a, ref = _make(CONFIGS[1]) + v = a.lazy[10:20] + np.testing.assert_array_equal(v.lazy.oindex[np.array([-1, -10])].result(), ref[[19, 10]]) + + def test_slice_then_oindex_write(self) -> None: + """Writes through the composed view land at the wrapped domain coordinates.""" + a, ref = _make(CONFIGS[1]) + v = a.lazy[10:20] + v.oindex[np.array([-1, -10])] = np.array([777, 888], dtype="i4") + expected = ref.copy() + expected[19] = 777 + expected[10] = 888 + np.testing.assert_array_equal(a[...], expected) + + def test_slice_then_oindex_out_of_bounds(self) -> None: + """Coordinates outside the view's ``[10, 20)`` domain raise before I/O: + a positive below the origin, a negative wrapping below the origin, and a + positive past the end.""" + a, _ = _make(CONFIGS[1]) + v = a.lazy[10:20] + for sel in (np.array([0]), np.array([-11]), np.array([20])): + with pytest.raises(BoundsCheckError): + _ = v.oindex[sel] + + +class TestLazyBoolScalar: + """`isinstance(True, int)` is True, so a bare boolean scalar would otherwise + pass the integer validators and silently read element 0/1. Reject it at the + boundary; boolean *arrays* (masks) stay valid.""" + + @pytest.mark.parametrize("accessor", ["basic", "oindex", "vindex"]) + @pytest.mark.parametrize("val", [True, False, np.True_, np.False_]) + def test_bool_scalar_rejected(self, accessor: str, val: Any) -> None: + a, _ = _make(CONFIGS[1]) + acc: Any = {"basic": a.lazy, "oindex": a.lazy.oindex, "vindex": a.lazy.vindex}[accessor] + with pytest.raises(IndexError): + _ = acc[val] + + def test_bool_array_still_valid(self) -> None: + """A boolean *array* is a mask, not a scalar, and remains a valid index.""" + a, ref = _make(CONFIGS[1]) + mask = np.zeros(24, dtype=bool) + mask[[3, 5, 7]] = True + np.testing.assert_array_equal(a.lazy.oindex[mask].result(), ref[[3, 5, 7]]) + + +class TestLazyMaskShape: + """A boolean mask must exactly match the view domain's shape on the + dimension(s) it consumes (NumPy parity: "boolean index did not match indexed + array"). Anything else raises IndexError at the boundary, before transform + construction — never a silent truncation.""" + + def test_correct_shape_masks_work(self) -> None: + """Exact-shape masks select as before, through oindex, vindex, and on a + non-zero-origin view (whose domain shape, not the base's, is the ruler).""" + a, ref = _make(CONFIGS[1]) # shape (24,) + mask = np.zeros(24, dtype=bool) + mask[[3, 5, 7]] = True + np.testing.assert_array_equal(a.lazy.oindex[mask].result(), ref[[3, 5, 7]]) + np.testing.assert_array_equal(a.lazy.vindex[mask].result(), ref[[3, 5, 7]]) + v = a.lazy[2:10] # domain [2, 10) -> shape (8,) + vmask = np.zeros(8, dtype=bool) + vmask[[3, 5]] = True # True positions are absolute coordinates + np.testing.assert_array_equal(v.lazy.oindex[vmask].result(), ref[[3, 5]]) + + @pytest.mark.parametrize("accessor", ["oindex", "vindex"]) + def test_under_length_mask_raises(self, accessor: str) -> None: + a, _ = _make(CONFIGS[1]) # shape (24,) + acc: Any = getattr(a.lazy, accessor) + with pytest.raises(IndexError, match="boolean index did not match"): + _ = acc[np.zeros(5, dtype=bool)] + + @pytest.mark.parametrize("accessor", ["oindex", "vindex"]) + @pytest.mark.parametrize("fill", [False, True]) + def test_over_length_mask_raises(self, accessor: str, fill: bool) -> None: + """Over-length masks raise the shape-mismatch error regardless of where + their True values fall (not a bounds error on the True positions).""" + a, _ = _make(CONFIGS[1]) + acc: Any = getattr(a.lazy, accessor) + with pytest.raises(IndexError, match="boolean index did not match"): + _ = acc[np.full(30, fill, dtype=bool)] + + @pytest.mark.parametrize("accessor", ["oindex", "vindex"]) + def test_wrong_ndim_mask_raises(self, accessor: str) -> None: + """A 2-D mask on a 1-D array is an IndexError (it consumes two + dimensions), not a ValueError from deeper in the transform layer.""" + a, _ = _make(CONFIGS[1]) + acc: Any = getattr(a.lazy, accessor) + with pytest.raises(IndexError, match="too many indices"): + _ = acc[np.zeros((4, 6), dtype=bool)] + + def test_view_mask_checked_against_view_domain(self) -> None: + """On a view, the mask ruler is the view's domain shape: the base + array's shape is the wrong length there.""" + a, _ = _make(CONFIGS[1]) + v = a.lazy[2:10] # domain shape (8,) + with pytest.raises(IndexError, match="boolean index did not match"): + _ = v.lazy.oindex[np.zeros(24, dtype=bool)] + + +_GET_FIELD_METHODS = [ + "get_basic_selection", + "get_orthogonal_selection", + "get_mask_selection", + "get_coordinate_selection", +] +_SET_FIELD_METHODS = [ + "set_basic_selection", + "set_orthogonal_selection", + "set_mask_selection", + "set_coordinate_selection", +] + + +class TestLazyFieldsOnView: + """`fields=` on a non-identity view routed through the legacy indexer, which + ignores the transform and reads/writes the wrong storage region. It must now + raise NotImplementedError before any storage access in all eight + get/set_*_selection methods (and the ``v['f0', :]`` sugar), leaving the base + array bit-identical.""" + + @staticmethod + def _struct_array() -> tuple[zarr.Array[Any], npt.NDArray[Any]]: + dt = np.dtype([("f0", "i4"), ("f1", "i4")]) + a = zarr.create_array(MemoryStore(), shape=(10,), chunks=(10,), dtype=dt, zarr_format=2) + base = np.zeros(10, dtype=dt) + base["f0"] = np.arange(10, dtype="i4") + base["f1"] = np.arange(10, dtype="i4") + 100 + a[...] = base + return a, base + + @staticmethod + def _selection_for(method: str) -> Any: + if "mask" in method: + return np.zeros(5, dtype=bool) # view shape is (5,) + if "coordinate" in method: + return (np.array([5, 6, 7], dtype=np.intp),) + return Ellipsis + + @pytest.mark.parametrize("method", _GET_FIELD_METHODS) + def test_get_fields_on_view_raises(self, method: str) -> None: + a, base = self._struct_array() + v = a.lazy[5:10] + with pytest.raises(NotImplementedError, match="field"): + getattr(v, method)(self._selection_for(method), fields="f0") + np.testing.assert_array_equal(a[...], base) + + @pytest.mark.parametrize("method", _SET_FIELD_METHODS) + def test_set_fields_on_view_raises(self, method: str) -> None: + a, _base = self._struct_array() + before = np.asarray(a[...]).copy() + v = a.lazy[5:10] + with pytest.raises(NotImplementedError, match="field"): + getattr(v, method)(self._selection_for(method), np.arange(3, dtype="i4"), fields="f0") + np.testing.assert_array_equal(a[...], before) + + def test_getitem_field_sugar_on_view_raises(self) -> None: + a, base = self._struct_array() + v = a.lazy[5:10] + selections: list[Any] = [("f0", slice(None)), "f0"] + for sel in selections: + with pytest.raises(NotImplementedError, match="field"): + _ = v[sel] + np.testing.assert_array_equal(a[...], base) + + def test_setitem_field_sugar_on_view_raises(self) -> None: + a, _base = self._struct_array() + before = np.asarray(a[...]).copy() + v = a.lazy[5:10] + field_sel: Any = ("f0", slice(None)) + with pytest.raises(NotImplementedError, match="field"): + v[field_sel] = np.arange(5, dtype="i4") + np.testing.assert_array_equal(a[...], before) + + def test_fields_on_identity_array_still_work(self) -> None: + """Field selection on the base (identity) array still routes to the legacy + path — the view guard must not block it. (Only the write path is checked + here; structured-dtype field *reads* are broken independently of this + change.)""" + a, _ = self._struct_array() + a.set_basic_selection(Ellipsis, np.arange(10, dtype="i4") + 1, fields="f0") + result: Any = np.asarray(a[...]) + np.testing.assert_array_equal(result["f0"], np.arange(10, dtype="i4") + 1) + + @pytest.mark.xfail( + reason="the eager fields= write path corrupts sibling fields (pre-existing " + "upstream bug in the legacy identity path, not introduced or fixed by the " + "lazy-view guard); strict xfail flips when the eager path is fixed", + strict=True, + ) + def test_fields_write_preserves_sibling_fields(self) -> None: + """A `fields='f0'` write must leave `f1` untouched. Documents the known + eager-path corruption (mirrors the TestKnownFancyIntBugs pinning style).""" + a, base = self._struct_array() + a.set_basic_selection(Ellipsis, np.arange(10, dtype="i4") + 1, fields="f0") + result: Any = np.asarray(a[...]) + np.testing.assert_array_equal(result["f1"], base["f1"]) + + +class TestFancyAfterFancy: + """Composing a fancy (orthogonal/vectorized) selection onto a view that + already carries an orthogonal ArrayMap axis is not supported: the reindexing + machinery would try to index a broadcast (singleton) axis of the existing + ArrayMap with the new coordinates. This used to leak a raw NumPy + ``IndexError`` (``index N is out of bounds for axis ... with size 1``) at + resolve time, which read like a user error; it now raises + ``NotImplementedError`` naming the limitation and the workaround at + composition time. Compositions that keep the fancy step on the existing + dependency axis (or on a correlated vindex view) still work — see + ``test_fancy_on_same_axis_still_works``.""" + + def test_oindex_then_oindex_other_axis_raises(self) -> None: + """``v.lazy.oindex[:, [5]]`` after ``v = a.lazy.oindex[[0, 2, 5], :]`` + raises at composition time, not a raw IndexError at resolve time.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + v = b.lazy.oindex[[0, 2, 5], :] + with pytest.raises(NotImplementedError, match="fancy"): + _ = v.lazy.oindex[:, [5]] + + def test_oindex_then_vindex_other_axis_raises(self) -> None: + """The vindex accessor closes the same gap: a coordinate selection over a + broadcast axis of an oindex view raises rather than crashing in numpy.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + v = b.lazy.oindex[[0, 2, 5], :] + with pytest.raises(NotImplementedError, match="fancy"): + _ = v.lazy.vindex[[0, 1], [5, 6]] + + def test_oindex_both_axes_then_oindex_raises(self) -> None: + """An outer-product view (two ArrayMaps) rejects any further fancy step: + the new axis is a broadcast axis of the *other* ArrayMap.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + b[...] = np.arange(48, dtype="i4").reshape(6, 8) + v = b.lazy.oindex[[0, 2, 5], [1, 3, 5]] + with pytest.raises(NotImplementedError, match="fancy"): + _ = v.lazy.oindex[[0, 1], :] + + def test_fancy_on_same_axis_still_works(self) -> None: + """A fancy step on the existing dependency axis is absorbed correctly and + must keep working; likewise a fancy step on a correlated vindex view.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + ref = np.arange(48, dtype="i4").reshape(6, 8) + b[...] = ref + rows = b.lazy.oindex[[0, 2, 5], :] + np.testing.assert_array_equal( + rows.lazy.oindex[[0, 1], :].result(), ref[[0, 2, 5], :][[0, 1], :] + ) + pts = b.lazy.vindex[[0, 2, 5], [1, 2, 3]] + expected = ref[[0, 2, 5], [1, 2, 3]] + np.testing.assert_array_equal(pts.lazy.oindex[[0, 1]].result(), expected[[0, 1]]) + np.testing.assert_array_equal(pts.lazy.vindex[[0, 1]].result(), expected[[0, 1]]) + + +class TestKnownFancyIntBugs: + """Strict-xfail pins for the int-on-fancy-picked-dim defect (see review notes): + integer indexing a dimension that an oindex/vindex selection created is + mis-lowered, crashing reads or mis-shaping results. These flip to failures + when the bug is fixed, forcing the pins to be removed.""" + + @pytest.mark.xfail( + reason="integer indexing on an oindex-picked dimension crashes in the " + "codec pipeline (ArrayMap not collapsed to ConstantMap)", + strict=True, + ) + def test_int_read_on_oindex_view(self) -> None: + """`rows[0]` on an oindex-created view must equal the picked row.""" + b = zarr.create_array({}, shape=(6, 8), chunks=(2, 3), dtype="i4") + ref = np.arange(48, dtype="i4").reshape(6, 8) + b[...] = ref + rows = b.lazy.oindex[np.array([1, 3]), slice(None)] + np.testing.assert_array_equal(rows[0], ref[[1, 3]][0]) + + def test_int_read_on_vindex_view_is_scalar(self) -> None: + """`pts[0]` on a vindex-created view must be a scalar, as in NumPy.""" + a = zarr.create_array({}, shape=(12,), chunks=(3,), dtype="i4") + a[...] = np.arange(12, dtype="i4") + pts = a.lazy.vindex[np.array([1, 3, 5])] + assert np.shape(pts[0]) == () + assert np.shape(pts.lazy[0].result()) == () + assert pts[0] == a[1] + + def test_block_selection_on_view_rejected(self) -> None: + """Block selection is ill-defined on a lazy view and must raise, not corrupt.""" + a, _ = _make(CONFIGS[3]) # 2d-unsharded + v = a.lazy[5:15] + with pytest.raises(NotImplementedError, match="block selection"): + _ = v.blocks[0] + with pytest.raises(NotImplementedError, match="block selection"): + v.blocks[0] = 0 + + def test_vindex_with_slice_rejected(self) -> None: + """vindex is coordinate-only; mixing a slice raises (parity with eager).""" + a, _ = _make(CONFIGS[3]) # 2d-unsharded + with pytest.raises(IndexError): + a.lazy.vindex[(np.array([0, 1], dtype=np.intp), slice(None))] + + def test_result_threads_prototype(self) -> None: + """``result(prototype=...)`` forwards the prototype rather than dropping it.""" + a, ref = _make(CONFIGS[3]) # 2d-unsharded + proto = default_buffer_prototype() + with mock.patch.object(type(a), "get_basic_selection", autospec=True) as gbs: + gbs.return_value = ref + a.result(prototype=proto) + assert gbs.call_args.kwargs.get("prototype") is proto + + +class TestLazyRandomizedRoundtrip: + """Model-based round-trips: apply random selections to both the zarr array and a + NumPy reference, then assert they stay equal. Parametrized over sharded and + unsharded grids with one body, so the same selections exercise chunk- and + shard-boundary read-modify-write (the pattern TensorStore's driver_testutil uses). + """ + + @pytest.mark.parametrize("cfg", RANDOM_CASES) + def test_basic_writes_track_numpy(self, cfg: Config) -> None: + """Random strided basic writes/reads match a NumPy model across boundaries.""" + rng = np.random.default_rng(0) + a, ref = _make(cfg) + ref = ref.copy() + for _ in range(25): + sel = tuple(_rand_slice(rng, s) for s in cfg.shape) + val = rng.integers(-9999, 9999, size=ref[sel].shape, dtype="i4") + ref[sel] = val + a.lazy[sel] = val + np.testing.assert_array_equal(a.lazy[sel][...], ref[sel]) + np.testing.assert_array_equal(a[...], ref) + + @pytest.mark.parametrize("cfg", RANDOM_CASES) + def test_vindex_writes_track_numpy(self, cfg: Config) -> None: + """Random coordinate (vindex) writes/reads match a NumPy model across boundaries.""" + rng = np.random.default_rng(1) + a, ref = _make(cfg) + ref = ref.copy() + for _ in range(25): + idx = _unique_coords(rng, cfg.shape, int(rng.integers(1, 7))) + val = rng.integers(-9999, 9999, size=idx[0].shape, dtype="i4") + ref[idx] = val + a.lazy.vindex[idx] = val + np.testing.assert_array_equal(a.lazy.vindex[idx][...], ref[idx]) + np.testing.assert_array_equal(a[...], ref) + + @pytest.mark.parametrize("cfg", RANDOM_CASES) + def test_oindex_single_axis_writes_track_numpy(self, cfg: Config) -> None: + """Random single-fancy-axis orthogonal writes/reads match a NumPy model.""" + rng = np.random.default_rng(2) + a, ref = _make(cfg) + ref = ref.copy() + size0 = cfg.shape[0] + for _ in range(25): + k = int(rng.integers(1, size0 + 1)) + idx0 = rng.choice(size0, size=k, replace=False).astype(np.intp) + sel: Any = (idx0, *(_rand_slice(rng, s) for s in cfg.shape[1:])) + val = rng.integers(-9999, 9999, size=ref[sel].shape, dtype="i4") + ref[sel] = val + a.lazy.oindex[sel] = val + np.testing.assert_array_equal(a.lazy.oindex[sel][...], ref[sel]) + np.testing.assert_array_equal(a[...], ref) + + +_VIEW_GUARDED_PROPERTIES = ( + "chunks", + "shards", + "read_chunk_sizes", + "write_chunk_sizes", + "cdata_shape", + "nchunks", + "nchunks_initialized", + "info", +) + +# method name -> call arguments +_VIEW_GUARDED_METHODS: dict[str, tuple[Any, ...]] = { + "nbytes_stored": (), + "info_complete": (), + "resize": ((24,),), + "append": (np.zeros((3,), dtype="i4"),), +} + + +class TestLazyViewGridGuards: + """Grid-describing members assume the array fills its chunk grid, so on a + non-identity lazy view they must raise instead of silently describing the + backing grid (a footgun for consumers that size reads off ``.chunks``).""" + + @staticmethod + def _array() -> zarr.Array[Any]: + a = zarr.create_array({}, shape=(12,), chunks=(3,), dtype="i4") + a[...] = np.arange(12, dtype="i4") + return a + + @pytest.mark.parametrize("name", _VIEW_GUARDED_PROPERTIES) + def test_property_raises_on_view(self, name: str) -> None: + """A grid-describing property works on the backing array but raises LazyViewError on a view.""" + a = self._array() + getattr(a, name) # backing (identity) array: no raise + with pytest.raises(LazyViewError): + getattr(a.lazy[2:10], name) + + @pytest.mark.parametrize("name", list(_VIEW_GUARDED_METHODS)) + def test_method_raises_on_view(self, name: str) -> None: + """A grid-describing/mutating method raises LazyViewError on a view.""" + a = self._array() + with pytest.raises(LazyViewError): + getattr(a.lazy[2:10], name)(*_VIEW_GUARDED_METHODS[name]) + + def test_size_and_nbytes_reflect_the_view(self) -> None: + """size / nbytes are logical members: they describe the view's extent, not the backing array.""" + a = self._array() # shape (12,), i4 (4 bytes) + view = a.lazy[2:10] # logical shape (8,) + assert view.shape == (8,) + assert view.size == 8 + assert view.nbytes == 8 * 4 + + +class TestChunkProjections: + """`chunk_projections` enumerates the stored chunks an (identity or view) array + projects onto, giving each one's stored region, the region of this array it maps + to, and whether it is partially covered.""" + + @staticmethod + def _array(shape: tuple[int, ...], chunks: tuple[int, ...]) -> zarr.Array[Any]: + a = zarr.create_array({}, shape=shape, chunks=chunks, dtype="i4") + a[...] = np.arange(int(np.prod(shape)), dtype="i4").reshape(shape) + return a + + def test_identity_tiles_the_whole_domain(self) -> None: + """On a full (identity) array, projections tile the domain exactly: one per chunk, none partial.""" + a = self._array((10,), (3,)) # chunks (3,3,3,1) + projs = list(a.chunk_projections()) + assert len(projs) == 4 + assert all(not p.is_partial for p in projs) + covered = np.zeros(10, dtype=bool) + for p in projs: + arr_sel: Any = p.array_selection + covered[arr_sel] = True + assert covered.all() # complete, and (bool assignment) each cell once + + def test_view_round_trip(self) -> None: + """Placing each projection's stored-chunk slice at its array_selection reconstructs the view.""" + a = self._array((10,), (3,)) + backing = np.asarray(a[...]) + view = a.lazy[2:9] # (7,) + out = np.empty(view.shape, dtype="i4") + for p in view.chunk_projections(): + offset = p.coord[0] * 3 # regular grid: chunk c starts at c*chunk_size + stored_chunk = backing[offset : offset + p.shape[0]] + arr_sel: Any = p.array_selection + chunk_sel: Any = p.chunk_selection + out[arr_sel] = stored_chunk[chunk_sel] + np.testing.assert_array_equal(out, backing[2:9]) + + def test_partial_flag_and_alignment(self) -> None: + """Boundary chunks a view only partially covers are flagged; a chunk-aligned view is not.""" + a = self._array((12,), (4,)) # chunks at [0,4) [4,8) [8,12) + partial = a.lazy[2:12] # first chunk partially covered + assert any(p.is_partial for p in partial.chunk_projections()) + assert not partial.is_chunk_aligned() + aligned = a.lazy[4:12] # exactly chunks 1 and 2 + assert all(not p.is_partial for p in aligned.chunk_projections()) + assert aligned.is_chunk_aligned() + + def test_sharded_write_unit(self) -> None: + """For a sharded array, unit='write' enumerates shards; unit='read' (inner chunks) is deferred.""" + a = zarr.create_array({}, shape=(12,), chunks=(2,), shards=(6,), dtype="i4") + a[...] = np.arange(12, dtype="i4") + shards = list(a.chunk_projections(unit="write")) + assert len(shards) == 2 # two shards of size 6 + assert all(p.shape == (6,) and not p.is_partial for p in shards) + with pytest.raises(NotImplementedError): + a.chunk_projections(unit="read") + + def test_multi_codec_sharded_read_unit_raises(self) -> None: + """A sharded array whose codec chain has more than one entry (a `ShardingCodec` + followed by a trailing bytes-bytes codec) is still sharded: `unit='read'` must + raise the same `NotImplementedError` as a plain single-codec sharded array.""" + with pytest.warns(ZarrUserWarning, match="sharding_indexed"): + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[ + ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec()]), + GzipCodec(), + ], + ) + a[...] = np.arange(12, dtype="i4") + shards = list(a.chunk_projections(unit="write")) + assert len(shards) == 2 # two shards of size 6 + assert all(p.shape == (6,) and not p.is_partial for p in shards) + with pytest.raises(NotImplementedError): + a.chunk_projections(unit="read") + + def test_invalid_unit_rejected(self) -> None: + """An unknown granularity is rejected eagerly.""" + a = self._array((6,), (2,)) + with pytest.raises(ValueError, match="unit"): + a.chunk_projections(unit="bogus") # type: ignore[arg-type] + + +class TestIsSharded: + """`_is_sharded` must key off the presence of a `ShardingCodec` in the codec + chain, not the chain's length: a valid sharded array may carry additional + codecs (e.g. a trailing bytes-bytes compressor after the sharding codec).""" + + def test_true_for_multi_codec_sharded_chain(self) -> None: + with pytest.warns(ZarrUserWarning, match="sharding_indexed"): + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[ + ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec()]), + GzipCodec(), + ], + ) + assert a._async_array._is_sharded is True + + def test_true_for_single_codec_sharded_chain(self) -> None: + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[ShardingCodec(chunk_shape=(2,), codecs=[BytesCodec()])], + ) + assert a._async_array._is_sharded is True + + def test_false_for_plain_bytes_chain(self) -> None: + a = zarr.create( + store={}, shape=(12,), chunks=(6,), dtype="i4", zarr_format=3, codecs=[BytesCodec()] + ) + assert a._async_array._is_sharded is False + + def test_false_for_plain_compressed_chain(self) -> None: + a = zarr.create( + store={}, + shape=(12,), + chunks=(6,), + dtype="i4", + zarr_format=3, + codecs=[BytesCodec(), GzipCodec()], + ) + assert a._async_array._is_sharded is False + + +class TestChunkCoverageParity: + """`chunk_projections`' `is_partial` (backed by `_covers_full_chunk`) must agree + with the write path's `_is_complete_chunk`: both derive coverage from the same + `iter_chunk_transforms` enumeration and must reach the same verdict — including + the case a naive selection-kind check misses: an integer index on a + chunk-extent-1 dimension covers that dimension exactly as its equivalent + length-1 slice does.""" + + @pytest.mark.parametrize( + "sel", + [ + (slice(None), slice(None)), # full slice, both dims -> aligned + (slice(1, 3), slice(None)), # partial slice on dim 0 (chunk extent 1) + (2, slice(None)), # int on a chunk-extent-1 dim -> aligned + (slice(2, 3), slice(None)), # equivalent length-1 slice -> aligned + (slice(None), 2), # int on a chunk-extent-4 dim -> partial + (slice(None), slice(2, 3)), # length-1 slice on the same dim -> partial + ], + ) + def test_is_partial_matches_is_complete_chunk(self, sel: tuple[Any, ...]) -> None: + """For every stored chunk a selection touches, `is_partial` (computed via + `chunk_projections`) must equal `not _is_complete_chunk(...)` computed + directly from the same `iter_chunk_transforms` enumeration.""" + from zarr_indexing.chunk_resolution import iter_chunk_transforms + + from zarr.core.array import _is_complete_chunk + + a = zarr.create_array({}, shape=(4, 4), chunks=(1, 4), dtype="i4") + a[...] = np.arange(16, dtype="i4").reshape(4, 4) + view = a.lazy[sel] + async_view = view._async_array + transform = async_view._transform.translate_domain_to( + (0,) * async_view._transform.input_rank + ) + chunk_grid = async_view._chunk_grid + + projections = {p.coord: p for p in view.chunk_projections()} + assert len(projections) > 0 + + seen: set[tuple[int, ...]] = set() + for chunk_coords, sub_transform, _out_indices in iter_chunk_transforms( + transform, chunk_grid + ): + spec = chunk_grid[chunk_coords] + if spec is None: + continue + seen.add(chunk_coords) + expected_partial = not _is_complete_chunk(sub_transform, spec) + assert projections[chunk_coords].is_partial == expected_partial + + assert seen == set(projections) + + def test_int_and_length1_slice_report_identically(self) -> None: + """An integer selection and its equivalent length-1 slice, on a + chunk-extent-1 dimension, must report the same is_partial/aligned verdict — + both cover the same data, so a read-modify-write plan should not depend on + which selection spelling was used.""" + a = zarr.create_array({}, shape=(4, 4), chunks=(1, 4), dtype="i4") + a[...] = np.arange(16, dtype="i4").reshape(4, 4) + + int_view = a.lazy[2, :] + slice_view = a.lazy[2:3, :] + + int_partial = [p.is_partial for p in int_view.chunk_projections()] + slice_partial = [p.is_partial for p in slice_view.chunk_projections()] + assert int_partial == slice_partial + assert int_view.is_chunk_aligned() == slice_view.is_chunk_aligned() + assert int_view.is_chunk_aligned() is True # chunk extent along dim 0 is 1 + + +class TestOutOfGridGuardrail: + """The transform I/O loops must raise loudly on an out-of-grid chunk + coordinate rather than silently `continue`. A read buffer is allocated with + `nd_buffer.empty()` (uninitialized), so a bug anywhere in transform + composition that yields an out-of-grid storage coordinate would otherwise + return garbage memory on reads and silently drop data on writes. + + The public indexing APIs already bounds-check selections before a transform + reaches chunk resolution, so an out-of-grid coordinate can only arise from a + transform built directly (simulating a composition bug upstream). Each test + hand-builds an `IndexTransform` whose `ArrayMap` storage coordinates + exceed the backing array's chunk grid. + """ + + @staticmethod + def _out_of_grid_transform() -> IndexTransform: + """For shape=(20,), chunks=(5,): storage coords 1 and 2 are in-grid (chunk + 0), but 23 and 24 exceed the array's extent and fall in chunk id 4, which + does not exist (the grid only has chunks 0-3).""" + idx = np.array([1, 2, 23, 24], dtype=np.intp) + return IndexTransform( + domain=IndexDomain.from_shape((4,)), output=(ArrayMap(index_array=idx),) + ) + + def test_get_selection_via_transform_raises(self) -> None: + """The read pipeline (`_get_selection_via_transform`) must raise instead + of silently returning whatever was in the uninitialized output buffer.""" + a = zarr.create_array({}, shape=(20,), chunks=(5,), dtype="i4") + a[...] = np.arange(20, dtype="i4") + transform = self._out_of_grid_transform() + with pytest.raises(IndexError, match="out of bounds"): + sync(a._async_array._get_selection_t(transform, prototype=default_buffer_prototype())) + + def test_set_selection_via_transform_raises(self) -> None: + """The write pipeline (`_set_selection_via_transform`) must raise instead + of silently dropping the out-of-grid entries -- and must not apply a + partial write for the in-grid entries either.""" + a = zarr.create_array({}, shape=(20,), chunks=(5,), dtype="i4") + original = np.arange(20, dtype="i4") + a[...] = original + transform = self._out_of_grid_transform() + value = np.array([100, 200, 300, 400], dtype="i4") + with pytest.raises(IndexError, match="out of bounds"): + sync( + a._async_array._set_selection_t( + transform, value, prototype=default_buffer_prototype() + ) + ) + # The batch is built (and validated) before any write is issued, so a + # rejected batch must leave the array untouched -- not even the in-grid + # entries (indices 1, 2) should have been written. + np.testing.assert_array_equal(a[...], original) + + def test_iter_chunk_projections_raises(self) -> None: + """`chunk_projections` (backed by `iter_chunk_projections`) must raise + instead of silently omitting the out-of-grid chunk from the enumeration.""" + a = zarr.create_array({}, shape=(20,), chunks=(5,), dtype="i4") + a[...] = np.arange(20, dtype="i4") + transform = self._out_of_grid_transform() + view = a._with_transform(transform) + with pytest.raises(IndexError, match="out of bounds"): + list(view.chunk_projections()) + + +class TestLazyArrayMapResolution: + """Chunk resolution of orthogonal (outer-product) and correlated (vectorized) + ArrayMaps: intersections and output selectors must preserve the distinct + semantics of the two flavours, across chunk boundaries, and for the + length-1 shapes that are degenerate under the shape-derived classifier. + """ + + def _make(self, shape: tuple[int, ...], chunks: tuple[int, ...]) -> tuple[Any, Any]: + a = zarr.create_array(store={}, shape=shape, chunks=chunks, dtype="i4") + data = np.arange(int(np.prod(shape)), dtype="i4").reshape(shape) + a[...] = data + return a, data + + def test_oindex_multiple_arrays_outer_product(self) -> None: + """Orthogonal selection with unequal index lengths equals ``np.ix_`` (read).""" + a, data = self._make((20, 30), (5, 10)) + rows = np.array([1, 11], dtype=np.intp) + cols = np.array([2, 12, 22], dtype=np.intp) + actual = a.lazy.oindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[np.ix_(rows, cols)]) + + def test_oindex_multiple_arrays_outer_product_write(self) -> None: + """Orthogonal outer-product write spanning several chunks (unsharded).""" + a, data = self._make((20, 30), (5, 10)) + rows = np.array([1, 11], dtype=np.intp) + cols = np.array([2, 12, 22], dtype=np.intp) + expected = data.copy() + val = (np.arange(6, dtype="i4").reshape(2, 3) + 1) * 100 + expected[np.ix_(rows, cols)] = val + a.lazy.oindex[rows, cols] = val + np.testing.assert_array_equal(a[...], expected) + + def test_oindex_length1_and_length3_axes(self) -> None: + """A length-1 orthogonal axis (all-singleton shape) must not be + misclassified as scalar/correlated: the outer product still holds.""" + a, data = self._make((6, 6), (2, 2)) + rows = np.array([2], dtype=np.intp) + cols = np.array([1, 3, 5], dtype=np.intp) + actual = a.lazy.oindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[np.ix_(rows, cols)]) + assert actual.shape == (1, 3) + + def test_vindex_length1_pair(self) -> None: + """Two length-1 correlated arrays (degenerate (1,) shape) stay a single + pointwise scatter, not an outer product.""" + a, data = self._make((6, 6), (2, 2)) + actual = a.lazy.vindex[np.array([2]), np.array([3])].result() + np.testing.assert_array_equal(actual, data[np.array([2]), np.array([3])]) + + def test_vindex_two_arrays_with_residual_slice(self) -> None: + """Vectorized selection with two correlated arrays plus a residual slice + dim (partial indexing of a 3-D array) matches NumPy fancy indexing.""" + a, data = self._make((4, 3, 5), (2, 2, 2)) + rows = np.array([1, 3], dtype=np.intp) + cols = np.array([2, 0], dtype=np.intp) + actual = a.lazy.vindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[rows, cols]) + assert actual.shape == (2, 5) + + def test_vindex_two_arrays_with_residual_slice_write(self) -> None: + """Write-through of the correlated + residual-slice case.""" + a, data = self._make((4, 3, 5), (2, 2, 2)) + rows = np.array([1, 3], dtype=np.intp) + cols = np.array([2, 0], dtype=np.intp) + expected = data.copy() + val = (np.arange(10, dtype="i4").reshape(2, 5) + 1) * 100 + expected[rows, cols] = val + a.lazy.vindex[rows, cols] = val + np.testing.assert_array_equal(a[...], expected) + + def test_vindex_partial_rank_bool_mask(self) -> None: + """A partial-rank boolean mask (rank < array ndim) leaves a residual slice + dim; the vectorized read matches NumPy.""" + a, data = self._make((4, 3, 5), (2, 2, 2)) + mask = np.zeros((4, 3), dtype=bool) + mask[1, 2] = mask[3, 0] = True + actual = a.lazy.vindex[mask].result() + np.testing.assert_array_equal(actual, data[mask]) + assert actual.shape == (2, 5) + + def test_basic_slice_after_oindex_independent_axis(self) -> None: + """Basic-slicing an axis the ArrayMap is independent of broadcasts the + map (does not slice its values); the slice narrows only that axis.""" + a, data = self._make((6, 6), (2, 2)) + cols = a.lazy.oindex[:, np.array([1, 3, 5], dtype=np.intp)] + ref = data[:, [1, 3, 5]] + np.testing.assert_array_equal(cols.lazy[2:4, :].result(), ref[2:4, :]) + + def test_basic_slice_after_oindex_fancy_axis(self) -> None: + """Basic-slicing the fancy axis of an ArrayMap slices the map's values.""" + a, data = self._make((6, 6), (2, 2)) + cols = a.lazy.oindex[:, np.array([1, 3, 5], dtype=np.intp)] + ref = data[:, [1, 3, 5]] + np.testing.assert_array_equal(cols.lazy[:, 0:2].result(), ref[:, 0:2]) + + def test_empty_oindex_chunk_projections(self) -> None: + """An empty fancy selection yields no chunk projections (no crash).""" + a = zarr.create_array(store={}, shape=(10,), chunks=(3,), dtype="i4") + assert list(a.lazy.oindex[np.array([], dtype=np.intp)].chunk_projections()) == [] + + def test_empty_oindex_read(self) -> None: + """An empty fancy selection reads back an empty array.""" + a, _ = self._make((10,), (3,)) + actual = a.lazy.oindex[np.array([], dtype=np.intp)].result() + assert actual.shape == (0,) + + +# --------------------------------------------------------------------------- +# Async-surface guards +# --------------------------------------------------------------------------- +# +# The async surface (``AsyncArray.getitem`` / ``setitem`` / selection methods and +# the ``AsyncOIndex`` / ``AsyncVIndex`` accessors) does not yet route selections +# through a view's index transform: it builds a legacy indexer from +# ``metadata.shape``, which for a lazy view silently reads or writes the *wrong* +# region. Until async transform routing lands, every such entry point must raise +# ``LazyViewError`` for a non-identity view, steering users to the synchronous +# ``Array`` surface (or ``.result()``). Identity arrays must be unaffected. + + +def _read_ops() -> list[Any]: + """Async read entry points that build a legacy indexer, keyed by name.""" + coord = (np.array([0, 1], dtype=np.intp),) + mask = np.zeros((10,), dtype=bool) + mask[0] = mask[2] = True + return [ + pytest.param(lambda av: av.getitem(slice(None)), id="getitem"), + pytest.param( + lambda av: av.get_orthogonal_selection((slice(0, 3),)), id="get_orthogonal_selection" + ), + pytest.param(lambda av: av.get_coordinate_selection(coord), id="get_coordinate_selection"), + pytest.param(lambda av: av.get_mask_selection(mask), id="get_mask_selection"), + pytest.param(lambda av: av.oindex.getitem((slice(0, 3),)), id="oindex.getitem"), + pytest.param(lambda av: av.vindex.getitem(coord), id="vindex.getitem"), + ] + + +class TestAsyncSurfaceGuards: + def _view(self) -> tuple[zarr.Array[Any], npt.NDArray[Any]]: + """A base array and a non-identity async view ``arr.lazy[2:12]``'s async_array.""" + a = zarr.create_array(MemoryStore(), shape=(24,), chunks=(4,), dtype="i4") + ref = np.arange(24, dtype="i4") + a[...] = ref + return a, ref + + @pytest.mark.parametrize("op", _read_ops()) + def test_read_entry_points_raise(self, op: Callable[[Any], Any]) -> None: + """Every async read entry point raises ``LazyViewError`` on a lazy view.""" + a, _ = self._view() + av = a.lazy[2:12]._async_array + assert not av._is_identity + with pytest.raises(LazyViewError): + sync(op(av)) + + def test_setitem_raises_and_leaves_base_unchanged(self) -> None: + """``AsyncArray.setitem`` on a view raises and writes nothing to the base.""" + a, ref = self._view() + av = a.lazy[5:10]._async_array + with pytest.raises(LazyViewError): + sync(av.setitem(slice(0, 5), 99)) + np.testing.assert_array_equal(a[...], ref) + + def test_from_array_view_explicit_kwargs_raises_and_writes_nothing(self) -> None: + """``from_array`` with a lazy-view source raises and leaves the target empty.""" + a, _ = self._view() + view = a.lazy[5:10] + target = MemoryStore() + with pytest.raises(LazyViewError): + zarr.from_array(target, data=view, chunks=(5,), shards=None, overwrite=False) + assert target._store_dict == {} + + def test_from_array_view_default_kwargs_raises_by_design(self) -> None: + """The default ('keep') kwarg path raises the same clear error, not by accident.""" + a, _ = self._view() + view = a.lazy[5:10] + target = MemoryStore() + with pytest.raises(LazyViewError): + zarr.from_array(target, data=view) + assert target._store_dict == {} + + def test_translate_by_async_view_getitem_raises(self) -> None: + """A view produced by ``AsyncArray.translate_by`` is guarded too.""" + a, _ = self._view() + av = a._async_array.translate_by((3,)) + assert not av._is_identity + with pytest.raises(LazyViewError): + sync(av.getitem(slice(None))) + + # --- identity regression guards: none of the above may affect eager arrays --- + + def test_identity_getitem_still_reads(self) -> None: + a, ref = self._view() + got = sync(a._async_array.getitem(slice(2, 12))) + np.testing.assert_array_equal(got, ref[2:12]) + + def test_identity_setitem_still_writes(self) -> None: + a, ref = self._view() + sync(a._async_array.setitem(slice(0, 5), 99)) + expected = ref.copy() + expected[0:5] = 99 + np.testing.assert_array_equal(a[...], expected) + + def test_identity_from_array_still_copies(self) -> None: + a, ref = self._view() + target = MemoryStore() + out = zarr.from_array(target, data=a) + np.testing.assert_array_equal(out[...], ref) diff --git a/tests/test_properties.py b/tests/test_properties.py index 33888bfd4e..fe6949dda1 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -2,7 +2,7 @@ import json import numbers from collections.abc import Generator -from typing import Any +from typing import Any, Literal import numpy as np import pytest @@ -15,24 +15,32 @@ import hypothesis.extra.numpy as npst import hypothesis.strategies as st -from hypothesis import assume, given, settings +from hypothesis import assume, example, given, settings +from hypothesis.stateful import RuleBasedStateMachine, invariant, rule +from tests.conftest import Expect from zarr.abc.store import Store from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata from zarr.core.sync import sync from zarr.testing.strategies import ( + IndexingProgram, + IndexMode, + IndexOperation, array_metadata, arrays, basic_indices, block_indices, block_test_arrays, complex_rectilinear_arrays, + indexing_programs, + numpy_array_indexers, numpy_arrays, - orthogonal_indices, + program_shapes, rectilinear_arrays, simple_arrays, stores, + windows, zarr_formats, ) @@ -116,33 +124,10 @@ def test_array_creates_implicit_groups(array): ) -# this decorator removes timeout; not ideal but it should avoid intermittent CI failures - - -@settings(deadline=None) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -@given(data=st.data()) -async def test_basic_indexing(data: st.DataObject) -> None: - zarray = data.draw(st.one_of(simple_arrays(), rectilinear_arrays())) - nparray = zarray[:] - indexer = data.draw(basic_indices(shape=nparray.shape)) - - # sync get - actual = zarray[indexer] - assert_array_equal(nparray[indexer], actual) - - # async get - async_zarray = zarray._async_array - actual = await async_zarray.getitem(indexer) - assert_array_equal(nparray[indexer], actual) - - # sync set - new_data = data.draw(numpy_arrays(shapes=st.just(actual.shape), dtype=nparray.dtype)) - zarray[indexer] = new_data - nparray[indexer] = new_data - assert_array_equal(nparray, zarray[:]) - - # TODO test async setitem? +# Eager basic/oindex/vindex/mask indexing is exercised comprehensively (read, +# out=, async read, and write) by the single oracle `test_indexing_parity` +# defined below. Special array shapes that the oracle's strategies don't cover +# keep dedicated tests (complex rectilinear arrays here; block indexing later). @settings(deadline=None) @@ -154,105 +139,6 @@ async def test_basic_indexing_complex_rectilinear(data: st.DataObject) -> None: assert_array_equal(nparray[indexer], zarray[indexer]) -@given(data=st.data()) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -async def test_oindex(data: st.DataObject) -> None: - # integer_array_indices can't handle 0-size dimensions. - zarray = data.draw( - st.one_of( - simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)), - rectilinear_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1, max_side=20)), - ) - ) - nparray = zarray[:] - zindexer, npindexer = data.draw(orthogonal_indices(shape=nparray.shape)) - - # sync get - actual = zarray.oindex[zindexer] - assert_array_equal(nparray[npindexer], actual) - - # async get - async_zarray = zarray._async_array - actual = await async_zarray.oindex.getitem(zindexer) - assert_array_equal(nparray[npindexer], actual) - - # sync get - assume(zarray.shards is None) # GH2834 - for idxr in npindexer: - if isinstance(idxr, np.ndarray) and idxr.size != np.unique(idxr).size: - # behaviour of setitem with repeated indices is not guaranteed in practice - assume(False) - new_data = data.draw(numpy_arrays(shapes=st.just(actual.shape), dtype=nparray.dtype)) - nparray[npindexer] = new_data - zarray.oindex[zindexer] = new_data - assert_array_equal(nparray, zarray[:]) - - # note: async oindex setitem not yet implemented - - -@given(data=st.data()) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -async def test_vindex(data: st.DataObject) -> None: - # integer_array_indices can't handle 0-size dimensions. - zarray = data.draw( - st.one_of( - simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)), - rectilinear_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=20)), - ) - ) - nparray = zarray[:] - indexer = data.draw( - npst.integer_array_indices( - shape=nparray.shape, result_shape=npst.array_shapes(min_side=1, max_dims=None) - ) - ) - - # sync get - actual = zarray.vindex[indexer] - assert_array_equal(nparray[indexer], actual) - - # async get - async_zarray = zarray._async_array - actual = await async_zarray.vindex.getitem(indexer) - assert_array_equal(nparray[indexer], actual) - - # sync set - # FIXME! - # when the indexer is such that a value gets overwritten multiple times, - # I think the output depends on chunking. - # new_data = data.draw(npst.arrays(shape=st.just(actual.shape), dtype=nparray.dtype)) - # nparray[indexer] = new_data - # zarray.vindex[indexer] = new_data - # assert_array_equal(nparray, zarray[:]) - - # note: async vindex setitem not yet implemented - - -@settings(deadline=None) -@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") -@given(data=st.data()) -def test_mask_indexing(data: st.DataObject) -> None: - zarray = data.draw(st.one_of(simple_arrays(), rectilinear_arrays())) - nparray = zarray[:] - mask = data.draw(npst.arrays(dtype=np.bool_, shape=st.just(nparray.shape))) - - expected = nparray[mask] - - # sync get, via both the dedicated method and the vindex interface - assert_array_equal(expected, zarray.get_mask_selection(mask)) - assert_array_equal(expected, zarray.vindex[mask]) - - # sync set, via both interfaces - assume(zarray.shards is None) # GH2834 - new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) - nparray[mask] = new_data - zarray.set_mask_selection(mask, new_data) - assert_array_equal(nparray, zarray[:]) - - zarray.vindex[mask] = new_data - assert_array_equal(nparray, zarray[:]) - - @settings(deadline=None) @pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") @given(data=st.data()) @@ -448,3 +334,735 @@ def test_array_metadata_meets_spec(meta: ArrayV2Metadata | ArrayV3Metadata) -> N assert serialized_complex_float_is_valid(asdict_dict["fill_value"]) elif dtype_native.kind in ("M", "m") and np.isnat(meta.fill_value): assert asdict_dict["fill_value"] == -9223372036854775808 + + +# The indexing modes and which Array method implements each. vindex/mask are +# "vectorized" — they scatter through a single flat index, so an out= buffer must +# be flat (number of selected points) rather than the multi-dimensional result. +_INDEX_MODES: tuple[IndexMode, ...] = ("basic", "oindex", "vindex", "mask") +# Modes that scatter through a flat index (so an out= buffer must be flat). Kept a +# tuple like _INDEX_MODES; membership is checked against it below. +_VECTORIZED_MODES: tuple[IndexMode, ...] = ("vindex", "mask") + + +def _get(target: zarr.Array, mode: IndexMode, zsel: Any, *, out: Any = None) -> Any: + """Read `zsel` from `target` via the get-method for `mode`.""" + if mode == "basic": + return target.get_basic_selection(zsel, out=out) + if mode == "oindex": + return target.get_orthogonal_selection(zsel, out=out) + if mode == "vindex": + return target.get_coordinate_selection(zsel, out=out) + if mode == "mask": + return target.get_mask_selection(zsel, out=out) + raise AssertionError(mode) + + +def _async_get(async_array: Any, mode: IndexMode, zsel: Any) -> Any: + """The async read coroutine for `mode` (vindex/mask share the vectorized accessor).""" + if mode == "basic": + return async_array.getitem(zsel) + if mode == "oindex": + return async_array.oindex.getitem(zsel) + return async_array.vindex.getitem(zsel) + + +def _setitem(zarray: zarr.Array, mode: IndexMode, zsel: Any, value: Any) -> None: + """Write `value` at `zsel` via the set-method for `mode`.""" + if mode == "basic": + zarray[zsel] = value + elif mode == "oindex": + zarray.oindex[zsel] = value + elif mode == "vindex": + zarray.vindex[zsel] = value + elif mode == "mask": + zarray.set_mask_selection(zsel, value) + else: + raise AssertionError(mode) + + +def _write_is_unambiguous(mode: IndexMode, zsel: Any, shape: tuple[int, ...]) -> bool: + """Whether a write to `zsel` targets each cell at most once. + + `zsel` must be the *raw zarr selection* (per-axis arrays/slices), NOT the + broadcast np.ix_-style form NumPy needs: for oindex the broadcast form tiles + each component across every axis, so the per-axis uniqueness check below would + trip on tiling instead of genuine duplicates and reject well-defined writes. + + Negative indices are normalized first (`[2, -2]` on length 4 both target cell + 2), then duplicate targets are detected. A repeated target makes the surviving + value depend on the order the writes are applied. NumPy happens to be + deterministic here (last-write-wins), but only as an implementation detail of + iterating the index serially in C order — it is not a guarantee. zarr writes + scalars into chunks and cannot in general promise a write order, so it cannot + be expected to match NumPy's choice. There is therefore no well-defined oracle + for such a write, and we reject it (a fundamental limitation, not a bug). + """ + sel = zsel if isinstance(zsel, tuple) else (zsel,) + if mode == "oindex": + # Independent axes: each fancy axis must select each index at most once. + for axis, s in enumerate(sel): + if isinstance(s, np.ndarray): + norm = np.where(s < 0, s + shape[axis], s) + if norm.size != np.unique(norm).size: + return False + return True + if mode == "vindex": + # Correlated coordinate arrays: a target is the tuple of normalized coords, + # so the write is well-defined iff no coordinate tuple repeats. + norm = [ + np.where(s < 0, s + shape[axis], s) + for axis, s in enumerate(sel) + if isinstance(s, np.ndarray) + ] + if len(norm) == 0: + return True + coords = np.stack([a.ravel() for a in np.broadcast_arrays(*norm)], axis=-1) + return len(coords) == len(np.unique(coords, axis=0)) + return True # basic / mask never target a cell twice + + +def _n_array_axes(zsel: Any) -> int: + """Number of integer-array (fancy) axes in a *raw* zarr selection (slices excluded). + + Must be the raw per-axis selection, not the broadcast numpy form: the latter + turns every axis into an array, so this would always return ``ndim``. + """ + sel = zsel if isinstance(zsel, tuple) else (zsel,) + return sum(isinstance(s, np.ndarray) for s in sel) + + +def _eligible(mode: IndexMode, shape: tuple[int, ...]) -> bool: + """Whether `mode` can be exercised on `shape`. + + Rank-0 arrays have no interesting selections. `basic` and `mask` handle + zero-length axes (a boolean mask of the array's shape simply selects nothing, + which zarr and NumPy agree on — verified against `set_mask_selection` too). + The integer-array fancy modes (`oindex`/`vindex`) cannot address an empty + axis, so they require every axis to be non-empty. + """ + if len(shape) == 0: + return False + if mode in ("basic", "mask"): + return True + return all(s > 0 for s in shape) + + +def assert_read_matches_numpy( + target: zarr.Array, + ref: np.ndarray[Any, Any], + mode: IndexMode, + zsel: Any, + npsel: Any, + *, + out_layout: Literal["eager", "transform"] = "eager", +) -> None: + """Assert `target`'s read of `zsel` (mode) matches `ref[npsel]`, with/without out=. + + Consumer-agnostic: `target` is any `zarr.Array` (an eager array or a lazy + view) and `ref` is the NumPy array it must behave like. The two consumers + disagree on the `out=` buffer shape for the vectorized vindex/mask modes, so + `out_layout` names which convention `target` implements: + + - `"eager"` (legacy path): vindex/mask scatter through a flat index, so the + buffer is flat — `(number of selected points,)`. + - `"transform"` (lazy views, Task 4): `out=` takes the broadcast selection + shape for every mode, so the buffer always has the result shape. + """ + expected = np.asarray(ref[npsel]) + actual = np.asarray(_get(target, mode, zsel)) + # dtype must match, not just values (catches silent dtype coercion); ndindex's + # assert_equal does the same. Guarded to ndim>=1 to avoid python-vs-numpy + # scalar dtype noise on 0-d results. + if expected.ndim >= 1: + assert actual.dtype == expected.dtype, (actual.dtype, expected.dtype) + assert_array_equal(actual, expected) + if expected.ndim >= 1 and expected.size > 0: + if out_layout == "eager" and mode in _VECTORIZED_MODES: + buf_shape: tuple[int, ...] = (expected.size,) + else: + buf_shape = expected.shape + buf = default_buffer_prototype().nd_buffer.empty(shape=buf_shape, dtype=expected.dtype) + _get(target, mode, zsel, out=buf) + assert_array_equal(np.asarray(buf.as_ndarray_like()).reshape(expected.shape), expected) + + +def _indexing_array(data: st.DataObject) -> zarr.Array: + """An eager array (regular or rectilinear) for the indexing-parity oracles.""" + return data.draw( + st.one_of( + # min_side=0 restores zero-extent-array coverage: the `_eligible` helper + # routes zero-size shapes to the modes that tolerate them (basic/mask). + simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=0)), + rectilinear_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=20)), + ) + ) + + +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +async def test_indexing_read_parity(data: st.DataObject) -> None: + """Every indexing *read* on an eager array matches NumPy (all modes). + + Covers {basic, oindex, vindex, mask}: sync read, read into an out= buffer, and + async read. No special-casing — reads are well-defined for every selection. + Consumer-agnostic via assert_read_matches_numpy, which test_lazy_view_indexing_parity + reuses for lazy views. + """ + zarray = _indexing_array(data) + nparray = zarray[:] + mode = data.draw(st.sampled_from(_INDEX_MODES)) + assume(_eligible(mode, nparray.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=nparray.shape)) + + assert_read_matches_numpy(zarray, nparray, mode, zsel, npsel) + expected = np.asarray(nparray[npsel]) + assert_array_equal(np.asarray(await _async_get(zarray._async_array, mode, zsel)), expected) + if mode == "mask": # a mask read may also be spelled vindex[mask]; they must agree + assert_array_equal(np.asarray(zarray.vindex[zsel]), expected) + + +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +async def test_indexing_write_parity(data: st.DataObject) -> None: + """Every *supported, well-defined* indexing write on an eager array matches NumPy. + + Two families of selections are filtered out with `assume` (they are not + failures to test here): + + - **Repeated coordinates** (vindex/oindex selecting a cell more than once): + the write order, and thus the final value, is undefined — NumPy has the same + ambiguity. This is a fundamental limitation, not a zarr bug. + - **Sharded orthogonal writes across >= 2 array axes**: unsupported by the + sharding partial-write codec; pinned at the pytest level by + `test_oindex_sharded_multiaxis_write_xfail` rather than hidden here. + """ + zarray = _indexing_array(data) + nparray = zarray[:] + mode = data.draw(st.sampled_from(_INDEX_MODES)) + assume(_eligible(mode, nparray.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=nparray.shape)) + if mode in ("oindex", "vindex"): + assume(_write_is_unambiguous(mode, zsel, nparray.shape)) + if mode == "oindex" and zarray.shards is not None: + assume(_n_array_axes(zsel) < 2) + + expected = np.asarray(nparray[npsel]) + new_data = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=nparray.dtype)) + _setitem(zarray, mode, zsel, new_data) + nparray[npsel] = new_data + assert_array_equal(nparray, zarray[:]) + if mode == "mask": # the vindex[mask] = ... spelling must agree with set_mask_selection + zarray.vindex[zsel] = new_data + assert_array_equal(nparray, zarray[:]) + + +@pytest.mark.xfail( + reason="GH2834: the sharding partial-write codec does not support orthogonal " + "writes spanning two or more array axes", + strict=True, +) +def test_oindex_sharded_multiaxis_write_xfail() -> None: + """Pin the one known-unsupported write: oindex across >= 2 array axes on a shard. + + Single-axis oindex, vindex, and mask writes to sharded arrays all work; only + the multi-array-axis orthogonal case is broken. Strict xfail so this flips to a + failure (prompting removal) if/when GH2834 is fixed. + """ + a = zarr.create_array({}, shape=(12, 12), chunks=(3, 3), shards=(6, 6), dtype="i4") + a[...] = np.arange(144, dtype="i4").reshape(12, 12) + i0, i1 = np.array([0, 5, 9]), np.array([1, 6, 10]) + a.set_orthogonal_selection((i0, i1), np.zeros((3, 3), dtype="i4")) + assert_array_equal(a.oindex[i0, i1], np.zeros((3, 3), dtype="i4")) + + +@pytest.mark.parametrize( + "case", + [ + Expect( + input=("vindex", (np.array([2, -2, 0, 1]),), (4,)), + output=False, + id="vindex-neg-collision", + ), + Expect(input=("vindex", (np.array([0, 1, 2]),), (4,)), output=True, id="vindex-distinct"), + Expect( + input=("oindex", (np.array([1, -3]),), (4,)), output=False, id="oindex-neg-collision" + ), + Expect( + input=("oindex", (np.array([0, 2]), np.array([1, 1])), (4, 4)), + output=False, + id="oindex-repeat", + ), + Expect( + input=("oindex", (np.array([0, 2]), np.array([1, 3])), (4, 4)), + output=True, + id="oindex-distinct", + ), + Expect(input=("basic", (slice(None),), (4,)), output=True, id="basic-always"), + ], + ids=lambda c: c.id, +) +def test_write_is_unambiguous( + case: Expect[tuple[IndexMode, Any, tuple[int, ...]], bool], +) -> None: + """A write is well-defined iff each target cell is hit at most once *after* + negative indices are normalized. + + Cases use the *raw zarr selection* (per-axis arrays) — the form real callers + pass — so `oindex-distinct` ([0,2] x [1,3]) is well-defined. See + `test_write_filter_rejects_broadcast_oindex` for why the broadcast form must not + be passed. Regression anchor for the CI fix: vindex `[2, -2, 0, 1]` on length 4 + maps `-2 -> 2`, so cell 2 is written twice (order-dependent) and must be rejected + even though the raw values are all distinct. + """ + mode, zsel, shape = case.input + assert _write_is_unambiguous(mode, zsel, shape) is case.output + + +def test_write_filter_rejects_broadcast_oindex() -> None: + """The write filter must be fed the raw per-axis selection, not the broadcast form. + + `numpy_array_indexers(mode="oindex")` returns a raw per-axis `zsel` and a broadcast + np.ix_-style `npsel`. A well-defined multi-axis write ([0,2] x [1,3]) is admitted + on the raw form but — because the broadcast form tiles each component across all + axes — is wrongly rejected on `npsel`. This pins the review finding that passing + `npsel` silently skipped essentially every multi-axis orthogonal write. + """ + zsel = (np.array([0, 2]), np.array([1, 3])) # raw per-axis, distinct -> well-defined + assert _write_is_unambiguous("oindex", zsel, (4, 4)) is True + # the fully-tiled form orthogonal_indices produces (np.broadcast_arrays over np.ix_) + broadcast = tuple(np.broadcast_arrays(*np.ix_(*zsel))) + assert _write_is_unambiguous("oindex", broadcast, (4, 4)) is False # tiling; do not pass + # _n_array_axes counts only genuine fancy axes on the raw form (slices excluded), + # so a single-fancy-axis oindex is correctly allowed on a sharded array. + assert _n_array_axes((slice(None), np.array([1, 3]))) == 1 + + +@st.composite +def _bounds_selection(draw: st.DrawFn, mode: IndexMode, shape: tuple[int, ...]) -> Any: + """An integer selection whose components may fall outside `[-size, size)`, to + probe out-of-bounds behavior. Returns `(zarr_sel, numpy_sel)`.""" + wide = {s: st.integers(-2 * s - 1, 2 * s) for s in set(shape)} + if mode == "basic": + sel = tuple(draw(wide[s]) for s in shape) + return sel, sel + n = draw(st.integers(1, 4)) + arrs = tuple( + np.array(draw(st.lists(wide[s], min_size=n, max_size=n)), dtype=np.intp) for s in shape + ) + if mode == "oindex": + return arrs, np.ix_(*arrs) # numpy outer-product oracle + return arrs, arrs # vindex: pointwise + + +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +def test_indexing_bounds_error_parity(data: st.DataObject) -> None: + """Out-of-bounds integer indexing raises in zarr iff it raises in NumPy. + + zarr's bounds errors (BoundsCheckError etc.) subclass IndexError, so the eager + methods must agree with NumPy on *whether* an index is in bounds. This is the + ndindex `check_same` error-parity idea, and it catches the silent-wrong-data + class — returning garbage where NumPy would raise. + """ + zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=3, min_side=1, max_side=6))) + nparray = zarray[:] + mode = data.draw(st.sampled_from(("basic", "oindex", "vindex"))) + zsel, npsel = data.draw(_bounds_selection(mode, nparray.shape)) + try: + expected = np.asarray(nparray[npsel]) + except IndexError: + with pytest.raises(IndexError): + _get(zarray, mode, zsel) + return + assert_array_equal(np.asarray(_get(zarray, mode, zsel)), expected) + + +class _IndexingStateMachine(RuleBasedStateMachine): + """Apply many random indexing writes to one array while a NumPy model tracks it + in lockstep; the model must match after every step. + + This is the TensorStore `driver_testutil` pattern with NumPy as the oracle: + a single array accumulates many writes, so it catches read-modify-write, + chunk-boundary, and shard-merge/persistence bugs that the single-shot + `test_indexing_write_parity` cannot see. Concrete geometries are defined by + subclasses; the rules are dtype-agnostic (fixed i4 — dtype breadth lives in the + property tests). + """ + + shape: tuple[int, ...] = () + chunks: tuple[int, ...] = () + shards: tuple[int, ...] | None = None + + def __init__(self) -> None: + super().__init__() + n = int(np.prod(self.shape, dtype=int)) + self.model = np.arange(n, dtype="i4").reshape(self.shape) + self.zarray = zarr.create_array( + {}, shape=self.shape, chunks=self.chunks, shards=self.shards, dtype="i4" + ) + self.zarray[...] = self.model + + @rule(data=st.data(), mode=st.sampled_from(_INDEX_MODES)) + def write(self, data: st.DataObject, mode: IndexMode) -> None: + if not _eligible(mode, self.shape): + return + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=self.shape)) + if mode in ("oindex", "vindex") and not _write_is_unambiguous(mode, zsel, self.shape): + return + if mode == "oindex" and self.shards is not None and _n_array_axes(zsel) >= 2: + return # GH2834, see test_oindex_sharded_multiaxis_write_xfail + expected = np.asarray(self.model[npsel]) + value = data.draw(numpy_arrays(shapes=st.just(expected.shape), dtype=self.model.dtype)) + self.model[npsel] = value + _setitem(self.zarray, mode, zsel, value) + + @rule(data=st.data(), mode=st.sampled_from(_INDEX_MODES)) + def read(self, data: st.DataObject, mode: IndexMode) -> None: + if not _eligible(mode, self.shape): + return + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=self.shape)) + assert_array_equal(np.asarray(_get(self.zarray, mode, zsel)), np.asarray(self.model[npsel])) + + @invariant() + def matches_model(self) -> None: + assert_array_equal(self.zarray[:], self.model) + + +class _RegularStateMachine(_IndexingStateMachine): + shape, chunks, shards = (12,), (3,), None + + +class _Sharded2DStateMachine(_IndexingStateMachine): + shape, chunks, shards = (8, 9), (2, 3), (4, 9) + + +class _ThreeDStateMachine(_IndexingStateMachine): + shape, chunks, shards = (4, 5, 6), (2, 3, 3), None + + +# Run these under the slow-hypothesis CI job (like the store stateful tests): a +# stateful run is a sequence of steps, so it is heavier than the single-shot tests. +TestIndexingStateMachineRegular = pytest.mark.slow_hypothesis(_RegularStateMachine.TestCase) +TestIndexingStateMachineSharded = pytest.mark.slow_hypothesis(_Sharded2DStateMachine.TestCase) +TestIndexingStateMachine3D = pytest.mark.slow_hypothesis(_ThreeDStateMachine.TestCase) + + +def _canonical_nonneg(sel: Any, shape: tuple[int, ...]) -> Any: + """NumPy-normalize a selection to its non-negative canonical form. + + Wraps negative integers and integer-array values, resolves slice bounds via + Python slice semantics, and expands ellipsis — producing the literal + spelling of the same NumPy selection, as the lazy path requires. + """ + items = sel if isinstance(sel, tuple) else (sel,) + n_explicit = sum(1 for s in items if s is not Ellipsis) + out: list[Any] = [] + dim = 0 + for s in items: + if s is Ellipsis: + for _ in range(len(shape) - n_explicit): + out.append(slice(0, shape[dim], 1)) + dim += 1 + elif isinstance(s, (int, np.integer)) and not isinstance(s, (bool, np.bool_)): + v = int(s) + out.append(v + shape[dim] if v < 0 else v) + dim += 1 + elif isinstance(s, slice): + start, stop, step = s.indices(shape[dim]) + if stop < start: # NumPy empty-by-reversal -> literal empty interval + stop = start + out.append(slice(start, stop, step)) + dim += 1 + elif isinstance(s, np.ndarray) and s.dtype != np.bool_: + out.append(np.where(s < 0, s + shape[dim], s)) + dim += 1 + else: # boolean mask or anything already canonical + out.append(s) + dim += 1 + return tuple(out) if isinstance(sel, tuple) else out[0] + + +@pytest.mark.skipif( + not hasattr(zarr.Array, "lazy"), reason="lazy indexing (Array.lazy) not available" +) +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +@given(data=st.data()) +async def test_lazy_view_indexing_parity(data: st.DataObject) -> None: + """The eager read oracle, applied to a re-zeroed lazy view (the lazy consumer). + + A view is just another `zarr.Array`, so it flows through the same + `assert_read_matches_numpy` harness. Views preserve their domain + (TensorStore semantics: coordinates are literal), so to compare against a + zero-based NumPy reference the view is explicitly re-zeroed with + `translate_to` — exercising domain translation on every example. Skipped + until `Array.lazy` exists, so the eager oracle can merge ahead of the lazy + feature. + """ + # min_side=0 exercises the `windows()` zero-extent branch and, via `_eligible`, + # zero-size lazy-view reads in the basic/mask modes that tolerate them. + zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=0))) + nparray = zarray[:] + window = data.draw(windows(shape=nparray.shape)) + view = zarray.lazy[window].translate_to((0,) * nparray.ndim) + vref = nparray[window] + mode = data.draw(st.sampled_from(_INDEX_MODES)) + assume(_eligible(mode, vref.shape)) + zsel, npsel = data.draw(numpy_array_indexers(mode=mode, shape=vref.shape)) + # The strategies draw NumPy-dialect selections. The lazy boundary wraps + # negatives against the (here zero-origin) view, matching NumPy, but positive + # indices are literal domain coordinates; canonicalizing to the equivalent + # non-negative form keeps the comparison valid for both rules. NumPy semantics + # are invariant under this, so the reference side (npsel) is unchanged. + # Lazy views take out= in the broadcast result shape for every mode (Task 4) + # — *except* when the composed transform is an identity, where the read + # methods fast-path to the legacy implementation and its flat-vindex out= + # convention applies. Mirror that dispatch here. + out_layout: Literal["eager", "transform"] = ( + "eager" if view._async_array._is_identity else "transform" + ) + assert_read_matches_numpy( + view, vref, mode, _canonical_nonneg(zsel, vref.shape), npsel, out_layout=out_layout + ) + + +# --------------------------------------------------------------------------- +# Composed indexing programs (Task 5): a differential oracle over sequences of +# composed indexing operations. Where the single-shot parity tests above check +# one selection at a time, these compose one to three operations through the real +# index-transform layer and check the composed result — and each of five +# execution recipes — against a NumPy model applied in lockstep. +# --------------------------------------------------------------------------- + +# Public program dialect -> the internal accessor mode used by _get/_setitem. +_PROGRAM_INDEX_MODE: dict[str, IndexMode] = { + "basic": "basic", + "orthogonal": "oindex", + "vectorized": "vindex", +} + + +def _program_selections(op: IndexOperation) -> tuple[Any, Any]: + """The ``(zarr_selection, numpy_selection)`` pair for one program step. + + Only ``orthogonal`` needs distinct spellings (per-axis vs the ``np.ix_`` + outer-product form); ``basic`` and ``vectorized`` index both sides alike. + """ + if op.mode == "orthogonal": + zsel, npsel = op.selection + return zsel, npsel + return op.selection, op.selection + + +def _apply_lazy_step(view: zarr.Array, mode: str, zsel: Any) -> zarr.Array: + """Compose one lazy indexing step onto ``view`` and re-zero the result. + + Re-zeroing (``translate_to``) keeps the next step's NumPy-dialect selection + valid against a zero-origin domain (positive coordinates stay literal, so they + match NumPy) while still exercising domain translation on every step. + """ + if mode == "basic": + out = view.lazy[zsel] + elif mode == "orthogonal": + out = view.lazy.oindex[zsel] + elif mode == "vectorized": + out = view.lazy.vindex[zsel] + else: + raise AssertionError(mode) + if out.ndim: + out = out.translate_to((0,) * out.ndim) + return out + + +def _compare_read(actual_raw: Any, expected: np.ndarray[Any, Any]) -> None: + """Assert a program read matches NumPy in scalar/array kind, shape, dtype, values.""" + actual = np.asarray(actual_raw) + # Scalar-versus-array kind: a zero-rank read must come back as a scalar, a + # ranked read as an array. np.ndim reads 0 for python/numpy scalars. + assert np.ndim(actual_raw) == expected.ndim, ("kind", np.ndim(actual_raw), expected.ndim) + if expected.ndim >= 1: # skip 0-d to avoid python-vs-numpy scalar dtype noise + assert actual.dtype == expected.dtype, (actual.dtype, expected.dtype) + assert actual.shape == expected.shape, (actual.shape, expected.shape) + assert_array_equal(actual, expected) + + +def _compose_base( + zarray: zarr.Array, nparray: np.ndarray[Any, Any], operations: tuple[IndexOperation, ...] +) -> tuple[zarr.Array, np.ndarray[Any, Any]]: + """Compose every operation *except the last* as lazy views, in lockstep with NumPy. + + Returns the composed (zero-origin) lazy base view and the matching NumPy + array; their shapes are asserted equal at each step. The prefix operations are + always basic (fancy steps only ever occur last), so the NumPy side stays a + view — which matters for the write executions, where assignment through it must + reach the root array. + """ + base = zarray + ref = nparray + for op in operations[:-1]: + zsel, npsel = _program_selections(op) + base = _apply_lazy_step(base, op.mode, _canonical_nonneg(zsel, base.shape)) + ref = ref[npsel] + assert base.shape == ref.shape, (base.shape, ref.shape) + return base, ref + + +def _run_program_read( + zarray: zarr.Array, nparray: np.ndarray[Any, Any], program: IndexingProgram +) -> None: + base, ref = _compose_base(zarray, nparray, program.operations) + last = program.operations[-1] + imode = _PROGRAM_INDEX_MODE[last.mode] + zsel, npsel = _program_selections(last) + czsel = _canonical_nonneg(zsel, base.shape) + expected = np.asarray(ref[npsel]) + + if program.execution == "eager_on_lazy": + # Apply the final step as an *eager* indexing read on the composed view. + _compare_read(_get(base, imode, czsel), expected) + return + + # materialize / out: compose the final step lazily too, then read the whole view. + view = _apply_lazy_step(base, last.mode, czsel) + _compare_read(view[...], expected) + if program.execution == "out" and expected.ndim >= 1 and expected.size > 0: + # A whole-view basic read: the transform path takes the full result shape + # (Task 4), for every mode, so the buffer is never the flat vectorized form. + buf = default_buffer_prototype().nd_buffer.empty(shape=expected.shape, dtype=expected.dtype) + view.get_basic_selection(Ellipsis, out=buf) + assert_array_equal(np.asarray(buf.as_ndarray_like()), expected) + + +def _run_program_write( + zarray: zarr.Array, nparray: np.ndarray[Any, Any], program: IndexingProgram +) -> None: + root = nparray.copy() # the NumPy model we mutate and compare against + base, sub = _compose_base(zarray, root, program.operations) + last = program.operations[-1] + imode = _PROGRAM_INDEX_MODE[last.mode] + zsel, npsel = _program_selections(last) + czsel = _canonical_nonneg(zsel, base.shape) + if last.mode in ("orthogonal", "vectorized"): + # A fancy write hitting a cell twice is order-dependent and has no oracle + # (same limitation NumPy has); program arrays are unsharded so GH2834 (the + # only other write restriction) cannot arise. + assume(_write_is_unambiguous(imode, czsel, base.shape)) + + result_shape = sub[npsel].shape + value: Any + if program.execution == "set_scalar": + value = np.array(-7, dtype=root.dtype) + else: # set_array: distinct (negative) values expose scatter/order bugs + value = (-(np.arange(int(np.prod(result_shape)), dtype="i8") + 1)).reshape(result_shape) + value = value.astype(root.dtype) + + sub[npsel] = value # writes through the basic-view chain into `root` + _setitem(base, imode, czsel, value) + assert_array_equal(zarray[:], root) + + +def _build_program_array(shape: tuple[int, ...]) -> tuple[zarr.Array, np.ndarray[Any, Any]]: + """A deterministic i4 array filled with ``arange``, chunked ~2-per-axis. + + Fixed dtype and a regular (unsharded) grid keep the program oracle focused on + composition and execution correctness — dtype breadth and rectilinear/sharded + grids are covered by the single-shot parity tests and the state machines. The + ~2-chunks-per-axis geometry makes composed selections routinely cross chunk + boundaries. + """ + chunks = tuple(max(1, (s + 1) // 2) for s in shape) + nparray = np.arange(int(np.prod(shape, dtype=int)), dtype="i4").reshape(shape) + zarray = zarr.create_array({}, shape=shape, chunks=chunks, dtype="i4") + zarray[...] = nparray + return zarray, nparray + + +@st.composite +def _indexing_scenarios(draw: st.DrawFn) -> tuple[tuple[int, ...], IndexingProgram]: + """Draw a ``(shape, program)`` scenario. Kept together so `@example` can pin both.""" + shape = draw(program_shapes) + program = draw(indexing_programs(shape=shape)) + return shape, program + + +def _orth(*axes: Any) -> tuple[tuple[np.ndarray[Any, Any], ...], Any]: + """Build an orthogonal selection ``(zarr_per_axis, numpy_ix_form)`` for `@example`.""" + arrs = tuple(np.asarray(a) for a in axes) + return arrs, np.ix_(*arrs) + + +@pytest.mark.skipif( + not hasattr(zarr.Array, "lazy"), reason="lazy indexing (Array.lazy) not available" +) +@settings(deadline=None) +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +# slice-then-oindex: the canonical composed-fancy regression. Reversing the Task 2 +# correlation predicate makes this example fail (mutation check). +@example( + scenario=( + (4, 4), + IndexingProgram( + ( + IndexOperation("basic", (slice(1, 4), slice(None))), + IndexOperation("orthogonal", _orth([0, 2], [1, 3])), + ), + "materialize", + ), + ) +) +# unequal orthogonal lengths: the outer product must not assume square selections. +@example( + scenario=( + (4, 4), + IndexingProgram((IndexOperation("orthogonal", _orth([0, 1, 2], [0, 1])),), "materialize"), + ) +) +# multidimensional vindex read into an out= buffer. +@example( + scenario=( + (4, 4), + IndexingProgram( + ( + IndexOperation( + "vectorized", (np.array([[0, 1], [2, 3]]), np.array([[1, 0], [3, 2]])) + ), + ), + "out", + ), + ) +) +# zero-rank result: a full integer selection collapses to a scalar. +@example( + scenario=( + (3, 3), + IndexingProgram((IndexOperation("basic", (1, 2)),), "materialize"), + ) +) +# empty result: a 0-length slice yields a zero-size read. +@example( + scenario=( + (4, 4), + IndexingProgram((IndexOperation("basic", (slice(0, 0), slice(None))),), "materialize"), + ) +) +@given(scenario=_indexing_scenarios()) +def test_indexing_program_parity(scenario: tuple[tuple[int, ...], IndexingProgram]) -> None: + """A composed indexing program matches NumPy under every execution recipe. + + Reads (`materialize`, `eager_on_lazy`, `out`) check scalar/array kind, shape, + dtype and values; writes (`set_scalar`, `set_array`) apply the equivalent + assignment to a NumPy root and compare the whole zarr array afterwards. The + generator composes at most one fancy step and always keeps it last — see + `indexing_programs` for the documented exclusions (fancy-after-fancy is + unsupported and raises `NotImplementedError`; basic-after-fancy is + known-broken; both are deliberately not generated). + """ + shape, program = scenario + zarray, nparray = _build_program_array(shape) + if program.execution in ("set_scalar", "set_array"): + _run_program_write(zarray, nparray, program) + else: + _run_program_read(zarray, nparray, program) diff --git a/uv.lock b/uv.lock index 7a0d713873..e82bbababe 100644 --- a/uv.lock +++ b/uv.lock @@ -6,6 +6,12 @@ resolution-markers = [ "python_full_version < '3.15'", ] +[manifest] +members = [ + "zarr", + "zarr-indexing", +] + [[package]] name = "aiobotocore" version = "3.7.0" @@ -607,6 +613,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ae/44/c1221527f6a71a01ec6fbad7fa78f1d50dfa02217385cf0fa3eec7087d59/click-8.3.3-py3-none-any.whl", hash = "sha256:a2bf429bb3033c89fa4936ffb35d5cb471e3719e1f3c8a7c3fff0b8314305613", size = 110502, upload-time = "2026-04-22T15:11:25.044Z" }, ] +[[package]] +name = "cloudpickle" +version = "3.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -781,6 +796,29 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a4/80/5e05de89ba61df072aab6f8a6ee3ffeec57db68a0a456825b3b4ce608426/cupy_cuda12x-14.1.1-cp314-cp314t-manylinux2014_x86_64.whl", hash = "sha256:238080487174268d0f09770fe518de7c5b206527bef5c6792aef7ba0626a1c48", size = 132635338, upload-time = "2026-06-01T04:53:35.229Z" }, ] +[[package]] +name = "dask" +version = "2026.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "cloudpickle" }, + { name = "fsspec" }, + { name = "packaging" }, + { name = "partd" }, + { name = "pyyaml" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d6/39/cbd21c9133d02b4e60899ed466ad5e876553ea68ebffee6209e7dcafc8d4/dask-2026.7.1.tar.gz", hash = "sha256:5727484427665f051e86bf87d021a64d6411141cdc8a20bfe3c1ad2968cc06b7", size = 11548794, upload-time = "2026-07-14T01:06:22.46Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/5f/7c22733da92b3a6cc4dddcaa8731089d213c2790bbc997e51c429a4e8f8b/dask-2026.7.1-py3-none-any.whl", hash = "sha256:985ffd6c5e9d7979ede515e84ae8d39b647d6aa64f77600f15714ff65f578fe6", size = 1496882, upload-time = "2026-07-14T01:06:20.341Z" }, +] + +[package.optional-dependencies] +array = [ + { name = "numpy" }, +] + [[package]] name = "defusedxml" version = "0.7.1" @@ -1271,6 +1309,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ce/62/b40b382fa0c66fee1478073eb8db352a4a6beda4a1adccf1df911d8c289c/librt-0.11.0-cp314-cp314t-win_arm64.whl", hash = "sha256:dee008f20b542e3cd162ba338a7f9ec0f6d23d395f66fe8aeeec3c9d067ea253", size = 102572, upload-time = "2026-05-10T18:17:06.809Z" }, ] +[[package]] +name = "locket" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2f/83/97b29fe05cb6ae28d2dbd30b81e2e402a3eed5f460c26e9eaa5895ceacf5/locket-1.0.0.tar.gz", hash = "sha256:5c0d4c052a8bbbf750e056a8e65ccd309086f4f0f18a2eac306a8dfa4112a632", size = 4350, upload-time = "2022-04-20T22:04:44.312Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/db/bc/83e112abc66cd466c6b83f99118035867cecd41802f8d044638aa78a106e/locket-1.0.0-py2.py3-none-any.whl", hash = "sha256:b6c819a722f7b6bd955b80781788e4a66a55628b858d347536b7e81325a3a5e3", size = 4398, upload-time = "2022-04-20T22:04:42.23Z" }, +] + [[package]] name = "markdown" version = "3.10.2" @@ -1986,6 +2033,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/90/96/04b8e52da071d28f5e21a805b19cb9390aa17a47462ac87f5e2696b9566d/paginate-0.5.7-py2.py3-none-any.whl", hash = "sha256:b885e2af73abcf01d9559fd5216b57ef722f8c42affbb63942377668e35c7591", size = 13746, upload-time = "2024-08-25T14:17:22.55Z" }, ] +[[package]] +name = "partd" +version = "1.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "locket" }, + { name = "toolz" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b2/3a/3f06f34820a31257ddcabdfafc2672c5816be79c7e353b02c1f318daa7d4/partd-1.4.2.tar.gz", hash = "sha256:d022c33afbdc8405c226621b015e8067888173d85f7f5ecebb3cafed9a20f02c", size = 21029, upload-time = "2024-05-06T19:51:41.945Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" }, +] + [[package]] name = "pathable" version = "0.5.0" @@ -3052,6 +3112,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6a/43/8bd850ee71a191bf072e31302c73a66be413fecdd98fdcd111ecbcce13ca/tomlkit-0.15.0-py3-none-any.whl", hash = "sha256:4dbc8f0fc024412b57ced8757ac7461305126a648ff8c2c807fcb8e133a78738", size = 41328, upload-time = "2026-05-10T07:38:23.517Z" }, ] +[[package]] +name = "toolz" +version = "1.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/11/d6/114b492226588d6ff54579d95847662fc69196bdeec318eb45393b24c192/toolz-1.1.0.tar.gz", hash = "sha256:27a5c770d068c110d9ed9323f24f1543e83b2f300a687b7891c1a6d56b697b5b", size = 52613, upload-time = "2025-10-17T04:03:21.661Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/12/5911ae3eeec47800503a238d971e51722ccea5feb8569b735184d5fcdbc0/toolz-1.1.0-py3-none-any.whl", hash = "sha256:15ccc861ac51c53696de0a5d6d4607f99c210739caf987b5d2054f3efed429d8", size = 58093, upload-time = "2025-10-17T04:03:20.435Z" }, +] + [[package]] name = "towncrier" version = "25.8.0" @@ -3390,6 +3459,7 @@ dependencies = [ { name = "numpy" }, { name = "packaging" }, { name = "typing-extensions" }, + { name = "zarr-indexing" }, ] [package.optional-dependencies] @@ -3461,6 +3531,9 @@ docs = [ { name = "s3fs" }, { name = "towncrier" }, ] +interop-tests = [ + { name = "dask", extra = ["array"] }, +] release = [ { name = "towncrier" }, ] @@ -3513,6 +3586,7 @@ requires-dist = [ { name = "typer", marker = "extra == 'cli'" }, { name = "typing-extensions", specifier = ">=4.14" }, { name = "universal-pathlib", marker = "extra == 'optional'" }, + { name = "zarr-indexing", editable = "packages/zarr-indexing" }, ] provides-extras = ["cast-value-rs", "cli", "gpu", "optional", "remote"] @@ -3567,6 +3641,7 @@ docs = [ { name = "s3fs", specifier = ">=2023.10.0" }, { name = "towncrier", specifier = "==25.8.0" }, ] +interop-tests = [{ name = "dask", extras = ["array"] }] release = [{ name = "towncrier", specifier = "==25.8.0" }] remote-tests = [ { name = "botocore" }, @@ -3602,3 +3677,21 @@ test = [ { name = "tomlkit", specifier = "==0.15.0" }, { name = "uv", specifier = "==0.11.26" }, ] + +[[package]] +name = "zarr-indexing" +source = { editable = "packages/zarr-indexing" } +dependencies = [ + { name = "numpy" }, +] + +[package.dev-dependencies] +test = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [{ name = "numpy", specifier = ">=2" }] + +[package.metadata.requires-dev] +test = [{ name = "pytest" }]