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
138
139
140
141
142
143
144
145
146
147
148

/// A macro which provides a native index implementation to the given type.
macro_rules! impl_native_index {
    ($t:ty) => {
        impl crate::index::Index for $t {
            fn is_trained(&self) -> bool {
                unsafe { faiss_Index_is_trained(self.inner_ptr()) != 0 }
            }

            fn ntotal(&self) -> u64 {
                unsafe { faiss_Index_ntotal(self.inner_ptr()) as u64 }
            }

            fn d(&self) -> u32 {
                unsafe { faiss_Index_d(self.inner_ptr()) as u32 }
            }

            fn metric_type(&self) -> crate::metric::MetricType {
                unsafe {
                    crate::metric::MetricType::from_code(
                        faiss_Index_metric_type(self.inner_ptr()) as u32
                    )
                    .unwrap()
                }
            }

            fn add(&mut self, x: &[f32]) -> Result<()> {
                unsafe {
                    let n = x.len() / self.d() as usize;
                    faiss_try(faiss_Index_add(self.inner_ptr(), n as i64, x.as_ptr()))?;
                    Ok(())
                }
            }

            fn add_with_ids(&mut self, x: &[f32], xids: &[crate::index::Idx]) -> Result<()> {
                unsafe {
                    let n = x.len() / self.d() as usize;
                    faiss_try(faiss_Index_add_with_ids(
                        self.inner_ptr(),
                        n as i64,
                        x.as_ptr(),
                        xids.as_ptr() as *const _,
                    ))?;
                    Ok(())
                }
            }
            fn train(&mut self, x: &[f32]) -> Result<()> {
                unsafe {
                    let n = x.len() / self.d() as usize;
                    faiss_try(faiss_Index_train(self.inner_ptr(), n as i64, x.as_ptr()))?;
                    Ok(())
                }
            }
            fn assign(
                &mut self,
                query: &[f32],
                k: usize,
            ) -> Result<crate::index::AssignSearchResult> {
                unsafe {
                    let nq = query.len() / self.d() as usize;
                    let mut out_labels = vec![Idx::none(); k * nq];
                    faiss_try(faiss_Index_assign(
                        self.inner_ptr(),
                        nq as idx_t,
                        query.as_ptr(),
                        out_labels.as_mut_ptr() as *mut _,
                        k as i64,
                    ))?;
                    Ok(crate::index::AssignSearchResult { labels: out_labels })
                }
            }
            fn search(&mut self, query: &[f32], k: usize) -> Result<crate::index::SearchResult> {
                unsafe {
                    let nq = query.len() / self.d() as usize;
                    let mut distances = vec![0_f32; k * nq];
                    let mut labels = vec![Idx::none(); k * nq];
                    faiss_try(faiss_Index_search(
                        self.inner_ptr(),
                        nq as idx_t,
                        query.as_ptr(),
                        k as idx_t,
                        distances.as_mut_ptr(),
                        labels.as_mut_ptr() as *mut _,
                    ))?;
                    Ok(crate::index::SearchResult { distances, labels })
                }
            }
            fn range_search(
                &mut self,
                query: &[f32],
                radius: f32,
            ) -> Result<crate::index::RangeSearchResult> {
                unsafe {
                    let nq = (query.len() / self.d() as usize) as idx_t;
                    let mut p_res: *mut FaissRangeSearchResult = ::std::ptr::null_mut();
                    faiss_try(faiss_RangeSearchResult_new(&mut p_res, nq))?;
                    faiss_try(faiss_Index_range_search(
                        self.inner_ptr(),
                        nq,
                        query.as_ptr(),
                        radius,
                        p_res,
                    ))?;
                    Ok(crate::index::RangeSearchResult { inner: p_res })
                }
            }

            fn reset(&mut self) -> Result<()> {
                unsafe {
                    faiss_try(faiss_Index_reset(self.inner_ptr()))?;
                    Ok(())
                }
            }

            fn remove_ids(&mut self, sel: &IdSelector) -> Result<usize> {
                unsafe {
                    let mut n_removed = 0;
                    faiss_try(faiss_Index_remove_ids(
                        self.inner_ptr(),
                        sel.inner_ptr(),
                        &mut n_removed,
                    ))?;
                    Ok(n_removed)
                }
            }
        }
    };
}

/// A macro which provides a Clone implementation to native index types.
macro_rules! impl_native_index_clone {
    ($t:ty) => {
        impl $t {
            /// Create an independent clone of this index.
            ///
            /// # Errors
            ///
            /// May result in a native error if the clone operation is not
            /// supported for the internal type of index.
            pub fn try_clone(&self) -> Result<Self> {
                unsafe {
                    let mut new_index_ptr = ::std::ptr::null_mut();
                    faiss_try(faiss_clone_index(self.inner_ptr(), &mut new_index_ptr))?;
                    Ok(crate::index::FromInnerPtr::from_inner_ptr(new_index_ptr))
                }
            }
        }
    };
}