从零实现一个向量搜索引擎

// Go 中可以通过 unsafe + 汇编实现 SIMD，这里给出纯 Go 版本作为参考
func L2SquaredDistance(a, b []float32) float32 {
    var sum float32
    for i := range a {
        diff := a[i] - b[i]
        sum += diff * diff
    }
    return sum
}

func InnerProduct(a, b []float32) float32 {
    var sum float32
    for i := range a {
        sum += a[i] * b[i]
    }
    return sum
}

func CosineDistance(a, b []float32) float32 {
    var dot, normA, normB float32
    for i := range a {
        dot += a[i] * b[i]
        normA += a[i] * a[i]
        normB += b[i] * b[i]
    }
    if normA == 0 || normB == 0 {
        return 1.0
    }
    return 1.0 - dot/float32(math.Sqrt(float64(normA)*float64(normB)))
}

// 启发式邻居选择：优先保留那些"不被已选邻居覆盖"的候选
func selectNeighborsHeuristic(candidates []NodeDist, m int, distFunc DistFunc) []uint32 {
    sort.Slice(candidates, func(i, j int) bool {
        return candidates[i].Dist < candidates[j].Dist
    })
    selected := make([]NodeDist, 0, m)
    for _, c := range candidates {
        if len(selected) >= m {
            break
        }
        // 检查 c 是否比所有已选邻居都更"有价值"
        good := true
        for _, s := range selected {
            if distFunc(c.Vec, s.Vec) < c.Dist {
                good = false
                break
            }
        }
        if good {
            selected = append(selected, c)
        }
    }
    ids := make([]uint32, len(selected))
    for i, s := range selected {
        ids[i] = s.ID
    }
    return ids
}

type PQIndex struct {
    M          int          // 子空间数量
    Ksub       int          // 每个子空间的聚类中心数，通常 256
    Dim        int          // 原始向量维度
    Centroids  [][]float32  // [M*Ksub][subDim] 聚类中心
    Codes      [][]byte     // [N][M] 压缩编码
}

// 构建查距离表：query 到每个子空间每个聚类中心的距离
func (pq *PQIndex) BuildDistTable(query []float32) [][]float32 {
    subDim := pq.Dim / pq.M
    table := make([][]float32, pq.M)
    for i := 0; i < pq.M; i++ {
        table[i] = make([]float32, pq.Ksub)
        qSub := query[i*subDim : (i+1)*subDim]
        for j := 0; j < pq.Ksub; j++ {
            centroid := pq.Centroids[i*pq.Ksub+j]
            table[i][j] = L2SquaredDistance(qSub, centroid)
        }
    }
    return table
}

// ADC 距离：用查表代替真实距离计算
func (pq *PQIndex) ADCDistance(table [][]float32, code []byte) float32 {
    var dist float32
    for i := 0; i < pq.M; i++ {
        dist += table[i][code[i]]
    }
    return dist
}

const (
    OpInsert byte = 0x01
    OpDelete byte = 0x02
)

type WAL struct {
    file *os.File
    mu   sync.Mutex
}

func OpenWAL(path string) (*WAL, error) {
    f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR|os.O_APPEND, 0644)
    if err != nil {
        return nil, err
    }
    return &WAL{file: f}, nil
}

func (w *WAL) AppendInsert(id uint64, vec []float32) error {
    w.mu.Lock()
    defer w.mu.Unlock()

    vecBytes := float32SliceToBytes(vec)
    payload := make([]byte, 1+8+len(vecBytes))
    payload[0] = OpInsert
    binary.LittleEndian.PutUint64(payload[1:9], id)
    copy(payload[9:], vecBytes)

    return w.writeRecord(payload)
}

func (w *WAL) AppendDelete(id uint64) error {
    w.mu.Lock()
    defer w.mu.Unlock()

    payload := make([]byte, 1+8)
    payload[0] = OpDelete
    binary.LittleEndian.PutUint64(payload[1:9], id)

    return w.writeRecord(payload)
}

func (w *WAL) writeRecord(payload []byte) error {
    length := uint32(len(payload))
    checksum := crc32.ChecksumIEEE(payload)

    header := make([]byte, 8)
    binary.LittleEndian.PutUint32(header[0:4], checksum)
    binary.LittleEndian.PutUint32(header[4:8], length)

    if _, err := w.file.Write(header); err != nil {
        return err
    }
    if _, err := w.file.Write(payload); err != nil {
        return err
    }
    return w.file.Sync()
}

