nordisk/knut

python/examples/social_graph.py

"""
Social-graph example
====================

Demonstrates the full round-trip:
  1. Register schemas on the server (Rust side — shown as comment).
  2. Insert a small social graph (vertices + edges) via DoPut.
  3. Fetch it back via DoGet.
  4. Build a NetworkX DiGraph and print basic statistics.

Prerequisite: graphar-flight server running locally on port 50051
with the following schemas registered:

    registry.register(
        "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name, n.age",
        Arc::new(Schema::new(vec![
            Field::new("_gar_id", DataType::Int64, false),
            Field::new("name",    DataType::Utf8,  true),
            Field::new("age",     DataType::Int64,  true),
        ])),
    );
    registry.register(
        "MATCH (a:Person)-[r:KNOWS]->(b:Person) RETURN a._gar_id AS src, b._gar_id AS dst, r.since",
        Arc::new(Schema::new(vec![
            Field::new("src",   DataType::Int64, false),
            Field::new("dst",   DataType::Int64, false),
            Field::new("since", DataType::Utf8,  true),
        ])),
    );
"""

import pandas as pd
import networkx as nx
from knut import KnutClient

VERTEX_QUERY = "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name, n.age"
EDGE_QUERY   = (
    "MATCH (a:Person)-[r:KNOWS]->(b:Person) "
    "RETURN a._gar_id AS src, b._gar_id AS dst, r.since"
)


def main() -> None:
    client = KnutClient(host="localhost", port=50051)

    # ── 1. Insert vertices ────────────────────────────────────────────────────
    persons = pd.DataFrame({
        "_gar_id": [0, 1, 2, 3],
        "name":    ["Alice", "Bob", "Carol", "Dave"],
        "age":     [30, 25, 35, 28],
    })
    n = client.insert_vertices("Person", persons)
    print(f"Inserted {n} vertices")

    # ── 2. Insert edges ───────────────────────────────────────────────────────
    knows = pd.DataFrame({
        "src":   [0, 0, 1, 2],
        "dst":   [1, 2, 3, 3],
        "since": ["2020-01-01", "2019-06-15", "2021-03-10", "2018-11-01"],
    })
    m = client.insert_edges("Person", "KNOWS", "Person", knows)
    print(f"Inserted {m} edges")

    # ── 3. Build NetworkX graph ───────────────────────────────────────────────
    # Arrow → pandas → NetworkX — no intermediate format conversion needed.
    G = client.to_networkx(
        vertex_query=VERTEX_QUERY,
        edge_query=EDGE_QUERY,
        node_id_col="_gar_id",
        src_col="src",
        dst_col="dst",
        directed=True,
    )

    # ── 4. Inspect ────────────────────────────────────────────────────────────
    print(f"\nGraph: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")

    print("\nNodes:")
    for node, attrs in G.nodes(data=True):
        print(f"  {node}: {attrs}")

    print("\nEdges:")
    for src, dst, attrs in G.edges(data=True):
        print(f"  {src} → {dst}  {attrs}")

    print("\nIn-degree centrality (top 3):")
    centrality = nx.in_degree_centrality(G)
    for node, score in sorted(centrality.items(), key=lambda x: -x[1])[:3]:
        name = G.nodes[node].get("name", node)
        print(f"  {name}: {score:.4f}")

    print("\nShortest path Alice → Dave:")
    try:
        path = nx.shortest_path(G, source=0, target=3)
        names = [G.nodes[n].get("name", n) for n in path]
        print("  " + " → ".join(names))
    except nx.NetworkXNoPath:
        print("  No path found")


if __name__ == "__main__":
    main()