什么是 ‘Vector Database Kernels’:利用 Go 手写 HNSW 索引实现亿级向量的毫秒级检索

向量数据库内核:利用 Go 手写 HNSW 索引实现亿级向量的毫秒级检索

1. 向量检索与向量数据库的崛起

在人工智能和机器学习日益普及的今天,我们处理的数据类型正在发生深刻的变化。传统的结构化数据,如数字和文本,已不再足以描述图像、音频、视频、自然语言的深层含义。为了捕捉这些复杂数据的高维语义信息,我们将其转化为向量(embeddings)。这些向量是高维空间中的点,它们之间的距离或相似度可以量化原始数据之间的语义关联。

向量检索(Vector Search),或称近似最近邻(Approximate Nearest Neighbor, ANN)搜索,旨在从海量向量数据集中快速找出与给定查询向量最相似的 K 个向量。这项技术是许多现代AI应用的核心基石,例如:

  • 推荐系统:为用户推荐相似的商品、电影或音乐。
  • 语义搜索:理解用户查询的意图,返回语义相关的文档或网页,而非仅仅关键词匹配。
  • 图像识别与检索:根据一张图片找到数据库中相似的图片。
  • 自然语言处理:问答系统、文本去重、抄袭检测。
  • 个性化广告:根据用户行为向量匹配广告向量。

随着向量数据规模的爆炸式增长,传统的数据库系统在处理高维向量的相似性查询时显得力不从心。暴力搜索(Brute-Force Search)的计算复杂度太高,无法满足实时性要求。因此,专门为向量检索而设计的向量数据库(Vector Database)应运而生。

向量数据库的核心在于其高效的索引结构和检索算法,这正是我们所说的“向量数据库内核”。一个高性能的内核能够在大规模数据集上实现毫秒级的查询响应。本讲座将深入探讨如何使用 Go 语言,从零开始构建一个 HNSW(Hierarchical Navigable Small World)索引,以应对亿级向量的毫秒级检索挑战。

选择 Go 语言有其独特的优势:

  • 并发模型:Go 的 Goroutine 和 Channel 机制天然支持高并发操作,非常适合构建高性能的服务。
  • 性能:Go 编译为原生机器码,运行效率接近 C/C++,同时拥有更快的开发速度和更好的安全性。
  • 内存管理:Go 的垃圾回收机制简化了内存管理,减少了内存泄漏的风险,但同时也需要开发者注意优化内存使用。
  • 生态系统:Go 拥有活跃的社区和丰富的标准库,易于构建网络服务和分布式系统。

我们的目标是理解 HNSW 算法的原理,并用 Go 语言将其实现为一个可扩展、高性能的向量索引内核,为亿级向量提供毫秒级检索能力。

2. 向量空间与相似性度量

在深入 HNSW 算法之前,我们首先需要理解向量以及如何衡量它们之间的“相似度”。

2.1 向量的表示

在数学上,一个向量是一个有方向和大小的量。在机器学习中,它通常表示为一个数字列表(或数组),每个数字代表一个维度上的特征值。例如,一个 3 维向量 v = [x, y, z] 可以看作是 3 维空间中的一个点。

// 示例:定义一个向量类型
type Vector []float32

// 假设我们有一个简单的函数来创建向量
func NewVector(dims int) Vector {
    return make(Vector, dims)
}

2.2 相似性度量

为了量化两个向量之间的相似性或距离,我们使用不同的度量标准。选择合适的度量标准对于检索结果的质量至关重要,它取决于向量的生成方式和应用场景。

2.2.1 欧氏距离 (Euclidean Distance, L2 Distance)

欧氏距离是多维空间中两点之间直线距离的泛化。它直观地反映了两点之间的“远近”。距离越小,相似度越高。

对于两个 $D$ 维向量 $A = [a_1, a_2, …, a_D]$ 和 $B = [b_1, b_2, …, bD]$,欧氏距离的计算公式为:
$$
d(A, B) = sqrt{sum
{i=1}^{D} (a_i – b_i)^2}
$$
在实际应用中,为了计算效率,常常使用欧氏距离的平方(L2 距离的平方),因为它避免了开方运算,且不影响距离的相对排序。

// L2DistanceSquared 计算两个向量的欧氏距离的平方
func L2DistanceSquared(v1, v2 Vector) float32 {
    var sum float32
    for i := 0; i < len(v1); i++ {
        diff := v1[i] - v2[i]
        sum += diff * diff
    }
    return sum
}

2.2.2 余弦相似度 (Cosine Similarity)

余弦相似度衡量的是两个向量在多维空间中方向的相似性,而非它们的大小。它计算的是两个向量夹角的余弦值。余弦值越接近 1,表示方向越一致,相似度越高;越接近 -1,表示方向越相反;为 0 则表示正交(不相关)。

对于两个 $D$ 维向量 $A$ 和 $B$,余弦相似度的计算公式为:
$$
text{similarity}(A, B) = frac{A cdot B}{|A| |B|} = frac{sum_{i=1}^{D} a_i bi}{sqrt{sum{i=1}^{D} ai^2} sqrt{sum{i=1}^{D} b_i^2}}
$$
在 HNSW 中,通常会将其转换为距离度量:$1 – text{similarity}(A, B)$。

// CosineSimilarity 计算两个向量的余弦相似度
func CosineSimilarity(v1, v2 Vector) float32 {
    var dotProduct float32
    var normV1, normV2 float32

    for i := 0; i < len(v1); i++ {
        dotProduct += v1[i] * v2[i]
        normV1 += v1[i] * v1[i]
        normV2 += v2[i] * v2[i]
    }

    if normV1 == 0 || normV2 == 0 {
        return 0 // 避免除以零
    }

    return dotProduct / (sqrt(normV1) * sqrt(normV2)) // 这里的sqrt需要math.Sqrtf或自己实现
}

// CosineDistance 转换为距离度量
func CosineDistance(v1, v2 Vector) float32 {
    return 1.0 - CosineSimilarity(v1, v2)
}

sqrt 函数可以从 math 包中引入 float32 版本,例如 float32(math.Sqrt(float64(normV1)))

2.2.3 内积 (Dot Product)

内积是两个向量对应元素乘积之和。在某些场景下,尤其是当向量已经经过归一化(即 $|A|=1, |B|=1$)时,内积与余弦相似度是等价的。内积越大,相似度越高。

$$
text{similarity}(A, B) = A cdot B = sum_{i=1}^{D} a_i b_i
$$
作为距离度量,通常用 $-A cdot B$ 或 $C – A cdot B$(其中 $C$ 是一个足够大的常数,确保距离非负)来表示。

// DotProduct 计算两个向量的内积
func DotProduct(v1, v2 Vector) float32 {
    var sum float32
    for i := 0; i < len(v1); i++ {
        sum += v1[i] * v2[i]
    }
    return sum
}

// DotProductDistance 转换为距离度量 (负内积,距离越小相似度越高)
func DotProductDistance(v1, v2 Vector) float32 {
    return -DotProduct(v1, v2)
}

