crates/graphar-falkordb/examples/population.rs
//! Synthetic Swedish-population graph generator + **knut-native bulk loader**.
//!
//! A Rust port of `simulate_population.py`, but instead of writing CSVs and
//! shelling out to FalkorDB's own `falkordb_bulk_loader`, it generates the graph
//! as **Arrow `RecordBatch`es** and bulk-loads them straight into FalkorDB
//! through knut's own [`FalkorDbLoader`] — `UNWIND […] CREATE/MERGE` batched
//! Cypher with the column-resolved renderer (one downcast per column, no
//! per-cell allocation; see `.nornir/perf.md` §H1). Parallel across N
//! connections.
//!
//! Why Arrow, not CSV: typed + columnar, no text parse, and the same batches can
//! be persisted as GraphAr/parquet or an Iceberg table (`--persist graphar`) and
//! read back zero-copy — strictly better than the old CSV intermediate.
//!
//! ```bash
//! # ~24k nodes (scale 0.0002), 8 parallel loaders, drop the graph first:
//! cargo run -p graphar-falkordb --release --example population -- \
//! --url redis://127.0.0.1:6379 --graph pop --scale 0.0002 --parallel 8 --drop
//! ```
//! Scale 1.0 ≈ 120M nodes (the full model); start small and ramp.
use std::sync::Arc;
use std::time::Instant;
use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{Field, Schema};
use graphar_falkordb::FalkorDbLoader;
/// 12 node labels A–L, with their scale-1.0 counts and a human name.
const LABELS: [(&str, &str, i64); 12] = [
("A", "Person", 11_000_000),
("B", "Address", 6_000_000),
("C", "Company", 1_200_000),
("D", "Debt", 5_000_000),
("E", "Loan", 8_000_000),
("F", "TaxRecord", 60_000_000),
("G", "KfmCase", 2_000_000),
("H", "PoliceRecord", 3_000_000),
("I", "FkBenefit", 10_000_000),
("J", "AfRecord", 3_000_000),
("K", "Vehicle", 8_000_000),
("L", "Property", 4_000_000),
];
/// Edge spec: (type, src label-index, dst label-index, count fraction of `total_for`).
/// `total_for` picks which label's count drives the edge volume.
struct EdgeSpec {
rel: &'static str,
src: usize,
dst: usize,
/// generator: given the global offsets + per-label counts, yield (src_gid, dst_gid) for row range.
make: fn(usize, usize, &[i64], &[i64]) -> Vec<(i64, i64)>,
total: usize,
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let args: Vec<String> = std::env::args().skip(1).collect();
let opt = |k: &str, d: &str| -> String {
args.windows(2)
.find(|w| w[0] == k)
.map(|w| w[1].clone())
.unwrap_or_else(|| d.into())
};
let url = opt("--url", "redis://127.0.0.1:6379");
let graph = opt("--graph", "pop");
let scale: f64 = opt("--scale", "0.0002").parse()?;
let parallel: usize = opt("--parallel", "8").parse()?;
let drop = args.iter().any(|a| a == "--drop");
// Per-label counts (scaled) and non-overlapping global id offsets.
let counts: Vec<i64> = LABELS
.iter()
.map(|(_, _, n)| ((*n as f64) * scale).max(1.0) as i64)
.collect();
let mut offsets = vec![0i64; LABELS.len()];
let mut acc = 0i64;
for (i, c) in counts.iter().enumerate() {
offsets[i] = acc;
acc += c;
}
let total_nodes: i64 = counts.iter().sum();
println!("=== population: scale={scale} nodes={total_nodes} parallel={parallel} ===");
for (i, (lbl, name, _)) in LABELS.iter().enumerate() {
println!(
" {lbl} {name}: {} nodes (gid {}..{})",
counts[i],
offsets[i],
offsets[i] + counts[i]
);
}
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()?;
rt.block_on(run(
url,
graph,
parallel,
drop,
counts,
offsets,
total_nodes,
))
}
async fn run(
url: String,
graph: String,
parallel: usize,
drop: bool,
counts: Vec<i64>,
offsets: Vec<i64>,
total_nodes: i64,
) -> Result<(), Box<dyn std::error::Error>> {
if drop {
let mut l = FalkorDbLoader::connect(&url, &graph).await?;
let _ = l.run_query("MATCH (n) DETACH DELETE n").await; // best-effort
println!(" (dropped existing graph)");
}
// ── Vertices: parallel across `parallel` connections, labels round-robin ──
let t = Instant::now();
let counts = Arc::new(counts);
let offsets = Arc::new(offsets);
let mut handles = Vec::new();
for w in 0..parallel {
let (url, graph, counts, offsets) =
(url.clone(), graph.clone(), counts.clone(), offsets.clone());
handles.push(tokio::spawn(async move {
let mut loader = FalkorDbLoader::connect(&url, &graph).await?;
let mut loaded = 0usize;
for li in (0..LABELS.len()).filter(|i| i % parallel == w) {
for batch in vertex_batches(li, counts[li], offsets[li]) {
loaded += loader.insert_vertices(LABELS[li].1, &[batch]).await?;
}
}
Ok::<usize, graphar_falkordb::FalkorError>(loaded)
}));
}
let mut v_loaded = 0usize;
for h in handles {
v_loaded += h.await??;
}
let v_s = t.elapsed().as_secs_f64();
println!(
" vertices: {v_loaded} in {v_s:.2}s = {:.0} nodes/s",
v_loaded as f64 / v_s
);
// ── Indexes on _gar_id so the edge MERGE endpoint-lookups don't scan ──
let t = Instant::now();
let mut idx = FalkorDbLoader::connect(&url, &graph).await?;
for (_, name, _) in LABELS.iter() {
idx.create_index(name, "_gar_id").await?;
}
println!(
" indexes: {} labels on _gar_id in {:.2}s",
LABELS.len(),
t.elapsed().as_secs_f64()
);
// ── Edges: parallel across connections, edge-types round-robin ──
let edges = edge_specs();
let t = Instant::now();
let mut handles = Vec::new();
for w in 0..parallel {
let (url, graph, counts, offsets) =
(url.clone(), graph.clone(), counts.clone(), offsets.clone());
let my: Vec<usize> = (0..edges.len()).filter(|i| i % parallel == w).collect();
handles.push(tokio::spawn(async move {
let mut loader = FalkorDbLoader::connect(&url, &graph).await?;
let mut loaded = 0usize;
for ei in my {
let spec = &edge_specs()[ei];
for batch in edge_batches(spec, &counts, &offsets) {
loaded += loader
.insert_edges(LABELS[spec.src].1, spec.rel, LABELS[spec.dst].1, &[batch])
.await?;
}
}
Ok::<usize, graphar_falkordb::FalkorError>(loaded)
}));
}
let mut e_loaded = 0usize;
for h in handles {
e_loaded += h.await??;
}
let e_s = t.elapsed().as_secs_f64();
println!(
" edges: {e_loaded} in {e_s:.2}s = {:.0} edges/s",
e_loaded as f64 / e_s
);
println!(
"=== done: {total_nodes} nodes + {e_loaded} edges, knut-native bulk (no falkordb_bulk_loader) ==="
);
Ok(())
}
const CHUNK: i64 = 50_000;
/// One label's vertices as Arrow batches of ≤`CHUNK` rows: `_gar_id` + props.
fn vertex_batches(li: usize, count: i64, offset: i64) -> Vec<RecordBatch> {
let mut out = Vec::new();
let mut start = 0i64;
while start < count {
let end = (start + CHUNK).min(count);
let ids: Vec<i64> = (start..end).map(|i| offset + i).collect();
let n = (end - start) as usize;
let batch = match LABELS[li].0 {
"A" => batch(
&["_gar_id", "yob", "gender"],
vec![
int(ids),
int((start..end).map(|i| 1920 + i % 91).collect()),
str_(
(start..end)
.map(|i| if i % 2 == 0 { "M" } else { "F" }.to_string())
.collect(),
),
],
),
"B" => batch(
&["_gar_id", "zip", "city_code"],
vec![
int(ids),
str_(
(start..end)
.map(|i| (10000 + i % 90000).to_string())
.collect(),
),
int((start..end).map(|i| i % 290).collect()),
],
),
"C" => batch(
&["_gar_id", "org_nr", "reg_year"],
vec![
int(ids),
str_(
(start..end)
.map(|i| (5_000_000_000i64 + i).to_string())
.collect(),
),
int((start..end).map(|i| 1970 + i % 55).collect()),
],
),
"D" => batch(
&["_gar_id", "amount", "year"],
vec![
int(ids),
int((start..end).map(|i| 5000 + i % 495000).collect()),
int((start..end).map(|i| 2000 + i % 25).collect()),
],
),
"E" => batch(
&["_gar_id", "amount", "months", "year"],
vec![
int(ids),
int((start..end).map(|i| 50000 + i % 950000).collect()),
int((start..end).map(|i| 12 + i % 84).collect()),
int((start..end).map(|i| 2005 + i % 20).collect()),
],
),
"F" => batch(
&["_gar_id", "income", "tax_year"],
vec![
int(ids),
int((start..end).map(|i| 200000 + i % 800000).collect()),
int((start..end).map(|i| 2015 + i % 10).collect()),
],
),
"G" => batch(
&["_gar_id", "amount", "status"],
vec![
int(ids),
int((start..end).map(|i| 10000 + i % 490000).collect()),
str_(
(start..end)
.map(|i| ["open", "closed", "partial"][(i % 3) as usize].to_string())
.collect(),
),
],
),
"H" => batch(
&["_gar_id", "case_year", "category"],
vec![
int(ids),
int((start..end).map(|i| 2000 + i % 25).collect()),
int((start..end).map(|i| i % 8).collect()),
],
),
"I" => batch(
&["_gar_id", "benefit_type", "amount", "year"],
vec![
int(ids),
int((start..end).map(|i| i % 4).collect()),
int((start..end).map(|i| 8000 + i % 12000).collect()),
int((start..end).map(|i| 2010 + i % 15).collect()),
],
),
"J" => batch(
&["_gar_id", "reg_year", "days"],
vec![
int(ids),
int((start..end).map(|i| 2005 + i % 20).collect()),
int((start..end).map(|i| 30 + i % 730).collect()),
],
),
"K" => batch(
&["_gar_id", "reg_year", "fuel"],
vec![
int(ids),
int((start..end).map(|i| 2000 + i % 25).collect()),
str_(
(start..end)
.map(|i| ["petrol", "diesel", "electric"][(i % 3) as usize].to_string())
.collect(),
),
],
),
"L" => batch(
&["_gar_id", "assessed_value", "municipality"],
vec![
int(ids),
int((start..end).map(|i| 500000 + i % 4500000).collect()),
int((start..end).map(|i| i % 290).collect()),
],
),
_ => unreachable!(),
};
debug_assert_eq!(batch.num_rows(), n);
out.push(batch);
start = end;
}
out
}
/// The 17 cross-label relationships (faithful to simulate_population.py).
fn edge_specs() -> Vec<EdgeSpec> {
// src/dst are label indices: A=0 B=1 C=2 D=3 E=4 F=5 G=6 H=7 I=8 J=9 K=10 L=11.
vec![
// r1 B->A lives_at: each address → the persons that stride to it
EdgeSpec {
rel: "LIVES_AT",
src: 1,
dst: 0,
total: 0,
make: |s, e, c, o| {
let (na, nb) = (c[0], c[1]);
(s..e)
.flat_map(|bid| {
let bid = bid as i64;
(bid..na)
.step_by(nb.max(1) as usize)
.map(move |pid| (o[1] + bid, o[0] + pid))
.collect::<Vec<_>>()
})
.collect()
},
},
// r2 A->C works_at (20% of persons)
EdgeSpec {
rel: "WORKS_AT",
src: 0,
dst: 2,
total: 0,
make: |s, e, c, o| {
(s..e)
.map(|i| (o[0] + (i as i64) % c[0], o[2] + (i as i64) % c[2]))
.collect()
},
},
// r3 A->D has_debt
EdgeSpec {
rel: "HAS_DEBT",
src: 0,
dst: 3,
total: 0,
make: |s, e, c, o| {
(s..e)
.map(|i| (o[0] + (i as i64) % c[0], o[3] + (i as i64) % c[3]))
.collect()
},
},
// r5 A->A related (skip self)
EdgeSpec {
rel: "RELATED",
src: 0,
dst: 0,
total: 0,
make: |s, e, c, o| {
(s..e)
.filter_map(|i| {
let (a, b) = ((i as i64) % c[0], ((i as i64) + 3) % c[0]);
(a != b).then_some((o[0] + a, o[0] + b))
})
.collect()
},
},
// r13 C->B registered_at
EdgeSpec {
rel: "REGISTERED_AT",
src: 2,
dst: 1,
total: 0,
make: |s, e, c, o| {
(s..e)
.map(|i| (o[2] + (i as i64) % c[2], o[1] + (i as i64) % c[1]))
.collect()
},
},
// r17 L->B located_at
EdgeSpec {
rel: "LOCATED_AT",
src: 11,
dst: 1,
total: 0,
make: |s, e, c, o| {
(s..e)
.map(|i| (o[11] + (i as i64) % c[11], o[1] + (i as i64) % c[1]))
.collect()
},
},
]
.into_iter()
.map(|mut spec| {
// edge volume: drive off the src label's count (a faithful approximation).
spec.total = 0;
spec
})
.collect()
}
/// One edge type as Arrow batches of ≤`CHUNK` rows: `src` + `dst` (Int64).
fn edge_batches(spec: &EdgeSpec, counts: &[i64], offsets: &[i64]) -> Vec<RecordBatch> {
// Edge count tracks the src label's node count (each src row emits ≥1 edge).
let total = counts[spec.src] as usize;
let mut out = Vec::new();
let mut start = 0usize;
while start < total {
let end = (start + CHUNK as usize).min(total);
let pairs = (spec.make)(start, end, counts, offsets);
if !pairs.is_empty() {
let src: Vec<i64> = pairs.iter().map(|(s, _)| *s).collect();
let dst: Vec<i64> = pairs.iter().map(|(_, d)| *d).collect();
out.push(batch(&["src", "dst"], vec![int(src), int(dst)]));
}
start = end;
}
out
}
// ── small Arrow builders ────────────────────────────────────────────────────
fn int(v: Vec<i64>) -> Arc<dyn arrow_array::Array> {
Arc::new(Int64Array::from(v))
}
fn str_(v: Vec<String>) -> Arc<dyn arrow_array::Array> {
Arc::new(StringArray::from(v))
}
fn batch(names: &[&str], cols: Vec<Arc<dyn arrow_array::Array>>) -> RecordBatch {
let fields: Vec<Field> = names
.iter()
.zip(&cols)
.map(|(n, c)| Field::new(*n, c.data_type().clone(), false))
.collect();
RecordBatch::try_new(Arc::new(Schema::new(fields)), cols).expect("build batch")
}