crates/knut-bifrost/src/codegen.rs
//! **rel2graph CODEGEN** — one [`MappingSpec`] projected onto executable
//! backends.
//!
//! bifröst **is** the rel2graph engine: a relational schema → a *mapping spec*
//! (nodes, merge keys, edges, directions, props) → materialised graph. Codegen
//! is *one tool* of that engine (its siblings: the live CDC listener, the
//! pipeline-config authoring, the schema→mapping inference in
//! [`crate::mapping::MappingSpec::from_arrow_schema`]). This module is the
//! **describe-first, execute-anywhere** half: it lowers a reviewed
//! [`MappingSpec`] into a single canonical [`GraphPlan`] and then *projects*
//! that plan onto a backend.
//!
//! ## The two-projections invariant
//!
//! A [`GraphPlan`] is projected two ways, and **both must agree exactly**:
//!
//! * [`GraphPlan::to_pyspark`] — a PySpark job that loads the schema into the
//! graph sink (the **Spark** backend projection). Lives here, pure Rust.
//! * a **thund [`Pipeline`]** the all-Rust engine runs on *any*
//! [`ExecBackend`] (Spark-Connect-SDP **or** native DataFusion) — lives in
//! `knut-thund` (`knut_thund::rel2graph`, feature `rel2graph`), because that
//! crate owns the execution IR. It projects from *this same* [`GraphPlan`].
//!
//! bifröst maps rel→graph **once**; thund picks the tech. Any divergence
//! between the two projections (a label, a merge key, an edge direction, the
//! batch size) is a bug — pinned by the `two_projections_agree` test in
//! `knut-thund` (which can see both projections without a dependency cycle:
//! `knut-pipelines` already depends on `knut-bifrost`, so bifröst must not
//! depend on thund; thund depending on bifröst is the acyclic direction).
//!
//! [`MappingSpec`]: crate::mapping::MappingSpec
//! [`Pipeline`]: https://docs.rs/knut-thund
//! [`ExecBackend`]: https://docs.rs/knut-thund
use std::fmt::Write as _;
use crate::error::{BifrostError, Result};
use crate::mapping::{MappingSpec, PipelineStep, SyncMode};
// ═══════════════════════════════════════════════════════════════════════════
// Options + backend/mode enums
// ═══════════════════════════════════════════════════════════════════════════
/// The relational **source plane** a schema is read from. In the intended
/// deployment Iceberg is the *data plane* and Postgres the *schema plane*
/// (the T1 catalog), so both are first-class.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SourceKind {
/// A Postgres (or Postgres-catalog-fronted) source.
Postgres,
/// An Iceberg table (the streaming-capable data plane).
Iceberg,
}
impl SourceKind {
/// The lower-case token used in generated artifacts / RON.
pub fn token(self) -> &'static str {
match self {
SourceKind::Postgres => "postgres",
SourceKind::Iceberg => "iceberg",
}
}
}
/// The graph **sink** a plan targets.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SinkKind {
/// FalkorDB (Cypher `UNWIND … MERGE`, the runtime sink bifröst already
/// drives in [`crate::sink`]).
FalkorDb,
/// GraphAr columnar graph files.
GraphAr,
}
impl SinkKind {
/// The lower-case token used in generated artifacts / RON.
pub fn token(self) -> &'static str {
match self {
SinkKind::FalkorDb => "falkordb",
SinkKind::GraphAr => "graphar",
}
}
}
/// Batch vs streaming. Streaming requires an Iceberg source (the changelog
/// scan) — enforced by [`GraphPlan::from_mapping`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PlanMode {
/// A bounded, idempotent full/incremental load.
Batch,
/// A continuous changelog-driven load (needs Iceberg).
Stream,
}
impl PlanMode {
/// The lower-case token used in generated artifacts / RON.
pub fn token(self) -> &'static str {
match self {
PlanMode::Batch => "batch",
PlanMode::Stream => "stream",
}
}
}
/// **Resolution tier** — where a node/edge mapping's authority *came from*, per
/// the rel2graph spec. Recorded on every plan element and surfaced in the
/// WARNINGS provenance block so a reviewer can see how confident each mapping
/// is (T1 catalog fact ▸ T4 inference proposal).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Tier {
/// **T1** — the upstream Postgres catalog (PK/FK constraints). The answer
/// for most schemas: Iceberg is the data plane, Postgres the schema plane.
T1PostgresCatalog,
/// **T2** — Iceberg identifier-field-ids (the table's declared identity).
T2IcebergFieldIds,
/// **T3** — `knut.*` TBLPROPERTIES on the table
/// (`knut.node.label`/`knut.node.key`/`knut.edge.*`). What a hand-authored
/// or reviewed [`MappingSpec`] encodes.
T3TblProperties,
/// **T4** — inclusion-dependency *inference* (the LEFT-ANTI-JOIN orphan
/// test). **Proposals only** — emitted COMMENTED-OUT with evidence, never
/// applied. A T4 element must never appear in the active plan.
T4InclusionDependency,
}
impl Tier {
/// A stable short token (`T1`…`T4`).
pub fn token(self) -> &'static str {
match self {
Tier::T1PostgresCatalog => "T1",
Tier::T2IcebergFieldIds => "T2",
Tier::T3TblProperties => "T3",
Tier::T4InclusionDependency => "T4",
}
}
/// A one-line human description of the tier.
pub fn describe(self) -> &'static str {
match self {
Tier::T1PostgresCatalog => "upstream Postgres catalog (PK/FK constraints)",
Tier::T2IcebergFieldIds => "Iceberg identifier-field-ids",
Tier::T3TblProperties => "knut.* TBLPROPERTIES on the table",
Tier::T4InclusionDependency => {
"inclusion-dependency inference (LEFT ANTI JOIN orphan test) — PROPOSAL only"
}
}
}
}
/// Knobs for [`GraphPlan::from_mapping`]. Defaults match the rel2graph spec
/// (1000-row UNWIND batches, single coalesced writer, Iceberg→FalkorDB batch).
#[derive(Debug, Clone)]
pub struct CodegenOptions {
/// The source plane the schema was read from.
pub source: SourceKind,
/// The graph sink to load.
pub sink: SinkKind,
/// Batch or streaming (streaming needs an Iceberg source).
pub mode: PlanMode,
/// Rows per `UNWIND … MERGE` batch (the idempotent write unit).
pub batch_rows: usize,
/// Writer concurrency — the `coalesce(N)` partition count the sink client
/// runs inside.
pub concurrency: usize,
/// The single checkpoint / storage location (one per query).
pub checkpoint: String,
/// The tier the *declared* mapping elements are attributed to. A reviewed
/// [`MappingSpec`] is T3 (its `knut.*` properties); a spec derived straight
/// from an Iceberg schema is T2. The live-catalog T1 resolution is a
/// worklist item (see the design doc).
pub declared_tier: Tier,
}
impl Default for CodegenOptions {
fn default() -> Self {
CodegenOptions {
source: SourceKind::Iceberg,
sink: SinkKind::FalkorDb,
mode: PlanMode::Batch,
batch_rows: 1000,
concurrency: 1,
checkpoint: "/tmp/rel2graph/checkpoint".to_string(),
declared_tier: Tier::T3TblProperties,
}
}
}
// ═══════════════════════════════════════════════════════════════════════════
// The canonical plan — the ONE object both projections derive from
// ═══════════════════════════════════════════════════════════════════════════
/// A node projection: one merge-key table → one graph node.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct NodePlan {
/// The fully-qualified source table.
pub table: String,
/// The graph node label (PascalCase singular).
pub label: String,
/// The merge key column — node identity.
pub key: String,
/// The non-key, non-FK scalar property columns, in **ordinal** order.
pub props: Vec<String>,
/// Where this node's authority came from.
pub tier: Tier,
}
/// An edge projection: a directed relationship `from → to`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EdgePlan {
/// The source table the edge rows are read from.
pub table: String,
/// The tail (child / "many") node label.
pub from_label: String,
/// The relationship type (SCREAMING_SNAKE).
pub rel: String,
/// The head (parent / "one") node label.
pub to_label: String,
/// The tail node's **merge-key property** — the property name the `from`
/// node is identified/MERGEd by in the graph (its `id_column`). This is the
/// property NAME on the node, which is NOT necessarily the column the value
/// is read from — see [`Self::from_col`].
pub from_key: String,
/// The head node's **merge-key property** — the property name the `to` node
/// is identified/MERGEd by in the graph (its `id_column`). For a self-ref /
/// renamed FK this differs from the column the value is read from
/// ([`Self::to_col`]): e.g. `Employee-REPORTS_TO->Employee` MERGEs on
/// `employee_id` but reads the endpoint value from the `manager_id` column.
pub to_key: String,
/// The column in `table` the **tail** endpoint VALUE is read from
/// (`src.column`). Usually equal to [`Self::from_key`]; differs only when the
/// tail FK column is named differently from the node's key.
pub from_col: String,
/// The column in `table` the **head** endpoint VALUE is read from
/// (`dst.column`, i.e. the FK column). Differs from [`Self::to_key`] for a
/// self-ref / renamed FK (`manager_id` read, MERGEd on `employee_id`).
pub to_col: String,
/// Edge property columns, ordinal order.
pub props: Vec<String>,
/// Where this edge's authority came from.
pub tier: Tier,
}
/// A **T4 proposal** — an inferred relationship the codegen must NOT apply,
/// only surface commented-out with its evidence. Carried through the plan so
/// both projections emit it identically (as a comment), never as live code.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct EdgeProposal {
/// The candidate edge (its `tier` is always [`Tier::T4InclusionDependency`]).
pub edge: EdgePlan,
/// Human-readable evidence (e.g. the orphan-row count of the LEFT ANTI
/// JOIN test) a reviewer confirms before promoting it.
pub evidence: String,
}
/// The canonical rel2graph plan: **the single source both projections read**.
///
/// Deterministic by construction (see [`GraphPlan::from_mapping`]):
/// tables schema-then-name, node props ordinal, nodes-before-edges, edges
/// sorted `(from_label, rel, to_label)`. No timestamps, no randomness, no
/// hashmap iteration order — the same input is the same bytes, always.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GraphPlan {
/// Source plane.
pub source: SourceKind,
/// Graph sink.
pub sink: SinkKind,
/// Batch/streaming.
pub mode: PlanMode,
/// Rows per `UNWIND … MERGE` batch.
pub batch_rows: usize,
/// Writer concurrency (`coalesce(N)`).
pub concurrency: usize,
/// The single checkpoint / storage location.
pub checkpoint: String,
/// Nodes, in canonical order.
pub nodes: Vec<NodePlan>,
/// Edges, in canonical order.
pub edges: Vec<EdgePlan>,
/// T4 inference proposals — emitted commented-out, never applied.
pub proposals: Vec<EdgeProposal>,
/// Provenance + honesty notes, closing every artifact.
pub warnings: Vec<String>,
}
impl GraphPlan {
/// **Lower a reviewed [`MappingSpec`] into the canonical plan.**
///
/// Walks tables in schema-then-name order; each `vertex` step becomes a
/// [`NodePlan`], each `edge`/`fk_edges` step a [`EdgePlan`] (direction =
/// `src → dst`, the child→parent FK direction). `custom` steps are opaque
/// and cannot be projected to a declarative graph load, so they are
/// SURFACED in [`warnings`](Self::warnings), never silently dropped.
///
/// Errors on the blocking gaps the spec names: a streaming plan on a
/// non-Iceberg source, or a node with no merge key.
pub fn from_mapping(spec: &MappingSpec, opts: &CodegenOptions) -> Result<GraphPlan> {
if opts.mode == PlanMode::Stream && opts.source != SourceKind::Iceberg {
return Err(BifrostError::Other(format!(
"MODE/SOURCE conflict: streaming needs an Iceberg source, got {:?}",
opts.source
)));
}
let mut nodes: Vec<NodePlan> = Vec::new();
let mut edges: Vec<EdgePlan> = Vec::new();
let mut warnings: Vec<String> = Vec::new();
// Tables schema-then-name (namespace before leaf), deterministically.
let mut tables: Vec<&crate::mapping::TableMapping> = spec.tables.iter().collect();
tables.sort_by(|a, b| table_sort_key(&a.table).cmp(&table_sort_key(&b.table)));
for tm in tables {
if tm.mode == SyncMode::Cdc && opts.mode == PlanMode::Batch {
warnings.push(format!(
"table `{}` declares CDC sync but the plan mode is batch — the changelog \
is read as a full insert load (see MODE)",
tm.table
));
}
for step in &tm.pipeline {
match step {
PipelineStep::Vertex { vertex } => {
if vertex.id_column.is_empty() {
return Err(BifrostError::Other(format!(
"node `{}` (table `{}`) has no merge key",
vertex.label, tm.table
)));
}
nodes.push(NodePlan {
table: tm.table.clone(),
label: vertex.label.clone(),
key: vertex.id_column.clone(),
props: vertex.properties.clone(),
tier: opts.declared_tier,
});
}
PipelineStep::Edge { edge: t } | PipelineStep::FkEdges { fk_edges: t } => {
edges.push(EdgePlan {
table: tm.table.clone(),
from_label: t.src.label.clone(),
rel: t.edge_type.clone(),
to_label: t.dst.label.clone(),
from_key: t
.src
.id_column
.clone()
.unwrap_or_else(|| t.src.column.clone()),
to_key: t
.dst
.id_column
.clone()
.unwrap_or_else(|| t.dst.column.clone()),
// The VALUE is always read from the endpoint's own
// `column` (the transform fabric reads `src.column`
// / `dst.column`); only the MERGE *property name*
// comes from `id_column`. Keeping these separate is
// what makes self-ref / renamed FKs correct.
from_col: t.src.column.clone(),
to_col: t.dst.column.clone(),
props: t.properties.clone(),
tier: opts.declared_tier,
});
}
PipelineStep::Sql { .. } => {
warnings.push(format!(
"table `{}` uses a `sql:` step — its output is decided by the query \
at run time and cannot be projected to a declarative graph load; wire \
it by hand on the target backend (NOT dropped)",
tm.table
));
}
PipelineStep::Custom { custom } => {
warnings.push(format!(
"table `{}` uses custom transform `{custom}` — an opaque Rust step \
that cannot be projected to a declarative graph load; wire it by \
hand on the target backend (NOT dropped)",
tm.table
));
}
}
}
}
// DETERMINISM: edges sorted (from_label, rel, to_label). Nodes keep the
// schema-then-name table order (one merge-key table → one node); props
// stay ordinal.
edges.sort_by(|a, b| {
(&a.from_label, &a.rel, &a.to_label).cmp(&(&b.from_label, &b.rel, &b.to_label))
});
// Provenance summary — the WARNINGS block leads with it.
warnings.insert(
0,
format!(
"provenance: {} node(s) + {} edge(s) resolved at tier {} ({}); {} T4 proposal(s) \
held back (commented-out, never applied)",
nodes.len(),
edges.len(),
opts.declared_tier.token(),
opts.declared_tier.describe(),
0
),
);
Ok(GraphPlan {
source: opts.source,
sink: opts.sink,
mode: opts.mode,
batch_rows: opts.batch_rows,
concurrency: opts.concurrency,
checkpoint: opts.checkpoint.clone(),
nodes,
edges,
proposals: Vec::new(),
warnings,
})
}
/// Attach T4 inference proposals (from an inclusion-dependency pass run
/// elsewhere). They are stamped [`Tier::T4InclusionDependency`] and will
/// only ever be emitted COMMENTED-OUT. Returns `self` for chaining.
pub fn with_proposals(mut self, proposals: impl IntoIterator<Item = EdgeProposal>) -> Self {
for mut p in proposals {
p.edge.tier = Tier::T4InclusionDependency;
self.proposals.push(p);
}
// Re-summarise the proposal count in the provenance line.
if let Some(first) = self.warnings.first_mut() {
if first.starts_with("provenance:") {
*first = format!(
"provenance: {} node(s) + {} edge(s) resolved; {} T4 proposal(s) held back \
(commented-out, never applied)",
self.nodes.len(),
self.edges.len(),
self.proposals.len()
);
}
}
self
}
/// A canonical, human-readable digest of the plan's **labels, merge keys,
/// edge directions and batch size** — the *contract* both projections must
/// honour. The two-projections-agree test compares what it can parse back
/// out of each artifact against this. Deterministic and stable.
pub fn contract(&self) -> String {
let mut s = String::new();
let _ = writeln!(s, "batch_rows={}", self.batch_rows);
for n in &self.nodes {
let _ = writeln!(s, "node {}={}", n.label, n.key);
}
for e in &self.edges {
let _ = writeln!(s, "edge {}-{}->{}", e.from_label, e.rel, e.to_label);
}
s
}
// ───────────────────────────────────────────────────────────────────────
// Projection (a): the PySpark job — the Spark backend
// ───────────────────────────────────────────────────────────────────────
/// **Projection (a): a PySpark job** loading the schema into the graph
/// sink. Obeys the rel2graph Python hard rules: a FAIR scheduler, indexes
/// created at startup, **nodes-before-edges**, the sink client constructed
/// *inside* the partition function, a `coalesce()` before
/// `foreachPartition`, and idempotent `UNWIND $rows AS r MERGE` in
/// `batch_rows`-sized chunks. Streaming wraps the same per-partition writer
/// in a single `foreachBatch` with one `checkpointLocation`.
///
/// Deterministic: same plan ⇒ byte-identical Python.
pub fn to_pyspark(&self) -> String {
let mut s = String::new();
let p = &mut s;
let _ = writeln!(
p,
"# ── rel2graph codegen · Spark backend projection ──────────────────"
);
let _ = writeln!(
p,
"# source={} sink={} mode={} batch_rows={} concurrency={}",
self.source.token(),
self.sink.token(),
self.mode.token(),
self.batch_rows,
self.concurrency
);
let _ = writeln!(
p,
"# GENERATED — one MappingSpec, two projections; edits belong upstream."
);
let _ = writeln!(p, "from pyspark.sql import SparkSession");
let _ = writeln!(p);
let _ = writeln!(p, "BATCH_ROWS = {}", self.batch_rows);
let _ = writeln!(p, "CHECKPOINT = {:?}", self.checkpoint);
let _ = writeln!(p);
// Sink client + per-batch MERGE helpers, built INSIDE the partition fn.
let _ = writeln!(p, "def _sink():");
let _ = writeln!(
p,
" # sink client constructed INSIDE the partition fn (never shipped from the driver)"
);
match self.sink {
SinkKind::FalkorDb => {
let _ = writeln!(p, " from falkordb import FalkorDB");
let _ = writeln!(
p,
" return FalkorDB(host='localhost', port=6379).select_graph('rel2graph')"
);
}
SinkKind::GraphAr => {
let _ = writeln!(p, " import pygraphar # columnar graph writer");
let _ = writeln!(p, " return pygraphar.Writer('rel2graph')");
}
}
let _ = writeln!(p);
// Node writer.
let _ = writeln!(p, "def _write_nodes(label, key, rows):");
let _ = writeln!(p, " g = _sink(); buf = []");
let _ = writeln!(p, " def flush():");
let _ = writeln!(p, " if not buf: return");
p.push_str(
" g.query('UNWIND $rows AS r MERGE (n:`' + label + '` {`' + key + '`: r[\"' + key + '\"]}) SET n += r', {'rows': buf})\n",
);
let _ = writeln!(p, " buf.clear()");
let _ = writeln!(p, " for row in rows:");
let _ = writeln!(p, " buf.append(row.asDict())");
let _ = writeln!(p, " if len(buf) >= BATCH_ROWS: flush()");
let _ = writeln!(p, " flush()");
let _ = writeln!(p);
// Edge writer.
let _ = writeln!(p, "def _write_edges(fl, rel, tl, fk, tk, rows):");
let _ = writeln!(p, " g = _sink(); buf = []");
let _ = writeln!(p, " def flush():");
let _ = writeln!(p, " if not buf: return");
p.push_str(
" g.query('UNWIND $rows AS r MERGE (a:`' + fl + '` {`' + fk + '`: r.src}) MERGE (b:`' + tl + '` {`' + tk + '`: r.dst}) MERGE (a)-[e:`' + rel + '`]->(b) SET e += r.props', {'rows': buf})\n",
);
let _ = writeln!(p, " buf.clear()");
let _ = writeln!(p, " for row in rows:");
let _ = writeln!(
p,
" d = row.asDict(); props = {{k: v for k, v in d.items() if k not in ('src', 'dst')}}"
);
let _ = writeln!(
p,
" buf.append({{'src': d['src'], 'dst': d['dst'], 'props': props}})"
);
let _ = writeln!(p, " if len(buf) >= BATCH_ROWS: flush()");
let _ = writeln!(p, " flush()");
let _ = writeln!(p);
// Startup indexes — one per node label+key.
let _ = writeln!(p, "def ensure_indexes(g):");
if self.nodes.is_empty() {
let _ = writeln!(p, " pass # no nodes");
}
for n in &self.nodes {
let _ = writeln!(
p,
" g.query('CREATE INDEX FOR (n:`{}`) ON (n.`{}`)') # {}",
n.label,
n.key,
n.tier.token()
);
}
let _ = writeln!(p);
// main(): FAIR scheduler, nodes before edges.
let _ = writeln!(p, "def main():");
let _ = writeln!(
p,
" spark = SparkSession.builder.config('spark.scheduler.mode', 'FAIR').getOrCreate()"
);
let _ = writeln!(p, " ensure_indexes(_sink())");
let stream = self.mode == PlanMode::Stream;
let reader = if stream { "readStream.table" } else { "table" };
let _ = writeln!(p, " # ── nodes (before edges) ──");
for n in &self.nodes {
let cols = node_select_cols(n);
let df = format!("df_node_{}", sanitize(&n.label));
let _ = writeln!(
p,
" {df} = spark.{reader}({:?}).selectExpr({cols})",
n.table
);
if stream {
let _ = writeln!(
p,
" {df}.writeStream.option('checkpointLocation', CHECKPOINT + '/node_{}').foreachBatch(",
sanitize(&n.label)
);
let _ = writeln!(
p,
" lambda b, _: b.coalesce({}).foreachPartition(lambda rows: _write_nodes({:?}, {:?}, rows))).start()",
self.concurrency, n.label, n.key
);
} else {
let _ = writeln!(
p,
" {df}.coalesce({}).foreachPartition(lambda rows: _write_nodes({:?}, {:?}, rows))",
self.concurrency, n.label, n.key
);
}
}
let _ = writeln!(p, " # ── edges ──");
for e in &self.edges {
let cols = edge_select_cols(e);
let df = format!(
"df_edge_{}_{}_{}",
sanitize(&e.from_label),
sanitize(&e.rel),
sanitize(&e.to_label)
);
let _ = writeln!(
p,
" {df} = spark.{reader}({:?}).selectExpr({cols})",
e.table
);
if stream {
let _ = writeln!(
p,
" {df}.writeStream.option('checkpointLocation', CHECKPOINT + '/edge_{}_{}_{}').foreachBatch(",
sanitize(&e.from_label),
sanitize(&e.rel),
sanitize(&e.to_label)
);
let _ = writeln!(
p,
" lambda b, _: b.coalesce({}).foreachPartition(lambda rows: _write_edges({:?}, {:?}, {:?}, {:?}, {:?}, rows))).start()",
self.concurrency, e.from_label, e.rel, e.to_label, e.from_key, e.to_key
);
} else {
let _ = writeln!(
p,
" {df}.coalesce({}).foreachPartition(lambda rows: _write_edges({:?}, {:?}, {:?}, {:?}, {:?}, rows))",
self.concurrency, e.from_label, e.rel, e.to_label, e.from_key, e.to_key
);
}
}
if stream {
let _ = writeln!(p, " spark.streams.awaitAnyTermination()");
}
let _ = writeln!(p);
let _ = writeln!(p, "if __name__ == '__main__':");
let _ = writeln!(p, " main()");
// T4 proposals — commented out with evidence, never applied.
self.emit_proposals_py(p);
// WARNINGS block closes the artifact.
self.emit_warnings(p, "#");
s
}
fn emit_proposals_py(&self, p: &mut String) {
if self.proposals.is_empty() {
return;
}
let _ = writeln!(p);
let _ = writeln!(
p,
"# ── T4 PROPOSALS (inclusion-dependency inference) — NOT applied ──"
);
let _ = writeln!(
p,
"# Confirm the evidence, then add upstream in the MappingSpec. Never auto-applied."
);
for pr in &self.proposals {
let e = &pr.edge;
let _ = writeln!(
p,
"# ({}:`{}`)-[:`{}`]->({}:`{}`) evidence: {}",
sanitize(&e.from_label),
e.from_label,
e.rel,
sanitize(&e.to_label),
e.to_label,
pr.evidence
);
}
}
/// The WARNINGS block, prefixed with `line_prefix` (e.g. `#` for Python,
/// `//` for RON). Closes every artifact — provenance first, then notes.
pub fn emit_warnings(&self, p: &mut String, line_prefix: &str) {
let _ = writeln!(p);
let _ = writeln!(p, "{line_prefix} ── WARNINGS ──");
if self.warnings.is_empty() {
let _ = writeln!(p, "{line_prefix} (none)");
}
for w in &self.warnings {
let _ = writeln!(p, "{line_prefix} {w}");
}
}
}
// ── helpers ─────────────────────────────────────────────────────────────────
/// Schema-then-name sort key: `(namespace, leaf)` so `a.z` precedes `b.a` on
/// namespace but `s.accounts` precedes `s.users` on leaf. Deterministic.
fn table_sort_key(table: &str) -> (String, String) {
match table.rsplit_once('.') {
Some((ns, leaf)) => (ns.to_string(), leaf.to_string()),
None => (String::new(), table.to_string()),
}
}
/// Python identifier-safe form of a label/rel for a variable name.
fn sanitize(s: &str) -> String {
s.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect()
}
/// `selectExpr(...)` column list for a node: key first, then ordinal props,
/// each backtick-quoted.
fn node_select_cols(n: &NodePlan) -> String {
let mut cols = vec![format!("\"`{}`\"", n.key)];
cols.extend(n.props.iter().map(|c| format!("\"`{c}`\"")));
cols.join(", ")
}
/// `selectExpr(...)` column list for an edge: `from_col AS src`, `to_col AS
/// dst`, then ordinal props (as themselves), each backtick-quoted.
///
/// The endpoint VALUES are read from the FK columns (`from_col`/`to_col`), NOT
/// the merge-key property names (`from_key`/`to_key`) — they differ for a
/// self-ref / renamed FK (`manager_id AS dst`, MERGEd on `employee_id`).
fn edge_select_cols(e: &EdgePlan) -> String {
let mut parts = vec![
format!("\"`{}` AS src\"", e.from_col),
format!("\"`{}` AS dst\"", e.to_col),
];
for c in &e.props {
parts.push(format!("\"`{c}`\""));
}
parts.join(", ")
}
// ═══════════════════════════════════════════════════════════════════════════
// Tests — canonical plan + the PySpark (Spark backend) projection. The
// *two-projections-agree* invariant is pinned in `knut-thund` (which can see
// the thund `Pipeline` projection too); here we pin determinism, ordering,
// provenance, and the Python hard rules.
// ═══════════════════════════════════════════════════════════════════════════
#[cfg(test)]
mod tests {
use super::*;
use crate::mapping::MappingSpec;
/// A two-table social schema: users (with a company FK) + companies, whose
/// declaration order is deliberately "wrong" (users before companies, edge
/// listed before the plain node) so the canonical ordering has work to do.
fn sample_spec() -> MappingSpec {
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] }
"#;
MappingSpec::from_yaml(yaml).expect("sample yaml parses")
}
#[test]
fn plan_is_canonical_and_deterministic() {
let spec = sample_spec();
let opts = CodegenOptions::default();
let a = GraphPlan::from_mapping(&spec, &opts).unwrap();
let b = GraphPlan::from_mapping(&spec, &opts).unwrap();
// Byte-identical PySpark on repeat — no timestamps/random/hashmap order.
assert_eq!(a.to_pyspark(), b.to_pyspark(), "codegen is deterministic");
// Tables schema-then-name: `companies` (leaf) sorts before `users`, so
// its node comes first even though it was declared last.
assert_eq!(a.nodes.len(), 2);
assert_eq!(a.nodes[0].label, "Company");
assert_eq!(a.nodes[1].label, "User");
// Props stay ordinal.
assert_eq!(a.nodes[1].props, vec!["name", "age"]);
// Edge: direction User -[WORKS_AT]-> Company, keys carried.
assert_eq!(a.edges.len(), 1);
let e = &a.edges[0];
assert_eq!(
(e.from_label.as_str(), e.rel.as_str(), e.to_label.as_str()),
("User", "WORKS_AT", "Company")
);
assert_eq!(e.from_key, "user_id");
assert_eq!(e.to_key, "company_id");
}
#[test]
fn pyspark_obeys_the_python_hard_rules() {
let plan = GraphPlan::from_mapping(&sample_spec(), &CodegenOptions::default()).unwrap();
let py = plan.to_pyspark();
assert!(py.contains("spark.scheduler.mode"), "FAIR scheduler set");
assert!(py.contains("'FAIR'"));
assert!(py.contains("ensure_indexes"), "indexes created at startup");
assert!(py.contains("CREATE INDEX FOR (n:`User`) ON (n.`user_id`)"));
assert!(
py.contains("def _sink():"),
"sink client built inside partition fn"
);
assert!(
py.contains(".coalesce(1).foreachPartition("),
"coalesce before foreachPartition"
);
assert!(
py.contains("UNWIND $rows AS r MERGE"),
"idempotent UNWIND MERGE"
);
assert!(py.contains("if len(buf) >= BATCH_ROWS:"), "batched writes");
assert!(py.contains("BATCH_ROWS = 1000"));
// REGRESSION GUARD: the Cypher `MERGE (n:… {`key`: …})` carries REAL
// SINGLE braces `{`/`}` in the emitted text. These are literal `{` in
// the source; emitting them via `format!`/`writeln!` (which read `{` as
// a placeholder) is a compile error that has bitten the kube-feature
// build — so they are built with a raw `push_str`, and this pins that
// the single-brace Cypher survives to the output verbatim.
assert!(
py.contains("MERGE (n:`' + label + '` {`' + key + '`: r[\"' + key + '\"]})"),
"node MERGE keeps its single-brace Cypher: {py}"
);
assert!(
py.contains("MERGE (a:`' + fl + '` {`' + fk + '`: r.src})"),
"edge MERGE keeps its single-brace Cypher"
);
assert!(
py.contains("{'rows': buf}"),
"the params dict uses single braces"
);
// And never the doubled braces a `format!` escape would leave behind.
assert!(!py.contains("{{"), "no doubled `{{` leaked into the output");
assert!(!py.contains("}}"), "no doubled `}}` leaked into the output");
// Nodes-before-edges: both node writers precede the edge writer in main().
let main = &py[py.find("def main()").unwrap()..];
let node_at = main.find("_write_nodes").unwrap();
let edge_at = main.find("_write_edges").unwrap();
assert!(node_at < edge_at, "nodes are loaded before edges");
// WARNINGS block closes the artifact, provenance first.
assert!(py.trim_end().contains("── WARNINGS ──"));
assert!(py.contains("provenance:"));
}
#[test]
fn streaming_needs_iceberg_and_wraps_one_foreachbatch() {
let spec = sample_spec();
// Stream on Postgres is refused up front.
let bad = CodegenOptions {
source: SourceKind::Postgres,
mode: PlanMode::Stream,
..Default::default()
};
assert!(
GraphPlan::from_mapping(&spec, &bad).is_err(),
"stream needs iceberg"
);
let ok = CodegenOptions {
source: SourceKind::Iceberg,
mode: PlanMode::Stream,
..Default::default()
};
let py = GraphPlan::from_mapping(&spec, &ok).unwrap().to_pyspark();
assert!(py.contains("readStream.table"), "streaming reader");
assert!(
py.contains("foreachBatch("),
"foreachBatch wraps the per-partition writer"
);
assert!(
py.contains("checkpointLocation"),
"one checkpointLocation per query"
);
assert!(py.contains("awaitAnyTermination"));
}
#[test]
fn custom_step_is_surfaced_not_dropped() {
let yaml = r#"
version: 1
tables:
- table: lake.x.t
pipeline:
- vertex: { label: T, id_column: id }
- custom: my_special_transform
"#;
let spec = MappingSpec::from_yaml(yaml).unwrap();
let plan = GraphPlan::from_mapping(&spec, &CodegenOptions::default()).unwrap();
assert_eq!(plan.nodes.len(), 1);
assert!(
plan.warnings
.iter()
.any(|w| w.contains("my_special_transform") && w.contains("NOT dropped")),
"the opaque custom step is surfaced in WARNINGS: {:?}",
plan.warnings
);
}
#[test]
fn t4_proposals_are_emitted_commented_out_never_applied() {
let plan = GraphPlan::from_mapping(&sample_spec(), &CodegenOptions::default())
.unwrap()
.with_proposals([EdgeProposal {
edge: EdgePlan {
table: "lake.social.users".into(),
from_label: "User".into(),
rel: "REFERRED_BY".into(),
to_label: "User".into(),
from_key: "user_id".into(),
to_key: "user_id".into(),
from_col: "user_id".into(),
to_col: "referrer_id".into(),
props: vec![],
tier: Tier::T1PostgresCatalog, // will be overwritten to T4
},
evidence: "LEFT ANTI JOIN found 0 orphan referrer_id rows (n=10000)".into(),
}]);
// Stamped T4 regardless of what the caller passed.
assert_eq!(plan.proposals[0].edge.tier, Tier::T4InclusionDependency);
let py = plan.to_pyspark();
assert!(py.contains("T4 PROPOSALS"), "proposals section present");
assert!(py.contains("REFERRED_BY"), "the proposal is shown");
assert!(py.contains("evidence:"), "with its evidence");
// But NEVER as an active edge writer.
assert!(
!py.contains("_write_edges(\"User\", \"REFERRED_BY\""),
"a T4 proposal is never emitted as a live edge load"
);
}
}