zarr-python design sketch: ChunkSource — decoupling lazy views from their chunk source (concat/virtual/manifest sources; TensorStore Driver analog)

2026-07-02-chunk-source-design.md

ChunkSource: decoupling lazy views from their chunk source

Version: 1 Date: 2026-07-02 Status: design sketch (pre-implementation; nothing here blocks #3906)

Related:

• Lazy-view chunk layout design ("Layers A–C"): gist f50c93bc5f365e6986dadba6e4aa1902 — this doc is the natural "Layer D".
• Discussion #4082 (L0–L5 ladder; chunk_projections is L4).
• TensorStore drivers: stack, cast, virtual_chunked, downsample.
• VirtualiZarr / kerchunk (manifest-backed arrays) — the nearest-term external constituency.

1. The observation

A lazy array is, abstractly, a pair:

lazy array = (selection, source of chunks)
           = (IndexTransform, ChunkSource)

Everything built in the lazy-indexing work — composition, domain preservation, literal coordinates, chunk resolution — lives on the transform side and never inspects what it is pointed at. The source is only asked one kind of question: "here is a transform whose output coordinates lie in your domain; read (or write) the cells it names."

This is exactly TensorStore's factorization. Every TensorStore object is Driver ∘ IndexTransform — verified empirically in our behavioral oracle, where repr(ts_view) shows precisely {'driver': ..., 'transform': {...}}. TensorStore's drivers form a DAG: storage leaves (zarr3, n5), adapters over other TensorStores (stack = lazy concatenation/overlay, cast, downsample), and synthetic sources (array, virtual_chunked). Lazy concatenation of two arrays is not a feature there; it is just another source.

Today zarr-python fuses the pair: AsyncArray holds _transform plus the tuple (store_path, metadata, codec_pipeline, chunk_grid, config). That tuple is the chunk source — passed as loose arguments through the one seam that already exists:

# src/zarr/core/array.py — the de-facto source interface, today
async def _get_selection_via_transform(
    store_path, metadata, config, transform, codec_pipeline, *,
    prototype, out=None, chunk_grid=None,
) -> NDArrayLikeOrScalar: ...

async def _set_selection_via_transform(...same shape...) -> None: ...

The proposal is reification, not redesign: name that tuple, put the two functions behind a protocol, and let AsyncArray hold (source, transform).

2. Goals / non-goals

Goals:

1. A minimal ChunkSource protocol such that the existing zarr array is one implementation and the lazy machinery is generic over it.
2. Storage-shaped composition: concat/stack, constant/virtual arrays, dtype-cast views, manifest-backed (kerchunk-style) arrays.
3. Keep the planning API (chunk_projections, #4082 L4) as an optional capability that composite sources propagate when they honestly can.
4. Behavior-neutral extraction first; the lazy test suite (oracle-parity, state machines, 2k+ tests) is the safety net.

Non-goals:

• General computation (map_blocks, reductions, rechunking-on-read): that is dask/cubed's graph. TensorStore draws the same line; downsample is as computational as its drivers get, and we defer even that.
• Public API commitment in the first pass — the protocol can prove itself privately (zarr.core) before moving to zarr.abc.
• Any change to the storage format or spec. This is entirely client-side.
• Implicit bounds / resize-aware sources (TensorStore's * bounds) — future.

3. The protocol

Coarse boundary, TensorStore-style: the source receives the whole transform and decomposes internally. (The alternative — generic machinery enumerates chunks and asks the source for one chunk at a time — forces every source to have chunks. Function-backed and constant sources do not.)

@runtime_checkable
class ChunkSource(Protocol):
    """Anything that can answer reads/writes for an index-transformed region.

    The lazy array is (transform, source). The transform's OUTPUT coordinates
    are the source's domain coordinates; its INPUT domain is the caller's
    buffer space and is zero-origin by contract (the Array layer normalizes
    via translate_domain_to before calling, as the I/O paths do today).
    """

    @property
    def domain(self) -> IndexDomain: ...       # extent answerable; origin 0 for leaves
    @property
    def dtype(self) -> ZDType[Any, Any]: ...

    async def read(
        self, transform: IndexTransform, *, prototype: BufferPrototype,
        out: NDBuffer | None = None,
    ) -> NDArrayLikeOrScalar: ...

    async def write(
        self, transform: IndexTransform, values: npt.ArrayLike, *,
        prototype: BufferPrototype,
    ) -> None: ...                             # raises SourceNotWritableError


class ChunkStructured(Protocol):
    """Optional capability: the source has enumerable stored units."""

    def chunk_projections(
        self, transform: IndexTransform, *, unit: Literal["read", "write"] = "read",
    ) -> Iterator[ChunkProjection]: ...

Contract notes:

• transform.output_rank == source.domain.ndim; every output coordinate the transform can produce must lie in source.domain (callers guarantee this — it is what the literal-coordinate bounds checks already enforce).
• Reads with out= fill positionally (input domain = buffer space).
• Write capability is dynamic: write raises SourceNotWritableError (new, subclassing PermissionError) rather than a static flag — mirrors how read-only stores behave, and lets concat-of-mixed-writability fail per-region with a precise message.
• Zarr-specific notions (metadata, chunk_grid, codecs, fill value) are not in the protocol. They remain on the leaf implementation, and the Array surface gates them exactly the way Layer C already gates grid members on views: present for a single-zarr-backed array, LazyViewError-style guard otherwise. The guards were retroactively the right call: a single chunks tuple was never going to survive composite sources.

4. The reference implementation is a rename

ZarrChunkSource = today's tuple, verbatim:

@dataclass(frozen=True)
class ZarrChunkSource:
    store_path: StorePath
    metadata: ArrayMetadata
    codec_pipeline: CodecPipeline
    chunk_grid: ChunkGrid
    config: ArrayConfig

    domain     # IndexDomain.from_shape(metadata.shape)
    read       # body of _get_selection_via_transform
    write      # body of _set_selection_via_transform
    chunk_projections  # body of chunk_partition.iter_chunk_projections

AsyncArray becomes (source, transform) structurally — _with_transform, the _is_identity fast path (transform vs source.domain), .lazy, guards, and translate_by/translate_to are all unchanged. The eager legacy fast path for identity arrays is untouched.

5. The unification: a chunked array is already a composite

iter_chunk_transforms(transform, chunk_grid) is a router: for each grid cell's box, transform.intersect(box) → restricted transform, translate → cell-local coordinates, delegate. A ConcatSource needs exactly those two operations over two irregular members instead of ~N regular cells. The chunk grid is a partition combinator; concat is the same combinator with fewer, fatter cells. One routing algorithm, two levels of the same tree.

6. Worked example: ConcatSource

@dataclass(frozen=True)
class ConcatSource:
    """Lazy concatenation along `axis` (TensorStore's `stack`, restricted)."""
    members: tuple[ChunkSource, ...]     # equal dtype; equal extents off-axis
    axis: int

    @property
    def domain(self) -> IndexDomain:     # off-axis extents shared; axis extents summed
        ...

    def _boxes(self) -> Iterator[tuple[ChunkSource, IndexDomain, tuple[int, ...]]]:
        # each member's box in concat coordinates + the shift into member-local
        offset = 0
        for m in self.members:
            box = member_extent_shifted_along_axis(m.domain, self.axis, offset)
            yield m, box, axis_shift(self.axis, -offset, self.domain.ndim)
            offset += m.domain.shape[self.axis]

    async def read(self, transform, *, prototype, out=None):
        out = out or allocate(transform.domain.shape, self.dtype, prototype)
        for member, box, shift in self._boxes():
            hit = transform.intersect(box)           # same op the chunk router uses
            if hit is None:
                continue
            restricted, out_indices = hit
            local = restricted.translate(shift)      # concat coords -> member coords
            await member.read(local, prototype=prototype,
                              out=out_view_for(out, restricted, out_indices))
        return out

    # write: identical routing, delegate member.write(local, values_view)

Eliding the out_indices bookkeeping (vectorized selections) — which is the point: it is the same bookkeeping _get_selection_via_transform already does per chunk, factored once.

chunk_projections (capability): implemented iff every member is ChunkStructured — delegate per member, shifting each projection's array_selection by the member's output offset. Projection identity across members is an open question (§10).

Acceptance test for the whole design: if ConcatSource needs anything beyond the public protocol — one peek into a member's internals — the boundary is wrong.

7. Source catalog (near-term seriousness ordering)

Source	Backs	read	write	chunk_projections	TensorStore analog
ZarrChunkSource	store+codecs (today's array)	✅	✅	✅	zarr3 driver
ConcatSource	other sources, one axis	✅ route	✅ route	✅ if members are	stack
ManifestSource	chunk manifest → byte ranges in foreign files	✅	❌	✅ (its whole point)	— (VirtualiZarr/kerchunk)
ConstantSource	a fill value	✅	❌	n/a	—
FunctionSource	f(domain_box) -> ndarray	✅	❌	n/a	virtual_chunked
CastSource	one source + dtype conversion	✅	⚠ lossless only	passthrough	cast
DownsampleSource	— deferred: computational	—	—	—	downsample

ManifestSource is the sleeper: it makes the VirtualiZarr/kerchunk pattern a first-class zarr-python citizen instead of a store-level trick, and it needs nothing but the protocol plus byte-range reads the stores already have.

8. Interactions with the existing design

• Layer C guards extend unchanged in spirit: members that require a single backing grid (metadata, chunk_grid, nchunks_initialized, resize, append, info*) are leaf-only; the guard message becomes capability-conditional instead of naming "the backing array".
• chunk_projections stays the public L4 planning API; it simply becomes a capability sources opt into. This is our deliberate divergence from TensorStore (which keeps partitioning driver-internal), and composition makes it more valuable: per-member grids are exactly when a single chunks tuple is meaningless and enumeration is the only honest answer.
• #4040 ChunkLayout is likewise a leaf capability; a composite source has layouts per member, not a layout.
• Effective granularity (#4082 future work) becomes a per-source question — which it always really was (each concat member can have different codecs).

9. Staging

1. Phase 0 — land #3906 as-is. The fusion is fine; the seam is clean.
2. Phase 1 — mechanical extraction of ZarrChunkSource (private, behavior-neutral, zero new features). Existing suite must pass untouched.
3. Phase 2 — ConcatSource behind zarr.experimental, plus the oracle habit: pin its semantics against tensorstore.stack empirically before freezing them (same method that de-risked the domain semantics).
4. Phase 3 — ManifestSource; protocol graduates to zarr.abc once two non-leaf sources have survived contact with it.

10. Open questions

1. Projection identity on composites: tag ChunkProjection with a source path (member=2), or re-key ("1/c/0/3")? Affects distributed planners.
2. Write-through-CastSource: forbid, or allow for lossless casts only?
3. Resize: sources are fixed-extent snapshots for now; a resized backing array invalidates open views. Revisit with implicit bounds.
4. Buffer prototypes: caller-owned (as sketched) or source-negotiated (GPU sources may care)?
5. Store vs source: a Store is key→bytes; a ChunkSource is domain→cells. Manifest blurs this — keep them distinct (a ManifestSource uses stores; it is not one).

11. Why this and not dask

The boundary rule of thumb: a ChunkSource may relocate, reinterpret, or synthesize cells; it may not combine them (no arithmetic across cells). Everything on the allowed list is storage-shaped — it terminates in bytes that already exist (or a constant/function) — so reads stay single-pass, write-through stays meaningful, and chunk_projections stays truthful. The moment cells combine, you need a task graph, and that is someone else's excellent library.