3.4 向量索引与检索优化（上）

向量索引和检索优化是RAG系统性能提升的关键技术。高效的索引策略和优化算法能够显著提升检索速度和准确性。

向量索引概述

基本概念

向量索引（Vector Indexing）是用于快速检索相似向量的数据结构，主要功能包括：

• 快速检索：在海量向量中快速找到相似向量
• 降维加速：通过降维技术提升检索效率
• 内存优化：减少内存占用和I/O开销
• 质量保证：确保检索结果的准确性

索引类型

主流索引算法

1. 精确索引算法

暴力搜索（Brute Force）

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

def brute_force_search(query_vector: np.ndarray, index_vectors: np.ndarray, top_k: int = 5):
    """暴力搜索"""
    # 计算所有向量的相似度
    similarities = cosine_similarity([query_vector], index_vectors)[0]
    
    # 获取top-k结果
    top_indices = np.argsort(similarities)[-top_k:][::-1]
    top_scores = similarities[top_indices]
    
    return list(zip(top_indices, top_scores))
​

特点：

• 简单易实现
• 准确性最高
• 时间复杂度O(n)
• 适合小规模数据

适用场景：

• 数据量小（<10k）
• 对准确性要求极高
• 毫秒级响应可接受

空间分区（Space Partitioning）

from sklearn.cluster import KMeans

class SpacePartitionIndex:
    """空间分区索引"""
    
    def __init__(self, n_clusters: int = 10):
        self.n_clusters = n_clusters
        self.kmeans = KMeans(n_clusters=n_clusters, random_state=42)
        self.index_vectors = None
        self.cluster_centers = None
    
    def fit(self, vectors: np.ndarray):
        """构建索引"""
        self.index_vectors = vectors
        self.kmeans.fit(vectors)
        self.cluster_centers = self.kmeans.cluster_centers_
    
    def search(self, query_vector: np.ndarray, top_k: int = 5):
        """搜索"""
        # 1. 找到最近的簇
        cluster_distances = np.linalg.norm(self.cluster_centers - query_vector, axis=1)
        nearest_clusters = np.argsort(cluster_distances)[:3]  # 搜索最近的3个簇
        
        # 2. 在这些簇中搜索
        candidates = []
        for cluster_id in nearest_clusters:
            cluster_indices = np.where(self.kmeans.labels_ == cluster_id)[0]
            cluster_vectors = self.index_vectors[cluster_indices]
            
            similarities = cosine_similarity([query_vector], cluster_vectors)[0]
            top_cluster_indices = np.argsort(similarities)[-top_k:][::-1]
            
            for idx in top_cluster_indices:
                original_idx = cluster_indices[idx]
                candidates.append((original_idx, similarities[idx]))
        
        # 3. 排序并返回top-k
        candidates.sort(key=lambda x: x[1], reverse=True)
        return candidates[:top_k]
​

特点：

• 基于聚类分区的搜索
• 准确性较好
• 速度优于暴力搜索
• 需要额外的聚类计算

2. 近似索引算法

HNSW（Hierarchical Navigable Small World）

import numpy as np
from sklearn.neighbors import NearestNeighbors

class HNSWIndex:
    """HNSW索引实现"""
    
    def __init__(self, M: int = 16, ef: int = 32):
        self.M = M  # 每个节点的最大连接数
        self.ef = ef  # 搜索时的候选节点数
        self.layers = []
        self.max_level = 0
        self.index_vectors = None
    
    def fit(self, vectors: np.ndarray):
        """构建HNSW索引"""
        self.index_vectors = vectors
        n_samples = len(vectors)
        
        # 1. 确定层数
        self.max_level = int(np.log2(n_samples)) - 1
        
        # 2. 构建层次结构
        for level in range(self.max_level, -1, -1):
            layer = self._build_layer(vectors, level)
            self.layers.append(layer)
    
    def _build_layer(self, vectors: np.ndarray, level: int):
        """构建单层"""
        layer = []
        n_samples = len(vectors)
        
        # 第一层：随机选择初始点
        if level == 0:
            entry_points = np.random.choice(n_samples, min(self.M * 2, n_samples), replace=False)
            for point in entry_points:
                neighbors = self._find_neighbors(vectors, point, min(self.M, n_samples - 1))
                layer.append({
                    'point': point,
                    'neighbors': neighbors.tolist()
                })
        else:
            # 其他层：使用下层的结果
            for point in range(n_samples):
                neighbors = self._find_neighbors(vectors, point, min(self.M, n_samples - 1))
                layer.append({
                    'point': point,
                    'neighbors': neighbors.tolist()
                })
        
        return layer
    
    def _find_neighbors(self, vectors: np.ndarray, point_idx: int, k: int):
        """找到k个最近邻"""
        distances = np.linalg.norm(vectors - vectors[point_idx], axis=1)
        # 排除自己
        distances[point_idx] = np.inf
        neighbor_indices = np.argsort(distances)[:k]
        return neighbor_indices
    
    def search(self, query_vector: np.ndarray, top_k: int = 5):
        """HNSW搜索"""
        # 从顶层开始搜索
        current_layer = self.layers[-1]
        candidates = []
        
        # 1. 在顶层搜索
        entry_points = [layer['point'] for layer in current_layer]
        distances = np.linalg.norm(self.index_vectors[entry_points] - query_vector, axis=1)
        current_point = entry_points[np.argmin(distances)]
        
        # 2. 层层向下搜索
        for layer in reversed(self.layers):
            # 在当前层搜索
            neighbors = self._get_layer_neighbors(layer, current_point)
            
            # 计算距离并选择最近的点
            distances = np.linalg.norm(self.index_vectors[neighbors] - query_vector, axis=1)
            current_point = neighbors[np.argmin(distances)]
        
        # 3. 在底层精确搜索
        distances = np.linalg.norm(self.index_vectors - query_vector, axis=1)
        top_indices = np.argsort(distances)[:top_k]
        
        return list(zip(top_indices, distances[top_indices]))
