falkor_load/graph_load_utils.py

"""Batch loaders for FalkorDB and Neo4j from staging parquet files."""

from __future__ import annotations

import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Callable, Iterator

import pyarrow.parquet as pq

BASE_DIR = Path(__file__).resolve().parent.parent
DEFAULT_STAGING_DIR = BASE_DIR / "clinical_graph_staging"
CONCEPTS_FILE = "concepts.parquet"
MRREL_FILE = "mrrel_edges.parquet"
HIER_FILE = "hier_edges.parquet"


def staging_paths(staging_dir: Path | None = None) -> dict[str, Path]:
    root = staging_dir or DEFAULT_STAGING_DIR
    return {
        "concepts": root / CONCEPTS_FILE,
        "mrrel": root / MRREL_FILE,
        "hier": root / HIER_FILE,
    }


def require_staging(staging_dir: Path | None = None) -> dict[str, Path]:
    paths = staging_paths(staging_dir)
    missing = [str(p) for p in paths.values() if not p.exists()]
    if missing:
        raise FileNotFoundError(
            "Missing staging parquet files:\n  "
            + "\n  ".join(missing)
            + "\nRun: uv run python3 build_clinical_graph.py --keep-staging"
        )
    return paths


def iter_parquet_batches(
    path: Path, batch_size: int, *, max_rows: int | None = None
) -> Iterator[list[dict]]:
    pf = pq.ParquetFile(path)
    sent = 0
    for batch in pf.iter_batches(batch_size=batch_size):
        df = batch.to_pandas()
        if max_rows is not None:
            remaining = max_rows - sent
            if remaining <= 0:
                return
            if len(df) > remaining:
                df = df.head(remaining)
        rows = df.where(df.notna(), None).to_dict(orient="records")
        if not rows:
            continue
        sent += len(rows)
        yield rows
        if max_rows is not None and sent >= max_rows:
            return


def _clean_str(value) -> str:
    if value is None:
        return ""
    return str(value).strip()


def concept_rows_for_load(rows: list[dict]) -> list[dict]:
    out: list[dict] = []
    for row in rows:
        aui = _clean_str(row.get("aui"))
        if not aui:
            continue
        out.append(
            {
                "aui": aui,
                "cui": _clean_str(row.get("cui")),
                "tui": _clean_str(row.get("tui")),
                "sab": _clean_str(row.get("sab")),
                "str": _clean_str(row.get("str")),
                "tty": _clean_str(row.get("tty")),
                "rank": int(row["rank"]) if row.get("rank") is not None else 0,
            }
        )
    return out


def mrrel_rows_for_load(rows: list[dict]) -> list[dict]:
    out: list[dict] = []
    for row in rows:
        src = _clean_str(row.get("src"))
        dst = _clean_str(row.get("dst"))
        if not src or not dst or src == dst:
            continue
        out.append(
            {
                "src": src,
                "dst": dst,
                "rel": _clean_str(row.get("rel")),
                "rela": _clean_str(row.get("rela")),
                "sab": _clean_str(row.get("sab")),
            }
        )
    return out


def hier_rows_for_load(rows: list[dict]) -> list[dict]:
    out: list[dict] = []
    for row in rows:
        src = _clean_str(row.get("src"))
        dst = _clean_str(row.get("dst"))
        if not src or not dst or src == dst:
            continue
        out.append(
            {
                "src": src,
                "dst": dst,
                "sab": _clean_str(row.get("sab")),
                "rela": _clean_str(row.get("rela")),
            }
        )
    return out


def parquet_row_count(path: Path) -> int:
    return pq.ParquetFile(path).metadata.num_rows


def estimate_batches(
    path: Path, batch_size: int, *, max_rows: int | None = None
) -> tuple[int, int]:
    """Return (row_count, num_batches)."""
    rows = parquet_row_count(path)
    if max_rows is not None:
        rows = min(rows, max_rows)
    batches = (rows + batch_size - 1) // batch_size if rows else 0
    return rows, batches


