mirror of
https://github.com/safishamsi/graphify.git
synced 2026-09-24 22:45:43 +00:00
perf(serve): trigram candidate prefilter to cut O(N) query latency
_score_nodes and _find_node scan every node per query (O(nodes x terms)), so query latency scales with total graph size regardless of where the answer lives. Add a lazily-built character-trigram index (cached on the graph object, auto-invalidated on hot-reload like _idf_cache) that narrows each query to a small candidate superset before the unchanged scoring loop. Results are byte-identical: the index is a pure candidate generator over the exact fields the scorer reads (norm_label, label_tokens, nid, source_file); a non-candidate node always scores 0, and IDF stays a whole-graph statistic. _find_node candidates are returned in graph iteration order so its exact/prefix/substring ordering, and matches[0], stay unchanged. A selectivity guard falls back to the full scan when a query term is too short to trigram or its rarest trigram is still common (broad terms like model/client), preserving a never-worse contract. The index builds eagerly at load and before a reloaded graph is swapped in, so neither the first query nor the first post-reload query pays the one-time build cost. Storage format is unchanged (in-memory index only). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
committed by
safishamsi
co-authored by
Claude Opus 4.8
parent
a4d09aefd8
commit
53edd27e6a
+129
-2
@@ -4,6 +4,7 @@ import json
|
||||
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 _build_server(graph_path: str):
|
||||
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
|
||||
|
||||
@@ -12,6 +12,10 @@ from graphify.serve import (
|
||||
_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"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user