faiss_next/index/
ivf.rs

1use std::ptr;
2
3use faiss_next_sys::{self, FaissIndex, FaissIndexIVF, FaissIndexIVFFlat};
4
5use crate::error::{check_return_code, Error, Result};
6use crate::index::native::InnerPtr;
7use crate::index::traits::{Index, IvfIndex};
8
9pub struct IndexIVFFlat {
10    inner: InnerPtr<FaissIndexIVFFlat>,
11}
12
13impl IndexIVFFlat {
14    pub fn new(quantizer: super::IndexFlat, nlist: usize) -> Result<Self> {
15        let d = quantizer.d();
16        let quantizer_ptr = quantizer.inner_ptr();
17
18        unsafe {
19            let mut inner = ptr::null_mut();
20            check_return_code(faiss_next_sys::faiss_IndexIVFFlat_new_with(
21                &mut inner,
22                quantizer_ptr as *mut FaissIndex,
23                d as usize,
24                nlist,
25            ))?;
26
27            std::mem::forget(quantizer);
28
29            Ok(Self {
30                inner: InnerPtr::new(inner)?,
31            })
32        }
33    }
34
35    pub fn from_index(index: super::IndexImpl) -> Result<Self> {
36        unsafe {
37            let ivf_ptr = faiss_next_sys::faiss_IndexIVFFlat_cast(index.inner_ptr());
38            if ivf_ptr.is_null() {
39                return Err(Error::invalid_cast(
40                    "IndexIVFFlat",
41                    "index is not an IVFFlat index",
42                ));
43            }
44            std::mem::forget(index);
45            Ok(Self {
46                inner: InnerPtr::new(ivf_ptr)?,
47            })
48        }
49    }
50
51    pub fn quantizer(&self) -> *mut FaissIndex {
52        unsafe { faiss_next_sys::faiss_IndexIVFFlat_quantizer(self.inner.as_ptr()) }
53    }
54}
55
56impl Index for IndexIVFFlat {
57    fn inner_ptr(&self) -> *mut FaissIndex {
58        self.inner.as_ptr() as *mut FaissIndex
59    }
60}
61
62impl IvfIndex for IndexIVFFlat {
63    fn nlist(&self) -> usize {
64        unsafe { faiss_next_sys::faiss_IndexIVFFlat_nlist(self.inner.as_ptr()) }
65    }
66
67    fn nprobe(&self) -> usize {
68        unsafe { faiss_next_sys::faiss_IndexIVFFlat_nprobe(self.inner.as_ptr()) }
69    }
70
71    fn set_nprobe(&mut self, nprobe: usize) {
72        unsafe { faiss_next_sys::faiss_IndexIVFFlat_set_nprobe(self.inner.as_ptr(), nprobe) }
73    }
74}
75
76impl Drop for IndexIVFFlat {
77    fn drop(&mut self) {
78        tracing::trace!("dropping IndexIVFFlat");
79        unsafe {
80            faiss_next_sys::faiss_IndexIVFFlat_free(self.inner.as_ptr());
81        }
82    }
83}
84
85unsafe impl Send for IndexIVFFlat {}
86unsafe impl Sync for IndexIVFFlat {}
87
88pub struct IndexIVF {
89    inner: InnerPtr<FaissIndexIVF>,
90}
91
92impl IndexIVF {
93    pub fn from_index(index: super::IndexImpl) -> Result<Self> {
94        unsafe {
95            let ivf_ptr = faiss_next_sys::faiss_IndexIVF_cast(index.inner_ptr());
96            if ivf_ptr.is_null() {
97                return Err(Error::invalid_cast("IndexIVF", "index is not an IVF index"));
98            }
99            std::mem::forget(index);
100            Ok(Self {
101                inner: InnerPtr::new(ivf_ptr)?,
102            })
103        }
104    }
105}
106
107impl Index for IndexIVF {
108    fn inner_ptr(&self) -> *mut FaissIndex {
109        self.inner.as_ptr() as *mut FaissIndex
110    }
111}
112
113impl IvfIndex for IndexIVF {
114    fn nlist(&self) -> usize {
115        unsafe { faiss_next_sys::faiss_IndexIVF_nlist(self.inner.as_ptr()) }
116    }
117
118    fn nprobe(&self) -> usize {
119        unsafe { faiss_next_sys::faiss_IndexIVF_nprobe(self.inner.as_ptr()) }
120    }
121
122    fn set_nprobe(&mut self, nprobe: usize) {
123        unsafe { faiss_next_sys::faiss_IndexIVF_set_nprobe(self.inner.as_ptr(), nprobe) }
124    }
125}
126
127impl Drop for IndexIVF {
128    fn drop(&mut self) {
129        tracing::trace!("dropping IndexIVF");
130        unsafe {
131            faiss_next_sys::faiss_IndexIVF_free(self.inner.as_ptr());
132        }
133    }
134}
135
136unsafe impl Send for IndexIVF {}
137unsafe impl Sync for IndexIVF {}
faiss_next/index/ivf.rs

faiss_next/index/
ivf.rs