我们的 HNSW 实现需要一个统一的距离计算接口,通常会要求距离越小表示越相似。因此,对于余弦相似度和内积,我们需要将其转换为距离形式。

// DistanceFunc 定义一个距离计算函数的接口
type DistanceFunc func(v1, v2 Vector) float32

3. 暴力搜索的局限性

在了解了向量和距离度量之后,最直观的检索方法是暴力搜索(Brute-Force Search),即计算查询向量与数据库中所有向量的距离,然后选取距离最近的 K 个。

假设数据库中有 $N$ 个 $D$ 维向量,查询一个向量需要进行 $N$ 次距离计算。每次距离计算的复杂度为 $O(D)$。因此,暴力搜索的总时间复杂度为 $O(N times D)$。

当 $N$ 达到亿级($10^8$)甚至更高,维度 $D$ 达到数百甚至上千时:

  • $N = 10^8$, $D = 768$ (BERT 模型的常见维度)
  • 总计算量约为 $10^8 times 768 approx 7.68 times 10^{10}$ 次浮点运算。
  • 即使现代 CPU 每秒能执行数十亿次浮点运算,完成一次查询也需要数十秒,这远不能满足毫秒级的实时检索要求。

因此,暴力搜索对于大规模向量数据集是不可行的,我们必须转向近似最近邻(ANN)搜索算法,HNSW 就是其中最先进、最广泛应用的算法之一。

4. HNSW 算法详解

HNSW (Hierarchical Navigable Small World) 算法是一种高效的近似最近邻搜索算法,它结合了图结构和多层跳表(skip list)的思想,能够在高维空间中实现快速且高质量的搜索。

4.1 核心思想:可导航小世界图

HNSW 的基础是“小世界图”(Small-World Graph)的概念。小世界图的特点是:

  1. 高聚类系数:图中的节点倾向于形成紧密的社群。
  2. 短平均路径长度:任意两个节点之间通常可以通过少量步骤连接起来。

“可导航小世界”(Navigable Small World, NSW)图在此基础上增加了“可导航”的特性,意味着我们可以通过一种贪婪的策略,从图中的任意一点出发,通过沿着距离目标更近的边移动,快速找到目标节点。

然而,原始的 NSW 算法在插入节点时会产生大量连接,导致内存消耗过大且搜索效率在大规模数据集上仍有瓶颈。HNSW 算法通过引入“分层结构”解决了这些问题。

4.2 HNSW 的分层结构

HNSW 将 NSW 图组织成一个多层结构,类似于一个概率跳表:

  • 顶层(Layer $L$):包含最少的节点,节点之间的连接稀疏且跳跃距离远。这一层用于快速定位到目标区域的大致范围。
  • 中间层(Layer $L-1, …, 1$):节点数量逐渐增多,连接密度逐渐增大,跳跃距离逐渐减小。
  • 底层(Layer $0$):包含所有节点,节点之间的连接最密集,用于精确搜索。

每个节点在图中的层数是随机确定的,通常服从指数衰减分布。这意味着大部分节点只存在于底层,而只有少数节点会出现在顶层。这种设计使得搜索可以从顶层快速缩小范围,然后逐层向下,在更密集的连接中进行局部精确搜索。

4.3 HNSW 的数据结构

HNSW 图由一系列节点组成。每个节点表示一个向量,并存储其在不同层上的连接。

// Node 表示 HNSW 图中的一个节点
type Node struct {
    ID        int        // 向量的唯一标识符
    Vector    Vector     // 实际的向量数据
    // Connections 存储不同层级的连接。
    // Connections[layer_idx] 是一个整数切片,存储该层连接到的其他节点的ID。
    Connections [][]int
    // 我们可以选择在Node中存储其最高的层级
    // maxLayer int
}

// HNSW 结构体
type HNSW struct {
    nodes        sync.Map          // 存储所有节点,key为节点ID,value为*Node
    entryPointID int               // 顶层的入口点ID
    maxLayer     int               // 当前HNSW中所有节点所达到的最高层级
    dims         int               // 向量维度
    distanceFunc DistanceFunc      // 距离计算函数
    rng          *rand.Rand        // 随机数生成器,用于确定节点层级

    // HNSW参数
    M            int     // 每个节点在构建时最多连接的邻居数量
    efConstruction int     // 构建时,在搜索每个新节点的邻居时,考虑的最近邻候选数量
    efSearch     int     // 搜索时,在搜索过程中维护的最近邻候选数量
    maxLayers    int     // HNSW图中允许的最大层数
    levelMultiplier float64 // 层级随机数生成器的参数

    // 为了并发安全,在AddVector和Search等操作中需要锁
    // 对于nodes的增删改,需要一个全局的写锁或对每个节点使用读写锁
    // 简单起见,这里先用sync.Map处理节点存储,它本身是并发安全的。
    // 但entryPointID和maxLayer的更新需要额外同步
    mu sync.RWMutex // 用于保护entryPointID和maxLayer的读写
}

// HNSWParams 封装 HNSW 的所有可配置参数
type HNSWParams struct {
    M               int      // 每个节点在构建时最多连接的邻居数量 (通常 10-20)
    EfConstruction  int      // 构建时,在搜索每个新节点的邻居时,考虑的最近邻候选数量 (通常 100-200)
    EfSearch        int      // 搜索时,在搜索过程中维护的最近邻候选数量 (通常 efConstruction 或更大)
    MaxLayers       int      // HNSW图中允许的最大层数 (通常 log2(N))
    LevelMultiplier float64  // 用于计算节点层级的指数分布参数 (通常 1/log(M))
    Seed            int64    // 随机数种子
}

// 默认参数
func DefaultHNSWParams() HNSWParams {
    return HNSWParams{
        M:               16,
        EfConstruction:  200,
        EfSearch:        200, // 通常与EfConstruction相同或更大
        MaxLayers:       10,  // 可以根据数据量调整,log2(N)
        LevelMultiplier: 1.0 / math.Log(16.0), // 1/log(M)
        Seed:            time.Now().UnixNano(),
    }
}

4.4 HNSW 图的构建 (AddVector)

向 HNSW 图中添加一个新向量是构建索引的关键过程。每次添加操作包括确定新节点的层级、在各层寻找邻居并建立连接。

插入过程概览:

  1. 生成随机层级 l:为新节点随机分配一个最高层级 l,通常使用指数分布。
  2. 初始化入口点 ep:如果图为空,新节点成为入口点。否则,从当前全局入口点 ep 开始搜索。
  3. 从顶层向下搜索:从 maxLayer 层(或新节点层级 l 和当前全局 maxLayer 中较小的一个)开始,逐层向下,使用贪婪搜索找到离查询向量最近的 efConstruction 个邻居。这些邻居将作为下一层的搜索入口点。
  4. l 层到 0 层建立连接:对于从 l 层到 0 层的每一层:
    • 在当前层中执行搜索,找到 efConstruction 个最近邻候选。
    • 从这些候选邻居中,选择 M 个最佳邻居与新节点连接。
    • 同时,对于被选中的邻居,如果它们的连接数未达到上限,也需要与新节点建立反向连接。
    • 这个选择邻居的过程涉及到复杂的启发式剪枝策略(selectNeighbors),以保证图的质量和连通性。

