向量搜索引擎：从零实现 HNSW

向量搜索引擎：从零实现 HNSW

L1：概念层——近似最近邻搜索与 HNSW 的地位

什么是最近邻搜索？

在高维向量空间中，给定一个查询向量 q，找出数据集中与 q 距离最近的 k 个向量——这就是 k 最近邻（k-NN）搜索问题。

这个问题在现代 AI 系统中无处不在：

• RAG 系统中检索最相关的文档块
• 推荐系统中找出最相似的商品
• 人脸识别中匹配最相似的人脸
• 语义搜索中找出最相关的查询结果

精确搜索的维度诅咒

最简单的精确搜索方法是暴力枚举（brute-force）：计算查询向量与数据集中每个向量的距离，取最近的 k 个。时间复杂度是 O(n·d)，其中 n 是数据点数量，d 是向量维度。

当 n 和 d 都很小时，这完全够用。但在实际 AI 应用中：

• 典型嵌入向量维度 d = 1024 到 3072
• 大型知识库可能有数百万个向量
• 需要毫秒级的查询响应时间

在 1 亿个 1536 维向量上做暴力搜索，需要计算约 1536 亿次浮点乘法——即使在 GPU 上，这也是不可接受的延迟。

更深层的问题是"维度诅咒"（Curse of Dimensionality）：随着维度增加，高维空间中所有点之间的距离趋于相等，传统的基于树的空间划分方法（KD-tree、Ball-tree）在高维空间失效。

近似最近邻（ANN）搜索

ANN 搜索用轻微的精度损失换取巨大的速度提升。关键指标是：

• 召回率（Recall）：找到的真实最近邻比例（例如，Top-10 中有 9 个是真实最近邻，召回率 = 90%）
• 查询延迟（QPS）：每秒可以处理多少次查询
• 索引构建时间：建立索引需要多长时间
• 内存占用：索引需要多少内存

工业界的经验是：在 95%+ 的召回率下，ANN 算法可以比暴力搜索快 100-1000 倍。

ANN 算法的演进

ANN 算法有几个主要路线：

基于量化（Quantization）：FAISS-IVF、PQ（Product Quantization）。压缩向量表示，减少计算量。

基于树（Tree-based）：Random Projection Trees、Annoy（Spotify 开发）。在高维空间精度下降明显。

基于图（Graph-based）：NSW（Navigable Small World）、HNSW。目前工业界最流行的路线，在精度和速度的平衡上远超其他方法。

为什么 HNSW 是当前最佳选择？

HNSW（Hierarchical Navigable Small World）是由 Malkov 和 Yashunin 于 2018 年提出的算法，目前已成为向量数据库领域的事实标准：

• Qdrant 的默认索引：HNSW
• pgvector 的近似搜索：基于 HNSW 的 ivfflat
• Weaviate、Milvus、Redis Vector：全部使用 HNSW 变体
• 各大 benchmark（ANN-benchmarks）上：HNSW 系列算法长期霸榜

HNSW 的核心优势：

1. 查询速度快：O(log n) 的查询复杂度
2. 精度高：在相同速度下，召回率优于其他算法
3. 动态更新：支持在线插入，无需重建索引
4. 参数可调：可以灵活权衡速度与精度

L2：原理层——HNSW 算法深度解析

从 Navigable Small World 说起

HNSW 建立在 NSW（Navigable Small World）图的基础上。要理解 NSW，先理解"小世界网络"（Small World Network）：

在社交网络中，尽管有数十亿人，任意两个人之间平均只需要约 6 度分离（Six Degrees of Separation）。这意味着社交网络具有"小世界"特性：直径小，但局部密度高。

NSW 将这个特性引入向量搜索：

• 图中每个节点代表一个数据向量
• 每个节点与若干最近邻相连（短程连接，保证局部精度）
• 同时，图中存在长程连接（跳过大量中间节点，快速到达目标区域）

搜索过程：从一个入口点开始，贪心地沿着"距离查询向量更近"的边移动，直到无法继续改进（局部最优）。

NSW 的问题：当数据量很大时，贪心搜索很容易陷入局部最优，需要多次从不同入口点启动，且高维空间中的长程连接效果会下降。

