SniperGold_ML/engine/storage.py

"""Persistent canonical storage (Layer 6).

Parquet with zstd level 3 and pinned pyarrow writer (spec 12.1, determinism
pins in spec 26). Physical Parquet bytes are integrity evidence only; the
logical content-id contract lives in engine/canonical.py.

Writer/reader version is pinned to pyarrow 25.0.0 (env check performed at
import time).
"""

import os

import pyarrow as pa
import pyarrow.parquet as pq

from . import versions

_PYARROW_PIN = versions.PYARROW_VERSION_PIN
_PA_ACTUAL = pa.__version__
if _PA_ACTUAL != _PYARROW_PIN:
    raise RuntimeError("pyarrow must be %s (pinned), found %s"
                       % (_PYARROW_PIN, _PA_ACTUAL))

TICK_SCHEMA = pa.schema([
    ("ts_ms", pa.int64()),
    ("bid_u", pa.int64()),
    ("ask_u", pa.int64()),
    ("spread_u", pa.int64()),
    ("vol", pa.int64()),
    ("src_line", pa.int64()),
    ("rec_ord", pa.int64()),
])

BAR_SCHEMA = pa.schema([
    ("bar_idx", pa.int64()),
    ("period_id", pa.int64()),
    ("start_ms", pa.int64()),
    ("end_ms", pa.int64()),
    ("open_u2", pa.int64()),
    ("high_u2", pa.int64()),
    ("low_u2", pa.int64()),
    ("close_u2", pa.int64()),
    ("tick_count", pa.int64()),
    ("vol_sum", pa.int64()),
    ("spread_min_u", pa.int64()),
    ("spread_max_u", pa.int64()),
    ("spread_sum_u", pa.int64()),
    ("spread_avg_u", pa.int64()),
    ("is_final", pa.bool_()),
    ("first_src_line", pa.int64()),
    ("last_src_line", pa.int64()),
])

TICK_FIELDS = [f.name for f in TICK_SCHEMA]
BAR_FIELDS = [f.name for f in BAR_SCHEMA]

# Canonical directory layout (spec 12.2)
def cert_path(out_root):
    return os.path.join(out_root, "certification", "source_certificate.json")


def chunkmap_path(out_root, source_id):
    return os.path.join(out_root, "chunkmap", "chunkmap-%s.json" % source_id)


def ticks_chunk_path(out_root, index):
    return os.path.join(out_root, "ticks", "chunk_%06d.parquet" % index)


def ticks_staging_path(out_root, index):
    return os.path.join(out_root, "staging", "chunk_%06d.parquet.tmp" % index)


def preservation_chunk_path(out_root, index):
    """Source-preservation sidecar for the G_TICKSTORY_MT5 ``last`` column
    (P3-DE-004 spec 8.6); dotted-grammar runs produce no sidecars."""
    return os.path.join(out_root, "source_preservation",
                        "chunk_%06d.jsonl" % index)



def malformed_chunk_path(out_root, index):
    return os.path.join(out_root, "malformed", "malformed_%06d.jsonl" % index)


def bar_part_path(out_root, timeframe, workload_index):
    return os.path.join(out_root, "bars", timeframe,
                        "wl_%03d.parquet" % workload_index)


def bar_dir(out_root, timeframe):
    return os.path.join(out_root, "bars", timeframe)


def state_dir(out_root):
    return os.path.join(out_root, "state")


def checkpoint_path(out_root):
    return os.path.join(out_root, "state", "checkpoint.json")


def journal_path(out_root):
    return os.path.join(out_root, "state", "commits.jsonl")


def lock_path(out_root):
    return os.path.join(out_root, "state", "lock")


def control_path(out_root):
    return os.path.join(out_root, "state", "control.json")


def progress_path(out_root):
    return os.path.join(out_root, "progress.json")


def evidence_path(out_root):
    return os.path.join(out_root, "evidence", "evidence.json")


def run_complete_path(out_root):
    return os.path.join(out_root, "RUN_COMPLETE.json")


def datasets_dir(out_root):
    return os.path.join(out_root, "datasets")


def verify_dir(out_root):
    return os.path.join(out_root, "verification")


def logs_dir(out_root):
    return os.path.join(out_root, "logs")


def _with_meta(schema, meta):
    md = dict(schema.metadata) if schema.metadata else {}
    md.update({k.encode(): str(v).encode() for k, v in meta.items()})
    return schema.with_metadata(md)


def _tick_meta(cfg, source_id, chunk_index):
    return {
        "schema_version": versions.SCHEMA_TICKS,
        "engine_version": versions.ENGINE_VERSION,
        "parser_version": versions.PARSER_VERSION,
        "algorithm_version": versions.ALGORITHM_VERSION,
        "source_id": source_id,
        "chunk_index": str(chunk_index),
        "price_scale": str(versions.PRICE_SCALE),
        "pyarrow_version": str(_PA_ACTUAL),
    }


def _bar_meta(cfg, timeframe, workload_index):
    return {
        "schema_version": versions.SCHEMA_BARS,
        "engine_version": versions.ENGINE_VERSION,
        "parser_version": versions.PARSER_VERSION,
        "algorithm_version": versions.ALGORITHM_VERSION,
        "timeframe": timeframe,
        "price_scale_bar": str(versions.PRICE_SCALE_BAR),
        "pyarrow_version": str(_PA_ACTUAL),
    }


def write_ticks_chunk(path, records, cfg, source_id, chunk_index):
    """Write canonical tick records (7-tuples in CTS_V1 order) to parquet."""
    arrays = [pa.array([r[i] for r in records], pa.int64()) for i in range(len(TICK_FIELDS))]
    table = pa.Table.from_arrays(arrays, schema=TICK_SCHEMA)
    table = table.replace_schema_metadata(_tick_meta(cfg, source_id, chunk_index))
    pq.write_table(table, path, compression="zstd", compression_level=3)
    return len(records)


def read_ticks_chunk(path):
    """Read canonical tick parquet into a list of 7-tuples in row order."""
    table = pq.read_table(path, columns=TICK_SCHEMA.names)
    cols = [table.column(i).to_pylist() for i in range(len(TICK_FIELDS))]
    rows = [tuple(c[i] for c in cols) for i in range(len(cols[0]))]
    return rows


def write_bar_part(path, rows, timeframe, workload_index, cfg):
    """Write CBS_V1 bar rows (tuples in BAR_SCHEMA order)."""
    n = len(BAR_FIELDS)
    arrays = []
    for i, f in enumerate(BAR_SCHEMA):
        if f.type == pa.bool_():
            arrays.append(pa.array([r[i] for r in rows], pa.bool_()))
        else:
            arrays.append(pa.array([r[i] for r in rows], pa.int64()))
    table = pa.Table.from_arrays(arrays, schema=BAR_SCHEMA)
    table = table.replace_schema_metadata(_bar_meta(cfg, timeframe, workload_index))
    pq.write_table(table, path, compression="zstd", compression_level=3)
    return len(rows)


def read_bar_part(path):
    table = pq.read_table(path, columns=BAR_SCHEMA.names)
    cols = [table.column(i).to_pylist() for i in range(len(BAR_FIELDS))]
    rows = [tuple(c[i] for c in cols) for i in range(len(cols[0]))]
    return rows