Add provider-specific Save/Load persistence to HnswCollection - #5
Merged
Conversation
The Microsoft.Extensions.VectorData abstraction has no persistence API, so expose Save(Stream)/Load(Stream) on the concrete HnswCollection. Consumers "break glass" by holding the concrete type to snapshot the records plus the backing HnswIndex graph to a single stream and restore them later. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
This PR adds provider-specific persistence to HnswCollection<TKey, TRecord> by exposing Save(Stream) / Load(Stream) so consumers can snapshot both the stored records and the underlying HnswIndex graph to a single stream and restore it later, since the Microsoft.Extensions.VectorData abstraction doesn’t define a persistence API.
Changes:
- Added
HnswCollection.Save(Stream)andHnswCollection.Load(Stream)using a versioned binary header plus JSON-serialized record payload and optionalHnswIndexbytes. - Added an xUnit test that seeds a collection, saves to a
MemoryStream, reloads into a new collection, and verifies Get + Search behavior after reload.
Show a summary per file
| File | Description |
|---|---|
| tests/Hnsw.Net.Tests/VectorStoreTests.cs | Adds a Save/Load round-trip test validating record retrieval and nearest-neighbor search after reload. |
| src/Hnsw.Net/VectorData/HnswCollection.cs | Implements provider-specific snapshot persistence via Save(Stream) and Load(Stream) with a custom stream format. |
Copilot's findings
Tip
Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Files reviewed: 2/2 changed files
- Comments generated: 3
The reflection-based Save/Load require runtime reflection and so are unusable under trimming or NativeAOT. Add overloads that take a source-generated JsonSerializerContext and serialize each key and record through the supplied JsonTypeInfo, so callers only need metadata for TKey and TRecord. Both the reflection and context overloads share one per-record binary framing, so a snapshot written by either is loadable by either. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Address three issues raised in review of the Save/Load persistence: - Validate payload lengths and verify the requested bytes were actually read, so a truncated or corrupt stream fails fast instead of surfacing as confusing JSON errors or oversized allocations. - Record the distance metric in the snapshot and reject a Load whose metric or vector dimension does not match the target collection's configured model. - Update the existing HnswCollectionData in place under its lock rather than swapping the instance in the shared dictionary, so the lock object stays stable for threads that already captured it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Fail fast when a snapshot claims an index but it has no vector for a record id, instead of silently storing an empty vector that would not be searchable. - Reject a payload length that exceeds the remaining bytes of a seekable stream before allocating, so a bogus length prefix cannot trigger an OutOfMemory. - Drop the misleading typeof(Name) suggestion from the missing-metadata error, which is not a valid expression for generic or nested types. - Fix the README JsonSerializerContext sample to include a body. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Validate source-gen Load deserialization results instead of null-forgiving, throwing InvalidDataException on null/wrong-typed JSON payloads. - After HnswIndex.Load, verify the index header dimension/metric match the snapshot header and fail fast on a mismatch. - Cap individual key/record payload length so a corrupt prefix on a non-seekable stream cannot trigger an OOM allocation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Reject negative dimension and nextId header values. - Fail fast when a record id is not less than nextId, which would otherwise surface later as a duplicate-id error on the next Upsert. - Build and validate the new state fully before mutating the live collection, so a corrupt snapshot no longer leaves it partially loaded. - Add the missing using directive to the README JsonSerializerContext sample. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- Enforce MaxPayloadLength in SaveCore so Save never emits a snapshot that Load would reject as oversized. - Establish/validate the record type in _collectionTypes before creating the per-collection data, closing a race where a concurrent GetCollection with a different TRecord could observe mixed-type data. - Use BinaryPrimitives little-endian writes in the corruption tests so they are stable on big-endian runtimes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- SaveCore validates the runtime key/record types and throws a clear InvalidOperationException instead of an opaque InvalidCastException when a collection name was accidentally used with mixed key/record types. - LoadCore rejects a non-empty snapshot that has no index, which would otherwise load records with empty vectors and silently return no search hits. - LoadCore fails fast on duplicate record keys or ids instead of silently overwriting earlier entries. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
A snapshot carrying an index with dimension 0 could build an HnswIndex whose Dimension is 0, leaving the collection in a broken state. Treat it as corrupt and fail fast during Load. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The fixed-size header/record reads in LoadCore previously surfaced a truncated stream as EndOfStreamException, inconsistent with the InvalidDataException used by ReadExact and the rest of Load. Wrap the reads and rethrow InvalidDataException so callers see a single corruption shape. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
A thorough self-review (plus a second-pass code review) surfaced three remaining gaps in the snapshot persistence contract: - Save now serializes and validates every entry (type and payload size) before writing any bytes, so a failure can no longer leave a partially-written snapshot on the destination stream. - Load normalizes the full range of exceptions a corrupt or malicious index region can raise from HnswIndex.Load (overflow, argument, I/O, etc.) to InvalidDataException, matching the rest of Load. HnswIndex.Load's own InvalidDataException still propagates unchanged. - Load also maps NotSupportedException from record deserialization to InvalidDataException alongside the existing JsonException handling. Adds tests for a corrupt index region and a malformed JSON payload. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Fail fast with a clear ArgumentException when Save receives a non-writable stream or Load receives a non-readable stream, instead of surfacing a later, less actionable NotSupportedException/IOException (which Load would otherwise mislabel as a deserialization failure). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
NextId starts at 0, so the connector never produces negative record ids. Treat a negative id in a snapshot as corruption and fail fast rather than loading state that violates the connector's id invariants. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Previously the collection type registry keyed only on TRecord, so the same collection name could in principle be associated with different TKey types, leaving records present but unretrievable via key equality. Track and validate the (TKey, TRecord) pair in both GetCollection and Load, with a clearer message. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Microsoft.Extensions.VectorData abstraction has no persistence API, so expose Save(Stream)/Load(Stream) on the concrete HnswCollection. Consumers "break glass" by holding the concrete type to snapshot the records plus the backing HnswIndex graph to a single stream and restore them later.