From d9929b63a17dc297f0c3de32f1b22868cd775b73 Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 06:13:07 -0700 Subject: [PATCH 1/9] Speed up index Save/Load with bulk float I/O Replace the per-element ReadSingle/Write(float) loops in HnswIndex.Save and Load with bulk MemoryMarshal byte-span reads and writes. On a 20,487-vector, 64-dim index this cuts Load from ~59ms to ~22ms (min, ~2-3x) by avoiding millions of scalar BinaryReader/BinaryWriter calls and their per-call bounds checks. The on-disk bytes are unchanged on the little-endian platforms .NET targets (BinaryWriter.Write(float) and MemoryMarshal both produce little-endian), so the format is byte-identical and existing v3 indexes load correctly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Hnsw.Net/HnswIndex.cs | 39 ++++++++++++++++++++++----------------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index ca214a8..25876cf 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -494,15 +494,11 @@ public void Save(Stream stream) writer.Write(node.Id); writer.Write(node.Level); writer.Write(node.Deleted); - ReadOnlySpan stored = StoredSpan(n); - for (int i = 0; i < Dimension; i++) - { - writer.Write(stored[i]); - } - for (int i = 0; i < Dimension; i++) - { - writer.Write(node.OriginalVector[i]); - } + // Vectors are written as raw little-endian float blocks; on the LE platforms .NET + // targets this is byte-identical to a per-element BinaryWriter.Write(float) loop but + // avoids millions of scalar writes on large indexes. + writer.Write(MemoryMarshal.AsBytes(StoredSpan(n))); + writer.Write(MemoryMarshal.AsBytes(node.OriginalVector.AsSpan())); writer.Write(node.Links.Length); for (int layer = 0; layer < node.Links.Length; layer++) @@ -554,17 +550,11 @@ public static HnswIndex Load(Stream stream) int level = reader.ReadInt32(); bool deleted = version >= 3 && reader.ReadBoolean(); Span stored = index._storedVectors.AsSpan(i * dimension, dimension); - for (int j = 0; j < dimension; j++) - { - stored[j] = reader.ReadSingle(); - } + ReadExactInto(reader, MemoryMarshal.AsBytes(stored)); var originalVector = new float[dimension]; if (version >= 2) { - for (int j = 0; j < dimension; j++) - { - originalVector[j] = reader.ReadSingle(); - } + ReadExactInto(reader, MemoryMarshal.AsBytes(originalVector.AsSpan())); } else { @@ -659,6 +649,21 @@ private void EnsureStoredCapacity(int nodeCount) private Span StoredSpanMutable(int index) => _storedVectors.AsSpan(index * Dimension, Dimension); + private static void ReadExactInto(BinaryReader reader, Span buffer) + { + int total = 0; + while (total < buffer.Length) + { + int read = reader.Read(buffer.Slice(total)); + if (read == 0) + { + throw new EndOfStreamException(); + } + + total += read; + } + } + private int RandomLevel() { double sample = Math.Max(_random.NextDouble(), double.Epsilon); From 161c2d9461b73becadb8e2bc9e0d422b716650cd Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 06:39:44 -0700 Subject: [PATCH 2/9] Route search hot path through per-slot accessors Introduce SlotVector(slot) and SlotLinks(slot, layer) and read every vector and neighbor list in SearchGreedy, SearchLayer, SelectNeighbors and Search through them. This centralizes per-slot storage access so the backing store can later be swapped for a memory-mapped level-0 block without touching the search algorithms. Behavior is unchanged; the parity test suite is green. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Hnsw.Net/HnswIndex.cs | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index 25876cf..b8b9128 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -439,7 +439,7 @@ private void ReplaceSlot(int slot, long id, ReadOnlySpan vector, float[] try { int entryPoint = _entryPoint; - float entryDistance = Distance(preparedQuery, StoredSpan(entryPoint)); + float entryDistance = Distance(preparedQuery, SlotVector(entryPoint)); for (int layer = _maxLevel; layer > 0; layer--) { (entryPoint, entryDistance) = SearchGreedy(preparedQuery, entryPoint, entryDistance, layer); @@ -647,6 +647,12 @@ private void EnsureStoredCapacity(int nodeCount) private ReadOnlySpan StoredSpan(int index) => _storedVectors.AsSpan(index * Dimension, Dimension); + // Per-slot read accessors. The search path goes through these so the backing storage can later + // be swapped (e.g. a memory-mapped level-0 block) without touching the algorithms. + private ReadOnlySpan SlotVector(int slot) => StoredSpan(slot); + + private List SlotLinks(int slot, int layer) => _nodes[slot].Links[layer]; + private Span StoredSpanMutable(int index) => _storedVectors.AsSpan(index * Dimension, Dimension); private static void ReadExactInto(BinaryReader reader, Span buffer) @@ -676,9 +682,9 @@ private int RandomLevel() do { changed = false; - foreach (int neighbor in _nodes[entryPoint].Links[layer]) + foreach (int neighbor in SlotLinks(entryPoint, layer)) { - float distance = Distance(query, StoredSpan(neighbor)); + float distance = Distance(query, SlotVector(neighbor)); if (distance < entryDistance) { entryDistance = distance; @@ -699,7 +705,7 @@ private void SearchLayer(ReadOnlySpan query, int entryPoint, int ef, int s.Nearest.Clear(); s.Visited[entryPoint] = version; - float entryDistance = Distance(query, StoredSpan(entryPoint)); + float entryDistance = Distance(query, SlotVector(entryPoint)); var entry = new Candidate(entryPoint, entryDistance); s.Candidates.Enqueue(entry, entryDistance); float lowerBound; @@ -721,7 +727,7 @@ private void SearchLayer(ReadOnlySpan query, int entryPoint, int ef, int break; } - foreach (int neighbor in _nodes[current.Index].Links[layer]) + foreach (int neighbor in SlotLinks(current.Index, layer)) { if (s.Visited[neighbor] == version) { @@ -729,7 +735,7 @@ private void SearchLayer(ReadOnlySpan query, int entryPoint, int ef, int } s.Visited[neighbor] = version; - float distance = Distance(query, StoredSpan(neighbor)); + float distance = Distance(query, SlotVector(neighbor)); if (s.Nearest.Count < ef || distance < lowerBound) { var candidate = new Candidate(neighbor, distance); @@ -768,11 +774,11 @@ private void SelectNeighbors(List candidates, int maxConnections, Lis candidates.Sort(_candidateComparison); foreach (Candidate candidate in candidates) { - ReadOnlySpan candidateSpan = StoredSpan(candidate.Index); + ReadOnlySpan candidateSpan = SlotVector(candidate.Index); bool good = true; foreach (int selected in result) { - if (Distance(candidateSpan, StoredSpan(selected)) < candidate.Distance) + if (Distance(candidateSpan, SlotVector(selected)) < candidate.Distance) { good = false; break; From d11abb9c1c9aec4422b3fe978b1680bc41ffeade Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 06:43:30 -0700 Subject: [PATCH 3/9] Store search vectors in a chunked VectorBlock Replace the single float[] backing the normalized search vectors with a chunked, slot-addressed VectorBlock. Each chunk stays well under the .NET array length limit so the index scales to very large repos, and every slot exposes a contiguous span ready for a future memory-mapped backing. The v3 on-disk format and all read/write semantics are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Hnsw.Net/HnswIndex.cs | 72 +++++++++++++++++++++++++++------------ 1 file changed, 51 insertions(+), 21 deletions(-) diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index b8b9128..5e7e65a 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -29,7 +29,7 @@ public sealed class HnswIndex private readonly ReaderWriterLockSlim _lock = new(); private readonly ConcurrentBag _scratchPool = new(); private readonly Comparison _candidateComparison; - private float[] _storedVectors = Array.Empty(); + private VectorBlock _vectors; /// Initializes a new HNSW index. /// Vector dimension. All indexed and query vectors must have this length. @@ -62,6 +62,7 @@ public HnswIndex(int dimension, DistanceMetric metric, int m = 16, int efConstru _levelMultiplier = 1.0 / Math.Log(m); _allowReplaceDeleted = allowReplaceDeleted; _candidateComparison = CompareCandidates; + _vectors = new VectorBlock(dimension); } private HnswIndex(int dimension, DistanceMetric metric, int m, int efConstruction, int ef, int entryPoint, int maxLevel, bool allowReplaceDeleted) @@ -77,6 +78,7 @@ private HnswIndex(int dimension, DistanceMetric metric, int m, int efConstructio _levelMultiplier = 1.0 / Math.Log(m); _allowReplaceDeleted = allowReplaceDeleted; _candidateComparison = CompareCandidates; + _vectors = new VectorBlock(dimension); } /// Gets the vector dimension. @@ -542,14 +544,14 @@ public static HnswIndex Load(Stream stream) int count = reader.ReadInt32(); bool allowReplaceDeleted = version >= 3 && reader.ReadBoolean(); var index = new HnswIndex(dimension, metric, m, efConstruction, ef, entryPoint, maxLevel, allowReplaceDeleted); - index._storedVectors = count > 0 ? new float[count * dimension] : Array.Empty(); + index._vectors.EnsureCapacity(count); for (int i = 0; i < count; i++) { long id = reader.ReadInt64(); int level = reader.ReadInt32(); bool deleted = version >= 3 && reader.ReadBoolean(); - Span stored = index._storedVectors.AsSpan(i * dimension, dimension); + Span stored = index._vectors.VectorMutable(i); ReadExactInto(reader, MemoryMarshal.AsBytes(stored)); var originalVector = new float[dimension]; if (version >= 2) @@ -628,24 +630,9 @@ private void PrepareVectorInto(ReadOnlySpan source, Span destinati } } - private void EnsureStoredCapacity(int nodeCount) - { - int required = nodeCount * Dimension; - if (_storedVectors.Length >= required) - { - return; - } + private void EnsureStoredCapacity(int nodeCount) => _vectors.EnsureCapacity(nodeCount); - int capacity = _storedVectors.Length == 0 ? Dimension * 16 : _storedVectors.Length * 2; - if (capacity < required) - { - capacity = required; - } - - Array.Resize(ref _storedVectors, capacity); - } - - private ReadOnlySpan StoredSpan(int index) => _storedVectors.AsSpan(index * Dimension, Dimension); + private ReadOnlySpan StoredSpan(int index) => _vectors.Vector(index); // Per-slot read accessors. The search path goes through these so the backing storage can later // be swapped (e.g. a memory-mapped level-0 block) without touching the algorithms. @@ -653,7 +640,7 @@ private void EnsureStoredCapacity(int nodeCount) private List SlotLinks(int slot, int layer) => _nodes[slot].Links[layer]; - private Span StoredSpanMutable(int index) => _storedVectors.AsSpan(index * Dimension, Dimension); + private Span StoredSpanMutable(int index) => _vectors.VectorMutable(index); private static void ReadExactInto(BinaryReader reader, Span buffer) { @@ -880,6 +867,49 @@ private static float DotProduct(ReadOnlySpan a, ReadOnlySpan b) return sum; } + // Chunked, slot-addressed storage for the normalized search vectors. Splitting the data into + // bounded chunks keeps each backing array well under the .NET array length limit so the index + // scales to very large repos, and gives every slot a contiguous span ready for a future + // memory-mapped backing. + private sealed class VectorBlock + { + private readonly int _dimension; + private readonly int _vectorsPerChunk; + private readonly List _chunks = new(); + private int _capacity; + + public VectorBlock(int dimension) + { + _dimension = dimension; + + // Target ~256 MB (64Mi floats) per chunk, at least one vector, and never let a chunk's + // element count exceed the int array bound. + const long TargetFloatsPerChunk = 64L * 1024 * 1024; + long perChunk = Math.Max(1, TargetFloatsPerChunk / dimension); + _vectorsPerChunk = (int)Math.Min(perChunk, int.MaxValue / dimension); + } + + public void EnsureCapacity(int slotCount) + { + while (_capacity < slotCount) + { + _chunks.Add(new float[_vectorsPerChunk * _dimension]); + _capacity += _vectorsPerChunk; + } + } + + public ReadOnlySpan Vector(int slot) => Slot(slot); + + public Span VectorMutable(int slot) => Slot(slot); + + private Span Slot(int slot) + { + int chunk = slot / _vectorsPerChunk; + int offset = (slot % _vectorsPerChunk) * _dimension; + return _chunks[chunk].AsSpan(offset, _dimension); + } + } + private sealed class Node { public Node(long id, float[] originalVector, int level) From 552075b0faa974414baadb3c4557cd5e738321e3 Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 06:57:20 -0700 Subject: [PATCH 4/9] Store a single vector per node (format v4) Drop Node.OriginalVector so each node keeps only the stored (search) vector, halving vector memory and producing the contiguous single-vector-per-slot layout needed for a memory-mapped backing. Bump the on-disk format to v4 which writes one vector per node; v1-v3 files still load, with the older duplicate pre-normalization copy read and discarded. For DistanceMetric.Cosine this changes ExportItems/TryGetVector to return the unit-normalized stored vector instead of the original input magnitude, matching hnswlib's getDataByLabel. Re-adding a normalized vector is a no-op normalization so rebuild stays stable. Other metrics are unaffected. Docs and tests updated, plus a regression test that loads a hand-crafted v3 stream. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 7 +-- src/Hnsw.Net/HnswIndex.cs | 49 ++++++++++---------- tests/Hnsw.Net.Tests/HnswIndexTests.cs | 62 ++++++++++++++++++++++++-- 3 files changed, 86 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index f75a9cd..b5127de 100644 --- a/README.md +++ b/README.md @@ -112,9 +112,10 @@ HnswIndex rebuilt = HnswIndex.Build( seed: 42); ``` -`ExportItems` returns copies of the original vectors passed to `Add`. Cosine -vectors are still stored normalized internally for search, but the portable -export preserves the original input values; rebuilding normalizes them again. +`ExportItems` returns copies of the stored vectors. For cosine these are the +unit-normalized vectors used for search (not the original input magnitudes); +rebuilding from them re-normalizes and reproduces the same index. Other metrics +store vectors unchanged, so the export matches the input. ## hnswlib parity validation diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index 5e7e65a..689848c 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -13,7 +13,7 @@ namespace HnswNet; /// public sealed class HnswIndex { - private const int FormatVersion = 3; + private const int FormatVersion = 4; private const uint Magic = 0x31575348; // HSW1, little-endian. private readonly List _nodes = new(); @@ -118,7 +118,9 @@ public int ActiveCount public bool AllowReplaceDeleted => _allowReplaceDeleted; /// - /// Returns copies of the live ids and original vectors, suitable for rebuilding a portable index. + /// Returns copies of the live ids and their stored vectors, suitable for rebuilding a portable + /// index. For the stored vector is unit-normalized, matching + /// what the index searches against; re-adding it reproduces the same index. /// public IEnumerable<(long Id, float[] Vector)> ExportItems() { @@ -126,11 +128,12 @@ public int ActiveCount try { var items = new List<(long, float[])>(_nodes.Count - _deletedCount); - foreach (Node node in _nodes) + for (int slot = 0; slot < _nodes.Count; slot++) { + Node node = _nodes[slot]; if (!node.Deleted) { - items.Add((node.Id, (float[])node.OriginalVector.Clone())); + items.Add((node.Id, _vectors.Vector(slot).ToArray())); } } @@ -175,7 +178,6 @@ public static HnswIndex Build( public void Add(long id, ReadOnlySpan vector) { ValidateVector(vector); - float[] originalVector = vector.ToArray(); _lock.EnterWriteLock(); try { @@ -189,7 +191,7 @@ public void Add(long id, ReadOnlySpan vector) int slot = _deletedSlots.Pop(); if (_nodes[slot].Deleted && slot != _entryPoint) { - ReplaceSlot(slot, id, originalVector, originalVector); + ReplaceSlot(slot, id, vector); return; } } @@ -197,8 +199,8 @@ public void Add(long id, ReadOnlySpan vector) int level = RandomLevel(); int newIndex = _nodes.Count; EnsureStoredCapacity(newIndex + 1); - PrepareVectorInto(originalVector, StoredSpanMutable(newIndex)); - var node = new Node(id, originalVector, level); + PrepareVectorInto(vector, StoredSpanMutable(newIndex)); + var node = new Node(id, level); _nodes.Add(node); _ids.Add(id, newIndex); @@ -307,7 +309,10 @@ public bool Contains(long id) } } - /// Gets a copy of the original vector for a live id. Returns false for unknown or deleted ids. + /// + /// Gets a copy of the stored vector for a live id. For this + /// is the unit-normalized vector. Returns false for unknown or deleted ids. + /// public bool TryGetVector(long id, out float[] vector) { _lock.EnterReadLock(); @@ -315,7 +320,7 @@ public bool TryGetVector(long id, out float[] vector) { if (_ids.TryGetValue(id, out int slot) && !_nodes[slot].Deleted) { - vector = (float[])_nodes[slot].OriginalVector.Clone(); + vector = _vectors.Vector(slot).ToArray(); return true; } @@ -369,7 +374,7 @@ private void LinkNode(int newIndex, int level, Scratch s) } } - private void ReplaceSlot(int slot, long id, ReadOnlySpan vector, float[] originalVector) + private void ReplaceSlot(int slot, long id, ReadOnlySpan vector) { Node node = _nodes[slot]; for (int layer = 0; layer < node.Links.Length; layer++) @@ -384,7 +389,6 @@ private void ReplaceSlot(int slot, long id, ReadOnlySpan vector, float[] _ids.Remove(node.Id); node.Id = id; - node.OriginalVector = originalVector; node.Deleted = false; _deletedCount--; _ids.Add(id, slot); @@ -500,7 +504,6 @@ public void Save(Stream stream) // targets this is byte-identical to a per-element BinaryWriter.Write(float) loop but // avoids millions of scalar writes on large indexes. writer.Write(MemoryMarshal.AsBytes(StoredSpan(n))); - writer.Write(MemoryMarshal.AsBytes(node.OriginalVector.AsSpan())); writer.Write(node.Links.Length); for (int layer = 0; layer < node.Links.Length; layer++) @@ -546,6 +549,10 @@ public static HnswIndex Load(Stream stream) var index = new HnswIndex(dimension, metric, m, efConstruction, ef, entryPoint, maxLevel, allowReplaceDeleted); index._vectors.EnsureCapacity(count); + // Formats 2 and 3 stored a second copy of the pre-normalized vector per node; it is read and + // discarded since only the (normalized) stored vector is retained now. + float[] discard = version is 2 or 3 ? new float[dimension] : Array.Empty(); + for (int i = 0; i < count; i++) { long id = reader.ReadInt64(); @@ -553,17 +560,12 @@ public static HnswIndex Load(Stream stream) bool deleted = version >= 3 && reader.ReadBoolean(); Span stored = index._vectors.VectorMutable(i); ReadExactInto(reader, MemoryMarshal.AsBytes(stored)); - var originalVector = new float[dimension]; - if (version >= 2) + if (discard.Length > 0) { - ReadExactInto(reader, MemoryMarshal.AsBytes(originalVector.AsSpan())); - } - else - { - stored.CopyTo(originalVector); + ReadExactInto(reader, MemoryMarshal.AsBytes(discard.AsSpan())); } - var node = new Node(id, originalVector, level) { Deleted = deleted }; + var node = new Node(id, level) { Deleted = deleted }; int layerCount = reader.ReadInt32(); if (layerCount != level + 1) { @@ -912,10 +914,9 @@ private Span Slot(int slot) private sealed class Node { - public Node(long id, float[] originalVector, int level) + public Node(long id, int level) { Id = id; - OriginalVector = originalVector; Level = level; Links = new List[level + 1]; for (int i = 0; i < Links.Length; i++) @@ -926,8 +927,6 @@ public Node(long id, float[] originalVector, int level) public long Id { get; set; } - public float[] OriginalVector { get; set; } - public int Level { get; } public bool Deleted { get; set; } diff --git a/tests/Hnsw.Net.Tests/HnswIndexTests.cs b/tests/Hnsw.Net.Tests/HnswIndexTests.cs index aa6b3d7..10ada17 100644 --- a/tests/Hnsw.Net.Tests/HnswIndexTests.cs +++ b/tests/Hnsw.Net.Tests/HnswIndexTests.cs @@ -1,3 +1,4 @@ +using System.Numerics.Tensors; using HnswNet; using Xunit; @@ -98,7 +99,7 @@ public void SaveLoadRoundTripProducesIdenticalResults() } [Fact] - public void ExportItemsRebuildsPortableIndexWithOriginalVectors() + public void ExportItemsRebuildsPortableIndexWithStoredVectors() { const int count = 400; const int dimension = 32; @@ -114,17 +115,19 @@ public void ExportItemsRebuildsPortableIndexWithOriginalVectors() index.Add(50_000 + i, vectors[i]); } + // Cosine stores unit-normalized vectors, so that is what is exported. + float[] expected0 = Normalize(vectors[0]); (long Id, float[] Vector)[] exported = index.ExportItems().ToArray(); Assert.Equal(count, exported.Length); - Assert.Equal(vectors[0], exported[0].Vector); + Assert.Equal(expected0, exported[0].Vector); exported[0].Vector[0] = 12345; - Assert.Equal(vectors[0][0], index.ExportItems().First().Vector[0]); + Assert.Equal(expected0[0], index.ExportItems().First().Vector[0]); using var stream = new MemoryStream(); index.Save(stream); stream.Position = 0; HnswIndex loaded = HnswIndex.Load(stream); - Assert.Equal(vectors[0], loaded.ExportItems().First().Vector); + Assert.Equal(expected0, loaded.ExportItems().First().Vector); HnswIndex rebuilt = HnswIndex.Build( dimension, @@ -142,6 +145,45 @@ public void ExportItemsRebuildsPortableIndexWithOriginalVectors() } } + [Fact] + public void LoadsLegacyV3FormatAndDiscardsDuplicateVector() + { + // Hand-craft a version-3 stream: a single DotProduct node (stored == original) so the loader + // must read and discard the second vector copy that older formats persisted. + const uint magic = 0x31575348; + float[] vector = [1f, 2f, 3f]; + using var stream = new MemoryStream(); + using (var writer = new BinaryWriter(stream, System.Text.Encoding.UTF8, leaveOpen: true)) + { + writer.Write(magic); + writer.Write(3); // version + writer.Write(vector.Length); // dimension + writer.Write((int)DistanceMetric.DotProduct); // metric + writer.Write(2); // m + writer.Write(10); // efConstruction + writer.Write(10); // ef + writer.Write(0); // entryPoint + writer.Write(0); // maxLevel + writer.Write(1); // count + writer.Write(false); // allowReplaceDeleted + + writer.Write(7L); // id + writer.Write(0); // level + writer.Write(false); // deleted + foreach (float f in vector) writer.Write(f); // stored vector + foreach (float f in vector) writer.Write(f); // original vector (discarded on load) + writer.Write(1); // layer count + writer.Write(0); // layer 0 link count + } + + stream.Position = 0; + HnswIndex loaded = HnswIndex.Load(stream); + Assert.Equal(1, loaded.Count); + Assert.True(loaded.TryGetVector(7, out float[] stored)); + Assert.Equal(vector, stored); + Assert.Equal(7, loaded.Search(vector, 1)[0].Id); + } + [Fact] public void HandlesEdgeCases() { @@ -174,6 +216,18 @@ private static float[][] RandomVectors(int count, int dimension, int seed) return vectors; } + private static float[] Normalize(float[] vector) + { + var copy = (float[])vector.Clone(); + float norm = MathF.Sqrt(TensorPrimitives.Dot(copy, copy)); + if (norm > 0) + { + TensorPrimitives.Divide(copy, norm, copy); + } + + return copy; + } + private static IReadOnlyList<(long Id, float Distance)> BruteForce( float[][] vectors, float[] query, From ddb7dc64e26eb0591be35de7884e500a115af744 Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 07:24:02 -0700 Subject: [PATCH 5/9] Add memory-mapped read path (LoadMapped) Restructure the v4 on-disk format into sections: a contiguous, fixed-stride vector block followed by a separate graph section. This lets LoadMapped memory- map the vector block (the bulk of the data) instead of reading it into the managed heap, while the graph is still loaded into RAM so search stays fast. The vector store is now an abstraction with a chunked heap implementation (build and default Load) and a MemoryMappedFile-backed implementation addressed with long offsets, so the vector section can exceed the .NET array length limit while each per-slot span stays within it. A mapped index is read-only (Add throws) and owns the mapping, so HnswIndex is now IDisposable. Formats v1-v3 still load via the previous interleaved reader. Adds LoadMapped round-trip/read-only/dispose tests (48 total) and README docs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- README.md | 16 ++ src/Hnsw.Net/Hnsw.Net.csproj | 1 + src/Hnsw.Net/HnswIndex.cs | 273 ++++++++++++++++++++----- tests/Hnsw.Net.Tests/HnswIndexTests.cs | 51 +++++ 4 files changed, 294 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index b5127de..6138007 100644 --- a/README.md +++ b/README.md @@ -117,6 +117,22 @@ unit-normalized vectors used for search (not the original input magnitudes); rebuilding from them re-normalizes and reproduces the same index. Other metrics store vectors unchanged, so the export matches the input. +## Memory-mapped loading + +For large indexes or cold-start-heavy scenarios, `LoadMapped` memory-maps the +vector section of a saved file instead of reading it into the managed heap. The +graph is still loaded into memory, so search stays fast while the bulk of the +data (the vectors) is paged in on demand by the OS: + +```csharp +using HnswIndex index = HnswIndex.LoadMapped("index.hnsw"); +IReadOnlyList<(long Id, float Distance)> hits = index.Search(query, 10); +``` + +The returned index owns the mapping and is read-only — it must be disposed, and +`Add` throws. To modify an index, load it with `Load` instead. `LoadMapped` +requires the current on-disk format; re-save older indexes first. + ## hnswlib parity validation `tests/Hnsw.Net.Tests/gen_parity.py` generates the committed diff --git a/src/Hnsw.Net/Hnsw.Net.csproj b/src/Hnsw.Net/Hnsw.Net.csproj index 9dd4fb6..6e0dbc3 100644 --- a/src/Hnsw.Net/Hnsw.Net.csproj +++ b/src/Hnsw.Net/Hnsw.Net.csproj @@ -6,6 +6,7 @@ enable true latest + true HnswNet Hnsw.Net true diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index 689848c..bb56662 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -1,4 +1,5 @@ using System.Collections.Concurrent; +using System.IO.MemoryMappedFiles; using System.Numerics; using System.Numerics.Tensors; using System.Runtime.CompilerServices; @@ -11,7 +12,7 @@ namespace HnswNet; /// An HNSW approximate-nearest-neighbor index over vectors. /// Builds and modifications are serialized; searches are thread-safe and may run concurrently. /// -public sealed class HnswIndex +public sealed class HnswIndex : IDisposable { private const int FormatVersion = 4; private const uint Magic = 0x31575348; // HSW1, little-endian. @@ -62,7 +63,7 @@ public HnswIndex(int dimension, DistanceMetric metric, int m = 16, int efConstru _levelMultiplier = 1.0 / Math.Log(m); _allowReplaceDeleted = allowReplaceDeleted; _candidateComparison = CompareCandidates; - _vectors = new VectorBlock(dimension); + _vectors = new HeapVectorBlock(dimension); } private HnswIndex(int dimension, DistanceMetric metric, int m, int efConstruction, int ef, int entryPoint, int maxLevel, bool allowReplaceDeleted) @@ -78,7 +79,7 @@ private HnswIndex(int dimension, DistanceMetric metric, int m, int efConstructio _levelMultiplier = 1.0 / Math.Log(m); _allowReplaceDeleted = allowReplaceDeleted; _candidateComparison = CompareCandidates; - _vectors = new VectorBlock(dimension); + _vectors = new HeapVectorBlock(dimension); } /// Gets the vector dimension. @@ -178,6 +179,11 @@ public static HnswIndex Build( public void Add(long id, ReadOnlySpan vector) { ValidateVector(vector); + if (_vectors.IsReadOnly) + { + throw new InvalidOperationException("This index was loaded with memory-mapped vectors and is read-only. Load it without mapping to modify it."); + } + _lock.EnterWriteLock(); try { @@ -494,17 +500,22 @@ public void Save(Stream stream) writer.Write(_nodes.Count); writer.Write(_allowReplaceDeleted); + // Vector section: one normalized vector per slot, contiguous and fixed-stride, so the + // read path can memory-map it and address slot s at base + s * Dimension * sizeof(float). + // Written as raw little-endian float blocks (byte-identical to a per-float loop on the + // LE platforms .NET targets, but without millions of scalar writes). + for (int n = 0; n < _nodes.Count; n++) + { + writer.Write(MemoryMarshal.AsBytes(StoredSpan(n))); + } + + // Graph section: per-node metadata and link lists, kept separate from the vectors. for (int n = 0; n < _nodes.Count; n++) { Node node = _nodes[n]; writer.Write(node.Id); writer.Write(node.Level); writer.Write(node.Deleted); - // Vectors are written as raw little-endian float blocks; on the LE platforms .NET - // targets this is byte-identical to a per-element BinaryWriter.Write(float) loop but - // avoids millions of scalar writes on large indexes. - writer.Write(MemoryMarshal.AsBytes(StoredSpan(n))); - writer.Write(node.Links.Length); for (int layer = 0; layer < node.Links.Length; layer++) { @@ -522,15 +533,98 @@ public void Save(Stream stream) } } - /// Loads an index saved by . + /// Loads an index saved by into managed memory. public static HnswIndex Load(Stream stream) { ArgumentNullException.ThrowIfNull(stream); using var reader = new BinaryReader(stream, System.Text.Encoding.UTF8, leaveOpen: true); + Header header = ReadHeader(reader); + var index = new HnswIndex(header.Dimension, header.Metric, header.M, header.EfConstruction, header.Ef, header.EntryPoint, header.MaxLevel, header.AllowReplaceDeleted); + index._vectors.EnsureCapacity(header.Count); + + if (header.Version <= 3) + { + ReadInterleavedBody(reader, index, header); + } + else + { + // Vector section first (one vector per slot), then the graph section. + for (int i = 0; i < header.Count; i++) + { + ReadExactInto(reader, MemoryMarshal.AsBytes(index._vectors.VectorMutable(i))); + } + + for (int i = 0; i < header.Count; i++) + { + long id = reader.ReadInt64(); + int level = reader.ReadInt32(); + bool deleted = reader.ReadBoolean(); + ReadNodeLinks(reader, index, i, id, level, deleted, header); + } + } + + return index; + } + + /// + /// Loads an index from a file, memory-mapping the vector section instead of reading it into the + /// managed heap. The graph is still loaded into memory. The returned index is read-only (calls to + /// throw) and owns the mapping, so it must be disposed. + /// Only the current on-disk format is supported; re-save older indexes first. + /// + public static HnswIndex LoadMapped(string path) + { + ArgumentNullException.ThrowIfNull(path); + Header header; + long vectorOffset; + HnswIndex index; + using (var stream = new FileStream(path, FileMode.Open, FileAccess.Read, FileShare.Read)) + using (var reader = new BinaryReader(stream, System.Text.Encoding.UTF8, leaveOpen: true)) + { + header = ReadHeader(reader); + if (header.Version != FormatVersion) + { + throw new InvalidDataException("Memory-mapped loading requires the current index format; re-save the index."); + } + + vectorOffset = stream.Position; + long vectorBytes = (long)header.Count * header.Dimension * sizeof(float); + index = new HnswIndex(header.Dimension, header.Metric, header.M, header.EfConstruction, header.Ef, header.EntryPoint, header.MaxLevel, header.AllowReplaceDeleted); + + stream.Seek(vectorOffset + vectorBytes, SeekOrigin.Begin); + for (int i = 0; i < header.Count; i++) + { + long id = reader.ReadInt64(); + int level = reader.ReadInt32(); + bool deleted = reader.ReadBoolean(); + ReadNodeLinks(reader, index, i, id, level, deleted, header); + } + } + + MemoryMappedFile? file = null; + MemoryMappedViewAccessor? view = null; + try + { + file = MemoryMappedFile.CreateFromFile(path, FileMode.Open, mapName: null, capacity: 0, MemoryMappedFileAccess.Read); + view = file.CreateViewAccessor(0, 0, MemoryMappedFileAccess.Read); + index._vectors = new MappedVectorBlock(file, view, vectorOffset, header.Dimension); + return index; + } + catch + { + view?.Dispose(); + file?.Dispose(); + throw; + } + } + + private static Header ReadHeader(BinaryReader reader) + { if (reader.ReadUInt32() != Magic) { throw new InvalidDataException("The stream is not an Hnsw.Net index."); } + int version = reader.ReadInt32(); if (version is < 1 or > FormatVersion) { @@ -546,54 +640,77 @@ public static HnswIndex Load(Stream stream) int maxLevel = reader.ReadInt32(); int count = reader.ReadInt32(); bool allowReplaceDeleted = version >= 3 && reader.ReadBoolean(); - var index = new HnswIndex(dimension, metric, m, efConstruction, ef, entryPoint, maxLevel, allowReplaceDeleted); - index._vectors.EnsureCapacity(count); + return new Header(version, dimension, metric, m, efConstruction, ef, entryPoint, maxLevel, count, allowReplaceDeleted); + } + // Reads the pre-v4 layout where each node's vector(s) are interleaved with its graph data. + private static void ReadInterleavedBody(BinaryReader reader, HnswIndex index, Header header) + { // Formats 2 and 3 stored a second copy of the pre-normalized vector per node; it is read and // discarded since only the (normalized) stored vector is retained now. - float[] discard = version is 2 or 3 ? new float[dimension] : Array.Empty(); - - for (int i = 0; i < count; i++) + float[] discard = header.Version is 2 or 3 ? new float[header.Dimension] : Array.Empty(); + for (int i = 0; i < header.Count; i++) { long id = reader.ReadInt64(); int level = reader.ReadInt32(); - bool deleted = version >= 3 && reader.ReadBoolean(); - Span stored = index._vectors.VectorMutable(i); - ReadExactInto(reader, MemoryMarshal.AsBytes(stored)); + bool deleted = header.Version >= 3 && reader.ReadBoolean(); + ReadExactInto(reader, MemoryMarshal.AsBytes(index._vectors.VectorMutable(i))); if (discard.Length > 0) { ReadExactInto(reader, MemoryMarshal.AsBytes(discard.AsSpan())); } - var node = new Node(id, level) { Deleted = deleted }; - int layerCount = reader.ReadInt32(); - if (layerCount != level + 1) - { - throw new InvalidDataException("Invalid layer count in Hnsw.Net index."); - } + ReadNodeLinks(reader, index, i, id, level, deleted, header); + } + } - for (int layer = 0; layer < layerCount; layer++) + private static void ReadNodeLinks(BinaryReader reader, HnswIndex index, int i, long id, int level, bool deleted, Header header) + { + var node = new Node(id, level) { Deleted = deleted }; + int layerCount = reader.ReadInt32(); + if (layerCount != level + 1) + { + throw new InvalidDataException("Invalid layer count in Hnsw.Net index."); + } + + for (int layer = 0; layer < layerCount; layer++) + { + int linkCount = reader.ReadInt32(); + for (int link = 0; link < linkCount; link++) { - int linkCount = reader.ReadInt32(); - for (int link = 0; link < linkCount; link++) - { - node.Links[layer].Add(reader.ReadInt32()); - } + node.Links[layer].Add(reader.ReadInt32()); } + } - index._ids.Add(id, i); - index._nodes.Add(node); - if (deleted) + index._ids.Add(id, i); + index._nodes.Add(node); + if (deleted) + { + index._deletedCount++; + if (header.AllowReplaceDeleted && i != header.EntryPoint) { - index._deletedCount++; - if (allowReplaceDeleted && i != entryPoint) - { - index._deletedSlots.Push(i); - } + index._deletedSlots.Push(i); } } + } - return index; + private readonly record struct Header( + int Version, + int Dimension, + DistanceMetric Metric, + int M, + int EfConstruction, + int Ef, + int EntryPoint, + int MaxLevel, + int Count, + bool AllowReplaceDeleted); + + /// Releases the memory mapping held by an index loaded via . + public void Dispose() + { + (_vectors as IDisposable)?.Dispose(); + _lock.Dispose(); } private void ValidateVector(ReadOnlySpan vector) @@ -869,18 +986,31 @@ private static float DotProduct(ReadOnlySpan a, ReadOnlySpan b) return sum; } - // Chunked, slot-addressed storage for the normalized search vectors. Splitting the data into - // bounded chunks keeps each backing array well under the .NET array length limit so the index - // scales to very large repos, and gives every slot a contiguous span ready for a future - // memory-mapped backing. - private sealed class VectorBlock + // Slot-addressed storage for the normalized search vectors. The read path goes through this + // abstraction so the backing store can be either the managed heap (build and default load) or a + // memory-mapped file (LoadMapped). + private abstract class VectorBlock + { + public abstract bool IsReadOnly { get; } + + public abstract void EnsureCapacity(int slotCount); + + public abstract ReadOnlySpan Vector(int slot); + + public abstract Span VectorMutable(int slot); + } + + // Chunked heap storage. Splitting the data into bounded chunks keeps each backing array well + // under the .NET array length limit so the index scales to very large repos, and gives every + // slot a contiguous span. + private sealed class HeapVectorBlock : VectorBlock { private readonly int _dimension; private readonly int _vectorsPerChunk; private readonly List _chunks = new(); private int _capacity; - public VectorBlock(int dimension) + public HeapVectorBlock(int dimension) { _dimension = dimension; @@ -891,7 +1021,9 @@ public VectorBlock(int dimension) _vectorsPerChunk = (int)Math.Min(perChunk, int.MaxValue / dimension); } - public void EnsureCapacity(int slotCount) + public override bool IsReadOnly => false; + + public override void EnsureCapacity(int slotCount) { while (_capacity < slotCount) { @@ -900,9 +1032,9 @@ public void EnsureCapacity(int slotCount) } } - public ReadOnlySpan Vector(int slot) => Slot(slot); + public override ReadOnlySpan Vector(int slot) => Slot(slot); - public Span VectorMutable(int slot) => Slot(slot); + public override Span VectorMutable(int slot) => Slot(slot); private Span Slot(int slot) { @@ -912,6 +1044,53 @@ private Span Slot(int slot) } } + // Read-only storage backed by a memory-mapped file. Each slot is served as a span straight off + // the mapped pages; the file is addressed with long offsets so the vector section can exceed the + // .NET array length limit while each per-slot span stays within it. + private sealed unsafe class MappedVectorBlock : VectorBlock, IDisposable + { + private readonly MemoryMappedFile _file; + private readonly MemoryMappedViewAccessor _view; + private readonly long _vectorOffset; + private readonly int _dimension; + private byte* _base; + + public MappedVectorBlock(MemoryMappedFile file, MemoryMappedViewAccessor view, long vectorOffset, int dimension) + { + _file = file; + _view = view; + _vectorOffset = vectorOffset; + _dimension = dimension; + byte* pointer = null; + view.SafeMemoryMappedViewHandle.AcquirePointer(ref pointer); + _base = pointer + view.PointerOffset; + } + + public override bool IsReadOnly => true; + + public override void EnsureCapacity(int slotCount) + { + } + + public override ReadOnlySpan Vector(int slot) + => new(_base + _vectorOffset + (long)slot * _dimension * sizeof(float), _dimension); + + public override Span VectorMutable(int slot) + => throw new NotSupportedException("Memory-mapped vectors are read-only."); + + public void Dispose() + { + if (_base != null) + { + _view.SafeMemoryMappedViewHandle.ReleasePointer(); + _base = null; + } + + _view.Dispose(); + _file.Dispose(); + } + } + private sealed class Node { public Node(long id, int level) diff --git a/tests/Hnsw.Net.Tests/HnswIndexTests.cs b/tests/Hnsw.Net.Tests/HnswIndexTests.cs index 10ada17..2ef8c30 100644 --- a/tests/Hnsw.Net.Tests/HnswIndexTests.cs +++ b/tests/Hnsw.Net.Tests/HnswIndexTests.cs @@ -69,6 +69,57 @@ public void TinyDataFindsExactNearestNeighbor(DistanceMetric metric) Assert.Equal(expected, index.Search(query, 1)[0].Id); } + [Fact] + public void LoadMappedProducesIdenticalResultsAndIsReadOnly() + { + const int count = 300; + const int dimension = 32; + float[][] vectors = RandomVectors(count, dimension, seed: 222); + float[][] queries = RandomVectors(8, dimension, seed: 333); + var index = new HnswIndex(dimension, DistanceMetric.Cosine, m: 16, efConstruction: 120, seed: 44) + { + Ef = 80, + }; + + for (int i = 0; i < vectors.Length; i++) + { + index.Add(10_000 + i, vectors[i]); + } + + string path = Path.GetTempFileName(); + try + { + using (var file = new FileStream(path, FileMode.Create, FileAccess.Write)) + { + index.Save(file); + } + + using (HnswIndex mapped = HnswIndex.LoadMapped(path)) + { + Assert.Equal(index.Count, mapped.Count); + foreach (float[] query in queries) + { + Assert.Equal(index.Search(query, 12), mapped.Search(query, 12)); + } + + Assert.True(mapped.TryGetVector(10_000, out float[] stored)); + Assert.Equal(Normalize(vectors[0]), stored); + Assert.Throws(() => mapped.Add(99_999, vectors[0])); + } + + // Dispose must release the mapping so the file is no longer locked. + File.Delete(path); + Assert.False(File.Exists(path)); + } + finally + { + if (File.Exists(path)) + { + File.Delete(path); + } + } + } + [Fact] public void SaveLoadRoundTripProducesIdenticalResults() { From 42481633d465c9e4e52e5f379e59e6dc0852ab8a Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 09:28:12 -0700 Subject: [PATCH 6/9] Add memory-mapped, lazy snapshot load to HnswCollection LoadCore reads every record payload off disk and deserializes all N records into objects up front, so a large collection pays its full size in load time and managed allocation before the first query. This adds a path-based load that memory-maps the file instead. - HnswIndex.LoadMapped(path, offset): map the vector section of an index that begins at a byte offset within a larger container file. - HnswCollection.Load(string [, offset], [context]): map the snapshot, read framing eagerly (keys materialized, record payloads located but not read), defer TRecord deserialization to first access from the mapped region, and map the embedded index's vectors. Records stay off the managed heap; the index owns the vectors, so the collection keeps none. The result is read-only. - Entry now supports lazy materialization (thread-safe, cached); the mapping lifetime is owned by HnswCollectionData and released on reload, collection deletion, or store disposal. - Framing corruption is still detected at load; record-payload JSON validity is validated on access for the mapped path (documented). Stream-based Load and its corruption tests are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Hnsw.Net/HnswIndex.cs | 14 +- src/Hnsw.Net/VectorData/HnswCollection.cs | 277 +++++++++++++++++- src/Hnsw.Net/VectorData/HnswCollectionData.cs | 115 +++++++- src/Hnsw.Net/VectorData/HnswVectorStore.cs | 20 +- tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs | 150 ++++++++++ 5 files changed, 566 insertions(+), 10 deletions(-) create mode 100644 tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index bb56662..ee3c056 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -572,15 +572,25 @@ public static HnswIndex Load(Stream stream) /// throw) and owns the mapping, so it must be disposed. /// Only the current on-disk format is supported; re-save older indexes first. /// - public static HnswIndex LoadMapped(string path) + public static HnswIndex LoadMapped(string path) => LoadMapped(path, 0); + + /// + /// Loads an index that begins at within a larger file, memory-mapping + /// its vector section. Use this when an index is embedded in a container format (for example a + /// collection snapshot). Behaves like otherwise: the result is + /// read-only, owns the mapping, and must be disposed. + /// + public static HnswIndex LoadMapped(string path, long baseOffset) { ArgumentNullException.ThrowIfNull(path); + ArgumentOutOfRangeException.ThrowIfNegative(baseOffset); Header header; long vectorOffset; HnswIndex index; using (var stream = new FileStream(path, FileMode.Open, FileAccess.Read, FileShare.Read)) using (var reader = new BinaryReader(stream, System.Text.Encoding.UTF8, leaveOpen: true)) { + stream.Seek(baseOffset, SeekOrigin.Begin); header = ReadHeader(reader); if (header.Version != FormatVersion) { @@ -706,7 +716,7 @@ private readonly record struct Header( int Count, bool AllowReplaceDeleted); - /// Releases the memory mapping held by an index loaded via . + /// Releases the memory mapping held by an index loaded via . public void Dispose() { (_vectors as IDisposable)?.Dispose(); diff --git a/src/Hnsw.Net/VectorData/HnswCollection.cs b/src/Hnsw.Net/VectorData/HnswCollection.cs index da46b1e..d6e394d 100644 --- a/src/Hnsw.Net/VectorData/HnswCollection.cs +++ b/src/Hnsw.Net/VectorData/HnswCollection.cs @@ -104,7 +104,11 @@ public override Task EnsureCollectionExistsAsync(CancellationToken cancellationT /// public override Task EnsureCollectionDeletedAsync(CancellationToken cancellationToken = default) { - _collections.TryRemove(Name, out _); + if (_collections.TryRemove(Name, out HnswCollectionData? data)) + { + data.Dispose(); + } + _collectionTypes.TryRemove(Name, out _); return Task.CompletedTask; } @@ -228,7 +232,7 @@ public override async Task UpsertAsync(IEnumerable records, Cancellatio long id = data.NextId++; index.Add(id, vector.Span); - data.Records[key] = new HnswCollectionData.Entry { Record = record, Id = id, Vector = vector }; + data.Records[key] = new HnswCollectionData.Entry(record, id, vector); data.IdToKey[id] = key; } } @@ -410,6 +414,46 @@ public void Load(Stream stream, JsonSerializerContext context) json => DeserializeWithTypeInfo(json, recordInfo)); } + /// + /// Provider-specific persistence that memory-maps a snapshot file instead of reading it into the managed + /// heap. Record payloads are deserialized lazily on first access and the backing vectors are mapped, so a + /// large collection loads with minimal time and allocation. The resulting collection is read-only (the + /// mapping is released when the store or collection data is disposed) and its records' JSON is validated + /// on access rather than at load time. + /// + /// The snapshot file path. + /// + /// This overload deserializes records by reflection. For trimmed or NativeAOT applications, use + /// with a source-generated context. + /// + [RequiresUnreferencedCode("Snapshot persistence serializes records by reflection and is incompatible with trimming. Use the JsonSerializerContext overload for trimming/NativeAOT.")] + [RequiresDynamicCode("Snapshot persistence serializes records by reflection and is incompatible with NativeAOT. Use the JsonSerializerContext overload for trimming/NativeAOT.")] + public void Load(string path) + => LoadMappedCore(path, 0, DeserializeByReflectionSpan, DeserializeByReflectionSpan); + + /// + /// Memory-mapping, AOT- and trimming-safe counterpart of , deserializing + /// records with the supplied source-generated . + /// + public void Load(string path, JsonSerializerContext context) + => Load(path, 0, context); + + /// + /// Memory-maps a snapshot that begins at within , for + /// use when a collection snapshot is embedded in a larger container file. + /// + public void Load(string path, long offset, JsonSerializerContext context) + { + ArgumentNullException.ThrowIfNull(context); + JsonTypeInfo keyInfo = ResolveTypeInfo(context, typeof(TKey)); + JsonTypeInfo recordInfo = ResolveTypeInfo(context, typeof(TRecord)); + LoadMappedCore( + path, + offset, + json => DeserializeWithTypeInfo(json, keyInfo), + json => DeserializeWithTypeInfo(json, recordInfo)); + } + // Shared binary framing for both the reflection and source-generated overloads. Records and keys are // serialized individually so the caller's context only needs metadata for TKey and TRecord. private void SaveCore(Stream stream, Func serializeKey, Func serializeRecord) @@ -620,7 +664,7 @@ or ArgumentException or InvalidOperationException or FormatException $"Corrupt Hnsw.Net collection snapshot: the index does not contain a vector for record id {id}."); } - if (!newRecords.TryAdd(key!, new HnswCollectionData.Entry { Record = record!, Id = id, Vector = vector })) + if (!newRecords.TryAdd(key!, new HnswCollectionData.Entry(record!, id, vector))) { throw new InvalidDataException( $"Corrupt Hnsw.Net collection snapshot: duplicate record key '{key}'."); @@ -633,6 +677,206 @@ or ArgumentException or InvalidOperationException or FormatException } } + InstallLoadedState(dimension, nextId, index, mappedRecords: null, newRecords, newIdToKey); + } + + private void LoadMappedCore(string path, long baseOffset, SpanReader deserializeKey, SpanReader deserializeRecord) + { + ArgumentNullException.ThrowIfNull(path); + ArgumentOutOfRangeException.ThrowIfNegative(baseOffset); + + int dimension; + DistanceMetric metric; + long nextId; + bool hasIndex; + long indexOffset; + // Framing only: key materialized eagerly (small), record payload located but not read yet. + var framed = new List<(TKey Key, long Id, long RecordOffset, int RecordLength)>(); + + using (var stream = new FileStream(path, FileMode.Open, FileAccess.Read, FileShare.Read)) + using (var reader = new BinaryReader(stream, Encoding.UTF8, leaveOpen: true)) + { + long fileLength = stream.Length; + stream.Seek(baseOffset, SeekOrigin.Begin); + try + { + if (reader.ReadUInt32() != SnapshotMagic) + { + throw new InvalidDataException("The file is not an Hnsw.Net collection snapshot."); + } + + if (reader.ReadInt32() != SnapshotVersion) + { + throw new InvalidDataException("Unsupported Hnsw.Net collection snapshot version."); + } + + dimension = reader.ReadInt32(); + metric = (DistanceMetric)reader.ReadInt32(); + nextId = reader.ReadInt64(); + + if (dimension < 0) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: negative vector dimension."); + } + + if (nextId < 0) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: negative next id."); + } + + int configured = _model.VectorProperty.Dimensions; + if (configured > 0 && dimension > 0 && dimension != configured) + { + throw new InvalidDataException( + $"The snapshot's vector dimension ({dimension}) does not match the collection's configured dimension ({configured})."); + } + + if (metric != _metric) + { + throw new InvalidDataException( + $"The snapshot's distance metric ({metric}) does not match the collection's configured metric ({_metric})."); + } + + int count = reader.ReadInt32(); + if (count < 0) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: negative record count."); + } + + for (int i = 0; i < count; i++) + { + long id = reader.ReadInt64(); + TKey key = deserializeKey(ReadExact(reader, reader.ReadInt32())); + int recordLength = reader.ReadInt32(); + if (recordLength < 0) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: negative payload length."); + } + + if (recordLength > MaxPayloadLength) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: payload length exceeds the maximum supported size."); + } + + long recordOffset = stream.Position; + if (recordLength > fileLength - recordOffset) + { + throw new InvalidDataException("Truncated Hnsw.Net collection snapshot."); + } + + stream.Seek(recordLength, SeekOrigin.Current); + framed.Add((key, id, recordOffset, recordLength)); + } + + hasIndex = reader.ReadBoolean(); + indexOffset = stream.Position; + } + catch (EndOfStreamException ex) + { + throw new InvalidDataException("Truncated Hnsw.Net collection snapshot.", ex); + } + catch (JsonException ex) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: a key payload is not valid JSON.", ex); + } + catch (NotSupportedException ex) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: a key payload could not be deserialized.", ex); + } + } + + if (framed.Count > 0 && !hasIndex) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: a non-empty collection has no index."); + } + + if (hasIndex && dimension == 0) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: an index is present but the vector dimension is zero."); + } + + HnswIndex? index = null; + MappedRecordFile? mapped = null; + try + { + try + { + index = hasIndex ? HnswIndex.LoadMapped(path, indexOffset) : null; + } + catch (Exception ex) when ( + ex is EndOfStreamException or IOException or OverflowException + or ArgumentException or InvalidOperationException or FormatException + && ex is not InvalidDataException) + { + throw new InvalidDataException("Corrupt Hnsw.Net collection snapshot: the index is invalid.", ex); + } + + if (index is not null && (index.Dimension != dimension || index.Metric != metric)) + { + throw new InvalidDataException( + $"Corrupt Hnsw.Net collection snapshot: the index header (dimension {index.Dimension}, metric {index.Metric}) " + + $"does not match the snapshot header (dimension {dimension}, metric {metric})."); + } + + mapped = new MappedRecordFile(path); + + var newRecords = new Dictionary(); + var newIdToKey = new Dictionary(); + foreach ((TKey key, long id, long recordOffset, int recordLength) in framed) + { + if (id < 0) + { + throw new InvalidDataException($"Corrupt Hnsw.Net collection snapshot: negative record id {id}."); + } + + if (id >= nextId) + { + throw new InvalidDataException( + $"Corrupt Hnsw.Net collection snapshot: record id {id} is not less than the next id ({nextId})."); + } + + if (index is not null && !index.Contains(id)) + { + throw new InvalidDataException( + $"Corrupt Hnsw.Net collection snapshot: the index does not contain a vector for record id {id}."); + } + + MappedRecordFile mappedRecords = mapped; + long offset = recordOffset; + int length = recordLength; + object Factory() => deserializeRecord(mappedRecords.Slice(offset, length))!; + + // The mapped index owns the vectors; the collection keeps none on the managed heap. + if (!newRecords.TryAdd(key!, new HnswCollectionData.Entry(Factory, id, ReadOnlyMemory.Empty))) + { + throw new InvalidDataException($"Corrupt Hnsw.Net collection snapshot: duplicate record key '{key}'."); + } + + if (!newIdToKey.TryAdd(id, key)) + { + throw new InvalidDataException($"Corrupt Hnsw.Net collection snapshot: duplicate record id {id}."); + } + } + + InstallLoadedState(dimension, nextId, index, mapped, newRecords, newIdToKey); + index = null; + mapped = null; + } + finally + { + mapped?.Dispose(); + index?.Dispose(); + } + } + + private void InstallLoadedState( + int dimension, + long nextId, + HnswIndex? index, + MappedRecordFile? mappedRecords, + Dictionary newRecords, + Dictionary newIdToKey) + { // Establish/validate the key and record types before touching _collections so a concurrent // GetCollection with a different TKey/TRecord can't slip in and observe mixed-type data for this name. (Type Key, Type Record) existingType = _collectionTypes.GetOrAdd(Name, (typeof(TKey), typeof(TRecord))); @@ -647,11 +891,15 @@ or ArgumentException or InvalidOperationException or FormatException HnswCollectionData data = _collections.GetOrAdd(Name, static _ => new HnswCollectionData()); lock (data.Lock) { + HnswIndex? supersededIndex = data.Index; + MappedRecordFile? supersededMapping = data.MappedRecords; + data.Records.Clear(); data.IdToKey.Clear(); data.Dimension = dimension; data.NextId = nextId; data.Index = index; + data.MappedRecords = mappedRecords; foreach ((object key, HnswCollectionData.Entry entry) in newRecords) { data.Records[key] = entry; @@ -661,9 +909,21 @@ or ArgumentException or InvalidOperationException or FormatException { data.IdToKey[id] = key!; } + + if (!ReferenceEquals(supersededIndex, index)) + { + supersededIndex?.Dispose(); + } + + if (!ReferenceEquals(supersededMapping, mappedRecords)) + { + supersededMapping?.Dispose(); + } } } + private delegate T SpanReader(ReadOnlySpan utf8Json); + // An individual key or record payload should never be huge; cap it so a corrupt length prefix on a // non-seekable stream cannot trigger a denial-of-service allocation before truncation is detected. private const int MaxPayloadLength = 128 * 1024 * 1024; @@ -706,6 +966,11 @@ private static T DeserializeWithTypeInfo(byte[] json, JsonTypeInfo typeInfo) ? value : throw new InvalidDataException($"Invalid Hnsw.Net collection snapshot: could not deserialize a '{typeof(T)}'."); + private static T DeserializeWithTypeInfo(ReadOnlySpan json, JsonTypeInfo typeInfo) + => JsonSerializer.Deserialize(json, typeInfo) is T value + ? value + : throw new InvalidDataException($"Invalid Hnsw.Net collection snapshot: could not deserialize a '{typeof(T)}'."); + [RequiresUnreferencedCode("Serializes by reflection.")] [RequiresDynamicCode("Serializes by reflection.")] private static byte[] SerializeByReflection(T value) @@ -717,6 +982,12 @@ private static T DeserializeByReflection(byte[] json) => JsonSerializer.Deserialize(json) ?? throw new InvalidDataException("Invalid Hnsw.Net collection snapshot."); + [RequiresUnreferencedCode("Deserializes by reflection.")] + [RequiresDynamicCode("Deserializes by reflection.")] + private static T DeserializeByReflectionSpan(ReadOnlySpan json) + => JsonSerializer.Deserialize(json) + ?? throw new InvalidDataException("Invalid Hnsw.Net collection snapshot."); + private HnswIndex GetOrCreateIndex(HnswCollectionData data, int dimension) { if (data.Index is null) diff --git a/src/Hnsw.Net/VectorData/HnswCollectionData.cs b/src/Hnsw.Net/VectorData/HnswCollectionData.cs index e157785..fae79cf 100644 --- a/src/Hnsw.Net/VectorData/HnswCollectionData.cs +++ b/src/Hnsw.Net/VectorData/HnswCollectionData.cs @@ -1,3 +1,5 @@ +using System.IO.MemoryMappedFiles; + namespace HnswNet; /// @@ -5,7 +7,7 @@ namespace HnswNet; /// collection name within a store. Holds the HNSW index, the record payloads, and the mapping between the /// caller's keys and the ids used by . /// -internal sealed class HnswCollectionData +internal sealed class HnswCollectionData : IDisposable { public readonly object Lock = new(); @@ -21,12 +23,117 @@ internal sealed class HnswCollectionData public long NextId; + /// + /// When the data was loaded by memory-mapping a snapshot, the mapping that backs the lazily-materialized + /// record payloads. Held for the lifetime of the data and released when it is replaced or disposed. + /// + public MappedRecordFile? MappedRecords; + + public void Dispose() + { + Index?.Dispose(); + MappedRecords?.Dispose(); + } + + /// + /// A record payload that is either materialized eagerly (mutation and stream-load paths) or deserialized + /// on first access from a memory-mapped snapshot region (the mmap load path). + /// public sealed class Entry { - public required object Record { get; init; } + private object? _record; + private Func? _factory; + + public Entry(object record, long id, ReadOnlyMemory vector) + { + _record = record; + Id = id; + Vector = vector; + } + + public Entry(Func recordFactory, long id, ReadOnlyMemory vector) + { + _factory = recordFactory; + Id = id; + Vector = vector; + } - public required long Id { get; init; } + public long Id { get; } + + public ReadOnlyMemory Vector { get; } + + /// The record, materialized on first access for mmap-backed entries. + public object Record + { + get + { + object? record = Volatile.Read(ref _record); + if (record is not null) + { + return record; + } + + lock (this) + { + if (_record is null) + { + _record = _factory!(); + _factory = null; + } + + return _record; + } + } + } + } +} + +/// +/// Owns a read-only memory mapping of a snapshot file and hands out spans over record payloads so they can be +/// deserialized lazily without copying the whole file onto the managed heap. Disposing releases the mapping; +/// spans must not be used afterward. +/// +internal sealed unsafe class MappedRecordFile : IDisposable +{ + private readonly MemoryMappedFile _file; + private readonly MemoryMappedViewAccessor _view; + private readonly byte* _base; + private bool _disposed; + + public MappedRecordFile(string path) + { + _file = MemoryMappedFile.CreateFromFile(path, FileMode.Open, mapName: null, capacity: 0, MemoryMappedFileAccess.Read); + try + { + _view = _file.CreateViewAccessor(0, 0, MemoryMappedFileAccess.Read); + byte* pointer = null; + _view.SafeMemoryMappedViewHandle.AcquirePointer(ref pointer); + _base = pointer + _view.PointerOffset; + } + catch + { + _view?.Dispose(); + _file.Dispose(); + throw; + } + } + + public ReadOnlySpan Slice(long offset, int length) + { + ObjectDisposedException.ThrowIf(_disposed, this); + return new ReadOnlySpan(_base + offset, length); + } + + public void Dispose() + { + if (_disposed) + { + return; + } - public required ReadOnlyMemory Vector { get; init; } + _disposed = true; + _view.SafeMemoryMappedViewHandle.ReleasePointer(); + _view.Dispose(); + _file.Dispose(); } } diff --git a/src/Hnsw.Net/VectorData/HnswVectorStore.cs b/src/Hnsw.Net/VectorData/HnswVectorStore.cs index edd4d78..b80130e 100644 --- a/src/Hnsw.Net/VectorData/HnswVectorStore.cs +++ b/src/Hnsw.Net/VectorData/HnswVectorStore.cs @@ -90,11 +90,29 @@ public override Task CollectionExistsAsync(string name, CancellationToken /// public override Task EnsureCollectionDeletedAsync(string name, CancellationToken cancellationToken = default) { - _collections.TryRemove(name, out _); + if (_collections.TryRemove(name, out HnswCollectionData? data)) + { + data.Dispose(); + } + _collectionTypes.TryRemove(name, out _); return Task.CompletedTask; } + /// + protected override void Dispose(bool disposing) + { + if (disposing) + { + foreach (HnswCollectionData data in _collections.Values) + { + data.Dispose(); + } + } + + base.Dispose(disposing); + } + /// public override object? GetService(Type serviceType, object? serviceKey = null) { diff --git a/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs b/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs new file mode 100644 index 0000000..e08a3bc --- /dev/null +++ b/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs @@ -0,0 +1,150 @@ +using Microsoft.Extensions.VectorData; +using HnswNet; +using Xunit; + +namespace Hnsw.Net.Tests; + +public partial class VectorStoreTests +{ + [Fact] + public async Task LoadMapped_FromPath_RoundTripsRecordsAndSearch() + { + HnswCollection collection = await SeedAsync(); + string path = Path.GetTempFileName(); + try + { + using (FileStream fs = File.Create(path)) + { + collection.Save(fs, SnapshotContext.Default); + } + + var store = new HnswVectorStore(); + HnswCollection reloaded = store.GetCollection("docs"); + reloaded.Load(path, SnapshotContext.Default); + + // Lazily materialized record payload. + Doc? fetched = await reloaded.GetAsync(2); + Assert.NotNull(fetched); + Assert.Equal("y axis", fetched!.Text); + + var results = new List>(); + await foreach (VectorSearchResult r in reloaded.SearchAsync(new float[] { 0.9f, 0.1f, 0f }, top: 3)) + { + results.Add(r); + } + + Assert.Equal(3, results.Count); + Assert.Equal(1, results[0].Record.Id); + + // Disposing the store must release the mapping so the file can be deleted (Windows locks maps). + store.Dispose(); + File.Delete(path); + Assert.False(File.Exists(path)); + } + finally + { + if (File.Exists(path)) + { + File.Delete(path); + } + } + } + + [Fact] + public async Task LoadMapped_ByReflection_RoundTrips() + { + HnswCollection collection = await SeedAsync(); + string path = Path.GetTempFileName(); + try + { + using (FileStream fs = File.Create(path)) + { + collection.Save(fs, SnapshotContext.Default); + } + + using var store = new HnswVectorStore(); + HnswCollection reloaded = store.GetCollection("docs"); + reloaded.Load(path); + + Doc? fetched = await reloaded.GetAsync(3); + Assert.NotNull(fetched); + Assert.Equal("z axis", fetched!.Text); + } + finally + { + File.Delete(path); + } + } + + [Fact] + public async Task LoadMapped_AtOffset_RoundTrips() + { + HnswCollection collection = await SeedAsync(); + string path = Path.GetTempFileName(); + const int prefix = 13; + try + { + using (FileStream fs = File.Create(path)) + { + fs.Write(new byte[prefix]); + collection.Save(fs, SnapshotContext.Default); + } + + using var store = new HnswVectorStore(); + HnswCollection reloaded = store.GetCollection("docs"); + reloaded.Load(path, prefix, SnapshotContext.Default); + + Doc? fetched = await reloaded.GetAsync(1); + Assert.NotNull(fetched); + Assert.Equal("x axis", fetched!.Text); + } + finally + { + File.Delete(path); + } + } + + [Fact] + public async Task LoadMapped_MatchesStreamLoadResults() + { + HnswCollection collection = await SeedAsync(); + string path = Path.GetTempFileName(); + try + { + using (FileStream fs = File.Create(path)) + { + collection.Save(fs, SnapshotContext.Default); + } + + var streamStore = new HnswVectorStore(); + HnswCollection streamLoaded = streamStore.GetCollection("docs"); + using (FileStream fs = File.OpenRead(path)) + { + streamLoaded.Load(fs, SnapshotContext.Default); + } + + using var mapStore = new HnswVectorStore(); + HnswCollection mapLoaded = mapStore.GetCollection("docs"); + mapLoaded.Load(path, SnapshotContext.Default); + + float[] query = { 0f, 0.2f, 0.9f }; + Assert.Equal(await CollectIdsAsync(streamLoaded, query), await CollectIdsAsync(mapLoaded, query)); + streamStore.Dispose(); + } + finally + { + File.Delete(path); + } + } + + private static async Task> CollectIdsAsync(HnswCollection collection, float[] query) + { + var ids = new List(); + await foreach (VectorSearchResult r in collection.SearchAsync(query, top: 3)) + { + ids.Add(r.Record.Id); + } + + return ids; + } +} From f87d38a1973fd011b66c7e7f06a721418cc4ca3c Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 11:20:14 -0700 Subject: [PATCH 7/9] Validate link indices and harden vector I/O endianness Validate every persisted link index (0 <= neighbor < Count) and bound link counts while loading. Unchecked indices reached MappedVectorBlock, which builds a span from a raw pointer with no bounds check, so a corrupt link could read arbitrary mapped memory or fault the process instead of throwing. Restore the explicit little-endian on-disk vector format that the pre-bulk-IO code guaranteed: bulk byte copy on little-endian hosts (the only platforms .NET supports) with a scalar fallback otherwise. LoadMapped reinterprets mapped bytes in place, so it now rejects big-endian platforms rather than returning silently wrong results. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Hnsw.Net/HnswIndex.cs | 71 +++++++++++++++++++++++--- tests/Hnsw.Net.Tests/HnswIndexTests.cs | 66 ++++++++++++++++++++++++ 2 files changed, 130 insertions(+), 7 deletions(-) diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index ee3c056..e51e189 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -1,3 +1,4 @@ +using System.Buffers.Binary; using System.Collections.Concurrent; using System.IO.MemoryMappedFiles; using System.Numerics; @@ -502,11 +503,9 @@ public void Save(Stream stream) // Vector section: one normalized vector per slot, contiguous and fixed-stride, so the // read path can memory-map it and address slot s at base + s * Dimension * sizeof(float). - // Written as raw little-endian float blocks (byte-identical to a per-float loop on the - // LE platforms .NET targets, but without millions of scalar writes). for (int n = 0; n < _nodes.Count; n++) { - writer.Write(MemoryMarshal.AsBytes(StoredSpan(n))); + WriteFloatsLittleEndian(writer, StoredSpan(n)); } // Graph section: per-node metadata and link lists, kept separate from the vectors. @@ -551,7 +550,7 @@ public static HnswIndex Load(Stream stream) // Vector section first (one vector per slot), then the graph section. for (int i = 0; i < header.Count; i++) { - ReadExactInto(reader, MemoryMarshal.AsBytes(index._vectors.VectorMutable(i))); + ReadFloatsLittleEndian(reader, index._vectors.VectorMutable(i)); } for (int i = 0; i < header.Count; i++) @@ -584,6 +583,15 @@ public static HnswIndex LoadMapped(string path, long baseOffset) { ArgumentNullException.ThrowIfNull(path); ArgumentOutOfRangeException.ThrowIfNegative(baseOffset); + + // The mapped vector section is read as host-endian floats straight off the pages, so it is + // only correct on little-endian platforms (the only ones .NET supports). Reject otherwise + // rather than return silently wrong results; the stream loader handles big-endian hosts. + if (!BitConverter.IsLittleEndian) + { + throw new PlatformNotSupportedException("Memory-mapped loading is only supported on little-endian platforms."); + } + Header header; long vectorOffset; HnswIndex index; @@ -664,10 +672,10 @@ private static void ReadInterleavedBody(BinaryReader reader, HnswIndex index, He long id = reader.ReadInt64(); int level = reader.ReadInt32(); bool deleted = header.Version >= 3 && reader.ReadBoolean(); - ReadExactInto(reader, MemoryMarshal.AsBytes(index._vectors.VectorMutable(i))); + ReadFloatsLittleEndian(reader, index._vectors.VectorMutable(i)); if (discard.Length > 0) { - ReadExactInto(reader, MemoryMarshal.AsBytes(discard.AsSpan())); + ReadFloatsLittleEndian(reader, discard.AsSpan()); } ReadNodeLinks(reader, index, i, id, level, deleted, header); @@ -686,9 +694,27 @@ private static void ReadNodeLinks(BinaryReader reader, HnswIndex index, int i, l for (int layer = 0; layer < layerCount; layer++) { int linkCount = reader.ReadInt32(); + + // Bound the count by the node total (a slot can have at most Count distinct neighbors) + // so a corrupt length cannot drive a huge allocation, then validate every neighbor slot. + // Unchecked indices would later reach MappedVectorBlock, which builds a span directly + // from a raw pointer with no bounds check, so an out-of-range link could read arbitrary + // mapped memory or fault the process instead of throwing. + if ((uint)linkCount > (uint)header.Count) + { + throw new InvalidDataException("Invalid link count in Hnsw.Net index."); + } + + List links = node.Links[layer]; for (int link = 0; link < linkCount; link++) { - node.Links[layer].Add(reader.ReadInt32()); + int neighbor = reader.ReadInt32(); + if ((uint)neighbor >= (uint)header.Count) + { + throw new InvalidDataException("Link index out of range in Hnsw.Net index."); + } + + links.Add(neighbor); } } @@ -786,6 +812,37 @@ private static void ReadExactInto(BinaryReader reader, Span buffer) } } + // The on-disk vector section is little-endian. On little-endian hosts (the only platforms .NET + // supports) this is a single bulk byte copy; the big-endian fallbacks keep the format portable. + private static void WriteFloatsLittleEndian(BinaryWriter writer, ReadOnlySpan values) + { + if (BitConverter.IsLittleEndian) + { + writer.Write(MemoryMarshal.AsBytes(values)); + return; + } + + Span scratch = stackalloc byte[sizeof(float)]; + foreach (float value in values) + { + BinaryPrimitives.WriteSingleLittleEndian(scratch, value); + writer.Write(scratch); + } + } + + private static void ReadFloatsLittleEndian(BinaryReader reader, Span destination) + { + Span bytes = MemoryMarshal.AsBytes(destination); + ReadExactInto(reader, bytes); + if (!BitConverter.IsLittleEndian) + { + for (int i = 0; i < bytes.Length; i += sizeof(float)) + { + bytes.Slice(i, sizeof(float)).Reverse(); + } + } + } + private int RandomLevel() { double sample = Math.Max(_random.NextDouble(), double.Epsilon); diff --git a/tests/Hnsw.Net.Tests/HnswIndexTests.cs b/tests/Hnsw.Net.Tests/HnswIndexTests.cs index 2ef8c30..a7c3e27 100644 --- a/tests/Hnsw.Net.Tests/HnswIndexTests.cs +++ b/tests/Hnsw.Net.Tests/HnswIndexTests.cs @@ -235,6 +235,72 @@ public void LoadsLegacyV3FormatAndDiscardsDuplicateVector() Assert.Equal(7, loaded.Search(vector, 1)[0].Id); } + [Fact] + public void RejectsOutOfRangeLinkIndexOnLoad() + { + byte[] payload = BuildV4StreamWithLink(neighbor: 5, count: 2, dimension: 3); + + Assert.Throws(() => HnswIndex.Load(new MemoryStream(payload))); + + string path = Path.Combine(Path.GetTempPath(), $"hnsw_corrupt_{Guid.NewGuid():N}.bin"); + try + { + File.WriteAllBytes(path, payload); + Assert.Throws(() => HnswIndex.LoadMapped(path)); + } + finally + { + File.Delete(path); + } + } + + // Hand-craft a current-format (v4) stream of `count` zero vectors where the first node carries a + // single layer-0 link to `neighbor`, used to verify out-of-range links are rejected at load. + private static byte[] BuildV4StreamWithLink(int neighbor, int count, int dimension) + { + const uint magic = 0x31575348; + const int version = 4; + using var stream = new MemoryStream(); + using (var writer = new BinaryWriter(stream, System.Text.Encoding.UTF8, leaveOpen: true)) + { + writer.Write(magic); + writer.Write(version); + writer.Write(dimension); + writer.Write((int)DistanceMetric.DotProduct); + writer.Write(2); // m + writer.Write(10); // efConstruction + writer.Write(10); // ef + writer.Write(0); // entryPoint + writer.Write(0); // maxLevel + writer.Write(count); + writer.Write(false); // allowReplaceDeleted + + for (int n = 0; n < count; n++) // vector section + { + for (int j = 0; j < dimension; j++) writer.Write(0f); + } + + for (int n = 0; n < count; n++) // graph section + { + writer.Write((long)n); // id + writer.Write(0); // level + writer.Write(false); // deleted + writer.Write(1); // layer count + if (n == 0) + { + writer.Write(1); // layer 0 link count + writer.Write(neighbor); // out-of-range neighbor + } + else + { + writer.Write(0); + } + } + } + + return stream.ToArray(); + } + [Fact] public void HandlesEdgeCases() { From cbf89c6bd00ba9e37392d9840c723cda0314ea55 Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 11:43:25 -0700 Subject: [PATCH 8/9] Harden Dispose, read-only contract, and mmap header validation - Make HnswIndex.Dispose idempotent (ReaderWriterLockSlim must not be disposed twice). - Enforce the read-only contract in MarkDeleted/UnmarkDeleted, not just Add. - Validate the vector section in LoadMapped: reject negative count/dimension, guard the size multiplication against overflow, and ensure it fits the file. - Dispose HnswCollectionData under its Lock and clear fields so teardown cannot race in-flight operations using unsafe mapped spans. - Use 'using' for stores in mmap tests so disposal cannot mask assertion failures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Hnsw.Net/HnswIndex.cs | 43 ++++++++++++++++--- src/Hnsw.Net/VectorData/HnswCollectionData.cs | 9 +++- tests/Hnsw.Net.Tests/HnswIndexTests.cs | 6 +++ tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs | 11 ++--- 4 files changed, 54 insertions(+), 15 deletions(-) diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index e51e189..234ea9b 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -32,6 +32,7 @@ public sealed class HnswIndex : IDisposable private readonly ConcurrentBag _scratchPool = new(); private readonly Comparison _candidateComparison; private VectorBlock _vectors; + private bool _disposed; /// Initializes a new HNSW index. /// Vector dimension. All indexed and query vectors must have this length. @@ -180,10 +181,7 @@ public static HnswIndex Build( public void Add(long id, ReadOnlySpan vector) { ValidateVector(vector); - if (_vectors.IsReadOnly) - { - throw new InvalidOperationException("This index was loaded with memory-mapped vectors and is read-only. Load it without mapping to modify it."); - } + ThrowIfReadOnly(); _lock.EnterWriteLock(); try @@ -249,6 +247,7 @@ public void Add(long id, ReadOnlySpan vector) /// public void MarkDeleted(long id) { + ThrowIfReadOnly(); _lock.EnterWriteLock(); try { @@ -279,6 +278,7 @@ public void MarkDeleted(long id) /// Restores a vector previously marked by , making it searchable again. public void UnmarkDeleted(long id) { + ThrowIfReadOnly(); _lock.EnterWriteLock(); try { @@ -606,7 +606,26 @@ public static HnswIndex LoadMapped(string path, long baseOffset) } vectorOffset = stream.Position; - long vectorBytes = (long)header.Count * header.Dimension * sizeof(float); + if (header.Count < 0 || header.Dimension < 0) + { + throw new InvalidDataException("The index header specifies a negative count or dimension."); + } + + long vectorBytes; + try + { + vectorBytes = checked((long)header.Count * header.Dimension * sizeof(float)); + } + catch (OverflowException) + { + throw new InvalidDataException("The index vector section size is invalid."); + } + + if (vectorOffset + vectorBytes > stream.Length) + { + throw new InvalidDataException("The index vector section extends beyond the end of the file."); + } + index = new HnswIndex(header.Dimension, header.Metric, header.M, header.EfConstruction, header.Ef, header.EntryPoint, header.MaxLevel, header.AllowReplaceDeleted); stream.Seek(vectorOffset + vectorBytes, SeekOrigin.Begin); @@ -745,10 +764,24 @@ private readonly record struct Header( /// Releases the memory mapping held by an index loaded via . public void Dispose() { + if (_disposed) + { + return; + } + + _disposed = true; (_vectors as IDisposable)?.Dispose(); _lock.Dispose(); } + private void ThrowIfReadOnly() + { + if (_vectors.IsReadOnly) + { + throw new InvalidOperationException("This index was loaded with memory-mapped vectors and is read-only. Load it without mapping to modify it."); + } + } + private void ValidateVector(ReadOnlySpan vector) { if (vector.Length != Dimension) diff --git a/src/Hnsw.Net/VectorData/HnswCollectionData.cs b/src/Hnsw.Net/VectorData/HnswCollectionData.cs index fae79cf..30d4dcc 100644 --- a/src/Hnsw.Net/VectorData/HnswCollectionData.cs +++ b/src/Hnsw.Net/VectorData/HnswCollectionData.cs @@ -31,8 +31,13 @@ internal sealed class HnswCollectionData : IDisposable public void Dispose() { - Index?.Dispose(); - MappedRecords?.Dispose(); + lock (Lock) + { + Index?.Dispose(); + MappedRecords?.Dispose(); + Index = null; + MappedRecords = null; + } } /// diff --git a/tests/Hnsw.Net.Tests/HnswIndexTests.cs b/tests/Hnsw.Net.Tests/HnswIndexTests.cs index a7c3e27..714d82e 100644 --- a/tests/Hnsw.Net.Tests/HnswIndexTests.cs +++ b/tests/Hnsw.Net.Tests/HnswIndexTests.cs @@ -105,6 +105,12 @@ public void LoadMappedProducesIdenticalResultsAndIsReadOnly() Assert.True(mapped.TryGetVector(10_000, out float[] stored)); Assert.Equal(Normalize(vectors[0]), stored); Assert.Throws(() => mapped.Add(99_999, vectors[0])); + Assert.Throws(() => mapped.MarkDeleted(10_000)); + Assert.Throws(() => mapped.UnmarkDeleted(10_000)); + + // Dispose must be idempotent (ReaderWriterLockSlim cannot be disposed twice). + mapped.Dispose(); + mapped.Dispose(); } // Dispose must release the mapping so the file is no longer locked. diff --git a/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs b/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs index e08a3bc..a126820 100644 --- a/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs +++ b/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs @@ -18,7 +18,7 @@ public async Task LoadMapped_FromPath_RoundTripsRecordsAndSearch() collection.Save(fs, SnapshotContext.Default); } - var store = new HnswVectorStore(); + using var store = new HnswVectorStore(); HnswCollection reloaded = store.GetCollection("docs"); reloaded.Load(path, SnapshotContext.Default); @@ -35,14 +35,10 @@ public async Task LoadMapped_FromPath_RoundTripsRecordsAndSearch() Assert.Equal(3, results.Count); Assert.Equal(1, results[0].Record.Id); - - // Disposing the store must release the mapping so the file can be deleted (Windows locks maps). - store.Dispose(); - File.Delete(path); - Assert.False(File.Exists(path)); } finally { + // store is disposed by `using` above, releasing the mapping so the file can be deleted (Windows locks maps). if (File.Exists(path)) { File.Delete(path); @@ -116,7 +112,7 @@ public async Task LoadMapped_MatchesStreamLoadResults() collection.Save(fs, SnapshotContext.Default); } - var streamStore = new HnswVectorStore(); + using var streamStore = new HnswVectorStore(); HnswCollection streamLoaded = streamStore.GetCollection("docs"); using (FileStream fs = File.OpenRead(path)) { @@ -129,7 +125,6 @@ public async Task LoadMapped_MatchesStreamLoadResults() float[] query = { 0f, 0.2f, 0.9f }; Assert.Equal(await CollectIdsAsync(streamLoaded, query), await CollectIdsAsync(mapLoaded, query)); - streamStore.Dispose(); } finally { From 27a65d03bbcbfd1ad60c14be2fd7ca0941d0ef88 Mon Sep 17 00:00:00 2001 From: Eric StJohn Date: Tue, 9 Jun 2026 11:52:11 -0700 Subject: [PATCH 9/9] Validate index header dimension/count in the shared loader - Reject a non-positive dimension and negative count in ReadHeader so both Load and LoadMapped fail with InvalidDataException instead of DivideByZeroException or a silently-empty index on corrupt input. - Dispose the partially-constructed index if creating the memory mapping fails. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/Hnsw.Net/HnswIndex.cs | 16 +++++++++++----- tests/Hnsw.Net.Tests/HnswIndexTests.cs | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 5 deletions(-) diff --git a/src/Hnsw.Net/HnswIndex.cs b/src/Hnsw.Net/HnswIndex.cs index 234ea9b..72d62f4 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -606,11 +606,6 @@ public static HnswIndex LoadMapped(string path, long baseOffset) } vectorOffset = stream.Position; - if (header.Count < 0 || header.Dimension < 0) - { - throw new InvalidDataException("The index header specifies a negative count or dimension."); - } - long vectorBytes; try { @@ -651,6 +646,7 @@ public static HnswIndex LoadMapped(string path, long baseOffset) { view?.Dispose(); file?.Dispose(); + index.Dispose(); throw; } } @@ -677,6 +673,16 @@ private static Header ReadHeader(BinaryReader reader) int maxLevel = reader.ReadInt32(); int count = reader.ReadInt32(); bool allowReplaceDeleted = version >= 3 && reader.ReadBoolean(); + if (dimension <= 0) + { + throw new InvalidDataException("The index header specifies a non-positive dimension."); + } + + if (count < 0) + { + throw new InvalidDataException("The index header specifies a negative count."); + } + return new Header(version, dimension, metric, m, efConstruction, ef, entryPoint, maxLevel, count, allowReplaceDeleted); } diff --git a/tests/Hnsw.Net.Tests/HnswIndexTests.cs b/tests/Hnsw.Net.Tests/HnswIndexTests.cs index 714d82e..40818bc 100644 --- a/tests/Hnsw.Net.Tests/HnswIndexTests.cs +++ b/tests/Hnsw.Net.Tests/HnswIndexTests.cs @@ -260,6 +260,24 @@ public void RejectsOutOfRangeLinkIndexOnLoad() } } + [Fact] + public void RejectsNonPositiveDimensionOnLoad() + { + byte[] payload = BuildV4StreamWithLink(neighbor: 0, count: 0, dimension: 0); + Assert.Throws(() => HnswIndex.Load(new MemoryStream(payload))); + + string path = Path.Combine(Path.GetTempPath(), $"hnsw_dim0_{Guid.NewGuid():N}.bin"); + try + { + File.WriteAllBytes(path, payload); + Assert.Throws(() => HnswIndex.LoadMapped(path)); + } + finally + { + File.Delete(path); + } + } + // Hand-craft a current-format (v4) stream of `count` zero vectors where the first node carries a // single layer-0 link to `neighbor`, used to verify out-of-range links are rejected at load. private static byte[] BuildV4StreamWithLink(int neighbor, int count, int dimension)