crates/knut-rel2graph-bench/src/harness.rs
//! Drive the bifröst rel2graph engine over generated data and summarise the
//! resulting graph — all through the engine's PUBLIC API, no live DB.
//!
//! The path mirrors what the real [`Listener`](knut_bifrost::Listener) does per
//! round, minus the source/sink: compile each table's declarative pipeline into
//! [`Transform`]s and apply them to the table's Arrow batches, merging every
//! step's output into one [`GraphMutation`]. That mutation — `VertexSet`s +
//! `EdgeSet`s — is the graph, at the exact seam the correctness bencher asserts
//! on (no FalkorDB required).
use std::collections::{BTreeMap, BTreeSet, HashSet};
use arrow_array::{Array, Int64Array, RecordBatch};
use knut_bifrost::{mapping::MappingSpec, mutation::GraphMutation};
use crate::generate::GeneratedDb;
/// Build the full [`GraphMutation`] by running every table's mapping pipeline
/// over its generated batches (append-only / full-load semantics).
///
/// This is a thin wrapper over the engine's canonical seam,
/// [`knut_bifrost::rel2graph::build_mutation`] — the SAME entrypoint knut-flight
/// and korp's `iceberg2graph` step call over real data. The bench differs only in
/// its [`TableSource`](knut_bifrost::rel2graph::TableSource): a seeded synthetic
/// `GeneratedDb` instead of a live Iceberg / parquet read. A compile or transform
/// failure is a bug in the fixed bench spec, so it panics rather than bubbling.
pub fn build_mutation(spec: &MappingSpec, db: &GeneratedDb) -> GraphMutation {
knut_bifrost::rel2graph::build_mutation(spec, db)
.expect("the bench mapping spec compiles and applies over generated data")
}
/// A correctness-focused digest of a [`GraphMutation`].
#[derive(Debug, Clone, Default)]
pub struct GraphSummary {
/// Node label → number of DISTINCT keys (the node count for that label).
pub node_counts: BTreeMap<String, usize>,
/// Node label → the set of key values (used to resolve edge endpoints).
pub node_keys: BTreeMap<String, HashSet<i64>>,
/// `(from_label, rel, to_label)` → count of edges whose BOTH endpoints are
/// non-null and resolve to an existing node.
pub edge_counts: BTreeMap<(String, String, String), usize>,
/// Every node label that appears (as a `VertexSet`).
pub labels: BTreeSet<String>,
/// Edges whose non-null endpoint value does NOT resolve to a node — an
/// integrity violation. Must be 0 for a consistent load.
pub dangling: usize,
/// Edge rows skipped because an endpoint was NULL (a null FK = "no edge",
/// standard SQL semantics — e.g. an org root's `REPORTS_TO`).
pub null_endpoint_rows: usize,
}
impl GraphSummary {
pub fn total_nodes(&self) -> usize {
self.node_counts.values().sum()
}
pub fn total_edges(&self) -> usize {
self.edge_counts.values().sum()
}
pub fn edges_for(&self, from: &str, rel: &str, to: &str) -> usize {
*self
.edge_counts
.get(&(from.to_string(), rel.to_string(), to.to_string()))
.unwrap_or(&0)
}
}
/// Read an Int64 column by name as `Vec<Option<i64>>` (a null stays `None`).
fn i64_opt(batch: &RecordBatch, col: &str) -> Vec<Option<i64>> {
let idx = batch
.schema()
.index_of(col)
.unwrap_or_else(|_| panic!("column `{col}` present in batch"));
let arr = batch
.column(idx)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap_or_else(|| panic!("column `{col}` is Int64"));
(0..arr.len())
.map(|i| {
if arr.is_null(i) {
None
} else {
Some(arr.value(i))
}
})
.collect()
}
/// Summarise a mutation: node key sets first (so edges can be resolved against
/// them), then every edge endpoint checked for null / dangling / resolved.
pub fn summarize(mutation: &GraphMutation) -> GraphSummary {
let mut s = GraphSummary::default();
// Pass 1 — vertices: accumulate distinct keys per label.
for vs in &mutation.vertices {
s.labels.insert(vs.label.clone());
let keys = s.node_keys.entry(vs.label.clone()).or_default();
for v in i64_opt(&vs.batch, &vs.id_column).into_iter().flatten() {
keys.insert(v);
}
}
for (label, keys) in &s.node_keys {
s.node_counts.insert(label.clone(), keys.len());
}
// Pass 2 — edges: resolve both endpoints against the node key sets.
for es in &mutation.edges {
let src = i64_opt(&es.batch, "src");
let dst = i64_opt(&es.batch, "dst");
let src_keys = s.node_keys.get(&es.src_label);
let dst_keys = s.node_keys.get(&es.dst_label);
let key = (
es.src_label.clone(),
es.edge_type.clone(),
es.dst_label.clone(),
);
let mut resolved = 0usize;
for (sv, dv) in src.iter().zip(dst.iter()) {
match (sv, dv) {
(Some(a), Some(b)) => {
let ok_src = src_keys.map(|k| k.contains(a)).unwrap_or(false);
let ok_dst = dst_keys.map(|k| k.contains(b)).unwrap_or(false);
if ok_src && ok_dst {
resolved += 1;
} else {
s.dangling += 1;
}
}
// A null on either endpoint is "no edge", not a dangling edge.
_ => s.null_endpoint_rows += 1,
}
}
*s.edge_counts.entry(key).or_default() += resolved;
}
s
}
// ── expectations computed straight from the generated relational data ──────────
// These read the SOURCE tables (not the mapping output), giving the correctness
// bencher an independent oracle to compare the GraphMutation against.
/// Distinct values of `col` in `table` — the expected node count for a table
/// whose PK is `col`.
pub fn distinct_count(db: &GeneratedDb, table: &str, col: &str) -> usize {
let mut seen = HashSet::new();
for b in db.batches(table) {
for v in i64_opt(b, col).into_iter().flatten() {
seen.insert(v);
}
}
seen.len()
}
/// Count of NON-NULL values of `col` in `table` — the FK cardinality, i.e. the
/// expected edge count for the relationship that FK column projects to.
pub fn nonnull_count(db: &GeneratedDb, table: &str, col: &str) -> usize {
let mut n = 0;
for b in db.batches(table) {
n += i64_opt(b, col).into_iter().flatten().count();
}
n
}