feat(zarr-indexing): factor chunk plans into a columnar GridPartition

[d-v-b]

🤖 AI text below 🤖

Summary

Chunk resolution in zarr-indexing built two IndexDomains, two IndexTransforms and a ChunkProjection per chunk, and for correlated (vindex) selections rescanned the whole index array once per candidate chunk. This PR lifts that construction from the chunk to the axis.

Restricting a transform to a chunk box distributes over output dimensions whenever each output map reads its own input axis (every basic and orthogonal selection). So each axis is resolved once against its grid into a table, and a projection is one row of each table:

• StridedSet: one output dimension read through a ConstantMap or DimensionMap, one row per touched chunk (chunk, chunk_start, chunk_extent, local_start, extent, origin, full).
• IndexedSet: an orthogonal ArrayMap, its coordinates grouped by chunk in CSR form (pointer, index, positions, derived local).
• JointSet: the correlated index arrays, sorted into chunks together once (chunk, pointer, index, positions, block_coordinates), since a chunk constrains all of them at once and they do not distribute.
• GridPartition: the tables plus row_shape; n_rows, chunk_coords() (vectorized) and __iter__ derive ChunkProjection rows on demand. Columns are read-only.

ChunkPlan.partition() memoizes the partition and ChunkPlan.projections() iterates it; there is no second mechanism. A hand-built diagonal (two DimensionMaps reading one input axis, which no selection produces) is rejected with ValueError; the whole-transform walk that main used for it produced wrong projections (a three-point diagonal yielded four projections covering six cells), so rejection is strictly better. This is the shape of TensorStore's IndexTransformGridPartition, with two simplifications noted in the design notes.

The projections a plan yields are unchanged for every transform a selection can produce: a parametrized evaluation oracle (both transforms agree pointwise, cells tile the request exactly once, chunk-local coordinates lie in the chunk, coverage is full exactly when the chunk's cells are each read once) runs across basic, strided, reversed, scalar, orthogonal, one-element, correlated, 2-D block, sorted 1-D, empty and rank-0 cases on fixed and varying grids, and an out-of-tree differential harness matched main field-for-field over ~24k random transforms during review.

Supporting changes: _prepare_correlated flattens a transform's correlated arrays once (shared by _intersect_general and the joint table); checked_affine gains an identity fast path (measured to halve the correlated object-row walk; a dtype-bound variant measured at noise was dropped in review); ArrayMap._with_affine shares the frozen array when a map is translated; IndexDomain._unchecked / IndexTransform._unchecked skip validation for objects derived from an already-valid transform (removing them measured +60% on basic and +130% on orthogonal walks).

Reviewed adversarially before merge-readiness (roborev, a correctness reviewer, a complexity reviewer); the last commit applies their findings and cuts ~860 lines, including everything that made a second mechanism.

Performance

Indexer-level cost of planning plus consuming, measured through a zarr indexer port on top of this package (legacy = today's zarr/core/indexing.py):

Case	legacy	projections (objects)	tables read directly
basic full read, 1331 chunks	0.94 ms	3.45 ms (was 21.8 ms)	0.58 ms
shard-like full read, 4096 chunks	2.89 ms	10.2 ms (was 68 ms)	1.72 ms
2-D coordinate selection, 10k points over 6304 chunks	10.9 ms	61 ms (was 472 ms)	6.6 ms

"was" is main. The object view is bounded by the five allocations per chunk it still makes; reading the tables is already faster than the legacy indexer, which is the form a consumer should move to. End-to-end zarr benchmarks with the port went from 1.43x to 1.12x of legacy from the projection-path work alone; consuming the tables in the codec pipeline is the follow-up that would take that below 1.0x, and is not part of this PR.

Documentation

The old narrative (intersect the whole transform with every candidate chunk) is retired everywhere it appeared, not just supplemented:

• chunk_resolution module docstring: explains the factored form, the three tables, the sum-versus-product cost, and that the whole-transform walk remains only for hand-built diagonals.
• Visual guide: a new final integrator section, A plan is a product of per-axis tables, with an executable snippet (docs/snippets/grid_partition.py, auto-discovered and run by test_doc_examples.py) that reads StridedSet, IndexedSet and JointSet tables off real plans and checks the plan's projections against the partition's rows.
• Integration boundaries: Reading the tables directly, a consumer that assembles a strided box from the tables with no projection materialized, including the literal-versus-positional origin detail.
• API index, landing page and design notes (TensorStore lineage: IndexTransformGridPartition; the performance caveat; the box/query split) point at the new section.
• The stale claim that the package's tests need zarr's ChunkGrid is corrected in pyproject.toml, the justfile, and the CI workflow: nothing in the package imports zarr; the suite runs from the repo root to share the pinned toolchain.

mkdocs build --strict passes.

Testing

• packages/zarr-indexing: 1331 passed, doctests and executable doc snippets included; pyright 0 errors; mypy below main's baseline; ruff clean; mkdocs build --strict passes.
• tests/test_indexing.py on unmodified zarr: 447 passed.
• A differential fuzzer (thousands of random selections per dialect, gathering and scattering through the port with codec-size chunk buffers, against NumPy and the legacy indexer) found no mismatches.

🤖 Generated with Claude Code