zarr_indexing.json

zarr_indexing.json

Lowering between canonical ndsel bodies and in-memory IndexTransforms.

This is the engine layer. Where messages.py is pure JSON→JSON and imposes no array constraints, this module converts a canonical ndsel transform body (spec section 4.3, as produced by zarr_indexing.messages.normalize_ndsel) into the numpy-backed IndexTransform the chunk engine runs on, and back.

Two engine constraints live here and only here:

• Finite bounds. An IndexDomain addresses a finite array, so a canonical body carrying a "-inf"/"+inf" bound cannot be lowered; from_json raises.
• Implicit bounds lower by value. The [n]-bracket implicit/explicit flag is a message-layer concern; the engine keeps only the integer value.

The index_array wire format (and the degenerate-collapse it documents)

ndsel and TensorStore both reject an output map that carries both input_dimension and index_array. The in-memory ArrayMap, however, records an input_dimension to pin the axis an orthogonal (oindex) array varies over. This module bridges the gap:

• On serialize (transform_to_canonical):
• An all-singleton index_array (size 1) selects a single coordinate regardless of input, so it is collapsed to a constant map {offset: offset + stride*value}. The size-1 input dimension stays in the domain, unconsumed — a valid transform. This makes a length-1 oindex selection round-trip behaviorally (an ArrayMap becomes a ConstantMap) rather than by object identity.

• Non-degenerate index_array maps are emitted without input_dimension.

• On load (transform_from_canonical): the in-memory input_dimension is reconstructed from the full-rank array's dependency axes (its non-singleton axes, see transform._array_map_dependency_axes). An array that solely owns a single non-singleton axis is orthogonal (input_dimension = that axis); arrays that share non-singleton axes, or vary over several, are correlated (vindex, input_dimension = None). A single 1-D array over a rank-1 domain is inherently ambiguous between the two flavours; it reconstructs as orthogonal, which is behaviorally identical for the single-array case.

index_transform_to_json / index_transform_from_json (and the *_domain_* variants) are these canonical converters under their historical names.

BoundJSON module-attribute

BoundJSON = int | str | list[IndexValueJSON]

IndexValueJSON module-attribute

IndexValueJSON = int | str

NestedIntList module-attribute

NestedIntList = list[Any]

index_transform_from_json module-attribute

index_transform_from_json = transform_from_canonical

index_transform_to_json module-attribute

index_transform_to_json = transform_to_canonical

IndexDomainJSON

Bases: TypedDict

Canonical JSON representation of an IndexDomain.

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

class IndexDomainJSON(TypedDict, total=False):
    """Canonical JSON representation of an IndexDomain."""

    input_inclusive_min: Required[list[BoundJSON]]
    input_exclusive_max: Required[list[BoundJSON]]
    input_labels: Required[list[str]]

input_exclusive_max instance-attribute

input_exclusive_max: Required[list[BoundJSON]]

input_inclusive_min instance-attribute

input_inclusive_min: Required[list[BoundJSON]]

input_labels instance-attribute

input_labels: Required[list[str]]

IndexTransformJSON

Bases: TypedDict

Canonical JSON representation of an IndexTransform (spec section 4.3).

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

class IndexTransformJSON(TypedDict, total=False):
    """Canonical JSON representation of an IndexTransform (spec section 4.3)."""

    input_rank: Required[int]
    input_inclusive_min: Required[list[BoundJSON]]
    input_exclusive_max: Required[list[BoundJSON]]
    input_labels: Required[list[str]]
    output: Required[list[OutputIndexMapJSON]]

input_exclusive_max instance-attribute

input_exclusive_max: Required[list[BoundJSON]]

input_inclusive_min instance-attribute

input_inclusive_min: Required[list[BoundJSON]]

input_labels instance-attribute

input_labels: Required[list[str]]

input_rank instance-attribute

input_rank: Required[int]

output instance-attribute

output: Required[list[OutputIndexMapJSON]]

OutputIndexMapJSON

Bases: TypedDict

Canonical JSON representation of a single output index map.

Exactly one of three forms (distinguished by which fields are present):

• {"offset": 5} — constant
• {"offset": 0, "stride": 1, "input_dimension": 0} — single_input_dimension
• {"offset": 0, "stride": 1, "index_array": [...], "index_array_bounds": ["-inf", "+inf"]} — index_array

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

