fix(indexing): plan every view read with its source grid

Remove Partition.result and partition-local source windows. Preserve the source grid in derived views and supply projections for every LazyArray reader call, including unpartitioned reads.

Assisted-by: Codex:GPT-6

Author: d-v-b

packages/zarr-indexing/changes/4349.feature.md

@@ -1 +1 @@
-Add `Partition.result()` to execute a planned partition independently with the same global transform and chunk projection supplied during parent assembly.
+Make every `LazyArray` read supply a chunk projection, including unpartitioned reads and independently executed partition views. Partition views retain the source grid and full source base shape, so indexing and repartitioning use the same coordinate frame as other views.

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

@@ -21,8 +21,10 @@ coordinates into its raw `Partition.view.array`, including for non-first
 partitions. `Partition.projection.chunk_transform` intentionally stays local to
 the selected chunk. During parent materialization (`view.result(parts=parts)`) the reader receives
 both frames in one `ReadContext`: the public global transform in `context.transform` and the
-same local plan in `context.projection`. Use `part.result()` to execute a
-partition independently with both frames. Direct `part.view.result()` calls
-resolve the general view with no projection.
+local plan in `context.projection`. Every view retains the source grid and plans
+its reads, so `part.view.result()` also supplies both frames. Its projection's
+result placement is relative to that part view, rather than the parent output.
+Further indexing and repartitioning use the same source-global coordinate frame.
+Even `unpartitioned()` reads carry a projection for the single source-wide cell.
 
 ::: zarr_indexing.lazy_array

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

@@ -12,8 +12,9 @@ ownership.
 `Reader.read_into(source, context, out)` receives a `ReadContext` whose
 `transform` maps zero-origin output-buffer coordinates to global coordinates in
 `source`, with `context.transform.domain.shape == out.shape`. Its optional
-`projection` is the existing plan for a partitioned read. Both parent
-assembly and independent `Partition.result()` calls supply that projection. The projection's
+`projection` describes one planned read. `LazyArray.result()` always supplies
+it, including for partition views and unpartitioned reads. Direct callers of
+the reader protocol may omit it when their reader supports that. The projection's
 `chunk_transform` remains chunk-local, its `cell_transform` describes result
 placement, and its `chunk_domain` describes the grid cell. The global read
 transform and the projection's chunk transform deliberately use different

packages/zarr-indexing/docs/guide/index.md

@@ -390,9 +390,10 @@ Within one `Partition`, the frames divide: `Partition.view.transform` is a
 different, global transform — it maps the part view directly into the raw
 wrapped source — while only `Partition.projection.chunk_transform` uses
 zero-origin chunk-local coordinates. Parent assembly passes both frames to
-the reader. Independently scheduled `part.result()` calls supply the same
-context. Calling `part.view.result()` instead resolves the view without the
-partition record, so its context has `projection=None`.
+the reader. Independently scheduled `part.view.result()` calls plan against the
+same source grid and supply both frames too. Their result placement is relative
+to the part view being read. A partition view can be indexed or repartitioned
+like any other view; its base shape remains the full source shape.
 
 | Projection field | What its output coordinates mean |
 | --- | --- |

packages/zarr-indexing/docs/snippets/chunk_projection.py

@@ -51,7 +51,7 @@ def indices_to_chunks(self, indices: NDArray[np.intp]) -> NDArray[np.intp]:
 ADVANCED_EXPECTED = image[[4, 1, 1]][:, 2:6]
 ADVANCED_RESULT = np.empty_like(ADVANCED_EXPECTED)
 for part in advanced.parts():
-    ADVANCED_RESULT[part.out_selection] = part.result()
+    ADVANCED_RESULT[part.out_selection] = part.view.result()
 
 np.testing.assert_array_equal(ADVANCED_RESULT, ADVANCED_EXPECTED)
 # --8<-- [end:advanced-projection]

packages/zarr-indexing/examples/lazy_indexing_dask/lazy_indexing_dask.py

@@ -66,7 +66,7 @@ def test_parts_as_tasks(source: zarr.Array) -> None:
     # places the returned blocks sequentially.
     @dask.delayed
     def read(part: object) -> np.ndarray:
-        return part.result()
+        return part.view.result()
 
     blocks = dask.compute(*[read(part) for part in parts], scheduler="threads")
 