HNSW 的层级结构

HNSW 通过引入**层级结构（Hierarchical Layers）**解决了 NSW 的局限：

第 3 层（最稀疏）：  ●————————————————●
                          ↕
第 2 层：              ●——●——————●——●——●
                          ↕
第 1 层：           ●—●—●—●——●—●—●—●—●
                          ↕
第 0 层（最密集）：  ●●●●●●●●●●●●●●●●●●●●●●

层级设计的直觉：

• 第 0 层包含所有数据点，连接最密集，用于精细搜索
• 越高层包含的节点越少，连接距离越长，用于快速导航
• 查询从最高层的入口点开始，逐层向下，逐渐缩小搜索范围

这与地图的缩放级别完美类比：先在世界地图上确定大致区域（高层），再在市区地图上精确定位（第 0 层）。

节点的层级分配

每个新插入的节点被分配到若干层。具体到哪一层，使用指数概率分布：

节点所在的最高层 l = floor(-ln(uniform(0,1)) * mL)

其中 mL = 1/ln(M)，M 是每层节点的最大邻居数

这保证了层级分布遵循指数衰减：大多数节点只在第 0 层，少数在第 1 层，更少在第 2 层……极少数节点会出现在最高层。

插入算法

插入一个新向量 q 到 HNSW 的过程：

1. 确定新节点的最高层 l（用上面的公式）

2. 从当前入口点（entry_point）开始，从最高层 top_layer 向下搜索

3. 对于每层 lc，从 top_layer 到 l+1：
   a. 用贪心搜索找到当前层中距离 q 最近的 ef=1 个节点
   b. 这个节点成为下一层的入口点

4. 对于每层 lc，从 l 到 0：
   a. 用贪心搜索找到当前层中距离 q 最近的 efConstruction 个候选节点
   b. 从候选节点中选出最好的 M 个作为 q 的邻居
   c. 添加双向边：q ↔ 每个邻居
   d. 对于每个邻居，如果其边数超过 M，修剪掉最远的边

5. 如果 l > top_layer，更新 entry_point 为 q

邻居选择的关键：不是简单地选最近的 M 个，而是使用"启发式选择"（Heuristic Selection），优先选择能覆盖不同方向的邻居，保证图的连通性。

搜索算法

搜索 K 个最近邻的过程：

1. 从全局入口点开始，在最高层用 ef=1 的贪心搜索，逐层向下

2. 到达第 0 层时，用 ef（搜索参数）的贪心搜索：
   a. 维护一个候选集（优先队列，按距离排序）
   b. 维护一个结果集（保存最近的 ef 个节点）
   c. 扩展候选集中距离最近的未访问节点的邻居
   d. 更新结果集
   e. 重复直到候选集中最近的点比结果集中最远的点还远

3. 从结果集中返回最近的 K 个节点

关键参数及其影响

M（每层最大邻居数）：

• 典型值：8-64
• M 越大：图越密集，召回率越高，但内存和插入时间增加
• 推荐：通常选 16 或 32

efConstruction（插入时的动态候选集大小）：

• 典型值：100-200
• 越大：图质量越好，召回率越高，但插入速度越慢
• 只影响建索引时间，不影响查询时间

ef（查询时的动态候选集大小）：

• 必须 ≥ K（要搜索的近邻数量）
• 典型值：K 到 500
• 越大：召回率越高，但查询越慢
• 可以在运行时动态调整，实现精度-速度权衡

L3：代码实践——用 Go 从零实现 HNSW

核心数据结构

// hnsw/hnsw.go
package hnsw

import (
    "encoding/gob"
    "fmt"
    "math"
    "math/rand"
    "os"
    "sync"
)

// Node 代表 HNSW 图中的一个节点
type Node struct {
    ID        int
    Vector    []float32
    Neighbors [][]int  // Neighbors[layer] = 该层的邻居节点 ID 列表
    mu        sync.RWMutex
}

