faiss-next 0.6.0

Rust bindings for Faiss (Facebook AI Similarity Search)
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
127
128
129
130
131
132
133
134
135
136
137

use std::ptr;

use faiss_next_sys::{self, FaissIndex, FaissIndexIVF, FaissIndexIVFFlat};

use crate::error::{check_return_code, Error, Result};
use crate::index::native::InnerPtr;
use crate::index::traits::{Index, IvfIndex};

pub struct IndexIVFFlat {
    inner: InnerPtr<FaissIndexIVFFlat>,
}

impl IndexIVFFlat {
    pub fn new(quantizer: super::IndexFlat, nlist: usize) -> Result<Self> {
        let d = quantizer.d();
        let quantizer_ptr = quantizer.inner_ptr();

        unsafe {
            let mut inner = ptr::null_mut();
            check_return_code(faiss_next_sys::faiss_IndexIVFFlat_new_with(
                &mut inner,
                quantizer_ptr as *mut FaissIndex,
                d as usize,
                nlist,
            ))?;

            std::mem::forget(quantizer);

            Ok(Self {
                inner: InnerPtr::new(inner)?,
            })
        }
    }

    pub fn from_index(index: super::IndexImpl) -> Result<Self> {
        unsafe {
            let ivf_ptr = faiss_next_sys::faiss_IndexIVFFlat_cast(index.inner_ptr());
            if ivf_ptr.is_null() {
                return Err(Error::invalid_cast(
                    "IndexIVFFlat",
                    "index is not an IVFFlat index",
                ));
            }
            std::mem::forget(index);
            Ok(Self {
                inner: InnerPtr::new(ivf_ptr)?,
            })
        }
    }

    pub fn quantizer(&self) -> *mut FaissIndex {
        unsafe { faiss_next_sys::faiss_IndexIVFFlat_quantizer(self.inner.as_ptr()) }
    }
}

impl Index for IndexIVFFlat {
    fn inner_ptr(&self) -> *mut FaissIndex {
        self.inner.as_ptr() as *mut FaissIndex
    }
}

impl IvfIndex for IndexIVFFlat {
    fn nlist(&self) -> usize {
        unsafe { faiss_next_sys::faiss_IndexIVFFlat_nlist(self.inner.as_ptr()) }
    }

    fn nprobe(&self) -> usize {
        unsafe { faiss_next_sys::faiss_IndexIVFFlat_nprobe(self.inner.as_ptr()) }
    }

    fn set_nprobe(&mut self, nprobe: usize) {
        unsafe { faiss_next_sys::faiss_IndexIVFFlat_set_nprobe(self.inner.as_ptr(), nprobe) }
    }
}

impl Drop for IndexIVFFlat {
    fn drop(&mut self) {
        tracing::trace!("dropping IndexIVFFlat");
        unsafe {
            faiss_next_sys::faiss_IndexIVFFlat_free(self.inner.as_ptr());
        }
    }
}

unsafe impl Send for IndexIVFFlat {}
unsafe impl Sync for IndexIVFFlat {}

pub struct IndexIVF {
    inner: InnerPtr<FaissIndexIVF>,
}

impl IndexIVF {
    pub fn from_index(index: super::IndexImpl) -> Result<Self> {
        unsafe {
            let ivf_ptr = faiss_next_sys::faiss_IndexIVF_cast(index.inner_ptr());
            if ivf_ptr.is_null() {
                return Err(Error::invalid_cast("IndexIVF", "index is not an IVF index"));
            }
            std::mem::forget(index);
            Ok(Self {
                inner: InnerPtr::new(ivf_ptr)?,
            })
        }
    }
}

impl Index for IndexIVF {
    fn inner_ptr(&self) -> *mut FaissIndex {
        self.inner.as_ptr() as *mut FaissIndex
    }
}

impl IvfIndex for IndexIVF {
    fn nlist(&self) -> usize {
        unsafe { faiss_next_sys::faiss_IndexIVF_nlist(self.inner.as_ptr()) }
    }

    fn nprobe(&self) -> usize {
        unsafe { faiss_next_sys::faiss_IndexIVF_nprobe(self.inner.as_ptr()) }
    }

    fn set_nprobe(&mut self, nprobe: usize) {
        unsafe { faiss_next_sys::faiss_IndexIVF_set_nprobe(self.inner.as_ptr(), nprobe) }
    }
}

impl Drop for IndexIVF {
    fn drop(&mut self) {
        tracing::trace!("dropping IndexIVF");
        unsafe {
            faiss_next_sys::faiss_IndexIVF_free(self.inner.as_ptr());
        }
    }
}

unsafe impl Send for IndexIVF {}
unsafe impl Sync for IndexIVF {}