查询一次的端到端:HNSW 遍历 + ZMQ 按需重算
30 秒导读: 这是全库工程含量最高的一章。上一章 01-storage-trick.md 讲了 LEANN 为省 97% 空间不存 embedding、只存图;省下来的代价,全在"查询时怎么办"这里还。本章端到端追一条 query:字符串 → query embedding → HNSW 图遍历 → 遍历中每碰到一个陌生节点,就把它的 id 通过 ZMQ 发回一个 Python embedding server,现场把原文重新编码成向量、算出距离、回传 → 整数 label 映射回 passage id →
SearchResult。
1. 这是什么(先建直觉)
一句话: 普通向量库搜索时,每个候选点的向量早就存在内存/磁盘里,直接算距离即可;LEANN 把这些向量删了,所以搜索走到某个点、需要它的向量时,得临时把这个点对应的那段原文重新喂给 embedding 模型,现场算出来。
这个"现场算"就叫 recompute(按需重算)。它由两个进程配合完成:
| 角色 | 进程 | 干什么 |
|---|---|---|
| 搜索引擎 | C++ faiss(mmap 进 Python) | 跑 HNSW 图遍历,决定"下一步要哪些节点的距离" |
| 重算服务 | Python embedding server(子进程) | 收到节点 id → 查原文 → 调模型算 embedding → 算距离 → 回传 |
两者之间用 ZMQ(ZeroMQ,一个轻量消息队列库)+ msgpack(二进制序列化)通信。
为什么值得单开一章? 因为这里有一个反直觉的架构:搜索的"内循环"跨了进程边界和语言边界。C++ 每走一步图,就得同步等 Python 把 embedding 算完。理解这条链路,才算真懂 LEANN 拿什么换了空间。
2. 顶层全景(一条 query 怎么流动)
先看整条链路。三个进程:发起搜索的主进程、跑图遍历的 C++ 索引(在主进程内,但是原生代码)、独立的 embedding server 子进程。
怎么读这张图: 从上到下是时间顺序;中间那段"遍历循环"里的 ①②③ 会反复来回几十次——图每扩展一层就来回一次,这是整个搜索最耗时的地方。
主进程 (Python) C++ faiss 索引 embedding server 子进程
─────────────── ───────────── ──────────────────────
LeannSearcher.search
│ query 字符串
│
├─(A) _ensure_server_running ──────────────────────────────▶ 启动/复用,拿到 zmq_port
│
├─(B) compute_query_embedding
│ "把 query 本身编码" ────── ZMQ 文本请求 ─────────────▶ _handle_text_embedding
│ query_embedding (1,D) ◀───────────────────────────── 返回向量
│
├─(C) backend_impl.search
│ 设 SearchParametersHNSW
│ set_zmq_port(zmq_port)
│ self._index.search(...) ──▶ ┌─────────────────────┐
│ │ HNSW 图遍历循环 │
│ │ │
│ │ ① 挑出要距离的 │
│ │ node_ids │
│ │ │ │
│ │ ② 发距离请求 ───────┼── ZMQ ──▶ _handle_distance_request
│ │ [ids, qvec] │ 取原文→compute_embeddings
│ │ │ →算 -dot / L2 →打包
│ │ ③ 收距离 ◀──────────┼─────────── 返回 [distances]
│ │ │ │
│ │ 用距离更新候选堆, │
│ │ 没收敛就回 ① 循环 │
│ └─────────┬───────────┘
│ 整数 labels + distances ◀────────────┘
│
├─(D) _id_map[label] → passage id 字 符串
│
└─(E) passage_manager.get_passage(id) → 组装 SearchResult 列表
各步落到源码:
| 步 | 干什么 | 源码位置 · 符号 |
|---|---|---|
| A | 起/复用 embedding server,拿 zmq_port | packages/leann-core/src/leann/api.py:1347-1356 LeannSearcher.search → _ensure_server_running |
| B | 把 query 字符串编码成向量 | api.py:1375-1380 compute_query_embedding |
| C | 设参数、set_zmq_port、调 C++ | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py:226-271 HNSWSearcher.search |
| ②③ | 遍历中按需重算距离 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py:211-271 _handle_distance_request |
| D | 整数 label → passage id | hnsw_backend.py:274-287 map_label / _id_map |
| E | id → 原文,组装 SearchResult | api.py:1437-1453 |
下面逐段拆。
3. (A)(B) 起服务 + 编码 query —— LeannSearcher.search
入口是 LeannSearcher.search(api.py:1240)。它先做一堆通用处理(BM25 纯关键词分支、混合检索权重、top_k 自动收缩到文档总数),向量检索的核心是这三步。
3.1 只有需要重算时才起 server
不是所有索引都要 recompute——如果索引是"非压缩"存了完整向量的,就不必起 embedding server。所以起 server 是条件式的:
# 示意,非源码:重点看「只有 effective_recompute 才起 server」
if effective_recompute: # 这个索引把向量删了,得现场重算
zmq_port = backend_impl._ensure_server_running(
meta_path_str, port=expected_zmq_port, ...
) # 拿到 server 监听的端口
真实实现见 api.py:1347-1356。effective_recompute 的来源:构造 LeannSearcher 时的配置(self.recompute_embeddings),search(recompute_embeddings=...) 这个每次调用的覆盖是 deprecated 的(api.py:1336-1344 会打警告)。
_ensure_server_running 在 searcher_base.py:62-102(BaseSearcher._ensure_server_running):它把 embedding_model、embedding_mode、distance_metric、passages_file(那个 .meta.json)交给 EmbeddingServerManager.start_server,后者负责起子进程或复用已在跑的,返回真正监听的端口 actual_port。一个细节:它会从 embedding_options 里滤掉所有 prompt 模板(searcher_base.py:82-86),因为模板已经在 compute_query_embedding 里贴到 query 上了,server 端再贴一次就双重模板了。
3.2 编码 query 本身
query 字符串要先变成向量,才能作为图遍历的"目标点"。这一步也走 embedding server(如果 recompute 开着):
# 示意,非源码
query_embedding = backend_impl.compute_query_embedding(
query,
use_server_if_available=effective_recompute, # 复用同一个 server
zmq_port=zmq_port,
query_template=query_template, # 某些模型要给 query 贴前缀
) # 得到 (1, D)
真实:api.py:1375-1380。落到 searcher_base.py:104-162 compute_query_embedding:
- 若有
query_template,先query = f"{query_template}{query}"(searcher_base.py:125-126)——例如 instruction 模型要求 query 带 "Represent this sentence:" 之类前缀。 - 若
use_server_if_available,调_compute_embedding_via_server([query], zmq_port)走 ZMQ(searcher_base.py:146)。 - server 挂了就 fallback:直接在主进程
import compute_embeddings本地算(searcher_base.py:149-162)。这是一处刻意的容错。
ZMQ 请求协议(_compute_embedding_via_server,searcher_base.py:164-194):一个 REQ socket,把 [query字符串] 用 msgpack 打包发过去,30 秒超时等回一个 [[float...]],转成 numpy。server 端识别"一个全是字符串的 list"就走 _handle_text_embedding(下面 §5)。
4. (C) HNSW 遍历 —— HNSWSearcher.search
拿到 query 向量后,交给 C++ 图搜索。Python 侧这层(hnsw_backend.py:174-289)主要是翻译参数和善后 label,真正的图遍历在 C++ faiss 里。
4.1 参数怎么翻译
LEANN 对外的四个搜索旋钮,映射到 faiss 的 SearchParametersHNSW:
| LEANN 参数 | 映射到 | 含义 · 源码 |
|---|---|---|
complexity | params.efSearch | 候选表大小,越大越准越慢 · hnsw_backend.py:229 |
beam_width | params.beam_size | 每轮并行扩展的路径数 · hnsw_backend.py:230 |
prune_ratio | params.pq_pruning_ratio | 用 PQ 近似距离预剪掉多少比例邻居 · hnsw_backend.py:243 |
pruning_strategy | local_prune + send_neigh_times_ratio | 剪枝候选怎么选,见下 · hnsw_backend.py:246-254 |
pruning_strategy 是三选一,映射不直观,单列一张表:
| 策略 | local_prune | send_neigh_times_ratio | 直觉 |
|---|---|---|---|
global(默认) | False | 0.0 | 用全局 PQ 队列大小选候选 |
local | True | 0.0 | 局部剪枝:排序后取最优候选 |
proportional | False | 1.0(任意 >1e-6) | 按"新邻居数量比例"决定发多少 |
见 hnsw_backend.py:246-254。还有一个只对 OpenAI cosine 模型的特判:关掉 check_relative_distance(hnsw_backend.py:235-240),避免所有分数挤在窄区间时提前终止遍历。
4.2 把 zmq_port 交给 C++
关键一句:C++ 遍历时要向 server 发请求,得知道端口。所以有两处都塞了 port——一处是对象方法 set_zmq_port,一处是搜索参数字段:
# 示意,非源码
if recompute_embeddings:
self._index.set_zmq_port(zmq_port) # 告诉 C++ 索引往哪发重算请求
params.zmq_port = zmq_port # 参数里也带一份
self._index.search(nq, query_ptr, top_k, dist_ptr, label_ptr, params)
真实:set_zmq_port 在 hnsw_backend.py:218-219,params.zmq_port 在 hnsw_backend.py:227-228,调 C++ self._index.search 在 hnsw_backend.py:264-271。从这一句 search 返回前,C++ 内部已经和 embedding server 来回了几十次——这些来回对 Python 是透明的,Python 只在最后拿到填好的 distances 和 labels 数组。
4.3 整数 label 映射回 passage id
C++ 只认整数节点号(0,1,2,…)。用户要的是 passage id(可能是文件路径+chunk 号那种字符串)。中间隔着一张 _id_map(从 <index>.ids.txt 每行一个 id 读进来,hnsw_backend.py:164-172):
# 示意,非源码
def map_label(x):
if 0 <= x < len(self._id_map):
return self._id_map[x] # 整数下标 → 原始 passage id 字符串
return str(x) # 越界就退回字符串数字
真实:hnsw_backend.py:274-287。没有 id_map 时直接把整数转字符串。返回 {"labels": [[str...]], "distances": ndarray}(hnsw_backend.py:289)。
5. 核心机制:遍历中怎么按需重算(embedding server 三类请求)
这是本章的心脏。C++ 遍历到需要距离的节点时,并不自己算——它没有向量。它把 node_ids + query 向量打包成 ZMQ 消息发给 Python server。server 端有一个 REP socket 循环,按消息形状分发到三个 handler。
5.1 服务主循环:一个带优雅关闭的 REP 循环
zmq_server_thread_with_shutdown(hnsw_embedding_server.py:161)绑定 tcp://*:{zmq_port},设了 1 秒收发超时(为了能定期检查 shutdown 事件),然后死循环 recv → 按形状分发 → send。
请求怎么区分? 靠 msgpack 解出来的结构形状(hnsw_embedding_server.py:371-396):
| 收到的形状 | 判为 | handler |
|---|---|---|
["__QUERY_MODEL__"] | 问模型名 | 直接回 [model_name] |
["str", "str", ...] 全字符串 | 文本编码 | _handle_text_embedding |
[[ids], [qvec]] 两个 list | 算距离 | _handle_distance_request |
| 其它 | 按 id 取 embedding | _handle_embedding_by_id |
REP socket 的铁律是收一个必须回一个,所以解包失败/handler 抛异常时,会送一个"安全兜底"响应(_build_safe_fallback,hnsw_embedding_server.py:180-193,例如距离请求就回一堆 1e9 的大距离),否则 REQ 端会永久卡死。
5.2 距离请求(遍历内循环的主力)—— _handle_distance_request
这是遍历中被调用最频繁的路径(hnsw_embedding_server.py:211-271)。核心四步:
# 示意,非源码:遍历中一次距离请求
node_ids = request[0] # C++ 要这些节点的距离
query_vec = np.array(request[1]) # 相对哪个 query
texts = [passages.get_passage(_map_node_id(nid))["text"] # ① 节点 id → 原文
for nid in node_ids]
embeddings = compute_embeddings(texts, model_name, ...) # ② 原文 → 现场重算向量
if distance_metric == "l2":
dist = np.sum((embeddings - query_vec)**2, axis=1) # ③a L2 距离
else:
dist = -np.dot(embeddings, query_vec) # ③b MIPS:负点积(越小越近)
rep_socket.send(msgpack.packb([dist.tolist()])) # ④ 回传
对应真实行:取原文 hnsw_embedding_server.py:228-241;重算 :248-253;算距离 :257-262(注意 MIPS 用 -dot 把"越大越相似"翻成"越小越近",配合 faiss 的最小堆);回传 :269。
几个务实的健壮性细节:
- 取原文先经
_map_node_id(:151-159)做整数→passage id 映射,和 §4.3 是同一张 id_map,只是在 server 侧独立加载(:132-149)。 - 某个节点原文缺失/查不到,不整批失败:该位置填哨兵大距离
1e9,只把找得到的算出来填回原位(found_indices机制,:227-265)。
5.3 另外两类请求
| handler | 触发 | 干什么 · 源码 |
|---|---|---|
_handle_text_embedding | query 编码(§3.2)、纯文本编码 | 直接 compute_embeddings(request) 回 [[float...]] · hnsw_embedding_server.py:195-209 |
_handle_embedding_by_id | C++ 要的是 embedding 本体而非距离 | 按 id 取原文→算向量→扁平化成 [[batch,dim],[flat floats]] 回传;有 NaN/Inf 检测 · :273-346 |
_handle_embedding_by_id 返回的是 [dims, flat_data] 两段:第一段是形状 [batch, dim],第二段是拉平的浮点数组,C++ 侧按 dim 切回矩阵。它还有个防呆:只把 found_indices 对应的槽位填上真向量,查不到的保持 0(:333-339)。
6. 为什么"pruned 索引"必须 recompute(不能关)
有一条硬约束值得单独讲:如果索引是 pruned/compact(把向量删了的),recompute_embeddings=False 会直接抛错,不给你降级。
# 真实:hnsw_backend.py:210-214
if not recompute_embeddings and self.is_pruned:
raise RuntimeError(
"Recompute is required for pruned/compact HNSW index. "
"Re-run search with --recompute, or rebuild with --no-recompute and --no-compact."
)
为什么不能软着陆? 因为 pruned 索引里根本没有向量数据可读——它只有图结构。不 recompute 就没有任何距离可算,搜索无从进行。这不是"慢一点"的问题,是"缺了必需的数据"。所以 LEANN 选择 fail fast,而不是返回错误结果。
is_pruned 从哪来?构造 HNSWSearcher 时,先看 meta 里显式的 is_pruned,否则默认取 is_recompute/is_compact(hnsw_backend.py:147-150)。它同时决定加载索引时的 hnsw_config.is_recompute(hnsw_backend.py:158-162,注释坦白说这个 C++ 字段"only for loading IIUC")。
对照关系:
| 索引怎么建的 | is_pruned | 能否 recompute=False |
|---|---|---|
| pruned / compact(默认,省空间) | True | 不能,强制重算 |
--no-recompute --no-compact(存全量向量) | False | 可以,直接读向量 |
这正是 01-storage-trick.md "省 97% 空间"的另一面:省下来的空间,在查询时用每个节点一次现场编码换回来。想不付这个查询代价,就得付回存储代价——二选一,LEANN 只是让默认选择偏向省空间。
7. (E) 组装 SearchResult
回到主进程(api.py:1437-1453)。C++ 给的是 (passage_id, distance) 对,用户要的是带原文和 metadata 的结构。逐个 id 去 passage_manager.get_passage(id) 捞原文:
# 示意,非源码
for string_id, dist in zip(labels[0], distances[0]):
p = passage_manager.get_passage(string_id) # id → {text, metadata}
results.append(SearchResult(
id=string_id, score=float(dist),
text=p["text"], metadata=p.get("metadata", {}),
))
真实:api.py:1441-1453。查不到的 id 记 error 但不中断(api.py:1466-1471)。之后按需过 metadata_filters(api.py:1473-1478)、写 query log(api.py:1485-1486),返回 list[SearchResult]。至此一条 query 走完全程。
8. 边界与坑
- 同步阻塞的内循环。 C++ 每轮遍历都要同步等 Python 把 embedding 算完再继续。embedding 计算(尤其无 GPU 时)是主要延迟来源;
complexity/beam_width越大,来回次数越多。这解释了为什么 LEANN 强调 warmup(api.py:1230warmup预热 embedding 路径)。 - REP/REQ 必须严格一收一回。 server 任何异常都得回个兜底响应(§5.1),否则客户端 30 秒超时前一直卡。这是把容错逻辑写满 handler 的原因。
- server 挂了只有 query 编码能 fallback。
compute_query_embedding能退回本地算(searcher_base.py:149-162);但遍历内循环的距离请求没有本地 fallback——server 一旦不可用,搜索就废了。 - daemon TTL。 server 可作 daemon 常驻,空闲超
daemon_ttl秒自动退出(hnsw_embedding_server.py:474-479),避免每次查询都付冷启动。
其它后端(IVF/DiskANN)对 recompute 的处理不同,见 03-backends-registry.md;embedding 到底怎么算(sentence-transformers / MLX / OpenAI / Ollama 多后端)见 04-embedding-and-chunking.md。
9. 代码地图(导航索引)
| 主题 | 文件路径 | 符号 |
|---|---|---|
| 搜索总入口 · 起 server · 编码 query | packages/leann-core/src/leann/api.py | LeannSearcher.search |
| SearchResult 组装 | packages/leann-core/src/leann/api.py | (api.py:1437-1453,SearchResult) |
| 起/复用 embedding server | packages/leann-core/src/leann/searcher_base.py | BaseSearcher._ensure_server_running |
| query 编码 + server fallback | packages/leann-core/src/leann/searcher_base.py | compute_query_embedding |
| ZMQ REQ 协议(客户端) | packages/leann-core/src/leann/searcher_base.py | _compute_embedding_via_server |
| HNSW 参数映射 + 调 C++ | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py | HNSWSearcher.search |
| pruned 强制重算约束 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py | hnsw_backend.py:210-214(RuntimeError) |
| 整数 label → passage id | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py | map_label / _id_map |
| ZMQ REP 主循环 + 分发 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py | zmq_server_thread_with_shutdown |
| 距离请求(遍历内循环主力) | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py | _handle_distance_request |
| 文本编码请求 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py | _handle_text_embedding |
| 按 id 取 embedding | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py | _handle_embedding_by_id |
| server 侧 id 映射 | packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py | _map_node_id |