packages/zarr-indexing/examples/lazy_indexing_numpy/lazy_indexing_numpy.py

@@ -105,7 +105,7 @@ def test_parts() -> None:
     # parts and placing them is what `result()` does.
     assembled = np.empty(view.shape, dtype=view.dtype)
     for part in parts:
-        assembled[part.out_selection] = part.result()
+        assembled[part.out_selection] = part.view.result()
     assert np.array_equal(assembled, view.result())
 
     # The partitioning is a read strategy, so a different one gives the same data.

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

@@ -30,10 +30,10 @@
 
 The part view's transform directly addresses its raw wrapped array. The paired
 projection deliberately retains the chunk-local frame. Parent materialization
-passes both in `ReadContext`, as does independent `part.result()` execution.
-Calling `part.view.result()` resolves the general view without its partition
-record and therefore uses `projection=None`. Use `part.result()` when scheduling
-partition reads.
+passes both in `ReadContext`. Every view retains the source's partitioning and
+plans its own reads, including views obtained from partitions. Calling
+`part.view.result()` therefore supplies the reader with the selected chunk's
+projection as well. Output placement is relative to the view being executed.
 
 The partitioning is discovered from the wrapped array at construction — first
 `read_chunk_sizes` (zarr's clipped per-axis sizes, sharding-aware), then
@@ -406,8 +406,8 @@ class Partition:
 
     Yielded by [`LazyArray.parts`][zarr_indexing.lazy_array.LazyArray.parts].
     The parts of a view tile it exactly and disjointly: assembling every
-    `part.result()` at its `out_selection` reproduces the whole view's
-    `result()`, with the same reader context on both execution paths. Parts can be
+    `part.view.result()` at its `out_selection` reproduces the whole view's
+    `result()`. Each view plans its reads using the retained source grid. Parts can be
     resolved concurrently when the source and reader permit it.
     Derived parts retain the same reader object; a shared stateful reader owns
     synchronization for concurrent calls.
@@ -448,12 +448,12 @@ class Partition:
     view
         A `LazyArray` covering exactly the cells of the view that live in this
         box. Its transform directly addresses its raw wrapped `array`; only the
-        projection's `chunk_transform` is chunk-local. Use `Partition.result()` to read
-        with this partition's projection through the selected reader. Named `view` rather than
+        projection's `chunk_transform` is chunk-local. The view retains the source
+        grid and reader and can be read or indexed like any other view. Named `view` rather than
         `array` because `LazyArray.array` is the opposite thing — the raw
         wrapped source — and the two sat next to each other meaning inverses.
     out_selection
-        Where `Partition.result()` belongs in an array of the whole view's shape — a
+        Where `view.result()` belongs in an array of the whole view's shape — a
         NumPy index tuple with one entry per dimension of the view, usable
         directly as `out[part.out_selection] = ...`.
     is_complete
@@ -472,7 +472,7 @@ class Partition:
     >>> view = LazyArray.from_numpy(source).with_parts((2, 2))
     >>> out = np.empty(view.shape, dtype=view.dtype)
     >>> for part in view.parts():
-    ...     out[part.out_selection] = part.result()
+    ...     out[part.out_selection] = part.view.result()
     >>> bool((out == source).all())
     True
     """
@@ -493,34 +493,6 @@ def is_complete(self) -> bool:
         """Whether the projection proves it covers the entire selected cell."""
         return self.projection.coverage == "full"
 
-    def result(self) -> Any:
-        """Materialize this partition with its existing projection.
-
-        Returns
-        -------
-        numpy.ndarray
-            The selected values in fresh system memory, with this part's view
-            shape and dtype. Assign the result at `out_selection` to assemble
-            the parent view. The reader receives the same global transform and
-            chunk projection as in parent materialization.
-
-        Notes
-        -----
-        This uses the part's reader and source without replanning. Backend
-        exceptions propagate unchanged. Concurrent calls require a source and
-        reader that support concurrent access.
-        """
-        # Share the view allocator so independent reads preserve masked-array output.
-        out = self.view._output_buffer(self.view.shape)  # pyright: ignore[reportPrivateUsage]
-        if math.prod(self.view.shape) != 0:
-            _invoke_reader(
-                self.view.reader,
-                self.view.array,
-                ReadContext(self.view.transform, self.projection),
-                out,
-            )
-        return out
-
 
 def _validate_prepared_parts(parts: Sequence[Partition], out_shape: tuple[int, ...]) -> None:
     """Require `parts` to address every output cell exactly once.
@@ -598,7 +570,7 @@ class LazyArray:
            [ 8, 10]])
     """
 
-    __slots__ = ("_array", "_part_owner", "_parts", "_reader", "_transform", "_window")
+    __slots__ = ("_array", "_part_owner", "_parts", "_reader", "_transform")
 
     def __init__(self, array: _WrappedArray) -> None:
         """Wrap `array` without reading it; parameters are documented on the class.
