Embeddings and Vector Search

// vsearch.go — exact, quantized and partitioned (IVF) nearest-neighbour search, with recall.
package main

import (
	"fmt"
	"math"
	"math/rand"
	"runtime"
	"sort"
	"sync"
	"time"
)

const (
	N   = 50000 // vectors
	Dim = 128
	K   = 10 // neighbours wanted
)

type Hit struct {
	ID    int32
	Score float32
}

// topK keeps the K best hits in a small sorted slice: for small K this beats a heap.
type topK []Hit

func (t *topK) push(id int32, score float32) {
	h := *t
	if len(h) == K && score <= h[K-1].Score {
		return
	}
	if len(h) < K {
		h = append(h, Hit{})
	}
	i := len(h) - 1
	for ; i > 0 && h[i-1].Score < score; i-- {
		h[i] = h[i-1]
	}
	h[i] = Hit{id, score}
	*t = h
}

func dot(a, b []float32) float32 {
	b = b[:len(a)]
	var s0, s1, s2, s3 float32
	for i := 0; i+4 <= len(a); i += 4 {
		s0 += a[i] * b[i]
		s1 += a[i+1] * b[i+1]
		s2 += a[i+2] * b[i+2]
		s3 += a[i+3] * b[i+3]
	}
	return s0 + s1 + s2 + s3
}

func dotInt8(a, b []int8) int32 {
	b = b[:len(a)]
	var s0, s1, s2, s3 int32
	for i := 0; i+4 <= len(a); i += 4 {
		s0 += int32(a[i]) * int32(b[i])
		s1 += int32(a[i+1]) * int32(b[i+1])
		s2 += int32(a[i+2]) * int32(b[i+2])
		s3 += int32(a[i+3]) * int32(b[i+3])
	}
	return s0 + s1 + s2 + s3
}

// quantize maps a component of a unit vector to 8 bits. Components of a 128-D unit vector are
// small (around 0.09), so they are scaled up first and clamped.
func quantize(x float32) int8 {
	return int8(max(-127, min(127, math.Round(float64(x)*400))))
}

func normalize(v []float32) {
	var ss float64
	for _, x := range v {
		ss += float64(x) * float64(x)
	}
	inv := float32(1 / math.Sqrt(ss))
	for i := range v {
		v[i] *= inv
	}
}

// Index stores every vector in ONE flat slice: contiguous, pointer-free, cache-friendly.
type Index struct {
	vecs  []float32 // N x Dim, unit length, so dot product = cosine similarity
	q8    []int8    // the same vectors at 8 bits per value
	cents []float32 // IVF: nlist centroids
	lists [][]int32 // IVF: vector IDs per centroid
}

func (ix *Index) vec(i int) []float32 { return ix.vecs[i*Dim : (i+1)*Dim] }

// Exact scans everything, splitting the range across goroutines.
func (ix *Index) Exact(q []float32, workers int) []Hit {
	parts := make([]topK, workers)
	var wg sync.WaitGroup
	chunk := (N + workers - 1) / workers
	for w := 0; w < workers; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			t := make(topK, 0, K)
			for i := w * chunk; i < min((w+1)*chunk, N); i++ {
				t.push(int32(i), dot(q, ix.vec(i)))
			}
			parts[w] = t // each goroutine writes its own slot
		}()
	}
	wg.Wait()
	merged := make(topK, 0, K)
	for _, p := range parts {
		for _, h := range p {
			merged.push(h.ID, h.Score)
		}
	}
	return merged
}

// Quantized scans int8 codes (4x less memory to read), then re-scores the best with float32.
func (ix *Index) Quantized(q []float32) []Hit {
	q8 := make([]int8, Dim)
	for i, x := range q {
		q8[i] = quantize(x)
	}
	// Pass 1: approximate scores for everything; keep a shortlist of 4K candidates.
	const shortlist = 4 * K
	ids := make([]int32, 0, shortlist)
	scores := make([]int32, 0, shortlist)
	for i := 0; i < N; i++ {
		s := dotInt8(q8, ix.q8[i*Dim:(i+1)*Dim])
		if len(ids) == shortlist && s <= scores[shortlist-1] {
			continue
		}
		if len(ids) < shortlist {
			ids, scores = append(ids, 0), append(scores, 0)
		}
		j := len(ids) - 1
		for ; j > 0 && scores[j-1] < s; j-- {
			ids[j], scores[j] = ids[j-1], scores[j-1]
		}
		ids[j], scores[j] = int32(i), s
	}
	// Pass 2: exact scores for the shortlist only.
	t := make(topK, 0, K)
	for _, id := range ids {
		t.push(id, dot(q, ix.vec(int(id))))
	}
	return t
}

