faiss_next/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use faiss_next_sys as sys;
4use std::ffi::CString;
5use std::ptr::{addr_of_mut, null_mut};
6use std::time::Instant;
7use tracing::trace;
8
9/// ## Metric Type
10/// please refer to <https://github.com/facebookresearch/faiss/wiki/MetricType-and-distances>
11pub use sys::FaissMetricType;
12
13/// Error
14#[derive(Debug, thiserror::Error)]
15pub enum Error {
16    #[error("{0}")]
17    NulErr(#[from] std::ffi::NulError),
18    #[error("faiss error,code={},message={}", .code, .message)]
19    Faiss { code: i32, message: String },
20}
21
22impl Error {
23    pub fn from_rc(rc: i32) -> Self {
24        match rc {
25            0 => unimplemented!(),
26            _ => {
27                let message = unsafe { std::ffi::CStr::from_ptr(sys::faiss_get_last_error()) };
28                let message = message
29                    .to_str()
30                    .unwrap_or("unknown error, failed to decode error message from bytes")
31                    .to_string();
32                Error::Faiss { code: rc, message }
33            }
34        }
35    }
36}
37
38pub type Result<T> = std::result::Result<T, Error>;
39
40macro_rules! faiss_rc {
41    ($blk: block) => {
42        match unsafe { $blk } {
43            0 => Ok(()),
44            rc @ _ => Err(Error::from_rc(rc)),
45        }
46    };
47}
48
49/// search result
50#[derive(Debug)]
51pub struct SearchResult {
52    /// - labels of search result
53    pub labels: Vec<i64>,
54    /// - distances of search result
55    pub distances: Vec<f32>,
56}
57
58/// Trait that can return inner pointer of index
59pub trait IndexInner {
60    /// return inner pointer
61    fn inner(&self) -> *mut sys::FaissIndex;
62}
63
64/// Index trait, all index should implement this trait
65pub trait Index: IndexInner {
66    /// add vectors to index, x.len() should be a multiple of d
67    fn add<T: AsRef<[f32]> + ?Sized>(&mut self, x: &T) -> Result<()> {
68        let x = x.as_ref();
69        faiss_rc!({
70            sys::faiss_Index_add(self.inner(), (x.len() / self.d()) as sys::idx_t, x.as_ptr())
71        })?;
72        trace!("add: index={:?}, x.len={}", self.inner(), x.as_ref().len());
73        Ok(())
74    }
75
76    /// search vector against index, `x.len()` should be a multiple of `d`, `k` means top k
77    fn search<T: AsRef<[f32]> + ?Sized>(&self, x: &T, k: usize) -> Result<SearchResult> {
78        let x = x.as_ref();
79        let n = (x.len() / self.d()) as sys::idx_t;
80        let mut labels = vec![0 as sys::idx_t; n as usize * k];
81        let mut distances = vec![0.0f32; n as usize * k];
82        let tm = Instant::now();
83        faiss_rc!({
84            {
85                sys::faiss_Index_search(
86                    self.inner(),
87                    n,
88                    x.as_ptr(),
89                    k as sys::idx_t,
90                    distances.as_mut_ptr(),
91                    labels.as_mut_ptr(),
92                )
93            }
94        })?;
95        trace!(
96            "search index d={}, n={}, k={}, tm={:?}",
97            self.d(),
98            n,
99            k,
100            tm.elapsed()
101        );
102        Ok(SearchResult { labels, distances })
103    }
104
105    /// return dimension of index
106    fn d(&self) -> usize {
107        unsafe { sys::faiss_Index_d(self.inner()) as usize }
108    }
109
110    // return whether index is trained
111    fn is_trained(&self) -> bool {
112        unsafe { sys::faiss_Index_is_trained(self.inner()) != 0 }
113    }
114
115    /// train index when some index impl is used
116    fn train<T: AsRef<[f32]> + ?Sized>(&mut self, x: &T) -> Result<()> {
117        let x = x.as_ref();
118        let n = (x.len() / self.d()) as sys::idx_t;
119        faiss_rc!({ sys::faiss_Index_train(self.inner(), n, x.as_ptr()) })?;
120        Ok(())
121    }
122
123    /// remove feature with [IDSelector](https://faiss.ai/cpp_api/struct/structfaiss_1_1IDSelector.html)
124    fn remove_ids(&mut self, sel: IDSelector) -> Result<usize> {
125        let mut n_removed = 0usize;
126        faiss_rc!({
127            sys::faiss_Index_remove_ids(self.inner(), sel.inner, addr_of_mut!(n_removed))
128        })?;
129        Ok(n_removed)
130    }
131
132    /// save index to disk
133    fn save<P: AsRef<str>>(&self, pth: P) -> Result<()> {
134        let pth = pth.as_ref();
135        let pth = CString::new(pth)?;
136        faiss_rc!({ sys::faiss_write_index_fname(self.inner() as *const _, pth.as_ptr()) })?;
137        todo!()
138    }
139
140    /// load index from disk
141    fn load<P: AsRef<str>>(pth: P) -> Result<CpuIndex> {
142        let pth = pth.as_ref();
143        let pth = CString::new(pth)?;
144        let mut inner = null_mut();
145        faiss_rc!({ sys::faiss_read_index_fname(pth.as_ptr(), 0, addr_of_mut!(inner)) })?;
146        Ok(CpuIndex { inner })
147    }
148}
149
150/// Select ID to delete feature in index [IDSelector](https://faiss.ai/cpp_api/struct/structfaiss_1_1IDSelector.html)
151pub struct IDSelector {
152    pub inner: *mut sys::FaissIDSelector,
153}
154
155impl IDSelector {
156    /// create a selector from batch ids
157    /// TODO: more id selector
158    pub fn batch(ids: &[i64]) -> Result<Self> {
159        let mut inner = null_mut();
160        faiss_rc!({
161            sys::faiss_IDSelectorBatch_new(addr_of_mut!(inner), ids.len(), ids.as_ptr())
162        })?;
163        Ok(Self {
164            inner: inner as *mut _,
165        })
166    }
167}
168
169impl Drop for IDSelector {
170    fn drop(&mut self) {
171        unsafe { sys::faiss_IDSelector_free(self.inner) }
172    }
173}
174
175/// Index use cpu
176/// # Examples
177/// ``` rust
178/// use faiss_next as faiss;
179/// use faiss::Index;
180/// let mut index = faiss::index_factory(128, "Flat", faiss::FaissMetricType::METRIC_L2).expect("failed to create index");
181/// index.add(&vec![0.0;128 * 128]).expect("failed to add feature");
182/// let ret = index.search(&vec![0.0;128], 1).expect("failed to search");
183/// ```
184#[derive(Debug)]
185pub struct CpuIndex {
186    pub inner: *mut sys::FaissIndex,
187}
188
189impl Drop for CpuIndex {
190    fn drop(&mut self) {
191        unsafe {
192            sys::faiss_Index_free(self.inner);
193        }
194        trace!("drop: index={:?}", self.inner);
195    }
196}
197
198impl IndexInner for CpuIndex {
199    fn inner(&self) -> *mut sys::FaissIndex {
200        self.inner
201    }
202}
203
204impl Index for CpuIndex {}
205
206impl CpuIndex {
207    /// create multi gpu index, `devices` is a list of gpu device id, `split` means split index on multiple gpu or not
208    #[cfg(feature = "gpu")]
209    pub fn to_multi_gpu(&self, devices: &[i32], shard: bool) -> Result<gpu::GpuIndex> {
210        let mut p_out = 0 as *mut _;
211        let providers = (0..devices.len())
212            .map(|_| -> Result<_> { gpu::GpuResourcesProvider::new() })
213            .collect::<Result<Vec<_>>>()?;
214        let mut options = 0 as *mut sys::FaissGpuClonerOptions;
215        faiss_rc!({ sys::faiss_GpuClonerOptions_new(addr_of_mut!(options)) })?;
216        if shard {
217            unsafe { sys::faiss_GpuMultipleClonerOptions_set_shard(options, 1) };
218        }
219        let providers_ = providers.iter().map(|p| p.inner).collect::<Vec<_>>();
220        faiss_rc!({
221            sys::faiss_index_cpu_to_gpu_multiple_with_options(
222                providers_.as_ptr(),
223                providers.len(),
224                devices.as_ptr(),
225                devices.len(),
226                self.inner,
227                options,
228                addr_of_mut!(p_out),
229            )
230        })?;
231        Ok(gpu::GpuIndex {
232            inner: p_out,
233            providers: providers,
234        })
235    }
236
237    /// create multi gpu index, `devices` is a list of gpu device id, `split` means split index on multiple gpu or not, cpu index will be dropped
238    #[cfg(feature = "gpu")]
239    pub fn into_multi_gpu(self, devices: &[i32], shard: bool) -> Result<gpu::GpuIndex> {
240        self.to_multi_gpu(devices, shard)
241    }
242
243    /// create gpu index, `device` is gpu device id
244    #[cfg(feature = "gpu")]
245    pub fn to_gpu(&self, device: i32) -> Result<gpu::GpuIndex> {
246        let mut p_out = 0 as *mut _;
247        let provider = gpu::GpuResourcesProvider::new()?;
248        trace!(?provider, "create gpu provider");
249        faiss_rc!({
250            sys::faiss_index_cpu_to_gpu(provider.inner, device, self.inner, addr_of_mut!(p_out))
251        })?;
252        trace!(
253            "into_gpu: from {:?} to index={:?}, device={}",
254            self.inner,
255            p_out,
256            device
257        );
258        Ok(gpu::GpuIndex {
259            inner: p_out,
260            providers: vec![provider],
261        })
262    }
263
264    /// create gpu index, `device` is gpu device id, cpu index will be dropped
265    #[cfg(feature = "gpu")]
266    pub fn into_gpu(self, device: i32) -> Result<gpu::GpuIndex> {
267        self.to_gpu(device)
268    }
269}
270
271impl Clone for CpuIndex {
272    fn clone(&self) -> Self {
273        let mut inner = null_mut();
274        faiss_rc!({ sys::faiss_clone_index(self.inner, addr_of_mut!(inner)) })
275            .unwrap_or_else(|_| panic!("failed to clone index with inner={:?}", self.inner));
276        Self { inner }
277    }
278}
279
280/// helper function to create cpu index, please refer to [doc](https://github.com/facebookresearch/faiss/wiki/The-index-factory) for details of `description` and `metric`
281pub fn index_factory(d: i32, description: &str, metric: FaissMetricType) -> Result<CpuIndex> {
282    let mut p_index = null_mut();
283    let description_ = CString::new(description)?;
284    faiss_rc! {{sys::faiss_index_factory(addr_of_mut!(p_index), d, description_.as_ptr(), metric)}}?;
285    trace!(
286        "index_factory: create index={:?}, description={}, metric={:?}",
287        p_index,
288        description,
289        metric
290    );
291    Ok(CpuIndex { inner: p_index })
292}
293
294///  gpu related module
295#[cfg(feature = "gpu")]
296pub mod gpu {
297    use super::{addr_of_mut, null_mut, sys, Error, Index, IndexInner};
298    use tracing::trace;
299
300    /// gpu index
301    pub struct GpuIndex {
302        /// - gpu resource provider of faiss
303        pub providers: Vec<GpuResourcesProvider>,
304        /// - raw pointer
305        pub inner: *mut sys::FaissGpuIndex,
306    }
307
308    /// gpu resource provider
309    #[derive(Debug)]
310    pub struct GpuResourcesProvider {
311        pub inner: *mut sys::FaissGpuResourcesProvider,
312    }
313
314    impl GpuResourcesProvider {
315        pub fn new() -> super::Result<Self> {
316            let mut inner = null_mut();
317            faiss_rc!({ sys::faiss_StandardGpuResources_new(addr_of_mut!(inner)) })?;
318            trace!(?inner, "create gpu provider");
319            Ok(Self { inner })
320        }
321    }
322
323    impl Drop for GpuResourcesProvider {
324        fn drop(&mut self) {
325            unsafe { sys::faiss_GpuResourcesProvider_free(self.inner) }
326            trace!("drop: gpu provider={:?}", self.inner);
327        }
328    }
329
330    impl IndexInner for GpuIndex {
331        fn inner(&self) -> *mut sys::FaissIndex {
332            self.inner
333        }
334    }
335
336    impl Index for GpuIndex {}
337
338    impl GpuIndex {
339        pub fn to_cpu(&self) -> super::Result<super::CpuIndex> {
340            let mut inner = null_mut();
341            faiss_rc!({ sys::faiss_index_gpu_to_cpu(self.inner, addr_of_mut!(inner)) })?;
342            Ok(super::CpuIndex { inner })
343        }
344
345        pub fn into_cpu(self) -> super::Result<super::CpuIndex> {
346            self.to_cpu()
347        }
348    }
349
350    impl Drop for GpuIndex {
351        fn drop(&mut self) {
352            unsafe {
353                sys::faiss_Index_free(self.inner);
354                trace!("drop: index={:?}", self.inner);
355            }
356        }
357    }
358}
faiss_next/lib.rs

faiss_next/
lib.rs
