Skip to content

zarr_indexing.chunk_resolution

zarr_indexing.chunk_resolution

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

Given an IndexTransform (which coordinates a user wants to access) and a ChunkGrid (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 output buffer?

The algorithm is:

  1. Enumerate candidate chunks — determine which chunks could possibly be touched by the transform's output coordinate ranges.

  2. Intersect — for each candidate chunk, call transform.intersect(chunk_domain) to restrict the transform to coordinates within that chunk. If the intersection is empty, skip it.

  3. Translate — shift the restricted transform to chunk-local coordinates via transform.translate(-chunk_origin).

  4. Project — pair the chunk-local storage transform with a transform back to the request's cells. Both use the same compact, zero-origin domain.

Sorted one-dimensional correlated array maps can be partitioned directly because every touched chunk owns a contiguous slice of the index array. That case bypasses candidate enumeration and repeated intersection.

The public result is a lazy, reusable ChunkPlan. Each ChunkProjection is source-independent: it identifies the chunk and expresses both sides of the gather without assuming NumPy selectors, a codec pipeline, or an execution scheduler.

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; iterating either the plan or projections() performs a fresh chunk walk.

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 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
@dataclass(frozen=True, slots=True)
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.

    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."""

    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)

    def __iter__(self) -> Iterator[ChunkProjection]:
        """Equivalent to `projections()`: each iteration performs a fresh chunk walk."""
        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(): each iteration performs a fresh chunk walk.

Source code in src/zarr_indexing/chunk_resolution.py 
160
161
162
def __iter__(self) -> Iterator[ChunkProjection]:
    """Equivalent to `projections()`: each iteration performs a fresh chunk walk."""
    return self.projections()

projections 

projections() -> Iterator[ChunkProjection]

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

Source code in src/zarr_indexing/chunk_resolution.py 
156
157
158
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)

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:

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 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
@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 
120
121
122
123
124
125
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}"
        )

plan_chunks 

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

Plan a transform against a caller-selected chunk grid.

Parameters:

Returns:

  • ChunkPlan –

    A reusable plan whose projections are computed lazily.

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 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
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.

    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)