Global 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::Filter() 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 Filter() 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/search.cc

@@ -336,68 +336,10 @@ struct BasicSearch {
     return UnifyResults(GetSubResults(node.tags, mapping), LogicOp::OR);
   }
 
-  void SearchKnnFlat(FlatVectorIndex<DocId>* vec_index, const AstKnnNode& knn,
-                     IndexResult&& sub_results) {
-    knn_distances_.reserve(sub_results.ApproximateSize());
-    auto cb = [&](auto* set) {
-      auto [dim, sim] = vec_index->Info();
-      for (DocId matched_doc : *set) {
-        float dist = VectorDistance(knn.vec.first.get(), vec_index->Get(matched_doc), dim, sim);
-        knn_distances_.emplace_back(dist, matched_doc);
-      }
-    };
-    visit(cb, sub_results.Borrowed());
-
-    size_t prefix_size = min(knn.limit, knn_distances_.size());
-    partial_sort(knn_distances_.begin(), knn_distances_.begin() + prefix_size,
-                 knn_distances_.end());
-    knn_distances_.resize(prefix_size);
-  }
-
-  void SearchKnnHnsw(HnswVectorIndex<DocId>* vec_index, const AstKnnNode& knn,
-                     IndexResult&& sub_results) {
-    if (indices_->GetAllDocs().size() == sub_results.ApproximateSize())  // TODO: remove approx size
-      knn_distances_ = vec_index->Knn(knn.vec.first.get(), knn.limit, knn.ef_runtime);
-    else
-      knn_distances_ =
-          vec_index->Knn(knn.vec.first.get(), knn.limit, knn.ef_runtime, sub_results.Take().first);
-  }
-
   // [KNN limit @field vec]: Compute distance from `vec` to all vectors keep closest `limit`
   IndexResult Search(const AstKnnNode& knn, string_view active_field) {
-    DCHECK(active_field.empty());
-    auto sub_results = SearchGeneric(*knn.filter, active_field);
-
-    auto* vec_index = GetIndex<BaseVectorIndex<DocId>>(knn.field);
-    if (!vec_index)
-      return IndexResult{};
-
-    // If vector dimension is 0, treat as placeholder/invalid - return empty results
-    // This allows tests to use dummy vector values like "<your_vector_blob>"
-    if (knn.vec.second == 0)
-      return IndexResult{};
-
-    if (auto [dim, _] = vec_index->Info(); dim != knn.vec.second) {
-      error_ =
-          absl::StrCat("Wrong vector index dimensions, got: ", knn.vec.second, ", expected: ", dim);
-      return IndexResult{};
-    }
-
-    knn_scores_.clear();
-    if (auto hnsw_index = dynamic_cast<HnswVectorIndex<DocId>*>(vec_index); hnsw_index)
-      SearchKnnHnsw(hnsw_index, knn, std::move(sub_results));
-    else
-      SearchKnnFlat(dynamic_cast<FlatVectorIndex<DocId>*>(vec_index), knn, std::move(sub_results));
-
-    vector<DocId> out(knn_distances_.size());
-    knn_scores_.reserve(knn_distances_.size());
-
-    for (size_t i = 0; i < knn_distances_.size(); i++) {
-      knn_scores_.emplace_back(knn_distances_[i].second, knn_distances_[i].first);
-      out[i] = knn_distances_[i].second;
-    }
-
-    return IndexResult{std::move(out)};
+    LOG(DFATAL) << "KNN node should not be searched in shard";
+    return IndexResult{};
   }
 
   // Determine node type and call specific search function
@@ -503,24 +445,15 @@ 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;
       }
+      case SchemaField::VECTOR: {
+        const auto& vparams = std::get<SchemaField::VectorParams>(field_info.special_params);
+        indices_[field_ident] = make_unique<ShardNoOpVectorIndex>(vparams);
+        break;
+      }
     }
   }
 }
@@ -666,14 +599,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;
+  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/main_service.cc

@@ -59,6 +59,7 @@ extern "C" {
 #include "server/multi_command_squasher.h"
 #include "server/namespaces.h"
 #include "server/script_mgr.h"
+#include "server/search/global_vector_index.h"
 #include "server/search/search_family.h"
 #include "server/server_state.h"
 #include "server/set_family.h"
@@ -1123,6 +1124,11 @@ void Service::Shutdown() {
 
   shard_set->PreShutdown();
   shard_set->Shutdown();
+
+#ifdef WITH_SEARCH
+  GlobalVectorIndexRegistry::Instance().Reset();
+#endif
+
   Transaction::Shutdown();
 
   pp_.AwaitFiberOnAll([](ProactorBase* pb) { ServerState::tlocal()->Destroy(); });

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)