Posiciones y evaluaciones: el join que no descarga 42 GB

"""Positions from the UCI games: normalized 4-field FEN, phase, deduplicated with a count.

``fen4`` is ``pieces turn castling en-passant`` (the key of
``Lichess/chess-position-evaluations``); move counters are dropped so the same position from
different games collapses into one row with ``n_seen`` occurrences. The first occurrence wins
the ``game_id``/``ply``/``last_move``/``result`` columns.
"""

from __future__ import annotations

from collections.abc import Iterator
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq
from pydantic import Field

from rukh.config import BaseConfig
from rukh.data.db import DuckDbConfig
from rukh.data.db import connect as db_connect
from rukh.data.manifest import FileHash, Manifest
from rukh.data.parallel import run_batches
from rukh.data.uci import resolve_workers, sha256_file
from rukh.paths import resolve

OPENING_MAX_PLY = 10
MIDDLEGAME_MIN_PIECES = 14
BATCH_GAMES = 5_000
POSITIONS_FILE = "positions.parquet"
PARTS_DIR = "parts"
POSITION_SCHEMA = pa.schema(
    [
        ("seq", pa.int64()),
        ("fen4", pa.string()),
        ("game_id", pa.int64()),
        ("ply", pa.int16()),
        ("last_move", pa.string()),
        ("result", pa.string()),
        ("phase", pa.string()),
    ]
)

class PositionsConfig(BaseConfig):
    """Which games to walk and how many distinct positions to keep."""

    uci_dir: str = "data/uci"
    month: str = "2025-01"
    out_dir: str = "data/positions"
    n_games: int = Field(default=300_000, ge=1)
    max_positions: int = Field(default=5_000_000, ge=1)
    workers: int = Field(default=0, ge=0)
    duckdb: DuckDbConfig = Field(default_factory=DuckDbConfig)

def fen4(board: object) -> str:
    """First four FEN fields, built without the move counters python-chess would format.

    The en-passant square is written only when the capture is legal, which is the convention
    of ``Lichess/chess-position-evaluations``.
    """
    import chess

    ep = chess.SQUARE_NAMES[board.ep_square] if board.has_legal_en_passant() else "-"  # type: ignore[attr-defined]
    turn = "w" if board.turn else "b"  # type: ignore[attr-defined]
    return f"{board.board_fen()} {turn} {board.castling_xfen()} {ep}"  # type: ignore[attr-defined]


def phase(ply: int, n_pieces: int) -> str:
    """``opening`` up to ply 10, ``middlegame`` with 14+ pieces afterwards, else ``endgame``."""
    if ply <= OPENING_MAX_PLY:
        return "opening"
    return "middlegame" if n_pieces >= MIDDLEGAME_MIN_PIECES else "endgame"

def walk_game(game_id: int, uci: str, result: str) -> Iterator[tuple[str, int, int, str, str, str]]:
    """Yield ``(fen4, game_id, ply, last_move, result, phase)`` after every move."""
    import chess

    board = chess.Board()
    for ply, token in enumerate(uci.split(), start=1):
        move = chess.Move.from_uci(token)
        board.push(move)
        yield fen4(board), game_id, ply, token, result, phase(ply, chess.popcount(board.occupied))


def walk_rows(rows: list[tuple[int, str, str]]) -> dict[str, list[object]]:
    """Walk a batch of ``(game_id, uci, result)``; module-level for the spawn pool."""
    out: dict[str, list[object]] = {name: [] for name in POSITION_SCHEMA.names if name != "seq"}
    for game_id, uci, result in rows:
        for fen, gid, ply, last_move, res, ph in walk_game(game_id, uci, result):
            out["fen4"].append(fen)
            out["game_id"].append(gid)
            out["ply"].append(ply)
            out["last_move"].append(last_move)
            out["result"].append(res)
            out["phase"].append(ph)
    return out

def iter_games(games: Path, n_games: int) -> Iterator[list[tuple[int, str, str]]]:
    """Batches of ``(game_id, uci, result)``, read from the parquet in row groups."""
    reader = pq.ParquetFile(games)
    remaining = n_games
    for batch in reader.iter_batches(batch_size=BATCH_GAMES, columns=["game_id", "uci", "result"]):
        if remaining <= 0:
            return
        rows = list(
            zip(
                batch.column("game_id").to_pylist(),
                batch.column("uci").to_pylist(),
                batch.column("result").to_pylist(),
                strict=True,
            )
        )[:remaining]
        remaining -= len(rows)
        yield rows