func (w *WAL) Replay(onInsert func(uint64, []float32), onDelete func(uint64)) error {
    w.file.Seek(0, io.SeekStart)
    reader := bufio.NewReader(w.file)

    for {
        header := make([]byte, 8)
        if _, err := io.ReadFull(reader, header); err != nil {
            if err == io.EOF || err == io.ErrUnexpectedEOF {
                return nil
            }
            return err
        }

        checksum := binary.LittleEndian.Uint32(header[0:4])
        length := binary.LittleEndian.Uint32(header[4:8])

        payload := make([]byte, length)
        if _, err := io.ReadFull(reader, payload); err != nil {
            return nil // 部分写入的记录，忽略
        }

        if crc32.ChecksumIEEE(payload) != checksum {
            return nil // CRC 不匹配，说明是不完整记录
        }

        switch payload[0] {
        case OpInsert:
            id := binary.LittleEndian.Uint64(payload[1:9])
            vec := bytesToFloat32Slice(payload[9:])
            onInsert(id, vec)
        case OpDelete:
            id := binary.LittleEndian.Uint64(payload[1:9])
            onDelete(id)
        }
    }
}

func (w *WAL) Truncate() error {
    w.mu.Lock()
    defer w.mu.Unlock()
    return w.file.Truncate(0)
}

func MmapFile(path string, size int) ([]byte, error) {
    f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0644)
    if err != nil {
        return nil, err
    }
    defer f.Close()

    // 确保文件足够大
    if err := f.Truncate(int64(size)); err != nil {
        return nil, err
    }

    data, err := syscall.Mmap(
        int(f.Fd()), 0, size,
        syscall.PROT_READ|syscall.PROT_WRITE,
        syscall.MAP_SHARED,
    )
    if err != nil {
        return nil, err
    }
    return data, nil
}

// 通过 mmap 读写向量，就像操作普通 byte slice 一样
func GetVector(mmapData []byte, id int, dim int) []float32 {
    offset := id * dim * 4
    return bytesToFloat32Slice(mmapData[offset : offset+dim*4])
}

func PostFilterSearch(engine *Engine, query []float32, k int, filter FilterFunc) []Result {
    // 多取一些候选，因为过滤会淘汰一部分
    overFetch := k * 10
    candidates := engine.HNSWSearch(query, overFetch)

    results := make([]Result, 0, k)
    for _, c := range candidates {
        if filter(c.Metadata) {
            results = append(results, c)
            if len(results) >= k {
                break
            }
        }
    }
    return results
}

func PreFilterSearch(engine *Engine, query []float32, k int, filter FilterFunc) []Result {
    // 第一步：元数据过滤，得到 bitmap
    bitmap := engine.MetadataIndex.Filter(filter)
    matchCount := bitmap.Count()

    if matchCount == 0 {
        return nil
    }

    // 如果匹配数量少，直接暴力搜索
    if matchCount < 1000 {
        return engine.BruteForceSearchWithBitmap(query, k, bitmap)
    }

    // 否则在 HNSW 中搜索，但跳过不在 bitmap 中的节点
    return engine.HNSWSearchWithBitmap(query, k, bitmap)
}

package main

import (
    "container/heap"
    "encoding/json"
    "fmt"
    "log"
    "math"
    "math/rand"
    "net/http"
    "sort"
    "sync"
    "sync/atomic"
    "time"
)

// --- 距离函数 ---

func l2Squared(a, b []float32) float32 {
    var sum float32
    for i := range a {
        d := a[i] - b[i]
        sum += d * d
    }
    return sum
}

// --- Min-Heap for topK ---

type Result struct {
    ID   uint64  `json:"id"`
    Dist float32 `json:"distance"`
}

