Skip to content

Commit

Permalink
sequential search with upper bound param (#335)
Browse files Browse the repository at this point in the history
* sequential search with upper bound param

* remove additional test
  • Loading branch information
friendlymatthew authored Jun 11, 2024
1 parent 286a6b8 commit abe9808
Showing 1 changed file with 34 additions and 0 deletions.
34 changes: 34 additions & 0 deletions pkg/hnsw/hnsw_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -476,6 +476,7 @@ func TestHnsw_InsertVector(t *testing.T) {
}
}
})

}

func TestHnsw_KnnSearch(t *testing.T) {
Expand Down Expand Up @@ -598,4 +599,37 @@ func TestHnsw_KnnSearch(t *testing.T) {
t.Fatalf("expected closest points to be %v, got %v", expected, got)
}
})

t.Run("sequential search with upper bound params", func(t *testing.T) {
h := NewHnsw(2, 12, 12, Point{0, 0})
for i := 1; i <= 8; i++ {
if err := h.InsertVector(Point{float32(i), float32(i + 1)}); err != nil {
t.Fatalf("failed to insert point: %v, err: %v", Point{float32(i), float32(i + 1)}, err)
}
}

found, err := h.KnnSearch(Point{float32(0), float32(0)}, 10)

if err != nil {
t.Fatalf("failed to find closest neighbors: %v", err)
}

if found.Len() != 9 {
t.Fatalf("expected to find 9 closest neighbors, got %v", found.Len())
}

expectedId := Id(0)

for found.IsEmpty() {
nnItem, err := found.PopItem()
if err != nil {
t.Fatalf("failed to pop item: %v, err: %v", found, err)
}
if expectedId != nnItem.id {
t.Fatalf("expected to find %v, got %v", expectedId, nnItem.id)
}

expectedId += 1
}
})
}

0 comments on commit abe9808

Please sign in to comment.