feat: add IndexTransform library for composable, lazy coordinate mappings#3906
feat: add IndexTransform library for composable, lazy coordinate mappings#3906d-v-b wants to merge 84 commits into
Conversation
…ings Add a new `src/zarr/core/transforms/` package implementing 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 explicitly resolved. Key types: - `IndexDomain` — rectangular region in N-dimensional integer space - `ConstantMap`, `DimensionMap`, `ArrayMap` — three representations of a set of storage coordinates (singleton, arithmetic progression, explicit enumeration) - `IndexTransform` — pairs an input domain with output maps (one per storage dim) - `compose(outer, inner)` — chain two transforms Key operations on IndexTransform: - `__getitem__`, `.oindex[]`, `.vindex[]` — indexing produces new transforms - `.intersect(domain)` — restrict to coordinates within a region (chunk resolution) - `.translate(shift)` — shift coordinates (make chunk-local) The transform library is standalone with no dependency on Array. Includes comprehensive test suite (143 tests covering all types, operations, composition, chunk resolution, and edge cases). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
github-actions
Bot
added
the
needs release notes
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #3906 +/- ##
==========================================
- Coverage 93.50% 92.90% -0.61%
==========================================
Files 90 98 +8
Lines 11981 13250 +1269
==========================================
+ Hits 11203 12310 +1107
- Misses 778 940 +162
🚀 New features to boost your workflow:
|
|
@d-v-b I'm new to zarr-python indexing, does My use case is mainly if I have an array of shape |
Add TypedDict definitions and conversion functions for serializing
IndexDomain, OutputIndexMap, and IndexTransform to/from JSON.
The JSON format follows TensorStore's conventions for interoperability:
- IndexDomain: input_inclusive_min, input_exclusive_max, input_labels
- OutputIndexMap: offset + optional stride/input_dimension/index_array
- IndexTransform: domain fields + output array
TypedDicts: IndexDomainJSON, OutputIndexMapJSON, IndexTransformJSON
Functions: index_domain_to_json, index_domain_from_json,
index_transform_to_json, index_transform_from_json
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
on this branch: >>> arr = create_array(store={}, shape=(100,100), dtype="uint8")
>>> arr.z[0]
<Array memory://136980013828096 shape=(100,) dtype=uint8 domain={ 0, [0, 100) }>
>>> arr.z[0].shape
(100,)the |
|
What is I have and I want to select the arrays for band 1 (second dim) but for all the times, would |
yeah it should! |
Merging this PR will degrade performance by 17.94%
Performance ChangesComparing 
Footnotes
|
|
I played with this a bit today and this type of slicing is really great: One thing that tripped me up -- probably just shows my naivety (PBCAK) -- is that e.g. data.z.shape is |
We should probably have a shape attribute there. Thanks for reporting that, and thanks for trying this branch out. Let me know if you find anything else we need to fix! |
|
I guess from my PoV it would be nice if I could just use Edit: Don't get me wrong, this is pretty great -- as an example, with this I can remove the dask.array wrapper here: instead, just check if my array is zarr and if so use |
I'm glad it's useful! I agree that this is probably how slicing should work by default. But that would be a big breaking change, so it's far down the road. The |
|
Makes sense. So maybe in followup _LazyIndexAccessor could have more of the array properties and support more of the numpy protocol, so that array.z could be used as a drop-in for array, but be lazy about slicing, etc. |
|
I like this a lot! |
I just wanted something really short. Since this is (IMO) a strictly better slicing API, I wanted to make accessing it as low-friction as possible. But I will totally bow to the will of the crowd here, we can use whatever people find intuitive.
I use |
|
I agree with @normanrz, the the length of |
|
what about |
|
Just for reference: xarray uses |
|
ping @ilan-gold for visibility |
ilan-gold
left a comment
ilan-gold
left a comment
There was a problem hiding this comment.
use data structures compatible with tensorstore. tensorstore's indexing machinery serializes to JSON. if we adopt the same patterns, we are closer to using tensorstore as an optional backend.
I'm sure zarrs could also support this but I wonder about JSON performance for things like vindex/oindex where you could easily have 10's of thousands of individual coordinates.
| - **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``. |
There was a problem hiding this comment.
I think it could be cool to have different kinds of composition like union and I don't see that anywhere intthe PR - it seems that only array.z[my_transform][my_other_transform_relative_to_the_first] is supported wheres it could be cool to have array.z[my_transform + my_other_transform]. Concretely, if I wanted (slice(0, 10), slice(0, 10)) and then also (slice(50, 60), slice(50, 60)) from an array, how would I do that? np.concatenate + np.arange to generate outer indexers is the only option AFAICT but that kind of stinks.
There was a problem hiding this comment.
that's a very cool idea, not sure about overloading the addition operator, but some method for combining indexing expressions seems nice. I'll see what we need for this
Resolve independent (oindex) and correlated (vindex) ArrayMaps by their dependency axes during chunk intersection instead of assuming >=2 arrays are correlated or raveling full-rank arrays to 1-D. - _intersect routes on input_dimension is None (correlated marker): orthogonal maps are filtered along their own dependency axis at full input rank; correlated maps are jointly masked over the shared broadcast block while residual DimensionMap dims are narrowed against the chunk bounds and preserved (previously a rank-1 domain was paired with untouched DimensionMap outputs, raising ValueError at read time for any vindex with >=2 arrays plus a residual slice dim, including partial-rank boolean masks). - The correlated intersection now emits a flat row-major scatter index covering (surviving points, residual slice block), and sub_transform_to_selections consumes it as the single out-selection. - _reindex_array applies basic selections only along an ArrayMap's dependency axes; selections on singleton (broadcast) axes narrow the domain without touching the array's values, fixing basic slicing of the non-fancy axis of an oindex view. - Degenerate length-1 selections (all-singleton shapes) use input_dimension as the tie-breaker so they are not misclassified. - iter_chunk_transforms returns early for empty fancy selections instead of crashing on chunk_ids.min(). Assisted-by: ClaudeCode:claude-fable-5
Normalize public selections at the Array boundary: wrap NumPy-style negatives against the current view's domain (positive indices stay literal domain coordinates), reject boolean scalar indices, and route basic-slice-then-fancy composition through the view's transform. Generalizes the coordinate-selection wrapping from d0abd4b into one helper used by the lazy accessors and the non-identity view methods; the transform layer stays literal. Also reject fields= on a non-identity view (NotImplementedError) in all eight get/set_*_selection methods before any storage access, fixing silent corruption where the legacy indexer ignored the transform. Assisted-by: ClaudeCode:claude-fable-5
A boolean mask must exactly match the view domain's shape on the dimensions it consumes (NumPy parity: 'boolean index did not match indexed array'); under/over-length masks previously truncated silently and a too-high-rank mask raised ValueError from the transform layer. Validate in _normalize_public_selection against the current view domain, raising IndexError before transform construction. Also pin the pre-existing eager fields= sibling-field corruption with a strict xfail so the identity fields test documents rather than hides it. Assisted-by: ClaudeCode:claude-fable-5
Zero-rank correlated reads (e.g. an integer index into a vindex-created view) collapsed a flat (1,) working buffer with as_scalar, which returns shape-(1,) data; collapse to 0-d first so the caller gets a true scalar. The caller-visible `out` contract for vectorized reads is now the broadcast selection shape rather than a flat buffer: a flat temporary is used internally and the reshaped result is copied into the supplied buffer. Efficiency: reshape the freshly-allocated contiguous working buffer as a view instead of np.array(...)-copying, and drop the now-redundant second reshape in get_coordinate_selection's transform branch. Assisted-by: ClaudeCode:claude-fable-5
Property-based differential oracle for the index-transform layer: an IndexingProgram composes one to three indexing operations (basic prefix, optional trailing orthogonal/vectorized step) on a small array and checks five execution recipes (materialize, eager-on-lazy, out=, scalar and array writes) against a NumPy model applied in lockstep. Deterministic @example cases pin slice-then-oindex, unequal orthogonal lengths, multidimensional vindex with out=, a zero-rank result, and an empty result. Also restores zero-extent coverage lost with the deleted pre-branch tests: _indexing_array and the lazy-view parity test draw min_side=0 again (mask mode verified to work on zero-size shapes and now admitted by _eligible), and assert_read_matches_numpy learns the transform-path out= convention (broadcast result shape) so lazy-view vindex reads with multidimensional results are checked correctly. Assisted-by: ClaudeCode:claude-fable-5
The async surface built legacy indexers from metadata.shape without consulting a lazy view's index transform, silently reading/writing the wrong region (getitem/setitem/selection methods, oindex/vindex accessors, and from_array's copy loop). Guard every such entry point with a new AsyncArray._require_identity_for_async helper that raises LazyViewError on non-identity views, and precheck from_array up front so no half-created target array is left behind. Steer users to the sync surface / .result(). Assisted-by: ClaudeCode:claude-fable-5
|
Following this closely as a downstream consumer and I'd like to help get it over the line — it's a big PR and the coordinate algebra is exactly the primitive a data loader wants under it. A few concrete ways I could contribute; happy to take whichever is most useful to you. Context: insitubatch plans a training batch as a set of chunk reads and gathers decoded chunks into one buffer — i.e. it hand-rolls a sample→chunk coordinate mapping today. Ways I could help:
Let me know what's actually useful vs. noise. |
…out-of-grid guardrail `_is_sharded` keyed off `len(codecs) == 1`, misclassifying a valid sharded array with a trailing bytes-bytes codec (e.g. an outer compressor) as unsharded, so `chunk_projections(unit="read")` silently returned shard-granularity projections instead of raising. `_covers_full_chunk` returned partial for any integer chunk_selection entry, diverging from the write path's `_is_complete_chunk` for a constant over a chunk dim of extent 1. The transform read/write loops and `iter_chunk_projections` silently `continue`d on an out-of-grid chunk coordinate instead of raising, risking uninitialized-buffer reads and silently dropped writes if a transform composition bug ever produced one. Assisted-by: ClaudeCode:claude-fable-5
…, 0-d pin, ChunkProjection docs Address five small code-review items on the lazy-indexing branch: - Convert branch-introduced bare collection truthiness checks to explicit len() comparisons (array.py, transforms/transform.py, testing/strategies.py, test_properties.py). - Convert branch-introduced RST-style ``double-backtick`` markup and :role:`target` references to plain markdown backticks, since docs render via mkdocs/mkdocstrings (transforms/*.py, chunk_partition.py, errors.py, array.py). - Add changes/3906.feature.md documenting the new .lazy accessor and chunk_projections. - Pin the intentional 0-d array iteration behavior change (raises TypeError eagerly at iter(), matching numpy, instead of the old silent empty sequence-protocol iteration). - Document zarr.ChunkProjection under docs/api/zarr/array.md, which was exported in zarr.__all__ but missing from the API docs. Assisted-by: ClaudeCode:claude-fable-5
…forms package Move the TensorStore-style index-transform algebra out of zarr core into a new numpy-only uv-workspace subpackage packages/zarr-transforms (import name zarr_transforms), mirroring packages/zarr-metadata. This is a pure move: behavior is byte-for-byte unchanged and the full suite stays green. The package depends on numpy + stdlib only and must not import zarr. Two couplings are broken: - errors: the canonical BoundsCheckError / VindexInvalidSelectionError class definitions now live in zarr_transforms.errors (both still subclass IndexError). zarr.errors re-exports the same objects, so zarr.errors.BoundsCheckError is zarr_transforms.errors.BoundsCheckError and every existing catch site is unaffected. - chunk grid: chunk_resolution is typed against new structural Protocols (ChunkGridLike / DimensionGridLike) in zarr_transforms.grid instead of importing zarr's concrete ChunkGrid, which satisfies them structurally. zarr now declares a runtime dependency on zarr-transforms, resolved to the in-tree package via a uv workspace source. The package tests are collected by the root test suite (they exercise chunk resolution against zarr's ChunkGrid, so they need zarr importable) rather than run in isolation. The package __init__ promotes the names the zarr integration layer consumes (iter_chunk_transforms, sub_transform_to_selections, selection_to_transform) to the public surface alongside the existing exports. Assisted-by: ClaudeCode:claude-opus-4.8
Assisted-by: ClaudeCode:claude-opus-4.8
The hatch run-coverage / run-coverage-html / run-hypothesis (and gputest run-coverage) scripts passed --source=src to coverage run, which excluded the index-transform algebra after its move to packages/zarr-transforms/src — it was measured before the move. Declare the measured source trees in [tool.coverage.run] source instead, so every coverage run invocation measures both src and the workspace package consistently, and drop the CLI flag from all four scripts. Assisted-by: ClaudeCode:claude-opus-4.8
Great idea, and thanks for offering your help. I'm working on moving the lazy indexing logic here out into a separate package in the zarr-python workspace. That PR should be much easier to land, since it wouldn't change anything about zarr python. Once that package is up and running, we can publish it on pypi and you can start trying it out. How does this sound?
I'm happy to prototype either tier if that's a useful thing to peel off. I think we can sidestep this for now, because the JSON thing is only relevant when / if we want to interchange array selections expressions over the wire. We don't need a performant JSON serialization inside Zarr-Python itself.
this is a shared need -- see #4028. I think the lazy-indexing-oriented API would look like this: you declare the N regions of the array you want to read, then you arrange those N regions into a concatenated array, and then you submit the request to fetch that array to a zarr backend that can efficiently do all the IO you need. This is just an idea right now, but I think it would be a very clean API. |
Applying an oindex/vindex selection to a view that already carries an
orthogonal ArrayMap could land a fancy index on a broadcast (singleton)
axis of that map, leaking a raw numpy IndexError ("index N is out of
bounds for axis ... with size 1") at resolve time — reading like a user
error rather than the implementation gap it is.
Add `_guard_fancy_after_fancy`, invoked from `_apply_oindex` and
`_apply_vindex`, which raises a clear NotImplementedError at composition
time when a fancy axis targets a non-dependency axis of an existing
ArrayMap. It names the limitation and the workaround (materialize via
.result() then index, or reorder so the fancy step is last). Compositions
that keep the fancy step on the dependency axis (or on a correlated vindex
view) are unaffected and still resolve correctly.
Adds TestFancyAfterFancy, documents the limitation in the lazy-indexing
user guide, and refreshes the now-stale generator/runner comments that
described this class as a silent bug.
Assisted-by: ClaudeCode:claude-fable-5
…anges
zarr-transforms is a hard runtime dependency of zarr with no CI. Add
`zarr-transforms.yml` (push/PR path-filtered test/ruff/pyright, mirroring
zarr-metadata; the test job runs from the repo root because the transform
tests import zarr) and `zarr-transforms-release.yml` (tag-triggered publish
on `zarr_transforms-v*`). Extend `check_changelogs.yml` to check the
package's changes directory.
Bump the package `requires-python` floor to >=3.12 (consistent with the
repo; nothing tested 3.11) and update its classifiers, ruff target, and
pyright pythonVersion to match. Extend the root hatch `--match v*` comment
to note the `zarr_transforms-v*` tags it also excludes.
Disclose the two deliberate eager-path changes in the 3906 changelog: the
Array repr `domain={...}` suffix and the 0-d iteration TypeError matching
NumPy. Fix a stale comment in array.py that claimed pop_fields yields []
for no fields (it now yields None).
Assisted-by: ClaudeCode:claude-fable-5
Mirrors zarr_metadata's importlib.metadata idiom; the release workflow's isolated-wheel check imports it. Assisted-by: ClaudeCode:claude-fable-5
…ansforms The ArrayMap (fancy) branch of iter_chunk_transforms enumerated the dense range(min_chunk, max_chunk + 1) bounding box and ran transform.intersect against every candidate chunk, making sparse fancy/vindex chunk resolution scale O(n_chunks) instead of O(n_touched). The exact touched chunk ids were already computed and then discarded in favor of min/max. Enumerate each fancy dimension's distinct touched chunk ids (np.unique) instead; the cartesian product then spans only touched-per-dimension combinations. Constant/Dimension dims keep their contiguous ranges. Semantics are unchanged: the dense range only ever added empty intersections, which intersect already skipped. Sparse vindex (2 far-apart coords) goes from 15.9x slower than eager at 1k chunks / 65.5x at 4k to ~1.1x at both, flat in grid size. Dense fancy selections are unchanged. Assisted-by: ClaudeCode:claude-fable-5
github-actions
Bot
removed
the
needs release notes
Vendored, unmodified, from ndsel branch fix/slice-origin-trunc (commit c132b4c1caa3205830ce35a42502363171f650a7). PROVENANCE.md records the source URL, commit SHA, and do-not-edit note. The corpus is the language-agnostic fixture set every ndsel implementation runs against. Assisted-by: ClaudeCode:claude-opus-4.8
Add a pure JSON->JSON message layer (messages.py: parse_ndsel,
normalize_ndsel, NdselError) implementing the ndsel draft wire format,
which adapts TensorStore's IndexTransform. It accepts all five kinds
(point/box/slice/points/transform), normalizes to the canonical transform
body (spec 4.3), and enforces the full error taxonomy. Verified against the
vendored conformance corpus.
Rework json.py into the engine/lowering layer: transform_{to,from}_canonical
(re-pointed from index_transform_{to,from}_json). Engine constraints live
here only (reject infinite bounds, implicit-lowers-by-value). Fix the oindex
wire format so index_array maps no longer carry input_dimension (rejected by
ndsel and TensorStore); degenerate all-singleton arrays collapse to constant
maps, and input_dimension is reconstructed from array dependency axes on load.
Cross-checked against a real tensorstore (importorskip) for a set of finite
canonical bodies.
Assisted-by: ClaudeCode:claude-opus-4.8
The pyright job added alongside zarr-metadata's config had never actually run before an external PR triggered it, and fails with 68 errors on this branch's HEAD. Downgrade reportUnknownVariableType/Argument/Member/ParameterType to warnings: the strict config was copied from zarr-metadata's JSON/dataclass code, but numpy's stubs return partially-unknown types even at fully-typed call sites, so this numpy-heavy package can't reasonably satisfy that family. CI only fails on errors, so warnings keep the signal visible without blocking the build. Fix the genuine findings in code: - reportUnnecessaryIsInstance (14 sites across transform.py, composition.py, json.py, chunk_resolution.py): replace tautological final isinstance arms (pyright proves them always-true once prior branches exhaust the union) with `else:` plus a comment naming the narrowed type. Two sites in composition.py and one in json.py also drop a `# pragma: no cover`-marked dead `raise TypeError` that followed. - reportPrivateUsage (2 sites): chunk_resolution.py's `chunk_grid._dimensions` (declared on the ChunkGridLike Protocol for structural typing against zarr's ChunkGrid) and json.py's cross-module import of transform.py's `_array_map_dependency_axes`. Both are suppressed with `# pyright: ignore[reportPrivateUsage]` plus a comment; renaming either is left as an open pre-publish API decision. Zero behavior change. Verified: pyright 0 errors (was 68); pytest 266 passed, 1 skipped; mypy clean; ruff clean. Assisted-by: ClaudeCode:claude-fable-5
Fixtures verified byte-identical to the merged c59bc556c; only the reference moves off the deleted feature branch. Assisted-by: ClaudeCode:claude-fable-5
…nsforms The fast path bound 'm' as ArrayMap before the general loop rebinds it to the OutputIndexMap union, which mypy rejects (and treats the ConstantMap arm as unreachable). packages/ is outside the pre-commit mypy scope, so #234's CI could not catch this. Assisted-by: ClaudeCode:claude-fable-5
…ndexing # Conflicts: # pyproject.toml # src/zarr/core/array.py # uv.lock
…indexing Pre-publish rename: the in-tree zarr-transforms workspace package is not yet published, so this is a safe, purely mechanical rename with zero behavior change. Distribution name is now zarr-indexing; import name is now zarr_indexing. Covers: the package directory and src layout, the package's own pyproject.toml (name, tag-pattern, wheel packages, towncrier package), the root pyproject.toml (workspace member, uv source, runtime dependency, coverage source, pytest testpaths, hatch version comment), all imports across src/zarr, tests, and the package's own tests, the zarr-transforms(-release).yml workflows (renamed to zarr-indexing(-release).yml with matching job names, path filters, tag pattern, environment names, artifact names, and smoke test), check_changelogs.yml's package path, changelog fragments (root changes/3906.feature.md and the package's own changes/), and the regenerated uv.lock. Assisted-by: ClaudeCode:claude-fable-5
…classes
Identity (non-view) arrays now render with the legacy 3.x repr with no
`domain={...}` suffix, byte-identical across `Array` and `AsyncArray`.
The suffix appears only on non-identity lazy views, and `AsyncArray`
gains it for views (reachable via `translate_by` and friends). Restores
the doctests and repr tests that had been rewritten to expect the
unconditional suffix on identity arrays.
Assisted-by: ClaudeCode:claude-opus-4.8
`get_coordinate_selection(out=...)` requires a flat `(n,)` buffer on eager/identity arrays but the broadcast selection shape on lazy views. Document both sides in the `out` parameter docstring and note that the broadcast-shape contract is the long-term direction. No behavior change. Assisted-by: ClaudeCode:claude-opus-4.8
…nning docs) - name the zarr-indexing changelog fragment by PR number (3906) so the integer-issue check passes - drop :class: sphinx roles from the indexing-program docstrings (mkdocs markdown; caught by main's new ci/lint_docs.py) - untrack docs/superpowers/ planning documents: upstream main gitignores that path (agent planning notes are local-only by convention), which also removes them from markdownlint's scope; files remain on disk Assisted-by: ClaudeCode:claude-fable-5
9 participants
My summary:
With dask maintenance on the decline, it's more important than ever that we give zarr-python users a dask-free way to do something very intuitive: index large zarr arrays without turning the whole thing into a numpy array first. This was discussed at length in #1603.
This PR, done with Claude, makes regular indexing go through a lazy indexing layer. The lazy indexing layer is based on abstractions defined in tensorstore. The basic idea is to explicitly model indexing an array as a transformation from some input coordinates to output coordinates, and to bind such a representation to our
Arrayclasses.Regular indexing via
.__getitem__is still immediate, but arrays have a new.zattribute that exposes the lazy indexing layer:Goals here:
mainare disparate ad-hoc copies of stuff from zarr-python 2.x. We can do better.Non-goals:
Claude's summary.
Add a new
src/zarr/core/transforms/package implementing TensorStore-inspiredindex 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 explicitly resolved.
Key types:
IndexDomain— rectangular region in N-dimensional integer spaceConstantMap,DimensionMap,ArrayMap— three representations of a set ofstorage coordinates (singleton, arithmetic progression, explicit enumeration)
IndexTransform— pairs an input domain with output maps (one per storage dim)compose(outer, inner)— chain two transformsKey operations on IndexTransform:
__getitem__,.oindex[],.vindex[]— indexing produces new transforms.intersect(domain)— restrict to coordinates within a region (chunk resolution).translate(shift)— shift coordinates (make chunk-local)The transform library is standalone with no dependency on Array.
Includes comprehensive test suite (143 tests covering all types, operations,
composition, chunk resolution, and edge cases).
Co-Authored-By: Claude Opus 4.6 (1M context) noreply@anthropic.com