zarr_indexing.chunk_resolution

zarr_indexing.chunk_resolution

Chunk resolution — mapping transforms to chunk-level I/O.

Given an IndexTransform (which coordinates a request reads) and one grid per storage dimension (how storage is divided into chunks), chunk resolution answers:

For each chunk, which storage coordinates does this transform touch,
and where do those values land in the request?

The public result is a lazy, reusable ChunkPlan whose rows are ChunkProjections. Each identifies a chunk and pairs a chunk-local transform with a transform back to the request's cells, over one shared zero-origin cell domain, without assuming NumPy selectors, a codec pipeline, or a scheduler.

The plan is computed in factored form, the GridPartition. Restricting a transform to a chunk box distributes over output dimensions whenever each output map reads its own input axis — every basic and orthogonal selection — so each axis is resolved once against its grid into a table:

• StridedSet — a ConstantMap or DimensionMap axis: one row per touched chunk, holding the chunk-local start, the extent, the request position of the first cell, and whether the row covers its chunk exactly once.
• IndexedSet — an orthogonal ArrayMap axis: its coordinates grouped by chunk in CSR form, with the request positions they fill.
• JointSet — the correlated (vindex) index arrays, which read the same input axes and so do not distribute: their points are sorted into chunks together, once.

A projection is one row of each table combined. Building the tables costs the sum of the touched chunks per axis rather than their product, rows are materialized only on request, and a consumer may read the tables directly instead. Two output maps that read one input axis through DimensionMaps (a diagonal, which no selection produces) have no factored form and are rejected with ValueError; an index array varying over an axis a DimensionMap also reads is rejected with NotImplementedError.

ChunkCoverage

ChunkCoverage = Literal['full', 'partial', 'unknown']

ChunkPlan dataclass

A reusable, lazy partition of an index transform over a chunk grid.

Construct plans with plan_chunks. The plan's factored form, a GridPartition, is built on first use and memoized; iterating the plan or projections() materializes fresh ChunkProjection rows from it.

Examples:

Row 1 of a (3, 4) array with (2, 2) chunks crosses two chunks, and the plan can be walked again after it is exhausted:

