ladybugdb

nx2lb.py

import networkx as nx
import real_ladybug
import csv
import tempfile
import os

def nx_to_ladybug_bulk(G: nx.Graph, db_path: str = "./ladybug_db"):
    """
    Convert a NetworkX graph to LadybugDB using bulk COPY FROM CSV.
    Much faster than row-by-row inserts for large graphs.
    """
    db = real_ladybug.Database(db_path)
    conn = real_ladybug.Connection(db)

    rel_type = "DIRECTED" if G.is_directed() else "CONNECTS"

    # --- Infer schemas ---
    node_prop_keys = sorted({k for _, d in G.nodes(data=True) for k in d})
    edge_prop_keys = sorted({k for _, _, d in G.edges(data=True) for k in d})

    with tempfile.TemporaryDirectory() as tmpdir:
        nodes_csv = os.path.join(tmpdir, "nodes.csv")
        edges_csv = os.path.join(tmpdir, "edges.csv")

        # --- Write nodes CSV ---
        node_headers = ["id"] + node_prop_keys
        with open(nodes_csv, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=node_headers)
            writer.writeheader()
            for node, data in G.nodes(data=True):
                row = {"id": str(node)}
                for k in node_prop_keys:
                    row[k] = str(data.get(k, ""))
                writer.writerow(row)

        # --- Write edges CSV ---
        edge_headers = ["src", "dst"] + edge_prop_keys
        with open(edges_csv, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=edge_headers)
            writer.writeheader()
            for src, dst, data in G.edges(data=True):
                row = {"src": str(src), "dst": str(dst)}
                for k in edge_prop_keys:
                    row[k] = str(data.get(k, ""))
                writer.writerow(row)

        # --- Create Node table ---
        node_cols = ", ".join(f"{k} STRING" for k in node_prop_keys)
        node_schema = f"id STRING, {node_cols}" if node_cols else "id STRING"
        conn.execute(f"CREATE NODE TABLE Node({node_schema}, PRIMARY KEY (id))")

        # --- Create Rel table ---
        edge_cols = ", ".join(f"{k} STRING" for k in edge_prop_keys)
        conn.execute(
            f"CREATE REL TABLE {rel_type}(FROM Node TO Node"
            + (f", {edge_cols}" if edge_cols else "")
            + ")"
        )

        # --- Bulk load ---
        conn.execute(f'COPY Node FROM "{nodes_csv}" (HEADER=TRUE)')
        conn.execute(f'COPY {rel_type} FROM "{edges_csv}" (HEADER=TRUE, FROM="src", TO="dst")')

    print(f"✅ Bulk imported {G.number_of_nodes()} nodes, {G.number_of_edges()} edges.")
    conn.close()
    return db


# --- Example usage ---
if __name__ == "__main__":
    G = nx.DiGraph()
    G.add_node("alice", age="30", role="engineer")
    G.add_node("bob",   age="25", role="designer")
    G.add_node("carol", age="28", role="manager")
    G.add_edge("alice", "bob",   since="2021", type="colleague")
    G.add_edge("carol", "alice", since="2019", type="manages")

    db = nx_to_ladybug_bulk(G, "./my_graph_db")

    # Query it back
    conn = real_ladybug.Connection(db)
    result = conn.execute("MATCH (a:Node)-[r]->(b:Node) RETURN a.id, b.id")
    while result.has_next():
        print(result.get_next())