DiskANN in DefraDB

DefraDB is a peer-to-peer document database with a GraphQL interface, and defradb.rs is its Rust implementation. It runs on servers, embedded devices and in the browser, on top of regolith, our LSM-tree storage engine. In October it got DISKANN, a vector index that keeps working when the data doesn't fit in memory:

type Doc {
embedding: [Float!] @index(vector: {dimensions: 768, alg: DISKANN})
}

Our other vector indexes (Flat, HNSW and IVF-PQ) store their structures as key-value records in the database, so they're transactional and replicate like everything else. But every record a search touches costs a database read, where an in-memory library pays a pointer chase. DiskANN is built to touch few records per search.

What DiskANN is

DiskANN is a graph index from Microsoft Research. Each vector is a node with at most R neighbours. A search walks greedily toward the query, keeping the best L candidates seen so far. A bigger L buys recall with more reads. Building runs the same walk, with list size L_build, to find a new node's neighbours, then prunes them. An alpha above 1 keeps some long edges so a walk gets close in few hops, which matters when each hop is a disk read.

Neighbours are ranked on compressed vectors. Product quantization (PQ) splits a vector into chunks and replaces each with the nearest of 256 trained centroids, giving a code of a few bytes. Full vectors are read only at the end, to re-rank the best candidates. Inserts and deletes follow the FreshDiskANN paper.

How we built it

Each record carries its neighbours' codes. Classic DiskANN keeps every PQ code in RAM and stores each full vector beside its node on disk. We store a node's own code plus its neighbours' ids and codes in one record of about 7 KB: 96 bytes per code plus 8 per id, for 67 edges on average at R = 64, since a node can hold up to 1.3·R edges between prunes. One read ranks every edge out of a node. LM-DiskANN (Pan, Sun and Yu, 2023) and libSQL's vector index work much the same way. A search keeps only the PQ codebooks and its candidate list in memory, and a walk stops after max(4·L, 32) nodes, so that list stays small at any corpus size.

The cost is disk. Each code is copied into every record that points at it, about 67 times on average, so the records come to 6.6 GiB at 1M vectors, more than twice the 2.9 GiB of raw float32 vectors.

Small indexes scan. Below 1,024 vectors a search compares the query with every vector. The write that crosses that threshold trains the quantizer once, on a random sample of up to 128 MiB, and builds the graph.

Inserts don't touch the graph. After training, an insert writes its vector and code to a pending list, so concurrent writers don't conflict. Searches walk the graph, scan the pending codes and re-rank both. A background sweep later merges pending vectors into the graph and removes deleted nodes, as in FreshDiskANN. The pending list is stored in the database, so it's transactional and survives a crash. On a 20k × 384 benchmark, an insert after training dropped from 2.11 ms to 0.041 ms. The merge still does the graph work later, and until it runs every search pays for the scan: with 2,000 vectors pending, a query took 0.76 ms against 0.62 ms after the merge. At 1M with 100,000 vectors pending, a query took 64 ms against 4.9 ms, and QPS fell from 1,233 to 81.

Two fixes bring that down. Summing PQ distances in four accumulators instead of one took the query to 49 ms. Then every search on a node began sharing an in-memory copy of the pending codes, caught up from a log as inserts commit, instead of scanning storage. That took it to 24 ms at 180 QPS. Keeping the four sums in registers rather than in an array on the stack then took it to 21 ms at 216 QPS, with recall unchanged throughout. The copy holds up to 64 MiB, about 640,000 pending vectors at 96-byte codes, and past that searches scan as before. The log is the pending list itself, so an insert still makes one pending write. With 16 or 64 concurrent writers, inserts run about 2.5% slower than before, with no conflicts.

Resident codes are opt-in. With diskann: {residentCodes: true}, each code is stored once in an in-memory array and records keep only neighbour ids. Records shrink 13× (517 MiB against 6,774 MiB at 1M), for 91 MiB of codes in memory plus an index from node id to slot.

