Writers

LazyArray.write(values) and assignment through a view are implemented by write_into, which writes through a transform using only basic integer/slice assignment on the source. Independent affine selections become one basic assignment. Other selections are scattered against the source's write grid: each touched cell is read once, updated in memory, and written back, so storage round trips are bounded by touched cells rather than selected elements. NumPy sources receive one fancy assignment instead, and a source that advertises no grid is written one element at a time without reading.

zarr_indexing.writer

Synchronous writes through coordinate transforms using basic source assignment.

write_into

write_into(
    source: Any,
    transform: IndexTransform,
    values: Any,
    write_grid: Sequence[DimensionGridLike] | None = None,
) -> None

Write broadcast values to precisely the cells addressed by transform.

The source must expose shape, a NumPy-compatible dtype, and basic integer/slice assignment. Values are snapshotted and converted before the first mutation, so source aliases and conversion failures are safe. Extra leading singleton value axes are accepted, as in NumPy assignment. Repeated source coordinates receive the last value in C order of the view.

Independent affine maps use one basic assignment. Other selections are scattered in bulk: a NumPy source receives one fancy assignment, and a readable source with a write_grid is written one grid cell at a time — the cell's touched hull is read once, updated in memory, and written back with one basic assignment, so storage round trips are bounded by touched cells. Without a grid, or for the few transforms the planner cannot factor (such as one input axis feeding two output maps), or for a source that cannot be read, the fallback is one integer assignment per element, which never reads. Backend assignment may itself read storage units as part of updating them.

Source assignment failures propagate and may leave preceding writes applied; this operation is not transactional.

Source code in src/zarr_indexing/writer.py

def write_into(
    source: Any,
    transform: IndexTransform,
    values: Any,
    write_grid: Sequence[DimensionGridLike] | None = None,
) -> None:
    """Write broadcast values to precisely the cells addressed by ``transform``.

    The source must expose ``shape``, a NumPy-compatible ``dtype``, and basic
    integer/slice assignment. Values are snapshotted and converted before the
    first mutation, so source aliases and conversion failures are safe. Extra
    leading singleton value axes are accepted, as in NumPy assignment. Repeated
    source coordinates receive the last value in C order of the view.

    Independent affine maps use one basic assignment. Other selections are
    scattered in bulk: a NumPy source receives one fancy assignment, and a
    readable source with a ``write_grid`` is written one grid cell at a time —
    the cell's touched hull is read once, updated in memory, and written back
    with one basic assignment, so storage round trips are bounded by touched
    cells. Without a grid, or for the few transforms the planner cannot factor
    (such as one input axis feeding two output maps), or for a source that
    cannot be read, the fallback is one integer assignment per element, which
    never reads. Backend assignment may itself read storage units as part of
    updating them.

    Source assignment failures propagate and may leave preceding writes
    applied; this operation is not transactional.
    """
    from zarr_indexing.lazy_array import LazyArray

    shape = transform.domain.shape
    source_shape = tuple(int(extent) for extent in source.shape)
    if transform.output_rank != len(source_shape):
        raise ValueError("transform output rank must match the source rank")
    if isinstance(values, LazyArray):
        # Its NumPy conversion protocol may discard MaskedArray metadata.
        values = values.result()
    if isinstance(source, np.ma.MaskedArray):
        prepared = np.ma.array(values, dtype=source.dtype, copy=True)
    else:
        prepared = np.array(values, dtype=source.dtype, copy=True)
    while prepared.ndim > len(shape) and prepared.shape[0] == 1:
        prepared = prepared[0]
    if isinstance(prepared, np.ma.MaskedArray):
        broadcast = np.ma.array(
            np.broadcast_to(prepared.data, shape),
            mask=np.broadcast_to(np.ma.getmaskarray(prepared), shape),
            copy=False,
        )
    else:
        broadcast = np.broadcast_to(prepared, shape)
    if any(extent == 0 for extent in shape):
        # Values were still converted and broadcast, so a mis-sized RHS is
        # reported even when the selection happens to be empty.
        return

    # Check every map's extremal outputs before any assignment. Python integer
    # arithmetic avoids overflow while computing bounds of affine maps.
    for extent, output in zip(source_shape, transform.output, strict=True):
        if isinstance(output, ConstantMap):
            low = high = output.offset
        else:
            if isinstance(output, DimensionMap):
                axis = output.input_dimension
                first = transform.domain.inclusive_min[axis]
                last = transform.domain.exclusive_max[axis] - 1
            else:
                first = int(output.index_array.min())
                last = int(output.index_array.max())
            endpoints = (
                output.offset + output.stride * first,
                output.offset + output.stride * last,
            )
            low, high = min(endpoints), max(endpoints)
        if low < 0 or high >= extent:
            raise IndexError("transform output coordinates are outside the source shape")

    if _write_affine(source, transform, broadcast):
        return
    if len(source_shape) == 0:
        # A zero-rank source has one cell; no grid or scatter applies.
        source[()] = broadcast[(0,) * len(shape)]
        return
    if isinstance(source, np.ndarray):
        _scatter_numpy(source, transform, broadcast)
        return
    if write_grid is not None and hasattr(source, "__getitem__"):
        try:
            # Materialized before any write: an unfactorable transform is
            # reported by planning, never after some cells were rewritten.
            projections = list(plan_chunks(transform, tuple(write_grid)))
        except (ValueError, NotImplementedError):
            projections = None
        if projections is not None:
            _scatter_planned(source, transform, broadcast, projections)
            return
    _scatter_elementwise(source, transform, broadcast)