diff --git a/README.md b/README.md index f75a9cd..6138007 100644 --- a/README.md +++ b/README.md @@ -112,9 +112,26 @@ 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. + +## 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 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 ca214a8..72d62f4 100644 --- a/src/Hnsw.Net/HnswIndex.cs +++ b/src/Hnsw.Net/HnswIndex.cs @@ -1,4 +1,6 @@ +using System.Buffers.Binary; using System.Collections.Concurrent; +using System.IO.MemoryMappedFiles; using System.Numerics; using System.Numerics.Tensors; using System.Runtime.CompilerServices; @@ -11,9 +13,9 @@ 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 = 3; + private const int FormatVersion = 4; private const uint Magic = 0x31575348; // HSW1, little-endian. private readonly List _nodes = new(); @@ -29,7 +31,8 @@ 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; + private bool _disposed; /// Initializes a new HNSW index. /// Vector dimension. All indexed and query vectors must have this length. @@ -62,6 +65,7 @@ public HnswIndex(int dimension, DistanceMetric metric, int m = 16, int efConstru _levelMultiplier = 1.0 / Math.Log(m); _allowReplaceDeleted = allowReplaceDeleted; _candidateComparison = CompareCandidates; + _vectors = new HeapVectorBlock(dimension); } private HnswIndex(int dimension, DistanceMetric metric, int m, int efConstruction, int ef, int entryPoint, int maxLevel, bool allowReplaceDeleted) @@ -77,6 +81,7 @@ private HnswIndex(int dimension, DistanceMetric metric, int m, int efConstructio _levelMultiplier = 1.0 / Math.Log(m); _allowReplaceDeleted = allowReplaceDeleted; _candidateComparison = CompareCandidates; + _vectors = new HeapVectorBlock(dimension); } /// Gets the vector dimension. @@ -116,7 +121,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() { @@ -124,11 +131,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())); } } @@ -173,7 +181,8 @@ public static HnswIndex Build( public void Add(long id, ReadOnlySpan vector) { ValidateVector(vector); - float[] originalVector = vector.ToArray(); + ThrowIfReadOnly(); + _lock.EnterWriteLock(); try { @@ -187,7 +196,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; } } @@ -195,8 +204,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); @@ -238,6 +247,7 @@ public void Add(long id, ReadOnlySpan vector) /// public void MarkDeleted(long id) { + ThrowIfReadOnly(); _lock.EnterWriteLock(); try { @@ -268,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 { @@ -305,7 +316,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(); @@ -313,7 +327,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; } @@ -367,7 +381,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++) @@ -382,7 +396,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); @@ -439,7 +452,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); @@ -488,22 +501,20 @@ 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). + for (int n = 0; n < _nodes.Count; n++) + { + WriteFloatsLittleEndian(writer, 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); - 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]); - } - writer.Write(node.Links.Length); for (int layer = 0; layer < node.Links.Length; layer++) { @@ -521,15 +532,132 @@ 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++) + { + ReadFloatsLittleEndian(reader, 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) => 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); + + // 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; + 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) + { + throw new InvalidDataException("Memory-mapped loading requires the current index format; re-save the index."); + } + + vectorOffset = stream.Position; + 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); + 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(); + index.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) { @@ -545,61 +673,119 @@ 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._storedVectors = count > 0 ? new float[count * dimension] : Array.Empty(); + if (dimension <= 0) + { + throw new InvalidDataException("The index header specifies a non-positive dimension."); + } - for (int i = 0; i < count; i++) + 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); + } + + // 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 = 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._storedVectors.AsSpan(i * dimension, dimension); - for (int j = 0; j < dimension; j++) + bool deleted = header.Version >= 3 && reader.ReadBoolean(); + ReadFloatsLittleEndian(reader, index._vectors.VectorMutable(i)); + if (discard.Length > 0) { - stored[j] = reader.ReadSingle(); - } - var originalVector = new float[dimension]; - if (version >= 2) - { - for (int j = 0; j < dimension; j++) - { - originalVector[j] = reader.ReadSingle(); - } - } - else - { - stored.CopyTo(originalVector); + ReadFloatsLittleEndian(reader, discard.AsSpan()); } - var node = new Node(id, originalVector, level) { Deleted = deleted }; - int layerCount = reader.ReadInt32(); - if (layerCount != level + 1) + ReadNodeLinks(reader, index, i, id, level, deleted, header); + } + } + + 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(); + + // 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 layer count in Hnsw.Net index."); + throw new InvalidDataException("Invalid link count in Hnsw.Net index."); } - for (int layer = 0; layer < layerCount; layer++) + List links = node.Links[layer]; + for (int link = 0; link < linkCount; link++) { - int linkCount = reader.ReadInt32(); - for (int link = 0; link < linkCount; link++) + int neighbor = reader.ReadInt32(); + if ((uint)neighbor >= (uint)header.Count) { - node.Links[layer].Add(reader.ReadInt32()); + throw new InvalidDataException("Link index out of range in Hnsw.Net index."); } + + links.Add(neighbor); } + } - 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() + { + 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) @@ -638,26 +824,63 @@ private void PrepareVectorInto(ReadOnlySpan source, Span destinati } } - private void EnsureStoredCapacity(int nodeCount) + private void EnsureStoredCapacity(int nodeCount) => _vectors.EnsureCapacity(nodeCount); + + 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. + private ReadOnlySpan SlotVector(int slot) => StoredSpan(slot); + + private List SlotLinks(int slot, int layer) => _nodes[slot].Links[layer]; + + private Span StoredSpanMutable(int index) => _vectors.VectorMutable(index); + + 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; + } + } + + // 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) { - int required = nodeCount * Dimension; - if (_storedVectors.Length >= required) + if (BitConverter.IsLittleEndian) { + writer.Write(MemoryMarshal.AsBytes(values)); return; } - int capacity = _storedVectors.Length == 0 ? Dimension * 16 : _storedVectors.Length * 2; - if (capacity < required) + Span scratch = stackalloc byte[sizeof(float)]; + foreach (float value in values) { - capacity = required; + BinaryPrimitives.WriteSingleLittleEndian(scratch, value); + writer.Write(scratch); } - - Array.Resize(ref _storedVectors, capacity); } - private ReadOnlySpan StoredSpan(int index) => _storedVectors.AsSpan(index * Dimension, Dimension); - - private Span StoredSpanMutable(int index) => _storedVectors.AsSpan(index * Dimension, Dimension); + 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() { @@ -671,9 +894,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; @@ -694,7 +917,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; @@ -716,7 +939,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) { @@ -724,7 +947,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); @@ -763,11 +986,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; @@ -869,12 +1092,116 @@ private static float DotProduct(ReadOnlySpan a, ReadOnlySpan b) return sum; } + // 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 HeapVectorBlock(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 override bool IsReadOnly => false; + + public override void EnsureCapacity(int slotCount) + { + while (_capacity < slotCount) + { + _chunks.Add(new float[_vectorsPerChunk * _dimension]); + _capacity += _vectorsPerChunk; + } + } + + public override ReadOnlySpan Vector(int slot) => Slot(slot); + + public override 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); + } + } + + // 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, 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++) @@ -885,8 +1212,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/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..30d4dcc 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,122 @@ 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() + { + lock (Lock) + { + Index?.Dispose(); + MappedRecords?.Dispose(); + Index = null; + MappedRecords = null; + } + } + + /// + /// 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/HnswIndexTests.cs b/tests/Hnsw.Net.Tests/HnswIndexTests.cs index aa6b3d7..40818bc 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; @@ -68,6 +69,63 @@ 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])); + 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. + File.Delete(path); + Assert.False(File.Exists(path)); + } + finally + { + if (File.Exists(path)) + { + File.Delete(path); + } + } + } + [Fact] public void SaveLoadRoundTripProducesIdenticalResults() { @@ -98,7 +156,7 @@ public void SaveLoadRoundTripProducesIdenticalResults() } [Fact] - public void ExportItemsRebuildsPortableIndexWithOriginalVectors() + public void ExportItemsRebuildsPortableIndexWithStoredVectors() { const int count = 400; const int dimension = 32; @@ -114,17 +172,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 +202,129 @@ 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 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); + } + } + + [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) + { + 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() { @@ -174,6 +357,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, diff --git a/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs b/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs new file mode 100644 index 0000000..a126820 --- /dev/null +++ b/tests/Hnsw.Net.Tests/VectorStoreMmapTests.cs @@ -0,0 +1,145 @@ +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); + } + + using 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); + } + 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); + } + } + } + + [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); + } + + using 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)); + } + 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; + } +}