4.4.1 随机层级生成

每个新节点被分配一个随机层级 l,其概率 $P(l) propto e^{-l / M_{L}}$,其中 $M_L$ 是一个常数(通常为 $1/ln M$)。这意味着大部分节点只会出现在底层,而只有少数节点会到达顶层。

// generateLevel 根据指数分布生成一个随机层级
func (h *HNSW) generateLevel() int {
    return int(-math.Log(h.rng.Float64()) * h.levelMultiplier)
}

4.4.2 searchLayer:层内贪婪搜索

这是 HNSW 搜索的核心。给定一个查询向量 q、一个入口点 ep、一个搜索层 lc 和一个候选集大小 efsearchLayer 函数会在当前层 lc 中执行贪婪搜索,找出 ef 个最近的节点。

它使用两个优先队列:

  • candidates (Min-Priority Queue):存储待访问的节点,按到查询向量的距离升序排列。
  • result (Max-Priority Queue):存储当前找到的 ef 个最近邻,按到查询向量的距离降序排列。
// searchLayer 在指定层级进行贪婪搜索
// q: 查询向量
// entryPoint: 当前层的搜索入口点
// searchLayer: 要搜索的层级
// ef: 返回的最近邻候选数量
// 返回值: 包含ef个最近邻的优先级队列
func (h *HNSW) searchLayer(q Vector, entryPoint *Node, currentLayer int, ef int) *PriorityQueue {
    // visited 集合用于避免重复访问节点
    visited := make(map[int]struct{})

    // candidates 存储待探索的节点,按距离q的距离升序排列 (Min-Priority Queue)
    candidates := NewPriorityQueue(true) // true for min-heap
    // result 存储找到的ef个最近邻,按距离q的距离降序排列 (Max-Priority Queue)
    result := NewPriorityQueue(false) // false for max-heap

    if entryPoint == nil {
        return result
    }

    dist := h.distanceFunc(q, entryPoint.Vector)
    candidates.Push(entryPoint.ID, dist)
    result.Push(entryPoint.ID, dist)
    visited[entryPoint.ID] = struct{}{}

    for candidates.Len() > 0 {
        curr := candidates.Pop() // 弹出距离最小的节点

        // 如果当前节点距离比结果队列中最远的节点还要远,且结果队列已满,则停止
        // 这是一个优化,但需要谨慎使用,特别是EfConstruction可能需要更远的探索
        if curr.Distance > result.Top().Distance && result.Len() == ef {
            break
        }

        currNode, ok := h.nodes.Load(curr.ID)
        if !ok {
            continue // 节点可能已被删除或不存在
        }

        node := currNode.(*Node)

        // 遍历当前节点在 currentLayer 的所有连接
        for _, neighborID := range node.Connections[currentLayer] {
            if _, seen := visited[neighborID]; !seen {
                visited[neighborID] = struct{}{}

                neighborNodeVal, ok := h.nodes.Load(neighborID)
                if !ok {
                    continue
                }
                neighborNode := neighborNodeVal.(*Node)

                neighborDist := h.distanceFunc(q, neighborNode.Vector)

                if result.Len() < ef || neighborDist < result.Top().Distance {
                    candidates.Push(neighborID, neighborDist)
                    result.Push(neighborID, neighborDist)

                    // 保持 result 队列大小不超过 ef
                    for result.Len() > ef {
                        result.Pop()
                    }
                }
            }
        }
    }
    return result
}

4.4.3 selectNeighbors:选择并剪枝邻居

searchLayer 找到 efConstruction 个候选邻居后,我们需要从这些候选邻居中选择 M 个最佳邻居来连接新节点。这个过程并非简单地选择最近的 M 个,而是要考虑图的连通性和质量,避免生成“劣质”的连接。HNSW 论文提出了两种策略:

  1. Simple Heuristic (简单启发式):选择距离最近的 M 个邻居。
  2. Extend Heuristic (扩展启发式):在选择 M 个邻居时,不仅考虑距离,还考虑这些邻居是否能够到达已经被选中的其他邻居。这个过程旨在确保新节点的邻居能够形成一个紧密的簇,同时保证图的连通性。

我们这里实现一个简化的版本,通常在实际应用中,会采用更复杂的 selectNeighbors 策略,例如 selectNeighborsHeuristic,它会选择 M 个邻居,并尝试避免选择与已选邻居非常近的邻居,以保持图的均匀性。

// selectNeighbors 从候选集中选择M个最佳邻居
// q: 查询向量
// candidates: 从searchLayer返回的候选邻居 (Max-Priority Queue)
// M: 要选择的邻居数量
// Returns: 选定的邻居ID列表
func (h *HNSW) selectNeighbors(q Vector, candidates *PriorityQueue, M int) []int {
    // 将 Max-PQ 转换为 Min-PQ,以便按距离升序处理
    minPQ := NewPriorityQueue(true)
    for candidates.Len() > 0 {
        item := candidates.Pop()
        minPQ.Push(item.ID, item.Distance)
    }

    selectedNeighbors := make([]int, 0, M)

    // 这里实现一个简化版本:直接选择距离最近的 M 个
    for i := 0; i < M && minPQ.Len() > 0; i++ {
        item := minPQ.Pop()
        selectedNeighbors = append(selectedNeighbors, item.ID)
    }

    // 复杂的启发式剪枝会在这里进行,以优化图结构
    // 例如,确保选择的邻居之间距离不会太近,以提高多样性

    return selectedNeighbors
}

4.4.4 AddVector 实现

现在我们可以将上述组件组合起来,实现 AddVector 方法。