>>> from zarr_indexing import IndexTransform
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
>>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
>>> [p.chunk_coords for p in plan]
[(0, 0), (0, 1)]
>>> [p.chunk_coords for p in plan.projections()]
[(0, 0), (0, 1)]

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class ChunkPlan:
    """A reusable, lazy partition of an index transform over a chunk grid.

    Construct plans with `plan_chunks`. The plan's factored form, a
    `GridPartition`, is built on first use and memoized; iterating the plan
    or `projections()` materializes fresh `ChunkProjection` rows from it.

    Examples
    --------
    Row 1 of a `(3, 4)` array with `(2, 2)` chunks crosses two chunks, and
    the plan can be walked again after it is exhausted:

    >>> from zarr_indexing import IndexTransform
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
    >>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
    >>> [p.chunk_coords for p in plan]
    [(0, 0), (0, 1)]
    >>> [p.chunk_coords for p in plan.projections()]
    [(0, 0), (0, 1)]
    """

    transform: IndexTransform
    """The composed request this plan partitions."""

    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` if two
        `DimensionMap`s read one input axis (a diagonal, which no selection
        produces), the one shape with no factored form.
        """
        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(self.partition())

    def __iter__(self) -> Iterator[ChunkProjection]:
        """Equivalent to `projections()`."""
        return self.projections()

dimension_grids instance-attribute

dimension_grids: tuple[DimensionGridLike, ...]

One grid per storage dimension, defining the chunk layout the plan walks.

transform instance-attribute

transform: IndexTransform

The composed request this plan partitions.

__init__

__init__(
    transform: IndexTransform,
    dimension_grids: tuple[DimensionGridLike, ...],
) -> None

__iter__

__iter__() -> Iterator[ChunkProjection]

Equivalent to projections().

Source code in src/zarr_indexing/chunk_resolution.py

def __iter__(self) -> Iterator[ChunkProjection]:
    """Equivalent to `projections()`."""
    return self.projections()

partition

partition() -> GridPartition

The plan in factored, columnar form: one table per axis plus a joint table.

Built once per plan and memoized. Raises ValueError if two DimensionMaps read one input axis (a diagonal, which no selection produces), the one shape with no factored form.

Source code in src/zarr_indexing/chunk_resolution.py

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` if two
    `DimensionMap`s read one input axis (a diagonal, which no selection
    produces), the one shape with no factored form.
    """
    cached = self._partition
    if cached is None:
        cached = _partition_transform(self.transform, self.dimension_grids)
        object.__setattr__(self, "_partition", cached)
    return cached

projections

projections() -> Iterator[ChunkProjection]

Return a fresh iterator over the chunks touched by this plan.

Source code in src/zarr_indexing/chunk_resolution.py

def projections(self) -> Iterator[ChunkProjection]:
    """Return a fresh iterator over the chunks touched by this plan."""
    return iter(self.partition())

ChunkProjection dataclass

One source-independent projection of a request through a chunk.

Both transforms share a synthetic input domain. chunk_transform maps that domain to chunk-local storage coordinates; cell_transform maps it to the original request domain.

Attributes:

• chunk_coords (tuple[int, ...]) –

  Coordinates of the selected cell in the caller's grid.

• chunk_domain (IndexDomain) –

  Bounds of that grid cell in global storage coordinates.

• chunk_transform (IndexTransform) –

  Mapping from the shared synthetic domain to chunk-local storage.

• cell_transform (IndexTransform) –

  Mapping from the shared synthetic domain to request coordinates.

• coverage (ChunkCoverage) –

  Whether the request is proven to cover the whole grid cell exactly once. Fancy selections are conservatively "unknown".

Examples:

Row 1 of a (3, 4) array with (2, 2) chunks touches only part of the first chunk, whose domain spans rows [0, 2) and columns [0, 2):

>>> from zarr_indexing import IndexTransform
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
>>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
>>> first = next(iter(plan))
>>> first.chunk_coords
(0, 0)
>>> first.chunk_domain.shape
(2, 2)
>>> first.coverage
'partial'

Source code in src/zarr_indexing/chunk_resolution.py

@dataclass(frozen=True, slots=True)
class ChunkProjection:
    """One source-independent projection of a request through a chunk.

    Both transforms share a synthetic input domain. ``chunk_transform`` maps
    that domain to chunk-local storage coordinates; ``cell_transform`` maps it
    to the original request domain.

    Attributes
    ----------
    chunk_coords
        Coordinates of the selected cell in the caller's grid.
    chunk_domain
        Bounds of that grid cell in global storage coordinates.
    chunk_transform
        Mapping from the shared synthetic domain to chunk-local storage.
    cell_transform
        Mapping from the shared synthetic domain to request coordinates.
    coverage
        Whether the request is proven to cover the whole grid cell exactly
        once. Fancy selections are conservatively ``"unknown"``.

    Examples
    --------
    Row 1 of a `(3, 4)` array with `(2, 2)` chunks touches only part of the
    first chunk, whose domain spans rows `[0, 2)` and columns `[0, 2)`:

    >>> from zarr_indexing import IndexTransform
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
    >>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
    >>> first = next(iter(plan))
    >>> first.chunk_coords
    (0, 0)
    >>> first.chunk_domain.shape
    (2, 2)
    >>> first.coverage
    'partial'
    """

    chunk_coords: tuple[int, ...]
    chunk_domain: IndexDomain
    chunk_transform: IndexTransform
    cell_transform: IndexTransform
    coverage: ChunkCoverage

    def __post_init__(self) -> None:
        if self.chunk_transform.domain != self.cell_transform.domain:
            raise ValueError(
                "chunk_transform and cell_transform must share an input domain; "
                f"got {self.chunk_transform.domain!r} and {self.cell_transform.domain!r}"
            )

cell_transform instance-attribute

cell_transform: IndexTransform

chunk_coords instance-attribute

chunk_coords: tuple[int, ...]

chunk_domain instance-attribute

chunk_domain: IndexDomain

chunk_transform instance-attribute

chunk_transform: IndexTransform

coverage instance-attribute

coverage: ChunkCoverage

__init__

__init__(
    chunk_coords: tuple[int, ...],
    chunk_domain: IndexDomain,
    chunk_transform: IndexTransform,
    cell_transform: IndexTransform,
    coverage: ChunkCoverage,
) -> None

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    if self.chunk_transform.domain != self.cell_transform.domain:
        raise ValueError(
            "chunk_transform and cell_transform must share an input domain; "
            f"got {self.chunk_transform.domain!r} and {self.cell_transform.domain!r}"
        )

GridPartition dataclass

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 n_rows == 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.

Take one 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]]
>>> [projection.chunk_transform.selection_repr for projection in partition][3]
'{ [0, 4) step 2, [0, 1) }'

Source code in src/zarr_indexing/chunk_resolution.py

@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 `n_rows` ``== 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.

    Take one 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]]
    >>> [projection.chunk_transform.selection_repr for projection in partition][3]
    '{ [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."""

    @property
    def n_rows(self) -> int:
        """The number of projections, as an exact integer (`len` raises above the platform limit)."""
        return math.prod(self.row_shape)

    def __len__(self) -> int:
        return self.n_rows

    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 = self.n_rows
        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 __iter__(self) -> Iterator[ChunkProjection]:
        """Materialize every row in order."""
        if self.n_rows == 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:
        domain = self.transform.domain
        if isinstance(axis, StridedSet):
            origin = (
                0 if axis.input_dimension is None else domain.inclusive_min[axis.input_dimension]
            )
            return _strided_piece(axis, row, origin)
        return _indexed_piece(axis, row, domain.ndim, domain.inclusive_min[axis.input_dimension])

    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()
        lo_all = self.transform.domain.inclusive_min
        pieces_per_set = [
            [
                _strided_piece(
                    cast("StridedSet", axis),
                    row,
                    0 if axis.input_dimension is None else lo_all[axis.input_dimension],
                    None if axis.input_dimension is None else slot_of.get(axis.input_dimension),
                )
                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",
        )

dimension_grids instance-attribute

dimension_grids: tuple[DimensionGridLike, ...]

One grid per storage dimension.

joint instance-attribute

joint: JointSet | None

The correlated index arrays' table, or None when there are none.

n_rows property

n_rows: int

The number of projections, as an exact integer (len raises above the platform limit).

row_shape instance-attribute

row_shape: tuple[int, ...]

Rows per table, joint last; the partition is walked in row-major order over it.

sets instance-attribute

sets: tuple[StridedSet | IndexedSet, ...]

Independent per-axis tables, in output-dimension order.

transform instance-attribute

transform: IndexTransform

The transform this partition factors.

__init__

__init__(
    transform: IndexTransform,
    dimension_grids: tuple[DimensionGridLike, ...],
    sets: tuple[StridedSet | IndexedSet, ...],
    joint: JointSet | None,
    row_shape: tuple[int, ...],
) -> None

__iter__

__iter__() -> Iterator[ChunkProjection]

Materialize every row in order.

Source code in src/zarr_indexing/chunk_resolution.py

def __iter__(self) -> Iterator[ChunkProjection]:
    """Materialize every row in order."""
    if self.n_rows == 0:
        return
    if self.joint is None:
        yield from self._iter_factorized()
    else:
        yield from self._iter_correlated()

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return self.n_rows

chunk_coords

chunk_coords() -> ndarray[Any, dtype[intp]]

Chunk coordinates of every row, shape (len(self), output rank), without materializing rows.

Source code in src/zarr_indexing/chunk_resolution.py

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 = self.n_rows
    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

IndexedSet dataclass

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. Columns are read-only NumPy arrays.

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])

