Global vector index

* Disable creation of vector index in shard
* Execution flow with/without knn node
* Add/Remove/Rebuild docs
* Use GlobalDocId for vector index

Author: mkaruza

src/core/search/ast_expr.cc

@@ -73,6 +73,10 @@ AstKnnNode::AstKnnNode(AstNode&& filter, AstKnnNode&& self) {
   this->filter = make_unique<AstNode>(std::move(filter));
 }
 
+bool AstKnnNode::HasFilterNode() const {
+  return filter == nullptr;
+}
+
 }  // namespace dfly::search
 
 namespace std {

src/core/search/ast_expr.h

@@ -114,6 +114,8 @@ struct AstKnnNode {
   OwnedFtVector vec;
   std::string score_alias;
   std::optional<float> ef_runtime;
+
+  bool HasFilterNode() const;
 };
 
 using NodeVariants =

src/core/search/base.h

@@ -16,10 +16,24 @@
 #include "absl/container/flat_hash_set.h"
 #include "base/pmr/memory_resource.h"
 #include "core/string_map.h"
+#include "server/tx_base.h"
 
 namespace dfly::search {
 
 using DocId = uint32_t;
+using GlobalDocId = uint64_t;
+
+inline GlobalDocId CreateGlobalDocId(ShardId shard_id, DocId local_doc_id) {
+  return ((uint64_t)shard_id << 32) | local_doc_id;
+}
+
+inline ShardId GlobalDocIdShardId(GlobalDocId id) {
+  return (id >> 32);
+}
+
+inline search::DocId GlobalDocIdLocalId(GlobalDocId id) {
+  return (id)&0xFFFF;
+}
 
 enum class VectorSimilarity { L2, IP, COSINE };
 

src/core/search/indices.cc

@@ -523,12 +523,12 @@ template <typename T> const float* FlatVectorIndex<T>::Get(T doc) const {
 }
 
 template <typename T> std::vector<T> FlatVectorIndex<T>::GetAllDocsWithNonNullValues() const {
-  std::vector<DocId> result;
+  std::vector<T> result;
 
   size_t num_vectors = entries_.size() / BaseVectorIndex<T>::dim_;
   result.reserve(num_vectors);
 
-  for (DocId id = 0; id < num_vectors; ++id) {
+  for (T id = 0; id < num_vectors; ++id) {
     // Check if the vector is not zero (all elements are 0)
     // TODO: Valid vector can contain 0s, we should use a better approach
     const float* vec = Get(id);
@@ -553,8 +553,9 @@ template <typename T> std::vector<T> FlatVectorIndex<T>::GetAllDocsWithNonNullVa
 }
 
 template struct FlatVectorIndex<DocId>;
+template struct FlatVectorIndex<GlobalDocId>;
 
-struct HnswlibAdapter {
+template <typename T> struct HnswlibAdapter {
   // Default setting of hnswlib/hnswalg
   constexpr static size_t kDefaultEfRuntime = 10;
 
@@ -564,34 +565,45 @@ struct HnswlibAdapter {
                100 /* seed*/} {
   }
 
-  void Add(const float* data, DocId id) {
-    if (world_.cur_element_count + 1 >= world_.max_elements_)
-      world_.resizeIndex(world_.cur_element_count * 2);
-    world_.addPoint(data, id);
+  void Add(const float* data, T id) {
+    while (true) {
+      try {
+        absl::ReaderMutexLock lock(&resize_mutex_);
+        world_.addPoint(data, id);
+        return;
+      } catch (const std::exception& e) {
+        std::string error_msg = e.what();
+        if (absl::StrContains(error_msg, "The number of elements exceeds the specified limit")) {
+          ResizeIfFull();
+          continue;
+        }
+        throw e;
+      }
+    }
   }
 
-  void Remove(DocId id) {
+  void Remove(T id) {
     try {
       world_.markDelete(id);
     } catch (const std::exception& e) {
     }
   }
 
-  vector<pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef) {
+  vector<pair<float, T>> Knn(float* target, size_t k, std::optional<size_t> ef) {
     world_.setEf(ef.value_or(kDefaultEfRuntime));
     return QueueToVec(world_.searchKnn(target, k));
   }
 
-  vector<pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef,
-                                 const vector<DocId>& allowed) {
+  vector<pair<float, T>> Knn(float* target, size_t k, std::optional<size_t> ef,
+                             const vector<T>& allowed) {
     struct BinsearchFilter : hnswlib::BaseFilterFunctor {
       virtual bool operator()(hnswlib::labeltype id) {
         return binary_search(allowed->begin(), allowed->end(), id);
       }
 
-      BinsearchFilter(const vector<DocId>* allowed) : allowed{allowed} {
+      BinsearchFilter(const vector<T>* allowed) : allowed{allowed} {
       }
-      const vector<DocId>* allowed;
+      const vector<T>* allowed;
     };
 
     world_.setEf(ef.value_or(kDefaultEfRuntime));
@@ -613,8 +625,8 @@ struct HnswlibAdapter {
     return visit([](auto& space) -> hnswlib::SpaceInterface<float>* { return &space; }, space_);
   }
 
-  template <typename Q> static vector<pair<float, DocId>> QueueToVec(Q queue) {
-    vector<pair<float, DocId>> out(queue.size());
+  template <typename Q> static vector<pair<float, T>> QueueToVec(Q queue) {
+    vector<pair<float, T>> out(queue.size());
     size_t idx = out.size();
     while (!queue.empty()) {
       out[--idx] = queue.top();
@@ -623,14 +635,37 @@ struct HnswlibAdapter {
     return out;
   }
 
+  void ResizeIfFull() {
+    {
+      absl::ReaderMutexLock lock(&resize_mutex_);
+      if (world_.getCurrentElementCount() < world_.getMaxElements() ||
+          (world_.allow_replace_deleted_ && world_.getDeletedCount() > 0)) {
+        return;
+      }
+    }
+    try {
+      absl::WriterMutexLock lock(&resize_mutex_);
+      if (world_.getCurrentElementCount() == world_.getMaxElements() &&
+          (!world_.allow_replace_deleted_ || world_.getDeletedCount() == 0)) {
+        auto max_elements = world_.getMaxElements();
+        world_.resizeIndex(max_elements * 2);
+        LOG(INFO) << "Resizing HNSW Index, current size: " << max_elements
+                  << ", expand by: " << max_elements * 2;
+      }
+    } catch (const std::exception& e) {
+      throw e;
+    }
+  }
+
   SpaceUnion space_;
   hnswlib::HierarchicalNSW<float> world_;
+  absl::Mutex resize_mutex_;
 };
 
 template <typename T>
 HnswVectorIndex<T>::HnswVectorIndex(const SchemaField::VectorParams& params,
                                     PMR_NS::memory_resource*)
-    : BaseVectorIndex<T>{params.dim, params.sim}, adapter_{make_unique<HnswlibAdapter>(params)} {
+    : BaseVectorIndex<T>{params.dim, params.sim}, adapter_{make_unique<HnswlibAdapter<T>>(params)} {
   DCHECK(params.use_hnsw);
   // TODO: Patch hnsw to use MR
 }
@@ -663,6 +698,7 @@ void HnswVectorIndex<T>::Remove(T id, const DocumentAccessor& doc, string_view f
 }
 
 template struct HnswVectorIndex<DocId>;
+template struct HnswVectorIndex<GlobalDocId>;
 
 GeoIndex::GeoIndex(PMR_NS::memory_resource* mr) : rtree_(make_unique<rtree>()) {
 }

src/core/search/indices.h

@@ -195,8 +195,9 @@ template <typename T> struct FlatVectorIndex : public BaseVectorIndex<T> {
 };
 
 extern template struct FlatVectorIndex<DocId>;
+extern template struct FlatVectorIndex<GlobalDocId>;
 
-struct HnswlibAdapter;
+template <typename T> struct HnswlibAdapter;
 
 template <typename T> struct HnswVectorIndex : public BaseVectorIndex<T> {
   HnswVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource* mr);
@@ -217,10 +218,11 @@ template <typename T> struct HnswVectorIndex : public BaseVectorIndex<T> {
   void AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) override;
 
  private:
-  std::unique_ptr<HnswlibAdapter> adapter_;
+  std::unique_ptr<HnswlibAdapter<T>> adapter_;
 };
 
 extern template struct HnswVectorIndex<DocId>;
+extern template struct HnswVectorIndex<GlobalDocId>;
 
 struct GeoIndex : public BaseIndex<DocId> {
   using point =

src/core/search/search.cc

@@ -503,24 +503,13 @@ void FieldIndices::CreateIndices(PMR_NS::memory_resource* mr) {
         indices_[field_ident] = make_unique<TagIndex>(mr, tparams);
         break;
       }
-      case SchemaField::VECTOR: {
-        unique_ptr<BaseVectorIndex<DocId>> vector_index;
-
-        DCHECK(holds_alternative<SchemaField::VectorParams>(field_info.special_params));
-        const auto& vparams = std::get<SchemaField::VectorParams>(field_info.special_params);
-
-        if (vparams.use_hnsw)
-          vector_index = make_unique<HnswVectorIndex<DocId>>(vparams, mr);
-        else
-          vector_index = make_unique<FlatVectorIndex<DocId>>(vparams, mr);
-
-        indices_[field_ident] = std::move(vector_index);
-        break;
-      }
       case SchemaField::GEO: {
         indices_[field_ident] = make_unique<GeoIndex>(mr);
         break;
       }
+      // We have global vector index
+      case SchemaField::VECTOR:
+        break;
     }
   }
 }
@@ -666,14 +655,21 @@ SearchResult SearchAlgorithm::Search(const FieldIndices* index, size_t cuttoff_l
   return bs.Search(*query_, cuttoff_limit);
 }
 
-optional<KnnScoreSortOption> SearchAlgorithm::GetKnnScoreSortOption() const {
-  DCHECK(query_);
-
-  // KNN query
-  if (auto* knn = get_if<AstKnnNode>(query_.get()); knn)
-    return KnnScoreSortOption{string_view{knn->score_alias}, knn->limit};
+bool SearchAlgorithm::IsKnnQuery() const {
+  return std::holds_alternative<AstKnnNode>(*query_);
+}
 
-  return nullopt;
+std::unique_ptr<AstNode> SearchAlgorithm::GetKnnNode() {
+  if (auto* knn = get_if<AstKnnNode>(query_.get()); knn) {
+    // Save knn score sort option
+    knn_score_sort_option_ = KnnScoreSortOption{string_view{knn->score_alias}, knn->limit};
+    auto node = std::move(query_);
+    if (!std::holds_alternative<AstStarNode>(*(knn)->filter))
+      query_.swap(knn->filter);
+    return node;
+  }
+  LOG(DFATAL) << "Should not reach here";
+  return nullptr;
 }
 
 void SearchAlgorithm::EnableProfiling() {

src/core/search/search.h

@@ -23,6 +23,7 @@ namespace dfly::search {
 
 struct AstNode;
 struct TextIndex;
+struct AstKnnNode;
 
 // Optional FILTER
 struct OptionalNumericFilter : public OptionalFilterBase {
@@ -201,14 +202,20 @@ class SearchAlgorithm {
   SearchResult Search(const FieldIndices* index,
                       size_t cuttoff_limit = std::numeric_limits<size_t>::max()) const;
 
-  // if enabled, return limit & alias for knn query
-  std::optional<KnnScoreSortOption> GetKnnScoreSortOption() const;
+  std::optional<KnnScoreSortOption> GetKnnScoreSortOption() const {
+    return knn_score_sort_option_;
+  }
+
+  bool IsKnnQuery() const;
+
+  std::unique_ptr<AstNode> GetKnnNode();
 
   void EnableProfiling();
 
  private:
   bool profiling_enabled_ = false;
   std::unique_ptr<AstNode> query_;
+  std::optional<KnnScoreSortOption> knn_score_sort_option_;
 };
 
 }  // namespace dfly::search

src/server/search/CMakeLists.txt

@@ -4,7 +4,7 @@ if (NOT WITH_SEARCH)
   return()
 endif()
 
-add_library(dfly_search_server aggregator.cc doc_accessors.cc doc_index.cc search_family.cc index_join.cc
+add_library(dfly_search_server aggregator.cc doc_accessors.cc doc_index.cc search_family.cc index_join.cc global_vector_index.cc
   ../cluster/coordinator.cc)
 target_link_libraries(dfly_search_server dfly_transaction dragonfly_lib dfly_facade redis_lib jsonpath TRDP::jsoncons)