def write_parts(games: Path, parts_dir: Path, cfg: PositionsConfig) -> int:
    """Walk the games into ``parts_dir/part-NNNN.parquet`` (all positions, not deduplicated)."""
    parts_dir.mkdir(parents=True, exist_ok=True)
    for old in parts_dir.glob("part-*.parquet"):
        old.unlink()
    workers = resolve_workers(cfg.workers)
    batches = iter_games(games, cfg.n_games)
    return _write_parts(parts_dir, run_batches(walk_rows, batches, workers), 0)


def _write_parts(parts_dir: Path, results: Iterator[dict[str, list[object]]], seq: int) -> int:
    for index, columns in enumerate(results):
        n = len(columns["fen4"])
        table = pa.table({"seq": list(range(seq, seq + n)), **columns}, schema=POSITION_SCHEMA)
        pq.write_table(table, parts_dir / f"part-{index:04d}.parquet", compression="zstd")
        seq += n
    return seq

def dedupe(
    parts_dir: Path, out: Path, max_positions: int, duckdb_cfg: DuckDbConfig | None = None
) -> int:
    """Group the parts by ``fen4`` (first occurrence wins, ``n_seen`` counted) with DuckDB."""
    glob = (parts_dir / "part-*.parquet").as_posix()
    con = db_connect(duckdb_cfg)
    try:
        con.execute(
            f"""
            COPY (
              SELECT fen4,
                     arg_min(game_id, seq) AS game_id,
                     arg_min(ply, seq) AS ply,
                     arg_min(last_move, seq) AS last_move,
                     arg_min(result, seq) AS result,
                     arg_min(phase, seq) AS phase,
                     count(*) AS n_seen
              FROM read_parquet('{glob}')
              GROUP BY fen4
              ORDER BY min(seq)
              LIMIT {int(max_positions)}
            ) TO '{out.as_posix()}' (FORMAT PARQUET, COMPRESSION ZSTD)
            """
        )
        row = con.execute(f"SELECT count(*) FROM read_parquet('{out.as_posix()}')").fetchone()
    finally:
        con.close()
    return int(row[0]) if row else 0

def month_parquet(uci_dir: Path, month: str) -> Path:
    year, mm = month.split("-")
    return uci_dir / f"year={year}" / f"month={mm}" / "games.parquet"


def run(cfg: PositionsConfig) -> Manifest:
    """Walk up to ``n_games`` of ``month`` and write ``out_dir/positions.parquet`` + manifest."""
    games = month_parquet(resolve(cfg.uci_dir), cfg.month)
    out_dir = resolve(cfg.out_dir)
    parts_dir = out_dir / PARTS_DIR
    n_positions = write_parts(games, parts_dir, cfg)
    out = out_dir / POSITIONS_FILE
    n_distinct = dedupe(parts_dir, out, cfg.max_positions, cfg.duckdb)
    manifest = Manifest(
        dataset="Lichess/standard-chess-games",
        months=[cfg.month],
        filters={
            "n_games": cfg.n_games,
            "max_positions": cfg.max_positions,
            "phase": {
                "opening_max_ply": OPENING_MAX_PLY,
                "middlegame_min_pieces": MIDDLEGAME_MIN_PIECES,
            },
        },
        counts={"positions": n_positions, "distinct": n_distinct},
        files=[FileHash(path=POSITIONS_FILE, sha256=sha256_file(out), bytes=out.stat().st_size)],
    )
    (out_dir / "manifest.json").write_text(manifest.model_dump_json(indent=2) + "\n", "utf-8")
    return manifest

"""One scoring convention for every multi-PV consumer (consolidation and DPO pairs).

``cp`` and ``mate`` in ``Lichess/chess-position-evaluations`` are from White's point of view.
A mate is worth ``MATE_SCORE - mate`` (so a mate in 1 beats a mate in 5) and a line is judged
from the side to move, which negates the score for Black.

Lines are ranked best first by: a mate *for the side to move*, then depth, then the score for
the side to move, then the move string as a tie-break. ``evals`` consolidates with the SQL
expressions below and ``pairs`` ranks with ``line_rank_key``; both orderings must stay
identical, which is why they live in one module.
"""

from __future__ import annotations

MATE_SCORE = 10_000
BEST_LINE_DOC = "deepest; mate for the mover first; then best cp for the side to move"

