Tests pass!

Author: ammario

encode.go

@@ -2,6 +2,7 @@ package hnsw
 
 import (
 	"bufio"
+	"cmp"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -43,6 +44,16 @@ func binaryRead(r io.Reader, data interface{}) (int, error) {
 		*v = string(s)
 		return len(s), err
 
+	case *[]float32:
+		var ln int
+		_, err := binaryRead(r, &ln)
+		if err != nil {
+			return 0, err
+		}
+
+		*v = make([]float32, ln)
+		return binary.Size(*v), binary.Read(r, byteOrder, *v)
+
 	case io.ReaderFrom:
 		n, err := v.ReadFrom(r)
 		return int(n), err
@@ -73,6 +84,12 @@ func binaryWrite(w io.Writer, data any) (int, error) {
 		}
 
 		return n + n2, nil
+	case []float32:
+		n, err := binaryWrite(w, len(v))
+		if err != nil {
+			return n, err
+		}
+		return n + binary.Size(v), binary.Write(w, byteOrder, v)
 
 	default:
 		sz := binary.Size(data)
@@ -113,7 +130,7 @@ const encodingVersion = 1
 // Export writes the graph to a writer.
 //
 // T must implement io.WriterTo.
-func (h *Graph[T]) Export(w io.Writer) error {
+func (h *Graph[K]) Export(w io.Writer) error {
 	distFuncName, ok := distanceFuncToName(h.Distance)
 	if !ok {
 		return fmt.Errorf("distance function %v must be registered with RegisterDistanceFunc", h.Distance)
@@ -134,24 +151,20 @@ func (h *Graph[T]) Export(w io.Writer) error {
 		return fmt.Errorf("encode number of layers: %w", err)
 	}
 	for _, layer := range h.layers {
-		_, err = binaryWrite(w, len(layer.Nodes))
+		_, err = binaryWrite(w, len(layer.nodes))
 		if err != nil {
 			return fmt.Errorf("encode number of nodes: %w", err)
 		}
-		for _, node := range layer.Nodes {
-			_, err = binaryWrite(w, node.Point)
+		for _, node := range layer.nodes {
+			_, err = multiBinaryWrite(w, node.ID, node.Vec, len(node.neighbors))
 			if err != nil {
-				return fmt.Errorf("encode node point: %w", err)
-			}
-
-			if _, err = binaryWrite(w, len(node.neighbors)); err != nil {
-				return fmt.Errorf("encode number of neighbors: %w", err)
+				return fmt.Errorf("encode node data: %w", err)
 			}
 
 			for neighbor := range node.neighbors {
 				_, err = binaryWrite(w, neighbor)
 				if err != nil {
-					return fmt.Errorf("encode neighbor %q: %w", neighbor, err)
+					return fmt.Errorf("encode neighbor %v: %w", neighbor, err)
 				}
 			}
 		}
@@ -164,7 +177,7 @@ func (h *Graph[T]) Export(w io.Writer) error {
 // T must implement io.ReaderFrom.
 // The imported graph does not have to match the exported graph's parameters (except for
 // dimensionality). The graph will converge onto the new parameters.
-func (h *Graph[T]) Import(r io.Reader) error {
+func (h *Graph[K]) Import(r io.Reader) error {
 	var (
 		version int
 		dist    string
@@ -195,44 +208,43 @@ func (h *Graph[T]) Import(r io.Reader) error {
 		return err
 	}
 
-	h.layers = make([]*layer[T], nLayers)
+	h.layers = make([]*layer[K], nLayers)
 	for i := 0; i < nLayers; i++ {
 		var nNodes int
 		_, err = binaryRead(r, &nNodes)
 		if err != nil {
 			return err
 		}
 
-		nodes := make(map[string]*layerNode[T], nNodes)
+		nodes := make(map[K]*layerNode[K], nNodes)
 		for j := 0; j < nNodes; j++ {
-			var point T
-			_, err = binaryRead(r, &point)
-			if err != nil {
-				return fmt.Errorf("decoding node %d: %w", j, err)
-			}
-
+			var id K
+			var vec Vector
 			var nNeighbors int
-			_, err = binaryRead(r, &nNeighbors)
+			_, err = multiBinaryRead(r, &id, &vec, &nNeighbors)
 			if err != nil {
-				return fmt.Errorf("decoding number of neighbors for node %d: %w", j, err)
+				return fmt.Errorf("decoding node %d: %w", j, err)
 			}
 
-			neighbors := make([]string, nNeighbors)
+			neighbors := make([]K, nNeighbors)
 			for k := 0; k < nNeighbors; k++ {
-				var neighbor string
+				var neighbor K
 				_, err = binaryRead(r, &neighbor)
 				if err != nil {
 					return fmt.Errorf("decoding neighbor %d for node %d: %w", k, j, err)
 				}
 				neighbors[k] = neighbor
 			}
 
-			node := &layerNode[T]{
-				vec:       point,
-				neighbors: make(map[string]*layerNode[T]),
+			node := &layerNode[K]{
+				Node: Node[K]{
+					ID:  id,
+					Vec: vec,
+				},
+				neighbors: make(map[K]*layerNode[K]),
 			}
 
-			nodes[point.ID()] = node
+			nodes[id] = node
 			for _, neighbor := range neighbors {
 				node.neighbors[neighbor] = nil
 			}
@@ -243,7 +255,7 @@ func (h *Graph[T]) Import(r io.Reader) error {
 				node.neighbors[id] = nodes[id]
 			}
 		}
-		h.layers[i] = &layer[T]{nodes: nodes}
+		h.layers[i] = &layer[K]{nodes: nodes}
 	}
 
 	return nil
@@ -253,8 +265,8 @@ func (h *Graph[T]) Import(r io.Reader) error {
 // changes to a file upon calls to Save. It is more convenient
 // but less powerful than calling Graph.Export and Graph.Import
 // directly.
-type SavedGraph[T Embeddable] struct {
-	*Graph[T]
+type SavedGraph[K cmp.Ordered] struct {
+	*Graph[K]
 	Path string
 }
 
@@ -265,7 +277,7 @@ type SavedGraph[T Embeddable] struct {
 //
 // It does not hold open a file descriptor, so SavedGraph can be forgotten
 // without ever calling Save.
-func LoadSavedGraph[T Embeddable](path string) (*SavedGraph[T], error) {
+func LoadSavedGraph[K cmp.Ordered](path string) (*SavedGraph[K], error) {
 	f, err := os.OpenFile(path, os.O_RDWR|os.O_CREATE, 0o600)
 	if err != nil {
 		return nil, err
@@ -276,15 +288,15 @@ func LoadSavedGraph[T Embeddable](path string) (*SavedGraph[T], error) {
 		return nil, err
 	}
 
-	g := NewGraph[T]()
+	g := NewGraph[K]()
 	if info.Size() > 0 {
 		err = g.Import(bufio.NewReader(f))
 		if err != nil {
 			return nil, fmt.Errorf("import: %w", err)
 		}
 	}
 
-	return &SavedGraph[T]{Graph: g, Path: path}, nil
+	return &SavedGraph[K]{Graph: g, Path: path}, nil
 }
 
 // Save writes the graph to the file.

encode_test.go

@@ -2,8 +2,7 @@ package hnsw
 
 import (
 	"bytes"
-	"math/rand"
-	"strconv"
+	"cmp"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -50,21 +49,21 @@ func Test_binaryWrite_string(t *testing.T) {
 	require.Empty(t, buf.Bytes())
 }
 
-func verifyGraphNodes[T Embeddable](t *testing.T, g *Graph[T]) {
+func verifyGraphNodes[K cmp.Ordered](t *testing.T, g *Graph[K]) {
 	for _, layer := range g.layers {
-		for _, node := range layer.Nodes {
+		for _, node := range layer.nodes {
 			for neighborKey, neighbor := range node.neighbors {
-				_, ok := layer.Nodes[neighbor.Point.ID()]
+				_, ok := layer.nodes[neighbor.ID]
 				if !ok {
 					t.Errorf(
-						"node %s has neighbor %s, but neighbor does not exist",
-						node.Point.ID(), neighbor.Point.ID(),
+						"node %v has neighbor %v, but neighbor does not exist",
+						node.ID, neighbor.ID,
 					)
 				}
 
-				if neighborKey != neighbor.Point.ID() {
-					t.Errorf("node %s has neighbor %s, but neighbor key is %s", node.Point.ID(),
-						neighbor.Point.ID(),
+				if neighborKey != neighbor.ID {
+					t.Errorf("node %v has neighbor %v, but neighbor key is %v", node.ID,
+						neighbor.ID,
 						neighborKey,
 					)
 				}
@@ -74,10 +73,10 @@ func verifyGraphNodes[T Embeddable](t *testing.T, g *Graph[T]) {
 }
 
 // requireGraphApproxEquals checks that two graphs are equal.
-func requireGraphApproxEquals[T Embeddable](t *testing.T, g1, g2 *Graph[T]) {
+func requireGraphApproxEquals[K cmp.Ordered](t *testing.T, g1, g2 *Graph[K]) {
 	require.Equal(t, g1.Len(), g2.Len())
-	a1 := Analyzer[T]{g1}
-	a2 := Analyzer[T]{g2}
+	a1 := Analyzer[K]{g1}
+	a2 := Analyzer[K]{g2}
 
 	require.Equal(
 		t,
@@ -119,11 +118,13 @@ func requireGraphApproxEquals[T Embeddable](t *testing.T, g1, g2 *Graph[T]) {
 }
 
 func TestGraph_ExportImport(t *testing.T) {
-	rng := rand.New(rand.NewSource(0))
-
-	g1 := newTestGraph[Vector]()
+	g1 := newTestGraph[int]()
 	for i := 0; i < 128; i++ {
-		g1.Add(MakeVector(strconv.Itoa(i), []float32{rng.Float32()}))
+		g1.Add(
+			Node[int]{
+				i, randFloats(1),
+			},
+		)
 	}
 
 	buf := &bytes.Buffer{}
@@ -132,7 +133,7 @@ func TestGraph_ExportImport(t *testing.T) {
 
 	// Don't use newTestGraph to ensure parameters
 	// are imported.
-	g2 := &Graph[Vector]{}
+	g2 := &Graph[int]{}
 	err = g2.Import(buf)
 	require.NoError(t, err)
 
@@ -157,17 +158,21 @@ func TestGraph_ExportImport(t *testing.T) {
 func TestSavedGraph(t *testing.T) {
 	dir := t.TempDir()
 
-	g1, err := LoadSavedGraph[Vector](dir + "/graph")
+	g1, err := LoadSavedGraph[int](dir + "/graph")
 	require.NoError(t, err)
 	require.Equal(t, 0, g1.Len())
 	for i := 0; i < 128; i++ {
-		g1.Add(MakeVector(strconv.Itoa(i), []float32{float32(i)}))
+		g1.Add(
+			Node[int]{
+				i, randFloats(1),
+			},
+		)
 	}
 
 	err = g1.Save()
 	require.NoError(t, err)
 
-	g2, err := LoadSavedGraph[Vector](dir + "/graph")
+	g2, err := LoadSavedGraph[int](dir + "/graph")
 	require.NoError(t, err)
 
 	requireGraphApproxEquals(t, g1.Graph, g2.Graph)
@@ -177,9 +182,13 @@ const benchGraphSize = 100
 
 func BenchmarkGraph_Import(b *testing.B) {
 	b.ReportAllocs()
-	g := newTestGraph[Vector]()
+	g := newTestGraph[int]()
 	for i := 0; i < benchGraphSize; i++ {
-		g.Add(MakeVector(strconv.Itoa(i), randFloats(100)))
+		g.Add(
+			Node[int]{
+				i, randFloats(256),
+			},
+		)
 	}
 
 	buf := &bytes.Buffer{}
@@ -192,7 +201,7 @@ func BenchmarkGraph_Import(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		b.StopTimer()
 		rdr := bytes.NewReader(buf.Bytes())
-		g := newTestGraph[Vector]()
+		g := newTestGraph[int]()
 		b.StartTimer()
 		err = g.Import(rdr)
 		require.NoError(b, err)
@@ -201,9 +210,13 @@ func BenchmarkGraph_Import(b *testing.B) {
 
 func BenchmarkGraph_Export(b *testing.B) {
 	b.ReportAllocs()
-	g := newTestGraph[Vector]()
+	g := newTestGraph[int]()
 	for i := 0; i < benchGraphSize; i++ {
-		g.Add(MakeVector(strconv.Itoa(i), randFloats(256)))
+		g.Add(
+			Node[int]{
+				i, randFloats(256),
+			},
+		)
 	}
 
 	var buf bytes.Buffer

example/readme/main.go

@@ -7,16 +7,16 @@ import (
 )
 
 func main() {
-	g := hnsw.NewGraph[hnsw.Vector]()
+	g := hnsw.NewGraph[int]()
 	g.Add(
-		hnsw.MakeVector("1", []float32{1, 1, 1}),
-		hnsw.MakeVector("2", []float32{1, -1, 0.999}),
-		hnsw.MakeVector("3", []float32{1, 0, -0.5}),
+		hnsw.MakeNode(1, []float32{1, 1, 1}),
+		hnsw.MakeNode(2, []float32{1, -1, 0.999}),
+		hnsw.MakeNode(3, []float32{1, 0, -0.5}),
 	)
 
 	neighbors := g.Search(
 		[]float32{0.5, 0.5, 0.5},
 		1,
 	)
-	fmt.Printf("best friend: %v\n", neighbors[0].Embedding())
+	fmt.Printf("best friend: %v\n", neighbors[0].Vec)
 }

graph.go

@@ -20,6 +20,10 @@ type Node[K cmp.Ordered] struct {
 	Vec Vector
 }
 
+func MakeNode[K cmp.Ordered](id K, vec Vector) Node[K] {
+	return Node[K]{ID: id, Vec: vec}
+}
+
 // layerNode is a node in a layer of the graph.
 type layerNode[K cmp.Ordered] struct {
 	Node[K]

graph_test.go

@@ -63,7 +63,7 @@ func Test_layerNode_search(t *testing.T) {
 
 	best := entry.search(2, 4, []float32{4}, EuclideanDistance)
 
-	require.Equal(t, 4, best[0].node.ID)
+	require.Equal(t, 5, best[0].node.ID)
 	require.Equal(t, 3, best[1].node.ID)
 	require.Len(t, best, 2)
 }