Merge remote-tracking branch 'upstream/main' into fix478

infiniflow · Mar 9, 2025 · 76fe7d2 · 76fe7d2
2 parents 79cde6d + 88a5a56
commit 76fe7d2
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/python/test_pysdk/test_knn.py b/python/test_pysdk/test_knn.py
@@ -1069,7 +1069,7 @@ def test_sparse_knn_with_index(self, check_data, suffix):
                                                index.IndexType.BMP,
                                                {"block_size": "8", "compress_type": "compress"}), ConflictType.Error)
 
-        table_obj.optimize("idx1", {"topk": "3"})
+        table_obj.optimize("idx1", {"topk": "3", "bp_reorder": ""})
 
         res, extra_result = (table_obj
                              .output(["*", "_row_id", "_similarity"])

diff --git a/src/storage/knn_index/sparse/bmp_alg.cppm b/src/storage/knn_index/sparse/bmp_alg.cppm
@@ -265,10 +265,13 @@ public:
     void Optimize(const BMPOptimizeOptions &options) {
         std::unique_lock lock(mtx_);
 
-        if (options.bp_reorder_) {
+        while (options.bp_reorder_) {
             SizeT block_size = this->block_fwd_.block_size();
             SizeT term_num = this->bm_ivt_.term_num();
             SizeT doc_num = this->doc_ids_.size() - this->doc_ids_.size() % block_size;
+            if (doc_num == 0) {
+                break;
+            }
 
             this->bm_ivt_ = BMPIvt<DataType, CompressType, BMPOwnMem::kTrue>(term_num);
             Vector<Pair<Vector<IdxType>, Vector<DataType>>> fwd = this->block_fwd_.GetFwd(doc_num, term_num);
@@ -293,6 +296,7 @@ public:
                 SparseVecRef<DataType, IdxType> doc((i32)indices.size(), indices.data(), data.data());
                 this->AddDoc(doc, doc_ids[i], false);
             }
+            break;
         }
         if (options.topk_ != 0) {
             SizeT term_num = this->bm_ivt_.term_num();

diff --git a/src/storage/knn_index/sparse/bmp_fwd.cppm b/src/storage/knn_index/sparse/bmp_fwd.cppm
@@ -441,6 +441,9 @@ public:
                 const auto &[term_id, block_size, block_offsets, scores] = iter.Value();
                 for (SizeT i = 0; i < block_size; ++i) {
                     BMPDocID doc_id = block_offsets[i] + block_id * block_size_;
+                    if (doc_id >= doc_num) {
+                        break;
+                    }
                     DataType score = scores[i];
                     fwd[doc_id].first.push_back(term_id);
                     fwd[doc_id].second.push_back(score);