falkor_load/build_falkor_graph.py

#!/usr/bin/env python3
"""Load clinical graph staging parquet into FalkorDB (GRAPH.BULK)."""

from __future__ import annotations

import argparse
import os
import sys
import time
from pathlib import Path

from .graph_bulk_falkor import (
    DEFAULT_BULK_BUFFER_MB,
    bulk_update_hier,
    ensure_falkor_graph_absent,
    export_hier_update_csv,
    falkor_bulk_load_graph,
    falkor_graph_counts,
    falkor_graph_redis_exists,
    verify_falkor_counts,
    wait_for_falkor_ready,
    warn_if_docker_memory_low,
)
from .graph_load_utils import (
    DEFAULT_STAGING_DIR,
    concept_rows_for_load,
    estimate_batches,
    hier_rows_for_load,
    mrrel_rows_for_load,
    parquet_row_count,
    require_staging,
    run_batched_load,
)

DEFAULT_FALKOR_GRAPH = "clinical_graph"
DEFAULT_BATCH_SIZE = 50_000
DEFAULT_EDGE_BATCH_SIZE = 10_000


def _print_load_plan(
    staging: dict[str, Path], batch_size: int, max_rows: int | None
) -> None:
    n_c = min(parquet_row_count(staging["concepts"]), max_rows or 10**18)
    n_r = min(parquet_row_count(staging["mrrel"]), max_rows or 10**18)
    n_h = min(parquet_row_count(staging["hier"]), max_rows or 10**18)
    _, b_c = estimate_batches(staging["concepts"], batch_size, max_rows=max_rows)
    _, b_r = estimate_batches(staging["mrrel"], batch_size, max_rows=max_rows)
    _, b_h = estimate_batches(staging["hier"], batch_size, max_rows=max_rows)
    total_batches = b_c + b_r + b_h
    print(
        f"\nLoad plan: {n_c:,} concepts + {n_r:,} MRREL + {n_h:,} HIER "
        f"≈ {total_batches:,} Cypher batches per backend (batch_size={batch_size:,})",
        flush=True,
    )
    print(
        "Note: Ladybug bulk COPY ~4 min. FalkorDB edges use GRAPH.BULK (fast). "
        "Use --falkor-method cypher for slow Cypher-only load.",
        flush=True,
    )
    print(
        "For quick profiling only: --max-rows 100000\n"
        "  python3 -m falkor_load.build_falkor_graph --rebuild\n",
        flush=True,
    )


def _falkor_concept_count(graph) -> int:
    """Count Concept nodes; graph Redis key must already exist."""
    result = graph.query("MATCH (c:Concept) RETURN count(c) AS n").result_set
    return int(result[0][0]) if result else 0


def _falkor_relation_counts(graph) -> tuple[int, int]:
    mr = graph.query("MATCH ()-[r:MRREL]->() RETURN count(r) AS n").result_set
    hi = graph.query("MATCH ()-[r:HIER_ISA]->() RETURN count(r) AS n").result_set
    return int(mr[0][0]) if mr else 0, int(hi[0][0]) if hi else 0


# Rough full-scale targets (for --resume-hier after OOM during bulk load).
_FULL_NODES = 6_500_000
_FULL_MRREL = 10_500_000
_FULL_HIER = 14_500_000


def _falkor_create_indexes(graph) -> None:
    """FalkorDB 1.6.x API: create_node_range_index(label, *properties)."""
    for prop in ("aui", "cui"):
        try:
            graph.create_node_range_index("Concept", prop)
            print(f"  index Concept.{prop} OK", flush=True)
        except Exception as exc:
            print(f"  index Concept.{prop}: {exc}", flush=True)


def _falkor_load_concepts_cypher(
    graph,
    staging: dict[str, Path],
    *,
    rebuild: bool,
    batch_size: int,
    max_rows: int | None,
) -> None:
    node_cypher = (
        """
        UNWIND $batch AS row
        CREATE (c:Concept {
            aui: row.aui, cui: row.cui, tui: row.tui, sab: row.sab,
            str: row.str, tty: row.tty, rank: row.rank
        })
        """
        if rebuild
        else """
        UNWIND $batch AS row
        MERGE (c:Concept {aui: row.aui})
        SET c.cui = row.cui, c.tui = row.tui, c.sab = row.sab,
            c.str = row.str, c.tty = row.tty, c.rank = row.rank
        """
    )

    def load_concepts(batch: list[dict]) -> None:
        graph.query(node_cypher, {"batch": batch})

    run_batched_load(
        "Falkor concepts",
        staging["concepts"],
        batch_size,
        concept_rows_for_load,
        load_concepts,
        max_rows=max_rows,
    )


