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

Signed-off-by: Shawn Wang <[email protected]>
zilliztech · Jan 23, 2025 · 862f75a · 862f75a
1 parent 0aa91b4
commit 862f75a
Show file tree

Hide file tree

Showing 6 changed files with 518 additions and 27 deletions.
diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h
@@ -178,6 +178,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
@@ -391,6 +391,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
@@ -377,6 +377,20 @@ class SparseInvertedIndexNode : public IndexNode {
                     sparse::SparseMetricType::METRIC_BM25);
                 index->SetBM25Params(k1, b, avgdl);
                 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, mmapped>(
+                        sparse::SparseMetricType::METRIC_BM25);
+                index->SetBM25Params(k1, b, avgdl);
+                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, mmapped>(
+                        sparse::SparseMetricType::METRIC_BM25);
+                index->SetBM25Params(k1, b, avgdl);
+                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, mmapped>(
                     sparse::SparseMetricType::METRIC_BM25);
@@ -400,6 +414,17 @@ class SparseInvertedIndexNode : public IndexNode {
                 auto index = new sparse::InvertedIndex<T, T, sparse::InvertedIndexAlgo::DAAT_MAXSCORE, mmapped>(
                     sparse::SparseMetricType::METRIC_IP);
                 return index;
+            } else if (cfg.inverted_index_algo.value() == "DAAT_BLOCKMAX_WAND") {
+                auto index = new sparse::InvertedIndex<T, T, sparse::InvertedIndexAlgo::DAAT_BLOCKMAX_WAND, mmapped>(
+                    sparse::SparseMetricType::METRIC_IP);
+                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, mmapped>(
+                        sparse::SparseMetricType::METRIC_IP);
+                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, mmapped>(
                     sparse::SparseMetricType::METRIC_IP);