diff --git a/packages/zarr-indexing/changes/292.bugfix.md b/packages/zarr-indexing/changes/292.bugfix.md new file mode 100644 index 0000000000..542be813d8 --- /dev/null +++ b/packages/zarr-indexing/changes/292.bugfix.md @@ -0,0 +1,14 @@ +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` 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 +the pairing; a further `.lazy` selection describes a different read and drops +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/changes/292.doc.md b/packages/zarr-indexing/changes/292.doc.md new file mode 100644 index 0000000000..aa2f71178d --- /dev/null +++ b/packages/zarr-indexing/changes/292.doc.md @@ -0,0 +1,7 @@ +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 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/292.feature.1.md b/packages/zarr-indexing/changes/292.feature.1.md new file mode 100644 index 0000000000..cd93e0d991 --- /dev/null +++ b/packages/zarr-indexing/changes/292.feature.1.md @@ -0,0 +1,13 @@ +`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. 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 +coordinate, which no cover slice can spell: `NoBasicSelectionError`. 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/292.feature.3.md b/packages/zarr-indexing/changes/292.feature.3.md new file mode 100644 index 0000000000..1a5d09152c --- /dev/null +++ b/packages/zarr-indexing/changes/292.feature.3.md @@ -0,0 +1,25 @@ +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 + `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 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 + 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. +Backend-specific acceptance is covered separately by the asyncio example. diff --git a/packages/zarr-indexing/changes/292.feature.4.md b/packages/zarr-indexing/changes/292.feature.4.md new file mode 100644 index 0000000000..fc4eefa081 --- /dev/null +++ b/packages/zarr-indexing/changes/292.feature.4.md @@ -0,0 +1,13 @@ +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. + +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. It also +rejects internally overlapping strided buffers, whose logical cells cannot +hold distinct result values. diff --git a/packages/zarr-indexing/changes/292.feature.5.md b/packages/zarr-indexing/changes/292.feature.5.md new file mode 100644 index 0000000000..160d558e30 --- /dev/null +++ b/packages/zarr-indexing/changes/292.feature.5.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 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. 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..89ffaf1eb5 --- /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 each cell's `chunk_domain` 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..870ead2035 100644 --- a/packages/zarr-indexing/docs/guide/integrations.md +++ b/packages/zarr-indexing/docs/guide/integrations.md @@ -70,6 +70,35 @@ 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. `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 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`. + ## 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..d782452caf --- /dev/null +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/README.md @@ -0,0 +1,48 @@ +# 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 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`, 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 +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 + +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..fc8cfb709c --- /dev/null +++ b/packages/zarr-indexing/examples/lazy_indexing_asyncio/lazy_indexing_asyncio.py @@ -0,0 +1,261 @@ +# /// 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, 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, + ReadContext, + numpy_reader, +) + + +class _NoSynchronousRead: + """A test reader proving the async adapter performs every source read.""" + + 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. + + 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 +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[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: zarr.AsyncArray[Any]) -> np.ndarray[Any, Any]: + """Materialize `view` by fetching every partition concurrently. + + 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(*(_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[Any]) -> None: + """Fetch a view's partitions concurrently through zarr.AsyncArray.""" + 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) + 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(cast(Any, 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_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 + 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, 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( + 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: 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( + (key, np.asanyarray(cell)) for key, cell in zip(missing.keys(), cells, strict=True) + ) + + 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) + for part in parts: + out[part.out_selection] = cache[cell_key(part)][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(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(cast(Any, 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_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 `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(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) + return await read_through(view, async_source) + + expected = np.asarray(source[:])[[30, 2, 2]][:, 4:10] + assert np.array_equal(asyncio.run(scenario()), expected) + + +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/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/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 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/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 9acfd28a21..99fea157b3 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, @@ -69,6 +73,7 @@ unit_step_reader, ) from zarr_indexing.transform import ( + BasicSelection, IndexTransform, ) @@ -77,6 +82,7 @@ __all__ = [ "ArrayMap", "BasicReader", + "BasicSelection", "BoundsCheckError", "ChunkCoverage", "ChunkGrid", @@ -95,6 +101,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 950a97b25e..1d5d53c992 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 @@ -145,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 @@ -168,12 +170,15 @@ numpy_reader, ) from zarr_indexing.transform import ( + BasicSelection, IndexTransform, ) if TYPE_CHECKING: from collections.abc import Callable, Iterator + from zarr_indexing.domain import IndexDomain + SelectFn = Callable[[Any, SelectionMode], "LazyArray"] __all__ = ["LazyArray", "Partition"] @@ -296,6 +301,64 @@ 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 _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. + + 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) + ), + 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, ...]: @@ -432,9 +495,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 @@ -446,9 +518,18 @@ 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 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. 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 +570,79 @@ 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) -> BasicSelection: + """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. 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) + ``` + + 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. The dedicated subclass of + `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. + + 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) -> 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 + `projection.chunk_domain`'s origin. For a consumer that caches decoded + 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. + + 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 +783,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 +815,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 +835,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 +853,7 @@ def _derive( view._parts = parts view._window = window view._reader = reader + view._projection = projection view._part_owner = _PartOwner() return view @@ -752,6 +924,7 @@ def with_reader(self, reader: Reader) -> LazyArray: self._parts, self._window, reader, + self._projection, ) # -- shape of the selection --------------------------------------------- @@ -1009,7 +1182,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. @@ -1050,9 +1225,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: @@ -1070,8 +1245,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) ) @@ -1084,6 +1268,12 @@ def parts(self) -> Iterator[Partition]: None, window, self._reader, + # `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, @@ -1133,7 +1323,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 +1351,123 @@ 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, 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): + 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`. + + 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, 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. + + 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") + 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 + # 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) + + 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 +1477,27 @@ 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) + 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 + ) return out written = 0 @@ -1195,7 +1515,12 @@ def result(self, *, parts: Sequence[Partition] | None = None) -> Any: _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/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..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) - if hi == lo: - 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) - 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) - if hi == lo: - return slice(0, 0, 1), DimensionMap(input_dimension=d, offset=-lo, stride=1) - first = checked_affine(m.offset, m.stride, lo) - 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), - 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/__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 4458d7922a..43b4066117 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]`). """ @@ -56,10 +70,13 @@ 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.reader import Reader, basic_reader +from zarr_indexing.output_map import DimensionMap +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, @@ -67,14 +84,18 @@ 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__ = [ "DEFAULT_DATA", "DEFAULT_PARTITIONINGS", "DEFAULT_SETTINGS", "ChainedIndexingStateMachine", + "ProjectionReader", "apply_selection", "outer_selection", + "projection_reader", "repartition", "state_machine_test", ] @@ -208,8 +229,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 @@ -249,6 +272,30 @@ 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 + fitting = [ + 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)) + ) + ] + # 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)) self.model = apply_selection(self.model, selection, mode) @@ -295,6 +342,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.""" @@ -313,7 +401,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)) @@ -367,12 +455,165 @@ 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. + + 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: + # 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 NoBasicSelectionError: + cell_lowered = False + assert not cell_lowered, ( + f"chunk_local_selection lowered a part whose source_selection " + f"refused to: {self.chain}" + ) + continue + slab = data[source_selection] + np.testing.assert_array_equal(slab, expected_block, 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_block, err_msg=str(self.chain) + ) + + +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 _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/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index a8a5963a26..eadc7e75fb 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, @@ -47,13 +51,29 @@ ) if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence import numpy.typing as npt 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 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. +""" + + @dataclass(frozen=True, slots=True) class _PointOutOfBounds(Exception): """Internal signal from the shared point kernel: one coordinate left the domain. @@ -448,6 +468,293 @@ def inverted(self) -> IndexTransform: output=tuple(inverse_output[dimension] for dimension in range(self.input_rank)), ) + 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 + 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 + `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. 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`. + + [`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 + 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 or None + One selector per output dimension, interleaved with a `None` for + each domain axis that neither an output map reads nor a constant + stands in for. + + Raises + ------ + 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 + source cell, which no slice spells), the selectors cannot produce + the domain's axes in increasing order and exactly once (basic + 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 + -------- + >>> 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)) + + 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() + Traceback (most recent call last): + ... + 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 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 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, " + "by an integer-indexed output dimension or a newaxis" + ) + 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: + 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" + ) + 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 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" + ) + d = output_map.input_dimension + fabricate_axes_before(d) + if d != next_axis: + 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 " + f"input dimension {next_axis} is due; basic indexing never " + "transposes or repeats axes" + ) + next_axis = d + 1 + 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, last = endpoints + if min(first, last) < 0: + 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" + ) + 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)) + # 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) + + 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) }'`. @@ -1201,7 +1508,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. """ @@ -1261,7 +1568,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. @@ -1456,6 +1763,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_doc_examples.py b/packages/zarr-indexing/tests/test_doc_examples.py index 240bd48531..95fd80d84f 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,33 @@ def test_documentation_example_executes(example: Path) -> None: runpy.run_path(str(example), run_name="__main__") +def _imported_modules(script: Path) -> tuple[str, ...]: + """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, None) for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: + modules[node.module] = 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, diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index ab94be1ae5..b07af1f7bf 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, @@ -2616,3 +2617,313 @@ 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'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()) + out = np.full(inner_view.shape, -1, dtype=inner_view.dtype) + for part in inner_view.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_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], :, :] + part = next(view.parts()) + assert not part.view.is_box + with pytest.raises(NoBasicSelectionError, match="ArrayMap"): + _ = part.source_selection + with pytest.raises(NoBasicSelectionError, 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_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() + 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()) + + # 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()) + 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_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)) + 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)) + 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_lazy_array_stateful.py b/packages/zarr-indexing/tests/test_lazy_array_stateful.py index 4f5dadc1e4..6ab4a25740 100644 --- a/packages/zarr-indexing/tests/test_lazy_array_stateful.py +++ b/packages/zarr-indexing/tests/test_lazy_array_stateful.py @@ -60,9 +60,34 @@ 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_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() + + readers = stateful._reader_set(LazyArray(np.arange(3)), (declared,)) + + assert readers == (basic_reader, declared) class OneDimensionalIndexing(ChainedIndexingStateMachine): diff --git a/packages/zarr-indexing/tests/test_transform.py b/packages/zarr-indexing/tests/test_transform.py index a13eaf6e28..19989c209f 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 ( @@ -395,6 +395,187 @@ 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", + ), + # 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() + # 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)) + + def test_rejects_an_array_map(self) -> None: + transform = IndexTransform.from_shape((10,)).oindex[[3, 1, 1]] + 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(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(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(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(NoBasicSelectionError, match="only a singleton axis"): + transform.as_basic_selection() + + 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),)) + 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),)) + 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(NoBasicSelectionError, match="count from the end"): + 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."""