Prototype prepared indexing execution for codec consumers

.github/workflows/zarr-indexing.yml

@@ -43,8 +43,8 @@ jobs:
         uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4
       - name: Set up Python ${{ matrix.python-version }}
         run: uv python install ${{ matrix.python-version }}
-      # The suite imports nothing from `zarr`; it runs against the repo-root
-      # environment, with this package as an editable overlay, to share the
+      # The suite runs against the repo-root environment, with this package
+      # as an editable overlay, for Zarr codec integration tests and 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

packages/zarr-indexing/benchmarks/README.md

@@ -26,3 +26,19 @@ 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.
+
+## Experimental execution
+
+The execution follow-up adds two scripts:
+
+```sh
+hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/execution.py
+hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/execution_io.py
+```
+
+`execution.py` compares Zarr indexers, declarative projections, and execution
+selectors, including retained rows, borrowed inputs, snapshots, and shard
+lowering. `execution_io.py` verifies and measures MemoryStore reads and writes
+through real codec pipelines. Neither establishes filesystem or cloud throughput.
+The prototype remains opt-in; planning wins alone do not justify replacing the
+existing indexers.

packages/zarr-indexing/benchmarks/execution.py

@@ -0,0 +1,181 @@
+"""Compare PR planner with existing Zarr indexers; no storage I/O.
+
+Grid construction and input selection allocation are excluded for both sides.
+Selection compilation/indexer construction and complete streaming walks are included.
+Run with the worktree src and packages/zarr-indexing/src on PYTHONPATH in Hatch.
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import statistics
+import time
+import tracemalloc
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import zarr.core.indexing as zi
+from zarr.core.chunk_grids import ChunkGrid
+
+from zarr_indexing import IndexTransform, plan_chunks
+from zarr_indexing._execution import execute_selection
+from zarr_indexing.grid import dimension_grids_from_chunks
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+def consume(iterator: Any) -> int:
+    return sum(1 for _ in iterator)
+
+
+def measure(op: Callable[[], Any], repeats: int = 31) -> dict[str, float]:
+    op()
+    samples = []
+    for _ in range(repeats):
+        start = time.perf_counter()
+        op()
+        samples.append((time.perf_counter() - start) * 1000)
+    tracemalloc.start()
+    result = op()
+    _, peak = tracemalloc.get_traced_memory()
+    tracemalloc.stop()
+    del result
+    return {"ms": statistics.median(samples), "peak_mib": peak / 2**20}
+
+
+def main() -> None:
+    i = np.arange(1000)
+    cases = [
+        ("basic_10000_chunks", (1000, 1000), (10, 10), (slice(None), slice(None)), "basic"),
+        (
+            "sorted_1M_points_10_chunks",
+            (1_000_000,),
+            (100_000,),
+            (np.arange(1_000_000),),
+            "orthogonal",
+        ),
+        (
+            "sorted_coordinate_1M_points_10_chunks",
+            (1_000_000,),
+            (100_000,),
+            (np.arange(1_000_000),),
+            "coordinate",
+        ),
+        ("correlated_dense", (1000, 1000), (10, 10), (i, i), "coordinate"),
+        ("correlated_sparse", (10000, 10000), (10, 10), (i * 9, i * 9), "coordinate"),
+        (
+            "independent_one_chunk",
+            (1000,) * 3,
+            (1000,) * 3,
+            (i[:, None], i[:, None], i[None, :]),
+            "coordinate",
+        ),
+        (
+            "independent_100_chunks",
+            (1000,) * 3,
+            (100,) * 3,
+            (i[:, None], i[:, None], i[None, :]),
+            "coordinate",
+        ),
+    ]
+    results = {
+        name: compare_case(shape, chunks, selection, mode)
+        for name, shape, chunks, selection, mode in cases
+    }
+    print(json.dumps(results, indent=2))
+
+
+def compare_case(
+    shape: tuple[int, ...], chunks: tuple[int, ...], selection: Any, mode: str
+) -> dict[str, Any]:
+    zg = ChunkGrid.from_sizes(shape, chunks)
+    pg = dimension_grids_from_chunks(chunks, shape)
+    cls = {
+        "basic": zi.BasicIndexer,
+        "orthogonal": zi.OrthogonalIndexer,
+        "coordinate": zi.CoordinateIndexer,
+    }[mode]
+
+    def baseline() -> Any:
+        return cls(selection, shape, zg)
+
+    def new() -> Any:
+        base = IndexTransform.from_shape(shape)
+        transform = (
+            base[selection]
+            if mode == "basic"
+            else (base.oindex[selection] if mode == "orthogonal" else base.vindex[selection])
+        )
+        return plan_chunks(transform, pg).partition()
+
+    old_coords = [tuple(p.chunk_coords) for p in baseline()]
+    partition = new()
+    new_coords = [tuple(p.chunk_coords) for p in partition]
+    assert old_coords == new_coords
+    expected_size = (
+        math.prod(shape)
+        if mode == "basic"
+        else (
+            selection[0].size
+            if mode == "orthogonal"
+            else math.prod(np.broadcast_shapes(*(s.shape for s in selection)))
+        )
+    )
+    assert sum(math.prod(p.cell_transform.domain.shape) for p in partition) == expected_size
+    # Alternate evaluation order across rounds to reduce temporal bias.
+    rounds = []
+
+    def immediate() -> Any:
+        return execute_selection(
+            selection,
+            shape,
+            zg._dimensions,
+            ownership="borrow",
+            mode={"basic": "basic", "orthogonal": "orthogonal", "coordinate": "vectorized"}[mode],
+        )
+
+    assert [tuple(p.chunk_coords) for p in immediate()] == old_coords
+    operations = {
+        "zarr_setup": baseline,
+        "new_setup": new,
+        "zarr_walk": lambda: consume(baseline()),
+        "new_walk": lambda: consume(new()),
+        "immediate_setup": immediate,
+        "immediate_walk": lambda: consume(immediate()),
+        "zarr_retained": lambda: list(baseline()),
+        "borrowed_retained": lambda: list(immediate()),
+        "snapshot_retained": lambda: list(
+            execute_selection(
+                selection,
+                shape,
+                zg._dimensions,
+                mode={"basic": "basic", "orthogonal": "orthogonal", "coordinate": "vectorized"}[
+                    mode
+                ],
+            )
+        ),
+        "shard_retained": lambda: list(immediate().lower("shard")),
+    }
+    for round_id in range(3):
+        names = list(operations)
+        if round_id % 2:
+            names.reverse()
+        rounds.append({key: measure(operations[key]) for key in names})
+    return {
+        "chunks": len(old_coords),
+        "elements": expected_size,
+        "metrics": {
+            key: {
+                "ms": statistics.median(r[key]["ms"] for r in rounds),
+                "peak_mib": max(r[key]["peak_mib"] for r in rounds),
+                "round_ms": [r[key]["ms"] for r in rounds],
+            }
+            for key in operations
+        },
+    }
+
+
+if __name__ == "__main__":
+    main()

packages/zarr-indexing/benchmarks/execution_io.py

@@ -0,0 +1,128 @@
+"""Measure actual in-memory Zarr codec reads/writes, including plan construction.
+
+Storage, source data, grids, and replacement data are prepared before timing.
+Both paths use the same array and codec pipeline. These are local MemoryStore
+measurements, not estimates of cloud or filesystem throughput.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import statistics
+import time
+import tracemalloc
+from typing import TYPE_CHECKING, Any, cast
+
+import numpy as np
+import zarr.api.asynchronous as za
+from zarr.core.buffer.core import default_buffer_prototype
+from zarr.core.indexing import BasicIndexer, CoordinateIndexer
+from zarr.storage import MemoryStore
+
+from zarr_indexing._execution import execute_selection
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from zarr.core.indexing import Indexer
+
+
+async def measure(operation: Callable[[], Awaitable[Any]]) -> dict[str, float]:
+    await operation()
+    samples = []
+    for _ in range(9):
+        start = time.perf_counter()
+        await operation()
+        samples.append((time.perf_counter() - start) * 1000)
+    tracemalloc.start()
+    result = await operation()
+    _, peak = tracemalloc.get_traced_memory()
+    tracemalloc.stop()
+    del result
+    return {"ms": statistics.median(samples), "peak_mib": peak / 2**20}
+
+
+async def run_case(case: str, sharded: bool) -> dict[str, Any]:
+    shape: tuple[int, ...]
+    chunks: tuple[int, ...]
+    selection: Any
+    if case == "basic":
+        shape, chunks = (128, 128), (16, 16)
+        selection, mode = (slice(1, 127), slice(1, 127)), "basic"
+    elif case == "sorted":
+        shape, chunks = (100000,), (10000,)
+        selection, mode = (np.arange(1, 99999),), "vectorized"
+    else:
+        shape, chunks = (64, 64, 64), (16, 16, 16)
+        i = np.arange(64)
+        selection, mode = (i[:, None], i[:, None], i[None, :]), "vectorized"
+    kwargs: dict[str, Any] = {}
+    if sharded:
+        kwargs["shards"] = tuple(c * 2 for c in chunks)
+    array = await za.create_array(
+        store=MemoryStore(), shape=shape, chunks=chunks, dtype="int64", **kwargs
+    )
+    source = np.arange(np.prod(shape), dtype=np.int64).reshape(shape)
+    await array.setitem(Ellipsis, source)
+    expected = source[selection]
+    replacement = expected + 1
+    grids = array._chunk_grid._dimensions
+    prototype = default_buffer_prototype()
+    cls = BasicIndexer if mode == "basic" else CoordinateIndexer
+
+    async def old_read() -> Any:
+        result = await array._get_selection(
+            cls(selection, shape, array._chunk_grid), prototype=prototype
+        )
+        # Zarr's public coordinate API restores sel_shape outside the codec
+        # entry point; include that view operation in the baseline.
+        return np.asarray(result).reshape(expected.shape)
+
+    async def new_read() -> Any:
+        plan = execute_selection(selection, shape, grids, mode=mode, ownership="borrow")
+        return await array._get_selection(
+            cast("Indexer", plan.lower("shard" if sharded else "numpy")), prototype=prototype
+        )
+
+    async def old_write() -> None:
+        await array._set_selection(
+            cls(selection, shape, array._chunk_grid),
+            replacement if mode == "basic" else replacement.reshape(-1),
+            prototype=prototype,
+        )
+
+    async def new_write() -> None:
+        plan = execute_selection(
+            selection, shape, grids, mode=mode, ownership="borrow", access="write"
+        )
+        await array._set_selection(
+            cast("Indexer", plan.lower("shard" if sharded else "numpy")),
+            replacement,
+            prototype=prototype,
+        )
+
+    np.testing.assert_array_equal(await old_read(), expected)
+    np.testing.assert_array_equal(await new_read(), expected)
+    results = {
+        "zarr_read": await measure(old_read),
+        "new_read": await measure(new_read),
+        "zarr_write": await measure(old_write),
+        "new_write": await measure(new_write),
+    }
+    expected_full = source.copy()
+    expected_full[selection] = replacement
+    np.testing.assert_array_equal(await array.getitem(Ellipsis), expected_full)
+    return results
+
+
+async def main() -> None:
+    results = {}
+    for case in ("basic", "sorted", "components"):
+        for sharded in (False, True):
+            results[case + ("_sharded" if sharded else "")] = await run_case(case, sharded)
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

packages/zarr-indexing/docs/api/grid.md

@@ -16,6 +16,12 @@ become `VaryingDimension` instances. `DimensionGridLike` remains the narrow
 protocol used by the chunk planner, while `EdgeDimensionGrid` is kept for
 explicit edge-based and coordinate-origin examples.
 
+`RegularDimensionGridLike` is an optional structural capability for fast paths.
+It adds `size` and `extent` to the grid protocol: chunk `k` must start at
+`k * size` and declare `size` elements, with `extent` clipping only the valid
+data in the last chunk. Zarr's existing uniform dimensions satisfy this
+capability without conversion to this package's concrete classes.
+
 Zarr's array implementation can later import these compact grid types from
 `zarr_indexing`; this package intentionally has no import dependency on Zarr.