Source code in src/zarr_indexing/chunk_resolution.py

@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. Columns are read-only NumPy arrays.

    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 __post_init__(self) -> None:
        _freeze(
            self.chunk,
            self.chunk_start,
            self.chunk_extent,
            self.pointer,
            self.index,
            self.positions,
        )

    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]))

chunk instance-attribute

chunk: ndarray[Any, dtype[intp]]

Chunk index along the axis, one per row, ascending.

chunk_extent instance-attribute

chunk_extent: ndarray[Any, dtype[intp]]

Data extent of each chunk.

chunk_start instance-attribute

chunk_start: ndarray[Any, dtype[intp]]

Storage origin of each chunk.

counts property

counts: ndarray[Any, dtype[intp]]

Entries per row.

index instance-attribute

index: ndarray[Any, dtype[intp]]

Index-array values grouped by chunk.

input_dimension instance-attribute

input_dimension: int

The request axis the index array varies over.

local property

local: ndarray[Any, dtype[intp]]

Chunk-local storage coordinate of every entry, grouped like index.

offset instance-attribute

offset: int

The map's affine offset: storage is offset + stride * index.

output_dimension instance-attribute

output_dimension: int

The storage axis this table describes.

pointer instance-attribute

pointer: ndarray[Any, dtype[intp]]

CSR row pointer: row i owns entries pointer[i] to pointer[i + 1].

positions instance-attribute

positions: ndarray[Any, dtype[intp]]

Positions along the request axis, grouped by chunk, ascending within a row.

stride instance-attribute

stride: int

The map's affine stride.

__init__

