fix(zarr-indexing): adapt the execution prototype to the trimmed planner

The planner now rejects diagonals with ValueError and has no shared
projection walk, so execute_transform factors the plan up front and lets
that rejection surface at construction instead of on first iteration.
_axis_plan reuses chunk_resolution's data-extent helper rather than
carrying a copy.

Assisted-by: ClaudeCode:claude-fable-5-1
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: d-v-b

packages/zarr-indexing/src/zarr_indexing/_axis_plan.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING, NamedTuple
 
 from zarr_indexing._affine import checked_affine
+from zarr_indexing.chunk_resolution import _data_size  # pyright: ignore[reportPrivateUsage]
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -21,11 +22,6 @@ class AxisRun(NamedTuple):
     position: int
 
 
-def data_size(grid: DimensionGridLike, chunk: int) -> int:
-    method = getattr(grid, "data_size", None)
-    return grid.chunk_size(chunk) if method is None else int(method(chunk))
-
-
 def axis_runs(start: int, stride: int, nitems: int, grid: DimensionGridLike) -> Iterator[AxisRun]:
     """Intersect an affine request with chunks, in request traversal order.
 
@@ -43,7 +39,7 @@ def axis_runs(start: int, stride: int, nitems: int, grid: DimensionGridLike) ->
         chunk = grid.index_to_chunk(coordinate)
         offset = grid.chunk_offset(chunk)
         local = coordinate - offset
-        extent = data_size(grid, chunk)
+        extent = _data_size(grid, chunk)
         if stride == 0:
             count = nitems
         else:

packages/zarr-indexing/src/zarr_indexing/_execution.py

@@ -18,12 +18,7 @@
 from zarr_indexing._axis_plan import axis_runs
 from zarr_indexing._selector import as_scalar_index
 from zarr_indexing.boundary import split_scalar_axes
-from zarr_indexing.chunk_resolution import (
-    ChunkPlan,
-    IndexedSet,
-    _shared_input_axis,  # pyright: ignore[reportPrivateUsage]
-    plan_chunks,
-)
+from zarr_indexing.chunk_resolution import ChunkPlan, IndexedSet, plan_chunks
 from zarr_indexing.errors import BoundsCheckError
 from zarr_indexing.grid import DimensionGridLike, RegularDimensionGridLike
 from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap
@@ -375,20 +370,18 @@ def execute_transform(
                 return _with_policy(sorted_plan, access, "snapshot", conflicts)
     _validate_storage_bounds(transform, grids)
     plan = plan_chunks(transform, grids)
-    # Prepare factored array grouping once. Affine diagonals intentionally use
-    # ChunkPlan's shared projection path rather than per-output-axis tables.
+    # Factor the plan once, up front: a transform the planner cannot factor
+    # (a diagonal) is rejected here rather than on first iteration.
+    partition = plan.partition()
     if (
         any(isinstance(m, ArrayMap) for m in transform.output)
-        and _shared_input_axis(transform) is None
+        and not partition.sets
+        and all(bool((joint.chunk_start >= 0).all()) for joint in partition.joint_sets)
     ):
-        partition = plan.partition()
-        if not partition.sets and all(
-            bool((joint.chunk_start >= 0).all()) for joint in partition.joint_sets
-        ):
-            # Column arithmetic is checked once by JointSet.local; nonnegative
-            # chunk origins make its final local subtraction safe in intp.
-            work = _ComponentWork(plan, tuple(joint.local for joint in partition.joint_sets))
-            return _with_policy(ExecutionPlan(domain.shape, work), access, "snapshot", conflicts)
+        # Column arithmetic is checked once by JointSet.local; nonnegative
+        # chunk origins make its final local subtraction safe in intp.
+        work = _ComponentWork(plan, tuple(joint.local for joint in partition.joint_sets))
+        return _with_policy(ExecutionPlan(domain.shape, work), access, "snapshot", conflicts)
     return _with_policy(ExecutionPlan(domain.shape, plan), access, "snapshot", conflicts)
 
 

packages/zarr-indexing/tests/test_execution.py

@@ -234,17 +234,17 @@ def test_scalar_coordinate_consumer_preserves_value_shape(consumer: Any) -> None
 
 
 @pytest.mark.parametrize("access", ["read", "write"])
-def test_diagonal_with_independent_gather_is_preparable(access: Any) -> None:
+def test_diagonal_is_rejected_before_iteration(access: Any) -> None:
     from zarr_indexing import ArrayMap
 
     transform = IndexTransform(
         IndexDomain.from_shape((3, 2)),
         (DimensionMap(0), DimensionMap(0), ArrayMap(np.array([[1, 0]]))),
     )
-    plan = execute_transform(
-        transform, dimension_grids_from_chunks((2, 2, 2), (3, 3, 2)), access=access
-    )
-    assert len(list(plan)) == 2
+    with pytest.raises(ValueError, match="diagonal"):
+        execute_transform(
+            transform, dimension_grids_from_chunks((2, 2, 2), (3, 3, 2)), access=access
+        )
 
 
 def test_last_write_lowering_removes_duplicate_destinations() -> None: