diff --git a/graph.go b/graph.go index 958705c..490b578 100644 --- a/graph.go +++ b/graph.go @@ -352,6 +352,10 @@ func (g *Graph[K]) Add(nodes ...Node[K]) { g.assertDims(vec) insertLevel := g.randomLevel() + if _, ok := g.Lookup(key); ok { + g.Delete(key) + } + // Create layers that don't exist yet. for insertLevel >= len(g.layers) { g.layers = append(g.layers, &layer[K]{}) @@ -388,7 +392,9 @@ func (g *Graph[K]) Add(nodes ...Node[K]) { // On subsequent layers, we use the elevator node to enter the graph // at the best point. if elevator != nil { - searchPoint = layer.nodes[*elevator] + if next, ok := layer.nodes[*elevator]; ok { + searchPoint = next + } } if g.Distance == nil { @@ -406,9 +412,6 @@ func (g *Graph[K]) Add(nodes ...Node[K]) { elevator = ptr(neighborhood[0].node.Key) if insertLevel >= i { - if _, ok := layer.nodes[key]; ok { - g.Delete(key) - } // Insert the new node into the layer. layer.nodes[key] = newNode for _, node := range neighborhood { @@ -464,7 +467,9 @@ func (h *Graph[K]) search(near Vector, k int) []SearchResult[K] { for layer := len(h.layers) - 1; layer >= 0; layer-- { searchPoint := h.layers[layer].entry() if elevator != nil { - searchPoint = h.layers[layer].nodes[*elevator] + if next, ok := h.layers[layer].nodes[*elevator]; ok { + searchPoint = next + } } // Descending hierarchies diff --git a/graph_test.go b/graph_test.go index 8211507..5fb9c84 100644 --- a/graph_test.go +++ b/graph_test.go @@ -165,6 +165,44 @@ func TestGraph_AddDelete(t *testing.T) { }) } +func TestGraph_AddReplacesExistingNode(t *testing.T) { + t.Parallel() + + const numNodes = 128 + g := newTestGraph[int]() + for i := 0; i < numNodes; i++ { + g.Add(MakeNode(i, Vector{float32(i)})) + } + + for i := 0; i < numNodes; i++ { + replacement := Vector{float32(i)} + g.Add(MakeNode(i, replacement)) + + require.Equal(t, numNodes, g.Len()) + got, ok := g.Lookup(i) + require.True(t, ok) + require.Equal(t, replacement, got) + } + + for i := 0; i < numNodes; i++ { + replacement := Vector{float32(numNodes + i)} + g.Add(MakeNode(i, replacement)) + + require.Equal(t, numNodes, g.Len()) + got, ok := g.Lookup(i) + require.True(t, ok) + require.Equal(t, replacement, got) + } + + for i := 0; i < numNodes; i++ { + replacement := Vector{float32(numNodes + i)} + results := g.Search(replacement, 1) + require.NotEmpty(t, results) + require.Equal(t, i, results[0].Key) + require.Equal(t, replacement, results[0].Value) + } +} + func Benchmark_HSNW(b *testing.B) { b.ReportAllocs()