__init__(
    output_dimension: int,
    input_dimension: int,
    offset: int,
    stride: int,
    chunk: ndarray[Any, dtype[intp]],
    chunk_start: ndarray[Any, dtype[intp]],
    chunk_extent: ndarray[Any, dtype[intp]],
    pointer: ndarray[Any, dtype[intp]],
    index: ndarray[Any, dtype[intp]],
    positions: ndarray[Any, dtype[intp]],
) -> None

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return int(self.chunk.size)

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    _freeze(
        self.chunk,
        self.chunk_start,
        self.chunk_extent,
        self.pointer,
        self.index,
        self.positions,
    )

run

run(row: int) -> slice

The slice of index / positions a row owns.

Source code in src/zarr_indexing/chunk_resolution.py

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]))

JointSet dataclass

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. Columns are read-only NumPy arrays.

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])

Source code in src/zarr_indexing/chunk_resolution.py

@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. Columns are
    read-only NumPy arrays.

    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 __post_init__(self) -> None:
        _freeze(
            self.chunk,
            self.chunk_start,
            self.chunk_extent,
            self.pointer,
            self.index,
            self.positions,
            self.block_coordinates,
        )

    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]))

block_coordinates instance-attribute

block_coordinates: ndarray[Any, dtype[intp]]

positions unravelled over broadcast_shape, shape (points, len(broadcast_axes)).

broadcast_axes instance-attribute

broadcast_axes: tuple[int, ...]

The request axes the arrays broadcast over.

broadcast_shape instance-attribute

broadcast_shape: tuple[int, ...]

The extent of those axes.

chunk instance-attribute

chunk: ndarray[Any, dtype[intp]]

Chunk coordinates on output_dimensions, shape (rows, k), lexicographic.

chunk_extent instance-attribute

chunk_extent: ndarray[Any, dtype[intp]]

Data extent of each chunk on output_dimensions, shape (rows, k).

chunk_start instance-attribute

chunk_start: ndarray[Any, dtype[intp]]

Storage origin of each chunk on output_dimensions, shape (rows, k).

counts property

counts: ndarray[Any, dtype[intp]]

Points per row.

index instance-attribute

index: ndarray[Any, dtype[intp]]

Index-array values per point and output dimension, shape (points, k).

local property

local: ndarray[Any, dtype[intp]]

Chunk-local storage coordinates of every point, shape (points, k).

offsets instance-attribute

offsets: tuple[int, ...]

Affine offset of each array's map, aligned with output_dimensions.

output_dimensions instance-attribute

output_dimensions: tuple[int, ...]

The storage axes read by correlated index arrays.

pointer instance-attribute

pointer: ndarray[Any, dtype[intp]]

CSR row pointer into index, positions and block_coordinates.

positions instance-attribute

positions: ndarray[Any, dtype[intp]]

Flat block position of each point, ascending within a row.

strides instance-attribute

strides: tuple[int, ...]

Affine stride of each array's map.

__init__

__init__(
    output_dimensions: tuple[int, ...],
    offsets: tuple[int, ...],
    strides: tuple[int, ...],
    broadcast_axes: tuple[int, ...],
    broadcast_shape: tuple[int, ...],
    chunk: ndarray[Any, dtype[intp]],
    chunk_start: ndarray[Any, dtype[intp]],
    chunk_extent: ndarray[Any, dtype[intp]],
    pointer: ndarray[Any, dtype[intp]],
    index: ndarray[Any, dtype[intp]],
    positions: ndarray[Any, dtype[intp]],
    block_coordinates: ndarray[Any, dtype[intp]],
) -> None

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return int(self.chunk.shape[0])

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    _freeze(
        self.chunk,
        self.chunk_start,
        self.chunk_extent,
        self.pointer,
        self.index,
        self.positions,
        self.block_coordinates,
    )

run

run(row: int) -> slice

The slice of the point arrays a row owns.

Source code in src/zarr_indexing/chunk_resolution.py

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]))

StridedSet dataclass

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 positions [origin[i], origin[i] + extent[i]) along the input axis, counted from the request domain's lower bound.

Columns are read-only NumPy arrays. extent and origin are measured along the request axis, whose bounds are arbitrary Python ints; they are intp unless a value does not fit, in which case they hold exact ints (dtype=object).

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])

Source code in src/zarr_indexing/chunk_resolution.py

