"""Fast FalkorDB import via falkordb-bulk-loader (GRAPH.BULK).""" from __future__ import annotations import csv import os import subprocess import time from pathlib import Path from timeit import default_timer as timer import duckdb from .graph_load_utils import CONCEPTS_FILE, HIER_FILE, MRREL_FILE BULK_CSV_DIR = "bulk_csv" # CSV files under bulk_csv/ are generated outputs for falkordb-bulk-insert / bulk-update. # Source of truth is always staging parquet (see graph_load_utils.CONCEPTS_FILE, etc.). def warn_if_docker_memory_low( container_name_part: str = "falkordb", *, min_gib: float = 8.0, ) -> None: """ Docker Desktop often caps containers at ~4 GiB even when compose sets mem_limit=10g. Prints a warning if the effective cgroup limit looks too small for a full clinical load. """ try: out = subprocess.check_output( [ "docker", "stats", "--no-stream", "--format", "{{.Name}}\t{{.MemUsage}}", ], text=True, stderr=subprocess.DEVNULL, ) except (subprocess.CalledProcessError, FileNotFoundError): return import re for line in out.splitlines(): if container_name_part not in line.lower(): continue m = re.search(r"/\s*([\d.]+)\s*(GiB|MiB)", line, re.I) if not m: continue val, unit = float(m.group(1)), m.group(2).lower() limit_gib = val if unit == "gib" else val / 1024 if limit_gib < min_gib: print( f"\n⚠ Docker reports FalkorDB memory LIMIT ≈ {limit_gib:.1f} GiB " f"(need ≥ {min_gib:.0f} GiB for nodes+MRREL+HIER).\n" " Docker Desktop → Settings → Resources → Memory: set to 14 GB or more, " "Apply & Restart, then recreate the container:\n" " docker compose -f docker-compose.graph-bench.yml up -d --force-recreate falkordb\n", flush=True, ) return def falkor_graph_redis_exists( graph_name: str, *, host: str = "localhost", port: int = 6379, ) -> bool: """True if the Falkor graph Redis key exists (without touching the graph).""" db = falkordb_connect(host=host, port=port) return bool(db.connection.exists(graph_name)) def ensure_falkor_graph_absent( graph_name: str, *, host: str = "localhost", port: int = 6379, ) -> None: """ Delete a Falkor graph and remove its Redis key(s). FalkorDB recreates an empty graph key if you run Cypher after GRAPH.DELETE, which makes falkordb-bulk-insert think the graph already exists. """ db = falkordb_connect(host=host, port=port) conn = db.connection graph = db.select_graph(graph_name) try: graph.delete() except Exception: pass if conn.exists(graph_name): conn.delete(graph_name) for key in conn.keys(f"{graph_name}*"): conn.delete(key) if conn.exists(graph_name): raise RuntimeError( f"Could not remove FalkorDB graph key '{graph_name}' from Redis; " "stop other clients using this graph and retry." ) def _sql_path(p: Path) -> str: return str(p).replace("'", "''") # DuckDB COPY emits RFC-4180 CSV; bulk loader defaults to QUOTE_NONE (breaks on commas/quotes). _BULK_CSV_COPY_OPTS = "(HEADER, DELIMITER ',', QUOTE '\"', ESCAPE '\"')" _BULK_INSERT_QUOTE = "0" # csv.QUOTE_MINIMAL — matches DuckDB output _BULK_INSERT_ESCAPECHAR = "none" # Smaller buffers reduce peak Redis memory during HIER_ISA (avoids OOM / disconnect). DEFAULT_BULK_BUFFER_MB = 48 DEFAULT_BULK_MAX_TOKEN_COUNT = 2048 # Pipe CSV for falkordb-bulk-update (uses QUOTE_NONE; no commas in fields). HIER_UPDATE_PIPE = "HIER_ISA_update.pipe" HIER_BULK_UPDATE_QUERY = ( "MATCH (a:Concept {aui: row[0]}), (b:Concept {aui: row[1]}) " "CREATE (a)-[:HIER_ISA {sab: row[2], rela: row[3]}]->(b)" ) # Only used with --split-bulk (slow fallback). Full loads use GRAPH.BULK for HIER_ISA. DEFAULT_HIER_UPDATE_TOKEN_MB = 32 DEFAULT_BULK_QUERY_TIMEOUT_MS = int( os.environ.get("FALKORDB_BULK_QUERY_TIMEOUT_MS", "1800000") ) # 30 min per batch def _bulk_socket_timeout() -> float | None: """None = wait indefinitely for FalkorDB to finish a large batch.""" raw = os.environ.get("FALKORDB_BULK_SOCKET_TIMEOUT_S", "0") return None if float(raw) <= 0 else float(raw) def falkordb_connect(host: str = "localhost", port: int = 6379): from falkordb import FalkorDB return FalkorDB(host=host, port=port, socket_timeout=_bulk_socket_timeout()) def _run_falkordb_bulk_insert_cli(argv: list[str]) -> int: """ Run ``falkordb-bulk-insert`` in-process with Redis socket_timeout disabled. The stock CLI uses redis-py's default 5s socket timeout. Each GRAPH.BULK buffer (~500k nodes) often takes longer, which surfaces as ``Timeout reading from socket`` or ``Timeout writing to socket`` — not OOM. """ import falkordb import redis from click.testing import CliRunner from falkordb_bulk_loader.bulk_insert import bulk_insert sock = _bulk_socket_timeout() orig_falkor = falkordb.FalkorDB.from_url orig_redis = redis.from_url def _with_timeout(**kwargs): kwargs.setdefault("socket_timeout", sock) kwargs.setdefault("socket_connect_timeout", 60.0) return kwargs def patched_redis_from_url(url, **kwargs): return orig_redis(url, **_with_timeout(**kwargs)) def patched_falkor_from_url(url, **kwargs): return orig_falkor(url, **_with_timeout(**kwargs)) falkordb.FalkorDB.from_url = patched_falkor_from_url redis.from_url = patched_redis_from_url try: result = CliRunner().invoke(bulk_insert, argv) if result.exception is not None: raise result.exception return int(result.exit_code) finally: falkordb.FalkorDB.from_url = orig_falkor redis.from_url = orig_redis def _utf8len(s: str) -> int: return len(s.encode("utf-8")) def _quote_bulk_update_cell(cell: str) -> str: cell = cell.strip() try: float(cell) except ValueError: if ( cell.lower() not in ("false", "true") and not (cell.startswith("[") and cell.endswith("]")) and not (cell.startswith('"') and cell.endswith('"')) and not (cell.startswith("'") and cell.endswith("'")) ): cell = f'"{cell}"' return cell def _sanitize_sql(expr: str) -> str: """Strip characters that break CSV / Cypher string literals.""" return ( f"replace(replace(replace(replace(coalesce({expr}, ''), " f"chr(10), ' '), chr(13), ' '), chr(9), ' '), chr(34), chr(39))" ) def validate_bulk_csv(path: Path, *, expected_columns: int) -> None: """Fail fast if CSV rows do not match the bulk loader's QUOTE_MINIMAL parsing.""" bad = 0 first_bad: tuple[int, int, list[str]] | None = None with path.open(newline="", encoding="utf-8") as fh: reader = csv.reader( fh, delimiter=",", skipinitialspace=True, quoting=csv.QUOTE_MINIMAL, ) header = next(reader) if len(header) != expected_columns: raise RuntimeError( f"{path.name}: header has {len(header)} columns, expected {expected_columns}" ) for line_no, row in enumerate(reader, start=2): if len(row) != expected_columns: bad += 1 if first_bad is None: first_bad = (line_no, len(row), row) if bad: line_no, ncol, row = first_bad or (0, 0, []) preview = ",".join(row[:6])[:120] raise RuntimeError( f"{path.name}: {bad:,} malformed row(s) for bulk import " f"(e.g. line {line_no}: {ncol} columns, expected {expected_columns}). " f"Preview: {preview!r}" ) def _copy_bulk_csv(con: duckdb.DuckDBPyConnection, select_sql: str, dest: Path) -> None: con.execute(f"COPY ({select_sql}) TO '{_sql_path(dest)}' {_BULK_CSV_COPY_OPTS}") def _log_parquet_source(label: str, parquet_path: Path) -> None: if not parquet_path.exists(): raise FileNotFoundError(f"Missing staging parquet for {label}: {parquet_path}") print(f" source {label}: {parquet_path}", flush=True) def export_concept_csv( staging_dir: Path, *, max_rows: int | None = None, ) -> Path: """Export Concept nodes for falkordb-bulk-insert (:ID schema header).""" bulk_dir = staging_dir / BULK_CSV_DIR bulk_dir.mkdir(parents=True, exist_ok=True) concept_csv = bulk_dir / "Concept.csv" limit = f"LIMIT {int(max_rows)}" if max_rows is not None else "" _log_parquet_source("Concept", staging_dir / CONCEPTS_FILE) con = duckdb.connect() print(" exporting → Concept.csv (overwrite) …", flush=True) _copy_bulk_csv( con, f""" SELECT trim(aui) AS "aui:ID(Concept)", {_sanitize_sql("trim(cui)")} AS "cui:STRING", {_sanitize_sql("trim(tui)")} AS "tui:STRING", {_sanitize_sql("trim(sab)")} AS "sab:STRING", {_sanitize_sql("trim(str)")} AS "str:STRING", {_sanitize_sql("trim(tty)")} AS "tty:STRING", coalesce(cast(rank AS BIGINT), 0) AS "rank:INT" FROM read_parquet('{_sql_path(staging_dir / CONCEPTS_FILE)}') WHERE trim(aui) <> '' {limit} """, concept_csv, ) validate_bulk_csv(concept_csv, expected_columns=7) con.close() return concept_csv def export_relation_csvs( staging_dir: Path, *, max_rows: int | None = None, ) -> dict[str, Path]: """Export MRREL / HIER_ISA staging parquet to schema CSVs for falkordb-bulk-insert.""" bulk_dir = staging_dir / BULK_CSV_DIR bulk_dir.mkdir(parents=True, exist_ok=True) mrrel_csv = bulk_dir / "MRREL.csv" hier_csv = bulk_dir / "HIER_ISA.csv" limit = f"LIMIT {int(max_rows)}" if max_rows is not None else "" _log_parquet_source("MRREL", staging_dir / MRREL_FILE) _log_parquet_source("HIER_ISA", staging_dir / HIER_FILE) con = duckdb.connect() print(" exporting → MRREL.csv (overwrite) …", flush=True) _copy_bulk_csv( con, f""" SELECT trim(src) AS ":START_ID(Concept)", trim(dst) AS ":END_ID(Concept)", {_sanitize_sql("trim(rel)")} AS "rel:STRING", {_sanitize_sql("trim(rela)")} AS "rela:STRING", {_sanitize_sql("trim(sab)")} AS "sab:STRING" FROM read_parquet('{_sql_path(staging_dir / MRREL_FILE)}') WHERE trim(src) <> '' AND trim(dst) <> '' AND trim(src) <> trim(dst) {limit} """, mrrel_csv, ) validate_bulk_csv(mrrel_csv, expected_columns=5) print(" exporting → HIER_ISA.csv (overwrite) …", flush=True) _copy_bulk_csv( con, f""" SELECT trim(src) AS ":START_ID(Concept)", trim(dst) AS ":END_ID(Concept)", {_sanitize_sql("trim(sab)")} AS "sab:STRING", {_sanitize_sql("trim(rela)")} AS "rela:STRING" FROM read_parquet('{_sql_path(staging_dir / HIER_FILE)}') WHERE trim(src) <> '' AND trim(dst) <> '' AND trim(src) <> trim(dst) {limit} """, hier_csv, ) validate_bulk_csv(hier_csv, expected_columns=4) con.close() return {"mrrel": mrrel_csv, "hier": hier_csv} def export_hier_update_csv( staging_dir: Path, *, max_rows: int | None = None, ) -> Path: """Pipe-separated CSV for falkordb-bulk-update (fast append to existing graph).""" bulk_dir = staging_dir / BULK_CSV_DIR bulk_dir.mkdir(parents=True, exist_ok=True) hier_pipe = bulk_dir / HIER_UPDATE_PIPE limit = f"LIMIT {int(max_rows)}" if max_rows is not None else "" _log_parquet_source("HIER_ISA", staging_dir / HIER_FILE) con = duckdb.connect() print(" exporting → HIER_ISA_update.pipe (overwrite) …", flush=True) con.execute( f""" COPY ( SELECT trim(src) AS src, trim(dst) AS dst, {_sanitize_sql("trim(sab)")} AS sab, {_sanitize_sql("trim(rela)")} AS rela FROM read_parquet('{_sql_path(staging_dir / HIER_FILE)}') WHERE trim(src) <> '' AND trim(dst) <> '' AND trim(src) <> trim(dst) {limit} ) TO '{_sql_path(hier_pipe)}' (HEADER, DELIMITER '|') """ ) con.close() return hier_pipe def bulk_update_hier( graph_name: str, hier_csv: Path, *, host: str = "localhost", port: int = 6379, max_token_mb: int = DEFAULT_HIER_UPDATE_TOKEN_MB, query_timeout_ms: int = DEFAULT_BULK_QUERY_TIMEOUT_MS, verbose: bool = True, ) -> None: """ Slow fallback: append HIER_ISA via batched UNWIND + MATCH + CREATE. Prefer a single ``falkordb-bulk-insert`` with HIER_ISA.csv (default) — native GRAPH.BULK is much faster than this path. """ wait_for_falkor_ready(graph_name, host=host, port=port) max_token_mb = max(4, int(max_token_mb)) unwound = " ".join(["UNWIND $rows AS", "row", HIER_BULK_UPDATE_QUERY]) max_token_bytes = max_token_mb * 1024 * 1024 - _utf8len(unwound) print( f" HIER_ISA bulk-update: {hier_csv.name} " f"(batch≤{max_token_mb}MB, query_timeout={query_timeout_ms // 1000}s, " f"socket_timeout={'none' if _bulk_socket_timeout() is None else _bulk_socket_timeout()}) …", flush=True, ) db = falkordb_connect(host=host, port=port) graph = db.select_graph(graph_name) graph.explain(" ".join(["CYPHER rows=[]", unwound])) start = timer() buffers_sent = 0 rels_created = 0 buffer_size = 0 rows_strs: list[str] = [] def emit_buffer() -> None: nonlocal buffers_sent, rels_created, buffer_size, rows_strs if not rows_strs: return rows_cypher = "".join(["CYPHER rows=[", ",".join(rows_strs), "]"]) command = " ".join([rows_cypher, unwound]) nbytes = _utf8len(command) buffers_sent += 1 if verbose: print( f" batch #{buffers_sent}: {nbytes / (1024 * 1024):.1f} MB, " f"{len(rows_strs):,} rows …", flush=True, ) t0 = timer() result = graph.query(command, timeout=query_timeout_ms) rels_created += int(result.relationships_created) if verbose: print( f" batch #{buffers_sent} done in {timer() - t0:.1f}s " f"(+{int(result.relationships_created):,} rels, total {rels_created:,})", flush=True, ) rows_strs = [] buffer_size = 0 with hier_csv.open(newline="", encoding="utf-8") as fh: next(fh) # header reader = csv.reader( fh, delimiter="|", skipinitialspace=True, quoting=csv.QUOTE_NONE, escapechar="\\", ) for row in reader: row_line = ",".join(_quote_bulk_update_cell(cell) for cell in row) next_line = f"[{row_line.strip()}]" added = _utf8len(next_line) + 1 if buffer_size + added > max_token_bytes and rows_strs: emit_buffer() rows_strs.append(next_line) buffer_size += added emit_buffer() print( f" HIER_ISA bulk-update complete: {rels_created:,} relationships in " f"{timer() - start:.1f}s ({buffers_sent} batches)", flush=True, ) def falkor_graph_counts( graph_name: str, *, host: str = "localhost", port: int = 6379, ) -> tuple[int, int, int]: """Return (concept_nodes, mrrel_edges, hier_edges) for a graph.""" wait_for_falkor_ready(graph_name, host=host, port=port) graph = falkordb_connect(host=host, port=port).select_graph(graph_name) n = graph.query("MATCH (c:Concept) RETURN count(c) AS n").result_set[0][0] mr = graph.query("MATCH ()-[r:MRREL]->() RETURN count(r) AS n").result_set[0][0] hi = graph.query("MATCH ()-[r:HIER_ISA]->() RETURN count(r) AS n").result_set[0][0] return int(n), int(mr), int(hi) def verify_falkor_counts( graph_name: str, *, host: str, port: int, expected_nodes: int, expected_mrrel: int, expected_hier: int, min_ratio: float = 0.98, ) -> tuple[bool, str]: """Check graph size against CSV/staging expectations (after OOM / disconnect).""" if not falkor_graph_redis_exists(graph_name, host=host, port=port): return False, "graph key missing in Redis" try: n, mr, hi = falkor_graph_counts(graph_name, host=host, port=port) except Exception as exc: return False, f"cannot query graph: {exc}" def ok(got: int, exp: int) -> bool: if exp <= 0: return got == 0 return got >= int(exp * min_ratio) parts = [ f"nodes {n:,}/{expected_nodes:,}", f"MRREL {mr:,}/{expected_mrrel:,}", f"HIER {hi:,}/{expected_hier:,}", ] if ok(n, expected_nodes) and ok(mr, expected_mrrel) and ok(hi, expected_hier): return True, "; ".join(parts) return False, "; ".join(parts) def wait_for_falkor_ready( graph_name: str, *, host: str = "localhost", port: int = 6379, timeout_s: float = 180.0, poll_s: float = 3.0, ) -> None: """Wait until FalkorDB accepts connections (e.g. after container restart).""" deadline = time.monotonic() + timeout_s last_err: Exception | None = None while time.monotonic() < deadline: try: db = falkordb_connect(host=host, port=port) db.connection.ping() db.select_graph(graph_name) return except Exception as exc: last_err = exc time.sleep(poll_s) raise RuntimeError( f"FalkorDB at {host}:{port} not ready after {timeout_s:.0f}s: {last_err}" ) def bulk_load_graph( graph_name: str, csv_paths: dict[str, Path], *, host: str = "localhost", port: int = 6379, skip_invalid_edges: bool = True, verbose: bool = True, max_buffer_mb: int = DEFAULT_BULK_BUFFER_MB, max_token_count: int = DEFAULT_BULK_MAX_TOKEN_COUNT, include_hier: bool = True, clear_graph: bool = False, ) -> tuple[int, int, int]: """ Import via falkordb-bulk-insert. Returns expected (n_concept, n_mrrel, n_hier) from CSV row counts. When include_hier=False, only Concept + MRREL are bulk-loaded (lower peak RAM). """ if clear_graph: ensure_falkor_graph_absent(graph_name, host=host, port=port) wait_for_falkor_ready(graph_name, host=host, port=port) try: from falkordb_bulk_loader.bulk_insert import bulk_insert # noqa: F401 except ImportError as exc: raise RuntimeError( "falkordb-bulk-loader not installed. Install: uv pip install falkordb-bulk-loader" ) from exc con = duckdb.connect() n_concept = int( con.execute( f"SELECT count(*) FROM read_csv('{_sql_path(csv_paths['concept'])}', header=true)" ).fetchone()[0] ) n_mrrel = int( con.execute( f"SELECT count(*) FROM read_csv('{_sql_path(csv_paths['mrrel'])}', header=true)" ).fetchone()[0] ) n_hier = ( int( con.execute( f"SELECT count(*) FROM read_csv('{_sql_path(csv_paths['hier'])}', header=true)" ).fetchone()[0] ) if include_hier else 0 ) con.close() bulk_argv = [ graph_name, "--server-url", f"redis://{host}:{port}", "--enforce-schema", "--id-type", "STRING", "-N", "Concept", str(csv_paths["concept"]), "-R", "MRREL", str(csv_paths["mrrel"]), "--max-buffer-size", str(max(8, int(max_buffer_mb))), "--max-token-count", str(max(256, int(max_token_count))), "--quote", _BULK_INSERT_QUOTE, "--escapechar", _BULK_INSERT_ESCAPECHAR, ] if include_hier: bulk_argv.extend(["-R", "HIER_ISA", str(csv_paths["hier"])]) if skip_invalid_edges: bulk_argv.append("--skip-invalid-edges") if verbose: bulk_argv.append("--verbose") label = ( f"{n_concept:,} Concept + {n_mrrel:,} MRREL + {n_hier:,} HIER_ISA" if include_hier else f"{n_concept:,} Concept + {n_mrrel:,} MRREL (HIER via Cypher next)" ) sock = _bulk_socket_timeout() print(f" GRAPH.BULK import: {label} …", flush=True) print( f" bulk buffers: max_buffer_size={max_buffer_mb}MB, max_token_count={max_token_count}, " f"redis_socket_timeout={'none' if sock is None else sock}", flush=True, ) print(f" argv: falkordb-bulk-insert {' '.join(bulk_argv)}", flush=True) exit_code = _run_falkordb_bulk_insert_cli(bulk_argv) if exit_code == 0: return n_concept, n_mrrel, n_hier # Bulk loader often dies on the last flush while FalkorDB finalizes matrices (OOM). print( " bulk loader exited non-zero — waiting for FalkorDB to restart and checking counts …", flush=True, ) time.sleep(15) wait_for_falkor_ready(graph_name, host=host, port=port, timeout_s=600.0) ok, detail = verify_falkor_counts( graph_name, host=host, port=port, expected_nodes=n_concept, expected_mrrel=n_mrrel, expected_hier=n_hier, ) if ok: print(f" graph load looks complete despite bulk exit: {detail}", flush=True) return n_concept, n_mrrel, n_hier hint = ( "Common causes: (1) Redis 5s socket timeout on huge GRAPH.BULK buffers — fixed in " "this wrapper via FALKORDB_BULK_SOCKET_TIMEOUT_S=0; retry --rebuild. " "(2) Docker OOM — use --split-bulk, give Falkor ≥10 GiB RAM, stop Neo4j during load. " ) if falkor_graph_redis_exists(graph_name, host=host, port=port): hint += "Try: build_falkor_neo4j_graph.py --falkor-only --resume-hier" else: hint += "Then: build_falkor_neo4j_graph.py --falkor-only --rebuild" raise RuntimeError( f"falkordb-bulk-insert exited with code {exit_code} ({detail}). {hint}" ) def falkor_bulk_load_graph( staging_dir: Path, *, graph_name: str, host: str, port: int, max_rows: int | None, max_buffer_mb: int = DEFAULT_BULK_BUFFER_MB, max_token_count: int = DEFAULT_BULK_MAX_TOKEN_COUNT, bulk_all: bool = True, ) -> tuple[int, int, int]: """ Export staging CSVs and load into FalkorDB via ``falkordb-bulk-insert``. Default (bulk_all=True): one GRAPH.BULK for Concept + MRREL + HIER_ISA (~7–10 min). Use bulk_all=False only with ``--split-bulk`` when Docker RAM is tight (<8 GiB). """ concept_csv = export_concept_csv(staging_dir, max_rows=max_rows) rel_csvs = export_relation_csvs(staging_dir, max_rows=max_rows) csv_paths = {"concept": concept_csv, **rel_csvs} exp_n, exp_mr, exp_hi_bulk = bulk_load_graph( graph_name, csv_paths, host=host, port=port, max_buffer_mb=max_buffer_mb, max_token_count=max_token_count, include_hier=bulk_all, clear_graph=True, ) if bulk_all: return exp_n, exp_mr, exp_hi_bulk con = duckdb.connect() exp_hi = int( con.execute( f"SELECT count(*) FROM read_csv('{_sql_path(csv_paths['hier'])}', header=true)" ).fetchone()[0] ) con.close() return exp_n, exp_mr, exp_hi # Backward-compatible alias falkor_bulk_load_edges = falkor_bulk_load_graph