grust-ladybug: bulk loads can skip the stored-key read-back for fresh rows

crates/grust-ladybug/src/bulk_load.rs

@@ -36,17 +36,35 @@ impl LadybugGraphStore {
         table: &str,
         nodes: &[&Node],
     ) -> Result<usize> {
-        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<usize> {
-        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::<Result<Vec<_>>>()?;
         let props = fresh
-            .values()
+            .iter()
             .map(|edge| props_to_string(&edge.props))
             .collect::<Result<Vec<_>>>()?;
         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(())
+        })
+    }
 }

crates/grust-ladybug/src/lib.rs

@@ -109,6 +109,14 @@ pub struct LadybugGraphStore {
     db: Arc<lbug::Database>,
     lock: Mutex<()>,
     schema: RwLock<Option<GraphSchema>>,
+    /// 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<std::collections::HashMap<String, std::collections::HashSet<String>>>,
 }
 
 #[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
@@ -145,13 +153,39 @@ 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()),
         })
     }
 
     pub fn in_memory() -> Result<Self> {
         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<PathBuf>) -> Result<Self> {
         Self::new(LadybugConfig {
             path: LadybugPath::Directory(path.into()),