def _load_hier_phase(
    graph,
    staging: dict[str, Path],
    *,
    graph_name: str,
    host: str,
    port: int,
    batch_size: int,
    max_rows: int | None,
    hier_method: str,
    hier_update_token_mb: int,
) -> None:
    if hier_method == "cypher":
        print(
            "  phase 2: HIER_ISA via Cypher UNWIND batches (slow) …",
            flush=True,
        )
        _falkor_load_hier_cypher(
            graph, staging, batch_size=batch_size, max_rows=max_rows
        )
        return

    print(
        "  phase 2: HIER_ISA via batched bulk-update (in-process, no Redis socket timeout) …",
        flush=True,
    )
    hier_pipe = export_hier_update_csv(staging["concepts"].parent, max_rows=max_rows)
    bulk_update_hier(
        graph_name,
        hier_pipe,
        host=host,
        port=port,
        max_token_mb=hier_update_token_mb,
    )


def _falkor_load_hier_cypher(
    graph,
    staging: dict[str, Path],
    *,
    batch_size: int,
    max_rows: int | None,
) -> None:
    from .graph_bulk_falkor import DEFAULT_BULK_QUERY_TIMEOUT_MS

    def load_hier(batch: list[dict]) -> None:
        graph.query(
            """
            UNWIND $batch AS row
            MATCH (a:Concept {aui: row.src}), (b:Concept {aui: row.dst})
            CREATE (a)-[:HIER_ISA {sab: row.sab, rela: row.rela}]->(b)
            """,
            {"batch": batch},
            timeout=DEFAULT_BULK_QUERY_TIMEOUT_MS,
        )

    edge_batch = min(batch_size, DEFAULT_EDGE_BATCH_SIZE)
    run_batched_load(
        "Falkor HIER_ISA (Cypher resume)",
        staging["hier"],
        edge_batch,
        hier_rows_for_load,
        load_hier,
        max_rows=max_rows,
        progress_every_batches=1,
    )


def _falkor_load_edges_cypher(
    graph,
    staging: dict[str, Path],
    *,
    batch_size: int,
    max_rows: int | None,
    hier_only: bool = False,
) -> None:
    if hier_only:
        _falkor_load_hier_cypher(
            graph, staging, batch_size=batch_size, max_rows=max_rows
        )
        return

    def load_mrrel(batch: list[dict]) -> None:
        graph.query(
            """
            UNWIND $batch AS row
            MATCH (a:Concept {aui: row.src}), (b:Concept {aui: row.dst})
            CREATE (a)-[:MRREL {rel: row.rel, rela: row.rela, sab: row.sab}]->(b)
            """,
            {"batch": batch},
        )

    edge_batch = min(batch_size, DEFAULT_EDGE_BATCH_SIZE)
    run_batched_load(
        "Falkor MRREL (Cypher)",
        staging["mrrel"],
        edge_batch,
        mrrel_rows_for_load,
        load_mrrel,
        max_rows=max_rows,
        progress_every_batches=1,
    )
    _falkor_load_hier_cypher(graph, staging, batch_size=batch_size, max_rows=max_rows)