// AddVector 将一个新向量及关联ID添加到HNSW索引中
func (h *HNSW) AddVector(id int, vector Vector) error {
    if len(vector) != h.dims {
        return fmt.Errorf("vector dimension mismatch: expected %d, got %d", h.dims, len(vector))
    }

    newNode := &Node{
        ID:        id,
        Vector:    vector,
        Connections: make([][]int, h.maxLayers), // 为所有可能的层级预分配空间
    }
    for i := range newNode.Connections {
        newNode.Connections[i] = make([]int, 0, h.M*2) // 预留M*2空间,因为邻居也可能连接回来
    }

    h.mu.Lock()
    currentMaxLayer := h.maxLayer
    currentEntryPointID := h.entryPointID
    h.mu.Unlock()

    // 1. 生成新节点的随机层级
    newLevel := h.generateLevel()
    if newLevel >= h.maxLayers {
        newLevel = h.maxLayers - 1 // 限制在最大层数内
    }

    // 2. 如果HNSW为空,新节点成为入口点
    if currentEntryPointID == -1 {
        h.mu.Lock()
        h.entryPointID = id
        h.maxLayer = newLevel
        h.mu.Unlock()
        h.nodes.Store(id, newNode)
        return nil
    }

    // 3. 从顶层向下搜索,找到各层的入口点
    // currentEntryPointNode 追踪当前搜索的入口节点
    var currentEntryPointNode *Node
    epVal, ok := h.nodes.Load(currentEntryPointID)
    if !ok {
        return fmt.Errorf("entry point node %d not found", currentEntryPointID)
    }
    currentEntryPointNode = epVal.(*Node)

    // 从最高层开始向下搜索,直到新节点的层级
    for level := currentMaxLayer; level > newLevel; level-- {
        // 在当前层找到一个更近的入口点
        // searchLayer返回的是Min-PQ,我们只需要最顶部的1个作为下一层的入口
        candidates := h.searchLayer(vector, currentEntryPointNode, level, 1)
        if candidates.Len() > 0 {
            epItem := candidates.Pop()
            epNodeVal, ok := h.nodes.Load(epItem.ID)
            if !ok {
                return fmt.Errorf("searchLayer returned invalid node %d", epItem.ID)
            }
            currentEntryPointNode = epNodeVal.(*Node)
        }
    }

    // 4. 在新节点的层级到第0层建立连接
    for level := min(newLevel, currentMaxLayer); level >= 0; level-- {
        // 在当前层进行搜索,找到efConstruction个最近邻
        candidates := h.searchLayer(vector, currentEntryPointNode, level, h.efConstruction)

        // 从候选集中选择M个最佳邻居
        selectedNeighborsIDs := h.selectNeighbors(vector, candidates, h.M)

        // 将新节点连接到选定的M个邻居
        newNode.Connections[level] = append(newNode.Connections[level], selectedNeighborsIDs...)

        // 反向连接:确保被选中的邻居也连接回新节点
        for _, neighborID := range selectedNeighborsIDs {
            neighborVal, ok := h.nodes.Load(neighborID)
            if !ok {
                continue
            }
            neighborNode := neighborVal.(*Node)

            // 这里需要并发安全的修改邻居节点的连接
            // 对于实际的生产级HNSW,Node的Connections字段应该有自己的锁,
            // 或者通过CAS操作进行无锁更新。这里简化为直接修改
            neighborNode.Connections[level] = append(neighborNode.Connections[level], id)

            // 如果邻居节点的连接数超过M,需要进行剪枝
            if len(neighborNode.Connections[level]) > h.M {
                h.pruneConnections(neighborNode, vector, level)
            }
        }

        // 更新下一层的入口点 (最接近新节点的一个)
        if candidates.Len() > 0 {
            epItem := candidates.Top() // 此时candidates是Min-PQ,top是最接近的
            epNodeVal, ok := h.nodes.Load(epItem.ID)
            if !ok {
                return fmt.Errorf("searchLayer returned invalid node %d", epItem.ID)
            }
            currentEntryPointNode = epNodeVal.(*Node)
        }
    }

    // 存储新节点
    h.nodes.Store(id, newNode)

    // 更新全局最高层级和入口点(如果新节点层级更高)
    h.mu.Lock()
    if newLevel > h.maxLayer {
        h.maxLayer = newLevel
        h.entryPointID = id // 新节点成为新的入口点
    }
    h.mu.Unlock()

    return nil
}

// pruneConnections 对节点的指定层连接进行剪枝,使其不超过M个
func (h *HNSW) pruneConnections(node *Node, queryVector Vector, level int) {
    // 简单的剪枝策略:保留M个距离最近的连接
    // 实际生产中会使用更复杂的启发式剪枝

    // 创建一个临时的优先级队列来存储当前连接,并按距离排序
    pq := NewPriorityQueue(true) // Min-PQ

    for _, neighborID := range node.Connections[level] {
        neighborVal, ok := h.nodes.Load(neighborID)
        if !ok {
            continue
        }
        neighborNode := neighborVal.(*Node)
        dist := h.distanceFunc(queryVector, neighborNode.Vector)
        pq.Push(neighborID, dist)
    }

    // 重新构建连接列表,只保留M个最近的
    node.Connections[level] = make([]int, 0, h.M)
    for i := 0; i < h.M && pq.Len() > 0; i++ {
        item := pq.Pop()
        node.Connections[level] = append(node.Connections[level], item.ID)
    }
}

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

4.5 HNSW 图的搜索 (Search)

HNSW 的搜索过程与插入过程中的 searchLayer 类似,也是一个从顶层向下逐层细化的过程。

搜索概览:

  1. 从顶层入口点开始:从 maxLayer 层开始,使用贪婪搜索找到离查询向量最近的节点。
  2. 逐层向下搜索:将上一步找到的最近节点作为下一层的入口点,重复贪婪搜索,直到第 0 层。在每一层,只返回一个最接近的节点作为下一层的入口。
  3. 在第 0 层执行扩展搜索:在第 0 层,我们不再只返回一个入口点,而是扩展搜索到 efSearch 个最近邻,以确保找到高质量的 K 个结果。
  4. 返回 K 个最近邻:从第 0 层找到的 efSearch 个候选邻居中,选择距离查询向量最近的 K 个作为最终结果。
// Search 执行K-近邻搜索
// q: 查询向量
// k: 要返回的最近邻数量
// 返回值: 包含k个最近邻的ID和距离的切片
func (h *HNSW) Search(q Vector, k int) ([]NearestNeighbor, error) {
    if len(q) != h.dims {
        return nil, fmt.Errorf("query vector dimension mismatch: expected %d, got %d", h.dims, len(q))
    }

    h.mu.RLock()
    currentMaxLayer := h.maxLayer
    currentEntryPointID := h.entryPointID
    h.mu.RUnlock()

    if currentEntryPointID == -1 {
        return []NearestNeighbor{}, nil // 索引为空
    }

    epVal, ok := h.nodes.Load(currentEntryPointID)
    if !ok {
        return nil, fmt.Errorf("entry point node %d not found", currentEntryPointID)
    }
    currentEntryPointNode := epVal.(*Node)

    // 1. 从顶层向下搜索,找到第0层的入口点
    // 在搜索层过程中,只关注找到一个最佳入口点,所以ef为1
    for level := currentMaxLayer; level > 0; level-- {
        candidates := h.searchLayer(q, currentEntryPointNode, level, 1)
        if candidates.Len() > 0 {
            epItem := candidates.Pop()
            epNodeVal, ok := h.nodes.Load(epItem.ID)
            if !ok {
                return nil, fmt.Errorf("searchLayer returned invalid node %d", epItem.ID)
            }
            currentEntryPointNode = epNodeVal.(*Node)
        } else {
            // 如果在某个层级没有找到候选,说明图可能不连通或入口点有问题
            // 此时,只能从当前入口点继续向下,或者回退到全局入口点
            // 这里简单处理为继续使用当前入口点
        }
    }

    // 2. 在第0层执行扩展搜索,找到efSearch个最近邻
    // 此时 ef 设置为 efSearch,以获取更广的搜索范围和更高质量的结果
    finalCandidates := h.searchLayer(q, currentEntryPointNode, 0, h.efSearch)

    // 3. 从最终候选集中选择k个最近邻
    results := make([]NearestNeighbor, 0, k)
    // finalCandidates 是 Max-PQ,所以需要弹出k次来获取最近的
    // 或者将其转换为Min-PQ再弹出

    // 将 Max-PQ 转换为 Min-PQ,以便按距离升序处理
    minPQ := NewPriorityQueue(true)
    for finalCandidates.Len() > 0 {
        item := finalCandidates.Pop()
        minPQ.Push(item.ID, item.Distance)
    }

    for i := 0; i < k && minPQ.Len() > 0; i++ {
        item := minPQ.Pop()
        results = append(results, NearestNeighbor{ID: item.ID, Distance: item.Distance})
    }

    return results, nil
}

