From bbf8f9d10f72b2c1a5b02d4ef43de075052ee30e Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 20:39:05 +0200 Subject: [PATCH 1/4] =?UTF-8?q?feat(zarr-indexing):=20LazyArray=20?= =?UTF-8?q?=E2=80=94=20generic=20lazy=20indexing=20over=20array-API=20arra?= =?UTF-8?q?ys?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wrap any array-API-like array (NumPy, zarr, CuPy, ...) and give it a `.lazy` accessor that composes IndexTransforms instead of reading data. `result()` resolves a composed view: chunk by chunk through `iter_chunk_transforms` when the wrapped array declares chunks, otherwise as a single pass of array operations. `__getitem__` stays eager so a LazyArray is a drop-in duck array. Selections use positional NumPy semantics, deliberately different from `zarr.Array.lazy`'s literal TensorStore dialect; `zarr_indexing.boundary` is the generic translation layer between the two. `zarr_indexing.grid` gains `EdgeDimensionGrid` and `dimension_grids_from_chunks` so a caller with no zarr chunk grid can still resolve chunk-aware. Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/+lazy-array.feature.md | 1 + packages/zarr-indexing/docs/api/boundary.md | 5 + packages/zarr-indexing/docs/api/index.md | 13 +- packages/zarr-indexing/docs/api/lazy_array.md | 5 + packages/zarr-indexing/docs/index.md | 47 ++ packages/zarr-indexing/mkdocs.yml | 2 + .../src/zarr_indexing/__init__.py | 10 +- .../src/zarr_indexing/boundary.py | 231 ++++++ .../zarr-indexing/src/zarr_indexing/grid.py | 198 ++++- .../src/zarr_indexing/lazy_array.py | 674 ++++++++++++++++++ .../zarr-indexing/tests/test_lazy_array.py | 378 ++++++++++ 11 files changed, 1559 insertions(+), 5 deletions(-) create mode 100644 packages/zarr-indexing/changes/+lazy-array.feature.md create mode 100644 packages/zarr-indexing/docs/api/boundary.md create mode 100644 packages/zarr-indexing/docs/api/lazy_array.md create mode 100644 packages/zarr-indexing/src/zarr_indexing/boundary.py create mode 100644 packages/zarr-indexing/src/zarr_indexing/lazy_array.py create mode 100644 packages/zarr-indexing/tests/test_lazy_array.py diff --git a/packages/zarr-indexing/changes/+lazy-array.feature.md b/packages/zarr-indexing/changes/+lazy-array.feature.md new file mode 100644 index 0000000000..05bbe31007 --- /dev/null +++ b/packages/zarr-indexing/changes/+lazy-array.feature.md @@ -0,0 +1 @@ +Added `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, CuPy, ...) TensorStore-style lazy indexing. `LazyArray(x).lazy[...]`, `.lazy.oindex[...]`, and `.lazy.vindex[...]` compose an `IndexTransform` and return a new view without reading data; `result()` materializes it. Selections use positional NumPy semantics (zero-based within the current view, negatives wrap, masks must match the view shape), which the new `zarr_indexing.boundary` module translates into the transform algebra's literal coordinates. `__getitem__` reads eagerly, so a `LazyArray` is a drop-in source for `dask.array.from_array`. Chunking is taken from an explicit `chunks=` argument, then `read_chunk_sizes`, then `chunks`; a chunked source resolves through `iter_chunk_transforms` so each touched chunk is read exactly once, and an unchunked source lowers to a single pass of array operations. `zarr_indexing.grid` gained `EdgeDimensionGrid`, a concrete `DimensionGridLike` built from explicit per-chunk sizes, and `dimension_grids_from_chunks`, which normalizes either chunk convention into one grid per axis. diff --git a/packages/zarr-indexing/docs/api/boundary.md b/packages/zarr-indexing/docs/api/boundary.md new file mode 100644 index 0000000000..4f9fa99531 --- /dev/null +++ b/packages/zarr-indexing/docs/api/boundary.md @@ -0,0 +1,5 @@ +--- +title: boundary +--- + +::: zarr_indexing.boundary diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index b9a58b70fa..827f32202b 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -28,7 +28,18 @@ and the wire format built on top of it. current codec pipeline expects) - [`zarr_indexing.grid`](grid.md) — `DimensionGridLike`, the Protocol describing the narrow chunk-grid surface chunk resolution consumes, so that - nothing here imports `zarr` + nothing here imports `zarr`, plus `EdgeDimensionGrid` and + `dimension_grids_from_chunks`, a concrete per-axis grid for callers with no + zarr grid to hand + +**Lazy arrays** + +- [`zarr_indexing.lazy_array`](lazy_array.md) — `LazyArray`, a wrapper that + gives any array-API-like array (NumPy, zarr, CuPy, …) a `.lazy` accessor for + TensorStore-style deferred indexing, resolving chunk-aware when the wrapped + array declares chunks +- [`zarr_indexing.boundary`](boundary.md) — the translation between NumPy's + positional dialect and the transform algebra's literal coordinates **The ndsel wire format** (see [the guide](../ndsel.md)) diff --git a/packages/zarr-indexing/docs/api/lazy_array.md b/packages/zarr-indexing/docs/api/lazy_array.md new file mode 100644 index 0000000000..1ad928d84d --- /dev/null +++ b/packages/zarr-indexing/docs/api/lazy_array.md @@ -0,0 +1,5 @@ +--- +title: lazy_array +--- + +::: zarr_indexing.lazy_array diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index d123cbd32b..1c1af51498 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -142,6 +142,53 @@ and translated into chunk-local coordinates — exactly what a codec pipeline needs to decode that chunk and scatter the result. `out_indices` carries the output scatter indices for array selections, and is `None` for basic indexing. +## Lazy views over any array + +[`LazyArray`](api/lazy_array.md) packages all of the above into a wrapper you +can put around an array you already have — NumPy, zarr, CuPy, anything with +`shape`, `dtype`, and `__getitem__`. Its `.lazy` accessor composes transforms +instead of reading, and `result()` resolves the composed view in one pass: + +```python +import numpy as np +import zarr + +from zarr_indexing import LazyArray + +arr = zarr.create_array({}, shape=(100, 80), chunks=(30, 40), dtype="int32") +arr[:] = np.arange(8000).reshape(100, 80) + +lazy = LazyArray(arr) +lazy.chunks # ((30, 30, 30, 10), (40, 40)) + +view = lazy.lazy[10:50, ::4].lazy.oindex[[3, 0, 0], :] +view.shape # (3, 20) +repr(view) # '' + +view.result()[:, :4] +# array([[1040, 1044, 1048, 1052], +# [ 800, 804, 808, 812], +# [ 800, 804, 808, 812]], dtype=int32) +``` + +Unlike the transform algebra underneath it — where indices are literal +coordinates in a view's own (possibly shifted) domain — `LazyArray` speaks +**positional NumPy semantics**: a view is re-zeroed, `-1` is the last element, +and boolean masks must match the view's shape. `.lazy.oindex` and `.lazy.vindex` +give the orthogonal and vectorized flavours: + +```python +LazyArray(np.arange(12).reshape(3, 4)).lazy.vindex[[0, 2], [1, 3]].result() +# array([ 1, 11]) +``` + +Chunking is discovered from the wrapped array (`read_chunk_sizes`, then +`chunks`) or declared with `chunks=`. When it is known, a view resolves through +[chunk resolution](api/chunk_resolution.md) so every touched chunk is read +exactly once; when it is not, the view lowers to a single pass of array +operations. Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a +`LazyArray` a drop-in source for `dask.array.from_array`. + ## Reference - [The ndsel wire format](ndsel.md) diff --git a/packages/zarr-indexing/mkdocs.yml b/packages/zarr-indexing/mkdocs.yml index d7261f32e1..ffd04c7c9d 100644 --- a/packages/zarr-indexing/mkdocs.yml +++ b/packages/zarr-indexing/mkdocs.yml @@ -22,6 +22,8 @@ nav: - ' zarr_indexing.composition': api/composition.md - ' zarr_indexing.chunk_resolution': api/chunk_resolution.md - ' zarr_indexing.grid': api/grid.md + - ' zarr_indexing.lazy_array': api/lazy_array.md + - ' zarr_indexing.boundary': api/boundary.md - ' zarr_indexing.json': api/json.md - ' zarr_indexing.messages': api/messages.md - ' zarr_indexing.errors': api/errors.md diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index effe38ca88..4d252a2103 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -28,7 +28,11 @@ ) from zarr_indexing.composition import compose from zarr_indexing.domain import IndexDomain -from zarr_indexing.grid import DimensionGridLike +from zarr_indexing.grid import ( + DimensionGridLike, + EdgeDimensionGrid, + dimension_grids_from_chunks, +) from zarr_indexing.json import ( IndexDomainJSON, IndexTransformJSON, @@ -40,6 +44,7 @@ transform_from_canonical, transform_to_canonical, ) +from zarr_indexing.lazy_array import LazyArray from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_indexing.transform import IndexTransform, selection_to_transform @@ -51,15 +56,18 @@ "ConstantMap", "DimensionGridLike", "DimensionMap", + "EdgeDimensionGrid", "IndexDomain", "IndexDomainJSON", "IndexTransform", "IndexTransformJSON", + "LazyArray", "NdselError", "OutputIndexMap", "OutputIndexMapJSON", "__version__", "compose", + "dimension_grids_from_chunks", "index_domain_from_json", "index_domain_to_json", "index_transform_from_json", diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py new file mode 100644 index 0000000000..1ff930f9ec --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -0,0 +1,231 @@ +"""The positional (NumPy) selection dialect, lowered onto the transform algebra. + +The transform algebra speaks **literal domain coordinates**: an index is a +point in the view's own coordinate system, which after `view = arr[10:50]` runs +from 10 to 49, and a negative index is genuinely negative rather than counted +from the end (TensorStore's convention — see `zarr_indexing.transform`). + +NumPy speaks **positions**: index 0 always means the first element of the thing +you are indexing, and `-1` means the last. This module is the translation layer +between the two. It validates a selection against the view's shape with NumPy +semantics, then shifts every coordinate by the domain's origin so the transform +layer sees literal coordinates. + +Note +---- +`zarr.Array` currently carries its own copy of this normalization, tuned to a +different boundary contract (`Array.lazy[...]` deliberately exposes the literal +dialect, so a view's coordinates keep their meaning across composition). This +module is the generic, zarr-free version used by `LazyArray`; consolidating +zarr's copy onto it is left to a follow-up. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np + +if TYPE_CHECKING: + from zarr_indexing.domain import IndexDomain + +SelectionMode = Literal["basic", "orthogonal", "vectorized"] + + +def _is_bool_scalar(sel: Any) -> bool: + """Return True for a Python or NumPy boolean scalar. + + `isinstance(True, int)` holds, so booleans have to be rejected before any + integer handling — `arr[True]` would otherwise silently mean `arr[1]`. + """ + return isinstance(sel, (bool, np.bool_)) + + +def _as_index_array(sel: Any) -> np.ndarray[Any, np.dtype[Any]] | None: + """Return `sel` as an ndarray if it is array-like, else None.""" + if isinstance(sel, np.ndarray): + return sel + if isinstance(sel, (list, tuple)): + arr = np.asarray(sel) + if arr.dtype.kind in "biu": + return arr + raise IndexError( + f"arrays used as indices must be of integer or boolean type; got dtype {arr.dtype}" + ) + return None + + +def _axes_consumed(sel: Any, mode: SelectionMode) -> int: + """How many axes of the view a single selection entry consumes.""" + if sel is None: + return 0 + arr = _as_index_array(sel) + # A multidimensional boolean mask consumes one axis per mask dimension. + # Orthogonal indexing is per-axis by construction, so a mask there is 1-D + # and consumes exactly one axis. + if mode == "vectorized" and arr is not None and arr.dtype == np.bool_: + return arr.ndim + return 1 + + +def _normalize_int(value: int, size: int, axis: int) -> int: + """Bounds-check a positional integer index, wrapping negatives NumPy-style.""" + idx = value + if idx < 0: + idx += size + if idx < 0 or idx >= size: + raise IndexError(f"index {value} is out of bounds for axis {axis} with size {size}") + return idx + + +def _normalize_slice(sel: slice, size: int, axis: int) -> tuple[int, int, int]: + """Resolve a positional slice to `(start, stop, step)` within `[0, size]`.""" + if sel.step is not None and not isinstance(sel.step, (int, np.integer)): + raise IndexError(f"slice step must be an integer; got {sel.step!r}") + step = 1 if sel.step is None else int(sel.step) + if step < 1: + raise IndexError( + f"slice step must be positive; got {step} for axis {axis}. Reversed " + "and zero-step slices are not supported." + ) + start, stop, step = sel.indices(size) + # `slice.indices` can hand back stop < start for an empty forward slice + # (e.g. `5:2`); the transform layer wants a canonical empty interval. + stop = max(stop, start) + return start, stop, step + + +def _normalize_int_array( + arr: np.ndarray[Any, np.dtype[Any]], size: int, axis: int +) -> np.ndarray[Any, np.dtype[np.intp]]: + """Bounds-check a positional integer index array, wrapping negatives.""" + if arr.dtype.kind not in "iu": + raise IndexError( + f"arrays used as indices must be of integer or boolean type; got dtype {arr.dtype}" + ) + # Cast before wrapping: an unsigned array cannot represent the intermediate + # negative values, and `intp` covers every index NumPy can address. + out = arr.astype(np.intp, copy=True) + if out.size > 0: + negative = out < 0 + if bool(negative.any()): + out = np.where(negative, out + size, out) + lo, hi = int(out.min()), int(out.max()) + if lo < 0 or hi >= size: + bad = lo if lo < 0 else hi + raise IndexError(f"index {bad} is out of bounds for axis {axis} with size {size}") + return out + + +def normalize_positional_selection( + selection: Any, + domain: IndexDomain, + mode: SelectionMode, +) -> Any: + """Translate a positional (NumPy-dialect) selection into literal coordinates. + + Positions are zero-based offsets into the *current view*; negatives wrap + from the end. The returned selection addresses the same cells in the + literal coordinate system `domain` uses, ready for + `zarr_indexing.transform.selection_to_transform`. + + Parameters + ---------- + selection + A NumPy-style selection: integers, slices, `Ellipsis`, integer arrays or + lists, or boolean arrays. + domain + The domain of the view being indexed. Its shape defines the positional + bounds and its origin the coordinate shift. + mode + Which selection dialect the entries follow: `"basic"` (integers and + slices), `"orthogonal"` (per-axis arrays, outer product), or + `"vectorized"` (correlated coordinate arrays or a single mask). + + Returns + ------- + tuple + The selection with every coordinate expressed in literal domain + coordinates. + + Raises + ------ + IndexError + If a boolean scalar is used as an index, a boolean mask does not match + the shape of the axes it covers, an index is out of bounds, or too many + indices are supplied. + """ + entries = selection if isinstance(selection, tuple) else (selection,) + shape = domain.shape + origin = domain.inclusive_min + ndim = domain.ndim + + for sel in entries: + if _is_bool_scalar(sel): + raise IndexError( + "boolean scalars are not valid indices; use a boolean array " + "matching the shape of the axes it selects" + ) + + n_ellipsis = sum(1 for sel in entries if sel is Ellipsis) + if n_ellipsis > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + consumed = sum(_axes_consumed(sel, mode) for sel in entries if sel is not Ellipsis) + if consumed > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, but {consumed} were indexed" + ) + + result: list[Any] = [] + axis = 0 + for sel in entries: + if sel is Ellipsis: + # Passed through rather than expanded: every mode of + # `selection_to_transform` expands an ellipsis (and pads short + # selections) to whole-axis slices itself, and `vectorized` mode + # rejects an explicit slice, so expanding here would turn a legal + # partial coordinate selection into an error. + result.append(Ellipsis) + axis += ndim - consumed + continue + if sel is None: + # newaxis: no axis of the view is consumed, and there is no + # coordinate to shift. The transform layer decides whether the mode + # accepts it. + result.append(None) + continue + + arr = _as_index_array(sel) + if arr is not None and arr.dtype == np.bool_: + n_axes = _axes_consumed(sel, mode) + expected = shape[axis : axis + n_axes] + if arr.shape != tuple(expected): + raise IndexError( + f"boolean index has shape {arr.shape} but the axes it " + f"covers have shape {tuple(expected)}" + ) + for offset, positions in enumerate(np.nonzero(arr)): + result.append(positions.astype(np.intp) + origin[axis + offset]) + axis += n_axes + continue + + if axis >= ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, " + f"but {consumed} were indexed" + ) + size = shape[axis] + if arr is not None: + result.append(_normalize_int_array(arr, size, axis) + origin[axis]) + elif isinstance(sel, slice): + start, stop, step = _normalize_slice(sel, size, axis) + result.append(slice(start + origin[axis], stop + origin[axis], step)) + elif isinstance(sel, (int, np.integer)): + result.append(_normalize_int(int(sel), size, axis) + origin[axis]) + else: + raise IndexError(f"unsupported selection type: {type(sel)!r}") + axis += 1 + + # Axes the selection did not mention are left to `selection_to_transform`, + # which pads them with whole-axis slices in every mode. + return tuple(result) diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py index de1dae2dfc..06d7b193c9 100644 --- a/packages/zarr-indexing/src/zarr_indexing/grid.py +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -1,18 +1,27 @@ -"""Structural typing for the chunk-grid surface used by chunk resolution. +"""Per-dimension chunk grids: the Protocol and a concrete implementation. `chunk_resolution` needs only a narrow slice of a chunk grid: the per-dimension mapping between storage indices and chunk coordinates, passed as one `DimensionGridLike` per storage dimension. Rather than import zarr's concrete grid types, we type against this Protocol; zarr's per-dimension grids satisfy it structurally, so no zarr import is needed here. + +`EdgeDimensionGrid` is a concrete implementation for callers that do not have a +zarr grid to hand — it is built from an explicit tuple of per-chunk sizes along +one axis, so a clipped ("edge") boundary chunk is expressed directly rather than +derived from a uniform chunk shape. `dimension_grids_from_chunks` builds one per +axis from either chunk convention. """ from __future__ import annotations -from typing import TYPE_CHECKING, Protocol +from typing import TYPE_CHECKING, Any, Protocol, cast + +import numpy as np if TYPE_CHECKING: - import numpy as np + from collections.abc import Sequence + import numpy.typing as npt @@ -23,3 +32,186 @@ def index_to_chunk(self, idx: int) -> int: ... def chunk_offset(self, chunk_ix: int) -> int: ... def chunk_size(self, chunk_ix: int) -> int: ... def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: ... + + +class EdgeDimensionGrid: + """A one-dimensional chunk grid described by an explicit tuple of chunk sizes. + + The sizes follow the dask `Array.chunks` convention for a single axis: one + entry per chunk giving that chunk's extent, with a possibly-smaller final + ("edge") chunk. Chunk offsets are the exclusive prefix sums, so lookups are + a binary search rather than a division — the grid does not have to be + regular. + + Parameters + ---------- + sizes + Per-chunk extents along this axis. Every entry must be positive. An + empty sequence describes a zero-length axis. + + Examples + -------- + >>> grid = EdgeDimensionGrid((3, 3, 1)) + >>> grid.index_to_chunk(4) + 1 + >>> grid.chunk_offset(2), grid.chunk_size(2) + (6, 1) + """ + + __slots__ = ("_offsets", "sizes") + + sizes: tuple[int, ...] + + def __init__(self, sizes: Sequence[int]) -> None: + normalized = tuple(int(s) for s in sizes) + for i, s in enumerate(normalized): + if s <= 0: + raise ValueError( + f"chunk sizes must be positive; got {s} at position {i} of {normalized}" + ) + self.sizes = normalized + # Exclusive prefix sums, length len(sizes) + 1: `_offsets[c]` is the + # first index of chunk `c`, and `_offsets[-1]` the total extent. + offsets: np.ndarray[Any, np.dtype[np.intp]] = np.zeros(len(normalized) + 1, dtype=np.intp) + if len(normalized) > 0: + np.cumsum(np.asarray(normalized, dtype=np.intp), out=offsets[1:]) + self._offsets = offsets + + @property + def num_chunks(self) -> int: + """The number of chunks along this axis.""" + return len(self.sizes) + + @property + def extent(self) -> int: + """The total number of indices this axis spans.""" + return int(self._offsets[-1]) + + def __repr__(self) -> str: + return f"EdgeDimensionGrid(sizes={self.sizes})" + + def __eq__(self, other: object) -> bool: + if not isinstance(other, EdgeDimensionGrid): + return NotImplemented + return self.sizes == other.sizes + + def __hash__(self) -> int: + return hash((type(self).__name__, self.sizes)) + + def index_to_chunk(self, idx: int) -> int: + """Return the index of the chunk holding index `idx`.""" + if idx < 0 or idx >= self.extent: + raise IndexError(f"index {idx} is out of bounds for an axis of extent {self.extent}") + return int(np.searchsorted(self._offsets, idx, side="right")) - 1 + + def chunk_offset(self, chunk_ix: int) -> int: + """Return the first index belonging to chunk `chunk_ix`.""" + if chunk_ix < 0 or chunk_ix >= len(self.sizes): + raise IndexError( + f"chunk index {chunk_ix} is out of bounds for {len(self.sizes)} chunks" + ) + return int(self._offsets[chunk_ix]) + + def chunk_size(self, chunk_ix: int) -> int: + """Return the extent of chunk `chunk_ix`.""" + if chunk_ix < 0 or chunk_ix >= len(self.sizes): + raise IndexError( + f"chunk index {chunk_ix} is out of bounds for {len(self.sizes)} chunks" + ) + return self.sizes[chunk_ix] + + def indices_to_chunks(self, indices: npt.NDArray[np.intp]) -> npt.NDArray[np.intp]: + """Return the chunk holding each index in `indices`, elementwise.""" + arr = np.asarray(indices, dtype=np.intp) + if arr.size > 0 and (int(arr.min()) < 0 or int(arr.max()) >= self.extent): + raise IndexError( + f"indices must lie in [0, {self.extent}); got [{int(arr.min())}, {int(arr.max())}]" + ) + return (np.searchsorted(self._offsets, arr, side="right") - 1).astype(np.intp) + + +def _uniform_chunk_sizes(extent: int, chunk: int) -> tuple[int, ...]: + """Expand a single uniform chunk length into per-chunk sizes, tail clipped.""" + if chunk <= 0: + raise ValueError(f"chunk shape entries must be positive; got {chunk}") + n_full, remainder = divmod(extent, chunk) + sizes = (chunk,) * n_full + if remainder > 0: + sizes = (*sizes, remainder) + return sizes + + +def dimension_grids_from_chunks( + chunks: Sequence[int] | Sequence[Sequence[int]], + shape: Sequence[int], +) -> tuple[EdgeDimensionGrid, ...]: + """Build one `EdgeDimensionGrid` per axis from either chunk convention. + + Two conventions are accepted, disambiguated by element type: + + - a **uniform chunk shape** — one integer per axis, the chunk length along + that axis. The trailing chunk is clipped to the array extent, matching how + a zarr regular chunk grid covers a shape that is not a multiple of the + chunk shape. + - **dask-convention per-axis sizes** — one sequence of per-chunk extents per + axis, e.g. `((3, 3, 1), (4,))`. Each sequence must sum to the + corresponding entry of `shape`. + + Parameters + ---------- + chunks + A uniform chunk shape or dask-convention per-axis chunk sizes. + shape + The array shape the grids cover. + + Returns + ------- + tuple of EdgeDimensionGrid + One grid per axis, in axis order. + + Raises + ------ + ValueError + If `chunks` has a different length than `shape`, mixes the two + conventions, or declares per-axis sizes that do not sum to `shape`. + + Examples + -------- + >>> grids = dimension_grids_from_chunks((3, 4), (7, 4)) + >>> (grids[0].sizes, grids[1].sizes) + ((3, 3, 1), (4,)) + """ + shape_t = tuple(int(s) for s in shape) + entries: tuple[Any, ...] = tuple(chunks) + if len(entries) != len(shape_t): + raise ValueError( + f"chunks must have one entry per dimension; got {len(entries)} " + f"entries for shape {shape_t}" + ) + + is_int = [isinstance(entry, (int, np.integer)) for entry in entries] + if len(entries) > 0 and all(is_int): + return tuple( + EdgeDimensionGrid(_uniform_chunk_sizes(extent, int(entry))) + for entry, extent in zip(entries, shape_t, strict=True) + ) + if any(is_int): + raise ValueError( + "chunks must be either a uniform chunk shape (one integer per " + "dimension) or per-axis chunk sizes (one sequence per dimension), " + f"not a mixture; got {entries!r}" + ) + + grids: list[EdgeDimensionGrid] = [] + for axis, (entry, extent) in enumerate(zip(entries, shape_t, strict=True)): + # Every remaining entry is a sequence: the all-integer case returned + # above and the mixed case raised. + sizes = tuple(int(s) for s in cast("Sequence[int]", entry)) + total = sum(sizes) + if total != extent: + raise ValueError( + f"per-axis chunk sizes for dimension {axis} sum to {total}, " + f"but the array extent is {extent}" + ) + grids.append(EdgeDimensionGrid(sizes)) + return tuple(grids) diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py new file mode 100644 index 0000000000..401e5c0ea3 --- /dev/null +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -0,0 +1,674 @@ +"""`LazyArray` — TensorStore-style lazy indexing over any array-API-like array. + +`LazyArray` wraps an existing array — NumPy, zarr, CuPy, anything with `shape`, +`dtype`, and `__getitem__` — and adds a `.lazy` accessor whose indexing +operations build up an [`IndexTransform`](transform.md) instead of reading data: + +```python +view = LazyArray(source).lazy[10:50, ::2].lazy.oindex[[3, 1, 1], :] +view.shape # known without touching the data +values = view.result() +``` + +Nothing is read until `result()` (or `__array__`, or an eager `__getitem__`). +Composition is free: a view of a view is still a single transform. + +Chunk awareness +--------------- +If the wrapped array is chunked, `LazyArray` resolves a view against the chunk +grid — each chunk is read exactly once with plain basic slicing, then the +selected coordinates are scattered into the output buffer. Chunking is +discovered in this order: + +1. an explicit `chunks=` argument; +2. `array.read_chunk_sizes` (zarr's clipped per-axis sizes, sharding-aware); +3. `array.chunks`, read as per-axis sizes if its entries are sequences and as a + uniform chunk shape if they are integers; +4. otherwise the array is treated as unchunked and the view is lowered to a + single pass of array operations (`take`, basic slicing, a flat gather). + +Both resolution strategies compute the same values; the chunked one exists to +keep reads chunk-aligned, the unchunked one to avoid the per-chunk bookkeeping +when there is nothing to align to. + +The positional dialect +---------------------- +Selections on `LazyArray` are **positional, NumPy-style**: index 0 is the first +element of the current view, `-1` is the last, boolean masks must match the +view's shape, and every index is bounds-checked against the view. + +This is deliberately different from `zarr.Array.lazy[...]`, which exposes the +**literal** TensorStore dialect: a zarr view keeps the coordinate system of the +array it came from, so after `v = arr.lazy[10:50]` the first element of `v` is +`v[10]` and a negative index is out of bounds rather than counted from the end. +That dialect is the right one for zarr, where a view's coordinates are meant to +stay comparable with the parent array's. `LazyArray` is a duck array — it has to +behave like the thing it wraps to be usable as a drop-in for NumPy or as a dask +source — so it re-zeroes its coordinates on every view and speaks positions. +`zarr_indexing.boundary` performs the translation between the two. + +Device caveat +------------- +Chunked resolution builds its output buffer with NumPy, because the chunk +resolver's scatter indices are host-side integer arrays. Wrapping a device array +that declares chunks therefore returns host memory. The unchunked path stays in +the wrapped array's own namespace. +""" + +from __future__ import annotations + +import itertools +import math +from typing import TYPE_CHECKING, Any, Protocol + +import numpy as np + +from zarr_indexing.boundary import SelectionMode, normalize_positional_selection +from zarr_indexing.chunk_resolution import ( + iter_chunk_transforms, + sub_transform_to_selections, +) +from zarr_indexing.grid import EdgeDimensionGrid, dimension_grids_from_chunks +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.transform import ( + IndexTransform, + _array_map_dependent_axis, # pyright: ignore[reportPrivateUsage] + selection_to_transform, +) + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + SelectFn = Callable[[Any, SelectionMode], "LazyArray"] + +__all__ = ["LazyArray"] + + +class ArrayLike(Protocol): + """The surface `LazyArray` needs from the array it wraps.""" + + @property + def shape(self) -> tuple[int, ...]: ... + @property + def dtype(self) -> Any: ... + def __getitem__(self, key: Any) -> Any: ... + + +# --------------------------------------------------------------------------- # +# Array-namespace helpers +# --------------------------------------------------------------------------- # + + +def _namespace(x: Any) -> Any: + """Return `x`'s array-API namespace, or None if it does not declare one.""" + getter = getattr(x, "__array_namespace__", None) + if getter is None: + return None + try: + return getter() + except Exception: # pragma: no cover - a namespace that refuses to load + return None + + +def _take(x: Any, indices: np.ndarray[Any, np.dtype[np.intp]], axis: int) -> Any: + """Gather along one axis, preferring the wrapped array's own namespace.""" + xp = _namespace(x) + take = getattr(xp, "take", None) + if take is not None: + return take(x, xp.asarray(indices), axis=axis) + return x[(slice(None),) * axis + (indices,)] + + +def _reshape(x: Any, shape: tuple[int, ...]) -> Any: + xp = _namespace(x) + reshape = getattr(xp, "reshape", None) + if reshape is not None: + return reshape(x, shape) + return x.reshape(shape) + + +def _transpose(x: Any, perm: tuple[int, ...]) -> Any: + xp = _namespace(x) + for name in ("permute_dims", "transpose"): + fn = getattr(xp, name, None) + if fn is not None: + return fn(x, perm) + return np.transpose(x, perm) + + +def _expand_dims(x: Any, axis: int) -> Any: + xp = _namespace(x) + fn = getattr(xp, "expand_dims", None) + if fn is not None: + return fn(x, axis=axis) + return np.expand_dims(x, axis=axis) + + +# --------------------------------------------------------------------------- # +# Chunk discovery +# --------------------------------------------------------------------------- # + + +def _read_chunk_attribute(array: Any, name: str) -> Any: + """Read a chunk-describing attribute, treating any failure as "absent". + + Chunk discovery inspects an object we did not write. A missing attribute is + the common case, but zarr raises an `AttributeError` subclass from + `read_chunk_sizes` on a lazy view, and other backends compute the attribute + lazily and may fail for their own reasons. Any failure here means "this + array does not advertise chunks", never a hard error — while a malformed + *value* still raises from `dimension_grids_from_chunks`. + """ + try: + return getattr(array, name, None) + except Exception: + return None + + +def _discover_grids( + array: Any, + shape: tuple[int, ...], + chunks: Sequence[int] | Sequence[Sequence[int]] | None, +) -> tuple[EdgeDimensionGrid, ...] | None: + """Resolve the chunk grids for `array`, or None if it is unchunked.""" + if chunks is not None: + return dimension_grids_from_chunks(chunks, shape) + declared = _read_chunk_attribute(array, "read_chunk_sizes") + if declared is None: + declared = _read_chunk_attribute(array, "chunks") + if declared is None: + return None + return dimension_grids_from_chunks(declared, shape) + + +# --------------------------------------------------------------------------- # +# Resolution +# --------------------------------------------------------------------------- # + + +def _is_identity_transform(transform: IndexTransform, shape: tuple[int, ...]) -> bool: + """True when `transform` maps every coordinate of `shape` to itself. + + Structural rather than an `==` against `IndexTransform.from_shape`: an + `ArrayMap` holds an ndarray, so equality on two transforms that both carry + one would try to take the truth value of an array. + """ + domain = transform.domain + if domain.inclusive_min != (0,) * len(shape) or domain.exclusive_max != shape: + return False + if len(transform.output) != len(shape): + return False + return all( + isinstance(m, DimensionMap) and m.input_dimension == i and m.offset == 0 and m.stride == 1 + for i, m in enumerate(transform.output) + ) + + +def _is_correlated(transform: IndexTransform) -> bool: + """True when the transform scatters through a single flat index (vindex).""" + return any(isinstance(m, ArrayMap) and m.input_dimension is None for m in transform.output) + + +def _dimension_map_coords( + m: DimensionMap, transform: IndexTransform +) -> np.ndarray[Any, np.dtype[np.intp]]: + """The storage coordinates a DimensionMap enumerates, in view order.""" + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = transform.domain.exclusive_max[d] + return (m.offset + m.stride * np.arange(lo, hi, dtype=np.intp)).astype(np.intp) + + +def _array_map_coords(m: ArrayMap) -> np.ndarray[Any, np.dtype[np.intp]]: + """The storage coordinates an ArrayMap enumerates, flattened.""" + return (m.offset + m.stride * m.index_array).astype(np.intp).reshape(-1) + + +def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: int) -> Any: + """Permute `result`'s axes into input-domain order, restoring dropped axes. + + `axis_input_dims[k]` is the input (domain) dimension that axis `k` of + `result` corresponds to. Axes are permuted so that they appear in increasing + domain-dimension order, and any domain dimension no output map depends on + (only reachable via `newaxis`) is reinserted as a singleton. + """ + if len(set(axis_input_dims)) != len(axis_input_dims): + raise NotImplementedError( + "resolving a transform whose output maps share an input dimension " + "(a diagonal view) is not supported" + ) + order = sorted(range(len(axis_input_dims)), key=lambda k: axis_input_dims[k]) + if order != list(range(len(order))): + result = _transpose(result, tuple(order)) + covered = sorted(axis_input_dims) + for dim in range(rank): + if dim not in covered: + result = _expand_dims(result, dim) + return result + + +def _resolve_unchunked(array: Any, transform: IndexTransform) -> Any: + """Lower a transform to one pass of array operations over an unchunked source.""" + if _is_correlated(transform): + return _resolve_unchunked_correlated(array, transform) + return _resolve_unchunked_orthogonal(array, transform) + + +def _resolve_unchunked_orthogonal(array: Any, transform: IndexTransform) -> Any: + """Basic slicing plus one `take` per fancy-indexed axis (an outer product). + + Orthogonal `ArrayMap`s vary over distinct input axes, so gathering them one + axis at a time is exact — a `take` along one storage axis leaves every other + axis's coordinates untouched. + """ + outputs = transform.output + gathered: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + for out_dim, m in enumerate(outputs): + if isinstance(m, ArrayMap): + gathered[out_dim] = _array_map_coords(m) + elif isinstance(m, DimensionMap) and m.stride < 0: + # A reversing map has no positive-step slice; gather it explicitly. + gathered[out_dim] = _dimension_map_coords(m, transform) + + result = array + for out_dim, coords in gathered.items(): + result = _take(result, coords, axis=out_dim) + + selection: list[Any] = [] + axis_input_dims: list[int] = [] + for out_dim, m in enumerate(outputs): + if isinstance(m, ConstantMap): + selection.append(m.offset) + continue + if out_dim in gathered: + selection.append(slice(None)) + else: + assert isinstance(m, DimensionMap) + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = transform.domain.exclusive_max[d] + selection.append(slice(m.offset + m.stride * lo, m.offset + m.stride * hi, m.stride)) + axis_input_dims.append( + _array_map_dependent_axis(m) if isinstance(m, ArrayMap) else m.input_dimension + ) + result = result[tuple(selection)] + return _restore_domain_axis_order(result, axis_input_dims, transform.input_rank) + + +def _resolve_unchunked_correlated(array: Any, transform: IndexTransform) -> Any: + """Flatten the correlated axes, gather the points once, reshape back. + + Correlated (`vindex`) `ArrayMap`s address a list of *points*, not an outer + product, so they cannot be gathered one axis at a time. The correlated + storage axes are moved to the front and flattened, the per-point coordinates + are converted to offsets into that flat axis with row-major strides, and a + single `take` collects them. + """ + outputs = transform.output + correlated_dims = [ + d for d, m in enumerate(outputs) if isinstance(m, ArrayMap) and m.input_dimension is None + ] + if any(isinstance(m, ArrayMap) and m.input_dimension is not None for m in outputs): + raise NotImplementedError( + "resolving a transform with both correlated and orthogonal ArrayMaps is not supported" + ) + + slice_input_dims = {m.input_dimension for m in outputs if isinstance(m, DimensionMap)} + broadcast_axes = [d for d in range(transform.input_rank) if d not in slice_input_dims] + broadcast_shape = tuple(transform.domain.shape[d] for d in broadcast_axes) + + # Gather any reversing slice axis first, then take the basic-slice cut. The + # correlated axes keep their full extent: their coordinates are absolute. + gathered: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = { + d: _dimension_map_coords(m, transform) + for d, m in enumerate(outputs) + if isinstance(m, DimensionMap) and m.stride < 0 + } + result = array + for out_dim, coords in gathered.items(): + result = _take(result, coords, axis=out_dim) + + selection: list[Any] = [] + residual_axis_dims: list[int] = [] + correlated_positions: list[int] = [] + residual_positions: list[int] = [] + axis = 0 + for out_dim, m in enumerate(outputs): + if isinstance(m, ConstantMap): + selection.append(m.offset) + continue + if out_dim in correlated_dims: + selection.append(slice(None)) + correlated_positions.append(axis) + elif out_dim in gathered: + selection.append(slice(None)) + residual_positions.append(axis) + assert isinstance(m, DimensionMap) + residual_axis_dims.append(m.input_dimension) + else: + assert isinstance(m, DimensionMap) + d = m.input_dimension + lo = transform.domain.inclusive_min[d] + hi = transform.domain.exclusive_max[d] + selection.append(slice(m.offset + m.stride * lo, m.offset + m.stride * hi, m.stride)) + residual_positions.append(axis) + residual_axis_dims.append(d) + axis += 1 + result = result[tuple(selection)] + + # Correlated axes to the front, in output order, so the flattening strides + # below match the order the coordinates are combined in. + perm = tuple(correlated_positions) + tuple(residual_positions) + if perm != tuple(range(len(perm))): + result = _transpose(result, perm) + + n_corr = len(correlated_dims) + corr_sizes = tuple(int(s) for s in result.shape[:n_corr]) + tail_shape = tuple(int(s) for s in result.shape[n_corr:]) + result = _reshape(result, (math.prod(corr_sizes), *tail_shape)) + + flat_index = np.zeros(math.prod(broadcast_shape), dtype=np.intp) + stride = 1 + for position in range(n_corr - 1, -1, -1): + m = outputs[correlated_dims[position]] + assert isinstance(m, ArrayMap) + flat_index = flat_index + _array_map_coords(m) * stride + stride *= corr_sizes[position] + + result = _take(result, flat_index, axis=0) + result = _reshape(result, broadcast_shape + tail_shape) + return _restore_domain_axis_order( + result, list(broadcast_axes) + residual_axis_dims, transform.input_rank + ) + + +def _resolve_chunked( + array: Any, + transform: IndexTransform, + grids: Sequence[EdgeDimensionGrid], + dtype: np.dtype[Any], +) -> np.ndarray[Any, np.dtype[Any]]: + """Resolve a transform chunk by chunk, reading each touched chunk once.""" + out_shape = transform.domain.shape + # A correlated transform scatters through a single flat index, so the + # working buffer is 1-D during the read and reshaped afterwards. + needs_flat_buffer = _is_correlated(transform) + buffer_shape = (math.prod(out_shape),) if needs_flat_buffer else out_shape + out = np.empty(buffer_shape, dtype=dtype) + + if math.prod(out_shape) > 0: + for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, grids): + block_selection = tuple( + slice(grid.chunk_offset(c), grid.chunk_offset(c) + grid.chunk_size(c)) + for grid, c in zip(grids, chunk_coords, strict=True) + ) + block = np.asarray(array[block_selection]) + chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_transform, out_indices) + selected = block[chunk_sel] + if len(drop_axes) > 0: + selected = selected.squeeze(axis=drop_axes) + out[out_sel] = selected + + if needs_flat_buffer: + out = out.reshape(out_shape) + return out + + +def _view_chunk_sizes( + grid: EdgeDimensionGrid, offset: int, stride: int, lo: int, hi: int +) -> tuple[int, ...]: + """Per-chunk extents of a strided slice, expressed in the view's own positions.""" + size = hi - lo + if size == 0: + return () + cuts = {0, size} + for chunk_ix in range(1, grid.num_chunks): + boundary = grid.chunk_offset(chunk_ix) + # First view position whose storage coordinate reaches this boundary. + position = -((offset + stride * lo - boundary) // stride) + if 0 < position < size: + cuts.add(int(position)) + return tuple(b - a for a, b in itertools.pairwise(sorted(cuts))) + + +# --------------------------------------------------------------------------- # +# The wrapper +# --------------------------------------------------------------------------- # + + +class LazyArray: + """A lazily-indexable view over an array-API-like array. + + Wrapping is cheap and non-destructive: the wrapped array is never copied and + never read at construction time. Indexing through `.lazy` composes an + `IndexTransform` and returns another `LazyArray`; `result()` materializes. + + Selections use the **positional NumPy dialect** — see the module docstring, + which also documents chunk discovery and how this differs from + `zarr.Array.lazy`. + + Parameters + ---------- + array + The array to wrap. It must expose `shape`, `dtype`, and `__getitem__` + with basic (integer/slice) indexing. + chunks + Optionally declare the chunking of `array`, overriding discovery. Either + a uniform chunk shape (one integer per dimension) or dask-convention + per-axis chunk sizes (one sequence per dimension). + + Examples + -------- + >>> import numpy as np + >>> source = np.arange(12).reshape(3, 4) + >>> view = LazyArray(source, chunks=(2, 2)).lazy[1:, ::2] + >>> view.shape + (2, 2) + >>> view.result() + array([[ 4, 6], + [ 8, 10]]) + """ + + __slots__ = ("_array", "_grids", "_transform") + + def __init__( + self, + array: ArrayLike, + *, + chunks: Sequence[int] | Sequence[Sequence[int]] | None = None, + ) -> None: + shape = tuple(int(s) for s in array.shape) + self._array = array + self._transform = IndexTransform.from_shape(shape) + self._grids = _discover_grids(array, shape, chunks) + + @classmethod + def _view( + cls, + array: ArrayLike, + transform: IndexTransform, + grids: tuple[EdgeDimensionGrid, ...] | None, + ) -> LazyArray: + """Build a view sharing `array` and `grids` but carrying a new transform.""" + view = cls.__new__(cls) + view._array = array + # Views re-zero their coordinate system: the positional dialect means a + # view's first element is at position 0 whatever it was sliced from. + view._transform = transform.translate_domain_to((0,) * transform.input_rank) + view._grids = grids + return view + + # -- array-API surface -------------------------------------------------- + + @property + def array(self) -> ArrayLike: + """The wrapped array.""" + return self._array + + @property + def transform(self) -> IndexTransform: + """The composed transform from this view's coordinates to storage.""" + return self._transform + + @property + def shape(self) -> tuple[int, ...]: + """The shape of this view — the transform's input domain, not the source's.""" + return self._transform.domain.shape + + @property + def ndim(self) -> int: + return self._transform.input_rank + + @property + def size(self) -> int: + return math.prod(self.shape) + + @property + def dtype(self) -> Any: + """The wrapped array's dtype; views never change it.""" + return self._array.dtype + + @property + def chunks(self) -> tuple[tuple[int, ...], ...] | None: + """Per-axis chunk sizes of this view, in dask's `Array.chunks` convention. + + The chunks of the wrapped array, cut down to what the view addresses. It + is `None` when the source declared no chunking, and also when the view + is not a plain basic-slice view: once an axis has been fancy-indexed, + its elements no longer arrive in contiguous runs and "the chunks of this + view" has no useful answer. Reordered or duplicated coordinates are + exactly the case where a chunk cannot be described by a size. + """ + if self._grids is None: + return None + transform = self._transform + per_input: dict[int, tuple[int, ...]] = {} + for out_dim, m in enumerate(transform.output): + if isinstance(m, ConstantMap): + continue + if not isinstance(m, DimensionMap) or m.stride <= 0: + return None + d = m.input_dimension + if d in per_input: + return None + per_input[d] = _view_chunk_sizes( + self._grids[out_dim], + m.offset, + m.stride, + transform.domain.inclusive_min[d], + transform.domain.exclusive_max[d], + ) + if len(per_input) != transform.input_rank: + return None + return tuple(per_input[d] for d in range(transform.input_rank)) + + # -- indexing ----------------------------------------------------------- + + @property + def lazy(self) -> _LazyIndexer: + """Lazy indexing: `lazy[...]`, `lazy.oindex[...]`, `lazy.vindex[...]`. + + Each returns a new `LazyArray` view; no data is read. + """ + return _LazyIndexer(self._select) + + def _select(self, selection: Any, mode: SelectionMode) -> LazyArray: + literal = normalize_positional_selection(selection, self._transform.domain, mode) + composed = selection_to_transform(literal, self._transform, mode) + return LazyArray._view(self._array, composed, self._grids) + + def __getitem__(self, selection: Any) -> Any: + """Read a basic selection **eagerly**, like `numpy.ndarray.__getitem__`. + + Eager, not lazy, so that a `LazyArray` is a drop-in duck array for + consumers (dask's `from_array`, `numpy.asarray`) that expect indexing to + produce data. Use `.lazy[...]` for the lazy form. + """ + return self._select(selection, "basic").result() + + def result(self) -> Any: + """Materialize this view. + + Returns + ------- + array + An array of shape `self.shape`. Chunked sources produce a NumPy + array (see the module docstring's device caveat); unchunked sources + produce whatever the wrapped array's namespace produces. A view with + a zero-rank domain returns a zero-dimensional array, not a scalar. + """ + transform = self._transform + if self._grids is None: + return _resolve_unchunked(self._array, transform) + return _resolve_chunked(self._array, transform, self._grids, np.dtype(self._array.dtype)) + + def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: + if copy is False: + raise ValueError( + "a LazyArray cannot be converted to a NumPy array without a " + "copy: materializing a view always produces new data" + ) + return np.asarray(self.result(), dtype=dtype) + + def __len__(self) -> int: + if self.ndim == 0: + raise TypeError("len() of unsized object") + return self.shape[0] + + def __repr__(self) -> str: + wrapped = type(self._array).__name__ + parts = [f"{wrapped} shape={self.shape} dtype={self.dtype}"] + if not _is_identity_transform(self._transform, tuple(int(s) for s in self._array.shape)): + parts.append(f"view={self._transform.selection_repr}") + return f"" + + +class _LazyIndexer: + """The `.lazy` accessor: builds views instead of reading data. + + Holds the owning view's bound `_select` rather than the view itself, so the + accessor classes never reach into another object's internals. + """ + + __slots__ = ("_select",) + + def __init__(self, select: SelectFn) -> None: + self._select = select + + def __getitem__(self, selection: Any) -> LazyArray: + """Basic (integer / slice / ellipsis) indexing, lazily.""" + return self._select(selection, "basic") + + @property + def oindex(self) -> _LazyOIndex: + """Orthogonal (outer-product) indexing, lazily.""" + return _LazyOIndex(self._select) + + @property + def vindex(self) -> _LazyVIndex: + """Vectorized (coordinate / mask) indexing, lazily.""" + return _LazyVIndex(self._select) + + +class _LazyOIndex: + """`lazy.oindex[...]` — one selection per axis, combined as an outer product.""" + + __slots__ = ("_select",) + + def __init__(self, select: SelectFn) -> None: + self._select = select + + def __getitem__(self, selection: Any) -> LazyArray: + return self._select(selection, "orthogonal") + + +class _LazyVIndex: + """`lazy.vindex[...]` — correlated coordinate arrays, or a single mask.""" + + __slots__ = ("_select",) + + def __init__(self, select: SelectFn) -> None: + self._select = select + + def __getitem__(self, selection: Any) -> LazyArray: + return self._select(selection, "vectorized") diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py new file mode 100644 index 0000000000..34ec109870 --- /dev/null +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -0,0 +1,378 @@ +"""Tests for `zarr_indexing.grid` grids and the `LazyArray` wrapper. + +The happy-path suite is a single oracle test: every selection case is applied +both to a `LazyArray` and to the NumPy array it wraps, and the results must +match. The case list is crossed with four source flavours — an unchunked NumPy +array, NumPy with each of the two declared chunk conventions, and a zarr array +whose chunking is auto-discovered — so the chunked and unchunked resolution +strategies are held to the same answers. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import numpy as np +import pytest + +from zarr_indexing import EdgeDimensionGrid, LazyArray, dimension_grids_from_chunks + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + +SHAPE = (7, 5, 4) +CHUNK_SHAPE = (3, 2, 3) +DASK_CHUNKS = ((3, 3, 1), (2, 2, 1), (3, 1)) + + +def reference() -> np.ndarray[Any, np.dtype[np.int64]]: + """The array every source flavour holds, and the oracle for every case.""" + return np.arange(int(np.prod(SHAPE)), dtype=np.int64).reshape(SHAPE) + + +def outer(ref: np.ndarray[Any, Any], selections: Sequence[Any]) -> np.ndarray[Any, Any]: + """NumPy oracle for orthogonal indexing: the outer product of per-axis selections.""" + axes = [np.arange(size)[sel] for size, sel in zip(ref.shape, selections, strict=True)] + return ref[np.ix_(*axes)] + + +# --------------------------------------------------------------------------- +# EdgeDimensionGrid +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("sizes", "expected_index_to_chunk", "expected_offsets", "expected_sizes"), + [ + # A clipped edge chunk. + ((3, 3, 1), [0, 0, 0, 1, 1, 1, 2], [0, 3, 6], [3, 3, 1]), + # A single chunk covering the whole axis. + ((4,), [0, 0, 0, 0], [0], [4]), + # A size-1 axis. + ((1,), [0], [0], [1]), + # Irregular sizes: the grid does not have to be regular. + ((1, 2, 1), [0, 1, 1, 2], [0, 1, 3], [1, 2, 1]), + ], +) +def test_edge_dimension_grid( + sizes: tuple[int, ...], + expected_index_to_chunk: list[int], + expected_offsets: list[int], + expected_sizes: list[int], +) -> None: + """All four `DimensionGridLike` methods agree with hand-computed values.""" + grid = EdgeDimensionGrid(sizes) + extent = sum(sizes) + + assert grid.num_chunks == len(sizes) + assert grid.extent == extent + assert [grid.index_to_chunk(i) for i in range(extent)] == expected_index_to_chunk + assert [grid.chunk_offset(c) for c in range(len(sizes))] == expected_offsets + assert [grid.chunk_size(c) for c in range(len(sizes))] == expected_sizes + np.testing.assert_array_equal( + grid.indices_to_chunks(np.arange(extent, dtype=np.intp)), + np.asarray(expected_index_to_chunk, dtype=np.intp), + ) + + +def test_edge_dimension_grid_rejects_nonpositive_size() -> None: + with pytest.raises(ValueError, match="chunk sizes must be positive"): + EdgeDimensionGrid((3, 0, 2)) + + +def test_edge_dimension_grid_rejects_out_of_bounds_index() -> None: + with pytest.raises(IndexError, match="out of bounds for an axis of extent 4"): + EdgeDimensionGrid((3, 1)).index_to_chunk(4) + + +def test_edge_dimension_grid_rejects_out_of_bounds_chunk() -> None: + with pytest.raises(IndexError, match="chunk index 2 is out of bounds"): + EdgeDimensionGrid((3, 1)).chunk_offset(2) + + +@pytest.mark.parametrize( + ("chunks", "shape", "expected"), + [ + # Uniform chunk shape, tail clipped. + ((3, 4), (7, 4), ((3, 3, 1), (4,))), + # Dask-convention per-axis sizes, passed through. + (((3, 3, 1), (2, 2)), (7, 4), ((3, 3, 1), (2, 2))), + # A chunk longer than the axis collapses to one clipped chunk. + ((10,), (4,), ((4,),)), + # A zero-length axis has no chunks at all. + ((3,), (0,), ((),)), + ], +) +def test_dimension_grids_from_chunks( + chunks: Any, shape: tuple[int, ...], expected: tuple[tuple[int, ...], ...] +) -> None: + grids = dimension_grids_from_chunks(chunks, shape) + assert tuple(grid.sizes for grid in grids) == expected + + +def test_dimension_grids_from_chunks_rejects_wrong_length() -> None: + with pytest.raises(ValueError, match="one entry per dimension"): + dimension_grids_from_chunks((3,), (7, 4)) + + +def test_dimension_grids_from_chunks_rejects_mixed_conventions() -> None: + with pytest.raises(ValueError, match="not a mixture"): + dimension_grids_from_chunks((3, (2, 2)), (7, 4)) + + +def test_dimension_grids_from_chunks_rejects_wrong_total() -> None: + with pytest.raises(ValueError, match="sum to 5, but the array extent is 7"): + dimension_grids_from_chunks(((3, 2), (4,)), (7, 4)) + + +# --------------------------------------------------------------------------- +# Sources +# --------------------------------------------------------------------------- + + +def make_source(flavour: str) -> LazyArray: + """Build a `LazyArray` over `reference()` with the requested chunk provenance.""" + data = reference() + if flavour == "numpy-unchunked": + return LazyArray(data) + if flavour == "numpy-dask-chunks": + return LazyArray(data, chunks=DASK_CHUNKS) + if flavour == "numpy-uniform-chunks": + return LazyArray(data, chunks=CHUNK_SHAPE) + if flavour == "zarr": + zarr = pytest.importorskip("zarr") + array = zarr.create_array({}, shape=SHAPE, chunks=CHUNK_SHAPE, dtype="int64") + array[:] = data + return LazyArray(array) + raise AssertionError(f"unknown source flavour {flavour!r}") + + +FLAVOURS = ["numpy-unchunked", "numpy-dask-chunks", "numpy-uniform-chunks", "zarr"] + + +@pytest.fixture(params=FLAVOURS) +def source(request: pytest.FixtureRequest) -> LazyArray: + return make_source(request.param) + + +# --------------------------------------------------------------------------- +# The oracle +# --------------------------------------------------------------------------- + +MASK = (reference() % 11) == 0 + +# (id, lazy view builder, NumPy oracle). Integer scalars are deliberately absent +# from the orthogonal cases: the transform algebra keeps an int-selected axis as +# a length-1 axis under `oindex`, which `np.ix_` cannot express. +CASES: list[tuple[str, Callable[[LazyArray], LazyArray], Callable[[Any], Any]]] = [ + ( + "basic-strided-and-int-drop", + lambda a: a.lazy[1:6:2, :, -1], + lambda r: r[1:6:2, :, -1], + ), + ("basic-ellipsis", lambda a: a.lazy[..., -2], lambda r: r[..., -2]), + ("basic-negative-scalar", lambda a: a.lazy[-3], lambda r: r[-3]), + ("basic-empty", lambda a: a.lazy[:, 2:2, :], lambda r: r[:, 2:2, :]), + ("basic-all-scalars", lambda a: a.lazy[-1, 0, 2], lambda r: r[-1, 0, 2]), + ( + "oindex-unsorted-duplicates-multi-axis", + lambda a: a.lazy.oindex[[4, 0, 0, 2], :, [3, 1]], + lambda r: outer(r, ([4, 0, 0, 2], slice(None), [3, 1])), + ), + ( + "oindex-negative-and-slice", + lambda a: a.lazy.oindex[:, [-1, 0], 1:4], + lambda r: outer(r, (slice(None), [-1, 0], slice(1, 4))), + ), + ( + "oindex-boolean-axis", + lambda a: a.lazy.oindex[np.array([True, False, True, False, False, False, True]), :, :], + lambda r: outer( + r, + (np.array([True, False, True, False, False, False, True]), slice(None), slice(None)), + ), + ), + ( + "vindex-coordinates", + lambda a: a.lazy.vindex[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])], + lambda r: r[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])], + ), + ( + "vindex-broadcast-pair", + lambda a: a.lazy.vindex[np.array([[0], [6]]), np.array([1, 4]), np.array([2, 0])], + lambda r: r[np.array([[0], [6]]), np.array([1, 4]), np.array([2, 0])], + ), + ( + "vindex-negative-coordinates", + lambda a: a.lazy.vindex[np.array([-1, -7]), np.array([-2, 0]), np.array([0, -1])], + lambda r: r[np.array([-1, -7]), np.array([-2, 0]), np.array([0, -1])], + ), + ("vindex-mask", lambda a: a.lazy.vindex[MASK], lambda r: r[MASK]), + ( + "compose-basic-then-oindex", + lambda a: a.lazy[1:6].lazy.oindex[[3, 0, 0], [4, 1], :], + lambda r: outer(r[1:6], ([3, 0, 0], [4, 1], slice(None))), + ), + ( + "compose-oindex-then-basic-other-axis", + lambda a: a.lazy.oindex[[4, 0, 2], :, :].lazy[:, 1:4, ::2], + lambda r: outer(r, ([4, 0, 2], slice(None), slice(None)))[:, 1:4, ::2], + ), + ( + "compose-basic-then-basic", + lambda a: a.lazy[2:, 1:].lazy[::2, -1], + lambda r: r[2:, 1:][::2, -1], + ), + ( + "compose-basic-then-vindex", + lambda a: a.lazy[1:6, :, 1:].lazy.vindex[ + np.array([0, 4]), np.array([2, 0]), np.array([1, 2]) + ], + lambda r: r[1:6, :, 1:][np.array([0, 4]), np.array([2, 0]), np.array([1, 2])], + ), +] + + +@pytest.mark.parametrize(("build", "oracle"), [c[1:] for c in CASES], ids=[c[0] for c in CASES]) +def test_selection_matches_numpy( + source: LazyArray, + build: Callable[[LazyArray], LazyArray], + oracle: Callable[[Any], Any], +) -> None: + """Every selection resolves to what NumPy computes positionally on the same data.""" + view = build(source) + expected = np.asarray(oracle(reference())) + + assert view.shape == expected.shape + assert view.ndim == expected.ndim + np.testing.assert_array_equal(np.asarray(view.result()), expected) + # `__getitem__` is eager, and `__array__` routes through `result()`. + np.testing.assert_array_equal(np.asarray(view), expected) + + +def test_eager_getitem_returns_data(source: LazyArray) -> None: + """`arr[...]` reads immediately, like `numpy.ndarray.__getitem__`.""" + np.testing.assert_array_equal(np.asarray(source[1:3, ::2, -1]), reference()[1:3, ::2, -1]) + + +# --------------------------------------------------------------------------- +# Attribute forwarding +# --------------------------------------------------------------------------- + + +def test_forwards_array_attributes(source: LazyArray) -> None: + assert source.shape == SHAPE + assert source.ndim == len(SHAPE) + assert source.size == int(np.prod(SHAPE)) + assert source.dtype == np.dtype("int64") + assert len(source) == SHAPE[0] + assert "LazyArray" in repr(source) + + +def test_view_shape_comes_from_the_transform(source: LazyArray) -> None: + """A view reports its own shape, not the wrapped array's.""" + view = source.lazy[1:6:2, :, -1] + assert view.shape == (3, 5) + assert view.ndim == 2 + assert view.size == 15 + assert view.dtype == source.dtype + assert "view=" in repr(view) + + +@pytest.mark.parametrize( + ("flavour", "expected"), + [ + ("numpy-unchunked", None), + ("numpy-dask-chunks", DASK_CHUNKS), + ("numpy-uniform-chunks", DASK_CHUNKS), + ("zarr", DASK_CHUNKS), + ], +) +def test_chunk_discovery(flavour: str, expected: tuple[tuple[int, ...], ...] | None) -> None: + """Declared, uniform, and auto-discovered chunkings all agree.""" + assert make_source(flavour).chunks == expected + + +def test_chunks_follow_a_basic_view() -> None: + """A basic-slice view reports its chunks in its own positions.""" + array = make_source("numpy-uniform-chunks") + # Storage chunk boundaries along axis 0 are 3 and 6; positions 0..4 of the + # view map to storage 1..5, so the boundary at 3 falls at position 2. + assert array.lazy[1:6, :, :].chunks == ((2, 3), (2, 2, 1), (3, 1)) + # Every other element of axis 0: storage 0, 2 | 4 | 6. + assert array.lazy[::2, :, :].chunks[0] == (2, 1, 1) + + +def test_chunks_are_undefined_for_fancy_views() -> None: + """Reordered/duplicated coordinates have no chunk-size description.""" + array = make_source("numpy-uniform-chunks") + assert array.lazy.oindex[[4, 0, 0], :, :].chunks is None + + +# --------------------------------------------------------------------------- +# dask interop +# --------------------------------------------------------------------------- + + +def test_dask_from_array_roundtrip() -> None: + """A `LazyArray` is a drop-in dask source — no translation ceremony.""" + da = pytest.importorskip("dask.array") + source = make_source("zarr") + + lazy = da.from_array(source) + np.testing.assert_array_equal(lazy.compute(), reference()) + + # The wrapper's own chunks can be handed straight back to dask. + chunked = da.from_array(source, chunks=source.chunks) + assert chunked.chunks == DASK_CHUNKS + np.testing.assert_array_equal(chunked[2:, ::2].compute(), reference()[2:, ::2]) + + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +def test_boolean_scalar_is_rejected() -> None: + with pytest.raises(IndexError, match="boolean scalars are not valid indices"): + make_source("numpy-uniform-chunks").lazy[True] + + +def test_mask_shape_must_match_the_view() -> None: + array = make_source("numpy-uniform-chunks") + with pytest.raises(IndexError, match="boolean index has shape"): + array.lazy.vindex[np.ones((2, 2, 2), dtype=bool)] + + +def test_scalar_index_out_of_bounds() -> None: + with pytest.raises(IndexError, match="index 7 is out of bounds for axis 0 with size 7"): + make_source("numpy-uniform-chunks").lazy[7] + + +def test_scalar_index_out_of_bounds_in_a_view() -> None: + """Bounds are the *view's*, not the wrapped array's.""" + array = make_source("numpy-uniform-chunks").lazy[1:4] + with pytest.raises(IndexError, match="index 3 is out of bounds for axis 0 with size 3"): + array.lazy[3] + + +def test_index_array_out_of_bounds() -> None: + array = make_source("numpy-uniform-chunks") + with pytest.raises(IndexError, match="index 99 is out of bounds for axis 0 with size 7"): + array.lazy.oindex[[0, 99], :, :] + + +def test_too_many_indices() -> None: + with pytest.raises(IndexError, match="too many indices"): + make_source("numpy-uniform-chunks").lazy[0, 0, 0, 0] + + +def test_chunks_that_do_not_sum_to_shape_are_rejected() -> None: + with pytest.raises(ValueError, match="sum to 6, but the array extent is 7"): + LazyArray(reference(), chunks=((3, 3), (2, 2, 1), (3, 1))) + + +def test_copy_false_conversion_is_rejected() -> None: + array = make_source("numpy-uniform-chunks") + with pytest.raises(ValueError, match="cannot be converted to a NumPy array without a copy"): + np.array(array, copy=False) From 59b279e98fad1a43666513e07a35e32ef5557aa4 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 20:40:31 +0200 Subject: [PATCH 2/4] docs(zarr-indexing): name the LazyArray changelog fragment after its PR Assisted-by: ClaudeCode:claude-fable-5 --- .../changes/{+lazy-array.feature.md => 267.feature.md} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename packages/zarr-indexing/changes/{+lazy-array.feature.md => 267.feature.md} (100%) diff --git a/packages/zarr-indexing/changes/+lazy-array.feature.md b/packages/zarr-indexing/changes/267.feature.md similarity index 100% rename from packages/zarr-indexing/changes/+lazy-array.feature.md rename to packages/zarr-indexing/changes/267.feature.md From 1ab135034fe635cf390e1f45ade48f3d107fa250 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 21:17:04 +0200 Subject: [PATCH 3/4] fix(zarr-indexing): NumPy-faithful scalar and advanced-index placement in LazyArray MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three defects the review found, all reachable from the documented idioms. 1. An integer index applied to an axis a fancy selection had already indexed was unresolvable: the composed ArrayMap became all-singleton but kept an `input_dimension` naming the axis the integer had just removed, which after renumbering aliased a live DimensionMap's axis. `_apply_basic_indexing` now collapses such a map to a ConstantMap — the long-deferred degenerate-collapse rule — and `array_map_dependent_axis` (promoted from a private helper, and exported) answers `None` for "depends on no axis" rather than falling back to the stale binding. 2. Scalar integers now drop their axis in every mode, as NumPy does: `oindex[0]` has the shape of `x[0]`, and `oindex[0, 1, 2]` and `vindex[0, 1, 2]` are zero-rank. They are split off at the boundary and applied as a basic step before the advanced one, rather than widened into length-1 index arrays that kept the axis. 3. A vectorized selection whose coordinate arrays are not leading (`vindex[..., i, j]`, `vindex[..., mask]`) resolved correctly unchunked but raised a shape mismatch on every real chunking: the domain put the gathered dimensions first while the per-chunk gather left them where NumPy puts them. `_apply_vindex` now follows NumPy's placement rule, `_intersect_correlated` computes scatter offsets against the real buffer layout, and the chunked resolver realigns each block's gathered axis — including for NumPy's integer-counts-as-advanced corner, which a ConstantMap in the chunk selection triggers. Also: both resolvers agree that a zero-rank result is a 0-d array; a discovered but unusable `chunks` attribute on a foreign object falls back to whole-array reads while an explicit `chunks=` stays strict; and a non-sequence chunk entry names the two conventions instead of leaking a TypeError. Tests: the oracle matrix gains scalar-after-fancy chains in both orders, non-leading vindex including the `[..., mask]` idiom, and a seeded 400-chain randomized equivalence sweep per source flavour. Verified offline over 36,000 chains across six seeds and four flavours with zero divergences. Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-indexing/changes/267.feature.md | 6 +- packages/zarr-indexing/docs/index.md | 23 ++ .../src/zarr_indexing/__init__.py | 7 +- .../src/zarr_indexing/boundary.py | 91 +++++ .../zarr-indexing/src/zarr_indexing/grid.py | 32 +- .../src/zarr_indexing/lazy_array.py | 121 +++++- .../src/zarr_indexing/transform.py | 174 ++++++--- .../zarr-indexing/tests/test_lazy_array.py | 344 +++++++++++++++++- 8 files changed, 724 insertions(+), 74 deletions(-) diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 05bbe31007..5f64b18c73 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -1 +1,5 @@ -Added `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, CuPy, ...) TensorStore-style lazy indexing. `LazyArray(x).lazy[...]`, `.lazy.oindex[...]`, and `.lazy.vindex[...]` compose an `IndexTransform` and return a new view without reading data; `result()` materializes it. Selections use positional NumPy semantics (zero-based within the current view, negatives wrap, masks must match the view shape), which the new `zarr_indexing.boundary` module translates into the transform algebra's literal coordinates. `__getitem__` reads eagerly, so a `LazyArray` is a drop-in source for `dask.array.from_array`. Chunking is taken from an explicit `chunks=` argument, then `read_chunk_sizes`, then `chunks`; a chunked source resolves through `iter_chunk_transforms` so each touched chunk is read exactly once, and an unchunked source lowers to a single pass of array operations. `zarr_indexing.grid` gained `EdgeDimensionGrid`, a concrete `DimensionGridLike` built from explicit per-chunk sizes, and `dimension_grids_from_chunks`, which normalizes either chunk convention into one grid per axis. +Added `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, CuPy, ...) TensorStore-style lazy indexing. `LazyArray(x).lazy[...]`, `.lazy.oindex[...]`, and `.lazy.vindex[...]` compose an `IndexTransform` and return a new view without reading data; `result()` materializes it. Selections use positional NumPy semantics (zero-based within the current view, negatives wrap, masks must match the view shape, scalar integers drop their axis, and a partial `vindex` places its gathered dimensions where NumPy would), which the new `zarr_indexing.boundary` module translates into the transform algebra's literal coordinates. `__getitem__` reads eagerly, so a `LazyArray` is a drop-in source for `dask.array.from_array`. Chunking is taken from an explicit `chunks=` argument, then `read_chunk_sizes`, then `chunks`; a chunked source resolves through `iter_chunk_transforms` so each touched chunk is read exactly once, and an unchunked source lowers to a single pass of array operations. `zarr_indexing.grid` gained `EdgeDimensionGrid`, a concrete `DimensionGridLike` built from explicit per-chunk sizes, and `dimension_grids_from_chunks`, which normalizes either chunk convention into one grid per axis. + +Fixed an integer index applied to an axis a previous orthogonal or vectorized selection had already indexed. The composed `ArrayMap` became all-singleton but kept an `input_dimension` naming the axis the integer had just removed, which after renumbering aliased a different axis. Such a map now collapses to a `ConstantMap` at composition time, and `array_map_dependent_axis` (promoted from a private helper) reports "no axis" as `None` instead of falling back to that stale binding. + +Fixed the domain layout and chunked resolution of a vectorized selection whose coordinate arrays are not on the leading axes — `vindex[..., i, j]`, `vindex[..., mask]`. The gathered dimensions now follow NumPy's placement rule (in the spot the advanced indices occupied when they are adjacent, leading when a slice separates them), and the per-chunk gather is realigned to match the scatter indices instead of raising a shape-mismatch `ValueError`. diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index 1c1af51498..89efe8416f 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -182,6 +182,29 @@ LazyArray(np.arange(12).reshape(3, 4)).lazy.vindex[[0, 2], [1, 3]].result() # array([ 1, 11]) ``` +Two further NumPy rules hold in every mode. A **scalar integer drops its axis** +— it is a basic index wherever it appears, applied before any advanced index +rather than broadcast against one — and **advanced indices land where NumPy puts +them**, next to the axes they replaced when they are adjacent and leading when a +slice separates them: + +```python +x = np.arange(12).reshape(3, 4) + +LazyArray(x).lazy.oindex[1].result() +# array([4, 5, 6, 7]) + +LazyArray(x).lazy.oindex[[2, 0], 1].result() +# array([9, 1]) + +LazyArray(x).lazy.vindex[..., [3, 0]].result() +# array([[ 3, 0], +# [ 7, 4], +# [11, 8]]) +``` + +Use a length-1 list (`oindex[[1]]`) where you want to keep an axis. + Chunking is discovered from the wrapped array (`read_chunk_sizes`, then `chunks`) or declared with `chunks=`. When it is known, a view resolves through [chunk resolution](api/chunk_resolution.md) so every touched chunk is read diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 4d252a2103..982f1e14da 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -47,7 +47,11 @@ from zarr_indexing.lazy_array import LazyArray from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap -from zarr_indexing.transform import IndexTransform, selection_to_transform +from zarr_indexing.transform import ( + IndexTransform, + array_map_dependent_axis, + selection_to_transform, +) __version__ = version("zarr-indexing") @@ -66,6 +70,7 @@ "OutputIndexMap", "OutputIndexMapJSON", "__version__", + "array_map_dependent_axis", "compose", "dimension_grids_from_chunks", "index_domain_from_json", diff --git a/packages/zarr-indexing/src/zarr_indexing/boundary.py b/packages/zarr-indexing/src/zarr_indexing/boundary.py index 1ff930f9ec..a2c35bf595 100644 --- a/packages/zarr-indexing/src/zarr_indexing/boundary.py +++ b/packages/zarr-indexing/src/zarr_indexing/boundary.py @@ -117,6 +117,97 @@ def _normalize_int_array( return out +def _expanded_axis_walk(entries: tuple[Any, ...], ndim: int, mode: SelectionMode) -> list[int]: + """The starting axis each entry addresses, with an ellipsis expanded. + + The returned list has one entry per element of `entries`; the value for an + `Ellipsis` (or a `newaxis`) is the axis it starts at, which is also the axis + the following entry resumes from once the skipped axes are accounted for. + """ + for sel in entries: + if _is_bool_scalar(sel): + raise IndexError( + "boolean scalars are not valid indices; use a boolean array " + "matching the shape of the axes it selects" + ) + if sum(1 for sel in entries if sel is Ellipsis) > 1: + raise IndexError("an index can only have a single ellipsis ('...')") + consumed = sum(_axes_consumed(sel, mode) for sel in entries if sel is not Ellipsis) + if consumed > ndim: + raise IndexError( + f"too many indices for array: array has {ndim} dimensions, but {consumed} were indexed" + ) + + axes: list[int] = [] + axis = 0 + for sel in entries: + axes.append(axis) + axis += (ndim - consumed) if sel is Ellipsis else _axes_consumed(sel, mode) + return axes + + +def split_scalar_axes( + selection: Any, + domain: IndexDomain, + mode: SelectionMode, +) -> tuple[tuple[Any, ...] | None, Any]: + """Peel scalar integer indices out of a fancy selection. + + In NumPy, a scalar integer is a *basic* index wherever it appears: it drops + its axis, and is applied before the advanced indices rather than broadcast + against them. `a[0, [1, 2], :]` is `a[0][[1, 2], :]`. Neither the orthogonal + nor the vectorized path of the transform algebra models that — both widen a + scalar into a length-1 index array, which keeps the axis — so the scalars + are split off here and applied as a separate basic step first. + + Parameters + ---------- + selection + A positional orthogonal or vectorized selection. + domain + The domain of the view being indexed. + mode + `"orthogonal"` or `"vectorized"`; controls how many axes each entry + covers, which decides where the scalars sit. + + Returns + ------- + tuple + `(basic_selection, remaining_selection)`. `basic_selection` is a + full-rank basic selection in **literal** domain coordinates that drops + the scalar axes, or `None` when the selection has no scalar entries (in + which case `remaining_selection` is `selection` unchanged). + + Raises + ------ + IndexError + If a boolean scalar is used as an index, an index is out of bounds, or + too many indices are supplied. + """ + entries = selection if isinstance(selection, tuple) else (selection,) + axes = _expanded_axis_walk(entries, domain.ndim, mode) + + scalar_axes: dict[int, int] = {} + remaining: list[Any] = [] + for sel, axis in zip(entries, axes, strict=True): + if isinstance(sel, (int, np.integer)) and not _is_bool_scalar(sel): + scalar_axes[axis] = _normalize_int(int(sel), domain.shape[axis], axis) + else: + remaining.append(sel) + + if len(scalar_axes) == 0: + return None, selection + + basic: list[Any] = [] + for axis in range(domain.ndim): + lo = domain.inclusive_min[axis] + if axis in scalar_axes: + basic.append(lo + scalar_axes[axis]) + else: + basic.append(slice(lo, domain.exclusive_max[axis])) + return tuple(basic), tuple(remaining) + + def normalize_positional_selection( selection: Any, domain: IndexDomain, diff --git a/packages/zarr-indexing/src/zarr_indexing/grid.py b/packages/zarr-indexing/src/zarr_indexing/grid.py index 06d7b193c9..079914e16b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/grid.py +++ b/packages/zarr-indexing/src/zarr_indexing/grid.py @@ -189,24 +189,34 @@ def dimension_grids_from_chunks( f"entries for shape {shape_t}" ) - is_int = [isinstance(entry, (int, np.integer)) for entry in entries] - if len(entries) > 0 and all(is_int): + _CONVENTIONS = ( + "chunks must be either a uniform chunk shape (one integer per " + "dimension) or per-axis chunk sizes (one sequence of integers per " + "dimension)" + ) + # `sum` rather than `all`/`any` over a comprehension: the latter narrows + # `entry` to Never on the fall-through branch, which is not what we mean. + n_int = sum(1 for entry in entries if isinstance(entry, (int, np.integer))) + if len(entries) > 0 and n_int == len(entries): return tuple( EdgeDimensionGrid(_uniform_chunk_sizes(extent, int(entry))) for entry, extent in zip(entries, shape_t, strict=True) ) - if any(is_int): - raise ValueError( - "chunks must be either a uniform chunk shape (one integer per " - "dimension) or per-axis chunk sizes (one sequence per dimension), " - f"not a mixture; got {entries!r}" - ) + if n_int > 0: + raise ValueError(f"{_CONVENTIONS}, not a mixture; got {entries!r}") grids: list[EdgeDimensionGrid] = [] for axis, (entry, extent) in enumerate(zip(entries, shape_t, strict=True)): - # Every remaining entry is a sequence: the all-integer case returned - # above and the mixed case raised. - sizes = tuple(int(s) for s in cast("Sequence[int]", entry)) + # Every remaining entry must be a sequence of sizes: the all-integer case + # returned above and the mixed case raised. Anything else — a float, a + # string — is neither convention, and should say so rather than leak a + # TypeError from the iteration. + try: + sizes = tuple(int(s) for s in cast("Sequence[int]", entry)) + except (TypeError, ValueError): + raise ValueError( + f"{_CONVENTIONS}; entry {entry!r} for dimension {axis} is neither" + ) from None total = sum(sizes) if total != extent: raise ValueError( diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index 401e5c0ea3..d0523727a6 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -47,6 +47,20 @@ source — so it re-zeroes its coordinates on every view and speaks positions. `zarr_indexing.boundary` performs the translation between the two. +Two more NumPy rules the dialect keeps, in every mode: + +- **A scalar integer drops its axis.** It is a basic index wherever it appears, + applied before any advanced index rather than broadcast against one. So + `lazy.oindex[0]` has the shape of `x[0]`, `lazy.oindex[0, [1, 2], :]` means + `x[0][numpy.ix_([1, 2], ...)]`, and `lazy.oindex[0, 1, 2]` and + `lazy.vindex[0, 1, 2]` are both zero-rank. Use a length-1 list to keep an + axis. +- **Advanced indices land where NumPy puts them.** For a `vindex` selection that + leaves some axes unindexed, the gathered dimensions sit where the coordinate + arrays sat when those arrays are adjacent, and lead when a slice separates + them — so `lazy.vindex[..., i, j]` has shape `(x.shape[0], *broadcast)`, + matching `x[..., i, j]`. + Device caveat ------------- Chunked resolution builds its output buffer with NumPy, because the chunk @@ -63,7 +77,11 @@ import numpy as np -from zarr_indexing.boundary import SelectionMode, normalize_positional_selection +from zarr_indexing.boundary import ( + SelectionMode, + normalize_positional_selection, + split_scalar_axes, +) from zarr_indexing.chunk_resolution import ( iter_chunk_transforms, sub_transform_to_selections, @@ -72,7 +90,7 @@ from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import ( IndexTransform, - _array_map_dependent_axis, # pyright: ignore[reportPrivateUsage] + array_map_dependent_axis, selection_to_transform, ) @@ -170,7 +188,15 @@ def _discover_grids( shape: tuple[int, ...], chunks: Sequence[int] | Sequence[Sequence[int]] | None, ) -> tuple[EdgeDimensionGrid, ...] | None: - """Resolve the chunk grids for `array`, or None if it is unchunked.""" + """Resolve the chunk grids for `array`, or None if it is unchunked. + + An explicit `chunks` argument is *our* API, so a malformed one raises. A + discovered `chunks` attribute is external input we merely parse: an + attribute that does not describe a chunking of `shape` means "this object + does not advertise chunks we understand", and the array is treated as + unchunked rather than rejected. Chunking is an optimization here — reading + the whole array is always a correct fallback. + """ if chunks is not None: return dimension_grids_from_chunks(chunks, shape) declared = _read_chunk_attribute(array, "read_chunk_sizes") @@ -178,7 +204,10 @@ def _discover_grids( declared = _read_chunk_attribute(array, "chunks") if declared is None: return None - return dimension_grids_from_chunks(declared, shape) + try: + return dimension_grids_from_chunks(declared, shape) + except (ValueError, TypeError): + return None # --------------------------------------------------------------------------- # @@ -250,8 +279,15 @@ def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: in def _resolve_unchunked(array: Any, transform: IndexTransform) -> Any: """Lower a transform to one pass of array operations over an unchunked source.""" if _is_correlated(transform): - return _resolve_unchunked_correlated(array, transform) - return _resolve_unchunked_orthogonal(array, transform) + result = _resolve_unchunked_correlated(array, transform) + else: + result = _resolve_unchunked_orthogonal(array, transform) + if isinstance(result, np.generic): + # Basic indexing every axis of a NumPy array yields a scalar; the + # chunked resolver's buffer is a zero-dimensional array. Agree on the + # array, which is what `result()` documents. + return np.asarray(result) + return result def _resolve_unchunked_orthogonal(array: Any, transform: IndexTransform) -> Any: @@ -288,9 +324,17 @@ def _resolve_unchunked_orthogonal(array: Any, transform: IndexTransform) -> Any: lo = transform.domain.inclusive_min[d] hi = transform.domain.exclusive_max[d] selection.append(slice(m.offset + m.stride * lo, m.offset + m.stride * hi, m.stride)) - axis_input_dims.append( - _array_map_dependent_axis(m) if isinstance(m, ArrayMap) else m.input_dimension - ) + if isinstance(m, ArrayMap): + axis = array_map_dependent_axis(m) + if axis is None: + raise NotImplementedError( + "resolving an orthogonal ArrayMap that varies over no input " + "dimension is not supported; such a map should have been " + "collapsed to a ConstantMap" + ) + axis_input_dims.append(axis) + else: + axis_input_dims.append(m.input_dimension) result = result[tuple(selection)] return _restore_domain_axis_order(result, axis_input_dims, transform.input_rank) @@ -382,6 +426,43 @@ def _resolve_unchunked_correlated(array: Any, transform: IndexTransform) -> Any: ) +def _gathered_axis_positions(chunk_sel: tuple[Any, ...]) -> tuple[int, int] | None: + """Where NumPy puts the gathered axis of `block[chunk_sel]`, and where it belongs. + + `chunk_sel` mixes integers (which drop an axis), slices (which keep one), + and index arrays (which together collapse to a single gathered axis). NumPy + inserts that gathered axis where the advanced indices sat when they are + adjacent, and first when a *slice* separates them. + + An integer counts as an advanced index for that rule — NumPy's own + documented surprise, `a[0, :, i]` having shape `(len(i), a.shape[1])` rather + than `(a.shape[1], len(i))` — so a `ConstantMap` sitting before a slice that + precedes the index arrays pulls the gathered axis to the front, where + neither the orthogonal output selections (which address it in place) nor the + correlated flat scatter (which addresses it first) expect it. + + Returns + ------- + tuple of int, or None + `(actual, in_place)` — where NumPy put the axis, and the position it + would occupy if the advanced indices had not been separated. `None` when + `chunk_sel` gathers nothing. + """ + array_positions = [i for i, sel in enumerate(chunk_sel) if isinstance(sel, np.ndarray)] + if len(array_positions) == 0: + return None + in_place = sum( + 1 for i, sel in enumerate(chunk_sel) if i < array_positions[0] and isinstance(sel, slice) + ) + advanced = [i for i, sel in enumerate(chunk_sel) if not isinstance(sel, slice)] + first, last = advanced[0], advanced[-1] + separated = any(isinstance(sel, slice) for i, sel in enumerate(chunk_sel) if first < i < last) + if separated: + return 0, in_place + actual = sum(1 for i, sel in enumerate(chunk_sel) if i < first and isinstance(sel, slice)) + return actual, in_place + + def _resolve_chunked( array: Any, transform: IndexTransform, @@ -407,6 +488,16 @@ def _resolve_chunked( selected = block[chunk_sel] if len(drop_axes) > 0: selected = selected.squeeze(axis=drop_axes) + # `chunk_sel` addresses the block in storage order, so NumPy may put + # the gathered axis somewhere the output selections do not expect: + # first for a correlated (flat) scatter, in place for an orthogonal + # one. Line the two up. + positions = _gathered_axis_positions(chunk_sel) + if positions is not None: + actual, in_place = positions + target = 0 if needs_flat_buffer else in_place + if actual != target: + selected = np.moveaxis(selected, actual, target) out[out_sel] = selected if needs_flat_buffer: @@ -573,8 +664,16 @@ def lazy(self) -> _LazyIndexer: return _LazyIndexer(self._select) def _select(self, selection: Any, mode: SelectionMode) -> LazyArray: - literal = normalize_positional_selection(selection, self._transform.domain, mode) - composed = selection_to_transform(literal, self._transform, mode) + transform = self._transform + if mode != "basic": + # NumPy applies scalar integers as basic indices before the advanced + # ones, dropping their axes. Split them into their own step. + scalar_selection, selection = split_scalar_axes(selection, transform.domain, mode) + if scalar_selection is not None: + transform = selection_to_transform(scalar_selection, transform, "basic") + transform = transform.translate_domain_to((0,) * transform.input_rank) + literal = normalize_positional_selection(selection, transform.domain, mode) + composed = selection_to_transform(literal, transform, mode) return LazyArray._view(self._array, composed, self._grids) def __getitem__(self, selection: Any) -> Any: diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index e1a3898b1d..1adbbf74f9 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -30,7 +30,7 @@ import math from dataclasses import dataclass -from typing import Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np @@ -38,6 +38,9 @@ from zarr_indexing.errors import BoundsCheckError, VindexInvalidSelectionError from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap +if TYPE_CHECKING: + from collections.abc import Sequence + @dataclass(frozen=True, slots=True) class IndexTransform: @@ -363,7 +366,14 @@ def _intersect_orthogonal( # axis, or `input_dimension` for a degenerate length-1 array). Filter # along that axis and keep the array at full input rank so the # singleton axes it broadcasts over are preserved. - d = _array_map_dependent_axis(m) + axis = array_map_dependent_axis(m) + if axis is None: + raise ValueError( + f"output[{out_dim}] is an ArrayMap that varies over no input " + "dimension; a map with no dependency axis should have been " + "collapsed to a ConstantMap" + ) + d = axis storage = m.offset + m.stride * m.index_array mask = (storage >= lo) & (storage < hi) # The array is singleton on every axis but `d`, so its mask reduces @@ -421,8 +431,6 @@ def _intersect_correlated( `out_indices` is the flat scatter index into the (row-major flattened) output buffer, of shape `(surviving_points,) + (residual slice sizes)`. """ - corr_maps = [cast("ArrayMap", transform.output[i]) for i in correlated_dims] - # Mixing correlated and orthogonal ArrayMaps in one transform is not produced # by any single selection and is not supported here. orthogonal_array_dims = [ @@ -436,10 +444,15 @@ def _intersect_correlated( "ArrayMaps is not supported" ) - # The broadcast (dependency) axes are shared by every correlated map; they are - # the leading axes of the domain, followed by the residual slice axes. - broadcast_axes = _array_map_dependency_axes(corr_maps[0].index_array) - broadcast_shape = tuple(corr_maps[0].index_array.shape[a] for a in broadcast_axes) + # The broadcast axes are exactly the input axes no `DimensionMap` binds: a + # correlated transform's input domain is its residual slice axes plus the + # collapsed broadcast block. Deriving them by complement rather than from the + # index array's non-singleton axes keeps this correct when a broadcast axis + # is itself size 1, and when NumPy's placement rule puts the broadcast block + # somewhere other than the front (see `_broadcast_insertion_point`). + bound_axes = {m.input_dimension for m in transform.output if isinstance(m, DimensionMap)} + broadcast_axes = tuple(a for a in range(transform.input_rank) if a not in bound_axes) + broadcast_shape = tuple(transform.domain.shape[a] for a in broadcast_axes) # Joint bounds mask over the broadcast block. combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None @@ -523,23 +536,35 @@ def _intersect_correlated( ) result = IndexTransform(domain=new_domain, output=tuple(new_output)) - # Flat scatter index into the row-major output buffer of shape - # (broadcast points, residual slice sizes...): flat = point * prod(slice) + - # (row-major offset within the slice block). - prod_slice = 1 - for _d, _lo, _hi, full, _m in slice_dims: - prod_slice *= full - out_indices: np.ndarray[Any, np.dtype[np.intp]] = (surviving * prod_slice).reshape( + # Flat scatter index into the caller's row-major output buffer, whose shape + # is the *input* domain's shape. The buffer is addressed positionally, so + # this assumes a zero-origin domain — the resolvers normalize with + # `translate_domain_to` before resolving. + # + # Each surviving point is a flat index into the broadcast block; unravel it + # to per-axis coordinates so the buffer stride of each broadcast axis is + # applied at its real position, wherever NumPy's placement rule put it. + domain_shape = transform.domain.shape + buffer_strides = [1] * len(domain_shape) + for axis in range(len(domain_shape) - 2, -1, -1): + buffer_strides[axis] = buffer_strides[axis + 1] * domain_shape[axis + 1] + + point_offsets = np.zeros(n_points, dtype=np.intp) + if len(broadcast_shape) > 0: + for axis, coords_along_axis in zip( + broadcast_axes, np.unravel_index(surviving, broadcast_shape), strict=True + ): + point_offsets = point_offsets + coords_along_axis.astype(np.intp) * buffer_strides[axis] + + out_indices: np.ndarray[Any, np.dtype[np.intp]] = point_offsets.reshape( (n_points,) + (1,) * n_slice ) - running = 1 - for j in range(n_slice - 1, -1, -1): - _d, nlo, nhi, full, _m = slice_dims[j] - coords = np.arange(nlo, nhi, dtype=np.intp) * running + for j in range(n_slice): + d, nlo, nhi, _full, _m = slice_dims[j] + coords = np.arange(nlo, nhi, dtype=np.intp) * buffer_strides[d] shape = [1] * (1 + n_slice) shape[1 + j] = coords.size out_indices = out_indices + coords.reshape(shape) - running *= full return (result, out_indices.astype(np.intp)) @@ -792,18 +817,31 @@ def _apply_basic_indexing(transform: IndexTransform, selection: Any) -> IndexTra raise RuntimeError(f"unexpected: dimension {d} not handled") else: # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) - new_arr = _reindex_array(m, normalized, transform.domain) - array_input_dim: int | None = None + old_axes = set(_array_map_dependency_axes(m.index_array)) if m.input_dimension is not None: - array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) - new_output.append( - ArrayMap( - index_array=new_arr, - offset=m.offset, - stride=m.stride, - input_dimension=array_input_dim, + old_axes.add(m.input_dimension) + new_arr = _reindex_array(m, normalized, transform.domain) + if len(old_axes) > 0 and old_axes <= dropped_dims and new_arr.size == 1: + # Degenerate collapse: every input axis this map varied over was + # consumed by an integer index, so it now designates exactly one + # storage coordinate. Keeping it as an ArrayMap would leave a + # dangling `input_dimension` pointing at an axis that no longer + # exists — and, after renumbering, aliasing a *different* axis. + new_output.append( + ConstantMap(offset=m.offset + m.stride * int(new_arr.reshape(-1)[0])) + ) + else: + array_input_dim: int | None = None + if m.input_dimension is not None: + array_input_dim = old_to_new_dim.get(m.input_dimension, m.input_dimension) + new_output.append( + ArrayMap( + index_array=new_arr, + offset=m.offset, + stride=m.stride, + input_dimension=array_input_dim, + ) ) - ) return IndexTransform(domain=new_domain, output=tuple(new_output)) @@ -821,18 +859,35 @@ def _array_map_dependency_axes(index_array: np.ndarray[Any, Any]) -> tuple[int, return tuple(axis for axis, size in enumerate(index_array.shape) if size != 1) -def _array_map_dependent_axis(m: ArrayMap) -> int: +def array_map_dependent_axis(m: ArrayMap) -> int | None: """Return the single input axis an orthogonal `ArrayMap` varies over. Normally this is the array's one non-singleton axis. A degenerate length-1 orthogonal selection normalizes to an all-singleton shape (its dependency axes are empty and indistinguishable by shape from a scalar), so - `input_dimension` breaks the tie — it records the axis the map binds. + `input_dimension` breaks the tie — it records the *live* axis the map binds. + + Returns + ------- + int or None + The axis the map varies over, or `None` when it varies over no input + axis at all — a correlated (`vindex`) map, or a map that has become + constant. `None` is a real answer, not an error: callers that need an + axis must handle it rather than reading a stale `input_dimension`. + Indexing operations collapse a map whose every dependency axis is + consumed by an integer index into a `ConstantMap`, so an + `input_dimension` surviving here always names a live axis. + + Raises + ------ + ValueError + If the map varies over more than one axis, which makes it correlated + rather than orthogonal. """ dep = _array_map_dependency_axes(m.index_array) if len(dep) == 1: return dep[0] - if m.input_dimension is not None: + if len(dep) == 0: return m.input_dimension raise ValueError( f"orthogonal ArrayMap must vary over exactly one axis; got dependency " @@ -1014,6 +1069,27 @@ def __getitem__(self, selection: Any) -> IndexTransform: return _apply_vindex(self._transform, selection) +def _broadcast_insertion_point(array_dims: Sequence[int], slice_dims: Sequence[int]) -> int: + """Where the broadcast dimensions land, as a count of leading slice dimensions. + + NumPy's advanced-indexing placement rule: when the advanced indices are all + next to each other in the index tuple, the broadcast dimensions are inserted + at the spot they occupied; when a slice separates them, they lead. So + `a[:, i, j]` has shape `(len(a), *broadcast)` while `a[i, :, j]` has shape + `(*broadcast, a.shape[1])`. + + Returns the number of slice dimensions that precede the broadcast block; `0` + means the broadcast dimensions lead. + """ + if len(array_dims) == 0: + return 0 + first, last = array_dims[0], array_dims[-1] + separated = any(first < d < last for d in slice_dims) + if separated: + return 0 + return sum(1 for d in slice_dims if d < first) + + def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: """Apply vectorized indexing to an IndexTransform. @@ -1092,27 +1168,30 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: broadcast_arrays = [] broadcast_shape = () - # Build new domain: broadcast dims first, then slice dims - new_inclusive_min: list[int] = [] - new_exclusive_max: list[int] = [] - - # Broadcast dimensions - for s in broadcast_shape: - new_inclusive_min.append(0) - new_exclusive_max.append(s) - # Slice dimensions (preserved-domain literal semantics, like basic indexing) slice_dim_params: dict[int, tuple[int, int, int]] = {} + slice_bounds: list[tuple[int, int]] = [] for old_dim in slice_dims: sel = processed[old_dim] assert isinstance(sel, slice) lo = transform.domain.inclusive_min[old_dim] hi = transform.domain.exclusive_max[old_dim] start, step, origin, size = _resolve_slice_ts(sel, old_dim, lo, hi) - new_inclusive_min.append(origin) - new_exclusive_max.append(origin + size) + slice_bounds.append((origin, origin + size)) slice_dim_params[old_dim] = (start, step, origin) + n_before = _broadcast_insertion_point(array_dims, slice_dims) + + # Build the new domain with NumPy's placement rule: the broadcast + # (correlated) dimensions sit where the advanced indices sat when those are + # adjacent, and lead when a slice separates them. + new_inclusive_min = [lo for lo, _ in slice_bounds[:n_before]] + new_exclusive_max = [hi for _, hi in slice_bounds[:n_before]] + new_inclusive_min.extend([0] * len(broadcast_shape)) + new_exclusive_max.extend(broadcast_shape) + new_inclusive_min.extend(lo for lo, _ in slice_bounds[n_before:]) + new_exclusive_max.extend(hi for _, hi in slice_bounds[n_before:]) + new_domain = IndexDomain( inclusive_min=tuple(new_inclusive_min), exclusive_max=tuple(new_exclusive_max), @@ -1139,7 +1218,9 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: # dependency axes derived from the shape coincide — the signature # of a pointwise scatter rather than an outer product. broadcast_arr = array_dim_to_broadcast[d] - full_arr = broadcast_arr.reshape(broadcast_shape + (1,) * len(slice_dims)) + full_arr = broadcast_arr.reshape( + (1,) * n_before + broadcast_shape + (1,) * (len(slice_dims) - n_before) + ) new_output.append( ArrayMap( index_array=full_arr, @@ -1152,7 +1233,8 @@ def _apply_vindex(transform: IndexTransform, selection: Any) -> IndexTransform: start, step, origin = slice_dim_params[d] new_offset = m.offset + m.stride * (start - step * origin) new_stride = m.stride * step - new_input_dim = n_broadcast_dims + slice_dims.index(d) + position = slice_dims.index(d) + new_input_dim = position if position < n_before else position + n_broadcast_dims new_output.append( DimensionMap( input_dimension=new_input_dim, offset=new_offset, stride=new_stride diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 34ec109870..25d7710d8e 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -15,7 +15,16 @@ import numpy as np import pytest -from zarr_indexing import EdgeDimensionGrid, LazyArray, dimension_grids_from_chunks +from zarr_indexing import ( + ArrayMap, + ConstantMap, + DimensionMap, + EdgeDimensionGrid, + IndexTransform, + LazyArray, + array_map_dependent_axis, + dimension_grids_from_chunks, +) if TYPE_CHECKING: from collections.abc import Callable, Sequence @@ -31,9 +40,31 @@ def reference() -> np.ndarray[Any, np.dtype[np.int64]]: def outer(ref: np.ndarray[Any, Any], selections: Sequence[Any]) -> np.ndarray[Any, Any]: - """NumPy oracle for orthogonal indexing: the outer product of per-axis selections.""" - axes = [np.arange(size)[sel] for size, sel in zip(ref.shape, selections, strict=True)] - return ref[np.ix_(*axes)] + """NumPy oracle for orthogonal indexing: the outer product of per-axis selections. + + Scalar integers are basic indices — NumPy applies them first and drops the + axis — so they are peeled off before the outer product is formed. + """ + scalars = tuple( + sel if isinstance(sel, (int, np.integer)) and not isinstance(sel, bool) else slice(None) + for sel in selections + ) + reduced = ref[scalars] + axes = [ + np.arange(size)[sel] + for size, sel in zip( + reduced.shape, + [ + s + for s in selections + if not (isinstance(s, (int, np.integer)) and not isinstance(s, bool)) + ], + strict=True, + ) + ] + if len(axes) == 0: + return reduced + return reduced[np.ix_(*axes)] # --------------------------------------------------------------------------- @@ -125,6 +156,31 @@ def test_dimension_grids_from_chunks_rejects_wrong_total() -> None: dimension_grids_from_chunks(((3, 2), (4,)), (7, 4)) +def test_dimension_grids_from_chunks_rejects_non_sequence_entry() -> None: + """A float entry is neither convention, and says so instead of raising TypeError.""" + with pytest.raises(ValueError, match="entry 3.5 for dimension 0 is neither"): + dimension_grids_from_chunks((3.5, (4,)), (7, 4)) + + +def test_array_map_dependent_axis_reports_no_axis() -> None: + """A map varying over nothing answers None rather than a stale binding.""" + correlated = ArrayMap(index_array=np.array([[1], [2]], dtype=np.intp)) + assert array_map_dependent_axis(correlated) == 0 + assert array_map_dependent_axis(ArrayMap(index_array=np.array([[3]], dtype=np.intp))) is None + + +def test_scalar_on_a_fancy_axis_collapses_to_a_constant() -> None: + """The degenerate-collapse rule: an all-singleton ArrayMap becomes a ConstantMap. + + Without it the map keeps an `input_dimension` naming an axis the integer + index just removed, which after renumbering aliases a different axis. + """ + view = IndexTransform.from_shape((7, 5)).oindex[np.array([3, 1]), slice(None)][0] + assert view.output[0] == ConstantMap(offset=3) + assert isinstance(view.output[1], DimensionMap) + assert view.domain.shape == (5,) + + # --------------------------------------------------------------------------- # Sources # --------------------------------------------------------------------------- @@ -160,6 +216,8 @@ def source(request: pytest.FixtureRequest) -> LazyArray: # --------------------------------------------------------------------------- MASK = (reference() % 11) == 0 +# A mask over the two trailing axes, for the `vindex[..., mask]` idiom. +TRAILING_MASK = (reference()[0] % 3) == 0 # (id, lazy view builder, NumPy oracle). Integer scalars are deliberately absent # from the orthogonal cases: the transform algebra keeps an int-selected axis as @@ -230,6 +288,102 @@ def source(request: pytest.FixtureRequest) -> LazyArray: ], lambda r: r[1:6, :, 1:][np.array([0, 4]), np.array([2, 0]), np.array([1, 2])], ), + # Scalar integers are basic indices in the positional dialect: they drop the + # axis, in every mode, exactly as NumPy does. + ("oindex-scalar-drops-axis", lambda a: a.lazy.oindex[0], lambda r: r[0]), + ( + "oindex-scalar-with-arrays", + lambda a: a.lazy.oindex[0, [1, 2], :], + lambda r: outer(r, (0, [1, 2], slice(None))), + ), + ( + "oindex-scalar-middle-axis", + lambda a: a.lazy.oindex[[3, 1], -1, :], + lambda r: outer(r, ([3, 1], -1, slice(None))), + ), + ("oindex-all-scalars", lambda a: a.lazy.oindex[0, 1, 2], lambda r: r[0, 1, 2]), + ("vindex-all-scalars", lambda a: a.lazy.vindex[0, 1, 2], lambda r: r[0, 1, 2]), + ( + "vindex-scalar-with-arrays", + lambda a: a.lazy.vindex[0, [1, 2], [3, 0]], + lambda r: r[0, [1, 2], [3, 0]], + ), + ( + "vindex-scalar-on-middle-axis", + lambda a: a.lazy.vindex[[1, 2], 0, [3, 0]], + lambda r: r[[1, 2], 0, [3, 0]], + ), + # Scalar applied to a previously fancy-indexed axis, both orders. + ( + "compose-oindex-then-scalar", + lambda a: a.lazy.oindex[[3, 1], :, :].lazy[0], + lambda r: outer(r, ([3, 1], slice(None), slice(None)))[0], + ), + ( + "compose-oindex-then-scalar-negative", + lambda a: a.lazy.oindex[[3, 1, 1], :, :].lazy[-1, 2], + lambda r: outer(r, ([3, 1, 1], slice(None), slice(None)))[-1, 2], + ), + ( + "compose-scalar-then-oindex", + lambda a: a.lazy[0].lazy.oindex[[3, 1], :], + lambda r: outer(r[0], ([3, 1], slice(None))), + ), + ( + "compose-vindex-then-scalar", + lambda a: a.lazy.vindex[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])].lazy[ + 1 + ], + lambda r: r[np.array([0, 6, 3]), np.array([1, 4, 0]), np.array([2, 0, 1])][1], + ), + # Partial vindex whose coordinate arrays are NOT on the leading axes: NumPy + # inserts the gathered axis where the (adjacent) advanced indices sat. + ( + "vindex-trailing-arrays", + lambda a: a.lazy.vindex[..., np.array([1, 4, 0]), np.array([2, 0, 1])], + lambda r: r[..., np.array([1, 4, 0]), np.array([2, 0, 1])], + ), + ( + "vindex-single-trailing-array", + lambda a: a.lazy.vindex[..., np.array([3, 0, 1])], + lambda r: r[..., np.array([3, 0, 1])], + ), + ( + "vindex-trailing-mask", + lambda a: a.lazy.vindex[..., TRAILING_MASK], + lambda r: r[..., TRAILING_MASK], + ), + ( + "vindex-leading-partial", + lambda a: a.lazy.vindex[np.array([1, 2, 2])], + lambda r: r[np.array([1, 2, 2])], + ), + ( + "compose-basic-then-vindex-trailing", + lambda a: a.lazy[2:, 1:].lazy.vindex[..., np.array([1, 3, 0])], + lambda r: r[2:, 1:][..., np.array([1, 3, 0])], + ), + # A ConstantMap sitting between a slice and the coordinate arrays: NumPy + # counts the integer as an advanced index, so the gathered axis moves to the + # front of the chunk block even though the arrays are trailing. + ( + "compose-vindex-trailing-then-scalar", + lambda a: a.lazy.vindex[..., np.array([3, 0])].lazy[2], + lambda r: r[..., np.array([3, 0])][2], + ), + # Two fancy axes, then a scalar on the first: the surviving ArrayMap ends up + # behind a ConstantMap and a slice, which is where NumPy's integer-counts-as- + # advanced rule bites. + ( + "compose-oindex-two-axes-then-scalar", + lambda a: a.lazy.oindex[[3, 1, 0], 3:5, [2, 0, 2]].lazy[0], + lambda r: outer(r, ([3, 1, 0], slice(3, 5), [2, 0, 2]))[0], + ), + ( + "compose-vindex-trailing-then-scalar-and-slice", + lambda a: a.lazy.vindex[..., np.array([3, 0, 1])].lazy[-1, 1:4], + lambda r: r[..., np.array([3, 0, 1])][-1, 1:4], + ), ] @@ -250,6 +404,133 @@ def test_selection_matches_numpy( np.testing.assert_array_equal(np.asarray(view), expected) +# --------------------------------------------------------------------------- +# Randomized chain sweep +# --------------------------------------------------------------------------- + + +def _random_basic(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: + selection: list[Any] = [] + for size in shape: + roll = rng.random() + if roll < 0.3: + selection.append(int(rng.integers(-size, size))) + elif roll < 0.7: + start = int(rng.integers(0, size)) + stop = int(rng.integers(start, size + 1)) + selection.append(slice(start, stop, int(rng.integers(1, 4)))) + else: + selection.append(slice(None)) + return tuple(selection) + + +def _random_oindex(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: + selection: list[Any] = [] + for size in shape: + roll = rng.random() + if roll < 0.25: + selection.append(int(rng.integers(-size, size))) + elif roll < 0.65: + count = int(rng.integers(1, 5)) + selection.append(rng.integers(-size, size, size=count).tolist()) + elif roll < 0.8: + mask = rng.random(size) < 0.5 + mask[int(rng.integers(0, size))] = True + selection.append(mask) + else: + start = int(rng.integers(0, size)) + selection.append(slice(start, size)) + return tuple(selection) + + +def _random_vindex(rng: np.random.Generator, shape: tuple[int, ...]) -> tuple[Any, ...]: + ndim = len(shape) + count = int(rng.integers(1, ndim + 1)) + trailing = bool(rng.random() < 0.5) + axes = range(ndim - count, ndim) if trailing else range(count) + sizes = [shape[axis] for axis in axes] + + entries: list[Any] + if rng.random() < 0.2: + # A single boolean mask spanning the whole covered block. + mask = rng.random(tuple(sizes)) < 0.5 + mask.flat[int(rng.integers(0, mask.size))] = True + entries = [mask] + else: + length = int(rng.integers(1, 5)) + entries = [ + int(rng.integers(-size, size)) + if rng.random() < 0.25 + else rng.integers(-size, size, size=length) + for size in sizes + ] + return (Ellipsis, *entries) if trailing else tuple(entries) + + +def _apply_oracle( + ref: np.ndarray[Any, Any], mode: str, selection: tuple[Any, ...] +) -> np.ndarray[Any, Any]: + if mode == "orthogonal": + return outer(ref, selection) + # NumPy's own semantics *are* basic and vectorized indexing. + return ref[selection] + + +def _apply_view(view: LazyArray, mode: str, selection: tuple[Any, ...]) -> LazyArray: + if mode == "basic": + return view.lazy[selection] + if mode == "orthogonal": + return view.lazy.oindex[selection] + return view.lazy.vindex[selection] + + +def _random_chain(rng: np.random.Generator) -> list[tuple[str, tuple[Any, ...]]]: + """A chain of 2-4 steps with at most one fancy step, in either order.""" + n_steps = int(rng.integers(2, 5)) + fancy_at = int(rng.integers(0, n_steps)) + fancy_mode = "orthogonal" if rng.random() < 0.5 else "vectorized" + + chain: list[tuple[str, tuple[Any, ...]]] = [] + running = reference() + for step in range(n_steps): + if running.ndim == 0 or running.size == 0: + break + mode = fancy_mode if step == fancy_at else "basic" + if mode == "basic": + selection = _random_basic(rng, running.shape) + elif mode == "orthogonal": + selection = _random_oindex(rng, running.shape) + else: + selection = _random_vindex(rng, running.shape) + chain.append((mode, selection)) + running = _apply_oracle(running, mode, selection) + return chain + + +@pytest.mark.parametrize( + "flavour", ["numpy-unchunked", "numpy-uniform-chunks", "numpy-dask-chunks"] +) +def test_random_chains_match_numpy(flavour: str) -> None: + """A seeded sweep of composed chains: both resolvers must match NumPy.""" + rng = np.random.default_rng(20260730) + source = make_source(flavour) + + for _ in range(400): + chain = _random_chain(rng) + expected = reference() + for mode, selection in chain: + expected = _apply_oracle(expected, mode, selection) + + view = source + for mode, selection in chain: + view = _apply_view(view, mode, selection) + + assert view.shape == expected.shape, f"{flavour}: {chain}" + np.testing.assert_array_equal( + np.asarray(view.result()), np.asarray(expected), err_msg=f"{flavour}: {chain}" + ) + + def test_eager_getitem_returns_data(source: LazyArray) -> None: """`arr[...]` reads immediately, like `numpy.ndarray.__getitem__`.""" np.testing.assert_array_equal(np.asarray(source[1:3, ::2, -1]), reference()[1:3, ::2, -1]) @@ -303,6 +584,61 @@ def test_chunks_follow_a_basic_view() -> None: assert array.lazy[::2, :, :].chunks[0] == (2, 1, 1) +def test_scalar_is_basic_even_when_a_slice_separates_it_from_the_arrays() -> None: + """A documented, deliberate departure from one NumPy corner. + + NumPy counts an integer as an *advanced* index for its placement rule, so + `x[0, :, i]` has shape `(len(i), x.shape[1])` — the arrays lead because the + slice separates the integer from them. The positional dialect instead treats + every scalar as a basic index applied first, which is the rule the rest of + the surface follows and the only one `oindex` can express, so the same + selection reads as `x[0][:, i]`. + """ + data = reference() + view = make_source("numpy-uniform-chunks").lazy.vindex[0, ..., np.array([3, 0])] + np.testing.assert_array_equal(np.asarray(view.result()), data[0][..., np.array([3, 0])]) + assert view.shape == data[0][..., np.array([3, 0])].shape + assert view.shape != data[0, ..., np.array([3, 0])].shape + + +def test_zero_dimensional_result_is_an_array(source: LazyArray) -> None: + """Both resolvers agree on the kind of a zero-rank result.""" + result = source.lazy[0, 1, 2].result() + assert isinstance(result, np.ndarray) + assert result.ndim == 0 + assert result[()] == reference()[0, 1, 2] + + +def test_malformed_discovered_chunks_are_ignored() -> None: + """A foreign object's unusable `chunks` falls back to whole-array reads. + + Discovery parses external input, so an attribute that does not describe a + chunking of the shape means "no chunks I understand", not an error. + """ + + class ForeignArray: + def __init__(self, data: np.ndarray[Any, Any], chunks: Any) -> None: + self._data = data + self.chunks = chunks + + @property + def shape(self) -> tuple[int, ...]: + return self._data.shape + + @property + def dtype(self) -> Any: + return self._data.dtype + + def __getitem__(self, key: Any) -> Any: + return self._data[key] + + data = reference() + for bogus in (((3, 3), (5,), (4,)), (3, 2), "nope", (3.5, 2, 2)): + wrapped = LazyArray(ForeignArray(data, bogus)) + assert wrapped.chunks is None + np.testing.assert_array_equal(np.asarray(wrapped.lazy[1:3].result()), data[1:3]) + + def test_chunks_are_undefined_for_fancy_views() -> None: """Reordered/duplicated coordinates have no chunk-size description.""" array = make_source("numpy-uniform-chunks") From c17785a1eafcdbc90c5ff33a13d5634594599c67 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Thu, 30 Jul 2026 21:44:10 +0200 Subject: [PATCH 4/4] refactor(zarr-indexing)!: LazyArray reads in parts, not chunks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Retires `chunks` from the wrapper's surface entirely — no property, no constructor argument. `chunks` is the *source's* word: discovery still reads `read_chunk_sizes`/`chunks` off the wrapped object, because that is the vocabulary those objects speak, but what the wrapper exposes is parts. The parts family: - `Partition` — one box of the base partitioning as it falls through the view: its base coordinates, a resolvable `LazyArray` covering exactly the view's cells inside it, where those cells belong in the result, and whether the view covers the whole box. - `LazyArray.parts()` — iterate them. They tile the view exactly and disjointly, and each resolves independently (and concurrently). - `LazyArray.with_parts()` — same base, same view, new boxes. Uniform shape, explicit per-axis sizes, or `None` for one whole-array part. Validated strictly: it is our API, unlike discovery, which parses and falls back. `result()` is now nothing but the assembly of that walk, so there is one partitioner and one lowering engine instead of two resolvers. The round-1 advanced-index placement fixes live in the lowering engine, which always returns its result in the view's own axis order; scattering a part then needs no realignment at all, and the correlated flat buffer is gone (scatter offsets are unravelled against the result's real shape). Dunders, each delegating rather than reimplementing: `__dask_tokenize__` (the wrapped array's token, the transform's canonical ndsel body, and the partitioning; dask imported lazily and never required), `__iter__`, `__bool__`, `__int__`, `__float__`, `__index__`. NumPy decides what a size-1 conversion means, exception types included. Deliberately absent: `__array_namespace__`, `__array_ufunc__`/`__array_function__`, the dask collection protocol — every non-indexing NumPy operation materializes through `__array__`, which the class docstring now says prominently. Tests: the partitioning invariant as a property — the seeded chain sweep runs each chain under four partitionings across five source flavours, including boxes deliberately straddling a wrapped zarr array's own chunks. Plus part-coverage (assembly equals `result()`, placements cover with no overlap), concurrent per-part resolution, completeness reporting, `with_parts` validation, and one test per dunder. Verified offline over 43,200 chain x partitioning evaluations with zero divergences. Assisted-by: ClaudeCode:claude-fable-5 --- packages/zarr-indexing/changes/267.feature.md | 8 +- packages/zarr-indexing/docs/api/index.md | 4 +- packages/zarr-indexing/docs/index.md | 45 +- .../src/zarr_indexing/__init__.py | 3 +- .../src/zarr_indexing/lazy_array.py | 627 ++++++++++++------ .../zarr-indexing/tests/test_lazy_array.py | 368 +++++++--- 6 files changed, 755 insertions(+), 300 deletions(-) diff --git a/packages/zarr-indexing/changes/267.feature.md b/packages/zarr-indexing/changes/267.feature.md index 5f64b18c73..045ee1fd12 100644 --- a/packages/zarr-indexing/changes/267.feature.md +++ b/packages/zarr-indexing/changes/267.feature.md @@ -1,5 +1,9 @@ -Added `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, CuPy, ...) TensorStore-style lazy indexing. `LazyArray(x).lazy[...]`, `.lazy.oindex[...]`, and `.lazy.vindex[...]` compose an `IndexTransform` and return a new view without reading data; `result()` materializes it. Selections use positional NumPy semantics (zero-based within the current view, negatives wrap, masks must match the view shape, scalar integers drop their axis, and a partial `vindex` places its gathered dimensions where NumPy would), which the new `zarr_indexing.boundary` module translates into the transform algebra's literal coordinates. `__getitem__` reads eagerly, so a `LazyArray` is a drop-in source for `dask.array.from_array`. Chunking is taken from an explicit `chunks=` argument, then `read_chunk_sizes`, then `chunks`; a chunked source resolves through `iter_chunk_transforms` so each touched chunk is read exactly once, and an unchunked source lowers to a single pass of array operations. `zarr_indexing.grid` gained `EdgeDimensionGrid`, a concrete `DimensionGridLike` built from explicit per-chunk sizes, and `dimension_grids_from_chunks`, which normalizes either chunk convention into one grid per axis. +Added `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, CuPy, ...) TensorStore-style lazy indexing. `LazyArray(x).lazy[...]`, `.lazy.oindex[...]`, and `.lazy.vindex[...]` compose an `IndexTransform` and return a new view without reading data; `result()` materializes it. Selections use positional NumPy semantics (zero-based within the current view, negatives wrap, masks must match the view shape, scalar integers drop their axis, and a partial `vindex` places its gathered dimensions where NumPy would), which the new `zarr_indexing.boundary` module translates into the transform algebra's literal coordinates. `__getitem__` reads eagerly, so a `LazyArray` is a drop-in source for `dask.array.from_array`; every other NumPy operation materializes the view through `__array__`. `__dask_tokenize__`, `__len__`, `__iter__`, `__bool__`, `__int__`, `__float__`, and `__index__` round out the surface, and a wrapper pickles as long as its base does. + +A read is broken up along a **partitioning** — a grid of boxes, discovered from the wrapped array's `read_chunk_sizes` or `chunks` (the source's vocabulary; the wrapper's own surface speaks only of parts). `LazyArray.parts()` walks those boxes as they fall through the view, yielding a `Partition` with the box's base coordinates, a resolvable `LazyArray` for the cells of the view inside it, where those cells belong in the result, and whether the box is fully covered. `result()` is the assembly of that walk and nothing else, so one partitioner and one lowering engine serve every read. `LazyArray.with_parts()` chooses a different partitioning — uniform box shape, explicit per-axis sizes, or `None` for a single whole-array part — without touching the data or the view, and never changes what `result()` returns. + +`zarr_indexing.grid` gained `EdgeDimensionGrid`, a concrete `DimensionGridLike` built from explicit per-part sizes, and `dimension_grids_from_chunks`, which normalizes either convention into one grid per axis. Fixed an integer index applied to an axis a previous orthogonal or vectorized selection had already indexed. The composed `ArrayMap` became all-singleton but kept an `input_dimension` naming the axis the integer had just removed, which after renumbering aliased a different axis. Such a map now collapses to a `ConstantMap` at composition time, and `array_map_dependent_axis` (promoted from a private helper) reports "no axis" as `None` instead of falling back to that stale binding. -Fixed the domain layout and chunked resolution of a vectorized selection whose coordinate arrays are not on the leading axes — `vindex[..., i, j]`, `vindex[..., mask]`. The gathered dimensions now follow NumPy's placement rule (in the spot the advanced indices occupied when they are adjacent, leading when a slice separates them), and the per-chunk gather is realigned to match the scatter indices instead of raising a shape-mismatch `ValueError`. +Fixed the domain layout and partitioned resolution of a vectorized selection whose coordinate arrays are not on the leading axes — `vindex[..., i, j]`, `vindex[..., mask]`. The gathered dimensions now follow NumPy's placement rule (in the spot the advanced indices occupied when they are adjacent, leading when a slice separates them), and the per-part gather is realigned to match the scatter indices instead of raising a shape-mismatch `ValueError`. diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index 827f32202b..09b6bb6b32 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -36,8 +36,8 @@ and the wire format built on top of it. - [`zarr_indexing.lazy_array`](lazy_array.md) — `LazyArray`, a wrapper that gives any array-API-like array (NumPy, zarr, CuPy, …) a `.lazy` accessor for - TensorStore-style deferred indexing, resolving chunk-aware when the wrapped - array declares chunks + TensorStore-style deferred indexing, plus `Partition` and the `parts()` / + `with_parts()` pair that decide the boxes a read is broken into - [`zarr_indexing.boundary`](boundary.md) — the translation between NumPy's positional dialect and the transform algebra's literal coordinates diff --git a/packages/zarr-indexing/docs/index.md b/packages/zarr-indexing/docs/index.md index 89efe8416f..82125d3c86 100644 --- a/packages/zarr-indexing/docs/index.md +++ b/packages/zarr-indexing/docs/index.md @@ -205,12 +205,45 @@ LazyArray(x).lazy.vindex[..., [3, 0]].result() Use a length-1 list (`oindex[[1]]`) where you want to keep an axis. -Chunking is discovered from the wrapped array (`read_chunk_sizes`, then -`chunks`) or declared with `chunks=`. When it is known, a view resolves through -[chunk resolution](api/chunk_resolution.md) so every touched chunk is read -exactly once; when it is not, the view lowers to a single pass of array -operations. Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a -`LazyArray` a drop-in source for `dask.array.from_array`. +Plain `lazy[...]` (no `.lazy`) reads eagerly, which makes a `LazyArray` a +drop-in source for `dask.array.from_array`. Everything else NumPy might do to it +— arithmetic, `numpy.sum`, `numpy.stack` — materializes the whole view through +`__array__` and hands back a plain NumPy array: laziness here is about indexing, +not about deferred compute. + +### Parts + +A read is broken up along a **partitioning**: a grid of boxes the wrapper walks +through [chunk resolution](api/chunk_resolution.md), so each box is read once +with plain basic slicing and the selection is applied to the block in memory. +`parts()` exposes that walk, projected through the view: + +```python +part = next(iter(lazy.lazy[0:35, 0:10].parts())) + +part.base_coords # (0, 0) — which box of the base grid +part.array.shape # (30, 10) — a LazyArray for this box's cells +part.out_selection # (slice(0, 30), slice(0, 10)) +part.is_complete # False — the view does not cover the whole box +``` + +Each `part.array` is an ordinary `LazyArray`, so parts resolve independently and +concurrently, and `result()` is nothing but the assembly of that walk. + +The partitioning is discovered from the wrapped array — `read_chunk_sizes`, then +`chunks`. Those are the *source's* words; the wrapper's own surface speaks only +of parts, and `with_parts` chooses a different one without touching the data or +the view: + +```python +view.with_parts((50, 20)) # uniform boxes, tail clipped +view.with_parts(((30, 30), ...)) # explicit per-axis sizes +view.with_parts(None) # one whole-array part; resolve in one shot +``` + +`result()` is identical under any of them — repartitioning changes how the read +is cut up, never what it returns. Boxes that deliberately straddle the source's +own (to bound peak memory, say) are legal; they cost I/O, not correctness. ## Reference diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 982f1e14da..c152d7d6c6 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -44,7 +44,7 @@ transform_from_canonical, transform_to_canonical, ) -from zarr_indexing.lazy_array import LazyArray +from zarr_indexing.lazy_array import LazyArray, Partition from zarr_indexing.messages import NdselError, normalize_ndsel, parse_ndsel from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_indexing.transform import ( @@ -69,6 +69,7 @@ "NdselError", "OutputIndexMap", "OutputIndexMapJSON", + "Partition", "__version__", "array_map_dependent_axis", "compose", diff --git a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py index d0523727a6..aa27f19c13 100644 --- a/packages/zarr-indexing/src/zarr_indexing/lazy_array.py +++ b/packages/zarr-indexing/src/zarr_indexing/lazy_array.py @@ -13,23 +13,44 @@ Nothing is read until `result()` (or `__array__`, or an eager `__getitem__`). Composition is free: a view of a view is still a single transform. -Chunk awareness ---------------- -If the wrapped array is chunked, `LazyArray` resolves a view against the chunk -grid — each chunk is read exactly once with plain basic slicing, then the -selected coordinates are scattered into the output buffer. Chunking is -discovered in this order: - -1. an explicit `chunks=` argument; -2. `array.read_chunk_sizes` (zarr's clipped per-axis sizes, sharding-aware); -3. `array.chunks`, read as per-axis sizes if its entries are sequences and as a - uniform chunk shape if they are integers; -4. otherwise the array is treated as unchunked and the view is lowered to a - single pass of array operations (`take`, basic slicing, a flat gather). - -Both resolution strategies compute the same values; the chunked one exists to -keep reads chunk-aligned, the unchunked one to avoid the per-chunk bookkeeping -when there is nothing to align to. +Parts +----- +A `LazyArray` carries a **partitioning** of the array it wraps: a grid of boxes +that a read is broken into. `parts()` walks those boxes as they fall through the +view, yielding a [`Partition`](#zarr_indexing.lazy_array.Partition) per box — +the base coordinates of the box, a `LazyArray` covering exactly the cells of the +view that live in it, and where those cells belong in the result. `result()` is +built on that walk and nothing else: + +```python +for part in view.parts(): + out[part.out_selection] = part.array.result() +``` + +so each box is read once, with plain basic slicing, and the selection is applied +to the block in memory. + +The partitioning is discovered from the wrapped array at construction — first +`read_chunk_sizes` (zarr's clipped per-axis sizes, sharding-aware), then +`chunks`, read as per-axis sizes if its entries are sequences and as a uniform +box shape if they are integers. Those attribute names are the *source's* +vocabulary; the wrapper's own surface speaks only of parts. An array that +advertises neither is a single whole-array part, and resolving it lowers the +whole view to one pass of array operations (`take`, basic slicing, a flat +gather) rather than materializing a block first. + +`with_parts` replaces the partitioning without touching the data or the view: + +```python +view.with_parts((64, 64)) # uniform boxes, tail clipped +view.with_parts(((3, 3, 1),)) # explicit per-axis sizes +view.with_parts(None) # one whole-array part; resolve in one shot +``` + +Repartitioning never changes what `result()` returns — only how the read is cut +up. Deliberately misaligning the parts with the source's own boxes is legal and +useful (to bound peak memory, or to batch small reads); it costs I/O, not +correctness. The positional dialect ---------------------- @@ -61,18 +82,33 @@ them — so `lazy.vindex[..., i, j]` has shape `(x.shape[0], *broadcast)`, matching `x[..., i, j]`. +Materializing on fallback +------------------------- +`LazyArray` implements `__array__` but deliberately implements neither +`__array_ufunc__`/`__array_function__` nor `__array_namespace__`. **Every NumPy +operation other than indexing therefore materializes the whole view**: +`numpy.sum(view)`, `view + 1`, `numpy.stack([view, view])` all convert through +`__array__` first and give you a plain NumPy array back. That is intentional — +laziness here is about *indexing*, not about building a deferred compute graph — +but it means a `LazyArray` is not a drop-in for arithmetic on a large array. Use +`.lazy[...]` to narrow first, or hand the wrapper to `dask.array.from_array` and +let dask own the compute graph. + Device caveat ------------- -Chunked resolution builds its output buffer with NumPy, because the chunk +Resolving a partitioned view builds its output buffer with NumPy, because the resolver's scatter indices are host-side integer arrays. Wrapping a device array -that declares chunks therefore returns host memory. The unchunked path stays in -the wrapped array's own namespace. +that advertises parts therefore returns host memory. A single whole-array part +lowers directly and stays in the wrapped array's own namespace. """ from __future__ import annotations -import itertools +import hashlib +import json import math +import operator +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Protocol import numpy as np @@ -87,6 +123,7 @@ sub_transform_to_selections, ) from zarr_indexing.grid import EdgeDimensionGrid, dimension_grids_from_chunks +from zarr_indexing.json import transform_to_canonical from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap from zarr_indexing.transform import ( IndexTransform, @@ -95,11 +132,15 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Sequence + from collections.abc import Callable, Iterator, Sequence SelectFn = Callable[[Any, SelectionMode], "LazyArray"] -__all__ = ["LazyArray"] +__all__ = ["LazyArray", "Partition"] + +# Above this many bytes, the no-dask token fallback describes an array +# structurally instead of digesting its contents. See `_wrapped_token`. +_TOKEN_DIGEST_LIMIT = 1 << 20 class ArrayLike(Protocol): @@ -163,19 +204,18 @@ def _expand_dims(x: Any, axis: int) -> Any: # --------------------------------------------------------------------------- # -# Chunk discovery +# Partition discovery # --------------------------------------------------------------------------- # -def _read_chunk_attribute(array: Any, name: str) -> Any: - """Read a chunk-describing attribute, treating any failure as "absent". +def _read_source_attribute(array: Any, name: str) -> Any: + """Read a partition-describing attribute, treating any failure as "absent". - Chunk discovery inspects an object we did not write. A missing attribute is - the common case, but zarr raises an `AttributeError` subclass from + Discovery inspects an object we did not write. A missing attribute is the + common case, but zarr raises an `AttributeError` subclass from `read_chunk_sizes` on a lazy view, and other backends compute the attribute lazily and may fail for their own reasons. Any failure here means "this - array does not advertise chunks", never a hard error — while a malformed - *value* still raises from `dimension_grids_from_chunks`. + array does not advertise a partitioning", never a hard error. """ try: return getattr(array, name, None) @@ -183,25 +223,18 @@ def _read_chunk_attribute(array: Any, name: str) -> Any: return None -def _discover_grids( - array: Any, - shape: tuple[int, ...], - chunks: Sequence[int] | Sequence[Sequence[int]] | None, -) -> tuple[EdgeDimensionGrid, ...] | None: - """Resolve the chunk grids for `array`, or None if it is unchunked. +def _discover_parts(array: Any, shape: tuple[int, ...]) -> tuple[EdgeDimensionGrid, ...] | None: + """Resolve the partitioning advertised by `array`, or None for one whole part. - An explicit `chunks` argument is *our* API, so a malformed one raises. A - discovered `chunks` attribute is external input we merely parse: an - attribute that does not describe a chunking of `shape` means "this object - does not advertise chunks we understand", and the array is treated as - unchunked rather than rejected. Chunking is an optimization here — reading - the whole array is always a correct fallback. + Discovery parses external input: an attribute that does not describe a + partitioning of `shape` means "this object does not advertise one I + understand", and the array is treated as unpartitioned rather than rejected. + A partitioning is an I/O strategy, so reading the whole array is always a + correct fallback. `with_parts` is *our* API, and validates strictly. """ - if chunks is not None: - return dimension_grids_from_chunks(chunks, shape) - declared = _read_chunk_attribute(array, "read_chunk_sizes") + declared = _read_source_attribute(array, "read_chunk_sizes") if declared is None: - declared = _read_chunk_attribute(array, "chunks") + declared = _read_source_attribute(array, "chunks") if declared is None: return None try: @@ -210,8 +243,13 @@ def _discover_grids( return None +def _whole_array_grids(shape: tuple[int, ...]) -> tuple[EdgeDimensionGrid, ...]: + """A partitioning with a single part covering the whole array.""" + return tuple(EdgeDimensionGrid((extent,) if extent > 0 else ()) for extent in shape) + + # --------------------------------------------------------------------------- # -# Resolution +# The lowering engine # --------------------------------------------------------------------------- # @@ -234,7 +272,7 @@ def _is_identity_transform(transform: IndexTransform, shape: tuple[int, ...]) -> def _is_correlated(transform: IndexTransform) -> bool: - """True when the transform scatters through a single flat index (vindex).""" + """True when the transform gathers a list of points rather than an outer product.""" return any(isinstance(m, ArrayMap) and m.input_dimension is None for m in transform.output) @@ -260,6 +298,10 @@ def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: in `result` corresponds to. Axes are permuted so that they appear in increasing domain-dimension order, and any domain dimension no output map depends on (only reachable via `newaxis`) is reinserted as a singleton. + + This is where NumPy's advanced-index placement rules are absorbed: whatever + order the gather produced, the lowered result always comes back in the + view's own axis order. """ if len(set(axis_input_dims)) != len(axis_input_dims): raise NotImplementedError( @@ -276,21 +318,25 @@ def _restore_domain_axis_order(result: Any, axis_input_dims: list[int], rank: in return result -def _resolve_unchunked(array: Any, transform: IndexTransform) -> Any: - """Lower a transform to one pass of array operations over an unchunked source.""" +def _lower(array: Any, transform: IndexTransform) -> Any: + """Lower a transform to one pass of array operations over `array`. + + The single lowering engine: every read, partitioned or not, ends here. The + result is always in the transform's own domain axis order and of exactly its + domain shape. + """ if _is_correlated(transform): - result = _resolve_unchunked_correlated(array, transform) + result = _lower_correlated(array, transform) else: - result = _resolve_unchunked_orthogonal(array, transform) + result = _lower_orthogonal(array, transform) if isinstance(result, np.generic): - # Basic indexing every axis of a NumPy array yields a scalar; the - # chunked resolver's buffer is a zero-dimensional array. Agree on the - # array, which is what `result()` documents. + # Basic indexing every axis of a NumPy array yields a scalar; `result()` + # documents a zero-dimensional array. return np.asarray(result) return result -def _resolve_unchunked_orthogonal(array: Any, transform: IndexTransform) -> Any: +def _lower_orthogonal(array: Any, transform: IndexTransform) -> Any: """Basic slicing plus one `take` per fancy-indexed axis (an outer product). Orthogonal `ArrayMap`s vary over distinct input axes, so gathering them one @@ -339,7 +385,7 @@ def _resolve_unchunked_orthogonal(array: Any, transform: IndexTransform) -> Any: return _restore_domain_axis_order(result, axis_input_dims, transform.input_rank) -def _resolve_unchunked_correlated(array: Any, transform: IndexTransform) -> Any: +def _lower_correlated(array: Any, transform: IndexTransform) -> Any: """Flatten the correlated axes, gather the points once, reshape back. Correlated (`vindex`) `ArrayMap`s address a list of *points*, not an outer @@ -426,100 +472,135 @@ def _resolve_unchunked_correlated(array: Any, transform: IndexTransform) -> Any: ) -def _gathered_axis_positions(chunk_sel: tuple[Any, ...]) -> tuple[int, int] | None: - """Where NumPy puts the gathered axis of `block[chunk_sel]`, and where it belongs. +# --------------------------------------------------------------------------- # +# Partitions +# --------------------------------------------------------------------------- # - `chunk_sel` mixes integers (which drop an axis), slices (which keep one), - and index arrays (which together collapse to a single gathered axis). NumPy - inserts that gathered axis where the advanced indices sat when they are - adjacent, and first when a *slice* separates them. - An integer counts as an advanced index for that rule — NumPy's own - documented surprise, `a[0, :, i]` having shape `(len(i), a.shape[1])` rather - than `(a.shape[1], len(i))` — so a `ConstantMap` sitting before a slice that - precedes the index arrays pulls the gathered axis to the front, where - neither the orthogonal output selections (which address it in place) nor the - correlated flat scatter (which addresses it first) expect it. +def _partition_out_selection( + sub_transform: IndexTransform, + out_indices: Any, + out_shape: tuple[int, ...], +) -> tuple[Any, ...]: + """Where a partition's values belong in an array of shape `out_shape`. - Returns - ------- - tuple of int, or None - `(actual, in_place)` — where NumPy put the axis, and the position it - would occupy if the advanced indices had not been separated. `None` when - `chunk_sel` gathers nothing. + A correlated sub-transform scatters through flat offsets into the row-major + result; unravelling them turns that into an index tuple for the result's own + shape, so callers never need a flat working buffer. """ - array_positions = [i for i, sel in enumerate(chunk_sel) if isinstance(sel, np.ndarray)] - if len(array_positions) == 0: - return None - in_place = sum( - 1 for i, sel in enumerate(chunk_sel) if i < array_positions[0] and isinstance(sel, slice) + _chunk_selection, out_selection, _drop_axes = sub_transform_to_selections( + sub_transform, out_indices ) - advanced = [i for i, sel in enumerate(chunk_sel) if not isinstance(sel, slice)] - first, last = advanced[0], advanced[-1] - separated = any(isinstance(sel, slice) for i, sel in enumerate(chunk_sel) if first < i < last) - if separated: - return 0, in_place - actual = sum(1 for i, sel in enumerate(chunk_sel) if i < first and isinstance(sel, slice)) - return actual, in_place - - -def _resolve_chunked( - array: Any, - transform: IndexTransform, - grids: Sequence[EdgeDimensionGrid], - dtype: np.dtype[Any], -) -> np.ndarray[Any, np.dtype[Any]]: - """Resolve a transform chunk by chunk, reading each touched chunk once.""" - out_shape = transform.domain.shape - # A correlated transform scatters through a single flat index, so the - # working buffer is 1-D during the read and reshaped afterwards. - needs_flat_buffer = _is_correlated(transform) - buffer_shape = (math.prod(out_shape),) if needs_flat_buffer else out_shape - out = np.empty(buffer_shape, dtype=dtype) - - if math.prod(out_shape) > 0: - for chunk_coords, sub_transform, out_indices in iter_chunk_transforms(transform, grids): - block_selection = tuple( - slice(grid.chunk_offset(c), grid.chunk_offset(c) + grid.chunk_size(c)) - for grid, c in zip(grids, chunk_coords, strict=True) - ) - block = np.asarray(array[block_selection]) - chunk_sel, out_sel, drop_axes = sub_transform_to_selections(sub_transform, out_indices) - selected = block[chunk_sel] - if len(drop_axes) > 0: - selected = selected.squeeze(axis=drop_axes) - # `chunk_sel` addresses the block in storage order, so NumPy may put - # the gathered axis somewhere the output selections do not expect: - # first for a correlated (flat) scatter, in place for an orthogonal - # one. Line the two up. - positions = _gathered_axis_positions(chunk_sel) - if positions is not None: - actual, in_place = positions - target = 0 if needs_flat_buffer else in_place - if actual != target: - selected = np.moveaxis(selected, actual, target) - out[out_sel] = selected - - if needs_flat_buffer: - out = out.reshape(out_shape) - return out - - -def _view_chunk_sizes( - grid: EdgeDimensionGrid, offset: int, stride: int, lo: int, hi: int -) -> tuple[int, ...]: - """Per-chunk extents of a strided slice, expressed in the view's own positions.""" - size = hi - lo - if size == 0: + if not _is_correlated(sub_transform): + return out_selection + if len(out_shape) == 0: return () - cuts = {0, size} - for chunk_ix in range(1, grid.num_chunks): - boundary = grid.chunk_offset(chunk_ix) - # First view position whose storage coordinate reaches this boundary. - position = -((offset + stride * lo - boundary) // stride) - if 0 < position < size: - cuts.add(int(position)) - return tuple(b - a for a, b in itertools.pairwise(sorted(cuts))) + flat = out_selection[0] + if isinstance(flat, slice): + flat = np.arange(flat.start, flat.stop, dtype=np.intp) + return tuple(np.unravel_index(np.asarray(flat, dtype=np.intp), out_shape)) + + +def _covers_whole_part(local_transform: IndexTransform, part_shape: tuple[int, ...]) -> bool: + """Whether a part-local transform addresses every cell of its part. + + Conservative: an `ArrayMap` could in principle enumerate a whole part, but + proving it costs more than the answer is worth, so a fancy-indexed axis + always reports incomplete. + """ + domain = local_transform.domain + for out_dim, m in enumerate(local_transform.output): + extent = part_shape[out_dim] + if isinstance(m, ConstantMap): + if extent != 1 or m.offset != 0: + return False + elif isinstance(m, DimensionMap): + if m.stride != 1: + return False + d = m.input_dimension + if m.offset + domain.inclusive_min[d] != 0: + return False + if m.offset + domain.exclusive_max[d] != extent: + return False + else: + return False + return True + + +@dataclass(frozen=True) +class Partition: + """One box of a `LazyArray`'s partitioning, as it falls through the view. + + Yielded by [`LazyArray.parts`][zarr_indexing.lazy_array.LazyArray.parts]. + The parts of a view tile it exactly and disjointly: assembling every + `array.result()` at its `out_selection` reproduces the view's `result()`, + and each part can be resolved independently and concurrently. + + Attributes + ---------- + base_coords + Which box of the base partitioning this is, one coordinate per dimension + of the wrapped array. + array + A `LazyArray` covering exactly the cells of the view that live in this + box. Resolving it reads the box once with basic slicing and applies the + selection to the block in memory. + out_selection + Where `array.result()` belongs in an array of the view's shape — a + NumPy index tuple, usable directly as `out[part.out_selection] = ...`. + is_complete + Whether the view covers the whole box. Useful to a writer deciding + between a blind overwrite and a read-modify-write. + """ + + base_coords: tuple[int, ...] + array: LazyArray + out_selection: tuple[Any, ...] + is_complete: bool + + +# --------------------------------------------------------------------------- # +# Tokenization +# --------------------------------------------------------------------------- # + + +def _wrapped_token(array: Any) -> Any: + """A deterministic token for the wrapped array. + + Determinism scope, in order of preference: the array's own + `__dask_tokenize__`; `dask.base.tokenize` when dask is importable (imported + lazily — this package never requires it); otherwise a local fallback that + digests the contents of a small array and, above `_TOKEN_DIGEST_LIMIT`, + falls back to a structural description that is stable across processes but + **not** sensitive to the array's contents. + """ + hook = getattr(array, "__dask_tokenize__", None) + if hook is not None: + try: + return hook() + except Exception: # pragma: no cover - a hook that refuses to run + pass + try: + # dask is an optional peer, never a dependency of this package, so it is + # imported here and its absence is ordinary. + from dask.base import tokenize # pyright: ignore[reportMissingImports] + except ImportError: + pass + else: + return tokenize(array) + + structural = ( + type(array).__qualname__, + tuple(int(s) for s in getattr(array, "shape", ())), + str(getattr(array, "dtype", "")), + ) + try: + contents = np.ascontiguousarray(array) + except Exception: + return structural + if contents.nbytes > _TOKEN_DIGEST_LIMIT: + return structural + return (*structural, hashlib.sha256(contents.tobytes()).hexdigest()) # --------------------------------------------------------------------------- # @@ -534,25 +615,23 @@ class LazyArray: never read at construction time. Indexing through `.lazy` composes an `IndexTransform` and returns another `LazyArray`; `result()` materializes. - Selections use the **positional NumPy dialect** — see the module docstring, - which also documents chunk discovery and how this differs from - `zarr.Array.lazy`. + Selections use the **positional NumPy dialect**, and reads are broken up + along a **partitioning** discovered from the wrapped array — see the module + docstring, which also covers how the dialect differs from `zarr.Array.lazy` + and why every non-indexing NumPy operation materializes the view. Parameters ---------- array The array to wrap. It must expose `shape`, `dtype`, and `__getitem__` - with basic (integer/slice) indexing. - chunks - Optionally declare the chunking of `array`, overriding discovery. Either - a uniform chunk shape (one integer per dimension) or dask-convention - per-axis chunk sizes (one sequence per dimension). + with basic (integer/slice) indexing. Its partitioning, if it advertises + one, is discovered here; use `with_parts` to choose a different one. Examples -------- >>> import numpy as np >>> source = np.arange(12).reshape(3, 4) - >>> view = LazyArray(source, chunks=(2, 2)).lazy[1:, ::2] + >>> view = LazyArray(source).with_parts((2, 2)).lazy[1:, ::2] >>> view.shape (2, 2) >>> view.result() @@ -560,35 +639,46 @@ class LazyArray: [ 8, 10]]) """ - __slots__ = ("_array", "_grids", "_transform") + __slots__ = ("_array", "_parts", "_transform", "_window") - def __init__( - self, - array: ArrayLike, - *, - chunks: Sequence[int] | Sequence[Sequence[int]] | None = None, - ) -> None: + def __init__(self, array: ArrayLike) -> None: shape = tuple(int(s) for s in array.shape) self._array = array + self._window: tuple[slice, ...] | None = None self._transform = IndexTransform.from_shape(shape) - self._grids = _discover_grids(array, shape, chunks) + self._parts = _discover_parts(array, shape) @classmethod - def _view( + def _derive( cls, array: ArrayLike, transform: IndexTransform, - grids: tuple[EdgeDimensionGrid, ...] | None, + parts: tuple[EdgeDimensionGrid, ...] | None, + window: tuple[slice, ...] | None, ) -> LazyArray: - """Build a view sharing `array` and `grids` but carrying a new transform.""" + """Build a wrapper sharing `array` but carrying a new transform or partitioning.""" view = cls.__new__(cls) view._array = array # Views re-zero their coordinate system: the positional dialect means a # view's first element is at position 0 whatever it was sliced from. view._transform = transform.translate_domain_to((0,) * transform.input_rank) - view._grids = grids + view._parts = parts + view._window = window return view + @property + def _base_shape(self) -> tuple[int, ...]: + """The shape of what this wrapper treats as its base array.""" + if self._window is None: + return tuple(int(s) for s in self._array.shape) + return tuple(s.stop - s.start for s in self._window) + + def _read_base(self) -> Any: + """The base array, materializing the window if this wrapper has one.""" + if self._window is None: + return self._array + return np.asarray(self._array[self._window]) + # -- array-API surface -------------------------------------------------- @property @@ -619,39 +709,97 @@ def dtype(self) -> Any: """The wrapped array's dtype; views never change it.""" return self._array.dtype - @property - def chunks(self) -> tuple[tuple[int, ...], ...] | None: - """Per-axis chunk sizes of this view, in dask's `Array.chunks` convention. - - The chunks of the wrapped array, cut down to what the view addresses. It - is `None` when the source declared no chunking, and also when the view - is not a plain basic-slice view: once an axis has been fancy-indexed, - its elements no longer arrive in contiguous runs and "the chunks of this - view" has no useful answer. Reordered or duplicated coordinates are - exactly the case where a chunk cannot be described by a size. + # -- partitioning ------------------------------------------------------- + + def with_parts(self, parts: Sequence[int] | Sequence[Sequence[int]] | None) -> LazyArray: + """Return the same view, read in different parts. + + The transform, the wrapped array, and therefore `result()` are all + unchanged; only the boxes the read is broken into differ. Cheap: nothing + is copied and nothing is read. + + Parameters + ---------- + parts + Either a uniform box shape (one integer per dimension of the base + array, with the trailing box clipped to the extent), dask-convention + per-axis sizes (one sequence of box extents per dimension, each + summing to the base extent), or `None` for a single part covering + the whole array — which makes `result()` lower the view in one shot + instead of assembling it block by block. + + Returns + ------- + LazyArray + The same view with a new partitioning. + + Raises + ------ + ValueError + If `parts` has the wrong length, mixes the two conventions, contains + a non-positive extent, or declares per-axis sizes that do not sum to + the base shape. Unlike discovery, this is our own API and validates + strictly. + + Examples + -------- + >>> import numpy as np + >>> view = LazyArray(np.arange(12).reshape(3, 4)) + >>> [part.base_coords for part in view.with_parts((2, 3)).parts()] + [(0, 0), (0, 1), (1, 0), (1, 1)] """ - if self._grids is None: - return None - transform = self._transform - per_input: dict[int, tuple[int, ...]] = {} - for out_dim, m in enumerate(transform.output): - if isinstance(m, ConstantMap): - continue - if not isinstance(m, DimensionMap) or m.stride <= 0: - return None - d = m.input_dimension - if d in per_input: - return None - per_input[d] = _view_chunk_sizes( - self._grids[out_dim], - m.offset, - m.stride, - transform.domain.inclusive_min[d], - transform.domain.exclusive_max[d], + grids = None if parts is None else dimension_grids_from_chunks(parts, self._base_shape) + return LazyArray._derive(self._array, self._transform, grids, self._window) + + def parts(self) -> Iterator[Partition]: + """Iterate the base partitioning, projected through this view. + + Yields one [`Partition`][zarr_indexing.lazy_array.Partition] per box the + view actually touches. The parts tile the view exactly and disjointly, + and each carries a `LazyArray` that can be resolved on its own — in + another thread, in another order, or not at all. + + A wrapper with no partitioning (see `with_parts`) yields a single part + covering the whole array. + + Yields + ------ + Partition + One per touched box, in the resolver's own order. + + Examples + -------- + >>> import numpy as np + >>> view = LazyArray(np.arange(12).reshape(3, 4)).with_parts((2, 2)) + >>> part = next(view.lazy[:, 1:].parts()) + >>> (part.base_coords, part.array.shape, part.is_complete) + ((0, 0), (2, 1), False) + """ + base_shape = self._base_shape + grids = self._parts if self._parts is not None else _whole_array_grids(base_shape) + out_shape = self.shape + rank = len(base_shape) + + for base_coords, local, out_indices in iter_chunk_transforms(self._transform, grids): + origin = tuple(grid.chunk_offset(c) for grid, c in zip(grids, base_coords, strict=True)) + extent = tuple(grid.chunk_size(c) for grid, c in zip(grids, base_coords, strict=True)) + if origin == (0,) * rank and extent == base_shape: + # The part is the whole base: lowering directly against the + # source beats materializing a block that is the source. + window = self._window + elif self._window is None: + window = tuple(slice(o, o + e) for o, e in zip(origin, extent, strict=True)) + else: + window = tuple( + slice(w.start + o, w.start + o + e) + for w, o, e in zip(self._window, origin, extent, strict=True) + ) + yield Partition( + base_coords=base_coords, + array=LazyArray._derive(self._array, local, None, window), + out_selection=_partition_out_selection(local, out_indices, out_shape), + is_complete=_covers_whole_part(local, extent), ) - if len(per_input) != transform.input_rank: - return None - return tuple(per_input[d] for d in range(transform.input_rank)) # -- indexing ----------------------------------------------------------- @@ -674,7 +822,7 @@ def _select(self, selection: Any, mode: SelectionMode) -> LazyArray: transform = transform.translate_domain_to((0,) * transform.input_rank) literal = normalize_positional_selection(selection, transform.domain, mode) composed = selection_to_transform(literal, transform, mode) - return LazyArray._view(self._array, composed, self._grids) + return LazyArray._derive(self._array, composed, self._parts, self._window) def __getitem__(self, selection: Any) -> Any: """Read a basic selection **eagerly**, like `numpy.ndarray.__getitem__`. @@ -688,18 +836,33 @@ def __getitem__(self, selection: Any) -> Any: def result(self) -> Any: """Materialize this view. + Assembles the view from its `parts()`, reading each box once. A wrapper + with a single whole-array part skips the buffer entirely and lowers + straight to array operations. + Returns ------- array - An array of shape `self.shape`. Chunked sources produce a NumPy - array (see the module docstring's device caveat); unchunked sources - produce whatever the wrapped array's namespace produces. A view with - a zero-rank domain returns a zero-dimensional array, not a scalar. + An array of shape `self.shape`, identical whatever partitioning is + in force. A partitioned view produces a NumPy array (see the module + docstring's device caveat); a single whole-array part produces + whatever the wrapped array's namespace produces. A view with a + zero-rank domain returns a zero-dimensional array, not a scalar. """ - transform = self._transform - if self._grids is None: - return _resolve_unchunked(self._array, transform) - return _resolve_chunked(self._array, transform, self._grids, np.dtype(self._array.dtype)) + if self._parts is None: + return _lower(self._read_base(), self._transform) + + out_shape = self.shape + out = np.empty(out_shape, dtype=np.dtype(self.dtype)) + if math.prod(out_shape) > 0: + for part in self.parts(): + value = np.asarray(part.array.result()) + if len(out_shape) == 0: + value = value.reshape(()) + out[part.out_selection] = value + return out + + # -- protocols ---------------------------------------------------------- def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: if copy is False: @@ -709,17 +872,61 @@ def __array__(self, dtype: Any = None, copy: bool | None = None) -> Any: ) return np.asarray(self.result(), dtype=dtype) + def __dask_tokenize__(self) -> Any: + """A deterministic token: the wrapped array, the view, and the parts. + + Two wrappers token-equal when they wrap the same data, address the same + cells, and read them in the same boxes. The view contributes its + canonical ndsel body, so transforms that differ only in representation + token alike. See `_wrapped_token` for the determinism scope of the + wrapped array's contribution; dask is imported lazily and is never a + requirement of this package. + """ + return ( + type(self).__qualname__, + _wrapped_token(self._array), + None if self._window is None else tuple((s.start, s.stop) for s in self._window), + json.dumps(transform_to_canonical(self._transform), sort_keys=True), + None if self._parts is None else tuple(grid.sizes for grid in self._parts), + ) + def __len__(self) -> int: if self.ndim == 0: raise TypeError("len() of unsized object") return self.shape[0] + def __iter__(self) -> Iterator[Any]: + """Iterate eagerly over the first axis, like a NumPy array. + + The rank check happens in `__iter__` itself rather than in the + generator, so `iter(view)` on a zero-rank view raises immediately as + NumPy's does, instead of waiting for the first `next`. + """ + if self.ndim == 0: + raise TypeError("iteration over a 0-d array") + return (self[position] for position in range(self.shape[0])) + + # NumPy's own conversions decide what a size-1 (or wrong-sized) view means, + # including which exception it raises, so these delegate rather than + # reimplement. Each materializes the view first. + def __bool__(self) -> bool: + return bool(self.result()) + + def __int__(self) -> int: + return int(self.result()) + + def __float__(self) -> float: + return float(self.result()) + + def __index__(self) -> int: + return operator.index(self.result()) + def __repr__(self) -> str: wrapped = type(self._array).__name__ - parts = [f"{wrapped} shape={self.shape} dtype={self.dtype}"] - if not _is_identity_transform(self._transform, tuple(int(s) for s in self._array.shape)): - parts.append(f"view={self._transform.selection_repr}") - return f"" + described = [f"{wrapped} shape={self.shape} dtype={self.dtype}"] + if not _is_identity_transform(self._transform, self._base_shape): + described.append(f"view={self._transform.selection_repr}") + return f"" class _LazyIndexer: diff --git a/packages/zarr-indexing/tests/test_lazy_array.py b/packages/zarr-indexing/tests/test_lazy_array.py index 25d7710d8e..646c4cf38c 100644 --- a/packages/zarr-indexing/tests/test_lazy_array.py +++ b/packages/zarr-indexing/tests/test_lazy_array.py @@ -10,6 +10,9 @@ from __future__ import annotations +import operator +import pickle +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any import numpy as np @@ -30,8 +33,8 @@ from collections.abc import Callable, Sequence SHAPE = (7, 5, 4) -CHUNK_SHAPE = (3, 2, 3) -DASK_CHUNKS = ((3, 3, 1), (2, 2, 1), (3, 1)) +PART_SHAPE = (3, 2, 3) +EXPLICIT_PARTS = ((3, 3, 1), (2, 2, 1), (3, 1)) def reference() -> np.ndarray[Any, np.dtype[np.int64]]: @@ -186,24 +189,38 @@ def test_scalar_on_a_fancy_axis_collapses_to_a_constant() -> None: # --------------------------------------------------------------------------- +def make_zarr_source() -> Any: + """A zarr array holding `reference()`, whose parts are auto-discovered.""" + zarr = pytest.importorskip("zarr") + array = zarr.create_array({}, shape=SHAPE, chunks=PART_SHAPE, dtype="int64") + array[:] = reference() + return array + + def make_source(flavour: str) -> LazyArray: - """Build a `LazyArray` over `reference()` with the requested chunk provenance.""" + """Build a `LazyArray` over `reference()` with the requested partitioning.""" data = reference() - if flavour == "numpy-unchunked": + if flavour == "numpy-whole": return LazyArray(data) - if flavour == "numpy-dask-chunks": - return LazyArray(data, chunks=DASK_CHUNKS) - if flavour == "numpy-uniform-chunks": - return LazyArray(data, chunks=CHUNK_SHAPE) + if flavour == "numpy-explicit-parts": + return LazyArray(data).with_parts(EXPLICIT_PARTS) + if flavour == "numpy-uniform-parts": + return LazyArray(data).with_parts(PART_SHAPE) if flavour == "zarr": - zarr = pytest.importorskip("zarr") - array = zarr.create_array({}, shape=SHAPE, chunks=CHUNK_SHAPE, dtype="int64") - array[:] = data - return LazyArray(array) + return LazyArray(make_zarr_source()) + if flavour == "zarr-misaligned": + # Parts that deliberately straddle the zarr array's own chunks. + return LazyArray(make_zarr_source()).with_parts((4, 3, 3)) raise AssertionError(f"unknown source flavour {flavour!r}") -FLAVOURS = ["numpy-unchunked", "numpy-dask-chunks", "numpy-uniform-chunks", "zarr"] +FLAVOURS = [ + "numpy-whole", + "numpy-explicit-parts", + "numpy-uniform-parts", + "zarr", + "zarr-misaligned", +] @pytest.fixture(params=FLAVOURS) @@ -507,15 +524,22 @@ def _random_chain(rng: np.random.Generator) -> list[tuple[str, tuple[Any, ...]]] return chain -@pytest.mark.parametrize( - "flavour", ["numpy-unchunked", "numpy-uniform-chunks", "numpy-dask-chunks"] -) +@pytest.mark.parametrize("flavour", FLAVOURS) def test_random_chains_match_numpy(flavour: str) -> None: - """A seeded sweep of composed chains: both resolvers must match NumPy.""" + """A seeded sweep of composed chains, under every partitioning. + + The partitioning invariant as a property: `result()` is identical whatever + boxes the read is broken into, including boxes deliberately misaligned with + the source's own. + """ rng = np.random.default_rng(20260730) source = make_source(flavour) + partitionings: list[Any] = [None, (2, 2, 2), (7, 5, 4), ((4, 3), (1, 3, 1), (3, 1))] + # Reads against a real store cost more per chain; the NumPy flavours carry + # the bulk of the sweep and exercise the identical code path. + n_chains = 120 if flavour.startswith("zarr") else 400 - for _ in range(400): + for _ in range(n_chains): chain = _random_chain(rng) expected = reference() for mode, selection in chain: @@ -529,6 +553,12 @@ def test_random_chains_match_numpy(flavour: str) -> None: np.testing.assert_array_equal( np.asarray(view.result()), np.asarray(expected), err_msg=f"{flavour}: {chain}" ) + for parts in partitionings: + np.testing.assert_array_equal( + np.asarray(view.with_parts(parts).result()), + np.asarray(expected), + err_msg=f"{flavour} parts={parts}: {chain}", + ) def test_eager_getitem_returns_data(source: LazyArray) -> None: @@ -560,28 +590,12 @@ def test_view_shape_comes_from_the_transform(source: LazyArray) -> None: assert "view=" in repr(view) -@pytest.mark.parametrize( - ("flavour", "expected"), - [ - ("numpy-unchunked", None), - ("numpy-dask-chunks", DASK_CHUNKS), - ("numpy-uniform-chunks", DASK_CHUNKS), - ("zarr", DASK_CHUNKS), - ], -) -def test_chunk_discovery(flavour: str, expected: tuple[tuple[int, ...], ...] | None) -> None: - """Declared, uniform, and auto-discovered chunkings all agree.""" - assert make_source(flavour).chunks == expected - - -def test_chunks_follow_a_basic_view() -> None: - """A basic-slice view reports its chunks in its own positions.""" - array = make_source("numpy-uniform-chunks") - # Storage chunk boundaries along axis 0 are 3 and 6; positions 0..4 of the - # view map to storage 1..5, so the boundary at 3 falls at position 2. - assert array.lazy[1:6, :, :].chunks == ((2, 3), (2, 2, 1), (3, 1)) - # Every other element of axis 0: storage 0, 2 | 4 | 6. - assert array.lazy[::2, :, :].chunks[0] == (2, 1, 1) +def test_the_wrapper_has_no_chunks_vocabulary() -> None: + """`chunks` is the *source's* word, never the wrapper's.""" + assert not hasattr(make_source("numpy-whole"), "chunks") + assert not hasattr(make_source("zarr"), "chunks") + with pytest.raises(TypeError): + LazyArray(reference(), chunks=PART_SHAPE) # type: ignore[call-arg] def test_scalar_is_basic_even_when_a_slice_separates_it_from_the_arrays() -> None: @@ -595,7 +609,7 @@ def test_scalar_is_basic_even_when_a_slice_separates_it_from_the_arrays() -> Non selection reads as `x[0][:, i]`. """ data = reference() - view = make_source("numpy-uniform-chunks").lazy.vindex[0, ..., np.array([3, 0])] + view = make_source("numpy-uniform-parts").lazy.vindex[0, ..., np.array([3, 0])] np.testing.assert_array_equal(np.asarray(view.result()), data[0][..., np.array([3, 0])]) assert view.shape == data[0][..., np.array([3, 0])].shape assert view.shape != data[0, ..., np.array([3, 0])].shape @@ -609,40 +623,240 @@ def test_zero_dimensional_result_is_an_array(source: LazyArray) -> None: assert result[()] == reference()[0, 1, 2] -def test_malformed_discovered_chunks_are_ignored() -> None: - """A foreign object's unusable `chunks` falls back to whole-array reads. +class ForeignArray: + """A minimal array-like whose advertised `chunks` we do not control.""" - Discovery parses external input, so an attribute that does not describe a - chunking of the shape means "no chunks I understand", not an error. - """ + def __init__(self, data: np.ndarray[Any, Any], chunks: Any) -> None: + self._data = data + self.chunks = chunks - class ForeignArray: - def __init__(self, data: np.ndarray[Any, Any], chunks: Any) -> None: - self._data = data - self.chunks = chunks + @property + def shape(self) -> tuple[int, ...]: + return self._data.shape - @property - def shape(self) -> tuple[int, ...]: - return self._data.shape + @property + def dtype(self) -> Any: + return self._data.dtype - @property - def dtype(self) -> Any: - return self._data.dtype + def __getitem__(self, key: Any) -> Any: + return self._data[key] - def __getitem__(self, key: Any) -> Any: - return self._data[key] +def test_malformed_discovered_parts_are_ignored() -> None: + """A foreign object's unusable `chunks` falls back to one whole-array part. + + Discovery parses external input, so an attribute that does not describe a + partitioning of the shape means "none I understand", not an error. + """ data = reference() for bogus in (((3, 3), (5,), (4,)), (3, 2), "nope", (3.5, 2, 2)): wrapped = LazyArray(ForeignArray(data, bogus)) - assert wrapped.chunks is None + assert len(list(wrapped.parts())) == 1 np.testing.assert_array_equal(np.asarray(wrapped.lazy[1:3].result()), data[1:3]) -def test_chunks_are_undefined_for_fancy_views() -> None: - """Reordered/duplicated coordinates have no chunk-size description.""" - array = make_source("numpy-uniform-chunks") - assert array.lazy.oindex[[4, 0, 0], :, :].chunks is None +def test_discovered_parts_come_from_the_source_vocabulary() -> None: + """`read_chunk_sizes` wins over `chunks`, and both are read as parts.""" + data = reference() + wrapped = LazyArray(ForeignArray(data, PART_SHAPE)) + assert [part.base_coords for part in wrapped.parts()][:3] == [(0, 0, 0), (0, 0, 1), (0, 1, 0)] + assert len(list(wrapped.parts())) == 3 * 3 * 2 + + +# --------------------------------------------------------------------------- +# Parts +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flavour", FLAVOURS) +@pytest.mark.parametrize( + "build", + [ + lambda a: a, + lambda a: a.lazy[1:6, :, 1:], + lambda a: a.lazy.oindex[[4, 0, 0], :, [3, 1]], + lambda a: a.lazy.vindex[..., np.array([1, 4, 0]), np.array([2, 0, 1])], + ], + ids=["identity", "basic", "oindex", "vindex"], +) +def test_parts_tile_the_view_exactly_and_disjointly( + flavour: str, build: Callable[[LazyArray], LazyArray] +) -> None: + """Assembling the parts reproduces `result()`; the placements cover with no overlap.""" + view = build(make_source(flavour)) + expected = np.asarray(view.result()) + + assembled = np.zeros(view.shape, dtype=view.dtype) + hits = np.zeros(view.shape, dtype=np.int64) + for part in view.parts(): + assembled[part.out_selection] = np.asarray(part.array.result()) + hits[part.out_selection] += 1 + + np.testing.assert_array_equal(assembled, expected) + np.testing.assert_array_equal(hits, np.ones(view.shape, dtype=np.int64)) + + +def test_parts_resolve_independently_and_concurrently() -> None: + """Each part's `array` is a standalone `LazyArray` with no shared mutable state.""" + view = make_source("zarr").lazy[1:7, :, 1:].with_parts((2, 2, 2)) + parts = list(view.parts()) + assert len(parts) > 1 + + with ThreadPoolExecutor(max_workers=4) as pool: + values = list(pool.map(lambda part: np.asarray(part.array.result()), parts)) + + assembled = np.zeros(view.shape, dtype=view.dtype) + for part, value in zip(parts, values, strict=True): + assembled[part.out_selection] = value + np.testing.assert_array_equal(assembled, np.asarray(view.result())) + + +def test_parts_report_completeness() -> None: + """`is_complete` distinguishes a fully-covered box from a partial one.""" + array = make_source("numpy-uniform-parts") + assert all(part.is_complete for part in array.parts()) + # Boxes along axis 2 are [0, 3) and [3, 4); dropping column 0 leaves the + # first partially covered and the second whole. + trimmed = {part.base_coords[2]: part.is_complete for part in array.lazy[:, :, 1:].parts()} + assert trimmed == {0: False, 1: True} + # A fancy axis is always reported incomplete. + assert not any(part.is_complete for part in array.lazy.oindex[[4, 0, 0], :, :].parts()) + + +def test_an_unpartitioned_wrapper_has_a_single_whole_array_part() -> None: + view = make_source("numpy-whole") + parts = list(view.parts()) + assert len(parts) == 1 + assert parts[0].base_coords == (0, 0, 0) + assert parts[0].array.shape == SHAPE + assert parts[0].is_complete + np.testing.assert_array_equal(np.asarray(parts[0].array.result()), reference()) + + +def test_with_parts_keeps_the_view_and_the_base() -> None: + view = make_source("zarr").lazy[1:6, ::2] + repartitioned = view.with_parts((2, 1, 4)) + assert repartitioned.shape == view.shape + assert repartitioned.array is view.array + assert repartitioned.transform == view.transform + # A different partitioning really is a different set of boxes. + assert [part.base_coords for part in repartitioned.parts()] != [ + part.base_coords for part in view.parts() + ] + np.testing.assert_array_equal(np.asarray(repartitioned.result()), np.asarray(view.result())) + + +def test_with_parts_none_forces_one_shot_resolution() -> None: + view = make_source("zarr").lazy.oindex[[4, 0, 0], :, :] + whole = view.with_parts(None) + assert len(list(whole.parts())) == 1 + np.testing.assert_array_equal(np.asarray(whole.result()), np.asarray(view.result())) + + +@pytest.mark.parametrize( + ("parts", "match"), + [ + ((3,), "one entry per dimension"), + ((3, (2, 2), 4), "not a mixture"), + (((3, 3), (2, 2, 1), (3, 1)), "sum to 6, but the array extent is 7"), + ((3, 0, 3), "chunk shape entries must be positive"), + ((3.5, 2, 2), "not a mixture"), + ((3.5, 2.5, 2.5), "for dimension 0 is neither"), + ], +) +def test_with_parts_validates_strictly(parts: Any, match: str) -> None: + """`with_parts` is our own API, so a malformed partitioning raises.""" + with pytest.raises(ValueError, match=match): + make_source("numpy-whole").with_parts(parts) + + +# --------------------------------------------------------------------------- +# Protocols +# --------------------------------------------------------------------------- + + +def test_dask_token_is_deterministic_and_discriminating() -> None: + """Same data, same view, same parts token alike; anything else differs.""" + data = reference() + base = LazyArray(data) + assert base.__dask_tokenize__() == LazyArray(reference()).__dask_tokenize__() + + tokens = { + "base": base.__dask_tokenize__(), + "view": base.lazy[1:3].__dask_tokenize__(), + "other view": base.lazy[2:4].__dask_tokenize__(), + "parts": base.with_parts((2, 2, 2)).__dask_tokenize__(), + "other data": LazyArray(data + 1).__dask_tokenize__(), + } + assert len({repr(token) for token in tokens.values()}) == len(tokens) + # Equivalent transforms reached different ways still token alike. + assert base.lazy[1:5].lazy[0:2].__dask_tokenize__() == base.lazy[1:3].__dask_tokenize__() + + +def test_iteration_yields_eager_slices(source: LazyArray) -> None: + rows = list(source.lazy[2:5]) + assert len(rows) == 3 + for row, expected in zip(rows, reference()[2:5], strict=True): + np.testing.assert_array_equal(np.asarray(row), expected) + + +def test_iteration_over_a_zero_dimensional_view_is_rejected() -> None: + with pytest.raises(TypeError, match="iteration over a 0-d array"): + iter(make_source("numpy-uniform-parts").lazy[0, 0, 0]) + + +def test_len_of_a_zero_dimensional_view_is_rejected() -> None: + with pytest.raises(TypeError, match="len\\(\\) of unsized object"): + len(make_source("numpy-uniform-parts").lazy[0, 0, 0]) + + +@pytest.mark.parametrize( + ("convert", "selection"), + [ + (bool, (1, 1, 1)), + (int, (1, 1, 1)), + (float, (1, 1, 1)), + (operator.index, (1, 1, 1)), + (bool, (1, 1, slice(0, 1))), + ], + ids=["bool-0d", "int-0d", "float-0d", "index-0d", "bool-size-1"], +) +def test_scalar_conversions_match_numpy(convert: Any, selection: Any) -> None: + """Size-1 conversions delegate to NumPy, values and all.""" + data = reference() + view = make_source("numpy-uniform-parts").lazy[selection] + assert convert(view) == convert(data[selection]) + + +@pytest.mark.parametrize( + ("convert", "selection", "error"), + [ + (bool, (slice(0, 2), 0, 0), ValueError), + (bool, (slice(0, 0), 0, 0), ValueError), + (int, (slice(0, 1), 0, 0), TypeError), + (float, (slice(0, 2), 0, 0), TypeError), + (operator.index, (slice(0, 1), 0, 0), TypeError), + ], + ids=["bool-many", "bool-empty", "int-1d", "float-many", "index-1d"], +) +def test_scalar_conversions_raise_what_numpy_raises( + convert: Any, selection: Any, error: type[Exception] +) -> None: + data = reference() + view = make_source("numpy-uniform-parts").lazy[selection] + with pytest.raises(error): + convert(view) + with pytest.raises(error): + convert(data[selection]) + + +def test_pickle_round_trip() -> None: + """A wrapper over a picklable base survives a round trip, view and parts intact.""" + view = LazyArray(reference()).with_parts((2, 2, 2)).lazy[1:6, ::2].lazy.oindex[[3, 0, 0], :, :] + restored = pickle.loads(pickle.dumps(view)) + assert restored.shape == view.shape + assert restored.__dask_tokenize__() == view.__dask_tokenize__() + np.testing.assert_array_equal(np.asarray(restored.result()), np.asarray(view.result())) # --------------------------------------------------------------------------- @@ -658,10 +872,11 @@ def test_dask_from_array_roundtrip() -> None: lazy = da.from_array(source) np.testing.assert_array_equal(lazy.compute(), reference()) - # The wrapper's own chunks can be handed straight back to dask. - chunked = da.from_array(source, chunks=source.chunks) - assert chunked.chunks == DASK_CHUNKS - np.testing.assert_array_equal(chunked[2:, ::2].compute(), reference()[2:, ::2]) + # dask chooses its own blocks; the wrapper reads each of them through its + # own parts, so the two partitionings need not agree. + blocked = da.from_array(source, chunks=(4, 3, 3)) + assert blocked.chunks == ((4, 3), (3, 2), (3, 1)) + np.testing.assert_array_equal(blocked[2:, ::2].compute(), reference()[2:, ::2]) # --------------------------------------------------------------------------- @@ -671,44 +886,39 @@ def test_dask_from_array_roundtrip() -> None: def test_boolean_scalar_is_rejected() -> None: with pytest.raises(IndexError, match="boolean scalars are not valid indices"): - make_source("numpy-uniform-chunks").lazy[True] + make_source("numpy-uniform-parts").lazy[True] def test_mask_shape_must_match_the_view() -> None: - array = make_source("numpy-uniform-chunks") + array = make_source("numpy-uniform-parts") with pytest.raises(IndexError, match="boolean index has shape"): array.lazy.vindex[np.ones((2, 2, 2), dtype=bool)] def test_scalar_index_out_of_bounds() -> None: with pytest.raises(IndexError, match="index 7 is out of bounds for axis 0 with size 7"): - make_source("numpy-uniform-chunks").lazy[7] + make_source("numpy-uniform-parts").lazy[7] def test_scalar_index_out_of_bounds_in_a_view() -> None: """Bounds are the *view's*, not the wrapped array's.""" - array = make_source("numpy-uniform-chunks").lazy[1:4] + array = make_source("numpy-uniform-parts").lazy[1:4] with pytest.raises(IndexError, match="index 3 is out of bounds for axis 0 with size 3"): array.lazy[3] def test_index_array_out_of_bounds() -> None: - array = make_source("numpy-uniform-chunks") + array = make_source("numpy-uniform-parts") with pytest.raises(IndexError, match="index 99 is out of bounds for axis 0 with size 7"): array.lazy.oindex[[0, 99], :, :] def test_too_many_indices() -> None: with pytest.raises(IndexError, match="too many indices"): - make_source("numpy-uniform-chunks").lazy[0, 0, 0, 0] - - -def test_chunks_that_do_not_sum_to_shape_are_rejected() -> None: - with pytest.raises(ValueError, match="sum to 6, but the array extent is 7"): - LazyArray(reference(), chunks=((3, 3), (2, 2, 1), (3, 1))) + make_source("numpy-uniform-parts").lazy[0, 0, 0, 0] def test_copy_false_conversion_is_rejected() -> None: - array = make_source("numpy-uniform-chunks") + array = make_source("numpy-uniform-parts") with pytest.raises(ValueError, match="cannot be converted to a NumPy array without a copy"): np.array(array, copy=False)