crates/knut-rel2graph-bench/src/generate.rs
//! The **scalable, seeded** synthetic-data generator.
//!
//! Parametrised by a `scale` factor and a `seed`. Node tables grow linearly with
//! `scale`; junction tables grow with their parents. Every row is produced by the
//! seeded [`Rng`] — no wall-clock, no thread RNG — so the same `(scale, seed)`
//! yields byte-for-byte the same relational instance on every run and machine.
//!
//! **Referential integrity is guaranteed by construction:** parents are generated
//! before children (schema order), every serial PK is `1..=count`, and every FK is
//! a uniform pick from the parent's known `1..=parent_count` range. A conformant
//! rel2graph mapping therefore has zero dangling endpoints — the correctness
//! bencher asserts exactly that.
use std::collections::{BTreeMap, HashSet};
use std::sync::Arc;
use arrow_array::{ArrayRef, Float64Array, Int64Array, RecordBatch, StringArray};
use arrow_schema::Schema;
use crate::prng::Rng;
use crate::schema::{Column, Fill, Table, schema};
/// A generated relational database instance: table name → its record batch(es).
#[derive(Debug, Clone)]
pub struct GeneratedDb {
pub scale: usize,
pub seed: u64,
tables: BTreeMap<String, Vec<RecordBatch>>,
}
impl GeneratedDb {
/// The batches generated for `table` (empty slice if unknown).
pub fn batches(&self, table: &str) -> &[RecordBatch] {
self.tables.get(table).map(|v| v.as_slice()).unwrap_or(&[])
}
/// Total rows generated for `table`.
pub fn row_count(&self, table: &str) -> usize {
self.batches(table).iter().map(|b| b.num_rows()).sum()
}
/// Every table name, in schema order.
pub fn table_names(&self) -> impl Iterator<Item = &str> {
self.tables.keys().map(|s| s.as_str())
}
/// Grand total of rows across all tables (a rough size metric).
pub fn total_rows(&self) -> usize {
self.tables.values().flatten().map(|b| b.num_rows()).sum()
}
}
/// The seeded synthetic database is a [`TableSource`](knut_bifrost::rel2graph::TableSource):
/// the bench feeds it to the engine's canonical
/// [`build_mutation`](knut_bifrost::rel2graph::build_mutation) exactly as a live
/// Iceberg / parquet read would feed an `ArrowTables`. Same engine, different data.
impl knut_bifrost::rel2graph::TableSource for GeneratedDb {
fn batches(&self, table: &str) -> &[RecordBatch] {
GeneratedDb::batches(self, table)
}
}
/// A stable per-table sub-seed so tables are independent streams and adding a
/// table never shifts another table's data.
fn table_seed(seed: u64, index: usize) -> u64 {
seed.wrapping_add((index as u64).wrapping_mul(0x100_0000_01B3)) ^ 0xA5A5_5A5A_C3C3_3C3C
}
/// Pick `k` **distinct** ids from `1..=count` (or all of them if `k >= count`).
fn distinct_picks(rng: &mut Rng, count: usize, k: usize) -> Vec<i64> {
let k = k.min(count);
let mut seen = HashSet::with_capacity(k);
let mut out = Vec::with_capacity(k);
// Bounded attempts: with k <= count this terminates quickly; the cap is a
// guard, not a correctness crutch.
let mut attempts = 0;
while out.len() < k && attempts < k * 16 + 32 {
attempts += 1;
let v = rng.range_i64(1, count as i64);
if seen.insert(v) {
out.push(v);
}
}
out
}
/// Generate the whole schema at `scale` with `seed`.
pub fn generate(scale: usize, seed: u64) -> GeneratedDb {
assert!(scale >= 1, "scale must be >= 1");
let tables = schema();
// Resolve node-table row counts first (parents before children).
let mut counts: BTreeMap<&'static str, usize> = BTreeMap::new();
for t in &tables {
if !t.junction {
counts.insert(t.name, t.rows_at(scale));
}
}
let mut out: BTreeMap<String, Vec<RecordBatch>> = BTreeMap::new();
for (idx, t) in tables.iter().enumerate() {
let mut rng = Rng::new(table_seed(seed, idx));
let arrow = Arc::new(t.arrow_schema());
let batch = if t.junction {
generate_junction(t, &arrow, &counts, &mut rng)
} else {
let n = *counts.get(t.name).unwrap();
generate_node(t, &arrow, n, &counts, &mut rng)
};
out.insert(t.name.to_string(), vec![batch]);
}
GeneratedDb {
scale,
seed,
tables: out,
}
}
/// Build one node/entity table of `n` rows, column by column in schema order.
fn generate_node(
t: &Table,
arrow: &Arc<Schema>,
n: usize,
counts: &BTreeMap<&'static str, usize>,
rng: &mut Rng,
) -> RecordBatch {
let arrays: Vec<ArrayRef> = t
.columns
.iter()
.map(|c| build_column(c, n, counts, rng))
.collect();
RecordBatch::try_new(arrow.clone(), arrays).expect("node batch schema matches arrays")
}
/// Build one column of `n` rows according to its [`Fill`] rule.
fn build_column(
c: &Column,
n: usize,
counts: &BTreeMap<&'static str, usize>,
rng: &mut Rng,
) -> ArrayRef {
match &c.fill {
Fill::PkSerial => {
let v: Vec<i64> = (1..=n as i64).collect();
Arc::new(Int64Array::from(v))
}
Fill::Fk { parent, null_prob } => {
let pc = *counts.get(parent).unwrap_or(&0);
let v: Vec<Option<i64>> = (0..n)
.map(|_| {
if pc == 0 || rng.chance(*null_prob) {
None
} else {
Some(rng.range_i64(1, pc as i64))
}
})
.collect();
Arc::new(Int64Array::from(v))
}
Fill::SelfRefManager { roots } => {
// Row i (1-based) reports to a uniformly chosen earlier row; the
// first `roots` rows are org roots with a NULL manager.
let v: Vec<Option<i64>> = (1..=n)
.map(|i| {
if i <= *roots {
None
} else {
Some(rng.range_i64(1, (i - 1) as i64))
}
})
.collect();
Arc::new(Int64Array::from(v))
}
Fill::IntRange(lo, hi) => {
let v: Vec<i64> = (0..n).map(|_| rng.range_i64(*lo, *hi)).collect();
Arc::new(Int64Array::from(v))
}
Fill::Money(lo, hi) => {
let v: Vec<f64> = (0..n).map(|_| rng.money(*lo, *hi)).collect();
Arc::new(Float64Array::from(v))
}
Fill::Name(prefix) => {
let v: Vec<String> = (0..n).map(|i| format!("{prefix}_{:06}", i + 1)).collect();
Arc::new(StringArray::from(v))
}
Fill::Enum(set) => {
let v: Vec<String> = (0..n)
.map(|_| set[rng.range_usize(0, set.len() - 1)].to_string())
.collect();
Arc::new(StringArray::from(v))
}
Fill::Epoch => {
let v: Vec<i64> = (0..n).map(|_| 19_000 + rng.range_i64(0, 3650)).collect();
Arc::new(Int64Array::from(v))
}
}
}
/// Build an M:N junction table. Its first two columns are the composite-PK FKs;
/// the remaining columns are edge properties. For each parent-1 id we fan out to
/// a handful of DISTINCT parent-2 ids, so `(fk1, fk2)` is unique — the composite
/// key the junction collapses to an edge on.
fn generate_junction(
t: &Table,
arrow: &Arc<Schema>,
counts: &BTreeMap<&'static str, usize>,
rng: &mut Rng,
) -> RecordBatch {
let Fill::Fk { parent: p1, .. } = &t.columns[0].fill else {
panic!("junction {} column 0 must be an FK", t.name);
};
let Fill::Fk { parent: p2, .. } = &t.columns[1].fill else {
panic!("junction {} column 1 must be an FK", t.name);
};
let c1 = *counts.get(p1).unwrap_or(&0);
let c2 = *counts.get(p2).unwrap_or(&0);
// Fan-out per parent-1 row: order_items 1..=5 products; product_suppliers
// 1..=3 suppliers. Distinguished by name to keep the shapes realistic.
let max_fanout = if t.name.contains("order_items") { 5 } else { 3 };
let mut fk1: Vec<i64> = Vec::new();
let mut fk2: Vec<i64> = Vec::new();
for id1 in 1..=c1 as i64 {
let k = rng.range_usize(1, max_fanout);
for pick in distinct_picks(rng, c2, k) {
fk1.push(id1);
fk2.push(pick);
}
}
let rows = fk1.len();
// Column order in the batch must match the schema: [fk1, fk2, props...].
let mut arrays: Vec<ArrayRef> = Vec::with_capacity(t.columns.len());
arrays.push(Arc::new(Int64Array::from(fk1)));
arrays.push(Arc::new(Int64Array::from(fk2)));
for c in &t.columns[2..] {
arrays.push(build_column(c, rows, counts, rng));
}
RecordBatch::try_new(arrow.clone(), arrays).expect("junction batch schema matches arrays")
}
#[cfg(test)]
mod tests {
use super::*;
/// Same `(scale, seed)` ⇒ identical row counts and identical first-column
/// values (determinism you can hang a benchmark on).
#[test]
fn generation_is_deterministic() {
let a = generate(2, 0xC0FFEE);
let b = generate(2, 0xC0FFEE);
for name in a.table_names() {
assert_eq!(a.row_count(name), b.row_count(name), "row count for {name}");
}
// A different seed changes the data (junction counts, at least).
let c = generate(2, 0xBEEF);
let ja = a.row_count("erp.sales.order_items");
let jc = c.row_count("erp.sales.order_items");
assert!(ja > 0 && jc > 0);
}
/// Node tables scale linearly.
#[test]
fn node_tables_scale_linearly() {
let s1 = generate(1, 7);
let s10 = generate(10, 7);
assert_eq!(
s1.row_count("erp.sales.orders") * 10,
s10.row_count("erp.sales.orders")
);
assert_eq!(
s1.row_count("erp.crm.customers") * 10,
s10.row_count("erp.crm.customers")
);
}
}