nordisk/knut

demo/flight_networkx_demo.py

#!/usr/bin/env python3
"""
Arrow Flight → networkx demo (N4).

Connects to a running ``knut-flight-server`` (the plain Arrow Flight server over
knut's rel2graph output), lists the advertised flights, pulls the ``nodes`` and
``edges`` datasets as Arrow tables, and rebuilds a ``networkx.MultiDiGraph`` from
them with the SAME loader the file-export demo (N3) uses. It then ASSERTS the
Flight-path graph agrees with the file-path graph / oracle — proving the two
delivery paths describe the identical graph.

Run:

    # 1. start the server (separate shell, or backgrounded):
    cargo run -p knut-flight --release --bin knut-flight-server -- --addr 0.0.0.0:50055 --scale 1
    # 2. (optional) produce the N3 file export to cross-check against:
    cargo run -p knut-rel2graph-bench --release --bin rel2graph-export -- 1 target/rel2graph-export
    # 3. run this demo:
    python demo/flight_networkx_demo.py [grpc://HOST:PORT] [EXPORT_DIR]
"""

from __future__ import annotations

import os
import sys

import networkx as nx
import pyarrow.flight as flight

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import knut_networkx as kn  # noqa: E402


def main() -> int:
    location = sys.argv[1] if len(sys.argv) > 1 else "grpc://localhost:50055"
    export_dir = sys.argv[2] if len(sys.argv) > 2 else "target/rel2graph-export"

    client = flight.FlightClient(location)
    print(f"== connected to knut-flight-server at {location} ==")

    # ── list_flights — discover what the server advertises ──────────────────────
    flights = list(client.list_flights())
    print(f"\n-- list_flights: {len(flights)} datasets --")
    names = []
    for fi in flights:
        name = fi.descriptor.command.decode() if fi.descriptor.command else "/".join(
            p.decode() for p in fi.descriptor.path
        )
        names.append(name)
        print(f"  {name:<40} rows={fi.total_records:<7} cols={len(fi.schema)}")
    assert "nodes" in names and "edges" in names, "server must advertise nodes + edges"

    # ── do_get: pull the aggregate nodes + edges as Arrow ───────────────────────
    nodes_table = client.do_get(flight.Ticket(b"nodes")).read_all()
    edges_table = client.do_get(flight.Ticket(b"edges")).read_all()
    print(f"\n-- do_get --")
    print(f"  nodes: {nodes_table.num_rows} rows, schema {nodes_table.schema.names}")
    print(f"  edges: {edges_table.num_rows} rows, schema {edges_table.schema.names}")

    # bonus: pull one source relational table over the SAME Flight surface.
    emp = client.do_get(flight.Ticket(b"table/erp.hr.employees")).read_all()
    print(f"  table/erp.hr.employees: {emp.num_rows} rows, schema {emp.schema.names}")

    # ── build the SAME MultiDiGraph from the Arrow tables ───────────────────────
    g = kn.build_multidigraph(
        kn.nodes_from_arrow(nodes_table),
        kn.edges_from_arrow(edges_table),
    )
    print(f"\n-- networkx.MultiDiGraph from Flight --")
    print(f"  nodes={g.number_of_nodes()} edges={g.number_of_edges()} "
          f"self-loops={nx.number_of_selfloops(g)} "
          f"wcc={nx.number_weakly_connected_components(g)}")

    by_label = kn.node_counts_by_label(g)
    by_type = kn.edge_counts_by_type(g)
    reports_to = [(u, v) for u, v, d in g.edges(data=True) if d["rel"] == "REPORTS_TO"]

    # ── ASSERT the Flight-path graph is self-consistent ─────────────────────────
    print("\n== asserting Flight-path graph ==")
    assert g.is_directed() and g.is_multigraph()
    assert g.number_of_nodes() == nodes_table.num_rows
    assert g.number_of_edges() == edges_table.num_rows
    assert nx.number_of_selfloops(g) == 0, "no self-loops (a7208dc REPORTS_TO fix)"
    assert len(reports_to) > 0 and all(u != v for u, v in reports_to)
    print("  OK — MultiDiGraph, counts match the Arrow tables, no self-loops.")

    # ── ASSERT the Flight path AGREES with the file path (N3), if present ───────
    if os.path.isdir(export_dir):
        f_nodes, f_edges, oracle = kn.load_json_export(export_dir)
        fg = kn.build_multidigraph(f_nodes, f_edges)
        print(f"\n== asserting Flight path == file path ({export_dir}) ==")
        assert g.number_of_nodes() == fg.number_of_nodes() == oracle["total_nodes"]
        assert g.number_of_edges() == fg.number_of_edges() == oracle["total_edges"]
        assert by_label == kn.node_counts_by_label(fg)
        assert by_label == oracle["node_counts"]
        oracle_edges = {(e["src"], e["rel"], e["dst"]): e["count"] for e in oracle["edge_counts"]}
        assert by_type == kn.edge_counts_by_type(fg) == oracle_edges
        assert nx.number_of_selfloops(g) == oracle["self_loops"]
        print("  OK — Flight and file paths build the identical graph, and both")
        print("       match knut's oracle (node/edge totals, per-label, per-type,")
        print("       self-ref REPORTS_TO with no self-loop).")
    else:
        print(f"\n(note: {export_dir} not found — skipping file-path cross-check; "
              "run rel2graph-export to enable it.)")

    print("\nALL ASSERTIONS PASSED.")
    return 0


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