sparse: add block-max wand and block-max maxscore algorithm

Signed-off-by: Shawn Wang <[email protected]>
zilliztech · Jan 14, 2025 · 9e6f9bd · 9e6f9bd
1 parent 59bfdc2
commit 9e6f9bd
Show file tree

Hide file tree

Showing 6 changed files with 494 additions and 26 deletions.
diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h
@@ -177,6 +177,7 @@ constexpr const char* PRQ_NUM = "nrq";      // for PRQ, number of redisual quant
 constexpr const char* INVERTED_INDEX_ALGO = "inverted_index_algo";
 constexpr const char* DROP_RATIO_BUILD = "drop_ratio_build";
 constexpr const char* DROP_RATIO_SEARCH = "drop_ratio_search";
+constexpr const char* BLOCKMAX_BLOCK_SIZE = "blockmax_block_size";
 }  // namespace indexparam
 
 using MetricType = std::string;

diff --git a/include/knowhere/sparse_utils.h b/include/knowhere/sparse_utils.h
@@ -386,6 +386,16 @@ class GrowableVectorView {
         return reinterpret_cast<const T*>(mmap_data_)[i];
     }
 
+    T&
+    back() {
+        return reinterpret_cast<T*>(mmap_data_)[size() - 1];
+    }
+
+    const T&
+    back() const {
+        return reinterpret_cast<const T*>(mmap_data_)[size() - 1];
+    }
+
     class iterator : public boost::iterator_facade<iterator, T, boost::random_access_traversal_tag, T&> {
      public:
         iterator() = default;

diff --git a/src/index/sparse/sparse_index_node.cc b/src/index/sparse/sparse_index_node.cc
@@ -364,6 +364,18 @@ class SparseInvertedIndexNode : public IndexNode {
                     new sparse::InvertedIndex<T, uint16_t, sparse::InvertedIndexAlgo::DAAT_WAND, true, mmapped>();
                 index->SetBM25Params(k1, b, avgdl, max_score_ratio);
                 return index;
+            } else if (cfg.inverted_index_algo.value() == "DAAT_BLOCKMAX_WAND") {
+                auto index = new sparse::InvertedIndex<T, uint16_t, sparse::InvertedIndexAlgo::DAAT_BLOCKMAX_WAND, true,
+                                                       mmapped>();
+                index->SetBM25Params(k1, b, avgdl, max_score_ratio);
+                index->SetBlockmaxBlockSize(cfg.blockmax_block_size.value());
+                return index;
+            } else if (cfg.inverted_index_algo.value() == "DAAT_BLOCKMAX_MAXSCORE") {
+                auto index = new sparse::InvertedIndex<T, uint16_t, sparse::InvertedIndexAlgo::DAAT_BLOCKMAX_MAXSCORE,
+                                                       true, mmapped>();
+                index->SetBM25Params(k1, b, avgdl, max_score_ratio);
+                index->SetBlockmaxBlockSize(cfg.blockmax_block_size.value());
+                return index;
             } else if (cfg.inverted_index_algo.value() == "DAAT_MAXSCORE") {
                 auto index =
                     new sparse::InvertedIndex<T, uint16_t, sparse::InvertedIndexAlgo::DAAT_MAXSCORE, true, mmapped>();
@@ -386,6 +398,16 @@ class SparseInvertedIndexNode : public IndexNode {
                 auto index =
                     new sparse::InvertedIndex<T, T, sparse::InvertedIndexAlgo::DAAT_MAXSCORE, false, mmapped>();
                 return index;
+            } else if (cfg.inverted_index_algo.value() == "DAAT_BLOCKMAX_WAND") {
+                auto index =
+                    new sparse::InvertedIndex<T, T, sparse::InvertedIndexAlgo::DAAT_BLOCKMAX_WAND, false, mmapped>();
+                index->SetBlockmaxBlockSize(cfg.blockmax_block_size.value());
+                return index;
+            } else if (cfg.inverted_index_algo.value() == "DAAT_BLOCKMAX_MAXSCORE") {
+                auto index = new sparse::InvertedIndex<T, T, sparse::InvertedIndexAlgo::DAAT_BLOCKMAX_MAXSCORE, false,
+                                                       mmapped>();
+                index->SetBlockmaxBlockSize(cfg.blockmax_block_size.value());
+                return index;
             } else if (cfg.inverted_index_algo.value() == "TAAT_NAIVE") {
                 auto index = new sparse::InvertedIndex<T, T, sparse::InvertedIndexAlgo::TAAT_NAIVE, false, mmapped>();
                 return index;