scripts/02_inspect_schema.py

#!/usr/bin/env python3
"""
Discover the ACTUAL schema of MSMARCO-XI before writing any pipeline against it.

Works with polars alone (pyarrow optional), so it runs on hosts without pyarrow.

  python scripts/02_inspect_schema.py --root $VOICERAG_ROOT

Writes $VOICERAG_ROOT/data/schema_report.json, which every downstream script
reads instead of hard-coding column names.

Resolution is three-tier:
  1. exact match against a candidate list (case-insensitive)
  2. heuristic: for `text`, the string field with the longest average value;
     for `text_en`, a field whose name starts with eng/english
  3. report it unresolved, and dump the raw struct so a human can decide
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.schema_utils import detect_layout, iter_passages, default_root  # noqa: E402

FIELD_CANDIDATES: dict[str, list[str]] = {
    "query_id":   ["query_id", "qid", "queryId", "id", "query_ids"],
    "query":      ["query", "question", "Query", "translated_query"],
    "query_en":   ["Eng_Query", "eng_query", "english_query", "query_en", "source_query"],
    "answer":     ["Answer", "answer", "answers", "wellFormedAnswers"],
    "answer_en":  ["Eng_Answer", "eng_answer", "english_answer", "answer_en"],
    "passages":   ["passages", "passage", "contexts", "documents"],
    "query_type": ["query_type", "queryType", "type"],
    "lang":       ["target_lang", "target_language", "tgt_lang", "language", "lang", "tgt"],
    "src_lang":   ["source_lang", "source_language", "src_lang", "source"],
    "meta":       ["meta", "metadata", "info"],
}

PASSAGE_CANDIDATES: dict[str, list[str]] = {
    "text": [
        "Translated_passages", "translated_passages",           # MSMARCO-XI
        "passage_text", "text", "translated_passage", "translated_passage_text",
        "translated_text", "passage", "content", "body", "Passage", "PassageText",
        "passage_text_translated", "tgt_passage_text",
    ],
    "text_en": [
        "English_passages", "english_passages",                  # MSMARCO-XI
        "Eng_passage_text", "eng_passage_text", "Eng_Passage_Text", "english_passage",
        "passage_text_en", "Eng_Passage", "eng_passage", "src_passage_text",
        "source_passage_text", "original_passage_text",
    ],
    "is_selected": ["is_selected", "isSelected", "label", "relevance"],
    "url": ["url", "URL", "source_url", "passage_url", "link"],
}


def find_parquets(root: Path) -> list[Path]:
    hf = root / "hf_cache"
    files = sorted(p for p in hf.rglob("*.parquet") if "MSMARCO" in str(p))
    return files or sorted(hf.rglob("*.parquet"))


def resolve_exact(candidates: dict[str, list[str]], present: set[str]) -> dict[str, str | None]:
    lower = {c.lower(): c for c in present}
    return {logical: next((lower[n.lower()] for n in names if n.lower() in lower), None)
            for logical, names in candidates.items()}


def sample_values(rows: list, key: str, limit: int = 200) -> list[str]:
    """Collect string values for one inner key, for EITHER passages layout.

    Handles struct-of-parallel-lists (values are List(String)) as well as
    list-of-structs (values are String).
    """
    out: list[str] = []
    for row in rows:
        items = row if isinstance(row, (list, tuple)) else [row]
        for it in items:
            if not isinstance(it, dict):
                continue
            v = it.get(key)
            if isinstance(v, str):
                out.append(v)
            elif isinstance(v, (list, tuple)):
                out.extend(x for x in v if isinstance(x, str))
            if len(out) >= limit:
                return out[:limit]
    return out


def heuristic_text_fields(rows: list, inner_keys: set[str],
                          taken: set[str]) -> tuple[str | None, str | None, dict]:
    """Fallback: identify the passage text field by value length, and the English
    counterpart by name prefix. Returns (text, text_en, stats)."""
    stats: dict[str, dict] = {}
    for k in sorted(inner_keys):
        vals = sample_values(rows, k)
        if not vals:
            continue
        stats[k] = {
            "n": len(vals),
            "mean_chars": round(sum(len(v) for v in vals) / len(vals), 1),
            "mean_words": round(sum(len(v.split()) for v in vals) / len(vals), 1),
            "sample": vals[0][:120],
        }
    # long-text fields, longest first
    long_fields = sorted((k for k, s in stats.items() if s["mean_chars"] >= 40),
                         key=lambda k: -stats[k]["mean_chars"])
    long_fields = [k for k in long_fields if k not in taken]
    en_like = [k for k in long_fields if k.lower().startswith(("eng", "english", "src", "source", "original"))]
    non_en = [k for k in long_fields if k not in en_like]
    return (non_en[0] if non_en else None,
            en_like[0] if en_like else None,
            stats)


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=None,
                    help="data root; defaults to $VOICERAG_ROOT or <repo>/../voicerag_data")
    ap.add_argument("--sample-rows", type=int, default=5)
    ap.add_argument("--max-files", type=int, default=30)
    args = ap.parse_args()

    root = (args.root.expanduser().resolve() if args.root else default_root())
    print(f"==> data root: {root}")
    import polars as pl
    try:
        import pyarrow.parquet as pq
    except ImportError:
        pq = None

    files = find_parquets(root)
    if not files:
        print(f"No parquet files under {root/'hf_cache'}. Run 01_download.py first.", file=sys.stderr)
        return 1

    total_bytes = sum(f.stat().st_size for f in files)
    print(f"==> {len(files)} parquet files, {total_bytes/2**30:,.1f} GiB")

    head = files[0]
    print(f"\n{'='*72}\nSCHEMA  ({head.name})\n{'='*72}")
    schema = pl.read_parquet_schema(head)
    for name, dtype in schema.items():
        print(f"  {name:22s} {dtype}")
    top = set(schema)

    mapping = resolve_exact(FIELD_CANDIDATES, top)
    print("\n-- top-level fields --")
    for logical, actual in mapping.items():
        print(f"  {'  ' if actual else '??'} {logical:12s} -> {actual}")

    # ---------------- passages: the part that actually matters ----------------
    pcol = mapping.get("passages")
    passage_map: dict[str, str | None] = {}
    inner_keys: set[str] = set()
    inner_stats: dict = {}
    heuristics_used: list[str] = []

    if pcol:
        print(f"\n{'='*72}\nPASSAGES COLUMN  ('{pcol}')\n{'='*72}")
        sample = pl.read_parquet(head, columns=[pcol], n_rows=max(args.sample_rows, 400))
        rows = sample[pcol].to_list()
        print(f"  dtype: {sample.schema[pcol]}")

        first = rows[0]
        layout = detect_layout(first)
        if isinstance(first, dict):
            inner_keys = set(first.keys())
        elif isinstance(first, (list, tuple)) and first and isinstance(first[0], dict):
            inner_keys = set(first[0].keys())
        print(f"  layout: {layout}")
        if layout == "struct_of_lists":
            n = {k: (len(v) if isinstance(v, (list, tuple)) else 1) for k, v in first.items()}
            print(f"  parallel list lengths in row 0: {n}")
            if len(set(n.values())) != 1:
                print("  !! lists are NOT equal length — positional alignment is unsafe")

        if inner_keys:
            print(f"\n  INNER KEYS: {sorted(inner_keys)}")
            passage_map = resolve_exact(PASSAGE_CANDIDATES, inner_keys)

            taken = {v for v in passage_map.values() if v}
            guess_text, guess_en, inner_stats = heuristic_text_fields(rows, inner_keys, taken)

            print("\n  -- per-key value stats --")
            for k, s in sorted(inner_stats.items(), key=lambda kv: -kv[1]["mean_chars"]):
                print(f"    {k:26s} mean {s['mean_chars']:>7.1f} chars / {s['mean_words']:>6.1f} words")
                print(f"    {'':26s} e.g. {s['sample'][:100]!r}")

            if passage_map.get("text") is None and guess_text:
                passage_map["text"] = guess_text
                heuristics_used.append(f"text <- {guess_text} (longest non-English string field)")
            if passage_map.get("text_en") is None and guess_en:
                passage_map["text_en"] = guess_en
                heuristics_used.append(f"text_en <- {guess_en} (long string field with eng/src prefix)")

            print("\n  -- resolved passage fields --")
            for logical, actual in passage_map.items():
                print(f"    {'  ' if actual else '??'} {logical:12s} -> {actual}")
            if heuristics_used:
                print("\n  !! resolved by HEURISTIC, please eyeball the stats above:")
                for h in heuristics_used:
                    print(f"       {h}")

    # ---------------- per-file profile ----------------
    print(f"\n{'='*72}\nPER-FILE PROFILE\n{'='*72}")
    langcol = mapping.get("lang")
    profile = []
    for f in files[: args.max_files]:
        try:
            n = pq.ParquetFile(f).metadata.num_rows if pq else pl.scan_parquet(f).select(pl.len()).collect().item()
            langs = []
            if langcol:
                try:
                    langs = pl.read_parquet(f, columns=[langcol], n_rows=5000)[langcol].unique().to_list()[:4]
                except Exception:
                    pass
            split = "validation" if "val" in f.name else "train"
            profile.append({"file": f.name, "split": split, "rows": n, "langs": langs})
            print(f"  {f.name:22s} {split:11s} {n:>10,} rows  {langs}")
        except Exception as exc:
            print(f"  {f.name:22s} ERROR {exc}")

    langs_train = {p["langs"][0] for p in profile if p["split"] == "train" and p["langs"]}
    langs_val = {p["langs"][0] for p in profile if p["split"] == "validation" and p["langs"]}
    print(f"\n  train languages      : {len(langs_train)}  {sorted(langs_train)}")
    print(f"  validation languages : {len(langs_val)}  {sorted(langs_val)}")
    only_val = sorted(langs_val - langs_train)
    if only_val:
        print(f"  !! validation-only (NO train split): {only_val}")

    # ---------------- passage lengths ----------------
    print(f"\n{'='*72}\nPASSAGE LENGTH (words)\n{'='*72}")
    lengths: list[int] = []
    if pcol and passage_map.get("text"):
        tkey = passage_map["text"]
        ekey = passage_map.get("text_en")
        skey = passage_map.get("is_selected")
        n_sel = 0
        try:
            for row in pl.read_parquet(head, columns=[pcol], n_rows=3000)[pcol]:
                for _i, txt, _en, sel, _u in iter_passages(row, tkey, ekey, skey, None):
                    if isinstance(txt, str):
                        lengths.append(len(txt.split()))
                        n_sel += int(sel or 0)
            print(f"  labelled positives in sample: {n_sel:,} of {len(lengths):,}")
        except Exception as exc:
            print(f"  could not extract: {exc}")
    if lengths:
        import statistics as st
        lengths.sort()
        q = lambda p: lengths[min(len(lengths) - 1, int(p * len(lengths)))]  # noqa: E731
        print(f"  n={len(lengths):,}  mean={st.mean(lengths):.1f}  median={q(.5)}")
        print(f"  p10={q(.1)} p25={q(.25)} p75={q(.75)} p90={q(.9)} p99={q(.99)} max={lengths[-1]}")
        print(f"\n  >> ~{st.mean(lengths):.0f} words/passage; a 256-token chunk spans "
              f"~{256/(st.mean(lengths)*1.4):.1f} passages")
    else:
        print("  (none extracted — see the passage mapping above)")

    out = root / "data" / "schema_report.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({
        "n_files": len(files),
        "total_bytes": total_bytes,
        # stored relative to --root so the report is portable across hosts
        "files": [str(f.relative_to(root)) if str(f).startswith(str(root))
                  else str(f) for f in files],
        "top_level_columns": sorted(top),
        "field_mapping": mapping,
        "passage_mapping": passage_map,
        "passage_inner_keys": sorted(inner_keys),
        "passage_inner_stats": inner_stats,
        "heuristics_used": heuristics_used,
        "profile": profile,
        "passage_word_lengths_sample": lengths[:5000],
    }, indent=2, ensure_ascii=False))
    print(f"\n==> wrote {out}")

    unresolved = [k for k, v in passage_map.items() if v is None and k in ("text", "text_en")]
    if unresolved:
        print(f"\n!! STILL UNRESOLVED: {unresolved}")
        print("   Send me the INNER KEYS and per-key value stats printed above.")
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())