大规模向量检索的分片策略：如何把亿级索引拆成可管理的子索引

大规模向量检索的分片策略如何把亿级索引拆成可管理的子索引一、深度引言与场景痛点当你的向量索引从百万级增长到亿级时单索引架构的第一个裂缝就会出现在查询延迟上。HNSW 索引构建需要 O(N log N) 时间查询需要 O(log N) 次距离计算——当 N 从 100 万变成 1 亿这个对数增长在绝对数值上也从 ~6 变成了 ~8.5。看起来差异不大但内存中 1 亿条 768 维向量约需 300GB单机的物理内存和 GPU 显存都扛不住。这就是向量索引单机墙——无论算法多高效当数据量突破单机内存/显存上限时你只有两条路要么堆更贵的硬件要么把数据拆开。堆硬件的天花板很快会到云服务器单机最大内存也就 TB 级而数据分片是唯一可无限横向扩展的方案。但分片不是把数据随便切成几份就完事了。核心挑战在于查询时需要检索所有分片然后合并结果所以分片数量直接决定了查询延迟的下限。分片策略需要在怎么切减少跨分片查询和切多少控制查询延迟之间找到最优平衡。二、底层机制与原理深度剖析向量索引的分片策略有三种主流方案随机分片Random最简单的方案按向量 ID 的哈希值分配。优点是负载绝对均衡缺点是完全忽略语义——每个查询都要扫全部分片没有任何剪枝空间。空间分片Spatial先对全量向量做 K-Means 聚类每个簇作为一个分片。查询时先判断查询向量属于哪个簇优先检索该分片。语义相近的数据放一起的代价是数据分布不均匀——热门领域的簇会比冷门领域大得多。分层分片Hierarchical两层架构——第一层是粗粒度的路由层由 K-Means 聚类中心组成数量少如 1000 个第二层是每个路由指向的实际分片。查询时先用路由层判断查询属于哪些簇通常 Top-3只检索这些簇对应的分片。flowchart TB subgraph 分层分片架构 Q[查询向量] -- ROUTE[路由层\nK-Means 聚类中心 (1000个)] ROUTE --|Top-3 相似簇| C1[簇 1] ROUTE --|Top-3 相似簇| C3[簇 3] ROUTE --|Top-3 相似簇| C7[簇 7] C1 -- S1[分片 1\nHNSW 索引\n~10万条] C3 -- S3[分片 3\nHNSW 索引\n~15万条] C7 -- S7[分片 7\nHNSW 索引\n~8万条] S1 --|局部 Top-K| MERGE[结果合并\n全局 Top-K 重排序] S3 -- MERGE S7 -- MERGE MERGE -- RESULT[最终结果] end subgraph 分片策略对比 RAND[随机分片\n• 绝对均衡\n• 全分片检索\n• 延迟 N × 单分片延迟] SPAT[空间分片\n• 语义聚合\n• 可能剪枝\n• 数据可能不均衡] HIER[分层分片\n• 路由剪枝\n• 延迟 K × 单分片延迟\n• 路由层有额外开销] end style ROUTE fill:#4A90D9,color:#fff style MERGE fill:#E8A838,color:#fff style RESULT fill:#D9534F,color:#fff style HIER fill:#5CB85C,color:#fff分层分片是亿级索引的最优方案。路由层通常只有几千个聚类中心可全部放内存剪枝后只检索 2-5 个分片将有效检索量降低了 99%。代价是召回率可能略微下降路由阶段漏掉的分片就找不到了。三、生产级代码实现下面实现一个分层分片的向量检索系统包括 K-Means 路由层构建、分片检索和结果合并。import heapq import logging import time from dataclasses import dataclass, field from typing import Any import numpy as np logger logging.getLogger(__name__) dataclass class ShardConfig: 分片配置。 num_routing_centroids: int 1024 # 路由层聚类中心数 max_shard_size: int 100_000 # 单分片最大向量数 top_k_routing: int 3 # 路由层返回的 Top-K 分片数 search_top_k_per_shard: int 20 # 每个分片检索的 Top-K dataclass class SearchResult: 检索结果。 vector_id: str score: float shard_id: int -1 metadata: dict | None None class HierarchicalVectorIndex: 分层分片向量索引。 两层架构 - 路由层K-Means 聚类中心用于快速剪枝 - 分片层每个分片是一个独立的向量索引简化版用 numpy 实现 def __init__(self, config: ShardConfig | None None): self.cfg config or ShardConfig() self._centroids: np.ndarray | None None self._shards: dict[int, dict[str, np.ndarray]] {} self._shard_metadata: dict[str, dict] {} self._total_vectors: int 0 def build(self, vectors: dict[str, np.ndarray]): 构建分层索引。 Args: vectors: {vector_id: numpy_array} 的全量向量字典 if len(vectors) self.cfg.num_routing_centroids: raise ValueError( f向量数量 ({len(vectors)}) 少于聚类中心数 f({self.cfg.num_routing_centroids}) ) logger.info(开始构建分层索引总向量数%d, len(vectors)) # 1. 训练 K-Means 路由层 all_ids list(vectors.keys()) all_vecs np.array([vectors[vid] for vid in all_ids], dtypenp.float32) self._train_routing_layer(all_vecs) # 2. 分配向量到分片 shard_assignments self._assign_to_shards(all_vecs, all_ids) # 3. 构建每个分片的索引 for shard_id, (ids, vecs) in shard_assignments.items(): self._shards[shard_id] { id_: vec for id_, vec in zip(ids, vecs) } self._total_vectors len(vectors) logger.info( 索引构建完成: %d 个分片, 每个分片 ~%d 条向量, len(self._shards), self._total_vectors // max(len(self._shards), 1), ) def _train_routing_layer(self, vectors: np.ndarray): 训练 K-Means 聚类中心简化版生产环境用 faiss/sklearn。 n len(vectors) k min(self.cfg.num_routing_centroids, n) # 简化 K-Means随机初始化 少量迭代 rng np.random.default_rng(42) indices rng.choice(n, k, replaceFalse) centroids vectors[indices].copy() for iteration in range(10): # 分配 distances self._compute_distances(vectors, centroids) labels np.argmin(distances, axis1) # 更新 new_centroids np.zeros_like(centroids) for j in range(k): mask labels j if mask.sum() 0: new_centroids[j] vectors[mask].mean(axis0) else: new_centroids[j] centroids[j] # 保持旧中心 shift np.linalg.norm(new_centroids - centroids) centroids new_centroids if shift 1e-4: break self._centroids centroids.astype(np.float32) logger.info(路由层训练完成: %d 个聚类中心, k) staticmethod def _compute_distances(a: np.ndarray, b: np.ndarray) - np.ndarray: 批量计算余弦距离矩阵。 a_norm a / (np.linalg.norm(a, axis1, keepdimsTrue) 1e-8) b_norm b / (np.linalg.norm(b, axis1, keepdimsTrue) 1e-8) return 1.0 - np.dot(a_norm, b_norm.T) def _assign_to_shards( self, vectors: np.ndarray, ids: list[str], ) - dict[int, tuple[list[str], np.ndarray]]: 分配向量到分片。 if self._centroids is None: raise ValueError(请先调用 _train_routing_layer()) distances self._compute_distances(vectors, self._centroids) labels np.argmin(distances, axis1) shard_data: dict[int, tuple[list[str], list[np.ndarray]]] {} for i, (vector_id, label) in enumerate(zip(ids, labels)): shard_id int(label) if shard_id not in shard_data: shard_data[shard_id] ([], []) shard_data[shard_id][0].append(vector_id) shard_data[shard_id][1].append(vectors[i]) return { shard_id: (ids_list, np.array(vecs_list, dtypenp.float32)) for shard_id, (ids_list, vecs_list) in shard_data.items() } def search( self, query_vector: np.ndarray, top_k: int 10, use_routing: bool True, ) - list[SearchResult]: 分层检索。 Args: query_vector: 查询向量 top_k: 返回结果数量 use_routing: 是否启用路由层剪枝关闭则全分片检索 if self._centroids is None or not self._shards: return [] query np.array(query_vector, dtypenp.float32).reshape(1, -1) if use_routing: # 路由选择 Top-N 分片 routing_scores 1.0 - self._compute_distances( query, self._centroids, )[0] top_shards np.argsort(routing_scores)[::-1][:self.cfg.top_k_routing] logger.debug( 路由选择 %d 个分片: %s, len(top_shards), top_shards.tolist(), ) else: top_shards list(self._shards.keys()) # 在各分片中检索 all_results: list[SearchResult] [] for shard_id in top_shards: shard self._shards.get(int(shard_id)) if shard is None: continue shard_results self._search_shard( query[0], shard, shard_id, self.cfg.search_top_k_per_shard, ) all_results.extend(shard_results) # 全局 Top-K 合并 all_results.sort(keylambda x: x.score, reverseTrue) return all_results[:top_k] def _search_shard( self, query: np.ndarray, shard: dict[str, np.ndarray], shard_id: int, top_k: int, ) - list[SearchResult]: 在单个分片中检索。 ids list(shard.keys()) vecs np.array([shard[vid] for vid in ids], dtypenp.float32) similarities np.dot(vecs, query) / ( np.linalg.norm(vecs, axis1) * np.linalg.norm(query) 1e-8 ) # Top-K简化版实际用 HNSW/IVF top_indices np.argsort(similarities)[-top_k:][::-1] return [ SearchResult( vector_idids[idx], scorefloat(similarities[idx]), shard_idshard_id, ) for idx in top_indices ] def insert(self, vector_id: str, vector: np.ndarray): 增量插入向量。 if self._centroids is None: raise ValueError(索引尚未构建) # 找到最近的路由中心 query np.array(vector, dtypenp.float32).reshape(1, -1) distances self._compute_distances(query, self._centroids)[0] shard_id int(np.argmin(distances)) if shard_id not in self._shards: self._shards[shard_id] {} self._shards[shard_id][vector_id] vector.astype(np.float32) self._total_vectors 1 # 检查分片是否需要分裂 if len(self._shards[shard_id]) self.cfg.max_shard_size * 1.2: logger.warning( 分片 %d 超过容量上限 (%d/%d)建议触发重平衡, shard_id, len(self._shards[shard_id]), self.cfg.max_shard_size, ) def get_stats(self) - dict: 获取索引统计。 shard_sizes {sid: len(shard) for sid, shard in self._shards.items()} if shard_sizes: return { total_vectors: self._total_vectors, num_shards: len(self._shards), num_centroids: len(self._centroids) if self._centroids is not None else 0, avg_shard_size: round(np.mean(list(shard_sizes.values())), 0), max_shard_size: max(shard_sizes.values()), min_shard_size: min(shard_sizes.values()), imbalance_ratio: round( max(shard_sizes.values()) / max(min(shard_sizes.values()), 1), 2, ), } return {total_vectors: 0, num_shards: 0} def demo(): import tempfile config ShardConfig( num_routing_centroids16, max_shard_size10_000, top_k_routing3, search_top_k_per_shard20, ) index HierarchicalVectorIndex(config) # 生成测试数据 rng np.random.default_rng(42) vectors { fvec_{i}: rng.random(128).astype(np.float32) * rng.choice([1, 5, 10]) for i in range(10_000) } # 构建索引 index.build(vectors) print(索引统计:, index.get_stats()) # 检索 query rng.random(128).astype(np.float32) # 分层检索路由剪枝 t0 time.monotonic() results_routed index.search(query, top_k10, use_routingTrue) t_routed (time.monotonic() - t0) * 1000 # 全量检索无路由剪枝作为对比 t0 time.monotonic() results_full index.search(query, top_k10, use_routingFalse) t_full (time.monotonic() - t0) * 1000 print(f\n分层检索: {t_routed:.1f}ms, 召回{len(results_routed)}条) for r in results_routed[:3]: print(f {r.vector_id}: {r.score:.4f} (shard{r.shard_id})) print(f\n全量检索: {t_full:.1f}ms) print(f加速比: {t_full/t_routed:.1f}x) # 检查重叠度 routed_ids {r.vector_id for r in results_routed} full_ids {r.vector_id for r in results_full} overlap routed_ids full_ids print(f\n分层检索与全量检索的重叠度: {len(overlap)}/{len(full_ids)} ({len(overlap)/len(full_ids)*100:.1f}%)) if __name__ __main__: demo()四、边界分析与架构权衡路由层精度与召回率的权衡路由层只检索 Top-3 分片理论上可能丢失正确答案在第 4 个分片的情况。top_k_routing越大召回率越高但延迟也越高。建议从top_k_routing3开始监控召回率必要时增加到 5。数据分布不均的对策K-Means 聚类天然会导致热门簇比冷门簇大很多。解决策略(1) 用 K-Means 的变体如 Constrained K-Means限制每簇大小(2) 对超大的分片做二次分裂。分片数量与查询延迟分片数量增加 路由层剪枝后仍需检索的分片数增加。当分片总数从 100 增长到 10,000 时即使路由剪枝 99%仍需检索 100 个分片。这是分层分片的局限——规模达到十亿级时需要三层甚至四层分片。增量更新的挑战向量索引的增量更新插入/删除单条向量在分片架构中比单索引更复杂。插入需要路由到正确的分片并更新该分片的索引结构如 HNSW 图删除需要更新所有分片。建议用增量 定期全量重建的策略。五、总结分层分片是突破向量索引单机墙的成熟方案。路由层做粗筛剪枝 99% 的不相关分片分片层做精排——这套两层架构把亿级检索的有效计算量控制在了百万级。核心工程决策是路由层的聚类中心数量通常 500-5000、Top-K 路由参数通常 3-5、以及分片的大小上限。三者的选择需要在召回率、延迟和存储成本之间做 trade-off没有万能的最优值需要在你的数据和查询分布上做 benchmark 才能确定。