diff --git a/crates/grust-ladybug/src/bulk_load.rs b/crates/grust-ladybug/src/bulk_load.rs index 8c133d0d..e1680c0b 100644 --- a/crates/grust-ladybug/src/bulk_load.rs +++ b/crates/grust-ladybug/src/bulk_load.rs @@ -36,17 +36,35 @@ impl LadybugGraphStore { table: &str, nodes: &[&Node], ) -> Result { - let existing = if Self::count(conn, &format!("MATCH (n:{table}) RETURN count(n);"))? == 0 { + let trusts_fresh = self.bulk_load_trusts_fresh_rows(); + let existing = if trusts_fresh + || 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()) { - self.write_node_locked(conn, node, table)?; - } else { - fresh.insert(node.id.as_str(), node); + if trusts_fresh { + // Already-copied ids (this store, this mode) are carried endpoints: + // skip them; everything else is new by the caller's word. + let copied = self + .fresh_copied_node_ids + .lock() + .map_err(|_| GrustError::Backend("Ladybug fresh-copy set poisoned".into()))?; + let seen = copied.get(table); + for node in nodes { + if !seen.is_some_and(|s| s.contains(node.id.as_str())) { + fresh.insert(node.id.as_str(), node); + } + } + } else { + for node in nodes { + if existing.contains(node.id.as_str()) { + self.write_node_locked(conn, node, table)?; + } else { + fresh.insert(node.id.as_str(), node); + } } } if fresh.is_empty() { @@ -63,6 +81,16 @@ impl LadybugGraphStore { ])?; self.copy_from_csv(conn, table, &batch)?; self.bulk_write_node_index(conn, table, &fresh)?; + if trusts_fresh { + let mut copied = self + .fresh_copied_node_ids + .lock() + .map_err(|_| GrustError::Backend("Ladybug fresh-copy set poisoned".into()))?; + copied + .entry(table.to_string()) + .or_default() + .extend(fresh.keys().map(|id| id.to_string())); + } Ok(nodes.len()) } @@ -76,38 +104,49 @@ impl LadybugGraphStore { to_table: &str, edges: &[&Edge], ) -> Result { - let existing = if Self::count( - conn, - &format!("MATCH ()-[r:{rel_table}]->() RETURN count(r);"), - )? == 0 + let trusts_fresh = self.bulk_load_trusts_fresh_rows(); + let existing = if trusts_fresh + || 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()); - if existing.contains(&(pair.0.to_string(), pair.1.to_string())) { - self.write_edge_locked(conn, edge, rel_table, from_table, to_table)?; - } else { - fresh.insert(pair, edge); + // Rows to copy, in input order. With the read-back on, one row per + // (from, to) pair, the last one winning, as a sequence of MERGEs + // would leave; with the caller vouching for fresh rows, every row, + // so parallel edges stay parallel edges. + let fresh: Vec<&Edge> = if trusts_fresh { + edges.to_vec() + } else { + let mut by_pair: BTreeMap<(&str, &str), &Edge> = BTreeMap::new(); + for edge in edges { + let pair = (edge.from.as_str(), edge.to.as_str()); + if existing.contains(&(pair.0.to_string(), pair.1.to_string())) { + self.write_edge_locked(conn, edge, rel_table, from_table, to_table)?; + } else { + by_pair.insert(pair, edge); + } } - } + by_pair.into_values().collect() + }; if fresh.is_empty() { return Ok(edges.len()); } let ids = fresh - .values() + .iter() .map(|edge| checked_edge_key(edge)) .collect::>>()?; let props = fresh - .values() + .iter() .map(|edge| props_to_string(&edge.props)) .collect::>>()?; let batch = string_batch(&[ - ("from", fresh.keys().map(|(from, _)| *from).collect()), - ("to", fresh.keys().map(|(_, to)| *to).collect()), + ("from", fresh.iter().map(|e| e.from.as_str()).collect()), + ("to", fresh.iter().map(|e| e.to.as_str()).collect()), ("id", ids.iter().map(String::as_str).collect()), ("props", props.iter().map(String::as_str).collect()), ])?; @@ -301,4 +340,71 @@ mod tests { Ok(()) }) } + + /// The chunked-load shape: every edge batch carries its endpoint nodes + /// again. In fresh-rows mode the carried nodes must be skipped, not fail + /// the COPY on a duplicate primary key, and every edge must land. + #[test] + fn fresh_rows_mode_skips_carried_endpoints_across_chunks() -> Result<()> { + let rt = tokio::runtime::Runtime::new().unwrap(); + rt.block_on(async { + let store = LadybugGraphStore::in_memory()?; + store.set_bulk_load_trusts_fresh_rows(true); + let node = |i: usize| Node::new("Thing", format!("n{i}"), props(&[])); + // nodes first, as the harness does, then edge chunks carrying endpoints + store + .put_graph(&Graph::new((0..4).map(node).collect(), vec![])) + .await?; + for (a, b) in [(0usize, 1usize), (1, 2), (2, 3), (0, 3)] { + let chunk = Graph::new( + vec![node(a), node(b)], + vec![Edge::new( + "REL", + format!("n{a}"), + format!("n{b}"), + props(&[]), + )], + ); + store.put_graph(&chunk).await?; + } + let mut total = 0; + for i in 0..4 { + total += store + .get_edges(EdgeQuery { + from: Some(NodeId::new(format!("n{i}"))), + ..Default::default() + }) + .await? + .len(); + } + assert_eq!(total, 4, "every edge landed exactly once"); + // A multigraph chunk: two parallel edges on one pair, distinct ids, + // must both land and both read back. + let multi = Graph::new( + vec![node(0), node(1)], + vec![ + Edge::new("REL", "n0", "n1", props(&[("t", json!(1))])).with_id("p1"), + Edge::new("REL", "n0", "n1", props(&[("t", json!(2))])).with_id("p2"), + ], + ); + store.put_graph(&multi).await?; + let out = store + .get_edges(EdgeQuery { + from: Some(NodeId::new("n0")), + to: Some(NodeId::new("n1")), + ..Default::default() + }) + .await?; + assert_eq!(out.len(), 3, "the original edge plus two parallel edges"); + for i in 0..4 { + assert!( + store + .get_node(&NodeId::new(format!("n{i}"))) + .await? + .is_some() + ); + } + Ok(()) + }) + } } diff --git a/crates/grust-ladybug/src/lib.rs b/crates/grust-ladybug/src/lib.rs index 6587fe3f..420485b3 100644 --- a/crates/grust-ladybug/src/lib.rs +++ b/crates/grust-ladybug/src/lib.rs @@ -109,6 +109,14 @@ pub struct LadybugGraphStore { db: Arc, lock: Mutex<()>, schema: RwLock>, + /// See [`LadybugGraphStore::set_bulk_load_trusts_fresh_rows`]. + bulk_load_trusts_fresh_rows: std::sync::atomic::AtomicBool, + /// Node ids this store copied under the fresh-rows mode, per table: a + /// later batch that carries an already-copied node (an edge batch + /// carrying its endpoints) skips it instead of failing the COPY on a + /// duplicate primary key. + fresh_copied_node_ids: + Mutex>>, } #[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)] @@ -145,6 +153,8 @@ impl LadybugGraphStore { db: Arc::new(db), lock: Mutex::new(()), schema: RwLock::new(None), + bulk_load_trusts_fresh_rows: std::sync::atomic::AtomicBool::new(false), + fresh_copied_node_ids: Mutex::new(std::collections::HashMap::new()), }) } @@ -152,6 +162,30 @@ impl LadybugGraphStore { Self::new(LadybugConfig::default()) } + /// Bulk loads (`put_graph`) normally read back every existing node id and + /// `(from, to)` pair of a target table before copying, so that rows which + /// already exist take the per-row upsert path. That read-back is a full + /// scan per call, and a caller that feeds one graph as many chunks pays it + /// on every chunk after the first: quadratic in the chunk count, hours at + /// a hundred million edges. A caller that knows its chunks carry no key the + /// store already holds — a fresh graph loaded once, in disjoint pieces — + /// can turn the read-back off. Every relationship row is then copied (a + /// repeated pair becomes a parallel edge). Node ids the store copied in + /// this mode are remembered in memory, per table, so a later batch that + /// carries an already-copied node — an edge batch carrying its endpoints — + /// skips it, props unchanged; a node id the store held before this mode + /// was entered is the caller's error and fails the COPY. Off by default; + /// single-row writes and non-bulk paths are unaffected. + pub fn set_bulk_load_trusts_fresh_rows(&self, trusts: bool) { + self.bulk_load_trusts_fresh_rows + .store(trusts, std::sync::atomic::Ordering::Relaxed); + } + + pub(crate) fn bulk_load_trusts_fresh_rows(&self) -> bool { + self.bulk_load_trusts_fresh_rows + .load(std::sync::atomic::Ordering::Relaxed) + } + pub fn open(path: impl Into) -> Result { Self::new(LadybugConfig { path: LadybugPath::Directory(path.into()),