def build_falkor(
    staging: dict[str, Path],
    *,
    host: str,
    port: int,
    graph_name: str,
    rebuild: bool,
    batch_size: int,
    max_rows: int | None,
    falkor_method: str,
    cypher_edges: bool,
    edges_only: bool,
    resume_hier: bool,
    bulk_buffer_mb: int,
    split_bulk: bool,
    hier_method: str,
    hier_update_token_mb: int,
) -> None:

    print(f"\n▶ FalkorDB {host}:{port} graph={graph_name}", flush=True)
    warn_if_docker_memory_low()
    wait_for_falkor_ready(graph_name, host=host, port=port)
    if falkor_method == "cypher":
        cypher_edges = True
        print(
            "  method: Cypher batches (reliable; edges can take many hours at full scale)",
            flush=True,
        )
    else:
        bulk_all = not split_bulk
        mode = (
            "GRAPH.BULK Concept+MRREL+HIER_ISA (fast, default)"
            if bulk_all
            else f"GRAPH.BULK nodes+MRREL only, then HIER via {hier_method} (--split-bulk)"
        )
        print(f"  method: {mode}", flush=True)
        if bulk_all:
            warn_if_docker_memory_low(min_gib=8.0)

    from .graph_bulk_falkor import falkordb_connect

    db = falkordb_connect(host=host, port=port)
    graph = db.select_graph(graph_name)

    if resume_hier:
        if not falkor_graph_redis_exists(graph_name, host=host, port=port):
            raise RuntimeError(
                "--resume-hier: no graph in Redis. Run a bulk load first or use --rebuild."
            )
        n_nodes = _falkor_concept_count(graph)
        n_mrrel, n_hier = _falkor_relation_counts(graph)
        print(
            f"  resume: {n_nodes:,} nodes, {n_mrrel:,} MRREL, {n_hier:,} HIER_ISA",
            flush=True,
        )
        if n_hier >= _FULL_HIER:
            print("  HIER_ISA already complete — nothing to do.", flush=True)
            return
        if n_nodes < _FULL_NODES or n_mrrel < _FULL_MRREL:
            raise RuntimeError(
                f"--resume-hier needs a partial bulk load (nodes+MRREL done, HIER missing). "
                f"Got {n_nodes:,} nodes, {n_mrrel:,} MRREL. Use --rebuild for a fresh load."
            )
        # Cypher / bulk-update HIER append is very slow. Full GRAPH.BULK reload is faster.
        print(
            "  --resume-hier: full GRAPH.BULK reload (Concept+MRREL+HIER, ~7–10 min) …",
            flush=True,
        )
        ensure_falkor_graph_absent(graph_name, host=host, port=port)
        graph = db.select_graph(graph_name)
        buf_mb = bulk_buffer_mb or DEFAULT_BULK_BUFFER_MB
        exp_n, exp_mr, exp_hi = falkor_bulk_load_graph(
            staging["concepts"].parent,
            graph_name=graph_name,
            host=host,
            port=port,
            max_rows=max_rows,
            max_buffer_mb=buf_mb,
            bulk_all=True,
        )
        print("  creating FalkorDB indexes on Concept.aui / Concept.cui …", flush=True)
        _falkor_create_indexes(graph)
        ok, detail = verify_falkor_counts(
            graph_name,
            host=host,
            port=port,
            expected_nodes=exp_n,
            expected_mrrel=exp_mr,
            expected_hier=exp_hi,
        )
        if not ok:
            raise RuntimeError(
                f"FalkorDB graph incomplete after resume reload: {detail}"
            )
        print(f"  FalkorDB after resume: {detail}", flush=True)
        return

    if rebuild and not edges_only:
        ensure_falkor_graph_absent(graph_name, host=host, port=port)
        print("  deleted existing graph (Redis key cleared)", flush=True)
        graph = db.select_graph(graph_name)
        n_existing = 0
    elif falkor_graph_redis_exists(graph_name, host=host, port=port):
        n_existing = _falkor_concept_count(graph)
    else:
        n_existing = 0

    if edges_only:
        if n_existing == 0:
            raise RuntimeError(
                "edges-only: no Concept nodes in graph. Load nodes first or drop --edges-only."
            )
        print(f"  edges-only: {n_existing:,} Concept nodes present", flush=True)
        if not cypher_edges:
            print(
                "  edges-only: GRAPH.BULK cannot append to an existing graph; "
                "using Cypher edge load (slow). For fast full reload use --rebuild without --edges-only.",
                flush=True,
            )
            cypher_edges = True
    elif n_existing > 0 and not rebuild:
        print(
            f"  graph already has {n_existing:,} Concept nodes — skip (use --rebuild)"
        )
        return

    if cypher_edges and not edges_only:
        _falkor_load_concepts_cypher(
            graph, staging, rebuild=rebuild, batch_size=batch_size, max_rows=max_rows
        )
        print("  creating FalkorDB indexes on Concept.aui / Concept.cui …", flush=True)
        _falkor_create_indexes(graph)

    if cypher_edges:
        if falkor_method != "cypher":
            print(
                "  WARNING: Cypher edge load is slow at full scale "
                "(MATCH per edge). Default is GRAPH.BULK.",
                flush=True,
            )
        _falkor_load_edges_cypher(
            graph, staging, batch_size=batch_size, max_rows=max_rows
        )
    else:
        # GRAPH.BULK (split by default: nodes+MRREL bulk, HIER Cypher — avoids OOM at end).
        if n_existing > 0:
            raise RuntimeError(
                f"GRAPH.BULK requires an empty graph (found {n_existing:,} Concept nodes). "
                "Use --rebuild to delete and reload, or --cypher-edges to append slowly."
            )
        print(
            "  tip: FalkorDB only, ≥10 GB — "
            "docker compose -f docker-compose.graph-bench.yml up -d --force-recreate falkordb",
            flush=True,
        )
        buf_mb = bulk_buffer_mb or DEFAULT_BULK_BUFFER_MB
        exp_n, exp_mr, exp_hi = falkor_bulk_load_graph(
            staging["concepts"].parent,
            graph_name=graph_name,
            host=host,
            port=port,
            max_rows=max_rows,
            max_buffer_mb=buf_mb,
            bulk_all=not split_bulk,
        )
        graph = db.select_graph(graph_name)
        print("  creating FalkorDB indexes on Concept.aui / Concept.cui …", flush=True)
        _falkor_create_indexes(graph)

        if split_bulk:
            n, mr, hi = falkor_graph_counts(graph_name, host=host, port=port)
            print(
                f"  after bulk phase: {n:,} nodes, {mr:,} MRREL, {hi:,} HIER_ISA",
                flush=True,
            )
            need_hier = hi < int(exp_hi * 0.98)
            if need_hier:
                _load_hier_phase(
                    graph,
                    staging,
                    graph_name=graph_name,
                    host=host,
                    port=port,
                    batch_size=batch_size,
                    max_rows=max_rows,
                    hier_method=hier_method,
                    hier_update_token_mb=hier_update_token_mb,
                )

        ok, detail = verify_falkor_counts(
            graph_name,
            host=host,
            port=port,
            expected_nodes=exp_n,
            expected_mrrel=exp_mr,
            expected_hier=exp_hi,
        )
        if not ok:
            raise RuntimeError(f"FalkorDB graph incomplete after load: {detail}")
        print(f"  FalkorDB load verified: {detail}", flush=True)

    n, mr, hi = falkor_graph_counts(graph_name, host=host, port=port)
    print(
        f"  FalkorDB ready: {n:,} Concept nodes, {mr:,} MRREL, {hi:,} HIER_ISA",
        flush=True,
    )


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--staging-dir", type=Path, default=DEFAULT_STAGING_DIR)
    p.add_argument("--falkor-only", action="store_true", help=argparse.SUPPRESS)
    p.add_argument(
        "--rebuild", action="store_true", help="Drop existing graph data first"
    )
    p.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        help=f"Rows per Cypher UNWIND batch (default: {DEFAULT_BATCH_SIZE:,})",
    )
    p.add_argument(
        "--max-rows", type=int, default=None, help="Limit rows per table (testing)"
    )
    p.add_argument(
        "--falkor-method",
        choices=("bulk", "cypher"),
        default="bulk",
        help="FalkorDB load: bulk=GRAPH.BULK (fast, default); cypher=all Cypher batches",
    )
    p.add_argument("--cypher-edges", action="store_true")
    p.add_argument("--edges-only", action="store_true")
    p.add_argument("--resume-hier", action="store_true")
    p.add_argument("--bulk-buffer-mb", type=int, default=48)
    p.add_argument("--split-bulk", action="store_true")
    p.add_argument("--bulk-all", action="store_true", help=argparse.SUPPRESS)
    p.add_argument(
        "--hier-method",
        choices=("bulk-update", "cypher"),
        default="bulk-update",
    )
    p.add_argument("--hier-update-token-mb", type=int, default=32)
    p.add_argument("--estimate-only", action="store_true")
    p.add_argument(
        "--falkor-host", default=os.environ.get("FALKORDB_HOST", "localhost")
    )
    p.add_argument(
        "--falkor-port", type=int, default=int(os.environ.get("FALKORDB_PORT", "6379"))
    )
    p.add_argument(
        "--falkor-graph",
        default=os.environ.get("FALKORDB_GRAPH", DEFAULT_FALKOR_GRAPH),
    )
    return p.parse_args()


def main() -> int:
    args = parse_args()
    t0 = time.time()

    try:
        staging = require_staging(args.staging_dir)
    except FileNotFoundError as exc:
        print(exc, file=sys.stderr)
        return 1

    _print_load_plan(staging, args.batch_size, args.max_rows)
    if args.estimate_only:
        return 0

    try:
        build_falkor(
            staging,
            host=args.falkor_host,
            port=args.falkor_port,
            graph_name=args.falkor_graph,
            rebuild=args.rebuild,
            batch_size=args.batch_size,
            max_rows=args.max_rows,
            falkor_method=args.falkor_method,
            cypher_edges=args.cypher_edges,
            edges_only=args.edges_only,
            resume_hier=args.resume_hier,
            bulk_buffer_mb=args.bulk_buffer_mb,
            split_bulk=args.split_bulk and not args.bulk_all,
            hier_method=args.hier_method,
            hier_update_token_mb=args.hier_update_token_mb,
        )
    except Exception as exc:
        print(f"FalkorDB load failed: {exc}", file=sys.stderr)
        return 1

    print(f"\nDone in {(time.time() - t0) / 60:.1f} min")
    return 0


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