SimurgArray/src/simurg/core/schemas.py

"""Canonical columnar data model (plan §5) as Apache Arrow schemas, stored as Parquet.

Schemas carry a ``simurg.schema`` metadata key with a version, so readers can refuse incompatible files.
Parquet files are written without timestamps so identical inputs give byte-identical files (plan §15).
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

from simurg.core.atomic import atomic_open
from simurg.formats.bpm import Manifest

LOCI_SCHEMA_ID = "simurg.loci/1"
INTENSITY_SCHEMA_ID = "simurg.sample_intensity/1"
CONTROLS_SCHEMA_ID = "simurg.sample_controls/1"
NORMALIZED_SCHEMA_ID = "simurg.sample_normalized/1"
TRANSFORMS_SCHEMA_ID = "simurg.sample_transforms/1"

LOCI_SCHEMA = pa.schema(
    [
        ("locus_idx", pa.int32()),
        ("name", pa.string()),
        ("ilmn_id", pa.string()),
        ("chrom", pa.string()),
        ("pos", pa.int64()),
        ("snp", pa.string()),
        ("ilmn_strand", pa.string()),
        ("source_strand", pa.string()),
        ("ref_strand", pa.string()),
        ("address_a", pa.int32()),
        ("address_b", pa.int32()),
        ("assay_type", pa.uint8()),
        ("norm_id", pa.uint8()),
        ("norm_bin_id", pa.uint8()),
        ("norm_lookup", pa.int32()),
        ("exp_clusters", pa.int8()),
        ("intensity_only", pa.bool_()),
        ("ploidy", pa.string()),
        ("genome_build", pa.string()),
    ],
    metadata={"simurg.schema": LOCI_SCHEMA_ID},
)

INTENSITY_SCHEMA = pa.schema(
    [("raw_x", pa.uint16()), ("raw_y", pa.uint16()), ("missing", pa.bool_())],
    metadata={"simurg.schema": INTENSITY_SCHEMA_ID},
)

CONTROLS_SCHEMA = pa.schema(
    [("address", pa.int32()), ("red", pa.uint16()), ("green", pa.uint16())],
    metadata={"simurg.schema": CONTROLS_SCHEMA_ID},
)

NORMALIZED_SCHEMA = pa.schema(
    [("norm_x", pa.float32()), ("norm_y", pa.float32()), ("theta", pa.float32()), ("r", pa.float32())],
    metadata={"simurg.schema": NORMALIZED_SCHEMA_ID},
)

TRANSFORMS_SCHEMA = pa.schema(
    [
        ("bin_index", pa.int32()),
        ("norm_bin_id", pa.uint8()),
        ("n_loci", pa.int32()),
        ("version", pa.int32()),
        ("offset_x", pa.float32()),
        ("offset_y", pa.float32()),
        ("scale_x", pa.float32()),
        ("scale_y", pa.float32()),
        ("shear", pa.float32()),
        ("theta", pa.float32()),
        ("cvx", pa.float32()),
        ("cvy", pa.float32()),
        ("nn12", pa.float32()),
        ("rr12", pa.float32()),
        ("taa", pa.float32()),
        ("tbb", pa.float32()),
    ],
    metadata={"simurg.schema": TRANSFORMS_SCHEMA_ID},
)

_WRITE_OPTS: dict[str, Any] = {"compression": "zstd", "use_dictionary": True, "write_statistics": True}


def _dict_cols(values: list[str]) -> pa.Array:
    return pa.array(values, type=pa.string())


def loci_table(m: Manifest) -> pa.Table:
    n = m.num_loci
    cols = {
        "locus_idx": pa.array(np.arange(n, dtype=np.int32)),
        "name": _dict_cols(m.names),
        "ilmn_id": _dict_cols(m.ilmn_ids),
        "chrom": _dict_cols(m.chroms),
        "pos": pa.array(m.map_infos.astype(np.int64)),
        "snp": _dict_cols(m.snps),
        "ilmn_strand": _dict_cols(m.ilmn_strands),
        "source_strand": _dict_cols(m.source_strands),
        "ref_strand": _dict_cols(m.ref_strands),
        "address_a": pa.array(m.address_a.astype(np.int32)),
        "address_b": pa.array(m.address_b.astype(np.int32)),
        "assay_type": pa.array(m.assay_types.astype(np.uint8)),
        "norm_id": pa.array(m.norm_ids.astype(np.uint8)),
        "norm_bin_id": pa.array(m.norm_bin_ids()),
        "norm_lookup": pa.array(m.normalization_lookups().astype(np.int32)),
        "exp_clusters": pa.array(m.exp_clusters.astype(np.int8)),
        "intensity_only": pa.array(m.intensity_only.astype(np.bool_)),
        "ploidy": _dict_cols(m.ploidies),
        "genome_build": _dict_cols(m.genome_builds),
    }
    meta = {"simurg.schema": LOCI_SCHEMA_ID, "manifest_name": m.manifest_name, "manifest_format": m.format}
    return pa.table(cols, schema=LOCI_SCHEMA.with_metadata(meta))


def with_metadata(table: pa.Table, extra: dict[str, Any]) -> pa.Table:
    meta = dict(table.schema.metadata or {})
    meta.update(
        {k.encode(): (v if isinstance(v, str) else json.dumps(v, sort_keys=True)).encode() for k, v in extra.items()}
    )
    return table.replace_schema_metadata(meta)


def intensity_table(raw_x: np.ndarray, raw_y: np.ndarray, missing: np.ndarray, meta: dict[str, Any]) -> pa.Table:
    t = pa.table(
        {
            "raw_x": pa.array(raw_x.astype(np.uint16)),
            "raw_y": pa.array(raw_y.astype(np.uint16)),
            "missing": pa.array(missing.astype(np.bool_)),
        },
        schema=INTENSITY_SCHEMA,
    )
    return with_metadata(t, meta)


def controls_table(addresses: np.ndarray, red: np.ndarray, green: np.ndarray, meta: dict[str, Any]) -> pa.Table:
    t = pa.table(
        {
            "address": pa.array(addresses.astype(np.int32)),
            "red": pa.array(red.astype(np.uint16)),
            "green": pa.array(green.astype(np.uint16)),
        },
        schema=CONTROLS_SCHEMA,
    )
    return with_metadata(t, meta)


def normalized_table(
    nx: np.ndarray, ny: np.ndarray, theta: np.ndarray, r: np.ndarray, meta: dict[str, Any]
) -> pa.Table:
    t = pa.table(
        {
            "norm_x": pa.array(nx, pa.float32()),
            "norm_y": pa.array(ny, pa.float32()),
            "theta": pa.array(theta, pa.float32()),
            "r": pa.array(r, pa.float32()),
        },
        schema=NORMALIZED_SCHEMA,
    )
    return with_metadata(t, meta)


def transforms_table(rows: list[dict[str, Any]], meta: dict[str, Any]) -> pa.Table:
    t = pa.Table.from_pylist(rows, schema=TRANSFORMS_SCHEMA)
    return with_metadata(t, meta)


def write_parquet(table: pa.Table, path: Path) -> None:
    """Atomic, deterministic Parquet write."""
    with atomic_open(Path(path), "wb") as fh:
        pq.write_table(table, fh, **_WRITE_OPTS)


def read_metadata(path: Path) -> dict[str, str]:
    md = pq.read_schema(path).metadata or {}
    return {k.decode(): v.decode() for k, v in md.items()}


def read_table(path: Path, expect_schema: str | None = None) -> pa.Table:
    t = pq.read_table(path)
    if expect_schema is not None:
        got = (t.schema.metadata or {}).get(b"simurg.schema", b"").decode()
        if got != expect_schema:
            raise ValueError(f"{path}: schema {got!r}, expected {expect_schema!r}")
    return t