// IVF looks only in the nprobe partitions whose centroids are closest to the query.
func (ix *Index) IVF(q []float32, nprobe int) []Hit {
	nlist := len(ix.lists)
	order := make([]Hit, nlist)
	for c := 0; c < nlist; c++ {
		order[c] = Hit{int32(c), dot(q, ix.cents[c*Dim:(c+1)*Dim])}
	}
	sort.Slice(order, func(a, b int) bool { return order[a].Score > order[b].Score })
	t := make(topK, 0, K)
	for _, c := range order[:nprobe] {
		for _, id := range ix.lists[c.ID] {
			t.push(id, dot(q, ix.vec(int(id))))
		}
	}
	return t
}

func recall(got, want []Hit) float64 {
	in := map[int32]bool{}
	for _, h := range want {
		in[h.ID] = true
	}
	n := 0
	for _, h := range got {
		if in[h.ID] {
			n++
		}
	}
	return float64(n) / float64(len(want))
}

func main() {
	rng := rand.New(rand.NewSource(3))
	// Synthetic "embeddings": 200 topics, each vector is its topic plus noise.
	const topics = 200
	centers := make([]float32, topics*Dim)
	for i := range centers {
		centers[i] = float32(rng.NormFloat64())
	}
	ix := &Index{vecs: make([]float32, N*Dim), q8: make([]int8, N*Dim)}
	for i := 0; i < N; i++ {
		c := rng.Intn(topics)
		v := ix.vec(i)
		for d := range v {
			v[d] = centers[c*Dim+d] + float32(rng.NormFloat64())*0.6
		}
		normalize(v)
		for d, x := range v {
			ix.q8[i*Dim+d] = quantize(x)
		}
	}
	// IVF build: sample centroids, assign every vector to its nearest one.
	const nlist = 224 // about sqrt(N)
	ix.cents = make([]float32, nlist*Dim)
	for c := 0; c < nlist; c++ {
		copy(ix.cents[c*Dim:(c+1)*Dim], ix.vec(rng.Intn(N)))
	}
	ix.lists = make([][]int32, nlist)
	for i := 0; i < N; i++ {
		best, bestS := 0, float32(-2)
		for c := 0; c < nlist; c++ {
			if s := dot(ix.vec(i), ix.cents[c*Dim:(c+1)*Dim]); s > bestS {
				best, bestS = c, s
			}
		}
		ix.lists[best] = append(ix.lists[best], int32(i))
	}

	queries := make([][]float32, 50)
	for i := range queries {
		queries[i] = append([]float32{}, ix.vec(rng.Intn(N))...)
		for d := range queries[i] {
			queries[i][d] += float32(rng.NormFloat64()) * 0.05
		}
		normalize(queries[i])
	}
	truth := make([][]Hit, len(queries))
	for i, q := range queries {
		truth[i] = ix.Exact(q, 1)
	}

	fmt.Printf("%d vectors x %d dims = %d MB as float32, %d MB as int8\n\n", N, Dim, N*Dim*4>>20, N*Dim>>20)
	fmt.Println("method                         ms/query   recall@10")
	run := func(name string, f func(q []float32) []Hit) {
		start := time.Now()
		r := 0.0
		for i, q := range queries {
			r += recall(f(q), truth[i])
		}
		fmt.Printf("%-28s %9.2f   %8.1f%%\n", name,
			float64(time.Since(start).Microseconds())/1000/float64(len(queries)), 100*r/float64(len(queries)))
	}
	workers := runtime.GOMAXPROCS(0)
	run("exact scan, 1 goroutine", func(q []float32) []Hit { return ix.Exact(q, 1) })
	run(fmt.Sprintf("exact scan, %d goroutines", workers), func(q []float32) []Hit { return ix.Exact(q, workers) })
	run("int8 scan + float re-score", ix.Quantized)
	for _, np := range []int{1, 4, 16, 64} {
		run(fmt.Sprintf("IVF, nprobe=%d of %d lists", np, nlist), func(q []float32) []Hit { return ix.IVF(q, np) })
	}
}