feox-ann 0.1.0

Dependency-free HNSW approximate nearest neighbor index with deterministic, reproducible builds
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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use crate::{AnnConfig, AnnIndex, AnnQuery};

fn sample_records(count: usize) -> Vec<(String, Vec<f32>)> {
    (0..count)
        .map(|i| {
            let angle = i as f32 * 0.037;
            (
                format!("record-{i:05}"),
                vec![
                    angle.sin(),
                    angle.cos(),
                    (angle * 2.0).sin(),
                    (angle * 2.0).cos(),
                ],
            )
        })
        .collect()
}

fn query_results(index: &AnnIndex, seed: f32) -> Vec<crate::AnnCandidate> {
    let angle = seed * 0.17;
    index
        .query(AnnQuery {
            vector: &[
                angle.sin(),
                angle.cos(),
                (angle * 2.0).sin(),
                (angle * 2.0).cos(),
            ],
            top_k: 10,
            ef_search: Some(64),
            filter: None,
        })
        .unwrap()
}

#[test]
fn bulk_load_is_thread_count_invariant() {
    let records = sample_records(500);

    let mut single = AnnIndex::new(AnnConfig::for_dimensions(4)).unwrap();
    single.bulk_load(records.clone(), 1).unwrap();

    let mut parallel = AnnIndex::new(AnnConfig::for_dimensions(4)).unwrap();
    parallel.bulk_load(records, 8).unwrap();

    assert_eq!(single.len(), parallel.len());
    for i in 0..25 {
        assert_eq!(
            query_results(&single, i as f32),
            query_results(&parallel, i as f32),
            "thread count must not change the graph"
        );
    }
}

#[test]
fn bulk_load_finds_true_neighbors() {
    let records = sample_records(1_000);
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(4)).unwrap();
    index.bulk_load(records, 4).unwrap();
    assert_eq!(index.len(), 1_000);

    let angle = 123_f32 * 0.037;
    let matches = index
        .query(AnnQuery {
            vector: &[
                angle.sin(),
                angle.cos(),
                (angle * 2.0).sin(),
                (angle * 2.0).cos(),
            ],
            top_k: 1,
            ef_search: Some(128),
            filter: None,
        })
        .unwrap();
    assert_eq!(matches[0].id, "record-00123");
    assert!(matches[0].score > 0.9999);
}

#[test]
fn bulk_load_upserts_duplicate_ids() {
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(2)).unwrap();
    index
        .bulk_load(
            vec![
                ("item".to_string(), vec![1.0, 0.0]),
                ("other".to_string(), vec![0.5, 0.5]),
                ("item".to_string(), vec![0.0, 1.0]),
            ],
            2,
        )
        .unwrap();

    assert_eq!(index.len(), 2);
    let matches = index
        .query(AnnQuery {
            vector: &[0.0, 1.0],
            top_k: 1,
            ef_search: Some(8),
            filter: None,
        })
        .unwrap();
    assert_eq!(matches[0].id, "item");
    assert!(matches[0].score > 0.99);
}

#[test]
fn concurrent_queries_share_the_index() {
    let records = sample_records(300);
    let mut index = AnnIndex::new(AnnConfig::for_dimensions(4)).unwrap();
    index.bulk_load(records, 4).unwrap();

    let expected: Vec<_> = (0..8).map(|i| query_results(&index, i as f32)).collect();
    std::thread::scope(|scope| {
        for (i, expected) in expected.iter().enumerate() {
            let index = &index;
            scope.spawn(move || {
                for _ in 0..50 {
                    assert_eq!(&query_results(index, i as f32), expected);
                }
            });
        }
    });
}