fithia2/libs/export/snapshot_export.py

"""Export feature snapshots + labels to Parquet with train/valid/test split."""
from __future__ import annotations

import json
import subprocess
import uuid
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from libs.common.logging import get_logger
from libs.common.time_utils import utc_now

logger = get_logger(__name__)

MANIFEST_FILENAME = "manifest.json"


def _get_git_commit_hash() -> str:
    """Return the current git commit hash (short), or 'unknown'."""
    try:
        result = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        return result.stdout.strip() or "unknown"
    except Exception:
        return "unknown"


def _temporal_split(
    rows: list[dict[str, Any]],
    split_policy: str = "temporal_70_15_15",
) -> dict[str, list[dict[str, Any]]]:
    """Split rows into train/valid/test by event_date (temporal order).

    Args:
        rows: List of row dicts that must have an "event_date" field.
        split_policy: E.g. "temporal_70_15_15" → 70% train, 15% valid, 15% test.

    Returns:
        Dict with keys "train", "valid", "test".
    """
    if not rows:
        return {"train": [], "valid": [], "test": []}

    parts = split_policy.replace("temporal_", "").split("_")
    if len(parts) != 3:
        raise ValueError(f"Invalid split_policy: {split_policy}")
    train_pct, valid_pct, _ = (int(p) for p in parts)

    sorted_rows = sorted(rows, key=lambda r: r.get("event_date", ""))
    n = len(sorted_rows)
    n_train = int(n * train_pct / 100)
    n_valid = int(n * valid_pct / 100)

    return {
        "train": sorted_rows[:n_train],
        "valid": sorted_rows[n_train : n_train + n_valid],
        "test": sorted_rows[n_train + n_valid :],
    }


def _rows_to_table(rows: list[dict[str, Any]]) -> pa.Table:
    """Convert list of dicts to a PyArrow Table."""
    if not rows:
        return pa.table({})
    # Collect all keys
    keys = list(rows[0].keys())
    arrays: dict[str, list[Any]] = {k: [] for k in keys}
    for row in rows:
        for k in keys:
            arrays[k].append(row.get(k))
    return pa.table({k: pa.array(v) for k, v in arrays.items()})


async def export_dataset_snapshot(
    session: AsyncSession,
    snapshot_id: str | None,
    split_policy: str,
    output_dir: str | Path,
    feature_version: str = "market_v1",
    label_version: str = "label-1.0.0",
    parser_version: str = "rule-1.0.0",
) -> dict[str, Any]:
    """Join FeatureSnapshot + EventLabel and export to Parquet.

    Args:
        session: Async DB session.
        snapshot_id: Unique ID for this snapshot (generated if None).
        split_policy: Temporal split policy string (e.g. "temporal_70_15_15").
        output_dir: Root directory for output files.
        feature_version: Snapshot name filter for FeatureSnapshot.
        label_version: Label version filter for EventLabel.
        parser_version: Parser version filter for Event.

    Returns:
        Manifest dict with metadata and row counts.
    """
    from libs.db.models import EventLabel, FeatureSnapshot

    if snapshot_id is None:
        snapshot_id = str(uuid.uuid4())

    out_path = Path(output_dir) / snapshot_id
    out_path.mkdir(parents=True, exist_ok=True)

    # Query: JOIN feature_snapshots + event_labels via event_id
    stmt = (
        select(FeatureSnapshot, EventLabel)
        .join(EventLabel, FeatureSnapshot.event_id == EventLabel.event_id)
        .where(FeatureSnapshot.snapshot_name == feature_version)
        .where(EventLabel.label_version == label_version)
        .where(EventLabel.label_status == "ok")
        .where(EventLabel.invalid_event_for_labeling.is_(False))
    )

    result = await session.execute(stmt)
    pairs = result.all()

    rows: list[dict[str, Any]] = []
    for fs, lbl in pairs:
        row: dict[str, Any] = {
            "event_id": fs.event_id,
            "snapshot_name": fs.snapshot_name,
            "snapshot_version": fs.snapshot_version,
            **fs.feature_json,
            "entry_convention": lbl.entry_convention,
            "reaction_date": lbl.reaction_date.isoformat() if lbl.reaction_date else None,
            "entry_date": lbl.entry_date.isoformat() if lbl.entry_date else None,
            "entry_price": float(lbl.entry_price) if lbl.entry_price else None,
            "fwd_return_1d": float(lbl.fwd_return_1d) if lbl.fwd_return_1d else None,
            "fwd_return_3d": float(lbl.fwd_return_3d) if lbl.fwd_return_3d else None,
            "fwd_return_5d": float(lbl.fwd_return_5d) if lbl.fwd_return_5d else None,
            "hit_pos_1r_within_3d": lbl.hit_pos_1r_within_3d,
            "hit_neg_1r_within_3d": lbl.hit_neg_1r_within_3d,
            "close_up_after_3d": lbl.close_up_after_3d,
            "close_up_after_5d": lbl.close_up_after_5d,
            "mfe_3d": float(lbl.mfe_3d) if lbl.mfe_3d else None,
            "mae_3d": float(lbl.mae_3d) if lbl.mae_3d else None,
            "mfe_5d": float(lbl.mfe_5d) if lbl.mfe_5d else None,
            "mae_5d": float(lbl.mae_5d) if lbl.mae_5d else None,
            "label_status": lbl.label_status,
            "label_version": lbl.label_version,
        }
        if "event_date" not in row:
            row["event_date"] = str(lbl.reaction_date) if lbl.reaction_date else ""
        rows.append(row)

    logger.info("snapshot_export_rows", snapshot_id=snapshot_id, total=len(rows))

    splits = _temporal_split(rows, split_policy)
    row_counts: dict[str, int] = {}

    for split_name, split_rows in splits.items():
        parquet_path = out_path / f"{split_name}.parquet"
        table = _rows_to_table(split_rows)
        pq.write_table(table, str(parquet_path))
        row_counts[split_name] = len(split_rows)
        logger.info("split_written", split=split_name, rows=len(split_rows), path=str(parquet_path))

    manifest: dict[str, Any] = {
        "snapshot_id": snapshot_id,
        "created_at_utc": utc_now().isoformat(),
        "code_commit_hash": _get_git_commit_hash(),
        "feature_version": feature_version,
        "parser_version": parser_version,
        "label_version": label_version,
        "split_policy": split_policy,
        "row_counts": row_counts,
        "total_rows": len(rows),
        "output_dir": str(out_path),
    }

    manifest_path = out_path / MANIFEST_FILENAME
    manifest_path.write_text(json.dumps(manifest, indent=2))
    logger.info("manifest_written", path=str(manifest_path), snapshot_id=snapshot_id)

    return manifest