125 lines
2.9 KiB
Go
125 lines
2.9 KiB
Go
package indexDB_test
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"runtime"
|
|
"testing"
|
|
"time"
|
|
|
|
"apigo.cc/go/cast"
|
|
"apigo.cc/go/indexDB"
|
|
"apigo.cc/go/rand"
|
|
)
|
|
|
|
func mockVector(dim int) []float32 {
|
|
v := make([]float32, dim)
|
|
for i := 0; i < dim; i++ {
|
|
v[i] = rand.Float[float32](0.0, 1.0)
|
|
}
|
|
return v
|
|
}
|
|
|
|
func reportMem(t *testing.T, msg string) {
|
|
var m runtime.MemStats
|
|
runtime.ReadMemStats(&m)
|
|
t.Logf("%s: Alloc = %v MiB, TotalAlloc = %v MiB, Sys = %v MiB, NumGC = %v",
|
|
msg, m.Alloc/1024/1024, m.TotalAlloc/1024/1024, m.Sys/1024/1024, m.NumGC)
|
|
}
|
|
|
|
func TestVectorPerformance(t *testing.T) {
|
|
fPath := "perf_fulltext"
|
|
vPath := "perf_vector"
|
|
defer os.RemoveAll(fPath)
|
|
defer os.RemoveAll(vPath)
|
|
|
|
dim := 128
|
|
count := 1000
|
|
|
|
embeddingFunc := func(text string) ([]float32, error) {
|
|
return mockVector(dim), nil
|
|
}
|
|
|
|
dbUnauth, err := indexDB.GetDB(fPath, vPath, embeddingFunc, nil)
|
|
if err != nil {
|
|
t.Fatalf("Failed to create engine: %v", err)
|
|
}
|
|
db := dbUnauth.Auth("_system")
|
|
|
|
reportMem(t, "Before adding docs")
|
|
start := time.Now()
|
|
|
|
for i := 0; i < count; i++ {
|
|
id := fmt.Sprintf("doc_%d", i)
|
|
text := fmt.Sprintf("This is document number %d with some content to index for fulltext search.", i)
|
|
meta := map[string]any{"index": i, "category": "test"}
|
|
err := db.Add(id, text, meta, nil)
|
|
if err != nil {
|
|
t.Fatalf("Failed to add doc %d: %v", i, err)
|
|
}
|
|
if (i+1)%200 == 0 {
|
|
t.Logf("Added %d docs...", i+1)
|
|
}
|
|
}
|
|
|
|
t.Logf("Added %d docs in %v", count, time.Since(start))
|
|
reportMem(t, "After adding docs")
|
|
|
|
// Search test
|
|
t.Log("Testing search performance...")
|
|
start = time.Now()
|
|
results, err := db.Search("", "document 500", 10, nil)
|
|
if err != nil {
|
|
t.Fatalf("Search failed: %v", err)
|
|
}
|
|
t.Logf("Search returned %d results in %v", len(results), time.Since(start))
|
|
|
|
for _, r := range results {
|
|
t.Logf("Result: ID=%s, Score=%f", r.ID, r.Score)
|
|
}
|
|
|
|
reportMem(t, "After search")
|
|
|
|
// Filter search
|
|
t.Log("Testing search with filter...")
|
|
filter := []indexDB.Condition{
|
|
{Field: "index", Operator: "gt", Value: 800},
|
|
}
|
|
start = time.Now()
|
|
results, err = db.Search("", "document", 10, filter)
|
|
if err != nil {
|
|
t.Fatalf("Filtered search failed: %v", err)
|
|
}
|
|
t.Logf("Filtered search returned %d results in %v", len(results), time.Since(start))
|
|
for _, r := range results {
|
|
idx := cast.To[float64](r.Metadata["index"])
|
|
if idx <= 800 {
|
|
t.Errorf("Result %s has index %v, expected > 800", r.ID, idx)
|
|
}
|
|
}
|
|
}
|
|
|
|
func BenchmarkSearch(b *testing.B) {
|
|
fPath := "bench_fulltext"
|
|
vPath := "bench_vector"
|
|
defer os.RemoveAll(fPath)
|
|
defer os.RemoveAll(vPath)
|
|
|
|
dim := 128
|
|
embeddingFunc := func(text string) ([]float32, error) {
|
|
return mockVector(dim), nil
|
|
}
|
|
|
|
dbUnauth, _ := indexDB.GetDB(fPath, vPath, embeddingFunc, nil)
|
|
db := dbUnauth.Auth("_system")
|
|
|
|
for i := 0; i < 500; i++ {
|
|
_ = db.Add(fmt.Sprintf("d%d", i), "content", nil, nil)
|
|
}
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
_, _ = db.Search("", "content", 10, nil)
|
|
}
|
|
}
|