demo/knut_networkx.py
"""
knut_networkx — the tiny loader that turns a knut rel2graph result into a
``networkx.MultiDiGraph``.
A rel2graph result is a set of **nodes** (each has a ``label``, an integer ``id``
that is unique only WITHIN its label, and a bag of properties) and **edges**
(each has ``src_label``/``rel``/``dst_label`` plus the endpoint ids ``src``/``dst``
and a bag of properties). Because an id is unique only within a label, the
networkx node identifier is the **pair** ``(label, id)`` — so ``Order 1`` and
``Employee 1`` are distinct nodes and never collide.
A ``MultiDiGraph`` (directed + multi) is the right target: it keeps direction
(``a -> b`` is not ``b -> a``) and preserves **parallel edges** — two different
rel-types between the same pair, or two rows of the same rel-type — which a plain
``DiGraph`` would silently collapse to one.
Two input paths feed the SAME builder, which is the whole point of N4's proof:
* **file path (N3)** — ``load_json_export(dir)`` reads the ``nodes.jsonl`` /
``edges.jsonl`` / ``oracle.json`` that the Rust ``rel2graph-export`` bin wrote.
* **Flight path (N4)** — ``nodes_from_arrow`` / ``edges_from_arrow`` turn the
Arrow tables pulled from ``knut-flight-server`` into the same node/edge dicts.
Both hand their dicts to :func:`build_multidigraph`, so a graph built from files
and a graph built from Flight are the same object shape and compare equal.
Depends only on ``networkx`` (+ the stdlib). The Arrow helpers accept an already
materialised ``pyarrow.Table`` — ``pyarrow`` is imported by the caller, not here.
"""
from __future__ import annotations
import json
import os
from typing import Any, Iterable
import networkx as nx
# ── node id ────────────────────────────────────────────────────────────────────
NodeId = tuple[str, int]
def node_id(label: str, id_: int) -> NodeId:
"""The networkx node identifier: the ``(label, id)`` pair (ids are unique
only within a label, so the label must be part of the key)."""
return (str(label), int(id_))
# ── file path (N3): JSON-lines + oracle ─────────────────────────────────────────
def _read_jsonl(path: str) -> list[dict[str, Any]]:
out: list[dict[str, Any]] = []
with open(path, "r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if line:
out.append(json.loads(line))
return out
def load_json_export(export_dir: str) -> tuple[list[dict], list[dict], dict]:
"""Load ``(nodes, edges, oracle)`` from a ``rel2graph-export`` directory."""
nodes = _read_jsonl(os.path.join(export_dir, "nodes.jsonl"))
edges = _read_jsonl(os.path.join(export_dir, "edges.jsonl"))
with open(os.path.join(export_dir, "oracle.json"), "r", encoding="utf-8") as fh:
oracle = json.load(fh)
return nodes, edges, oracle
# ── Flight path (N4): Arrow tables → the same node/edge dicts ────────────────────
def nodes_from_arrow(table: "Any") -> list[dict]:
"""Normalise a pyarrow node table (columns ``label``, ``id``, + optional
property columns) into the node-dict shape :func:`build_multidigraph` wants."""
out = []
for row in table.to_pylist():
label = row.pop("label")
id_ = row.pop("id")
out.append({"label": label, "id": id_, "props": row})
return out
def edges_from_arrow(table: "Any") -> list[dict]:
"""Normalise a pyarrow edge table (columns ``src_label``, ``rel``,
``dst_label``, ``src``, ``dst``, + optional property columns) into the
edge-dict shape :func:`build_multidigraph` wants."""
out = []
for row in table.to_pylist():
rec = {
"src_label": row.pop("src_label"),
"rel": row.pop("rel"),
"dst_label": row.pop("dst_label"),
"src": row.pop("src"),
"dst": row.pop("dst"),
}
rec["props"] = row
out.append(rec)
return out
# ── the one builder both paths share ─────────────────────────────────────────────
def build_multidigraph(nodes: Iterable[dict], edges: Iterable[dict]) -> nx.MultiDiGraph:
"""Build a ``MultiDiGraph`` from node/edge dicts.
* node identifier = ``(label, id)``; node attributes = ``label`` + the
node's ``props``.
* edge = ``(src_label, src) -> (dst_label, dst)`` with attributes ``rel`` +
the edge's ``props``. Parallel edges (same endpoints, different rel or row)
are preserved — ``MultiDiGraph`` assigns each its own key.
"""
g = nx.MultiDiGraph()
for n in nodes:
nid = node_id(n["label"], n["id"])
g.add_node(nid, label=str(n["label"]), **n.get("props", {}))
for e in edges:
src = node_id(e["src_label"], e["src"])
dst = node_id(e["dst_label"], e["dst"])
# Endpoints may not have been in the node list (defensive): add_edge
# creates them, but tag the label so the by-label census stays honest.
if src not in g:
g.add_node(src, label=str(e["src_label"]))
if dst not in g:
g.add_node(dst, label=str(e["dst_label"]))
g.add_edge(src, dst, rel=str(e["rel"]), **e.get("props", {}))
return g
# ── census helpers the demos assert on ───────────────────────────────────────────
def node_counts_by_label(g: nx.MultiDiGraph) -> dict[str, int]:
counts: dict[str, int] = {}
for _, data in g.nodes(data=True):
counts[data["label"]] = counts.get(data["label"], 0) + 1
return counts
def edge_counts_by_type(g: nx.MultiDiGraph) -> dict[tuple[str, str, str], int]:
"""``(src_label, rel, dst_label) -> count``, read back off the built graph."""
counts: dict[tuple[str, str, str], int] = {}
for u, v, data in g.edges(data=True):
key = (u[0], data["rel"], v[0])
counts[key] = counts.get(key, 0) + 1
return counts