diff --git a/docs/superpowers/plans/2026-07-15-tensorstore-index-transform-corrections.md b/docs/superpowers/plans/2026-07-15-tensorstore-index-transform-corrections.md new file mode 100644 index 0000000000..438fbea082 --- /dev/null +++ b/docs/superpowers/plans/2026-07-15-tensorstore-index-transform-corrections.md @@ -0,0 +1,319 @@ +# TensorStore-Aligned Index Transform Corrections Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make transform-based Zarr indexing preserve eager NumPy/Zarr semantics across lazy composition, orthogonal and vectorized indexing, zero-rank arrays, caller-provided output buffers, and writes. + +**Architecture:** Normalize every `ArrayMap.index_array` over the complete transform input domain, following TensorStore. Derive independent versus correlated indexing from array dependency axes, preserve those axes through composition and chunk intersection, and keep public Zarr negative-index normalization at the array-selection boundary. Verify behavior with focused regressions plus generated indexing programs compared against NumPy. + +**Tech Stack:** Python 3.12+, NumPy, pytest, Hypothesis, Zarr transform and codec pipeline internals. + +## Global Constraints + +- Keep the public indexing API unchanged. +- Use full-input-rank index arrays; do not add a transform-wide indexing-mode flag. +- Preserve TensorStore-style literal coordinates in the general `IndexTransform` API while applying NumPy-style negative wraparound at the public Zarr array boundary. +- Never silently bypass a lazy view transform. +- Keep negative slice steps unsupported. +- Follow test-first red/green/refactor cycles. + +--- + +### Task 1: Normalize ArrayMap Dependency Axes + +**Files:** +- Modify: `src/zarr/core/transforms/transform.py` +- Modify: `src/zarr/core/transforms/composition.py` +- Modify: `src/zarr/core/transforms/json.py` +- Test: `tests/test_transforms/test_transform.py` +- Test: `tests/test_transforms/test_composition.py` +- Test: `tests/test_transforms/test_json.py` + +**Interfaces:** +- Consumes: `IndexTransform(domain, output)`, `_apply_oindex`, `_apply_vindex`, `compose`. +- Produces: full-input-rank `ArrayMap.index_array` values and an internal dependency-axis classifier used by later tasks. + +- [ ] **Step 1: Write failing transform-construction tests** + +Add assertions that orthogonal maps have shapes `(2, 1)` and `(1, 3)`, while two vectorized maps share shape `(2,)`. Assert JSON round-tripping preserves singleton axes. + +```python +def test_oindex_multiple_arrays_preserves_independent_axes() -> 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 result.output[0].index_array.shape == (2, 1) + assert result.output[1].index_array.shape == (1, 3) + + +def test_vindex_multiple_arrays_preserves_shared_axes() -> None: + t = IndexTransform.from_shape((10, 20)) + result = t.vindex[np.array([1, 3]), np.array([2, 4])] + assert result.domain.shape == (2,) + assert result.output[0].index_array.shape == (2,) + assert result.output[1].index_array.shape == (2,) +``` + +- [ ] **Step 2: Run the focused tests and confirm the orthogonal shape assertion fails** + +Run: + +```bash +.venv/bin/python -m pytest tests/test_transforms/test_transform.py tests/test_transforms/test_composition.py tests/test_transforms/test_json.py -q +``` + +Expected: the new orthogonal singleton-axis assertions fail against squeezed arrays. + +- [ ] **Step 3: Implement full-rank normalization** + +Update `_apply_oindex` to reshape each public one-dimensional integer array with one varying axis and singleton axes elsewhere. Update `_apply_vindex` to align every broadcast array with the broadcast dimensions plus singleton slice dimensions. Add a private helper with this contract: + +```python +def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, ...]: + """Return input axes on which a normalized index array varies.""" + return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) +``` + +Preserve full rank in basic indexing, composition, and JSON conversion. Reject an `ArrayMap` whose rank differs from its containing transform input rank, except during a clearly defined compatibility normalization at construction. + +- [ ] **Step 4: Run focused transform tests and refactor without changing behavior** + +Run the command from Step 2. Expected: PASS. + +- [ ] **Step 5: Commit Task 1** + +```bash +git add src/zarr/core/transforms tests/test_transforms +git commit -m "fix(indexing): preserve array map dependency axes" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 2: Resolve Independent and Correlated Array Maps Correctly + +**Files:** +- Modify: `src/zarr/core/transforms/transform.py` +- Modify: `src/zarr/core/transforms/chunk_resolution.py` +- Test: `tests/test_transforms/test_transform.py` +- Test: `tests/test_transforms/test_chunk_resolution.py` +- Test: `tests/test_lazy_indexing.py` + +**Interfaces:** +- Consumes: normalized full-rank `ArrayMap` values and `_array_map_dependency_axes` from Task 1. +- Produces: chunk intersections and output selectors that preserve outer-product or pairwise semantics. + +- [ ] **Step 1: Write failing orthogonal resolution tests** + +Create a chunked 2-D array and compare a lazy orthogonal selection with unequal index lengths to `np.ix_`. Cover both reads and writes across more than one chunk. + +```python +def test_lazy_oindex_multiple_arrays_outer_product(arr, data) -> None: + rows = np.array([1, 11]) + cols = np.array([2, 12, 22]) + actual = arr.lazy.oindex[rows, cols].result() + np.testing.assert_array_equal(actual, data[np.ix_(rows, cols)]) +``` + +- [ ] **Step 2: Run focused tests and confirm failure in intersection or placement** + +```bash +.venv/bin/python -m pytest tests/test_transforms/test_chunk_resolution.py tests/test_lazy_indexing.py -q +``` + +Expected: the new multi-array orthogonal regression fails. + +- [ ] **Step 3: Implement dependency-aware intersection and selectors** + +Replace the `len(array_dims) >= 2` correlation test with dependency analysis. Jointly mask maps only when they share the same non-singleton dependency axes. Filter independent maps along their own axes and retain singleton axes. In `sub_transform_to_selections`, emit shared scatter indices only for correlated maps; build orthogonal selectors for independent maps. + +- [ ] **Step 4: Verify orthogonal and vectorized behavior** + +Run: + +```bash +.venv/bin/python -m pytest tests/test_transforms/test_transform.py tests/test_transforms/test_chunk_resolution.py tests/test_lazy_indexing.py -q +``` + +Expected: PASS, including existing vectorized cases. + +- [ ] **Step 5: Commit Task 2** + +```bash +git add src/zarr/core/transforms tests/test_transforms tests/test_lazy_indexing.py +git commit -m "fix(indexing): distinguish orthogonal array maps" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 3: Preserve Lazy Views and Normalize Public Indices + +**Files:** +- Modify: `src/zarr/core/array.py` +- Modify: `src/zarr/core/transforms/transform.py` +- Test: `tests/test_lazy_indexing.py` +- Test: `tests/test_indexing.py` + +**Interfaces:** +- Consumes: `selection_to_transform(selection, transform, mode)`. +- Produces: array-boundary selection normalization relative to the current visible view and transform-aware fallback behavior. + +- [ ] **Step 1: Write failing composed-view and validation tests** + +Cover `arr.lazy[10:20].oindex[[0]]`, the corresponding write, `arr.lazy[-1]`, negative integer arrays, overly negative indices, positive out-of-range indices, and invalid Boolean masks. Assert writes by comparing the complete root array with NumPy. + +- [ ] **Step 2: Confirm the regressions fail for the reviewed reasons** + +```bash +.venv/bin/python -m pytest tests/test_lazy_indexing.py tests/test_indexing.py -q +``` + +Expected: composed advanced operations access the wrong offset or validation differs from eager Zarr. + +- [ ] **Step 3: Add a public-selection normalization boundary** + +Implement a helper in `array.py` that normalizes integer scalars and arrays against `self.shape`, validates bounds and mask shapes, and returns a selection suitable for `selection_to_transform`. Use it in `_LazyIndexAccessor`, `_LazyOIndex`, and `_LazyVIndex` before composing transforms. + +For advanced access on a non-identity view, compose with `self._transform` instead of creating a legacy indexer against storage. Allow the legacy path only where it is guaranteed to operate on the identity transform. Raise `NotImplementedError` for transformed structured-field operations that cannot preserve field-aware semantics. + +- [ ] **Step 4: Run focused lazy and eager indexing tests** + +Run the command from Step 2. Expected: PASS. + +- [ ] **Step 5: Commit Task 3** + +```bash +git add src/zarr/core/array.py src/zarr/core/transforms/transform.py tests/test_lazy_indexing.py tests/test_indexing.py +git commit -m "fix(indexing): compose advanced lazy selections" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 4: Correct Zero-Rank and Caller-Provided Buffer Handling + +**Files:** +- Modify: `src/zarr/core/array.py` +- Test: `tests/test_indexing.py` +- Test: `tests/test_lazy_indexing.py` + +**Interfaces:** +- Consumes: normalized transforms and chunk output selectors. +- Produces: scalar-shaped zero-rank I/O and safe multidimensional vectorized `out` behavior. + +- [ ] **Step 1: Write failing result-kind and `out` tests** + +Assert that `arr[...]` and `arr.lazy[...].result()` for a zero-dimensional array have shape `()` rather than `(1,)`. Add multidimensional coordinate arrays and an `NDBuffer` whose shape equals the broadcast selection shape; assert returned values and buffer contents match NumPy. + +- [ ] **Step 2: Run tests and observe the exact shape/placement failures** + +```bash +.venv/bin/python -m pytest tests/test_indexing.py::test_get_basic_selection_0d tests/test_indexing.py::test_get_selection_out tests/test_lazy_indexing.py -q +``` + +Expected: strict zero-rank shape assertions and multidimensional `out` regression fail. + +- [ ] **Step 3: Implement correlation-aware flattening** + +Define flattening as `bool(correlated_array_maps)` rather than `all(...)`. Keep `buffer_shape == ()` for a zero-rank identity transform. When `out` is supplied for correlated multidimensional indexing, read into a flat temporary, reshape to `out_shape`, and copy into `out`; return the caller-visible result without converting backend buffers unnecessarily. + +- [ ] **Step 4: Verify focused buffer tests** + +Run the command from Step 2. Expected: PASS. + +- [ ] **Step 5: Commit Task 4** + +```bash +git add src/zarr/core/array.py tests/test_indexing.py tests/test_lazy_indexing.py +git commit -m "fix(indexing): preserve transform result shapes" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 5: Add Stateful Property-Based Indexing Programs + +**Files:** +- Modify: `src/zarr/testing/strategies.py` +- Modify: `tests/test_properties.py` + +**Interfaces:** +- Consumes: array shapes and existing basic/orthogonal/vectorized strategies. +- Produces: `indexing_programs(shape)` yielding valid operation sequences and execution modes for NumPy/Zarr differential testing. + +- [ ] **Step 1: Add the program model and failing composed examples** + +Define small immutable operation records: + +```python +@dataclass(frozen=True) +class IndexOperation: + mode: Literal["basic", "orthogonal", "vectorized"] + selection: Any + + +@dataclass(frozen=True) +class IndexingProgram: + operations: tuple[IndexOperation, ...] + execution: Literal["materialize", "eager_on_lazy", "out", "set_scalar", "set_array"] +``` + +Generate one to three operations while carrying the visible shape forward. Keep sizes small enough for 300 Hypothesis examples. Add deterministic `@example` cases for slice-then-oindex, unequal orthogonal lengths, multidimensional vindex with `out`, and a zero-rank result. + +- [ ] **Step 2: Run the new property with deterministic regression examples enabled** + +```bash +.venv/bin/python -m pytest tests/test_properties.py -k indexing_program -q +``` + +Expected on the fixed tree: PASS. Temporarily reverse the Task 2 correlation predicate in the working tree, rerun the slice-then-oindex deterministic example, and confirm it fails before restoring the Task 2 code. This mutation check proves the generated-program oracle detects the original semantic error. + +- [ ] **Step 3: Implement the NumPy/Zarr program runner** + +For reads, compare exact shape, scalar-versus-array kind, dtype, values, and `out`. For writes, retain the NumPy root array, apply the equivalent indexed assignment, and compare the entire Zarr root array after the write. Use `np.ix_` to implement NumPy orthogonal semantics and normal NumPy indexing for vectorized semantics. + +- [ ] **Step 4: Run property tests under the CI profile** + +```bash +HYPOTHESIS_PROFILE=ci .venv/bin/python -m pytest --run-slow-hypothesis tests/test_properties.py -k "indexing_program or basic_indexing or oindex or vindex" -q +``` + +Expected: PASS with 300 deterministic examples per property. + +- [ ] **Step 5: Commit Task 5** + +```bash +git add src/zarr/testing/strategies.py tests/test_properties.py +git commit -m "test(indexing): generate composed indexing programs" -m "Assisted-by: Codex:GPT-5" +``` + +### Task 6: Full Verification and Review + +**Files:** +- Verify all modified production and test files. + +**Interfaces:** +- Consumes: Tasks 1-5. +- Produces: evidence that the branch is ready for review. + +- [ ] **Step 1: Run the complete affected test modules** + +```bash +.venv/bin/python -m pytest tests/test_transforms tests/test_lazy_indexing.py tests/test_indexing.py tests/test_array.py -q +``` + +Expected: PASS with only existing documented xfails. + +- [ ] **Step 2: Run the focused Hypothesis suite** + +```bash +HYPOTHESIS_PROFILE=ci .venv/bin/python -m pytest --run-slow-hypothesis tests/test_properties.py -q +``` + +Expected: PASS. + +- [ ] **Step 3: Run formatting and static checks** + +Use the repository's configured commands for Ruff, mypy, and pre-commit on modified files. Expected: no new errors. + +- [ ] **Step 4: Inspect the final branch diff** + +```bash +git diff --check origin/main...HEAD +git status --short +``` + +Expected: no whitespace errors and no uncommitted implementation changes. + +- [ ] **Step 5: Request a branch code review** + +Run the branch-review workflow and address any correctness findings before declaring completion. diff --git a/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md b/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md new file mode 100644 index 0000000000..b2edb86e6f --- /dev/null +++ b/docs/superpowers/specs/2026-07-15-tensorstore-index-transform-corrections-design.md @@ -0,0 +1,125 @@ +# TensorStore-Aligned Index Transform Corrections + +## Goal + +Correct lazy and transform-based indexing so that basic, orthogonal, and vectorized selections preserve NumPy/Zarr semantics through composition, chunk resolution, reads, writes, and serialization. + +## Scope + +This change fixes five reviewed regressions: + +1. Legacy advanced-indexing fallbacks bypass a lazy view's transform. +2. Multiple independent orthogonal index arrays are mistaken for correlated vectorized arrays. +3. Lazy selections do not consistently normalize negative indices or validate bounds. +4. Zero-dimensional identity transforms are incorrectly flattened. +5. Multidimensional vectorized reads cannot safely write directly into a caller-provided `out` buffer. + +The public indexing API remains unchanged. Structured-dtype field selection remains on the legacy codec path, but it must either operate on the correct composed storage selection or reject unsupported transformed access explicitly; it must never silently access the wrong storage region. + +## Transform Representation + +Follow TensorStore's normalized index-transform model. Every `ArrayMap.index_array` is interpreted over the complete input domain of its containing `IndexTransform`. + +- The index array rank equals the transform input rank. +- An axis on which an index map does not vary is represented as a singleton/broadcast axis. +- Independent orthogonal maps vary along different input dimensions. +- Correlated vectorized maps vary along the same broadcast input dimensions. +- Zero-rank transforms use zero-rank index arrays where applicable. + +For an orthogonal selection with row indices of length 2 and column indices of length 3, the normalized maps have shapes `(2, 1)` and `(1, 3)`. For a pairwise vectorized selection of length 2, both maps have shape `(2,)`. + +Correlation must be derived from the maps' input-dimension dependencies. Code must not infer vectorized semantics merely from the presence of two or more `ArrayMap` instances. + +## Construction and Composition + +Basic, orthogonal, and vectorized selection constructors normalize their inputs before producing transforms: + +- Python and NumPy integer scalars use normal Zarr negative-index wraparound. +- Integer arrays normalize negative elements and reject values outside the selected view's bounds. +- Boolean masks must match the dimensions they consume. +- Slice steps remain positive-only, consistent with current Zarr behavior. +- Orthogonal integer arrays remain one-dimensional at the public API boundary, then become full-rank broadcast arrays in the normalized transform. +- Vectorized integer arrays are broadcast together before being expanded to full input rank. + +Composition preserves full-rank dependency information. When an outer transform indexes an inner `ArrayMap`, the resulting array is evaluated and broadcast over the outer input domain without squeezing dependency axes. Composition may simplify a map to `ConstantMap` or `DimensionMap` only when the resulting mapping is provably equivalent. + +## Intersection and Chunk Resolution + +Intersection classifies maps by their dependency axes: + +- Correlated maps are filtered with one joint mask so a coordinate survives only when all mapped storage coordinates are inside the chunk. +- Independent maps are filtered per dependency axis and retain outer-product semantics. +- Mixed `ArrayMap`, `DimensionMap`, and `ConstantMap` transforms retain their input-domain axis ordering. + +Chunk resolution produces output selectors consistent with the normalized input domain. Correlated maps share scatter coordinates; independent maps produce orthogonal selectors. The resolver must not rely on array-map count as a proxy for correlation. + +## Array Read and Write Paths + +All non-field lazy reads and writes use the composed `IndexTransform` path. Eager advanced operations may continue using legacy indexers when operating on an identity transform, but a transformed view must not fall back to storage-relative indexing that ignores its transform. + +Structured-field operations continue using the existing field-aware codec path. If a transformed structured-field selection cannot be expressed correctly through that path, raise a clear `NotImplementedError` instead of reading or writing the wrong region. + +Buffer handling follows transform semantics: + +- Flatten only when the transform contains correlated array maps whose chunk resolution emits flat scatter indices. +- Never treat an empty output-map tuple as vectorized; zero-dimensional identity reads and writes retain shape `()`. +- When a multidimensional vectorized read supplies `out`, allocate a flat temporary buffer, perform scatter reads into it, reshape it to the selection shape, and copy it into `out`. +- Writes validate broadcasting against the visible selection domain before any vectorized flattening. + +## Serialization + +JSON round-tripping preserves full-rank index-array shapes, including singleton dependency axes. No transform-wide `orthogonal` or `vectorized` mode flag is added. Existing JSON that already satisfies the rank invariant continues to load unchanged. + +If compatibility with branch-local JSON fixtures containing lower-rank arrays is required, the loader may normalize them only when their dependency alignment is unambiguous. Ambiguous lower-rank arrays must be rejected rather than guessed. + +## Error Handling + +- Scalar and array indices outside explicit view bounds raise `IndexError` before chunk access. +- Invalid boolean mask shapes raise `IndexError` during transform construction. +- Incompatible assignment shapes raise `ValueError` before storage mutation. +- Unsupported transformed structured-field access raises `NotImplementedError` with an explanation. +- All validation is relative to the current lazy view domain, not the root array shape. + +## Test Strategy + +Add focused regression tests before implementation and verify that each fails for the reviewed reason: + +- A sliced lazy view followed by advanced orthogonal read and write accesses the composed storage region. +- Multiple orthogonal index arrays with unequal lengths produce an outer-product result across chunks. +- Multiple vectorized arrays retain pairwise/broadcast semantics. +- Negative scalar and array indices work relative to a lazy view; overly negative and positive out-of-range values raise. +- Invalid Boolean mask shapes raise. +- Zero-dimensional reads return a scalar and zero-dimensional writes persist it. +- Multidimensional coordinate selection with a caller-provided `out` buffer matches NumPy. +- Transform and JSON unit tests assert full-rank singleton-axis normalization and round-trip preservation. + +After focused tests pass, run the complete transform, lazy-indexing, indexing, and array test modules, followed by the repository's configured static checks for the modified files. + +### Stateful Indexing Properties + +Extend the property-based test architecture from single eager selections to indexing programs. A generated example consists of: + +1. A NumPy array and equivalent chunked Zarr array. +2. A valid sequence of one or more indexing operations chosen from basic, orthogonal, and vectorized modes. +3. An execution mode chosen from lazy materialization, eager access on a lazy view, caller-provided `out`, scalar assignment, and array assignment. + +Apply the same program to a NumPy reference and to Zarr. Each operation is relative to the visible domain produced by the preceding operations. The generator must preserve whether index arrays are independent (`oindex`) or correlated (`vindex`) and must deliberately cover unequal orthogonal index lengths, broadcast-compatible vectorized shapes, negative indices, zero-rank results, empty results, and selections spanning multiple chunks. + +Read properties compare all of the following where applicable: + +- result shape; +- scalar versus array result kind; +- dtype; +- element values; +- contents of a caller-provided `out` buffer. + +Write properties compare the entire NumPy and Zarr backing arrays after the operation, not merely the selected view. This detects writes that produce the expected selected values while mutating the wrong storage region. + +Keep focused example-based regressions for each reviewed bug. The stateful properties complement those tests by exploring compositions and shape interactions; they do not replace precise regression cases. + +## Non-Goals + +- Negative slice steps. +- New public indexing modes or APIs. +- General TensorStore dimension labels, implicit bounds, or arbitrary-origin domains. +- Unrelated restructuring of the legacy indexer or codec pipeline. diff --git a/docs/user-guide/lazy_indexing.md b/docs/user-guide/lazy_indexing.md new file mode 100644 index 0000000000..19951e0922 --- /dev/null +++ b/docs/user-guide/lazy_indexing.md @@ -0,0 +1,294 @@ +# 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 every index is a **literal coordinate** +in that domain. + +```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 +w = a.translate_by((-10,)) # a view of `a` labeled [-10, 2) +print(w.lazy[-1].result()) # -1 is just another index: base cell 9 +``` + +### `-1` is just another index + +Because indices are literal coordinates, a negative index is *not* "from the +end" — it names the coordinate `-1`, which your domain may or may not contain. +On the translated view above, `-1` was perfectly valid. On a fresh array (domain +`[0, n)`), it is out of bounds — in **every** syntactic form: integers, slice +bounds, and index arrays are treated identically. + +```python exec="true" session="lazy" source="above" result="ansi" +for make in (lambda: a.lazy[-1], lambda: a.lazy[-3:], lambda: a.lazy.oindex[[-1]]): + try: + make() + except IndexError as e: + print(type(e).__name__, "-", str(e).split(";")[0]) +``` + +To select from the end, say what you mean literally: + +```python exec="true" session="lazy" source="above" result="ansi" +n = a.shape[0] +print(a.lazy[n - 3 :].result()) # the last three elements +``` + +### 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 domain coordinates. (The base array's ordinary `a[...]` API is +unchanged: it keeps full NumPy semantics, negatives and all. The literal rules +apply to *views* and to the `.lazy` accessor.) NumPy-style zero-based access to +a view's data 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 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 are not from-the-end** — in any form (integer, slice + bound, index array). They name literal coordinates, which fresh arrays' + domains (starting at 0) do not contain. Use `shape[dim] - k`, or translate + the domain so negative coordinates exist. +- **No clamping**: out-of-range slice bounds raise; 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]`). +- `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 e4e757e630..10f47b4ff5 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/storage.md 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/api/synchronous.py b/src/zarr/api/synchronous.py index 8386427b3f..d6775b9104 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -1138,20 +1138,20 @@ def from_array( ... ) >>> arr2 = zarr.from_array(store, data=arr, overwrite=True) >>> arr2 - + >>> asyncio.run(store.clear()) # Remove files generated by test Create an array from an existing NumPy array: >>> import numpy as np >>> zarr.from_array({}, data=np.arange(10000, dtype="i4").reshape(100, 100)) - + Create an array from any array-like object: >>> arr3 = zarr.from_array({}, data=[[1, 2], [3, 4]]) >>> arr3 - + >>> arr3[...] array([[1, 2], [3, 4]]) @@ -1160,7 +1160,7 @@ def from_array( >>> arr4 = zarr.from_array({}, data=[[1, 2], [3, 4]]) >>> arr5 = zarr.from_array({}, data=arr4, write_data=False) >>> arr5 - + >>> arr5[...] array([[0, 0], [0, 0]]) """ diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 977520b12e..f27ef4195c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -42,6 +42,7 @@ from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, + ChunkSpec, _is_rectilinear_chunks, as_regular_shape, guess_chunks, @@ -55,6 +56,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 +106,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, @@ -130,9 +136,16 @@ parse_node_type_array, ) from zarr.core.sync import sync +from zarr.core.transforms.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr.core.transforms.output_map import ArrayMap +from zarr.core.transforms.transform import ( + IndexTransform, + selection_to_transform, +) from zarr.errors import ( ArrayNotFoundError, ChunkNotFoundError, + LazyViewError, MetadataValidationError, ZarrDeprecationWarning, ZarrUserWarning, @@ -329,6 +342,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 @@ -365,6 +379,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( @@ -778,6 +793,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 @@ -796,7 +840,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, ...]: @@ -807,6 +851,88 @@ 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 :meth:`_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`." + ) + + @property + def _is_sharded(self) -> bool: + """Whether the array stores inner chunks inside shards (a sharding codec).""" + from zarr.codecs.sharding import ShardingCodec + + codecs: tuple[Codec, ...] = getattr(self.metadata, "codecs", ()) + return len(codecs) == 1 and isinstance(codecs[0], ShardingCodec) + + 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 @@ -822,6 +948,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 @@ -848,6 +975,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 @@ -879,7 +1007,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 @@ -895,6 +1023,7 @@ def shards(self) -> tuple[int, ...] | None: tuple[int, ...]: The shard shape of the Array. """ + self._require_identity("shards") return self.metadata.shards @property @@ -906,7 +1035,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, ...]: @@ -1072,6 +1203,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 @@ -1128,6 +1260,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 @@ -1212,6 +1345,7 @@ async def example(): result = asyncio.run(example()) ``` """ + self._require_identity("nchunks_initialized") return await _nchunks_initialized(self) async def _nshards_initialized(self) -> int: @@ -1253,6 +1387,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( @@ -1568,6 +1703,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, @@ -1658,6 +1829,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, ...]: @@ -1679,6 +1851,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: @@ -1750,6 +1923,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: @@ -1769,6 +1943,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( @@ -1831,6 +2006,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, @@ -2166,7 +2384,7 @@ def cdata_shape(self) -> tuple[int, ...]: When sharding is used, this counts inner chunks (not shards) per dimension. """ - return self.async_array._chunk_grid_shape + return self.async_array.cdata_shape @property def _chunk_grid_shape(self) -> tuple[int, ...]: @@ -2244,6 +2462,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: """ @@ -2827,14 +3066,19 @@ 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, + 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, + ) ) - ) + 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, @@ -2936,8 +3180,16 @@ 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)) + 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 + 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, @@ -3064,12 +3316,23 @@ 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 + 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. + 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, @@ -3182,10 +3445,22 @@ 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) + 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. + 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, @@ -3270,12 +3545,28 @@ 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 + 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, @@ -3359,8 +3650,25 @@ 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)) + 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, @@ -3447,16 +3755,47 @@ 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 + 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) + 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}" + ) + sel_normalized = tuple( + np.where(np.asarray(s) < 0, np.asarray(s) + hi, s) + for s, hi in zip( + sel_normalized, self._async_array._transform.domain.exclusive_max, strict=True + ) + ) + transform = selection_to_transform( + sel_normalized, self._async_array._transform, "vectorized" + ) + out_array = sync( + self._async_array._get_selection_t(transform, out=out, prototype=prototype) + ) + # Reshape to the broadcast shape of the coordinate arrays + sel_tuple = sel_normalized + sel_arrays = [np.asarray(s) for s in sel_tuple] + sel_shape = np.broadcast_shapes(*(s.shape for s in sel_arrays)) + if hasattr(out_array, "shape") and sel_shape != (): + out_array = cast("NDArrayLikeOrScalar", np.array(out_array).reshape(sel_shape)) return out_array def set_coordinate_selection( @@ -3538,30 +3877,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 an empty fields list/tuple to None (pop_fields yields [] for + # no fields); keep the exact-emptiness check rather than a falsy one. + if fields == [] or fields == (): + fields = None + 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}" + ) + sel_normalized = tuple( + np.where(np.asarray(s) < 0, np.asarray(s) + hi, s) + for s, hi in zip( + sel_normalized, self._async_array._transform.domain.exclusive_max, strict=True + ) + ) + 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, @@ -3660,6 +4026,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( @@ -3760,6 +4134,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)) @@ -3789,6 +4170,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 @@ -3892,7 +4286,11 @@ def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: return type(self)(new_array) def __repr__(self) -> str: - return f"" + t = self._async_array._transform + return ( + f"" + ) @property def info(self) -> Any: @@ -4005,6 +4403,65 @@ async def _shards_initialized( type SerializerLike = dict[str, JSON] | ArrayBytesCodec | Literal["auto"] +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]: + 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: + 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]: + 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: + 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]: + 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: + new_t = selection_to_transform(selection, self._array._async_array._transform, "basic") + self._array._with_transform(new_t)[...] = value + + @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 @@ -5329,17 +5786,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 :func:`_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 :func:`_is_complete_chunk`. + """ return ArraySpec( shape=spec.codec_shape, dtype=metadata.dtype, @@ -5349,6 +5807,317 @@ 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.core.transforms.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 :class:`ChunkSpec` (the caller looks + it up once and shares it with :func:`_array_spec_from_chunk_spec`). + """ + from zarr.core.transforms.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 coordinate ArrayMap, input_dimension None — even a + # single one, e.g. vindex[idx_2d]) scatters through a single flat index, so + # the 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. + needs_flat_buffer = _is_vectorized_transform(transform) + buffer_shape = (product(out_shape),) if needs_flat_buffer else out_shape + + # Setup output buffer. Vectorized reads scatter through a flat index, so the + # caller must supply a flat out buffer (matching the eager path); otherwise it + # is the full multi-dimensional out_shape. + 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 != buffer_shape: + raise ValueError( + f"shape of out argument doesn't match. Expected {buffer_shape}, got {out.shape}" + ) + 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: + continue + 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 missing_info: + 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 + if out_shape == (): + return out_buffer.as_scalar() + out_result = out_buffer.as_ndarray_like() + # Reshape if we flattened for array indexing + if needs_flat_buffer and hasattr(out_result, "reshape"): + out_result = np.array(out_result).reshape(out_shape) + 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: + continue + 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, @@ -5912,9 +6681,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..712921928d --- /dev/null +++ b/src/zarr/core/chunk_partition.py @@ -0,0 +1,104 @@ +"""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.core.transforms.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from zarr.core.chunk_grids import ChunkGrid + from zarr.core.transforms.transform import IndexTransform + +__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``. + + Only a per-dimension full ``0:size:1`` slice is a full cover; an integer or + array selection touches a strict subset, so the chunk is partial. + """ + if len(chunk_selection) != len(shape): + return False + for sel, size in zip(chunk_selection, shape, strict=True): + if not isinstance(sel, slice): + return False + start, stop, step = sel.indices(size) + if not (start == 0 and stop == size and step == 1): + 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: + continue + 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/group.py b/src/zarr/core/group.py index 52eaa3e144..f3ecd8b11b 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -1939,7 +1939,7 @@ def __setitem__(self, key: str, value: Any) -> None: >>> group = zarr.group() >>> group["foo"] = np.array(zarr.zeros((10,))) >>> group["foo"] - + """ self._sync(self._async_group.setitem(key, value)) @@ -2259,7 +2259,7 @@ def arrays(self) -> Generator[tuple[str, AnyArray], None]: >>> group = zarr.group() >>> subarray = group.create_array("subarray", dtype="i1", shape=(10,), chunks=(10,)) >>> list(group.arrays()) - [('subarray', )] + [('subarray', )] """ for name, async_array in self._sync_iter(self._async_group.arrays()): yield name, Array(async_array) @@ -2288,7 +2288,7 @@ def array_values(self) -> Generator[AnyArray, None]: >>> group = zarr.group() >>> subarray = group.create_array("subarray", dtype="i1", shape=(10,), chunks=(10,)) >>> list(group.array_values()) - [] + [] """ for _, array in self.arrays(): yield array diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index f6eb495cd9..26eb06c996 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1442,8 +1442,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/core/transforms/__init__.py b/src/zarr/core/transforms/__init__.py new file mode 100644 index 0000000000..ec98e02d4d --- /dev/null +++ b/src/zarr/core/transforms/__init__.py @@ -0,0 +1,46 @@ +"""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 +""" + +from zarr.core.transforms.composition import compose +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.json import ( + IndexDomainJSON, + IndexTransformJSON, + OutputIndexMapJSON, + index_domain_from_json, + index_domain_to_json, + index_transform_from_json, + index_transform_to_json, +) +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.core.transforms.transform import IndexTransform + +__all__ = [ + "ArrayMap", + "ConstantMap", + "DimensionMap", + "IndexDomain", + "IndexDomainJSON", + "IndexTransform", + "IndexTransformJSON", + "OutputIndexMap", + "OutputIndexMapJSON", + "compose", + "index_domain_from_json", + "index_domain_to_json", + "index_transform_from_json", + "index_transform_to_json", +] diff --git a/src/zarr/core/transforms/chunk_resolution.py b/src/zarr/core/transforms/chunk_resolution.py new file mode 100644 index 0000000000..e57d71d408 --- /dev/null +++ b/src/zarr/core/transforms/chunk_resolution.py @@ -0,0 +1,223 @@ +"""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. + +``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.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform + +if TYPE_CHECKING: + from collections.abc import Iterator + + from zarr.core.chunk_grids import ChunkGrid + +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 iter_chunk_transforms( + transform: IndexTransform, + chunk_grid: ChunkGrid, +) -> 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. + """ + dim_grids = chunk_grid._dimensions + + # Enumerate all possible chunks via cartesian product of per-dim chunk ranges + # For each candidate chunk, intersect the transform with the chunk domain. + # The transform.intersect method handles both orthogonal and vectorized cases. + chunk_ranges: list[range] = [] + 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_ranges.append(range(c, c + 1)) + 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_ranges.append(range(first, last + 1)) + elif isinstance(m, ArrayMap): + storage = m.offset + m.stride * m.index_array + flat = storage.ravel().astype(np.intp) + chunk_ids = dg.indices_to_chunks(flat) + first = int(chunk_ids.min()) + last = int(chunk_ids.max()) + chunk_ranges.append(range(first, last + 1)) + + import itertools + + for chunk_coords_tuple in itertools.product(*chunk_ranges): + 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) + + chunk_sel: list[int | slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + out_sel: list[slice | np.ndarray[tuple[int, ...], np.dtype[np.intp]]] = [] + + # Single-pass build for the basic / single-array / vectorized cases. + # ConstantMap dims are dropped (no out_sel entry). + n_array_maps = 0 + + 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 + n_array_maps += 1 + if m.offset == 0 and m.stride == 1: + chunk_sel.append(m.index_array) + else: + chunk_sel.append((m.offset + m.stride * m.index_array).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, len(m.index_array))) + + # Vectorized: ≥2 correlated ArrayMaps scatter through a single shared index. + if out_indices is not None and n_array_maps >= 2: + out_sel = [out_indices] + + return tuple(chunk_sel), tuple(out_sel), () diff --git a/src/zarr/core/transforms/composition.py b/src/zarr/core/transforms/composition.py new file mode 100644 index 0000000000..9d07bd3324 --- /dev/null +++ b/src/zarr/core/transforms/composition.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.core.transforms.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) + + if isinstance(inner_map, ArrayMap): + return _compose_array(outer, inner_map) + + raise TypeError(f"Unknown output map type: {type(inner_map)}") # pragma: no cover + + +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, + ) + + if isinstance(outer_map, ArrayMap): + return ArrayMap( + index_array=outer_map.index_array, + offset=offset_i + stride_i * outer_map.offset, + stride=stride_i * outer_map.stride, + ) + + raise TypeError(f"Unknown output map type: {type(outer_map)}") # pragma: no cover + + +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/src/zarr/core/transforms/domain.py b/src/zarr/core/transforms/domain.py new file mode 100644 index 0000000000..bca6488ad1 --- /dev/null +++ b/src/zarr/core/transforms/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/src/zarr/core/transforms/json.py b/src/zarr/core/transforms/json.py new file mode 100644 index 0000000000..79fe6e58b8 --- /dev/null +++ b/src/zarr/core/transforms/json.py @@ -0,0 +1,167 @@ +"""JSON serialization for index transforms. + +Defines TypedDict types matching TensorStore's JSON representation of +IndexTransform and IndexDomain, plus conversion functions. + +The JSON format follows TensorStore's conventions for interoperability:: + + { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [100, 200], + "input_labels": ["x", "y"], + "output": [ + {"offset": 5}, + {"offset": 10, "stride": 2, "input_dimension": 1}, + {"offset": 0, "stride": 1, "index_array": [[1, 2, 0]]} + ] + } +""" + +from __future__ import annotations + +from typing import Required, TypedDict + +import numpy as np + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.core.transforms.transform import IndexTransform + +# --------------------------------------------------------------------------- +# TypedDict definitions (JSON shapes) +# --------------------------------------------------------------------------- + + +class IndexDomainJSON(TypedDict, total=False): + """JSON representation of an IndexDomain.""" + + input_inclusive_min: Required[list[int]] + input_exclusive_max: Required[list[int]] + input_labels: list[str] + + +class OutputIndexMapJSON(TypedDict, total=False): + """JSON representation of a single output index map. + + Exactly one of three forms: + - ``{"offset": 5}`` — constant + - ``{"offset": 0, "stride": 1, "input_dimension": 0}`` — dimension + - ``{"offset": 0, "stride": 1, "index_array": [...]}`` — array + """ + + offset: int + stride: int + input_dimension: int + index_array: list[int] | list[list[int]] + + +class IndexTransformJSON(TypedDict, total=False): + """JSON representation of an IndexTransform.""" + + input_inclusive_min: Required[list[int]] + input_exclusive_max: Required[list[int]] + input_labels: list[str] + output: Required[list[OutputIndexMapJSON]] + + +# --------------------------------------------------------------------------- +# IndexDomain serialization +# --------------------------------------------------------------------------- + + +def index_domain_to_json(domain: IndexDomain) -> IndexDomainJSON: + """Convert an IndexDomain to its JSON representation.""" + result: IndexDomainJSON = { + "input_inclusive_min": list(domain.inclusive_min), + "input_exclusive_max": list(domain.exclusive_max), + } + if domain.labels is not None: + result["input_labels"] = list(domain.labels) + return result + + +def index_domain_from_json(data: IndexDomainJSON) -> IndexDomain: + """Construct an IndexDomain from its JSON representation.""" + return IndexDomain( + inclusive_min=tuple(data["input_inclusive_min"]), + exclusive_max=tuple(data["input_exclusive_max"]), + labels=tuple(data["input_labels"]) if "input_labels" in data else None, + ) + + +# --------------------------------------------------------------------------- +# OutputIndexMap serialization +# --------------------------------------------------------------------------- + + +def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON: + """Convert an output index map to its JSON representation.""" + if isinstance(m, ConstantMap): + result: OutputIndexMapJSON = {"offset": m.offset} + return result + + if isinstance(m, DimensionMap): + result = {"offset": m.offset, "input_dimension": m.input_dimension} + if m.stride != 1: + result["stride"] = m.stride + return result + + if isinstance(m, ArrayMap): + result = {"offset": m.offset, "index_array": m.index_array.tolist()} + if m.stride != 1: + result["stride"] = m.stride + # Present only for orthogonal (oindex) arrays; omitted for vectorized. + if m.input_dimension is not None: + result["input_dimension"] = m.input_dimension + return result + + raise TypeError(f"Unknown output map type: {type(m)}") + + +def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap: + """Construct an output index map from its JSON representation.""" + if "index_array" in data: + return ArrayMap( + index_array=np.asarray(data["index_array"], dtype=np.intp), + offset=data.get("offset", 0), + stride=data.get("stride", 1), + input_dimension=data.get("input_dimension"), + ) + + if "input_dimension" in data: + return DimensionMap( + input_dimension=data["input_dimension"], + offset=data.get("offset", 0), + stride=data.get("stride", 1), + ) + + # Constant map: only offset present + return ConstantMap(offset=data.get("offset", 0)) + + +# --------------------------------------------------------------------------- +# IndexTransform serialization +# --------------------------------------------------------------------------- + + +def index_transform_to_json(transform: IndexTransform) -> IndexTransformJSON: + """Convert an IndexTransform to its JSON representation.""" + result: IndexTransformJSON = { + "input_inclusive_min": list(transform.domain.inclusive_min), + "input_exclusive_max": list(transform.domain.exclusive_max), + "output": [output_index_map_to_json(m) for m in transform.output], + } + if transform.domain.labels is not None: + result["input_labels"] = list(transform.domain.labels) + return result + + +def index_transform_from_json(data: IndexTransformJSON) -> IndexTransform: + """Construct an IndexTransform from its JSON representation.""" + domain = IndexDomain( + inclusive_min=tuple(data["input_inclusive_min"]), + exclusive_max=tuple(data["input_exclusive_max"]), + labels=tuple(data["input_labels"]) if "input_labels" in data else None, + ) + output = tuple(output_index_map_from_json(m) for m in data["output"]) + return IndexTransform(domain=domain, output=output) diff --git a/src/zarr/core/transforms/output_map.py b/src/zarr/core/transforms/output_map.py new file mode 100644 index 0000000000..c53402be4d --- /dev/null +++ b/src/zarr/core/transforms/output_map.py @@ -0,0 +1,98 @@ +"""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). + + ``input_dimension`` records which single input dimension indexes the array, + binding it the way :class:`DimensionMap` is bound. It distinguishes the two + flavours of multi-array fancy indexing: + + - **orthogonal** (``oindex``): each array indexes a *distinct* input + dimension, so ``input_dimension`` is set; the result is their outer + product. + - **vectorized** (``vindex``): the arrays are correlated and jointly + indexed by the same (possibly multi-dimensional) input range, so + ``input_dimension`` is ``None``. + + This binding is what lets chunk resolution tell an outer product from a + pointwise scatter when more than one ``ArrayMap`` is present. + """ + + index_array: npt.NDArray[np.intp] + offset: int = 0 + stride: int = 1 + input_dimension: int | None = None + + +OutputIndexMap = ConstantMap | DimensionMap | ArrayMap diff --git a/src/zarr/core/transforms/transform.py b/src/zarr/core/transforms/transform.py new file mode 100644 index 0000000000..473d15d0ef --- /dev/null +++ b/src/zarr/core/transforms/transform.py @@ -0,0 +1,1075 @@ +"""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.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +from zarr.errors import BoundsCheckError, VindexInvalidSelectionError + + +@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: + 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}") + elif isinstance(m, 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}]") + elif isinstance(m, 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, + ) + ) + elif isinstance(m, 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]). + + For orthogonal transforms (ConstantMap, DimensionMap, independent ArrayMaps), + each dimension is intersected independently and the input domain is narrowed. + + For vectorized transforms (correlated ArrayMaps), all array dimensions + must be checked simultaneously — a point survives only if ALL its + coordinates fall within the output domain. + + 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 ArrayMaps (vindex) are jointly indexed (input_dimension is None); + # >= 2 of them means a vectorized intersection. Orthogonal ArrayMaps (oindex) + # each bind a distinct input dimension and are handled per-dimension below. + array_dims = [i for i, m in enumerate(transform.output) if isinstance(m, ArrayMap)] + vectorized_dims = [ + i for i in array_dims if cast("ArrayMap", transform.output[i]).input_dimension is None + ] + if len(vectorized_dims) >= 2: + return _intersect_vectorized(transform, output_domain, vectorized_dims) + + # Orthogonal: intersect each output dimension independently. Multiple + # ArrayMaps bound to distinct input dimensions form an outer product, so each + # array dimension'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 + input_lo = new_min[d] + input_hi = new_max[d] + if input_lo >= input_hi: + return None + + # Find input range that produces storage coords in [lo, hi) + 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 + + new_min[d] = new_input_lo + new_max[d] = new_input_hi + new_output.append(m) + + elif isinstance(m, ArrayMap): + storage = m.offset + m.stride * m.index_array + mask = (storage >= lo) & (storage < hi) + survivors = np.nonzero(mask.ravel())[0].astype(np.intp) + if survivors.size == 0: + return None + filtered = m.index_array.ravel()[survivors] + new_output.append( + ArrayMap( + index_array=filtered, + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) + # Narrow this array's own input dimension to the surviving count. + if m.input_dimension is not None: + new_min[m.input_dimension] = 0 + new_max[m.input_dimension] = 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_vectorized( + transform: IndexTransform, + output_domain: IndexDomain, + array_dims: list[int], +) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]] | None] | None: + """Intersect a vectorized transform with an output domain. + + All ArrayMap outputs are correlated — a point survives only if ALL its + storage coordinates fall within the output domain. + """ + # Compute storage coords per array dim and check bounds simultaneously + masks: list[np.ndarray[Any, np.dtype[np.bool_]]] = [] + + for out_dim in array_dims: + m = transform.output[out_dim] + assert isinstance(m, ArrayMap) + storage = m.offset + m.stride * m.index_array + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + masks.append((storage >= lo) & (storage < hi)) + + # A point survives only if it's in-bounds on ALL array dims + combined_mask = masks[0] + for mask in masks[1:]: + combined_mask = combined_mask & mask + + if not np.any(combined_mask): + return None + + surviving = np.nonzero(combined_mask.ravel())[0].astype(np.intp) + + # Build new output maps + new_output: list[OutputIndexMap] = [] + for out_dim, m in enumerate(transform.output): + if isinstance(m, ArrayMap): + filtered = m.index_array.ravel()[surviving] + new_output.append( + ArrayMap( + index_array=filtered, + offset=m.offset, + stride=m.stride, + input_dimension=m.input_dimension, + ) + ) + elif isinstance(m, ConstantMap): + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + if lo <= m.offset < hi: + new_output.append(m) + else: + return None + elif isinstance(m, DimensionMap): + new_output.append(m) + + new_domain = IndexDomain.from_shape((len(surviving),)) + result = IndexTransform(domain=new_domain, output=tuple(new_output)) + return (result, surviving) + + +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( + arr: np.ndarray[Any, np.dtype[np.intp]], + 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). When input dimensions are dropped (int), sliced, + or inserted (newaxis), the array must be updated accordingly. + """ + # 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: + # Convert absolute domain coordinate to 0-based array index + array_idx = sel - domain.inclusive_min[old_dim] + idx.append(array_idx) + old_dim += 1 + elif isinstance(sel, slice): + if old_dim < arr.ndim: + 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)) + 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) + + +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 + elif isinstance(sel, slice): + 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") + elif isinstance(m, ArrayMap): + new_arr = _reindex_array(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 _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 + elif 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) + 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_output.append( + ArrayMap( + index_array=dim_array[d], + offset=m.offset, + stride=m.stride, + # Each oindex array indexes its own input dimension; this + # binding is what marks the selection orthogonal (outer + # product) rather than vectorized. + input_dimension=old_to_new_dim[d], + ) + ) + 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") + elif isinstance(m, ArrayMap): + 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 arrays: + 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: + new_output.append( + ArrayMap( + index_array=array_dim_to_broadcast[d], + 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 + ) + ) + elif isinstance(m, ArrayMap): + 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 = ( + "; negative indices are literal coordinates in lazy indexing, not from-the-end " + "(use `shape[dim] - k`, or materialize with `result()`, for NumPy semantics)" +) + + +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/src/zarr/errors.py b/src/zarr/errors.py index 781bebe534..d3bbb5ff74 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -9,6 +9,7 @@ "ContainsGroupError", "DataTypeValidationError", "GroupNotFoundError", + "LazyViewError", "MetadataValidationError", "NegativeStepError", "NodeTypeValidationError", @@ -151,6 +152,19 @@ class BoundsCheckError(IndexError): ... class ArrayIndexError(IndexError): ... +class LazyViewError(NotImplementedError): + """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. + """ + + class ChunkNotFoundError(BaseZarrError): """ Raised when a chunk that was expected to exist in storage was not retrieved successfully. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index f2c83677cc..60c99e92a1 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -27,6 +27,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 @@ -598,6 +599,76 @@ 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}") + + @st.composite def block_indices( draw: st.DrawFn, *, chunk_sizes: tuple[tuple[int, ...], ...] @@ -606,17 +677,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. @@ -629,7 +700,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. @@ -665,7 +736,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,)) ) @@ -690,9 +761,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()): @@ -738,10 +809,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: @@ -770,7 +841,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 0d6d2d5906..991120e9a0 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -1977,7 +1977,8 @@ def test_array_repr(store: Store) -> None: shape = (2, 3, 4) dtype = "uint8" arr = zarr.create_array(store, shape=shape, dtype=dtype) - assert str(arr) == f"" + domain = "{ [0, 2), [0, 3), [0, 4) }" + assert str(arr) == f"" class UnknownObjectCodecDtype(VariableLengthUTF8): diff --git a/tests/test_lazy_indexing.py b/tests/test_lazy_indexing.py new file mode 100644 index 0000000000..67403d1ea6 --- /dev/null +++ b/tests/test_lazy_indexing.py @@ -0,0 +1,785 @@ +"""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 + +import zarr +from zarr.core.buffer import default_buffer_prototype +from zarr.errors import BoundsCheckError, LazyViewError +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. + + Bounds are spelled literally: lazy coordinates are literal, so the + from-the-end `1:-1` spelling is rejected on the lazy path (see + TestLazyErrors.test_negative_slice_bounds_are_literal). + """ + 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_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]`` uses literal coordinates; ``-1`` is out of the domain.""" + 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]) + with pytest.raises(BoundsCheckError): + v[..., -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_flat_out_buffer(self) -> None: + """vindex with a multi-dim result and out= on a view uses a flat out buffer. + + Vectorized indexing scatters through a single flat index, so (as in the + eager path) the out buffer must be flat (shape = number of points). + """ + 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] + buf = default_buffer_prototype().nd_buffer.empty( + shape=(expected.size,), dtype=np.dtype("i4") + ) + v.get_coordinate_selection((i0, i1), out=buf) + np.testing.assert_array_equal( + np.asarray(buf.as_ndarray_like()).reshape(expected.shape), 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 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_is_literal(self) -> None: + """The lazy accessor copies TensorStore: indices are literal coordinates. + + Negative indices are absolute (origin 0), not from-the-end, so they fall + outside the ``[0, N)`` domain and raise — unlike the NumPy-normalizing + eager/view-method paths. Out-of-bounds array values raise cleanly rather + than silently wrapping. + """ + a, _ = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + with pytest.raises(IndexError): + _ = a.lazy[-1][...] + for sel in (np.array([-1], dtype=np.intp), np.array([24], dtype=np.intp)): + with pytest.raises(IndexError, match="out of bounds"): + _ = a.lazy.oindex[(sel,)][...] + with pytest.raises(IndexError, match="out of bounds"): + _ = a.lazy.vindex[(sel,)][...] + + def test_negative_slice_bounds_are_literal(self) -> None: + """Negative slice bounds are literal coordinates on the lazy path, like + every other index form — never from-the-end — so they raise. + + Consistency rule: a lazy selection is a declaration in literal + coordinates, and view domains start at 0, so a negative value is never + in bounds regardless of syntactic form (integer, slice bound, or index + array). + """ + a, _ = _make(CONFIGS[1]) # 1d-unsharded, shape (24,) + for sel in (slice(-3, None), slice(-24, None)): + with pytest.raises(BoundsCheckError, match="not contained"): + a.lazy[sel] + with pytest.raises(BoundsCheckError, match="not contained"): + a.lazy.oindex[(sel,)] + # bounds that RESOLVE reversed (stop < start), e.g. [1, -1) or [0, -1), + # are invalid intervals — TensorStore's a[5:2] case — not empties + for sel in (slice(None, -1), slice(1, -1)): + with pytest.raises(IndexError, match="interval"): + a.lazy[sel] + # ... on views too + v = a.lazy[2:10] + with pytest.raises(BoundsCheckError, match="not contained"): + v.lazy[-2:] + # ... and for writes + with pytest.raises(BoundsCheckError, match="not contained"): + a.lazy[-3:] = 0 + + 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_one_dialect_view_methods_use_domain_coordinates(self) -> None: + """A view has ONE coordinate system: its preserved domain. Eager methods + (`v[...]`, `v[k]`) use the same literal coordinates as `.lazy` — + matching TensorStore, where indexing a view of [2, 10) with 0 or -1 is + out of bounds. Zero-based NumPy-style access is spelled explicitly: + materialize (`result()` / `np.asarray`) or re-zero with `translate_to`. + """ + a, ref = _make(CONFIGS[1]) + v = a.lazy[2:10] + np.testing.assert_array_equal(v[2:5], ref[2:5]) + np.testing.assert_array_equal(v[3], ref[3]) + bad_reads: list[Callable[[], Any]] = [ + lambda: v[-1], + lambda: v[0:3], + lambda: v[-3:], + ] + for bad in bad_reads: + 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: + """All three literal-coordinate rejections raise BoundsCheckError (an + IndexError), with one message shape naming the valid range.""" + a, _ = _make(CONFIGS[1]) + triggers: list[Callable[[], Any]] = [ + lambda: a.lazy[-1], + lambda: a.lazy[-3:], + lambda: a.lazy.oindex[(np.array([-1], 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 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]) + + @pytest.mark.xfail( + reason="vindex-created views return shape-(1,) data for integer reads " + "where NumPy returns a scalar", + strict=True, + ) + 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]) == () + + 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_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] diff --git a/tests/test_properties.py b/tests/test_properties.py index 33888bfd4e..c71a713ed9 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -16,23 +16,27 @@ import hypothesis.extra.numpy as npst import hypothesis.strategies as st from hypothesis import assume, 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 ( + IndexMode, array_metadata, arrays, basic_indices, block_indices, block_test_arrays, complex_rectilinear_arrays, + numpy_array_indexers, numpy_arrays, - orthogonal_indices, rectilinear_arrays, simple_arrays, stores, + windows, zarr_formats, ) @@ -116,33 +120,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 +135,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 +330,463 @@ 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 not norm: + 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; the fancy modes + (oindex/vindex/mask via integer/boolean arrays) can't handle zero-size axes. + """ + if len(shape) == 0: + return False + return mode == "basic" or 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 +) -> 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 `out=` buffer + is flat for the vectorized vindex/mask modes (which scatter through a flat + index) and full-shape otherwise. + """ + 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: + buf_shape = (expected.size,) if mode in _VECTORIZED_MODES else 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( + simple_arrays(shapes=npst.array_shapes(max_dims=4)), + 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. + """ + zarray = data.draw(simple_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1))) + 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 (negatives from-the-end); + # lazy-path coordinates are literal, so canonicalize to the equivalent + # non-negative form for the view. NumPy semantics are invariant under this, + # so the reference side (npsel) is unchanged. + assert_read_matches_numpy(view, vref, mode, _canonical_nonneg(zsel, vref.shape), npsel) diff --git a/tests/test_transforms/__init__.py b/tests/test_transforms/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_transforms/test_chunk_resolution.py b/tests/test_transforms/test_chunk_resolution.py new file mode 100644 index 0000000000..893866740d --- /dev/null +++ b/tests/test_transforms/test_chunk_resolution.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.chunk_grids import ChunkGrid, FixedDimension +from zarr.core.transforms.chunk_resolution import iter_chunk_transforms, sub_transform_to_selections +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform + + +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 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 == () diff --git a/tests/test_transforms/test_composition.py b/tests/test_transforms/test_composition.py new file mode 100644 index 0000000000..e96616933d --- /dev/null +++ b/tests/test_transforms/test_composition.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr.core.transforms.composition import compose +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.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/tests/test_transforms/test_domain.py b/tests/test_transforms/test_domain.py new file mode 100644 index 0000000000..5a222a548c --- /dev/null +++ b/tests/test_transforms/test_domain.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import pytest + +from zarr.core.transforms.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/tests/test_transforms/test_json.py b/tests/test_transforms/test_json.py new file mode 100644 index 0000000000..db582a9561 --- /dev/null +++ b/tests/test_transforms/test_json.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.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.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform + + +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]} + 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(self) -> None: + domain = IndexDomain.from_shape((5,)) + json = index_domain_to_json(domain) + assert "input_labels" not in json + 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], + } + 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_omitted(self) -> None: + """stride=1 is the default and should be omitted from JSON.""" + m = DimensionMap(input_dimension=0) + json = output_index_map_to_json(m) + assert "stride" not in json + assert json == {"offset": 0, "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) + assert json == {"offset": 2, "stride": 3, "index_array": [1, 5, 9]} + 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_omitted(self) -> None: + arr = np.array([0, 1, 2], dtype=np.intp) + m = ArrayMap(index_array=arr) + json = output_index_map_to_json(m) + assert "stride" not in json + 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) + + +class TestIndexTransformJSON: + def test_identity(self) -> None: + t = IndexTransform.from_shape((10, 20)) + json = index_transform_to_json(t) + assert json == { + "input_inclusive_min": [0, 0], + "input_exclusive_max": [10, 20], + "output": [ + {"offset": 0, "input_dimension": 0}, + {"offset": 0, "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) + restored = index_transform_from_json(json) + assert isinstance(restored.output[0], ArrayMap) + np.testing.assert_array_equal(restored.output[0].index_array, idx) + assert isinstance(restored.output[1], DimensionMap) + + 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: + """Verify the JSON matches TensorStore's format exactly.""" + json: IndexTransformJSON = { + "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 diff --git a/tests/test_transforms/test_output_map.py b/tests/test_transforms/test_output_map.py new file mode 100644 index 0000000000..57dbebaf44 --- /dev/null +++ b/tests/test_transforms/test_output_map.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import numpy as np + +from zarr.core.transforms.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/tests/test_transforms/test_tensorstore_parity.py b/tests/test_transforms/test_tensorstore_parity.py new file mode 100644 index 0000000000..4670cf5928 --- /dev/null +++ b/tests/test_transforms/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.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.transform import IndexTransform +from zarr.errors import BoundsCheckError + + +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/tests/test_transforms/test_transform.py b/tests/test_transforms/test_transform.py new file mode 100644 index 0000000000..8e3810465c --- /dev/null +++ b/tests/test_transforms/test_transform.py @@ -0,0 +1,528 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from zarr.core.transforms.domain import IndexDomain +from zarr.core.transforms.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr.core.transforms.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) + np.testing.assert_array_equal(result.output[0].index_array, idx) + 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) + + +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) + + +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