mirror of
https://github.com/syntrex-lab/gomcp.git
synced 2026-04-27 13:26:21 +02:00
cc7956d835
- QJL (1-bit) approximate filter for 2.3x fast search - PolarQuant (4-bit/8-bit) compressed storage with PQDropFloat64 memory reclamation (15x heap reduction) - Two-Phase SearchQJL with fallback to CompressedSimilarity - Edge Deployment prototype (pq_attention.cu) for LLaMA 1.5M token context
421 lines
13 KiB
Go
421 lines
13 KiB
Go
package vectorstore
|
|
|
|
import (
|
|
"fmt"
|
|
"math"
|
|
"math/rand"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// --- QJL Core Algorithm Tests ---
|
|
|
|
func TestQJL_Quantize_Deterministic(t *testing.T) {
|
|
proj := NewQJLProjection(256, 128, 42)
|
|
vec := randomVector(128, 1)
|
|
|
|
sig1 := proj.Quantize(vec)
|
|
sig2 := proj.Quantize(vec)
|
|
|
|
assert.Equal(t, sig1, sig2, "identical input must produce identical signature")
|
|
}
|
|
|
|
func TestQJL_Quantize_DifferentSeeds(t *testing.T) {
|
|
vec := randomVector(128, 1)
|
|
|
|
proj1 := NewQJLProjection(256, 128, 42)
|
|
proj2 := NewQJLProjection(256, 128, 99)
|
|
|
|
sig1 := proj1.Quantize(vec)
|
|
sig2 := proj2.Quantize(vec)
|
|
|
|
// Different seeds should produce different projections → different signatures.
|
|
assert.NotEqual(t, sig1, sig2, "different seeds must produce different signatures")
|
|
}
|
|
|
|
func TestQJL_HammingSimilarity_Identical(t *testing.T) {
|
|
proj := NewQJLProjection(256, 128, 42)
|
|
vec := randomVector(128, 1)
|
|
|
|
sig := proj.Quantize(vec)
|
|
sim := HammingSimilarity(sig, sig, 256)
|
|
|
|
assert.InDelta(t, 1.0, sim, 0.001, "identical signatures → similarity = 1.0")
|
|
}
|
|
|
|
func TestQJL_HammingSimilarity_Orthogonal(t *testing.T) {
|
|
proj := NewQJLProjection(512, 128, 42) // More bits for better approximation.
|
|
v1 := make([]float64, 128)
|
|
v2 := make([]float64, 128)
|
|
v1[0] = 1.0
|
|
v2[1] = 1.0
|
|
|
|
sig1 := proj.Quantize(v1)
|
|
sig2 := proj.Quantize(v2)
|
|
|
|
sim := HammingSimilarity(sig1, sig2, 512)
|
|
// Orthogonal vectors → ~0.5 Hamming similarity (random coin-flip).
|
|
assert.InDelta(t, 0.5, sim, 0.1, "orthogonal vectors → similarity ≈ 0.5, got %.4f", sim)
|
|
}
|
|
|
|
func TestQJL_HammingSimilarity_Opposite(t *testing.T) {
|
|
proj := NewQJLProjection(512, 128, 42)
|
|
v1 := randomVector(128, 1)
|
|
v2 := make([]float64, 128)
|
|
for i := range v1 {
|
|
v2[i] = -v1[i]
|
|
}
|
|
|
|
sig1 := proj.Quantize(v1)
|
|
sig2 := proj.Quantize(v2)
|
|
|
|
sim := HammingSimilarity(sig1, sig2, 512)
|
|
// Opposite vectors → ~0.0 Hamming similarity.
|
|
assert.InDelta(t, 0.0, sim, 0.1, "opposite vectors → similarity ≈ 0.0, got %.4f", sim)
|
|
}
|
|
|
|
func TestQJL_PreservesOrdering(t *testing.T) {
|
|
// For well-separated vectors, QJL Hamming ordering should match cosine ordering.
|
|
proj := NewQJLProjection(512, 128, 42)
|
|
|
|
query := randomVector(128, 1)
|
|
close := perturbVector(query, 0.1, 2) // ~10% perturbation = close
|
|
far := perturbVector(query, 0.9, 3) // ~90% perturbation = far
|
|
|
|
cosClose := CosineSimilarity(query, close)
|
|
cosFar := CosineSimilarity(query, far)
|
|
require.Greater(t, cosClose, cosFar, "sanity: cosine(query, close) > cosine(query, far)")
|
|
|
|
sigQ := proj.Quantize(query)
|
|
sigClose := proj.Quantize(close)
|
|
sigFar := proj.Quantize(far)
|
|
|
|
hamClose := HammingSimilarity(sigQ, sigClose, 512)
|
|
hamFar := HammingSimilarity(sigQ, sigFar, 512)
|
|
|
|
assert.Greater(t, hamClose, hamFar,
|
|
"QJL must preserve ordering: hamming(query,close)=%.4f > hamming(query,far)=%.4f",
|
|
hamClose, hamFar)
|
|
}
|
|
|
|
func TestQJL_EstimatedCosineSimilarity(t *testing.T) {
|
|
// Hamming sim = 1.0 → estimated cosine = cos(0) = 1.0
|
|
assert.InDelta(t, 1.0, EstimatedCosineSimilarity(1.0), 0.001)
|
|
// Hamming sim = 0.5 → estimated cosine = cos(π/2) = 0.0
|
|
assert.InDelta(t, 0.0, EstimatedCosineSimilarity(0.5), 0.001)
|
|
// Hamming sim = 0.0 → estimated cosine = cos(π) = -1.0
|
|
assert.InDelta(t, -1.0, EstimatedCosineSimilarity(0.0), 0.001)
|
|
}
|
|
|
|
func TestQJL_MemoryReduction(t *testing.T) {
|
|
proj := NewQJLProjection(256, 128, 42)
|
|
vec := randomVector(128, 1)
|
|
sig := proj.Quantize(vec)
|
|
|
|
float64Bytes := 128 * 8 // 1024 bytes
|
|
qjlBytes := len(sig) * 8 // 4 * 8 = 32 bytes
|
|
reduction := float64(float64Bytes) / float64(qjlBytes)
|
|
|
|
assert.Equal(t, 4, len(sig), "256 bits → 4 uint64 words")
|
|
assert.Equal(t, 32, qjlBytes, "256-bit signature = 32 bytes")
|
|
assert.InDelta(t, 32.0, reduction, 0.5, "expected 32x memory reduction")
|
|
}
|
|
|
|
func TestQJL_HammingSimilarity_LengthMismatch(t *testing.T) {
|
|
a := QJLSignature{0xFF}
|
|
b := QJLSignature{0xFF, 0x00}
|
|
assert.Equal(t, 0.0, HammingSimilarity(a, b, 128))
|
|
}
|
|
|
|
func TestQJL_HammingSimilarity_ZeroBits(t *testing.T) {
|
|
a := QJLSignature{0xFF}
|
|
b := QJLSignature{0xFF}
|
|
assert.Equal(t, 0.0, HammingSimilarity(a, b, 0))
|
|
}
|
|
|
|
// --- Store QJL Integration Tests ---
|
|
|
|
func TestStore_QJLEnabled_Default(t *testing.T) {
|
|
s := New(nil)
|
|
assert.True(t, s.QJLEnabled(), "QJL should be enabled by default")
|
|
}
|
|
|
|
func TestStore_QJLDisabled_Explicit(t *testing.T) {
|
|
s := New(&Config{QJLProjections: -1})
|
|
assert.False(t, s.QJLEnabled(), "QJL should be disabled with -1")
|
|
}
|
|
|
|
func TestStore_SearchQJL_MatchesExact(t *testing.T) {
|
|
s := New(&Config{QJLProjections: 512, QJLVectorDim: 3, QJLSeed: 42})
|
|
|
|
s.Add(&IntentRecord{ID: "r1", Text: "read data", Vector: []float64{1.0, 0.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "w1", Text: "write data", Vector: []float64{0.0, 1.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "e1", Text: "exec code", Vector: []float64{0.0, 0.0, 1.0}})
|
|
|
|
// SearchQJL should return the same top-1 as Search for well-separated vectors.
|
|
exact := s.Search([]float64{1.0, 0.0, 0.0}, 1)
|
|
qjl := s.SearchQJL([]float64{1.0, 0.0, 0.0}, 1)
|
|
|
|
require.Len(t, exact, 1)
|
|
require.Len(t, qjl, 1)
|
|
assert.Equal(t, exact[0].Record.ID, qjl[0].Record.ID, "QJL top-1 should match exact top-1")
|
|
assert.InDelta(t, exact[0].Similarity, qjl[0].Similarity, 0.001)
|
|
}
|
|
|
|
func TestStore_SearchQJL_Fallback_NoQJL(t *testing.T) {
|
|
s := New(&Config{QJLProjections: -1})
|
|
|
|
s.Add(&IntentRecord{ID: "r1", Vector: []float64{1.0, 0.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "w1", Vector: []float64{0.0, 1.0, 0.0}})
|
|
|
|
results := s.SearchQJL([]float64{1.0, 0.0, 0.0}, 1)
|
|
require.Len(t, results, 1)
|
|
assert.Equal(t, "r1", results[0].Record.ID, "fallback should still work")
|
|
}
|
|
|
|
func TestStore_SearchQJL_Stats(t *testing.T) {
|
|
s := New(&Config{QJLProjections: 256, QJLVectorDim: 128})
|
|
s.Add(&IntentRecord{Route: "read", Verdict: "ALLOW", Entropy: 3.0, Vector: randomVector(128, 1)})
|
|
|
|
stats := s.GetStats()
|
|
assert.True(t, stats.QJLEnabled)
|
|
assert.Equal(t, 256, stats.QJLProjections)
|
|
assert.Equal(t, 256, stats.QJLBitsPerVec)
|
|
assert.Equal(t, 32, stats.QJLBytesPerVec)
|
|
}
|
|
|
|
func TestStore_QJL_LRU_Eviction(t *testing.T) {
|
|
s := New(&Config{Capacity: 3, QJLProjections: 64, QJLVectorDim: 3})
|
|
|
|
s.Add(&IntentRecord{ID: "a", Vector: []float64{1.0, 0.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "b", Vector: []float64{0.0, 1.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "c", Vector: []float64{0.0, 0.0, 1.0}})
|
|
s.Add(&IntentRecord{ID: "d", Vector: []float64{0.5, 0.5, 0.0}}) // Evicts "a"
|
|
|
|
assert.Equal(t, 3, s.Count())
|
|
assert.Nil(t, s.Get("a"))
|
|
|
|
// SearchQJL should still work after eviction.
|
|
results := s.SearchQJL([]float64{0.5, 0.5, 0.0}, 1)
|
|
require.Len(t, results, 1)
|
|
assert.Equal(t, "d", results[0].Record.ID)
|
|
}
|
|
|
|
func TestStore_SearchQJL_EmptyStore(t *testing.T) {
|
|
s := New(nil)
|
|
assert.Nil(t, s.SearchQJL([]float64{1.0}, 5))
|
|
}
|
|
|
|
func TestStore_SearchQJL_EmptyQuery(t *testing.T) {
|
|
s := New(nil)
|
|
s.Add(&IntentRecord{ID: "r1", Vector: []float64{1.0, 0.0}})
|
|
assert.Nil(t, s.SearchQJL(nil, 5))
|
|
}
|
|
|
|
// --- PolarQuant Store Integration Tests ---
|
|
|
|
func TestStore_PQEnabled(t *testing.T) {
|
|
// PQ disabled by default.
|
|
s1 := New(nil)
|
|
assert.False(t, s1.PQEnabled(), "PQ should be disabled by default")
|
|
|
|
// PQ enabled with 4-bit.
|
|
s2 := New(&Config{PQBitsPerDim: 4})
|
|
assert.True(t, s2.PQEnabled(), "PQ should be enabled with 4-bit")
|
|
}
|
|
|
|
func TestStore_PQ_Stats(t *testing.T) {
|
|
s := New(&Config{QJLProjections: 256, QJLVectorDim: 128, PQBitsPerDim: 4, PQSeed: 7})
|
|
s.Add(&IntentRecord{Route: "read", Vector: randomVector(128, 1)})
|
|
|
|
stats := s.GetStats()
|
|
assert.True(t, stats.PQEnabled)
|
|
assert.Equal(t, 4, stats.PQBitsPerDim)
|
|
assert.Equal(t, 68, stats.PQBytesPerVec) // 64 data + 4 radius
|
|
assert.Greater(t, stats.PQCompressionRate, 14.0)
|
|
}
|
|
|
|
func TestStore_SearchQJL_WithPQ_MatchesExact(t *testing.T) {
|
|
// Full TurboQuant pipeline: QJL filter → PQ compressed rerank.
|
|
s := New(&Config{QJLProjections: 512, QJLVectorDim: 3, PQBitsPerDim: 8, PQSeed: 7})
|
|
|
|
s.Add(&IntentRecord{ID: "r1", Vector: []float64{1.0, 0.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "w1", Vector: []float64{0.0, 1.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "e1", Vector: []float64{0.0, 0.0, 1.0}})
|
|
|
|
results := s.SearchQJL([]float64{0.9, 0.1, 0.0}, 1)
|
|
require.Len(t, results, 1)
|
|
assert.Equal(t, "r1", results[0].Record.ID, "PQ rerank should still pick the closest vector")
|
|
}
|
|
|
|
func TestStore_PQ_LRU_Eviction(t *testing.T) {
|
|
s := New(&Config{Capacity: 3, QJLProjections: 64, QJLVectorDim: 3, PQBitsPerDim: 4, PQSeed: 7})
|
|
|
|
s.Add(&IntentRecord{ID: "a", Vector: []float64{1.0, 0.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "b", Vector: []float64{0.0, 1.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "c", Vector: []float64{0.0, 0.0, 1.0}})
|
|
s.Add(&IntentRecord{ID: "d", Vector: []float64{0.5, 0.5, 0.0}}) // Evicts "a"
|
|
|
|
assert.Equal(t, 3, s.Count())
|
|
assert.Nil(t, s.Get("a"))
|
|
|
|
// SearchQJL with PQ should still work after eviction.
|
|
results := s.SearchQJL([]float64{0.5, 0.5, 0.0}, 1)
|
|
require.Len(t, results, 1)
|
|
assert.Equal(t, "d", results[0].Record.ID)
|
|
}
|
|
|
|
func TestStore_SearchQJL_PQOnly_NoBrokenFallback(t *testing.T) {
|
|
// PQ enabled, QJL disabled → should fallback to brute-force (no PQ rerank without QJL)
|
|
s := New(&Config{QJLProjections: -1, PQBitsPerDim: 4, QJLVectorDim: 3, PQSeed: 7})
|
|
|
|
s.Add(&IntentRecord{ID: "r1", Vector: []float64{1.0, 0.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "w1", Vector: []float64{0.0, 1.0, 0.0}})
|
|
|
|
results := s.SearchQJL([]float64{1.0, 0.0, 0.0}, 1)
|
|
require.Len(t, results, 1)
|
|
assert.Equal(t, "r1", results[0].Record.ID, "fallback should work with PQ but no QJL")
|
|
}
|
|
|
|
func TestStore_PQ_DropFloat64(t *testing.T) {
|
|
s := New(&Config{PQBitsPerDim: 4, PQDropFloat64: true})
|
|
s.Add(&IntentRecord{ID: "r1", Vector: []float64{1.0, 0.0, 0.0}})
|
|
s.Add(&IntentRecord{ID: "r2", Vector: []float64{0.0, 1.0, 0.0}})
|
|
|
|
// Original float64 vector should be nil'd out.
|
|
assert.Nil(t, s.Get("r1").Vector, "Vector should be dropped to save memory")
|
|
|
|
stats := s.GetStats()
|
|
assert.True(t, stats.PQDropFloat64, "Stats should reflect drop float64 true")
|
|
|
|
// Search should still work via compressed similarity fallback in brute-force searchLocked.
|
|
results := s.Search([]float64{1.0, 0.0, 0.0}, 1)
|
|
require.Len(t, results, 1)
|
|
assert.Equal(t, "r1", results[0].Record.ID, "Search should work via PQ fallback when vector is dropped")
|
|
}
|
|
|
|
// --- Benchmarks ---
|
|
|
|
func BenchmarkSearch_BruteForce(b *testing.B) {
|
|
s := New(&Config{Capacity: 10000, QJLProjections: -1, QJLVectorDim: 128})
|
|
populateStore(s, 1000, 128)
|
|
query := randomVector(128, 999)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
s.Search(query, 10)
|
|
}
|
|
}
|
|
|
|
func BenchmarkSearchQJL_TwoPhase(b *testing.B) {
|
|
s := New(&Config{Capacity: 10000, QJLProjections: 256, QJLVectorDim: 128, QJLSeed: 42})
|
|
populateStore(s, 1000, 128)
|
|
query := randomVector(128, 999)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
s.SearchQJL(query, 10)
|
|
}
|
|
}
|
|
|
|
func BenchmarkSearchTurboQuant_Full(b *testing.B) {
|
|
// Full TurboQuant: QJL filter + PolarQuant compressed rerank.
|
|
s := New(&Config{
|
|
Capacity: 10000,
|
|
QJLProjections: 256,
|
|
QJLVectorDim: 128,
|
|
QJLSeed: 42,
|
|
PQBitsPerDim: 4,
|
|
PQSeed: 7,
|
|
})
|
|
populateStore(s, 1000, 128)
|
|
query := randomVector(128, 999)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
s.SearchQJL(query, 10)
|
|
}
|
|
}
|
|
|
|
func BenchmarkQJL_Quantize(b *testing.B) {
|
|
proj := NewQJLProjection(256, 128, 42)
|
|
vec := randomVector(128, 1)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
proj.Quantize(vec)
|
|
}
|
|
}
|
|
|
|
func BenchmarkHammingSimilarity_256bit(b *testing.B) {
|
|
proj := NewQJLProjection(256, 128, 42)
|
|
v1 := randomVector(128, 1)
|
|
v2 := randomVector(128, 2)
|
|
sig1 := proj.Quantize(v1)
|
|
sig2 := proj.Quantize(v2)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
HammingSimilarity(sig1, sig2, 256)
|
|
}
|
|
}
|
|
|
|
func BenchmarkCosineSimilarity_128dim(b *testing.B) {
|
|
v1 := randomVector(128, 1)
|
|
v2 := randomVector(128, 2)
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
CosineSimilarity(v1, v2)
|
|
}
|
|
}
|
|
|
|
// --- Helpers ---
|
|
|
|
func randomVector(dim int, seed int64) []float64 {
|
|
rng := rand.New(rand.NewSource(seed))
|
|
vec := make([]float64, dim)
|
|
for i := range vec {
|
|
vec[i] = rng.NormFloat64()
|
|
}
|
|
// L2-normalize.
|
|
var norm float64
|
|
for _, v := range vec {
|
|
norm += v * v
|
|
}
|
|
norm = math.Sqrt(norm)
|
|
for i := range vec {
|
|
vec[i] /= norm
|
|
}
|
|
return vec
|
|
}
|
|
|
|
func perturbVector(v []float64, noise float64, seed int64) []float64 {
|
|
rng := rand.New(rand.NewSource(seed))
|
|
perturbed := make([]float64, len(v))
|
|
for i := range v {
|
|
perturbed[i] = v[i] + noise*rng.NormFloat64()
|
|
}
|
|
// L2-normalize.
|
|
var norm float64
|
|
for _, val := range perturbed {
|
|
norm += val * val
|
|
}
|
|
norm = math.Sqrt(norm)
|
|
for i := range perturbed {
|
|
perturbed[i] /= norm
|
|
}
|
|
return perturbed
|
|
}
|
|
|
|
func populateStore(s *Store, n, dim int) {
|
|
for i := 0; i < n; i++ {
|
|
s.Add(&IntentRecord{
|
|
ID: fmt.Sprintf("vec-%d", i),
|
|
Vector: randomVector(dim, int64(i)),
|
|
Route: "test",
|
|
})
|
|
}
|
|
}
|