Skip to content

Commit 23342f1

Browse files
author
Jack Moffitt
committed
Import diskann-garnet
1 parent 6323474 commit 23342f1

14 files changed

Lines changed: 5487 additions & 40 deletions

File tree

Cargo.lock

Lines changed: 1396 additions & 39 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ members = [
1616
"diskann-providers",
1717
"diskann-disk",
1818
"diskann-label-filter",
19+
"diskann-garnet",
1920
# Infrastructure
2021
"diskann-benchmark-runner",
2122
"diskann-benchmark-core",
@@ -35,7 +36,7 @@ default-members = [
3536
resolver = "3"
3637

3738
[workspace.package]
38-
version = "0.47.0" # Obeying semver
39+
version = "0.47.0"
3940
description = "DiskANN is a fast approximate nearest neighbor search library for high dimensional data"
4041
authors = ["Microsoft"]
4142
documentation = "https://github.com/microsoft/DiskANN"

diskann-garnet/Cargo.toml

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
[package]
2+
name = "diskann-garnet"
3+
version = "1.0.25"
4+
edition = "2024"
5+
6+
[lib]
7+
crate-type = ["cdylib"]
8+
9+
[dependencies]
10+
bytemuck.workspace = true
11+
crossbeam = "0.8.4"
12+
dashmap = { workspace = true, features = ["inline"] }
13+
diskann.workspace = true
14+
diskann-quantization.workspace = true
15+
diskann-providers.workspace = true
16+
diskann-vector.workspace = true
17+
foldhash = "0.2.0"
18+
thiserror.workspace = true
19+
tokio.workspace = true
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
<?xml version="1.0" encoding="utf-8"?>
2+
<!-- Version History
3+
| 1.0.0 | Sep 2025 | Initial version |
4+
-->
5+
<package>
6+
<metadata>
7+
<id>diskann-garnet</id>
8+
<version>1.0.25</version>
9+
<readme>docs/README.md</readme>
10+
<authors>DiskANN Team</authors>
11+
<owners>Encyclopedia@service.microsoft.com</owners>
12+
<projectUrl>https://msdata.visualstudio.com/CosmosDB/_git/diskann-garnet</projectUrl>
13+
<description>DiskANN FFI and Data Provider for Garnet</description>
14+
<copyright>© Microsoft Corporation. All rights reserved.</copyright>
15+
<tags>microsoft diskann garnet</tags>
16+
<title />
17+
</metadata>
18+
<files>
19+
<file src="linux/libdiskann_garnet.so" target="runtimes/linux-x64/native/" />
20+
<file src="windows/diskann_garnet.dll" target="runtimes/win-x64/native/" />
21+
<file src="windows/diskann_garnet.pdb" target="runtimes/win-x64/native/" />
22+
<file src="docs/README.md" target="docs/" />
23+
</files>
24+
</package>

diskann-garnet/docs/ffi-design.rs

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
/// Create a new empty index
2+
/// Takes the params of VADD (see: https://redis.io/docs/latest/commands/vadd/), maps to a reasonable interpretation
3+
///
4+
/// (context % 4) == 0, xxx_callbacks add 0/1/2/3 depending on data stored
5+
///
6+
/// Expectation is any state necessary to recover an index is stored via read/write callbacks - including quantizers.
7+
///
8+
/// reduce_dims == 0 to indicate no reduction requested (and can be ignored even if provided, if that is reasonable).
9+
///
10+
/// quant_type needs option that map from NoQuant, Bin, and Q8 (as that's what provided in Redis) in addition to any custom index. They don't need to be exact, just reasonable.
11+
///
12+
/// Returning a single pointer conceal all the generics behind an opaque handle
13+
#[unsafe(no_mangle)]
14+
extern "C" fn create_index(
15+
context: u64,
16+
dimensions: u32,
17+
reduce_dims: u32,
18+
quant_type: SomeCStyleEnumeration,
19+
build_exploration_factor: u32,
20+
num_links: u32,
21+
read_callback: unsafe extern "C" fn(u64, *const u8, usize, *mut u8, usize) -> i32,
22+
write_callback: unsafe extern "C" fn(u64, *const u8, usize, *const u8, usize) -> bool,
23+
delete_callback: unsafe extern "C" fn(u64, *const u8, usize) -> bool,
24+
) -> *mut c_void;
25+
26+
/// Drop a previously created index
27+
///
28+
/// Not called if any other operation against the index may be in flight or started.
29+
#[unsafe(no_mangle)]
30+
extern "C" fn drop_index(
31+
context: u64,
32+
index: *const c_void
33+
);
34+
35+
/// Insert a vector into an index.
36+
///
37+
/// Returns true if the vector is added, false if it is not.
38+
///
39+
/// False may result from the vector already being in the index, or writes failing.
40+
///
41+
/// Note that insert has to be aware of quantizer weirdness, if buffering has to happen it happens here. If we transition from not-quantizing to quantizing, it also has to happen here.
42+
///
43+
/// For now, attribute_data/attribute_len can be ignored - just want space for them.
44+
#[unsafe(no_mangle)]
45+
extern "C" fn insert(
46+
context: u64,
47+
index: *const c_void,
48+
id_data: *const u8,
49+
id_len: usize,
50+
vector_data: *const f32,
51+
vector_len: usize,
52+
attribute_data: *const u8,
53+
attribute_len: usize
54+
) -> bool;
55+
56+
/// Update attribute data on a vector already in the index.
57+
///
58+
/// To implement VSETATTR (https://redis.io/docs/latest/commands/vsetattr/).
59+
///
60+
/// We can skip implementing this for now since we don't need filters yet, just needs to be spec'd.
61+
///
62+
/// Return true if vector was in index and attribute was updated (even if attribute did not change), false otherwise.
63+
#[unsafe(no_mangle)]
64+
extern "C" fn set_attribute(
65+
context: u64,
66+
index: *const c_void,
67+
id_data: *const u8,
68+
id_len: usize,
69+
attribute_data: *const u8,
70+
attribute_len: usize
71+
) -> bool;
72+
73+
/// Find similar vectors, takes parameters of VSIM (https://redis.io/docs/latest/commands/vsim/) and maps to a reasonable interpretation.
74+
///
75+
/// Works with vector values.
76+
///
77+
/// vector_data is unquantized, vector_len will always match dimensions from create_index.
78+
///
79+
/// delta will be [0, 1].
80+
///
81+
/// Maximum number of results is indicated by output_distances_len, elements are i32 length prefixed in byte blobs in output_ids.
82+
///
83+
/// distances are [0, 1].
84+
///
85+
/// Returns number of results, sets continuation to non-zero if there are more to fetch.
86+
///
87+
/// Filtering can be ignored for now, just reserving space in FFI.
88+
///
89+
/// Various search & effort values can be ignored for now, but will eventually be mapped to something sensible. Exist for compat with Redis.
90+
#[unsafe(no_mangle)]
91+
extern "C" fn search_vector(
92+
context: u64,
93+
index_ptr: *const c_void,
94+
vector_data: *const f32,
95+
vector_len: usize,
96+
delta: float,
97+
search_exploration_factor: i32,
98+
filter_data: *const u8,
99+
filter_len: usize,
100+
max_filtering_effort: usize,
101+
output_ids: *mut u8,
102+
output_ids_len: usize,
103+
output_distances: *mut f32,
104+
output_distances_len: usize,
105+
continuation: *mut c_void,
106+
) -> i32;
107+
108+
109+
/// Find similar vectors, takes parameters of VSIM (https://redis.io/docs/latest/commands/vsim/) and maps to a reasonable interpretation.
110+
///
111+
/// Works with item id
112+
///
113+
/// delta will be [0, 1].
114+
///
115+
/// Maximum number of results is indicated by output_distances_len, elements are i32 length prefixed in byte blobs in output_ids.
116+
///
117+
/// distances are [0, 1].
118+
///
119+
/// Returns number of results.
120+
///
121+
/// Filtering can be ignored for now, just reserving space in FFI.
122+
///
123+
/// Various search & effort values can be ignored for now, but will eventually be mapped to something sensible. Exist for compat with Redis.
124+
#[unsafe(no_mangle)]
125+
extern "C" fn search_element(
126+
context: u64,
127+
index_ptr: *const c_void,
128+
id_data: *const u8,
129+
id_length: usize,
130+
delta: float,
131+
search_exploration_factor: i32,
132+
filter_data: *const u8,
133+
filter_len: usize,
134+
max_filtering_effort: i32,
135+
output_ids: *mut u8,
136+
output_ids_len: usize,
137+
output_distances: *mut f32,
138+
output_distances_len: usize,
139+
continuation: *mut c_void
140+
) -> i32;
141+
142+
/// Continues fetching results if not all were available after a call to search_xxx
143+
///
144+
/// Returns number of results placed in output_xxx
145+
///
146+
/// Sets new_continuation to non-zero if even more results are available.
147+
#[unsafe(no_mangle)]
148+
extern "C" fn continue_search(
149+
context: u64,
150+
index_ptr: *const c_void,
151+
continuation: usize,
152+
output_ids: *mut u8,
153+
output_ids_len: usize,
154+
output_distances: *mut f32,
155+
output_distances_len: usize,
156+
new_continuation: *mut c_void
157+
) -> i32;
158+
159+
/// Remove vector from index.
160+
///
161+
/// For implementing VREM (https://redis.io/docs/latest/commands/vrem/).
162+
///
163+
/// Returns true if element was removed from index.
164+
#[unsafe(no_mangle)]
165+
extern "C" fn delete(
166+
context: u64,
167+
index_ptr: *const c_void,
168+
vector_data: *const u8,
169+
vector_len: usize
170+
) -> bool;
171+
172+
/// Return number of vectors stored in index.
173+
///
174+
/// Equivalent to VCARD (https://redis.io/docs/latest/commands/vcard/) can be approximate, must be fast.
175+
#[unsafe(no_mangle)]
176+
extern "C" fn card(
177+
context: u64,
178+
index_ptr: *const c_void,
179+
) -> u64;
180+
181+
182+
/// Check if a vector exists in the index.
183+
///
184+
/// For implementing VISMEMBER - checks whether a vector with the given id is present in the index.
185+
///
186+
/// Returns true if the vector exists in the index, false otherwise.
187+
#[unsafe(no_mangle)]
188+
extern "C" fn has_vector(
189+
context: u64,
190+
index_ptr: *const c_void,
191+
id_data: *const u8,
192+
id_len: usize,
193+
) -> bool;
194+
195+
// To inspect neighbor lists and vector data, Garnet just has to be aware of the format - not a big deal, no need for FFI.

diskann-garnet/src/alloc.rs

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
use diskann_quantization::alloc::{AllocatorCore, AllocatorError, GlobalAllocator};
2+
use std::ptr::NonNull;
3+
4+
/// Custom allocator that over-aligns to 8 bytes. This is needed since Garnet will hand us byte slices for f32 data
5+
/// that may be unaligned, so we need an allocator to make owned, aligned byte containers.
6+
#[derive(Debug, Clone, Copy)]
7+
pub(crate) struct AlignToEight;
8+
9+
unsafe impl AllocatorCore for AlignToEight {
10+
#[inline]
11+
fn allocate(&self, layout: std::alloc::Layout) -> Result<NonNull<[u8]>, AllocatorError> {
12+
// Bump up the alignment.
13+
let layout = layout.align_to(8).map_err(|_| AllocatorError)?;
14+
GlobalAllocator.allocate(layout)
15+
}
16+
17+
#[inline]
18+
unsafe fn deallocate(&self, ptr: NonNull<[u8]>, layout: std::alloc::Layout) {
19+
// Lint: The given `layout` **should** be the same as that passed to `allocate`,
20+
// which must have succeeded for the pointer to be valid in the first place.
21+
#[allow(clippy::expect_used)]
22+
let layout = layout.align_to(8).expect("invalid layout provided");
23+
unsafe { GlobalAllocator.deallocate(ptr, layout) }
24+
}
25+
}

0 commit comments

Comments
 (0)