Merge db_v4 into db_v4_gql_avoid_collect

Author: ljeub-pometry

.github/workflows/stress-test.yml

@@ -21,10 +21,8 @@ jobs:
         uses: Swatinem/rust-cache@v2
         with:
           cache-all-crates: true
-      - name: Install maturin
-        run: pip install maturin==1.8.3
-      - name: Build raphtory
-        run: make install-python
+      - uses: ./.github/actions/setup_rust
+        name: Setup Rust
       - name: Set up pnpm
         uses: pnpm/action-setup@v4
         with:
@@ -33,9 +31,12 @@ jobs:
         with:
           k6-version: '1.0.0'
       - name: Run stress test
+        env:
+          RUST_BACKTRACE: 1
         run: |
+          cargo build --package raphtory-graphql --bin raphtory-graphql --profile=build-fast
+          ./target/build-fast/raphtory-graphql server --work-dir graphs &
           cd graphql-bench
-          raphtory server &
           make stress-test
       - name: Upload k6 report
         if: always()

Cargo.lock

@@ -78,7 +78,7 @@ reqwest = { version = "0.12.8", default-features = false, features = [
 boxcar = "0.2.14"
 iter-enum = { version = "1.2.0", features = ["rayon"] }
 serde = { version = "1.0.197", features = ["derive", "rc"] }
-serde_json = "1.0.114"
+serde_json = { version = "1.0.114", features = ["float_roundtrip"] }
 pyo3 = { version = "0.27.2", features = ["multiple-pymethods", "chrono"] }
 pyo3-build-config = "0.27.2"
 pyo3-arrow = "0.15.0"

Cargo.toml

@@ -84,7 +84,7 @@ build-python:
 	cd python && maturin develop -r --extras=dev
 
 debug-python:
-	cd python && maturin develop --profile=debug --extras=dev
+	cd python && maturin develop --profile=dev --extras=dev
 
 # Testing
 python-test:

Makefile

@@ -187,7 +187,7 @@ where
         // VIDs in the resolver may not be initialised yet, need to double-check the node actually exists!
         let nodes = self.storage().nodes();
         let (page_id, pos) = nodes.resolve_pos(vid);
-        let node_page = nodes.segments().get(page_id)?;
+        let node_page = nodes.get_segment(page_id)?;
 
         if pos.0 < node_page.num_nodes() {
             Some(vid)
@@ -357,16 +357,16 @@ where
         self.graph
     }
 
-    pub fn resize_chunks_to_vid(&mut self, vid: VID) {
-        let (chunks_needed, _) = self.graph.storage.nodes().resolve_pos(vid);
-        self.graph.storage().nodes().grow(chunks_needed + 1);
+    pub fn resize_segments_to_vid(&mut self, vid: VID) {
+        let (segment_id, _) = self.graph.storage.nodes().resolve_pos(vid);
+        self.graph.storage().nodes().grow(segment_id + 1);
         std::mem::take(&mut self.nodes);
         self.nodes = self.graph.storage.nodes().write_locked();
     }
 
-    pub fn resize_chunks_to_eid(&mut self, eid: EID) {
-        let (chunks_needed, _) = self.graph.storage.edges().resolve_pos(eid);
-        self.graph.storage().edges().grow(chunks_needed + 1);
+    pub fn resize_segments_to_eid(&mut self, eid: EID) {
+        let (segment_id, _) = self.graph.storage.edges().resolve_pos(eid);
+        self.graph.storage().edges().grow(segment_id + 1);
         std::mem::take(&mut self.edges);
         self.edges = self.graph.storage.edges().write_locked();
     }

db4-graph/src/lib.rs

@@ -16,8 +16,7 @@ use raphtory_api::core::{
 use storage::{
     api::{edges::EdgeSegmentOps, graph_props::GraphPropSegmentOps, nodes::NodeSegmentOps},
     error::StorageError,
-    pages::resolve_pos,
-    persist::{config::ConfigOps, strategy::PersistenceStrategy},
+    persist::strategy::PersistenceStrategy,
     resolver::GIDResolverOps,
     wal::{GraphReplay, TransactionID, LSN},
     ES, GS, NS,
@@ -33,7 +32,7 @@ where
     fn replay_add_edge(
         &mut self,
         lsn: LSN,
-        transaction_id: TransactionID,
+        _transaction_id: TransactionID,
         t: EventTime,
         src_name: Option<GID>,
         src_id: VID,
@@ -44,14 +43,7 @@ where
         layer_id: usize,
         props: Vec<(String, usize, Prop)>,
     ) -> Result<(), StorageError> {
-        let node_max_page_len = self.graph().extension().config().max_node_page_len();
-        let edge_max_page_len = self.graph().extension().config().max_edge_page_len();
-
-        // 1. Insert prop ids into edge meta.
-        // No need to validate props again since they are already validated before
-        // being logged to the WAL.
-
-        // 2. Insert node ids into resolver.
+        // Insert node ids into resolver.
         if let Some(src_name) = src_name.as_ref() {
             self.graph()
                 .logical_to_physical
@@ -64,16 +56,16 @@ where
                 .set(dst_name.as_ref(), dst_id)?;
         }
 
-        // 4. Grab src writer and add edge data.
-        let (src_segment_id, src_pos) = resolve_pos(src_id, node_max_page_len);
-        let resize_vid = VID::from(src_id.index() + 1);
-        self.resize_chunks_to_vid(resize_vid); // Create enough segments.
+        // Grab src writer and add edge data.
+        let (src_segment_id, src_pos) = self.graph().storage().nodes().resolve_pos(src_id);
+        self.resize_segments_to_vid(src_id); // Create enough segments.
 
         let segment = self
             .graph()
             .storage()
             .nodes()
             .get_or_create_segment(src_segment_id);
+
         let immut_lsn = segment.immut_lsn();
 
         // Replay this entry only if it doesn't exist in immut.
@@ -112,16 +104,16 @@ where
             drop(src_writer);
         }
 
-        // 5. Grab dst writer and add edge data.
-        let (dst_segment_id, dst_pos) = resolve_pos(dst_id, node_max_page_len);
-        let resize_vid = VID::from(dst_id.index() + 1);
-        self.resize_chunks_to_vid(resize_vid);
+        // Grab dst writer and add edge data.
+        let (dst_segment_id, dst_pos) = self.graph().storage().nodes().resolve_pos(dst_id);
+        self.resize_segments_to_vid(dst_id);
 
         let segment = self
             .graph()
             .storage()
             .nodes()
             .get_or_create_segment(dst_segment_id);
+
         let immut_lsn = segment.immut_lsn();
 
         // Replay this entry only if it doesn't exist in immut.
@@ -157,24 +149,26 @@ where
             drop(dst_writer);
         }
 
-        // 6. Grab edge writer and add temporal props & metadata.
-        let (edge_segment_id, edge_pos) = resolve_pos(eid, edge_max_page_len);
-        let resize_eid = EID::from(eid.index() + 1);
-        self.resize_chunks_to_eid(resize_eid);
+        // Grab edge writer and add temporal props & metadata.
+        let (edge_segment_id, edge_pos) = self.graph().storage().edges().resolve_pos(eid);
+        self.resize_segments_to_eid(eid);
 
         let segment = self
             .graph()
             .storage()
             .edges()
             .get_or_create_segment(edge_segment_id);
+
         let immut_lsn = segment.immut_lsn();
 
         // Replay this entry only if it doesn't exist in immut.
         if immut_lsn < lsn {
             let edge_meta = self.graph().edge_meta();
 
+            // Insert prop ids into edge meta.
             for (prop_name, prop_id, prop_value) in &props {
                 let prop_mapper = edge_meta.temporal_prop_mapper();
+
                 match prop_mapper.get_dtype(*prop_id) {
                     None => {
                         prop_mapper.set_id_and_dtype(
@@ -194,7 +188,7 @@ where
                 }
             }
 
-            // 3. Insert layer id into the layer meta of both edge and node.
+            // Insert layer id into the layer meta of both edge and node.
             let node_meta = self.graph().node_meta();
 
             edge_meta
@@ -233,4 +227,95 @@ where
 
         Ok(())
     }
+
+    fn replay_add_node(
+        &mut self,
+        lsn: LSN,
+        _transaction_id: TransactionID,
+        t: EventTime,
+        node_name: Option<GID>,
+        node_id: VID,
+        node_type_and_id: Option<(String, usize)>,
+        props: Vec<(String, usize, Prop)>,
+    ) -> Result<(), StorageError> {
+        // Insert node id into resolver.
+        if let Some(ref name) = node_name {
+            self.graph()
+                .logical_to_physical
+                .set(name.as_ref(), node_id)?;
+        }
+
+        // Resolve segment and check LSN.
+        let (segment_id, pos) = self.graph().storage().nodes().resolve_pos(node_id);
+        self.resize_segments_to_vid(node_id);
+
+        let segment = self
+            .graph()
+            .storage()
+            .nodes()
+            .get_or_create_segment(segment_id);
+
+        let immut_lsn = segment.immut_lsn();
+
+        // Replay this entry only if it doesn't exist in immut.
+        if immut_lsn < lsn {
+            let node_meta = self.graph().node_meta();
+
+            for (prop_name, prop_id, prop_value) in &props {
+                let prop_mapper = node_meta.temporal_prop_mapper();
+                match prop_mapper.get_dtype(*prop_id) {
+                    None => {
+                        prop_mapper.set_id_and_dtype(
+                            prop_name.as_str(),
+                            *prop_id,
+                            prop_value.dtype(),
+                        );
+                    }
+                    Some(old_dtype) => {
+                        let dtype = prop_value.dtype();
+                        let mut unified = false;
+                        let new_dtype = unify_types(&old_dtype, &dtype, &mut unified)?;
+                        if unified {
+                            prop_mapper.set_dtype(*prop_id, new_dtype);
+                        }
+                    }
+                }
+            }
+
+            // Set node type metadata early to prevent issues with borrowing node_writer.
+            if let Some((ref node_type, node_type_id)) = node_type_and_id {
+                node_meta
+                    .node_type_meta()
+                    .set_id(node_type.as_str(), node_type_id);
+            }
+
+            let mut node_writer = self.nodes.get_mut(segment_id).unwrap().writer();
+
+            if !node_writer.has_node(pos, STATIC_GRAPH_LAYER_ID) {
+                node_writer.increment_seg_num_nodes();
+            }
+
+            if let Some(name) = node_name {
+                node_writer.store_node_id(pos, STATIC_GRAPH_LAYER_ID, name);
+            }
+
+            if let Some((_, node_type_id)) = node_type_and_id {
+                node_writer.store_node_type(pos, STATIC_GRAPH_LAYER_ID, node_type_id);
+            }
+
+            // Add the node with its timestamp and props.
+            node_writer.add_props(
+                t,
+                pos,
+                STATIC_GRAPH_LAYER_ID,
+                props
+                    .into_iter()
+                    .map(|(_, prop_id, prop_value)| (prop_id, prop_value)),
+            );
+
+            node_writer.mut_segment.set_lsn(lsn);
+        }
+
+        Ok(())
+    }
 }

db4-graph/src/replay.rs

@@ -33,6 +33,9 @@ iter-enum = { workspace = true, features = ["rayon"] }
 chrono = { workspace = true, optional = true }
 clap.workspace = true
 tracing.workspace = true
+dashmap.workspace = true
+lock_api.workspace = true
+once_cell.workspace = true
 
 [dev-dependencies]
 proptest.workspace = true

db4-storage/Cargo.toml

@@ -126,6 +126,8 @@ pub trait LockedESegment: Send + Sync + std::fmt::Debug {
         &'a self,
         layer_ids: &'b LayerIds,
     ) -> impl ParallelIterator<Item = Self::EntryRef<'a>> + Sync + 'a;
+
+    fn num_edges(&self) -> u32;
 }
 
 pub trait EdgeEntryOps<'a>: Send + Sync {

db4-storage/src/api/edges.rs

@@ -110,7 +110,7 @@ pub trait NodeSegmentOps: Send + Sync + Debug + 'static {
 
     fn entry(&self, pos: impl Into<LocalPOS>) -> Self::Entry<'_>;
 
-    fn locked(self: &Arc<Self>) -> Self::ArcLockedSegment;
+    fn locked(&self) -> Self::ArcLockedSegment;
 
     fn flush(&self) -> Result<(), StorageError>;
 
@@ -129,7 +129,7 @@ pub trait NodeSegmentOps: Send + Sync + Debug + 'static {
     fn nodes_counter(&self) -> &AtomicU32;
 
     fn increment_num_nodes(&self, max_page_len: u32) {
-        increment_and_clamp(self.nodes_counter(), max_page_len);
+        increment_and_clamp(self.nodes_counter(), 1, max_page_len);
     }
 
     fn num_nodes(&self) -> u32 {

db4-storage/src/api/nodes.rs

@@ -81,10 +81,9 @@ pub type GraphTProps<'a> = GenericTProps<'a, MemGraphPropRef<'a>>;
 pub mod error {
     use std::{path::PathBuf, sync::Arc};
 
+    use crate::resolver::mapping_resolver::InvalidNodeId;
     use raphtory_api::core::{entities::properties::prop::PropError, utils::time::ParseTimeError};
-    use raphtory_core::entities::{
-        graph::logical_to_physical::InvalidNodeId, properties::props::MetadataError,
-    };
+    use raphtory_core::entities::properties::props::MetadataError;
 
     #[derive(thiserror::Error, Debug)]
     pub enum StorageError {