crates/knut-thund/src/rel2graph.rs
//! **rel2graph → thund**: project a bifröst [`GraphPlan`] into a real thund
//! [`Pipeline`] that any [`ExecBackend`] runs.
//!
//! This is the *second* projection of the two-projections invariant (the first,
//! the PySpark job, lives in `knut_bifrost::codegen`). bifröst maps a relational
//! schema to a graph **once** — the canonical [`GraphPlan`] (labels, merge keys,
//! edge directions, props, batch size). Here that ONE plan becomes a thund
//! [`Pipeline`]: one [`Dataset`] per node/edge (an [`OutputType::Sink`] carrying
//! the graph target in its `knut.*` [`Dataset::properties`]) fed by a projection
//! [`Flow`] over the source table. The resulting pipeline is a normal thund IR —
//! it `validate()`s, round-trips through the RON DSL, and is handed to
//! [`SparkBackend`](crate::backend::spark) (→ SDP) or
//! [`NativeBackend`](crate::backend::native) (→ DataFusion). "bifrost mapping →
//! thund → Spark|Native" is one path; thund picks the tech.
//!
//! **What is a projection vs a worklist item.** The graph *shape* (which nodes,
//! which merge keys, which directed edges, which columns, in canonical order)
//! projects fully and losslessly, and that shape is exactly what must agree with
//! the PySpark job — pinned by [`tests::two_projections_agree`]. The graph
//! *sink write itself* (a backend that turns an `OutputType::Sink` +
//! `knut.node.*`/`knut.edge.*` properties into `UNWIND … MERGE`) is a backend
//! capability not yet built — see the design doc's worklist. Until then the
//! Pipeline is the runnable, reviewable, diffable **description**; the PySpark
//! job is the executable Spark path.
//!
//! [`GraphPlan`]: knut_bifrost::GraphPlan
//! [`ExecBackend`]: crate::backend::ExecBackend
use std::collections::BTreeMap;
use knut_bifrost::{GraphPlan, PlanMode};
use crate::error::Result;
use crate::ir::{Dataset, Flow, Materialization, OutputType, Pipeline};
/// The property key namespace the graph target is carried under on each
/// [`Dataset`] — read by a (future) graph-sink `ExecBackend`, and by the
/// two-projections-agree test.
pub mod prop {
/// Node label.
pub const NODE_LABEL: &str = "knut.node.label";
/// Node merge-key column.
pub const NODE_KEY: &str = "knut.node.key";
/// Node property columns (comma-joined, ordinal).
pub const NODE_PROPS: &str = "knut.node.props";
/// Source table the node rows come from.
pub const NODE_TABLE: &str = "knut.node.table";
/// Edge relationship type.
pub const EDGE_REL: &str = "knut.edge.rel";
/// Edge tail (child) label.
pub const EDGE_FROM: &str = "knut.edge.from";
/// Edge head (parent) label.
pub const EDGE_TO: &str = "knut.edge.to";
/// The tail node's **merge-key property** — the property the `from` node is
/// MERGEd by (its `id_column`). NOT necessarily the column the value is read
/// from; see [`EDGE_FROM_COL`].
pub const EDGE_FROM_KEY: &str = "knut.edge.from_key";
/// The head node's **merge-key property** — the property the `to` node is
/// MERGEd by (its `id_column`). Differs from [`EDGE_TO_COL`] for a self-ref /
/// renamed FK (MERGE on `employee_id`, value read from `manager_id`).
pub const EDGE_TO_KEY: &str = "knut.edge.to_key";
/// The column in the edge table the **tail** endpoint VALUE is read from
/// (`src.column`). The projection flow projects this column (renamed to
/// `src` by the graph-sink backend). Usually equals [`EDGE_FROM_KEY`].
pub const EDGE_FROM_COL: &str = "knut.edge.from_col";
/// The column in the edge table the **head** endpoint VALUE is read from
/// (the FK column, `dst.column`). The projection flow projects THIS column,
/// not [`EDGE_TO_KEY`] — the two differ for a self-ref / renamed FK.
pub const EDGE_TO_COL: &str = "knut.edge.to_col";
/// Source table the edge rows come from.
pub const EDGE_TABLE: &str = "knut.edge.table";
/// Rows per `UNWIND … MERGE` batch (mirrors the PySpark `BATCH_ROWS`).
pub const BATCH_ROWS: &str = "knut.graph.batch_rows";
/// The resolution tier (`T1`…`T4`) the element was resolved at.
pub const TIER: &str = "knut.graph.tier";
/// The graph sink token (`falkordb` / `graphar`).
pub const SINK: &str = "knut.graph.sink";
/// The source plane token (`postgres` / `iceberg`).
pub const SOURCE: &str = "knut.graph.source";
}
/// **Project a [`GraphPlan`] into a thund [`Pipeline`].**
///
/// One [`Dataset`] per node (`node__<Label>`) and per edge
/// (`edge__<From>__<Rel>__<To>`), each an [`OutputType::Sink`] whose
/// `knut.*` properties carry the graph target, fed by a projection [`Flow`]
/// reading the source table's needed columns. Nodes precede edges (declaration
/// order follows the plan's canonical order). Streaming plans mark datasets
/// [`Materialization::Incremental`]; batch plans [`Materialization::Full`].
///
/// The returned pipeline is `validate()`d. T4 proposals are **not** projected —
/// they are never applied (they remain on [`GraphPlan::proposals`] for a
/// reviewer, and are emitted commented-out only by the PySpark projection).
pub fn to_thund_pipeline(plan: &GraphPlan) -> Result<Pipeline> {
let mut p = Pipeline::new("rel2graph");
p.storage = Some(plan.checkpoint.clone());
let mat = match plan.mode {
PlanMode::Stream => Materialization::Incremental,
PlanMode::Batch => Materialization::Full,
};
let sink_fmt = plan.sink.token();
let batch = plan.batch_rows.to_string();
for n in &plan.nodes {
let ds_name = node_dataset(&n.label);
let mut props: BTreeMap<String, String> = BTreeMap::new();
props.insert(prop::NODE_LABEL.into(), n.label.clone());
props.insert(prop::NODE_KEY.into(), n.key.clone());
props.insert(prop::NODE_PROPS.into(), n.props.join(","));
props.insert(prop::NODE_TABLE.into(), n.table.clone());
props.insert(prop::BATCH_ROWS.into(), batch.clone());
props.insert(prop::TIER.into(), n.tier.token().into());
props.insert(prop::SINK.into(), sink_fmt.into());
props.insert(prop::SOURCE.into(), plan.source.token().into());
let ds = Dataset::new(ds_name.clone(), OutputType::Sink)
.with_format(format!("{sink_fmt}-node"))
.with_properties(props);
let ds = with_materialization(ds, mat);
p = p.with_dataset(ds);
// Projection flow: read the source table, prune to key + props.
let mut cols = vec![n.key.clone()];
cols.extend(n.props.iter().cloned());
p = p.with_flow(
Flow::batch(format!("flow_{ds_name}"), ds_name, [n.table.clone()])
.with_projection(cols),
);
}
for e in &plan.edges {
let ds_name = edge_dataset(&e.from_label, &e.rel, &e.to_label);
let mut props: BTreeMap<String, String> = BTreeMap::new();
props.insert(prop::EDGE_REL.into(), e.rel.clone());
props.insert(prop::EDGE_FROM.into(), e.from_label.clone());
props.insert(prop::EDGE_TO.into(), e.to_label.clone());
props.insert(prop::EDGE_FROM_KEY.into(), e.from_key.clone());
props.insert(prop::EDGE_TO_KEY.into(), e.to_key.clone());
props.insert(prop::EDGE_FROM_COL.into(), e.from_col.clone());
props.insert(prop::EDGE_TO_COL.into(), e.to_col.clone());
props.insert(prop::EDGE_TABLE.into(), e.table.clone());
props.insert(prop::BATCH_ROWS.into(), batch.clone());
props.insert(prop::TIER.into(), e.tier.token().into());
props.insert(prop::SINK.into(), sink_fmt.into());
props.insert(prop::SOURCE.into(), plan.source.token().into());
let ds = Dataset::new(ds_name.clone(), OutputType::Sink)
.with_format(format!("{sink_fmt}-edge"))
.with_properties(props);
let ds = with_materialization(ds, mat);
p = p.with_dataset(ds);
// Projection flow: read the edge table, prune to the two endpoint
// VALUE columns (the FK columns `from_col`/`to_col`, NOT the merge-key
// property names) + edge props. For a self-ref / renamed FK the value
// lives in a different column than the node key (`manager_id` read, node
// MERGEd on `employee_id`); projecting the merge key here would read the
// wrong column and collapse the edge to a self-loop.
let mut cols = vec![e.from_col.clone(), e.to_col.clone()];
cols.extend(e.props.iter().cloned());
p = p.with_flow(
Flow::batch(format!("flow_{ds_name}"), ds_name, [e.table.clone()])
.with_projection(cols),
);
}
p.validate()?;
Ok(p)
}
/// Project a [`GraphPlan`] into a thund [`Pipeline`] rendered as **RON** — the
/// `pipeline.ron` artifact, a first-class thund DSL spec (requires the `dsl`
/// feature, which `rel2graph` enables). Round-trips through
/// [`crate::authoring::dsl::from_ron`].
pub fn to_pipeline_ron(plan: &GraphPlan) -> Result<String> {
let p = to_thund_pipeline(plan)?;
crate::authoring::dsl::to_ron(&p)
}
// ── helpers ─────────────────────────────────────────────────────────────────
fn with_materialization(ds: Dataset, mat: Materialization) -> Dataset {
match mat {
Materialization::Incremental => ds.incremental(),
Materialization::Full => ds, // Full is the default
}
}
/// A stable, identifier-safe dataset name for a node label.
fn node_dataset(label: &str) -> String {
format!("node__{}", sanitize(label))
}
/// A stable, identifier-safe dataset name for a directed edge.
fn edge_dataset(from: &str, rel: &str, to: &str) -> String {
format!(
"edge__{}__{}__{}",
sanitize(from),
sanitize(rel),
sanitize(to)
)
}
fn sanitize(s: &str) -> String {
s.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect()
}
// ═══════════════════════════════════════════════════════════════════════════
// Tests — THE two-projections-must-agree invariant lives here, where BOTH
// projections are visible: the PySpark job (bifröst) and the thund Pipeline.
// ═══════════════════════════════════════════════════════════════════════════
#[cfg(test)]
mod tests {
use super::*;
use knut_bifrost::MappingSpec;
use knut_bifrost::codegen::CodegenOptions;
fn sample_plan(opts: &CodegenOptions) -> GraphPlan {
// Declaration order is deliberately unsorted so the canonical ordering
// is exercised on both projections identically.
let yaml = r#"
version: 1
tables:
- table: lake.social.users
pipeline:
- vertex: { label: User, id_column: user_id, properties: [name, age] }
- fk_edges:
type: WORKS_AT
src: { label: User, column: user_id }
dst: { label: Company, column: company_id }
properties: [since]
- table: lake.social.companies
pipeline:
- vertex: { label: Company, id_column: company_id, properties: [name] }
"#;
let spec = MappingSpec::from_yaml(yaml).unwrap();
GraphPlan::from_mapping(&spec, opts).unwrap()
}
/// (label, key) node facts parsed back out of the PySpark job's startup
/// index statements — the Spark projection's ground truth.
fn nodes_from_pyspark(py: &str) -> Vec<(String, String)> {
let mut out = Vec::new();
for line in py.lines() {
// ... CREATE INDEX FOR (n:`Label`) ON (n.`key`) ...
if let Some(rest) = line.split_once("CREATE INDEX FOR (n:`") {
let (label, tail) = rest.1.split_once("`) ON (n.`").unwrap();
let key = tail.split_once("`)").unwrap().0;
out.push((label.to_string(), key.to_string()));
}
}
out.sort();
out
}
/// (from, rel, to) edge facts parsed back out of the PySpark job's
/// `_write_edges("From", "REL", "To", …)` calls.
fn edges_from_pyspark(py: &str) -> Vec<(String, String, String)> {
let mut out = Vec::new();
for line in py.lines() {
// Skip the writer's *definition* (`def _write_edges(fl, rel, tl…)`),
// whose args are unquoted params, not the call sites.
if line.trim_start().starts_with("def ") {
continue;
}
if let Some((_, rest)) = line.split_once("_write_edges(") {
// Args are Rust {:?} string literals: "From", "REL", "To", ...
let args: Vec<&str> = rest.split(", ").collect();
let unq = |s: &str| s.trim().trim_matches('"').to_string();
out.push((unq(args[0]), unq(args[1]), unq(args[2])));
}
}
out.sort();
out
}
fn batch_from_pyspark(py: &str) -> String {
py.lines()
.find_map(|l| l.strip_prefix("BATCH_ROWS = "))
.unwrap()
.to_string()
}
fn nodes_from_pipeline(p: &Pipeline) -> Vec<(String, String)> {
let mut out: Vec<(String, String)> = p
.datasets
.iter()
.filter_map(|d| {
Some((
d.properties.get(prop::NODE_LABEL)?.clone(),
d.properties.get(prop::NODE_KEY)?.clone(),
))
})
.collect();
out.sort();
out
}
fn edges_from_pipeline(p: &Pipeline) -> Vec<(String, String, String)> {
let mut out: Vec<(String, String, String)> = p
.datasets
.iter()
.filter_map(|d| {
Some((
d.properties.get(prop::EDGE_FROM)?.clone(),
d.properties.get(prop::EDGE_REL)?.clone(),
d.properties.get(prop::EDGE_TO)?.clone(),
))
})
.collect();
out.sort();
out
}
fn batch_from_pipeline(p: &Pipeline) -> String {
p.datasets
.iter()
.find_map(|d| d.properties.get(prop::BATCH_ROWS).cloned())
.unwrap()
}
/// **THE invariant.** The PySpark job and the thund Pipeline are two
/// projections of the SAME [`GraphPlan`], and they must agree on labels,
/// merge keys, edge directions, and batch size. We derive each set
/// INDEPENDENTLY — parsed from the generated Python on one side, read from
/// the Pipeline's `knut.*` properties on the other — and assert equality.
/// A divergence in either emitter (a renamed label, a flipped edge, a
/// different batch) fails here.
#[test]
fn two_projections_agree() {
let plan = sample_plan(&CodegenOptions::default());
let py = plan.to_pyspark();
let pipe = to_thund_pipeline(&plan).unwrap();
assert_eq!(
nodes_from_pyspark(&py),
nodes_from_pipeline(&pipe),
"node labels+keys agree across the two projections"
);
assert_eq!(
edges_from_pyspark(&py),
edges_from_pipeline(&pipe),
"edge directions (from,rel,to) agree across the two projections"
);
assert_eq!(
batch_from_pyspark(&py),
batch_from_pipeline(&pipe),
"batch size agrees across the two projections"
);
// And both agree with the plan's own canonical contract.
let contract = plan.contract();
assert!(contract.contains("node User=user_id"));
assert!(contract.contains("node Company=company_id"));
assert!(contract.contains("edge User-WORKS_AT->Company"));
assert!(contract.contains("batch_rows=1000"));
// Non-empty, so the assertions above are not vacuously true.
assert_eq!(nodes_from_pyspark(&py).len(), 2);
assert_eq!(edges_from_pyspark(&py).len(), 1);
}
/// The thund Pipeline projection is a REAL, valid thund IR: it validates,
/// its graph target rides in `knut.*` properties, nodes precede edges, and
/// it round-trips through the RON DSL losslessly (so `pipeline.ron` is a
/// first-class thund spec an `ExecBackend` consumes).
#[test]
fn pipeline_is_real_thund_ir_and_round_trips() {
let plan = sample_plan(&CodegenOptions::default());
let pipe = to_thund_pipeline(&plan).unwrap();
pipe.validate()
.expect("the projected pipeline is valid thund IR");
// Nodes-before-edges in declaration order.
let names: Vec<&str> = pipe.datasets.iter().map(|d| d.name.as_str()).collect();
assert_eq!(
names,
[
"node__Company",
"node__User",
"edge__User__WORKS_AT__Company"
]
);
// Every dataset is a graph Sink carrying its target.
assert!(
pipe.datasets
.iter()
.all(|d| d.output_type == OutputType::Sink)
);
// The projection flows prune to exactly the needed columns.
let user_flow = pipe
.flows
.iter()
.find(|f| f.target == "node__User")
.unwrap();
assert_eq!(user_flow.projection, vec!["user_id", "name", "age"]);
assert_eq!(user_flow.reads, vec!["lake.social.users"]);
// Round-trip through the RON DSL (the pipeline.ron artifact).
let ron = to_pipeline_ron(&plan).unwrap();
let back = crate::authoring::dsl::from_ron(&ron).unwrap();
assert_eq!(back, pipe, "pipeline.ron round-trips losslessly");
// The graph target genuinely reached the RON.
assert!(ron.contains("knut.node.label"));
assert!(ron.contains("WORKS_AT"));
}
/// A streaming plan projects to incremental datasets (and still validates /
/// round-trips), matching the PySpark `foreachBatch` streaming path.
#[test]
fn streaming_plan_projects_incremental() {
use knut_bifrost::{PlanMode, SourceKind};
let opts = CodegenOptions {
source: SourceKind::Iceberg,
mode: PlanMode::Stream,
..Default::default()
};
let plan = sample_plan(&opts);
let pipe = to_thund_pipeline(&plan).unwrap();
assert!(
pipe.datasets
.iter()
.all(|d| d.materialization == Materialization::Incremental),
"streaming plan → incremental datasets"
);
// Still agrees with the Spark projection's streaming form.
assert!(plan.to_pyspark().contains("foreachBatch("));
}
}