#!/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 /../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())