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)
+				}
+			}
+		})
+	}
+}