// NearestNeighbor 结构体用于存储搜索结果
type NearestNeighbor struct {
    ID       int
    Distance float32
}

4.6 优先级队列 (PriorityQueue) 实现

HNSW 算法大量依赖于优先级队列来管理候选节点和结果集。Go 标准库 container/heap 提供了实现堆的接口,我们可以基于此构建优先级队列。

import (
    "container/heap"
    "fmt"
    "math"
    "math/rand"
    "sync"
    "time"
)

// Item 是优先级队列中的元素
type Item struct {
    ID       int
    Distance float32 // 距离,作为优先级
    index    int     // 堆中元素的索引,用于更新
}

// PriorityQueue 实现 heap.Interface 接口
type PriorityQueue struct {
    items []*Item
    isMin bool // true for min-heap, false for max-heap
}

// NewPriorityQueue 创建一个新的优先级队列
func NewPriorityQueue(isMin bool) *PriorityQueue {
    pq := &PriorityQueue{
        items: make([]*Item, 0),
        isMin: isMin,
    }
    heap.Init(pq)
    return pq
}

func (pq *PriorityQueue) Len() int { return len(pq.items) }

func (pq *PriorityQueue) Less(i, j int) bool {
    if pq.isMin {
        return pq.items[i].Distance < pq.items[j].Distance // Min-heap: 距离越小优先级越高
    }
    return pq.items[i].Distance > pq.items[j].Distance // Max-heap: 距离越大优先级越高
}

func (pq *PriorityQueue) Swap(i, j int) {
    pq.items[i], pq.items[j] = pq.items[j], pq.items[i]
    pq.items[i].index = i
    pq.items[j].index = j
}

func (pq *PriorityQueue) Push(x interface{}) {
    item := x.(*Item)
    item.index = len(pq.items)
    pq.items = append(pq.items, item)
}

func (pq *PriorityQueue) Pop() interface{} {
    old := pq.items
    n := len(old)
    item := old[n-1]
    item.index = -1 // for safety
    pq.items = old[0 : n-1]
    return item
}

// Push 将一个新元素推入队列
func (pq *PriorityQueue) Push(id int, distance float32) {
    item := &Item{ID: id, Distance: distance}
    heap.Push(pq, item)
}

// Pop 弹出队列中优先级最高的元素
func (pq *PriorityQueue) Pop() *Item {
    if pq.Len() == 0 {
        return nil
    }
    return heap.Pop(pq).(*Item)
}

// Top 返回队列中优先级最高的元素,不弹出
func (pq *PriorityQueue) Top() *Item {
    if pq.Len() == 0 {
        return nil
    }
    return pq.items[0]
}

4.7 HNSW 参数的含义与影响

HNSW 算法的性能和召回率高度依赖于其参数的配置。

参数 含义 影响 典型值
M 每个节点在构建时最多连接的邻居数量 越大,召回率越高,搜索更准确,但内存和构建时间增加 10-20
efConstruction 构建时,在搜索每个新节点的邻居时,考虑的最近邻候选数量 越大,图的质量越高,召回率越高,但构建时间增加 100-200
efSearch 搜索时,在搜索过程中维护的最近邻候选数量 越大,召回率越高,搜索更准确,但查询延迟增加 100-500
maxLayers HNSW图中允许的最大层数 越大,层级结构越深,搜索路径可能更短,但内存占用略增 $log_2 N$ 的数量级
levelMultiplier 用于计算节点层级的指数分布参数 影响高层节点的数量和分布,通常设置为 $1/ln(M)$ $1/ln(M)$

5. Go 语言实现 HNSW

在上一节,我们已经详细讲解了 HNSW 的算法原理和 Go 语言的核心代码结构。现在,我们来更完整地组织这些代码,以构建一个可用的 HNSW 索引。

package hnsw

import (
    "container/heap"
    "fmt"
    "math"
    "math/rand"
    "sync"
    "time"
)

// Vector 类型定义
type Vector []float32

// DistanceFunc 定义距离计算函数的接口
type DistanceFunc func(v1, v2 Vector) float32

// L2DistanceSquared 计算两个向量的欧氏距离的平方
func L2DistanceSquared(v1, v2 Vector) float32 {
    var sum float32
    for i := 0; i < len(v1); i++ {
        diff := v1[i] - v2[i]
        sum += diff * diff
    }
    return sum
}

// CosineDistance 计算两个向量的余弦距离 (1 - CosineSimilarity)
func CosineDistance(v1, v2 Vector) float32 {
    var dotProduct float32
    var normV1, normV2 float32

    for i := 0; i < len(v1); i++ {
        dotProduct += v1[i] * v2[i]
        normV1 += v1[i] * v1[i]
        normV2 += v2[i] * v2[i]
    }

    if normV1 == 0 || normV2 == 0 {
        return 1.0 // 避免除以零,零向量与任何向量都不相似
    }

    similarity := dotProduct / (float32(math.Sqrt(float64(normV1))) * float32(math.Sqrt(float64(normV2))))
    return 1.0 - similarity
}

// Node 表示 HNSW 图中的一个节点
type Node struct {
    ID          int        // 向量的唯一标识符
    Vector      Vector     // 实际的向量数据
    Connections [][]int    // Connections[layer_idx] 是一个整数切片,存储该层连接到的其他节点的ID
    mu          sync.RWMutex // 用于保护Connections的并发访问
}

// HNSWParams 封装 HNSW 的所有可配置参数
type HNSWParams struct {
    M               int      // 每个节点在构建时最多连接的邻居数量
    EfConstruction  int      // 构建时,在搜索每个新节点的邻居时,考虑的最近邻候选数量
    EfSearch        int      // 搜索时,在搜索过程中维护的最近邻候选数量
    MaxLayers       int      // HNSW图中允许的最大层数
    LevelMultiplier float64  // 用于计算节点层级的指数分布参数 (通常 1/log(M))
    Seed            int64    // 随机数种子
}

// DefaultHNSWParams 提供一组默认参数
func DefaultHNSWParams() HNSWParams {
    return HNSWParams{
        M:               16,
        EfConstruction:  200,
        EfSearch:        200,
        MaxLayers:       10, // 通常根据数据量 N 调整, 例如 log2(N)
        LevelMultiplier: 1.0 / math.Log(16.0), // 1/log(M)
        Seed:            time.Now().UnixNano(),
    }
}