@@ -619,7 +591,6 @@ def __init__(self, array: _WrappedArray) -> None:
             )
         shape = tuple(int(s) for s in array.shape)
         self._array = array
-        self._window: tuple[slice, ...] | None = None
         self._transform = IndexTransform.from_shape(shape)
         self._parts = _discover_parts(array, shape)
         self._reader = basic_reader
@@ -640,7 +611,6 @@ def _derive(
         array: _WrappedArray,
         transform: IndexTransform,
         parts: tuple[DimensionGrid, ...] | None,
-        window: tuple[slice, ...] | None,
         reader: Reader,
     ) -> LazyArray:
         """Build a wrapper sharing `array` but carrying a new transform or partitioning."""
@@ -650,17 +620,14 @@ def _derive(
         # view's first element is at position 0 whatever it was sliced from.
         view._transform = transform.translate_domain_to((0,) * transform.input_rank)
         view._parts = parts
-        view._window = window
         view._reader = reader
         view._part_owner = _PartOwner()
         return view
 
     @property
     def _base_shape(self) -> tuple[int, ...]:
         """The shape of what this wrapper treats as its base array."""
-        if self._window is None:
-            return tuple(int(s) for s in self._array.shape)
-        return tuple(s.stop - s.start for s in self._window)
+        return tuple(int(s) for s in self._array.shape)
 
     # -- array-like surface -------------------------------------------------
 
@@ -674,11 +641,8 @@ def base_shape(self) -> tuple[int, ...]:
         """The shape the partitioning is expressed in — not this view's shape.
 
         `with_parts` and `with_parts_per_axis` describe boxes of the array being
-        read, not of the view reading it, so a narrowed view still partitions
-        the extents named here. For a part's own `array`, this is the part's
-        box, which is why the same call means different sizes there. Without
-        somewhere to read it, the frame in force could only be inferred from an
-        error message.
+        read, not of the view reading it. All derived views, including partition
+        views, retain the full source shape as their partitioning frame.
         """
         return self._base_shape
 
@@ -720,7 +684,6 @@ def with_reader(self, reader: Reader) -> LazyArray:
             self._array,
             self._transform,
             self._parts,
-            self._window,
             reader,
         )
 
@@ -949,9 +912,8 @@ def _part_entries(parts: Sequence[Any], method: str) -> tuple[Any, ...]:
     def unpartitioned(self) -> LazyArray:
         """Return the same view, read in one pass.
 
-        `result()` still allocates its owned output buffer first, then calls the
-        reader once with the whole projected transform. `parts()` still yields a
-        single part covering everything.
+        The source is treated as a single grid cell. Nonempty reads still pass
+        its projection to the reader; empty reads make no reader calls.
 
         Returns
         -------
@@ -961,7 +923,7 @@ def unpartitioned(self) -> LazyArray:
         return self._with_grids(None)
 
     def _with_grids(self, grids: tuple[DimensionGrid, ...] | None) -> LazyArray:
-        return LazyArray._derive(self._array, self._transform, grids, self._window, self._reader)
+        return LazyArray._derive(self._array, self._transform, grids, self._reader)
 
     def parts(self) -> Iterator[Partition]:
         """Iterate the base partitioning, projected through this view.
@@ -972,7 +934,7 @@ def parts(self) -> Iterator[Partition]:
 
         Yields one [`Partition`][zarr_indexing.lazy_array.Partition] per box the
         view actually touches. The parts tile the view exactly and disjointly,