​

特点：

• 层次化结构
• 搜索速度快
• 准确性较好
• 内存占用适中

适用场景：

• 中大规模数据
• 对速度要求高
• 可以接受一定的精度损失

LSH（Locality-Sensitive Hashing）

import numpy as np
from sklearn.preprocessing import normalize

class LSHIndex:
    """LSH索引实现"""
    
    def __init__(self, n_planes: int = 10, n_tables: int = 5):
        self.n_planes = n_planes
        self.n_tables = n_tables
        self.planes = []
        self.hash_tables = [{} for _ in range(n_tables)]
        self.index_vectors = None
    
    def fit(self, vectors: np.ndarray):
        """构建LSH索引"""
        self.index_vectors = vectors
        
        # 1. 生成随机平面
        for _ in range(self.n_planes):
            plane = np.random.randn(vectors.shape[1])
            self.planes.append(plane)
        
        # 2. 构建哈希表
        for i, vector in enumerate(vectors):
            for table_idx in range(self.n_tables):
                # 计算哈希值
                hash_key = self._compute_hash(vector, table_idx)
                
                # 添加到哈希表
                if hash_key not in self.hash_tables[table_idx]:
                    self.hash_tables[table_idx][hash_key] = []
                self.hash_tables[table_idx][hash_key].append(i)
    
    def _compute_hash(self, vector: np.ndarray, table_idx: int):
        """计算哈希值"""
        # 选择部分平面
        selected_planes = self.planes[table_idx * 2:(table_idx + 1) * 2]
        
        # 计算投影
        projections = []
        for plane in selected_planes:
            projection = np.dot(vector, plane)
            projections.append(1 if projection > 0 else 0)
        
        # 生成哈希键
        return tuple(projections)
    
    def search(self, query_vector: np.ndarray, top_k: int = 5):
        """LSH搜索"""
        # 1. 在所有表中查找候选点
        candidates = set()
        for table_idx in range(self.n_tables):
            hash_key = self._compute_hash(query_vector, table_idx)
            if hash_key in self.hash_tables[table_idx]:
                candidates.update(self.hash_tables[table_idx][hash_key])
        
        # 2. 计算相似度
        if not candidates:
            return []
        
        candidate_vectors = self.index_vectors[list(candidates)]
        similarities = cosine_similarity([query_vector], candidate_vectors)[0]
        
        # 3. 获取top-k结果
        top_indices = np.argsort(similarities)[-top_k:][::-1]
        top_scores = similarities[top_indices]
        
        return [(list(candidates)[i], top_scores[j]) for j, i in enumerate(top_indices)]
​

特点：

• 基于哈希的近似搜索
• 构建速度快
• 内存占用小
• 准确性一般

适用场景：

• 超大规模数据
• 对内存敏感
• 可以接受一定的精度损失

检索优化策略

1. 查询优化

查询预处理

