From a9894ebb7facaa81a79bbb2ff78fca33b775c59c Mon Sep 17 00:00:00 2001
From: Mohammad-Palla <77429079+Mohammad-Palla@users.noreply.github.com>
Date: Thu, 27 Aug 2026 18:01:13 +0100
Subject: [PATCH] perf(cluster): call graspologic_native.leiden() directly,
 skip graspologic's import chain

The Leiden path imported graspologic, which drags in a heavy chain (umap / pynndescent /
numba JIT) at import time. Call the underlying graspologic_native.leiden() binding directly
with the same arguments the wrapper uses (seed 42, randomness 0.001, modularity, unit edge
weights), falling back to the graspologic wrapper and then NetworkX Louvain when the native
binding is absent. Clustering output is unchanged.
---
 graphify/cluster.py | 83 +++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 81 insertions(+), 2 deletions(-)

diff --git a/graphify/cluster.py b/graphify/cluster.py
index 68221070..db476b5b 100644
--- a/graphify/cluster.py
+++ b/graphify/cluster.py
@@ -19,11 +19,86 @@ def _suppress_output():
     return contextlib.redirect_stdout(io.StringIO())
 
 
+def _native_leiden(stable: nx.Graph, resolution: float) -> dict[str, int] | None:
+    """Call graspologic_native.leiden() directly, bypassing graspologic's own
+    package import.
+
+    graspologic.partition.leiden() is a thin wrapper around exactly this
+    native (Rust) call. Importing the *package* — as opposed to the native
+    extension module it depends on — pulls in graspologic.layouts, which
+    imports umap, which imports pynndescent, which numba-JIT-compiles at
+    import time for a layout algorithm this function never calls: measured
+    at 7-19s of one-time import cost against a ~1s native call and a ~1.4s
+    full round trip (conversion + call + map-back) — see the "third update"
+    in GRAPHIFY_BUILD_PERF.md for the measurements this is based on.
+
+    Returns None (the caller falls through to the graspologic.partition.leiden
+    path, then to the networkx Louvain fallback) if graspologic_native isn't
+    installed, or if `stable` isn't the plain undirected, non-multigraph
+    input leiden actually supports — the same shape check
+    graspologic.partition.leiden itself makes before calling the same native
+    function.
+    """
+    try:
+        import graspologic_native as gn
+    except ImportError:
+        return None
+
+    if stable.is_directed() or stable.is_multigraph():
+        return None
+
+    # graspologic_native identifies nodes by their string form; two DISTINCT
+    # node objects that happen to stringify the same way would silently merge
+    # under it (this is exactly what graspologic.partition.leiden's own
+    # _IdentityMapper guards against). Graphify's own node IDs are already
+    # unique strings by construction — extractors/resolution.py's
+    # _disambiguate_colliding_node_ids salts any two distinct nodes that would
+    # otherwise share a string id before the graph is ever built — so this is
+    # a defensive check on an assumption that should never actually trip, not
+    # an expected path. One pass over the nodes, cheaper than an
+    # _IdentityMapper-style dict-store-per-edge-endpoint.
+    id_to_node: dict[str, object] = {}
+    for node in stable.nodes():
+        key = str(node)
+        existing = id_to_node.get(key)
+        if existing is not None and existing != node:
+            return None  # let graspologic.partition.leiden's own check handle/raise on this
+        id_to_node[key] = node
+
+    edges = [
+        (str(u), str(v), float(attrs.get("weight", 1.0)))
+        for u, v, attrs in stable.edges(data=True)
+    ]
+
+    try:
+        old_stderr = sys.stderr
+        try:
+            sys.stderr = io.StringIO()
+            with _suppress_output():
+                _quality, native_partitions = gn.leiden(
+                    edges=edges,
+                    starting_communities=None,
+                    resolution=resolution,
+                    randomness=0.001,
+                    iterations=1,
+                    use_modularity=True,
+                    seed=42,
+                    trials=1,
+                )
+        finally:
+            sys.stderr = old_stderr
+    except Exception:
+        return None
+
+    return {id_to_node[node_id]: community for node_id, community in native_partitions.items()}
+
+
 def _partition(G: nx.Graph, resolution: float = 1.0) -> dict[str, int]:
     """Run community detection. Returns {node_id: community_id}.
 
-    Tries Leiden (graspologic) first — best quality.
-    Falls back to Louvain (built into networkx) if graspologic is not installed.
+    Tries Leiden (graspologic_native directly, then graspologic) first — best
+    quality. Falls back to Louvain (built into networkx) if neither is
+    installed.
 
     resolution > 1.0 → more, smaller communities.
     resolution < 1.0 → fewer, larger communities.
@@ -44,6 +119,10 @@ def _partition(G: nx.Graph, resolution: float = 1.0) -> dict[str, int]:
     for src, tgt, attrs in edge_rows:
         stable.add_edge(src, tgt, **attrs)
 
+    native_result = _native_leiden(stable, resolution)
+    if native_result is not None:
+        return native_result
+
     try:
         from graspologic.partition import leiden
         lsig = inspect.signature(leiden).parameters