crates/knut-bifrost/src/rel2graph.rs
//! **The canonical offline rel2graph engine seam.**
//!
//! bifröst *is* the rel2graph engine: a [`MappingSpec`](crate::MappingSpec) is
//! compiled to a [`Transform`](crate::Transform) fabric and applied to a table's
//! Arrow batches to produce a [`GraphMutation`](crate::GraphMutation) (the
//! `VertexSet`s + `EdgeSet`s that ARE the graph). The live path runs this per
//! round inside the [`Listener`](crate::Listener), reading from a
//! [`Source`](crate::source) and writing to a [`GraphSink`](crate::sink::GraphSink).
//!
//! This module exposes the SAME work with **no source and no sink** — you hand it
//! the per-table Arrow batches, it hands you the `GraphMutation`. It is the one
//! entrypoint every DB-free consumer of the engine shares:
//!
//! * the [correctness + timing bench](../../knut_rel2graph_bench/index.html) drives
//! it over a seeded synthetic database;
//! * [`knut-flight`](../../knut_flight/index.html) drives it over an Iceberg /
//! parquet export to serve a REAL graph as Arrow Flight;
//! * the networkx file-export drives it to dump a node-list + edge-list + oracle;
//! * **korp's `iceberg2graph` demo step** calls it to turn the tables it just read
//! from the lakehouse into the graph it hands the viewer — the same graph the
//! PySpark projection ([`GraphPlan::to_pyspark`](crate::GraphPlan::to_pyspark))
//! and the thund projection (`knut_thund::rel2graph::to_thund_pipeline`) would
//! build, since all three derive from ONE reviewed [`MappingSpec`].
//!
//! The point of the seam is that swapping the data behind it never touches the
//! engine: whatever yields per-table [`RecordBatch`]es — the bench generator, a
//! parquet directory, a Spark-Connect Iceberg read — implements [`TableSource`]
//! and feeds the identical [`build_mutation`].
use std::collections::BTreeMap;
use arrow_array::RecordBatch;
use crate::{
error::Result, mapping::MappingSpec, mutation::GraphMutation, registry::TransformRegistry,
transform::TableContext,
};
/// A source of Arrow batches keyed by fully-qualified table name — the seam the
/// offline rel2graph engine reads from.
///
/// Implemented by anything that can produce a table's rows as Arrow: the bench's
/// `GeneratedDb`, the in-memory [`ArrowTables`] (a parquet-directory load or an
/// Iceberg scan collected into memory), or a bespoke adapter. [`build_mutation`]
/// asks only for `batches(table)`; an unknown table returns an empty slice and is
/// skipped, so a source may legitimately carry a superset of the spec's tables.
pub trait TableSource {
/// The batches for `table`, or an empty slice if the source has no such table.
fn batches(&self, table: &str) -> &[RecordBatch];
}
/// **Run the bifröst rel2graph engine offline over `source`.**
///
/// Compiles every table's declarative pipeline once, applies the resulting
/// [`Transform`](crate::Transform)s to that table's batches, and merges every
/// step's output into one [`GraphMutation`]. This mirrors exactly what the live
/// [`Listener`](crate::Listener) does per round, minus the source read and the
/// sink write — so the mutation this returns is the mutation a live load would
/// MERGE, at the same seam the correctness oracle asserts on.
///
/// Tables the `source` does not carry are skipped (empty batches), so a partial
/// source yields a partial-but-consistent graph.
///
/// ```
/// use knut_bifrost::rel2graph::{build_mutation, ArrowTables};
/// use knut_bifrost::MappingSpec;
/// # use arrow_array::{RecordBatch, Int64Array};
/// # use arrow_schema::{Schema, Field, DataType};
/// # use std::sync::Arc;
/// let spec = MappingSpec::from_yaml(r#"
/// version: 1
/// tables:
/// - table: lake.hr.employees
/// pipeline:
/// - vertex: { label: Employee, id_column: employee_id, properties: [] }
/// "#).unwrap();
/// # let schema = Arc::new(Schema::new(vec![Field::new("employee_id", DataType::Int64, false)]));
/// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1i64, 2, 3]))]).unwrap();
/// let mut src = ArrowTables::new();
/// src.insert("lake.hr.employees", vec![batch]);
/// let mutation = build_mutation(&spec, &src).unwrap();
/// assert_eq!(mutation.vertices.len(), 1);
/// ```
pub fn build_mutation(spec: &MappingSpec, source: &dyn TableSource) -> Result<GraphMutation> {
build_mutation_reported(spec, source).map(|(m, _)| m)
}
/// **What a `build_mutation` run actually read.**
///
/// `build_mutation` skips a spec table the source does not carry, which is the right
/// default (a partial source yields a partial-but-consistent graph) and is also
/// exactly how an **empty graph comes back looking like a successful build**: point
/// the source at the wrong namespace and every table is "skipped", the mutation is
/// empty, and the call returns `Ok`.
///
/// This report is the discrimination. It is carried BESIDE the mutation rather than
/// folded into it, because the mutation is what a sink MERGEs and a report is not.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct BuildReport {
/// Spec tables the source carried rows for, in spec order.
pub tables_mapped: Vec<String>,
/// Spec tables the source carried NO rows for — absent, or present with zero
/// rows. These are the ones that silently vanish from an unreported build.
pub tables_empty: Vec<String>,
/// Total input rows read across every mapped table.
pub rows_read: usize,
/// Vertex rows in the resulting mutation.
pub vertices: usize,
/// Edge rows in the resulting mutation.
pub edges: usize,
}
impl BuildReport {
/// **The honest one-liner a caller uses instead of trusting `Ok`.**
///
/// `Ok(self)` when at least one spec table contributed rows AND the mutation is
/// non-empty; otherwise a typed [`crate::error::BifrostError::Source`] naming
/// every table that came back empty. A build that mapped nothing is not a graph.
pub fn require_non_empty(self) -> Result<Self> {
if self.rows_read == 0 || (self.vertices == 0 && self.edges == 0) {
return Err(crate::error::BifrostError::Source(format!(
"the rel2graph build read {} row(s) from {} table(s) and produced {} vertex \
row(s) + {} edge row(s) — that is an EMPTY graph, not a built one. Tables \
the source carried nothing for: {:?}. Check the catalog, the namespace and \
the table names in the mapping spec.",
self.rows_read,
self.tables_mapped.len(),
self.vertices,
self.edges,
self.tables_empty,
)));
}
Ok(self)
}
}
/// [`build_mutation`] plus a [`BuildReport`] saying what it read. Same work, same
/// result — this is the one implementation and `build_mutation` delegates to it, so
/// the counts can never describe a different run than the mutation.
pub fn build_mutation_reported(
spec: &MappingSpec,
source: &dyn TableSource,
) -> Result<(GraphMutation, BuildReport)> {
let registry = TransformRegistry::new();
let mut mutation = GraphMutation::default();
let mut report = BuildReport::default();
for tm in &spec.tables {
let batches = source.batches(&tm.table);
let rows: usize = batches.iter().map(RecordBatch::num_rows).sum();
if batches.is_empty() || rows == 0 {
report.tables_empty.push(tm.table.clone());
continue;
}
report.tables_mapped.push(tm.table.clone());
report.rows_read += rows;
let ctx = TableContext {
table: tm.table.clone(),
};
let transforms = tm.compile(®istry)?;
for transform in &transforms {
let m = transform.apply(&ctx, batches)?;
mutation.merge(m);
}
}
report.vertices = mutation.vertices.iter().map(|s| s.batch.num_rows()).sum();
report.edges = mutation.edges.iter().map(|s| s.batch.num_rows()).sum();
Ok((mutation, report))
}
/// An in-memory [`TableSource`]: a map of fully-qualified table name → Arrow
/// batches.
///
/// This is the adapter that turns "some Arrow I already have" into engine input.
/// Fill it from a parquet directory, an Iceberg scan, a Spark-Connect read, or by
/// hand, then drive [`build_mutation`] over real data with no bench and no live
/// FalkorDB.
#[derive(Debug, Clone, Default)]
pub struct ArrowTables {
tables: BTreeMap<String, Vec<RecordBatch>>,
}
impl ArrowTables {
/// An empty source.
pub fn new() -> Self {
Self::default()
}
/// Add (or replace) a table's batches. Returns nothing — see
/// [`with_table`](Self::with_table) for the builder form.
pub fn insert(&mut self, table: impl Into<String>, batches: Vec<RecordBatch>) {
self.tables.insert(table.into(), batches);
}
/// Builder form of [`insert`](Self::insert): add a table, return `self`.
pub fn with_table(mut self, table: impl Into<String>, batches: Vec<RecordBatch>) -> Self {
self.insert(table, batches);
self
}
/// The table names this source carries, in sorted order.
pub fn table_names(&self) -> impl Iterator<Item = &str> {
self.tables.keys().map(String::as_str)
}
/// Whether the source carries no tables.
pub fn is_empty(&self) -> bool {
self.tables.is_empty()
}
/// The number of tables carried.
pub fn len(&self) -> usize {
self.tables.len()
}
}
impl TableSource for ArrowTables {
fn batches(&self, table: &str) -> &[RecordBatch] {
self.tables.get(table).map(Vec::as_slice).unwrap_or(&[])
}
}
/// A shared reference to a [`TableSource`] is itself a [`TableSource`], so
/// `&db`/`&tables` can be passed where an owned source is not convenient.
impl<T: TableSource + ?Sized> TableSource for &T {
fn batches(&self, table: &str) -> &[RecordBatch] {
(**self).batches(table)
}
}
#[cfg(test)]
mod tests {
use super::*;
use arrow_array::Int64Array;
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
fn id_batch(col: &str, ids: &[i64]) -> RecordBatch {
let schema = Arc::new(Schema::new(vec![Field::new(col, DataType::Int64, false)]));
RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(ids.to_vec()))]).unwrap()
}
#[test]
fn arrow_tables_source_drives_the_engine() {
let spec = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: lake.hr.employees
pipeline:
- vertex: { label: Employee, id_column: employee_id, properties: [] }
"#,
)
.unwrap();
let src = ArrowTables::new().with_table(
"lake.hr.employees",
vec![id_batch("employee_id", &[1, 2, 3])],
);
assert_eq!(src.len(), 1);
let mutation = build_mutation(&spec, &src).unwrap();
assert_eq!(
mutation.vertices.len(),
1,
"one VertexSet for the one table"
);
assert_eq!(mutation.vertices[0].label, "Employee");
assert_eq!(mutation.vertices[0].batch.num_rows(), 3);
}
#[test]
fn missing_tables_are_skipped_not_errors() {
let spec = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: present.table
pipeline:
- vertex: { label: A, id_column: id, properties: [] }
- table: absent.table
pipeline:
- vertex: { label: B, id_column: id, properties: [] }
"#,
)
.unwrap();
// Source carries only `present.table`; `absent.table` yields no batches.
let src = ArrowTables::new().with_table("present.table", vec![id_batch("id", &[10, 20])]);
let mutation = build_mutation(&spec, &src).unwrap();
assert_eq!(
mutation.vertices.len(),
1,
"absent table skipped, present table mapped"
);
assert_eq!(mutation.vertices[0].label, "A");
}
}