impl node_types for disk graph

python/tests/test_diskgraph.py

@@ -4,10 +4,9 @@
 import tempfile
 import os
 
-
 def test_disk_graph():
     curr_dir = os.path.dirname(os.path.abspath(__file__))
-    rsc_dir = os.path.join(curr_dir, "..","..", "pometry-storage-private", "resources")
+    rsc_dir = os.path.join(curr_dir, "..", "..", "pometry-storage-private", "resources")
     rsc_dir = os.path.normpath(rsc_dir)
     print("rsc_dir:", rsc_dir + "/netflowsorted/nft_sorted")
 
@@ -66,6 +65,7 @@ def test_disk_graph():
             read_chunk_size,
             concurrent_files,
             num_threads,
+            None,
             print_result=False,
         )
 
@@ -91,3 +91,65 @@ def test_disk_graph():
         "Page Rank", algorithms.pagerank, g.layer("netflow"), 100, print_result=False
     )
     assert len(list(actual.get_all_with_names())) == 1624
+
+def test_disk_graph_type_filter():
+    curr_dir = os.path.dirname(os.path.abspath(__file__))
+    rsc_dir = os.path.join(curr_dir, "..", "..", "pometry-storage-private", "resources")
+    rsc_dir = os.path.normpath(rsc_dir)
+    print("rsc_dir:", rsc_dir + "/netflowsorted/nft_sorted")
+
+    graph_dir = tempfile.TemporaryDirectory()
+    layer_parquet_cols = [
+        {
+            "parquet_dir": rsc_dir + "/netflowsorted/nft_sorted",
+            "layer": "netflow",
+            "src_col": "src",
+            "dst_col": "dst",
+            "time_col": "epoch_time",
+        }
+    ]
+
+    chunk_size = 268_435_456
+    num_threads = 4
+    t_props_chunk_size = int(chunk_size / 8)
+    read_chunk_size = 4_000_000
+    concurrent_files = 1
+
+    g = DiskGraph.load_from_parquets(
+        graph_dir.name,
+        layer_parquet_cols,
+        rsc_dir + "/netflowsorted/props/props.parquet",
+        chunk_size,
+        t_props_chunk_size,
+        read_chunk_size,
+        concurrent_files,
+        num_threads,
+        "node_type"
+    )
+
+    assert g.count_nodes() == 1619
+    assert g.layer("netflow").count_edges() == 2018
+    assert g.earliest_time == 7257619
+    assert g.latest_time == 7343970
+
+    assert len(g.nodes.type_filter(["A"]).name.collect()) == 785
+    assert len(g.nodes.type_filter([""]).name.collect()) == 0
+    assert len(g.nodes.type_filter(["A", "B"]).name.collect()) == 1619
+
+    neighbor_names = g.nodes.type_filter(["A"]).neighbours.name.collect()
+    total_length = sum(len(names) for names in neighbor_names)
+    assert total_length == 2056
+
+    assert g.nodes.type_filter([]).name.collect() == []
+
+    neighbor_names = g.nodes.type_filter(["A"]).neighbours.type_filter(["B"]).name.collect()
+    total_length = sum(len(names) for names in neighbor_names)
+    assert total_length == 1023
+
+    assert g.node("Comp175846").neighbours.type_filter(["A"]).name.collect() == ["Comp844043"]
+    assert g.node("Comp175846").neighbours.type_filter(["B"]).name.collect() == []
+    assert g.node("Comp175846").neighbours.type_filter([]).name.collect() == []
+    assert g.node("Comp175846").neighbours.type_filter(["A", "B"]).name.collect() == ["Comp844043"]
+
+    neighbor_names = g.node("Comp175846").neighbours.neighbours.name.collect()
+    assert len(neighbor_names) == 193

raphtory-api/Cargo.toml

@@ -17,3 +17,19 @@ edition.workspace = true
 [dependencies]
 serde = { workspace = true, features = ["derive"] }
 chrono.workspace = true