# SQL counterparts used by ``rukh.data.evals`` (``sign`` is ``mover_sign`` of the FEN turn).
SCORE_WHITE_SQL = f"""
CASE WHEN mate IS NOT NULL
     THEN CASE WHEN mate > 0 THEN {MATE_SCORE} - mate ELSE -{MATE_SCORE} - mate END
     ELSE cp END
"""
RANK_ORDER_SQL = (
    "(mate IS NOT NULL AND sign * mate > 0) DESC, depth DESC, sign * score_white DESC, move"
)

def mover_sign(turn: str) -> int:
    """``+1`` when White is to move, ``-1`` for Black."""
    return 1 if turn == "w" else -1


def score_white(cp: int | None, mate: int | None) -> int | None:
    """Centipawns from White's point of view; a mate in ``n`` is ``+/-(10000 - n)``."""
    if mate is not None:
        return MATE_SCORE - mate if mate > 0 else -MATE_SCORE - mate
    return None if cp is None else int(cp)


def score_mover(cp: int | None, mate: int | None, turn: str) -> int | None:
    """``score_white`` seen from the side to move: unchanged for White, negated for Black."""
    white = score_white(cp, mate)
    return None if white is None else mover_sign(turn) * white


def line_rank_key(
    cp: int | None, mate: int | None, depth: int | None, move: str, turn: str
) -> tuple[bool, int, int, str]:
    """Sort key (ascending) reproducing ``RANK_ORDER_SQL``: the best line sorts first."""
    sign = mover_sign(turn)
    white = score_white(cp, mate)
    if white is None:
        raise ValueError("a line without cp and without mate cannot be ranked")
    mates_for_mover = mate is not None and sign * mate > 0
    # DuckDB orders NULL depths last, so an unknown depth ranks below any known one.
    return (not mates_for_mover, -(depth if depth is not None else -1), -sign * white, move)

"""Join the sampled positions with ``Lichess/chess-position-evaluations`` (streaming, resumable).

Each remote file ``data_00NN.parquet`` is semi-joined with the local positions by ``fen`` and
written as ``data/evals/part-NN.parquet``; an existing part is skipped, so an interrupted run
resumes where it stopped. ``consolidate`` then keeps one row per FEN with the best line and
every line as ``pvs``; the ordering and the mate convention come from ``rukh.data.scoring``,
shared with ``rukh.data.pairs``.
"""

from __future__ import annotations

from pathlib import Path

from pydantic import Field

from rukh.config import BaseConfig
from rukh.data.db import DuckDbConfig, connect
from rukh.data.manifest import FileHash, Manifest
from rukh.data.scoring import BEST_LINE_DOC, MATE_SCORE, RANK_ORDER_SQL, SCORE_WHITE_SQL
from rukh.data.uci import sha256_file
from rukh.paths import resolve

EVAL_FILE = "positions-eval.parquet"

__all__ = ["MATE_SCORE", "EvalsConfig", "consolidate", "fetch_part", "parse_files", "run"]


class EvalsConfig(BaseConfig):
    """Remote dataset, local positions and the output directory."""

    dataset: str = "Lichess/chess-position-evaluations"
    remote_pattern: str = "hf://datasets/{dataset}/data/data_{index:04d}.parquet"
    n_files: int = Field(default=20, ge=1)
    positions: str = "data/positions/positions.parquet"
    out_dir: str = "data/evals"
    duckdb: DuckDbConfig = Field(default_factory=DuckDbConfig)

def parse_files(spec: str, n_files: int) -> list[int]:
    """``"0-19"``, ``"3"`` or ``"0,2,5-7"`` into a sorted list of file indices."""
    indices: set[int] = set()
    for part in spec.split(","):
        part = part.strip()
        if not part:
            continue
        if "-" in part:
            lo, hi = part.split("-", 1)
            indices.update(range(int(lo), int(hi) + 1))
        else:
            indices.add(int(part))
    bad = [i for i in indices if i < 0 or i >= n_files]
    if bad:
        raise ValueError(f"file index out of range 0-{n_files - 1}: {bad}")
    return sorted(indices)


def _source(cfg: EvalsConfig, index: int) -> str:
    """Remote parquet for file ``index`` (tests monkeypatch this to a local file)."""
    return cfg.remote_pattern.format(dataset=cfg.dataset, index=index)


def part_path(out_dir: Path, index: int) -> Path:
    return out_dir / f"part-{index:02d}.parquet"