// HNSW 结构体
type HNSW struct {
    nodes        sync.Map          // 存储所有节点,key为节点ID,value为*Node
    entryPointID int               // 顶层的入口点ID
    maxLayer     int               // 当前HNSW中所有节点所达到的最高层级
    dims         int               // 向量维度
    distanceFunc DistanceFunc      // 距离计算函数
    rng          *rand.Rand        // 随机数生成器,用于确定节点层级

    params HNSWParams // HNSW参数

    mu sync.RWMutex // 用于保护entryPointID和maxLayer的读写
}

// NewHNSW 创建并初始化一个新的HNSW索引
func NewHNSW(dims int, distFunc DistanceFunc, params HNSWParams) *HNSW {
    if params.M <= 0 {
        params.M = DefaultHNSWParams().M
    }
    if params.EfConstruction <= 0 {
        params.EfConstruction = DefaultHNSWParams().EfConstruction
    }
    if params.EfSearch <= 0 {
        params.EfSearch = DefaultHNSWParams().EfSearch
    }
    if params.MaxLayers <= 0 {
        params.MaxLayers = DefaultHNSWParams().MaxLayers
    }
    if params.LevelMultiplier <= 0 {
        params.LevelMultiplier = DefaultHNSWParams().LevelMultiplier
    }
    if params.Seed == 0 {
        params.Seed = DefaultHNSWParams().Seed
    }

    return &HNSW{
        nodes:        sync.Map{},
        entryPointID: -1, // -1 表示索引为空
        maxLayer:     -1,
        dims:         dims,
        distanceFunc: distFunc,
        rng:          rand.New(rand.NewSource(params.Seed)),
        params:       params,
    }
}

// generateLevel 根据指数分布生成一个随机层级
func (h *HNSW) generateLevel() int {
    return int(-math.Log(h.rng.Float64()) * h.params.LevelMultiplier)
}

// searchLayer 在指定层级进行贪婪搜索
// q: 查询向量
// entryPoint: 当前层的搜索入口点
// currentLayer: 要搜索的层级
// ef: 返回的最近邻候选数量
// 返回值: 包含ef个最近邻的优先级队列 (Max-Priority Queue)
func (h *HNSW) searchLayer(q Vector, entryPoint *Node, currentLayer int, ef int) *PriorityQueue {
    visited := make(map[int]struct{})

    candidates := NewPriorityQueue(true) // Min-heap for candidates to explore
    result := NewPriorityQueue(false)    // Max-heap for the top 'ef' nearest neighbors

    if entryPoint == nil {
        return result
    }

    dist := h.distanceFunc(q, entryPoint.Vector)
    candidates.Push(entryPoint.ID, dist)
    result.Push(entryPoint.ID, dist)
    visited[entryPoint.ID] = struct{}{}

    for candidates.Len() > 0 {
        currItem := candidates.Pop() // 弹出距离最小的节点

        // Optimization: if current candidate is farther than the farthest in result, and result is full, stop
        // This is critical for performance.
        if result.Len() == ef && currItem.Distance > result.Top().Distance {
            break
        }

        currNodeVal, ok := h.nodes.Load(currItem.ID)
        if !ok {
            continue
        }
        currNode := currNodeVal.(*Node)

        currNode.mu.RLock() // Read lock for connections
        connections := currNode.Connections[currentLayer]
        currNode.mu.RUnlock()

        for _, neighborID := range connections {
            if _, seen := visited[neighborID]; !seen {
                visited[neighborID] = struct{}{}

                neighborNodeVal, ok := h.nodes.Load(neighborID)
                if !ok {
                    continue
                }
                neighborNode := neighborNodeVal.(*Node)

                neighborDist := h.distanceFunc(q, neighborNode.Vector)

                if result.Len() < ef || neighborDist < result.Top().Distance {
                    candidates.Push(neighborID, neighborDist)
                    result.Push(neighborID, neighborDist)

                    for result.Len() > ef {
                        result.Pop()
                    }
                }
            }
        }
    }
    return result
}

// selectNeighbors 从候选集中选择M个最佳邻居 (简化版)
func (h *HNSW) selectNeighbors(q Vector, candidates *PriorityQueue, M int) []int {
    // Convert Max-PQ to Min-PQ to get elements in ascending distance order
    minPQ := NewPriorityQueue(true)
    for candidates.Len() > 0 {
        item := candidates.Pop()
        minPQ.Push(item.ID, item.Distance)
    }

    selectedNeighbors := make([]int, 0, M)
    for i := 0; i < M && minPQ.Len() > 0; i++ {
        item := minPQ.Pop()
        selectedNeighbors = append(selectedNeighbors, item.ID)
    }
    return selectedNeighbors
}

// pruneConnections 对节点的指定层连接进行剪枝,使其不超过M个
func (h *HNSW) pruneConnections(node *Node, queryVector Vector, level int) {
    node.mu.Lock()
    defer node.mu.Unlock()

    if len(node.Connections[level]) <= h.params.M {
        return // No need to prune
    }

    pq := NewPriorityQueue(true) // Min-PQ to find M closest neighbors

    // Collect distances for current connections
    for _, neighborID := range node.Connections[level] {
        neighborVal, ok := h.nodes.Load(neighborID)
        if !ok {
            continue
        }
        neighborNode := neighborVal.(*Node)
        dist := h.distanceFunc(queryVector, neighborNode.Vector)
        pq.Push(neighborID, dist)
    }

    // Rebuild connections with only the M closest
    newConnections := make([]int, 0, h.params.M)
    for i := 0; i < h.params.M && pq.Len() > 0; i++ {
        item := pq.Pop()
        newConnections = append(newConnections, item.ID)
    }
    node.Connections[level] = newConnections
}

