Merge pull request #7 from oramasearch/feat/implements-pq-residual

feat: implements PQ Residual

Cargo.toml

@@ -22,6 +22,8 @@ log = "0.4.22"
 thiserror = "2.0.3"
 env_logger = "0.11.5"
 serde = { version = "1.0.215", features = ["derive"] }
+numpy = "0.23.0"
+num-traits = "0.2.19"
 
 [dev-dependencies]
 criterion = { version = "0.5", features = ["html_reports"] }

src/errors.rs

@@ -2,6 +2,24 @@ use crate::pq::CodeType;
 use ndarray::Ix;
 use thiserror::Error;
 
+#[derive(Error, Debug)]
+pub enum PQResidualError {
+    #[error("Must use at least one product quantizer")]
+    MissingProductQuantizer,
+
+    #[error("Missing dataset name")]
+    MissingDatasetName,
+
+    #[error("Model not trained. Call fit() before calling this method")]
+    ModelNotTrained,
+
+    #[error("I/O error: {0}")]
+    Io(#[from] std::io::Error),
+
+    #[error("PQ Residual error: {0}")]
+    Pq(#[from] PQError),
+}
+
 #[derive(Error, Debug)]
 pub enum PQError {
     #[error("Number of clusters (ks) must be between 1 and 2**32 - 1. Got {0}")]

src/lib.rs

@@ -1,3 +1,4 @@
 mod errors;
 pub mod pq;
+mod pq_residual;
 mod utils;

src/pq.rs

@@ -4,20 +4,20 @@ use log::{info, trace, warn};
 use ndarray::parallel::prelude::*;
 use ndarray::{s, Array2, Array3, Axis};
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub enum CodeType {
     U8,
     U16,
     U32,
 }
 
 pub struct PQ {
-    m: usize,
-    ks: u32,
-    code_dtype: CodeType,
-    codewords: Option<Array3<f32>>,
-    ds: Option<Vec<usize>>,
-    dim: Option<usize>,
+    pub m: usize,
+    pub code_dtype: CodeType,
+    pub ks: u32,
+    pub codewords: Option<Array3<f32>>,
+    pub ds: Option<Vec<usize>>,
+    pub dim: Option<usize>,
 }
 
 impl PQ {