One build fix: a back-edge to a node already at R edges re-pruned it on every insert, and pruning had grown to three quarters of build time. Letting nodes grow to 1.3·R before pruning back to R, as Microsoft's implementation does, fixed it.

The numbers

We ran every index on the Cohere 1M dataset from VectorDBBench: 1M Wikipedia passages at 768 dimensions, cosine distance, top 10 for 1,000 queries. Our exact top 10 matches the published ground truth on every query. The box is an AWS i4i.2xlarge: 4 cores with 2 hyperthreads each, 61 GiB RAM, and local NVMe that does a 4 KiB random read in 0.11 ms at queue depth one.

Our HNSW stores each node, full vector included, as a regolith record and reads one per neighbour it visits: an HNSW inside a key-value store. For scale, the last two rows are usearch, a fast in-memory HNSW library, on the same box and data.

Each index runs in a fresh process. Cold is the first pass over the queries after dropping the OS page cache. Warm means the whole store was read into the page cache, then 2,000 warm-up queries from outside the test set ran before measuring. Latency is one query at a time, and QPS is 8 threads with one search in flight each. The 1M graphs were built in one pass, with the quantizer trained on a 128 MiB sample of the whole corpus.

A live database trains on the first 1,024 vectors instead, and everything after goes through pending lists and merges. At 200k vectors, recall at L = 100 was 0.977 against 0.982 for the one-pass build, whether the graph merged after every insert or in batches of 1,024 or 16,384. At 1M, merging in batches of 16,384, it was 0.979 against 0.981, with warm latency at 5.2 against 4.9 ms and the same QPS (1,239 against 1,233).

DiskANN: R 64, L_build 100, alpha 1.2, PQ 96 bytes per vector (the schema defaults). HNSW: M 16, ef_construction 128. IVF-PQ: 4,000 lists, PQ 96 bytes, 10 re-rank candidates per result. usearch 2: ef_construction 128 (M 16) and 200 (M 32), f32 vectors, index 4.3 and 4.6 GiB. Cold DiskANN at L = 100 is the median of four runs (27 to 32 ms). Other rows are single runs, and warm repeated within 0.3 ms. Server profile, with regolith's 512 MiB block cache.

At about 0.98 recall, warm DiskANN answers in 4.9 ms against 56 ms for our HNSW, with 12× the QPS. A DiskANN query at L = 100 misses regolith's block cache about 190 times and an HNSW query at ef = 256 misses about 3,000 times, once per neighbour it scores. usearch is still faster: 2.3 ms at 0.988 recall, against 4.9 ms for DiskANN here and 3.8 ms on the embedded profile. Both need their data in RAM for that, 4.6 GiB for usearch and the whole store for DiskANN.

A full preload is a lot to ask. Starting cold with only the 2,000 warm-up queries, DiskANN at L = 100 was still at 15.5 ms and resident codes at 9.0 ms. Records of 7 KB spread over 6.6 GiB take far more than 2,000 queries to pull in.

Cold, each hop waits for its record before choosing the next, because there is no beam search yet to overlap the reads. At L = 100 a query reads about 106 graph records, then the full vectors of its 100 best candidates for re-ranking, which classic DiskANN avoids by storing vectors beside nodes. At the device, a cold query issued 161 reads, about 0.19 ms each at 30 ms a query, against the drive's 0.11 ms for a single 4 KiB read and 0.12 ms for 8 KiB. Resident codes issued 103: their 517 MiB of records mostly get cached during the pass, so later queries go to disk mainly for full vectors.

Less memory than data

To take RAM away we ran each index in a cgroup capped at 2 GiB, 1 GiB and 512 MiB, swap off. The cap counts the page cache too, so it covers everything the process can keep in memory. Every store is bigger than every cap: about 3 GiB for HNSW and IVF-PQ, and 6.6 GiB of DiskANN records. Both DiskANN layouts share one 12.8 GiB store. Rerun in separate stores of 9.6 and 3.4 GiB, no latency or QPS cell below moved by more than 7%. All runs use regolith's embedded profile, which has no block cache, and warm up on the 2,000 unseen queries first.

