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

.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

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

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.

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()

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.

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

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.