faiss_next/index/
scalar_quantizer.rs

1use std::ptr;
2
3use faiss_next_sys::{self, FaissIndex, FaissIndexScalarQuantizer, FaissQuantizerType};
4
5use crate::error::{check_return_code, Error, Result};
6use crate::index::native::InnerPtr;
7use crate::index::traits::Index;
8use crate::metric::MetricType;
9
10pub struct IndexScalarQuantizer {
11    inner: InnerPtr<FaissIndexScalarQuantizer>,
12}
13
14impl IndexScalarQuantizer {
15    pub fn new(d: u32, qtype: QuantizerType, metric: MetricType) -> Result<Self> {
16        unsafe {
17            let mut inner = ptr::null_mut();
18            check_return_code(faiss_next_sys::faiss_IndexScalarQuantizer_new_with(
19                &mut inner,
20                d as i64,
21                qtype.as_native(),
22                metric.as_native(),
23            ))?;
24            Ok(Self {
25                inner: InnerPtr::new(inner)?,
26            })
27        }
28    }
29
30    pub fn from_index(index: super::IndexImpl) -> Result<Self> {
31        unsafe {
32            let sq_ptr = faiss_next_sys::faiss_IndexScalarQuantizer_cast(index.inner_ptr());
33            if sq_ptr.is_null() {
34                return Err(Error::invalid_cast(
35                    "IndexScalarQuantizer",
36                    "index is not a ScalarQuantizer index",
37                ));
38            }
39            std::mem::forget(index);
40            Ok(Self {
41                inner: InnerPtr::new(sq_ptr)?,
42            })
43        }
44    }
45}
46
47impl Index for IndexScalarQuantizer {
48    fn inner_ptr(&self) -> *mut FaissIndex {
49        self.inner.as_ptr() as *mut FaissIndex
50    }
51}
52
53impl Drop for IndexScalarQuantizer {
54    fn drop(&mut self) {
55        tracing::trace!("dropping IndexScalarQuantizer");
56        unsafe {
57            faiss_next_sys::faiss_IndexScalarQuantizer_free(self.inner.as_ptr());
58        }
59    }
60}
61
62unsafe impl Send for IndexScalarQuantizer {}
63unsafe impl Sync for IndexScalarQuantizer {}
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub enum QuantizerType {
67    Qt8bit,
68    Qt4bit,
69    Qt8bitUniform,
70    Qt4bitUniform,
71    QtFp16,
72    Qt8bitDirect,
73    Qt6bit,
74}
75
76impl QuantizerType {
77    pub fn as_native(&self) -> FaissQuantizerType {
78        match self {
79            QuantizerType::Qt8bit => FaissQuantizerType::QT_8bit,
80            QuantizerType::Qt4bit => FaissQuantizerType::QT_4bit,
81            QuantizerType::Qt8bitUniform => FaissQuantizerType::QT_8bit_uniform,
82            QuantizerType::Qt4bitUniform => FaissQuantizerType::QT_4bit_uniform,
83            QuantizerType::QtFp16 => FaissQuantizerType::QT_fp16,
84            QuantizerType::Qt8bitDirect => FaissQuantizerType::QT_8bit_direct,
85            QuantizerType::Qt6bit => FaissQuantizerType::QT_6bit,
86        }
87    }
88}
faiss_next/index/scalar_quantizer.rs

faiss_next/index/
scalar_quantizer.rs
