implement the binary quantization in meilisearch

milli/src/search/hybrid.rs

@@ -190,7 +190,7 @@ impl<'a> Search<'a> {
             return Ok(return_keyword_results(self.limit, self.offset, keyword_results));
         };
         // no embedder, no semantic search
-        let Some(SemanticSearch { vector, embedder_name, embedder }) = semantic else {
+        let Some(SemanticSearch { vector, embedder_name, embedder, quantized }) = semantic else {
             return Ok(return_keyword_results(self.limit, self.offset, keyword_results));
         };
 
@@ -212,7 +212,7 @@ impl<'a> Search<'a> {
         };
 
         search.semantic =
-            Some(SemanticSearch { vector: Some(vector_query), embedder_name, embedder });
+            Some(SemanticSearch { vector: Some(vector_query), embedder_name, embedder, quantized });
 
         // TODO: would be better to have two distinct functions at this point
         let vector_results = search.execute()?;

milli/src/search/mod.rs

@@ -32,6 +32,7 @@ pub struct SemanticSearch {
     vector: Option<Vec<f32>>,
     embedder_name: String,
     embedder: Arc<Embedder>,
+    quantized: bool,
 }
 
 pub struct Search<'a> {
@@ -89,9 +90,10 @@ impl<'a> Search<'a> {
         &mut self,
         embedder_name: String,
         embedder: Arc<Embedder>,
+        quantized: bool,
         vector: Option<Vec<f32>>,
     ) -> &mut Search<'a> {
-        self.semantic = Some(SemanticSearch { embedder_name, embedder, vector });
+        self.semantic = Some(SemanticSearch { embedder_name, embedder, quantized, vector });
         self
     }
 
@@ -206,7 +208,7 @@ impl<'a> Search<'a> {
             degraded,
             used_negative_operator,
         } = match self.semantic.as_ref() {
-            Some(SemanticSearch { vector: Some(vector), embedder_name, embedder }) => {
+            Some(SemanticSearch { vector: Some(vector), embedder_name, embedder, quantized }) => {
                 execute_vector_search(
                     &mut ctx,
                     vector,
@@ -219,6 +221,7 @@ impl<'a> Search<'a> {
                     self.limit,
                     embedder_name,
                     embedder,
+                    *quantized,
                     self.time_budget.clone(),
                     self.ranking_score_threshold,
                 )?

milli/src/search/new/mod.rs

@@ -320,6 +320,7 @@ fn get_ranking_rules_for_vector<'ctx>(
     target: &[f32],
     embedder_name: &str,
     embedder: &Embedder,
+    quantized: bool,
 ) -> Result<Vec<BoxRankingRule<'ctx, PlaceholderQuery>>> {
     // query graph search
 
@@ -347,6 +348,7 @@ fn get_ranking_rules_for_vector<'ctx>(
                         limit_plus_offset,
                         embedder_name,
                         embedder,
+                        quantized,
                     )?;
                     ranking_rules.push(Box::new(vector_sort));
                     vector = true;
@@ -576,6 +578,7 @@ pub fn execute_vector_search(
     length: usize,
     embedder_name: &str,
     embedder: &Embedder,
+    quantized: bool,
     time_budget: TimeBudget,
     ranking_score_threshold: Option<f64>,
 ) -> Result<PartialSearchResult> {
@@ -591,6 +594,7 @@ pub fn execute_vector_search(
         vector,
         embedder_name,
         embedder,
+        quantized,
     )?;
 
     let mut placeholder_search_logger = logger::DefaultSearchLogger;

milli/src/search/new/vector_sort.rs

@@ -16,6 +16,7 @@ pub struct VectorSort<Q: RankingRuleQueryTrait> {
     limit: usize,
     distribution_shift: Option<DistributionShift>,
     embedder_index: u8,
+    quantized: bool,
 }
 
 impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
@@ -26,6 +27,7 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
         limit: usize,
         embedder_name: &str,
         embedder: &Embedder,
+        quantized: bool,
     ) -> Result<Self> {
         let embedder_index = ctx
             .index
@@ -41,6 +43,7 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
             limit,
             distribution_shift: embedder.distribution(),
             embedder_index,
+            quantized,
         })
     }
 
@@ -49,16 +52,15 @@ impl<Q: RankingRuleQueryTrait> VectorSort<Q> {
         ctx: &mut SearchContext<'_>,
         vector_candidates: &RoaringBitmap,
     ) -> Result<()> {
-        let readers: std::result::Result<Vec<_>, _> =
-            ctx.index.arroy_readers(ctx.txn, self.embedder_index).collect();
-        let readers = readers?;
+        let readers: Vec<_> =
+            ctx.index.arroy_readers(self.embedder_index, self.quantized).collect();
 
         let target = &self.target;
         let mut results = Vec::new();
 
         for reader in readers.iter() {
             let nns_by_vector =
-                reader.nns_by_vector(ctx.txn, target, self.limit, None, Some(vector_candidates))?;
+                reader.nns_by_vector(ctx.txn, target, self.limit, Some(vector_candidates))?;
             results.extend(nns_by_vector.into_iter());
         }
         results.sort_unstable_by_key(|(_, distance)| OrderedFloat(*distance));

milli/src/search/similar.rs

@@ -18,6 +18,7 @@ pub struct Similar<'a> {
     embedder_name: String,
     embedder: Arc<Embedder>,
     ranking_score_threshold: Option<f64>,
+    quantized: bool,
 }
 
 impl<'a> Similar<'a> {
@@ -29,6 +30,7 @@ impl<'a> Similar<'a> {
         rtxn: &'a heed::RoTxn<'a>,
         embedder_name: String,
         embedder: Arc<Embedder>,
+        quantized: bool,
     ) -> Self {
         Self {
             id,
@@ -40,6 +42,7 @@ impl<'a> Similar<'a> {
             embedder_name,
             embedder,
             ranking_score_threshold: None,
+            quantized,
         }
     }
 
@@ -67,10 +70,7 @@ impl<'a> Similar<'a> {
                 .get(self.rtxn, &self.embedder_name)?
                 .ok_or_else(|| crate::UserError::InvalidEmbedder(self.embedder_name.to_owned()))?;
 
-        let readers: std::result::Result<Vec<_>, _> =
-            self.index.arroy_readers(self.rtxn, embedder_index).collect();
-
-        let readers = readers?;
+        let readers: Vec<_> = self.index.arroy_readers(embedder_index, self.quantized).collect();
 
         let mut results = Vec::new();
 
@@ -79,7 +79,6 @@ impl<'a> Similar<'a> {
                 self.rtxn,
                 self.id,
                 self.limit + self.offset + 1,
-                None,
                 Some(&universe),
             )?;
             if let Some(mut nns_by_item) = nns_by_item {