class OutputIndexMapJSON(TypedDict, total=False):
    """Canonical JSON representation of a single output index map.

    Exactly one of three forms (distinguished by which fields are present):

    - `{"offset": 5}` — constant
    - `{"offset": 0, "stride": 1, "input_dimension": 0}` — single_input_dimension
    - `{"offset": 0, "stride": 1, "index_array": [...],
       "index_array_bounds": ["-inf", "+inf"]}` — index_array
    """

    offset: int
    stride: int
    input_dimension: int
    index_array: NestedIntList
    index_array_bounds: list[IndexValueJSON]

index_array instance-attribute

index_array: NestedIntList

index_array_bounds instance-attribute

index_array_bounds: list[IndexValueJSON]

input_dimension instance-attribute

input_dimension: int

offset instance-attribute

offset: int

stride instance-attribute

stride: int

index_domain_from_json

index_domain_from_json(
    data: IndexDomainJSON,
) -> IndexDomain

Construct an IndexDomain from its canonical JSON representation.

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

def index_domain_from_json(data: IndexDomainJSON) -> IndexDomain:
    """Construct an IndexDomain from its canonical JSON representation."""
    inclusive_min = tuple(
        _lower_bound(b, f"input_inclusive_min[{i}]")
        for i, b in enumerate(data["input_inclusive_min"])
    )
    exclusive_max = tuple(
        _lower_bound(b, f"input_exclusive_max[{i}]")
        for i, b in enumerate(data["input_exclusive_max"])
    )
    labels = _lower_labels(list(data["input_labels"]))
    return IndexDomain(inclusive_min=inclusive_min, exclusive_max=exclusive_max, labels=labels)

index_domain_to_json

index_domain_to_json(
    domain: IndexDomain,
) -> IndexDomainJSON

Convert an IndexDomain to its canonical JSON representation.

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

def index_domain_to_json(domain: IndexDomain) -> IndexDomainJSON:
    """Convert an IndexDomain to its canonical JSON representation."""
    return {
        "input_inclusive_min": list(domain.inclusive_min),
        "input_exclusive_max": list(domain.exclusive_max),
        "input_labels": _emit_labels(domain.labels, domain.ndim),
    }

output_index_map_from_json

output_index_map_from_json(
    data: OutputIndexMapJSON,
) -> OutputIndexMap

Construct an output index map from its canonical JSON representation.

An index_array map's input_dimension is reconstructed from the array's dependency axes in isolation (single non-singleton axis → orthogonal). The transform-level loader classifies globally; use it when several maps may share axes.

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

def output_index_map_from_json(data: OutputIndexMapJSON) -> OutputIndexMap:
    """Construct an output index map from its canonical JSON representation.

    An `index_array` map's `input_dimension` is reconstructed from the array's
    dependency axes in isolation (single non-singleton axis → orthogonal). The
    transform-level loader classifies globally; use it when several maps may
    share axes.
    """
    if "index_array" in data:
        arr = np.asarray(data["index_array"], dtype=np.intp)
        return ArrayMap(
            index_array=arr,
            offset=data.get("offset", 0),
            stride=data.get("stride", 1),
            input_dimension=_solo_dependency_axis(arr),
        )

    if "input_dimension" in data:
        return DimensionMap(
            input_dimension=data["input_dimension"],
            offset=data.get("offset", 0),
            stride=data.get("stride", 1),
        )

    return ConstantMap(offset=data.get("offset", 0))

output_index_map_to_json

output_index_map_to_json(
    m: OutputIndexMap,
) -> OutputIndexMapJSON

Convert an output index map to its canonical JSON representation.

A degenerate all-singleton ArrayMap collapses to a constant map; a non-degenerate one is emitted without input_dimension (see the module docstring on the wire format).

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

def output_index_map_to_json(m: OutputIndexMap) -> OutputIndexMapJSON:
    """Convert an output index map to its canonical JSON representation.

    A degenerate all-singleton `ArrayMap` collapses to a `constant` map; a
    non-degenerate one is emitted without `input_dimension` (see the module
    docstring on the wire format).
    """
    if isinstance(m, ConstantMap):
        return {"offset": m.offset}

    if isinstance(m, DimensionMap):
        return {"offset": m.offset, "stride": m.stride, "input_dimension": m.input_dimension}

    # m: ArrayMap (OutputIndexMap = ConstantMap | DimensionMap | ArrayMap)
    if m.index_array.size == 1:
        value = int(m.index_array.reshape(-1)[0])
        return {"offset": m.offset + m.stride * value}
    return {
        "offset": m.offset,
        "stride": m.stride,
        "index_array": m.index_array.tolist(),
        "index_array_bounds": ["-inf", "+inf"],
    }