-        and each can be resolved with `Partition.result()`: in another thread,
+        and each view can be resolved with `part.view.result()`: in another thread,
         in another order, or not at all. Those views share this
         view's reader, and `LazyArray` does not serialize calls, so a stateful
         reader must synchronize its own mutable state.
@@ -993,48 +955,16 @@ def parts(self) -> Iterator[Partition]:
         >>> (part.base_coords, part.view.shape, part.is_complete)
         ((0, 0), (2, 1), False)
         """
-        base_shape = self._base_shape
-        grids = self._parts if self._parts is not None else _whole_array_grids(base_shape)
-        rank = len(base_shape)
-
-        if self._window is None:
-            plan_transform = self._transform
-        else:
-            plan_transform = self._transform.translate(tuple(-item.start for item in self._window))
-
-        for projection in plan_chunks(plan_transform, grids):
-            base_coords = projection.chunk_coords
-            local = projection.chunk_transform
-            origin = tuple(grid.chunk_offset(c) for grid, c in zip(grids, base_coords, strict=True))
-            extent = tuple(grid.data_size(c) for grid, c in zip(grids, base_coords, strict=True))
-            if origin == (0,) * rank and extent == base_shape:
-                # The part is the whole base: lowering directly against the
-                # source beats materializing a block that is the source.
-                window = self._window
-            elif self._window is None:
-                window = tuple(slice(o, o + e) for o, e in zip(origin, extent, strict=True))
-            else:
-                window = tuple(
-                    slice(w.start + o, w.start + o + e)
-                    for w, o, e in zip(self._window, origin, extent, strict=True)
-                )
-            # The global box, computed from the origin directly rather than from
-            # `window`: a part covering the whole base carries no window (so
-            # nothing is pre-materialized) but still sits somewhere concrete.
-            if self._window is None:
-                global_origin = origin
-            else:
-                global_origin = tuple(
-                    w.start + o for w, o in zip(self._window, origin, strict=True)
-                )
+        grids = self._parts if self._parts is not None else _whole_array_grids(self._base_shape)
+        for projection in plan_chunks(self._transform, grids):
+            domain = projection.chunk_domain
             yield Partition(
                 projection=projection,
-                box=tuple((o, o + e) for o, e in zip(global_origin, extent, strict=True)),
+                box=tuple(zip(domain.inclusive_min, domain.exclusive_max, strict=True)),
                 view=LazyArray._derive(
                     self._array,
-                    local.translate(global_origin),
-                    None,
-                    window,
+                    projection.chunk_transform.translate(domain.inclusive_min),
+                    self._parts,
                     self._reader,
                 ),
                 out_selection=_partition_out_selection(projection.cell_transform),
@@ -1068,7 +998,7 @@ def _select(self, selection: Any, mode: SelectionMode) -> LazyArray:
             composed = transform[literal]
         else:
             composed = transform.select(literal, mode)
-        return LazyArray._derive(self._array, composed, self._parts, self._window, self._reader)
+        return LazyArray._derive(self._array, composed, self._parts, self._reader)
 
     def __getitem__(self, selection: Any) -> Any:
         """Read a basic selection eagerly, like `numpy.ndarray.__getitem__`.
@@ -1086,6 +1016,8 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any:
         partition is read through the selected reader directly into its
         rectangular destination, or into an owned dense temporary before fancy
         placement. Empty views allocate without reading the source.
+        Every reader call includes a chunk projection, even when the source is
+        treated as a single grid cell by `unpartitioned()`.
 
         Parameters
         ----------
@@ -1128,10 +1060,6 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any:
         if size == 0:
             return out
 
-        if prepared_parts is None and self._parts is None:
-            _invoke_reader(self._reader, self._array, ReadContext(self._transform), out)
-            return out
-
         written = 0
         selected_parts = self.parts() if prepared_parts is None else prepared_parts
         for part in selected_parts:

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

@@ -48,7 +48,10 @@ class ReadContext:
     """Maps zero-origin output-buffer coordinates to global coordinates in the source."""
 
     projection: ChunkProjection | None = None
-    """The partition plan when this read is one part of a partitioned view, else `None`."""
+    """The read plan, always supplied by `LazyArray` execution.
+
+    Direct reader callers may omit it if their reader supports unplanned reads.
+    """
 
 
 class Reader(Protocol):