faiss_next/
search_params.rs

1use std::ptr;
2
3use faiss_next_sys::{self, FaissSearchParameters, FaissSearchParametersIVF};
4
5use crate::error::{check_return_code, Result};
6
7pub trait SearchParams {
8    fn as_ptr(&self) -> *const FaissSearchParameters;
9}
10
11pub struct SearchParameters {
12    ptr: *mut FaissSearchParameters,
13}
14
15impl SearchParameters {
16    pub fn new() -> Result<Self> {
17        unsafe {
18            let mut ptr: *mut FaissSearchParameters = ptr::null_mut();
19            check_return_code(faiss_next_sys::faiss_SearchParameters_new(
20                &mut ptr,
21                ptr::null_mut(),
22            ))?;
23            Ok(Self { ptr })
24        }
25    }
26}
27
28impl SearchParams for SearchParameters {
29    fn as_ptr(&self) -> *const FaissSearchParameters {
30        self.ptr
31    }
32}
33
34impl Default for SearchParameters {
35    fn default() -> Self {
36        Self::new().expect("Failed to create SearchParameters")
37    }
38}
39
40impl Drop for SearchParameters {
41    fn drop(&mut self) {
42        if !self.ptr.is_null() {
43            unsafe {
44                faiss_next_sys::faiss_SearchParameters_free(self.ptr);
45            }
46        }
47    }
48}
49
50pub struct SearchParametersIvf {
51    ptr: *mut FaissSearchParametersIVF,
52}
53
54impl SearchParametersIvf {
55    pub fn new() -> Result<Self> {
56        unsafe {
57            let mut ptr: *mut FaissSearchParametersIVF = ptr::null_mut();
58            check_return_code(faiss_next_sys::faiss_SearchParametersIVF_new(&mut ptr))?;
59            Ok(Self { ptr })
60        }
61    }
62
63    pub fn with_params(nprobe: usize, max_codes: usize) -> Result<Self> {
64        unsafe {
65            let mut ptr: *mut FaissSearchParametersIVF = ptr::null_mut();
66            check_return_code(faiss_next_sys::faiss_SearchParametersIVF_new_with(
67                &mut ptr,
68                ptr::null_mut(),
69                nprobe,
70                max_codes,
71            ))?;
72            Ok(Self { ptr })
73        }
74    }
75
76    pub fn nprobe(&self) -> usize {
77        unsafe { faiss_next_sys::faiss_SearchParametersIVF_nprobe(self.ptr) }
78    }
79
80    pub fn set_nprobe(&mut self, nprobe: usize) {
81        unsafe { faiss_next_sys::faiss_SearchParametersIVF_set_nprobe(self.ptr, nprobe) }
82    }
83
84    pub fn max_codes(&self) -> usize {
85        unsafe { faiss_next_sys::faiss_SearchParametersIVF_max_codes(self.ptr) }
86    }
87
88    pub fn set_max_codes(&mut self, max_codes: usize) {
89        unsafe { faiss_next_sys::faiss_SearchParametersIVF_set_max_codes(self.ptr, max_codes) }
90    }
91}
92
93impl SearchParams for SearchParametersIvf {
94    fn as_ptr(&self) -> *const FaissSearchParameters {
95        self.ptr as *const FaissSearchParameters
96    }
97}
98
99impl Default for SearchParametersIvf {
100    fn default() -> Self {
101        Self::new().expect("Failed to create SearchParametersIvf")
102    }
103}
104
105impl Drop for SearchParametersIvf {
106    fn drop(&mut self) {
107        if !self.ptr.is_null() {
108            unsafe {
109                faiss_next_sys::faiss_SearchParametersIVF_free(self.ptr);
110            }
111        }
112    }
113}
faiss_next/search_params.rs

faiss_next/
search_params.rs