def fetch_part(cfg: EvalsConfig, index: int, positions: Path, out_dir: Path) -> tuple[Path, bool]:
    """Semi-join one remote file with the positions; ``(path, fetched)`` (False = existed)."""
    target = part_path(out_dir, index)
    if target.is_file():
        return target, False
    tmp = target.with_suffix(".parquet.tmp")
    con = connect(cfg.duckdb)
    try:
        con.execute(
            f"""
            COPY (
              SELECT e.fen, e.line, e.depth, e.knodes, e.cp, e.mate
              FROM read_parquet('{_source(cfg, index)}') e
              SEMI JOIN read_parquet('{positions.as_posix()}') p ON e.fen = p.fen4
            ) TO '{tmp.as_posix()}' (FORMAT PARQUET, COMPRESSION ZSTD)
            """
        )
    finally:
        con.close()
    tmp.replace(target)
    return target, True

def consolidate(
    out_dir: Path, positions: Path, duckdb_cfg: DuckDbConfig | None = None
) -> dict[str, int]:
    """Best line per FEN plus ``pvs``, joined back with the positions' metadata."""
    parts = (out_dir / "part-*.parquet").as_posix()
    target = out_dir / EVAL_FILE
    con = connect(duckdb_cfg)
    try:
        con.execute(
            f"""
            CREATE TEMP TABLE lines AS
            SELECT fen,
                   split_part(line, ' ', 1) AS move,
                   depth, knodes, cp, mate,
                   CASE WHEN split_part(fen, ' ', 2) = 'w' THEN 1 ELSE -1 END AS sign,
                   {SCORE_WHITE_SQL} AS score_white
            FROM read_parquet('{parts}')
            WHERE line IS NOT NULL AND line <> ''
            """
        )
        con.execute(
            f"""
            CREATE TEMP TABLE ranked AS
            SELECT *, sign * score_white AS score_mover,
                   row_number() OVER (PARTITION BY fen ORDER BY {RANK_ORDER_SQL}) AS rk
            FROM lines
            """
        )
        con.execute(
            f"""
            COPY (
              SELECT p.fen4 AS fen, p.game_id, p.ply, p.last_move, p.result, p.phase, p.n_seen,
                     b.move AS best_move, b.cp, b.mate, b.depth, b.knodes,
                     pv.pvs
              FROM read_parquet('{positions.as_posix()}') p
              JOIN ranked b ON b.fen = p.fen4 AND b.rk = 1
              JOIN (
                SELECT fen,
                       list(struct_pack(move := move, cp := cp, mate := mate, depth := depth)
                            ORDER BY rk) AS pvs
                FROM ranked GROUP BY fen
              ) pv ON pv.fen = p.fen4
              ORDER BY p.game_id, p.ply
            ) TO '{target.as_posix()}' (FORMAT PARQUET, COMPRESSION ZSTD)
            """
        )
        total = con.execute(
            f"SELECT count(*) FROM read_parquet('{positions.as_posix()}')"
        ).fetchone()
        covered = con.execute(
            f"SELECT count(*) FROM read_parquet('{target.as_posix()}')"
        ).fetchone()
        n_lines = con.execute("SELECT count(*) FROM lines").fetchone()
    finally:
        con.close()
    return {
        "positions": int(total[0]) if total else 0,
        "evaluated": int(covered[0]) if covered else 0,
        "lines": int(n_lines[0]) if n_lines else 0,
    }

def run(cfg: EvalsConfig, files: str | None = None) -> Manifest:
    """Fetch the requested parts (skipping existing ones), consolidate and write the manifest."""
    positions = resolve(cfg.positions)
    out_dir = resolve(cfg.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    indices = parse_files(files or f"0-{cfg.n_files - 1}", cfg.n_files)
    fetched: list[int] = []
    for index in indices:
        _, was_fetched = fetch_part(cfg, index, positions, out_dir)
        if was_fetched:
            fetched.append(index)
    counts = consolidate(out_dir, positions, cfg.duckdb)
    present = sorted(int(p.stem.split("-")[1]) for p in out_dir.glob("part-*.parquet"))
    target = out_dir / EVAL_FILE
    manifest = Manifest(
        dataset=cfg.dataset,
        months=[],
        filters={
            "positions": Path(cfg.positions).as_posix(),
            "files_present": present,
            "files_fetched_now": fetched,
            "n_files": cfg.n_files,
            "coverage": round(counts["evaluated"] / counts["positions"], 4)
            if counts["positions"]
            else 0.0,
            "best_line": BEST_LINE_DOC,
            "mate_score": MATE_SCORE,
        },
        counts=counts,
        files=[FileHash(path=EVAL_FILE, sha256=sha256_file(target), bytes=target.stat().st_size)],
    )
    (out_dir / "manifest.json").write_text(manifest.model_dump_json(indent=2) + "\n", "utf-8")
    return manifest