crates/knut-bifrost/tests/robot_rel2graph_graph.rs
//! **ROBOT — the bifröst rel2graph bridge, asserted on the graph it emits.**
//!
//! [`build_mutation`] is the seam every DB-free consumer of the engine shares
//! (the bench, `knut-flight`, korp's `iceberg2graph` step). Its doctest asserts
//! `mutation.vertices.len() == 1`. That is a shape check: it passes on a
//! `VertexSet` carrying zero rows, and it says nothing at all about the EDGES —
//! which are the only part of a relational→graph bridge that can be wrong in an
//! interesting way.
//!
//! So this suite reads the mutation's Arrow batches back out and asserts the
//! **applied output**: the exact node id list per label, the exact `(src, dst)`
//! endpoint pairs of every relationship, the property columns that rode along
//! under their renamed names, and the endpoint LABELS and id-columns the sink
//! will `MATCH` on.
//!
//! Off the identity throughout: ids are 101.. and 7.. (never 0/1, which is what
//! a positional fallback would produce), the fan-out is 3/2/1 (never uniform,
//! which cannot distinguish a keyed join from a positional zip), one department
//! has no employees at all, and the rows span **two** RecordBatches so a
//! per-batch bug cannot hide.
//!
//! Run: `cargo test -p knut-bifrost --test robot_rel2graph_graph -- --nocapture`
use std::sync::Arc;
use arrow_array::{Array, Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use knut_bifrost::mapping::{FilterValue, RowFilter, apply_row_filter};
use knut_bifrost::{ArrowTables, GraphMutation, MappingSpec, build_mutation};
// ── fixture ─────────────────────────────────────────────────────────────────
/// Employees, split over TWO batches (4 + 2) so a per-batch id or offset bug
/// shows up as a wrong endpoint rather than a wrong count.
fn employees() -> Vec<RecordBatch> {
let schema = Arc::new(Schema::new(vec![
Field::new("emp_id", DataType::Int64, false),
Field::new("dept_id", DataType::Int64, false),
Field::new("full_name", DataType::Utf8, false),
Field::new("hired", DataType::Int64, false),
]));
let mk = |emp: Vec<i64>, dept: Vec<i64>, name: Vec<&str>, hired: Vec<i64>| {
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int64Array::from(emp)),
Arc::new(Int64Array::from(dept)),
Arc::new(StringArray::from(name)),
Arc::new(Int64Array::from(hired)),
],
)
.unwrap()
};
vec![
mk(
vec![101, 102, 103, 104],
vec![7, 7, 8, 7],
vec!["astrid", "bengt", "cecilia", "dag"],
vec![2011, 2014, 2016, 2018],
),
mk(
vec![105, 106],
vec![8, 9],
vec!["elin", "folke"],
vec![2020, 2023],
),
]
}
/// Four departments — 10 has NOBODY, so a bridge that emitted an edge per
/// department, or a cartesian product, cannot pass.
fn departments() -> Vec<RecordBatch> {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("dept_id", DataType::Int64, false),
Field::new("dept_name", DataType::Utf8, false),
Field::new("country", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![7i64, 8, 9, 10])),
Arc::new(StringArray::from(vec!["Bygg", "Data", "Juridik", "Arkiv"])),
Arc::new(StringArray::from(vec!["SE", "SE", "NO", "SE"])),
],
)
.map(|b| vec![b])
.unwrap()
}
const SPEC: &str = r#"
version: 1
tables:
- table: lake.hr.employees
pipeline:
- vertex:
label: Employee
id_column: emp_id
properties: [full_name]
rename: { full_name: name }
- fk_edges:
type: WORKS_IN
src: { label: Employee, column: emp_id }
dst: { label: Department, column: dept_id }
properties: [hired]
rename: { hired: since }
- table: lake.hr.departments
pipeline:
- vertex:
label: Department
id_column: dept_id
properties: [dept_name]
"#;
/// Read one Int64 column of a mutation batch out as strings, in row order.
fn ints(batch: &RecordBatch, col: &str) -> Vec<i64> {
let a = batch
.column_by_name(col)
.unwrap_or_else(|| panic!("column `{col}` missing; have {:?}", names(batch)));
a.as_any()
.downcast_ref::<Int64Array>()
.unwrap_or_else(|| panic!("column `{col}` is {:?}, not Int64", a.data_type()))
.iter()
.map(|v| v.expect("no nulls in the fixture"))
.collect()
}
fn strs(batch: &RecordBatch, col: &str) -> Vec<String> {
batch
.column_by_name(col)
.unwrap_or_else(|| panic!("column `{col}` missing; have {:?}", names(batch)))
.as_any()
.downcast_ref::<StringArray>()
.expect("Utf8")
.iter()
.map(|v| v.expect("no nulls").to_string())
.collect()
}
fn names(batch: &RecordBatch) -> Vec<String> {
batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect()
}
/// All rows of every `VertexSet` carrying `label`, as (id, extra-columns).
fn vertex_ids(m: &GraphMutation, label: &str, id_col: &str) -> Vec<i64> {
m.vertices
.iter()
.filter(|v| v.label == label)
.flat_map(|v| ints(&v.batch, id_col))
.collect()
}
/// Every `(src, dst)` endpoint pair of every `EdgeSet` of `edge_type`.
fn edge_pairs(m: &GraphMutation, edge_type: &str) -> Vec<(i64, i64)> {
m.edges
.iter()
.filter(|e| e.edge_type == edge_type)
.flat_map(|e| {
ints(&e.batch, "src")
.into_iter()
.zip(ints(&e.batch, "dst"))
.collect::<Vec<_>>()
})
.collect()
}
fn full_source() -> ArrowTables {
ArrowTables::new()
.with_table("lake.hr.employees", employees())
.with_table("lake.hr.departments", departments())
}
// ── the suite ───────────────────────────────────────────────────────────────
/// **The bridge's edges land on the ids the FK column says.**
///
/// Asserted as the exact ordered `(emp_id, dept_id)` pair list across BOTH input
/// batches — not `edges.len()`, which a positional zip, a per-batch id reset, or
/// a cartesian product would each satisfy.
#[test]
fn robot_fk_edges_carry_the_exact_endpoint_pairs() {
let m = build_mutation(&MappingSpec::from_yaml(SPEC).expect("spec"), &full_source())
.expect("build_mutation");
// ── Node ids, exactly. Two batches of employees must both be present. ──
assert_eq!(
vertex_ids(&m, "Employee", "emp_id"),
vec![101, 102, 103, 104, 105, 106],
"every employee id, across both input batches"
);
assert_eq!(
vertex_ids(&m, "Department", "dept_id"),
vec![7, 8, 9, 10],
"every department, including the empty one"
);
// ── Endpoint pairs, exactly. ──
let pairs = edge_pairs(&m, "WORKS_IN");
assert_eq!(
pairs,
vec![(101, 7), (102, 7), (103, 8), (104, 7), (105, 8), (106, 9)],
"the FK column IS the destination endpoint — a positional zip would give \
(101,7),(102,7),(103,8),(104,7) then restart at (105,7),(106,7)"
);
// 4 employees × 6 departments would be 24; a per-batch restart would still be 6.
assert_eq!(
pairs.len(),
6,
"one edge per employee row, no cartesian blow-up"
);
// ── The fan-out is 3/2/1/0 — non-uniform, so this can tell a join from a zip. ──
let mut fan = std::collections::BTreeMap::<i64, usize>::new();
for (_, d) in &pairs {
*fan.entry(*d).or_default() += 1;
}
assert_eq!(
fan,
[(7, 3), (8, 2), (9, 1)].into_iter().collect(),
"departments 7/8/9 take 3/2/1 employees and 10 takes NONE"
);
assert!(
!fan.contains_key(&10),
"the empty department must attract no edge: {fan:?}"
);
// ── The sink's MATCH shape: labels and id columns, applied. ──
let e = m
.edges
.iter()
.find(|e| e.edge_type == "WORKS_IN")
.expect("a WORKS_IN edge set");
assert_eq!(e.src_label, "Employee");
assert_eq!(e.dst_label, "Department");
assert_eq!(
(e.src_id_column.as_str(), e.dst_id_column.as_str()),
("emp_id", "dept_id"),
"the endpoint id columns the sink will MATCH on — defaulting to the \
carrier column name, since neither endpoint set `id_column`"
);
// ── Renamed properties rode along under their NEW names, with their values. ──
let ebatch = &e.batch;
assert_eq!(
names(ebatch),
vec!["src", "dst", "since"],
"the endpoint columns become src/dst and `hired` lands as `since`"
);
assert_eq!(
ints(ebatch, "since"),
vec![2011, 2014, 2016, 2018],
"the first batch's edge properties, by value — a rename that dropped the \
column would leave the schema check above as the only witness"
);
let v = m
.vertices
.iter()
.find(|v| v.label == "Employee")
.expect("an Employee vertex set");
assert_eq!(v.id_column, "emp_id");
assert_eq!(
names(&v.batch),
vec!["emp_id", "name"],
"`full_name` was renamed to `name`, and `dept_id`/`hired` were projected away"
);
assert_eq!(
strs(&v.batch, "name"),
vec!["astrid", "bengt", "cecilia", "dag"]
);
eprintln!(
"[robot-bifrost] {} vertex sets / {} edge sets · pairs {pairs:?}",
m.vertices.len(),
m.edges.len()
);
}
/// **The non-vacuous anchor.** A source that carries the employees table MUST
/// produce edges; a source that does not MUST produce none — and the two must
/// differ.
///
/// `build_mutation` skips a table the source lacks, which is the documented
/// "partial source yields a partial-but-consistent graph" behaviour. The failure
/// it hides is a source wired to the wrong table name: every step silently
/// vanishes and `build_mutation` returns `Ok` with an empty mutation.
#[test]
fn robot_a_missing_table_yields_no_edges_and_differs_from_the_full_run() {
let spec = MappingSpec::from_yaml(SPEC).expect("spec");
// ── must produce OUTPUT ──
let full = build_mutation(&spec, &full_source()).expect("full");
let full_pairs = edge_pairs(&full, "WORKS_IN");
assert_eq!(full_pairs.len(), 6);
assert!(!full.is_empty());
// ── must produce NONE: the departments table only. ──
let partial = ArrowTables::new().with_table("lake.hr.departments", departments());
let m = build_mutation(&spec, &partial).expect("partial");
assert_eq!(
edge_pairs(&m, "WORKS_IN"),
Vec::<(i64, i64)>::new(),
"no employees table → no WORKS_IN edges at all"
);
assert!(m.edges.is_empty(), "and no edge sets: {:?}", m.edges.len());
assert!(
vertex_ids(&m, "Employee", "emp_id").is_empty(),
"and no Employee nodes"
);
// But it is not simply empty — the departments DID come through, so this
// asserts the skip is scoped to the missing table.
assert_eq!(vertex_ids(&m, "Department", "dept_id"), vec![7, 8, 9, 10]);
// ── and nothing at all when the source is empty. ──
let none = build_mutation(&spec, &ArrowTables::new()).expect("empty");
assert!(none.is_empty(), "an empty source builds an empty mutation");
// ── the three cases must be distinct. ──
assert_ne!(full.edges.len(), m.edges.len());
assert_ne!(m.vertices.len(), none.vertices.len());
}
/// **The row filter selects, and its miss-case selects nothing.**
///
/// `apply_row_filter` is the predicate the Iceberg planner pushes down; when the
/// planner cannot prune, this is the residual that runs. Asserted on the exact
/// surviving ids — a count would not notice a filter that kept the wrong three.
#[test]
fn robot_row_filter_keeps_exactly_the_matching_departments() {
let batches = departments();
// ── must produce OUTPUT: three of four departments are Swedish. ──
let se = apply_row_filter(&batches, &RowFilter::eq("country", "SE")).expect("filter SE");
let kept: Vec<i64> = se.iter().flat_map(|b| ints(b, "dept_id")).collect();
assert_eq!(
kept,
vec![7, 8, 10],
"exactly the SE departments — 9 (Juridik/NO) must go"
);
assert_eq!(
se.iter()
.flat_map(|b| strs(b, "country"))
.collect::<Vec<_>>(),
vec!["SE", "SE", "SE"],
"and nothing non-SE survived"
);
// ── must produce NONE: a value nothing carries. ──
let dk = apply_row_filter(&batches, &RowFilter::eq("country", "DK")).expect("filter DK");
let none: Vec<i64> = dk.iter().flat_map(|b| ints(b, "dept_id")).collect();
assert!(none.is_empty(), "no Danish department exists, got {none:?}");
// ── and they differ — the whole point. ──
assert_ne!(
kept, none,
"a filter that kept everything would collapse these"
);
// ── IN over two values, exact. ──
let both = apply_row_filter(
&batches,
&RowFilter::is_in(
"country",
vec![FilterValue::from("NO"), FilterValue::from("DK")],
),
)
.expect("filter IN");
assert_eq!(
both.iter()
.flat_map(|b| ints(b, "dept_id"))
.collect::<Vec<_>>(),
vec![9],
"IN (NO, DK) selects only Juridik"
);
// ── a column that does not exist is an ERROR, never a silent empty set. ──
let err = apply_row_filter(&batches, &RowFilter::eq("nope", "SE"));
let msg = format!(
"{}",
err.expect_err("a missing filter column must be an error")
);
assert!(
msg.contains("nope") && msg.contains("not in the batch schema"),
"the error must name the column: {msg}"
);
}
/// **A self-key edge must emit BOTH endpoints from its one input column.**
///
/// `project_edge` exists specifically because the generic projection would drop
/// one side of a `src.column == dst.column` link (a rename map is keyed by
/// source column and can hold one mapping). Asserted on the emitted pairs: a
/// dropped endpoint would leave the batch one column short, and a shared array
/// would still have to hold the same values in both.
#[test]
fn robot_a_self_keyed_edge_emits_src_and_dst_from_one_column() {
let spec = MappingSpec::from_yaml(
r#"
version: 1
tables:
- table: lake.hr.alias
pipeline:
- edge:
type: SAME_AS
src: { label: Employee, column: emp_id }
dst: { label: Employee, column: emp_id }
properties: [confidence]
"#,
)
.expect("self-key spec");
let batch = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("emp_id", DataType::Int64, false),
Field::new("confidence", DataType::Int64, false),
])),
vec![
Arc::new(Int64Array::from(vec![101i64, 105])),
Arc::new(Int64Array::from(vec![90i64, 75])),
],
)
.unwrap();
let m = build_mutation(
&spec,
&ArrowTables::new().with_table("lake.hr.alias", vec![batch]),
)
.expect("build");
let e = m.edges.first().expect("one SAME_AS edge set");
assert_eq!(
names(&e.batch),
vec!["src", "dst", "confidence"],
"both endpoints must be emitted from the single `emp_id` column"
);
assert_eq!(
edge_pairs(&m, "SAME_AS"),
vec![(101, 101), (105, 105)],
"and both carry the same value per row"
);
assert_eq!(e.src_label, "Employee");
assert_eq!(e.dst_label, "Employee");
}