@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 positions ``[origin[i], origin[i] + extent[i])``
    along the input axis, counted from the request domain's lower bound.

    Columns are read-only NumPy arrays. `extent` and `origin` are measured
    along the request axis, whose bounds are arbitrary Python ints; they are
    ``intp`` unless a value does not fit, in which case they hold exact ints
    (``dtype=object``).

    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]]
    """Position of the row's first cell along the request axis (``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 __post_init__(self) -> None:
        _freeze(
            self.chunk,
            self.chunk_start,
            self.chunk_extent,
            self.local_start,
            self.extent,
            self.origin,
            self.full,
        )

    def __len__(self) -> int:
        return int(self.chunk.size)

chunk instance-attribute

chunk: ndarray[Any, dtype[intp]]

Chunk index along the axis, one per row, ascending.

chunk_extent instance-attribute

chunk_extent: ndarray[Any, dtype[intp]]

Data extent of each chunk (clipped at the array boundary).

chunk_start instance-attribute

chunk_start: ndarray[Any, dtype[intp]]

Storage origin of each chunk.

extent instance-attribute

extent: ndarray[Any, dtype[intp]]

Cells the row selects along the request axis (1 for a constant).

full instance-attribute

full: ndarray[Any, dtype[bool_]]

Whether the row covers its chunk's data extent exactly once, in order.

input_dimension instance-attribute

input_dimension: int | None

The request axis the map reads, or None for a constant.

local_start instance-attribute

local_start: ndarray[Any, dtype[intp]]

Chunk-local storage coordinate of the row's first cell.

origin instance-attribute

origin: ndarray[Any, dtype[intp]]

Position of the row's first cell along the request axis (0 for a constant).

output_dimension instance-attribute

output_dimension: int

The storage axis this table describes.

stride instance-attribute

stride: int

Storage step per request cell; 0 for a constant.

__init__

__init__(
    output_dimension: int,
    input_dimension: int | None,
    stride: int,
    chunk: ndarray[Any, dtype[intp]],
    chunk_start: ndarray[Any, dtype[intp]],
    chunk_extent: ndarray[Any, dtype[intp]],
    local_start: ndarray[Any, dtype[intp]],
    extent: ndarray[Any, dtype[intp]],
    origin: ndarray[Any, dtype[intp]],
    full: ndarray[Any, dtype[bool_]],
) -> None

__len__

__len__() -> int

Source code in src/zarr_indexing/chunk_resolution.py

def __len__(self) -> int:
    return int(self.chunk.size)

__post_init__

__post_init__() -> None

Source code in src/zarr_indexing/chunk_resolution.py

def __post_init__(self) -> None:
    _freeze(
        self.chunk,
        self.chunk_start,
        self.chunk_extent,
        self.local_start,
        self.extent,
        self.origin,
        self.full,
    )

plan_chunks

plan_chunks(
    transform: IndexTransform,
    dimension_grids: Sequence[DimensionGridLike],
) -> ChunkPlan

Plan a transform against a caller-selected chunk grid.

Parameters:

• transform (IndexTransform) –

  Mapping from the request domain to storage coordinates.

• dimension_grids (Sequence[DimensionGridLike]) –

  One storage grid per transform output dimension.

Returns:

• ChunkPlan –

  A reusable plan whose projections are computed lazily; its partition() is the factored form they are derived from.

Examples:

Row 1 of a (3, 4) array with (2, 2) chunks touches the two chunks in the top grid row, each contributing a (2, 2) chunk domain:

>>> from zarr_indexing import IndexTransform
>>> from zarr_indexing.grid import dimension_grids_from_chunks
>>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
>>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
>>> [p.chunk_coords for p in plan]
[(0, 0), (0, 1)]
>>> [p.chunk_domain.shape for p in plan]
[(2, 2), (2, 2)]

Source code in src/zarr_indexing/chunk_resolution.py

def plan_chunks(
    transform: IndexTransform,
    dimension_grids: Sequence[DimensionGridLike],
) -> ChunkPlan:
    """Plan a transform against a caller-selected chunk grid.

    Parameters
    ----------
    transform
        Mapping from the request domain to storage coordinates.
    dimension_grids
        One storage grid per transform output dimension.

    Returns
    -------
    ChunkPlan
        A reusable plan whose projections are computed lazily; its
        `partition()` is the factored form they are derived from.

    Examples
    --------
    Row 1 of a `(3, 4)` array with `(2, 2)` chunks touches the two chunks in
    the top grid row, each contributing a `(2, 2)` chunk domain:

    >>> from zarr_indexing import IndexTransform
    >>> from zarr_indexing.grid import dimension_grids_from_chunks
    >>> grids = dimension_grids_from_chunks((2, 2), shape=(3, 4))
    >>> plan = plan_chunks(IndexTransform.from_shape((3, 4))[1, :], grids)
    >>> [p.chunk_coords for p in plan]
    [(0, 0), (0, 1)]
    >>> [p.chunk_domain.shape for p in plan]
    [(2, 2), (2, 2)]
    """
    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}"
        )
    return ChunkPlan(transform=transform, dimension_grids=grids)