Use BF for HNSW and PQ+Refine for DiskANN when Filter Ratio is High (#…

…861)

Signed-off-by: Patrick Weizhi Xu <[email protected]>

knowhere/common/Heap.h

@@ -0,0 +1,62 @@
+// Copyright (C) 2019-2023 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+#pragma once
+
+#include <optional>
+#include <queue>
+#include <utility>
+
+namespace knowhere {
+
+// Maintain intermediate top-k results via maxheap
+// TODO: this naive implementation might be optimzed later
+//     1. Based on top-k and pushed element count to swtich strategy
+//     2. Combine `pop` and `push` operation to `replace`
+template <typename DisT, typename IdT>
+class ResultMaxHeap {
+ public:
+    explicit ResultMaxHeap(size_t k) : k_(k) {}
+
+    inline std::optional<std::pair<DisT, IdT>>
+    Pop() {
+        if (pq.empty()) {
+            return std::nullopt;
+        }
+        std::optional<std::pair<DisT, IdT>> res = pq.top();
+        pq.pop();
+        return res;
+    }
+
+    inline void
+    Push(DisT dis, IdT id) {
+        if (pq.size() < k_) {
+            pq.emplace(dis, id);
+            return;
+        }
+
+        if (dis < pq.top().first) {
+            pq.pop();
+            pq.emplace(dis, id);
+        }
+    }
+
+    inline size_t
+    Size() {
+        return pq.size();
+    }
+
+ private:
+    size_t k_;
+    std::priority_queue<std::pair<DisT, IdT>> pq;
+};
+
+}  // namespace knowhere

knowhere/index/vector_index/IndexDiskANN.cpp

@@ -403,7 +403,7 @@ IndexDiskANN<T>::Query(const DatasetPtr& dataset_ptr, const Config& config, cons
         futures.push_back(pool_->push([&, index = row]() {
             pq_flash_index_->cached_beam_search(query + (index * dim), k, query_conf.search_list_size,
                                                 p_id + (index * k), p_dist + (index * k), query_conf.beamwidth, false,
-                                                nullptr, nullptr, bitset);
+                                                nullptr, nullptr, bitset, query_conf.filter_threshold);
         }));
     }
 

knowhere/index/vector_index/IndexDiskANNConfig.cpp

@@ -37,6 +37,7 @@ static constexpr const char* kAioMaxnr = "aio_maxnr";
 
 static constexpr const char* kK = "k";
 static constexpr const char* kBeamwidth = "beamwidth";
+static constexpr const char* kFilterThreshold = "filter_threshold";
 
 static constexpr const char* kRadius = "radius";
 static constexpr const char* kRangeFilter = "range_filter";
@@ -65,6 +66,8 @@ static constexpr std::optional<uint32_t> kDiskPqBytesMaxValue = std::nullopt;
 static constexpr uint32_t kSearchListSizeMaxValue = 200;
 static constexpr uint32_t kBeamwidthMinValue = 1;
 static constexpr uint32_t kBeamwidthMaxValue = 128;
+static constexpr float kFilterThresholdMinValue = -1;
+static constexpr float kFilterThresholdMaxValue = 1;
 static constexpr uint64_t kKMinValue = 1;
 static constexpr std::optional<uint64_t> kKMaxValue = std::nullopt;
 static constexpr uint64_t kAioMaxnrMinValue = 1;