type MaxHeap []Result

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.(Result)) }
func (h *MaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

// --- HNSW 简化实现 ---

const (
    hnswM             = 16
    hnswMmax0         = 32
    hnswEfConstruction = 200
    hnswMl            = 1.0 / math.Ln2 // 1/ln(M)
)

type hnswNode struct {
    id      uint64
    vec     []float32
    level   int
    friends [][]uint32 // friends[layer] = list of neighbor indices
}

type hnswIndex struct {
    nodes      []*hnswNode
    entryPoint int
    maxLevel   int
    idToIdx    map[uint64]int
}

func newHNSW() *hnswIndex {
    return &hnswIndex{
        entryPoint: -1,
        maxLevel:   -1,
        idToIdx:    make(map[uint64]int),
    }
}

func (h *hnswIndex) randomLevel() int {
    level := 0
    for rand.Float64() < 0.5 && level < 8 {
        level++
    }
    return level
}

func (h *hnswIndex) searchLayer(query []float32, ep int, ef int, layer int) []Result {
    visited := make(map[int]bool)
    visited[ep] = true

    candidates := &MaxHeap{{ID: uint64(ep), Dist: l2Squared(query, h.nodes[ep].vec)}}
    heap.Init(candidates)

    results := &MaxHeap{{ID: uint64(ep), Dist: l2Squared(query, h.nodes[ep].vec)}}
    heap.Init(results)

    for candidates.Len() > 0 {
        current := heap.Pop(candidates).(Result)

        if results.Len() >= ef {
            worstResult := (*results)[0]
            if current.Dist > worstResult.Dist {
                break
            }
        }

        idx := int(current.ID)
        if idx >= len(h.nodes) || layer >= len(h.nodes[idx].friends) {
            continue
        }
        for _, neighborIdx := range h.nodes[idx].friends[layer] {
            ni := int(neighborIdx)
            if visited[ni] {
                continue
            }
            visited[ni] = true

            dist := l2Squared(query, h.nodes[ni].vec)

            if results.Len() < ef {
                heap.Push(results, Result{ID: uint64(ni), Dist: dist})
                heap.Push(candidates, Result{ID: uint64(ni), Dist: dist})
            } else if dist < (*results)[0].Dist {
                heap.Pop(results)
                heap.Push(results, Result{ID: uint64(ni), Dist: dist})
                heap.Push(candidates, Result{ID: uint64(ni), Dist: dist})
            }
        }
    }

    out := make([]Result, results.Len())
    for i := results.Len() - 1; i >= 0; i-- {
        out[i] = heap.Pop(results).(Result)
    }
    return out
}

func (h *hnswIndex) selectNeighbors(candidates []Result, m int) []uint32 {
    sort.Slice(candidates, func(i, j int) bool {
        return candidates[i].Dist < candidates[j].Dist
    })
    n := m
    if n > len(candidates) {
        n = len(candidates)
    }
    result := make([]uint32, n)
    for i := 0; i < n; i++ {
        result[i] = uint32(candidates[i].ID)
    }
    return result
}

func (h *hnswIndex) insert(id uint64, vec []float32) {
    level := h.randomLevel()
    idx := len(h.nodes)
    node := &hnswNode{
        id:      id,
        vec:     vec,
        level:   level,
        friends: make([][]uint32, level+1),
    }
    for i := range node.friends {
        node.friends[i] = make([]uint32, 0, hnswM)
    }
    h.nodes = append(h.nodes, node)
    h.idToIdx[id] = idx

    if h.entryPoint == -1 {
        h.entryPoint = idx
        h.maxLevel = level
        return
    }

    ep := h.entryPoint

    // 从最高层贪心下降到 level+1
    for l := h.maxLevel; l > level; l-- {
        results := h.searchLayer(vec, ep, 1, l)
        if len(results) > 0 {
            ep = int(results[0].ID)
        }
    }

    // 从 min(level, maxLevel) 到第 0 层，搜索并连边
    topLayer := level
    if topLayer > h.maxLevel {
        topLayer = h.maxLevel
    }
    for l := topLayer; l >= 0; l-- {
        candidates := h.searchLayer(vec, ep, hnswEfConstruction, l)
        maxM := hnswM
        if l == 0 {
            maxM = hnswMmax0
        }
        neighbors := h.selectNeighbors(candidates, maxM)
        node.friends[l] = neighbors

        // 添加反向边
        for _, nIdx := range neighbors {
            ni := int(nIdx)
            if l < len(h.nodes[ni].friends) {
                h.nodes[ni].friends[l] = append(h.nodes[ni].friends[l], uint32(idx))
                // 裁剪超过 maxM 的邻居
                if len(h.nodes[ni].friends[l]) > maxM {
                    friendDists := make([]Result, len(h.nodes[ni].friends[l]))
                    for fi, fIdx := range h.nodes[ni].friends[l] {
                        friendDists[fi] = Result{
                            ID:   uint64(fIdx),
                            Dist: l2Squared(h.nodes[ni].vec, h.nodes[int(fIdx)].vec),
                        }
                    }
                    h.nodes[ni].friends[l] = h.selectNeighbors(friendDists, maxM)
                }
            }
        }

        if len(candidates) > 0 {
            ep = int(candidates[0].ID)
        }
    }

    if level > h.maxLevel {
        h.maxLevel = level
        h.entryPoint = idx
    }
}

func (h *hnswIndex) search(query []float32, k int, ef int) []Result {
    if h.entryPoint == -1 {
        return nil
    }

    ep := h.entryPoint
    for l := h.maxLevel; l > 0; l-- {
        results := h.searchLayer(query, ep, 1, l)
        if len(results) > 0 {
            ep = int(results[0].ID)
        }
    }

    candidates := h.searchLayer(query, ep, ef, 0)

    // 把内部索引转换为外部 ID
    for i := range candidates {
        candidates[i].ID = h.nodes[int(candidates[i].ID)].id
    }

    if len(candidates) > k {
        candidates = candidates[:k]
    }
    return candidates
}

// --- 向量引擎 ---

const bruteForceThreshold = 1000

type VectorEngine struct {
    mu       sync.RWMutex
    hnsw     *hnswIndex
    vectors  map[uint64][]float32
    nextID   uint64
    dim      int
    count    int64
    efSearch int
}

func NewVectorEngine(dim int) *VectorEngine {
    return &VectorEngine{
        hnsw:     newHNSW(),
        vectors:  make(map[uint64][]float32),
        dim:      dim,
        efSearch: 64,
    }
}

func (e *VectorEngine) Insert(vec []float32) uint64 {
    e.mu.Lock()
    defer e.mu.Unlock()

    id := e.nextID
    e.nextID++

    copied := make([]float32, len(vec))
    copy(copied, vec)
    e.vectors[id] = copied

    e.hnsw.insert(id, copied)
    atomic.AddInt64(&e.count, 1)
    return id
}

func (e *VectorEngine) BatchInsert(vecs [][]float32) []uint64 {
    e.mu.Lock()
    defer e.mu.Unlock()

    ids := make([]uint64, len(vecs))
    for i, vec := range vecs {
        id := e.nextID
        e.nextID++

        copied := make([]float32, len(vec))
        copy(copied, vec)
        e.vectors[id] = copied

        e.hnsw.insert(id, copied)
        atomic.AddInt64(&e.count, 1)
        ids[i] = id
    }
    return ids
}

func (e *VectorEngine) Search(query []float32, k int) []Result {
    e.mu.RLock()
    defer e.mu.RUnlock()

    n := atomic.LoadInt64(&e.count)

    // 向量数少时用暴力搜索
    if n < bruteForceThreshold {
        return e.bruteForceSearch(query, k)
    }
    return e.hnsw.search(query, k, e.efSearch)
}

func (e *VectorEngine) bruteForceSearch(query []float32, k int) []Result {
    h := &MaxHeap{}
    heap.Init(h)

    for id, vec := range e.vectors {
        dist := l2Squared(query, vec)
        if h.Len() < k {
            heap.Push(h, Result{ID: id, Dist: dist})
        } else if dist < (*h)[0].Dist {
            heap.Pop(h)
            heap.Push(h, Result{ID: id, Dist: dist})
        }
    }

    results := make([]Result, h.Len())
    for i := h.Len() - 1; i >= 0; i-- {
        results[i] = heap.Pop(h).(Result)
    }
    return results
}

func (e *VectorEngine) Stats() map[string]interface{} {
    return map[string]interface{}{
        "count":     atomic.LoadInt64(&e.count),
        "dimension": e.dim,
        "ef_search": e.efSearch,
    }
}

// --- HTTP API ---

type InsertRequest struct {
    Vectors [][]float32 `json:"vectors"`
}

type InsertResponse struct {
    IDs []uint64 `json:"ids"`
}

type SearchRequest struct {
    Vector []float32 `json:"vector"`
    TopK   int       `json:"top_k"`
}

type SearchResponse struct {
    Results []Result `json:"results"`
    Latency string   `json:"latency"`
}

func main() {
    dim := 128
    engine := NewVectorEngine(dim)

    http.HandleFunc("/insert", func(w http.ResponseWriter, r *http.Request) {
        if r.Method != http.MethodPost {
            http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
            return
        }
        var req InsertRequest
        if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
            http.Error(w, err.Error(), http.StatusBadRequest)
            return
        }
        for _, v := range req.Vectors {
            if len(v) != dim {
                http.Error(w, fmt.Sprintf("expected dim %d", dim), http.StatusBadRequest)
                return
            }
        }
        ids := engine.BatchInsert(req.Vectors)
        json.NewEncoder(w).Encode(InsertResponse{IDs: ids})
    })

    http.HandleFunc("/search", func(w http.ResponseWriter, r *http.Request) {
        if r.Method != http.MethodPost {
            http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
            return
        }
        var req SearchRequest
        if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
            http.Error(w, err.Error(), http.StatusBadRequest)
            return
        }
        if len(req.Vector) != dim {
            http.Error(w, fmt.Sprintf("expected dim %d", dim), http.StatusBadRequest)
            return
        }
        if req.TopK <= 0 {
            req.TopK = 10
        }

        start := time.Now()
        results := engine.Search(req.Vector, req.TopK)
        latency := time.Since(start)

        json.NewEncoder(w).Encode(SearchResponse{
            Results: results,
            Latency: latency.String(),
        })
    })

    http.HandleFunc("/stats", func(w http.ResponseWriter, r *http.Request) {
        json.NewEncoder(w).Encode(engine.Stats())
    })

    addr := ":8080"
    log.Printf("vector engine listening on %s (dim=%d)", addr, dim)
    log.Fatal(http.ListenAndServe(addr, nil))
}

