crates/knut-rel2graph-bench/src/schema.rs
//! The test relational schema — a Northwind/ERP-shaped set of **12 tables** that
//! exercises every rel2graph correctness rule, plus the single source of truth
//! the mapping, the generator and the correctness expectations all read from.
//!
//! ## Coverage of the mapping rules (see the bench design doc)
//! * **table → node label / PK → node key** — the 10 entity tables.
//! * **FK → edge (many→one)** — every `*_id` FK column becomes a directed edge.
//! * **M:N junction table → EDGE, not a spurious vertex** — `order_items` and
//! `product_suppliers` have **no vertex step** at all; each is projected only
//! to edges. A conformant engine must NOT emit an `OrderItem`/`ProductSupplier`
//! node label.
//! * **composite key** — both junction tables key on a *pair* of FKs
//! (`(order_id, product_id)`, `(product_id, supplier_id)`).
//! * **self-referential FK** — `employees.manager_id → employees.employee_id`
//! becomes an `Employee-[:REPORTS_TO]->Employee` self-loop; the org roots carry
//! a NULL manager (a null FK is "no edge", standard SQL semantics).
use arrow_schema::{DataType, Field, Schema};
use knut_bifrost::{
mapping::{MappingSpec, PipelineStep, TableMapping},
transform::{EdgeTransform, EndpointSpec, VertexTransform},
};
use std::sync::Arc;
/// How the generator fills one column with deterministic, FK-consistent data.
#[derive(Debug, Clone)]
pub enum Fill {
/// Serial primary key `1..=row_count` (unique by construction).
PkSerial,
/// A foreign key: a uniform pick from the parent table's serial id range.
/// `null_prob` lets a nullable FK (a self-ref root) be absent.
Fk {
/// The parent table this FK references (qualified name).
parent: &'static str,
/// Probability the value is NULL (0.0 for a mandatory FK).
null_prob: f64,
},
/// A **self-referential** FK: points at an *earlier* row of the same table
/// (a valid manager already generated), or NULL for the first `roots` rows.
SelfRefManager {
/// The first `roots` employees have no manager (NULL).
roots: usize,
},
/// A uniform integer in `[lo, hi]`.
IntRange(i64, i64),
/// A price-like `f64` in `[lo, hi)` rounded to cents.
Money(f64, f64),
/// A synthetic name `"<prefix>_<00042>"`.
Name(&'static str),
/// A pick from a fixed set of category-like strings.
Enum(&'static [&'static str]),
/// A monotone-ish epoch-day integer.
Epoch,
}
/// One column: its Arrow field plus how to fill it.
#[derive(Debug, Clone)]
pub struct Column {
pub name: &'static str,
pub dtype: DataType,
pub nullable: bool,
pub fill: Fill,
}
/// A vertex projection for a node table.
#[derive(Debug, Clone)]
pub struct VertexModel {
pub label: &'static str,
pub key: &'static str,
pub props: &'static [&'static str],
}
/// An edge projection: `(src_label)-[rel]->(dst_label)`, both endpoint values
/// read from columns of *this* table.
#[derive(Debug, Clone)]
pub struct EdgeModel {
pub rel: &'static str,
pub src_label: &'static str,
/// Column in this table holding the src node's key value.
pub src_col: &'static str,
/// The src node's own key column (its `id_column`).
pub src_key: &'static str,
pub dst_label: &'static str,
/// Column in this table holding the dst node's key value (the FK).
pub dst_col: &'static str,
/// The dst node's own key column.
pub dst_key: &'static str,
pub props: &'static [&'static str],
}
/// A table in the schema — the single source of truth.
#[derive(Debug, Clone)]
pub struct Table {
/// Fully-qualified name (`erp.sales.orders`).
pub name: &'static str,
/// Base row count at scale 1 (node tables). Junction tables derive their
/// count from their parents and set this to 0.
pub base_rows: usize,
/// `true` for an M:N junction table (no vertex — edges only).
pub junction: bool,
/// The vertex step, if this table is an entity (`None` for a junction).
pub vertex: Option<VertexModel>,
/// The edges projected out of this table's FK columns.
pub edges: &'static [EdgeModel],
/// All columns (Arrow schema + generator fill rules).
pub columns: Vec<Column>,
}
impl Table {
/// The node label produced by this table, if any.
pub fn label(&self) -> Option<&'static str> {
self.vertex.as_ref().map(|v| v.label)
}
/// Build the Arrow schema for this table's rows.
pub fn arrow_schema(&self) -> Schema {
Schema::new(
self.columns
.iter()
.map(|c| Field::new(c.name, c.dtype.clone(), c.nullable))
.collect::<Vec<_>>(),
)
}
/// The row count this table generates at `scale` (node tables scale
/// linearly; junctions are computed by the generator and reported as 0
/// here — call the generator for their true counts).
pub fn rows_at(&self, scale: usize) -> usize {
self.base_rows * scale
}
/// Lower this table to a bifröst [`TableMapping`] (a vertex step, if any,
/// then one `fk_edges` step per edge).
pub fn to_mapping(&self) -> TableMapping {
let mut pipeline: Vec<PipelineStep> = Vec::new();
if let Some(v) = &self.vertex {
pipeline.push(PipelineStep::Vertex {
vertex: VertexTransform {
label: v.label.to_string(),
id_column: v.key.to_string(),
properties: v.props.iter().map(|s| s.to_string()).collect(),
rename: Default::default(),
},
});
}
for e in self.edges {
pipeline.push(PipelineStep::FkEdges {
fk_edges: EdgeTransform {
edge_type: e.rel.to_string(),
src: EndpointSpec {
label: e.src_label.to_string(),
column: e.src_col.to_string(),
id_column: Some(e.src_key.to_string()),
},
dst: EndpointSpec {
label: e.dst_label.to_string(),
column: e.dst_col.to_string(),
id_column: Some(e.dst_key.to_string()),
},
properties: e.props.iter().map(|s| s.to_string()).collect(),
rename: Default::default(),
key: Vec::new(),
replace: None,
},
});
}
TableMapping {
table: self.name.to_string(),
mode: Default::default(),
pipeline,
columns: Vec::new(),
filter: None,
}
}
}
// ── column helpers ────────────────────────────────────────────────────────────
fn i64c(name: &'static str, fill: Fill) -> Column {
Column {
name,
dtype: DataType::Int64,
nullable: false,
fill,
}
}
fn i64c_null(name: &'static str, fill: Fill) -> Column {
Column {
name,
dtype: DataType::Int64,
nullable: true,
fill,
}
}
fn f64c(name: &'static str, fill: Fill) -> Column {
Column {
name,
dtype: DataType::Float64,
nullable: false,
fill,
}
}
fn strc(name: &'static str, fill: Fill) -> Column {
Column {
name,
dtype: DataType::Utf8,
nullable: false,
fill,
}
}
const SEGMENTS: &[&str] = &["retail", "wholesale", "enterprise", "gov"];
const STATUSES: &[&str] = &["new", "picking", "shipped", "delivered", "returned"];
const METHODS: &[&str] = &["card", "invoice", "wire", "paypal"];
const COUNTRIES: &[&str] = &["SE", "NO", "DK", "FI", "DE", "US"];
/// **The 12-table ERP schema.** Node tables first (parents before children so
/// the generator can pick FK targets from already-known id ranges), junction
/// tables last.
pub fn schema() -> Vec<Table> {
vec![
// 1. Category (leaf reference table, no FKs)
Table {
name: "erp.ref.categories",
base_rows: 8,
junction: false,
vertex: Some(VertexModel {
label: "Category",
key: "category_id",
props: &["name", "description"],
}),
edges: &[],
columns: vec![
i64c("category_id", Fill::PkSerial),
strc("name", Fill::Name("cat")),
strc("description", Fill::Name("desc")),
],
},
// 2. Supplier
Table {
name: "erp.ref.suppliers",
base_rows: 20,
junction: false,
vertex: Some(VertexModel {
label: "Supplier",
key: "supplier_id",
props: &["name", "country", "rating"],
}),
edges: &[],
columns: vec![
i64c("supplier_id", Fill::PkSerial),
strc("name", Fill::Name("supp")),
strc("country", Fill::Enum(COUNTRIES)),
i64c("rating", Fill::IntRange(1, 5)),
],
},
// 3. Shipper
Table {
name: "erp.ref.shippers",
base_rows: 4,
junction: false,
vertex: Some(VertexModel {
label: "Shipper",
key: "shipper_id",
props: &["name", "phone"],
}),
edges: &[],
columns: vec![
i64c("shipper_id", Fill::PkSerial),
strc("name", Fill::Name("ship")),
strc("phone", Fill::Name("tel")),
],
},
// 4. Product → Category (FK edge IN_CATEGORY)
Table {
name: "erp.catalog.products",
base_rows: 120,
junction: false,
vertex: Some(VertexModel {
label: "Product",
key: "product_id",
props: &["name", "price", "in_stock"],
}),
edges: &[EdgeModel {
rel: "IN_CATEGORY",
src_label: "Product",
src_col: "product_id",
src_key: "product_id",
dst_label: "Category",
dst_col: "category_id",
dst_key: "category_id",
props: &[],
}],
columns: vec![
i64c("product_id", Fill::PkSerial),
i64c(
"category_id",
Fill::Fk {
parent: "erp.ref.categories",
null_prob: 0.0,
},
),
strc("name", Fill::Name("prod")),
f64c("price", Fill::Money(5.0, 500.0)),
i64c("in_stock", Fill::IntRange(0, 1000)),
],
},
// 5. Customer
Table {
name: "erp.crm.customers",
base_rows: 60,
junction: false,
vertex: Some(VertexModel {
label: "Customer",
key: "customer_id",
props: &["name", "segment", "created_at"],
}),
edges: &[],
columns: vec![
i64c("customer_id", Fill::PkSerial),
strc("name", Fill::Name("cust")),
strc("segment", Fill::Enum(SEGMENTS)),
i64c("created_at", Fill::Epoch),
],
},
// 6. Address → Customer (FK edge BELONGS_TO)
Table {
name: "erp.crm.addresses",
base_rows: 90,
junction: false,
vertex: Some(VertexModel {
label: "Address",
key: "address_id",
props: &["city", "zip", "country"],
}),
edges: &[EdgeModel {
rel: "BELONGS_TO",
src_label: "Address",
src_col: "address_id",
src_key: "address_id",
dst_label: "Customer",
dst_col: "customer_id",
dst_key: "customer_id",
props: &[],
}],
columns: vec![
i64c("address_id", Fill::PkSerial),
i64c(
"customer_id",
Fill::Fk {
parent: "erp.crm.customers",
null_prob: 0.0,
},
),
strc("city", Fill::Name("city")),
strc("zip", Fill::Name("zip")),
strc("country", Fill::Enum(COUNTRIES)),
],
},
// 7. Employee → Employee (SELF-REF edge REPORTS_TO, nullable at roots)
Table {
name: "erp.hr.employees",
base_rows: 24,
junction: false,
vertex: Some(VertexModel {
label: "Employee",
key: "employee_id",
props: &["name", "title", "hired_at"],
}),
edges: &[EdgeModel {
rel: "REPORTS_TO",
src_label: "Employee",
src_col: "employee_id",
src_key: "employee_id",
dst_label: "Employee",
dst_col: "manager_id",
dst_key: "employee_id",
props: &[],
}],
columns: vec![
i64c("employee_id", Fill::PkSerial),
i64c_null("manager_id", Fill::SelfRefManager { roots: 3 }),
strc("name", Fill::Name("emp")),
strc("title", Fill::Name("title")),
i64c("hired_at", Fill::Epoch),
],
},
// 8. Order → Customer / Employee / Shipper (three FK edges)
Table {
name: "erp.sales.orders",
base_rows: 300,
junction: false,
vertex: Some(VertexModel {
label: "Order",
key: "order_id",
props: &["order_date", "status"],
}),
edges: &[
EdgeModel {
rel: "PLACED_BY",
src_label: "Order",
src_col: "order_id",
src_key: "order_id",
dst_label: "Customer",
dst_col: "customer_id",
dst_key: "customer_id",
props: &[],
},
EdgeModel {
rel: "HANDLED_BY",
src_label: "Order",
src_col: "order_id",
src_key: "order_id",
dst_label: "Employee",
dst_col: "employee_id",
dst_key: "employee_id",
props: &[],
},
EdgeModel {
rel: "SHIPPED_VIA",
src_label: "Order",
src_col: "order_id",
src_key: "order_id",
dst_label: "Shipper",
dst_col: "shipper_id",
dst_key: "shipper_id",
props: &[],
},
],
columns: vec![
i64c("order_id", Fill::PkSerial),
i64c(
"customer_id",
Fill::Fk {
parent: "erp.crm.customers",
null_prob: 0.0,
},
),
i64c(
"employee_id",
Fill::Fk {
parent: "erp.hr.employees",
null_prob: 0.0,
},
),
i64c(
"shipper_id",
Fill::Fk {
parent: "erp.ref.shippers",
null_prob: 0.0,
},
),
i64c("order_date", Fill::Epoch),
strc("status", Fill::Enum(STATUSES)),
],
},
// 9. Payment → Order (FK edge PAYS)
Table {
name: "erp.sales.payments",
base_rows: 280,
junction: false,
vertex: Some(VertexModel {
label: "Payment",
key: "payment_id",
props: &["amount", "method", "paid_at"],
}),
edges: &[EdgeModel {
rel: "PAYS",
src_label: "Payment",
src_col: "payment_id",
src_key: "payment_id",
dst_label: "Order",
dst_col: "order_id",
dst_key: "order_id",
props: &[],
}],
columns: vec![
i64c("payment_id", Fill::PkSerial),
i64c(
"order_id",
Fill::Fk {
parent: "erp.sales.orders",
null_prob: 0.0,
},
),
f64c("amount", Fill::Money(10.0, 5000.0)),
strc("method", Fill::Enum(METHODS)),
i64c("paid_at", Fill::Epoch),
],
},
// 10. Review → Product / Customer (two FK edges)
Table {
name: "erp.crm.reviews",
base_rows: 200,
junction: false,
vertex: Some(VertexModel {
label: "Review",
key: "review_id",
props: &["rating", "body"],
}),
edges: &[
EdgeModel {
rel: "REVIEWS",
src_label: "Review",
src_col: "review_id",
src_key: "review_id",
dst_label: "Product",
dst_col: "product_id",
dst_key: "product_id",
props: &[],
},
EdgeModel {
rel: "WRITTEN_BY",
src_label: "Review",
src_col: "review_id",
src_key: "review_id",
dst_label: "Customer",
dst_col: "customer_id",
dst_key: "customer_id",
props: &[],
},
],
columns: vec![
i64c("review_id", Fill::PkSerial),
i64c(
"product_id",
Fill::Fk {
parent: "erp.catalog.products",
null_prob: 0.0,
},
),
i64c(
"customer_id",
Fill::Fk {
parent: "erp.crm.customers",
null_prob: 0.0,
},
),
i64c("rating", Fill::IntRange(1, 5)),
strc("body", Fill::Name("rev")),
],
},
// 11. JUNCTION: order_items — composite PK (order_id, product_id),
// collapses to Order-[:CONTAINS]->Product. NO VERTEX.
Table {
name: "erp.sales.order_items",
base_rows: 0,
junction: true,
vertex: None,
edges: &[EdgeModel {
rel: "CONTAINS",
src_label: "Order",
src_col: "order_id",
src_key: "order_id",
dst_label: "Product",
dst_col: "product_id",
dst_key: "product_id",
props: &["quantity", "unit_price"],
}],
columns: vec![
i64c(
"order_id",
Fill::Fk {
parent: "erp.sales.orders",
null_prob: 0.0,
},
),
i64c(
"product_id",
Fill::Fk {
parent: "erp.catalog.products",
null_prob: 0.0,
},
),
i64c("quantity", Fill::IntRange(1, 20)),
f64c("unit_price", Fill::Money(5.0, 500.0)),
],
},
// 12. JUNCTION: product_suppliers — composite PK (product_id, supplier_id),
// collapses to Product-[:SUPPLIED_BY]->Supplier. NO VERTEX.
Table {
name: "erp.catalog.product_suppliers",
base_rows: 0,
junction: true,
vertex: None,
edges: &[EdgeModel {
rel: "SUPPLIED_BY",
src_label: "Product",
src_col: "product_id",
src_key: "product_id",
dst_label: "Supplier",
dst_col: "supplier_id",
dst_key: "supplier_id",
props: &["supply_price", "lead_days"],
}],
columns: vec![
i64c(
"product_id",
Fill::Fk {
parent: "erp.catalog.products",
null_prob: 0.0,
},
),
i64c(
"supplier_id",
Fill::Fk {
parent: "erp.ref.suppliers",
null_prob: 0.0,
},
),
f64c("supply_price", Fill::Money(2.0, 400.0)),
i64c("lead_days", Fill::IntRange(1, 90)),
],
},
]
}
/// Assemble the full [`MappingSpec`] (version 1) from the schema.
pub fn mapping_spec() -> MappingSpec {
MappingSpec {
version: 1,
tables: schema().iter().map(|t| t.to_mapping()).collect(),
}
}
/// Convenience: the schema's tables as `Arc`-shared Arrow schemas keyed by name.
pub fn arrow_schemas() -> Vec<(&'static str, Arc<Schema>)> {
schema()
.iter()
.map(|t| (t.name, Arc::new(t.arrow_schema())))
.collect()
}