faiss_next/
factory.rs

1use std::ffi::CString;
2use std::ptr;
3
4use faiss_next_sys;
5
6use crate::error::{check_return_code, Result};
7use crate::index::{Index, IndexImpl};
8use crate::metric::MetricType;
9
10pub fn index_factory(d: u32, description: &str, metric: MetricType) -> Result<IndexImpl> {
11    let c_description = CString::new(description)?;
12    let mut inner = ptr::null_mut();
13
14    unsafe {
15        check_return_code(faiss_next_sys::faiss_index_factory(
16            &mut inner,
17            d as i32,
18            c_description.as_ptr(),
19            metric.as_native(),
20        ))?;
21
22        IndexImpl::from_raw(inner)
23    }
24}
25
26pub struct IndexBuilder {
27    d: u32,
28    description: String,
29    metric: MetricType,
30    nprobe: Option<usize>,
31    verbose: bool,
32}
33
34impl IndexBuilder {
35    pub fn new(d: u32) -> Self {
36        Self {
37            d,
38            description: "Flat".to_string(),
39            metric: MetricType::L2,
40            nprobe: None,
41            verbose: false,
42        }
43    }
44
45    pub fn description(mut self, desc: impl Into<String>) -> Self {
46        self.description = desc.into();
47        self
48    }
49
50    pub fn flat(self) -> Self {
51        self.description("Flat")
52    }
53
54    pub fn ivf_flat(self, nlist: usize) -> Self {
55        self.description(format!("IVF{},Flat", nlist))
56    }
57
58    pub fn ivf_pq(self, nlist: usize, m: usize, nbits: usize) -> Self {
59        self.description(format!("IVF{},PQ{}x{}", nlist, m, nbits))
60    }
61
62    pub fn pq(self, m: usize, nbits: usize) -> Self {
63        self.description(format!("PQ{}x{}", m, nbits))
64    }
65
66    pub fn hnsw(self, m: usize) -> Self {
67        self.description(format!("HNSW{}", m))
68    }
69
70    pub fn lsh(self, nbits: usize) -> Self {
71        self.description(format!("LSH{}", nbits))
72    }
73
74    pub fn metric(mut self, metric: MetricType) -> Self {
75        self.metric = metric;
76        self
77    }
78
79    pub fn l2(self) -> Self {
80        self.metric(MetricType::L2)
81    }
82
83    pub fn ip(self) -> Self {
84        self.metric(MetricType::InnerProduct)
85    }
86
87    pub fn nprobe(mut self, nprobe: usize) -> Self {
88        self.nprobe = Some(nprobe);
89        self
90    }
91
92    pub fn verbose(mut self, verbose: bool) -> Self {
93        self.verbose = verbose;
94        self
95    }
96
97    pub fn build(self) -> Result<IndexImpl> {
98        let mut index = index_factory(self.d, &self.description, self.metric)?;
99
100        if self.verbose {
101            index.set_verbose(true);
102        }
103
104        Ok(index)
105    }
106}
faiss_next/factory.rs

faiss_next/
factory.rs
