perf(serve): trigram candidate prefilter to cut O(N) query latency on large graphs (byte-identical results)

graphify/serve.py

@@ -4,6 +4,7 @@
 import math
 import re
 import sys
+from array import array
 from pathlib import Path
 import networkx as nx
 from networkx.readwrite import json_graph
@@ -135,6 +136,109 @@ def _compute_idf(G: nx.Graph, terms: list[str]) -> dict[str, float]:
     return {t: cache.get(t, math.log(1 + N)) for t in terms}
 
 
+def _trigrams(text: str) -> set[str]:
+    """Character trigrams of `text`; for <3-char text the whole string is the key."""
+    if len(text) < 3:
+        return {text} if text else set()
+    return {text[i:i + 3] for i in range(len(text) - 2)}
+
+
+def _node_search_text(data: dict, nid: str) -> str:
+    """Concatenate every field _score_nodes / _find_node match a query against, so
+    one trigram index over this text is a complete candidate generator for both.
+
+    - `norm_label` and `source_file` feed _score_nodes' per-term substring tiers.
+    - `label_tokens` (the space-joined token form) feeds _find_node's
+      `term in label_tokens` branch, where a multi-word `term` can span a token
+      boundary that punctuation hides in `norm_label` (e.g. query "foo bar" matches
+      label "foo.bar" only via its tokenized form).
+    - `nid` feeds the whole-query `joined == nid_lower` tier.
+
+    NUL separators stop a trigram from spanning two fields (a query never contains
+    NUL, so a cross-field trigram can never be a real match).
+    """
+    norm_label = data.get("norm_label") or _strip_diacritics(data.get("label") or "").lower()
+    label_tokens = " ".join(_search_tokens(data.get("label") or ""))
+    source = (data.get("source_file") or "").lower()
+    return "\x00".join((norm_label, label_tokens, str(nid).lower(), source))
+
+
+def _get_trigram_index(G: nx.Graph) -> dict:
+    """Lazily build and cache a trigram -> node-position postings map on the graph.
+
+    Cached on `G.graph` so it auto-invalidates when _maybe_reload() swaps in a
+    fresh graph object, exactly like `_idf_cache`. `set_cache` memoizes per-trigram
+    id-sets across queries within one graph generation.
+    """
+    idx = G.graph.get("_trigram_index")
+    if idx is not None:
+        return idx
+    ids = list(G.nodes())
+    postings: dict[str, array] = {}
+    for i, nid in enumerate(ids):
+        for g in _trigrams(_node_search_text(G.nodes[nid], nid)):
+            bucket = postings.get(g)
+            if bucket is None:
+                bucket = array("i")
+                postings[g] = bucket
+            bucket.append(i)
+    idx = {"ids": ids, "postings": postings, "set_cache": {}}
+    G.graph["_trigram_index"] = idx
+    return idx
+
+
+def _trigram_candidates(G: nx.Graph, needles: list[str], *, guard_frac: float = 0.10) -> list[str] | None:
+    """Node IDs whose text could contain any `needle` as a substring, via the
+    trigram index — a *superset* the caller then re-scores with the exact predicates.
+
+    Returns candidates in graph-iteration order (so order-sensitive callers like
+    _find_node stay byte-identical to a full scan), or **None** when the index isn't
+    worth it — a needle is too short to trigram, or its rarest trigram is still
+    common enough that the candidate set would approach the whole graph. The caller
+    falls back to the full scan, preserving the never-worse contract. The guard is
+    cheap: postings-length lookups only, no set intersection.
+    """
+    idx = _get_trigram_index(G)
+    ids, postings, set_cache = idx["ids"], idx["postings"], idx["set_cache"]
+    n = len(ids)
+    if n == 0:
+        return []
+    needles = [s for s in needles if s]
+    thresh = int(n * guard_frac)
+    for s in needles:
+        tgs = _trigrams(s)
+        if not tgs or any(len(g) < 3 for g in tgs):
+            return None  # too short to trigram-filter
+        present = [len(postings[g]) for g in tgs if g in postings]
+        if not present:
+            continue  # this needle matches nothing — contributes no candidates
+        if min(present) > thresh:
+            return None  # rarest trigram still too common -> not worth the index
+    cand: set[int] = set()
+    for s in needles:
+        sets: list[set] | None = []
+        for g in _trigrams(s):
+            bucket = postings.get(g)
+            if bucket is None:
+                sets = None  # a trigram absent everywhere -> needle matches nothing
+                break
+            cached = set_cache.get(g)
+            if cached is None:
+                cached = set(bucket)
+                set_cache[g] = cached
+            sets.append(cached)
+        if not sets:
+            continue
+        sets.sort(key=len)  # intersect smallest-first
+        hit = set(sets[0])
+        for other in sets[1:]:
+            hit &= other
+            if not hit:
+                break
+        cand |= hit
+    return [ids[i] for i in sorted(cand)]
+
+
 def _score_nodes(G: nx.Graph, terms: list[str]) -> list[tuple[float, str]]:
     scored = []
     norm_terms = [tok for t in terms for tok in _search_tokens(t)]
