From 46c6033aa3d76727cffece70c5a815a4f7a907c6 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 2 Sep 2026 18:48:10 +0200 Subject: [PATCH 01/18] feat(zarr-indexing): factor chunk plans into a columnar GridPartition Restricting a transform to a chunk box distributes over output dimensions whenever each output map reads its own input axis, which is every basic and orthogonal selection. Chunk resolution therefore no longer intersects the whole transform with every candidate chunk; it resolves each axis once against its grid into a table (StridedSet / IndexedSet), sorts correlated (vindex) index arrays into chunks once into a JointSet, and derives each ChunkProjection as one row of each table. ChunkPlan.partition() and partition_transform() expose the factored form, so a consumer can read the tables directly instead of materializing an object graph per chunk. The projections a plan yields are unchanged; the general whole-transform walk remains for hand-built diagonals, which have no factored form. Along the way: _intersect_general reuses a precomputed _CorrelatedBlock and accepts survivor positions; checked_affine has identity and dtype-bounded fast paths; ArrayMap._with_affine shares frozen index arrays on translate; IndexDomain._unchecked / IndexTransform._unchecked skip validation for objects derived from an already-valid transform. Assisted-by: ClaudeCode:claude-fable-5-1 --- packages/zarr-indexing/README.md | 3 + .../src/zarr_indexing/__init__.py | 14 +- .../src/zarr_indexing/_affine.py | 34 +- .../src/zarr_indexing/chunk_resolution.py | 1081 ++++++++++++++++- .../zarr-indexing/src/zarr_indexing/domain.py | 17 + .../src/zarr_indexing/output_map.py | 15 + .../src/zarr_indexing/transform.py | 161 ++- .../tests/test_chunk_resolution.py | 158 ++- 8 files changed, 1359 insertions(+), 124 deletions(-) diff --git a/packages/zarr-indexing/README.md b/packages/zarr-indexing/README.md index 7e2cec10dd..110ebcb722 100644 --- a/packages/zarr-indexing/README.md +++ b/packages/zarr-indexing/README.md @@ -25,6 +25,9 @@ Key types: - `ChunkPlan` and `ChunkProjection` — lazily partition a selection over a caller-selected grid and pair each chunk-local transform with its placement in the request, without binding a storage backend or scheduler +- `GridPartition` — the plan's factored, columnar form: one table per axis + (`StridedSet`, `IndexedSet`) plus a `JointSet` for correlated index arrays, + from which projections are derived on demand - `ConstantMap`, `DimensionMap`, `ArrayMap` — the three ways a single output dimension can depend on the input - `compose` — chain two transforms into one diff --git a/packages/zarr-indexing/src/zarr_indexing/__init__.py b/packages/zarr-indexing/src/zarr_indexing/__init__.py index 9acfd28a21..fec26b871f 100644 --- a/packages/zarr-indexing/src/zarr_indexing/__init__.py +++ b/packages/zarr-indexing/src/zarr_indexing/__init__.py @@ -18,7 +18,9 @@ backends use an explicit `Reader` adapter. `plan_chunks` projects a transform through a caller-selected chunk grid without -coupling the result to a storage backend or scheduler. `selection_to_transform` +coupling the result to a storage backend or scheduler; `ChunkPlan.partition` +(or `partition_transform`) exposes the plan's factored, columnar form as a +`GridPartition` of `StridedSet`, `IndexedSet` and `JointSet` tables. `selection_to_transform` is also exported for consumers starting with a NumPy-style selection. The `DimensionGridLike` Protocol describes the narrow grid surface chunk resolution consumes without importing zarr. @@ -30,6 +32,11 @@ ChunkCoverage, ChunkPlan, ChunkProjection, + GridPartition, + IndexedSet, + JointSet, + StridedSet, + partition_transform, plan_chunks, ) from zarr_indexing.domain import IndexDomain @@ -89,10 +96,13 @@ "DimensionMap", "EdgeDimensionGrid", "FixedDimension", + "GridPartition", "IndexDomain", "IndexDomainJSON", "IndexTransform", "IndexTransformJSON", + "IndexedSet", + "JointSet", "LazyArray", "NdselError", "NumPyReader", @@ -101,6 +111,7 @@ "Partition", "ReadContext", "Reader", + "StridedSet", "UnitStepReader", "VaryingDimension", "VindexInvalidSelectionError", @@ -111,6 +122,7 @@ "numpy_reader", "output_index_map_from_json", "parse_ndsel", + "partition_transform", "plan_chunks", "unit_step_reader", ] diff --git a/packages/zarr-indexing/src/zarr_indexing/_affine.py b/packages/zarr-indexing/src/zarr_indexing/_affine.py index d1846ffe2d..eaa2157fc1 100644 --- a/packages/zarr-indexing/src/zarr_indexing/_affine.py +++ b/packages/zarr-indexing/src/zarr_indexing/_affine.py @@ -14,6 +14,20 @@ def _fits_intp(value: int) -> bool: return _INTP_INFO.min <= value <= _INTP_INFO.max +_DTYPE_LIMITS: dict[np.dtype[Any], tuple[bool, int]] = {} + + +def _dtype_limits(dtype: np.dtype[Any]) -> tuple[bool, int]: + """Whether every value of an integer ``dtype`` fits ``np.intp``, and the dtype's largest magnitude.""" + limits = _DTYPE_LIMITS.get(dtype) + if limits is None: + info = np.iinfo(dtype) + fits = int(info.min) >= _INTP_INFO.min and int(info.max) <= _INTP_INFO.max + limits = (fits, max(-int(info.min), int(info.max))) + _DTYPE_LIMITS[dtype] = limits + return limits + + @overload def checked_affine(offset: int, stride: int, coordinates: int) -> int: ... @@ -37,6 +51,14 @@ def checked_affine( NumPy performs fixed-width arithmetic. The common representable case then uses an ``np.intp`` fast path whose multiplication and addition were proven safe; cancellation cases use exact object arithmetic. + + Two shortcuts avoid scanning the array at all. The identity affine + (``offset == 0 and stride == 1``) of an array whose dtype fits ``np.intp`` + cannot overflow, so it is returned as-is (re-typed). And when the dtype's + own bounds already prove ``offset + stride * value`` representable for + every value the dtype can hold, the fixed-width path is taken directly. + Chunk resolution builds many small maps, where the scan's cost is the + call overhead rather than the elements. """ offset = int(offset) stride = int(stride) @@ -49,8 +71,16 @@ def checked_affine( if coordinates.size == 0: return np.empty(coordinates.shape, dtype=np.intp) - coordinate_min = int(np.min(coordinates)) - coordinate_max = int(np.max(coordinates)) + dtype_fits, dtype_bound = _dtype_limits(coordinates.dtype) + if offset == 0 and stride == 1: + if dtype_fits: + return np.asarray(coordinates, dtype=np.intp) + elif abs(offset) + abs(stride) * dtype_bound <= _INTP_INFO.max: + intp_coordinates = coordinates.astype(np.intp, copy=False) + return np.asarray(offset + stride * intp_coordinates, dtype=np.intp) + + coordinate_min = int(coordinates.min()) + coordinate_max = int(coordinates.max()) product_at_min = stride * coordinate_min product_at_max = stride * coordinate_max mapped_at_min = offset + product_at_min diff --git a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py index af148685a0..0df720fc54 100644 --- a/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py +++ b/packages/zarr-indexing/src/zarr_indexing/chunk_resolution.py @@ -29,20 +29,31 @@ source-independent: it identifies the chunk and expresses both sides of the gather without assuming NumPy selectors, a codec pipeline, or an execution scheduler. + +Behind the walk sits its factored form, the `GridPartition`: one table per +output dimension the transform reads independently (`StridedSet`, +`IndexedSet`) and one `JointSet` for correlated index arrays. A projection is +one row of each table combined, so the tables cost the *sum* of the touched +chunks per axis rather than their product, and a vectorized consumer can read +them directly without materializing a projection per chunk. """ from __future__ import annotations -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Literal +import itertools +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np from zarr_indexing._affine import checked_affine from zarr_indexing.domain import IndexDomain -from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap +from zarr_indexing.output_map import ArrayMap, ConstantMap, DimensionMap, OutputIndexMap from zarr_indexing.transform import ( IndexTransform, + _intersect_dimension_map, # pyright: ignore[reportPrivateUsage] + _intersect_general, # pyright: ignore[reportPrivateUsage] + _prepare_correlated, # pyright: ignore[reportPrivateUsage] ) if TYPE_CHECKING: @@ -130,7 +141,8 @@ class ChunkPlan: """A reusable, lazy partition of an index transform over a chunk grid. Construct plans with `plan_chunks`; iterating either the plan or - `projections()` performs a fresh chunk walk. + `projections()` performs a fresh chunk walk. `partition()` exposes the + factored form the walk is derived from (see `GridPartition`). Examples -------- @@ -153,9 +165,28 @@ class ChunkPlan: dimension_grids: tuple[DimensionGridLike, ...] """One grid per storage dimension, defining the chunk layout the plan walks.""" + # Memoized factored form; derived state, excluded from identity (see + # `IndexDomain._shape` for the same pattern). + _partition: GridPartition | None = field(default=None, init=False, repr=False, compare=False) + + def partition(self) -> GridPartition: + """The plan in factored, columnar form: one table per axis plus a joint table. + + Built once per plan and memoized. Raises `ValueError` for the one + structure that has no factored form, a transform whose output maps + share an input axis (a diagonal); `projections()` still walks those. + """ + cached = self._partition + if cached is None: + cached = partition_transform(self.transform, self.dimension_grids) + object.__setattr__(self, "_partition", cached) + return cached + def projections(self) -> Iterator[ChunkProjection]: """Return a fresh iterator over the chunks touched by this plan.""" - return _iter_chunk_projections(self.transform, self.dimension_grids) + if _partitionable(self.transform): + return iter(self.partition()) + return _iter_general_projections(self.transform, self.dimension_grids) def __iter__(self) -> Iterator[ChunkProjection]: """Equivalent to `projections()`: each iteration performs a fresh chunk walk.""" @@ -253,6 +284,55 @@ def _iter_sorted_1d_array_map( start = stop +def _group_points_by_chunk( + chunk_ids: Sequence[np.ndarray[Any, np.dtype[np.intp]]], +) -> list[tuple[tuple[int, ...], np.ndarray[Any, np.dtype[np.intp]]]]: + """Partition the points of a correlated block by the chunk each lands in. + + ``chunk_ids`` holds, per correlated output dimension, the chunk index of + every block point. Returns one ``(chunk_coords, positions)`` pair per + touched chunk, in lexicographic chunk order, with ``positions`` (flat + indices into the block) ascending — the same order `np.nonzero` would + give. Costs ``O(points log points)`` regardless of grid size, so a + selection scattered over many chunks is not rescanned once per chunk. + """ + n = int(chunk_ids[0].size) + if n == 0: + return [] + keys: np.ndarray[Any, np.dtype[np.intp]] | None + if len(chunk_ids) == 1: + keys = np.asarray(chunk_ids[0], dtype=np.intp) + else: + # Mixed-radix key with the first dimension most significant, so sorting + # the keys sorts the chunk coordinates lexicographically. + keys = np.zeros(n, dtype=np.intp) + multiplier = 1 + for ids in reversed(chunk_ids): + radix = int(ids.max()) + 1 + if multiplier * radix >= 2**62: + keys = None + break + keys += np.asarray(ids, dtype=np.intp) * multiplier + multiplier *= radix + if keys is None: + stacked = np.stack( + [np.asarray(ids, dtype=np.intp).ravel() for ids in chunk_ids], axis=1 + ) + _, inverse = np.unique(stacked, axis=0, return_inverse=True) + keys = np.asarray(inverse, dtype=np.intp).reshape(-1) + order = np.argsort(keys, kind="stable") + sorted_keys = keys[order] + boundaries = np.flatnonzero(sorted_keys[1:] != sorted_keys[:-1]) + 1 + starts = [0, *boundaries.tolist()] + ends = [*starts[1:], n] + groups: list[tuple[tuple[int, ...], np.ndarray[Any, np.dtype[np.intp]]]] = [] + for start, stop in zip(starts, ends, strict=True): + first = order[start] + coords = tuple(int(ids[first]) for ids in chunk_ids) + groups.append((coords, order[start:stop])) + return groups + + def _iter_chunk_transform_results( transform: IndexTransform, dim_grids: Sequence[DimensionGridLike], @@ -300,14 +380,17 @@ def _iter_chunk_transform_results( # resolution scale with grid size instead of with the number of selected # coordinates. # - Correlated (vindex) `ArrayMap` dims share one *joint* slot holding the - # distinct chunk-coordinate tuples the points actually land in. The - # cartesian product of their per-dimension distinct sets would include + # distinct chunk-coordinate tuples the points actually land in, found by + # sorting the points once (`_group_points_by_chunk`). The cartesian + # product of their per-dimension distinct sets would include # combinations no point touches — quadratic in the number of selected # points for a diagonal selection — while the joint distinct set is - # bounded by the point count (see zarr-python gh-4174). + # bounded by the point count (see zarr-python gh-4174). The same sort + # hands each chunk its surviving points, so the per-chunk intersection + # never rescans the whole selection. structure = transform.index_array_structure + block = _prepare_correlated(transform) if structure == "general" else None correlated_dims: list[int] = [] - correlated_chunk_ids: list[np.ndarray[Any, np.dtype[np.intp]]] = [] slot_dims: list[tuple[int, ...]] = [] slot_candidates: list[Sequence[tuple[int, ...]]] = [] for out_dim, m in enumerate(transform.output): @@ -324,29 +407,10 @@ def _iter_chunk_transform_results( dim_hi = transform.domain.exclusive_max[d] if dim_lo >= dim_hi: return # empty domain - first_storage = checked_affine(m.offset, m.stride, dim_lo) - if m.stride > 0: - s_min = first_storage - s_max = checked_affine(m.offset, m.stride, dim_hi - 1) - elif m.stride < 0: - s_min = checked_affine(m.offset, m.stride, dim_hi - 1) - s_max = first_storage - else: - s_min = s_max = first_storage - first = dg.index_to_chunk(s_min) - last = dg.index_to_chunk(s_max) slot_dims.append((out_dim,)) - point_count = dim_hi - dim_lo - chunk_count = last - first + 1 - if point_count < chunk_count: - steps = np.arange(point_count, dtype=np.intp) - storage = checked_affine(first_storage, m.stride, steps) - chunk_ids = dg.indices_to_chunks(storage) - slot_candidates.append([(int(c),) for c in np.unique(chunk_ids)]) - else: - slot_candidates.append([(c,) for c in range(first, last + 1)]) - else: - # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap). + slot_candidates.append(_dimension_map_candidates(m, dim_lo, dim_hi, dg)) + elif block is None: + # m: ArrayMap with orthogonal structure. # Storage coordinates were already computed for a correlated 1-D map. storage = ( array_map_1d[1] @@ -356,33 +420,27 @@ def _iter_chunk_transform_results( if storage.size == 0: # Empty fancy selection: no coordinates, so no chunks are touched. return - # Keep the index-array shape: correlated maps broadcast against each - # other below, and raveling first would lose the singleton axes. chunk_ids = dg.indices_to_chunks(storage) - if structure == "orthogonal": - slot_dims.append((out_dim,)) - slot_candidates.append([(int(c),) for c in np.unique(chunk_ids)]) - else: - # Every index array of a general transform joins one joint - # slot: their chunk ids broadcast over the shared block, so the - # distinct tuples enumerate only combinations some point - # actually touches. - correlated_dims.append(out_dim) - correlated_chunk_ids.append(chunk_ids) - - if len(correlated_dims) == 1: - slot_dims.append((correlated_dims[0],)) - slot_candidates.append([(int(c),) for c in np.unique(correlated_chunk_ids[0])]) - elif len(correlated_dims) >= 2: - # Group the points jointly: distinct rows of the per-point chunk - # coordinates, O(points log points) regardless of grid size. - broadcast = np.broadcast_arrays(*correlated_chunk_ids) - stacked = np.stack([b.ravel() for b in broadcast], axis=1) - joint = np.unique(stacked, axis=0) + slot_dims.append((out_dim,)) + slot_candidates.append([(int(c),) for c in np.unique(chunk_ids)]) + else: + correlated_dims.append(out_dim) + + joint_slot: int | None = None + joint_positions: dict[tuple[int, ...], np.ndarray[Any, np.dtype[np.intp]]] = {} + if correlated_dims: + assert block is not None + chunk_ids_per_dim = [ + dim_grids[out_dim].indices_to_chunks(block.flat_storage[out_dim]) + for out_dim in correlated_dims + ] + groups = _group_points_by_chunk(chunk_ids_per_dim) + if not groups: + return + joint_slot = len(slot_dims) slot_dims.append(tuple(correlated_dims)) - slot_candidates.append([tuple(int(c) for c in row) for row in joint]) - - import itertools + slot_candidates.append([coords for coords, _ in groups]) + joint_positions = dict(groups) output_rank = len(transform.output) for combo in itertools.product(*slot_candidates): @@ -410,7 +468,13 @@ def _iter_chunk_transform_results( ) # Intersect transform with chunk domain - result = transform.intersect(chunk_domain) + result: tuple[IndexTransform, _OutIndices] | None + if block is not None and joint_slot is not None: + result = _intersect_general( + transform, chunk_domain, block=block, positions=joint_positions[combo[joint_slot]] + ) + else: + result = transform.intersect(chunk_domain) if result is None: continue @@ -422,6 +486,31 @@ def _iter_chunk_transform_results( yield (chunk_coords, local, surviving) +def _dimension_map_candidates( + m: DimensionMap, dim_lo: int, dim_hi: int, dg: DimensionGridLike +) -> Sequence[tuple[int, ...]]: + """The chunks a nonempty `DimensionMap` over input `[dim_lo, dim_hi)` can touch, ascending.""" + first_storage = checked_affine(m.offset, m.stride, dim_lo) + if m.stride > 0: + s_min = first_storage + s_max = checked_affine(m.offset, m.stride, dim_hi - 1) + elif m.stride < 0: + s_min = checked_affine(m.offset, m.stride, dim_hi - 1) + s_max = first_storage + else: + s_min = s_max = first_storage + first = dg.index_to_chunk(s_min) + last = dg.index_to_chunk(s_max) + point_count = dim_hi - dim_lo + chunk_count = last - first + 1 + if point_count < chunk_count: + steps = np.arange(point_count, dtype=np.intp) + storage = checked_affine(first_storage, m.stride, steps) + chunk_ids = dg.indices_to_chunks(storage) + return [(int(c),) for c in np.unique(chunk_ids)] + return [(c,) for c in range(first, last + 1)] + + def _covers_whole_chunk(transform: IndexTransform, chunk_shape: tuple[int, ...]) -> bool: """Whether an affine chunk-local transform bijects onto every chunk cell.""" domain = transform.domain @@ -556,11 +645,18 @@ def _cell_transform( return _orthogonal_cell_transform(original, restricted, survivors) -def _iter_chunk_projections( +def _iter_general_projections( transform: IndexTransform, dim_grids: Sequence[DimensionGridLike], ) -> Iterator[ChunkProjection]: - """Convert private intersection results into public paired projections.""" + """Convert private intersection results into public paired projections. + + The general walk: intersect the whole transform with every candidate + chunk. `GridPartition` covers every transform a selection can produce; + this remains for hand-built diagonals, which have no factored form. + """ + if any(size == 0 for size in transform.domain.shape): + return for chunk_coords, chunk_transform, survivors in _iter_chunk_transform_results( transform, dim_grids ): @@ -593,3 +689,870 @@ def _iter_chunk_projections( cell_transform=cell_transform, coverage=coverage, ) + + +# --------------------------------------------------------------------------- # +# Grid partition: the factored form of a plan +# --------------------------------------------------------------------------- # +# +# Restricting a transform to a chunk box distributes over output dimensions +# whenever each output map reads its own input axis: the domain is a product +# of intervals, each map depends on one of them, and the box is a product. The +# chunks such a transform touches are then the cartesian product of the chunks +# each axis touches, and the restriction to any one of them is the product of +# one-dimensional restrictions. `GridPartition` stores those one-dimensional +# restrictions as tables (`StridedSet`, `IndexedSet`); a `ChunkProjection` is +# one row of each table, combined. Correlated index arrays (`vindex`) read the +# same input axes, so they do not distribute; they are sorted into chunks once +# and kept in a single `JointSet`. This is the structure TensorStore's +# `IndexTransformGridPartition` uses (strided sets and index-array sets). + + +def _factorizable(transform: IndexTransform) -> bool: + """True when every output map binds its own input axis and none is correlated.""" + if transform.index_array_structure == "general": + return False + seen: set[int] = set() + for m in transform.output: + if isinstance(m, DimensionMap): + axis = m.input_dimension + elif isinstance(m, ArrayMap): + dependent = m.dependent_axis + if dependent is None: + return False + axis = dependent + else: + continue + if axis in seen: + return False + seen.add(axis) + return True + + +def _correlated_partitionable(transform: IndexTransform) -> bool: + """True when a general transform's index arrays vary only over its broadcast axes.""" + if transform.index_array_structure != "general": + return False + bound = {m.input_dimension for m in transform.output if isinstance(m, DimensionMap)} + return all( + not any(axis in bound for axis in m.dependency_axes) + for m in transform.output + if isinstance(m, ArrayMap) + ) + + +def _partitionable(transform: IndexTransform) -> bool: + return _factorizable(transform) or _correlated_partitionable(transform) + + +def _int_column(values: Sequence[int]) -> np.ndarray[Any, np.dtype[np.intp]]: + """An ``intp`` column, or an object column of Python ints when a value does not fit. + + Request coordinates are unbounded Python ints in the transform algebra; a + domain near the ``intp`` limit (see `IndexDomain`) still partitions, its + table just carries exact ints. + """ + try: + return np.array(values, dtype=np.intp) + except OverflowError: + return cast("np.ndarray[Any, np.dtype[np.intp]]", np.array(values, dtype=object)) + + +def _chunk_bounds( + dg: DimensionGridLike, chunks: np.ndarray[Any, np.dtype[np.intp]] +) -> tuple[np.ndarray[Any, np.dtype[np.intp]], np.ndarray[Any, np.dtype[np.intp]]]: + """Storage origin and data extent of each chunk, one grid call per touched chunk.""" + n = int(chunks.size) + starts = np.fromiter((dg.chunk_offset(int(c)) for c in chunks), dtype=np.intp, count=n) + extents = np.fromiter((_data_size(dg, int(c)) for c in chunks), dtype=np.intp, count=n) + return starts, extents + + +@dataclass(frozen=True, slots=True) +class StridedSet: + """One output dimension read through a `ConstantMap` or `DimensionMap`, one row per chunk. + + Row ``i`` is the map restricted to chunk ``chunk[i]`` and re-based to + chunk-local, zero-origin coordinates: the chunk-local map is + `DimensionMap(input_dimension, offset=local_start[i], stride=stride)` over + ``[0, extent[i])`` (a `ConstantMap(local_start[i])` for a constant), and + its cells are request coordinates ``[origin[i], origin[i] + extent[i])`` + along the input axis. + + Examples + -------- + >>> from zarr_indexing import IndexTransform, plan_chunks + >>> from zarr_indexing.grid import dimension_grids_from_chunks + >>> grids = dimension_grids_from_chunks((4,), shape=(10,)) + >>> (axis,) = plan_chunks(IndexTransform.from_shape((10,))[1:9:2], grids).partition().sets + >>> axis.chunk.tolist(), axis.local_start.tolist(), axis.extent.tolist(), axis.origin.tolist() + ([0, 1], [1, 1], [2, 2], [0, 2]) + """ + + output_dimension: int + """The storage axis this table describes.""" + + input_dimension: int | None + """The request axis the map reads, or `None` for a constant.""" + + stride: int + """Storage step per request cell; ``0`` for a constant.""" + + chunk: np.ndarray[Any, np.dtype[np.intp]] + """Chunk index along the axis, one per row, ascending.""" + + chunk_start: np.ndarray[Any, np.dtype[np.intp]] + """Storage origin of each chunk.""" + + chunk_extent: np.ndarray[Any, np.dtype[np.intp]] + """Data extent of each chunk (clipped at the array boundary).""" + + local_start: np.ndarray[Any, np.dtype[np.intp]] + """Chunk-local storage coordinate of the row's first cell.""" + + extent: np.ndarray[Any, np.dtype[np.intp]] + """Cells the row selects along the request axis (``1`` for a constant).""" + + origin: np.ndarray[Any, np.dtype[np.intp]] + """Request coordinate of the row's first cell (``0`` for a constant).""" + + full: np.ndarray[Any, np.dtype[np.bool_]] + """Whether the row covers its chunk's data extent exactly once, in order.""" + + def __len__(self) -> int: + return int(self.chunk.size) + + def chunk_map(self, row: int, input_dimension: int | None = None) -> OutputIndexMap: + """The chunk-local output map of a row, optionally with a renumbered input axis.""" + offset = int(self.local_start[row]) + if self.input_dimension is None: + return ConstantMap(offset=offset) + axis = self.input_dimension if input_dimension is None else input_dimension + return DimensionMap(input_dimension=axis, offset=offset, stride=self.stride) + + def cell_map(self, row: int) -> DimensionMap | None: + """The map from the row's zero-origin cells back to request coordinates.""" + if self.input_dimension is None: + return None + return DimensionMap(input_dimension=self.input_dimension, offset=int(self.origin[row])) + + +@dataclass(frozen=True, slots=True) +class IndexedSet: + """One output dimension read through an orthogonal `ArrayMap`, one row per chunk. + + The map's coordinates are grouped by chunk in CSR form: row ``i`` owns + ``index[pointer[i]:pointer[i + 1]]`` (the index-array values, in request + order) and ``positions[pointer[i]:pointer[i + 1]]`` (their positions along + the request axis, ascending). `local` gives the same values as chunk-local + storage coordinates. + + Examples + -------- + >>> import numpy as np + >>> from zarr_indexing import IndexTransform, plan_chunks + >>> from zarr_indexing.grid import dimension_grids_from_chunks + >>> grids = dimension_grids_from_chunks((4,), shape=(10,)) + >>> transform = IndexTransform.from_shape((10,)).oindex[np.array([9, 1, 2, 8])] + >>> (axis,) = plan_chunks(transform, grids).partition().sets + >>> axis.chunk.tolist(), axis.pointer.tolist() + ([0, 2], [0, 2, 4]) + >>> axis.local.tolist(), axis.positions.tolist() + ([1, 2, 1, 0], [1, 2, 0, 3]) + """ + + output_dimension: int + """The storage axis this table describes.""" + + input_dimension: int + """The request axis the index array varies over.""" + + offset: int + """The map's affine offset: storage is ``offset + stride * index``.""" + + stride: int + """The map's affine stride.""" + + chunk: np.ndarray[Any, np.dtype[np.intp]] + """Chunk index along the axis, one per row, ascending.""" + + chunk_start: np.ndarray[Any, np.dtype[np.intp]] + """Storage origin of each chunk.""" + + chunk_extent: np.ndarray[Any, np.dtype[np.intp]] + """Data extent of each chunk.""" + + pointer: np.ndarray[Any, np.dtype[np.intp]] + """CSR row pointer: row ``i`` owns entries ``pointer[i]`` to ``pointer[i + 1]``.""" + + index: np.ndarray[Any, np.dtype[np.intp]] + """Index-array values grouped by chunk.""" + + positions: np.ndarray[Any, np.dtype[np.intp]] + """Positions along the request axis, grouped by chunk, ascending within a row.""" + + def __len__(self) -> int: + return int(self.chunk.size) + + @property + def counts(self) -> np.ndarray[Any, np.dtype[np.intp]]: + """Entries per row.""" + return np.diff(self.pointer) + + @property + def local(self) -> np.ndarray[Any, np.dtype[np.intp]]: + """Chunk-local storage coordinate of every entry, grouped like `index`.""" + storage = checked_affine(self.offset, self.stride, self.index) + return storage - np.repeat(self.chunk_start, self.counts) + + def run(self, row: int) -> slice: + """The slice of `index` / `positions` a row owns.""" + return slice(int(self.pointer[row]), int(self.pointer[row + 1])) + + +@dataclass(frozen=True, slots=True) +class JointSet: + """The correlated index arrays of a transform, grouped by the chunk each point lands in. + + Correlated (`vindex`) arrays read the same input axes, so a chunk + constrains all of them at once; they are sorted into chunks together. + Row ``i`` is one touched chunk, `chunk[i]` its coordinates on the + `output_dimensions`, and CSR range ``pointer[i]:pointer[i + 1]`` its + points: `index` holds their index-array values per output dimension, + `positions` their flat positions in the request's broadcast block, and + `block_coordinates` those positions unravelled over the block. + + Examples + -------- + >>> import numpy as np + >>> from zarr_indexing import IndexTransform, plan_chunks + >>> from zarr_indexing.grid import dimension_grids_from_chunks + >>> grids = dimension_grids_from_chunks((3, 4), shape=(7, 9)) + >>> transform = IndexTransform.from_shape((7, 9)).vindex[ + ... np.array([0, 6, 6, 1]), np.array([8, 0, 1, 8]) + ... ] + >>> joint = plan_chunks(transform, grids).partition().joint + >>> joint.chunk.tolist(), joint.pointer.tolist(), joint.positions.tolist() + ([[0, 2], [2, 0]], [0, 2, 4], [0, 3, 1, 2]) + """ + + output_dimensions: tuple[int, ...] + """The storage axes read by correlated index arrays.""" + + offsets: tuple[int, ...] + """Affine offset of each array's map, aligned with `output_dimensions`.""" + + strides: tuple[int, ...] + """Affine stride of each array's map.""" + + broadcast_axes: tuple[int, ...] + """The request axes the arrays broadcast over.""" + + broadcast_shape: tuple[int, ...] + """The extent of those axes.""" + + chunk: np.ndarray[Any, np.dtype[np.intp]] + """Chunk coordinates on `output_dimensions`, shape ``(rows, k)``, lexicographic.""" + + chunk_start: np.ndarray[Any, np.dtype[np.intp]] + """Storage origin of each chunk on `output_dimensions`, shape ``(rows, k)``.""" + + chunk_extent: np.ndarray[Any, np.dtype[np.intp]] + """Data extent of each chunk on `output_dimensions`, shape ``(rows, k)``.""" + + pointer: np.ndarray[Any, np.dtype[np.intp]] + """CSR row pointer into `index`, `positions` and `block_coordinates`.""" + + index: np.ndarray[Any, np.dtype[np.intp]] + """Index-array values per point and output dimension, shape ``(points, k)``.""" + + positions: np.ndarray[Any, np.dtype[np.intp]] + """Flat block position of each point, ascending within a row.""" + + block_coordinates: np.ndarray[Any, np.dtype[np.intp]] + """`positions` unravelled over `broadcast_shape`, shape ``(points, len(broadcast_axes))``.""" + + def __len__(self) -> int: + return int(self.chunk.shape[0]) + + @property + def counts(self) -> np.ndarray[Any, np.dtype[np.intp]]: + """Points per row.""" + return np.diff(self.pointer) + + @property + def local(self) -> np.ndarray[Any, np.dtype[np.intp]]: + """Chunk-local storage coordinates of every point, shape ``(points, k)``.""" + storage = np.stack( + [ + checked_affine(offset, stride, self.index[:, column]) + for column, (offset, stride) in enumerate( + zip(self.offsets, self.strides, strict=True) + ) + ], + axis=1, + ) + return storage - np.repeat(self.chunk_start, self.counts, axis=0) + + def run(self, row: int) -> slice: + """The slice of the point arrays a row owns.""" + return slice(int(self.pointer[row]), int(self.pointer[row + 1])) + + +def _strided_set( + transform: IndexTransform, out_dim: int, m: ConstantMap | DimensionMap, dg: DimensionGridLike +) -> StridedSet: + domain = transform.domain + if isinstance(m, ConstantMap): + c = dg.index_to_chunk(checked_affine(m.offset, 0, 0)) + chunks = np.array([c], dtype=np.intp) + starts, extents = _chunk_bounds(dg, chunks) + local = m.offset - starts + return StridedSet( + output_dimension=out_dim, + input_dimension=None, + stride=0, + chunk=chunks, + chunk_start=starts, + chunk_extent=extents, + local_start=local, + extent=np.ones(1, dtype=np.intp), + origin=np.zeros(1, dtype=np.intp), + full=(extents == 1) & (local == 0), + ) + k = m.input_dimension + lo = domain.inclusive_min[k] + hi = domain.exclusive_max[k] + stride = m.stride + unit = abs(stride) == 1 + rows: list[tuple[int, int, int, int, int, int, bool]] = [] + # Exact Python-int arithmetic per touched chunk: the values can exceed + # np.intp before cancellation (a large-origin domain), and the number of + # touched chunks along one axis is a sum, not a product. + for (c,) in _dimension_map_candidates(m, lo, hi, dg): + c_start = dg.chunk_offset(c) + c_extent = _data_size(dg, c) + narrowed = _intersect_dimension_map(m, lo, hi, c_start, c_start + c_extent) + if narrowed is None: + continue + nlo, nhi = narrowed + extent = nhi - nlo + # Chunk-local, then re-based to a zero-origin input axis. + local_start = m.offset - c_start + stride * nlo + if unit: + last = local_start + stride * (extent - 1) + full = min(local_start, last) == 0 and max(local_start, last) == c_extent - 1 + else: + full = False + rows.append((c, c_start, c_extent, local_start, extent, nlo, full)) + columns = list(zip(*rows, strict=True)) if rows else [()] * 7 + return StridedSet( + output_dimension=out_dim, + input_dimension=k, + stride=stride, + chunk=_int_column(columns[0]), + chunk_start=_int_column(columns[1]), + chunk_extent=_int_column(columns[2]), + local_start=_int_column(columns[3]), + extent=_int_column(columns[4]), + origin=_int_column(columns[5]), + full=np.array(columns[6], dtype=np.bool_), + ) + + +def _indexed_set(out_dim: int, m: ArrayMap, dg: DimensionGridLike) -> IndexedSet: + dependent = m.dependent_axis + assert dependent is not None + flat = m.index_array.reshape(-1) + n = int(flat.size) + storage = checked_affine(m.offset, m.stride, flat) + # Probe the extreme coordinates with the scalar lookup first: a grid's + # scalar error names the offending coordinate, where the vectorized one + # reports a range. + dg.index_to_chunk(int(storage.min())) + dg.index_to_chunk(int(storage.max())) + chunk_ids = dg.indices_to_chunks(storage) + # Already-sorted coordinates (the common case) need no sort. + if bool((chunk_ids[1:] >= chunk_ids[:-1]).all()): + positions = np.arange(n, dtype=np.intp) + sorted_ids = chunk_ids + index = flat + else: + positions = np.argsort(chunk_ids, kind="stable") + sorted_ids = chunk_ids[positions] + index = flat[positions] + boundaries = np.flatnonzero(sorted_ids[1:] != sorted_ids[:-1]) + 1 + pointer = np.concatenate([[0], boundaries, [n]]).astype(np.intp) + chunks = np.asarray(sorted_ids[pointer[:-1]], dtype=np.intp) + starts, extents = _chunk_bounds(dg, chunks) + return IndexedSet( + output_dimension=out_dim, + input_dimension=dependent, + offset=m.offset, + stride=m.stride, + chunk=chunks, + chunk_start=starts, + chunk_extent=extents, + pointer=pointer, + index=np.asarray(index, dtype=np.intp), + positions=np.asarray(positions, dtype=np.intp), + ) + + +def _chunk_keys( + chunk_ids: Sequence[np.ndarray[Any, np.dtype[np.intp]]], +) -> np.ndarray[Any, np.dtype[np.intp]]: + """One sortable key per point whose order is the lexicographic chunk order.""" + n = int(chunk_ids[0].size) + if len(chunk_ids) == 1: + return np.asarray(chunk_ids[0], dtype=np.intp) + # Mixed-radix key with the first dimension most significant. + keys = np.zeros(n, dtype=np.intp) + multiplier = 1 + for ids in reversed(chunk_ids): + radix = int(ids.max()) + 1 + if multiplier * radix >= 2**62: + stacked = np.stack([np.asarray(i, dtype=np.intp).ravel() for i in chunk_ids], axis=1) + _, inverse = np.unique(stacked, axis=0, return_inverse=True) + return np.asarray(inverse, dtype=np.intp).reshape(-1) + keys += np.asarray(ids, dtype=np.intp) * multiplier + multiplier *= radix + return keys + + +def _joint_set(transform: IndexTransform, dim_grids: Sequence[DimensionGridLike]) -> JointSet: + block = _prepare_correlated(transform) + dims = block.correlated_dims + chunk_ids = [dim_grids[d].indices_to_chunks(block.flat_storage[d]) for d in dims] + n_points = int(chunk_ids[0].size) + keys = _chunk_keys(chunk_ids) + positions = np.argsort(keys, kind="stable") + sorted_keys = keys[positions] + boundaries = np.flatnonzero(sorted_keys[1:] != sorted_keys[:-1]) + 1 + pointer = np.concatenate([[0], boundaries, [n_points]]).astype(np.intp) + first = positions[pointer[:-1]] + chunk = np.stack([np.asarray(ids, dtype=np.intp)[first] for ids in chunk_ids], axis=1) + index = np.stack([block.flat_index[d][positions] for d in dims], axis=1) + if len(block.broadcast_shape) > 0: + block_coordinates = np.stack( + np.unravel_index(positions, block.broadcast_shape), axis=1 + ).astype(np.intp) + else: + block_coordinates = np.empty((n_points, 0), dtype=np.intp) + starts = np.empty_like(chunk) + extents = np.empty_like(chunk) + for column, d in enumerate(dims): + starts[:, column], extents[:, column] = _chunk_bounds(dim_grids[d], chunk[:, column]) + maps = [cast("ArrayMap", transform.output[d]) for d in dims] + return JointSet( + output_dimensions=dims, + offsets=tuple(m.offset for m in maps), + strides=tuple(m.stride for m in maps), + broadcast_axes=block.broadcast_axes, + broadcast_shape=block.broadcast_shape, + chunk=chunk, + chunk_start=starts, + chunk_extent=extents, + pointer=pointer, + index=index, + positions=np.asarray(positions, dtype=np.intp), + block_coordinates=block_coordinates, + ) + + +# One table row's contribution to a projection: (chunk index, chunk start, +# chunk data extent, bound input axis or None, restricted extent of that axis, +# the rebased chunk-local map, the cell map for the axis, whole-chunk cover). +_AxisPiece = tuple[int, int, int, int | None, int, OutputIndexMap, OutputIndexMap | None, bool] + + +def _slot(slot_of: dict[int, int], axis: StridedSet | IndexedSet) -> int | None: + """The restricted-domain axis a residual table's input axis maps to, if it has one.""" + if axis.input_dimension is None: + return None + return slot_of.get(axis.input_dimension) + + +def _strided_piece(axis: StridedSet, row: int, input_dimension: int | None = None) -> _AxisPiece: + return ( + int(axis.chunk[row]), + int(axis.chunk_start[row]), + int(axis.chunk_extent[row]), + axis.input_dimension, + int(axis.extent[row]), + axis.chunk_map(row, input_dimension), + axis.cell_map(row), + bool(axis.full[row]), + ) + + +def _indexed_piece(axis: IndexedSet, row: int, rank: int, origin: int) -> _AxisPiece: + run = axis.run(row) + k = axis.input_dimension + count = run.stop - run.start + shape = (1,) * k + (count,) + (1,) * (rank - k - 1) + c_start = int(axis.chunk_start[row]) + return ( + int(axis.chunk[row]), + c_start, + int(axis.chunk_extent[row]), + k, + count, + ArrayMap( + index_array=axis.index[run].reshape(shape), + offset=axis.offset - c_start, + stride=axis.stride, + ), + ArrayMap(index_array=axis.positions[run].reshape(shape), offset=origin), + False, + ) + + +@dataclass(frozen=True, slots=True) +class GridPartition: + """A plan in factored form: per-axis tables whose product is the chunk walk. + + `sets` holds one `StridedSet` or `IndexedSet` per output dimension the + transform reads independently, in output-dimension order; `joint` holds + the correlated index arrays, if any. A projection is one row of each + table, so the partition has ``prod(row_shape)`` rows, walked in row-major + order over `row_shape` (the joint table last). Rows are materialized into + `ChunkProjection` objects only on request; a vectorized consumer can read + the tables directly. + + Build one with `partition_transform`, or take it from `ChunkPlan.partition`. + + Examples + -------- + `arr[1:6:2, 5:]` on a `(7, 9)` array with `(3, 4)` chunks touches two + chunks along each axis, so the partition has four rows: + + >>> from zarr_indexing import IndexTransform, plan_chunks + >>> from zarr_indexing.grid import dimension_grids_from_chunks + >>> grids = dimension_grids_from_chunks((3, 4), shape=(7, 9)) + >>> partition = plan_chunks(IndexTransform.from_shape((7, 9))[1:6:2, 5:], grids).partition() + >>> partition.row_shape, len(partition) + ((2, 2), 4) + >>> partition.chunk_coords().tolist() + [[0, 1], [0, 2], [1, 1], [1, 2]] + >>> partition[3].chunk_transform.selection_repr + '{ [0, 4) step 2, [0, 1) }' + """ + + transform: IndexTransform + """The transform this partition factors.""" + + dimension_grids: tuple[DimensionGridLike, ...] + """One grid per storage dimension.""" + + sets: tuple[StridedSet | IndexedSet, ...] + """Independent per-axis tables, in output-dimension order.""" + + joint: JointSet | None + """The correlated index arrays' table, or `None` when there are none.""" + + row_shape: tuple[int, ...] + """Rows per table, `joint` last; the partition is walked in row-major order over it.""" + + def __len__(self) -> int: + return int(np.prod(self.row_shape, dtype=np.intp)) if self.row_shape else 1 + + def chunk_coords(self) -> np.ndarray[Any, np.dtype[np.intp]]: + """Chunk coordinates of every row, shape ``(len(self), output rank)``, without materializing rows.""" + n_rows = len(self) + out = np.empty((n_rows, self.transform.output_rank), dtype=np.intp) + if n_rows == 0 or not self.row_shape: + return out + indices = np.unravel_index(np.arange(n_rows, dtype=np.intp), self.row_shape) + for axis, table_rows in zip(self.sets, indices, strict=False): + out[:, axis.output_dimension] = axis.chunk[table_rows] + if self.joint is not None: + joint_rows = indices[len(self.sets)] + out[:, list(self.joint.output_dimensions)] = self.joint.chunk[joint_rows] + return out + + def __getitem__(self, row: int) -> ChunkProjection: + """Materialize one row (a negative index counts from the end).""" + n_rows = len(self) + if row < 0: + row += n_rows + if not 0 <= row < n_rows: + raise IndexError(f"row {row} is out of range for a partition of {n_rows} rows") + indices = tuple(int(i) for i in np.unravel_index(row, self.row_shape)) + return self._materialize(indices) + + def __iter__(self) -> Iterator[ChunkProjection]: + """Materialize every row in order.""" + if len(self) == 0: + return + if self.joint is None: + yield from self._iter_factorized() + else: + yield from self._iter_correlated() + + # -- factorized assembly ------------------------------------------------ + + def _pieces(self, axis: StridedSet | IndexedSet, row: int) -> _AxisPiece: + if isinstance(axis, StridedSet): + return _strided_piece(axis, row) + domain = self.transform.domain + return _indexed_piece(axis, row, domain.ndim, domain.inclusive_min[axis.input_dimension]) + + def _materialize(self, indices: tuple[int, ...]) -> ChunkProjection: + if self.joint is None: + pieces = [self._pieces(axis, row) for axis, row in zip(self.sets, indices, strict=True)] + return self._factorized_projection(pieces, *self._unbound()) + n_lead, slot_of = self._slots() + residual = [ + _strided_piece(cast("StridedSet", axis), row, _slot(slot_of, axis)) + for axis, row in zip(self.sets, indices[:-1], strict=True) + ] + return self._correlated_projection(residual, indices[-1], n_lead, slot_of) + + def _unbound(self) -> tuple[list[OutputIndexMap | None], bool]: + """Cell maps for request axes no output reads, and whether they permit full coverage. + + A whole-chunk cover must biject onto the chunk, so an unread axis can + only be a singleton. + """ + domain = self.transform.domain + bound = {axis.input_dimension for axis in self.sets if axis.input_dimension is not None} + cell_maps: list[OutputIndexMap | None] = [None] * domain.ndim + unbound_ok = True + for axis in range(domain.ndim): + if axis not in bound: + cell_maps[axis] = DimensionMap( + input_dimension=axis, offset=domain.inclusive_min[axis] + ) + unbound_ok = unbound_ok and domain.shape[axis] <= 1 + return cell_maps, unbound_ok + + def _iter_factorized(self) -> Iterator[ChunkProjection]: + pieces_per_set = [ + [self._pieces(axis, row) for row in range(len(axis))] for axis in self.sets + ] + base_cell_maps, unbound_ok = self._unbound() + for combo in itertools.product(*pieces_per_set): + yield self._factorized_projection(list(combo), base_cell_maps, unbound_ok) + + def _factorized_projection( + self, + pieces: list[_AxisPiece], + base_cell_maps: list[OutputIndexMap | None], + unbound_ok: bool, + ) -> ChunkProjection: + domain = self.transform.domain + shape = list(domain.shape) + cell_maps = list(base_cell_maps) + chunk_maps: list[OutputIndexMap] = [] + chunk_coords: list[int] = [] + chunk_min: list[int] = [] + chunk_max: list[int] = [] + has_array = False + full = unbound_ok + for c, c_start, c_extent, k, extent, chunk_map, cell_map, piece_full in pieces: + chunk_coords.append(c) + chunk_min.append(c_start) + chunk_max.append(c_start + c_extent) + chunk_maps.append(chunk_map) + if isinstance(chunk_map, ArrayMap): + has_array = True + if k is not None: + shape[k] = extent + cell_maps[k] = cell_map + full = full and piece_full + synthetic = IndexDomain._unchecked((0,) * domain.ndim, tuple(shape)) # pyright: ignore[reportPrivateUsage] + if has_array: + coverage: ChunkCoverage = "unknown" + elif full: + coverage = "full" + else: + coverage = "partial" + return ChunkProjection( + chunk_coords=tuple(chunk_coords), + chunk_domain=IndexDomain._unchecked(tuple(chunk_min), tuple(chunk_max)), # pyright: ignore[reportPrivateUsage] + chunk_transform=IndexTransform._unchecked(synthetic, tuple(chunk_maps)), # pyright: ignore[reportPrivateUsage] + cell_transform=IndexTransform._unchecked( # pyright: ignore[reportPrivateUsage] + synthetic, tuple(cast("list[OutputIndexMap]", cell_maps)) + ), + coverage=coverage, + ) + + # -- correlated assembly ------------------------------------------------ + + def _slots(self) -> tuple[int, dict[int, int]]: + """Where each residual slice axis sits in the restricted domain. + + The restricted domain is the collapsed points axis (if the broadcast + block has any axis) followed by the residual slice axes in + input-dimension order. + """ + joint = self.joint + assert joint is not None + n_lead = 1 if len(joint.broadcast_shape) > 0 else 0 + slice_axes = sorted( + m.input_dimension for m in self.transform.output if isinstance(m, DimensionMap) + ) + return n_lead, {axis: n_lead + slot for slot, axis in enumerate(slice_axes)} + + def _iter_correlated(self) -> Iterator[ChunkProjection]: + joint = self.joint + assert joint is not None + n_lead, slot_of = self._slots() + pieces_per_set = [ + [ + _strided_piece(cast("StridedSet", axis), row, _slot(slot_of, axis)) + for row in range(len(axis)) + ] + for axis in self.sets + ] + for residual in itertools.product(*pieces_per_set): + for row in range(len(joint)): + yield self._correlated_projection(list(residual), row, n_lead, slot_of) + + def _correlated_projection( + self, + residual: list[_AxisPiece], + row: int, + n_lead: int, + slot_of: dict[int, int], + ) -> ChunkProjection: + joint = self.joint + assert joint is not None + transform = self.transform + domain = transform.domain + rank = domain.ndim + lo_all = domain.inclusive_min + output_rank = transform.output_rank + n_slice = len(slot_of) + run = joint.run(row) + n_points = run.stop - run.start + points_shape = (n_points,) if n_lead else () + corr_shape = points_shape + (1,) * n_slice + + chunk_coords = [0] * output_rank + chunk_min = [0] * output_rank + chunk_max = [0] * output_rank + chunk_maps: list[OutputIndexMap | None] = [None] * output_rank + extents = [0] * n_slice + slice_origin = [0] * n_slice + for out_dim, (c, c_start, c_extent, k, extent, chunk_map, cell_map, _full) in zip( + (axis.output_dimension for axis in self.sets), residual, strict=True + ): + chunk_coords[out_dim] = c + chunk_min[out_dim] = c_start + chunk_max[out_dim] = c_start + c_extent + chunk_maps[out_dim] = chunk_map + if k is not None: + slot = slot_of[k] - n_lead + extents[slot] = extent + slice_origin[slot] = cast("DimensionMap", cell_map).offset + for column, out_dim in enumerate(joint.output_dimensions): + c_start = int(joint.chunk_start[row, column]) + chunk_coords[out_dim] = int(joint.chunk[row, column]) + chunk_min[out_dim] = c_start + chunk_max[out_dim] = c_start + int(joint.chunk_extent[row, column]) + chunk_maps[out_dim] = ArrayMap( + index_array=joint.index[run, column].reshape(corr_shape), + offset=joint.offsets[column] - c_start, + stride=joint.strides[column], + ) + shape = points_shape + tuple(extents) + synthetic = IndexDomain._unchecked((0,) * (n_lead + n_slice), shape) # pyright: ignore[reportPrivateUsage] + + # One cell map per request axis, materialized over the whole restricted + # block exactly as unravelling the flat scatter offsets would give. + cell_maps: list[OutputIndexMap] = [] + for axis in range(rank): + slot_index = slot_of.get(axis) + if slot_index is None: + column = joint.broadcast_axes.index(axis) + values = joint.block_coordinates[run, column].reshape(corr_shape) + else: + slot = slot_index - n_lead + extent = extents[slot] + values = ( + np.arange(extent, dtype=np.intp) + (slice_origin[slot] - lo_all[axis]) + ).reshape((1,) * (n_lead + slot) + (extent,) + (1,) * (n_slice - slot - 1)) + cell_maps.append( + ArrayMap(index_array=np.broadcast_to(values, shape), offset=lo_all[axis]) + ) + return ChunkProjection( + chunk_coords=tuple(chunk_coords), + chunk_domain=IndexDomain._unchecked(tuple(chunk_min), tuple(chunk_max)), # pyright: ignore[reportPrivateUsage] + chunk_transform=IndexTransform._unchecked( # pyright: ignore[reportPrivateUsage] + synthetic, tuple(cast("list[OutputIndexMap]", chunk_maps)) + ), + cell_transform=IndexTransform._unchecked(synthetic, tuple(cell_maps)), # pyright: ignore[reportPrivateUsage] + coverage="unknown", + ) + + +def partition_transform( + transform: IndexTransform, + dimension_grids: Sequence[DimensionGridLike], +) -> GridPartition: + """Factor a transform over a chunk grid into per-axis tables. + + Parameters + ---------- + transform + Mapping from the request domain to storage coordinates. + dimension_grids + One storage grid per transform output dimension. + + Returns + ------- + GridPartition + The factored plan; iterate it for `ChunkProjection` rows. + + Raises + ------ + ValueError + If two output maps read the same input axis (a diagonal), which has + no factored form; `plan_chunks` still walks such transforms. + + Examples + -------- + >>> from zarr_indexing import IndexTransform + >>> from zarr_indexing.grid import dimension_grids_from_chunks + >>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4)) + >>> partition = partition_transform(IndexTransform.from_shape((3, 4))[1, :], grids) + >>> len(partition), [p.chunk_coords for p in partition] + (2, [(0, 0), (0, 1)]) + """ + grids = tuple(dimension_grids) + if len(grids) != transform.output_rank: + raise ValueError( + "dimension_grids must have one entry per transform output dimension; " + f"got {len(grids)} grids for output rank {transform.output_rank}" + ) + if not _partitionable(transform): + raise ValueError( + "the transform has no factored form: two of its output maps read the same " + "input axis; iterate plan_chunks(...) instead" + ) + if any(size == 0 for size in transform.domain.shape): + return GridPartition( + transform=transform, dimension_grids=grids, sets=(), joint=None, row_shape=(0,) + ) + correlated = transform.index_array_structure == "general" + sets: list[StridedSet | IndexedSet] = [] + for out_dim, m in enumerate(transform.output): + if isinstance(m, ArrayMap): + if correlated: + continue + sets.append(_indexed_set(out_dim, m, grids[out_dim])) + else: + sets.append(_strided_set(transform, out_dim, m, grids[out_dim])) + joint = _joint_set(transform, grids) if correlated else None + row_shape = tuple(len(axis) for axis in sets) + ((len(joint),) if joint is not None else ()) + if any(rows == 0 for rows in row_shape): + row_shape = (0,) + return GridPartition( + transform=transform, + dimension_grids=grids, + sets=tuple(sets), + joint=joint, + row_shape=row_shape, + ) diff --git a/packages/zarr-indexing/src/zarr_indexing/domain.py b/packages/zarr-indexing/src/zarr_indexing/domain.py index a353b81254..b42bd04ccc 100644 --- a/packages/zarr-indexing/src/zarr_indexing/domain.py +++ b/packages/zarr-indexing/src/zarr_indexing/domain.py @@ -78,6 +78,23 @@ def __post_init__(self) -> None: f"Got {len(self.labels)} labels for {len(self.inclusive_min)} dimensions." ) + @classmethod + def _unchecked( + cls, inclusive_min: tuple[int, ...], exclusive_max: tuple[int, ...] + ) -> IndexDomain: + """Build an unlabeled domain from bounds the caller has already established. + + Skips `__post_init__`. For internal producers that derive bounds from an + already-valid domain — chunk resolution builds two domains per chunk, + and re-validating them was a measurable share of a plan's cost. + """ + domain = object.__new__(cls) + object.__setattr__(domain, "inclusive_min", inclusive_min) + object.__setattr__(domain, "exclusive_max", exclusive_max) + object.__setattr__(domain, "labels", None) + object.__setattr__(domain, "_shape", None) + return domain + @classmethod def from_shape(cls, shape: tuple[int, ...]) -> IndexDomain: """Create a domain with origin at zero.""" diff --git a/packages/zarr-indexing/src/zarr_indexing/output_map.py b/packages/zarr-indexing/src/zarr_indexing/output_map.py index 3ee7250efa..16ad7aeb9b 100644 --- a/packages/zarr-indexing/src/zarr_indexing/output_map.py +++ b/packages/zarr-indexing/src/zarr_indexing/output_map.py @@ -223,6 +223,21 @@ def __reduce__(self) -> tuple[object, tuple[object, int, int]]: (self.index_array, self.offset, self.stride), ) + def _with_affine(self, offset: int, stride: int) -> ArrayMap: + """This map's coordinates under a different affine adjustment. + + The frozen index array is shared rather than copied: it is already + owned by immutable bytes and read-only, so the ownership invariant + `__post_init__` establishes holds for the new map too. Chunk + resolution translates every restricted map once per chunk, and + re-copying the array there dominated the cost of small selections. + """ + new = object.__new__(ArrayMap) + object.__setattr__(new, "index_array", self.index_array) + object.__setattr__(new, "offset", offset) + object.__setattr__(new, "stride", stride) + return new + def __eq__(self, other: object) -> bool: """Value equality, comparing index arrays element-wise. diff --git a/packages/zarr-indexing/src/zarr_indexing/transform.py b/packages/zarr-indexing/src/zarr_indexing/transform.py index a8a5963a26..b8fab332ff 100644 --- a/packages/zarr-indexing/src/zarr_indexing/transform.py +++ b/packages/zarr-indexing/src/zarr_indexing/transform.py @@ -177,6 +177,19 @@ def output_rank(self) -> int: """Number of output dimensions — one per output map.""" return len(self.output) + @classmethod + def _unchecked(cls, domain: IndexDomain, output: tuple[OutputIndexMap, ...]) -> IndexTransform: + """Build a transform whose maps the caller has already fitted to `domain`. + + Skips `__post_init__`. For internal producers such as chunk resolution, + which derive every map from an already-valid transform and build two + transforms per chunk. + """ + transform = object.__new__(cls) + object.__setattr__(transform, "domain", domain) + object.__setattr__(transform, "output", output) + return transform + @classmethod def identity(cls, domain: IndexDomain) -> IndexTransform: """The identity transform over `domain`: every result cell reads the source at its own address.""" @@ -545,11 +558,7 @@ def translate(self, shift: tuple[int, ...]) -> IndexTransform: else: # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap) new_output.append( - ArrayMap( - index_array=m.index_array, - offset=m.offset + s, - stride=m.stride, - ) + m._with_affine(m.offset + s, m.stride) # pyright: ignore[reportPrivateUsage] ) return IndexTransform(domain=self.domain, output=tuple(new_output)) @@ -1026,9 +1035,72 @@ def _intersect_orthogonal( return (result, out_indices) +@dataclass(frozen=True, slots=True) +class _CorrelatedBlock: + """The correlated index arrays of a general transform, flattened over the broadcast block. + + Built once per transform by `_prepare_correlated` so that intersecting the + transform with many chunks does not re-broadcast and re-scan the arrays for + each chunk: `flat_storage` holds the storage coordinate of every block + point per correlated output dimension, `flat_index` the raw index values. + """ + + correlated_dims: tuple[int, ...] + broadcast_axes: tuple[int, ...] + broadcast_shape: tuple[int, ...] + flat_index: dict[int, np.ndarray[Any, np.dtype[np.intp]]] + flat_storage: dict[int, np.ndarray[Any, np.dtype[np.intp]]] + + +def _prepare_correlated(transform: IndexTransform) -> _CorrelatedBlock: + """Flatten a general transform's index arrays over its broadcast block.""" + correlated_dims = tuple(i for i, m in enumerate(transform.output) if isinstance(m, ArrayMap)) + + # 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) + + flat_index: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + flat_storage: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} + for out_dim in correlated_dims: + arr_map = cast("ArrayMap", transform.output[out_dim]) + if any(a not in broadcast_axes for a in arr_map.dependency_axes): + # Reachable only by hand-building a transform: no selection binds + # the same input axis to both a slice map and an index array. + raise NotImplementedError( + "intersecting a transform whose index array varies over an " + "input dimension also bound by a slice map is not supported" + ) + arr = arr_map.index_array + # Index arrays are singleton on every non-broadcast axis, so they + # collapse (C-order) to the broadcast block. A map may also be singleton + # along a block axis it does not vary over (an orthogonal member, or a + # leftover broadcast axis), so the collapsed array is broadcast up to the + # full block rather than reshaped. + block = arr.reshape(tuple(arr.shape[a] for a in broadcast_axes)) + flat = np.ascontiguousarray(np.broadcast_to(block, broadcast_shape)).reshape(-1) + flat_index[out_dim] = np.asarray(flat, dtype=np.intp) + flat_storage[out_dim] = checked_affine(arr_map.offset, arr_map.stride, flat) + return _CorrelatedBlock( + correlated_dims=correlated_dims, + broadcast_axes=broadcast_axes, + broadcast_shape=broadcast_shape, + flat_index=flat_index, + flat_storage=flat_storage, + ) + + def _intersect_general( transform: IndexTransform, output_domain: IndexDomain, + block: _CorrelatedBlock | None = None, + positions: np.ndarray[Any, np.dtype[np.intp]] | None = None, ) -> tuple[IndexTransform, np.ndarray[Any, np.dtype[np.intp]]] | None: """Intersect a transform with any index-array structure, pointwise. @@ -1051,49 +1123,33 @@ def _intersect_general( rank 0: the block either survives whole or the intersection is empty. The result keeps only the residual slice axes and `out_indices` loses its leading points axis, so the sub-transform's rank still matches the view's. - """ - correlated_dims = [i for i, m in enumerate(transform.output) if isinstance(m, ArrayMap)] - # 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) - - for out_dim in correlated_dims: - arr_map = cast("ArrayMap", transform.output[out_dim]) - if any(a not in broadcast_axes for a in arr_map.dependency_axes): - # Reachable only by hand-building a transform: no selection binds - # the same input axis to both a slice map and an index array. - raise NotImplementedError( - "intersecting a transform whose index array varies over an " - "input dimension also bound by a slice map is not supported" - ) - - # Joint bounds mask over the broadcast block. - combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None - for out_dim in correlated_dims: - cm = cast("ArrayMap", transform.output[out_dim]) - storage = checked_affine(cm.offset, cm.stride, cm.index_array) - lo = output_domain.inclusive_min[out_dim] - hi = output_domain.exclusive_max[out_dim] - mask = (storage >= lo) & (storage < hi) - combined = mask if combined is None else (combined & mask) - assert combined is not None - # Index arrays are singleton on every non-broadcast axis, so the mask - # collapses (C-order) to the broadcast block. A map may also be singleton - # along a block axis it does not vary over (an orthogonal member, or a - # leftover broadcast axis), so the collapsed mask is broadcast up to the - # full block rather than reshaped. - combined_block = combined.reshape(tuple(combined.shape[a] for a in broadcast_axes)) - combined_bcast = np.broadcast_to(combined_block, broadcast_shape) - surviving = np.nonzero(combined_bcast.reshape(-1))[0].astype(np.intp) - if surviving.size == 0: + `block` is the transform's flattened correlated arrays (`_prepare_correlated`), + reused across many chunks by chunk resolution. `positions` names the block + points that land in `output_domain`, ascending; a caller that has already + partitioned the points by chunk supplies it so the intersection costs + ``O(surviving points)`` rather than a scan of the whole block. + """ + if block is None: + block = _prepare_correlated(transform) + correlated_dims = block.correlated_dims + broadcast_axes = block.broadcast_axes + broadcast_shape = block.broadcast_shape + + if positions is None: + # Joint bounds mask over the broadcast block. + combined: np.ndarray[Any, np.dtype[np.bool_]] | None = None + for out_dim in correlated_dims: + storage = block.flat_storage[out_dim] + lo = output_domain.inclusive_min[out_dim] + hi = output_domain.exclusive_max[out_dim] + mask = (storage >= lo) & (storage < hi) + combined = mask if combined is None else (combined & mask) + assert combined is not None + positions = np.flatnonzero(combined).astype(np.intp) + if positions.size == 0: return None + surviving = positions # Intersect residual (slice / constant) dimensions independently. Slice dims # are ordered by input dimension so their flat-buffer strides are row-major. @@ -1118,14 +1174,9 @@ def _intersect_general( n_points = int(surviving.size) n_slice = len(slice_dims) - corr_values: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = {} - for out_dim in correlated_dims: - arr = cast("ArrayMap", transform.output[out_dim]).index_array - block = arr.reshape(tuple(arr.shape[a] for a in broadcast_axes)) - corr_values[out_dim] = np.asarray( - np.ascontiguousarray(np.broadcast_to(block, broadcast_shape)).reshape(-1)[surviving], - dtype=np.intp, - ) + corr_values: dict[int, np.ndarray[Any, np.dtype[np.intp]]] = { + out_dim: block.flat_index[out_dim][surviving] for out_dim in correlated_dims + } # A rank-0 broadcast block contributes no axis: the leading `(n_points,)` of # the domain, of every index array, and of `out_indices` is present only when @@ -1150,7 +1201,7 @@ def _intersect_general( corr = cast("ArrayMap", m) new_output.append( ArrayMap( - index_array=corr_values[out_dim].reshape(corr_shape).astype(np.intp), + index_array=corr_values[out_dim].reshape(corr_shape), offset=corr.offset, stride=corr.stride, ) diff --git a/packages/zarr-indexing/tests/test_chunk_resolution.py b/packages/zarr-indexing/tests/test_chunk_resolution.py index fbd779f463..4ad1b9e304 100644 --- a/packages/zarr-indexing/tests/test_chunk_resolution.py +++ b/packages/zarr-indexing/tests/test_chunk_resolution.py @@ -465,7 +465,7 @@ def test_sorted_coordinates_bypass_intersection(self, monkeypatch: pytest.Monkey ] == [[0, 1], [2, 3], [4, 5]] def test_unsorted_coordinates_use_intersection(self, monkeypatch: pytest.MonkeyPatch) -> None: - """Unsorted coordinates retain the general intersection path.""" + """Unsorted coordinates are grouped by chunk without whole-transform intersections.""" transform = IndexTransform.from_shape((12,)).vindex[np.array([9, 0, 4], dtype=np.intp)] grid = ChunkGrid(dimensions=(FixedDimension(size=4, extent=12),)) calls = _count_intersect_calls(monkeypatch) @@ -473,7 +473,7 @@ def test_unsorted_coordinates_use_intersection(self, monkeypatch: pytest.MonkeyP projections = list(plan_chunks(transform, grid.dimensions)) assert [projection.chunk_coords for projection in projections] == [(0,), (1,), (2,)] - assert calls["n"] == 3 + assert calls["n"] == 0 class CountingUnitGrid: @@ -544,8 +544,8 @@ def test_sparse_one_dimensional_selection_skips_the_dense_span( @pytest.mark.parametrize( ("mode", "expected_coords", "expected_calls"), [ - ("orthogonal", [(0, 0), (0, 999), (999, 0), (999, 999)], 4), - ("correlated", [(0, 0), (999, 999)], 2), + ("orthogonal", [(0, 0), (0, 999), (999, 0), (999, 999)], 0), + ("correlated", [(0, 0), (999, 999)], 0), ], ) def test_sparse_two_dimensional_selection_uses_only_touched_combinations( @@ -555,7 +555,12 @@ def test_sparse_two_dimensional_selection_uses_only_touched_combinations( expected_coords: list[tuple[int, int]], expected_calls: int, ) -> None: - """Orthogonal points use their outer product; correlated points remain paired.""" + """Orthogonal points use their outer product; correlated points remain paired. + + Neither walk intersects the whole transform per chunk: the orthogonal + axes are resolved once each, and correlated points are sorted into + their chunks up front. + """ base = IndexTransform.from_shape((4000, 4000)) first = np.array([1, 3997], dtype=np.intp) second = np.array([2, 3998], dtype=np.intp) @@ -578,7 +583,7 @@ def test_sparse_two_dimensional_selection_uses_only_touched_combinations( def test_correlated_diagonal_scales_with_points_not_their_product( self, monkeypatch: pytest.MonkeyPatch ) -> None: - """Fifty diagonal points require fifty, rather than 2500, intersections.""" + """Fifty diagonal points touch fifty chunks, not 2500, and need no re-intersection.""" n_points = 50 coordinates = np.arange(n_points, dtype=np.intp) * 8 transform = IndexTransform.from_shape((4000, 4000)).vindex[coordinates, coordinates] @@ -595,4 +600,143 @@ def test_correlated_diagonal_scales_with_points_not_their_product( assert sorted(projection.chunk_coords for projection in projections) == [ (2 * index, 2 * index) for index in range(n_points) ] - assert calls["n"] == n_points + assert calls["n"] == 0 + + +# --------------------------------------------------------------------------- +# GridPartition: the factored form +# --------------------------------------------------------------------------- + + +def _partition_cases() -> list[tuple[str, IndexTransform, tuple[Any, ...]]]: + base = IndexTransform.from_shape((7, 9, 5)) + fixed = dimension_grids_from_chunks((3, 4, 2), shape=(7, 9, 5)) + varying = ( + VaryingDimension(edges=(2, 5), extent=7), + VaryingDimension(edges=(1, 4, 4), extent=9), + FixedDimension(size=2, extent=5), + ) + return [ + ("identity", base, fixed), + ("strided", base[1:6:2, 5:, ::3], fixed), + ("strided varying", base[1:6:2, 5:, ::3], varying), + ("scalars", base[3, :, 4], fixed), + ("reversed", base[::-1, 8:1:-1, :], fixed), + ("empty", base[3:3, :, :], fixed), + ("oindex arrays", base.oindex[np.array([6, 0, 0, 2]), :, np.array([4, 1])], fixed), + ("oindex mixed", base.oindex[np.array([1, 5]), 2, 1:5:2], varying), + ("oindex one element", base.oindex[np.array([2]), :, :], fixed), + ("vindex", base.vindex[np.array([0, 6, 6, 1]), np.array([8, 0, 1, 8]), :], fixed), + ("vindex varying", base.vindex[np.array([0, 6, 6, 1]), np.array([8, 0, 1, 8]), 2], varying), + ( + "vindex 2-d block", + base.vindex[ + np.array([[0, 6], [3, 1]]), np.array([[8, 0], [2, 8]]), np.array([[4, 0], [1, 1]]) + ], + fixed, + ), + ( + "vindex 1-d sorted", + IndexTransform.from_shape((20,)).vindex[np.array([1, 5, 9, 17])], + dimension_grids_from_chunks((4,), shape=(20,)), + ), + ("rank 0", IndexTransform.from_shape(()), ()), + ] + + +@pytest.mark.parametrize("case", _partition_cases(), ids=lambda case: case[0]) +def test_partition_rows_match_general_walk( + case: tuple[str, IndexTransform, tuple[Any, ...]], +) -> None: + """Every partition row equals the projection the whole-transform walk produces.""" + _, transform, grids = case + partition = plan_chunks(transform, grids).partition() + rows = list(partition) + general = list(chunk_resolution._iter_general_projections(transform, grids)) + assert rows == general + assert len(partition) == len(rows) + assert [partition[i] for i in range(len(partition))] == rows + if rows: + assert partition[-1] == rows[-1] + + +@pytest.mark.parametrize("case", _partition_cases(), ids=lambda case: case[0]) +def test_partition_chunk_coords_are_vectorized_rows( + case: tuple[str, IndexTransform, tuple[Any, ...]], +) -> None: + """`chunk_coords` reads the tables without materializing a row per chunk.""" + _, transform, grids = case + partition = plan_chunks(transform, grids).partition() + coords = partition.chunk_coords() + assert coords.shape == (len(partition), transform.output_rank) + assert coords.tolist() == [list(p.chunk_coords) for p in partition] + + +def test_partition_tables_describe_chunk_local_coordinates() -> None: + """The columnar tables carry exactly what each row's chunk transform maps to.""" + transform = IndexTransform.from_shape((7, 9)).oindex[np.array([6, 0, 0, 2]), 5:] + grids = dimension_grids_from_chunks((3, 4), shape=(7, 9)) + partition = plan_chunks(transform, grids).partition() + indexed, strided = partition.sets + assert isinstance(indexed, chunk_resolution.IndexedSet) + assert isinstance(strided, chunk_resolution.StridedSet) + # rows are in chunk order; the array [6, 0, 0, 2] lands in chunks 0, 0, 0, 2 + assert indexed.chunk.tolist() == [0, 2] + assert indexed.pointer.tolist() == [0, 3, 4] + assert indexed.local.tolist() == [0, 0, 2, 0] + assert indexed.positions.tolist() == [1, 2, 3, 0] + assert strided.chunk.tolist() == [1, 2] + assert strided.local_start.tolist() == [1, 0] + assert strided.extent.tolist() == [3, 1] + assert strided.origin.tolist() == [5, 8] # literal request coordinates: the axis is [5, 9) + assert strided.full.tolist() == [False, True] + for row, projection in enumerate(partition): + table_rows = np.unravel_index(row, partition.row_shape) + run = indexed.run(int(table_rows[0])) + storage = projection.chunk_transform.apply_many( + np.array(list(np.ndindex(*projection.chunk_transform.domain.shape))) + ) + assert sorted(set(storage[:, 0].tolist())) == sorted(set(indexed.local[run].tolist())) + + +def test_joint_set_groups_points_by_chunk() -> None: + transform = IndexTransform.from_shape((7, 9)).vindex[ + np.array([0, 6, 6, 1]), np.array([8, 0, 1, 8]) + ] + grids = dimension_grids_from_chunks((3, 4), shape=(7, 9)) + joint = plan_chunks(transform, grids).partition().joint + assert joint is not None + assert joint.chunk.tolist() == [[0, 2], [2, 0]] + assert joint.pointer.tolist() == [0, 2, 4] + assert joint.positions.tolist() == [0, 3, 1, 2] + assert joint.local.tolist() == [[0, 0], [1, 0], [0, 0], [0, 1]] + + +def test_partition_rejects_diagonal_but_plan_still_walks() -> None: + """Two maps reading one input axis have no factored form; the general walk covers them.""" + transform = IndexTransform( + domain=IndexDomain.from_shape((4,)), + output=(DimensionMap(input_dimension=0), DimensionMap(input_dimension=0)), + ) + grids = dimension_grids_from_chunks((2, 2), shape=(4, 4)) + plan = plan_chunks(transform, grids) + with pytest.raises(ValueError, match="no factored form"): + plan.partition() + assert [p.chunk_coords for p in plan] == [(0, 0), (1, 1)] + + +def test_partition_is_memoized_on_the_plan() -> None: + plan = plan_chunks( + IndexTransform.from_shape((6,)), dimension_grids_from_chunks((2,), shape=(6,)) + ) + assert plan.partition() is plan.partition() + + +def test_partition_getitem_bounds() -> None: + partition = plan_chunks( + IndexTransform.from_shape((6,)), dimension_grids_from_chunks((2,), shape=(6,)) + ).partition() + assert len(partition) == 3 + assert partition[-1].chunk_coords == (2,) + with pytest.raises(IndexError): + partition[3] From 7895f2d82068ce0d1a850d020dbadbb22286ac5a Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 2 Sep 2026 18:48:32 +0200 Subject: [PATCH 02/18] docs(zarr-indexing): changelog fragment for the grid partition Assisted-by: ClaudeCode:claude-fable-5-1 --- packages/zarr-indexing/changes/316.feature.md | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 packages/zarr-indexing/changes/316.feature.md diff --git a/packages/zarr-indexing/changes/316.feature.md b/packages/zarr-indexing/changes/316.feature.md new file mode 100644 index 0000000000..c25a99d626 --- /dev/null +++ b/packages/zarr-indexing/changes/316.feature.md @@ -0,0 +1,8 @@ +Chunk plans now have a factored, columnar form. `ChunkPlan.partition()` and +`partition_transform()` return a `GridPartition`: one `StridedSet` or +`IndexedSet` table per output dimension the transform reads independently, +plus a `JointSet` for correlated (`vindex`) index arrays. A `ChunkProjection` +is one row of each table and is derived on demand, so planning costs the sum +of the touched chunks per axis rather than their product, and a vectorized +consumer can read the tables without materializing a projection per chunk. +The projections a plan yields are unchanged. From 30b8fa7ab45987906b9a5c6eb897e7641359df13 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 2 Sep 2026 18:54:55 +0200 Subject: [PATCH 03/18] docs(zarr-indexing): say why the suite runs from the repo root The package and its tests import nothing from zarr; the old comments claimed the chunk-resolution tests needed zarr's ChunkGrid, which stopped being true once the package grew its own grids. The real reason is the shared pinned test toolchain. Assisted-by: ClaudeCode:claude-fable-5-1 --- .github/workflows/zarr-indexing.yml | 9 ++++----- packages/zarr-indexing/justfile | 10 ++++------ packages/zarr-indexing/pyproject.toml | 17 +++++++++-------- 3 files changed, 17 insertions(+), 19 deletions(-) diff --git a/.github/workflows/zarr-indexing.yml b/.github/workflows/zarr-indexing.yml index f0b0cec65d..3b106e16aa 100644 --- a/.github/workflows/zarr-indexing.yml +++ b/.github/workflows/zarr-indexing.yml @@ -43,11 +43,10 @@ jobs: uses: extractions/setup-just@53165ef7e734c5c07cb06b3c8e7b647c5aa16db3 # v4 - name: Set up Python ${{ matrix.python-version }} run: uv python install ${{ matrix.python-version }} - # The transform tests exercise chunk resolution against zarr's ChunkGrid, - # so they run against the repo-root environment (which provides `zarr`) - # with this package as an editable overlay rather than in package - # isolation. The recipes carry that invocation; this step only fixes the - # interpreter the matrix asked for. + # The suite imports nothing from `zarr`; it runs against the repo-root + # environment, with this package as an editable overlay, to share the + # parent project's pinned test toolchain. The recipes carry that + # invocation; this step only fixes the interpreter the matrix asked for. - name: Sync test dependency group run: uv sync --project ../.. --group test --python ${{ matrix.python-version }} - name: Run pytest diff --git a/packages/zarr-indexing/justfile b/packages/zarr-indexing/justfile index 1b7164f647..44874bc0be 100644 --- a/packages/zarr-indexing/justfile +++ b/packages/zarr-indexing/justfile @@ -5,12 +5,10 @@ default: @just --list -# The chunk-resolution tests exercise this package against zarr's ChunkGrid, so -# they need an environment that has both `zarr` and this package installed. -# `zarr` is deliberately not a dependency of this package, and the repo is not -# a uv workspace, so run against the repo-root environment (which provides -# `zarr`) with this package layered in as an editable overlay — the same -# invocation CI uses. +# Nothing here imports `zarr`; the suite runs against the repo-root +# environment so it shares the parent project's pinned test toolchain +# (hypothesis, ruff, pyright). The repo is not a uv workspace, so this package +# is layered in as an editable overlay — the same invocation CI uses. # Run the test suite; extra args are passed to pytest test *args: uv run --project ../.. --group test --with-editable . python -m pytest tests src/zarr_indexing {{ args }} diff --git a/packages/zarr-indexing/pyproject.toml b/packages/zarr-indexing/pyproject.toml index 60f69a1dbb..953a432f07 100644 --- a/packages/zarr-indexing/pyproject.toml +++ b/packages/zarr-indexing/pyproject.toml @@ -48,14 +48,15 @@ Changelog = "https://github.com/zarr-developers/zarr-python/blob/main/packages/z Documentation = "https://zarr-indexing.readthedocs.io/" [dependency-groups] -# The transform tests exercise chunk resolution against zarr's ChunkGrid -# (tests/test_chunk_resolution.py) and are collected by the parent zarr-python -# test suite, which already has zarr installed. `zarr` is intentionally NOT -# listed here to avoid a workspace dependency cycle; run these tests from the -# repo root (`uv run pytest packages/zarr-indexing/tests`), not in isolation. -# `hypothesis` arrives via the `testing` extra, which is what -# `zarr_indexing.testing` needs; the repo-root `test` group pins the exact -# version CI runs against. Bump the two together. +# The package and its tests import nothing from `zarr`: chunk resolution +# consumes the `DimensionGridLike` protocol and the tests use this package's +# own grids (`zarr_indexing.grid`). The suite is nevertheless run from the +# repo-root environment (`just test`, the invocation CI uses) so that it sees +# the same pinned toolchain — hypothesis, ruff, pyright — as the parent +# project; the repo is not a uv workspace, so this package is layered in as an +# editable overlay. `hypothesis` arrives via the `testing` extra, which is +# what `zarr_indexing.testing` needs; the repo-root `test` group pins the +# exact version CI runs against. Bump the two together. test = ["pytest", "hypothesis>=6.160.0"] docs = [ # Pins match the zarr-python docs environment in the repo-root From ad48b0a010e856d3ff0cf7b29b45d54502ea7fd9 Mon Sep 17 00:00:00 2001 From: Davis Vann Bennett Date: Wed, 2 Sep 2026 19:15:36 +0200 Subject: [PATCH 04/18] docs(zarr-indexing): document the grid partition and retire the per-chunk narrative The module docstring described intersecting the whole transform with every candidate chunk as "the algorithm"; that walk is now the fallback for hand-built diagonals only. It now explains the factored form and its three tables, and why they cost the sum of the touched chunks per axis. The visual guide gains a final integrator section, "A plan is a product of per-axis tables", with an executable snippet that reads the StridedSet, IndexedSet and JointSet tables off real plans and checks the plan's projections against the partition's rows. Integration boundaries gains "Reading the tables directly", a consumer that assembles a strided box from the tables with no projection materialized. The API index, landing page and design notes (TensorStore lineage, the performance caveat, and the box/query split) point at the new section. Assisted-by: ClaudeCode:claude-fable-5-1 --- packages/zarr-indexing/docs/api/index.md | 9 +- packages/zarr-indexing/docs/design-notes.md | 29 +++-- packages/zarr-indexing/docs/guide/index.md | 77 ++++++++++++- .../zarr-indexing/docs/guide/integrations.md | 26 +++++ packages/zarr-indexing/docs/index.md | 4 +- .../docs/snippets/grid_partition.py | 102 ++++++++++++++++++ .../src/zarr_indexing/chunk_resolution.py | 63 ++++++----- 7 files changed, 265 insertions(+), 45 deletions(-) create mode 100644 packages/zarr-indexing/docs/snippets/grid_partition.py diff --git a/packages/zarr-indexing/docs/api/index.md b/packages/zarr-indexing/docs/api/index.md index 674c701a0f..dee436c1ba 100644 --- a/packages/zarr-indexing/docs/api/index.md +++ b/packages/zarr-indexing/docs/api/index.md @@ -11,9 +11,9 @@ links: [Lazy views compose](../guide/index.md#lazy-views-compose), then open [`zarr_indexing.lazy_array`](lazy_array.md) for `LazyArray`. - **Integrate a chunked source:** finish - [One cell domain, two projections](../guide/index.md#one-cell-domain-two-projections), + [A plan is a product of per-axis tables](../guide/index.md#a-plan-is-a-product-of-per-axis-tables), then open [`zarr_indexing.chunk_resolution`](chunk_resolution.md) for - `plan_chunks`. Start with + `plan_chunks` and `GridPartition`. Start with [Coordinates are addresses](../guide/index.md#coordinates-are-addresses) if literal coordinates are unfamiliar. @@ -36,7 +36,10 @@ and the wire format built on top of it. - [`zarr_indexing.chunk_resolution`](chunk_resolution.md) — `plan_chunks`, which lazily projects a request through a caller-selected grid, - plus the reusable `ChunkPlan` and paired-transform `ChunkProjection` values + the reusable `ChunkPlan` and paired-transform `ChunkProjection` values, and + the plan's factored form: `GridPartition` (from `ChunkPlan.partition` or + `partition_transform`), holding one `StridedSet` or `IndexedSet` table per + axis and a `JointSet` for correlated index arrays - [`zarr_indexing.grid`](grid.md) — `DimensionGridLike`, the Protocol describing the narrow chunk-grid surface chunk resolution consumes, so that nothing here imports `zarr`, plus `EdgeDimensionGrid` and diff --git a/packages/zarr-indexing/docs/design-notes.md b/packages/zarr-indexing/docs/design-notes.md index 78b352d376..1050354f0a 100644 --- a/packages/zarr-indexing/docs/design-notes.md +++ b/packages/zarr-indexing/docs/design-notes.md @@ -35,6 +35,16 @@ about the deliberately matching semantics: discriminator, and `tests/test_ndsel_tensorstore.py` loads our bodies into `tensorstore.IndexTransform(json=...)` and round-trips them back through our engine layer. +- **Chunk partitioning.** Both factor a transform over a grid before visiting + any cell. TensorStore's `IndexTransformGridPartition` holds *strided sets* + (one per affine grid dimension) and *index array sets* (correlated array + dimensions grouped into connected sets, their points partitioned by grid + cell), and derives a per-cell transform on iteration. + [`GridPartition`](api/chunk_resolution.md) is the same structure — + `StridedSet` and `IndexedSet` tables per axis, a `JointSet` for correlated + arrays — and derives each `ChunkProjection` from one row of each table + ([the guide](guide/index.md#a-plan-is-a-product-of-per-axis-tables) shows + the tables). The representations differ in one place: index arrays. Both models want an index array at the transform's full input rank, with singleton axes for the dimensions @@ -85,11 +95,14 @@ safe, `partial` proves it is not, and `unknown` conservatively covers fancy selections whose duplicates would require additional work to classify. The comparison also runs the other way. TensorStore is a mature, heavily -optimized C++ system whose performance this library cannot approach: resolution -here is Python-level bookkeeping over NumPy, and the per-part overhead is -significant. This library is small and depends on nothing beyond NumPy, so the -algebra can be adopted by a Python project that wants the model without the C++ -runtime. +optimized C++ system whose performance this library cannot approach. Planning +here is vectorized per axis, but each materialized `ChunkProjection` is +Python-level bookkeeping over NumPy — two domains, two transforms and the +projection itself — so the per-part overhead of the object view is +significant; a consumer that reads the partition's tables directly pays only +for the copies it makes. This library is small and depends on nothing beyond +NumPy, so the algebra can be adopted by a Python project that wants the model +without the C++ runtime. ## Bounding-box selections vs query selections @@ -146,7 +159,11 @@ soon as any stride exceeds 1. The two also behave differently under partitioning: a box touches a regularly-spaced run of parts, in increasing order, each at most once — a stride larger than a part's extent skips parts outright, so the run is not contiguous — while a query can touch any subset of -them, in any order, more than once. +them, in any order, more than once. The +[partition](guide/index.md#a-plan-is-a-product-of-per-axis-tables) makes the +split structural: a box factors into one `StridedSet` per axis, an orthogonal +query adds `IndexedSet`s, and only correlated (`vindex`) arrays need the +`JointSet`. [`LazyArray`](api/lazy_array.md) exposes the category directly: diff --git a/packages/zarr-indexing/docs/guide/index.md b/packages/zarr-indexing/docs/guide/index.md index 42e47530fb..f4108930f3 100644 --- a/packages/zarr-indexing/docs/guide/index.md +++ b/packages/zarr-indexing/docs/guide/index.md @@ -8,8 +8,9 @@ through those stages. The first four sections are for anyone indexing arrays: coordinates, transforms, composition, and result axes. **If you are using lazy indexing rather than building a storage backend, you can stop after section four.** -The last two sections are for integrators: they turn a request into a chunk -plan and pair each chunk read with its place in the result. +The last three sections are for integrators: they turn a request into a chunk +plan, pair each chunk read with its place in the result, and show the per-axis +tables the plan is built from. Throughout, one division of labor holds: the transform answers **which values?** and is independent of the backend; the reader answers **how do I @@ -327,7 +328,9 @@ each projection are the next section's subject.) The plan describes work but does not perform it. It contains no array source, storage backend, codec pipeline, buffer, or scheduler. A Zarr reader, a task queue, or a viewport can consume the same logical plan and decide independently -how and when to fetch its two chunks. +how and when to fetch its two chunks. Nor does it hold two projection objects: +it holds one small table per axis, from which the projections are derived — +[the last section](#a-plan-is-a-product-of-per-axis-tables) shows them. On the wrapper, this partitioning is called **parts**: `with_parts(shape)` gives a `LazyArray` a grid of uniform boxes to divide its reads along @@ -435,6 +438,74 @@ The paired representation preserves information that a bounding box or local selector discards: exact request order, duplicate destinations, and the correspondence between every request position and its chunk-local source cell. +## A plan is a product of per-axis tables {#a-plan-is-a-product-of-per-axis-tables} + +Look again at the two projections of `image[1, :]`. Each chunk transform has +one map per source axis, and every one of those maps came from restricting +the request's map for *that axis alone* to *that axis's* chunk: the fixed row +`ConstantMap(1)` lands in row-chunk 0 whatever the column chunk is, and the +column slice meets column-chunk 0 as local columns `0:2` and column-chunk 1 as +local columns `0:2` whatever the row chunk is. Restricting a transform to a +chunk distributes over axes whenever each output map reads its own request +axis — which every basic and orthogonal selection satisfies. So the plan does +not intersect the whole transform with every chunk. It resolves each axis +once, into a table with one row per chunk that axis touches, and a projection +is one row of each table combined. + +```text +image[1, :] over 2-by-2 chunks + +axis 0 (rows): ConstantMap(1) axis 1 (columns): DimensionMap +row | chunk start local extent row | chunk start local_start extent origin full + 0 | 0 0 1 1 0 | 0 0 0 2 0 yes + 1 | 1 2 0 2 2 yes + +row_shape (1, 2): 1 x 2 = 2 projections +projection (0, 0) = axis-0 row 0 x axis-1 row 0 -> chunk (0, 0), local (1, 0:2), request columns 0:2 +projection (0, 1) = axis-0 row 0 x axis-1 row 1 -> chunk (0, 1), local (1, 0:2), request columns 2:4 +``` + +`ChunkPlan.partition()` returns this factored form, a `GridPartition`. Its +`sets` hold one table per source axis, in axis order; `row_shape` is the +number of rows in each; and a projection is addressed by one row index per +table, walked in row-major order. The executable example reads the two tables +above off the plan, checks that the plan's projections are exactly the +partition's rows, and evaluates row `(0, 1)` on both of its transforms: + +```python +--8<-- "snippets/grid_partition.py:strided-tables" +``` + +There are three kinds of table, matching the three map kinds and the one +arrangement that does not factor: + +| Table | Holds | One row per | +| --- | --- | --- | +| `StridedSet` | A `ConstantMap` or `DimensionMap` axis: chunk-local `local_start`, `extent`, the request `origin` of the first cell, and `full`, whether the row covers its chunk exactly once | touched chunk along that axis | +| `IndexedSet` | An orthogonal `ArrayMap` axis (`.oindex`): its coordinates grouped by chunk in CSR form — `pointer`, `index`, and the request `positions` they fill, with `local` for the chunk-local coordinates | touched chunk along that axis | +| `JointSet` | The correlated arrays of a `.vindex` selection, which read the same request axes and so do not distribute: a chunk constrains all of them at once, so their points are sorted into chunks together, once | touched chunk, addressed on all correlated axes | + +The gather from the previous section, `oindex[[4, 1, 1], 2:6]` over 3-by-4 +chunks, groups rows `1, 1` into chunk 0 and row `4` into chunk 1 while +remembering that row `4` fills request position 0. A `vindex` selection keeps +its points paired in the joint table: + +```python +--8<-- "snippets/grid_partition.py:indexed-and-joint" +``` + +Two properties follow from the factoring. Building the tables costs the *sum* +of the chunks touched per axis, never their product, and correlated points +cost one sort; a selection over a million chunks is described by three short +tables. And projections are derived from rows only when asked for — by +iterating, by `partition[row]`, or not at all: `chunk_coords()` lists every +chunk the plan touches without materializing a row, and a consumer can read +the columns directly, as [Integration boundaries](integrations.md#reading-the-tables-directly) +shows. The only transform with no factored form is a hand-built diagonal, two +output maps reading one request axis; `partition()` raises `ValueError` for +it, and `plan_chunks` still walks it by intersecting the whole transform with +each chunk. + ---