Skip to content

Commit

Permalink
Knnsearch benchmark (#334)
Browse files Browse the repository at this point in the history
* add benchmark

* powers of two
  • Loading branch information
friendlymatthew authored Jun 14, 2024
1 parent 2015b11 commit 36e81d3
Showing 1 changed file with 45 additions and 0 deletions.
45 changes: 45 additions & 0 deletions pkg/hnsw/hnsw_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@ package hnsw

import (
"errors"
"fmt"
"math"
"reflect"
"testing"
)
Expand Down Expand Up @@ -661,3 +663,46 @@ func TestHnsw_KnnSearch(t *testing.T) {
}
})
}

func generatePoints(numPoints int) []Point {
points := make([]Point, numPoints)

for i := 0; i < numPoints; i++ {
points[i] = Point{float32(i), float32(i), float32(i)}
}

return points
}

func BenchmarkHnsw_KnnSearch(b *testing.B) {

for i := 0; i <= 20; i++ {
numPoints := int(math.Pow(2, float64(i)))
b.Run(fmt.Sprintf("knnsearch %d_points", numPoints), func(b *testing.B) {
h := NewHnsw(3, 10, 12, Point{0, 0, 0})

points := generatePoints(numPoints)

for _, point := range points {
if len(point) != 3 {
b.Fatalf("expected point of dim 3, got dim: %v", len(point))
}

if err := h.InsertVector(point); err != nil {
b.Fatalf("failed to insert point: %v, err: %v", point, err)
}
}

q := Point{float32(numPoints / 2), float32(numPoints / 2), float32(numPoints / 2)}
b.ResetTimer()

for i := 0; i < b.N; i++ {
_, err := h.KnnSearch(q, 5)
if err != nil {
b.Fatalf("failed to perform knnsearch: %v", err)
}
}

})
}
}

0 comments on commit 36e81d3

Please sign in to comment.