// AddVector 将一个新向量及关联ID添加到HNSW索引中
func (h *HNSW) AddVector(id int, vector Vector) error {
    if len(vector) != h.dims {
        return fmt.Errorf("vector dimension mismatch: expected %d, got %d", h.dims, len(vector))
    }

    newNode := &Node{
        ID:        id,
        Vector:    vector,
        Connections: make([][]int, h.params.MaxLayers),
    }
    for i := range newNode.Connections {
        // Pre-allocate space for connections in each layer (M*2 for bidirectional)
        newNode.Connections[i] = make([]int, 0, h.params.M*2) 
    }

    h.mu.RLock()
    currentMaxLayer := h.maxLayer
    currentEntryPointID := h.entryPointID
    h.mu.RUnlock()

    // 1. 生成新节点的随机层级
    newLevel := h.generateLevel()
    if newLevel >= h.params.MaxLayers {
        newLevel = h.params.MaxLayers - 1 // Clamp to MaxLayers
    }

    // 2. 如果HNSW为空,新节点成为入口点
    if currentEntryPointID == -1 {
        h.mu.Lock()
        h.entryPointID = id
        h.maxLayer = newLevel
        h.mu.Unlock()
        h.nodes.Store(id, newNode)
        return nil
    }

    // 3. 从顶层向下搜索,找到各层的入口点
    var currentEntryPointNode *Node
    epVal, ok := h.nodes.Load(currentEntryPointID)
    if !ok {
        return fmt.Errorf("entry point node %d not found", currentEntryPointID)
    }
    currentEntryPointNode = epVal.(*Node)

    // Search from the current top layer down to the new node's level (or current max layer if lower)
    for level := currentMaxLayer; level > newLevel; level-- {
        candidates := h.searchLayer(vector, currentEntryPointNode, level, 1) // Only need 1 closest for entry point
        if candidates.Len() > 0 {
            epItem := candidates.Pop()
            epNodeVal, ok := h.nodes.Load(epItem.ID)
            if !ok {
                return fmt.Errorf("searchLayer returned invalid node %d", epItem.ID)
            }
            currentEntryPointNode = epNodeVal.(*Node)
        }
    }

    // 4. 在新节点的层级到第0层建立连接
    // Iterate downwards from the minimum of newLevel and currentMaxLayer
    // This ensures we don't try to connect to non-existent higher layers
    for level := min(newLevel, currentMaxLayer); level >= 0; level-- {
        // Find efConstruction nearest neighbors in the current layer
        candidates := h.searchLayer(vector, currentEntryPointNode, level, h.params.EfConstruction)

        // Select M best neighbors (simplified: M closest)
        selectedNeighborsIDs := h.selectNeighbors(vector, candidates, h.params.M)

        // Add connections from newNode to selected neighbors
        newNode.mu.Lock()
        newNode.Connections[level] = append(newNode.Connections[level], selectedNeighborsIDs...)
        newNode.mu.Unlock()

        // Add reverse connections from selected neighbors to newNode
        for _, neighborID := range selectedNeighborsIDs {
            neighborVal, ok := h.nodes.Load(neighborID)
            if !ok {
                continue
            }
            neighborNode := neighborVal.(*Node)

            neighborNode.mu.Lock()
            neighborNode.Connections[level] = append(neighborNode.Connections[level], id)
            neighborNode.mu.Unlock() // Unlock before pruning, as prune might take a lock itself

            // Prune if neighbor's connections exceed M
            h.pruneConnections(neighborNode, neighborNode.Vector, level)
        }

        // Update entry point for the next lower layer (closest from current candidates)
        if candidates.Len() > 0 {
            // Candidates is a Max-PQ, so its Top() is the farthest.
            // To get the closest from the candidates, we'd need to convert it to Min-PQ
            // or iterate through its elements. For simplicity here, we assume the entryPointNode
            // is updated based on the actual closest neighbor found in the previous search.
            // A more rigorous approach would pick the closest from `candidates` again.
            // For now, we'll keep `currentEntryPointNode` as is for the next iteration.
        }
    }

    // Store the new node
    h.nodes.Store(id, newNode)

    // Update global maxLayer and entryPointID if new node is in a higher layer
    h.mu.Lock()
    if newLevel > h.maxLayer {
        h.maxLayer = newLevel
        h.entryPointID = id // New node becomes the new entry point
    }
    h.mu.Unlock()

    return nil
}

// NearestNeighbor 结构体用于存储搜索结果
type NearestNeighbor struct {
    ID       int
    Distance float32
}

// Search 执行K-近邻搜索
// q: 查询向量
// k: 要返回的最近邻数量
// 返回值: 包含k个最近邻的ID和距离的切片
func (h *HNSW) Search(q Vector, k int) ([]NearestNeighbor, error) {
    if len(q) != h.dims {
        return nil, fmt.Errorf("query vector dimension mismatch: expected %d, got %d", h.dims, len(q))
    }
    if k <= 0 {
        return []NearestNeighbor{}, nil
    }

    h.mu.RLock()
    currentMaxLayer := h.maxLayer
    currentEntryPointID := h.entryPointID
    h.mu.RUnlock()

    if currentEntryPointID == -1 {
        return []NearestNeighbor{}, nil // Index is empty
    }

    epVal, ok := h.nodes.Load(currentEntryPointID)
    if !ok {
        return nil, fmt.Errorf("entry point node %d not found", currentEntryPointID)
    }
    currentEntryPointNode := epVal.(*Node)

    // 1. 从顶层向下搜索,找到第0层的入口点
    // In search layers, we only need to find a single best entry point, so ef is 1.
    for level := currentMaxLayer; level > 0; level-- {
        candidates := h.searchLayer(q, currentEntryPointNode, level, 1)
        if candidates.Len() > 0 {
            epItem := candidates.Pop()
            epNodeVal, ok := h.nodes.Load(epItem.ID)
            if !ok {
                return nil, fmt.Errorf("searchLayer returned invalid node %d", epItem.ID)
            }
            currentEntryPointNode = epNodeVal.(*Node)
        }
    }

    // 2. 在第0层执行扩展搜索,找到efSearch个最近邻
    // Here, ef is set to efSearch to get a wider search scope and higher quality results.
    finalCandidates := h.searchLayer(q, currentEntryPointNode, 0, h.params.EfSearch)

    // 3. 从最终候选集中选择k个最近邻
    results := make([]NearestNeighbor, 0, k)

    // Convert Max-PQ to Min-PQ to retrieve the closest `k` items
    minPQ := NewPriorityQueue(true)
    for finalCandidates.Len() > 0 {
        item := finalCandidates.Pop()
        minPQ.Push(item.ID, item.Distance)
    }

    for i := 0; i < k && minPQ.Len() > 0; i++ {
        item := minPQ.Pop()
        results = append(results, NearestNeighbor{ID: item.ID, Distance: item.Distance})
    }

    return results, nil
}

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

// --- PriorityQueue Implementation (as defined in 4.6) ---
// Item, PriorityQueue, NewPriorityQueue, Len, Less, Swap, Push, Pop, Top methods would be here.
// For brevity, it's not repeated, but it should be part of the hnsw package.

上述 Go 语言实现包含了 HNSW 的核心逻辑,包括节点结构、HNSW 主结构、参数配置、距离函数、优先级队列、以及 AddVectorSearch 方法。
需要注意的是,Node 结构体中的 Connections 字段在并发写入时需要额外的同步保护。这里为每个 Node 添加了 sync.RWMutex 来保护其连接列表。sync.Map 保证了 nodes 集合本身的并发安全,但单个 Node 内部字段的并发修改仍需手动处理。

6. 性能优化与亿级向量挑战

实现 HNSW 只是第一步,要应对亿级向量的毫秒级检索,还需要进行大量的性能优化和系统设计。

6.1 内存管理

亿级向量意味着巨大的内存占用。例如,一个 768 维的 float32 向量需要 768 * 4 = 3072 字节。1 亿个向量就是 3072 * 10^8 ≈ 300 GB。这还不包括 HNSW 图本身的连接信息。

  • 数据类型选择float32 通常足够满足精度要求,相比 float64 可以节省一半内存。
  • 紧凑存储:确保 Node 结构体尽可能紧凑,减少内存碎片。
  • MMap 技术:对于海量向量数据,可以考虑使用内存映射文件(mmap)。将向量数据存储在磁盘上,通过 mmap 映射到进程地址空间。操作系统会按需将数据页载入内存,有效管理内存使用,并利用操作系统的页缓存机制。Go 语言可以通过 syscall 包或第三方库实现 mmap
  • 连接优化:HNSW 的连接信息也可能占用大量内存。可以考虑:
    • 动态调整 M:在构建索引时,对于不同层级的节点,M 可以不是固定值。
    • 存储相对索引:如果节点 ID 是连续的,可以存储相对偏移量。
    • 稀疏连接压缩:对于连接较多的情况,考虑压缩存储。

