Fix and add regression tests for issue #2724 related to Vector Index corruption (#2737)

Author: JKamsker

LiteDB.Tests/Engine/DropCollection_Tests.cs

@@ -3,6 +3,7 @@
 using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Threading.Tasks;
 using FluentAssertions;
 using LiteDB;
 using LiteDB.Engine;
@@ -212,6 +213,56 @@ public void DropCollection_WithVectorIndex_ReclaimsTrackedPages()
             }
         }
 
+        /// <summary>
+        /// Regression test for LiteDB issue #2724.
+        /// </summary>
+        /// <remarks>
+        /// https://github.com/litedb-org/LiteDB/issues/2724
+        /// </remarks>
+        [Fact]
+        public void DropCollection_WithVectorIndex_AfterParallelUpserts_DoesNotThrow()
+        {
+            using var file = new TempFile();
+
+            const int documentCount = 64;
+            const int updateCount = 256;
+            const ushort dimensions = 32;
+
+            using (var db = DatabaseFactory.Create(
+                TestDatabaseType.Disk,
+                $"Filename={file.Filename};Connection=Shared"))
+            {
+                var collection = db.GetCollection<VectorDocument>("docs");
+                var options = new VectorIndexOptions(dimensions, VectorDistanceMetric.Cosine);
+
+                collection.EnsureIndex(VectorIndexName, x => x.Embedding, options);
+
+                var initial = Enumerable.Range(1, documentCount)
+                    .Select(i => new VectorDocument
+                    {
+                        Id = i,
+                        Embedding = CreateLargeVector(i, dimensions)
+                    })
+                    .ToList();
+
+                collection.Insert(initial);
+
+                Parallel.For(0, updateCount, i =>
+                {
+                    var id = (i % documentCount) + 1;
+                    collection.Upsert(new VectorDocument
+                    {
+                        Id = id,
+                        Embedding = CreateLargeVector(i + 5000, dimensions)
+                    });
+                });
+
+                Action drop = () => db.DropCollection("docs");
+
+                drop.Should().NotThrow();
+            }
+        }
+
         private static Dictionary<PageType, int> CountPagesByType(string filename)
         {
             var counts = new Dictionary<PageType, int>();

LiteDB.Tests/Query/VectorIndex_Tests.cs

@@ -111,6 +111,80 @@ private static float[] CreateVector(Random random, int dimensions)
             return vector;
         }
 
