Support file-based serialize/deserialize for DiskANN indices using Save/Load traits#1079
Support file-based serialize/deserialize for DiskANN indices using Save/Load traits#1079suhasjs wants to merge 28 commits into
Conversation
…borrow for Writer::finish() since it calls std::io::Write::flush() that needs a mutable borrow
…ave::Results<Handle>
…w directly impls for Value<'static>
… wherever needed + modify writer return type to Result<>
… required for their Load/Save impls)
… diskann-providers with a new SingleUseWriteProvider (also for reads with SingleUseReadProvider)
…ad_from_bin<> to only use one read since VectorDataIterator already has the metadata needed for
…ghborProviderAsync
hildebrandmw
left a comment
hildebrandmw
left a comment
There was a problem hiding this comment.
Thanks Suhas! My comments this round are mainly about enum handling. I think we can do this more cleanly by having enums "handle themselves" rather than lifting it into a somewhat fragile interaction in the Save/Load traits.
Other than that, the code inside diskann-record could use considerable hardening and testing.
Do we want to plan out how plugins like VFS would work in this PR, or save that for the future? I ask because VFS is used pretty heavily in tests, so having a replacement that's compliant with this new system will be required. And it will be a good stress test that we can expand the backend serialization to other formats in the future.
| /// load that the tag's presence matches the corresponding [`Load::IS_ENUM`]. | ||
| fn variant(&self) -> Option<Cow<'_, str>> { | ||
| None | ||
| } |
There was a problem hiding this comment.
After thinking about this for longer than I care to admit, I think we can get rid of enum handling entirely at the trait level. If we consider three cases:
enum UnitEnum {
A,
B,
C,
}
enum InlineStruct {
A(f32),
B {
x: String,
y: f32,
}
}
enum Struct {
A(A),
B(B),
}
struct A {
field: f32,
}
struct B {
x: String,
y: f32,
}The serialized representations in JSON could be
{
"$version": "0.0.0",
"A": null,
}
{
$version": "0.0.0",
"B": {
"x": "hello world",
"y": 1,
}
}
{
"$version": "0.0.0",
"A": {
"$version": "0.0.1",
"0": 1,
}
}For versions 1 and 2, there is just a top level version, which makes sense since the structure of these enums evolves in lockstep with the outer enum. For Version 3, this allows nested objects in enums to have distinct versions naturally.
There was a problem hiding this comment.
$variant is a little more explicit in how it handles enums. But since most of our enums fall into one of the three variants you described, I've gone ahead and removed $variant handling.
| enum ErrorInner { | ||
| Light(Kind), | ||
| Heavy(anyhow::Error), | ||
| } |
There was a problem hiding this comment.
While I still really like the Light/Heavy idea, if we're honest it's probably overkill and will likely not get implemented ever.
Instead, the main goal I think is to differentiate crucial errors vs "recoverable" errors, where the latter is used for a probing API. In that context, we'd have something like:
struct Error {
inner: anyhow::Error,
recoverable: bool,
}Most of the constructors would imply critical errors, with probing APIs using an additional (_recoverable) API. This keeps normal code simple and the probing API possible and reasonably efficient.
There was a problem hiding this comment.
It did seem overkill. I've removed it in the latest commit.
| "handle references file {:?} which is not registered in the manifest", | ||
| key, | ||
| ))); | ||
| } |
There was a problem hiding this comment.
Probably worth adding a
if key.contains("..") || key_as_path.is_absolute() {
return Err(/* ... */);
}to keep files inside the directory.
There was a problem hiding this comment.
Good catch. Added it.
| * Licensed under the MIT license. | ||
| */ | ||
|
|
||
| //! Adapters that expose a single in-hand [`Read`]/[`Write`] target as a |
There was a problem hiding this comment.
Should we go ahead and start updating the exiting code to use io::Read/io::Write directly so we don't need this thing?
There was a problem hiding this comment.
By 'existing code', do you mean all code that uses Storage{Read/Write}Provider traits? Not sure I understood your ask here.
There was a problem hiding this comment.
Most of the code that currently takes Storage{Read/Write}Provider does so by taking the provider and a file name and immediately opening the necessary file. These could all be pretty easily migrated to taking std::io::Read and std::io::Write directly, moving the file open up one level and avoiding the need for this hacky one-off provider.
|
|
||
| pub fn get(&self, key: &str) -> Option<&Value<'a>> { | ||
| self.record.get(key) | ||
| } |
There was a problem hiding this comment.
In the vein of moving enums to handling themselves, if we add:
impl<'a> Record<'a> {
fn keys(&self) -> Keys<'_> {
Keys::new(self.record.key())
}
}
struct Keys<'a> {
keys: std::collections::hash_map::Keys<'a, Cow<'a, str>, Value<'a>>,
}
impl<'a> Iterator for Keys<'a> {
// iterate
}
impl ExactSizeIterator for Keys<'_> {}Then enums could have a way of inspecting the keys for loading.
… that wraps anyhow::Error along with a recoverable bool to allow probing APIUs
…-> throws an error now
Resolve conflicts in Cargo.toml (versions + members), diskann-vector/Cargo.toml (dev-deps), diskann-providers/src/storage/bin.rs (single-read load_from_bin), diskann-providers/src/index/wrapped_async.rs (combine tests), diskann/src/graph/index.rs (combine impls + tests), and port FastMemoryVectorProviderAsync / MemoryVectorProviderAsync save/load impls to main's T: VectorRepr (formerly Data: GraphDataType in this branch).
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #1079 +/- ##
==========================================
- Coverage 89.46% 89.25% -0.22%
==========================================
Files 482 493 +11
Lines 91082 93018 +1936
==========================================
+ Hits 81491 83021 +1530
- Misses 9591 9997 +406
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
hildebrandmw
left a comment
hildebrandmw
left a comment
There was a problem hiding this comment.
A few more drive-by comments.
|
|
||
| Keys starting with `$` are reserved for infrastructure (`$version`, `$variant`, | ||
| `$handle`). Since Rust identifiers can't start with `$`, the `save_fields!` macro is | ||
| inherently safe — no runtime check needed in the macro path. |
There was a problem hiding this comment.
- This struct can be created without any issues:
pub struct Test {
version: String,
handle: String,
variant: String,
}
- The choice of programming language shouldn’t determine the field names in a public contract.
| "$version": { | ||
| "major": 0, |
There was a problem hiding this comment.
version object repeats everywhere. It will bloat this file. Is it a concern? Should we consider tracking version at high level objects only?
| } | ||
| ``` | ||
|
|
||
| The lifetime `'a` allows the save path to borrow from the data being saved (zero-copy |
There was a problem hiding this comment.
Is there a measurable benefit of introducing a lifetime here? I would say that config load/save is outside of hot path or at least its contribution is negligible to overall RAM/CPU usage, since config is tiny.
|
|
||
| impl ContextInner { | ||
| pub(super) fn new(metadata: &Path, dir: &Path) -> Result<Self> { | ||
| let file = std::fs::File::open(metadata).map_err(|e| { |
There was a problem hiding this comment.
we should not code directly against the file system. Instead, use Storage{Read/Write}Provider abstraction. It allows to write directly to remote store and provides better testability of code.
|
|
||
| ## Remaining Work | ||
|
|
||
| ### Value::Bytes Wiring |
There was a problem hiding this comment.
are we planning to store bytes in a configuration file?
| "$version": { | ||
| "major": 0, | ||
| "minor": 0, | ||
| "patch": 0 |
There was a problem hiding this comment.
I think major/minor/patch is overkill. Each object should just have a single version number that is not tied to the overall version of diskann.
…/Deserialize<> for Version; updated sample_output
5 participants
Reference Issues/PRs
Initial attempt at #737
What does this implement/fix? Briefly explain your changes.
Introduces two new traits (
diskann_record::save::Saveanddiskann_record::load::Load)and a new creatediskann-recordto support file-based serialization + deserialization. This PR also implements the traits forDiskANNIndex(and all nested structs) for a successful prototype.Sample output can be found here.
Any other comments?
This PR is meant to be in a draft state until we're happy with the abstractions and the results.