fogfish · fogfish · Aug 17, 2024 · Aug 17, 2024
diff --git a/go.mod b/go.mod
@@ -1,19 +1,18 @@
 module github.com/fogfish/hnsw
 
-go 1.21.3
+go 1.22.2
 
 require (
 	github.com/bits-and-blooms/bitset v1.13.0
 	github.com/fogfish/faults v0.2.0
 	github.com/fogfish/guid/v2 v2.0.4
 	github.com/fogfish/it/v2 v2.0.1
 	github.com/kelindar/binary v1.0.19
-	github.com/kshard/vector v0.0.4
-
+	github.com/kshard/vector v0.1.1
 )
 
 require (
 	github.com/chewxy/math32 v1.10.1 // indirect
 	github.com/fogfish/golem/pure v0.10.1 // indirect
-	golang.org/x/sys v0.18.0 // indirect
+	golang.org/x/sys v0.24.0 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -14,11 +14,11 @@ github.com/fogfish/it/v2 v2.0.1 h1:vu3kV2xzYDPHoMHMABxXeu5CoMcTfRc4gkWkzOUkRJY=
 github.com/fogfish/it/v2 v2.0.1/go.mod h1:h5FdKaEQT4sUEykiVkB8VV4jX27XabFVeWhoDZaRZtE=
 github.com/kelindar/binary v1.0.19 h1:DNyQCtKjkLhBh9pnP49OWREddLB0Mho+1U/AOt/Qzxw=
 github.com/kelindar/binary v1.0.19/go.mod h1:/twdz8gRLNMffx0U4UOgqm1LywPs6nd9YK2TX52MDh8=
-github.com/kshard/vector v0.0.4 h1:1CAo52yGHZGx0cb3O0UYdHqzNCH4G589BaJTQC7nWIQ=
-github.com/kshard/vector v0.0.4/go.mod h1:l5c902GqrnE4/LvJAmSCbtGorXeMe9si5SHtwckr8jc=
+github.com/kshard/vector v0.1.1 h1:4sz566fGYEyYdDBDLUJHBMrb7XKdR4UJVsxBEy+UpPM=
+github.com/kshard/vector v0.1.1/go.mod h1:gHkE5jmnnl1T7hX5rFHuY7lWbg35XzWieRl+qsnihXU=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w=
 github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
-golang.org/x/sys v0.18.0 h1:DBdB3niSjOA/O0blCZBqDefyWNYveAYMNF1Wum0DYQ4=
-golang.org/x/sys v0.18.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
+golang.org/x/sys v0.24.0 h1:Twjiwq9dn6R1fQcyiK+wQyHWfaz/BJB+YIpzU/Cv3Xg=
+golang.org/x/sys v0.24.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
diff --git a/hnsw.go b/hnsw.go
@@ -114,7 +114,7 @@ func (h *HNSW[Vector]) Head() Vector { return h.heap[h.head].Vector }
 // Return current level
 func (h *HNSW[Vector]) Level() int { return h.level }
 
-// Return current size
+// Return number of vectors in the data structure
 func (h *HNSW[Vector]) Size() int { return len(h.heap) }
 
 // Calculate distance between two vectors using defined surface distance function.

diff --git a/hnsw_test.go b/hnsw_test.go
@@ -23,53 +23,96 @@ const (
 )
 
 var (
-	rnd = rand.NewSource(0x211111111)
+	rnd     = rand.NewSource(0x211111111)
+	vectors = rndVectors()
 )
 
+func sut(surface surface.Surface[surface.F32]) *hnsw.HNSW[vector.VF32] {
+	return hnsw.New(
+		vector.SurfaceVF32(surface),
+		hnsw.WithRandomSource(rnd),
+		hnsw.WithM0(64),
+	)
+}
+
+func TestInsert(t *testing.T) {
+	for _, df := range []surface.Surface[surface.F32]{
+		surface.Euclidean(),
+		surface.Cosine(),
+	} {
+		index := sut(df)
+		for i, v := range vectors {
+			index.Insert(vector.VF32{Key: uint32(i), Vec: v})
+		}
+
+		for _, q := range nodes(index) {
+			seq := index.Search(q, 1, 100)
+			if seq[0].Key != q.Key {
+				t.Errorf("Not found %v in %v", q, seq)
+			}
+		}
+	}
+}
+
+func TestUpdate(t *testing.T) {
+	for _, df := range []surface.Surface[surface.F32]{
+		surface.Euclidean(),
+		surface.Cosine(),
+	} {
+		index := sut(df)
+		for i, v := range vectors {
+			index.Insert(vector.VF32{Key: uint32(i), Vec: v})
+		}
+
+		// Update vectors (set new key)
+		for i, v := range vectors {
+			key := uint32((1 << 31) | i)
+			index.Insert(vector.VF32{Key: key, Vec: v})
+		}
+
+		// Check that key is updated
+		for _, n := range nodes(index) {
+			if (n.Key & (1 << 31)) == 0 {
+				t.Errorf("Not updated %v", n.Key)
+			}
+		}
+	}
+}
+
+//------------------------------------------------------------------------------
+
 func random() float32 {
 again:
-	f := float32(rnd.Int63()) / (1 << 63)
+	f := float64(rnd.Int63()) / (1 << 63)
 	if f == 1 {
 		goto again // resample; this branch is taken O(never)
 	}
-	return f
+	return float32(f)
 }
 
-func vRand() surface.F32 {
+func rndVector() surface.F32 {
 	v := make(surface.F32, d)
 	for i := 0; i < d; i++ {
-		v[i] = random()
+		v[i] = 2*random() - 1
 	}
 	return v
 }
 
-func TestHNSW(t *testing.T) {
+func rndVectors() []surface.F32 {
 	vecs := make([]surface.F32, n)
 	for i := 0; i < n; i++ {
-		vecs[i] = vRand()
-	}
-
-	index := hnsw.New(
-		vector.SurfaceVF32(surface.Euclidean()),
-		hnsw.WithRandomSource(rnd),
-	)
-
-	for i, v := range vecs {
-		index.Insert(vector.VF32{Key: uint32(i), Vec: v})
+		vecs[i] = rndVector()
 	}
+	return vecs
+}
 
+func nodes(index *hnsw.HNSW[vector.VF32]) []vector.VF32 {
 	nodes := make([]vector.VF32, 0)
 	index.ForAll(0,
 		func(rank int, vector vector.VF32, edges []vector.VF32) error {
 			nodes = append(nodes, vector)
 			return nil
 		},
 	)
-
-	for _, q := range nodes {
-		seq := index.Search(q, 4, 100)
-		if seq[0].Key != q.Key {
-			t.Errorf("Not found %v in %v", q, seq)
-		}
-	}
+	return nodes
 }
diff --git a/insert.go b/insert.go
@@ -15,6 +15,10 @@ import (
 	"github.com/fogfish/hnsw/internal/types"
 )
 
+// anything in the range (-eps; +eps) is considered as similar item
+const minEps = -1e-6
+const maxEps = +1e-6
+
 // generate random float from random source generator
 func (h *HNSW[Vector]) rand() float64 {
 again:
@@ -87,6 +91,14 @@ func (h *HNSW[Vector]) Insert(v Vector) {
 		for i := w.Len() - 1; i >= 0; i-- {
 			candidate := w.Deq()
 			edges[i] = candidate.Addr
+
+			// Consider the update
+			if minEps < candidate.Distance && candidate.Distance < maxEps {
+				if h.surface.Equal(h.heap[candidate.Addr].Vector, v) {
+					h.heap[candidate.Addr].Vector = v
+					return
+				}
+			}
 		}
 		node.Connections[lvl] = edges
 	}