diff --git a/.github/workflows/zarr-indexing.yml b/.github/workflows/zarr-indexing.yml index 3b106e16aa..142b91618a 100644 --- a/.github/workflows/zarr-indexing.yml +++ b/.github/workflows/zarr-indexing.yml @@ -43,8 +43,8 @@ jobs: uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 - name: Set up Python ${{ matrix.python-version }} run: uv python install ${{ matrix.python-version }} - # The suite imports nothing from `zarr`; it runs against the repo-root - # environment, with this package as an editable overlay, to share the + # The suite runs against the repo-root environment, with this package + # as an editable overlay, for Zarr codec integration tests and the # parent project's pinned test toolchain. The recipes carry that # invocation; this step only fixes the interpreter the matrix asked for. - name: Sync test dependency group diff --git a/packages/zarr-indexing/benchmarks/README.md b/packages/zarr-indexing/benchmarks/README.md index 28ef3e6a16..c2fb0af8dc 100644 --- a/packages/zarr-indexing/benchmarks/README.md +++ b/packages/zarr-indexing/benchmarks/README.md @@ -26,3 +26,19 @@ These scripts measure planning rather than codec or storage throughput. Repeat measurements with alternating operation order before interpreting small timing differences. Preserve raw benchmark output as an experiment artifact rather than accumulating successive result tables in this README. + +## Experimental execution + +The execution follow-up adds two scripts: + +```sh +hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/execution.py +hatch run test.py3.12-minimal:python packages/zarr-indexing/benchmarks/execution_io.py +``` + +`execution.py` compares Zarr indexers, declarative projections, and execution +selectors, including retained rows, borrowed inputs, snapshots, and shard +lowering. `execution_io.py` verifies and measures MemoryStore reads and writes +through real codec pipelines. Neither establishes filesystem or cloud throughput. +The prototype remains opt-in; planning wins alone do not justify replacing the +existing indexers. diff --git a/packages/zarr-indexing/benchmarks/execution.py b/packages/zarr-indexing/benchmarks/execution.py new file mode 100644 index 0000000000..4d9fee7029 --- /dev/null +++ b/packages/zarr-indexing/benchmarks/execution.py @@ -0,0 +1,181 @@ +"""Compare PR planner with existing Zarr indexers; no storage I/O. + +Grid construction and input selection allocation are excluded for both sides. +Selection compilation/indexer construction and complete streaming walks are included. +Run with the worktree src and packages/zarr-indexing/src on PYTHONPATH in Hatch. +""" + +from __future__ import annotations + +import json +import math +import statistics +import time +import tracemalloc +from typing import TYPE_CHECKING, Any + +import numpy as np +import zarr.core.indexing as zi +from zarr.core.chunk_grids import ChunkGrid + +from zarr_indexing import IndexTransform, plan_chunks +from zarr_indexing._execution import execute_selection +from zarr_indexing.grid import dimension_grids_from_chunks + +if TYPE_CHECKING: + from collections.abc import Callable + + +def consume(iterator: Any) -> int: + return sum(1 for _ in iterator) + + +def measure(op: Callable[[], Any], repeats: int = 31) -> dict[str, float]: + op() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + op() + samples.append((time.perf_counter() - start) * 1000) + tracemalloc.start() + result = op() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + del result + return {"ms": statistics.median(samples), "peak_mib": peak / 2**20} + + +def main() -> None: + i = np.arange(1000) + cases = [ + ("basic_10000_chunks", (1000, 1000), (10, 10), (slice(None), slice(None)), "basic"), + ( + "sorted_1M_points_10_chunks", + (1_000_000,), + (100_000,), + (np.arange(1_000_000),), + "orthogonal", + ), + ( + "sorted_coordinate_1M_points_10_chunks", + (1_000_000,), + (100_000,), + (np.arange(1_000_000),), + "coordinate", + ), + ("correlated_dense", (1000, 1000), (10, 10), (i, i), "coordinate"), + ("correlated_sparse", (10000, 10000), (10, 10), (i * 9, i * 9), "coordinate"), + ( + "independent_one_chunk", + (1000,) * 3, + (1000,) * 3, + (i[:, None], i[:, None], i[None, :]), + "coordinate", + ), + ( + "independent_100_chunks", + (1000,) * 3, + (100,) * 3, + (i[:, None], i[:, None], i[None, :]), + "coordinate", + ), + ] + results = { + name: compare_case(shape, chunks, selection, mode) + for name, shape, chunks, selection, mode in cases + } + print(json.dumps(results, indent=2)) + + +def compare_case( + shape: tuple[int, ...], chunks: tuple[int, ...], selection: Any, mode: str +) -> dict[str, Any]: + zg = ChunkGrid.from_sizes(shape, chunks) + pg = dimension_grids_from_chunks(chunks, shape) + cls = { + "basic": zi.BasicIndexer, + "orthogonal": zi.OrthogonalIndexer, + "coordinate": zi.CoordinateIndexer, + }[mode] + + def baseline() -> Any: + return cls(selection, shape, zg) + + def new() -> Any: + base = IndexTransform.from_shape(shape) + transform = ( + base[selection] + if mode == "basic" + else (base.oindex[selection] if mode == "orthogonal" else base.vindex[selection]) + ) + return plan_chunks(transform, pg).partition() + + old_coords = [tuple(p.chunk_coords) for p in baseline()] + partition = new() + new_coords = [tuple(p.chunk_coords) for p in partition] + assert old_coords == new_coords + expected_size = ( + math.prod(shape) + if mode == "basic" + else ( + selection[0].size + if mode == "orthogonal" + else math.prod(np.broadcast_shapes(*(s.shape for s in selection))) + ) + ) + assert sum(math.prod(p.cell_transform.domain.shape) for p in partition) == expected_size + # Alternate evaluation order across rounds to reduce temporal bias. + rounds = [] + + def immediate() -> Any: + return execute_selection( + selection, + shape, + zg._dimensions, + ownership="borrow", + mode={"basic": "basic", "orthogonal": "orthogonal", "coordinate": "vectorized"}[mode], + ) + + assert [tuple(p.chunk_coords) for p in immediate()] == old_coords + operations = { + "zarr_setup": baseline, + "new_setup": new, + "zarr_walk": lambda: consume(baseline()), + "new_walk": lambda: consume(new()), + "immediate_setup": immediate, + "immediate_walk": lambda: consume(immediate()), + "zarr_retained": lambda: list(baseline()), + "borrowed_retained": lambda: list(immediate()), + "snapshot_retained": lambda: list( + execute_selection( + selection, + shape, + zg._dimensions, + mode={"basic": "basic", "orthogonal": "orthogonal", "coordinate": "vectorized"}[ + mode + ], + ) + ), + "shard_retained": lambda: list(immediate().lower("shard")), + } + for round_id in range(3): + names = list(operations) + if round_id % 2: + names.reverse() + rounds.append({key: measure(operations[key]) for key in names}) + return { + "chunks": len(old_coords), + "elements": expected_size, + "metrics": { + key: { + "ms": statistics.median(r[key]["ms"] for r in rounds), + "peak_mib": max(r[key]["peak_mib"] for r in rounds), + "round_ms": [r[key]["ms"] for r in rounds], + } + for key in operations + }, + } + + +if __name__ == "__main__": + main() diff --git a/packages/zarr-indexing/benchmarks/execution_io.py b/packages/zarr-indexing/benchmarks/execution_io.py new file mode 100644 index 0000000000..124d942fe0 --- /dev/null +++ b/packages/zarr-indexing/benchmarks/execution_io.py @@ -0,0 +1,128 @@ +"""Measure actual in-memory Zarr codec reads/writes, including plan construction. + +Storage, source data, grids, and replacement data are prepared before timing. +Both paths use the same array and codec pipeline. These are local MemoryStore +measurements, not estimates of cloud or filesystem throughput. +""" + +from __future__ import annotations + +import asyncio +import json +import statistics +import time +import tracemalloc +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +import zarr.api.asynchronous as za +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.indexing import BasicIndexer, CoordinateIndexer +from zarr.storage import MemoryStore + +from zarr_indexing._execution import execute_selection + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable + + from zarr.core.indexing import Indexer + + +async def measure(operation: Callable[[], Awaitable[Any]]) -> dict[str, float]: + await operation() + samples = [] + for _ in range(9): + start = time.perf_counter() + await operation() + samples.append((time.perf_counter() - start) * 1000) + tracemalloc.start() + result = await operation() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + del result + return {"ms": statistics.median(samples), "peak_mib": peak / 2**20} + + +async def run_case(case: str, sharded: bool) -> dict[str, Any]: + shape: tuple[int, ...] + chunks: tuple[int, ...] + selection: Any + if case == "basic": + shape, chunks = (128, 128), (16, 16) + selection, mode = (slice(1, 127), slice(1, 127)), "basic" + elif case == "sorted": + shape, chunks = (100000,), (10000,) + selection, mode = (np.arange(1, 99999),), "vectorized" + else: + shape, chunks = (64, 64, 64), (16, 16, 16) + i = np.arange(64) + selection, mode = (i[:, None], i[:, None], i[None, :]), "vectorized" + kwargs: dict[str, Any] = {} + if sharded: + kwargs["shards"] = tuple(c * 2 for c in chunks) + array = await za.create_array( + store=MemoryStore(), shape=shape, chunks=chunks, dtype="int64", **kwargs + ) + source = np.arange(np.prod(shape), dtype=np.int64).reshape(shape) + await array.setitem(Ellipsis, source) + expected = source[selection] + replacement = expected + 1 + grids = array._chunk_grid._dimensions + prototype = default_buffer_prototype() + cls = BasicIndexer if mode == "basic" else CoordinateIndexer + + async def old_read() -> Any: + result = await array._get_selection( + cls(selection, shape, array._chunk_grid), prototype=prototype + ) + # Zarr's public coordinate API restores sel_shape outside the codec + # entry point; include that view operation in the baseline. + return np.asarray(result).reshape(expected.shape) + + async def new_read() -> Any: + plan = execute_selection(selection, shape, grids, mode=mode, ownership="borrow") + return await array._get_selection( + cast("Indexer", plan.lower("shard" if sharded else "numpy")), prototype=prototype + ) + + async def old_write() -> None: + await array._set_selection( + cls(selection, shape, array._chunk_grid), + replacement if mode == "basic" else replacement.reshape(-1), + prototype=prototype, + ) + + async def new_write() -> None: + plan = execute_selection( + selection, shape, grids, mode=mode, ownership="borrow", access="write" + ) + await array._set_selection( + cast("Indexer", plan.lower("shard" if sharded else "numpy")), + replacement, + prototype=prototype, + ) + + np.testing.assert_array_equal(await old_read(), expected) + np.testing.assert_array_equal(await new_read(), expected) + results = { + "zarr_read": await measure(old_read), + "new_read": await measure(new_read), + "zarr_write": await measure(old_write), + "new_write": await measure(new_write), + } + expected_full = source.copy() + expected_full[selection] = replacement + np.testing.assert_array_equal(await array.getitem(Ellipsis), expected_full) + return results + + +async def main() -> None: + results = {} + for case in ("basic", "sorted", "components"): + for sharded in (False, True): + results[case + ("_sharded" if sharded else "")] = await run_case(case, sharded) + print(json.dumps(results, indent=2)) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/packages/zarr-indexing/docs/api/grid.md b/packages/zarr-indexing/docs/api/grid.md index b7c376eb85..f5609fb3d3 100644 --- a/packages/zarr-indexing/docs/api/grid.md +++ b/packages/zarr-indexing/docs/api/grid.md @@ -16,6 +16,12 @@ become `VaryingDimension` instances. `DimensionGridLike` remains the narrow protocol used by the chunk planner, while `EdgeDimensionGrid` is kept for explicit edge-based and coordinate-origin examples. +`RegularDimensionGridLike` is an optional structural capability for fast paths. +It adds `size` and `extent` to the grid protocol: chunk `k` must start at +`k * size` and declare `size` elements, with `extent` clipping only the valid +data in the last chunk. Zarr's existing uniform dimensions satisfy this +capability without conversion to this package's concrete classes. + Zarr's array implementation can later import these compact grid types from `zarr_indexing`; this package intentionally has no import dependency on Zarr. diff --git a/packages/zarr-indexing/docs/guide/execution.md b/packages/zarr-indexing/docs/guide/execution.md new file mode 100644 index 0000000000..66883f6062 --- /dev/null +++ b/packages/zarr-indexing/docs/guide/execution.md @@ -0,0 +1,70 @@ +# Experimental execution consumers + +This follow-up implements a private, opt-in execution experiment in +`zarr_indexing._execution`. It does not replace Zarr's default indexers or change +the public transform algebra. Start with [From a selection to chunk +operations](selection-flow.md) for coordinate spaces, partitioning, and the paired +source/result correspondence. + +## Preparation and lowering + +`execute_selection(selection, shape, grids, mode=...)` interprets literal +coordinates and prepares an `ExecutionPlan`. Basic, orthogonal, and vectorized +modes choose the selection mapping. Orthogonal scalar integers are applied as a +basic selection first, dropping their result axes before the remaining orthogonal +selection. This differs intentionally from the low-level transform `oindex`, +which retains scalar axes. The resulting selectors already have the reduced +rank; they do not need the legacy indexer's later `drop_axes` squeeze. +Negative scalar coordinates still raise: this frontend does not normalize +NumPy-style negative indices. Vectorized mixed-scalar semantics remain those of +the transform algebra; this is not a universal Zarr Indexer replacement. + +`execute_transform(transform, grids)` accepts an existing immutable mapping. +Both entry points share affine-axis planning with declarative tables. Large +basic axes remain implicit; small axes cache at most 128 pieces. Sorted arrays +use a structural regular-grid capability, and connected-component tables lower +directly where supported. Other transforms use declarative projections. + +Iterating a plan lowers it for the NumPy consumer. Each row has `chunk_coords`, +`chunk_selection`, `out_selection`, and `is_complete_chunk`. For a read: + +```text +result[out_selection] = decoded_chunk[chunk_selection] +``` + +`plan.lower("numpy").operations()` also provides `value_shape` and +`selector_kind`, distinguishing basic selectors from paired broadcast coordinate +arrays. `plan.lower("shard")` adapts to the current shard indexer. That adaptation +can expand compact selections into flat coordinates; passing compact plans +through the shard boundary remains future work. + +## Ownership and writes + +`ownership="snapshot"` is the default. `ownership="borrow"` permits borrowing +caller arrays, which must remain unchanged throughout every use of the plan and +its iterators. It does not guarantee zero copies on every planning path. + +Writers prepare with `access="write"`. By default, repeated destinations raise; +`conflicts="last"` explicitly retains the last value in row-major request order +before dispatch. A write plan that discards overwritten values is not a read +plan. Reads preserve repeated positions. + +Completeness refers to the **valid data extent**, matching declarative coverage +and Zarr's existing indexers. A full boundary write can skip its read: Zarr's +merge helper fills the rest of a declared codec buffer when the selected slab +is smaller. A singleton selected once is full regardless of stride magnitude. +The execution adapter emits a true flag only for layouts the current codec's +merge shortcut can safely consume; reverse and general coordinate layouts stay +conservative. Coverage alone does not prove the values are in codec-buffer order. + +Preparation validates supported bounds and write policies, but does not provide +transactional I/O or protect against mutation of borrowed arrays. + +## Verification boundary + +`tests/test_indexing_execution.py` lives in the package suite, which CI runs in +the repo-root environment with Zarr installed. It skips only when Zarr itself is +unavailable. Both codec pipelines run against v2, v3, and sharded layouts, +including orthogonal scalar removal. Boundary-write tests count storage reads, +so marking a full boundary cell partial cannot silently add read-modify-write. +The package runtime still has no dependency on Zarr. diff --git a/packages/zarr-indexing/justfile b/packages/zarr-indexing/justfile index 44874bc0be..a448ba238d 100644 --- a/packages/zarr-indexing/justfile +++ b/packages/zarr-indexing/justfile @@ -5,8 +5,8 @@ default: @just --list -# Nothing here imports `zarr`; the suite runs against the repo-root -# environment so it shares the parent project's pinned test toolchain +# The suite runs against the repo-root environment for Zarr codec integration +# tests and the parent project's pinned test toolchain # (hypothesis, ruff, pyright). The repo is not a uv workspace, so this package # is layered in as an editable overlay — the same invocation CI uses. # Run the test suite; extra args are passed to pytest diff --git a/packages/zarr-indexing/mkdocs.yml b/packages/zarr-indexing/mkdocs.yml index a541e5d599..be64319c35 100644 --- a/packages/zarr-indexing/mkdocs.yml +++ b/packages/zarr-indexing/mkdocs.yml @@ -35,6 +35,7 @@ nav: - guide/index.md - Indexing patterns: guide/patterns.md - From a selection to chunk operations: guide/selection-flow.md + - Experimental execution consumers: guide/execution.md - Integration boundaries: guide/integrations.md - Examples: - Lazy indexing a NumPy array: examples/lazy_indexing_numpy.md diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index 953a432f07..c39d53e1ac 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -48,10 +48,9 @@ Changelog = "https://github.com/zarr-developers/zarr-python/blob/main/packages/z Documentation = "https://zarr-indexing.readthedocs.io/" [dependency-groups] -# The package and its tests import nothing from `zarr`: chunk resolution -# consumes the `DimensionGridLike` protocol and the tests use this package's -# own grids (`zarr_indexing.grid`). The suite is nevertheless run from the -# repo-root environment (`just test`, the invocation CI uses) so that it sees +# The package imports nothing from `zarr`: chunk resolution consumes the +# `DimensionGridLike` protocol. Codec integration tests optionally import Zarr. +# The suite runs from the repo-root environment (`just test`, as in CI), providing # the same pinned toolchain — hypothesis, ruff, pyright — as the parent # project; the repo is not a uv workspace, so this package is layered in as an # editable overlay. `hypothesis` arrives via the `testing` extra, which is diff --git a/packages/zarr-indexing/src/zarr_indexing/_axis_plan.py b/packages/zarr-indexing/src/zarr_indexing/_axis_plan.py new file mode 100644 index 0000000000..4943773a36 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/_axis_plan.py @@ -0,0 +1,51 @@ +"""Shared implicit affine-axis planning for tables and immediate selectors.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NamedTuple + +from zarr_indexing._affine import checked_affine +from zarr_indexing.chunk_resolution import _data_size # pyright: ignore[reportPrivateUsage] + +if TYPE_CHECKING: + from collections.abc import Iterator + + from zarr_indexing.grid import DimensionGridLike + + +class AxisRun(NamedTuple): + chunk: int + chunk_start: int + data_extent: int + local_start: int + nitems: int + position: int + + +def axis_runs(start: int, stride: int, nitems: int, grid: DimensionGridLike) -> Iterator[AxisRun]: + """Intersect an affine request with chunks, in request traversal order. + + Source endpoints are checked before the first run. Request positions and + counts retain Python integer precision; only storage coordinates must fit + intp. Stride zero keeps repetitions symbolic in a single run. + """ + if nitems == 0: + return + grid.index_to_chunk(checked_affine(start, stride, 0)) + grid.index_to_chunk(checked_affine(start, stride, nitems - 1)) + position = 0 + while position < nitems: + coordinate = start + position * stride + chunk = grid.index_to_chunk(coordinate) + offset = grid.chunk_offset(chunk) + local = coordinate - offset + extent = _data_size(grid, chunk) + if stride == 0: + count = nitems + else: + count = min( + nitems - position, + (extent - 1 - local) // stride + 1 if stride > 0 else local // -stride + 1, + ) + yield AxisRun(chunk, offset, extent, local, count, position) + position += count diff --git a/packages/zarr-indexing/src/zarr_indexing/_execution.py b/packages/zarr-indexing/src/zarr_indexing/_execution.py new file mode 100644 index 0000000000..debc83127f --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/_execution.py @@ -0,0 +1,698 @@ +"""Internal selector-execution prototype; not a public API. + +Prepared work shares literal-coordinate semantics with IndexTransform. Inputs +are snapshotted by default; borrowing, access intent, and duplicate-write policy +are explicit. NumPy and shard consumers lower the same work to their required +selector layout. Neither path changes Zarr's default indexers. +""" + +from __future__ import annotations + +import itertools +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast + +import numpy as np + +from zarr_indexing._affine import checked_affine +from zarr_indexing._axis_plan import axis_runs +from zarr_indexing._selector import as_scalar_index +from zarr_indexing.boundary import split_scalar_axes +from zarr_indexing.chunk_resolution import ChunkPlan, IndexedSet, plan_chunks +from zarr_indexing.errors import BoundsCheckError +from zarr_indexing.grid import DimensionGridLike, RegularDimensionGridLike +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import ( + IndexTransform, + _normalize_basic_selection, # pyright: ignore[reportPrivateUsage] + _positional_slice, # pyright: ignore[reportPrivateUsage] + _resolve_slice_ts, # pyright: ignore[reportPrivateUsage] +) + +type Selector = int | slice | np.ndarray[Any, Any] + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + + +class ExecutionChunk(NamedTuple): + """Legacy four-field codec row produced by a named consumer's lowering.""" + + chunk_coords: tuple[int, ...] + chunk_selection: tuple[Selector, ...] + out_selection: tuple[Selector, ...] + is_complete_chunk: bool + + +@dataclass(frozen=True, slots=True) +class ExecutionPlan: + """Prepared semantic work, independent of a consumer's selector layout. + + Snapshot ownership is the default. Borrowing is an explicit caller promise + to leave arrays unchanged throughout every use of the plan and its iterators. + Writers must prepare with access='write'; duplicate coordinates are rejected + unless conflicts='last' explicitly requests request-order last-write-wins. + """ + + shape: tuple[int, ...] + work: _BasicWork | _SortedWork | _ComponentWork | ChunkPlan + access: Literal["read", "write"] = "read" + ownership: Literal["snapshot", "borrow"] = "snapshot" + conflicts: Literal["error", "last"] = "error" + drop_axes: tuple[int, ...] = () + + def __iter__(self) -> Iterator[ExecutionChunk]: + return iter(self.lower()) + + def lower(self, consumer: Literal["numpy", "shard"] = "numpy") -> LoweredPlan: + """Choose a consumer; shard lowering may materialize paired coordinates.""" + if consumer not in ("numpy", "shard"): + raise ValueError(f"unknown execution consumer: {consumer}") + return LoweredPlan(self, consumer) + + +@dataclass(frozen=True, slots=True) +class LoweredOperation: + """Selectors and their selected-value shape for one named consumer. + + Coordinate selectors are paired and broadcast to value_shape. Basic + selectors use NumPy slice/integer semantics. row is the legacy four-field + tuple consumed by Zarr's current codec pipeline. + """ + + row: ExecutionChunk + value_shape: tuple[int, ...] + selector_kind: Literal["basic", "paired"] + + +@dataclass(frozen=True, slots=True) +class LoweredPlan: + plan: ExecutionPlan + consumer: Literal["numpy", "shard"] + + @property + def shape(self) -> tuple[int, ...]: + return self.plan.shape + + @property + def drop_axes(self) -> tuple[int, ...]: + return self.plan.drop_axes + + def operations(self) -> Iterator[LoweredOperation]: + return _lower(self.plan, self.consumer) + + def __iter__(self) -> Iterator[ExecutionChunk]: + return _consumer_rows(self.plan, self.consumer) + + +class _BasicAxis(NamedTuple): + start: int + step: int + nitems: int + grid: DimensionGridLike + scalar: bool = False + + +@dataclass(frozen=True, slots=True) +class _BasicWork: + axes: tuple[_BasicAxis, ...] + + +@dataclass(frozen=True, slots=True) +class _SortedWork: + coordinates: np.ndarray[Any, Any] + chunk_size: int + first: int + cuts: np.ndarray[Any, Any] + + +@dataclass(frozen=True, slots=True) +class _ComponentWork: + plan: ChunkPlan + local: tuple[np.ndarray[Any, Any], ...] + + +def _axis_rows(axis: _BasicAxis) -> Iterator[tuple[int, int | slice, slice | None, bool]]: + start, step, count, grid, scalar = axis + for run in axis_runs(start, step, count, grid): + yield ( + run.chunk, + run.local_start if scalar else _positional_slice(run.local_start, run.nitems, step), + None if scalar else slice(run.position, run.position + run.nitems), + (step == 1 or run.nitems == 1) + and run.local_start == 0 + and run.nitems == run.data_extent, + ) + + +def _basic_rows(axes: tuple[_BasicAxis, ...]) -> Iterator[ExecutionChunk]: + if any(axis.nitems == 0 for axis in axes): + return + # Cache only small axes. Large axes stay implicit, including before the + # first result; a long dimension must not be drained by itertools.product. + pools: list[tuple[tuple[int, int | slice, slice | None, bool], ...] | None] = [] + for axis in axes: + span = ( + abs( + axis.grid.index_to_chunk(axis.start + (axis.nitems - 1) * axis.step) + - axis.grid.index_to_chunk(axis.start) + ) + + 1 + ) + pools.append(tuple(_axis_rows(axis)) if min(span, axis.nitems) <= 128 else None) + combinations = ( + itertools.product(*cast("list[tuple[Any, ...]]", pools)) + if all(pool is not None for pool in pools) + else _implicit_product(axes, pools) + ) + for pieces in combinations: + yield ExecutionChunk( + tuple(piece[0] for piece in pieces), + tuple(piece[1] for piece in pieces), + tuple(piece[2] for piece in pieces if piece[2] is not None), + all(piece[3] for piece in pieces), + ) + + +def _implicit_product( + axes: tuple[_BasicAxis, ...], + pools: list[tuple[Any, ...] | None], + dimension: int = 0, + prefix: tuple[Any, ...] = (), +) -> Iterator[tuple[Any, ...]]: + if dimension == len(axes): + yield prefix + else: + pool = pools[dimension] + for piece in pool if pool is not None else _axis_rows(axes[dimension]): + yield from _implicit_product(axes, pools, dimension + 1, (*prefix, piece)) + + +def _basic_plan(shape: tuple[int, ...], axes: tuple[_BasicAxis, ...]) -> ExecutionPlan: + # Validate all axes before handing any work to a writer. A late failure + # must not leave earlier chunks of an invalid selection modified. + if all(shape): + for axis in axes: + if axis.nitems: + axis.grid.index_to_chunk(axis.start) + axis.grid.index_to_chunk(axis.start + (axis.nitems - 1) * axis.step) + return ExecutionPlan(shape, _BasicWork(axes)) + + +def _sorted_plan( + coordinates: np.ndarray[Any, Any], grids: tuple[DimensionGridLike, ...] +) -> ExecutionPlan | None: + if ( + len(grids) != 1 + or not isinstance(grids[0], RegularDimensionGridLike) + or coordinates.dtype != np.dtype(np.intp) + or coordinates.ndim != 1 + or coordinates.size == 0 + ): + return None + grid = grids[0] + if coordinates[0] < 0 or coordinates[-1] >= grid.extent or grid.size == 0: + return None + first = int(coordinates[0]) // grid.size + last = int(coordinates[-1]) // grid.size + # Sparse or unordered selections use the shared component/table planner. + if (last - first + 1) * coordinates.size.bit_length() >= coordinates.size: + return None + if not bool((coordinates[:-1] <= coordinates[1:]).all()): + return None + # Only internal boundaries: the last chunk's end may exceed intp. + edges = np.arange(first + 1, last + 1, dtype=np.intp) * grid.size + cuts = np.searchsorted(coordinates, edges) + cuts.setflags(write=False) + return ExecutionPlan(coordinates.shape, _SortedWork(coordinates, grid.size, first, cuts)) + + +def _sorted_rows( + coordinates: np.ndarray[Any, Any], + chunk_size: int, + first: int, + cuts: np.ndarray[Any, Any], +) -> Iterator[ExecutionChunk]: + start = 0 + for relative in range(cuts.size + 1): + stop = int(cuts[relative]) if relative < cuts.size else coordinates.size + if stop > start: + chunk = first + relative + yield ExecutionChunk( + (chunk,), + (coordinates[start:stop] - chunk * chunk_size,), + (slice(start, stop),), + False, + ) + start = stop + + +def execute_selection( + selection: Any, + shape: tuple[int, ...], + dimension_grids: Sequence[DimensionGridLike], + *, + mode: str = "basic", + ownership: Literal["snapshot", "borrow"] = "snapshot", + access: Literal["read", "write"] = "read", + conflicts: Literal["error", "last"] = "error", +) -> ExecutionPlan: + """Compile literal-coordinate selections directly to execution selectors. + + Snapshot ownership is the default, independent of optimizer dispatch. + ownership='borrow' permits borrowing; callers must leave inputs unchanged + for every use of the plan. This is a literal-coordinate frontend, not a + NumPy/Zarr selection-normalization API. + """ + grids = tuple(dimension_grids) + if len(grids) != len(shape): + raise ValueError("dimension_grids must have one entry per storage dimension") + if any(size < 0 for size in shape): + raise ValueError("shape dimensions must be nonnegative") + if mode == "basic": + normalized = _normalize_basic_selection(selection, len(shape)) + if all(sel is not None for sel in normalized): + axes: list[_BasicAxis] = [] + out_shape: list[int] = [] + for dim, (sel, size, grid) in enumerate(zip(normalized, shape, grids, strict=True)): + if isinstance(sel, int): + if not 0 <= sel < size: + raise BoundsCheckError(f"index {sel} is out of bounds for dimension {dim}") + axes.append(_BasicAxis(sel, 1, 1, grid, True)) + else: + assert isinstance(sel, slice) + start, step, _origin, count = _resolve_slice_ts(sel, dim, 0, size) + axes.append(_BasicAxis(start, step, count, grid)) + out_shape.append(count) + return _with_policy( + _basic_plan(tuple(out_shape), tuple(axes)), access, ownership, conflicts + ) + elif mode in ("orthogonal", "vectorized"): + items: tuple[Any, ...] = selection if isinstance(selection, tuple) else (selection,) + if ( + len(shape) == len(items) == 1 + and isinstance(items[0], np.ndarray) + and items[0].ndim == 1 + and items[0].size > 0 + and 0 <= items[0][0] <= items[0][-1] < shape[0] + ): + coordinates = items[0] + if ownership == "snapshot": + coordinates = coordinates.copy() + coordinates.setflags(write=False) + sorted_plan = _sorted_plan(coordinates, grids) + if sorted_plan is not None: + return _with_policy(sorted_plan, access, ownership, conflicts) + else: + raise ValueError(f"unknown indexing mode: {mode}") + base = IndexTransform.from_shape(shape) + if mode == "orthogonal": + # Keep scalar-axis removal at the execution boundary; oindex's + # transform algebra intentionally retains scalar axes as length one. + items = selection if isinstance(selection, tuple) else (selection,) + for item in items: + scalar = as_scalar_index(item) + if scalar is not None and scalar < 0: + raise BoundsCheckError("negative scalar is outside the literal source domain") + scalars, selection = split_scalar_axes(selection, base.domain, "orthogonal") + if scalars is not None: + base = base[scalars] + transform = ( + base[selection] + if mode == "basic" + else base.oindex[selection] + if mode == "orthogonal" + else base.vindex[selection] + ) + return _with_policy(execute_transform(transform, grids), access, ownership, conflicts) + + +def execute_transform( + transform: IndexTransform, + dimension_grids: Sequence[DimensionGridLike], + *, + access: Literal["read", "write"] = "read", + conflicts: Literal["error", "last"] = "error", +) -> ExecutionPlan: + """Lower an existing immutable transform through the same execution paths.""" + grids = tuple(dimension_grids) + if len(grids) != transform.output_rank: + raise ValueError("dimension_grids must have one entry per storage dimension") + domain = transform.domain + axes: list[_BasicAxis] = [] + input_axes: list[int] = [] + for m, grid in zip(transform.output, grids, strict=True): + if isinstance(m, ConstantMap): + axes.append(_BasicAxis(m.offset, 1, 1, grid, True)) + elif isinstance(m, DimensionMap) and m.stride != 0: + axis = m.input_dimension + input_axes.append(axis) + axes.append( + _BasicAxis( + m.offset + m.stride * domain.inclusive_min[axis], + m.stride, + domain.shape[axis], + grid, + ) + ) + else: + break + else: + if input_axes == list(range(domain.ndim)): + return _with_policy( + _basic_plan(domain.shape, tuple(axes)), access, "snapshot", conflicts + ) + if transform.input_rank == transform.output_rank == 1: + (m,) = transform.output + if isinstance(m, ArrayMap) and m.offset == 0 and m.stride == 1: + sorted_plan = _sorted_plan(m.index_array, grids) + if sorted_plan is not None and sorted_plan.shape == domain.shape: + return _with_policy(sorted_plan, access, "snapshot", conflicts) + _validate_storage_bounds(transform, grids) + plan = plan_chunks(transform, grids) + # Factor the plan once, up front: a transform the planner cannot factor + # (a diagonal) is rejected here rather than on first iteration. + partition = plan.partition() + if ( + any(isinstance(m, ArrayMap) for m in transform.output) + and not partition.sets + and all(bool((joint.chunk_start >= 0).all()) for joint in partition.joint_sets) + ): + # Column arithmetic is checked once by JointSet.local; nonnegative + # chunk origins make its final local subtraction safe in intp. + work = _ComponentWork(plan, tuple(joint.local for joint in partition.joint_sets)) + return _with_policy(ExecutionPlan(domain.shape, work), access, "snapshot", conflicts) + return _with_policy(ExecutionPlan(domain.shape, plan), access, "snapshot", conflicts) + + +def _coordinates(transform: IndexTransform, origins: tuple[int, ...]) -> tuple[Selector, ...]: + """Broadcast coordinate selectors in synthetic-axis order without expansion.""" + domain = transform.domain + selectors: list[Selector] = [] + for m, origin in zip(transform.output, origins, strict=True): + if isinstance(m, ConstantMap): + values = np.full((1,) * domain.ndim, m.offset - origin, dtype=np.intp) + elif isinstance(m, DimensionMap): + axis = m.input_dimension + shape = tuple(domain.shape[k] if k == axis else 1 for k in range(domain.ndim)) + positions = np.arange(domain.shape[axis], dtype=np.intp).reshape(shape) + values = checked_affine( + m.offset + m.stride * domain.inclusive_min[axis] - origin, m.stride, positions + ) + else: + values = checked_affine(m.offset - origin, m.stride, m.index_array) + selectors.append(np.broadcast_to(values, domain.shape)) + return tuple(selectors) + + +def _general_rows(plan: ChunkPlan) -> Iterator[ExecutionChunk]: + transform = plan.transform + for projection in plan: + yield ExecutionChunk( + projection.chunk_coords, + _coordinates(projection.chunk_transform, (0,) * transform.output_rank), + _coordinates(projection.cell_transform, transform.domain.inclusive_min), + # Arrays may repeat coordinates; only the direct basic path proves + # full data-extent coverage with a codec-compatible value layout. + False, + ) + + +def _flat_selectors( + selection: tuple[Selector, ...], shape: tuple[int, ...] +) -> tuple[Selector, ...]: + if not selection: + return () + if all(isinstance(sel, np.ndarray) for sel in selection): + arrays = cast("tuple[np.ndarray[Any, Any], ...]", selection) + else: + rank = sum(not isinstance(sel, int) for sel in selection) + basic_arrays: list[np.ndarray[Any, Any]] = [] + axis = 0 + for sel, extent in zip(selection, shape, strict=True): + if isinstance(sel, int): + values = np.asarray(sel, dtype=np.intp).reshape((1,) * rank) + else: + values = ( + np.arange(*sel.indices(extent), dtype=np.intp) + if isinstance(sel, slice) + else sel + ) + values = values.reshape((1,) * axis + (values.size,) + (1,) * (rank - axis - 1)) + axis += 1 + basic_arrays.append(values) + arrays = tuple(basic_arrays) + return tuple(array.reshape(-1) for array in np.broadcast_arrays(*arrays)) + + +def _shard_rows(plan: ExecutionPlan, rows: Iterator[ExecutionChunk]) -> Iterator[ExecutionChunk]: + for row in rows: + if not row.out_selection: + # The shard coordinate indexer returns a length-one vector for a + # 0-D array selector. Integer selectors preserve a scalar result. + yield ExecutionChunk( + row.chunk_coords, + tuple( + int(sel.item()) if isinstance(sel, np.ndarray) and sel.ndim == 0 else sel + for sel in row.chunk_selection + ), + (), + row.is_complete_chunk, + ) + continue + if all(isinstance(sel, np.ndarray) and sel.ndim <= 1 for sel in row.chunk_selection) and ( + all(isinstance(sel, np.ndarray) and sel.ndim <= 1 for sel in row.out_selection) + or (len(row.out_selection) == 1 and isinstance(row.out_selection[0], slice)) + ): + # The existing shard indexer already produces this flat value + # shape. In particular, retain sorted runs' output slices. + yield row + continue + if any( + isinstance(sel, np.ndarray) + or (isinstance(sel, slice) and sel.step is not None and sel.step < 0) + for sel in row.chunk_selection + ): + shape = _chunk_shape(plan, row.chunk_coords) + yield ExecutionChunk( + row.chunk_coords, + _flat_selectors(row.chunk_selection, shape), + _flat_selectors(row.out_selection, plan.shape), + False, + ) + else: + yield row + + +def _validate_storage_bounds( + transform: IndexTransform, grids: tuple[DimensionGridLike, ...] +) -> None: + if 0 in transform.domain.shape: + return + for m, grid in zip(transform.output, grids, strict=True): + if isinstance(m, ConstantMap): + bounds = (m.offset, m.offset) + elif isinstance(m, DimensionMap): + axis = m.input_dimension + bounds = ( + checked_affine(m.offset, m.stride, transform.domain.inclusive_min[axis]), + checked_affine(m.offset, m.stride, transform.domain.exclusive_max[axis] - 1), + ) + else: + mapped = checked_affine(m.offset, m.stride, m.index_array) + bounds = (int(mapped.min()), int(mapped.max())) + grid.index_to_chunk(min(bounds)) + grid.index_to_chunk(max(bounds)) + + +def _with_policy( + plan: ExecutionPlan, + access: Literal["read", "write"], + ownership: Literal["snapshot", "borrow"], + conflicts: Literal["error", "last"], +) -> ExecutionPlan: + if access not in ("read", "write"): + raise ValueError(f"unknown access intent: {access}") + if ownership not in ("snapshot", "borrow"): + raise ValueError(f"unknown ownership policy: {ownership}") + if conflicts not in ("error", "last"): + raise ValueError(f"unknown conflict policy: {conflicts}") + result = ( + plan + if (plan.access, plan.ownership, plan.conflicts) == (access, ownership, conflicts) + else replace(plan, access=access, ownership=ownership, conflicts=conflicts) + ) + if access == "write" and conflicts == "error": + _validate_unique_writes(result) + return result + + +def _validate_unique_writes(plan: ExecutionPlan) -> None: + if 0 in plan.shape or isinstance(plan.work, _BasicWork): + return + work = plan.work + if isinstance(work, _SortedWork): + unique = not bool((work.coordinates[1:] == work.coordinates[:-1]).any()) + else: + source_plan = work.plan if isinstance(work, _ComponentWork) else work + transform = source_plan.transform + referenced: set[int] = set() + for m in transform.output: + if isinstance(m, DimensionMap) and m.stride != 0: + referenced.add(m.input_dimension) + elif isinstance(m, ArrayMap) and m.stride != 0: + referenced.update(m.dependency_axes) + unique = all(size <= 1 or axis in referenced for axis, size in enumerate(plan.shape)) + if unique and transform.index_array_structure != "general": + for axis, size in enumerate(plan.shape): + if size <= 1 or any( + isinstance(m, DimensionMap) and m.input_dimension == axis and m.stride != 0 + for m in transform.output + ): + continue + columns = [ + m.index_array.reshape(-1) + for m in transform.output + if isinstance(m, ArrayMap) and m.dependent_axis == axis and m.stride != 0 + ] + unique &= ( + bool(columns) and np.unique(np.stack(columns, axis=1), axis=0).shape[0] == size + ) + elif unique and any(isinstance(m, ArrayMap) for m in transform.output): + partition = source_plan.partition() + for table in partition.sets: + if isinstance(table, IndexedSet): + unique &= table.stride != 0 and np.unique(table.index).size == table.index.size + for joint in partition.joint_sets: + active_columns = [i for i, stride in enumerate(joint.strides) if stride != 0] + values = joint.index[:, active_columns] + unique &= np.unique(values, axis=0).shape[0] == values.shape[0] + if not unique: + raise ValueError("duplicate writes require conflicts='last'") + + +def _raw_rows(plan: ExecutionPlan) -> Iterator[ExecutionChunk]: + work = plan.work + if isinstance(work, _BasicWork): + return _basic_rows(work.axes) + if isinstance(work, _SortedWork): + return _sorted_rows(work.coordinates, work.chunk_size, work.first, work.cuts) + if isinstance(work, _ComponentWork): + return _component_rows(work) + return _general_rows(work) + + +def _chunk_shape(plan: ExecutionPlan, coords: tuple[int, ...]) -> tuple[int, ...]: + work = plan.work + if isinstance(work, _SortedWork): + return (work.chunk_size,) + grids = ( + tuple(axis.grid for axis in work.axes) + if isinstance(work, _BasicWork) + else work.plan.dimension_grids + if isinstance(work, _ComponentWork) + else work.dimension_grids + ) + return tuple(grid.chunk_size(c) for grid, c in zip(grids, coords, strict=True)) + + +def _ordered_write_rows( + plan: ExecutionPlan, rows: Iterator[ExecutionChunk] +) -> Iterator[ExecutionChunk]: + for row in rows: + chunk = _flat_selectors(row.chunk_selection, _chunk_shape(plan, row.chunk_coords)) + out = _flat_selectors(row.out_selection, plan.shape) + if not out: + yield row + continue + positions = cast("tuple[np.ndarray[Any, Any], ...]", out) + order = np.lexsort(positions[::-1]) + if not chunk: + # A scalar target repeated over a request has one final value. + yield ExecutionChunk( + row.chunk_coords, (), tuple(int(p[order[-1]]) for p in positions), False + ) + else: + # Eliminate duplicate destinations ourselves: a backend's repeated + # advanced-assignment order must not define our conflict policy. + destinations = np.stack([cast("np.ndarray[Any, Any]", c)[order] for c in chunk], axis=1) + _, reversed_positions = np.unique(destinations[::-1], axis=0, return_index=True) + order = order[np.sort(order.size - 1 - reversed_positions)] + yield ExecutionChunk( + row.chunk_coords, + tuple(cast("np.ndarray[Any, Any]", c)[order] for c in chunk), + tuple(p[order] for p in positions), + False, + ) + + +def _consumer_rows( + plan: ExecutionPlan, consumer: Literal["numpy", "shard"] +) -> Iterator[ExecutionChunk]: + rows = _raw_rows(plan) + if ( + plan.access == "write" + and plan.conflicts == "last" + and not isinstance(plan.work, _BasicWork) + ): + rows = _ordered_write_rows(plan, rows) + return _shard_rows(plan, rows) if consumer == "shard" else rows + + +def _lower(plan: ExecutionPlan, consumer: Literal["numpy", "shard"]) -> Iterator[LoweredOperation]: + for row in _consumer_rows(plan, consumer): + if all(isinstance(sel, np.ndarray) for sel in row.out_selection) and row.out_selection: + shape = np.broadcast_shapes( + *(cast("np.ndarray[Any, Any]", sel).shape for sel in row.out_selection) + ) + else: + shape = tuple( + len(range(*sel.indices(size))) + for sel, size in zip(row.out_selection, plan.shape, strict=True) + if isinstance(sel, slice) + ) + kind: Literal["basic", "paired"] = ( + "paired" if any(isinstance(sel, np.ndarray) for sel in row.chunk_selection) else "basic" + ) + yield LoweredOperation(row, shape, kind) + + +def _component_rows(work: _ComponentWork) -> Iterator[ExecutionChunk]: + """Lower factored coordinate columns without constructing transform pairs.""" + partition = work.plan.partition() + joints = partition.joint_sets + domain = work.plan.transform.domain + if 0 in domain.shape: + return + referenced = {axis for joint in joints for axis in joint.broadcast_axes} + unread = tuple(axis for axis in range(domain.ndim) if axis not in referenced) + slots: list[int | None] = [] + lead = 0 + for joint in joints: + slots.append(lead if joint.broadcast_axes else None) + lead += bool(joint.broadcast_axes) + rank = lead + len(unread) + for rows in itertools.product(*(range(len(joint)) for joint in joints)): + chunk: list[Selector] = [0] * work.plan.transform.output_rank + out: list[Selector] = [0] * domain.ndim + coords = [0] * len(chunk) + shape = [1] * rank + for joint, local, slot, row in zip(joints, work.local, slots, rows, strict=True): + run = joint.run(row) + component_shape = [1] * rank + if slot is not None: + shape[slot] = component_shape[slot] = run.stop - run.start + for column, dimension in enumerate(joint.output_dimensions): + coords[dimension] = int(joint.chunk[row, column]) + chunk[dimension] = local[run, column].reshape(component_shape) + for column, axis in enumerate(joint.broadcast_axes): + out[axis] = joint.block_coordinates[run, column].reshape(component_shape) + for i, axis in enumerate(unread): + component_shape = [1] * rank + shape[lead + i] = component_shape[lead + i] = domain.shape[axis] + out[axis] = np.arange(domain.shape[axis], dtype=np.intp).reshape(component_shape) + if any(domain.shape[axis] > 1 for axis in unread): + chunk = [ + np.broadcast_to(cast("np.ndarray[Any, Any]", sel), tuple(shape)) for sel in chunk + ] + yield ExecutionChunk(tuple(coords), tuple(chunk), tuple(out), False) diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py index 20ad3f95c2..34a6186a8e 100644 --- a/packages/zarr-indexing/src/zarr_indexing/grid.py +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -61,6 +61,21 @@ def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.int ... +@runtime_checkable +class RegularDimensionGridLike(DimensionGridLike, Protocol): + """A regular-grid capability, implemented structurally across packages. + + Chunk k starts at k * size and declares size elements; extent clips only + the final chunk's valid data. Implementations must obey these invariants. + """ + + @property + def size(self) -> int: ... + + @property + def extent(self) -> int: ... + + def _bounded_indices(indices: npt.NDArray[np.intp], extent: int) -> npt.NDArray[np.intp]: """Normalize a vector lookup and enforce the scalar grid bounds.""" arr = np.asarray(indices, dtype=np.intp) diff --git a/packages/zarr-indexing/tests/test_execution.py b/packages/zarr-indexing/tests/test_execution.py new file mode 100644 index 0000000000..0e0e4fc7de --- /dev/null +++ b/packages/zarr-indexing/tests/test_execution.py @@ -0,0 +1,348 @@ +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest +from hypothesis import given +from hypothesis import strategies as st + +from zarr_indexing import DimensionMap, IndexDomain, IndexTransform +from zarr_indexing._execution import execute_selection, execute_transform +from zarr_indexing.grid import dimension_grids_from_chunks + + +@pytest.mark.parametrize( + ("shape", "chunks", "selection", "mode"), + [ + ((7, 9), (3, 4), (slice(None), slice(None)), "basic"), + ((7, 9), (3, 4), (slice(1, 7, 2), 3), "basic"), + ((7, 9), (3, 4), (slice(None, None, -1), slice(1, 8, 2)), "basic"), + ((7, 9), ((2, 5), (4, 5)), (slice(2, 7), slice(None)), "basic"), + ((7, 9), (3, 4), (slice(2, 2), slice(None)), "basic"), + ((7, 9), (3, 4), (None, Ellipsis), "basic"), + ((), (), (), "basic"), + ((1000,), (100,), (np.repeat(np.arange(1000), 2),), "vectorized"), + ((7,), (3,), (np.array([6, 0, 6, 2]),), "vectorized"), + ((7,), (3,), (np.array([], dtype=np.intp),), "vectorized"), + ((1000,), (100,), (np.r_[np.arange(100), np.arange(900, 1000)],), "vectorized"), + ( + (7, 9, 5), + (3, 4, 2), + (np.array([6, 0])[:, None], np.array([8, 2])[:, None], np.array([4, 0, 2])[None, :]), + "vectorized", + ), + ], +) +def test_execution_reads_and_writes( + shape: tuple[int, ...], chunks: tuple[Any, ...], selection: Any, mode: str +) -> None: + source = np.arange(np.prod(shape), dtype=np.int64).reshape(shape) + grids = dimension_grids_from_chunks(chunks, shape) + base = IndexTransform.from_shape(shape) + transform = base[selection] if mode == "basic" else base.vindex[selection] + expected = source[selection] + for execution in ( + execute_selection(selection, shape, grids, mode=mode), + execute_transform(transform, grids), + ): + result = np.empty(execution.shape, dtype=source.dtype) + written = source.copy() + replacements = np.arange(expected.size).reshape(expected.shape) + 10000 + for row in execution: + bounds = tuple( + slice(g.chunk_offset(c), g.chunk_offset(c) + g.data_size(c)) + for g, c in zip(grids, row.chunk_coords, strict=True) + ) + result[row.out_selection] = source[bounds][row.chunk_selection] + target = written[bounds] if bounds else written + target[row.chunk_selection] = replacements[row.out_selection] + np.testing.assert_array_equal(result, expected) + expected_write = source.copy() + expected_write[selection] = replacements + np.testing.assert_array_equal(written, expected_write) + assert [r.chunk_coords for r in execution] == [r.chunk_coords for r in execution] + + +def test_boundary_chunk_has_complete_data_extent() -> None: + grids = dimension_grids_from_chunks((3,), (7,)) + rows = list(execute_selection(slice(None), (7,), grids)) + assert [row.is_complete_chunk for row in rows] == [True, True, True] + + +def test_execution_rejects_grid_rank_mismatch() -> None: + with pytest.raises(ValueError, match="one entry"): + execute_selection(slice(None), (7,), ()) + + +def test_execution_rejects_unknown_mode() -> None: + with pytest.raises(ValueError, match="unknown indexing mode"): + execute_selection( + slice(None), (7,), dimension_grids_from_chunks((3,), (7,)), mode="invalid" + ) + + +def test_execution_rejects_negative_shape() -> None: + with pytest.raises(ValueError, match="nonnegative"): + execute_selection(slice(None), (-1,), dimension_grids_from_chunks((3,), (7,))) + + +def test_execution_rejects_out_of_bounds_integer() -> None: + with pytest.raises(IndexError, match="out of bounds"): + execute_selection(7, (7,), dimension_grids_from_chunks((3,), (7,))) + + +def test_execution_sorted_coordinates_check_source_bounds() -> None: + with pytest.raises(IndexError): + execute_selection( + np.arange(1000), (100,), dimension_grids_from_chunks((100,), (1000,)), mode="vectorized" + ) + + +def test_execution_rejects_zero_slice_step() -> None: + with pytest.raises(IndexError, match="step must not be zero"): + execute_selection(slice(None, None, 0), (7,), dimension_grids_from_chunks((3,), (7,))) + + +def test_execution_validates_grid_bounds_before_returning_work() -> None: + with pytest.raises(IndexError): + execute_selection(slice(None), (100,), dimension_grids_from_chunks((10,), (50,))) + + +def test_execution_rejects_repeated_unread_input_axes() -> None: + transform = IndexTransform(IndexDomain.from_shape((2, 3)), (DimensionMap(0),)) + with pytest.raises(ValueError, match="duplicate writes"): + execute_transform(transform, dimension_grids_from_chunks((2,), (2,)), access="write") + + +def test_declarative_execution_retains_snapshot() -> None: + coordinates = np.arange(1000) + transform = IndexTransform.from_shape((1000,)).vindex[coordinates] + plan = execute_transform(transform, dimension_grids_from_chunks((100,), (1000,))) + coordinates[:] = 0 + first = next(iter(plan)) + np.testing.assert_array_equal(first.chunk_selection[0], np.arange(100)) + + +def test_immediate_execution_accepts_readonly_sorted_coordinates() -> None: + coordinates = np.arange(1000) + coordinates.setflags(write=False) + plan = execute_selection( + coordinates, (1000,), dimension_grids_from_chunks((100,), (1000,)), mode="vectorized" + ) + assert len(list(plan)) == 10 + np.testing.assert_array_equal(coordinates, np.arange(1000)) + + +@pytest.mark.parametrize( + ("offset", "stride", "values", "extent"), + [ + (2**63, -1, [2**63 - 1, 2**63 - 2], 3), + (-(2**63) + 1, 2, [2**62, 2**62 + 2], 6), + ], +) +def test_lowering_preserves_exact_affine_cancellation( + offset: int, stride: int, values: list[int], extent: int +) -> None: + from zarr_indexing import ArrayMap + + transform = IndexTransform( + IndexDomain.from_shape((2,)), + (ArrayMap(np.array(values, dtype=np.intp), offset=offset, stride=stride),), + ) + plan = execute_transform( + transform, dimension_grids_from_chunks((3,), (extent,)), access="write" + ) + result = np.zeros(extent, dtype=np.int64) + for row in plan: + result[row.chunk_coords[0] * 3 : row.chunk_coords[0] * 3 + 3][row.chunk_selection] = ( + np.array([10, 20])[row.out_selection] + ) + assert result[transform.apply((0,))[0]] == 10 + assert result[transform.apply((1,))[0]] == 20 + + +@pytest.mark.parametrize("ownership", ["snapshot", "borrow"]) +def test_explicit_ownership(ownership: Any) -> None: + coordinates = np.arange(1000) + plan = execute_selection( + coordinates, + (1000,), + dimension_grids_from_chunks((100,), (1000,)), + mode="vectorized", + ownership=ownership, + ) + assert plan.ownership == ownership + if ownership == "snapshot": + coordinates[:] = 0 + np.testing.assert_array_equal(next(iter(plan)).chunk_selection[0], np.arange(100)) + + +@pytest.mark.parametrize(("access", "conflicts"), [("read", "error"), ("write", "last")]) +def test_repeated_unread_axes_have_explicit_access_policy(access: Any, conflicts: Any) -> None: + transform = IndexTransform(IndexDomain.from_shape((2, 3)), (DimensionMap(0),)) + plan = execute_transform( + transform, dimension_grids_from_chunks((2,), (2,)), access=access, conflicts=conflicts + ) + source = np.array([10, 20]) + for consumer in ("numpy", "shard"): + result = np.empty(plan.shape, dtype=np.int64) + written = source.copy() + for op in plan.lower(consumer).operations(): + row = op.row + result[row.out_selection] = source[row.chunk_selection] + if access == "write": + written[row.chunk_selection] = np.arange(6).reshape(2, 3)[row.out_selection] + if access == "read": + np.testing.assert_array_equal(result, [[10, 10, 10], [20, 20, 20]]) + else: + np.testing.assert_array_equal(written, [2, 5]) + + +def test_write_rejects_duplicate_sorted_coordinates() -> None: + with pytest.raises(ValueError, match="duplicate writes"): + execute_selection( + np.repeat(np.arange(1000), 2), + (1000,), + dimension_grids_from_chunks((100,), (1000,)), + mode="vectorized", + access="write", + ) + + +def test_first_basic_result_keeps_large_axes_implicit() -> None: + import tracemalloc + + plan = execute_selection(Ellipsis, (10**9, 2), dimension_grids_from_chunks((1, 1), (10**9, 2))) + tracemalloc.start() + first = next(iter(plan)) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + assert first.chunk_coords == (0, 0) + assert peak < 100_000 + + +@pytest.mark.parametrize("consumer", ["numpy", "shard"]) +def test_scalar_coordinate_consumer_preserves_value_shape(consumer: Any) -> None: + from zarr_indexing import ArrayMap + + transform = IndexTransform(IndexDomain.from_shape(()), (ArrayMap(np.array(4)),)) + plan = execute_transform(transform, dimension_grids_from_chunks((1,), (8,))) + (operation,) = plan.lower(consumer).operations() + result = np.empty((), dtype=np.int64) + result[operation.row.out_selection] = np.array([4])[operation.row.chunk_selection] + assert result == 4 + assert operation.value_shape == () + + +@pytest.mark.parametrize("access", ["read", "write"]) +def test_diagonal_is_rejected_before_iteration(access: Any) -> None: + from zarr_indexing import ArrayMap + + transform = IndexTransform( + IndexDomain.from_shape((3, 2)), + (DimensionMap(0), DimensionMap(0), ArrayMap(np.array([[1, 0]]))), + ) + with pytest.raises(ValueError, match="diagonal"): + execute_transform( + transform, dimension_grids_from_chunks((2, 2, 2), (3, 3, 2)), access=access + ) + + +def test_last_write_lowering_removes_duplicate_destinations() -> None: + values = np.repeat(np.arange(1000), 2) + plan = execute_selection( + values, + (1000,), + dimension_grids_from_chunks((100,), (1000,)), + mode="vectorized", + access="write", + conflicts="last", + ) + for row in plan: + destinations = row.chunk_selection[0] + assert np.unique(destinations).size == np.size(destinations) + np.testing.assert_array_equal(np.asarray(row.out_selection[0]) % 2, 1) + + +def test_empty_coordinate_plan_emits_no_io() -> None: + plan = execute_selection( + np.array([], dtype=np.intp), + (5,), + dimension_grids_from_chunks((2,), (5,)), + mode="vectorized", + access="write", + ) + assert list(plan) == [] + assert list(plan.lower("shard")) == [] + + +def test_execution_rejects_unknown_access() -> None: + with pytest.raises(ValueError, match="access intent"): + execute_selection(Ellipsis, (1,), dimension_grids_from_chunks((1,), (1,)), access="bad") # type: ignore[arg-type] + + +def test_execution_rejects_unknown_ownership() -> None: + with pytest.raises(ValueError, match="ownership policy"): + execute_selection(Ellipsis, (1,), dimension_grids_from_chunks((1,), (1,)), ownership="bad") # type: ignore[arg-type] + + +def test_execution_rejects_unknown_conflicts() -> None: + with pytest.raises(ValueError, match="conflict policy"): + execute_selection(Ellipsis, (1,), dimension_grids_from_chunks((1,), (1,)), conflicts="bad") # type: ignore[arg-type] + + +def test_execution_rejects_unknown_consumer() -> None: + plan = execute_selection(Ellipsis, (1,), dimension_grids_from_chunks((1,), (1,))) + with pytest.raises(ValueError, match="consumer"): + plan.lower("bad") # type: ignore[arg-type] + + +def test_orthogonal_scalar_rejects_negative_literal_coordinate() -> None: + with pytest.raises(IndexError, match="negative scalar"): + execute_selection( + (-1, [1, 2]), (4, 4), dimension_grids_from_chunks((2, 2), (4, 4)), mode="orthogonal" + ) + + +@given( + points=st.lists(st.tuples(st.integers(0, 4), st.integers(0, 6)), max_size=40), + chunk_sizes=st.tuples(st.integers(1, 5), st.integers(1, 7)), +) +def test_execution_coordinate_roundtrip_and_last_write( + points: list[tuple[int, int]], chunk_sizes: tuple[int, int] +) -> None: + """Both consumers preserve request order and last writes for arbitrary gathers.""" + shape = (5, 7) + source = np.arange(35).reshape(shape) + coordinates = np.array(points, dtype=np.intp).reshape(-1, 2) + selection = (coordinates[:, 0], coordinates[:, 1]) + grids = dimension_grids_from_chunks(chunk_sizes, shape) + values = np.arange(len(points)) + 100 + expected = source.copy() + for point, value in zip(points, values, strict=True): + expected[point] = value + for consumer in ("numpy", "shard"): + plan = execute_selection(selection, shape, grids, mode="vectorized") + result = np.empty(len(points), dtype=source.dtype) + covered = np.zeros(len(points), dtype=np.intp) + for row in plan.lower(consumer): + bounds = tuple( + slice(g.chunk_offset(c), g.chunk_offset(c) + g.data_size(c)) + for g, c in zip(grids, row.chunk_coords, strict=True) + ) + result[row.out_selection] = source[bounds][row.chunk_selection] + covered[row.out_selection] += 1 + np.testing.assert_array_equal(result, source[selection]) + np.testing.assert_array_equal(covered, np.ones(len(points), dtype=np.intp)) + written = source.copy() + plan = execute_selection( + selection, shape, grids, mode="vectorized", access="write", conflicts="last" + ) + for row in plan.lower(consumer): + bounds = tuple( + slice(g.chunk_offset(c), g.chunk_offset(c) + g.data_size(c)) + for g, c in zip(grids, row.chunk_coords, strict=True) + ) + written[bounds][row.chunk_selection] = values[row.out_selection] + np.testing.assert_array_equal(written, expected) diff --git a/packages/zarr-indexing/tests/test_indexing_execution.py b/packages/zarr-indexing/tests/test_indexing_execution.py new file mode 100644 index 0000000000..74dc6d0ae4 --- /dev/null +++ b/packages/zarr-indexing/tests/test_indexing_execution.py @@ -0,0 +1,101 @@ +"""Exercise the optional indexing prototype through real Zarr codec pipelines.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, cast + +import numpy as np +import pytest + +from zarr_indexing import _execution as execution + +zarr = pytest.importorskip("zarr") + +if TYPE_CHECKING: + from zarr.core.indexing import Indexer + +pytestmark = pytest.mark.asyncio + + +@pytest.mark.parametrize("pipeline", ["BatchedCodecPipeline", "FusedCodecPipeline"]) +@pytest.mark.parametrize("layout", ["v2", "v3", "sharded"]) +@pytest.mark.parametrize( + "case", ["basic", "integer", "reverse", "sorted", "components", "orthogonal"] +) +async def test_execution_codec_read_write(pipeline: str, layout: str, case: str) -> None: + from zarr.core.buffer.core import default_buffer_prototype + + shape: tuple[int, ...] + chunks: tuple[int, ...] + selection: Any + if case == "sorted": + shape, chunks = (1003,), (100,) + selection, mode = (np.arange(1, 1003),), "vectorized" + else: + shape, chunks = (7, 9, 5), (3, 4, 2) + selection, mode = { + "basic": ((slice(1, 7, 2), slice(None), slice(1, 5)), "basic"), + "integer": ((2, slice(1, 8, 2), slice(None)), "basic"), + "reverse": ((slice(None, None, -1), slice(None), slice(None)), "basic"), + "orthogonal": ((3, np.array([1, 2]), slice(None)), "orthogonal"), + "components": ( + ( + np.array([6, 0])[:, None], + np.array([8, 2])[:, None], + np.array([4, 0, 2])[None, :], + ), + "vectorized", + ), + }[case] + source = np.arange(np.prod(shape), dtype=np.int64).reshape(shape) + kwargs: dict[str, Any] = {"zarr_format": 2 if layout == "v2" else 3} + if layout == "sharded": + kwargs["shards"] = tuple(c * 2 for c in chunks) + with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline." + pipeline}): + array = zarr.create_array( + store=zarr.storage.MemoryStore(), shape=shape, chunks=chunks, dtype="int64", **kwargs + ) + array[:] = source + async_array = array._async_array + # The pipeline processes shard-sized buffers for sharded arrays. + grids = async_array._chunk_grid._dimensions + plan = execution.execute_selection(selection, shape, grids, mode=mode) + if layout == "sharded": + plan = plan.lower("shard") + indexer = cast("Indexer", plan) + prototype = default_buffer_prototype() + result = await async_array._get_selection(indexer, prototype=prototype) + np.testing.assert_array_equal(result, source[selection]) + replacement = np.arange(np.prod(plan.shape)).reshape(plan.shape) + 10000 + write_plan = execution.execute_selection(selection, shape, grids, mode=mode, access="write") + if layout == "sharded": + write_plan = write_plan.lower("shard") + await async_array._set_selection( + cast("Indexer", write_plan), replacement, prototype=prototype + ) + expected = source.copy() + expected[selection] = replacement + np.testing.assert_array_equal(await async_array.getitem(Ellipsis), expected) + + +@pytest.mark.parametrize("pipeline", ["BatchedCodecPipeline", "FusedCodecPipeline"]) +@pytest.mark.parametrize("step", [1, 2]) +async def test_boundary_complete_write_skips_read(pipeline: str, step: int) -> None: + from zarr.core.buffer.core import default_buffer_prototype + + store = zarr.storage.LoggingStore(zarr.storage.MemoryStore()) + with zarr.config.set({"codec_pipeline.path": "zarr.core.codec_pipeline." + pipeline}): + array = zarr.create_array(store=store, shape=(7,), chunks=(3,), dtype="int64") + array[:] = np.arange(7) + plan = execution.execute_selection( + slice(6, 7, step), (7,), array._async_array._chunk_grid._dimensions, access="write" + ) + store.counter.clear() + await array._async_array._set_selection( + cast("Indexer", plan), np.array([99]), prototype=default_buffer_prototype() + ) + assert store.counter["get"] == 0 + assert store.counter["get_sync"] == 0 + np.testing.assert_array_equal( + await array._async_array.getitem(Ellipsis), [0, 1, 2, 3, 4, 5, 99] + )