@@ -206,8 +209,10 @@ from_json(const Config& config, DiskANNPrepareConfig& prep_conf) {
 
 void
 to_json(Config& config, const DiskANNQueryConfig& query_conf) {
-    config =
-        Config{{kK, query_conf.k}, {kSearchListSize, query_conf.search_list_size}, {kBeamwidth, query_conf.beamwidth}};
+    config = Config{{kK, query_conf.k},
+                    {kSearchListSize, query_conf.search_list_size},
+                    {kBeamwidth, query_conf.beamwidth},
+                    {kFilterThreshold, query_conf.filter_threshold}};
 }
 
 void
@@ -218,6 +223,8 @@ from_json(const Config& config, DiskANNQueryConfig& query_conf) {
                                       std::max(kSearchListSizeMaxValue, static_cast<uint32_t>(10 * query_conf.k)),
                                       query_conf.search_list_size);
     CheckNumericParamAndSet<uint32_t>(config, kBeamwidth, kBeamwidthMinValue, kBeamwidthMaxValue, query_conf.beamwidth);
+    CheckNumericParamAndSet<float>(config, kFilterThreshold, kFilterThresholdMinValue, kFilterThresholdMaxValue,
+                                   query_conf.filter_threshold);
 }
 
 void

knowhere/index/vector_index/IndexDiskANNConfig.h

@@ -92,6 +92,10 @@ struct DiskANNQueryConfig {
     // slightly higher total number of IO requests to SSD per query. For the highest query throughput with a fixed SSD
     // IOps rating, use W=1. For best latency, use W=4,8 or higher complexity search.
     uint32_t beamwidth = 8;
+    // The threshold which determines when to switch to PQ + Refine strategy based on the number of bits set. The
+    // value should be in range of [0.0, 1.0] which means when greater or equal to x% of the bits are set,
+    // use PQ + Refine. Default to -1.0f, negative vlaues will use dynamic threshold calculator given topk.
+    float filter_threshold = -1.0f;
 
     static DiskANNQueryConfig
     Get(const Config& config);

thirdparty/DiskANN/include/pq_flash_index.h

@@ -3,6 +3,7 @@
 
 #pragma once
 #include <cassert>
+#include <optional>
 #include <sstream>
 #include <stack>
 #include <string>
@@ -98,8 +99,9 @@ namespace diskann {
         const T *query, const _u64 k_search, const _u64 l_search, _s64 *res_ids,
         float *res_dists, const _u64 beam_width,
         const bool use_reorder_data = false, QueryStats *stats = nullptr,
-        const knowhere::feder::diskann::FederResultUniq& feder = nullptr,
-        faiss::BitsetView bitset_view = nullptr);
+        const knowhere::feder::diskann::FederResultUniq &feder = nullptr,
+        faiss::BitsetView                             bitset_view = nullptr,
+        const float                                   filter_ratio = -1.0f);
 
 
     DISKANN_DLLEXPORT _u32 range_search(const T *query1, const double range,
@@ -144,6 +146,23 @@ namespace diskann {
                          : sector_buf + (node_id % nnodes_per_sector) * max_node_len;
     }
 
+    inline void copy_vec_base_data(T *des, const int64_t des_idx, void *src);
+
+    // Init thread data and returns query norm if avaialble.
+    // If there is no value, there is nothing to do with the given query
+    std::optional<float> init_thread_data(ThreadData<T> &data, const T *query1);
+
+    // Brute force search for the given query. Use beam search rather than
+    // sending whole bunch of requests at once to avoid all threads sending I/O
+    // requests and the time overlaps.
+    // The beam width is adjusted in the function.
+    void brute_force_beam_search(
+        ThreadData<T> &data, const float query_norm, const _u64 k_search,
+        _s64 *indices, float *distances, const _u64 beam_width_param, IOContext &ctx,
+        QueryStats                                      *stats,
+        const knowhere::feder::diskann::FederResultUniq &feder,
+        faiss::BitsetView                             bitset_view);
+
     // index info
     // nhood of node `i` is in sector: [i / nnodes_per_sector]
     // offset in sector: [(i % nnodes_per_sector) * max_node_len]

thirdparty/DiskANN/include/utils.h

@@ -33,7 +33,6 @@ typedef HANDLE FileHandle;
 typedef int FileHandle;
 #endif
 
-#include "utils.h"
 #include "logger.h"
 #include "cached_io.h"
 #include "ann_exception.h"