# 插入两条向量
curl -s -X POST http://localhost:8080/insert \
  -d '{"vectors": [[1.0, 0.5, 0.3, ...], [0.2, 0.8, 0.1, ...]]}' \
  | jq .

# 查询最近的 5 个
curl -s -X POST http://localhost:8080/search \
  -d '{"vector": [1.0, 0.5, 0.3, ...], "top_k": 5}' \
  | jq .

# 查看统计信息
curl -s http://localhost:8080/stats | jq .

func (e *VectorEngine) BatchInsertWithWAL(wal *WAL, vecs [][]float32) ([]uint64, error) {
    // 先全部写入 WAL
    ids := make([]uint64, len(vecs))
    for i, vec := range vecs {
        id := atomic.AddUint64(&e.nextID, 1) - 1
        ids[i] = id
        if err := wal.AppendInsert(id, vec); err != nil {
            return nil, err
        }
    }
    // 一次 fsync
    wal.file.Sync()

    // 再批量更新索引
    e.mu.Lock()
    defer e.mu.Unlock()
    for i, vec := range vecs {
        copied := make([]float32, len(vec))
        copy(copied, vec)
        e.vectors[ids[i]] = copied
        e.hnsw.insert(ids[i], copied)
        atomic.AddInt64(&e.count, 1)
    }

    return ids, nil
}

var resultPool = sync.Pool{
    New: func() interface{} {
        h := make(MaxHeap, 0, 128)
        return &h
    },
}

func (e *VectorEngine) SearchPooled(query []float32, k int) []Result {
    e.mu.RLock()
    defer e.mu.RUnlock()

    h := resultPool.Get().(*MaxHeap)
    *h = (*h)[:0]
    defer resultPool.Put(h)

    // ... 搜索逻辑复用 h 而非每次 make
    return e.hnsw.search(query, k, e.efSearch)
}

func BenchmarkSearch(b *testing.B) {
    engine := NewVectorEngine(128)
    // 插入 100K 随机向量
    for i := 0; i < 100000; i++ {
        vec := randomVector(128)
        engine.Insert(vec)
    }
    query := randomVector(128)

    b.ResetTimer()
    b.RunParallel(func(pb *testing.PB) {
        for pb.Next() {
            engine.Search(query, 10)
        }
    })
}