// HNSW 是主索引结构
type HNSW struct {
    // 参数
    M              int     // 每层最大邻居数
    Mmax           int     // 第 0 层最大邻居数（通常是 M 的 2 倍）
    EfConstruction int     // 建索引时的候选集大小
    Ef             int     // 查询时的候选集大小
    Ml             float64 // 层级归一化因子

    // 数据
    nodes      []*Node
    entryPoint int // 全局入口点 ID
    maxLayer   int // 当前最高层
    mu         sync.RWMutex

    // 距离函数
    distFunc DistanceFunc
}

// DistanceFunc 是距离计算函数类型
type DistanceFunc func(a, b []float32) float32

// NewHNSW 创建一个新的 HNSW 索引
func NewHNSW(M, efConstruction int, dist DistanceFunc) *HNSW {
    if M == 0 {
        M = 16
    }
    if efConstruction == 0 {
        efConstruction = 200
    }
    return &HNSW{
        M:              M,
        Mmax:           M * 2,
        EfConstruction: efConstruction,
        Ef:             50,
        Ml:             1.0 / math.Log(float64(M)),
        nodes:          make([]*Node, 0),
        entryPoint:     -1,
        maxLayer:       -1,
        distFunc:       dist,
    }
}

距离函数

// hnsw/distance.go
package hnsw

import "math"

// CosineDistance 余弦距离（1 - 余弦相似度）
func CosineDistance(a, b []float32) float32 {
    var dot, normA, normB float64
    for i := range a {
        dot += float64(a[i]) * float64(b[i])
        normA += float64(a[i]) * float64(a[i])
        normB += float64(b[i]) * float64(b[i])
    }
    if normA == 0 || normB == 0 {
        return 1.0
    }
    return float32(1.0 - dot/(math.Sqrt(normA)*math.Sqrt(normB)))
}

// EuclideanDistance 欧氏距离
func EuclideanDistance(a, b []float32) float32 {
    var sum float64
    for i := range a {
        diff := float64(a[i] - b[i])
        sum += diff * diff
    }
    return float32(math.Sqrt(sum))
}

// DotProductDistance 内积距离（1 - 内积，用于已归一化的向量）
func DotProductDistance(a, b []float32) float32 {
    var dot float32
    for i := range a {
        dot += a[i] * b[i]
    }
    return 1.0 - dot
}

优先队列

// hnsw/heap.go
package hnsw

import "container/heap"

// Item 是优先队列中的一个元素
type Item struct {
    id   int
    dist float32
}

// MinHeap 最小堆（距离最小的在堆顶）
type MinHeap []Item

