diff --git a/Cargo.toml b/Cargo.toml index 0d6c3a99..870f37b9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,7 +39,7 @@ categories = ["database", "data-structures"] [workspace.dependencies] async-trait = "0.1.89" helix-db = "=3.0.0" -lbug = "0.20.2" +lbug = "0.20.4" lancedb = "0.30.0" redis = "1.6.0" r2d2 = "0.8.10" diff --git a/crates/grust-ladybug/Cargo.toml b/crates/grust-ladybug/Cargo.toml index 69369ab8..03a27bdd 100644 --- a/crates/grust-ladybug/Cargo.toml +++ b/crates/grust-ladybug/Cargo.toml @@ -15,6 +15,7 @@ default = ["arrow"] arrow = ["dep:arrow", "lbug/arrow", "dep:grust-arrow"] [dependencies] +tempfile = { workspace = true } grust-arrow = { version = "0.20.0", path = "../grust-arrow", default-features = false, features = ["arrow-55"], optional = true } arrow = { version = "55", optional = true } async-trait = { workspace = true } @@ -25,3 +26,7 @@ serde_json = { workspace = true } [dev-dependencies] tempfile = { workspace = true } tokio = { workspace = true } + +[[example]] +name = "copy_bench" +required-features = ["arrow"] diff --git a/crates/grust-ladybug/examples/copy_bench.rs b/crates/grust-ladybug/examples/copy_bench.rs new file mode 100644 index 00000000..7dcd354f --- /dev/null +++ b/crates/grust-ladybug/examples/copy_bench.rs @@ -0,0 +1,294 @@ +//! Times the ways the `lbug` crate can bulk-load a graph, on one synthetic +//! graph, so the adapter's load path can be chosen from a measurement: +//! +//! A. what `grust-ladybug` does today: register Arrow batches as a scratch +//! table and `COPY t FROM (MATCH (x:scratch) RETURN …)`; +//! B. `COPY t FROM scratch` over the registered table directly, the shape +//! of Python's `COPY t FROM $df` (may be rejected: reported, not fatal); +//! C. `COPY t FROM 'file.csv'`, the engine's native file path; +//! D. C again with INT64 primary keys and no properties, the CSR-parquet +//! shape the maintainer loads cit-Patents from in seconds; +//! E. the CSR registration (`create_arrow_rel_table_csr`) with INT64 keys, +//! copied through MATCH and directly. +//! +//! VARIANTS=ABCDE cargo run --release -p grust-ladybug --example copy_bench -- [nodes] [edges] +//! +//! Prints one line per variant: rows, seconds, rows per second, or the +//! engine's error text. +use std::{fs::File, io::Write, sync::Arc, time::Instant}; + +use arrow::{ + array::{ArrayRef, Int64Array, StringArray, UInt64Array}, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, +}; + +fn arg(i: usize, default: usize) -> usize { + std::env::args() + .nth(i) + .and_then(|s| s.parse().ok()) + .unwrap_or(default) +} + +fn strings(name: &str, values: &[String]) -> (Field, ArrayRef) { + ( + Field::new(name, DataType::Utf8, false), + Arc::new(StringArray::from_iter_values( + values.iter().map(String::as_str), + )), + ) +} + +fn ints(name: &str, values: &[i64]) -> (Field, ArrayRef) { + ( + Field::new(name, DataType::Int64, false), + Arc::new(Int64Array::from(values.to_vec())), + ) +} + +/// CSR node offsets: the engine wants UINT64 row offsets, not key values. +fn uints(name: &str, values: &[i64]) -> (Field, ArrayRef) { + ( + Field::new(name, DataType::UInt64, false), + Arc::new(UInt64Array::from( + values.iter().map(|v| *v as u64).collect::>(), + )), + ) +} + +fn batch(cols: Vec<(Field, ArrayRef)>) -> RecordBatch { + let (fields, arrays): (Vec<_>, Vec<_>) = cols.into_iter().unzip(); + RecordBatch::try_new(Arc::new(Schema::new(fields)), arrays).expect("batch") +} + +fn q(conn: &lbug::Connection<'_>, sql: &str) -> Result<(), String> { + conn.query(sql).map(drop).map_err(|e| e.to_string()) +} + +/// Runs `f`, prints the variant's rate or the error it returned. +fn time(label: &str, rows: usize, f: impl FnOnce() -> Result<(), String>) { + let wanted = std::env::var("VARIANTS").unwrap_or_else(|_| "ABCDE".into()); + if !wanted.contains(&label[..1]) { + println!("{label:58} skipped (VARIANTS={wanted})"); + return; + } + let t = Instant::now(); + match f() { + Ok(()) => { + let s = t.elapsed().as_secs_f64(); + println!( + "{label:58} {rows:>9} rows {s:8.2} s {:>10.0} rows/s", + rows as f64 / s + ); + } + Err(e) => println!("{label:58} ERROR {}", e.lines().next().unwrap_or("")), + } +} + +fn main() { + let n = arg(1, 1_000_000); + let m = arg(2, 5_000_000); + let dir = tempfile::tempdir().expect("tempdir"); + let db = lbug::Database::new(dir.path().join("db"), lbug::SystemConfig::default()).expect("db"); + let conn = lbug::Connection::new(&db).expect("conn"); + println!( + "lbug {} · {n} nodes, {m} edges · db in {}", + env!("CARGO_PKG_VERSION"), + dir.path().display() + ); + + // ---- the data: string ids like the adapter's, and their integer twins + let ids: Vec = (0..n).map(|i| format!("n{i}")).collect(); + let empty: Vec = vec!["{}".to_string(); n]; + let from_i: Vec = (0..m).map(|i| ((i * 7919) % n) as i64).collect(); + let to_i: Vec = (0..m).map(|i| ((i * 104_729 + 13) % n) as i64).collect(); + let from_s: Vec = from_i.iter().map(|i| format!("n{i}")).collect(); + let to_s: Vec = to_i.iter().map(|i| format!("n{i}")).collect(); + let eids: Vec = (0..m).map(|i| format!("e{i}")).collect(); + let eprops: Vec = vec!["{}".to_string(); m]; + + // ---- string-keyed tables, the adapter's schema + for t in ["N", "N2", "N3"] { + q( + &conn, + &format!("CREATE NODE TABLE {t}(id STRING, props STRING, PRIMARY KEY(id));"), + ) + .unwrap(); + } + for (e, t) in [("E", "N"), ("E2", "N2"), ("E3", "N3")] { + q( + &conn, + &format!("CREATE REL TABLE {e}(FROM {t} TO {t}, id STRING, props STRING);"), + ) + .unwrap(); + } + + // A. registered Arrow + COPY FROM (MATCH …): the adapter today + let node_batch = batch(vec![strings("id", &ids), strings("props", &empty)]); + time("A nodes: arrow table, COPY FROM (MATCH …)", n, || { + conn.create_arrow_table("s_nodes", &[node_batch.clone()]) + .map_err(|e| e.to_string())?; + let r = q( + &conn, + "COPY N FROM (MATCH (x:s_nodes) RETURN x.id, x.props);", + ); + conn.drop_arrow_table("s_nodes") + .map_err(|e| e.to_string())?; + r + }); + let rel_batch = batch(vec![ + strings("from", &from_s), + strings("to", &to_s), + strings("id", &eids), + strings("props", &eprops), + ]); + time("A edges: arrow rel table, COPY FROM (MATCH …)", m, || { + conn.create_arrow_rel_table("s_rels", &[rel_batch.clone()], "N", "N") + .map_err(|e| e.to_string())?; + let r = q( + &conn, + "COPY E FROM (MATCH (a:N)-[r:s_rels]->(b:N) RETURN a.id, b.id, r.id, r.props);", + ); + conn.drop_arrow_table("s_rels").map_err(|e| e.to_string())?; + r + }); + + // B. COPY t FROM directly + time("B nodes: arrow table, COPY FROM s_nodes", n, || { + conn.create_arrow_table("s_nodes", &[node_batch.clone()]) + .map_err(|e| e.to_string())?; + let r = q(&conn, "COPY N2 FROM s_nodes;"); + conn.drop_arrow_table("s_nodes") + .map_err(|e| e.to_string())?; + r + }); + time("B edges: arrow rel table, COPY FROM s_rels", m, || { + conn.create_arrow_rel_table("s_rels", &[rel_batch.clone()], "N2", "N2") + .map_err(|e| e.to_string())?; + let r = q(&conn, "COPY E2 FROM s_rels;"); + conn.drop_arrow_table("s_rels").map_err(|e| e.to_string())?; + r + }); + + // C. CSV files, string keys + let nodes_csv = dir.path().join("nodes.csv"); + let rels_csv = dir.path().join("rels.csv"); + { + let mut f = File::create(&nodes_csv).unwrap(); + for id in &ids { + writeln!(f, "{id},{{}}").unwrap(); + } + let mut f = File::create(&rels_csv).unwrap(); + for i in 0..m { + writeln!(f, "{},{},{},{{}}", from_s[i], to_s[i], eids[i]).unwrap(); + } + } + time("C nodes: COPY FROM 'nodes.csv' (STRING keys)", n, || { + q(&conn, &format!("COPY N3 FROM '{}';", nodes_csv.display())) + }); + time("C edges: COPY FROM 'rels.csv' (STRING keys)", m, || { + q(&conn, &format!("COPY E3 FROM '{}';", rels_csv.display())) + }); + + // D. CSV files, INT64 keys, no properties: the CSR-parquet shape + q(&conn, "CREATE NODE TABLE NI(id INT64, PRIMARY KEY(id));").unwrap(); + for e in ["EI", "EI2", "EI3"] { + q(&conn, &format!("CREATE REL TABLE {e}(FROM NI TO NI);")).unwrap(); + } + let nodes_i_csv = dir.path().join("nodes_i.csv"); + let rels_i_csv = dir.path().join("rels_i.csv"); + { + let mut f = File::create(&nodes_i_csv).unwrap(); + for i in 0..n { + writeln!(f, "{i}").unwrap(); + } + let mut f = File::create(&rels_i_csv).unwrap(); + for i in 0..m { + writeln!(f, "{},{}", from_i[i], to_i[i]).unwrap(); + } + } + time("D nodes: COPY FROM 'nodes_i.csv' (INT64 keys)", n, || { + q(&conn, &format!("COPY NI FROM '{}';", nodes_i_csv.display())) + }); + time( + "D edges: COPY FROM 'rels_i.csv' (INT64 keys, no props)", + m, + || q(&conn, &format!("COPY EI FROM '{}';", rels_i_csv.display())), + ); + + // E. CSR registration with INT64 keys: indices sorted by source, indptr offsets + let mut order: Vec = (0..m).collect(); + order.sort_by_key(|&i| from_i[i]); + let indices: Vec = order.iter().map(|&i| to_i[i]).collect(); + let mut indptr: Vec = vec![0; n + 1]; + for &i in &order { + indptr[from_i[i] as usize + 1] += 1; + } + for i in 0..n { + indptr[i + 1] += indptr[i]; + } + let indices_batch = batch(vec![uints("to", &indices)]); + let indptr_batch = batch(vec![uints("offset", &indptr)]); + time( + "E edges: CSR arrow rel table, COPY FROM (MATCH …)", + m, + || { + conn.create_arrow_rel_table_csr( + "s_csr", + &[indices_batch.clone()], + &[indptr_batch.clone()], + "NI", + "NI", + "to", + ) + .map_err(|e| e.to_string())?; + let r = q( + &conn, + "COPY EI2 FROM (MATCH (a:NI)-[r:s_csr]->(b:NI) RETURN a.id, b.id);", + ); + conn.drop_arrow_table("s_csr").map_err(|e| e.to_string())?; + r + }, + ); + time("E edges: CSR arrow rel table, COPY FROM s_csr", m, || { + conn.create_arrow_rel_table_csr( + "s_csr", + &[indices_batch.clone()], + &[indptr_batch.clone()], + "NI", + "NI", + "to", + ) + .map_err(|e| e.to_string())?; + let r = q(&conn, "COPY EI3 FROM s_csr;"); + conn.drop_arrow_table("s_csr").map_err(|e| e.to_string())?; + r + }); + + // sanity: what landed + for (t, want) in [ + ("N", n), + ("N2", n), + ("N3", n), + ("NI", n), + ("E", m), + ("E2", m), + ("E3", m), + ("EI", m), + ("EI2", m), + ("EI3", m), + ] { + let pat = if t.starts_with('E') { + format!("MATCH ()-[r:{t}]->() RETURN count(r);") + } else { + format!("MATCH (x:{t}) RETURN count(x);") + }; + let got = conn + .query(&pat) + .ok() + .and_then(|mut r| r.next()) + .map(|row| format!("{:?}", row[0])) + .unwrap_or_else(|| "?".into()); + println!("count {t:4} = {got} (want {want})"); + } +} diff --git a/crates/grust-ladybug/src/bulk_load.rs b/crates/grust-ladybug/src/bulk_load.rs index 19df632b..8c133d0d 100644 --- a/crates/grust-ladybug/src/bulk_load.rs +++ b/crates/grust-ladybug/src/bulk_load.rs @@ -1,32 +1,32 @@ -//! Bulk loading through registered Arrow tables. +//! Bulk loading through the engine's file `COPY`. //! //! One `MERGE` per node or edge costs the engine a full statement each, //! measured at about 50 ms of CPU for a relationship `MERGE … SET` on a 4-vCPU -//! host, so a 200,000-edge slice took five hours. The `lbug` crate can register -//! Arrow record batches as tables; `COPY … FROM (MATCH …)` over such a table -//! loads the same 100,000 nodes in a quarter of a second and 112,000 edges in -//! eight. +//! host, so a 200,000-edge slice took five hours. Registering the rows as an +//! Arrow table and running `COPY … FROM (MATCH …)` over it fixed the nodes +//! (a million rows a second) but not the edges: that probe is superlinear in +//! the engine, 4,100 edges a second at 500,000 and unfinished after two hours +//! at five million, which is why cit-Patents never loaded. The engine's own +//! `COPY … FROM 'file.csv'` takes the same rows, STRING keys and JSON props +//! included, at 1.3 million edges a second (five million in 3.9 s; 2.2 million +//! nodes a second), so every bulk write goes through a temporary CSV file. +//! The `lbug` crate has no equivalent of Python's `COPY … FROM $df`, and the +//! binder rejects `COPY … FROM `. //! //! `put_graph` keeps its upsert semantics: rows whose node id, or whose //! `(from, to)` pair for a relationship table, already exists go through the //! per-row `MERGE` path exactly as before; every other row is copied. Within //! one load the last row for an id or pair wins, which is what a sequence of -//! `MERGE`s would leave behind. +//! `MERGE`s would leave behind. An empty target table skips the read-back of +//! existing keys, so a fresh load costs nothing beyond the copy. -use std::{ - collections::{BTreeMap, HashSet}, - sync::atomic::AtomicU64, -}; +use std::collections::{BTreeMap, HashSet}; use arrow::record_batch::RecordBatch; use grust_core::prelude::*; use super::{LadybugGraphStore, ladybug_error, props_to_string}; -/// A registered Arrow table's name is unique per process so two loads on the -/// same connection never collide. -static SCRATCH_TABLES: AtomicU64 = AtomicU64::new(0); - impl LadybugGraphStore { /// Write `nodes` into `table`, copying rows whose id is new and merging /// the rest. Returns how many rows were written. @@ -36,7 +36,11 @@ impl LadybugGraphStore { table: &str, nodes: &[&Node], ) -> Result { - let existing = self.existing_node_ids(conn, table)?; + let existing = if Self::count(conn, &format!("MATCH (n:{table}) RETURN count(n);"))? == 0 { + HashSet::new() + } else { + self.existing_node_ids(conn, table)? + }; let mut fresh: BTreeMap<&str, &Node> = BTreeMap::new(); for node in nodes { if existing.contains(node.id.as_str()) { @@ -57,7 +61,7 @@ impl LadybugGraphStore { ("id", ids.clone()), ("props", props.iter().map(String::as_str).collect()), ])?; - self.copy_from_arrow_nodes(conn, table, batch)?; + self.copy_from_csv(conn, table, &batch)?; self.bulk_write_node_index(conn, table, &fresh)?; Ok(nodes.len()) } @@ -72,7 +76,15 @@ impl LadybugGraphStore { to_table: &str, edges: &[&Edge], ) -> Result { - let existing = self.existing_edge_pairs(conn, rel_table, from_table, to_table)?; + let existing = if Self::count( + conn, + &format!("MATCH ()-[r:{rel_table}]->() RETURN count(r);"), + )? == 0 + { + HashSet::new() + } else { + self.existing_edge_pairs(conn, rel_table, from_table, to_table)? + }; let mut fresh: BTreeMap<(&str, &str), &Edge> = BTreeMap::new(); for edge in edges { let pair = (edge.from.as_str(), edge.to.as_str()); @@ -99,35 +111,59 @@ impl LadybugGraphStore { ("id", ids.iter().map(String::as_str).collect()), ("props", props.iter().map(String::as_str).collect()), ])?; - let scratch = scratch_table_name(&self.config.table_prefix); - conn.create_arrow_rel_table(&scratch, &[batch], from_table, to_table) - .map_err(ladybug_error)?; - let copy = Self::exec( - conn, - &format!( - "COPY {rel_table} FROM (MATCH (a:{from_table})-[r:{scratch}]->(b:{to_table}) RETURN a.id, b.id, r.id, r.props);" - ), - ); - conn.drop_arrow_table(&scratch).map_err(ladybug_error)?; - copy?; + // Positional: FROM, TO, then the rel table's own columns. + self.copy_from_csv(conn, rel_table, &batch)?; Ok(edges.len()) } - fn copy_from_arrow_nodes( + /// Writes `batch` to a temporary CSV and runs the engine's `COPY` over it, + /// columns positional. The file lives only for the statement. + fn copy_from_csv( &self, conn: &lbug::Connection<'_>, table: &str, - batch: RecordBatch, + batch: &RecordBatch, ) -> Result<()> { - let scratch = scratch_table_name(&self.config.table_prefix); - conn.create_arrow_table(&scratch, &[batch]) - .map_err(ladybug_error)?; - let copy = Self::exec( + let io = |err: std::io::Error| { + GrustError::Backend(format!("Ladybug bulk copy file error: {err}")) + }; + let dir = tempfile::Builder::new() + .prefix("grust-ladybug-copy-") + .tempdir() + .map_err(io)?; + let path = dir.path().join("rows.csv"); + { + let file = std::fs::File::create(&path).map_err(io)?; + let mut writer = arrow::csv::WriterBuilder::new() + .with_header(false) + .build(file); + writer.write(batch).map_err(|err| { + GrustError::Serialization(format!("Ladybug bulk copy CSV error: {err}")) + })?; + } + let path = path.display().to_string(); + if path.contains('\'') { + return Err(GrustError::Backend(format!( + "Ladybug bulk copy: temp path has a quote: {path}" + ))); + } + Self::exec( conn, - &format!("COPY {table} FROM (MATCH (n:{scratch}) RETURN n.id, n.props);"), - ); - conn.drop_arrow_table(&scratch).map_err(ladybug_error)?; - copy + &format!( + "COPY {table} FROM '{path}' (HEADER=false, DELIM=',', QUOTE='\"', ESCAPE='\"');" + ), + ) + } + + /// One integer from a `RETURN count(…)` query. + fn count(conn: &lbug::Connection<'_>, query: &str) -> Result { + let mut rows = conn.query(query).map_err(ladybug_error)?; + match rows.next().as_deref() { + Some([lbug::Value::Int64(n), ..]) => Ok(*n), + other => Err(GrustError::Serialization(format!( + "Ladybug count returned {other:?}" + ))), + } } /// The metadata index gets one row per new node, copied the same way. @@ -145,17 +181,7 @@ impl LadybugGraphStore { ("label", labels), ("table_name", vec![table; fresh.len()]), ])?; - let scratch = scratch_table_name(&self.config.table_prefix); - conn.create_arrow_table(&scratch, &[batch]) - .map_err(ladybug_error)?; - let copy = Self::exec( - conn, - &format!( - "COPY {node_index} FROM (MATCH (n:{scratch}) RETURN n.id, n.kind, n.label, n.table_name);" - ), - ); - conn.drop_arrow_table(&scratch).map_err(ladybug_error)?; - copy + self.copy_from_csv(conn, &node_index, &batch) } fn existing_node_ids( @@ -192,13 +218,87 @@ impl LadybugGraphStore { } } -fn scratch_table_name(prefix: &str) -> String { - let n = SCRATCH_TABLES.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - format!("{prefix}_arrow_in_{}_{n}", std::process::id()) -} - /// A record batch of non-null UTF-8 columns. fn string_batch(columns: &[(&str, Vec<&str>)]) -> Result { grust_arrow::v55::string_batch(columns) .map_err(|err| GrustError::Serialization(format!("Ladybug Arrow batch error: {err}"))) } + +#[cfg(test)] +mod tests { + use super::super::LadybugGraphStore; + use grust_core::prelude::*; + use serde_json::json; + + fn props(pairs: &[(&str, serde_json::Value)]) -> Props { + pairs + .iter() + .map(|(k, v)| (k.to_string(), Value::from(v.clone()))) + .collect() + } + + /// Values that break naive CSV: commas, double quotes, newlines, a + /// backslash, a quote-only string, an empty string, unicode. + #[test] + fn bulk_copy_round_trips_hostile_props_and_then_upserts() -> Result<()> { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let store = LadybugGraphStore::in_memory()?; + let nasty = props(&[ + ("comma", json!("a,b")), + ("quote", json!("say \"hi\"")), + ("newline", json!("line1\nline2")), + ("backslash", json!("c:\\path")), + ("only_quote", json!("\"")), + ("empty", json!("")), + ("unicode", json!("naïve 日本 🚀")), + ("nested", json!({"k": ["x,y", "\"z\""]})), + ]); + let nodes: Vec = (0..3) + .map(|i| Node::new("Thing", format!("n{i}"), nasty.clone())) + .collect(); + let edges = vec![ + Edge::new("REL", "n0", "n1", nasty.clone()), + Edge::new("REL", "n1", "n2", props(&[("plain", json!(1))])), + ]; + let report = store.put_graph(&Graph::new(nodes, edges)).await?; + assert_eq!((report.nodes, report.edges), (3, 2)); + // Reads add the id into the props; compare the stored values only. + let mut got = store.get_node(&NodeId::new("n1")).await?.expect("n1"); + got.props.remove("id"); + assert_eq!(got.props, nasty); + let out = store + .get_edges(EdgeQuery { + from: Some(NodeId::new("n0")), + ..Default::default() + }) + .await?; + assert_eq!(out.len(), 1); + assert_eq!(out[0].props, nasty); + + // A second load of the same ids takes the merge path (the tables + // are no longer empty) and must not duplicate anything. + // A put carries its edges' endpoints, as every caller's batches do. + let again = Graph::new( + vec![ + Node::new("Thing", "n0", nasty.clone()), + Node::new("Thing", "n1", props(&[("v", json!(2))])), + ], + vec![Edge::new("REL", "n0", "n1", props(&[("v", json!(2))]))], + ); + store.put_graph(&again).await?; + let mut got = store.get_node(&NodeId::new("n1")).await?.expect("n1"); + got.props.remove("id"); + assert_eq!(got.props, props(&[("v", json!(2))])); + let out = store + .get_edges(EdgeQuery { + from: Some(NodeId::new("n0")), + ..Default::default() + }) + .await?; + assert_eq!(out.len(), 1, "upsert must not add a parallel edge"); + assert_eq!(out[0].props, props(&[("v", json!(2))])); + Ok(()) + }) + } +}