feat: flush Friends to disk

pkg/hnsw/friends.go

@@ -1,7 +1,9 @@
 package hnsw
 
 import (
+	"encoding/binary"
 	"errors"
+	"fmt"
 	"math"
 )
 
@@ -63,6 +65,45 @@ func (v *Friends) GetFriendsAtLevel(level int) (*DistHeap, error) {
 	return v.friends[level], nil
 }
 
+func (v *Friends) Flush(numNeighbors int) ([]byte, error) {
+	if len(v.friends) == 0 {
+		panic("no levels to be found")
+	}
+
+	// for every neighbor, we're going to serialize
+	// +-------+-------------------------+
+	// | level |      Id
+
+	buf := make([]byte, (4+1)*numNeighbors)
+
+	level0 := v.friends[0]
+	copyLevel0 := level0.Clone()
+
+	for i := 0; i < numNeighbors; i++ {
+		if copyLevel0.IsEmpty() {
+			// write out max values here
+			continue
+		}
+
+		closestItem, err := copyLevel0.PopMinItem()
+		if err != nil {
+			return []byte{}, fmt.Errorf("failed to find closest item in friends: %v", err)
+		}
+
+		closestId := closestItem.id
+		closestIdMaxLevel, ok := v.maxLevels[closestId]
+
+		if !ok {
+			return []byte{}, fmt.Errorf("failed to find id %v in maxLevels map", closestId)
+		}
+
+		buf[i*(1+4)] = byte(closestIdMaxLevel)
+		binary.BigEndian.PutUint32(buf[i*(1+4)+1:], uint32(closestId))
+	}
+
+	return buf, nil
+}
+
 func EuclidDistance(p0, p1 Point) float32 {
 	var sum float32
 

pkg/hnsw/friends_test.go

@@ -2,6 +2,7 @@ package hnsw
 
 import (
 	"math"
+	"reflect"
 	"testing"
 )
 
@@ -139,3 +140,40 @@ func TestVector_EuclidDistance(t *testing.T) {
 		}
 	})
 }
+
+func TestFriends_Flush(t *testing.T) {
+	t.Run("flush single friend", func(t *testing.T) {
+		f := NewFriends(3)
+
+		f.InsertFriendsAtLevel(2, 1, 4)
+
+		buf, err := f.Flush(1)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !reflect.DeepEqual(buf, []byte{2, 0, 0, 0, 1}) {
+			t.Fatalf("expected %v, got %v", []byte{2, 0, 0, 0, 1}, buf)
+		}
+	})
+
+	t.Run("flushes 8 friends exactly", func(t *testing.T) {
+		f := NewFriends(4)
+		f.InsertFriendsAtLevel(2, 1, 1)
+		f.InsertFriendsAtLevel(3, 2, 2)
+		f.InsertFriendsAtLevel(1, 3, 3)
+		f.InsertFriendsAtLevel(0, 4, 4)
+		f.InsertFriendsAtLevel(4, 5, 5)
+		f.InsertFriendsAtLevel(2, 6, 6)
+		f.InsertFriendsAtLevel(0, 7, 7)
+		f.InsertFriendsAtLevel(2, 8, 8)
+
+		buf, err := f.Flush(8)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !reflect.DeepEqual(buf, []byte{2, 0, 0, 0, 1, 3, 0, 0, 0, 2, 1, 0, 0, 0, 3, 0, 0, 0, 0, 4, 4, 0, 0, 0, 5, 2, 0, 0, 0, 6, 0, 0, 0, 0, 7, 2, 0, 0, 0, 8}) {
+			t.Fatalf("expected %v, got %v", []byte{2, 0, 0, 0, 1, 3, 0, 0, 0, 2, 1, 0, 0, 0, 3, 0, 0, 0, 0, 4, 4, 0, 0, 0, 5, 2, 0, 0, 0, 6, 0, 0, 0, 0, 7, 2, 0, 0, 0, 8}, buf)
+		}
+	})
+}