+        /// <summary>
+        /// Regression test for LiteDB issue #2724.
+        /// </summary>
+        /// <remarks>
+        /// https://github.com/litedb-org/LiteDB/issues/2724
+        /// </remarks>
+        [Fact]
+        public void VectorIndex_ShouldMaintainBidirectionalEdges()
+        {
+            const int count = 256;
+            const int dimensions = 8;
+
+            using var db = DatabaseFactory.Create();
+            var collection = db.GetCollection<VectorDocument>("docs");
+
+            collection.EnsureIndex(
+                "embedding_idx",
+                BsonExpression.Create("$.Embedding"),
+                new VectorIndexOptions(dimensions, VectorDistanceMetric.Cosine));
+
+            var random = new Random(123);
+
+            for (var i = 1; i <= count; i++)
+            {
+                collection.Insert(new VectorDocument
+                {
+                    Id = i,
+                    Embedding = CreateVector(random, dimensions)
+                });
+            }
+
+            db.Checkpoint();
+
+            InspectVectorIndex(db, "docs", (snapshot, collation, metadata) =>
+            {
+                CountNodes(snapshot, metadata.Root).Should().Be(count);
+
+                var visited = new HashSet<PageAddress>();
+                var queue = new Queue<PageAddress>();
+                queue.Enqueue(metadata.Root);
+
+                while (queue.Count > 0)
+                {
+                    var address = queue.Dequeue();
+                    if (!visited.Add(address))
+                    {
+                        continue;
+                    }
+
+                    var node = snapshot.GetPage<VectorIndexPage>(address.PageID).GetNode(address.Index);
+
+                    for (var level = 0; level < node.LevelCount; level++)
+                    {
+                        foreach (var neighbor in node.GetNeighbors(level))
+                        {
+                            if (neighbor.IsEmpty)
+                            {
+                                continue;
+                            }
+
+                            queue.Enqueue(neighbor);
+
+                            var neighborNode = snapshot.GetPage<VectorIndexPage>(neighbor.PageID).GetNode(neighbor.Index);
+
+                            (level < neighborNode.LevelCount).Should().BeTrue();
+                            neighborNode.GetNeighbors(level).Should().Contain(address);
+                        }
+                    }
+                }
+
+                return true;
+            }).Should().BeTrue();
+        }
+
         private static float[] ReadExternalVector(DataService dataService, PageAddress start, int dimensions, out int blocksRead)
         {
             var totalBytes = dimensions * sizeof(float);

LiteDB/Engine/Services/VectorIndexService.cs

@@ -293,11 +293,16 @@ private void Insert(VectorIndexMetadata metadata, PageAddress dataBlock, float[]
                     candidates.Where(x => x.Address != newAddress).ToList(),
                     VectorIndexNode.MaxNeighborsPerLevel);
 
-                node.SetNeighbors(level, selected.Select(x => x.Address).ToList());
+                var selectedAddresses = selected.Select(x => x.Address).ToList();
 
-                foreach (var neighbor in selected)
+                node.SetNeighbors(level, selectedAddresses);
+
+                foreach (var neighbor in selectedAddresses)
                 {
-                    this.EnsureBidirectional(metadata, neighbor.Address, newAddress, level, vectorCache);
+                    if (!this.EnsureBidirectional(metadata, neighbor, newAddress, level, vectorCache))
+                    {
+                        node.RemoveNeighbor(level, neighbor);
+                    }
                 }
 
                 if (selected.Count > 0)
@@ -420,10 +425,11 @@ private List<NodeDistance> SearchLayer(
             return this.SelectNeighbors(results, Math.Max(1, maxResults));
         }
 
-        private void EnsureBidirectional(VectorIndexMetadata metadata, PageAddress source, PageAddress target, int level, Dictionary<PageAddress, float[]> vectorCache)
+        private bool EnsureBidirectional(VectorIndexMetadata metadata, PageAddress source, PageAddress target, int level, Dictionary<PageAddress, float[]> vectorCache)
         {
             var node = this.GetNode(source);
-            var neighbors = node.GetNeighbors(level).ToList();
+            var neighbors = node.GetNeighbors(level).Where(x => !x.IsEmpty).ToList();
+            var before = neighbors.ToList();
 
             if (!neighbors.Contains(target))
             {
@@ -432,6 +438,33 @@ private void EnsureBidirectional(VectorIndexMetadata metadata, PageAddress sourc
 
             var pruned = this.PruneNeighbors(metadata, source, neighbors, vectorCache);
             node.SetNeighbors(level, pruned);
+
+            foreach (var removed in before)
+            {
+                if (!pruned.Contains(removed))
+                {
+                    this.RemoveBackLink(removed, source, level);
+                }
+            }
+
+            return pruned.Contains(target);
+        }
+
+        private void RemoveBackLink(PageAddress source, PageAddress target, int level)
+        {
+            if (source.IsEmpty || target.IsEmpty)
+            {
+                return;
+            }
+
+            var node = this.GetNode(source);
+
+            if (level < 0 || level >= node.LevelCount)
+            {
+                return;
+            }
+
+            node.RemoveNeighbor(level, target);
         }
 
         private IReadOnlyList<PageAddress> PruneNeighbors(VectorIndexMetadata metadata, PageAddress source, List<PageAddress> neighbors, Dictionary<PageAddress, float[]> vectorCache)