From e1e62033b5151a46571a7994e91e3d9478c5c980 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 10:41:12 +0200 Subject: [PATCH 01/16] feat(zarr-indexing): lower parts to basic selections, carry projections, add result_into Public API for consumers that plan reads here but perform them through their own I/O layer (the rio-tiler / zarr.AsyncArray pattern): - IndexTransform.as_basic_selection() lowers a box transform to the int/slice tuple that reads exactly its cells at exactly its domain shape, producing collapsed-constant singleton axes as length-1 slices and refusing queries, broadcasts, and axis permutations with ValueError. - Partition.source_selection / Partition.chunk_local_selection expose a part's read as the wrapped array's own basic selection and as the grid-cell-relative equivalent, so an async consumer's loop is out[part.out_selection] = await src.getitem(part.source_selection) and a decoded-chunk cache keyed on base_coords needs no coordinate arithmetic. - part.view.result() now passes the paired ChunkProjection to the reader instead of projection=None, matching the parent view's partitioned read; with_reader and repartitioning keep the pairing, a new selection drops it. - LazyArray.result_into(out, *, parts=None) is the non-allocating form of result(): the caller's validated buffer is filled in place and returned. ChainedIndexingStateMachine gained an invariant running both documented assembly loops literally under plain NumPy semantics, reader-free. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+part-lowering.feature.md | 19 ++ .../changes/+part-view-projection.bugfix.md | 8 + .../changes/+result-into.feature.md | 6 + .../src/zarr_indexing/lazy_array.py | 207 +++++++++++++++-- .../src/zarr_indexing/testing/stateful.py | 57 +++++ .../src/zarr_indexing/transform.py | 148 ++++++++++++ .../zarr-indexing/tests/test_lazy_array.py | 210 ++++++++++++++++++ .../zarr-indexing/tests/test_transform.py | 105 +++++++++ 8 files changed, 747 insertions(+), 13 deletions(-) create mode 100644 packages/zarr-indexing/changes/+part-lowering.feature.md create mode 100644 packages/zarr-indexing/changes/+part-view-projection.bugfix.md create mode 100644 packages/zarr-indexing/changes/+result-into.feature.md diff --git a/packages/zarr-indexing/changes/+part-lowering.feature.md b/packages/zarr-indexing/changes/+part-lowering.feature.md new file mode 100644 index 0000000000..927d687789 --- /dev/null +++ b/packages/zarr-indexing/changes/+part-lowering.feature.md @@ -0,0 +1,19 @@ +Box selections and partitions now lower to backend-native basic selections, +for consumers that plan reads here but fetch through their own I/O layer (an +async store, an HTTP range endpoint): + +- `IndexTransform.as_basic_selection()` converts a box transform to a tuple of + integers and slices such that `source[selection]` reads exactly the + transform's cells, at exactly its domain shape. A collapsed + single-coordinate gather keeps its singleton axis through a length-1 slice; + queries, broadcasts, and transposed or repeated axes raise `ValueError` + instead of guessing a slab. +- `Partition.source_selection` is that lowering of a part's global read, so an + async consumer's whole loop is + `out[part.out_selection] = await source.getitem(part.source_selection)`. +- `Partition.chunk_local_selection` is the same read relative to + `projection.chunk_domain`'s origin, for decoded-chunk caches keyed on + `base_coords`. + +`ChainedIndexingStateMachine` gained an invariant that runs both documented +assembly loops literally under plain NumPy semantics, with no reader involved. diff --git a/packages/zarr-indexing/changes/+part-view-projection.bugfix.md b/packages/zarr-indexing/changes/+part-view-projection.bugfix.md new file mode 100644 index 0000000000..c5e742bf45 --- /dev/null +++ b/packages/zarr-indexing/changes/+part-view-projection.bugfix.md @@ -0,0 +1,8 @@ +Resolving a partition's view on its own (`part.view.result()`) now hands the +reader the same `ReadContext` the parent's partitioned `result()` passes: the +paired `ChunkProjection` rides along instead of arriving as `projection=None`. +A custom reader keyed on `projection.chunk_coords` — a decoded-chunk cache — +now behaves identically on both paths, which the dask example's +task-per-partition pattern relies on. `with_reader` and repartitioning keep +the pairing; a further `.lazy` selection describes a different read and drops +it. diff --git a/packages/zarr-indexing/changes/+result-into.feature.md b/packages/zarr-indexing/changes/+result-into.feature.md new file mode 100644 index 0000000000..576b94dd1f --- /dev/null +++ b/packages/zarr-indexing/changes/+result-into.feature.md @@ -0,0 +1,6 @@ +Added `LazyArray.result_into(out, *, parts=None)`: the non-allocating form of +`result()`. The caller's writable buffer is validated against the view's shape +and dtype, filled in place — every cell written exactly once — and returned. A +view into a larger array qualifies, so a part's block can land directly in its +final slot, and a `numpy.ma` masked buffer keeps a masked source's mask. +`result()` itself is unchanged and always allocates. diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 950a97b25e..c99b8ade45 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -133,8 +133,10 @@ Ownership --------- `result()` always allocates fresh system memory before reading through the -selected reader. A `numpy.ma` source keeps its mask by receiving a masked -output buffer; other source-specific array types do not survive materializing. +selected reader. `result_into(out)` is the non-allocating form: the caller's +buffer is validated against the view's shape and dtype, filled in place, and +returned. A `numpy.ma` source keeps its mask by receiving a masked output +buffer; other source-specific array types do not survive materializing. """ from __future__ import annotations @@ -446,9 +448,14 @@ class Partition: 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. Resolving the view reads - the box once through its selected reader. 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. + the box once through its selected reader, handing it a `ReadContext` + that carries this partition's `projection` — the same context the + parent view's partitioned `result()` passes, so a reader keyed on + `chunk_coords` behaves identically on either path. A further `.lazy` + selection describes a different read and drops that pairing. 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 `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 @@ -489,6 +496,68 @@ def is_complete(self) -> bool: """Whether the projection proves it covers the entire selected cell.""" return self.projection.coverage == "full" + @property + def source_selection(self) -> tuple[int | slice, ...]: + """The basic selection on the raw wrapped array that reads this part. + + Lowered from `view.transform` by + [`IndexTransform.as_basic_selection`][zarr_indexing.transform.IndexTransform.as_basic_selection], + so the coordinates are global to `view.array` and + `view.array[part.source_selection]` reads the same block + `view.result()` reads, in the same order. This is the request to hand + to a consumer-owned I/O layer whose vocabulary is a basic selection + rather than a transform — for one driving reads through an async + store, assembly is one line per part: + + ```python + out[part.out_selection] = await source.getitem(part.source_selection) + ``` + + Defined for box-shaped parts (`view.is_box`); a query part has no + slab to request and raises `ValueError`, as does the one degenerate + box with no basic-selection spelling (an axis restored by repetition, + reached by gathering a collapsed constant with duplicates). A consumer + mixing selection kinds catches `ValueError` — or checks `view.is_box` + for the common case — and falls back to `view.result()`. A reversing + view lowers to a negative-step slice. + + Examples + -------- + >>> import numpy as np + >>> view = LazyArray.from_numpy(np.arange(12).reshape(3, 4)).with_parts((2, 2)) + >>> [part.source_selection for part in view.lazy[1:, ::2].parts()] + [(slice(1, 2, 1), slice(0, 1, 2)), (slice(1, 2, 1), slice(2, 3, 2)), \ +(slice(2, 3, 1), slice(0, 1, 2)), (slice(2, 3, 1), slice(2, 3, 2))] + """ + return self.view.transform.as_basic_selection() + + @property + def chunk_local_selection(self) -> tuple[int | slice, ...]: + """The same read as `source_selection`, relative to the part's grid cell. + + Lowered from `projection.chunk_transform`, so coordinate 0 per axis is + `projection.chunk_domain`'s origin. For a consumer that caches decoded + cells keyed on `base_coords`, this is the selection to apply to a + cached cell: `cell[part.chunk_local_selection]` yields the same values + as `view.array[part.source_selection]`, and both land at + `out_selection` in the request buffer. + + Defined exactly when `source_selection` is: the two lower the same + maps, so a part that refuses one spelling refuses both. + + Examples + -------- + >>> import numpy as np + >>> view = LazyArray.from_numpy(np.arange(12).reshape(3, 4)).with_parts((2, 2)) + >>> part = next(view.lazy[1:, :2].parts()) + >>> (part.base_coords, part.source_selection, part.chunk_local_selection) + ((0, 0), (slice(1, 2, 1), slice(0, 2, 1)), (slice(1, 2, 1), slice(0, 2, 1))) + >>> part = next(view.lazy[2:, :2].parts()) + >>> (part.base_coords, part.source_selection, part.chunk_local_selection) + ((1, 0), (slice(2, 3, 1), slice(0, 2, 1)), (slice(0, 1, 1), slice(0, 2, 1))) + """ + return self.projection.chunk_transform.as_basic_selection() + def _validate_prepared_parts(parts: Sequence[Partition], out_shape: tuple[int, ...]) -> None: """Require `parts` to address every output cell exactly once. @@ -629,7 +698,15 @@ class LazyArray: [ 8, 10]]) """ - __slots__ = ("_array", "_part_owner", "_parts", "_reader", "_transform", "_window") + __slots__ = ( + "_array", + "_part_owner", + "_parts", + "_projection", + "_reader", + "_transform", + "_window", + ) def __init__(self, array: _WrappedArray) -> None: """Wrap `array` without reading it; parameters are documented on the class. @@ -653,6 +730,7 @@ def __init__(self, array: _WrappedArray) -> None: self._transform = IndexTransform.from_shape(shape) self._parts = _discover_parts(array, shape) self._reader = basic_reader + self._projection: ChunkProjection | None = None self._part_owner = _PartOwner() @classmethod @@ -672,8 +750,16 @@ def _derive( parts: tuple[DimensionGrid, ...] | None, window: tuple[slice, ...] | None, reader: Reader, + projection: ChunkProjection | None = None, ) -> LazyArray: - """Build a wrapper sharing `array` but carrying a new transform or partitioning.""" + """Build a wrapper sharing `array` but carrying a new transform or partitioning. + + `projection` is the paired partition plan when the wrapper is a + `Partition.view`; it rides along so an unpartitioned `result()` hands + the reader the same `ReadContext` the parent's partitioned read would. + A new selection describes a different read, so `_select` leaves it at + the default `None`. + """ view = cls.__new__(cls) view._array = array # Views re-zero their coordinate system: the positional dialect means a @@ -682,6 +768,7 @@ def _derive( view._parts = parts view._window = window view._reader = reader + view._projection = projection view._part_owner = _PartOwner() return view @@ -752,6 +839,7 @@ def with_reader(self, reader: Reader) -> LazyArray: self._parts, self._window, reader, + self._projection, ) # -- shape of the selection --------------------------------------------- @@ -1009,7 +1097,9 @@ 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._window, self._reader, self._projection + ) def parts(self) -> Iterator[Partition]: """Iterate the base partitioning, projected through this view. @@ -1084,6 +1174,7 @@ def parts(self) -> Iterator[Partition]: None, window, self._reader, + projection, ), out_selection=_partition_out_selection(projection.cell_transform), _owner=self._part_owner, @@ -1133,7 +1224,9 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any: Every result starts as a fresh system-memory buffer. Each touched 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. + placement. Empty views allocate without reading the source. To fill a + buffer the caller already owns instead, use + [`result_into`][zarr_indexing.lazy_array.LazyArray.result_into]. Parameters ---------- @@ -1159,6 +1252,86 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any: If this library's own partition walk fails to cover the view — a bug in zarr-indexing, never a consequence of the caller's input. """ + prepared_parts = self._prepared_parts(parts) + return self._read_into_buffer(self._output_buffer(self.shape), prepared_parts) + + def result_into( + self, out: np.ndarray[Any, Any], *, parts: Sequence[Partition] | None = None + ) -> Any: + """Materialize this view into a buffer the caller owns. + + The non-allocating form of + [`result`][zarr_indexing.lazy_array.LazyArray.result]: `out` is + validated, filled in place, and returned, and no output buffer is + allocated here — a consumer assembling many views into one array, or + holding a pool of reusable tile buffers, decides where results live. + (A part whose placement is fancy still gathers through an owned dense + temporary before scattering, exactly as `result()` does.) + + Every cell of `out` is overwritten exactly once. If a reader raises + midway, `out` is left partially written. + + Parameters + ---------- + out + A writable `numpy.ndarray` of exactly `self.shape` and this view's + dtype. A view into a larger array qualifies, so a part's result + can land directly in its slot: + `part.view.result_into(final[part.out_selection])` for a + rectangular part. A `numpy.ma` masked buffer (what `result()` + would allocate for a masked source) also qualifies. + parts + A reusable sequence previously returned by this exact view's + `parts()` method, exactly as for `result`. + + Returns + ------- + numpy.ndarray + The `out` object that was passed in, filled. + + Raises + ------ + TypeError + If `out` is not a `numpy.ndarray`. + ValueError + If `out` has the wrong shape or dtype or is read-only, or if + supplied parts were prepared by another view or do not tile this + view exactly. + AssertionError + If this library's own partition walk fails to cover the view — a + bug in zarr-indexing, never a consequence of the caller's input. + + Examples + -------- + >>> import numpy as np + >>> source = np.arange(12).reshape(3, 4) + >>> view = LazyArray.from_numpy(source).lazy[1:, ::2] + >>> out = np.empty(view.shape, dtype=view.dtype) + >>> returned = view.result_into(out) + >>> returned is out + True + >>> out + array([[ 4, 6], + [ 8, 10]]) + """ + out_shape = self.shape + if not isinstance(cast(object, out), np.ndarray): + raise TypeError(f"out must be a numpy.ndarray, got {type(out).__name__}") + if tuple(out.shape) != out_shape: + raise ValueError( + f"out has shape {tuple(out.shape)}, but this view has shape {out_shape}" + ) + if out.dtype != np.dtype(self.dtype): + raise ValueError( + f"out has dtype {out.dtype}, but this view has dtype {np.dtype(self.dtype)}" + ) + if not out.flags.writeable: + raise ValueError("out is read-only; result_into writes every cell of it") + prepared_parts = self._prepared_parts(parts) + return self._read_into_buffer(out, prepared_parts) + + def _prepared_parts(self, parts: Sequence[Partition] | None) -> tuple[Partition, ...] | None: + """Validate a caller-supplied partition plan against this view.""" prepared_parts = None if parts is None else tuple(parts) if prepared_parts is not None and any( # Module-private provenance deliberately crosses the two public @@ -1168,16 +1341,24 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any: ): raise ValueError("prepared parts do not belong to this view") - out_shape = self.shape if prepared_parts is not None: - _validate_prepared_parts(prepared_parts, out_shape) - out = self._output_buffer(out_shape) + _validate_prepared_parts(prepared_parts, self.shape) + return prepared_parts + + def _read_into_buffer(self, out: Any, prepared_parts: tuple[Partition, ...] | None) -> Any: + """Read this view into `out`, part by part, and return `out`.""" + out_shape = self.shape size = math.prod(out_shape) 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) + # A partition view carries its paired projection, so resolving it + # alone hands the reader the same ReadContext the parent's + # partitioned read would. + _invoke_reader( + self._reader, self._array, ReadContext(self._transform, self._projection), out + ) return out written = 0 diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 4458d7922a..32741b15b6 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -57,6 +57,7 @@ def make_source(self, data): from hypothesis.stateful import RuleBasedStateMachine, initialize, invariant, precondition, rule from zarr_indexing.lazy_array import LazyArray +from zarr_indexing.output_map import ConstantMap from zarr_indexing.reader import Reader, basic_reader from zarr_indexing.testing.strategies import ( basic_selections, @@ -367,6 +368,62 @@ def parts_tile_the_view(self) -> None: hits, np.ones(self.view.shape, dtype=np.int64), err_msg=str(self.chain) ) + @invariant() + def box_parts_lower_to_basic_selections(self) -> None: + """The consumer-owned-I/O assembly, run literally. + + [`Partition`][zarr_indexing.lazy_array.Partition] documents + `out[part.out_selection] = source[part.source_selection]` as the + assembly for a consumer fetching box parts through its own I/O layer, + and `cell[part.chunk_local_selection]` as the equivalent read from a + cached grid cell. Both selections are applied here under plain NumPy + semantics to the values the source holds and checked against the + model; a query part must refuse to lower instead of guessing a slab. + + No reader runs in this invariant: it proves the lowered selections + alone carry the read, which is exactly what a consumer that plans here + but fetches elsewhere relies on. + """ + data = np.asarray(type(self).data) + for part in self.view.parts(): + try: + source_selection = part.source_selection + except ValueError: + # Refusal is the documented answer for a query part, and for + # the one box shape with no basic spelling: an axis restored + # by repetition, reached by gathering a collapsed constant + # with duplicates. Either way the properties must agree. + assert not part.view.is_box or any( + isinstance(m, ConstantMap) for m in part.view.transform.output + ), f"a plain box part refused to lower: {self.chain}" + for attribute in ("source_selection", "chunk_local_selection"): + lowered = True + try: + getattr(part, attribute) + except ValueError: + lowered = False + assert not lowered, ( + f"{attribute} lowered a part whose paired selection " + f"refused to: {self.chain}" + ) + continue + expected = self.model[part.out_selection] + slab = data[source_selection] + np.testing.assert_array_equal(slab, expected, err_msg=str(self.chain)) + cell = data[ + tuple( + slice(lo, hi) + for lo, hi in zip( + part.projection.chunk_domain.inclusive_min, + part.projection.chunk_domain.exclusive_max, + strict=True, + ) + ) + ] + np.testing.assert_array_equal( + cell[part.chunk_local_selection], expected, err_msg=str(self.chain) + ) + def _reader_set(view: LazyArray, declared: Sequence[Reader] | None) -> tuple[Reader, ...]: """The readers `choose_reader` draws from, `basic_reader` always among them.""" diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index a8a5963a26..ff3dc263ac 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -448,6 +448,154 @@ def inverted(self) -> IndexTransform: output=tuple(inverse_output[dimension] for dimension in range(self.input_rank)), ) + def as_basic_selection(self) -> tuple[int | slice, ...]: + """Lower a box transform to the NumPy basic selection it describes. + + Returns one selector per output dimension such that + `source[transform.as_basic_selection()]` under NumPy indexing + semantics reads exactly the cells this transform addresses, in domain + order: the result has shape `domain.shape` exactly, and its element at + position `p` holds the source value at + `apply(domain.inclusive_min + p)`. + + A `DimensionMap` lowers to a `slice`. A `ConstantMap` usually lowers + to an `int`, which drops its axis exactly as the constant map carries + no input dimension — but a singleton domain axis that no output map + references (a single-coordinate fancy index collapses to a constant at + construction, leaving its axis behind: `oindex[:, [2]]`) is instead + produced by lowering a constant to a length-1 slice, so the result + keeps the domain's shape. + + This is the boundary for consumers that plan reads here but perform + them through their own I/O layer — an async store, an HTTP range + request — whose request vocabulary is a basic selection rather than a + transform. A reversing map lowers to a negative-step slice, so a + backend restricted to ascending unit-step reads should inspect the + steps (or keep such selections out of its plans; compare + `UnitStepReader`). + + Returns + ------- + tuple of int or slice + One selector per output dimension. + + Raises + ------ + ValueError + If the transform cannot be expressed as a basic selection: an + output map is an `ArrayMap` (a query is a lookup table, not a + slab), a `DimensionMap` has stride 0 (a broadcast repeats one + source cell, which no slice spells), the selectors cannot produce + the domain's axes in increasing order and exactly once (basic + indexing never transposes, repeats, or fabricates axes), an + unreferenced domain axis has extent other than 1, or a selected + coordinate is negative (NumPy would count it from the end of the + source). + + Examples + -------- + >>> IndexTransform.from_shape((10, 20))[2:8, ::2].as_basic_selection() + (slice(2, 8, 1), slice(0, 19, 2)) + >>> IndexTransform.from_shape((10, 20))[3, 4:6].as_basic_selection() + (3, slice(4, 6, 1)) + + The collapsed single-coordinate gather keeps its axis: + + >>> IndexTransform.from_shape((10, 20)).oindex[slice(None), [2]].as_basic_selection() + (slice(0, 10, 1), slice(2, 3, 1)) + + A query has no basic-selection spelling: + + >>> IndexTransform.from_shape((10,)).oindex[[3, 1, 1]].as_basic_selection() + Traceback (most recent call last): + ... + ValueError: cannot lower to a basic selection: output[0] is an ArrayMap; \ +a query selection is a lookup table, not a slab + """ + for output_dimension, output_map in enumerate(self.output): + # Named before the axis bookkeeping below, which would otherwise + # blame a query's gathered axis for being unreferenced. + if isinstance(output_map, ArrayMap): + raise ValueError( # noqa: TRY004 - valid map, no basic spelling + f"cannot lower to a basic selection: output[{output_dimension}] " + "is an ArrayMap; a query selection is a lookup table, not a slab" + ) + referenced = {m.input_dimension for m in self.output if isinstance(m, DimensionMap)} + for axis, extent in enumerate(self.domain.shape): + if axis not in referenced and extent != 1: + raise ValueError( + f"cannot lower to a basic selection: no output map references " + f"input dimension {axis}, whose extent is {extent}; only a " + "singleton axis can be produced by an integer-indexed output " + "dimension" + ) + selection: list[int | slice] = [] + # The next domain axis a selector has to produce. Slices produce axes + # in selector order, so walking the outputs left to right must meet the + # domain axes in increasing order. + next_axis = 0 + for output_dimension, output_map in enumerate(self.output): + if isinstance(output_map, ConstantMap): + if output_map.offset < 0: + raise ValueError( + f"cannot lower to a basic selection: output[{output_dimension}] " + f"addresses coordinate {output_map.offset}, which NumPy would " + "count from the end of the source" + ) + if next_axis < self.input_rank and next_axis not in referenced: + # This constant stands in for the singleton axis: a + # length-1 slice keeps the axis where an int would drop it. + selection.append(slice(output_map.offset, output_map.offset + 1, 1)) + next_axis += 1 + else: + selection.append(output_map.offset) + continue + assert isinstance(output_map, DimensionMap) # ArrayMap raised above + if output_map.stride == 0: + raise ValueError( + f"cannot lower to a basic selection: output[{output_dimension}] " + "has stride 0, which repeats one source cell along an axis; " + "no slice spells a broadcast" + ) + d = output_map.input_dimension + if d != next_axis: + raise ValueError( + "cannot lower to a basic selection: the selectors must produce " + f"the domain's axes in increasing order and exactly once, but " + f"output[{output_dimension}] produces input dimension {d} where " + f"input dimension {next_axis} is due; basic indexing never " + "transposes, repeats, or fabricates axes" + ) + next_axis = d + 1 + lo = self.domain.inclusive_min[d] + hi = self.domain.exclusive_max[d] + if hi <= lo: + selection.append(slice(0, 0, 1)) + continue + first = checked_affine(output_map.offset, output_map.stride, lo) + last = checked_affine(output_map.offset, output_map.stride, hi - 1) + if min(first, last) < 0: + raise ValueError( + f"cannot lower to a basic selection: output[{output_dimension}] " + f"addresses coordinate {min(first, last)}, which NumPy would " + "count from the end of the source" + ) + if output_map.stride > 0: + selection.append(slice(first, last + 1, output_map.stride)) + else: + # Descending: the stop sits one step past the final coordinate, + # and a stop below 0 has no literal spelling — None walks to + # the front edge instead of wrapping. + stop = last + output_map.stride + selection.append(slice(first, stop if stop >= 0 else None, output_map.stride)) + if next_axis != self.input_rank: + raise ValueError( + "cannot lower to a basic selection: no selector is left to produce " + f"input dimension {next_axis}, so a basic read cannot yield that " + "axis of the domain" + ) + return tuple(selection) + @property def selection_repr(self) -> str: """Compact domain string, e.g. `'{ [2, 8), [0, 10) }'`. diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index ab94be1ae5..a3eb1cd90e 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2616,3 +2616,213 @@ def test_fancy_composition_over_an_empty_axis() -> None: scalar = composed.lazy.vindex[..., np.array(1)] assert scalar.shape == (2, 0) assert np.asarray(scalar.result()).shape == (2, 0) + + +# --------------------------------------------------------------------------- +# Backend-native selection lowering +# --------------------------------------------------------------------------- + + +def test_box_part_selections_carry_the_read() -> None: + """The lowered selections agree with the documented assemblies. + + `out[part.out_selection] = data[part.source_selection]` reproduces the + view — the loop a consumer runs when it plans here but fetches through its + own I/O layer — and `chunk_local_selection` reads the same values from the + part's grid cell, the loop a decoded-chunk cache runs. + """ + data = reference() + view = LazyArray.from_numpy(data).with_parts((3, 2, 3)) + cases = ( + view, + view.lazy[1:6, ::2, 1:], + view.lazy[::-1, 2, 1::2], + # The single-coordinate gather collapses to a constant and keeps its + # axis through a length-1 slice. + view.lazy.oindex[:, [2], :], + view.lazy[5, 1, 2], + view.unpartitioned().lazy[2:6, :, ::2], + ) + for case in cases: + expected = np.asarray(case.result()) + out = np.full(case.shape, -1, dtype=case.dtype) + for part in case.parts(): + slab = data[part.source_selection] + cell = data[ + tuple( + slice(lo, hi) + for lo, hi in zip( + part.projection.chunk_domain.inclusive_min, + part.projection.chunk_domain.exclusive_max, + strict=True, + ) + ) + ] + np.testing.assert_array_equal(cell[part.chunk_local_selection], slab) + out[part.out_selection] = slab + np.testing.assert_array_equal(out, expected) + + +def test_part_of_part_selections_stay_global_and_cell_local() -> None: + """A nested part addresses the raw source; its cell stays window-relative.""" + data = reference() + view = LazyArray.from_numpy(data).with_parts((3, 2, 3)).lazy[1:6, 1:4, :] + outer = next(view.parts()) + inner_view = outer.view.with_parts((2, 1, 2)) + expected = np.asarray(outer.view.result()) + boxed = data[tuple(slice(lo, hi) for lo, hi in outer.box)] + out = np.full(inner_view.shape, -1, dtype=inner_view.dtype) + for part in inner_view.parts(): + slab = data[part.source_selection] + cell = boxed[ + tuple( + slice(lo, hi) + for lo, hi in zip( + part.projection.chunk_domain.inclusive_min, + part.projection.chunk_domain.exclusive_max, + strict=True, + ) + ) + ] + np.testing.assert_array_equal(cell[part.chunk_local_selection], slab) + out[part.out_selection] = slab + np.testing.assert_array_equal(out, expected) + + +def test_query_part_selections_refuse_to_lower() -> None: + """A query part has no slab; both spellings must say so, not guess.""" + view = LazyArray.from_numpy(reference()).with_parts((3, 2, 3)).lazy.oindex[[6, 1, 1], :, :] + part = next(view.parts()) + assert not part.view.is_box + with pytest.raises(ValueError, match="ArrayMap"): + _ = part.source_selection + with pytest.raises(ValueError, match="ArrayMap"): + _ = part.chunk_local_selection + + +# --------------------------------------------------------------------------- +# Part views carry their projection +# --------------------------------------------------------------------------- + + +def test_a_part_view_resolved_alone_carries_its_projection() -> None: + """`part.view.result()` hands the reader the same context the parent does. + + A custom reader keyed on `projection.chunk_coords` (a decoded-chunk cache) + must behave identically whether the consumer calls `view.result()` or + resolves each part's view itself. + """ + reader = RecordingReader() + view = LazyArray(np.arange(8)).with_reader(reader).with_parts((4,)) + parts = list(view.parts()) + for part in parts: + np.testing.assert_array_equal( + np.asarray(part.view.result()), np.arange(8)[part.source_selection] + ) + assert [context.projection for context in reader.contexts] == [ + part.projection for part in parts + ] + # The context transform stays source-global even though the projection's + # chunk_transform is chunk-local. + assert reader.contexts[1].transform.apply((0,)) == (4,) + + +def test_reader_and_partitioning_swaps_keep_a_part_views_projection() -> None: + """`with_reader` and `unpartitioned` return the same view, pairing intact.""" + reader = RecordingReader() + view = LazyArray(np.arange(8)).with_parts((4,)) + part = next(view.parts()) + part.view.with_reader(reader).result() + part.view.with_reader(reader).unpartitioned().result() + assert [context.projection for context in reader.contexts] == [part.projection] * 2 + + +def test_a_new_selection_on_a_part_view_drops_the_projection() -> None: + """A further selection describes a different read, so the pairing ends.""" + reader = RecordingReader() + view = LazyArray(np.arange(8)).with_reader(reader).with_parts((4,)) + part = next(view.parts()) + part.view.lazy[1:3].result() + assert reader.contexts[-1].projection is None + + +# --------------------------------------------------------------------------- +# result_into +# --------------------------------------------------------------------------- + + +def test_result_into_fills_the_callers_buffer() -> None: + """`result_into` fills and returns `out` for the shapes `result` covers.""" + data = reference() + for view in ( + LazyArray.from_numpy(data).with_parts((3, 2, 3)).lazy[1:6, ::2, :], + # Fancy placement scatters through owned temporaries into `out`. + LazyArray.from_numpy(data).lazy.oindex[[6, 1, 1], :, :], + LazyArray.from_numpy(data).lazy[5, 1, 2], + # An empty view returns `out` untouched without reading. + LazyArray.from_numpy(data).lazy[1:1], + ): + expected = np.asarray(view.result()) + out = np.full(view.shape, -1, dtype=view.dtype) + returned = view.result_into(out) + assert returned is out + np.testing.assert_array_equal(out, expected) + + # A view into a larger array is a valid destination, so a part's result + # can land directly in its final slot. + view = LazyArray.from_numpy(data).with_parts((3, 2, 3)).lazy[1:6, 1:4, :] + final = np.full((10, *view.shape[1:]), -1, dtype=view.dtype) + assert view.result_into(final[2:7]) is not final + np.testing.assert_array_equal(final[2:7], np.asarray(view.result())) + assert np.all(final[:2] == -1) + assert np.all(final[7:] == -1) + + # Prepared parts are honored exactly as in `result`. + view = LazyArray.from_numpy(data).with_parts((3, 2, 3)) + parts = tuple(view.parts()) + out = np.empty(view.shape, dtype=view.dtype) + assert view.result_into(out, parts=parts) is out + np.testing.assert_array_equal(out, data) + + # A masked buffer — what result() would allocate for a masked source — + # keeps the source's mask. + masked_source = np.ma.masked_array(data, mask=data % 5 == 0) + masked_view = LazyArray.from_numpy(masked_source) + out_masked = np.ma.masked_all(masked_view.shape, dtype=masked_view.dtype) + masked_view.result_into(out_masked) + np.testing.assert_array_equal(np.ma.getmaskarray(out_masked), masked_source.mask) + np.testing.assert_array_equal(out_masked.compressed(), masked_source.compressed()) + + +def test_result_into_rejects_a_non_array() -> None: + view = LazyArray.from_numpy(reference()) + with pytest.raises(TypeError, match="must be a numpy.ndarray"): + view.result_into(cast("Any", [[0]])) + + +def test_result_into_rejects_the_wrong_shape() -> None: + view = LazyArray.from_numpy(reference()) + with pytest.raises(ValueError, match="has shape"): + view.result_into(np.empty((1, 2, 3), dtype=view.dtype)) + + +def test_result_into_rejects_the_wrong_dtype() -> None: + view = LazyArray.from_numpy(reference()) + with pytest.raises(ValueError, match="has dtype"): + view.result_into(np.empty(view.shape, dtype=np.float32)) + + +def test_result_into_rejects_a_read_only_buffer() -> None: + view = LazyArray.from_numpy(reference()) + out = np.empty(view.shape, dtype=view.dtype) + out.flags.writeable = False + with pytest.raises(ValueError, match="read-only"): + view.result_into(out) + + +def test_result_into_rejects_foreign_parts() -> None: + view = LazyArray.from_numpy(reference()).with_parts((3, 2, 3)) + other = LazyArray.from_numpy(reference()).with_parts((3, 2, 3)) + out = np.empty(view.shape, dtype=view.dtype) + with pytest.raises(ValueError, match="do not belong"): + view.result_into(out, parts=tuple(other.parts())) diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index a13eaf6e28..ce90428cf9 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -395,6 +395,111 @@ def test_inverted_rejects_input_labels_that_cannot_be_preserved(self) -> None: transform.inverted() +def _pointwise_read(transform: IndexTransform, source: np.ndarray) -> np.ndarray: + """Read through `transform` one point at a time — the oracle.""" + domain = transform.domain + if transform.input_rank == 0: + return np.asarray(source[transform.apply(())]) + axes = [ + np.arange(lo, hi) for lo, hi in zip(domain.inclusive_min, domain.exclusive_max, strict=True) + ] + if any(axis.size == 0 for axis in axes): + return np.empty(domain.shape, dtype=source.dtype) + mesh = np.meshgrid(*axes, indexing="ij") + points = np.stack([m.ravel() for m in mesh], axis=1) + coordinates = transform.apply_many(points) + return source[tuple(coordinates.T)].reshape(domain.shape) + + +class TestAsBasicSelection: + @pytest.mark.parametrize( + "transform", + [ + pytest.param(IndexTransform.from_shape((10,)), id="identity"), + pytest.param(IndexTransform.from_shape((10,))[2:8], id="slice"), + pytest.param(IndexTransform.from_shape((10,))[1:9:3], id="strided"), + pytest.param(IndexTransform.from_shape((6,))[::-1], id="reversed"), + pytest.param(IndexTransform.from_shape((10,))[::-3], id="reversed-strided"), + pytest.param(IndexTransform.from_shape((10,))[4:4], id="empty"), + pytest.param(IndexTransform.from_shape((5, 7))[3, 1:6:2], id="int-and-slice"), + pytest.param(IndexTransform.from_shape((5, 7))[2:5][3:5], id="composed-literal"), + pytest.param(IndexTransform.from_shape((5, 7))[3, 2], id="rank-zero"), + pytest.param( + # The single-coordinate gather collapses to a constant at + # construction, leaving a singleton domain axis no map + # references; the constant lowers to a length-1 slice. + IndexTransform.from_shape((5, 7)).oindex[slice(None), [2]], + id="collapsed-single-gather", + ), + pytest.param( + IndexTransform.from_shape((5, 7)).oindex[[3], [2]], + id="two-collapsed-gathers", + ), + ], + ) + def test_selection_reproduces_the_transform(self, transform: IndexTransform) -> None: + """`source[selection]` has the domain's shape and its exact values.""" + source = np.arange(35).reshape(5, 7) if transform.output_rank == 2 else np.arange(10) + selection = transform.as_basic_selection() + assert len(selection) == transform.output_rank + read = np.asarray(source[selection]) + assert read.shape == transform.domain.shape + np.testing.assert_array_equal(read, _pointwise_read(transform, source)) + + def test_rejects_an_array_map(self) -> None: + transform = IndexTransform.from_shape((10,)).oindex[[3, 1, 1]] + with pytest.raises(ValueError, match="is an ArrayMap"): + transform.as_basic_selection() + + def test_rejects_a_zero_stride(self) -> None: + transform = IndexTransform( + IndexDomain.from_shape((3,)), (DimensionMap(0, offset=2, stride=0),) + ) + with pytest.raises(ValueError, match="stride 0"): + transform.as_basic_selection() + + def test_rejects_a_transposed_dimension_order(self) -> None: + transform = IndexTransform( + IndexDomain.from_shape((2, 3)), (DimensionMap(1), DimensionMap(0)) + ) + with pytest.raises(ValueError, match="increasing order"): + transform.as_basic_selection() + + def test_rejects_a_repeated_input_dimension(self) -> None: + transform = IndexTransform( + IndexDomain.from_shape((2,)), (DimensionMap(0), DimensionMap(0, offset=1)) + ) + with pytest.raises(ValueError, match="increasing order"): + transform.as_basic_selection() + + def test_rejects_an_unreferenced_non_singleton_dimension(self) -> None: + """An axis restored by repetition — a duplicate gather of a constant.""" + transform = IndexTransform.from_shape((5,)).oindex[[3]].oindex[np.array([0, 0])] + with pytest.raises(ValueError, match="only a singleton axis"): + transform.as_basic_selection() + + def test_rejects_a_singleton_axis_with_no_selector_left(self) -> None: + transform = IndexTransform(IndexDomain.from_shape((2, 1)), (DimensionMap(0),)) + with pytest.raises(ValueError, match="no selector is left"): + transform.as_basic_selection() + + def test_rejects_a_newaxis(self) -> None: + """`None` inserts an axis basic int/slice selectors cannot produce.""" + transform = IndexTransform.from_shape((5,))[None] + with pytest.raises(ValueError, match="increasing order"): + transform.as_basic_selection() + + def test_rejects_a_negative_constant_coordinate(self) -> None: + transform = IndexTransform(IndexDomain.from_shape(()), (ConstantMap(-1),)) + with pytest.raises(ValueError, match="count from the end"): + transform.as_basic_selection() + + def test_rejects_a_negative_slice_coordinate(self) -> None: + transform = IndexTransform(IndexDomain.from_shape((2,)), (DimensionMap(0, offset=-3),)) + with pytest.raises(ValueError, match="count from the end"): + transform.as_basic_selection() + + class TestIndexTransformBasicIndexing: def test_slice_identity(self) -> None: """slice(None) on identity transform is a no-op.""" From d16f22c82a9b0fa14e37d59934944b3dafbd8386 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 10:41:38 +0200 Subject: [PATCH 02/16] docs(zarr-indexing): add asyncio integration example and consumer-owned I/O guide MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A peer to the Dask example: parts() driven by asyncio.gather against zarr.AsyncArray, with three loops — one fetch per part through Partition.source_selection, a decoded-chunk cache keyed on base_coords and placed with chunk_local_selection, and the ValueError fallback for query parts. The integrations guide gains a matching 'Consumer-owned I/O' section with a tested snippet. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+asyncio-example.doc.md | 6 + .../docs/examples/lazy_indexing_asyncio.md | 13 ++ .../zarr-indexing/docs/guide/integrations.md | 22 ++ .../docs/snippets/integrations.py | 16 ++ .../examples/lazy_indexing_asyncio/README.md | 44 ++++ .../lazy_indexing_asyncio.py | 193 ++++++++++++++++++ packages/zarr-indexing/mkdocs.yml | 1 + 7 files changed, 295 insertions(+) create mode 100644 packages/zarr-indexing/changes/+asyncio-example.doc.md create mode 100644 packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md create mode 100644 packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md create mode 100644 packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py diff --git a/packages/zarr-indexing/changes/+asyncio-example.doc.md b/packages/zarr-indexing/changes/+asyncio-example.doc.md new file mode 100644 index 0000000000..55199590fa --- /dev/null +++ b/packages/zarr-indexing/changes/+asyncio-example.doc.md @@ -0,0 +1,6 @@ +Added an asyncio integration example (`examples/lazy_indexing_asyncio/`), +peer to the Dask example: `parts()` driven by `asyncio.gather` against +`zarr.AsyncArray`, using `Partition.source_selection` for per-part fetches, a +decoded-chunk cache keyed on `base_coords` placed with +`chunk_local_selection`, and the `ValueError` fallback for query parts. The +integrations guide gained a matching "Consumer-owned I/O" section. diff --git a/packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md b/packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md new file mode 100644 index 0000000000..cc31abb587 --- /dev/null +++ b/packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md @@ -0,0 +1,13 @@ +--8<-- "lazy_indexing_asyncio/README.md" + +`LazyArray` owns the indexing-derived plan — which grid cells a view touches, +what to request from each, and where each block lands — while the consumer's +event loop owns concurrency and I/O. The projection pair travels with each +partition, so a cache keyed on `base_coords` and sliced with +`chunk_local_selection` needs no coordinate arithmetic of its own. + +## Source Code + +```python +--8<-- "lazy_indexing_asyncio/lazy_indexing_asyncio.py" +``` diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index 03661ed43c..ec45c26647 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -70,6 +70,28 @@ view = LazyArray(source).with_reader(unit_step_reader) A strided selection then over-reads its cover by the stride factor, which the partitioning above bounds by one part. +### Consumer-owned I/O: parts as backend requests + +Both regimes above still read *through* the wrapper. A consumer with its own +I/O layer — an async store, an HTTP endpoint, a connection pool — can instead +use the wrapper purely as a planner: every box-shaped part lowers to a +backend-native basic selection with +[`Partition.source_selection`][zarr_indexing.lazy_array.Partition], and its +paired `out_selection` places whatever comes back: + +```python +--8<-- "snippets/integrations.py:consumer-owned-io" +``` + +[`Partition.chunk_local_selection`][zarr_indexing.lazy_array.Partition] is the +same read relative to the part's grid cell, for consumers caching decoded +chunks keyed on `base_coords`. A query part (an `oindex`/`vindex` gather) has +no slab spelling and raises `ValueError`, so mixed consumers fall back to +`part.view.result()` for those parts. The +[asyncio example](../examples/lazy_indexing_asyncio.md) drives all three +loops — gather-per-part, decoded-chunk cache, and the query fallback — with +`asyncio.gather` over `zarr.AsyncArray`. + ## napari-like consumer This is a **napari-like consumer**, not a napari integration. It models the diff --git a/packages/zarr-indexing/docs/snippets/integrations.py b/packages/zarr-indexing/docs/snippets/integrations.py index 5699c0d319..d573beeaa9 100644 --- a/packages/zarr-indexing/docs/snippets/integrations.py +++ b/packages/zarr-indexing/docs/snippets/integrations.py @@ -165,3 +165,19 @@ def materialize(view: LazyArray) -> Any: for key in slab_source.keys ) # --8<-- [end:dense-box-repartition] + + +# --8<-- [start:consumer-owned-io] +recorder = RecordingArray(np.arange(100).reshape(10, 10), chunks=(4, 4)) +planner = LazyArray(recorder) +view = planner.lazy[1:9:2, 4:] + +consumer_io = np.arange(100).reshape(10, 10) # stands in for the consumer's I/O layer +out = np.empty(view.shape, dtype=view.dtype) +for part in view.parts(): + # Each box part lowers to a backend-native basic selection; nothing + # reads through the wrapper or its reader. + out[part.out_selection] = consumer_io[part.source_selection] +assert (out == consumer_io[1:9:2, 4:]).all() +assert recorder.keys == [] # the planner's own source was never read +# --8<-- [end:consumer-owned-io] diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md new file mode 100644 index 0000000000..07e62f3f92 --- /dev/null +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md @@ -0,0 +1,44 @@ +# Lazy Indexing with asyncio + +This example demonstrates using `zarr_indexing.LazyArray` as a chunk *planner* +while an async I/O layer — here `zarr.AsyncArray` — performs every read. The +package deliberately contains no scheduler; `parts()` exposes the partition +structure and this example shows `asyncio.gather` driving it. + +The example shows how to: + +- Lower each box-shaped partition to a backend-native request with + `part.source_selection` — a tuple of integers and slices in the wrapped + array's own coordinates — and fetch all partitions concurrently, assembling + each block with `out[part.out_selection] = await source.getitem(part.source_selection)` +- Build a decoded-chunk cache keyed on `part.base_coords`: fetch each touched + grid cell (`projection.chunk_domain`) once, then serve every overlapping view + from the cache with `part.chunk_local_selection` +- Fall back for query partitions (`oindex`/`vindex`/mask selections), whose + gathers have no single-slab spelling: `source_selection` raises `ValueError` + and the part resolves through its own `part.view.result()` instead + +The async side only needs one method — `async def getitem(selection)` accepting +a basic selection — so the same loop drives an HTTP range endpoint or any other +async source. Note that a reversing view (`lazy[::-1]`) lowers to a +negative-step slice, which Zarr's basic selections reject; inspect the slice +steps if your backend only walks forward. + +## Running the Example + +The script declares its dependencies inline +([PEP 723](https://peps.python.org/pep-0723/)), so the easiest way to run it is +with [uv](https://docs.astral.sh/uv/), which installs them automatically: + +```bash +cd packages/zarr-indexing +uv run --with-editable . examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py +``` + +Alternatively, run it with plain Python, in which case you must first install +`zarr`, `zarr-indexing`, `numpy`, and `pytest` yourself: + +```bash +cd packages/zarr-indexing +python examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py +``` diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py new file mode 100644 index 0000000000..a60fb0e445 --- /dev/null +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py @@ -0,0 +1,193 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = [ +# "zarr @ git+https://github.com/zarr-developers/zarr-python.git@main", +# "zarr-indexing>=0.1", +# "numpy==2.4.3", +# "pytest==9.0.2" +# ] +# /// +# + +""" +Demonstrate driving zarr_indexing.LazyArray's partition plan with asyncio +""" + +import asyncio +import sys +from typing import Any, Protocol + +import numpy as np +import pytest +import zarr +import zarr.api.asynchronous + +from zarr_indexing import LazyArray, Partition + + +class AsyncSource(Protocol): + """The surface the async loop needs: one awaitable basic-selection read. + + `zarr.AsyncArray` satisfies it; so does anything else that can serve a + tuple of integers and ascending slices — an HTTP tile endpoint, an fsspec + wrapper, a database. The planner never sees this object. + """ + + async def getitem(self, selection: Any) -> Any: ... + + +@pytest.fixture +def store() -> dict[str, Any]: + """A fresh in-memory store, shared by the sync and async handles.""" + return {} + + +@pytest.fixture +def source(store: dict[str, Any]) -> zarr.Array: + """A chunked Zarr array, created synchronously and shared with the async side.""" + array = zarr.create_array(store=store, shape=(40, 30), chunks=(10, 10), dtype="i4") + array[:] = np.arange(40 * 30).reshape(40, 30) + return array + + +async def read_through(view: LazyArray, source: AsyncSource) -> np.ndarray: + """Materialize `view` by fetching every partition concurrently. + + The wrapper plans; the async source fetches. Each box partition lowers to + the basic selection `part.source_selection`, is fetched through the + caller's own I/O layer, and lands at `part.out_selection` — no reader, no + thread pool, and no scheduler inside zarr-indexing. + """ + parts = tuple(view.parts()) + blocks = await asyncio.gather(*(source.getitem(part.source_selection) for part in parts)) + out = np.empty(view.shape, dtype=view.dtype) + for part, block in zip(parts, blocks, strict=True): + out[part.out_selection] = block + return out + + +def test_parts_with_asyncio_gather(store: dict[str, Any], source: zarr.Array) -> None: + """Fetch a view's partitions concurrently through zarr.AsyncArray.""" + view = LazyArray(source).lazy[5:35, 3:27] + + async def scenario() -> np.ndarray: + async_source = await zarr.api.asynchronous.open_array(store=store) + return await read_through(view, async_source) + + result = asyncio.run(scenario()) + assert result.shape == (30, 24) + assert np.array_equal(result, source[5:35, 3:27]) + + # Strided and integer selections lower the same way; a scalar axis lowers + # to an integer and drops, exactly as it does in the output placement. + decimated = LazyArray(source).lazy[::4, 7] + + async def decimated_scenario() -> np.ndarray: + async_source = await zarr.api.asynchronous.open_array(store=store) + return await read_through(decimated, async_source) + + assert np.array_equal(asyncio.run(decimated_scenario()), source[::4, 7]) + + +def test_decoded_chunk_cache(store: dict[str, Any], source: zarr.Array) -> None: + """Key a decoded-chunk cache on `base_coords`; place with `chunk_local_selection`. + + A tile server reads many overlapping views of the same array. Fetching + whole chunks once and slicing every view out of the cached cells turns + N overlapping requests into one fetch per chunk. `projection.chunk_domain` + is the cell to fetch, `base_coords` is its cache key, and + `chunk_local_selection` is the view's read relative to the cell's origin. + """ + cache: dict[tuple[int, ...], np.ndarray] = {} + + def cell_selection(part: Partition) -> tuple[slice, ...]: + domain = part.projection.chunk_domain + return tuple( + slice(lo, hi) for lo, hi in zip(domain.inclusive_min, domain.exclusive_max, strict=True) + ) + + async def fetch_missing(parts: tuple[Partition, ...], async_source: AsyncSource) -> None: + missing = { + part.base_coords: cell_selection(part) + for part in parts + if part.base_coords not in cache + } + cells = await asyncio.gather( + *(async_source.getitem(selection) for selection in missing.values()) + ) + cache.update(zip(missing.keys(), cells, strict=True)) + + async def read_cached(view: LazyArray, async_source: AsyncSource) -> np.ndarray: + parts = tuple(view.parts()) + await fetch_missing(parts, async_source) + out = np.empty(view.shape, dtype=view.dtype) + for part in parts: + out[part.out_selection] = cache[part.base_coords][part.chunk_local_selection] + return out + + async def scenario() -> tuple[np.ndarray, np.ndarray]: + async_source = await zarr.api.asynchronous.open_array(store=store) + first = await read_cached(LazyArray(source).lazy[5:25, 3:27], async_source) + # The second view overlaps the first, so most cells are already cached. + second = await read_cached(LazyArray(source).lazy[15:35, ::2], async_source) + return first, second + + first, second = asyncio.run(scenario()) + assert np.array_equal(first, source[5:25, 3:27]) + assert np.array_equal(second, source[15:35, ::2]) + # Every cached cell was fetched at most once: the first view touches 9 + # chunks, and of the 9 the second touches only 3 are new. + assert len(cache) == 12 + + +def test_query_parts_fall_back(store: dict[str, Any], source: zarr.Array) -> None: + """A query part refuses to lower; the wrapper's own reader is the fallback. + + A gather (`oindex`, `vindex`, a mask) has no single-slab spelling, so + `source_selection` raises `ValueError` instead of guessing one. A consumer + mixing selection kinds catches that and resolves the part through + `part.view.result()`, which reads through the wrapped array's reader. + """ + view = LazyArray(source).lazy.oindex[[30, 2, 2], 4:10] + + async def scenario() -> np.ndarray: + async_source = await zarr.api.asynchronous.open_array(store=store) + parts = tuple(view.parts()) + out = np.empty(view.shape, dtype=view.dtype) + + async def resolve(part: Partition) -> tuple[Partition, np.ndarray]: + try: + selection = part.source_selection + except ValueError: + # The gathered axis needs a lookup, not a slab: read this part + # through the wrapper (synchronously here; a real consumer + # might push it to a thread, or fetch the part's bounding box + # and gather in memory). + return part, np.asarray(part.view.result()) + return part, np.asarray(await async_source.getitem(selection)) + + for part, block in await asyncio.gather(*(resolve(part) for part in parts)): + out[part.out_selection] = block + return out + + assert np.array_equal(asyncio.run(scenario()), source.oindex[[30, 2, 2], 4:10]) + + +if __name__ == "__main__": + # Run the example with printed output, and a dummy pytest configuration file specified. + # Without the dummy configuration file, at test time pytest will attempt to use the + # configuration file in the project root, which will error because Zarr is using some + # plugins that are not installed in this example. + sys.exit( + pytest.main( + [ + "-s", + __file__, + f"-c {__file__}", + # Suppress: "PytestAssertRewriteWarning: Module already imported so + # cannot be rewritten; zarr" + "-W", + "ignore::pytest.PytestAssertRewriteWarning", + ] + ) + ) diff --git a/packages/zarr-indexing/mkdocs.yml b/packages/zarr-indexing/mkdocs.yml index d97e63b150..d752b9361c 100644 --- a/packages/zarr-indexing/mkdocs.yml +++ b/packages/zarr-indexing/mkdocs.yml @@ -38,6 +38,7 @@ nav: - Examples: - Lazy indexing a NumPy array: examples/lazy_indexing_numpy.md - Lazy indexing with Dask: examples/lazy_indexing_dask.md + - Lazy indexing with asyncio: examples/lazy_indexing_asyncio.md - System-memory chunk cache: examples/system_memory_chunk_cache.md - The ndsel wire format: ndsel.md - Design notes: design-notes.md From edc35b5a5baeb5cb09331ab403ea853cff7d396b Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 16:54:00 +0200 Subject: [PATCH 03/16] fix(zarr-indexing): name partition cells in the source's own coordinates A view that partitions a window plans over that window, so its parts named their cell in window coordinates while their transform addressed the source. Three consequences, all of them silent: - `chunk_domain` promises global storage coordinates, so the documented decoded-cell read `cell[part.chunk_local_selection]` fetched the wrong region for a part of a part. It is now translated back onto the source. - A part of a part carried that window-relative projection into its own `result()`, where a reader keyed on `chunk_coords` would fetch a chunk the part does not read. Before the pairing existed such a read refused loudly; it refuses again, and the partitioned path takes its context from the part view so both paths agree by construction rather than by coincidence. - `view.result(parts=view.parts())` on an unpartitioned view read through a freshly synthesized whole-base projection instead of the view's own. `base_coords` counts cells of the partitioned base, so it keys a decoded-chunk cache only alongside the grid that produced it; the asyncio example keys on the cell's global origin instead, and the docs say which is which. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+asyncio-example.doc.md | 2 +- .../changes/+part-lowering.feature.md | 7 +- .../changes/+part-view-projection.bugfix.md | 8 +- .../docs/examples/lazy_indexing_asyncio.md | 2 +- .../zarr-indexing/docs/guide/integrations.md | 5 +- .../examples/lazy_indexing_asyncio/README.md | 6 +- .../lazy_indexing_asyncio.py | 21 +++-- .../src/zarr_indexing/lazy_array.py | 93 ++++++++++++++----- .../zarr-indexing/tests/test_lazy_array.py | 46 ++++++++- 9 files changed, 150 insertions(+), 40 deletions(-) diff --git a/packages/zarr-indexing/changes/+asyncio-example.doc.md b/packages/zarr-indexing/changes/+asyncio-example.doc.md index 55199590fa..1c0fd483ff 100644 --- a/packages/zarr-indexing/changes/+asyncio-example.doc.md +++ b/packages/zarr-indexing/changes/+asyncio-example.doc.md @@ -1,6 +1,6 @@ Added an asyncio integration example (`examples/lazy_indexing_asyncio/`), peer to the Dask example: `parts()` driven by `asyncio.gather` against `zarr.AsyncArray`, using `Partition.source_selection` for per-part fetches, a -decoded-chunk cache keyed on `base_coords` placed with +decoded-chunk cache keyed on each cell's `chunk_domain` origin and placed with `chunk_local_selection`, and the `ValueError` fallback for query parts. The integrations guide gained a matching "Consumer-owned I/O" section. diff --git a/packages/zarr-indexing/changes/+part-lowering.feature.md b/packages/zarr-indexing/changes/+part-lowering.feature.md index 927d687789..4ce8aa4120 100644 --- a/packages/zarr-indexing/changes/+part-lowering.feature.md +++ b/packages/zarr-indexing/changes/+part-lowering.feature.md @@ -12,8 +12,11 @@ async store, an HTTP range endpoint): async consumer's whole loop is `out[part.out_selection] = await source.getitem(part.source_selection)`. - `Partition.chunk_local_selection` is the same read relative to - `projection.chunk_domain`'s origin, for decoded-chunk caches keyed on - `base_coords`. + `projection.chunk_domain`'s origin, for decoded-chunk caches. The domain + names its cell in the source's own coordinates even for a view partitioning + a window of it, so the cached cell is the same read either way; + `base_coords` counts cells of the partitioned base, so it keys such a cache + only together with the grid that produced it. `ChainedIndexingStateMachine` gained an invariant that runs both documented assembly loops literally under plain NumPy semantics, with no reader involved. diff --git a/packages/zarr-indexing/changes/+part-view-projection.bugfix.md b/packages/zarr-indexing/changes/+part-view-projection.bugfix.md index c5e742bf45..29d43def61 100644 --- a/packages/zarr-indexing/changes/+part-view-projection.bugfix.md +++ b/packages/zarr-indexing/changes/+part-view-projection.bugfix.md @@ -5,4 +5,10 @@ A custom reader keyed on `projection.chunk_coords` — a decoded-chunk cache — now behaves identically on both paths, which the dask example's task-per-partition pattern relies on. `with_reader` and repartitioning keep the pairing; a further `.lazy` selection describes a different read and drops -it. +it. So does a further partitioning of a part, whose cells are counted in the +part's own box and would name the wrong chunk of the source — such a read +pairs with no projection, as it did before, rather than with a misleading one. + +Reusing a plan (`view.result(parts=view.parts())`) on an unpartitioned view +now reads through the same context as the plain call, instead of through a +freshly synthesized whole-base projection. diff --git a/packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md b/packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md index cc31abb587..89ffaf1eb5 100644 --- a/packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md +++ b/packages/zarr-indexing/docs/examples/lazy_indexing_asyncio.md @@ -3,7 +3,7 @@ `LazyArray` owns the indexing-derived plan — which grid cells a view touches, what to request from each, and where each block lands — while the consumer's event loop owns concurrency and I/O. The projection pair travels with each -partition, so a cache keyed on `base_coords` and sliced with +partition, so a cache keyed on each cell's `chunk_domain` and sliced with `chunk_local_selection` needs no coordinate arithmetic of its own. ## Source Code diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index ec45c26647..f6972abdaf 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -85,7 +85,10 @@ paired `out_selection` places whatever comes back: [`Partition.chunk_local_selection`][zarr_indexing.lazy_array.Partition] is the same read relative to the part's grid cell, for consumers caching decoded -chunks keyed on `base_coords`. A query part (an `oindex`/`vindex` gather) has +chunks. `projection.chunk_domain` locates that cell in the source, so it is +what to fetch and a sound cache key; `base_coords` counts cells of whatever +base the view partitions, so it keys a cache only alongside the grid that +produced it. A query part (an `oindex`/`vindex` gather) has no slab spelling and raises `ValueError`, so mixed consumers fall back to `part.view.result()` for those parts. The [asyncio example](../examples/lazy_indexing_asyncio.md) drives all three diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md index 07e62f3f92..17e4aad65f 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md @@ -11,9 +11,9 @@ The example shows how to: `part.source_selection` — a tuple of integers and slices in the wrapped array's own coordinates — and fetch all partitions concurrently, assembling each block with `out[part.out_selection] = await source.getitem(part.source_selection)` -- Build a decoded-chunk cache keyed on `part.base_coords`: fetch each touched - grid cell (`projection.chunk_domain`) once, then serve every overlapping view - from the cache with `part.chunk_local_selection` +- Build a decoded-chunk cache keyed on each touched grid cell's global origin: + fetch the cell (`projection.chunk_domain`) once, then serve every overlapping + view from the cache with `part.chunk_local_selection` - Fall back for query partitions (`oindex`/`vindex`/mask selections), whose gathers have no single-slab spelling: `source_selection` raises `ValueError` and the part resolves through its own `part.view.result()` instead diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py index a60fb0e445..6733241734 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py @@ -90,16 +90,25 @@ async def decimated_scenario() -> np.ndarray: def test_decoded_chunk_cache(store: dict[str, Any], source: zarr.Array) -> None: - """Key a decoded-chunk cache on `base_coords`; place with `chunk_local_selection`. + """Cache decoded cells; place each view's share with `chunk_local_selection`. A tile server reads many overlapping views of the same array. Fetching whole chunks once and slicing every view out of the cached cells turns N overlapping requests into one fetch per chunk. `projection.chunk_domain` - is the cell to fetch, `base_coords` is its cache key, and - `chunk_local_selection` is the view's read relative to the cell's origin. + is the cell to fetch, and `chunk_local_selection` is the view's read + relative to the cell's origin. + + The cache is keyed on that domain's origin rather than on `base_coords`, + which counts cells of whatever base its view partitions: two views sharing + a source but not a grid — or a part re-partitioned into smaller boxes — + number their cells differently, while the origin names one region of the + source however it was reached. """ cache: dict[tuple[int, ...], np.ndarray] = {} + def cell_key(part: Partition) -> tuple[int, ...]: + return part.projection.chunk_domain.inclusive_min + def cell_selection(part: Partition) -> tuple[slice, ...]: domain = part.projection.chunk_domain return tuple( @@ -108,9 +117,7 @@ def cell_selection(part: Partition) -> tuple[slice, ...]: async def fetch_missing(parts: tuple[Partition, ...], async_source: AsyncSource) -> None: missing = { - part.base_coords: cell_selection(part) - for part in parts - if part.base_coords not in cache + cell_key(part): cell_selection(part) for part in parts if cell_key(part) not in cache } cells = await asyncio.gather( *(async_source.getitem(selection) for selection in missing.values()) @@ -122,7 +129,7 @@ async def read_cached(view: LazyArray, async_source: AsyncSource) -> np.ndarray: await fetch_missing(parts, async_source) out = np.empty(view.shape, dtype=view.dtype) for part in parts: - out[part.out_selection] = cache[part.base_coords][part.chunk_local_selection] + out[part.out_selection] = cache[cell_key(part)][part.chunk_local_selection] return out async def scenario() -> tuple[np.ndarray, np.ndarray]: diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index c99b8ade45..045ac82a8c 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -147,7 +147,7 @@ import operator import uuid from collections.abc import Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, Any, Protocol, cast import numpy as np @@ -161,6 +161,7 @@ ChunkProjection, plan_chunks, ) +from zarr_indexing.domain import IndexDomain from zarr_indexing.grid import DimensionGrid, FixedDimension, dimension_grids_from_chunks from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.reader import ( @@ -298,6 +299,18 @@ class _PartOwner: """Opaque identity shared only by one view and the parts it prepared.""" +def _translated_domain(domain: IndexDomain, window: tuple[slice, ...]) -> IndexDomain: + """`domain`, expressed in the coordinates `window` was cut from.""" + return IndexDomain( + inclusive_min=tuple( + w.start + lo for w, lo in zip(window, domain.inclusive_min, strict=True) + ), + exclusive_max=tuple( + w.start + hi for w, hi in zip(window, domain.exclusive_max, strict=True) + ), + ) + + def _partition_out_selection( cell_transform: IndexTransform, ) -> tuple[Any, ...]: @@ -434,9 +447,18 @@ class Partition: domain, mapping each selected cell to chunk-local storage and request coordinates respectively. This is the authoritative placement model; `base_coords` and `is_complete` are conveniences derived from it. + `chunk_domain` locates the cell in the wrapped array's own global + coordinates, while `chunk_coords` names it in the grid the producing + view partitions — the two differ once a view partitions a window of + the source rather than the whole of it. base_coords - Which box of the base partitioning this is, one coordinate per dimension - of the wrapped array. + Which box of the base partitioning this is, one coordinate per + dimension of the wrapped array. A cell coordinate in the producing + view's grid, so it identifies a chunk of the source only for a view + that partitions the whole source (`base_shape` equals the source's + shape) with the source's own grid. A cache keyed on it must therefore + be keyed on that grid too, or on `projection.chunk_domain`, which is + global. box The box itself, in the global storage coordinates of the wrapped array: one `[inclusive_min, exclusive_max)` interval per dimension. It @@ -448,11 +470,15 @@ class Partition: 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. Resolving the view reads - the box once through its selected reader, handing it a `ReadContext` - that carries this partition's `projection` — the same context the - parent view's partitioned `result()` passes, so a reader keyed on - `chunk_coords` behaves identically on either path. A further `.lazy` - selection describes a different read and drops that pairing. Named + the box once through its selected reader, handing it exactly the + `ReadContext` the parent view's partitioned `result()` passes, so a + reader keyed on `chunk_coords` behaves identically on either path. + That context carries this partition's `projection` when the parent + partitions the source itself; a part of a part partitions a window, + whose cell coordinates would name the wrong chunk of the source, so it + pairs with no projection and a projection-keyed reader refuses it + rather than reading the wrong cell. A further `.lazy` selection + describes a different read and drops the pairing too. 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. @@ -537,10 +563,13 @@ def chunk_local_selection(self) -> tuple[int | slice, ...]: Lowered from `projection.chunk_transform`, so coordinate 0 per axis is `projection.chunk_domain`'s origin. For a consumer that caches decoded - cells keyed on `base_coords`, this is the selection to apply to a - cached cell: `cell[part.chunk_local_selection]` yields the same values - as `view.array[part.source_selection]`, and both land at - `out_selection` in the request buffer. + cells, this is the selection to apply to a cached cell: with `cell` + holding `view.array[projection.chunk_domain]` — the domain is global, + so that read is the same whatever narrowed the view — + `cell[part.chunk_local_selection]` yields the same values as + `view.array[part.source_selection]`, and both land at `out_selection` + in the request buffer. `base_coords` keys such a cache only for a view + partitioning the whole source; see its own note. Defined exactly when `source_selection` is: the two lower the same maps, so a part that refuses one spelling refuses both. @@ -1140,9 +1169,9 @@ def parts(self) -> Iterator[Partition]: 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 + for planned in plan_chunks(plan_transform, grids): + base_coords = planned.chunk_coords + local = planned.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: @@ -1160,8 +1189,17 @@ def parts(self) -> Iterator[Partition]: # `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: + projection = planned global_origin = origin else: + # A windowed view partitions its window, so the walk names the + # cell in window coordinates. `chunk_domain` promises global + # storage coordinates, and `chunk_local_selection` is only the + # cell's own read if the cell it names is the one the source + # holds, so translate it back onto the source. + projection = replace( + planned, chunk_domain=_translated_domain(planned.chunk_domain, self._window) + ) global_origin = tuple( w.start + o for w, o in zip(self._window, origin, strict=True) ) @@ -1174,7 +1212,12 @@ def parts(self) -> Iterator[Partition]: None, window, self._reader, - projection, + # `chunk_coords` addresses this view's own grid, which is + # the source's own only when nothing narrowed the base. + # Handing a reader that keys on it a cell coordinate of a + # grid over a window would fetch the wrong chunk, so a + # part of a part pairs with no projection at all. + projection if self._window is None else None, ), out_selection=_partition_out_selection(projection.cell_transform), _owner=self._part_owner, @@ -1352,10 +1395,13 @@ def _read_into_buffer(self, out: Any, prepared_parts: tuple[Partition, ...] | No if size == 0: return out - if prepared_parts is None and self._parts is None: - # A partition view carries its paired projection, so resolving it - # alone hands the reader the same ReadContext the parent's - # partitioned read would. + if self._parts is None: + # An unpartitioned view's walk is the single part that is the view + # itself, so reading it directly is the same read — and the only + # one that keeps the projection this view was paired with. Going + # through the loop would hand the reader the walk's freshly + # synthesized whole-base projection instead, making a supplied + # `parts=` plan read differently from the plain call. _invoke_reader( self._reader, self._array, ReadContext(self._transform, self._projection), out ) @@ -1376,7 +1422,12 @@ def _read_into_buffer(self, out: Any, prepared_parts: tuple[Partition, ...] | No _invoke_reader( self._reader, self._array, - ReadContext(part.view.transform, part.projection), + # The part view's own pairing, not `part.projection`: the two + # agree wherever the projection describes a read of the source + # itself, and taking it from the view is what makes resolving + # the part here and resolving it alone the same read by + # construction rather than by coincidence. + ReadContext(part.view.transform, part.view._projection), destination, ) if not direct: diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index a3eb1cd90e..ea92b6dc8b 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2664,17 +2664,21 @@ def test_box_part_selections_carry_the_read() -> None: def test_part_of_part_selections_stay_global_and_cell_local() -> None: - """A nested part addresses the raw source; its cell stays window-relative.""" + """A nested part's cell is named in the source's own coordinates. + + A view partitioning a window plans over that window, but `chunk_domain` + describes the source, so the decoded-cell loop reads the same cell from + the raw array whether or not anything narrowed the view first. + """ data = reference() view = LazyArray.from_numpy(data).with_parts((3, 2, 3)).lazy[1:6, 1:4, :] outer = next(view.parts()) inner_view = outer.view.with_parts((2, 1, 2)) expected = np.asarray(outer.view.result()) - boxed = data[tuple(slice(lo, hi) for lo, hi in outer.box)] out = np.full(inner_view.shape, -1, dtype=inner_view.dtype) for part in inner_view.parts(): slab = data[part.source_selection] - cell = boxed[ + cell = data[ tuple( slice(lo, hi) for lo, hi in zip( @@ -2737,6 +2741,42 @@ def test_reader_and_partitioning_swaps_keep_a_part_views_projection() -> None: assert [context.projection for context in reader.contexts] == [part.projection] * 2 +def test_a_part_of_a_part_pairs_with_no_projection() -> None: + """A window's cell coordinates would name the wrong chunk of the source. + + A part view partitions its own box, so its parts' `chunk_coords` count + cells of that box. Handing those to a reader keyed on `chunk_coords` + alongside the raw source would fetch a different chunk than the one the + part reads, so the pairing stops at the first level and such a reader + refuses the read instead. + """ + reader = RecordingReader() + view = LazyArray(np.arange(24).reshape(4, 6)).with_parts((2, 3)) + outer = list(view.parts())[3] + assert outer.base_coords == (1, 1) + inner = next(outer.view.with_parts((1, 3)).parts()) + # The cell it counts is its own box's first, not the source's. + assert inner.base_coords == (0, 0) + # ... while the domain that names the cell stays global, as does the read. + assert inner.projection.chunk_domain.inclusive_min == (2, 3) + assert inner.source_selection == (slice(2, 3, 1), slice(3, 6, 1)) + inner.view.with_reader(reader).result() + assert reader.contexts[-1].projection is None + + +def test_prepared_parts_read_through_the_same_context_as_the_plain_call() -> None: + """The `parts=` reuse path is an optimization, not a different read.""" + reader = RecordingReader() + view = LazyArray(np.arange(8)).with_parts((4,)) + part = list(view.parts())[1] + part_view = part.view.with_reader(reader) + np.testing.assert_array_equal( + np.asarray(part_view.result(parts=list(part_view.parts()))), + np.asarray(part_view.result()), + ) + assert [context.projection for context in reader.contexts] == [part.projection] * 2 + + def test_a_new_selection_on_a_part_view_drops_the_projection() -> None: """A further selection describes a different read, so the pairing ends.""" reader = RecordingReader() From 53879a848db1f9eb01586c783b5546ec719c3328 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 16:58:01 +0200 Subject: [PATCH 04/16] fix(zarr-indexing): reject result_into buffers that cannot hold the result MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `result()` allocates, so two buffers it could never produce were unreachable until `result_into` let the caller supply one, and both failed silently: - A plain ndarray for a `numpy.ma` source passed the dtype check and came back holding the values beneath the mask, presented as data. - A buffer overlapping the wrapped array had each part overwrite cells the parts after it still had to read, so the result was wrong in a way that depended on the partitioning. Both are now refused. The overlap test asks the cheap bounds question first and only pays for an exact answer once memory is known to overlap. `result_into`'s own docstring recommended `final[part.out_selection]` as a destination with only a parenthetical hedge, though NumPy hands back a copy for a fancy `out_selection` — which this method would fill and return, leaving `final` untouched. It cannot tell a copy from a view, so the doc now branches on the selection instead of hedging. Translating a windowed cell's domain keeps its labels (roborev, edc35b5). Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+result-into.feature.md | 5 ++ .../src/zarr_indexing/lazy_array.py | 68 ++++++++++++++++--- .../zarr-indexing/tests/test_lazy_array.py | 18 +++++ 3 files changed, 82 insertions(+), 9 deletions(-) diff --git a/packages/zarr-indexing/changes/+result-into.feature.md b/packages/zarr-indexing/changes/+result-into.feature.md index 576b94dd1f..ba10c447c8 100644 --- a/packages/zarr-indexing/changes/+result-into.feature.md +++ b/packages/zarr-indexing/changes/+result-into.feature.md @@ -4,3 +4,8 @@ and dtype, filled in place — every cell written exactly once — and returned. view into a larger array qualifies, so a part's block can land directly in its final slot, and a `numpy.ma` masked buffer keeps a masked source's mask. `result()` itself is unchanged and always allocates. + +Validation rejects what `result()`'s own allocation made impossible: a plain +buffer for a masked source, which would silently present the values beneath +the mask as data, and a buffer sharing memory with the wrapped array, where +each part would overwrite cells the parts after it still have to read. diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 045ac82a8c..4a1e4737a5 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -161,7 +161,6 @@ ChunkProjection, plan_chunks, ) -from zarr_indexing.domain import IndexDomain from zarr_indexing.grid import DimensionGrid, FixedDimension, dimension_grids_from_chunks from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.reader import ( @@ -177,6 +176,8 @@ if TYPE_CHECKING: from collections.abc import Callable, Iterator + from zarr_indexing.domain import IndexDomain + SelectFn = Callable[[Any, SelectionMode], "LazyArray"] __all__ = ["LazyArray", "Partition"] @@ -299,9 +300,28 @@ class _PartOwner: """Opaque identity shared only by one view and the parts it prepared.""" +def _overlaps(out: np.ndarray[Any, Any], array: Any) -> bool: + """Whether `out` and `array` can be writing and reading the same memory. + + Only NumPy sources can be compared this way; anything else reaches its + data through its own machinery, where sharing memory with a result buffer + is not expressible. + """ + if not isinstance(array, np.ndarray): + return False + # The cheap bounds test first: an exact answer can cost real work, and it + # is only needed once the two are known to occupy overlapping memory. + return bool(np.may_share_memory(out, array)) and bool(np.shares_memory(out, array)) + + def _translated_domain(domain: IndexDomain, window: tuple[slice, ...]) -> IndexDomain: - """`domain`, expressed in the coordinates `window` was cut from.""" - return IndexDomain( + """`domain`, expressed in the coordinates `window` was cut from. + + Moving a region does not rename its axes, so `replace` carries everything + but the bounds — the labels among them — through unchanged. + """ + return replace( + domain, inclusive_min=tuple( w.start + lo for w, lo in zip(window, domain.inclusive_min, strict=True) ), @@ -1318,11 +1338,24 @@ def result_into( ---------- out A writable `numpy.ndarray` of exactly `self.shape` and this view's - dtype. A view into a larger array qualifies, so a part's result - can land directly in its slot: - `part.view.result_into(final[part.out_selection])` for a - rectangular part. A `numpy.ma` masked buffer (what `result()` - would allocate for a masked source) also qualifies. + dtype, not overlapping the wrapped array. A view into a larger + array qualifies, so a part's result can land directly in its slot: + + ```python + if all(isinstance(s, slice) for s in part.out_selection): + part.view.result_into(final[part.out_selection]) + else: + final[part.out_selection] = part.view.result() + ``` + + The branch is the whole contract: `final[part.out_selection]` is a + writable view only while every selector is a slice, and NumPy hands + back a *copy* for a fancy `out_selection` — which this method would + dutifully fill and return, leaving `final` untouched. It cannot + tell the two apart; only the caller knows what `final` was. + + A masked source requires a `numpy.ma` buffer — what `result()` + would allocate for it — since a plain one drops the mask. parts A reusable sequence previously returned by this exact view's `parts()` method, exactly as for `result`. @@ -1337,7 +1370,8 @@ def result_into( TypeError If `out` is not a `numpy.ndarray`. ValueError - If `out` has the wrong shape or dtype or is read-only, or if + If `out` has the wrong shape or dtype, is read-only, is unmasked + for a masked source, or shares memory with the wrapped array; or if supplied parts were prepared by another view or do not tile this view exactly. AssertionError @@ -1370,6 +1404,22 @@ def result_into( ) if not out.flags.writeable: raise ValueError("out is read-only; result_into writes every cell of it") + if isinstance(self._array, np.ma.MaskedArray) and not isinstance(out, np.ma.MaskedArray): + # dtypes match, so nothing above rejects a plain buffer, and the + # reader would write the values under the mask into it as if they + # were data. + raise ValueError( # noqa: TRY004 - an ndarray, just not one this view fits + "out must be a numpy.ma.MaskedArray for a masked source; a plain " + "ndarray cannot carry the mask, and the values beneath it are not data" + ) + if _overlaps(out, self._array): + # Parts are read in walk order, so a destination overlapping the + # source is read after earlier parts have already overwritten the + # cells it covers — silently, and differently per partitioning. + raise ValueError( + "out shares memory with the wrapped array; reading a view into " + "the array it reads from would overwrite cells later parts read" + ) prepared_parts = self._prepared_parts(parts) return self._read_into_buffer(out, prepared_parts) diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index ea92b6dc8b..18fe703569 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2860,6 +2860,24 @@ def test_result_into_rejects_a_read_only_buffer() -> None: view.result_into(out) +def test_result_into_rejects_an_unmasked_buffer_for_a_masked_source() -> None: + data = reference() + view = LazyArray.from_numpy(np.ma.masked_array(data, mask=data % 5 == 0)) + with pytest.raises(ValueError, match="MaskedArray for a masked source"): + view.result_into(np.empty(view.shape, dtype=view.dtype)) + + +def test_result_into_rejects_a_buffer_that_overlaps_the_source() -> None: + """Reading into the source overwrites cells later parts still have to read.""" + data = reference() + view = LazyArray.from_numpy(data).with_parts((3, 2, 3)).lazy[::-1] + with pytest.raises(ValueError, match="shares memory"): + view.result_into(data) + # A disjoint slice of the same buffer is refused for the same reason. + with pytest.raises(ValueError, match="shares memory"): + LazyArray.from_numpy(data).lazy[:1].result_into(data[1:2]) + + def test_result_into_rejects_foreign_parts() -> None: view = LazyArray.from_numpy(reference()).with_parts((3, 2, 3)) other = LazyArray.from_numpy(reference()).with_parts((3, 2, 3)) From 413a520deeef0a65e0362964f69dbfba3e571635 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 17:01:40 +0200 Subject: [PATCH 05/16] feat(zarr-indexing): lower fabricated axes to the newaxis that made them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `LazyArray._select` admits NumPy's `None`, so `lazy[None, 2:8]` is a box view like any other — but its domain carries an axis no output map reads, and `as_basic_selection` had no spelling for one. Every part of such a view raised "the selectors must produce the domain's axes in increasing order", and since `is_box` is `True` throughout, a consumer following the documented precheck met an uncaught `ValueError` on a selection NumPy itself spells natively. An unreferenced axis is one nothing reads, which is what `None` does, so it now lowers back to `None` — except where a constant is due in its position, which keeps the axis as a length-1 slice, as before. Selections gain `None` alongside ints and slices, and the trailing-axis refusal becomes an assertion: a referenced axis is produced by the map that references it, and an unreferenced one is a proven singleton. `is_box` is documented as necessary but not sufficient — broadcasts and axes restored by repetition still refuse — with `ValueError` the decider. A `None` selector, like a negative step, is one a backend narrower than NumPy may reject. Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/lazy_array.py | 20 ++--- .../src/zarr_indexing/transform.py | 81 +++++++++++++------ .../zarr-indexing/tests/test_lazy_array.py | 11 +++ .../zarr-indexing/tests/test_transform.py | 26 +++--- 4 files changed, 96 insertions(+), 42 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 4a1e4737a5..0722e0b76f 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -543,7 +543,7 @@ def is_complete(self) -> bool: return self.projection.coverage == "full" @property - def source_selection(self) -> tuple[int | slice, ...]: + def source_selection(self) -> tuple[int | slice | None, ...]: """The basic selection on the raw wrapped array that reads this part. Lowered from `view.transform` by @@ -559,13 +559,15 @@ def source_selection(self) -> tuple[int | slice, ...]: out[part.out_selection] = await source.getitem(part.source_selection) ``` - Defined for box-shaped parts (`view.is_box`); a query part has no - slab to request and raises `ValueError`, as does the one degenerate - box with no basic-selection spelling (an axis restored by repetition, - reached by gathering a collapsed constant with duplicates). A consumer - mixing selection kinds catches `ValueError` — or checks `view.is_box` - for the common case — and falls back to `view.result()`. A reversing - view lowers to a negative-step slice. + A query part has no slab to request and raises `ValueError`, as do the + two degenerate boxes with no basic-selection spelling: an axis + restored by repetition (reached by gathering a collapsed constant with + duplicates) and an axis broadcast from one cell. `view.is_box` is + therefore necessary but not sufficient — catching `ValueError` is what + decides it — and a consumer mixing selection kinds falls back to + `view.result()` for the parts that refuse. A reversing view lowers to + a negative-step slice and a fabricated axis to `None`, which a backend + narrower than NumPy may not accept. Examples -------- @@ -578,7 +580,7 @@ def source_selection(self) -> tuple[int | slice, ...]: return self.view.transform.as_basic_selection() @property - def chunk_local_selection(self) -> tuple[int | slice, ...]: + def chunk_local_selection(self) -> tuple[int | slice | None, ...]: """The same read as `source_selection`, relative to the part's grid cell. Lowered from `projection.chunk_transform`, so coordinate 0 per axis is diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index ff3dc263ac..a67bb6e523 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -448,10 +448,11 @@ def inverted(self) -> IndexTransform: output=tuple(inverse_output[dimension] for dimension in range(self.input_rank)), ) - def as_basic_selection(self) -> tuple[int | slice, ...]: + def as_basic_selection(self) -> tuple[int | slice | None, ...]: """Lower a box transform to the NumPy basic selection it describes. - Returns one selector per output dimension such that + Returns one selector per output dimension, plus a `None` for each + domain axis no output map reads, such that `source[transform.as_basic_selection()]` under NumPy indexing semantics reads exactly the cells this transform addresses, in domain order: the result has shape `domain.shape` exactly, and its element at @@ -464,20 +465,24 @@ def as_basic_selection(self) -> tuple[int | slice, ...]: references (a single-coordinate fancy index collapses to a constant at construction, leaving its axis behind: `oindex[:, [2]]`) is instead produced by lowering a constant to a length-1 slice, so the result - keeps the domain's shape. + keeps the domain's shape. An unreferenced axis with no constant to + stand in for it is one nothing reads — a `None`/newaxis in the + selection that made it — and lowers back to `None`. This is the boundary for consumers that plan reads here but perform them through their own I/O layer — an async store, an HTTP range request — whose request vocabulary is a basic selection rather than a - transform. A reversing map lowers to a negative-step slice, so a - backend restricted to ascending unit-step reads should inspect the - steps (or keep such selections out of its plans; compare - `UnitStepReader`). + transform. Such a backend may accept less than NumPy does: a reversing + map lowers to a negative-step slice and a fabricated axis to `None`, + neither of which Zarr's own basic selections take, so a consumer whose + backend is stricter should inspect the lowered selectors (or keep such + views out of its plans; compare `UnitStepReader`). Returns ------- - tuple of int or slice - One selector per output dimension. + tuple of int or slice or None + One selector per output dimension, interleaved with a `None` for + each domain axis no output map reads. Raises ------ @@ -487,10 +492,10 @@ def as_basic_selection(self) -> tuple[int | slice, ...]: slab), a `DimensionMap` has stride 0 (a broadcast repeats one source cell, which no slice spells), the selectors cannot produce the domain's axes in increasing order and exactly once (basic - indexing never transposes, repeats, or fabricates axes), an - unreferenced domain axis has extent other than 1, or a selected - coordinate is negative (NumPy would count it from the end of the - source). + indexing never transposes or repeats axes), an unreferenced domain + axis has extent other than 1 (an axis restored by repetition, + where a newaxis would give extent 1), or a selected coordinate is + negative (NumPy would count it from the end of the source). Examples -------- @@ -504,6 +509,17 @@ def as_basic_selection(self) -> tuple[int | slice, ...]: >>> IndexTransform.from_shape((10, 20)).oindex[slice(None), [2]].as_basic_selection() (slice(0, 10, 1), slice(2, 3, 1)) + An axis the selection fabricated lowers back to the newaxis that made it: + + >>> IndexTransform.from_shape((10, 20))[:, None].as_basic_selection() + (slice(0, 10, 1), None, slice(0, 20, 1)) + + unless a constant is due where the axis is, in which case that + constant keeps it and no newaxis is needed: + + >>> IndexTransform.from_shape((10, 20))[None, 3].as_basic_selection() + (slice(3, 4, 1), slice(0, 20, 1)) + A query has no basic-selection spelling: >>> IndexTransform.from_shape((10,)).oindex[[3, 1, 1]].as_basic_selection() @@ -526,14 +542,30 @@ def as_basic_selection(self) -> tuple[int | slice, ...]: raise ValueError( f"cannot lower to a basic selection: no output map references " f"input dimension {axis}, whose extent is {extent}; only a " - "singleton axis can be produced by an integer-indexed output " - "dimension" + "singleton axis can be produced without reading a source axis, " + "by an integer-indexed output dimension or a newaxis" ) - selection: list[int | slice] = [] + selection: list[int | slice | None] = [] # The next domain axis a selector has to produce. Slices produce axes # in selector order, so walking the outputs left to right must meet the # domain axes in increasing order. next_axis = 0 + + def fabricate_axes_before(axis: int) -> None: + """Spell every unreferenced axis due before `axis` as a newaxis. + + An axis no output map reads is one NumPy inserts rather than + reads, which is what `None` does. The extent check above already + proved each is a singleton, and a constant standing in for one is + preferred to this — a length-1 slice reads the coordinate the + constant names, where a newaxis would need the constant's own + selector to survive as well. + """ + nonlocal next_axis + while next_axis < axis and next_axis not in referenced: + selection.append(None) + next_axis += 1 + for output_dimension, output_map in enumerate(self.output): if isinstance(output_map, ConstantMap): if output_map.offset < 0: @@ -558,13 +590,14 @@ def as_basic_selection(self) -> tuple[int | slice, ...]: "no slice spells a broadcast" ) d = output_map.input_dimension + fabricate_axes_before(d) if d != next_axis: raise ValueError( "cannot lower to a basic selection: the selectors must produce " f"the domain's axes in increasing order and exactly once, but " f"output[{output_dimension}] produces input dimension {d} where " f"input dimension {next_axis} is due; basic indexing never " - "transposes, repeats, or fabricates axes" + "transposes or repeats axes" ) next_axis = d + 1 lo = self.domain.inclusive_min[d] @@ -588,12 +621,14 @@ def as_basic_selection(self) -> tuple[int | slice, ...]: # the front edge instead of wrapping. stop = last + output_map.stride selection.append(slice(first, stop if stop >= 0 else None, output_map.stride)) - if next_axis != self.input_rank: - raise ValueError( - "cannot lower to a basic selection: no selector is left to produce " - f"input dimension {next_axis}, so a basic read cannot yield that " - "axis of the domain" - ) + # Trailing axes no output map reads sit past the last selector, where + # the same newaxis spelling puts them at the end of the result. + fabricate_axes_before(self.input_rank) + assert next_axis == self.input_rank, ( + f"input dimension {next_axis} of {self.input_rank} went unproduced; " + "every referenced axis is produced by the map that references it, and " + "every unreferenced one was proven a singleton above" + ) return tuple(selection) @property diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 18fe703569..0986785b8d 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2693,6 +2693,17 @@ def test_part_of_part_selections_stay_global_and_cell_local() -> None: np.testing.assert_array_equal(out, expected) +def test_parts_of_a_newaxis_view_lower_like_any_other_box() -> None: + """A view can fabricate axes, so its parts have to lower with them.""" + data = reference() + view = LazyArray.from_numpy(data).with_parts((3, 2, 3)).lazy[None, 1:6, :, None, ::2] + out = np.full(view.shape, -1, dtype=view.dtype) + for part in view.parts(): + out[part.out_selection] = data[part.source_selection] + np.testing.assert_array_equal(out, np.asarray(view.result())) + np.testing.assert_array_equal(out, data[None, 1:6, :, None, ::2]) + + def test_query_part_selections_refuse_to_lower() -> None: """A query part has no slab; both spellings must say so, not guess.""" view = LazyArray.from_numpy(reference()).with_parts((3, 2, 3)).lazy.oindex[[6, 1, 1], :, :] diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index ce90428cf9..e6a13e38b9 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -435,13 +435,25 @@ class TestAsBasicSelection: IndexTransform.from_shape((5, 7)).oindex[[3], [2]], id="two-collapsed-gathers", ), + # An axis nothing reads is one the selection fabricated, in every + # position it can hold relative to the axes that are read. + pytest.param(IndexTransform.from_shape((5, 7))[None], id="leading-newaxis"), + pytest.param(IndexTransform.from_shape((5, 7))[:, None], id="interior-newaxis"), + pytest.param(IndexTransform.from_shape((10,))[:, None], id="trailing-newaxis"), + pytest.param(IndexTransform.from_shape((10,))[None, None], id="stacked-newaxis"), + pytest.param(IndexTransform.from_shape((5, 7))[None, 3], id="newaxis-and-int"), + pytest.param( + IndexTransform.from_shape((5, 7)).oindex[slice(None), [2]][None], + id="newaxis-and-collapsed-gather", + ), ], ) def test_selection_reproduces_the_transform(self, transform: IndexTransform) -> None: """`source[selection]` has the domain's shape and its exact values.""" source = np.arange(35).reshape(5, 7) if transform.output_rank == 2 else np.arange(10) selection = transform.as_basic_selection() - assert len(selection) == transform.output_rank + # One selector per output dimension, plus one per fabricated axis. + assert len(selection) - sum(item is None for item in selection) == transform.output_rank read = np.asarray(source[selection]) assert read.shape == transform.domain.shape np.testing.assert_array_equal(read, _pointwise_read(transform, source)) @@ -478,16 +490,10 @@ def test_rejects_an_unreferenced_non_singleton_dimension(self) -> None: with pytest.raises(ValueError, match="only a singleton axis"): transform.as_basic_selection() - def test_rejects_a_singleton_axis_with_no_selector_left(self) -> None: + def test_a_trailing_singleton_axis_lowers_to_a_newaxis(self) -> None: + """No output map reads it, so nothing but a newaxis can produce it.""" transform = IndexTransform(IndexDomain.from_shape((2, 1)), (DimensionMap(0),)) - with pytest.raises(ValueError, match="no selector is left"): - transform.as_basic_selection() - - def test_rejects_a_newaxis(self) -> None: - """`None` inserts an axis basic int/slice selectors cannot produce.""" - transform = IndexTransform.from_shape((5,))[None] - with pytest.raises(ValueError, match="increasing order"): - transform.as_basic_selection() + assert transform.as_basic_selection() == (slice(0, 2, 1), None) def test_rejects_a_negative_constant_coordinate(self) -> None: transform = IndexTransform(IndexDomain.from_shape(()), (ConstantMap(-1),)) From a527676fb636b50a0d10a16cc12b048f5b4cddb1 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 17:05:13 +0200 Subject: [PATCH 06/16] test(zarr-indexing): skip examples by what they import, tighten the lowering invariant MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The example runner guarded on the script's *name* — `if "dask" in stem` — so the asyncio example, which needs `zarr`, was the first to need a guard and not get one: in an environment without zarr (which the suite otherwise supports, and which pyproject deliberately leaves out of the test group) it failed instead of skipping. The guard now reads the imports out of the script, so it holds for whatever an example needs next. The lowering invariant excused any refusing box that contained a `ConstantMap`, which is most of them — an integer-indexed box like `[3, 4:6]` could start refusing and the invariant would stay green. It now names the two shapes that may refuse: an axis restored by repetition, and one broadcast from a single cell. That is also the whole set now that a fabricated axis lowers, so the check is tight in both directions rather than loose in one and wrong in the other. The paired-refusal check no longer re-evaluates the property that already raised on its way in. Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/testing/stateful.py | 55 +++++++++++++------ .../zarr-indexing/tests/test_doc_examples.py | 22 +++++++- 2 files changed, 57 insertions(+), 20 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 32741b15b6..5875db4662 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -57,7 +57,7 @@ def make_source(self, data): from hypothesis.stateful import RuleBasedStateMachine, initialize, invariant, precondition, rule from zarr_indexing.lazy_array import LazyArray -from zarr_indexing.output_map import ConstantMap +from zarr_indexing.output_map import DimensionMap from zarr_indexing.reader import Reader, basic_reader from zarr_indexing.testing.strategies import ( basic_selections, @@ -68,6 +68,7 @@ def make_source(self, data): if TYPE_CHECKING: from zarr_indexing.boundary import SelectionMode + from zarr_indexing.transform import IndexTransform __all__ = [ "DEFAULT_DATA", @@ -389,23 +390,27 @@ def box_parts_lower_to_basic_selections(self) -> None: try: source_selection = part.source_selection except ValueError: - # Refusal is the documented answer for a query part, and for - # the one box shape with no basic spelling: an axis restored - # by repetition, reached by gathering a collapsed constant - # with duplicates. Either way the properties must agree. - assert not part.view.is_box or any( - isinstance(m, ConstantMap) for m in part.view.transform.output - ), f"a plain box part refused to lower: {self.chain}" - for attribute in ("source_selection", "chunk_local_selection"): - lowered = True - try: - getattr(part, attribute) - except ValueError: - lowered = False - assert not lowered, ( - f"{attribute} lowered a part whose paired selection " - f"refused to: {self.chain}" - ) + # Refusal is the documented answer for a query part and for + # the two degenerate boxes: an axis restored by repetition + # (gathering a collapsed constant with duplicates) and an axis + # broadcast from one cell. Every other box must lower — an + # integer-indexed one included, which is the shape a laxer + # check would quietly excuse. + assert not part.view.is_box or _is_degenerate_box(part.view.transform), ( + f"a plain box part refused to lower: {self.chain}" + ) + # The paired spellings lower the same maps, so refusing one + # and answering the other would leave a consumer holding a + # cell selection that no source selection matches. + cell_lowered = True + try: + _ = part.chunk_local_selection + except ValueError: + cell_lowered = False + assert not cell_lowered, ( + f"chunk_local_selection lowered a part whose source_selection " + f"refused to: {self.chain}" + ) continue expected = self.model[part.out_selection] slab = data[source_selection] @@ -425,6 +430,20 @@ def box_parts_lower_to_basic_selections(self) -> None: ) +def _is_degenerate_box(transform: IndexTransform) -> bool: + """Whether `transform` is one of the boxes with no basic-selection spelling. + + Either an axis no output map reads but that a length-1 slice or a newaxis + cannot produce because its extent is not 1 — an axis restored by + repetition — or an axis broadcast from a single source cell by a stride-0 + map. Both name more values than the cells they read, which no slab does. + """ + referenced = {m.input_dimension for m in transform.output if isinstance(m, DimensionMap)} + return any( + axis not in referenced and extent != 1 for axis, extent in enumerate(transform.domain.shape) + ) or any(isinstance(m, DimensionMap) and m.stride == 0 for m in transform.output) + + def _reader_set(view: LazyArray, declared: Sequence[Reader] | None) -> tuple[Reader, ...]: """The readers `choose_reader` draws from, `basic_reader` always among them.""" readers = list(declared) if declared is not None else [view.reader] diff --git a/packages/zarr-indexing/tests/test_doc_examples.py b/packages/zarr-indexing/tests/test_doc_examples.py index 240bd48531..561a2fdd3a 100644 --- a/packages/zarr-indexing/tests/test_doc_examples.py +++ b/packages/zarr-indexing/tests/test_doc_examples.py @@ -22,6 +22,7 @@ from __future__ import annotations +import ast import re import runpy import subprocess @@ -131,11 +132,28 @@ def test_documentation_example_executes(example: Path) -> None: runpy.run_path(str(example), run_name="__main__") +def _imported_modules(script: Path) -> tuple[str, ...]: + """The top-level modules `script` imports, in source order. + + Examples declare their own dependencies (they are PEP 723 scripts run + outside this environment), and the suite is expected to run without the + optional ones installed. Reading the imports keeps the skip rule tied to + what a script actually needs rather than to what its name suggests. + """ + modules: dict[str, None] = {} + for node in ast.walk(ast.parse(script.read_text())): + if isinstance(node, ast.Import): + modules.update((alias.name.split(".")[0], None) for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: + modules[node.module.split(".")[0]] = None + return tuple(modules) + + @pytest.mark.parametrize("script", CLI_EXAMPLES, ids=lambda path: path.stem) def test_cli_example_runs_as_a_subprocess(script: Path) -> None: """The CLI examples exit 0 when run the way their READMEs instruct.""" - if "dask" in script.stem: - pytest.importorskip("dask.array") + for module in _imported_modules(script): + pytest.importorskip(module) completed = subprocess.run( [sys.executable, str(script)], capture_output=True, From 02e4e30aabd393491e77fe655f6e916ca7dd106c Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 17:07:29 +0200 Subject: [PATCH 07/16] refactor(zarr-indexing): one answer for the coordinates a DimensionMap reaches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `as_basic_selection` computed the first and last coordinate a map addresses over its domain with the same four lines the reader's two slab-pushers each carry, so planning a read through the public lowering and executing one through `BasicReader` asked the same question of two implementations. A change to either — a clamping rule, an overflow guard — would have silently moved the planned request away from the executed read, with both files' doctests still passing. `DimensionMap.endpoints` is now that question, asked once. It returns `None` for an empty interval, which is the case with no coordinate to name and the reason the pair was never just two `checked_affine` calls. The Returns section no longer promises a `None` per unread axis, since a constant due in that position stands in for one (roborev, 413a520). Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/output_map.py | 24 +++++++++++++++++++ .../zarr-indexing/src/zarr_indexing/reader.py | 12 +++++----- .../src/zarr_indexing/transform.py | 17 ++++++------- 3 files changed, 39 insertions(+), 14 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py index 3ee7250efa..8880127d5b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/output_map.py +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -130,6 +130,30 @@ class DimensionMap: stride: int = 1 """The output-coordinate step per unit input step; negative walks backward, zero repeats `offset`.""" + def endpoints(self, inclusive_min: int, exclusive_max: int) -> tuple[int, int] | None: + """The first and last output coordinates reached over a domain interval. + + `None` for an empty interval, which reaches no coordinate at all — the + one case with nothing to name, and the reason this is not two calls to + `checked_affine` at each site that needs the pair. A descending map + reports them in the order it walks them, so `first` may exceed `last`. + + Examples + -------- + >>> DimensionMap(input_dimension=0, offset=2, stride=3).endpoints(0, 4) + (2, 11) + >>> DimensionMap(input_dimension=0, offset=9, stride=-2).endpoints(0, 3) + (9, 5) + >>> DimensionMap(input_dimension=0).endpoints(4, 4) is None + True + """ + if exclusive_max <= inclusive_min: + return None + return ( + checked_affine(self.offset, self.stride, inclusive_min), + checked_affine(self.offset, self.stride, exclusive_max - 1), + ) + def to_json(self) -> OutputIndexMapJSON: """Convert to the canonical wire form: the `single_input_dimension` map. diff --git a/packages/zarr-indexing/src/zarr_indexing/reader.py b/packages/zarr-indexing/src/zarr_indexing/reader.py index 8d47d51d21..dad2296591 100644 --- a/packages/zarr-indexing/src/zarr_indexing/reader.py +++ b/packages/zarr-indexing/src/zarr_indexing/reader.py @@ -495,10 +495,10 @@ def _push_slice_for_dimension_map( d = m.input_dimension lo = transform.domain.inclusive_min[d] hi = max(transform.domain.exclusive_max[d], lo) - if hi == lo: + endpoints = m.endpoints(lo, hi) + if endpoints is None: return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) - first = checked_affine(m.offset, m.stride, lo) - last = checked_affine(m.offset, m.stride, hi - 1) + first, last = endpoints if m.stride > 0: return ( slice(first, last + 1, m.stride), @@ -531,15 +531,15 @@ def _push_unit_slice_for_dimension_map( d = m.input_dimension lo = transform.domain.inclusive_min[d] hi = max(transform.domain.exclusive_max[d], lo) - if hi == lo: + endpoints = m.endpoints(lo, hi) + if endpoints is None: return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) - first = checked_affine(m.offset, m.stride, lo) + first, last = endpoints if m.stride == 0: return ( slice(first, first + 1, 1), DimensionMap(input_dimension=d, offset=0, stride=0), ) - last = checked_affine(m.offset, m.stride, hi - 1) origin = min(first, last) return ( slice(origin, max(first, last) + 1, 1), diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index a67bb6e523..7123c94d0b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -452,8 +452,8 @@ def as_basic_selection(self) -> tuple[int | slice | None, ...]: """Lower a box transform to the NumPy basic selection it describes. Returns one selector per output dimension, plus a `None` for each - domain axis no output map reads, such that - `source[transform.as_basic_selection()]` under NumPy indexing + domain axis no output map reads and no constant stands in for, such + that `source[transform.as_basic_selection()]` under NumPy indexing semantics reads exactly the cells this transform addresses, in domain order: the result has shape `domain.shape` exactly, and its element at position `p` holds the source value at @@ -482,7 +482,8 @@ def as_basic_selection(self) -> tuple[int | slice | None, ...]: ------- tuple of int or slice or None One selector per output dimension, interleaved with a `None` for - each domain axis no output map reads. + each domain axis that neither an output map reads nor a constant + stands in for. Raises ------ @@ -600,13 +601,13 @@ def fabricate_axes_before(axis: int) -> None: "transposes or repeats axes" ) next_axis = d + 1 - lo = self.domain.inclusive_min[d] - hi = self.domain.exclusive_max[d] - if hi <= lo: + endpoints = output_map.endpoints( + self.domain.inclusive_min[d], self.domain.exclusive_max[d] + ) + if endpoints is None: selection.append(slice(0, 0, 1)) continue - first = checked_affine(output_map.offset, output_map.stride, lo) - last = checked_affine(output_map.offset, output_map.stride, hi - 1) + first, last = endpoints if min(first, last) < 0: raise ValueError( f"cannot lower to a basic selection: output[{output_dimension}] " From ff805de1872aab220d04a0af2ea30308e7e50eee Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 17:09:22 +0200 Subject: [PATCH 08/16] test(zarr-indexing): skip on the exact module an example imports Truncating `dask.array` to `dask` asked a coarser question than the guard it replaced: a distribution can be installed while the submodule an example needs is not importable, and the example would then run and fail rather than skip. Keeping the dotted path restores that precision (roborev, a527676). Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-indexing/tests/test_doc_examples.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/packages/zarr-indexing/tests/test_doc_examples.py b/packages/zarr-indexing/tests/test_doc_examples.py index 561a2fdd3a..95fd80d84f 100644 --- a/packages/zarr-indexing/tests/test_doc_examples.py +++ b/packages/zarr-indexing/tests/test_doc_examples.py @@ -133,19 +133,24 @@ def test_documentation_example_executes(example: Path) -> None: def _imported_modules(script: Path) -> tuple[str, ...]: - """The top-level modules `script` imports, in source order. + """The modules `script` imports, dotted paths and all, in source order. Examples declare their own dependencies (they are PEP 723 scripts run outside this environment), and the suite is expected to run without the optional ones installed. Reading the imports keeps the skip rule tied to what a script actually needs rather than to what its name suggests. + + Submodules are kept whole: a distribution can be installed while the + submodule an example needs is not importable — `dask` without + `dask.array`, whose extra dependencies are their own install — and only + the path the script actually imports answers that. """ modules: dict[str, None] = {} for node in ast.walk(ast.parse(script.read_text())): if isinstance(node, ast.Import): - modules.update((alias.name.split(".")[0], None) for alias in node.names) + modules.update((alias.name, None) for alias in node.names) elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: - modules[node.module.split(".")[0]] = None + modules[node.module] = None return tuple(modules) From d1fb56f05ff27c02d79d619c5652134babcc10b7 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 18:29:09 +0200 Subject: [PATCH 09/16] test(zarr-indexing): reach the states that hid this PR's defects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The state machine explored hard in the region where the code was already right. Its view was always the one built from the source — nothing ever set it to a part's view — so a view with a window, and every bug that needs one, sat outside the reachable state space. Its selections never carried `None`, so an axis no source axis backs was never drawn. And every reader it drew answers from `context.transform` alone, so a projection describing a different read than the transform beside it could not change a single value. Three additions, one per gap: - `descend_into_a_part` follows a part's view and boxes it again, reaching the part-of-a-part whose cell coordinates and transform count from different origins. The model follows by the documented assembly. - `fabricates_an_axis` draws a basic selection carrying `None`. - `ProjectionReader` reads each part out of the cell its `chunk_domain` names, the way a decoded-chunk cache does, and joins `basic_reader` as a reader every machine draws from — so both halves of a `ReadContext` now face the same NumPy model. Against this PR's first two commits each addition fails: the descent alone trips `box_parts_lower_to_basic_selections` with no newaxis involved, and resolving a nested part view returns the model's values through `basic_reader` while disagreeing through `ProjectionReader` — the value-level statement of the bug that check could not previously make. 1500 examples over 14 steps pass on three sources here. An explicit per-axis partitioning describes the source's extents, so it is skipped where a descent has narrowed the base it would have to sum to. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+stateful-part-descent.feature.md | 21 +++ .../src/zarr_indexing/testing/__init__.py | 6 + .../src/zarr_indexing/testing/stateful.py | 165 +++++++++++++++++- .../src/zarr_indexing/testing/strategies.py | 22 +++ .../tests/test_lazy_array_stateful.py | 16 +- 5 files changed, 219 insertions(+), 11 deletions(-) create mode 100644 packages/zarr-indexing/changes/+stateful-part-descent.feature.md diff --git a/packages/zarr-indexing/changes/+stateful-part-descent.feature.md b/packages/zarr-indexing/changes/+stateful-part-descent.feature.md new file mode 100644 index 0000000000..69c1b362bc --- /dev/null +++ b/packages/zarr-indexing/changes/+stateful-part-descent.feature.md @@ -0,0 +1,21 @@ +`ChainedIndexingStateMachine` gained two rules and a reader, each covering a +state it could not previously reach: + +- `descend_into_a_part` continues the chain from one part's view and boxes it + again. A part's view is documented as resolvable on its own, and it bases + its boxes on its own window rather than on the source, so following one + reaches the part-of-a-part — where a projection's cell coordinates and its + view's transform are counted from different origins. +- `fabricates_an_axis` draws a basic selection carrying `None` + (`newaxis_selections`, also exported), which adds a domain axis no source + axis backs. +- `ProjectionReader` reads each part by fetching the cell `chunk_domain` names + and gathering it with `chunk_transform`, the way a decoded-chunk cache does. + It joins `basic_reader` as a reader every machine draws from, so a + `ReadContext` whose projection describes a different read than its transform + is caught by the same NumPy model as everything else — no reader that + ignores the projection can see that. + +Subclasses inherit all three. A partitioning declared as explicit per-axis +sizes describes the source's extents, so it is skipped where a descent has +narrowed the base it would have to sum to. diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py index e98d051900..aa87784d40 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/__init__.py @@ -27,14 +27,17 @@ def make_source(self, data): DEFAULT_PARTITIONINGS, DEFAULT_SETTINGS, ChainedIndexingStateMachine, + ProjectionReader, apply_selection, outer_selection, + projection_reader, repartition, state_machine_test, ) from zarr_indexing.testing.strategies import ( basic_selections, masks, + newaxis_selections, orthogonal_selections, slice_selections, vectorized_selections, @@ -45,11 +48,14 @@ def make_source(self, data): "DEFAULT_PARTITIONINGS", "DEFAULT_SETTINGS", "ChainedIndexingStateMachine", + "ProjectionReader", "apply_selection", "basic_selections", "masks", + "newaxis_selections", "orthogonal_selections", "outer_selection", + "projection_reader", "repartition", "slice_selections", "state_machine_test", diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 5875db4662..d189a4f245 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -38,9 +38,23 @@ def make_source(self, data): The `choose_reader` rule draws a reader and applies it to the view, so the execution strategy becomes part of the chain. Every reader listed by a subclass -must preserve the NumPy model for its source. The universal `basic_reader` is -always exercised, even when a subclass lists only specialized readers; with no -declared readers it is the sole strategy drawn. +must preserve the NumPy model for its source. Two universal readers are always +exercised, even when a subclass lists only specialized ones: `basic_reader`, +which answers from the read's transform, and `ProjectionReader`, which answers +from its projection instead — fetching whole grid cells and gathering out of +them, the way a decoded-chunk cache does. Between them both halves of a +`ReadContext` are checked against the model. A projection describing a +different read than the transform it arrives with is invisible to every reader +that ignores it, which is how such a pairing can be wrong while every value +assertion passes. + +`descend_into_a_part` continues the chain from one part's view and boxes it +again. A part's view is documented as resolvable on its own, so it must satisfy +everything a view satisfies; and because it bases its boxes on its own window +rather than on the source, following one reaches the part-of-a-part — where a +projection's cell coordinates and its view's transform are counted from +different origins, and where a cell named relative to the wrong one still reads +plausible values. Requires the `testing` extra (`pip install zarr-indexing[testing]`). """ @@ -58,9 +72,10 @@ def make_source(self, data): from zarr_indexing.lazy_array import LazyArray from zarr_indexing.output_map import DimensionMap -from zarr_indexing.reader import Reader, basic_reader +from zarr_indexing.reader import ReadContext, Reader, basic_reader from zarr_indexing.testing.strategies import ( basic_selections, + newaxis_selections, orthogonal_selections, slice_selections, vectorized_selections, @@ -68,6 +83,7 @@ def make_source(self, data): if TYPE_CHECKING: from zarr_indexing.boundary import SelectionMode + from zarr_indexing.domain import IndexDomain from zarr_indexing.transform import IndexTransform __all__ = [ @@ -75,8 +91,10 @@ def make_source(self, data): "DEFAULT_PARTITIONINGS", "DEFAULT_SETTINGS", "ChainedIndexingStateMachine", + "ProjectionReader", "apply_selection", "outer_selection", + "projection_reader", "repartition", "state_machine_test", ] @@ -210,8 +228,10 @@ class ChainedIndexingStateMachine(RuleBasedStateMachine): mid-chain does, and spends the whole step budget on indexing. readers Execution strategies `choose_reader` may draw. Every listed reader must - preserve the model for the source. `basic_reader` is always included; - `None` means the reader already carried by the constructed view. + preserve the model for the source. `basic_reader` and + `projection_reader` are always included, since between them they check + both halves of a `ReadContext`; `None` means the reader already carried + by the constructed view. """ data: ClassVar[Any] = DEFAULT_DATA @@ -251,6 +271,26 @@ def _indexable(self) -> bool: """ return self.model.ndim > 0 and self.model.size > 0 + def _fitting_partitionings(self) -> list[Any]: + """The declared partitionings that describe this view's own base. + + `partitionings` is written against the source's shape, and a part + views its own window as its base, so the explicit per-axis spelling — + whose sizes must sum to each extent — does not survive a descent. The + uniform spellings do: a box wider than the extent is one box. + """ + base = self.view.base_shape + return [ + boxes + for boxes in type(self).partitionings + if boxes is None + or not any(isinstance(entry, Sequence) for entry in boxes) + or ( + len(boxes) == len(base) + and all(sum(sizes) == extent for sizes, extent in zip(boxes, base, strict=True)) + ) + ] + def _step(self, mode: SelectionMode, selection: tuple[Any, ...]) -> None: self.chain.append((mode, selection)) self.model = apply_selection(self.model, selection, mode) @@ -297,6 +337,47 @@ def slices_only(self, data: st.DataObject) -> None: """ self._step("orthogonal", data.draw(slice_selections(self.model.shape))) + @precondition(lambda self: self._indexable()) + @rule(data=st.data()) + def fabricates_an_axis(self, data: st.DataObject) -> None: + """A basic step carrying `None`, which adds an axis no source axis backs. + + Its own rule for the reason `slices_only` is: a domain axis that no + output map reads is a distinct shape for everything downstream — the + lowering to a basic selection most of all, which has to produce that + axis without reading one — and folding the spelling into `basic` would + leave it drawn a fraction of the time. + """ + self._step("basic", data.draw(newaxis_selections(self.model.shape))) + + @precondition(lambda self: self.model.size > 0) + @rule(data=st.data()) + def descend_into_a_part(self, data: st.DataObject) -> None: + """Continue the chain from one part's view, then re-box it. + + `Partition.view` is documented as resolvable on its own — in another + thread, in another order, or not at all — so everything this machine + checks of a view has to hold of one. Following a part also re-bases + the partitioning: the new view boxes its own window rather than the + source, and boxing it again reaches the part-of-a-part, whose cell + coordinates and whose transform are counted from different origins. + Nothing else here reaches that state, which is why the re-boxing is + part of this rule rather than left to `repartition` — that one waits + for a chain to run out of axes, and the state is most interesting + while there are axes left. + + The model follows by the documented assembly: a part's values are the + model's at the part's `out_selection`. + """ + parts = list(self.view.parts()) + part = parts[data.draw(st.integers(0, len(parts) - 1))] + self.model = self.model[part.out_selection] + self.view = part.view + self.chain.append(("part", part.base_coords)) + boxes = data.draw(st.sampled_from(self._fitting_partitionings())) + self.view = repartition(self.view, boxes) + self.chain.append(("parts", boxes)) + @rule(data=st.data()) def choose_reader(self, data: st.DataObject) -> None: """Read the rest of the chain through another conforming strategy.""" @@ -315,7 +396,7 @@ def repartition(self, data: st.DataObject) -> None: a rank-0 view read through every partitioning is exactly the state a collapsed correlated selection reaches. """ - parts = data.draw(st.sampled_from(list(type(self).partitionings))) + parts = data.draw(st.sampled_from(self._fitting_partitionings())) self.view = repartition(self.view, parts) self.chain.append(("parts", parts)) @@ -444,11 +525,79 @@ def _is_degenerate_box(transform: IndexTransform) -> bool: ) or any(isinstance(m, DimensionMap) and m.stride == 0 for m in transform.output) +def _domain_points(domain: IndexDomain) -> np.ndarray[Any, np.dtype[np.intp]]: + """Every coordinate in `domain`, one row per point, in C order.""" + axes = [ + np.arange(lo, hi, dtype=np.intp) + for lo, hi in zip(domain.inclusive_min, domain.exclusive_max, strict=True) + ] + if not axes: + # A rank-zero domain holds exactly one point, which has no coordinates. + return np.zeros((1, 0), dtype=np.intp) + mesh = np.meshgrid(*axes, indexing="ij") + return np.stack([axis.ravel() for axis in mesh], axis=1) + + +class ProjectionReader: + """Read each part by gathering it out of the whole grid cell it lives in. + + Every other reader answers from `context.transform` alone, so a projection + describing a different read than the transform it arrives with cannot + change what they return, and no assertion about values can see it. This + one fetches the cell `projection.chunk_domain` names and gathers the + request out of it with `projection.chunk_transform` — what a decoded-chunk + cache does — which puts the pairing of the two under the same NumPy model + as everything else. + + Deliberately naive: whole cells, gathered pointwise. It is a contract + check, not a strategy to copy for performance. + """ + + def read_into(self, source: Any, context: ReadContext, out: np.ndarray[Any, Any], /) -> None: + """Fill `out` through the projection rather than the transform.""" + projection = context.projection + if projection is None: + # An unpartitioned read is paired with no cell, so there is + # nothing here that `basic_reader` does not already check. + basic_reader.read_into(source, context, out) + return + domain = projection.chunk_domain + cell = np.asanyarray( + source[ + tuple( + slice(lo, hi) + for lo, hi in zip(domain.inclusive_min, domain.exclusive_max, strict=True) + ) + ] + ) + chunk_points = projection.chunk_transform.apply_many( + _domain_points(projection.chunk_transform.domain) + ) + values = cell[tuple(chunk_points.T)] + # The projection's synthetic domain and the buffer's domain enumerate + # the same cells in the same order and at the same shape, which is the + # pairing `ChunkProjection` documents; a reshape that fails has caught + # that promise breaking. + out[...] = values.reshape(out.shape) + + +projection_reader = ProjectionReader() +"""The shared `ProjectionReader`; readers are stateless, so one serves every view.""" + + def _reader_set(view: LazyArray, declared: Sequence[Reader] | None) -> tuple[Reader, ...]: - """The readers `choose_reader` draws from, `basic_reader` always among them.""" + """The readers `choose_reader` draws from. + + `basic_reader` and `projection_reader` are always among them: neither + knows anything about a particular source beyond the slab reads every + source already serves, and between them they check both halves of a + `ReadContext` against the model. + """ readers = list(declared) if declared is not None else [view.reader] if all(reader is not basic_reader for reader in readers): readers.insert(0, basic_reader) + if all(not isinstance(reader, ProjectionReader) for reader in readers): + readers.insert(1, projection_reader) unique: list[Reader] = [] for reader in readers: if all(reader is not existing for existing in unique): diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py index 63a3d351a9..4591a09852 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/strategies.py @@ -37,6 +37,7 @@ def test_my_array_slices_like_numpy(selection): "basic_selections", "empty_masks", "masks", + "newaxis_selections", "orthogonal_selections", "slice_selections", "vectorized_selections", @@ -138,6 +139,27 @@ def basic_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ... return _entries(shape, _basic_entry) +@st.composite +def newaxis_selections(draw: st.DrawFn, shape: tuple[int, ...]) -> tuple[Any, ...]: + """Basic selections carrying `None`, the selector that fabricates an axis. + + NumPy's `None` belongs to basic indexing and `LazyArray` admits it, but it + is the one basic selector that adds an axis of the view no axis of the + source backs — a domain axis no output map reads, which everything + downstream must then produce out of nothing. `basic_selections` promises + one entry per axis, so the spelling that breaks that count is drawn here + instead of widening it. + + A selection carries one or two of them, in any position: leading, between + two axes, or trailing, each of which lands the fabricated axis somewhere + different in the result. + """ + entries = list(draw(basic_selections(shape))) + for _ in range(draw(st.integers(1, 2))): + entries.insert(draw(st.integers(0, len(entries))), None) + return tuple(entries) + + def orthogonal_selections(shape: tuple[int, ...]) -> st.SearchStrategy[tuple[Any, ...]]: """Orthogonal (`oindex`) selections: an outer product of per-axis choices. diff --git a/packages/zarr-indexing/tests/test_lazy_array_stateful.py b/packages/zarr-indexing/tests/test_lazy_array_stateful.py index 4f5dadc1e4..fe0fa97122 100644 --- a/packages/zarr-indexing/tests/test_lazy_array_stateful.py +++ b/packages/zarr-indexing/tests/test_lazy_array_stateful.py @@ -60,9 +60,19 @@ def test_reader_set_deduplicates_by_identity_without_hashing() -> None: readers = stateful._reader_set(LazyArray(np.arange(3)), (first, first, second)) assert readers[0] is basic_reader - assert readers[1] is first - assert readers[2] is second - assert len(readers) == 3 + assert readers[1] is stateful.projection_reader + assert readers[2] is first + assert readers[3] is second + assert len(readers) == 4 + + +def test_reader_set_keeps_a_declared_projection_reader() -> None: + """The universal pair is added, not duplicated over a subclass's own.""" + declared = stateful.ProjectionReader() + + readers = stateful._reader_set(LazyArray(np.arange(3)), (declared,)) + + assert readers == (basic_reader, declared) class OneDimensionalIndexing(ChainedIndexingStateMachine): From 4b397808529ab867c13ba7dc1a11de9678b44c8f Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 18:43:57 +0200 Subject: [PATCH 10/16] test(zarr-indexing): leave a descent something to draw from MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `partitionings` describes the source's extents, so a subclass declaring nothing but explicit per-axis sizes has none that fit a part's narrowed base — and `sampled_from` on the resulting empty list raises from inside Hypothesis rather than saying so. Falling back to leaving the view boxed as it already is states the same thing and keeps the run going (roborev, d1fb56f). Assisted-by: ClaudeCode:claude-fable-5 --- .../src/zarr_indexing/testing/stateful.py | 6 +++++- .../tests/test_lazy_array_stateful.py | 15 +++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index d189a4f245..77432cf63f 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -280,7 +280,7 @@ def _fitting_partitionings(self) -> list[Any]: uniform spellings do: a box wider than the extent is one box. """ base = self.view.base_shape - return [ + fitting = [ boxes for boxes in type(self).partitionings if boxes is None @@ -290,6 +290,10 @@ def _fitting_partitionings(self) -> list[Any]: and all(sum(sizes) == extent for sizes, extent in zip(boxes, base, strict=True)) ) ] + # A subclass may declare nothing but per-axis sizes, none of which can + # fit a narrowed base. Leaving the view boxed as it already is says so, + # where drawing from nothing would raise from inside Hypothesis. + return fitting or [None] def _step(self, mode: SelectionMode, selection: tuple[Any, ...]) -> None: self.chain.append((mode, selection)) diff --git a/packages/zarr-indexing/tests/test_lazy_array_stateful.py b/packages/zarr-indexing/tests/test_lazy_array_stateful.py index fe0fa97122..6ab4a25740 100644 --- a/packages/zarr-indexing/tests/test_lazy_array_stateful.py +++ b/packages/zarr-indexing/tests/test_lazy_array_stateful.py @@ -66,6 +66,21 @@ def test_reader_set_deduplicates_by_identity_without_hashing() -> None: assert len(readers) == 4 +def test_only_per_axis_partitionings_still_leave_a_descent_something_to_draw() -> None: + """A narrowed base can fit none of them, and drawing from nothing raises.""" + + class PerAxisOnly(ChainedIndexingStateMachine): + data = np.arange(30, dtype=np.int64) + partitionings: ClassVar[tuple[Any, ...]] = (((7, 8, 15),),) + + machine = PerAxisOnly() + assert machine._fitting_partitionings() == [((7, 8, 15),)] + machine.view = stateful.repartition(machine.view, ((7, 8, 15),)) + machine.view = next(machine.view.parts()).view + assert machine.view.base_shape == (7,) + assert machine._fitting_partitionings() == [None] + + def test_reader_set_keeps_a_declared_projection_reader() -> None: """The universal pair is added, not duplicated over a subclass's own.""" declared = stateful.ProjectionReader() From fd68127a6eccab8abb196fa9b58ef3dc5d9ee973 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 21:38:57 +0200 Subject: [PATCH 11/16] feat(zarr-indexing): give the lowering refusal a name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `as_basic_selection` is a partial function, and its refusal shared a type with everything else that can go wrong: the documented consumer fallback was `except ValueError`, which also catches a genuine defect in the lowering — silently degrading every part to the slow path with nothing reporting it. Refusals now raise `NoBasicSelectionError`, exported at the package root and subclassing `ValueError` so existing catch sites keep working. The asyncio example, the stateful invariant, and the error tests all catch the precise type now; the invariant in particular no longer excuses a bug that happens to raise `ValueError`. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+asyncio-example.doc.md | 3 +- .../+no-basic-selection-error.feature.md | 6 ++++ .../changes/+part-lowering.feature.md | 3 +- .../zarr-indexing/docs/guide/integrations.md | 3 +- .../examples/lazy_indexing_asyncio/README.md | 3 +- .../lazy_indexing_asyncio.py | 7 ++-- .../src/zarr_indexing/__init__.py | 7 +++- .../zarr-indexing/src/zarr_indexing/errors.py | 34 +++++++++++++++++-- .../src/zarr_indexing/lazy_array.py | 14 ++++---- .../src/zarr_indexing/testing/stateful.py | 5 +-- .../src/zarr_indexing/transform.py | 24 +++++++------ .../zarr-indexing/tests/test_lazy_array.py | 5 +-- .../zarr-indexing/tests/test_transform.py | 16 ++++----- 13 files changed, 92 insertions(+), 38 deletions(-) create mode 100644 packages/zarr-indexing/changes/+no-basic-selection-error.feature.md diff --git a/packages/zarr-indexing/changes/+asyncio-example.doc.md b/packages/zarr-indexing/changes/+asyncio-example.doc.md index 1c0fd483ff..e8a108d2eb 100644 --- a/packages/zarr-indexing/changes/+asyncio-example.doc.md +++ b/packages/zarr-indexing/changes/+asyncio-example.doc.md @@ -2,5 +2,6 @@ Added an asyncio integration example (`examples/lazy_indexing_asyncio/`), peer to the Dask example: `parts()` driven by `asyncio.gather` against `zarr.AsyncArray`, using `Partition.source_selection` for per-part fetches, a decoded-chunk cache keyed on each cell's `chunk_domain` origin and placed with -`chunk_local_selection`, and the `ValueError` fallback for query parts. The +`chunk_local_selection`, and the `NoBasicSelectionError` fallback for query +parts. The integrations guide gained a matching "Consumer-owned I/O" section. diff --git a/packages/zarr-indexing/changes/+no-basic-selection-error.feature.md b/packages/zarr-indexing/changes/+no-basic-selection-error.feature.md new file mode 100644 index 0000000000..0609c75910 --- /dev/null +++ b/packages/zarr-indexing/changes/+no-basic-selection-error.feature.md @@ -0,0 +1,6 @@ +`as_basic_selection` refusals now raise `NoBasicSelectionError`, a dedicated +`ValueError` subclass exported at the package root. The documented consumer +fallback (`except NoBasicSelectionError: part.view.result()`) can no longer +silently absorb a genuine defect in the lowering, which bare `except +ValueError` did. Existing catch sites keep working: the subclass is caught by +`except ValueError` unchanged. diff --git a/packages/zarr-indexing/changes/+part-lowering.feature.md b/packages/zarr-indexing/changes/+part-lowering.feature.md index 4ce8aa4120..a98f9320d5 100644 --- a/packages/zarr-indexing/changes/+part-lowering.feature.md +++ b/packages/zarr-indexing/changes/+part-lowering.feature.md @@ -6,7 +6,8 @@ async store, an HTTP range endpoint): integers and slices such that `source[selection]` reads exactly the transform's cells, at exactly its domain shape. A collapsed single-coordinate gather keeps its singleton axis through a length-1 slice; - queries, broadcasts, and transposed or repeated axes raise `ValueError` + queries, broadcasts, and transposed or repeated axes raise + `NoBasicSelectionError` (a `ValueError` subclass) instead of guessing a slab. - `Partition.source_selection` is that lowering of a part's global read, so an async consumer's whole loop is diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index f6972abdaf..badbb6ae1d 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -89,7 +89,8 @@ chunks. `projection.chunk_domain` locates that cell in the source, so it is what to fetch and a sound cache key; `base_coords` counts cells of whatever base the view partitions, so it keys a cache only alongside the grid that produced it. A query part (an `oindex`/`vindex` gather) has -no slab spelling and raises `ValueError`, so mixed consumers fall back to +no slab spelling and raises `NoBasicSelectionError` (a `ValueError` +subclass), so mixed consumers fall back to `part.view.result()` for those parts. The [asyncio example](../examples/lazy_indexing_asyncio.md) drives all three loops — gather-per-part, decoded-chunk cache, and the query fallback — with diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md index 17e4aad65f..749b4833d4 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md @@ -15,7 +15,8 @@ The example shows how to: fetch the cell (`projection.chunk_domain`) once, then serve every overlapping view from the cache with `part.chunk_local_selection` - Fall back for query partitions (`oindex`/`vindex`/mask selections), whose - gathers have no single-slab spelling: `source_selection` raises `ValueError` + gathers have no single-slab spelling: `source_selection` raises + `NoBasicSelectionError` and the part resolves through its own `part.view.result()` instead The async side only needs one method — `async def getitem(selection)` accepting diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py index 6733241734..781a7a0c9e 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py @@ -22,7 +22,7 @@ import zarr import zarr.api.asynchronous -from zarr_indexing import LazyArray, Partition +from zarr_indexing import LazyArray, NoBasicSelectionError, Partition class AsyncSource(Protocol): @@ -154,6 +154,9 @@ def test_query_parts_fall_back(store: dict[str, Any], source: zarr.Array) -> Non `source_selection` raises `ValueError` instead of guessing one. A consumer mixing selection kinds catches that and resolves the part through `part.view.result()`, which reads through the wrapped array's reader. + Catching the dedicated subclass — not bare `ValueError` — keeps a genuine + defect in the lowering loud instead of silently degrading every part to + the fallback path. """ view = LazyArray(source).lazy.oindex[[30, 2, 2], 4:10] @@ -165,7 +168,7 @@ async def scenario() -> np.ndarray: async def resolve(part: Partition) -> tuple[Partition, np.ndarray]: try: selection = part.source_selection - except ValueError: + except NoBasicSelectionError: # The gathered axis needs a lookup, not a slab: read this part # through the wrapper (synchronously here; a real consumer # might push it to a thread, or fetch the part's bounding box diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 9acfd28a21..42a41d6f90 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -33,7 +33,11 @@ plan_chunks, ) from zarr_indexing.domain import IndexDomain -from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError +from zarr_indexing.errors import ( + BoundsCheckError, + NoBasicSelectionError, + VindexInvalidSelectionError, +) from zarr_indexing.grid import ( ChunkGrid, ChunkSpec, @@ -95,6 +99,7 @@ "IndexTransformJSON", "LazyArray", "NdselError", + "NoBasicSelectionError", "NumPyReader", "OutputIndexMap", "OutputIndexMapJSON", diff --git a/packages/zarr-indexing/src/zarr_indexing/errors.py b/packages/zarr-indexing/src/zarr_indexing/errors.py index efd3b1ecd6..29c29313bc 100644 --- a/packages/zarr-indexing/src/zarr_indexing/errors.py +++ b/packages/zarr-indexing/src/zarr_indexing/errors.py @@ -1,7 +1,8 @@ -"""Canonical index-error types raised by the transform algebra. +"""Canonical error types raised by the transform algebra. -Both subclass the built-in `IndexError`, so an `except IndexError` catch site -keeps working unchanged whichever library raised. +The index errors subclass the built-in `IndexError` and the lowering refusal +subclasses `ValueError`, so a catch site written against the built-in keeps +working unchanged whichever library raised. `zarr.errors` defines classes of the same names, and they are *not* these objects: `zarr.errors.BoundsCheckError is BoundsCheckError` is false. Catching @@ -13,10 +14,37 @@ __all__ = [ "BoundsCheckError", + "NoBasicSelectionError", "VindexInvalidSelectionError", ] +class NoBasicSelectionError(ValueError): + """The transform has no basic-selection spelling. + + `IndexTransform.as_basic_selection` is a partial function: integers, + slices, and `None` spell exactly the reads that take one source axis per + result axis, in increasing order, each an arithmetic progression of + nonnegative coordinates. A transform outside that set — a query's lookup + table, a stride-0 broadcast, transposed or repeated axes, a negative + coordinate — is refused with this error rather than approximated. + + Distinct from the `ValueError` a genuine defect would raise, so a + consumer's fallback (`except NoBasicSelectionError: use the reader`) + cannot silently absorb a bug in the lowering itself. Subclasses + `ValueError`, so catch sites written before this class existed keep + working. + + Examples + -------- + >>> from zarr_indexing import IndexTransform + >>> IndexTransform.from_shape((10,)).oindex[[3, 1, 1]].as_basic_selection() + Traceback (most recent call last): + ... + zarr_indexing.errors.NoBasicSelectionError: ... + """ + + class VindexInvalidSelectionError(IndexError): """A wrapper `vindex` selection contained a slice. diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 0722e0b76f..ab5d835d5a 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -559,13 +559,15 @@ def source_selection(self) -> tuple[int | slice | None, ...]: out[part.out_selection] = await source.getitem(part.source_selection) ``` - A query part has no slab to request and raises `ValueError`, as do the - two degenerate boxes with no basic-selection spelling: an axis - restored by repetition (reached by gathering a collapsed constant with - duplicates) and an axis broadcast from one cell. `view.is_box` is - therefore necessary but not sufficient — catching `ValueError` is what + A query part has no slab to request and raises + [`NoBasicSelectionError`][zarr_indexing.errors.NoBasicSelectionError], + as do the two degenerate boxes with no basic-selection spelling: an + axis restored by repetition (reached by gathering a collapsed constant + with duplicates) and an axis broadcast from one cell. `view.is_box` is + therefore necessary but not sufficient — catching the error is what decides it — and a consumer mixing selection kinds falls back to - `view.result()` for the parts that refuse. A reversing view lowers to + `view.result()` for the parts that refuse. The dedicated subclass of + `ValueError` keeps that fallback from absorbing a genuine defect. A reversing view lowers to a negative-step slice and a fabricated axis to `None`, which a backend narrower than NumPy may not accept. diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 77432cf63f..98d179a00a 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -70,6 +70,7 @@ def make_source(self, data): from hypothesis import strategies as st from hypothesis.stateful import RuleBasedStateMachine, initialize, invariant, precondition, rule +from zarr_indexing.errors import NoBasicSelectionError from zarr_indexing.lazy_array import LazyArray from zarr_indexing.output_map import DimensionMap from zarr_indexing.reader import ReadContext, Reader, basic_reader @@ -474,7 +475,7 @@ def box_parts_lower_to_basic_selections(self) -> None: for part in self.view.parts(): try: source_selection = part.source_selection - except ValueError: + except NoBasicSelectionError: # Refusal is the documented answer for a query part and for # the two degenerate boxes: an axis restored by repetition # (gathering a collapsed constant with duplicates) and an axis @@ -490,7 +491,7 @@ def box_parts_lower_to_basic_selections(self) -> None: cell_lowered = True try: _ = part.chunk_local_selection - except ValueError: + except NoBasicSelectionError: cell_lowered = False assert not cell_lowered, ( f"chunk_local_selection lowered a part whose source_selection " diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 7123c94d0b..56148efc63 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -37,7 +37,11 @@ from zarr_indexing._selector import as_scalar_index, require_index from zarr_indexing.boundary import validate_advanced_selection from zarr_indexing.domain import IndexDomain -from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError +from zarr_indexing.errors import ( + BoundsCheckError, + NoBasicSelectionError, + VindexInvalidSelectionError, +) from zarr_indexing.output_map import ( ArrayMap, ConstantMap, @@ -487,7 +491,7 @@ def as_basic_selection(self) -> tuple[int | slice | None, ...]: Raises ------ - ValueError + NoBasicSelectionError If the transform cannot be expressed as a basic selection: an output map is an `ArrayMap` (a query is a lookup table, not a slab), a `DimensionMap` has stride 0 (a broadcast repeats one @@ -526,21 +530,21 @@ def as_basic_selection(self) -> tuple[int | slice | None, ...]: >>> IndexTransform.from_shape((10,)).oindex[[3, 1, 1]].as_basic_selection() Traceback (most recent call last): ... - ValueError: cannot lower to a basic selection: output[0] is an ArrayMap; \ -a query selection is a lookup table, not a slab + zarr_indexing.errors.NoBasicSelectionError: cannot lower to a basic \ +selection: output[0] is an ArrayMap; a query selection is a lookup table, not a slab """ for output_dimension, output_map in enumerate(self.output): # Named before the axis bookkeeping below, which would otherwise # blame a query's gathered axis for being unreferenced. if isinstance(output_map, ArrayMap): - raise ValueError( # noqa: TRY004 - valid map, no basic spelling + raise NoBasicSelectionError( f"cannot lower to a basic selection: output[{output_dimension}] " "is an ArrayMap; a query selection is a lookup table, not a slab" ) referenced = {m.input_dimension for m in self.output if isinstance(m, DimensionMap)} for axis, extent in enumerate(self.domain.shape): if axis not in referenced and extent != 1: - raise ValueError( + raise NoBasicSelectionError( f"cannot lower to a basic selection: no output map references " f"input dimension {axis}, whose extent is {extent}; only a " "singleton axis can be produced without reading a source axis, " @@ -570,7 +574,7 @@ def fabricate_axes_before(axis: int) -> None: for output_dimension, output_map in enumerate(self.output): if isinstance(output_map, ConstantMap): if output_map.offset < 0: - raise ValueError( + raise NoBasicSelectionError( f"cannot lower to a basic selection: output[{output_dimension}] " f"addresses coordinate {output_map.offset}, which NumPy would " "count from the end of the source" @@ -585,7 +589,7 @@ def fabricate_axes_before(axis: int) -> None: continue assert isinstance(output_map, DimensionMap) # ArrayMap raised above if output_map.stride == 0: - raise ValueError( + raise NoBasicSelectionError( f"cannot lower to a basic selection: output[{output_dimension}] " "has stride 0, which repeats one source cell along an axis; " "no slice spells a broadcast" @@ -593,7 +597,7 @@ def fabricate_axes_before(axis: int) -> None: d = output_map.input_dimension fabricate_axes_before(d) if d != next_axis: - raise ValueError( + raise NoBasicSelectionError( "cannot lower to a basic selection: the selectors must produce " f"the domain's axes in increasing order and exactly once, but " f"output[{output_dimension}] produces input dimension {d} where " @@ -609,7 +613,7 @@ def fabricate_axes_before(axis: int) -> None: continue first, last = endpoints if min(first, last) < 0: - raise ValueError( + raise NoBasicSelectionError( f"cannot lower to a basic selection: output[{output_dimension}] " f"addresses coordinate {min(first, last)}, which NumPy would " "count from the end of the source" diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 0986785b8d..4aed2029ab 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -30,6 +30,7 @@ FixedDimension, IndexTransform, LazyArray, + NoBasicSelectionError, ReadContext, VaryingDimension, dimension_grids_from_chunks, @@ -2709,9 +2710,9 @@ def test_query_part_selections_refuse_to_lower() -> None: view = LazyArray.from_numpy(reference()).with_parts((3, 2, 3)).lazy.oindex[[6, 1, 1], :, :] part = next(view.parts()) assert not part.view.is_box - with pytest.raises(ValueError, match="ArrayMap"): + with pytest.raises(NoBasicSelectionError, match="ArrayMap"): _ = part.source_selection - with pytest.raises(ValueError, match="ArrayMap"): + with pytest.raises(NoBasicSelectionError, match="ArrayMap"): _ = part.chunk_local_selection diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index e6a13e38b9..9fc9f8e00f 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -6,7 +6,7 @@ import pytest from zarr_indexing.domain import IndexDomain -from zarr_indexing.errors import BoundsCheckError +from zarr_indexing.errors import BoundsCheckError, NoBasicSelectionError from zarr_indexing.lazy_array import LazyArray from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import ( @@ -460,34 +460,34 @@ def test_selection_reproduces_the_transform(self, transform: IndexTransform) -> def test_rejects_an_array_map(self) -> None: transform = IndexTransform.from_shape((10,)).oindex[[3, 1, 1]] - with pytest.raises(ValueError, match="is an ArrayMap"): + with pytest.raises(NoBasicSelectionError, match="is an ArrayMap"): transform.as_basic_selection() def test_rejects_a_zero_stride(self) -> None: transform = IndexTransform( IndexDomain.from_shape((3,)), (DimensionMap(0, offset=2, stride=0),) ) - with pytest.raises(ValueError, match="stride 0"): + with pytest.raises(NoBasicSelectionError, match="stride 0"): transform.as_basic_selection() def test_rejects_a_transposed_dimension_order(self) -> None: transform = IndexTransform( IndexDomain.from_shape((2, 3)), (DimensionMap(1), DimensionMap(0)) ) - with pytest.raises(ValueError, match="increasing order"): + with pytest.raises(NoBasicSelectionError, match="increasing order"): transform.as_basic_selection() def test_rejects_a_repeated_input_dimension(self) -> None: transform = IndexTransform( IndexDomain.from_shape((2,)), (DimensionMap(0), DimensionMap(0, offset=1)) ) - with pytest.raises(ValueError, match="increasing order"): + with pytest.raises(NoBasicSelectionError, match="increasing order"): transform.as_basic_selection() def test_rejects_an_unreferenced_non_singleton_dimension(self) -> None: """An axis restored by repetition — a duplicate gather of a constant.""" transform = IndexTransform.from_shape((5,)).oindex[[3]].oindex[np.array([0, 0])] - with pytest.raises(ValueError, match="only a singleton axis"): + with pytest.raises(NoBasicSelectionError, match="only a singleton axis"): transform.as_basic_selection() def test_a_trailing_singleton_axis_lowers_to_a_newaxis(self) -> None: @@ -497,12 +497,12 @@ def test_a_trailing_singleton_axis_lowers_to_a_newaxis(self) -> None: def test_rejects_a_negative_constant_coordinate(self) -> None: transform = IndexTransform(IndexDomain.from_shape(()), (ConstantMap(-1),)) - with pytest.raises(ValueError, match="count from the end"): + with pytest.raises(NoBasicSelectionError, match="count from the end"): transform.as_basic_selection() def test_rejects_a_negative_slice_coordinate(self) -> None: transform = IndexTransform(IndexDomain.from_shape((2,)), (DimensionMap(0, offset=-3),)) - with pytest.raises(ValueError, match="count from the end"): + with pytest.raises(NoBasicSelectionError, match="count from the end"): transform.as_basic_selection() From 0a7522ad9f314f79cf773e42dda0ab7af3a5db55 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 21:45:22 +0200 Subject: [PATCH 12/16] feat(zarr-indexing): total decomposition into cover and residual MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `as_basic_selection` is partial by nature — integers and slices spell only diagonal, order-preserving reads — but its refusal was a dead end, and the reader privately owned the fact that it need not be: `_decompose` already factored every transform into an ascending cover plus a block-local residual. That machinery moves to transform.py and becomes public: `IndexTransform.decompose()` returns `(cover, residual)` with the value law (resolving the residual against `source[cover]` reads exactly the transform's cells) and the composition law (the cover, read as the transform it denotes, chained onto the residual, reads cell for cell as the original — the factorization is inverted by `compose`). Queries factor into their bounding interval plus a block-local gather, so a consumer with its own I/O layer can keep every part on it. `decompose_unit_step()` is the contiguous-cover variant `UnitStepReader` reads through; the readers now call the public methods, so the planned request and the executed read share one implementation. The shared walk also gained the negative-coordinate refusal, which previously produced a cover slice NumPy would have quietly wrapped. The stateful invariant now holds every part — query parts included — to the factorization law against the NumPy model, and the decompose tests run the composition law over every transform shape the lowering tests cover, including the ones as_basic_selection refuses. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+decompose.feature.md | 12 + .../zarr-indexing/docs/guide/integrations.md | 4 +- .../src/zarr_indexing/lazy_array.py | 6 +- .../zarr-indexing/src/zarr_indexing/reader.py | 119 +--------- .../src/zarr_indexing/testing/stateful.py | 23 +- .../src/zarr_indexing/transform.py | 219 +++++++++++++++++- .../zarr-indexing/tests/test_transform.py | 70 ++++++ 7 files changed, 331 insertions(+), 122 deletions(-) create mode 100644 packages/zarr-indexing/changes/+decompose.feature.md diff --git a/packages/zarr-indexing/changes/+decompose.feature.md b/packages/zarr-indexing/changes/+decompose.feature.md new file mode 100644 index 0000000000..30d0814865 --- /dev/null +++ b/packages/zarr-indexing/changes/+decompose.feature.md @@ -0,0 +1,12 @@ +`IndexTransform.decompose()` — the total counterpart of +`as_basic_selection`: every transform factors into an ascending basic cover +plus an in-memory residual, `(cover, residual)`, such that resolving the +residual against `source[cover]` reads exactly the transform's cells (and the +cover, read as the transform it denotes, composes with the residual back to +the original). Queries decompose into their bounding interval plus a +block-local gather, so a consumer can keep even query parts on its own I/O +path. `decompose_unit_step()` is the variant whose cover is contiguous and +ascending, the factorization `UnitStepReader` reads through; both live where +the readers now source their own decomposition, so the planned request and the +executed read cannot drift apart. The one refusal is a negative output +coordinate, which no cover slice can spell: `NoBasicSelectionError`. diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index badbb6ae1d..ae2c68d5f7 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -91,7 +91,9 @@ base the view partitions, so it keys a cache only alongside the grid that produced it. A query part (an `oindex`/`vindex` gather) has no slab spelling and raises `NoBasicSelectionError` (a `ValueError` subclass), so mixed consumers fall back to -`part.view.result()` for those parts. The +`part.view.result()` for those parts — or stay on their own I/O path with +`part.view.transform.decompose()`, which factors *any* transform into a basic +cover to fetch plus a residual to resolve in memory. The [asyncio example](../examples/lazy_indexing_asyncio.md) drives all three loops — gather-per-part, decoded-chunk cache, and the query fallback — with `asyncio.gather` over `zarr.AsyncArray`. diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index ab5d835d5a..d56c4af28d 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -567,7 +567,11 @@ def source_selection(self) -> tuple[int | slice | None, ...]: therefore necessary but not sufficient — catching the error is what decides it — and a consumer mixing selection kinds falls back to `view.result()` for the parts that refuse. The dedicated subclass of - `ValueError` keeps that fallback from absorbing a genuine defect. A reversing view lowers to + `ValueError` keeps that fallback from absorbing a genuine defect. A + consumer that must stay on its own I/O path even for query parts can + instead fetch the cover half of + [`view.transform.decompose()`][zarr_indexing.transform.IndexTransform.decompose] + and finish the gather in memory. A reversing view lowers to a negative-step slice and a fabricated axis to `None`, which a backend narrower than NumPy may not accept. diff --git a/packages/zarr-indexing/src/zarr_indexing/reader.py b/packages/zarr-indexing/src/zarr_indexing/reader.py index dad2296591..7a73afa64e 100644 --- a/packages/zarr-indexing/src/zarr_indexing/reader.py +++ b/packages/zarr-indexing/src/zarr_indexing/reader.py @@ -4,19 +4,14 @@ import math from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, Protocol - -if TYPE_CHECKING: - from collections.abc import Callable +from typing import Any, Final, Protocol import numpy as np from zarr_indexing._affine import checked_affine from zarr_indexing.chunk_resolution import ChunkProjection # noqa: TC001 (runtime annotation) -from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap -from zarr_indexing.transform import ( - IndexTransform, -) +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import IndexTransform # noqa: TC001 (doctest runtime) __all__ = [ "BasicReader", @@ -126,7 +121,7 @@ class BasicReader: def read_into(self, source: Any, context: ReadContext, out: Any, /) -> None: """Read one transform through a positive-slice slab and residual lowering.""" transform = context.transform - key, residual = _decompose_basic(transform) + key, residual = transform.decompose() block = np.asanyarray(source[key]) out[...] = _lower(block, residual) @@ -154,7 +149,7 @@ class NumPyReader: def read_into(self, source: Any, context: ReadContext, out: Any, /) -> None: """Read one transform through a narrowed slab into `out`.""" transform = context.transform - key, residual = _decompose_basic(transform) + key, residual = transform.decompose() block = np.asanyarray(source[key]) out[...] = _lower(block, residual) @@ -196,7 +191,7 @@ class UnitStepReader: def read_into(self, source: Any, context: ReadContext, out: Any, /) -> None: """Read one transform through an ascending unit-step slab into `out`.""" transform = context.transform - key, residual = _decompose_unit_step(transform) + key, residual = transform.decompose_unit_step() block = np.asanyarray(source[key]) out[...] = _lower(block, residual) @@ -481,105 +476,3 @@ def _lower_general(array: Any, transform: IndexTransform) -> Any: return _restore_domain_axis_order( result, list(broadcast_axes) + residual_axis_dims, transform.domain.shape ) - - -def _push_slice_for_dimension_map( - m: DimensionMap, transform: IndexTransform -) -> tuple[slice, DimensionMap]: - """The positive-step slice covering a `DimensionMap`, and its block-local map. - - A negative step is read forwards and reversed by the residual: a source is - only ever asked for a slice that walks upwards, which is the one form every - array-like agrees on. - """ - d = m.input_dimension - lo = transform.domain.inclusive_min[d] - hi = max(transform.domain.exclusive_max[d], lo) - endpoints = m.endpoints(lo, hi) - if endpoints is None: - return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) - first, last = endpoints - if m.stride > 0: - return ( - slice(first, last + 1, m.stride), - DimensionMap(input_dimension=d, offset=-lo, stride=1), - ) - if m.stride == 0: - return ( - slice(first, first + 1, 1), - DimensionMap(input_dimension=d, offset=0, stride=0), - ) - # Descending: the block holds the same coordinates in ascending order, so - # the residual walks it backwards from the last block position. - return ( - slice(last, first + 1, -m.stride), - DimensionMap(input_dimension=d, offset=hi - 1, stride=-1), - ) - - -def _push_unit_slice_for_dimension_map( - m: DimensionMap, transform: IndexTransform -) -> tuple[slice, DimensionMap]: - """The unit-step slice covering a `DimensionMap`, and its block-local map. - - Strides and reversals stay in the residual: the source is only ever asked - for a contiguous ascending slice, and the original stride is replayed - against the in-memory block. The cover therefore over-reads a strided - selection by its stride factor, which is the price of a source that - accepts nothing but `slice(start, stop, 1)`. - """ - d = m.input_dimension - lo = transform.domain.inclusive_min[d] - hi = max(transform.domain.exclusive_max[d], lo) - endpoints = m.endpoints(lo, hi) - if endpoints is None: - return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) - first, last = endpoints - if m.stride == 0: - return ( - slice(first, first + 1, 1), - DimensionMap(input_dimension=d, offset=0, stride=0), - ) - origin = min(first, last) - return ( - slice(origin, max(first, last) + 1, 1), - DimensionMap(input_dimension=d, offset=m.offset - origin, stride=m.stride), - ) - - -def _decompose_basic(transform: IndexTransform) -> tuple[tuple[slice, ...], IndexTransform]: - return _decompose(transform, _push_slice_for_dimension_map) - - -def _decompose_unit_step(transform: IndexTransform) -> tuple[tuple[slice, ...], IndexTransform]: - return _decompose(transform, _push_unit_slice_for_dimension_map) - - -def _decompose( - transform: IndexTransform, - push_dimension_map: Callable[[DimensionMap, IndexTransform], tuple[slice, DimensionMap]], -) -> tuple[tuple[slice, ...], IndexTransform]: - key: list[slice] = [] - residual: list[OutputIndexMap] = [] - for output_map in transform.output: - if isinstance(output_map, ConstantMap): - coordinate = checked_affine(output_map.offset, 0, 0) - key.append(slice(coordinate, coordinate + 1, 1)) - residual.append(ConstantMap(offset=0)) - elif isinstance(output_map, DimensionMap): - pushed, local = push_dimension_map(output_map, transform) - key.append(pushed) - residual.append(local) - else: - coordinates = checked_affine( - output_map.offset, output_map.stride, output_map.index_array - ) - if coordinates.size == 0: - key.append(slice(0, 0, 1)) - local_index = coordinates - else: - origin = int(coordinates.min()) - key.append(slice(origin, int(coordinates.max()) + 1, 1)) - local_index = checked_affine(-origin, 1, coordinates) - residual.append(ArrayMap(index_array=local_index)) - return tuple(key), IndexTransform(domain=transform.domain, output=tuple(residual)) diff --git a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py index 98d179a00a..43b4066117 100644 --- a/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py +++ b/packages/zarr-indexing/src/zarr_indexing/testing/stateful.py @@ -467,12 +467,24 @@ def box_parts_lower_to_basic_selections(self) -> None: semantics to the values the source holds and checked against the model; a query part must refuse to lower instead of guessing a slab. - No reader runs in this invariant: it proves the lowered selections - alone carry the read, which is exactly what a consumer that plans here - but fetches elsewhere relies on. + Every part — query parts included — must also satisfy the + factorization law: `transform.decompose()` yields a cover and a + residual whose resolution against `data[cover]` reproduces the model, + the read a consumer performs when it fetches a bounding slab through + its own I/O layer and finishes the gather in memory. + + Apart from resolving that residual, no reader runs in this invariant: + it proves the lowered selections alone carry the read, which is + exactly what a consumer that plans here but fetches elsewhere relies + on. """ data = np.asarray(type(self).data) for part in self.view.parts(): + expected_block = self.model[part.out_selection] + cover, residual = part.view.transform.decompose() + decomposed = np.empty(expected_block.shape, dtype=data.dtype) + basic_reader.read_into(data[cover], ReadContext(residual), decomposed) + np.testing.assert_array_equal(decomposed, expected_block, err_msg=str(self.chain)) try: source_selection = part.source_selection except NoBasicSelectionError: @@ -498,9 +510,8 @@ def box_parts_lower_to_basic_selections(self) -> None: f"refused to: {self.chain}" ) continue - expected = self.model[part.out_selection] slab = data[source_selection] - np.testing.assert_array_equal(slab, expected, err_msg=str(self.chain)) + np.testing.assert_array_equal(slab, expected_block, err_msg=str(self.chain)) cell = data[ tuple( slice(lo, hi) @@ -512,7 +523,7 @@ def box_parts_lower_to_basic_selections(self) -> None: ) ] np.testing.assert_array_equal( - cell[part.chunk_local_selection], expected, err_msg=str(self.chain) + cell[part.chunk_local_selection], expected_block, err_msg=str(self.chain) ) diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 56148efc63..bdad6faec5 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -51,7 +51,7 @@ ) if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence import numpy.typing as npt @@ -473,6 +473,10 @@ def as_basic_selection(self) -> tuple[int | slice | None, ...]: stand in for it is one nothing reads — a `None`/newaxis in the selection that made it — and lowers back to `None`. + [`decompose`][zarr_indexing.transform.IndexTransform.decompose] is + the total counterpart: every transform factors into a basic cover + plus an in-memory residual, so a refusal here is not a dead end. + This is the boundary for consumers that plan reads here but perform them through their own I/O layer — an async store, an HTTP range request — whose request vocabulary is a basic selection rather than a @@ -636,6 +640,105 @@ def fabricate_axes_before(axis: int) -> None: ) return tuple(selection) + def decompose(self) -> tuple[tuple[slice, ...], IndexTransform]: + """Factor this transform into a basic cover and an in-memory residual. + + The total counterpart of + [`as_basic_selection`][zarr_indexing.transform.IndexTransform.as_basic_selection]: + where that method is defined only for reads that *are* a basic + selection, every transform factors as one basic read followed by an + in-memory rearrangement. Returns `(cover, residual)` such that + resolving `residual` against the block `source[cover]` reads exactly + the cells this transform addresses, at exactly its domain shape — so a + consumer whose I/O layer speaks basic selections can fetch the cover + through it and finish any gather, reversal, or broadcast in memory. + + The cover is one ascending slice per output dimension: a + `DimensionMap` covers its arithmetic progression (a descending map is + read forwards and reversed by the residual), a `ConstantMap` covers + its single coordinate, and an `ArrayMap` covers the range from its + smallest to its largest coordinate — the price of a slab vocabulary is + over-reading a sparse gather's bounding interval. The residual keeps + this transform's domain and rewrites each output map to block-local + coordinates, so the factorization inverts by composition: the cover, + read as the diagonal transform it denotes, chained onto the residual, + reads cell for cell as `self` (the maps may spell offsets + differently; the reads are identical). + + Returns + ------- + tuple + `(cover, residual)`: a tuple of ascending slices to request from + the source, and the `IndexTransform` to resolve against the + returned block. + + Raises + ------ + NoBasicSelectionError + If an output coordinate is negative — NumPy would count it from + the end of the source, so no cover slice can spell it. This is + the one transform shape with no factorization; every other + transform, including queries and broadcasts, decomposes. + + Examples + -------- + A descending strided read: the cover walks forward, the residual + reverses. + + >>> import numpy as np + >>> source = np.arange(10) + >>> cover, residual = IndexTransform.from_shape((10,))[::-3].decompose() + >>> cover + (slice(0, 10, 3),) + >>> block = source[cover] + >>> lo, hi = residual.domain.inclusive_min[0], residual.domain.exclusive_max[0] + >>> [int(block[residual.apply((p,))]) for p in range(lo, hi)] + [9, 6, 3, 0] + + A query decomposes into its bounding interval and a block-local + lookup, where `as_basic_selection` refuses: + + >>> cover, residual = IndexTransform.from_shape((10,)).oindex[[7, 2, 2]].decompose() + >>> cover + (slice(2, 8, 1),) + >>> block = source[cover] + >>> [int(block[residual.apply((p,))]) for p in range(3)] + [7, 2, 2] + """ + return _decompose_basic(self) + + def decompose_unit_step(self) -> tuple[tuple[slice, ...], IndexTransform]: + """Factor as `decompose` does, with a contiguous ascending cover. + + The variant of + [`decompose`][zarr_indexing.transform.IndexTransform.decompose] for a + source that accepts nothing but `slice(start, stop, 1)` — compare + [`UnitStepReader`][zarr_indexing.reader.UnitStepReader], which reads + through exactly this factorization. Strides and reversals stay in the + residual, so the cover over-reads a strided selection by its stride + factor: the price of a unit-step request vocabulary. + + Returns + ------- + tuple + `(cover, residual)` exactly as `decompose` returns them, with + every cover slice contiguous and ascending. + + Raises + ------ + NoBasicSelectionError + If an output coordinate is negative, exactly as for `decompose`. + + Examples + -------- + >>> cover, residual = IndexTransform.from_shape((10,))[1:9:3].decompose_unit_step() + >>> cover + (slice(1, 8, 1),) + >>> residual.output + (DimensionMap(input_dimension=0, offset=0, stride=3),) + """ + return _decompose_unit_step(self) + @property def selection_repr(self) -> str: """Compact domain string, e.g. `'{ [2, 8), [0, 10) }'`. @@ -1644,6 +1747,120 @@ def _reshape_to_axis( return flat.reshape(shape) +# --------------------------------------------------------------------------- # +# Cover-and-residual decomposition +# --------------------------------------------------------------------------- # + + +def _push_slice_for_dimension_map( + m: DimensionMap, transform: IndexTransform +) -> tuple[slice, DimensionMap]: + """The positive-step slice covering a `DimensionMap`, and its block-local map. + + A negative step is read forwards and reversed by the residual: a source is + only ever asked for a slice that walks upwards, which is the one form every + array-like agrees on. + """ + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = max(transform.domain.exclusive_max[d], lo) + endpoints = m.endpoints(lo, hi) + if endpoints is None: + return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) + first, last = endpoints + if m.stride > 0: + return ( + slice(first, last + 1, m.stride), + DimensionMap(input_dimension=d, offset=-lo, stride=1), + ) + if m.stride == 0: + return ( + slice(first, first + 1, 1), + DimensionMap(input_dimension=d, offset=0, stride=0), + ) + # Descending: the block holds the same coordinates in ascending order, so + # the residual walks it backwards from the last block position. + return ( + slice(last, first + 1, -m.stride), + DimensionMap(input_dimension=d, offset=hi - 1, stride=-1), + ) + + +def _push_unit_slice_for_dimension_map( + m: DimensionMap, transform: IndexTransform +) -> tuple[slice, DimensionMap]: + """The unit-step slice covering a `DimensionMap`, and its block-local map. + + Strides and reversals stay in the residual: the source is only ever asked + for a contiguous ascending slice, and the original stride is replayed + against the in-memory block. The cover therefore over-reads a strided + selection by its stride factor, which is the price of a source that + accepts nothing but `slice(start, stop, 1)`. + """ + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = max(transform.domain.exclusive_max[d], lo) + endpoints = m.endpoints(lo, hi) + if endpoints is None: + return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) + first, last = endpoints + if m.stride == 0: + return ( + slice(first, first + 1, 1), + DimensionMap(input_dimension=d, offset=0, stride=0), + ) + origin = min(first, last) + return ( + slice(origin, max(first, last) + 1, 1), + DimensionMap(input_dimension=d, offset=m.offset - origin, stride=m.stride), + ) + + +def _decompose_basic(transform: IndexTransform) -> tuple[tuple[slice, ...], IndexTransform]: + return _decompose(transform, _push_slice_for_dimension_map) + + +def _decompose_unit_step(transform: IndexTransform) -> tuple[tuple[slice, ...], IndexTransform]: + return _decompose(transform, _push_unit_slice_for_dimension_map) + + +def _decompose( + transform: IndexTransform, + push_dimension_map: Callable[[DimensionMap, IndexTransform], tuple[slice, DimensionMap]], +) -> tuple[tuple[slice, ...], IndexTransform]: + key: list[slice] = [] + residual: list[OutputIndexMap] = [] + for output_map in transform.output: + if isinstance(output_map, ConstantMap): + coordinate = checked_affine(output_map.offset, 0, 0) + key.append(slice(coordinate, coordinate + 1, 1)) + residual.append(ConstantMap(offset=0)) + elif isinstance(output_map, DimensionMap): + pushed, local = push_dimension_map(output_map, transform) + key.append(pushed) + residual.append(local) + else: + coordinates = checked_affine( + output_map.offset, output_map.stride, output_map.index_array + ) + if coordinates.size == 0: + key.append(slice(0, 0, 1)) + local_index = coordinates + else: + origin = int(coordinates.min()) + key.append(slice(origin, int(coordinates.max()) + 1, 1)) + local_index = checked_affine(-origin, 1, coordinates) + residual.append(ArrayMap(index_array=local_index)) + for output_dimension, entry in enumerate(key): + if entry.start < 0: + raise NoBasicSelectionError( + f"cannot cover with a basic selection: output[{output_dimension}] " + f"addresses coordinate {entry.start}, which NumPy would count " + "from the end of the source" + ) + return tuple(key), IndexTransform(domain=transform.domain, output=tuple(residual)) + + class _OIndexHelper: """Helper that provides orthogonal (outer) indexing via `transform.oindex[...]`.""" diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index 9fc9f8e00f..19989c209f 100644 --- a/packages/zarr-indexing/tests/test_transform.py +++ b/packages/zarr-indexing/tests/test_transform.py @@ -506,6 +506,76 @@ def test_rejects_a_negative_slice_coordinate(self) -> None: transform.as_basic_selection() +class TestDecompose: + @pytest.mark.parametrize( + "transform", + [ + pytest.param(IndexTransform.from_shape((10,)), id="identity"), + pytest.param(IndexTransform.from_shape((10,))[1:9:3], id="strided"), + pytest.param(IndexTransform.from_shape((6,))[::-1], id="reversed"), + pytest.param(IndexTransform.from_shape((10,))[4:4], id="empty"), + pytest.param(IndexTransform.from_shape((5, 7))[3, 1:6:2], id="int-and-slice"), + pytest.param(IndexTransform.from_shape((5, 7))[3, 2], id="rank-zero"), + pytest.param(IndexTransform.from_shape((10,)).oindex[[7, 2, 2]], id="query-gather"), + pytest.param( + IndexTransform.from_shape((5, 7)).oindex[slice(None), [2]], + id="collapsed-single-gather", + ), + pytest.param( + IndexTransform( + IndexDomain.from_shape((3,)), (DimensionMap(0, offset=2, stride=0),) + ), + id="broadcast", + ), + pytest.param( + IndexTransform(IndexDomain.from_shape((2, 3)), (DimensionMap(1), DimensionMap(0))), + id="transposed", + ), + pytest.param(IndexTransform.from_shape((5,))[None], id="newaxis"), + pytest.param( + IndexTransform.from_shape((5,)).oindex[[3]].oindex[np.array([0, 0])], + id="axis-restored-by-repetition", + ), + ], + ) + def test_cover_and_residual_reproduce_the_transform(self, transform: IndexTransform) -> None: + """The factorization law, checked both ways. + + Value law: resolving the residual against `source[cover]` reads + exactly what the transform reads. Composition law: the cover, read as + the diagonal transform it denotes, chained onto the residual, is the + original transform — `decompose` is inverted by `compose`. Every + transform decomposes, including the shapes `as_basic_selection` + refuses. + """ + source = np.arange(35).reshape(5, 7) if transform.output_rank == 2 else np.arange(10) + cover, residual = transform.decompose() + assert len(cover) == transform.output_rank + assert all(entry.step >= 1 for entry in cover) + block = np.asarray(source[cover]) + np.testing.assert_array_equal( + _pointwise_read(residual, block), _pointwise_read(transform, source) + ) + cover_transform = IndexTransform( + IndexDomain.from_shape(block.shape), + tuple( + DimensionMap(axis, offset=entry.start, stride=entry.step) + for axis, entry in enumerate(cover) + ), + ) + composed = residual.compose(cover_transform) + assert composed.domain == transform.domain + np.testing.assert_array_equal( + _pointwise_read(composed, source), _pointwise_read(transform, source) + ) + + def test_rejects_a_negative_coordinate(self) -> None: + """The one shape with no cover: NumPy would wrap the coordinate.""" + transform = IndexTransform(IndexDomain.from_shape((2,)), (DimensionMap(0, offset=-3),)) + with pytest.raises(NoBasicSelectionError, match="count from the end"): + transform.decompose() + + class TestIndexTransformBasicIndexing: def test_slice_identity(self) -> None: """slice(None) on identity transform is a no-op.""" From e407df38d9c6a925baa77b72c2b8da1f6a31bbbb Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 13 Aug 2026 21:58:37 +0200 Subject: [PATCH 13/16] feat(zarr-indexing): name the lowering's output vocabulary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `tuple[int | slice | None, ...]` was spelled out at six signature sites and had already drifted once (`source_selection` said `tuple[int | slice, ...]` until the newaxis lowering widened it). It is now the public alias `BasicSelection`, defined beside the lowering and exported at the root. A plain alias, deliberately not a NewType: the tuple's validity is relative to the array it is applied to — the same value can be a correct source_selection and a wrong chunk_local_selection — so a nominal brand would assert a provenance the type system cannot carry, and no signature inside or outside the package could enforce it. The name is the destination contract's: `zarr.AsyncArray.getitem` types its selection parameter `BasicSelection`, and NumPy calls the dialect basic indexing. The asyncio example's `AsyncSource` protocol now uses it in the one place the type appears in parameter position. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+basic-selection-alias.feature.md | 6 ++++++ .../lazy_indexing_asyncio.py | 4 ++-- .../src/zarr_indexing/__init__.py | 2 ++ .../src/zarr_indexing/lazy_array.py | 5 +++-- .../src/zarr_indexing/transform.py | 20 ++++++++++++++++--- 5 files changed, 30 insertions(+), 7 deletions(-) create mode 100644 packages/zarr-indexing/changes/+basic-selection-alias.feature.md diff --git a/packages/zarr-indexing/changes/+basic-selection-alias.feature.md b/packages/zarr-indexing/changes/+basic-selection-alias.feature.md new file mode 100644 index 0000000000..162b5a8d81 --- /dev/null +++ b/packages/zarr-indexing/changes/+basic-selection-alias.feature.md @@ -0,0 +1,6 @@ +`BasicSelection` — the public alias for the lowering's output vocabulary, +`tuple[int | slice | None, ...]`, exported at the package root. The name is +the destination contract's: it is what `zarr.AsyncArray.getitem` calls its +selection parameter, and what NumPy calls basic indexing. Use it to annotate a +consumer-owned I/O protocol (`async def getitem(self, selection: +BasicSelection)`), as the asyncio example now does. diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py index 781a7a0c9e..6a353bfbd7 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py @@ -22,7 +22,7 @@ import zarr import zarr.api.asynchronous -from zarr_indexing import LazyArray, NoBasicSelectionError, Partition +from zarr_indexing import BasicSelection, LazyArray, NoBasicSelectionError, Partition class AsyncSource(Protocol): @@ -33,7 +33,7 @@ class AsyncSource(Protocol): wrapper, a database. The planner never sees this object. """ - async def getitem(self, selection: Any) -> Any: ... + async def getitem(self, selection: BasicSelection) -> Any: ... @pytest.fixture diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 42a41d6f90..99fea157b3 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -73,6 +73,7 @@ unit_step_reader, ) from zarr_indexing.transform import ( + BasicSelection, IndexTransform, ) @@ -81,6 +82,7 @@ __all__ = [ "ArrayMap", "BasicReader", + "BasicSelection", "BoundsCheckError", "ChunkCoverage", "ChunkGrid", diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index d56c4af28d..a6254f7890 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -170,6 +170,7 @@ numpy_reader, ) from zarr_indexing.transform import ( + BasicSelection, IndexTransform, ) @@ -543,7 +544,7 @@ def is_complete(self) -> bool: return self.projection.coverage == "full" @property - def source_selection(self) -> tuple[int | slice | None, ...]: + def source_selection(self) -> BasicSelection: """The basic selection on the raw wrapped array that reads this part. Lowered from `view.transform` by @@ -586,7 +587,7 @@ def source_selection(self) -> tuple[int | slice | None, ...]: return self.view.transform.as_basic_selection() @property - def chunk_local_selection(self) -> tuple[int | slice | None, ...]: + def chunk_local_selection(self) -> BasicSelection: """The same read as `source_selection`, relative to the part's grid cell. Lowered from `projection.chunk_transform`, so coordinate 0 per axis is diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index bdad6faec5..2fb671e5e8 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -58,6 +58,20 @@ from zarr_indexing.json import IndexTransformJSON +type BasicSelection = tuple[int | slice | None, ...] +"""A selection in NumPy's basic-indexing dialect: one selector per axis. + +The request vocabulary of `source[selection]` under basic indexing — also the +parameter type of `zarr.AsyncArray.getitem`. Denotationally a product of +arithmetic progressions: an `int` reads one coordinate and drops its axis, a +`slice` reads an arithmetic progression, and `None` fabricates an axis no +source axis backs. Produced by +[`IndexTransform.as_basic_selection`][zarr_indexing.transform.IndexTransform.as_basic_selection] +and the `Partition` lowering properties; which cells a value denotes is +relative to the array it is applied to. +""" + + @dataclass(frozen=True, slots=True) class _PointOutOfBounds(Exception): """Internal signal from the shared point kernel: one coordinate left the domain. @@ -452,7 +466,7 @@ def inverted(self) -> IndexTransform: output=tuple(inverse_output[dimension] for dimension in range(self.input_rank)), ) - def as_basic_selection(self) -> tuple[int | slice | None, ...]: + def as_basic_selection(self) -> BasicSelection: """Lower a box transform to the NumPy basic selection it describes. Returns one selector per output dimension, plus a `None` for each @@ -1492,7 +1506,7 @@ def _intersect_general( return (result, out_indices.astype(np.intp)) -def _normalize_basic_selection(selection: Any, ndim: int) -> tuple[int | slice | None, ...]: +def _normalize_basic_selection(selection: Any, ndim: int) -> BasicSelection: """Normalize a selection to a tuple of int, slice, or None (newaxis), expanding ellipsis and padding with slice(None) as needed. """ @@ -1552,7 +1566,7 @@ def _positional_slice(pos: int, size: int, step: int) -> slice: def _reindex_array( m: ArrayMap, - normalized: tuple[int | slice | None, ...], + normalized: BasicSelection, domain: IndexDomain, ) -> np.ndarray[Any, np.dtype[np.intp]]: """Apply basic indexing operations to an ArrayMap's index_array. From 141a1b76441d93bf2fccd230c18d459c5f3f8a08 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 14 Aug 2026 18:22:00 +0200 Subject: [PATCH 14/16] fix(zarr-indexing): reject overlapping result buffers Reject destination views whose logical elements share storage while preserving valid reversed, transposed, and strided outputs. Add regressions for zero-stride and nonzero-stride overlap. Assisted-by: Codex:gpt-5 --- .../src/zarr_indexing/lazy_array.py | 50 ++++++++++++++++--- .../zarr-indexing/tests/test_lazy_array.py | 31 ++++++++++++ 2 files changed, 73 insertions(+), 8 deletions(-) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index a6254f7890..1d5d53c992 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -315,6 +315,33 @@ def _overlaps(out: np.ndarray[Any, Any], array: Any) -> bool: return bool(np.may_share_memory(out, array)) and bool(np.shares_memory(out, array)) +def _has_internal_overlap(out: np.ndarray[Any, Any]) -> bool: + """Whether distinct logical cells of ``out`` may share storage. + + NumPy exposes exact overlap checks between two arrays, but not within one + strided array. Prove the common layouts disjoint by walking axes from the + smallest absolute stride outward: each next axis must begin beyond the + complete byte span reachable through the axes already seen. Standard + slices, reversals, and transposes satisfy this proof. Ambiguous exotic + layouts are rejected along with layouts that definitely overlap; accepting + one would make ``result_into`` silently return a buffer whose logical cells + cannot hold distinct result values. + """ + if out.size <= 1 or out.dtype.itemsize == 0: + return False + byte_span = out.dtype.itemsize + axes = sorted( + (abs(stride), extent) + for extent, stride in zip(out.shape, out.strides, strict=True) + if extent > 1 + ) + for stride, extent in axes: + if stride < byte_span: + return True + byte_span += (extent - 1) * stride + return False + + def _translated_domain(domain: IndexDomain, window: tuple[slice, ...]) -> IndexDomain: """`domain`, expressed in the coordinates `window` was cut from. @@ -553,8 +580,8 @@ def source_selection(self) -> BasicSelection: `view.array[part.source_selection]` reads the same block `view.result()` reads, in the same order. This is the request to hand to a consumer-owned I/O layer whose vocabulary is a basic selection - rather than a transform — for one driving reads through an async - store, assembly is one line per part: + rather than a transform. When that backend accepts NumPy's complete + basic-selection dialect, assembly is one line per part: ```python out[part.out_selection] = await source.getitem(part.source_selection) @@ -1347,8 +1374,9 @@ def result_into( ---------- out A writable `numpy.ndarray` of exactly `self.shape` and this view's - dtype, not overlapping the wrapped array. A view into a larger - array qualifies, so a part's result can land directly in its slot: + dtype, with distinct storage for its logical cells and not + overlapping the wrapped array. A view into a larger array + qualifies, so a part's result can land directly in its slot: ```python if all(isinstance(s, slice) for s in part.out_selection): @@ -1379,10 +1407,11 @@ def result_into( TypeError If `out` is not a `numpy.ndarray`. ValueError - If `out` has the wrong shape or dtype, is read-only, is unmasked - for a masked source, or shares memory with the wrapped array; or if - supplied parts were prepared by another view or do not tile this - view exactly. + If `out` has the wrong shape or dtype, is read-only, has an + internally overlapping memory layout, is unmasked for a masked + source, or shares memory with the wrapped array; or if supplied + parts were prepared by another view or do not tile this view + exactly. AssertionError If this library's own partition walk fails to cover the view — a bug in zarr-indexing, never a consequence of the caller's input. @@ -1413,6 +1442,11 @@ def result_into( ) if not out.flags.writeable: raise ValueError("out is read-only; result_into writes every cell of it") + if _has_internal_overlap(out): + raise ValueError( + "out has overlapping elements; distinct cells of this view need " + "distinct destination storage" + ) if isinstance(self._array, np.ma.MaskedArray) and not isinstance(out, np.ma.MaskedArray): # dtypes match, so nothing above rejects a plain buffer, and the # reader would write the values under the mask into it as if they diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 4aed2029ab..b07af1f7bf 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -2845,6 +2845,16 @@ def test_result_into_fills_the_callers_buffer() -> None: np.testing.assert_array_equal(np.ma.getmaskarray(out_masked), masked_source.mask) np.testing.assert_array_equal(out_masked.compressed(), masked_source.compressed()) + # Non-contiguous layouts with distinct cells remain valid destinations. + view = LazyArray.from_numpy(data).lazy[1:6, 1:4, :] + strided_outputs = ( + np.empty(view.shape, dtype=view.dtype)[::-1], + np.empty(view.shape[::-1], dtype=view.dtype).transpose(2, 1, 0), + ) + for strided_out in strided_outputs: + assert view.result_into(strided_out) is strided_out + np.testing.assert_array_equal(strided_out, np.asarray(view.result())) + def test_result_into_rejects_a_non_array() -> None: view = LazyArray.from_numpy(reference()) @@ -2872,6 +2882,27 @@ def test_result_into_rejects_a_read_only_buffer() -> None: view.result_into(out) +def test_result_into_rejects_a_zero_stride_buffer() -> None: + view = LazyArray.from_numpy(np.arange(4)) + storage = np.empty(1, dtype=view.dtype) + out = np.lib.stride_tricks.as_strided(storage, shape=view.shape, strides=(0,), writeable=True) + with pytest.raises(ValueError, match="overlapping elements"): + view.result_into(out) + + +def test_result_into_rejects_a_nonzero_stride_overlapping_buffer() -> None: + view = LazyArray.from_numpy(np.arange(4).reshape(2, 2)) + storage = np.empty(3, dtype=view.dtype) + out = np.lib.stride_tricks.as_strided( + storage, + shape=view.shape, + strides=(view.dtype.itemsize, view.dtype.itemsize), + writeable=True, + ) + with pytest.raises(ValueError, match="overlapping elements"): + view.result_into(out) + + def test_result_into_rejects_an_unmasked_buffer_for_a_masked_source() -> None: data = reference() view = LazyArray.from_numpy(np.ma.masked_array(data, mask=data % 5 == 0)) From 06612286e2f2491c202cc4c9879fa669d08edc8c Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 14 Aug 2026 18:22:31 +0200 Subject: [PATCH 15/16] fix(zarr-indexing): adapt lazy reads to AsyncArray Normalize new-axis, negative-step, and query selections into an ascending cover plus an in-memory residual. Exercise the adapter with real AsyncArray reads and include the example in strict typechecking. Assisted-by: Codex:gpt-5 --- .../zarr-indexing/docs/guide/integrations.md | 13 +- .../examples/lazy_indexing_asyncio/README.md | 17 +- .../lazy_indexing_asyncio.py | 162 ++++++++++++------ packages/zarr-indexing/justfile | 2 +- packages/zarr-indexing/pyproject.toml | 1 + .../src/zarr_indexing/transform.py | 12 +- 6 files changed, 136 insertions(+), 71 deletions(-) diff --git a/packages/zarr-indexing/docs/guide/integrations.md b/packages/zarr-indexing/docs/guide/integrations.md index ae2c68d5f7..870ead2035 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -88,12 +88,13 @@ same read relative to the part's grid cell, for consumers caching decoded chunks. `projection.chunk_domain` locates that cell in the source, so it is what to fetch and a sound cache key; `base_coords` counts cells of whatever base the view partitions, so it keys a cache only alongside the grid that -produced it. A query part (an `oindex`/`vindex` gather) has -no slab spelling and raises `NoBasicSelectionError` (a `ValueError` -subclass), so mixed consumers fall back to -`part.view.result()` for those parts — or stay on their own I/O path with -`part.view.transform.decompose()`, which factors *any* transform into a basic -cover to fetch plus a residual to resolve in memory. The +produced it. A query part (an `oindex`/`vindex` gather) has no slab spelling +and raises `NoBasicSelectionError` (a `ValueError` subclass), so the AsyncArray +adapter stays on its own I/O path with `part.view.transform.decompose()`. That +factors the transform into an ascending basic cover to fetch plus a residual +to resolve in memory. The same path handles NumPy's newaxis and negative-step +slices, which Zarr's narrower basic-selection dialect rejects even though +`source_selection` can spell them. The [asyncio example](../examples/lazy_indexing_asyncio.md) drives all three loops — gather-per-part, decoded-chunk cache, and the query fallback — with `asyncio.gather` over `zarr.AsyncArray`. diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md index 749b4833d4..d782452caf 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md @@ -7,23 +7,26 @@ structure and this example shows `asyncio.gather` driving it. The example shows how to: -- Lower each box-shaped partition to a backend-native request with +- Lower each compatible box-shaped partition to a backend-native request with `part.source_selection` — a tuple of integers and slices in the wrapped array's own coordinates — and fetch all partitions concurrently, assembling each block with `out[part.out_selection] = await source.getitem(part.source_selection)` +- Normalize the two NumPy basic selectors that `zarr.AsyncArray.getitem` does + not accept: a newaxis (`None`) and a negative-step slice. Those parts fetch + their ascending cover through Zarr and apply the residual transform in + memory. - Build a decoded-chunk cache keyed on each touched grid cell's global origin: fetch the cell (`projection.chunk_domain`) once, then serve every overlapping view from the cache with `part.chunk_local_selection` - Fall back for query partitions (`oindex`/`vindex`/mask selections), whose gathers have no single-slab spelling: `source_selection` raises - `NoBasicSelectionError` - and the part resolves through its own `part.view.result()` instead + `NoBasicSelectionError`, so the adapter fetches their ascending cover and + applies the residual gather in memory The async side only needs one method — `async def getitem(selection)` accepting -a basic selection — so the same loop drives an HTTP range endpoint or any other -async source. Note that a reversing view (`lazy[::-1]`) lowers to a -negative-step slice, which Zarr's basic selections reject; inspect the slice -steps if your backend only walks forward. +ascending basic slices — so the same adapter drives an HTTP range endpoint or +any other async source with Zarr's selection dialect. A backend with a wider +dialect can take the direct `source_selection` path for more parts. ## Running the Example diff --git a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py index 6a353bfbd7..fc8cfb709c 100644 --- a/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py @@ -15,25 +15,69 @@ import asyncio import sys -from typing import Any, Protocol +from typing import Any, cast import numpy as np import pytest import zarr import zarr.api.asynchronous +from zarr.core.indexing import BasicSelection as ZarrBasicSelection -from zarr_indexing import BasicSelection, LazyArray, NoBasicSelectionError, Partition +from zarr_indexing import ( + BasicSelection, + LazyArray, + NoBasicSelectionError, + Partition, + ReadContext, + numpy_reader, +) -class AsyncSource(Protocol): - """The surface the async loop needs: one awaitable basic-selection read. +class _NoSynchronousRead: + """A test reader proving the async adapter performs every source read.""" - `zarr.AsyncArray` satisfies it; so does anything else that can serve a - tuple of integers and ascending slices — an HTTP tile endpoint, an fsspec - wrapper, a database. The planner never sees this object. - """ + def read_into( + self, + _source: Any, + _context: ReadContext, + _out: np.ndarray[Any, Any], + /, + ) -> None: + raise AssertionError("the AsyncArray integration performed a synchronous read") + + +def _as_zarr_selection(selection: BasicSelection) -> ZarrBasicSelection | None: + """Narrow a NumPy basic selection to the dialect AsyncArray accepts. - async def getitem(self, selection: BasicSelection) -> Any: ... + Zarr accepts integers and positive-step slices, but not NumPy's ``None`` + newaxis or negative-step slices. Returning ``None`` selects the cover and + residual path below; no backend exception is used for feature detection. + """ + if any( + item is None or (isinstance(item, slice) and item.step is not None and item.step < 1) + for item in selection + ): + return None + return cast("ZarrBasicSelection", selection) + + +async def _read_part(part: Partition, source: zarr.AsyncArray[Any]) -> np.ndarray[Any, Any]: + """Fetch one part through AsyncArray, normalizing its narrower dialect.""" + try: + direct = _as_zarr_selection(part.source_selection) + except NoBasicSelectionError: + direct = None + if direct is not None: + return np.asanyarray(await source.getitem(direct)) + + # A cover contains one ascending slice per source axis, which is always a + # Zarr basic selection. The residual restores a newaxis or reversal and + # performs any query gather against the fetched NumPy block. + cover, residual = part.view.transform.decompose() + block = np.asanyarray(await source.getitem(cover)) + out = np.empty(part.view.shape, dtype=part.view.dtype) + numpy_reader.read_into(block, ReadContext(residual), out) + return out @pytest.fixture @@ -43,32 +87,32 @@ def store() -> dict[str, Any]: @pytest.fixture -def source(store: dict[str, Any]) -> zarr.Array: +def source(store: dict[str, Any]) -> zarr.Array[Any]: """A chunked Zarr array, created synchronously and shared with the async side.""" array = zarr.create_array(store=store, shape=(40, 30), chunks=(10, 10), dtype="i4") array[:] = np.arange(40 * 30).reshape(40, 30) return array -async def read_through(view: LazyArray, source: AsyncSource) -> np.ndarray: +async def read_through(view: LazyArray, source: zarr.AsyncArray[Any]) -> np.ndarray[Any, Any]: """Materialize `view` by fetching every partition concurrently. - The wrapper plans; the async source fetches. Each box partition lowers to - the basic selection `part.source_selection`, is fetched through the - caller's own I/O layer, and lands at `part.out_selection` — no reader, no - thread pool, and no scheduler inside zarr-indexing. + The wrapper plans; the async source fetches. Compatible box partitions use + `part.source_selection` directly. Other partitions fetch an ascending + cover and apply the residual NumPy selection in memory. Every block lands + at `part.out_selection` — no thread pool or scheduler inside zarr-indexing. """ parts = tuple(view.parts()) - blocks = await asyncio.gather(*(source.getitem(part.source_selection) for part in parts)) + blocks = await asyncio.gather(*(_read_part(part, source) for part in parts)) out = np.empty(view.shape, dtype=view.dtype) for part, block in zip(parts, blocks, strict=True): out[part.out_selection] = block return out -def test_parts_with_asyncio_gather(store: dict[str, Any], source: zarr.Array) -> None: +def test_parts_with_asyncio_gather(store: dict[str, Any], source: zarr.Array[Any]) -> None: """Fetch a view's partitions concurrently through zarr.AsyncArray.""" - view = LazyArray(source).lazy[5:35, 3:27] + view = LazyArray(cast(Any, source)).lazy[5:35, 3:27] async def scenario() -> np.ndarray: async_source = await zarr.api.asynchronous.open_array(store=store) @@ -80,7 +124,7 @@ async def scenario() -> np.ndarray: # Strided and integer selections lower the same way; a scalar axis lowers # to an integer and drops, exactly as it does in the output placement. - decimated = LazyArray(source).lazy[::4, 7] + decimated = LazyArray(cast(Any, source)).lazy[::4, 7] async def decimated_scenario() -> np.ndarray: async_source = await zarr.api.asynchronous.open_array(store=store) @@ -89,7 +133,29 @@ async def decimated_scenario() -> np.ndarray: assert np.array_equal(asyncio.run(decimated_scenario()), source[::4, 7]) -def test_decoded_chunk_cache(store: dict[str, Any], source: zarr.Array) -> None: +def test_asyncarray_dialect_is_normalized(store: dict[str, Any], source: zarr.Array[Any]) -> None: + """New axes and reversals take the cover-and-residual path through Zarr.""" + data = np.asarray(source[:]) + planner = LazyArray(cast(Any, source)).with_reader(_NoSynchronousRead()) + views_and_expected = ( + (planner.lazy[None, 5:35, 3:27], data[None, 5:35, 3:27]), + (planner.lazy[35:4:-3, ::-2], data[35:4:-3, ::-2]), + (planner.lazy[5:10, :, None], data[5:10, :, None]), + ) + + async def scenario() -> tuple[np.ndarray[Any, Any], ...]: + async_source = await zarr.api.asynchronous.open_array(store=store) + return tuple( + await asyncio.gather( + *(read_through(view, async_source) for view, _expected in views_and_expected) + ) + ) + + for result, (_view, expected) in zip(asyncio.run(scenario()), views_and_expected, strict=True): + np.testing.assert_array_equal(result, expected) + + +def test_decoded_chunk_cache(store: dict[str, Any], source: zarr.Array[Any]) -> None: """Cache decoded cells; place each view's share with `chunk_local_selection`. A tile server reads many overlapping views of the same array. Fetching @@ -115,16 +181,22 @@ def cell_selection(part: Partition) -> tuple[slice, ...]: slice(lo, hi) for lo, hi in zip(domain.inclusive_min, domain.exclusive_max, strict=True) ) - async def fetch_missing(parts: tuple[Partition, ...], async_source: AsyncSource) -> None: + async def fetch_missing( + parts: tuple[Partition, ...], async_source: zarr.AsyncArray[Any] + ) -> None: missing = { cell_key(part): cell_selection(part) for part in parts if cell_key(part) not in cache } cells = await asyncio.gather( *(async_source.getitem(selection) for selection in missing.values()) ) - cache.update(zip(missing.keys(), cells, strict=True)) + cache.update( + (key, np.asanyarray(cell)) for key, cell in zip(missing.keys(), cells, strict=True) + ) - async def read_cached(view: LazyArray, async_source: AsyncSource) -> np.ndarray: + async def read_cached( + view: LazyArray, async_source: zarr.AsyncArray[Any] + ) -> np.ndarray[Any, Any]: parts = tuple(view.parts()) await fetch_missing(parts, async_source) out = np.empty(view.shape, dtype=view.dtype) @@ -134,9 +206,9 @@ async def read_cached(view: LazyArray, async_source: AsyncSource) -> np.ndarray: async def scenario() -> tuple[np.ndarray, np.ndarray]: async_source = await zarr.api.asynchronous.open_array(store=store) - first = await read_cached(LazyArray(source).lazy[5:25, 3:27], async_source) + first = await read_cached(LazyArray(cast(Any, source)).lazy[5:25, 3:27], async_source) # The second view overlaps the first, so most cells are already cached. - second = await read_cached(LazyArray(source).lazy[15:35, ::2], async_source) + second = await read_cached(LazyArray(cast(Any, source)).lazy[15:35, ::2], async_source) return first, second first, second = asyncio.run(scenario()) @@ -147,40 +219,26 @@ async def scenario() -> tuple[np.ndarray, np.ndarray]: assert len(cache) == 12 -def test_query_parts_fall_back(store: dict[str, Any], source: zarr.Array) -> None: - """A query part refuses to lower; the wrapper's own reader is the fallback. +def test_query_parts_fetch_cover_asynchronously( + store: dict[str, Any], source: zarr.Array[Any] +) -> None: + """A query part fetches its cover asynchronously and gathers in memory. A gather (`oindex`, `vindex`, a mask) has no single-slab spelling, so - `source_selection` raises `ValueError` instead of guessing one. A consumer - mixing selection kinds catches that and resolves the part through - `part.view.result()`, which reads through the wrapped array's reader. - Catching the dedicated subclass — not bare `ValueError` — keeps a genuine - defect in the lowering loud instead of silently degrading every part to - the fallback path. + `source_selection` raises `NoBasicSelectionError` instead of guessing one. + The adapter catches that dedicated error, fetches the part's ascending + cover through AsyncArray, and applies its residual transform in memory. """ - view = LazyArray(source).lazy.oindex[[30, 2, 2], 4:10] + view = ( + LazyArray(cast(Any, source)).with_reader(_NoSynchronousRead()).lazy.oindex[[30, 2, 2], 4:10] + ) async def scenario() -> np.ndarray: async_source = await zarr.api.asynchronous.open_array(store=store) - parts = tuple(view.parts()) - out = np.empty(view.shape, dtype=view.dtype) - - async def resolve(part: Partition) -> tuple[Partition, np.ndarray]: - try: - selection = part.source_selection - except NoBasicSelectionError: - # The gathered axis needs a lookup, not a slab: read this part - # through the wrapper (synchronously here; a real consumer - # might push it to a thread, or fetch the part's bounding box - # and gather in memory). - return part, np.asarray(part.view.result()) - return part, np.asarray(await async_source.getitem(selection)) - - for part, block in await asyncio.gather(*(resolve(part) for part in parts)): - out[part.out_selection] = block - return out + return await read_through(view, async_source) - assert np.array_equal(asyncio.run(scenario()), source.oindex[[30, 2, 2], 4:10]) + expected = np.asarray(source[:])[[30, 2, 2]][:, 4:10] + assert np.array_equal(asyncio.run(scenario()), expected) if __name__ == "__main__": diff --git a/packages/zarr-indexing/justfile b/packages/zarr-indexing/justfile index 1b7164f647..a484ca65b0 100644 --- a/packages/zarr-indexing/justfile +++ b/packages/zarr-indexing/justfile @@ -30,7 +30,7 @@ lint: # Type-check the package sources, documentation Python, and their contract tests typecheck: - uv run --group test --with pyright pyright + uv run --project ../.. --group test --with-editable . --with pyright pyright # Run everything CI runs for this package check: lint typecheck test test-tensorstore docs-check diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index 60f69a1dbb..ec39aa0062 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -131,6 +131,7 @@ filterwarnings = [ include = [ "src", "docs/snippets", + "examples/lazy_indexing_asyncio", "tests/test_doc_examples.py", ] enableExperimentalFeatures = true diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index 2fb671e5e8..eadc7e75fb 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -61,11 +61,13 @@ type BasicSelection = tuple[int | slice | None, ...] """A selection in NumPy's basic-indexing dialect: one selector per axis. -The request vocabulary of `source[selection]` under basic indexing — also the -parameter type of `zarr.AsyncArray.getitem`. Denotationally a product of -arithmetic progressions: an `int` reads one coordinate and drops its axis, a -`slice` reads an arithmetic progression, and `None` fabricates an axis no -source axis backs. Produced by +The request vocabulary of `source[selection]` under NumPy basic indexing. +It is deliberately wider than `zarr.AsyncArray.getitem`'s selection type: +Zarr rejects `None` and negative-step slices, so an integration with that +backend must normalize those two forms before making the request. +Denotationally this is a product of arithmetic progressions: an `int` reads one +coordinate and drops its axis, a `slice` reads an arithmetic progression, and +`None` fabricates an axis no source axis backs. Produced by [`IndexTransform.as_basic_selection`][zarr_indexing.transform.IndexTransform.as_basic_selection] and the `Partition` lowering properties; which cells a value denotes is relative to the array it is applied to. From 7d160bad98eb447be073eda03af535dfb251c05e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Fri, 14 Aug 2026 18:23:01 +0200 Subject: [PATCH 16/16] chore(zarr-indexing): number PR 292 changelog fragments Rename every fragment to the PR-numbered convention required by the changelog validator and update the entries to describe the AsyncArray and result-buffer fixes. Assisted-by: Codex:gpt-5 --- .../zarr-indexing/changes/+basic-selection-alias.feature.md | 6 ------ .../changes/+no-basic-selection-error.feature.md | 6 ------ .../{+part-view-projection.bugfix.md => 292.bugfix.md} | 2 +- .../changes/{+asyncio-example.doc.md => 292.doc.md} | 4 ++-- .../changes/{+decompose.feature.md => 292.feature.1.md} | 3 ++- packages/zarr-indexing/changes/292.feature.2.md | 6 ++++++ .../changes/{+part-lowering.feature.md => 292.feature.3.md} | 6 ++++-- .../changes/{+result-into.feature.md => 292.feature.4.md} | 4 +++- .../{+stateful-part-descent.feature.md => 292.feature.5.md} | 2 +- packages/zarr-indexing/changes/292.feature.md | 6 ++++++ 10 files changed, 25 insertions(+), 20 deletions(-) delete mode 100644 packages/zarr-indexing/changes/+basic-selection-alias.feature.md delete mode 100644 packages/zarr-indexing/changes/+no-basic-selection-error.feature.md rename packages/zarr-indexing/changes/{+part-view-projection.bugfix.md => 292.bugfix.md} (91%) rename packages/zarr-indexing/changes/{+asyncio-example.doc.md => 292.doc.md} (70%) rename packages/zarr-indexing/changes/{+decompose.feature.md => 292.feature.1.md} (82%) create mode 100644 packages/zarr-indexing/changes/292.feature.2.md rename packages/zarr-indexing/changes/{+part-lowering.feature.md => 292.feature.3.md} (83%) rename packages/zarr-indexing/changes/{+result-into.feature.md => 292.feature.4.md} (86%) rename packages/zarr-indexing/changes/{+stateful-part-descent.feature.md => 292.feature.5.md} (92%) create mode 100644 packages/zarr-indexing/changes/292.feature.md diff --git a/packages/zarr-indexing/changes/+basic-selection-alias.feature.md b/packages/zarr-indexing/changes/+basic-selection-alias.feature.md deleted file mode 100644 index 162b5a8d81..0000000000 --- a/packages/zarr-indexing/changes/+basic-selection-alias.feature.md +++ /dev/null @@ -1,6 +0,0 @@ -`BasicSelection` — the public alias for the lowering's output vocabulary, -`tuple[int | slice | None, ...]`, exported at the package root. The name is -the destination contract's: it is what `zarr.AsyncArray.getitem` calls its -selection parameter, and what NumPy calls basic indexing. Use it to annotate a -consumer-owned I/O protocol (`async def getitem(self, selection: -BasicSelection)`), as the asyncio example now does. diff --git a/packages/zarr-indexing/changes/+no-basic-selection-error.feature.md b/packages/zarr-indexing/changes/+no-basic-selection-error.feature.md deleted file mode 100644 index 0609c75910..0000000000 --- a/packages/zarr-indexing/changes/+no-basic-selection-error.feature.md +++ /dev/null @@ -1,6 +0,0 @@ -`as_basic_selection` refusals now raise `NoBasicSelectionError`, a dedicated -`ValueError` subclass exported at the package root. The documented consumer -fallback (`except NoBasicSelectionError: part.view.result()`) can no longer -silently absorb a genuine defect in the lowering, which bare `except -ValueError` did. Existing catch sites keep working: the subclass is caught by -`except ValueError` unchanged. diff --git a/packages/zarr-indexing/changes/+part-view-projection.bugfix.md b/packages/zarr-indexing/changes/292.bugfix.md similarity index 91% rename from packages/zarr-indexing/changes/+part-view-projection.bugfix.md rename to packages/zarr-indexing/changes/292.bugfix.md index 29d43def61..542be813d8 100644 --- a/packages/zarr-indexing/changes/+part-view-projection.bugfix.md +++ b/packages/zarr-indexing/changes/292.bugfix.md @@ -1,6 +1,6 @@ Resolving a partition's view on its own (`part.view.result()`) now hands the reader the same `ReadContext` the parent's partitioned `result()` passes: the -paired `ChunkProjection` rides along instead of arriving as `projection=None`. +paired `ChunkProjection` now rides along instead of arriving as `projection=None`. A custom reader keyed on `projection.chunk_coords` — a decoded-chunk cache — now behaves identically on both paths, which the dask example's task-per-partition pattern relies on. `with_reader` and repartitioning keep diff --git a/packages/zarr-indexing/changes/+asyncio-example.doc.md b/packages/zarr-indexing/changes/292.doc.md similarity index 70% rename from packages/zarr-indexing/changes/+asyncio-example.doc.md rename to packages/zarr-indexing/changes/292.doc.md index e8a108d2eb..aa2f71178d 100644 --- a/packages/zarr-indexing/changes/+asyncio-example.doc.md +++ b/packages/zarr-indexing/changes/292.doc.md @@ -2,6 +2,6 @@ Added an asyncio integration example (`examples/lazy_indexing_asyncio/`), peer to the Dask example: `parts()` driven by `asyncio.gather` against `zarr.AsyncArray`, using `Partition.source_selection` for per-part fetches, a decoded-chunk cache keyed on each cell's `chunk_domain` origin and placed with -`chunk_local_selection`, and the `NoBasicSelectionError` fallback for query -parts. The +`chunk_local_selection`, and an ascending-cover fallback for queries, new +axes, and negative-step slices that `AsyncArray.getitem` cannot take directly. The integrations guide gained a matching "Consumer-owned I/O" section. diff --git a/packages/zarr-indexing/changes/+decompose.feature.md b/packages/zarr-indexing/changes/292.feature.1.md similarity index 82% rename from packages/zarr-indexing/changes/+decompose.feature.md rename to packages/zarr-indexing/changes/292.feature.1.md index 30d0814865..cd93e0d991 100644 --- a/packages/zarr-indexing/changes/+decompose.feature.md +++ b/packages/zarr-indexing/changes/292.feature.1.md @@ -5,7 +5,8 @@ residual against `source[cover]` reads exactly the transform's cells (and the cover, read as the transform it denotes, composes with the residual back to the original). Queries decompose into their bounding interval plus a block-local gather, so a consumer can keep even query parts on its own I/O -path. `decompose_unit_step()` is the variant whose cover is contiguous and +path. The AsyncArray adapter also uses this factorization for NumPy selectors +that Zarr rejects directly. `decompose_unit_step()` is the variant whose cover is contiguous and ascending, the factorization `UnitStepReader` reads through; both live where the readers now source their own decomposition, so the planned request and the executed read cannot drift apart. The one refusal is a negative output diff --git a/packages/zarr-indexing/changes/292.feature.2.md b/packages/zarr-indexing/changes/292.feature.2.md new file mode 100644 index 0000000000..b4ecf9c389 --- /dev/null +++ b/packages/zarr-indexing/changes/292.feature.2.md @@ -0,0 +1,6 @@ +`as_basic_selection` refusals now raise `NoBasicSelectionError`, a dedicated +`ValueError` subclass exported at the package root. A consumer can catch it to +fetch the transform's basic cover and apply its residual without silently +absorbing a genuine defect in the lowering, which bare `except ValueError` +would do. Existing catch sites keep working: the subclass is caught by +`except ValueError` unchanged. diff --git a/packages/zarr-indexing/changes/+part-lowering.feature.md b/packages/zarr-indexing/changes/292.feature.3.md similarity index 83% rename from packages/zarr-indexing/changes/+part-lowering.feature.md rename to packages/zarr-indexing/changes/292.feature.3.md index a98f9320d5..1a5d09152c 100644 --- a/packages/zarr-indexing/changes/+part-lowering.feature.md +++ b/packages/zarr-indexing/changes/292.feature.3.md @@ -10,8 +10,9 @@ async store, an HTTP range endpoint): `NoBasicSelectionError` (a `ValueError` subclass) instead of guessing a slab. - `Partition.source_selection` is that lowering of a part's global read, so an - async consumer's whole loop is - `out[part.out_selection] = await source.getitem(part.source_selection)`. + async consumer can hand compatible selectors directly to its backend. + Backends narrower than NumPy normalize the remaining selectors at their + integration boundary. - `Partition.chunk_local_selection` is the same read relative to `projection.chunk_domain`'s origin, for decoded-chunk caches. The domain names its cell in the source's own coordinates even for a view partitioning @@ -21,3 +22,4 @@ async store, an HTTP range endpoint): `ChainedIndexingStateMachine` gained an invariant that runs both documented assembly loops literally under plain NumPy semantics, with no reader involved. +Backend-specific acceptance is covered separately by the asyncio example. diff --git a/packages/zarr-indexing/changes/+result-into.feature.md b/packages/zarr-indexing/changes/292.feature.4.md similarity index 86% rename from packages/zarr-indexing/changes/+result-into.feature.md rename to packages/zarr-indexing/changes/292.feature.4.md index ba10c447c8..fc4eefa081 100644 --- a/packages/zarr-indexing/changes/+result-into.feature.md +++ b/packages/zarr-indexing/changes/292.feature.4.md @@ -8,4 +8,6 @@ final slot, and a `numpy.ma` masked buffer keeps a masked source's mask. Validation rejects what `result()`'s own allocation made impossible: a plain buffer for a masked source, which would silently present the values beneath the mask as data, and a buffer sharing memory with the wrapped array, where -each part would overwrite cells the parts after it still have to read. +each part would overwrite cells the parts after it still have to read. It also +rejects internally overlapping strided buffers, whose logical cells cannot +hold distinct result values. diff --git a/packages/zarr-indexing/changes/+stateful-part-descent.feature.md b/packages/zarr-indexing/changes/292.feature.5.md similarity index 92% rename from packages/zarr-indexing/changes/+stateful-part-descent.feature.md rename to packages/zarr-indexing/changes/292.feature.5.md index 69c1b362bc..160d558e30 100644 --- a/packages/zarr-indexing/changes/+stateful-part-descent.feature.md +++ b/packages/zarr-indexing/changes/292.feature.5.md @@ -16,6 +16,6 @@ state it could not previously reach: is caught by the same NumPy model as everything else — no reader that ignores the projection can see that. -Subclasses inherit all three. A partitioning declared as explicit per-axis +Subclasses inherit all three NumPy-semantic checks. A partitioning declared as explicit per-axis sizes describes the source's extents, so it is skipped where a descent has narrowed the base it would have to sum to. diff --git a/packages/zarr-indexing/changes/292.feature.md b/packages/zarr-indexing/changes/292.feature.md new file mode 100644 index 0000000000..c09de45d37 --- /dev/null +++ b/packages/zarr-indexing/changes/292.feature.md @@ -0,0 +1,6 @@ +`BasicSelection` — the public alias for the lowering's output vocabulary, +`tuple[int | slice | None, ...]`, exported at the package root. The name is +NumPy's basic-indexing contract. It is wider than +`zarr.AsyncArray.getitem`'s same-named selection type because NumPy also +accepts `None` and negative-step slices; the asyncio example narrows or +normalizes those forms at the backend boundary.