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"