transform_from_canonical

transform_from_canonical(
    data: IndexTransformJSON,
) -> IndexTransform

Construct an IndexTransform from a canonical (or canonicalizable) body.

The body is first run through the message layer (normalize_ndsel) so that omitted fields — identity output, default bounds/labels — are filled and validated, then lowered to the engine representation. index_array maps' input_dimension values are reconstructed by global dependency-axis ownership (see the module docstring).

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

def transform_from_canonical(data: IndexTransformJSON) -> IndexTransform:
    """Construct an IndexTransform from a canonical (or canonicalizable) body.

    The body is first run through the message layer (`normalize_ndsel`) so that
    omitted fields — identity `output`, default bounds/labels — are filled and
    validated, then lowered to the engine representation. `index_array` maps'
    `input_dimension` values are reconstructed by global dependency-axis
    ownership (see the module docstring).
    """
    body = normalize_ndsel({"kind": "transform", **data})

    inclusive_min = tuple(
        _lower_bound(b, f"input_inclusive_min[{i}]")
        for i, b in enumerate(body["input_inclusive_min"])
    )
    exclusive_max = tuple(
        _lower_bound(b, f"input_exclusive_max[{i}]")
        for i, b in enumerate(body["input_exclusive_max"])
    )
    domain = IndexDomain(
        inclusive_min=inclusive_min,
        exclusive_max=exclusive_max,
        labels=_lower_labels(body["input_labels"]),
    )

    output_raw: list[dict[str, Any]] = body["output"]

    # Classify index_array maps globally: an axis owned by exactly one array map
    # (and the map's sole non-singleton axis) marks that map orthogonal; shared
    # or multiple non-singleton axes mark the maps correlated (vindex).
    array_axes: dict[int, tuple[int, ...]] = {}
    axis_owners: Counter[int] = Counter()
    for i, om in enumerate(output_raw):
        if "index_array" in om:
            arr = np.asarray(om["index_array"], dtype=np.intp)
            dep = _array_map_dependency_axes(arr)
            array_axes[i] = dep
            axis_owners.update(dep)

    output: list[OutputIndexMap] = []
    for i, om in enumerate(output_raw):
        if "index_array" in om:
            dep = array_axes[i]
            input_dim = dep[0] if len(dep) == 1 and axis_owners[dep[0]] == 1 else None
            output.append(
                ArrayMap(
                    index_array=np.asarray(om["index_array"], dtype=np.intp),
                    offset=om.get("offset", 0),
                    stride=om.get("stride", 1),
                    input_dimension=input_dim,
                )
            )
        elif "input_dimension" in om:
            output.append(
                DimensionMap(
                    input_dimension=om["input_dimension"],
                    offset=om.get("offset", 0),
                    stride=om.get("stride", 1),
                )
            )
        else:
            output.append(ConstantMap(offset=om.get("offset", 0)))

    return IndexTransform(domain=domain, output=tuple(output))

transform_to_canonical

transform_to_canonical(
    transform: IndexTransform,
) -> IndexTransformJSON

Convert an IndexTransform to its canonical ndsel transform body.

The result is fully explicit (spec section 4.3): input_rank, fully written bounds and labels, and an explicit output with offset/stride present on every affine and array map.

Source code in packages/zarr-indexing/src/zarr_indexing/json.py

def transform_to_canonical(transform: IndexTransform) -> IndexTransformJSON:
    """Convert an IndexTransform to its canonical ndsel transform body.

    The result is fully explicit (spec section 4.3): `input_rank`, fully written
    bounds and labels, and an explicit `output` with `offset`/`stride` present
    on every affine and array map.
    """
    return {
        "input_rank": transform.domain.ndim,
        "input_inclusive_min": list(transform.domain.inclusive_min),
        "input_exclusive_max": list(transform.domain.exclusive_max),
        "input_labels": _emit_labels(transform.domain.labels, transform.domain.ndim),
        "output": [output_index_map_to_json(m) for m in transform.output],
    }