6.2 并发与并行

Go 语言以其高效的并发模型而闻名,这对于构建高性能的向量数据库至关重要。

  • Goroutine 与 Channel
    • 构建时并发AddVector 操作可以并发执行。多个 Goroutine 可以同时添加向量,但需要解决并发冲突:
      • HNSW 结构体中的 entryPointIDmaxLayer 需要通过 sync.RWMutex 保护。
      • 每个 NodeConnections 字段也需要 sync.RWMutex 保护,以允许并发读写其连接列表。
      • sync.Map 已经提供了对节点集合的并发安全访问。
    • 搜索时并发Search 操作通常是读操作,可以并发执行。searchLayer 函数不需要写锁,可以被多个 Goroutine 同时调用。
  • 批处理:对于大规模数据导入,将向量分批次添加可以减少锁竞争,提高吞吐量。

6.3 距离计算优化

距离计算是 HNSW 算法中最频繁执行的操作,其效率直接影响整体性能。

  • SIMD 指令:现代 CPU 支持单指令多数据(Single Instruction, Multiple Data, SIMD)指令集(如 SSE、AVX)。这些指令允许 CPU 同时对多个数据元素执行相同的操作,大幅加速向量运算。Go 语言可以通过以下方式利用 SIMD:
    • Cgo 调用:编写 C/C++ 代码并使用 SIMD 内联函数,然后通过 Cgo 在 Go 中调用。
    • Go 汇编:直接编写 Go 汇编代码,利用特定 CPU 架构的 SIMD 指令。
    • 向量库:使用已经优化好的 Go 向量库(如果存在)。
  • 预计算与缓存:对于频繁查询的向量或其归一化值,可以进行预计算并缓存,避免重复计算。

6.4 网络与分布式

亿级向量往往需要分布式部署才能承载。

  • 分片 (Sharding):将总数据集分成多个子集,每个子集存储在一个 HNSW 实例中,并部署到不同的服务器节点上。查询时,将查询向量发送到所有分片,然后聚合结果。
  • 负载均衡:使用负载均衡器(如 Nginx, Envoy, F5)将查询请求分发到不同的 HNSW 服务实例。
  • gRPC/REST API:为 HNSW 内核提供标准的 gRPC 或 RESTful API,方便外部服务调用。gRPC 性能更高,更适合服务间通信。
  • 数据同步与高可用
    • 复制 (Replication):每个分片可以有多个副本,提供高可用性和读扩展性。
    • 一致性哈希:用于将向量均匀地分配到各个分片,并处理节点的增删。
    • Leader-Follower 架构:对于写操作,可以由 Leader 节点处理,然后同步到 Follower 节点。

6.5 亿级向量的策略

  • 多级索引/过滤:在 HNSW 之前,可以增加一个粗粒度过滤层。例如,先通过属性过滤缩小搜索范围,或者使用更快的粗粒度 ANN 算法(如 LSH 或量化索引)快速筛选出一小部分候选集,再在这些候选集上运行 HNSW 进行精确搜索。
  • GPU 加速:对于极端性能要求,可以通过 Cgo 接口调用 CUDA 库,利用 GPU 的大规模并行计算能力加速距离计算。
  • 动态更新与删除:HNSW 原始算法对动态更新和删除支持不够友好。实现动态更新通常需要:
    • 惰性删除:标记节点为删除状态,在搜索时跳过。定期进行垃圾回收和重建。
    • 增量更新:如果只是少量更新,可以重新添加节点并重建其邻居连接。
    • 对于频繁的更新和删除,可能需要考虑其他支持动态性的 ANN 算法或更复杂的 HNSW 变体。

7. 构建完整的向量数据库内核

一个生产级的向量数据库内核不仅仅是 HNSW 索引。它需要包含更多功能来支持实际应用。

  • 数据管理
    • 向量存储:除了 HNSW 索引,还需要持久化存储原始向量数据。可以考虑 MMap 文件、RocksDB、LevelDB 或其他键值存储。
    • 元数据存储:与向量关联的文本、标签、时间戳等元数据也需要存储。
  • API 层:提供易于使用的 gRPC 或 HTTP/REST API 接口,供客户端应用接入。
  • 持久化与恢复
    • HNSW 索引本身需要序列化到磁盘,以便服务重启后快速恢复。这涉及图结构(节点、连接)和参数的保存。
    • 定期备份机制。
  • 监控与可观测性
    • 集成 Prometheus 等监控系统,暴露 HNSW 的性能指标(查询 QPS、延迟、召回率、索引大小等)。
    • 日志记录:记录重要的操作和错误信息,便于故障排查。
  • 容错与高可用
    • 副本机制:多副本部署,当主节点故障时自动切换到备用节点。
    • 数据一致性:确保多副本之间的数据一致性。
  • 过滤功能
    • 前置过滤 (Pre-filtering):在向量搜索之前,根据元数据条件过滤掉不符合要求的向量,只对剩余向量进行 HNSW 搜索。
    • 后置过滤 (Post-filtering):先进行 HNSW 搜索得到 K 个最近邻,再对这些结果进行元数据过滤。通常需要更大的 efSearch 来保证过滤后仍有足够多的结果。
  • 索引重建与优化:随着数据量的变化或参数调整,可能需要定期重建 HNSW 索引以保持最佳性能和召回率。

8. 实践案例与未来展望

HNSW 索引作为向量数据库的核心,在多个领域展现出强大的能力:

  • 推荐系统:为用户生成一个兴趣向量,在商品或内容向量库中快速找到相似项。
  • 图像与视频搜索:将图像/视频编码为向量,实现“以图搜图”或根据语义内容搜索。
  • 智能问答与聊天机器人:将用户问题和知识库内容转化为向量,快速匹配最相关的答案。
  • 药物发现与材料科学:在化学分子或材料特性向量空间中寻找相似结构或性质。

未来,向量数据库和 HNSW 算法将持续演进:

  • 动态性增强:更高效地支持实时插入、更新和删除操作,减少索引重建的需求。
  • 混合搜索:更紧密地结合向量搜索与传统关键词搜索、结构化数据过滤,提供更丰富、更精确的查询能力。
  • 硬件加速:进一步利用 GPU、FPGA 甚至专用 AI 芯片(如 Google TPU)加速向量计算。
  • 新算法:研究和应用如 DiskANN、SPANN 等更适合磁盘存储或特定工作负载的 ANN 算法。
  • 标准化与生态:向量数据库领域将逐步走向标准化,形成更完善的工具链和生态系统。

通过 Go 语言手写 HNSW 索引,我们不仅深入理解了其工作原理,也为构建高性能、可扩展的向量数据库内核奠定了基础。这对于应对未来AI应用中日益增长的向量检索需求至关重要。

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注