+dashmap = { workspace = true }
+rustc-hash = { workspace = true }
+lock_api = { workspace = true }
+parking_lot = { workspace = true }
+pyo3 = { workspace = true, optional = true }
+rayon = { workspace = true }
+rand = { workspace = true }
+quickcheck = { workspace = true }
+quickcheck_macros = { workspace = true }
+
+[features]
+default = []
+# Enables generating the pyo3 python bindings
+python = [
+    "dep:pyo3",
+]

raphtory-api/src/core/storage/arc_str.rs

@@ -0,0 +1,121 @@
+use serde::{Deserialize, Serialize};
+use std::{
+    borrow::Borrow,
+    cmp::Ordering,
+    fmt,
+    fmt::{Display, Formatter},
+    ops::Deref,
+    sync::Arc,
+};
+
+#[derive(Clone, Debug, Eq, Ord, Hash, Serialize, Deserialize)]
+pub struct ArcStr(pub Arc<str>);
+
+impl Display for ArcStr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        Display::fmt(&self.0, f)
+    }
+}
+
+impl<T: Into<Arc<str>>> From<T> for ArcStr {
+    fn from(value: T) -> Self {
+        ArcStr(value.into())
+    }
+}
+
+impl From<ArcStr> for String {
+    fn from(value: ArcStr) -> Self {
+        value.to_string()
+    }
+}
+
+impl From<&ArcStr> for String {
+    fn from(value: &ArcStr) -> Self {
+        value.clone().into()
+    }
+}
+
+impl Deref for ArcStr {
+    type Target = Arc<str>;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl Borrow<str> for ArcStr {
+    #[inline]
+    fn borrow(&self) -> &str {
+        self.0.borrow()
+    }
+}
+
+impl<T> AsRef<T> for ArcStr
+where
+    T: ?Sized,
+    <ArcStr as Deref>::Target: AsRef<T>,
+{
+    fn as_ref(&self) -> &T {
+        self.deref().as_ref()
+    }
+}
+
+impl<T: Borrow<str> + ?Sized> PartialEq<T> for ArcStr {
+    fn eq(&self, other: &T) -> bool {
+        <ArcStr as Borrow<str>>::borrow(self).eq(other.borrow())
+    }
+}
+
+impl<T: Borrow<str>> PartialOrd<T> for ArcStr {
+    fn partial_cmp(&self, other: &T) -> Option<Ordering> {
+        <ArcStr as Borrow<str>>::borrow(self).partial_cmp(other.borrow())
+    }
+}
+
+pub trait OptionAsStr<'a> {
+    fn as_str(self) -> Option<&'a str>;
+}
+
+impl<'a, O: AsRef<str> + 'a> OptionAsStr<'a> for &'a Option<O> {
+    fn as_str(self) -> Option<&'a str> {
+        self.as_ref().map(|s| s.as_ref())
+    }
+}
+
+impl<'a, O: AsRef<str> + 'a> OptionAsStr<'a> for Option<&'a O> {
+    fn as_str(self) -> Option<&'a str> {
+        self.map(|s| s.as_ref())
+    }
+}
+
+#[cfg(test)]
+mod test_arc_str {
+    use crate::core::storage::arc_str::{ArcStr, OptionAsStr};
+    use std::sync::Arc;
+
+    #[test]
+    fn can_compare_with_str() {
+        let test: ArcStr = "test".into();
+        assert_eq!(test, "test");
+        assert_eq!(test, "test".to_string());
+        assert_eq!(test, Arc::from("test"));
+        assert_eq!(&test, &"test".to_string())
+    }
+
+    #[test]
+    fn test_option_conv() {
+        let test: Option<ArcStr> = Some("test".into());
+
+        let opt_str = test.as_str();
+        assert_eq!(opt_str, Some("test"));
+
+        let test_ref = test.as_ref();
+        let opt_str = test_ref.as_str();
+        assert_eq!(opt_str, Some("test"));
+
+        let test = Some("test".to_string());
+        let opt_str = test.as_str();
+        assert_eq!(opt_str, Some("test"));
+    }
+}