diff --git a/.github/workflows/zarr-indexing.yml b/.github/workflows/zarr-indexing.yml index f0b0cec65d..3b106e16aa 100644 --- a/.github/workflows/zarr-indexing.yml +++ b/.github/workflows/zarr-indexing.yml @@ -43,11 +43,10 @@ jobs: uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 - name: Set up Python ${{ matrix.python-version }} run: uv python install ${{ matrix.python-version }} - # The transform tests exercise chunk resolution against zarr's ChunkGrid, - # so they run against the repo-root environment (which provides `zarr`) - # with this package as an editable overlay rather than in package - # isolation. The recipes carry that invocation; this step only fixes the - # interpreter the matrix asked for. + # The suite imports nothing from `zarr`; it runs against the repo-root + # environment, with this package as an editable overlay, to share the + # parent project's pinned test toolchain. The recipes carry that + # invocation; this step only fixes the interpreter the matrix asked for. - name: Sync test dependency group run: uv sync --project ../.. --group test --python ${{ matrix.python-version }} - name: Run pytest diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index 7e2cec10dd..4328de1cc1 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -25,6 +25,9 @@ Key types: - `ChunkPlan` and `ChunkProjection` — lazily partition a selection over a caller-selected grid and pair each chunk-local transform with its placement in the request, without binding a storage backend or scheduler +- `GridPartition` — the plan's factored, columnar form: one table per axis + (`StridedSet`, `IndexedSet`) plus `joint_sets` for connected index-array components, + from which projections are derived on demand - `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single output dimension can depend on the input - `compose` — chain two transforms into one diff --git a/packages/zarr-indexing/benchmarks/README.md b/packages/zarr-indexing/benchmarks/README.md new file mode 100644 index 0000000000..28ef3e6a16 --- /dev/null +++ b/packages/zarr-indexing/benchmarks/README.md @@ -0,0 +1,28 @@ +# Chunk planning benchmarks + +Run from the repository root, with `zarr-indexing` installed in the selected +Hatch environment: + +```sh +hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/chunk_planning.py --repeats 9 +``` + +`chunk_planning.py` measures construction, projection traversal, local-column +access, and chunk-coordinate enumeration. See the script for workloads and output fields. + +Use the same interpreter, dependencies, inputs, and script revision when +comparing checkouts. Put the intended checkout's package source on `PYTHONPATH`; +installed package metadata is also required for version lookup. Record both git +revisions and the interpreter/dependency versions alongside saved output. + +Compare complete operations, not only constructor times: plans can defer work +until iteration. Measure allocation separately from elapsed time. Distinguish +streaming consumption from retaining all projections or coordinates, and report +whether input arrays and transform construction are included. Repeated local +column access measures cache reuse, which trades allocation against retained +memory. Bounded coordinate batches avoid constructing the full coordinate array. + +These scripts measure planning rather than codec or storage throughput. Repeat +measurements with alternating operation order before interpreting small timing +differences. Preserve raw benchmark output as an experiment artifact rather than +accumulating successive result tables in this README. diff --git a/packages/zarr-indexing/benchmarks/chunk_planning.py b/packages/zarr-indexing/benchmarks/chunk_planning.py new file mode 100644 index 0000000000..934d68dcbc --- /dev/null +++ b/packages/zarr-indexing/benchmarks/chunk_planning.py @@ -0,0 +1,126 @@ +"""Reproducible planning measurements; no storage I/O or timing assertions. + +Run with the package on PYTHONPATH in the repository's Hatch test environment. +Use the same interpreter and this script with the old/new package paths to compare. +The connected-components probe is deliberately not a production planner: it +measures independent table construction, without defining a new public row API. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +import tracemalloc +from typing import TYPE_CHECKING, Any + +import numpy as np + +from zarr_indexing import ArrayMap, IndexDomain, IndexTransform, plan_chunks +from zarr_indexing.grid import dimension_grids_from_chunks + +if TYPE_CHECKING: + from collections.abc import Callable + + +def measure(operation: Callable[[], Any], repeats: int) -> dict[str, float]: + samples = [] + for _ in range(repeats): + start = time.perf_counter() + operation() + samples.append(time.perf_counter() - start) + tracemalloc.start() + result = operation() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + del result + return {"median_ms": statistics.median(samples) * 1000, "peak_mib": peak / 2**20} + + +def consume(values: Any) -> int: + return sum(1 for _ in values) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + results: dict[str, Any] = {} + indices = np.arange(1000) + cases = [ + ("box", IndexTransform.from_shape((1000, 1000)), (10, 10), (1000, 1000)), + ( + "sorted", + IndexTransform.from_shape((1_000_000,)).oindex[np.arange(1_000_000)], + (100_000,), + (1_000_000,), + ), + ( + "correlated", + IndexTransform.from_shape((10_000, 10_000)).vindex[indices * 9, indices * 9], + (10, 10), + (10_000, 10_000), + ), + ( + "correlated_slab", + IndexTransform.from_shape((1000, 1000, 1000)).vindex[indices, indices, :], + (1000, 1000, 1000), + (1000, 1000, 1000), + ), + ] + for name, transform, chunks, shape in cases: + grids = dimension_grids_from_chunks(chunks, shape=shape) + results[name] = measure( + lambda transform=transform, grids=grids: consume(plan_chunks(transform, grids)), + args.repeats, + ) + + transform = IndexTransform.from_shape((10_000,)).oindex[np.arange(10_000)] + grids = dimension_grids_from_chunks((10,), shape=(10_000,)) + (table,) = plan_chunks(transform, grids).partition().sets + + def read_local_rows() -> None: + for row in range(len(table)): + table.local[table.run(row)] + + results["local_rows"] = measure(read_local_rows, args.repeats) + part = plan_chunks( + IndexTransform.from_shape((100, 100, 100)), + dimension_grids_from_chunks((1, 1, 1), shape=(100, 100, 100)), + ).partition() + results["all_coordinates"] = measure(part.chunk_coords, args.repeats) + + # TensorStore factors connected components in the input/grid dependency + # graph. Here outputs 0 and 1 depend on u, while output 2 depends on v. + # A single joint block unnecessarily expands these independent components. + i = np.arange(1000) + transform = IndexTransform( + IndexDomain.from_shape((1000, 1000)), + (ArrayMap(i[:, None]), ArrayMap(i[:, None]), ArrayMap(i[None, :])), + ) + grids = dimension_grids_from_chunks((1000, 1000, 1000), shape=(1000, 1000, 1000)) + results["component_partition"] = measure( + lambda: plan_chunks(transform, grids).partition(), args.repeats + ) + results["component_walk"] = measure( + lambda: consume(plan_chunks(transform, grids)), args.repeats + ) + base = IndexTransform.from_shape((1000, 1000, 1000)) + results["component_vindex_compile"] = measure( + lambda: base.vindex[i[:, None], i[:, None], i[None, :]], args.repeats + ) + left = IndexTransform(IndexDomain.from_shape((1000,)), (ArrayMap(i), ArrayMap(i))) + right = IndexTransform(IndexDomain.from_shape((1000,)), (ArrayMap(i),)) + results["independent_components_probe"] = measure( + lambda: ( + plan_chunks(left, grids[:2]).partition(), + plan_chunks(right, grids[2:]).partition(), + ), + args.repeats, + ) + print(json.dumps(results, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/packages/zarr-indexing/changes/4310.feature.md b/packages/zarr-indexing/changes/4310.feature.md new file mode 100644 index 0000000000..96191c0cc7 --- /dev/null +++ b/packages/zarr-indexing/changes/4310.feature.md @@ -0,0 +1,20 @@ +Chunk plans now have a factored, columnar form. `ChunkPlan.partition()` +returns a `GridPartition`: one `StridedSet` or `IndexedSet` table per output +dimension the transform reads independently, plus `joint_sets`, one `JointSet` +per connected component of index arrays. Arrays that share request axes are +sorted into chunks together; independent ones stay in separate tables, so +`(u, v) -> (a[u], b[u], c[v])` stores 3,000 index values rather than +3,000,000. A `ChunkProjection` is one row of each table and is derived on +demand, so planning costs the sum of the touched chunks per axis rather than +their product, and a vectorized consumer can read the tables (`chunk_coords()`, +the CSR columns, the memoized read-only `local` coordinates) without +materializing a projection per chunk. + +The projections a plan yields map the same cells to the same storage as +before. Correlated projections now keep residual slices symbolic and index +arrays compact along axes they do not vary over, and `vindex` compilation +preserves singleton axes for the same reason. + +A hand-built transform in which two output maps read one input axis (a +diagonal, which no selection produces) is rejected with `ValueError`; the +whole-transform walk that previously served it produced wrong projections. diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index 674c701a0f..a5b7163bbd 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -11,9 +11,9 @@ links: [Lazy views compose](../guide/index.md#lazy-views-compose), then open [`zarr_indexing.lazy_array`](lazy_array.md) for `LazyArray`. - **Integrate a chunked source:** finish - [One cell domain, two projections](../guide/index.md#one-cell-domain-two-projections), + [A plan is a product of per-axis tables](../guide/index.md#a-plan-is-a-product-of-per-axis-tables), then open [`zarr_indexing.chunk_resolution`](chunk_resolution.md) for - `plan_chunks`. Start with + `plan_chunks` and `GridPartition`. Start with [Coordinates are addresses](../guide/index.md#coordinates-are-addresses) if literal coordinates are unfamiliar. @@ -36,7 +36,10 @@ and the wire format built on top of it. - [`zarr_indexing.chunk_resolution`](chunk_resolution.md) — `plan_chunks`, which lazily projects a request through a caller-selected grid, - plus the reusable `ChunkPlan` and paired-transform `ChunkProjection` values + the reusable `ChunkPlan` and paired-transform `ChunkProjection` values, and + the plan's factored form: `GridPartition` (from `ChunkPlan.partition`), + holding one `StridedSet` or `IndexedSet` table per axis and `joint_sets` + for connected index-array components - [`zarr_indexing.grid`](grid.md) — `DimensionGridLike`, the Protocol describing the narrow chunk-grid surface chunk resolution consumes, so that nothing here imports `zarr`, plus `EdgeDimensionGrid` and diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index 78b352d376..59eef4dacd 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -35,27 +35,21 @@ about the deliberately matching semantics: discriminator, and `tests/test_ndsel_tensorstore.py` loads our bodies into `tensorstore.IndexTransform(json=...)` and round-trips them back through our engine layer. - -The representations differ in one place: index arrays. Both models want an index -array at the transform's full input rank, with singleton axes for the dimensions -a map does not vary over. TensorStore enforces it — its JSON parser rejects a -rank-1 array over a rank-2 domain outright, with `Index array for output -dimension 0 has rank 1 but must have rank 2` (checked against tensorstore -0.1.84) — while our loader is the more permissive of the two and also accepts a -lower-rank array that broadcasts against the input domain. That is a -compatibility affordance, not a difference in the model: ndsel leaves index-array -rank to [the engine layer](ndsel.md#lowering-to-a-transform), and everything the -algebra builds itself is at full rank. - -The reason full rank matters here is that we *derive* meaning from those -singletons rather than merely tolerating them: an array full-sized on one axis -and singleton elsewhere is orthogonal, and one varying over several shared axes -is vectorized, so the distinction is readable off the shape — and the shape is -the *only* place it lives. An earlier `ArrayMap.input_dimension` field pinned -the orthogonal axis redundantly and was retired: the one shape it disambiguated -(a single-coordinate array, all axes singleton) is now normalized away at -construction, collapsed to the `ConstantMap` it equals, exactly as -[the serializer](api/json.md) has always collapsed it on the wire. +- **Chunk partitioning.** Both factor a transform over a grid before visiting + any cell, rather than intersecting the whole transform with each chunk. + TensorStore's `IndexTransformGridPartition` holds strided sets and index + array sets and derives a per-cell transform on iteration; + [`GridPartition`](api/chunk_resolution.md) also partitions index-array + dependency components. Its `StridedSet` is per storage axis, where TensorStore's + is per input dimension and spans every grid axis reading it (which is how + TensorStore factors diagonals). `joint_sets` holds one `JointSet` per + connected index-array component, including independent single-array + components within a mixed request. Both derive the per-chunk transforms + from the partition ([the guide](guide/index.md#a-plan-is-a-product-of-per-axis-tables) + shows the tables). TensorStore keeps strided sets implicit, while this + library materializes their per-axis rows for vectorized consumers. Diagonals are rejected here; supporting them needs a strided set per + *input* dimension spanning every storage axis that reads it, TensorStore's + representation. Four deliberate differences: @@ -85,11 +79,14 @@ safe, `partial` proves it is not, and `unknown` conservatively covers fancy selections whose duplicates would require additional work to classify. The comparison also runs the other way. TensorStore is a mature, heavily -optimized C++ system whose performance this library cannot approach: resolution -here is Python-level bookkeeping over NumPy, and the per-part overhead is -significant. This library is small and depends on nothing beyond NumPy, so the -algebra can be adopted by a Python project that wants the model without the C++ -runtime. +optimized C++ system whose performance this library cannot approach. Independent strided planning +here is per axis, but each materialized `ChunkProjection` is +Python-level bookkeeping over NumPy — two domains, two transforms and the +projection itself — so the per-part overhead of the object view is +significant; a consumer that reads the partition's tables directly pays no +per-chunk object construction. This library is small and depends on nothing beyond +NumPy, so the algebra can be adopted by a Python project that wants the model +without the C++ runtime. ## Bounding-box selections vs query selections @@ -277,15 +274,20 @@ rewritten in place. Resolution classifies the result by structure (`index_array_structure`): pure per-axis outer products keep the orthogonal resolvers, and everything else — correlated maps, mixtures, index arrays sharing an input axis (a diagonal gather, reachable only by hand-building a -transform) — takes the pointwise path that collapses the joint block. +transform) — takes the general reader/intersection path. Chunk planning +factors index arrays into connected dependency components before flattening, +so independent groups do not expand one another. Vectorized selection preserves +broadcast singletons to retain those dependencies. Three limits remain, all intentional and all expected to be lifted: -- **Affine diagonals.** A hand-built transform in which an *index array* and a - *slice map* bind the same input dimension, or two slice maps share one, is - rejected at resolution with `NotImplementedError`. No selection dialect - produces one; supporting them means lowering the slice maps into the joint - block too. *Planned.* +- **Affine diagonals.** A hand-built transform in which two output maps read + one input dimension — two slice maps, or a slice map and an orthogonal index + array — is rejected at planning with `ValueError`; a correlated index array + varying over a dimension a slice map also reads is rejected with + `NotImplementedError`. No selection dialect produces either. Supporting them + needs a strided set per *input* dimension spanning all dependent storage + axes, TensorStore's connected-component representation. *Planned.* - **Finite explicit bounds only.** `IndexDomain` has no implicit or unbounded dimensions; the message layer will normalize a body with `"-inf"`/`"+inf"` bounds, but the engine layer refuses to lower one into a transform. @@ -294,3 +296,8 @@ Three limits remain, all intentional and all expected to be lifted: dimension labels and the wire format round-trips them, but indexing operations build new domains without them, so a label does not survive a slice. *Planned.* + +## Selection to chunk operations + +[The selection-flow guide](guide/selection-flow.md) traces normalization, +partitioning, and paired chunk-local/request projections with executable examples. diff --git a/packages/zarr-indexing/docs/guide/index.md b/packages/zarr-indexing/docs/guide/index.md index 42e47530fb..68192f58cb 100644 --- a/packages/zarr-indexing/docs/guide/index.md +++ b/packages/zarr-indexing/docs/guide/index.md @@ -8,8 +8,9 @@ through those stages. The first four sections are for anyone indexing arrays: coordinates, transforms, composition, and result axes. **If you are using lazy indexing rather than building a storage backend, you can stop after section four.** -The last two sections are for integrators: they turn a request into a chunk -plan and pair each chunk read with its place in the result. +The last three sections are for integrators: they turn a request into a chunk +plan, pair each chunk read with its place in the result, and show the per-axis +tables the plan is built from. Throughout, one division of labor holds: the transform answers **which values?** and is independent of the backend; the reader answers **how do I @@ -435,6 +436,83 @@ The paired representation preserves information that a bounding box or local selector discards: exact request order, duplicate destinations, and the correspondence between every request position and its chunk-local source cell. +## A plan is a product of per-axis tables {#a-plan-is-a-product-of-per-axis-tables} + +Look again at the two projections of `image[1, :]`. Each chunk transform has +one map per source axis, and every one of those maps came from restricting +the request's map for *that axis alone* to *that axis's* chunk: the fixed row +`ConstantMap(1)` lands in row-chunk 0 whatever the column chunk is, and the +column slice meets column-chunk 0 as local columns `0:2` and column-chunk 1 as +local columns `0:2` whatever the row chunk is. Restricting a transform to a +chunk distributes over axes whenever each output map reads its own request +axis — which every basic and orthogonal selection satisfies. So the plan does +not intersect the whole transform with every chunk. It resolves each axis +once, into a table with one row per chunk that axis touches, and a projection +is one row of each table combined. + +```text +image[1, :] over 2-by-2 chunks + +axis 0 (rows): ConstantMap(1) axis 1 (columns): DimensionMap +row | chunk start local extent row | chunk start local_start extent origin full + 0 | 0 0 1 1 0 | 0 0 0 2 0 yes + 1 | 1 2 0 2 2 yes + +row_shape (1, 2): 1 x 2 = 2 projections +projection (0, 0) = axis-0 row 0 x axis-1 row 0 -> chunk (0, 0), local (1, 0:2), request columns 0:2 +projection (0, 1) = axis-0 row 0 x axis-1 row 1 -> chunk (0, 1), local (1, 0:2), request columns 2:4 +``` + +`ChunkPlan.partition()` returns this factored form, a `GridPartition`. Its +`sets` hold one table per source axis the request reads independently, in +axis order, and `joint_sets` one table per connected group of index arrays; `row_shape` is the +number of rows in each; and the plan walks the rows in row-major order over +it. The executable example reads the two tables above off the plan, checks +that the plan's projections are exactly the partition's rows, and evaluates +the second row on both of its transforms: + +```python +--8<-- "snippets/grid_partition.py:strided-tables" +``` + +There are three kinds of table, matching the three map kinds and the one +arrangement that does not factor. A `StridedSet` holds a `ConstantMap` or +`DimensionMap` axis; an `IndexedSet` holds an orthogonal `ArrayMap` axis +(`.oindex`) with its coordinates grouped by chunk; and each `JointSet` holds one connected group of index arrays. Arrays +sharing request axes are sorted together; independent groups are separate +tables in `joint_sets`. Singleton broadcast axes retain this independence. The +[API reference](../api/chunk_resolution.md) documents every column. + +The gather from the previous section, `oindex[[4, 1, 1], 2:6]` over 3-by-4 +chunks, groups rows `1, 1` into chunk 0 and row `4` into chunk 1 while +remembering that row `4` fills request position 0. A `vindex` selection keeps +its points paired in the joint table: + +```python +--8<-- "snippets/grid_partition.py:indexed-and-joint" +``` + +Building independent strided tables costs the *sum* of touched chunks per +axis. Index-array grouping also scales with the selected point count; the +planner flattens each connected component separately. Independent components +do not expand into their Cartesian product until their rows are iterated. And projections are derived from rows only when asked for: +`chunk_coords()` lists every chunk the plan touches without materializing a +row, and a consumer can read the columns directly, as +[Integration boundaries](integrations.md#reading-the-tables-directly) shows. +For example, `(u, v) -> (a[u], b[u], c[v])` has two independent components. +The first two storage axes share `u`; the third reads `v`. With 1,000 values +per input axis, planning stores 3,000 index values rather than 3,000,000. +Each projection has one synthetic axis per nonconstant component, followed +by residual request axes, including axes no output reads. + +```python +--8<-- "snippets/grid_partition.py:components" +``` + +A hand-built affine diagonal — two `DimensionMap`s reading one request axis, +which no selection produces — would need a table spanning several storage +axes; the plan rejects it with `ValueError`. + ---