def run_batched_load(
    label: str,
    path: Path,
    batch_size: int,
    row_mapper: Callable[[list[dict]], list[dict]],
    loader: Callable[[list[dict]], None],
    *,
    max_rows: int | None = None,
    progress_every_batches: int = 10,
) -> int:
    expected_rows, expected_batches = estimate_batches(
        path, batch_size, max_rows=max_rows
    )
    print(
        f"  {label}: {expected_rows:,} rows (~{expected_batches:,} batches @ {batch_size:,})",
        flush=True,
    )
    total = 0
    batch_num = 0
    t0 = time.time()
    last_report = t0
    for raw in iter_parquet_batches(path, batch_size, max_rows=max_rows):
        batch = row_mapper(raw)
        if not batch:
            continue
        loader(batch)
        total += len(batch)
        batch_num += 1
        now = time.time()
        if (
            batch_num % progress_every_batches == 0
            or now - last_report >= 30
            or batch_num == expected_batches
        ):
            elapsed = now - t0
            rate = total / elapsed if elapsed > 0 else 0.0
            pct = 100.0 * total / expected_rows if expected_rows else 100.0
            eta_s = (expected_rows - total) / rate if rate > 0 else 0
            eta_m = eta_s / 60
            print(
                f"  {label}: {total:,}/{expected_rows:,} ({pct:.1f}%) "
                f"batch {batch_num}/{expected_batches} "
                f"@ {rate:,.0f} rows/s, ETA ~{eta_m:.0f} min",
                flush=True,
            )
            last_report = now
    elapsed = time.time() - t0
    rate = total / elapsed if elapsed > 0 else 0.0
    print(
        f"  {label}: done {total:,} rows in {elapsed / 60:.1f} min ({rate:,.0f} rows/s)",
        flush=True,
    )
    return total


def _partition_row_groups(path: Path, n_workers: int) -> list[list[int]]:
    pf = pq.ParquetFile(path)
    n = max(1, min(n_workers, pf.num_row_groups or 1))
    buckets: list[list[int]] = [[] for _ in range(n)]
    for i, rg in enumerate(range(pf.num_row_groups)):
        buckets[i % n].append(rg)
    return [b for b in buckets if b]


def _iter_row_group_batches(
    path: Path,
    row_group_ids: list[int],
    batch_size: int,
    *,
    max_rows: int | None = None,
) -> Iterator[list[dict]]:
    pf = pq.ParquetFile(path)
    sent = 0
    for rg in row_group_ids:
        table = pf.read_row_group(rg)
        df = table.to_pandas()
        if max_rows is not None:
            remaining = max_rows - sent
            if remaining <= 0:
                return
            if len(df) > remaining:
                df = df.head(remaining)
        for start in range(0, len(df), batch_size):
            chunk = df.iloc[start : start + batch_size]
            rows = chunk.where(chunk.notna(), None).to_dict(orient="records")
            if rows:
                sent += len(rows)
                yield rows
                if max_rows is not None and sent >= max_rows:
                    return


def run_parallel_parquet_load(
    label: str,
    path: Path,
    batch_size: int,
    row_mapper: Callable[[list[dict]], list[dict]],
    loader_factory: Callable[[], Callable[[list[dict]], None]],
    *,
    workers: int | None = None,
    max_rows: int | None = None,
) -> int:
    """
    Load parquet in parallel: one DB connection per worker, row-groups split across workers.
    """
    n_workers = workers or min(8, max(1, (os.cpu_count() or 4) - 1))
    expected_rows, _ = estimate_batches(path, batch_size, max_rows=max_rows)
    partitions = _partition_row_groups(path, n_workers)
    print(
        f"  {label}: {expected_rows:,} rows, {len(partitions)} workers "
        f"(batch_size={batch_size:,})",
        flush=True,
    )

    progress_lock = threading.Lock()
    progress = {"rows": 0, "batches": 0}
    t0 = time.time()
    last_report = t0

    def worker(part_id: int, row_groups: list[int]) -> int:
        load_fn = loader_factory()
        local = 0
        try:
            for raw in _iter_row_group_batches(
                path, row_groups, batch_size, max_rows=max_rows
            ):
                batch = row_mapper(raw)
                if not batch:
                    continue
                load_fn(batch)
                local += len(batch)
                with progress_lock:
                    progress["rows"] += len(batch)
                    progress["batches"] += 1
        finally:
            closer = getattr(load_fn, "close", None)
            if callable(closer):
                closer()
        return local

    total = 0
    with ThreadPoolExecutor(max_workers=len(partitions)) as pool:
        futures = {pool.submit(worker, i, rg): i for i, rg in enumerate(partitions)}
        for fut in as_completed(futures):
            total += fut.result()
            now = time.time()
            with progress_lock:
                rows = progress["rows"]
                batches = progress["batches"]
            if now - last_report >= 15:
                elapsed = now - t0
                rate = rows / elapsed if elapsed > 0 else 0.0
                pct = 100.0 * rows / expected_rows if expected_rows else 100.0
                eta_m = (expected_rows - rows) / rate / 60 if rate > 0 else 0
                print(
                    f"  {label}: {rows:,}/{expected_rows:,} ({pct:.1f}%) "
                    f"{batches:,} batches @ {rate:,.0f} rows/s, ETA ~{eta_m:.0f} min",
                    flush=True,
                )
                last_report = now

    elapsed = time.time() - t0
    rate = total / elapsed if elapsed > 0 else 0.0
    print(
        f"  {label}: done {total:,} rows in {elapsed / 60:.1f} min ({rate:,.0f} rows/s)",
        flush=True,
    )
    return total