cortexm/index/nsg.py
Rama main · réplica de lectura del repositorio original.
"""NSG (Navigating Spreading-out Graph) index backend.
Alternative to the default quadrant page-clustered tree. NSG is more
compact and has better search efficiency at high recall on high-dim
vectors. Trade-off: build is slower (kNN pass), but query latency is
lower at the same recall target.
This module mirrors the pattern from ``cortexm.accel.QuadrantANN`` —
when the Rust wheel is installed (``pip install ./rust/quadrant``), the
hot paths route through the compiled ``NsgIndex``; otherwise a pure-numpy
implementation produces bit-identical recall behavior (it uses the same
MRNG pruning rule, the same medoid-selection logic, and the same greedy
best-first search). The numpy fallback exists so that the NSG backend
ships in production even before the wheel is built for a target platform.
Algorithm references:
* Fu, Xiang, Wang, Huang — "Fast Approximate Nearest Neighbor Search
with Navigable Spreading-out Graphs", VLDB 2019.
* The MRNG (Monotonic Relative Neighborhood Graph) prune rule keeps
an edge (p, q) iff no other neighbor r of p satisfies
dist(p, r) < dist(p, q) AND dist(q, r) < dist(p, q).
The numpy fallback uses cosine distance = 1 - dot(a, b) on unit vectors
(which is monotonic in squared Euclidean distance for unit vectors),
matching the Rust side exactly so recall ordering is preserved.
"""
from __future__ import annotations
import os
from typing import Any
import numpy as np
# --- optional Rust wheel ---------------------------------------------------
try: # pragma: no cover
import quadrant as _quadrant # type: ignore
except Exception: # pragma: no cover
_quadrant = None
_env = os.environ.get("CONTEXTM_RUST", "auto")
NSG_RUST_AVAILABLE = _quadrant is not None and hasattr(_quadrant, "NsgIndex")
NSG_RUST_ENABLED = (_env == "1") or (_env == "auto" and NSG_RUST_AVAILABLE)
def _sim_to_dist(sim: float) -> float:
"""Convert cosine similarity to MRNG distance.
For unit vectors, ``1 - dot(a, b)`` is monotonic in squared
Euclidean distance (``|a-b|^2 = 2 - 2*dot(a, b)``), so it preserves
the prune-rule ordering the Rust side relies on.
"""
return 1.0 - float(sim)
class _NsgNumpyFallback:
"""Pure-numpy NSG implementation — same recall behavior as Rust.
Build is O(N^2 D) (brute kNN), search is greedy best-first with
frontier width ``ef_search``. Intended for fallback when the wheel
is not available; the Rust path is preferred for any non-trivial N.
"""
def __init__(self, k: int = 200, ef_search: int = 64) -> None:
self.k_build = int(k)
self.ef_search = int(ef_search)
self.vectors: np.ndarray | None = None
self.dims: int = 0
self.n: int = 0
self.edges: list[np.ndarray] = []
self.nav_node: int = 0
self.n_edges: int = 0
# ids map node index -> external string id
self.ids: list[str] = []
def _medoid(self) -> int:
"""Pick the node with max mean cosine sim to all others.
Deterministic and avoids RNG. For N <= 1024 we use the full
O(N^2 D) pass; for larger N we sample 256 candidates × 1024
eval rows, mirroring the Rust crate's strategy.
"""
if self.n == 0:
return 0
if self.n <= 1024:
cand = np.arange(self.n)
eval_set = np.arange(self.n)
else:
step_c = max(1, self.n // 256)
step_e = max(1, self.n // 1024)
cand = np.arange(0, self.n, step_c)[:256]
eval_set = np.arange(0, self.n, step_e)[:1024]
eval_vecs = self.vectors[eval_set]
best = int(cand[0])
best_mean = -np.inf
for c in cand:
sims = self.vectors[c] @ eval_vecs.T
mask = eval_set != c
m = float(sims[mask].mean()) if mask.any() else float(sims.mean())
if m > best_mean:
best_mean = m
best = int(c)
return best
def _knn(self, p: int, k: int) -> np.ndarray:
"""kNN of vector p (excluding itself), ascending distance."""
sims = self.vectors @ self.vectors[p]
sims[p] = -np.inf # exclude self
# Take top-k by sim (highest sim = closest). argpartition + sort
# gives O(N) partition + O(k log k) sort — same total ordering
# the Rust side produces via heap selection.
kk = min(k, self.n - 1)
if kk <= 0:
return np.empty(0, dtype=np.int64)
part = np.argpartition(-sims, kk - 1)[:kk]
order = part[np.argsort(-sims[part])]
return order
def build(self, vectors: np.ndarray, ids: list[str] | None = None) -> None:
v = np.ascontiguousarray(vectors, dtype=np.float32)
if v.ndim != 2:
raise ValueError(f"vectors must be 2-D, got shape {v.shape}")
n, d = v.shape
if n == 0:
raise ValueError("empty corpus")
# Normalize rows so cosine == dot product. We DO NOT mutate the
# caller's array; this is the index's own storage.
norms = np.linalg.norm(v, axis=1, keepdims=True)
norms = np.where(norms == 0, 1.0, norms)
self.vectors = (v / norms).astype(np.float32)
self.dims = int(d)
self.n = int(n)
self.nav_node = self._medoid()
if ids is None:
self.ids = [str(i) for i in range(n)]
else:
if len(ids) != n:
raise ValueError(
f"len(ids)={len(ids)} != n_vectors={n}")
self.ids = list(ids)
# Build pipeline (mirrors Fu et al. VLDB 2019 + the Rust crate):
# 1. kNN graph (full, pre-prune).
# 2. MRNG prune rule.
# 3. Tree-traversal connectivity pass — BFS from the medoid; for
# any unreachable node i, add an edge from the nearest visited
# kNN of i to i. This guarantees a greedy search starting from
# the medoid can reach every node (without it, well-separated
# clusters can leave the search stuck in the medoid's cluster).
k_eff = min(self.k_build, self.n - 1)
if k_eff <= 0:
k_eff = 1
# 1) kNN graph.
knn_graph: list[np.ndarray] = [self._knn(i, k_eff) for i in range(n)]
# 2) MRNG prune.
self.edges = []
for i in range(n):
knn = knn_graph[i]
kept: list[int] = []
pv = self.vectors[i]
for q in knn:
q = int(q)
qv = self.vectors[q]
sim_pq = float(pv @ qv)
d_pq = _sim_to_dist(sim_pq)
redundant = False
for r in kept:
rv = self.vectors[r]
sim_pr = float(pv @ rv)
d_pr = _sim_to_dist(sim_pr)
if d_pr < d_pq:
sim_qr = float(qv @ rv)
d_qr = _sim_to_dist(sim_qr)
if d_qr < d_pq:
redundant = True
break
if not redundant:
kept.append(q)
self.edges.append(np.asarray(kept, dtype=np.int64))
# 3) Tree-traversal connectivity pass.
visited = np.zeros(self.n, dtype=bool)
visited[self.nav_node] = True
self._bfs_mark(self.nav_node, visited)
# Iterate until every node is reachable.
while True:
unvisited = np.where(~visited)[0]
if len(unvisited) == 0:
break
orphan = int(unvisited[0])
# Nearest visited neighbor via kNN list (ascending distance).
added = None
for cand in knn_graph[orphan]:
cand = int(cand)
if visited[cand]:
added = cand
break
# Pathological: no visited kNN — fall back to medoid edge.
src = added if added is not None else self.nav_node
# Idempotent: don't double-add.
if orphan not in self.edges[src]:
self.edges[src] = np.append(self.edges[src], orphan)
visited[orphan] = True
self._bfs_mark(orphan, visited)
self.n_edges = sum(len(e) for e in self.edges)
def _bfs_mark(self, start: int, visited: np.ndarray) -> None:
"""Mark every node reachable from `start` via `self.edges` as visited."""
stack = [start]
while stack:
node = stack.pop()
for nb in self.edges[node]:
nb = int(nb)
if not visited[nb]:
visited[nb] = True
stack.append(nb)
def search(self, query: np.ndarray, k: int = 10) -> list[tuple[str, float]]:
if self.vectors is None or self.n == 0:
return []
q = np.ascontiguousarray(query, dtype=np.float32).ravel()
if q.shape[0] != self.dims:
raise ValueError(
f"query len {q.shape[0]} != dims {self.dims}")
nq = float(np.linalg.norm(q))
if nq == 0.0:
nq = 1.0
q = q / nq
ef = max(self.ef_search, k, 1)
# Greedy best-first from nav_node — same algorithm as Rust side:
# min-heap on dist = max-heap on -sim; frontier = top-ef by sim;
# neighbor enters candidate-expansion queue iff it survives the
# frontier. Terminates when best candidate is worse than the
# frontier's worst.
import heapq
visited = np.zeros(self.n, dtype=bool)
start = int(self.nav_node)
start_sim = float(self.vectors[start] @ q)
candidates: list[tuple[float, int]] = [] # min-heap on dist (1-sim)
frontier: list[tuple[float, int]] = [] # min-heap on sim (worst at root)
visited[start] = True
heapq.heappush(candidates, (1.0 - start_sim, start))
heapq.heappush(frontier, (start_sim, start))
while candidates:
c_dist, c_id = heapq.heappop(candidates)
c_sim = 1.0 - c_dist
if len(frontier) == ef and c_sim < frontier[0][0]:
break
for nb in self.edges[c_id]:
nb = int(nb)
if visited[nb]:
continue
visited[nb] = True
s = float(self.vectors[nb] @ q)
# Admit to frontier iff (not full) OR (better than worst)
admit = (len(frontier) < ef) or (s > frontier[0][0])
if admit:
heapq.heappush(frontier, (s, nb))
if len(frontier) > ef:
heapq.heappop(frontier)
heapq.heappush(candidates, (1.0 - s, nb))
# Return top-k sorted desc by sim, as (id_str, sim) tuples.
out = sorted(frontier, key=lambda t: -t[0])[:k]
return [(self.ids[i], float(s)) for s, i in out]
def stats(self) -> dict:
avg_deg = (self.n_edges / self.n) if self.n else 0.0
return {
"mode": "numpy-fallback",
"n_vectors": self.n,
"dims": self.dims,
"k_build": self.k_build,
"n_edges": self.n_edges,
"avg_degree": round(float(avg_deg), 4),
"nav_node": self.nav_node,
"ef_search": self.ef_search,
}
class NsgBackend:
"""NSG (Navigating Spreading-out Graph) index backend.
Alternative to the default quadrant page-clustered tree. NSG is more
compact and has better search efficiency at high recall on high-dim
vectors. Trade-off: build is slower (kNN pass), but query latency
is lower at the same recall target.
Mirrors the pattern of ``cortexm.accel.QuadrantANN``: if the Rust
wheel (``quadrant.NsgIndex``) is installed and enabled, the hot
paths route through it; otherwise a pure-numpy NSG produces the same
recall ordering (the brute kNN pass uses the same selection rule,
the MRNG prune is bit-identical, the search uses the same greedy
best-first walk with the same termination criterion).
"""
def __init__(self, dims: int, k: int = 200, ef_search: int = 64) -> None:
self.dims = int(dims)
self.k_build = int(k)
self.ef_search = int(ef_search)
self._rust_idx: Any = None
self._np_idx: _NsgNumpyFallback | None = None
self._ids: list[str] = []
self._vectors: np.ndarray | None = None
self._mode: str = "uninitialized"
def _use_rust(self) -> bool:
return NSG_RUST_AVAILABLE and NSG_RUST_ENABLED
def build(self, vectors: np.ndarray, ids: list[str] | None = None) -> None:
v = np.ascontiguousarray(vectors, dtype=np.float32)
if v.ndim != 2:
raise ValueError(f"vectors must be 2-D, got shape {v.shape}")
n, d = v.shape
if d != self.dims:
raise ValueError(
f"vectors have dims={d}, expected {self.dims}")
if n == 0:
raise ValueError("empty corpus")
self._vectors = v.copy()
if ids is None:
self._ids = [str(i) for i in range(n)]
else:
if len(ids) != n:
raise ValueError(f"len(ids)={len(ids)} != n_vectors={n}")
self._ids = list(ids)
if self._use_rust():
# The Rust wheel owns the vectors and adjacency lists; we
# keep a reference to the numpy storage so `search()` can
# resolve ids when callers pass through the Python surface.
self._rust_idx = _quadrant.NsgIndex.build(v, self.k_build)
self._np_idx = None
self._mode = "rust"
else: # pragma: no cover
self._rust_idx = None
self._np_idx = _NsgNumpyFallback(k=self.k_build, ef_search=self.ef_search)
self._np_idx.build(v, self._ids)
self._mode = "numpy-fallback"
def search(self, query: np.ndarray, k: int = 10) -> list[tuple[str, float]]:
if self._mode == "uninitialized":
raise RuntimeError("NsgBackend.build() must be called first")
q = np.ascontiguousarray(query, dtype=np.float32).ravel()
if q.shape[0] != self.dims:
raise ValueError(
f"query len {q.shape[0]} != dims {self.dims}")
if self._mode == "rust":
ids_u32, sims = self._rust_idx.search(q, k, self.ef_search)
return [(self._ids[int(i)], float(s))
for i, s in zip(ids_u32, sims)]
return self._np_idx.search(q, k) # pragma: no cover
def stats(self) -> dict:
if self._mode == "uninitialized":
return {"mode": "uninitialized", "dims": self.dims,
"k_build": self.k_build, "ef_search": self.ef_search}
if self._mode == "rust":
return {
"mode": "rust",
"dims": self.dims,
"k_build": self.k_build,
"ef_search": self.ef_search,
"detail": self._rust_idx.stats(),
"n_vectors": int(self._rust_idx.n_vectors()),
"n_edges": int(self._rust_idx.n_edges()),
"nav_node": int(self._rust_idx.nav_node()),
}
return self._np_idx.stats() # pragma: no cover
def __repr__(self) -> str:
return (f"NsgBackend(dims={self.dims}, k={self.k_build}, "
f"ef_search={self.ef_search}, mode={self._mode})")
def nsg_status() -> dict[str, Any]:
"""Report whether the Rust NSG wheel is wired and active."""
return {
"rust_available": bool(NSG_RUST_AVAILABLE),
"rust_enabled": bool(NSG_RUST_ENABLED),
"env": _env,
"hint": ("" if NSG_RUST_AVAILABLE else
"build with: pip install ./rust/quadrant "
"(requires cargo + maturin) — numpy fallback active"),
}