Disk reads are counted at the device and averaged over the run. We didn't count them for HNSW or usearch. p99 stays within 1.4× of p50 in every DiskANN and IVF-PQ cell. HNSW warmed on the test queries themselves, which can only flatter it (at 0.4 to 0.6 s a query we didn't rerun it), and its 2 GiB cell used the server profile. usearch served a separate M 32 build, saved to a 3.1 GiB file, through view(), which maps the file instead of loading it. It warmed on 500 unseen queries, and its QPS is one pass over 8 threads. Uncapped, the mapped index answers in 2.5 ms at 0.985 recall.

Both HNSWs collapse once their scattered reads go to disk. Ours is 7 to 15× slower than with the store cached on the same profile, and usearch goes from 2.5 ms to 330 ms at 2 GiB and 859 ms at 512 MiB. DiskANN touches about 200 small records a query and stays usable. IVF-PQ degrades least, since it reads a few lists front to back. At matching recall, about 0.96, DiskANN at L = 50 answers in a quarter of IVF-PQ's time with nearly 4× its QPS, at every cap.

Resident codes do best when memory is short. At 2 GiB their records stay cached, and a query needs half the reads and half the time of the default layout. The lead shrinks with the cap, to about 10% at 512 MiB, where the 205 MiB process shares the same limit. Warm, they lose about 1 ms a query: each hop looks up 64 neighbour codes in an in-memory hash index at about 110 ns each, where the default layout already has them in the record it just read. So the default suits stores that fit in memory, and resident codes suit stores that don't.

Memory

On the server profile every index peaks at 570 to 750 MiB resident, about 512 MiB of it regolith's block cache. To see the index itself we reopened the stores on the embedded profile, which has none. Peak RSS here leaves out the OS page cache.

Resident codes add about 140 MiB at 1M, for the codes and their id index. Most of DiskANN's 66 MiB is regolith's SSTable metadata: loading SSTable indexes and filters accounts for 36 MiB of a 45 MiB heap peak, 16 MiB of it pinned filters. Under 10 MiB is left for everything else, the 1,000 test queries included. The metadata grows with the store, which is why 200k vectors take 22 MiB. The DiskANN store holds both layouts in 12.8 GiB against 3 GiB for HNSW and IVF-PQ, so those rows aren't like for like. In separate stores of 9.6 and 3.4 GiB, the two layouts peak at 55 and 176 MiB.

DiskANN and HNSW are also faster warm on the embedded profile than on the server one: 3.8 against 4.9 ms for DiskANN at L = 100, and 38 against 56 ms for HNSW, both preloaded. IVF-PQ goes the other way, 54 against 48 ms. For the graph indexes, once everything is in the page cache, a read through regolith's block cache costs more than one the page cache serves. We haven't profiled why.

Building it

The benchmark builds each index in memory, one insert at a time on one task, then copies it into regolith, so these are harness build times. At 1M, DiskANN took 94 minutes with a 37.6 GiB peak, HNSW 80 minutes at 24.6 GiB, and IVF-PQ 17 minutes at 24.6 GiB. The peaks come from the harness holding the corpus, the in-memory index and copies of both. usearch built in 4 to 9 minutes on 8 threads. Training the quantizer and building the first graph at 1,024 vectors took 4.0 s, and the write that crosses the threshold waits for it.

A live database builds through merges instead, here one every 16,384 inserts. On main a live 100k build took 345 s. A stack of seven PRs in review (#2140 up to #2196) brings that to 231 s, with recall unchanged at 0.981. With the stack, a live 1M build takes 40 to 42 minutes.

The code is in crates/db/src/index/vector/engine/diskann/ in defradb.rs. The benchmark is crates/db/tests/vector/vector_disk_cost.rs (#2054, merged in 36e21d5), with the exact commands in its doc comment, and tools/vector-bench has the corpus converter, the memory cap and the usearch reference.