@@ -144,7 +248,16 @@ def _score_nodes(G: nx.Graph, terms: list[str]) -> list[tuple[float, str]]:
     # Weight the full-query bonus by the rarest constituent term so a specific
     # multi-word label still outweighs common-token noise; floor at 1.0.
     joined_w = max((idf.get(t, 1.0) for t in norm_terms), default=1.0)
-    for nid, data in G.nodes(data=True):
+    # Trigram prefilter: score only nodes whose text could match a term, falling
+    # back to the whole graph when the index isn't selective. The result is
+    # identical either way — the per-node scoring below is unchanged and a
+    # non-candidate node always scores 0. (IDF above stays a whole-graph statistic.)
+    candidate_ids = _trigram_candidates(G, norm_terms + ([joined] if joined else []))
+    node_iter = (
+        G.nodes(data=True) if candidate_ids is None
+        else ((nid, G.nodes[nid]) for nid in candidate_ids)
+    )
+    for nid, data in node_iter:
         norm_label = data.get("norm_label") or _strip_diacritics(data.get("label") or "").lower()
         bare_label = norm_label.rstrip("()")
         # Tokenized form of the label (punctuation stripped, same transform as the
@@ -462,7 +575,14 @@ def _find_node(G: nx.Graph, label: str) -> list[str]:
     exact: list[str] = []
     prefix: list[str] = []
     substring: list[str] = []
-    for nid, d in G.nodes(data=True):
+    # Trigram prefilter (graph-iteration order preserved so exact/prefix/substring
+    # ordering — and thus matches[0] — is byte-identical to the full scan).
+    candidate_ids = _trigram_candidates(G, [term])
+    node_iter = (
+        G.nodes(data=True) if candidate_ids is None
+        else ((nid, G.nodes[nid]) for nid in candidate_ids)
+    )
+    for nid, d in node_iter:
         norm_label = d.get("norm_label") or _strip_diacritics(d.get("label") or "").lower()
         bare_label = norm_label.rstrip("()")
         label_tokens = " ".join(_search_tokens(d.get("label") or ""))
@@ -530,6 +650,11 @@ def _build_server(graph_path: str):
 
     G = _load_graph(graph_path)
     communities = _communities_from_graph(G)
+    # Build the trigram query index eagerly so the first query doesn't pay the
+    # one-time build (a few seconds on a large graph). serve exists to answer
+    # queries, so absorbing the cost at startup beats spiking a random first
+    # request; it's cached on G.graph, so queries just reuse it.
+    _get_trigram_index(G)
 
     # Hot-reload state: mtime+size key lets us detect graph.json changes without
     # polling. Initialised from the file stat at startup so the first tool call
@@ -562,6 +687,8 @@ def _maybe_reload() -> None:
                 new_G = _load_graph(graph_path)
             except SystemExit:
                 return  # keep serving stale graph on transient read error
+            _get_trigram_index(new_G)  # warm before exposing, so the first
+            # post-reload query is fast too (same rationale as startup)
             G = new_G
             communities = _communities_from_graph(new_G)
             _reload_state["mtime_ns"], _reload_state["size"] = key

tests/test_serve.py

@@ -12,6 +12,10 @@
     _bfs,
     _dfs,
     _find_node,
+    _trigrams,
+    _node_search_text,
+    _get_trigram_index,
+    _trigram_candidates,
     _filter_graph_by_context,
     _infer_context_filters,
     _query_terms,
@@ -130,6 +134,107 @@ def test_find_node_matches_full_punctuated_unicode_label():
     assert _find_node(G, "Skill /auditar — Auditoría inquisitiva de enlaces") == ["n1"]
 
 
