diskann-providers 0.57.0

DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases.
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

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use diskann::neighbor::{self, Neighbor};
use diskann_utils::views::MatrixView;

/// Compute the ground truth for a small dataset.
///
/// Counter intuitively, this function puts nearest neighbors **at the end** of the
/// vector rather than the beginning.
///
/// This allows filtering by `is_match` to be much more efficient because it decreases
/// the number of elements that have to be moved.
pub fn groundtruth<T, F>(data: MatrixView<T>, query: &[T], f: F) -> Vec<Neighbor<u32>>
where
    F: Fn(&[T], &[T]) -> f32,
{
    let mut results: Vec<_> = data
        .row_iter()
        .enumerate()
        .map(|(i, row)| Neighbor::new(i as u32, f(row, query)))
        .collect();

    results.sort_unstable_by(neighbor::ord::reverse(neighbor::ord::fast_distance));
    results
}

/// Decide if `neighbor` is the next nearest neighbor in the groundtruth slice, using
/// the `margin` factor to accommodate ties. If it *is* a match, return the position
/// of the matching entry in `groundtruth`.
///
/// # Panics
///
/// Panics if the `neighbor.id` cannot be found in `groundtruth`.
#[cfg(test)]
pub fn is_match(
    groundtruth: &[Neighbor<u32>],
    neighbor: Neighbor<u32>,
    margin: f32,
) -> Option<usize> {
    for i in (0..groundtruth.len()).rev() {
        // Check if the distance matches.
        let gt = groundtruth[i];
        if (gt.distance() - neighbor.distance()).abs() > margin {
            return None;
        }
        if gt.id() == neighbor.id() {
            return Some(i);
        }
    }
    panic!(
        "could not find neighbor {:?}. Remaining: {:?}",
        neighbor, groundtruth
    );
}

/// Asserts that the top-k results exactly match the ground truth
///
/// For each of the top-k results, this function verifies that both the distance and ID
/// match exactly with what's expected in the ground truth.
pub fn assert_top_k_exactly_match(
    query_id: usize,
    gt: &[Neighbor<u32>],
    ids: &[u32],
    distances: &[f32],
    top_k: usize,
) {
    for i in 0..top_k {
        let neighbor = gt[gt.len() - 1 - i];
        assert_eq!(
            *neighbor.distance(),
            distances[i],
            "failed on query {} for result {}",
            query_id,
            i
        );
        assert_eq!(
            *neighbor.id(),
            ids[i],
            "failed on query {} for result {}",
            query_id,
            i
        );
    }
}

/// Asserts that the range results exactly match the ground truth
///
/// For each of the range results, this function verifies that both the distance and ID
/// match exactly with what's expected in the ground truth.
#[cfg(test)]
pub fn assert_range_results_exactly_match(
    query_id: usize,
    gt: &[Neighbor<u32>],
    ids: &[u32],
    radius: f32,
    inner_radius: Option<f32>,
) {
    let gt_ids = if let Some(inner_radius) = inner_radius {
        gt.iter()
            .filter(|nbh| *nbh.distance() >= inner_radius && *nbh.distance() <= radius)
            .map(|nbh| *nbh.id())
            .collect::<Vec<_>>()
    } else {
        gt.iter()
            .filter(|nbh| *nbh.distance() <= radius)
            .map(|nbh| *nbh.id())
            .collect::<Vec<_>>()
    };
    if ids.iter().any(|id| !gt_ids.contains(id)) {
        panic!(
            "query {}: found ids {:?} in range search with radius {}, inner radius {}, but expected {:?}",
            query_id,
            ids,
            radius,
            inner_radius.unwrap_or(f32::MIN),
            gt_ids
        );
    }
}