diff --git a/internal/store/memory_bench_test.go b/internal/store/memory_bench_test.go new file mode 100644 index 0000000..ef6ffdd --- /dev/null +++ b/internal/store/memory_bench_test.go @@ -0,0 +1,77 @@ +package store + +import ( + "context" + "fmt" + "math" + "math/rand" + "path/filepath" + "testing" +) + +// benchDim is the resident embedder's width (multilingual-e5-small, 384), so +// the per-row decode cost the benchmark measures is the real one. +const benchDim = 384 + +// seedMemVectors fills a fresh store with n L2-normalized rows carrying a meta +// blob the size recall actually stores — the note text plus its type — because +// the cost this benchmark exists to measure is unmarshalling that blob for +// every row when only topK survivors need it. +func seedMemVectors(tb testing.TB, n int) *MemoryStore { + tb.Helper() + path := filepath.Join(tb.TempDir(), "mem_bench.db") + st, err := Open(context.Background(), path) + if err != nil { + tb.Fatalf("Open: %v", err) + } + tb.Cleanup(func() { _ = st.Close() }) + m := st.VectorMemory() + + rng := rand.New(rand.NewSource(1)) + ctx := context.Background() + for i := 0; i < n; i++ { + if err := m.Insert(ctx, fmt.Sprintf("note:%d", i), randUnitVec(rng, benchDim), map[string]string{ + "type": "note", + "text": fmt.Sprintf("заметка номер %d о том, что надо не забыть сделать на неделе", i), + }); err != nil { + tb.Fatalf("Insert %d: %v", i, err) + } + } + return m +} + +func randUnitVec(rng *rand.Rand, dim int) []float32 { + v := make([]float32, dim) + var norm float64 + for i := range v { + f := rng.NormFloat64() + v[i] = float32(f) + norm += f * f + } + norm = math.Sqrt(norm) + for i := range v { + v[i] = float32(float64(v[i]) / norm) + } + return v +} + +// BenchmarkMemoryStoreSearch measures one recall query against a store of n +// rows. Row counts bracket the documented scale: 1000 is a plausible today, +// 10000 is the "thousands, not millions" ceiling the type doc claims a full +// scan is fine at. +func BenchmarkMemoryStoreSearch(b *testing.B) { + for _, n := range []int{1000, 10000} { + b.Run(fmt.Sprintf("rows=%d", n), func(b *testing.B) { + m := seedMemVectors(b, n) + q := randUnitVec(rand.New(rand.NewSource(2)), benchDim) + ctx := context.Background() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + if _, err := m.Search(ctx, q, 10); err != nil { + b.Fatal(err) + } + } + }) + } +}