+# --- trigram candidate prefilter (the trigram index that shrinks the O(N) scan) ---
+
+
+def _force_full_scan(monkeypatch):
+    """Disable the prefilter so a call exercises the original full-node scan."""
+    monkeypatch.setattr("graphify.serve._trigram_candidates", lambda *a, **k: None)
+
+
+def _make_big_graph(n: int = 150) -> nx.Graph:
+    """A graph large enough that the selectivity guard lets the fast-path fire for
+    rare terms and fall back for common ones. Most labels share the 'item'/'node'
+    stem (common), plus a few distinctive rare labels and one punctuated label."""
+    G = nx.Graph()
+    for i in range(n):
+        G.add_node(f"id{i}", label=f"item node {i}", source_file=f"pkg/item_{i}.py")
+    G.add_node("rareA", label="ZebraQuokkaWidget", source_file="zoo/zqw.py")
+    G.add_node("rareB", label="MarmosetGadget handler", source_file="zoo/marmoset.py")
+    G.add_node("punct", label="Foo.Bar:Baz", source_file="pkg/foobar.py")
+    return G
+
+
+def test_trigrams_basic():
+    assert _trigrams("foobar") == {"foo", "oob", "oba", "bar"}
+    assert _trigrams("ab") == {"ab"}        # <3 chars -> whole string is the key
+    assert _trigrams("") == set()
+
+
+def test_node_search_text_includes_all_matched_fields():
+    G = _make_big_graph()
+    text = _node_search_text(G.nodes["punct"], "punct")
+    # norm_label, the tokenized label (label_tokens), nid, and source are all present,
+    # NUL-separated so trigrams can't span fields.
+    parts = text.split("\x00")
+    assert parts[0] == "foo.bar:baz"          # norm_label (punctuation kept)
+    assert parts[1] == "foo bar baz"          # label_tokens (tokenized)
+    assert parts[2] == "punct"                # nid
+    assert parts[3] == "pkg/foobar.py"        # source_file
+
+
+def test_trigram_candidates_fast_path_fires_for_rare_term():
+    G = _make_big_graph()
+    cand = _trigram_candidates(G, ["zebraquokkawidget"])
+    assert cand is not None                   # selective -> fast-path used
+    assert "rareA" in cand
+    assert len(cand) < G.number_of_nodes()    # a real shrink, not the whole graph
+
+
+def test_trigram_candidates_falls_back_on_common_term():
+    G = _make_big_graph()
+    # 'item' is in the label of every one of the 150 'item node N' nodes -> the
+    # rarest trigram is still common -> guard returns None (full-scan fallback).
+    assert _trigram_candidates(G, ["item"]) is None
+
+
+def test_trigram_candidates_falls_back_on_short_token():
+    G = _make_big_graph()
+    assert _trigram_candidates(G, ["ab"]) is None   # <3 chars -> can't trigram-filter
+
+
+def test_score_nodes_prefilter_is_identical_to_full_scan(monkeypatch):
+    G = _make_big_graph()
+    queries = ["zebraquokkawidget", "marmosetgadget handler", "foo bar baz",
+               "item", "node 42", "nonexistentxyz"]
+    for q in queries:
+        terms = _query_terms(q)
+        fast = _score_nodes(G, terms)
+        _force_full_scan(monkeypatch)
+        full = _score_nodes(G, terms)
+        monkeypatch.undo()
+        assert fast == full, f"prefilter diverged from full scan for {q!r}"
+
+
+def test_find_node_prefilter_is_identical_to_full_scan(monkeypatch):
+    G = _make_big_graph()
+    # includes the punctuated label, exercised via its tokenized (label_tokens) form
+    for label in ["ZebraQuokkaWidget", "MarmosetGadget handler", "Foo Bar Baz",
+                  "item node 7", "missing"]:
+        fast = _find_node(G, label)
+        _force_full_scan(monkeypatch)
+        full = _find_node(G, label)
+        monkeypatch.undo()
+        assert fast == full, f"_find_node prefilter diverged (order!) for {label!r}"
+
+
+def test_find_node_label_tokens_branch_covered_by_index():
+    # "foo bar baz" matches label "Foo.Bar:Baz" only via the tokenized label_tokens
+    # form (the dotted/colon norm_label never contains the spaced query). The index
+    # must surface this node as a candidate, or the prefilter would silently drop it.
+    G = _make_big_graph()
+    assert _find_node(G, "Foo Bar Baz") == ["punct"]
+
+
+def test_trigram_index_cached_and_rebuilt_per_graph():
+    G = _make_big_graph()
+    idx1 = _get_trigram_index(G)
+    assert idx1 is _get_trigram_index(G)            # cached on the same graph object
+    assert G.graph["_trigram_index"] is idx1
+    G2 = _make_big_graph()
+    assert _get_trigram_index(G2) is not idx1       # a fresh graph rebuilds (reload safety)
+
+
 def test_query_terms_strips_search_punctuation():
     assert _query_terms("what calls extract?") == ["what", "calls", "extract"]