faiss_next/
id_selector.rs

1use std::ptr;
2
3use faiss_next_sys::{
4    self, FaissIDSelector, FaissIDSelectorAnd, FaissIDSelectorBatch, FaissIDSelectorNot,
5    FaissIDSelectorOr, FaissIDSelectorRange, FaissIDSelectorXOr,
6};
7
8use crate::error::{check_return_code, Result};
9use crate::idx::Idx;
10
11pub trait IDSelector {
12    fn as_ptr(&self) -> *const FaissIDSelector;
13}
14
15pub struct IDSelectorRange {
16    ptr: *mut FaissIDSelectorRange,
17}
18
19impl IDSelectorRange {
20    pub fn new(imin: i64, imax: i64) -> Result<Self> {
21        unsafe {
22            let mut ptr: *mut FaissIDSelectorRange = ptr::null_mut();
23            check_return_code(faiss_next_sys::faiss_IDSelectorRange_new(
24                &mut ptr, imin, imax,
25            ))?;
26            Ok(Self { ptr })
27        }
28    }
29
30    pub fn imin(&self) -> i64 {
31        unsafe { faiss_next_sys::faiss_IDSelectorRange_imin(self.ptr) }
32    }
33
34    pub fn imax(&self) -> i64 {
35        unsafe { faiss_next_sys::faiss_IDSelectorRange_imax(self.ptr) }
36    }
37}
38
39impl IDSelector for IDSelectorRange {
40    fn as_ptr(&self) -> *const FaissIDSelector {
41        self.ptr as *const FaissIDSelector
42    }
43}
44
45impl Drop for IDSelectorRange {
46    fn drop(&mut self) {
47        if !self.ptr.is_null() {
48            unsafe {
49                faiss_next_sys::faiss_IDSelectorRange_free(self.ptr);
50            }
51        }
52    }
53}
54
55pub struct IDSelectorBatch {
56    ptr: *mut FaissIDSelectorBatch,
57}
58
59impl IDSelectorBatch {
60    pub fn new(indices: &[i64]) -> Result<Self> {
61        unsafe {
62            let mut ptr: *mut FaissIDSelectorBatch = ptr::null_mut();
63            check_return_code(faiss_next_sys::faiss_IDSelectorBatch_new(
64                &mut ptr,
65                indices.len(),
66                indices.as_ptr(),
67            ))?;
68            Ok(Self { ptr })
69        }
70    }
71
72    pub fn from_idx(indices: &[Idx]) -> Result<Self> {
73        let ids: Vec<i64> = indices.iter().map(|&id| id.as_repr()).collect();
74        Self::new(&ids)
75    }
76}
77
78impl IDSelector for IDSelectorBatch {
79    fn as_ptr(&self) -> *const FaissIDSelector {
80        self.ptr as *const FaissIDSelector
81    }
82}
83
84impl Drop for IDSelectorBatch {
85    fn drop(&mut self) {
86        if !self.ptr.is_null() {
87            unsafe {
88                faiss_next_sys::faiss_IDSelector_free(self.ptr as *mut FaissIDSelector);
89            }
90        }
91    }
92}
93
94pub struct IDSelectorNot {
95    ptr: *mut FaissIDSelectorNot,
96}
97
98impl IDSelectorNot {
99    pub fn new<S: IDSelector>(sel: &S) -> Result<Self> {
100        unsafe {
101            let mut ptr: *mut FaissIDSelectorNot = ptr::null_mut();
102            check_return_code(faiss_next_sys::faiss_IDSelectorNot_new(
103                &mut ptr,
104                sel.as_ptr(),
105            ))?;
106            Ok(Self { ptr })
107        }
108    }
109}
110
111impl IDSelector for IDSelectorNot {
112    fn as_ptr(&self) -> *const FaissIDSelector {
113        self.ptr as *const FaissIDSelector
114    }
115}
116
117impl Drop for IDSelectorNot {
118    fn drop(&mut self) {
119        if !self.ptr.is_null() {
120            unsafe {
121                faiss_next_sys::faiss_IDSelector_free(self.ptr as *mut FaissIDSelector);
122            }
123        }
124    }
125}
126
127pub struct IDSelectorAnd {
128    ptr: *mut FaissIDSelectorAnd,
129}
130
131impl IDSelectorAnd {
132    pub fn new<L: IDSelector, R: IDSelector>(lhs: &L, rhs: &R) -> Result<Self> {
133        unsafe {
134            let mut ptr: *mut FaissIDSelectorAnd = ptr::null_mut();
135            check_return_code(faiss_next_sys::faiss_IDSelectorAnd_new(
136                &mut ptr,
137                lhs.as_ptr(),
138                rhs.as_ptr(),
139            ))?;
140            Ok(Self { ptr })
141        }
142    }
143}
144
145impl IDSelector for IDSelectorAnd {
146    fn as_ptr(&self) -> *const FaissIDSelector {
147        self.ptr as *const FaissIDSelector
148    }
149}
150
151impl Drop for IDSelectorAnd {
152    fn drop(&mut self) {
153        if !self.ptr.is_null() {
154            unsafe {
155                faiss_next_sys::faiss_IDSelector_free(self.ptr as *mut FaissIDSelector);
156            }
157        }
158    }
159}
160
161pub struct IDSelectorOr {
162    ptr: *mut FaissIDSelectorOr,
163}
164
165impl IDSelectorOr {
166    pub fn new<L: IDSelector, R: IDSelector>(lhs: &L, rhs: &R) -> Result<Self> {
167        unsafe {
168            let mut ptr: *mut FaissIDSelectorOr = ptr::null_mut();
169            check_return_code(faiss_next_sys::faiss_IDSelectorOr_new(
170                &mut ptr,
171                lhs.as_ptr(),
172                rhs.as_ptr(),
173            ))?;
174            Ok(Self { ptr })
175        }
176    }
177}
178
179impl IDSelector for IDSelectorOr {
180    fn as_ptr(&self) -> *const FaissIDSelector {
181        self.ptr as *const FaissIDSelector
182    }
183}
184
185impl Drop for IDSelectorOr {
186    fn drop(&mut self) {
187        if !self.ptr.is_null() {
188            unsafe {
189                faiss_next_sys::faiss_IDSelector_free(self.ptr as *mut FaissIDSelector);
190            }
191        }
192    }
193}
194
195pub struct IDSelectorXOr {
196    ptr: *mut FaissIDSelectorXOr,
197}
198
199impl IDSelectorXOr {
200    pub fn new<L: IDSelector, R: IDSelector>(lhs: &L, rhs: &R) -> Result<Self> {
201        unsafe {
202            let mut ptr: *mut FaissIDSelectorXOr = ptr::null_mut();
203            check_return_code(faiss_next_sys::faiss_IDSelectorXOr_new(
204                &mut ptr,
205                lhs.as_ptr(),
206                rhs.as_ptr(),
207            ))?;
208            Ok(Self { ptr })
209        }
210    }
211}
212
213impl IDSelector for IDSelectorXOr {
214    fn as_ptr(&self) -> *const FaissIDSelector {
215        self.ptr as *const FaissIDSelector
216    }
217}
218
219impl Drop for IDSelectorXOr {
220    fn drop(&mut self) {
221        if !self.ptr.is_null() {
222            unsafe {
223                faiss_next_sys::faiss_IDSelector_free(self.ptr as *mut FaissIDSelector);
224            }
225        }
226    }
227}
faiss_next/id_selector.rs

faiss_next/
id_selector.rs
