crates/knut-bifrost/tests/first_load_bulk_falkordb.rs
//! **First load in bulk, then MERGE** — against a real FalkorDB, and measured.
//!
//! The same N rows go to two graphs: `bulk` (an empty label: the first load is a
//! verbatim `GRAPH.BULK`) and `merge` (the label already holds a node, so the
//! same rows are `UNWIND … MERGE`d, on an indexed key). Both must end as the
//! same graph — N nodes, the same properties, no `_gar_id` — and a second round
//! into `bulk` (updates + new ids) must land as MERGE: updated, not duplicated.
//! The wall time of each first load is printed.
//!
//! ```bash
//! KNUT_FALKORDB_IMAGE=localhost/nordisk-falkordb-valkey:v4.20.1 KNUT_BULK_ROWS=200000 \
//! cargo test -p knut-bifrost --release --test first_load_bulk_falkordb -- --ignored --nocapture
//! ```
use std::sync::Arc;
use std::time::{Duration, Instant};
use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use knut_bifrost::sink::{FalkorDbSink, GraphSink};
use knut_bifrost::{GraphMutation, mutation::{EdgeSet, VertexSet}};
use testcontainers::core::ContainerPort;
use testcontainers::{GenericImage, ImageExt, runners::AsyncRunner};
fn people(ids: std::ops::Range<i64>, tag: &str) -> VertexSet {
let ids: Vec<i64> = ids.collect();
VertexSet {
label: "Person".into(),
id_column: "id".into(),
batch: RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("age", DataType::Int64, false),
])),
vec![
Arc::new(Int64Array::from(ids.clone())),
Arc::new(StringArray::from(ids.iter().map(|i| format!("{tag}{i}")).collect::<Vec<_>>())),
Arc::new(Int64Array::from(ids.iter().map(|i| i % 90).collect::<Vec<_>>())),
],
)
.unwrap(),
}
}
fn companies(n: i64) -> VertexSet {
let ids: Vec<i64> = (0..n).collect();
VertexSet {
label: "Company".into(),
id_column: "id".into(),
batch: RecordBatch::try_new(
Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false), Field::new("name", DataType::Utf8, false)])),
vec![
Arc::new(Int64Array::from(ids.clone())),
Arc::new(StringArray::from(ids.iter().map(|i| format!("c{i}")).collect::<Vec<_>>())),
],
)
.unwrap(),
}
}
/// Person i WORKS_AT company i % c, `since` = i; plus two duplicate rows (the
/// last must win) and two rows to companies that do not exist (MERGE makes them).
fn works_at(n: i64, c: i64, tag: i64) -> EdgeSet {
let mut src: Vec<i64> = (0..n).collect();
let mut dst: Vec<i64> = (0..n).map(|i| i % c).collect();
let mut since: Vec<i64> = (0..n).map(|i| i + tag).collect();
src.extend([0, 1, n - 1, n - 2]);
dst.extend([0, 1, c + 1, c + 2]);
since.extend([-1, -2, 7, 8]);
EdgeSet {
src_label: "Person".into(),
edge_type: "WORKS_AT".into(),
dst_label: "Company".into(),
src_id_column: "id".into(),
dst_id_column: "id".into(),
edge_key_props: Vec::new(),
replace_outgoing: false,
batch: RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("src", DataType::Int64, false),
Field::new("dst", DataType::Int64, false),
Field::new("since", DataType::Int64, false),
])),
vec![Arc::new(Int64Array::from(src)), Arc::new(Int64Array::from(dst)), Arc::new(Int64Array::from(since))],
)
.unwrap(),
}
}
fn round(sets: Vec<VertexSet>) -> GraphMutation {
GraphMutation { vertices: sets, ..Default::default() }
}
fn ensure_docker_host() {
if std::env::var("DOCKER_HOST").is_err() {
let uid = std::process::Command::new("id")
.arg("-u")
.output()
.ok()
.and_then(|o| String::from_utf8(o.stdout).ok())
.and_then(|s| s.trim().parse::<u32>().ok())
.unwrap_or(1000);
let sock = format!("/run/user/{uid}/podman/podman.sock");
if std::path::Path::new(&sock).exists() {
unsafe { std::env::set_var("DOCKER_HOST", format!("unix://{sock}")) };
}
}
}
async fn scalar(conn: &mut redis::aio::MultiplexedConnection, graph: &str, q: &str) -> i64 {
fn walk(v: &redis::Value) -> Option<i64> {
match v {
redis::Value::Int(i) => Some(*i),
redis::Value::Array(items) | redis::Value::Set(items) => items.iter().find_map(walk),
redis::Value::BulkString(b) => String::from_utf8_lossy(b).parse().ok(),
_ => None,
}
}
let v: redis::Value = redis::cmd("GRAPH.QUERY").arg(graph).arg(q).query_async(conn).await.unwrap();
walk(&v).unwrap_or(-1)
}
async fn indexed(conn: &mut redis::aio::MultiplexedConnection, graph: &str) -> bool {
let v: redis::Value = redis::cmd("GRAPH.QUERY").arg(graph).arg("CALL db.indexes()").query_async(conn).await.unwrap();
let text = format!("{v:?}");
text.contains("Person") && text.contains("id")
}
#[tokio::test]
#[ignore = "needs podman and a FalkorDB image; run with --ignored"]
async fn the_first_load_is_bulk_then_merge_and_both_make_the_same_graph() {
ensure_docker_host();
let n: i64 = std::env::var("KNUT_BULK_ROWS").ok().and_then(|v| v.parse().ok()).unwrap_or(50_000);
let image = std::env::var("KNUT_FALKORDB_IMAGE").unwrap_or_else(|_| "localhost/knut-falkordb-valkey9:latest".into());
let (name, tag) = image.rsplit_once(':').unwrap_or((image.as_str(), "latest"));
let falkor = GenericImage::new(name, tag)
.with_exposed_port(ContainerPort::Tcp(6379))
.with_startup_timeout(Duration::from_secs(120))
.start()
.await
.expect("start falkordb (build the image first, or set KNUT_FALKORDB_IMAGE)");
let url = format!("redis://127.0.0.1:{}", falkor.get_host_port_ipv4(6379).await.unwrap());
let mut conn = redis::Client::open(url.clone()).unwrap().get_multiplexed_async_connection().await.unwrap();
// ── `merge`: one node already there, so the label is not empty ─────────
let mut merge = FalkorDbSink::connect(&url, "merge").await.unwrap();
merge.apply(&round(vec![people(-1..0, "seed")])).await.unwrap();
let t = Instant::now();
merge.apply(&round(vec![people(0..n, "p")])).await.unwrap();
let merge_ms = t.elapsed().as_millis();
// ── `bulk`: an empty graph ─────────────────────────────────────────────
let mut bulk = FalkorDbSink::connect(&url, "bulk").await.unwrap();
let t = Instant::now();
let r = bulk.apply(&round(vec![people(0..n, "p")])).await.unwrap();
let bulk_ms = t.elapsed().as_millis();
assert_eq!(r.vertices_written as i64, n);
eprintln!(
"first load of {n} nodes: GRAPH.BULK {bulk_ms} ms, indexed UNWIND…MERGE {merge_ms} ms ({:.1}x)",
merge_ms as f64 / bulk_ms.max(1) as f64
);
// ── the same graph ─────────────────────────────────────────────────────
assert_eq!(scalar(&mut conn, "bulk", "MATCH (n:Person) RETURN count(n)").await, n);
assert_eq!(scalar(&mut conn, "merge", "MATCH (n:Person) WHERE n.id >= 0 RETURN count(n)").await, n);
for g in ["bulk", "merge"] {
assert_eq!(scalar(&mut conn, g, "MATCH (n:Person) WHERE n._gar_id IS NOT NULL RETURN count(n)").await, 0, "{g}: no _gar_id");
assert_eq!(
scalar(&mut conn, g, &format!("MATCH (n:Person {{id: {}}}) WHERE n.name = 'p{}' AND n.age = {} RETURN count(n)", n - 1, n - 1, (n - 1) % 90)).await,
1,
"{g}: properties"
);
assert!(indexed(&mut conn, g).await, "{g}: Person(id) is indexed");
}
// ── EDGES: the first load of a relationship type ────────────────────────
let c = (n / 10).max(3);
for (g, sink) in [("bulk", &mut bulk), ("merge", &mut merge)] {
sink.apply(&round(vec![companies(c)])).await.unwrap();
assert_eq!(scalar(&mut conn, g, "MATCH (n:Company) RETURN count(n)").await, c);
}
// `merge` already has one WORKS_AT edge, so its load is MERGE throughout.
merge
.apply(&GraphMutation {
edges: vec![EdgeSet { batch: works_at(3, c, 0).batch.slice(0, 1), ..works_at(3, c, 0) }],
..Default::default()
})
.await
.unwrap();
let edges = |sink_n| GraphMutation { edges: vec![works_at(sink_n, c, 0)], ..Default::default() };
let t = Instant::now();
let rm = merge.apply(&edges(n)).await.unwrap();
let merge_edges_ms = t.elapsed().as_millis();
let t = Instant::now();
let rb = bulk.apply(&edges(n)).await.unwrap();
let bulk_edges_ms = t.elapsed().as_millis();
eprintln!(
"first load of {n} edges: GRAPH.BULK {bulk_edges_ms} ms ({} written), indexed MERGE {merge_edges_ms} ms ({} written) ({:.1}x)",
rb.edges_written,
rm.edges_written,
merge_edges_ms as f64 / bulk_edges_ms.max(1) as f64
);
for g in ["bulk", "merge"] {
assert_eq!(scalar(&mut conn, g, "MATCH ()-[e:WORKS_AT]->() RETURN count(e)").await, n + 2, "{g}: n edges + 2 to new companies, duplicates collapsed");
assert_eq!(scalar(&mut conn, g, "MATCH (:Person {id: 0})-[e:WORKS_AT]->(:Company {id: 0}) WHERE e.since = -1 RETURN count(e)").await, 1, "{g}: the last duplicate wins");
assert_eq!(scalar(&mut conn, g, &format!("MATCH (:Person {{id: {}}})-[e:WORKS_AT]->(:Company {{id: {}}}) WHERE e.since = 7 RETURN count(e)", n - 1, c + 1)).await, 1, "{g}: an edge to a company not loaded");
assert_eq!(scalar(&mut conn, g, "MATCH (n:Company) RETURN count(n)").await, c + 2, "{g}: MERGE made the 2 missing companies");
assert_eq!(scalar(&mut conn, g, &format!("MATCH (:Person {{id: {}}})-[e:WORKS_AT]->(:Company {{id: {}}}) WHERE e.since = {} RETURN count(e)", n / 2, (n / 2) % c, n / 2)).await, 1, "{g}: properties");
}
// The second edge round into `bulk` is a MERGE: updated, not duplicated.
bulk.apply(&GraphMutation { edges: vec![works_at(5, c, 1000)], ..Default::default() }).await.unwrap();
// works_at(5, …) is persons 0..5 (already there) + its 4 extra rows: two of
// them duplicates of existing edges, two NEW edges (4 -> c+1, 3 -> c+2).
assert_eq!(scalar(&mut conn, "bulk", "MATCH ()-[e:WORKS_AT]->() RETURN count(e)").await, n + 4, "2 new edges, no duplicates");
assert_eq!(scalar(&mut conn, "bulk", "MATCH (:Person {id: 4})-[e:WORKS_AT]->() WHERE e.since = 1004 RETURN count(e)").await, 1, "updated");
// ── the second round into `bulk` is a MERGE: updated, not duplicated ───
let r = bulk.apply(&round(vec![people(n - 10..n + 10, "q")])).await.unwrap();
assert_eq!(r.vertices_written, 20);
assert_eq!(scalar(&mut conn, "bulk", "MATCH (n:Person) RETURN count(n)").await, n + 10, "10 updated, 10 new");
assert_eq!(scalar(&mut conn, "bulk", &format!("MATCH (n:Person {{id: {}}}) WHERE n.name = 'q{}' RETURN count(n)", n - 1, n - 1)).await, 1);
}