sparse: make the distance returned by RangeSearch and AnnIterator to …

…be the raw instead of the quantized distance

Signed-off-by: Buqian Zheng <zhengbuqian@gmail.com>

src/index/sparse/sparse_index_node.cc

@@ -122,6 +122,47 @@ class SparseInvertedIndexNode : public IndexNode {
         return GenResultDataSet(nq, k, p_id.release(), p_dist.release());
     }
 
+ private:
+    class RefineIterator : public IndexIterator {
+     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)
+            : IndexIterator(true, refine_ratio),
+              index_(index),
+              query_(std::move(query)),
+              precomputed_it_(precomputed_it),
+              computer_(computer) {
+        }
+
+     protected:
+        // returns n_rows / 10 DistId for the first time to create a large enough window for refinement.
+        void
+        next_batch(std::function<void(const std::vector<DistId>&)> batch_handler) override {
+            std::vector<DistId> dists;
+            auto num = first_return_ ? (index_->n_rows() / 10) : 1;
+            first_return_ = false;
+            for (size_t i = 0; i < num && precomputed_it_->HasNext(); ++i) {
+                auto [id, dist] = precomputed_it_->Next();
+                dists.emplace_back(id, dist);
+            }
+            batch_handler(dists);
+        }
+
+        float
+        raw_distance(int64_t id) override {
+            return index_->GetRawDistance(id, query_, computer_);
+        }
+
+     private:
+        const sparse::BaseInvertedIndex<T>* index_;
+        sparse::SparseRow<T> query_;
+        const sparse::DocValueComputer<T> computer_;
+        std::shared_ptr<PrecomputedDistanceIterator> precomputed_it_;
+        bool first_return_ = true;
+    };
+
+ public:
     // TODO: for now inverted index and wand use the same impl for AnnIterator.
     [[nodiscard]] expected<std::vector<IndexNode::IteratorPtr>>
     AnnIterator(const DataSetPtr dataset, std::unique_ptr<Config> config, const BitsetView& bitset) const override {
@@ -142,13 +183,23 @@ class SparseInvertedIndexNode : public IndexNode {
         auto computer = computer_or.value();
         auto drop_ratio_search = cfg.drop_ratio_search.value_or(0.0f);
 
+        const bool has_quantization = index_->HasQuantization() || drop_ratio_search > 0;
+
         auto vec = std::vector<std::shared_ptr<IndexNode::iterator>>(nq, nullptr);
         std::vector<folly::Future<folly::Unit>> futs;
         futs.reserve(nq);
         for (int i = 0; i < nq; ++i) {
             futs.emplace_back(search_pool_->push([&, i]() {
-                vec[i] = std::make_shared<PrecomputedDistanceIterator>(
+                auto it = std::make_shared<PrecomputedDistanceIterator>(
                     index_->GetAllDistances(queries[i], drop_ratio_search, bitset, computer), true);
+                if (!has_quantization) {
+                    vec[i] = it;
+                } else {
+                    sparse::SparseRow<T> query_copy(queries[i]);
+                    auto refine_it = std::make_shared<RefineIterator>(index_, std::move(query_copy), it, computer);
+                    refine_it->initialize();
+                    vec[i] = std::move(refine_it);
+                }
             }));
         }
         WaitAllSuccess(futs);
@@ -372,6 +423,11 @@ class SparseInvertedIndexNodeCC : public SparseInvertedIndexNode<T, use_wand> {
     expected<std::vector<IndexNode::IteratorPtr>>
     AnnIterator(const DataSetPtr dataset, std::unique_ptr<Config> cfg, const BitsetView& bitset) const override {
         ReadPermission permission(*this);
+        // Always uses PrecomputedDistanceIterator for SparseInvertedIndexNodeCC:
+        // If we want to use RefineIterator, it needs to get another ReadPermission when calling
+        // index_->GetRawDistance(). If an Add task is added in between, there will be a deadlock.
+        auto config = static_cast<const knowhere::SparseInvertedIndexConfig&>(*cfg);
+        config.drop_ratio_search = 0.0f;
         return SparseInvertedIndexNode<T, use_wand>::AnnIterator(dataset, std::move(cfg), bitset);
     }
 

src/index/sparse/sparse_inverted_index.h

@@ -60,12 +60,18 @@ class BaseInvertedIndex {
     GetAllDistances(const SparseRow<T>& query, float drop_ratio_search, const BitsetView& bitset,
                     const DocValueComputer<T>& computer) const = 0;
 
+    virtual float
+    GetRawDistance(const label_t id, const SparseRow<T>& query, const DocValueComputer<T>& computer) const = 0;
+
     virtual void
     GetVectorById(const label_t id, SparseRow<T>& output) const = 0;
 
     virtual expected<DocValueComputer<T>>
     GetDocValueComputer(const SparseInvertedIndexConfig& cfg) const = 0;
 
+    virtual bool
+    HasQuantization() const = 0;
+
     [[nodiscard]] virtual size_t
     size() const = 0;
 
@@ -435,6 +441,7 @@ class InvertedIndex : public BaseInvertedIndex<T> {
         }
     }
 
+    // Returned distances are inaccurate based on the drop_ratio.
     std::vector<float>
     GetAllDistances(const SparseRow<T>& query, float drop_ratio_search, const BitsetView& bitset,
                     const DocValueComputer<T>& computer) const override {
@@ -459,6 +466,12 @@ class InvertedIndex : public BaseInvertedIndex<T> {
         return distances;
     }
 
+    float
+    GetRawDistance(const label_t id, const SparseRow<T>& query, const DocValueComputer<T>& computer) const override {
+        T doc_sum = bm25 ? bm25_params_->row_sums.at(id) : 0;
+        return query.dot(raw_data_[id], computer, doc_sum);
+    }
+
     void
     GetVectorById(const label_t id, SparseRow<T>& output) const override {
         output = raw_data_[id];
@@ -502,6 +515,11 @@ class InvertedIndex : public BaseInvertedIndex<T> {
         return n_cols_internal();
     }
 
+    [[nodiscard]] virtual bool
+    HasQuantization() const override {
+        return drop_during_build_;
+    }
+
  private:
     size_t
     n_rows_internal() const {

tests/ut/test_sparse.cc

@@ -222,18 +222,24 @@ TEST_CASE("Test Mem Sparse Index With Float Vector", "[float metrics]") {
         REQUIRE(iterators_or.has_value());
         auto& iterators = iterators_or.value();
         REQUIRE(iterators.size() == (size_t)nq);
-        // verify the distances are monotonic decreasing, as INDEX_SPARSE_INVERTED_INDEX and INDEX_SPARSE_WAND
-        // performs exausitive search for iterator.
+
+        int count = 0;
+        int out_of_order = 0;
         for (int i = 0; i < nq; ++i) {
             auto& iter = iterators[i];
             float prev_dist = std::numeric_limits<float>::max();
             while (iter->HasNext()) {
                 auto [id, dist] = iter->Next();
                 REQUIRE(!bitset.test(id));
-                REQUIRE(prev_dist >= dist);
+                count++;
+                if (prev_dist < dist) {
+                    out_of_order++;
+                }
                 prev_dist = dist;
             }
         }
+        // less than 5% of the distances are out of order.
+        REQUIRE(out_of_order * 20 <= count);
     }
 
     SECTION("Test Sparse Range Search") {