class QueryOptimizer:
    """查询优化器"""
    
    def __init__(self):
        self.query_cache = {}
    
    def optimize_query(self, query: dict) -> dict:
        """优化查询"""
        # 查询预处理
        optimized_query = query.copy()
        
        # 1. 参数验证和标准化
        optimized_query = self._validate_parameters(optimized_query)
        
        # 2. 查询缓存
        query_hash = self._hash_query(optimized_query)
        if query_hash in self.query_cache:
            return self.query_cache[query_hash]
        
        # 3. 查询重写
        optimized_query = self._rewrite_query(optimized_query)
        
        # 4. 缓存结果
        self.query_cache[query_hash] = optimized_query
        
        return optimized_query
    
    def _validate_parameters(self, query: dict) -> dict:
        """验证查询参数"""
        # 标准化参数
        if 'top_k' not in query or query['top_k'] <= 0:
            query['top_k'] = 10
        
        if 'similarity_threshold' not in query:
            query['similarity_threshold'] = 0.5
        
        if 'search_strategy' not in query:
            query['search_strategy'] = 'auto'
        
        return query
    
    def _rewrite_query(self, query: dict) -> dict:
        """重写查询"""
        # 基于查询模式的重写
        if query.get('query_type') == 'semantic':
            # 语义查询优化
            query['use_semantic_filter'] = True
            query['expand_keywords'] = True
        elif query.get('query_type') == 'keyword':
            # 关键词查询优化
            query['use_keyword_boost'] = True
            query['keyword_weight'] = 0.7
        
        return query
​

批量查询优化

def batch_query_optimization(queries: list, batch_size: int = 100):
    """批量查询优化"""
    results = []
    
    for i in range(0, len(queries), batch_size):
        batch_queries = queries[i:i + batch_size]
        
        # 批量处理
        batch_results = []
        for query in batch_queries:
            # 预处理查询
            optimized_query = query_optimizer.optimize_query(query)
            
            # 执行查询
            result = vector_index.search(optimized_query)
            batch_results.append(result)
        
        results.extend(batch_results)
    
    return results
​

2. 索引优化

索引参数调优

def optimize_index_parameters(vectors: np.ndarray, query_pattern: str):
    """优化索引参数"""
    n_samples, n_dimensions = vectors.shape
    
    # 基于数据规模和查询模式选择参数
    if query_pattern == 'speed':
        # 速度优先模式
        params = {
            'index_type': 'HNSW',
            'M': 32,  # 连接数
            'ef': 64,  # 搜索宽度
            'n_clusters': min(100, n_samples // 10)
        }
    elif query_pattern == 'accuracy':
        # 准确性优先模式
        params = {
            'index_type': 'HNSW',
            'M': 64,  # 连接数
            'ef': 128,  # 搜索宽度
            'n_clusters': min(200, n_samples // 5)
        }
    elif query_pattern == 'balanced':
        # 平衡模式
        params = {
            'index_type': 'HNSW',
            'M': 24,  # 连接数
            'ef': 48,  # 搜索宽度
            'n_clusters': min(150, n_samples // 8)
        }
    else:
        # 默认模式
        params = {
            'index_type': 'HNSW',
            'M': 16,  # 连接数
            'ef': 32,  # 搜索宽度
            'n_clusters': min(100, n_samples // 10)
        }
    
    return params
​

索引重建策略

class IndexManager:
    """索引管理器"""
    
    def __init__(self, index_type: str = 'HNSW'):
        self.index_type = index_type
        self.index = None
        self.data_buffer = []
        self.rebuild_threshold = 1000  # 数据量达到阈值时重建索引
    
    def add_data(self, vectors: np.ndarray):
        """添加数据"""
        self.data_buffer.extend(vectors.tolist())
        
        # 检查是否需要重建索引
        if len(self.data_buffer) >= self.rebuild_threshold:
            self.rebuild_index()
    
    def rebuild_index(self):
        """重建索引"""
        if not self.data_buffer:
            return
        
        # 合并新数据
        if self.index is not None:
            all_vectors = np.vstack([self.index_vectors, np.array(self.data_buffer)])
        else:
            all_vectors = np.array(self.data_buffer)
        
        # 重建索引
        if self.index_type == 'HNSW':
            self.index = HNSWIndex()
            self.index.fit(all_vectors)
        elif self.index_type == 'LSH':
            self.index = LSHIndex()
            self.index.fit(all_vectors)
        
        # 清空缓冲区
        self.data_buffer = []
    
    def search(self, query_vector: np.ndarray, top_k: int = 5):
        """搜索"""
        # 如果有未处理的数据，先处理
        if self.data_buffer and len(self.data_buffer) >= self.rebuild_threshold:
            self.rebuild_index()
        
        return self.index.search(query_vector, top_k)
​

最佳实践总结

1. 索引选择原则

• 数据规模：根据数据量选择合适的索引算法
• 查询模式：根据查询频率和复杂度选择索引
• 准确性要求：根据准确性要求选择索引类型
• 资源限制：根据内存和计算资源选择索引

2. 性能优化策略

• 参数调优：根据实际数据特点调整索引参数
• 缓存机制：使用缓存减少重复计算
• 批量处理：使用批量处理提升吞吐量
• 内存管理：合理使用内存和存储资源