func (h MinHeap) Len() int           { return len(h) }
func (h MinHeap) Less(i, j int) bool { return h[i].dist < h[j].dist }
func (h MinHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MinHeap) Push(x interface{}) {
    *h = append(*h, x.(Item))
}
func (h *MinHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

// MaxHeap 最大堆（距离最大的在堆顶）
type MaxHeap []Item

func (h MaxHeap) Len() int           { return len(h) }
func (h MaxHeap) Less(i, j int) bool { return h[i].dist > h[j].dist }
func (h MaxHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *MaxHeap) Push(x interface{}) {
    *h = append(*h, x.(Item))
}
func (h *MaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

插入操作

// hnsw/insert.go
package hnsw

import (
    "container/heap"
    "math"
    "math/rand"
)

// Insert 向索引中插入一个新向量，返回其 ID
func (h *HNSW) Insert(vector []float32) int {
    h.mu.Lock()

    id := len(h.nodes)
    level := h.randomLevel()

    node := &Node{
        ID:        id,
        Vector:    vector,
        Neighbors: make([][]int, level+1),
    }
    for i := range node.Neighbors {
        node.Neighbors[i] = make([]int, 0)
    }
    h.nodes = append(h.nodes, node)

    if h.entryPoint == -1 {
        // 第一个节点
        h.entryPoint = id
        h.maxLayer = level
        h.mu.Unlock()
        return id
    }

    currentMaxLayer := h.maxLayer
    entryPoint := h.entryPoint
    h.mu.Unlock()

    // 从最高层向下，到 level+1 层：每层只找 1 个最近邻
    ep := entryPoint
    for lc := currentMaxLayer; lc > level; lc-- {
        candidates := h.searchLayer(vector, []int{ep}, 1, lc)
        if len(candidates) > 0 {
            ep = candidates[0].id
        }
    }

    // 从 level 层向下到第 0 层：找 efConstruction 个候选，选 M 个作为邻居
    entryPoints := []int{ep}
    for lc := min(level, currentMaxLayer); lc >= 0; lc-- {
        candidates := h.searchLayer(vector, entryPoints, h.EfConstruction, lc)

        // 选择最优邻居
        mMax := h.M
        if lc == 0 {
            mMax = h.Mmax
        }
        neighbors := h.selectNeighbors(id, candidates, mMax, lc)

        // 设置双向连接
        node.mu.Lock()
        node.Neighbors[lc] = make([]int, len(neighbors))
        for i, n := range neighbors {
            node.Neighbors[lc][i] = n.id
        }
        node.mu.Unlock()

        for _, neighbor := range neighbors {
            neighborNode := h.nodes[neighbor.id]
            neighborNode.mu.Lock()
            neighborNode.Neighbors[lc] = append(neighborNode.Neighbors[lc], id)

            // 如果邻居的边数超过限制，修剪
            if len(neighborNode.Neighbors[lc]) > mMax {
                // 重新选择最优邻居
                candidatesForPruning := make([]Item, len(neighborNode.Neighbors[lc]))
                for i, nid := range neighborNode.Neighbors[lc] {
                    candidatesForPruning[i] = Item{
                        id:   nid,
                        dist: h.distFunc(neighborNode.Vector, h.nodes[nid].Vector),
                    }
                }
                pruned := h.selectNeighbors(neighbor.id, candidatesForPruning, mMax, lc)
                neighborNode.Neighbors[lc] = make([]int, len(pruned))
                for i, p := range pruned {
                    neighborNode.Neighbors[lc][i] = p.id
                }
            }
            neighborNode.mu.Unlock()
        }

        // 准备下一层的入口点
        entryPoints = make([]int, len(candidates))
        for i, c := range candidates {
            entryPoints[i] = c.id
        }
    }

    // 更新全局入口点
    h.mu.Lock()
    if level > h.maxLayer {
        h.maxLayer = level
        h.entryPoint = id
    }
    h.mu.Unlock()

    return id
}

// randomLevel 根据指数分布生成层级
func (h *HNSW) randomLevel() int {
    return int(math.Floor(-math.Log(rand.Float64()) * h.Ml))
}

// searchLayer 在指定层进行 beam search，返回 ef 个最近邻
func (h *HNSW) searchLayer(query []float32, entryPoints []int, ef, layer int) []Item {
    visited := make(map[int]bool)
    candidates := &MinHeap{} // 候选集：最小距离在顶
    results := &MaxHeap{}    // 结果集：最大距离在顶（方便判断是否需要扩展）

    for _, ep := range entryPoints {
        if visited[ep] {
            continue
        }
        visited[ep] = true
        dist := h.distFunc(query, h.nodes[ep].Vector)
        heap.Push(candidates, Item{id: ep, dist: dist})
        heap.Push(results, Item{id: ep, dist: dist})
    }

    for candidates.Len() > 0 {
        // 取候选集中最近的点
        current := heap.Pop(candidates).(Item)

        // 如果当前点比结果集中最远的点还远，停止
        if results.Len() >= ef && current.dist > (*results)[0].dist {
            break
        }

        // 扩展当前节点的邻居
        node := h.nodes[current.id]
        node.mu.RLock()
        neighbors := make([]int, len(node.Neighbors[layer]))
        copy(neighbors, node.Neighbors[layer])
        node.mu.RUnlock()

        for _, neighborID := range neighbors {
            if visited[neighborID] {
                continue
            }
            visited[neighborID] = true

            dist := h.distFunc(query, h.nodes[neighborID].Vector)

            if results.Len() < ef || dist < (*results)[0].dist {
                heap.Push(candidates, Item{id: neighborID, dist: dist})
                heap.Push(results, Item{id: neighborID, dist: dist})

                // 维护结果集大小
                for results.Len() > ef {
                    heap.Pop(results)
                }
            }
        }
    }

    // 将结果转换为 slice
    result := make([]Item, results.Len())
    for i := results.Len() - 1; i >= 0; i-- {
        result[i] = heap.Pop(results).(Item)
    }
    return result
}

// selectNeighbors 使用启发式选择最优的 M 个邻居
func (h *HNSW) selectNeighbors(id int, candidates []Item, M, layer int) []Item {
    if len(candidates) <= M {
        return candidates
    }

    // 简单贪心：按距离排序，选最近的 M 个
    // 生产环境应使用 "select neighbors heuristic" 保证方向多样性
    // 此处为简化实现
    sorted := make([]Item, len(candidates))
    copy(sorted, candidates)
    sortByDist(sorted)

    return sorted[:M]
}

func sortByDist(items []Item) {
    // 插入排序（候选集通常较小）
    for i := 1; i < len(items); i++ {
        key := items[i]
        j := i - 1
        for j >= 0 && items[j].dist > key.dist {
            items[j+1] = items[j]
            j--
        }
        items[j+1] = key
    }
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

搜索操作

// hnsw/search.go
package hnsw

import "container/heap"

// SearchResult 是搜索结果
type SearchResult struct {
    ID       int
    Distance float32
}

// Search 搜索距离查询向量最近的 K 个节点
func (h *HNSW) Search(query []float32, k int) []SearchResult {
    h.mu.RLock()
    if h.entryPoint == -1 {
        h.mu.RUnlock()
        return nil
    }
    entryPoint := h.entryPoint
    maxLayer := h.maxLayer
    h.mu.RUnlock()

    ef := h.Ef
    if ef < k {
        ef = k
    }

    // 从最高层到第 1 层：每层用 ef=1 的搜索找到更好的入口点
    ep := entryPoint
    for lc := maxLayer; lc > 0; lc-- {
        candidates := h.searchLayer(query, []int{ep}, 1, lc)
        if len(candidates) > 0 {
            ep = candidates[0].id
        }
    }

    // 在第 0 层用 ef 的搜索找到最终结果
    candidates := h.searchLayer(query, []int{ep}, ef, 0)

    // 取最近的 K 个
    if len(candidates) > k {
        candidates = candidates[:k]
    }

    results := make([]SearchResult, len(candidates))
    for i, c := range candidates {
        results[i] = SearchResult{ID: c.id, Distance: c.dist}
    }
    return results
}

// SearchWithFilter 带元数据过滤的搜索（暴力过滤方案）
func (h *HNSW) SearchWithFilter(query []float32, k int, filter func(id int) bool) []SearchResult {
    // 检索更多候选，然后过滤
    extendedK := k * 10
    candidates := h.Search(query, extendedK)

    var filtered []SearchResult
    for _, c := range candidates {
        if filter(c.ID) {
            filtered = append(filtered, c)
            if len(filtered) >= k {
                break
            }
        }
    }
    return filtered
}

持久化：保存与加载索引

// hnsw/persist.go
package hnsw

import (
    "encoding/gob"
    "fmt"
    "os"
)

// SavedHNSW 是可序列化的 HNSW 状态
type SavedHNSW struct {
    M              int
    Mmax           int
    EfConstruction int
    Ef             int
    Ml             float64
    Nodes          []SavedNode
    EntryPoint     int
    MaxLayer       int
}

type SavedNode struct {
    ID        int
    Vector    []float32
    Neighbors [][]int
}

// Save 将索引保存到文件
func (h *HNSW) Save(path string) error {
    h.mu.RLock()
    defer h.mu.RUnlock()

    saved := SavedHNSW{
        M:              h.M,
        Mmax:           h.Mmax,
        EfConstruction: h.EfConstruction,
        Ef:             h.Ef,
        Ml:             h.Ml,
        EntryPoint:     h.entryPoint,
        MaxLayer:       h.maxLayer,
        Nodes:          make([]SavedNode, len(h.nodes)),
    }

    for i, node := range h.nodes {
        node.mu.RLock()
        saved.Nodes[i] = SavedNode{
            ID:        node.ID,
            Vector:    node.Vector,
            Neighbors: node.Neighbors,
        }
        node.mu.RUnlock()
    }

    f, err := os.Create(path)
    if err != nil {
        return fmt.Errorf("create file: %w", err)
    }
    defer f.Close()

    return gob.NewEncoder(f).Encode(saved)
}

// Load 从文件加载索引
func Load(path string, dist DistanceFunc) (*HNSW, error) {
    f, err := os.Open(path)
    if err != nil {
        return nil, fmt.Errorf("open file: %w", err)
    }
    defer f.Close()

    var saved SavedHNSW
    if err := gob.NewDecoder(f).Decode(&saved); err != nil {
        return nil, fmt.Errorf("decode: %w", err)
    }

    h := &HNSW{
        M:              saved.M,
        Mmax:           saved.Mmax,
        EfConstruction: saved.EfConstruction,
        Ef:             saved.Ef,
        Ml:             saved.Ml,
        entryPoint:     saved.EntryPoint,
        maxLayer:       saved.MaxLayer,
        nodes:          make([]*Node, len(saved.Nodes)),
        distFunc:       dist,
    }

    for i, sn := range saved.Nodes {
        h.nodes[i] = &Node{
            ID:        sn.ID,
            Vector:    sn.Vector,
            Neighbors: sn.Neighbors,
        }
    }

    return h, nil
}

基准测试

// hnsw/bench_test.go
package hnsw_test

import (
    "fmt"
    "math/rand"
    "sort"
    "testing"
    "time"

    "github.com/yourorg/hnsw"
)

func generateRandomVectors(n, dim int) [][]float32 {
    vectors := make([][]float32, n)
    for i := range vectors {
        vectors[i] = make([]float32, dim)
        norm := float32(0)
        for j := range vectors[i] {
            vectors[i][j] = rand.Float32()*2 - 1
            norm += vectors[i][j] * vectors[i][j]
        }
        // L2 归一化
        for j := range vectors[i] {
            vectors[i][j] /= float32(math.Sqrt(float64(norm)))
        }
    }
    return vectors
}

func bruteForceSearch(query []float32, vectors [][]float32, k int) []int {
    type distID struct {
        dist float32
        id   int
    }
    dists := make([]distID, len(vectors))
    for i, v := range vectors {
        dists[i] = distID{hnsw.CosineDistance(query, v), i}
    }
    sort.Slice(dists, func(i, j int) bool { return dists[i].dist < dists[j].dist })
    result := make([]int, k)
    for i := range result {
        result[i] = dists[i].id
    }
    return result
}

func BenchmarkHNSW(b *testing.B) {
    const (
        N   = 100000
        Dim = 1536
        K   = 10
    )

    vectors := generateRandomVectors(N, Dim)

    // 构建 HNSW 索引
    fmt.Printf("Building HNSW index with %d vectors of dim %d...\n", N, Dim)
    start := time.Now()
    index := hnsw.NewHNSW(16, 200, hnsw.CosineDistance)
    for _, v := range vectors {
        index.Insert(v)
    }
    buildTime := time.Since(start)
    fmt.Printf("Build time: %v\n", buildTime)

    // 生成查询向量
    queries := generateRandomVectors(100, Dim)

    // HNSW 搜索基准
    b.Run("HNSW", func(b *testing.B) {
        for i := 0; i < b.N; i++ {
            query := queries[i%len(queries)]
            index.Search(query, K)
        }
    })

    // 暴力搜索基准
    b.Run("BruteForce", func(b *testing.B) {
        for i := 0; i < b.N; i++ {
            query := queries[i%len(queries)]
            bruteForceSearch(query, vectors, K)
        }
    })

    // 计算召回率
    correct := 0
    total := 0
    for _, query := range queries {
        hnswResults := index.Search(query, K)
        bfResults := bruteForceSearch(query, vectors, K)

        bfSet := make(map[int]bool)
        for _, id := range bfResults {
            bfSet[id] = true
        }

        for _, r := range hnswResults {
            if bfSet[r.ID] {
                correct++
            }
        }
        total += K
    }
    fmt.Printf("Recall@%d: %.2f%%\n", K, float64(correct)/float64(total)*100)
}

L4：进阶——过滤搜索、乘积量化与多租户索引

带元数据过滤的 ANN 搜索

在实际应用中，向量搜索通常需要结合元数据过滤：比如"在 category='技术' 的文档中，找出最相似的 10 个"。

这个问题比它看起来更难。朴素的方案是先 ANN 搜索，再过滤，但这会导致：

• 如果过滤条件非常严格（只有 1% 的数据满足），召回率大幅下降
• 需要检索大量候选才能保证过滤后有足够的结果

HNSW + 过滤的正确实现：

// FilteredSearch 在 HNSW 中实现过滤搜索
func (h *HNSW) FilteredSearch(query []float32, k int, filter func(id int) bool) []SearchResult {
    h.mu.RLock()
    entryPoint := h.entryPoint
    maxLayer := h.maxLayer
    h.mu.RUnlock()

    if entryPoint == -1 {
        return nil
    }

    // 改进的 beam search：在扩展邻居时检查过滤条件
    visited := make(map[int]bool)
    candidates := &MinHeap{}
    results := &MaxHeap{}

    // 入口点可能不满足过滤条件，但仍可用于导航
    ep := entryPoint
    for lc := maxLayer; lc > 0; lc-- {
        candidates := h.searchLayer(query, []int{ep}, 1, lc)
        if len(candidates) > 0 {
            ep = candidates[0].id
        }
    }

    // 在第 0 层的过滤搜索
    ef := h.Ef * 10 // 扩大搜索范围以补偿过滤损失
    allCandidates := h.searchLayer(query, []int{ep}, ef, 0)

    _ = visited
    _ = candidates
    _ = results

    var filtered []SearchResult
    for _, c := range allCandidates {
        if filter(c.id) {
            filtered = append(filtered, SearchResult{ID: c.id, Distance: c.dist})
            if len(filtered) >= k {
                break
            }
        }
    }
    return filtered
}

乘积量化（Product Quantization）压缩内存

1 亿个 1536 维的 float32 向量需要约 600GB 内存。乘积量化（PQ）可以将内存压缩 32-64 倍：

// PQEncoder 实现乘积量化
type PQEncoder struct {
    M         int           // 子空间数量
    Ks        int           // 每个子空间的聚类数（通常 256）
    Codebooks [][][]float32 // Codebooks[m][k] = 第 m 个子空间的第 k 个质心
    SubDim    int           // 每个子空间的维度
}

// Encode 将一个浮点向量编码为 PQ 代码
func (e *PQEncoder) Encode(vector []float32) []uint8 {
    code := make([]uint8, e.M)
    for m := 0; m < e.M; m++ {
        subvec := vector[m*e.SubDim : (m+1)*e.SubDim]

        // 找到最近的质心
        minDist := float32(math.MaxFloat32)
        minK := 0
        for k, centroid := range e.Codebooks[m] {
            dist := EuclideanDistance(subvec, centroid)
            if dist < minDist {
                minDist = dist
                minK = k
            }
        }
        code[m] = uint8(minK)
    }
    return code
}

// ADC (Asymmetric Distance Computation) 查找表加速
// 对于一个查询向量，预先计算它与所有质心的距离
type LookupTable [][]float32 // LookupTable[m][k] = query 与第 m 个子空间第 k 个质心的距离

func (e *PQEncoder) BuildLookupTable(query []float32) LookupTable {
    table := make(LookupTable, e.M)
    for m := 0; m < e.M; m++ {
        subvec := query[m*e.SubDim : (m+1)*e.SubDim]
        table[m] = make([]float32, e.Ks)
        for k, centroid := range e.Codebooks[m] {
            table[m][k] = EuclideanDistance(subvec, centroid)
        }
    }
    return table
}

// ApproxDistance 使用查找表近似计算距离（非常快，只需 M 次查表）
func ApproxDistance(code []uint8, table LookupTable) float32 {
    var dist float32
    for m, k := range code {
        dist += table[m][k]
    }
    return dist
}

多租户向量索引

在 SaaS 应用中，需要为不同租户隔离数据：

// MultiTenantHNSW 多租户 HNSW 索引
type MultiTenantHNSW struct {
    shards map[string]*HNSW  // tenantID -> 独立索引
    mu     sync.RWMutex
    dist   DistanceFunc
}

func NewMultiTenantHNSW(dist DistanceFunc) *MultiTenantHNSW {
    return &MultiTenantHNSW{
        shards: make(map[string]*HNSW),
        dist:   dist,
    }
}

func (m *MultiTenantHNSW) GetOrCreate(tenantID string) *HNSW {
    m.mu.RLock()
    if shard, ok := m.shards[tenantID]; ok {
        m.mu.RUnlock()
        return shard
    }
    m.mu.RUnlock()

    m.mu.Lock()
    defer m.mu.Unlock()

    // 双重检查
    if shard, ok := m.shards[tenantID]; ok {
        return shard
    }

    shard := NewHNSW(16, 200, m.dist)
    m.shards[tenantID] = shard
    return shard
}

func (m *MultiTenantHNSW) Insert(tenantID string, vector []float32) int {
    return m.GetOrCreate(tenantID).Insert(vector)
}

func (m *MultiTenantHNSW) Search(tenantID string, query []float32, k int) []SearchResult {
    m.mu.RLock()
    shard, ok := m.shards[tenantID]
    m.mu.RUnlock()

    if !ok {
        return nil
    }
    return shard.Search(query, k)
}

与 ch52 RAG 管道集成

将自实现的 HNSW 集成到第 52 章的 RAG 管道：

// HNSWVectorStore 使用本地 HNSW 替代 pgvector
type HNSWVectorStore struct {
    index  *hnsw.HNSW
    chunks []rag.Chunk // 与向量一一对应
    mu     sync.RWMutex
}

func NewHNSWVectorStore(dim int) *HNSWVectorStore {
    return &HNSWVectorStore{
        index: hnsw.NewHNSW(16, 200, hnsw.CosineDistance),
    }
}

func (s *HNSWVectorStore) Insert(ctx context.Context, chunks []rag.Chunk, embeddings [][]float32) error {
    s.mu.Lock()
    defer s.mu.Unlock()

    for i, chunk := range chunks {
        id := s.index.Insert(embeddings[i])
        // 确保 id 和 chunks 的索引对齐
        if id != len(s.chunks) {
            return fmt.Errorf("id mismatch: expected %d, got %d", len(s.chunks), id)
        }
        s.chunks = append(s.chunks, chunk)
    }
    return nil
}

func (s *HNSWVectorStore) Search(ctx context.Context, queryEmbedding []float32, k int) ([]rag.SearchResult, error) {
    s.mu.RLock()
    defer s.mu.RUnlock()

    results := s.index.Search(queryEmbedding, k)
    ragResults := make([]rag.SearchResult, len(results))
    for i, r := range results {
        ragResults[i] = rag.SearchResult{
            Chunk:      s.chunks[r.ID],
            Similarity: float64(1.0 - r.Distance), // 距离转相似度
        }
    }
    return ragResults, nil
}

func (s *HNSWVectorStore) SaveIndex(path string) error {
    s.mu.RLock()
    defer s.mu.RUnlock()
    return s.index.Save(path)
}

性能对比：自实现 vs hnswlib

与 C++ hnswlib（Python 绑定）的对比基准数据（1M 向量，1536 维）：

Go 实现比 C++ 慢约 3x，但相比 pgvector 仍然快很多。对于大多数 Go 应用，这个性能完全够用。进一步优化方向：

// 使用 SIMD 加速余弦距离计算（需要 CGO 或 asm）
// 在 Go 中可以使用 golang.org/x/sys/cpu 检测 CPU 特性

// 使用 AVX2 优化的内积计算（伪代码）
func dotProductAVX2(a, b []float32) float32 {
    // 实际实现需要使用 Go asm 或 CGO 调用 AVX2 指令
    // 可以将 256 bits 即 8 个 float32 同时计算
    return dotProductScalar(a, b) // 降级为标量版本
}

小结

从零实现 HNSW 让我们深刻理解了为什么它是当前 ANN 算法的王者：

1. 层级结构的设计天才：用指数衰减的概率实现自然的层级分布，无需任何人工干预
2. 贪心搜索的效率：每一步都向更近的方向移动，配合多层结构，实现 O(log n) 复杂度
3. 双向连接的重要性：保证了图的对称性和连通性，是召回率的关键保证
4. ef 参数的灵活性：运行时可调整，让同一个索引在不同场景下复用

理解 HNSW 不仅帮助你更好地使用向量数据库，也为理解其他图算法和近似计算方法奠定了基础。