embedvec 0.8.0

Fast, lightweight, in-process vector database with HNSW indexing, E8/H4 lattice quantization (up to 24.8x compression), metadata filtering, and PyO3 bindings
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

// Measure the *resident* memory footprint of an in-memory embedvec index.
//
// embedvec keeps the HNSW graph + vector cache in RAM, so this reports the real
// per-vector cost (vectors + graph + metadata + allocator overhead), not just
// the on-disk vector bytes.
//
//   cargo run --release --example scale_memory -- <raw|h4|e8> <n> [dim]
//
// Example: cargo run --release --example scale_memory -- h4 1000000

use embedvec::{Distance, EmbedVec, Quantization};
use std::time::Instant;

fn rss_bytes() -> u64 {
    memory_stats::memory_stats()
        .map(|m| m.physical_mem as u64)
        .unwrap_or(0)
}

#[tokio::main]
async fn main() {
    let args: Vec<String> = std::env::args().collect();
    let mode = args.get(1).map(|s| s.as_str()).unwrap_or("h4");
    let n: usize = args.get(2).and_then(|s| s.parse().ok()).unwrap_or(100_000);
    let dim: usize = args.get(3).and_then(|s| s.parse().ok()).unwrap_or(768);

    let quant = match mode {
        "raw" | "none" => Quantization::None,
        "h4" => Quantization::h4_default(),
        "e8" => Quantization::e8_default(),
        other => {
            eprintln!("unknown mode {other}; use raw|h4|e8");
            return;
        }
    };

    let base = rss_bytes();

    let mut db = EmbedVec::builder()
        .dimension(dim)
        .metric(Distance::Cosine)
        .m(16)
        .ef_construction(100)
        .quantization(quant)
        .build()
        .await
        .unwrap();

    // Deterministic pseudo-random vectors, generated on the fly to avoid holding
    // a second copy of the whole dataset in RAM while we measure.
    let mut state: u64 = 0x9E3779B97F4A7C15;
    let mut next = || {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        (state as f32 / u64::MAX as f32) * 2.0 - 1.0
    };

    let t = Instant::now();
    let mut v = vec![0.0f32; dim];
    for i in 0..n {
        for x in v.iter_mut() {
            *x = next();
        }
        db.add_internal(&v, serde_json::json!({ "i": i })).unwrap();
    }
    let build = t.elapsed();

    let peak = rss_bytes();
    let used = peak.saturating_sub(base);
    let vec_bytes = db.storage.read().memory_bytes();

    println!(
        "mode={mode:>4} n={n:>9} dim={dim}  build={build:>8.1?}  \
         RSS_total={:.2} GB  RSS_used={:.2} GB  bytes/vec(RSS)={:.0}  vec_bytes/vec={:.0}",
        peak as f64 / 1e9,
        used as f64 / 1e9,
        used as f64 / n as f64,
        vec_bytes as f64 / n as f64,
    );
}