sparse: quantize float to uint16_t for BM25 inverted index

To reduce the memory usage of the inverted index, when BM25 metric is used, quantize term frequency from float to uint16_t. All values that exceeds the maximum of uint16_t is quantized to the maximum of uint16_t. Signed-off-by: Shawn Wang <[email protected]>
zilliztech · Jan 3, 2025 · eadb481 · eadb481
1 parent 7d5fe8b
commit eadb481
Show file tree

Hide file tree

Showing 4 changed files with 214 additions and 130 deletions.
diff --git a/src/index/sparse/sparse_index_node.cc b/src/index/sparse/sparse_index_node.cc
@@ -19,7 +19,6 @@
 #include "knowhere/config.h"
 #include "knowhere/dataset.h"
 #include "knowhere/expected.h"
-#include "knowhere/feature.h"
 #include "knowhere/index/index_factory.h"
 #include "knowhere/index/index_node.h"
 #include "knowhere/log.h"
@@ -125,7 +124,7 @@ class SparseInvertedIndexNode : public IndexNode {
      public:
         RefineIterator(const sparse::BaseInvertedIndex<T>* index, sparse::SparseRow<T>&& query,
                        std::shared_ptr<PrecomputedDistanceIterator> precomputed_it,
-                       const sparse::DocValueComputer<T>& computer, const float refine_ratio = 0.5f)
+                       const sparse::DocValueComputer<float>& computer, const float refine_ratio = 0.5f)
             : IndexIterator(true, refine_ratio),
               index_(index),
               query_(std::move(query)),
@@ -155,7 +154,7 @@ class SparseInvertedIndexNode : public IndexNode {
      private:
         const sparse::BaseInvertedIndex<T>* index_;
         sparse::SparseRow<T> query_;
-        const sparse::DocValueComputer<T> computer_;
+        const sparse::DocValueComputer<float> computer_;
         std::shared_ptr<PrecomputedDistanceIterator> precomputed_it_;
         bool first_return_ = true;
     };
@@ -251,7 +250,7 @@ class SparseInvertedIndexNode : public IndexNode {
             return index_or.error();
         }
         index_ = index_or.value();
-        return index_->Load(reader);
+        return index_->Load(reader, 0, "");
     }
 
     Status
@@ -327,7 +326,8 @@ class SparseInvertedIndexNode : public IndexNode {
     expected<sparse::BaseInvertedIndex<T>*>
     CreateIndex(const SparseInvertedIndexConfig& cfg) const {
         if (IsMetricType(cfg.metric_type.value(), metric::BM25)) {
-            auto idx = new sparse::InvertedIndex<T, use_wand, true, mmapped>();
+            // quantize float to uint16_t when BM25 metric type is used.
+            auto idx = new sparse::InvertedIndex<T, uint16_t, use_wand, true, mmapped>();
             if (!cfg.bm25_k1.has_value() || !cfg.bm25_b.has_value() || !cfg.bm25_avgdl.has_value()) {
                 return expected<sparse::BaseInvertedIndex<T>*>::Err(
                     Status::invalid_args, "BM25 parameters k1, b, and avgdl must be set when building/loading");
@@ -339,7 +339,7 @@ class SparseInvertedIndexNode : public IndexNode {
             idx->SetBM25Params(k1, b, avgdl, max_score_ratio);
             return idx;
         } else {
-            return new sparse::InvertedIndex<T, use_wand, false, mmapped>();
+            return new sparse::InvertedIndex<T, T, use_wand, false, mmapped>();
         }
     }