crates/knut-rel2graph-bench/src/export.rs
//! **networkx bridge (N3)** — dump one rel2graph result as a stable, language-
//! neutral node-list + edge-list + correctness ORACLE that Python's networkx can
//! load directly into a [`networkx.MultiDiGraph`].
//!
//! The rel2graph output is a [`GraphMutation`](knut_bifrost::GraphMutation):
//! `VertexSet`s (nodes: label + id column + property columns) and `EdgeSet`s
//! (edges: `src`/`dst` id columns + `src_label`/`edge_type`/`dst_label` + property
//! columns, incl. the composite-key junction edges and the self-referential
//! `Employee-[:REPORTS_TO]->Employee`). Rather than invent a wire format we emit
//! two JSON-lines files plus one oracle JSON:
//!
//! * `nodes.jsonl` — one object per node: `{"label","id","props":{…}}`. Node ids
//! are unique only WITHIN a label, so a networkx node id must be the
//! `(label, id)` PAIR — the loader keys on exactly that.
//! * `edges.jsonl` — one object per RESOLVED edge:
//! `{"src_label","rel","dst_label","src","dst","props":{…}}`. Rows whose
//! endpoint is NULL (a null FK = "no edge", standard SQL — e.g. an org root's
//! `REPORTS_TO`) are dropped here, so the edge-list count equals the oracle's
//! resolved-edge count exactly. `MultiDiGraph` preserves the PARALLEL edges
//! (two distinct rel-types, or two `edge_key_props` combinations, between the
//! same endpoint pair) that a plain `DiGraph` would collapse.
//! * `oracle.json` — the [`GraphSummary`](crate::GraphSummary) totals + per-label
//! node counts + per-`(src,rel,dst)` edge counts, plus `self_loops` (edges whose
//! two endpoints are the SAME node — must be 0: proves the a7208dc self-ref fix
//! did NOT collapse `REPORTS_TO` into every employee pointing at itself). The
//! Python demo ASSERTS the graph it built matches every one of these — a
//! correctness demo, not a pretty print.
//!
//! Only the three Arrow column types the bench schema uses are read:
//! `Int64` / `Float64` / `Utf8`. Anything else serialises as JSON `null` with the
//! column still named, which is enough for the demo and keeps this leaf tiny.
use std::collections::HashSet;
use std::fs::{self, File};
use std::io::{self, BufWriter, Write};
use std::path::Path;
use arrow_array::{Array, Float64Array, Int64Array, RecordBatch, StringArray};
use knut_bifrost::mutation::GraphMutation;
use serde_json::{Map, Value, json};
use crate::GraphSummary;
use crate::generate::GeneratedDb;
use crate::schema::schema;
/// Write the generated **source tables** to `out_dir` as one `<table>.parquet`
/// per table — the on-disk export that stands in for an Iceberg / object-store
/// landing zone. Feeds `knut-flight`'s live-data seam (`Catalog::from_parquet_dir`)
/// and the export bin's `--parquet` mode, so the "read parquet → rel2graph → serve"
/// path exercises the SAME engine over real files rather than the in-RAM generator.
///
/// Returns the paths written (schema order). A table with no rows is skipped.
pub fn write_source_parquet(db: &GeneratedDb, out_dir: &Path) -> io::Result<Vec<String>> {
use parquet::arrow::ArrowWriter;
fs::create_dir_all(out_dir)?;
let mut written = Vec::new();
for t in schema() {
let batches = db.batches(t.name);
if batches.is_empty() {
continue;
}
let path = out_dir.join(format!("{}.parquet", t.name));
let file = File::create(&path)?;
let mut writer = ArrowWriter::try_new(file, batches[0].schema(), None)
.map_err(|e| io::Error::other(e.to_string()))?;
for b in batches {
writer
.write(b)
.map_err(|e| io::Error::other(e.to_string()))?;
}
writer
.close()
.map_err(|e| io::Error::other(e.to_string()))?;
written.push(path.display().to_string());
}
Ok(written)
}
/// What [`export_networkx`] wrote — paths + the headline counts, so a caller
/// (the `rel2graph-export` bin) can print a one-line receipt.
#[derive(Debug, Clone)]
pub struct ExportManifest {
pub nodes_path: String,
pub edges_path: String,
pub oracle_path: String,
pub nodes_written: usize,
pub edges_written: usize,
}
/// Read Arrow cell `row` of `col` as a JSON value (`Int64`/`Float64`/`Utf8`;
/// null → `Value::Null`; any other dtype → `Value::Null`).
fn cell(col: &dyn Array, row: usize) -> Value {
if col.is_null(row) {
return Value::Null;
}
if let Some(a) = col.as_any().downcast_ref::<Int64Array>() {
return json!(a.value(row));
}
if let Some(a) = col.as_any().downcast_ref::<Float64Array>() {
return json!(a.value(row));
}
if let Some(a) = col.as_any().downcast_ref::<StringArray>() {
return json!(a.value(row));
}
Value::Null
}
/// Read one Int64 cell as `Option<i64>` (null → `None`). Panics if the column is
/// not Int64 — `src`/`dst`/id columns are Int64 by construction in this schema.
fn i64_cell(col: &dyn Array, row: usize) -> Option<i64> {
if col.is_null(row) {
return None;
}
let a = col
.as_any()
.downcast_ref::<Int64Array>()
.expect("endpoint / id column is Int64");
Some(a.value(row))
}
/// Every column of `batch` except those in `skip`, as a JSON object for row `row`.
fn props(batch: &RecordBatch, row: usize, skip: &[&str]) -> Value {
let schema = batch.schema();
let mut m = Map::new();
for (i, f) in schema.fields().iter().enumerate() {
let name = f.name().as_str();
if skip.contains(&name) {
continue;
}
m.insert(name.to_string(), cell(batch.column(i).as_ref(), row));
}
Value::Object(m)
}
/// Export `mutation` (+ its independently-computed `summary` oracle) as the three
/// files described in the module docs, under `out_dir` (created if absent).
pub fn export_networkx(
mutation: &GraphMutation,
summary: &GraphSummary,
out_dir: &Path,
) -> io::Result<ExportManifest> {
fs::create_dir_all(out_dir)?;
let nodes_path = out_dir.join("nodes.jsonl");
let edges_path = out_dir.join("edges.jsonl");
let oracle_path = out_dir.join("oracle.json");
// ── nodes.jsonl — dedupe on (label, id) so the count == distinct keys ──────
let mut nw = BufWriter::new(File::create(&nodes_path)?);
let mut seen: HashSet<(String, i64)> = HashSet::new();
let mut nodes_written = 0usize;
for vs in &mutation.vertices {
let schema = vs.batch.schema();
let id_idx = schema
.index_of(&vs.id_column)
.expect("vertex id column present in batch");
let id_col = vs.batch.column(id_idx);
for row in 0..vs.batch.num_rows() {
let Some(id) = i64_cell(id_col.as_ref(), row) else {
continue; // a null PK is not a node
};
if !seen.insert((vs.label.clone(), id)) {
continue; // already emitted this node
}
let obj = json!({
"label": vs.label,
"id": id,
"props": props(&vs.batch, row, &[&vs.id_column]),
});
writeln!(nw, "{obj}")?;
nodes_written += 1;
}
}
nw.flush()?;
// ── edges.jsonl — one line per RESOLVED (both endpoints non-null) edge ─────
let mut ew = BufWriter::new(File::create(&edges_path)?);
let mut edges_written = 0usize;
let mut self_loops = 0usize;
for es in &mutation.edges {
let schema = es.batch.schema();
let src_idx = schema.index_of("src").expect("edge batch has `src`");
let dst_idx = schema.index_of("dst").expect("edge batch has `dst`");
let src_col = es.batch.column(src_idx);
let dst_col = es.batch.column(dst_idx);
for row in 0..es.batch.num_rows() {
let (Some(src), Some(dst)) = (
i64_cell(src_col.as_ref(), row),
i64_cell(dst_col.as_ref(), row),
) else {
continue; // null endpoint = "no edge"
};
// A true self-loop: same label AND same key value on both ends.
if es.src_label == es.dst_label && src == dst {
self_loops += 1;
}
let obj = json!({
"src_label": es.src_label,
"rel": es.edge_type,
"dst_label": es.dst_label,
"src": src,
"dst": dst,
"props": props(&es.batch, row, &["src", "dst"]),
});
writeln!(ew, "{obj}")?;
edges_written += 1;
}
}
ew.flush()?;
// ── oracle.json — the independent correctness oracle the demo asserts on ───
let node_counts: Map<String, Value> = summary
.node_counts
.iter()
.map(|(k, v)| (k.clone(), json!(v)))
.collect();
let edge_counts: Vec<Value> = summary
.edge_counts
.iter()
.map(|((s, r, d), c)| json!({ "src": s, "rel": r, "dst": d, "count": c }))
.collect();
let oracle = json!({
"total_nodes": summary.total_nodes(),
"total_edges": summary.total_edges(),
"dangling": summary.dangling,
"null_endpoint_rows": summary.null_endpoint_rows,
"self_loops": self_loops,
"node_counts": Value::Object(node_counts),
"edge_counts": edge_counts,
});
fs::write(&oracle_path, serde_json::to_vec_pretty(&oracle)?)?;
Ok(ExportManifest {
nodes_path: nodes_path.display().to_string(),
edges_path: edges_path.display().to_string(),
oracle_path: oracle_path.display().to_string(),
nodes_written,
edges_written,
})
}