faiss_next/
result.rs

1use crate::idx::Idx;
2
3#[derive(Debug, Clone)]
4pub struct SearchResult {
5    pub distances: Vec<f32>,
6    pub labels: Vec<Idx>,
7}
8
9impl SearchResult {
10    pub fn new(distances: Vec<f32>, labels: Vec<Idx>) -> Self {
11        Self { distances, labels }
12    }
13
14    pub fn len(&self) -> usize {
15        self.labels.len()
16    }
17
18    pub fn is_empty(&self) -> bool {
19        self.labels.is_empty()
20    }
21
22    pub fn n_queries(&self) -> usize {
23        if self.labels.is_empty() {
24            0
25        } else {
26            self.distances.len() / self.labels.len() * self.labels.len()
27        }
28    }
29
30    pub fn iter(&self) -> impl Iterator<Item = (Idx, f32)> + '_ {
31        self.labels
32            .iter()
33            .copied()
34            .zip(self.distances.iter().copied())
35    }
36
37    pub fn labels(&self) -> impl Iterator<Item = Option<u64>> + '_ {
38        self.labels.iter().map(|&l| l.get())
39    }
40
41    pub fn get(&self, query_idx: usize, k: usize) -> Option<(Vec<Idx>, Vec<f32>)> {
42        let start = query_idx * k;
43        let end = start + k;
44        if end <= self.labels.len() && end <= self.distances.len() {
45            Some((
46                self.labels[start..end].to_vec(),
47                self.distances[start..end].to_vec(),
48            ))
49        } else {
50            None
51        }
52    }
53
54    pub fn first_label(&self) -> Option<u64> {
55        self.labels.first().and_then(|&l| l.get())
56    }
57
58    pub fn first_distance(&self) -> Option<f32> {
59        self.distances.first().copied()
60    }
61}
62
63#[derive(Debug, Clone)]
64pub struct RangeSearchResult {
65    pub labels: Vec<Idx>,
66    pub distances: Vec<f32>,
67    pub lims: Vec<usize>,
68}
69
70impl RangeSearchResult {
71    pub fn new(labels: Vec<Idx>, distances: Vec<f32>, lims: Vec<usize>) -> Self {
72        Self {
73            labels,
74            distances,
75            lims,
76        }
77    }
78
79    pub fn n_queries(&self) -> usize {
80        self.lims.len().saturating_sub(1)
81    }
82
83    pub fn total_results(&self) -> usize {
84        self.labels.len()
85    }
86
87    pub fn get(&self, query_idx: usize) -> Option<(Vec<Idx>, Vec<f32>)> {
88        let start = *self.lims.get(query_idx)?;
89        let end = *self.lims.get(query_idx + 1)?;
90        Some((
91            self.labels[start..end].to_vec(),
92            self.distances[start..end].to_vec(),
93        ))
94    }
95
96    pub fn iter(&self) -> RangeSearchResultIter<'_> {
97        RangeSearchResultIter {
98            result: self,
99            query_idx: 0,
100        }
101    }
102}
103
104pub struct RangeSearchResultIter<'a> {
105    result: &'a RangeSearchResult,
106    query_idx: usize,
107}
108
109impl<'a> Iterator for RangeSearchResultIter<'a> {
110    type Item = (Vec<Idx>, Vec<f32>);
111
112    fn next(&mut self) -> Option<Self::Item> {
113        let result = self.result.get(self.query_idx)?;
114        self.query_idx += 1;
115        Some(result)
116    }
117}
118
119#[derive(Debug, Clone)]
120pub struct BinarySearchResult {
121    pub distances: Vec<i32>,
122    pub labels: Vec<Idx>,
123}
124
125impl BinarySearchResult {
126    pub fn new(distances: Vec<i32>, labels: Vec<Idx>) -> Self {
127        Self { distances, labels }
128    }
129
130    pub fn len(&self) -> usize {
131        self.labels.len()
132    }
133
134    pub fn is_empty(&self) -> bool {
135        self.labels.is_empty()
136    }
137
138    pub fn iter(&self) -> impl Iterator<Item = (Idx, i32)> + '_ {
139        self.labels
140            .iter()
141            .copied()
142            .zip(self.distances.iter().copied())
143    }
144
145    pub fn get(&self, query_idx: usize, k: usize) -> Option<(Vec<Idx>, Vec<i32>)> {
146        let start = query_idx * k;
147        let end = start + k;
148        if end <= self.labels.len() && end <= self.distances.len() {
149            Some((
150                self.labels[start..end].to_vec(),
151                self.distances[start..end].to_vec(),
152            ))
153        } else {
154            None
155        }
156    }
157}
faiss_next/result.rs

faiss_next/
result.rs
