Repository: neo4j-labs/graph (language: rust)

Issue:
Python: Surface wcc



Discussion / hints:

Fix patch:
diff --git a/crates/algos/src/afforest.rs b/crates/algos/src/afforest.rs
index ad4ae2b4..914c1275 100644
--- a/crates/algos/src/afforest.rs
+++ b/crates/algos/src/afforest.rs
@@ -25,7 +25,7 @@ impl<NI: Idx> UnionFind<NI> for Afforest<NI> {
 
         while p1 != p2 {
             let high = NI::max(p1, p2);
-            let low = p1 + (p2 - high);
+            let low = p1 + p2 - high;
             let p_high = self.find(high);
 
             if p_high == low
diff --git a/crates/algos/src/wcc.rs b/crates/algos/src/wcc.rs
index 365482f5..fb96d6e4 100644
--- a/crates/algos/src/wcc.rs
+++ b/crates/algos/src/wcc.rs
@@ -42,29 +42,33 @@ pub use crate::dss::DisjointSetStruct;
 #[cfg_attr(feature = "clap", derive(clap::Args))]
 pub struct WccConfig {
     /// Number of nodes to be processed in batch by a single thread.
-    #[cfg_attr(feature = "clap", clap(long, default_value_t = WccConfig::default().chunk_size))]
+    #[cfg_attr(feature = "clap", clap(long, default_value_t = WccConfig::DEFAULT_CHUNK_SIZE))]
     pub chunk_size: usize,
 
     /// Number of relationships of each node to sample during subgraph linking.
-    #[cfg_attr(feature = "clap", clap(long, default_value_t = WccConfig::default().neighbor_rounds))]
+    #[cfg_attr(feature = "clap", clap(long, default_value_t = WccConfig::DEFAULT_NEIGHBOR_ROUNDS))]
     pub neighbor_rounds: usize,
 
     /// Number of samples to draw from the DSS to find the largest component.
-    #[cfg_attr(feature = "clap", clap(long, default_value_t = WccConfig::default().sampling_size))]
+    #[cfg_attr(feature = "clap", clap(long, default_value_t = WccConfig::DEFAULT_SAMPLING_SIZE))]
     pub sampling_size: usize,
 }
 
 impl Default for WccConfig {
     fn default() -> Self {
         Self {
-            chunk_size: 16384,
-            neighbor_rounds: 2,
-            sampling_size: 1024,
+            chunk_size: WccConfig::DEFAULT_CHUNK_SIZE,
+            neighbor_rounds: WccConfig::DEFAULT_NEIGHBOR_ROUNDS,
+            sampling_size: WccConfig::DEFAULT_SAMPLING_SIZE,
         }
     }
 }
 
 impl WccConfig {
+    pub const DEFAULT_CHUNK_SIZE: usize = 16384;
+    pub const DEFAULT_NEIGHBOR_ROUNDS: usize = 2;
+    pub const DEFAULT_SAMPLING_SIZE: usize = 1024;
+
     pub fn new(chunk_size: usize, neighbor_rounds: usize, sampling_size: usize) -> Self {
         Self {
             chunk_size,
diff --git a/crates/unladen_swallow/README.md b/crates/unladen_swallow/README.md
index 248a2f5d..4d15001a 100644
--- a/crates/unladen_swallow/README.md
+++ b/crates/unladen_swallow/README.md
@@ -27,20 +27,20 @@ source .env/bin/activate
 Build in debug mode.
 
 ```
-maturin develop --cargo-extra-args="--features=extension-module"
+maturin develop
 ```
 
 Build in release mode.
 
 ```
-maturin develop --release --cargo-extra-args="--features=extension-module"
+maturin develop --release
 ```
 
 Rebuild the extension in release mode 2 seconds after the last file change.
 This is an optional step.
 
 ```
-cargo watch --shell 'maturin develop --release --cargo-extra-args="--features=extension-module"' --delay 2
+cargo watch --shell 'maturin develop --release' --delay 2
 ```
 
 ### Testing
@@ -62,4 +62,7 @@ isort tests
 
 # Verify with https://pypi.org/project/flake8/
 flake8 tests
+
+# Very types using http://mypy-lang.org
+mypy .
 ```
diff --git a/crates/unladen_swallow/requirements/dev.txt b/crates/unladen_swallow/requirements/dev.txt
index e7311579..238362a7 100644
--- a/crates/unladen_swallow/requirements/dev.txt
+++ b/crates/unladen_swallow/requirements/dev.txt
@@ -4,6 +4,7 @@ flake8
 isort
 maturin
 mypy
+numpy
 pytest
 pytest-annotate
 tox
diff --git a/crates/unladen_swallow/src/graphs/digraph.rs b/crates/unladen_swallow/src/graphs/digraph.rs
index 33516c82..70a117c3 100644
--- a/crates/unladen_swallow/src/graphs/digraph.rs
+++ b/crates/unladen_swallow/src/graphs/digraph.rs
@@ -1,6 +1,6 @@
 use super::{Graph, Layout, PyGraph};
-use crate::page_rank::PageRankResult;
-use graph::{page_rank::PageRankConfig, prelude::DirectedCsrGraph};
+use crate::{page_rank::PageRankResult, wcc::WccResult};
+use graph::{page_rank::PageRankConfig, prelude::DirectedCsrGraph, wcc::WccConfig};
 use numpy::PyArray1;
 use pyo3::{prelude::*, types::PyList};
 use std::path::PathBuf;
@@ -114,6 +114,24 @@ impl DiGraph {
         let config = PageRankConfig::new(max_iterations, tolerance, damping_factor);
         crate::page_rank::page_rank(py, self.inner.g(), config)
     }
+
+    /// Run Weakly Connected Compontents on this graph.
+    #[args(
+        "*",
+        chunk_size = "WccConfig::DEFAULT_CHUNK_SIZE",
+        neighbor_rounds = "WccConfig::DEFAULT_NEIGHBOR_ROUNDS",
+        sampling_size = "WccConfig::DEFAULT_SAMPLING_SIZE"
+    )]
+    pub fn wcc(
+        &self,
+        py: Python<'_>,
+        chunk_size: usize,
+        neighbor_rounds: usize,
+        sampling_size: usize,
+    ) -> WccResult {
+        let config = WccConfig::new(chunk_size, neighbor_rounds, sampling_size);
+        WccResult::new(crate::wcc::wcc(py, self.inner.g(), config))
+    }
 }
 
 impl std::fmt::Debug for DiGraph {
diff --git a/crates/unladen_swallow/src/lib.rs b/crates/unladen_swallow/src/lib.rs
index 83f9d034..94b4e818 100644
--- a/crates/unladen_swallow/src/lib.rs
+++ b/crates/unladen_swallow/src/lib.rs
@@ -8,6 +8,7 @@ use pyo3_log::{Caching, Logger};
 
 mod graphs;
 mod page_rank;
+mod wcc;
 
 type GResult<T> = std::result::Result<T, GError>;
 
@@ -34,6 +35,7 @@ fn unladen_swallow(py: Python, m: &PyModule) -> PyResult<()> {
 
     graphs::register(py, m)?;
     page_rank::register(py, m)?;
+    wcc::register(py, m)?;
 
     Ok(())
 }
diff --git a/crates/unladen_swallow/src/wcc.rs b/crates/unladen_swallow/src/wcc.rs
new file mode 100644
index 00000000..818cb79d
--- /dev/null
+++ b/crates/unladen_swallow/src/wcc.rs
@@ -0,0 +1,87 @@
+use crate::graphs::{NumpyType, SharedSlice};
+use graph::prelude::{
+    wcc_afforest as graph_wcc, Components, DirectedDegrees, DirectedNeighbors, Graph as GraphTrait,
+    Idx, WccConfig,
+};
+use numpy::PyArray1;
+use pyo3::prelude::*;
+use std::hash::Hash;
+use std::marker::PhantomData;
+use std::time::{Duration, Instant};
+
+pub(crate) fn register(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+    m.add_class::<WccResult>()?;
+    Ok(())
+}
+
+pub(crate) fn wcc<NI, G, C>(py: Python<'_>, graph: &G, config: C) -> WccRes<NI>
+where
+    NI: Idx + Hash + NumpyType,
+    G: GraphTrait<NI> + DirectedDegrees<NI> + DirectedNeighbors<NI> + Sync,
+    C: Into<Option<WccConfig>> + Send,
+{
+    py.allow_threads(move || inner_wcc(graph, config))
+}
+
+fn inner_wcc<NI, G>(graph: &G, config: impl Into<Option<WccConfig>>) -> WccRes<NI>
+where
+    NI: Idx + Hash + NumpyType,
+    G: GraphTrait<NI> + DirectedDegrees<NI> + DirectedNeighbors<NI> + Sync,
+{
+    let config = config.into().unwrap_or_default();
+    let start = Instant::now();
+    let components = graph_wcc(graph, config).to_vec();
+    let micros = start.elapsed().as_micros().min(u64::MAX as _) as _;
+    let components = SharedSlice::from_vec(components);
+    WccRes {
+        components,
+        micros,
+        _phantom: PhantomData,
+    }
+}
+
+pub struct WccRes<NI> {
+    components: SharedSlice,
+    micros: u64,
+    _phantom: PhantomData<NI>,
+}
+
+#[pyclass]
+#[derive(Clone)]
+pub struct WccResult {
+    components: SharedSlice,
+    #[pyo3(get)]
+    micros: u64,
+}
+
+impl WccResult {
+    pub(crate) fn new(result: WccRes<u32>) -> Self {
+        Self {
+            components: result.components,
+            micros: result.micros,
+        }
+    }
+}
+
+impl std::fmt::Debug for WccResult {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("WccResult")
+            .field(
+                "components",
+                &format!("[... {} values]", self.components.len()),
+            )
+            .field("took", &Duration::from_micros(self.micros))
+            .finish()
+    }
+}
+
+#[pymethods]
+impl WccResult {
+    pub fn components<'py>(&self, py: Python<'py>) -> PyResult<&'py PyArray1<u32>> {
+        self.components.clone().into_numpy(py)
+    }
+
+    fn __repr__(&self) -> String {
+        format!("{:?}", self)
+    }
+}
diff --git a/crates/unladen_swallow/unladen_swallow.pyi b/crates/unladen_swallow/unladen_swallow.pyi
index 5f50f785..6d14ba7e 100644
--- a/crates/unladen_swallow/unladen_swallow.pyi
+++ b/crates/unladen_swallow/unladen_swallow.pyi
@@ -81,6 +81,10 @@ class DiGraph:
         self, *, max_iterations: int, tolerance: float, damping_factor: float
     ) -> PageRankResult:
         """Run Page Rank on this graph."""
+    def wcc(
+        self, *, chunk_size: int, neighbor_rounds: int, sampling_size: int
+    ) -> WccResult:
+        """Run Weakly Connected Components on this graph."""
 
 class Graph:
     """
@@ -135,3 +139,12 @@ class PageRankResult:
         pass
     def __repr__(self) -> str:
         pass
+
+class WccResult:
+    def components(self) -> npt.NDArray[np.uint32]:
+        pass
+    @property
+    def micros(self) -> float:
+        pass
+    def __repr__(self) -> str:
+        pass


Test patch:
diff --git a/crates/unladen_swallow/tests/wcc_test.py b/crates/unladen_swallow/tests/wcc_test.py
new file mode 100644
index 00000000..153e78c1
--- /dev/null
+++ b/crates/unladen_swallow/tests/wcc_test.py
@@ -0,0 +1,20 @@
+import pytest
+
+from unladen_swallow import DiGraph
+
+
+def test_wcc(g: DiGraph):
+    wcc = g.wcc()
+
+    assert wcc.micros > 0
+
+    components = wcc.components()
+
+    assert len(components) == 1 << 8
+    for component in components:
+        assert component >= 0
+        assert component < g.node_count()
+
+def test_config_must_be_kwargs(g: DiGraph):
+    with pytest.raises(TypeError):
+        g.wcc(42, 1.0, 0.1)