crates/graphar-flight/tests/integration_doput_parallel.rs
//! **Flight SQL DoPut: serial vs parallel bulk insert.**
//!
//! `do_put_statement_update` used to call the loader's SERIAL
//! `insert_vertices` / `insert_edges`, which await one `UNWIND … CREATE` at a
//! time over the single link — the client idles for a full round-trip between
//! windows, so a DoPut ingested at about one core no matter how big the box. The
//! in-process bulk loader had been on the concurrent path (`insert_*_par`,
//! `buffer_unordered` over cloned multiplexed connections) for a while. Same
//! crate, same renderer, two very different speeds for identical work.
//!
//! This test drives BOTH through the real Flight SQL wire and asserts the two
//! things a perf change has to prove:
//!
//! * **byte-identity** — the graph a parallel DoPut builds is byte-for-byte the
//! graph the serial DoPut builds (every `_gar_id`, every property, every edge).
//! Concurrency may change the time, never the data.
//! * **the number** — the wall time of each, so the speedup is measured, not
//! claimed.
//!
//! ## Prerequisites
//!
//! A reachable FalkorDB. Either export `KNUT_FALKORDB_URL` (e.g.
//! `redis://127.0.0.1:6379`) at an already-running one, or let the test start a
//! container itself:
//!
//! ```bash
//! systemctl --user start podman.socket
//! podman build -t knut-falkordb-valkey9 -f containers/falkordb-valkey9/Containerfile .
//! ```
//!
//! ## Run
//!
//! ```bash
//! cargo test -p graphar-flight --test integration_doput_parallel -- --ignored --nocapture
//! ```
use std::sync::Arc;
use std::time::{Duration, Instant};
use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_flight::FlightDescriptor;
use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::flight_service_client::FlightServiceClient;
use arrow_flight::sql::{Any, CommandStatementUpdate, DoPutUpdateResult};
use arrow_schema::{DataType, Field, Schema};
use futures::{StreamExt, TryStreamExt};
use graphar_flight::{SchemaRegistry, ServerOptions, run_server_with};
use prost::Message;
use testcontainers::{GenericImage, ImageExt, core::ContainerPort, runners::AsyncRunner};
/// Rows per DoPut. Large enough that the per-window round-trip dominates (that
/// is the whole cost the parallel dispatch hides), small enough to stay a test.
const ROWS: i64 = 40_000;
/// Vertex label under test.
const LABEL: &str = "Person";
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}")) };
}
}
}
fn free_port() -> u16 {
std::net::TcpListener::bind("127.0.0.1:0")
.unwrap()
.local_addr()
.unwrap()
.port()
}
/// The payload both DoPuts send — one batch, three columns.
fn people(rows: i64) -> RecordBatch {
let ids: Vec<i64> = (0..rows).collect();
let names: Vec<String> = ids.iter().map(|i| format!("person_{i}")).collect();
let cities: Vec<String> = ids
.iter()
.map(|i| ["Stockholm", "Göteborg", "Malmö", "Uppsala"][(i % 4) as usize].to_string())
.collect();
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("city", DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(ids)),
Arc::new(StringArray::from(names)),
Arc::new(StringArray::from(cities)),
],
)
.expect("people batch")
}
/// Start a Flight SQL server against `graph` with an explicit DoPut concurrency,
/// and return a raw Flight client connected to it.
///
/// A raw [`FlightServiceClient`] rather than the SQL wrapper: the stock
/// `FlightSqlServiceClient::execute_update` sends the command with **no record
/// batches**, and this test's whole subject is the batch-carrying DoPut.
async fn start_server(
redis_url: &str,
graph: &'static str,
put_concurrency: usize,
) -> FlightServiceClient<tonic::transport::Channel> {
let port = free_port();
let addr = format!("127.0.0.1:{port}").parse().unwrap();
let opts = ServerOptions {
put_concurrency: Some(put_concurrency),
..ServerOptions::default()
};
let url = redis_url.to_string();
tokio::spawn(async move {
run_server_with(&url, graph, Arc::new(SchemaRegistry::new()), addr, opts)
.await
.unwrap();
});
tokio::time::sleep(Duration::from_millis(400)).await;
let channel = tonic::transport::Channel::from_shared(format!("http://127.0.0.1:{port}"))
.unwrap()
.connect()
.await
.expect("connect flight channel");
FlightServiceClient::new(channel)
}
/// Send `batch` as a `CommandStatementUpdate` DoPut carrying `VERTICES:<LABEL>`,
/// and return `(rows_reported, wall_time)`.
async fn do_put_vertices(
client: &mut FlightServiceClient<tonic::transport::Channel>,
batch: RecordBatch,
) -> (i64, Duration) {
let cmd = CommandStatementUpdate {
query: format!("VERTICES:{LABEL}"),
transaction_id: None,
};
let descriptor = FlightDescriptor::new_cmd(Any::pack(&cmd).unwrap().encode_to_vec());
let stream = FlightDataEncoderBuilder::new()
.with_flight_descriptor(Some(descriptor))
.build(futures::stream::iter(vec![Ok(batch)]))
.map(|r| r.expect("encode flight data"));
let t = Instant::now();
let mut resp = client
.do_put(stream)
.await
.expect("DoPut accepted")
.into_inner();
let result = resp
.try_next()
.await
.expect("DoPut result")
.expect("one PutResult");
let elapsed = t.elapsed();
let decoded = DoPutUpdateResult::decode(&*result.app_metadata).expect("decode update result");
(decoded.record_count, elapsed)
}
/// Delete `graph` if it exists, so a re-run always starts from empty (a leftover
/// graph would make the two dumps differ for reasons that have nothing to do
/// with concurrency).
async fn reset_graph(redis_url: &str, graph: &str) {
let client = redis::Client::open(redis_url).unwrap();
let mut conn = client.get_multiplexed_async_connection().await.unwrap();
let _: std::result::Result<redis::Value, _> = redis::cmd("GRAPH.DELETE")
.arg(graph)
.query_async(&mut conn)
.await;
}
/// Read a graph back as a deterministic, canonical dump: one `String` per node,
/// carrying every property, ordered by `_gar_id`. Two graphs are identical iff
/// these vectors are equal — no per-property assertion needed, and a mismatch
/// can be reported as the FIRST differing row instead of a megabyte of debug.
async fn dump_graph(redis_url: &str, graph: &str) -> Vec<String> {
/// Rows per page. FalkorDB caps a single result set at `RESULTSET_SIZE`
/// (10000 by default), so a one-shot dump would silently truncate — and a
/// truncated comparison is not a byte-identity proof. Page instead of
/// mutating the server's global config.
const PAGE: usize = 5_000;
let client = redis::Client::open(redis_url).unwrap();
let mut conn = client.get_multiplexed_async_connection().await.unwrap();
let mut out: Vec<String> = Vec::new();
loop {
// One STRING column per node so the reply is trivially comparable — the
// typed/compact reply shapes would compare FalkorDB's encoding, not data.
let (_header, rows, _stats): (Vec<String>, Vec<Vec<String>>, Vec<String>) =
redis::cmd("GRAPH.QUERY")
.arg(graph)
.arg(format!(
"MATCH (n:{LABEL}) \
RETURN toString(n._gar_id) + '|' + toString(n.id) + '|' + n.name + '|' + n.city \
AS row ORDER BY n._gar_id SKIP {} LIMIT {PAGE}",
out.len()
))
.query_async(&mut conn)
.await
.expect("dump query");
if rows.is_empty() {
break;
}
out.extend(rows.into_iter().map(|mut r| r.remove(0)));
}
out
}
/// The node count, read straight from the engine (never capped by
/// `RESULTSET_SIZE` — it is one aggregate row).
async fn node_count(redis_url: &str, graph: &str) -> i64 {
let client = redis::Client::open(redis_url).unwrap();
let mut conn = client.get_multiplexed_async_connection().await.unwrap();
let (_h, rows, _s): (Vec<String>, Vec<Vec<i64>>, Vec<String>) = redis::cmd("GRAPH.QUERY")
.arg(graph)
.arg(format!("MATCH (n:{LABEL}) RETURN count(n)"))
.query_async(&mut conn)
.await
.expect("count query");
rows.first().and_then(|r| r.first()).copied().unwrap_or(0)
}
/// Report the first place two dumps diverge — a byte-identity failure must name
/// WHERE, not print the whole graph.
fn first_difference(a: &[String], b: &[String]) -> Option<String> {
if a.len() != b.len() {
return Some(format!(
"row COUNT differs: serial {} vs parallel {}",
a.len(),
b.len()
));
}
a.iter()
.zip(b)
.enumerate()
.find(|(_, (x, y))| x != y)
.map(|(i, (x, y))| format!("row {i} differs:\n serial = {x:?}\n parallel = {y:?}"))
}
#[tokio::test(flavor = "multi_thread")]
#[ignore = "needs a live FalkorDB (KNUT_FALKORDB_URL) or podman; run with --ignored"]
async fn doput_parallel_is_byte_identical_to_serial_and_faster() {
ensure_docker_host();
// Prefer an already-running FalkorDB; otherwise bring one up.
let (redis_url, _guard) = match std::env::var("KNUT_FALKORDB_URL") {
Ok(u) if !u.is_empty() => (u, None),
_ => {
let image = std::env::var("KNUT_FALKORDB_IMAGE")
.unwrap_or_else(|_| "localhost/knut-falkordb-valkey9:latest".into());
let (name, tag) = image.rsplit_once(':').unwrap();
let c = GenericImage::new(name, tag)
.with_exposed_port(ContainerPort::Tcp(6379))
.with_startup_timeout(Duration::from_secs(120))
.start()
.await
.expect("start falkordb");
let port = c.get_host_port_ipv4(6379).await.unwrap();
(format!("redis://127.0.0.1:{port}"), Some(c))
}
};
let batch = people(ROWS);
let concurrency = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
// A re-run must start from empty graphs.
reset_graph(&redis_url, "doput_serial").await;
reset_graph(&redis_url, "doput_parallel").await;
// ── SERIAL: put_concurrency = 1 → the pre-fix one-round-trip-at-a-time ──
let mut serial = start_server(&redis_url, "doput_serial", 1).await;
let (serial_rows, serial_t) = do_put_vertices(&mut serial, batch.clone()).await;
// ── PARALLEL: the shipped default — up to `concurrency` windows in flight ─
let mut par = start_server(&redis_url, "doput_parallel", concurrency).await;
let (par_rows, par_t) = do_put_vertices(&mut par, batch.clone()).await;
// ── the row contract ─────────────────────────────────────────────────────
assert_eq!(serial_rows, ROWS, "serial DoPut reported every row");
assert_eq!(
par_rows, serial_rows,
"the parallel DoPut reports the SAME row count as the serial one"
);
// ── BYTE-IDENTITY: the two graphs must be indistinguishable ─────────────
assert_eq!(
node_count(&redis_url, "doput_serial").await,
ROWS,
"serial graph node count"
);
assert_eq!(
node_count(&redis_url, "doput_parallel").await,
ROWS,
"parallel graph node count"
);
let serial_dump = dump_graph(&redis_url, "doput_serial").await;
let par_dump = dump_graph(&redis_url, "doput_parallel").await;
assert_eq!(
serial_dump.len(),
ROWS as usize,
"the serial dump has every node"
);
if let Some(diff) = first_difference(&serial_dump, &par_dump) {
panic!(
"the graph a PARALLEL DoPut builds must be byte-identical to the serial one \
(same _gar_id order, same properties) — concurrency changes the time, not the \
data.\n{diff}"
);
}
// ── the IN-PROCESS reference: the same loader call the DoPut now makes, ──
// driven directly, with no Flight wire in between. This is the number the
// DoPut path was said to be ~100x behind; measuring it here says how much
// of any gap the transport owns versus the insert path.
reset_graph(&redis_url, "doput_inproc").await;
let loader = graphar_falkordb::FalkorDbLoader::connect(&redis_url, "doput_inproc")
.await
.expect("connect loader");
let t = Instant::now();
let inproc_rows = loader
.insert_vertices_par(LABEL, std::slice::from_ref(&batch), concurrency)
.await
.expect("in-process bulk insert");
let inproc_t = t.elapsed();
assert_eq!(
inproc_rows as i64, ROWS,
"in-process insert reported every row"
);
assert_eq!(
node_count(&redis_url, "doput_inproc").await,
ROWS,
"in-process graph node count"
);
let inproc_dump = dump_graph(&redis_url, "doput_inproc").await;
if let Some(diff) = first_difference(&serial_dump, &inproc_dump) {
panic!(
"the IN-PROCESS bulk insert must build the byte-identical graph too — \
the DoPut path is the same loader call over a wire.\n{diff}"
);
}
// ── the number ───────────────────────────────────────────────────────────
let speedup = serial_t.as_secs_f64() / par_t.as_secs_f64().max(f64::MIN_POSITIVE);
println!(
"DOPUT rows={ROWS} concurrency={concurrency} \
serial_s={:.6} parallel_s={:.6} inprocess_s={:.6} \
serial_rows_per_s={:.0} parallel_rows_per_s={:.0} inprocess_rows_per_s={:.0} \
speedup_x={speedup:.2} doput_vs_inprocess_x={:.2}",
serial_t.as_secs_f64(),
par_t.as_secs_f64(),
inproc_t.as_secs_f64(),
ROWS as f64 / serial_t.as_secs_f64(),
ROWS as f64 / par_t.as_secs_f64(),
ROWS as f64 / inproc_t.as_secs_f64(),
par_t.as_secs_f64() / inproc_t.as_secs_f64().max(f64::MIN_POSITIVE),
);
assert!(
par_t < serial_t,
"the parallel DoPut must not be SLOWER than the serial one \
(serial {serial_t:?}, parallel {par_t:?})"
);
// Clean up so a re-run starts from an empty graph.
let client = redis::Client::open(redis_url.as_str()).unwrap();
let mut conn = client.get_multiplexed_async_connection().await.unwrap();
for g in ["doput_serial", "doput_parallel", "doput_inproc"] {
let _: redis::Value = redis::cmd("GRAPH.DELETE")
.arg(g)
.query_async(&mut conn)
.await
.unwrap();
}
}