WIP 2

Author: mkaruza

src/core/search/indices.cc

@@ -509,14 +509,15 @@ FlatVectorIndex::FlatVectorIndex(const SchemaField::VectorParams& params, ShardI
   DCHECK(!params.use_hnsw);
   entries_.resize(shard_set_size);
   for (size_t i = 0; i < shard_set_size; i++) {
-    entries_[i].resize(params.capacity * params.dim);
+    entries_[i].reserve(params.capacity * params.dim);
   }
 }
 
 void FlatVectorIndex::AddVector(GlobalDocId id,
                                 const typename BaseVectorIndex<GlobalDocId>::VectorPtr& vector) {
   auto shard_id = search::GlobalDocIdShardId(id);
   auto shard_doc_id = search::GlobalDocIdLocalId(id);
+  DCHECK_LE(shard_doc_id * BaseVectorIndex<GlobalDocId>::dim_, entries_[shard_id].size());
   if (shard_doc_id * BaseVectorIndex<GlobalDocId>::dim_ == entries_[shard_id].size()) {
     unique_lock<util::fb2::SharedMutex> lock{shard_vector_locks_[shard_id]};
     entries_[shard_id].resize((shard_doc_id + 1) * BaseVectorIndex<GlobalDocId>::dim_);

src/core/search/search.h

@@ -202,7 +202,7 @@ class SearchAlgorithm {
   SearchResult Search(const FieldIndices* index,
                       size_t cuttoff_limit = std::numeric_limits<size_t>::max()) const;
 
-  std::optional<KnnScoreSortOption> GetKnnScoreSortOption() const {
+  const std::optional<KnnScoreSortOption>& GetKnnScoreSortOption() const {
     return knn_score_sort_option_;
   }
 

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/doc_index.cc

@@ -9,6 +9,7 @@
 #include <memory>
 #include <optional>
 #include <queue>
+#include <string_view>
 
 #include "absl/strings/str_cat.h"
 #include "base/logging.h"
@@ -382,11 +383,9 @@ void ShardDocIndex::AddDocToGlobalVectorIndex(std::string_view index_name,
   for (const auto& [field_ident, field_info] : base_->schema.fields) {
     if (field_info.type == search::SchemaField::VECTOR &&
         !(field_info.flags & search::SchemaField::NOINDEX)) {
-      if (auto vector_info = accessor->GetVector(field_ident); vector_info && vector_info->first) {
-        auto global_index =
-            GlobalVectorIndexRegistry::Instance().GetVectorIndex(index_name, field_info.short_name);
-        global_index->Add(global_id, *accessor, field_ident);
-      }
+      GlobalVectorIndexRegistry::Instance()
+          .GetVectorIndex(index_name, field_info.short_name)
+          ->Add(global_id, *accessor, field_ident);
     }
   }
 }
@@ -411,18 +410,18 @@ void ShardDocIndex::RemoveDocFromGlobalVectorIndex(std::string_view index_name,
 void ShardDocIndex::RebuildGlobalVectorIndices(std::string_view index_name, const OpArgs& op_args) {
   auto cb = [this, index_name](string_view key, const BaseAccessor& doc) {
     auto local_id = key_index_.Find(key);
+
     if (!local_id)
       return;
+
     GlobalDocId global_id = search::CreateGlobalDocId(EngineShard::tlocal()->shard_id(), *local_id);
 
     for (const auto& [field_ident, field_info] : base_->schema.fields) {
       if (field_info.type == search::SchemaField::VECTOR &&
           !(field_info.flags & search::SchemaField::NOINDEX)) {
-        if (auto vector_info = doc.GetVector(field_ident); vector_info && vector_info->first) {
-          auto global_index = GlobalVectorIndexRegistry::Instance().GetVectorIndex(
-              index_name, field_info.short_name);
-          global_index->Add(global_id, doc, field_ident);
-        }
+        GlobalVectorIndexRegistry::Instance()
+            .GetVectorIndex(index_name, field_info.short_name)
+            ->Add(global_id, doc, field_ident);
       }
     }
   };

src/server/search/global_vector_index.cc

@@ -66,10 +66,10 @@ std::vector<std::pair<float, search::GlobalDocId>> GlobalVectorIndex::SearchKnnF
     return index->Knn(knn->vec.first.get());
 }
 
-std::vector<SearchResult> GlobalVectorIndex::Search(const search::AstKnnNode* knn_node,
-                                                    const std::vector<SearchResult>& filter_docs,
-                                                    const CommandContext& cmd_cntx,
-                                                    const SearchParams& params) {
+std::vector<SearchResult> GlobalVectorIndex::Search(
+    const search::AstKnnNode* knn_node, const std::vector<SearchResult>& filter_docs,
+    const std::optional<search::KnnScoreSortOption>& knn_score_option,
+    const CommandContext& cmd_cntx, const SearchParams& params) {
   std::vector<SearchResult> results(1);
 
   std::optional<std::vector<search::GlobalDocId>> filter_docs_global_ids = std::nullopt;
@@ -105,7 +105,8 @@ std::vector<SearchResult> GlobalVectorIndex::Search(const search::AstKnnNode* kn
     return results;
   }
 
-  std::vector<SerializedSearchDoc> knn_result_docs(knn_results.size());
+  std::vector<SerializedSearchDoc> knn_result_docs;
+  knn_result_docs.reserve(knn_results.size());
 
   // Group by shard with minimal allocations
   std::vector<std::vector<std::pair<float, search::DocId>>> shard_doc_ids(shard_size);
@@ -148,28 +149,56 @@ std::vector<SearchResult> GlobalVectorIndex::Search(const search::AstKnnNode* kn
     // Cache schema reference to avoid repeated lookups
     const auto& schema = index->GetInfo().base_index.schema;
 
+    auto& sort_option = params.sort_option;
+    bool need_fetch_sort_field = false;
+
+    if (sort_option) {
+      need_fetch_sort_field = !params.sort_option->IsSame(*knn_score_option);
+    }
+
     // Optimize serialization based on query type
     if (params.ShouldReturnAllFields()) {
       // Full serialization for full queries
       for (const auto& [score, doc_id] : shard_requests) {
         if (auto entry = index->LoadEntry(doc_id, t->GetOpArgs(es))) {
           auto& [key, accessor] = *entry;
           auto fields = accessor->Serialize(schema);
+
+          search::SortableValue sort_score = std::monostate{};
+
+          if (need_fetch_sort_field) {
+            sort_score = fields[params.sort_option->field.Name()];
+          }
+
           docs_for_shard.push_back(
-              {doc_id, std::string{key}, std::move(fields), score, std::monostate{}});
+              {doc_id, std::string{key}, std::move(fields), score, sort_score});
         }
       }
     } else {
       // Selective field serialization
       const auto& return_fields = params.return_fields.value_or(std::vector<FieldReference>{});
+      bool sort_field_in_return_fields = false;
+
+      if (need_fetch_sort_field) {
+        for (auto& field_reference : return_fields) {
+          sort_field_in_return_fields = field_reference.Name() == sort_option->field.Name();
+        }
+      }
+
       for (const auto& [score, doc_id] : shard_requests) {
         if (auto entry = index->LoadEntry(doc_id, t->GetOpArgs(es))) {
           auto& [key, accessor] = *entry;
           auto fields = return_fields.empty() ? SearchDocData{}
                                               // NOCONTENT query - no fields needed
                                               : accessor->Serialize(schema, return_fields);
+
+          search::SortableValue sort_score = std::monostate{};
+          if (sort_field_in_return_fields && need_fetch_sort_field) {
+            sort_score = fields[params.sort_option->field.Name()];
+          }
+
           docs_for_shard.push_back(
-              {doc_id, std::string{key}, std::move(fields), score, std::monostate{}});
+              {doc_id, std::string{key}, std::move(fields), score, sort_score});
         }
       }
     }
@@ -182,6 +211,12 @@ std::vector<SearchResult> GlobalVectorIndex::Search(const search::AstKnnNode* kn
                            std::make_move_iterator(docs.end()));
   }
 
+  auto cb = [](SerializedSearchDoc& l, SerializedSearchDoc& r) {
+    return l.knn_score > r.knn_score;
+  };
+  partial_sort(knn_result_docs.begin(), knn_result_docs.begin() + knn_result_docs.size() / 2,
+               knn_result_docs.end(), cb);
+
   results[0].total_hits = knn_results.size();
   results[0].docs = std::move(knn_result_docs);
   return results;
@@ -230,6 +265,11 @@ std::shared_ptr<GlobalVectorIndex> GlobalVectorIndexRegistry::GetVectorIndex(
   return it != indices_.end() ? it->second : nullptr;
 }
 
+void GlobalVectorIndexRegistry::Reset() {
+  std::unique_lock<std::shared_mutex> lock(registry_mutex_);
+  indices_.clear();
+}
+
 std::string GlobalVectorIndexRegistry::MakeKey(std::string_view index_name,
                                                std::string_view field_name) const {
   return absl::StrCat(index_name, ":", field_name);

src/server/search/global_vector_index.h

@@ -21,6 +21,8 @@
 
 namespace dfly {
 
+struct KnnScoreSortOption;
+
 class GlobalVectorIndex {
  public:
   GlobalVectorIndex(const search::SchemaField::VectorParams& params, std::string_view index_name,
@@ -31,9 +33,10 @@ class GlobalVectorIndex {
   bool Add(search::GlobalDocId id, const search::DocumentAccessor& doc, std::string_view field);
   void Remove(search::GlobalDocId id, const search::DocumentAccessor& doc, std::string_view field);
 
-  std::vector<SearchResult> Search(const search::AstKnnNode* knn,
-                                   const std::vector<SearchResult>& filter_docs,
-                                   const CommandContext& cmd_cntx, const SearchParams& params);
+  std::vector<SearchResult> Search(
+      const search::AstKnnNode* knn, const std::vector<SearchResult>& filter_docs,
+      const std::optional<search::KnnScoreSortOption>& knn_score_option,
+      const CommandContext& cmd_cntx, const SearchParams& params);
 
   std::pair<size_t, search::VectorSimilarity> Info() const;
 
@@ -70,6 +73,9 @@ class GlobalVectorIndexRegistry {
   std::shared_ptr<GlobalVectorIndex> GetVectorIndex(std::string_view index_name,
                                                     std::string_view field_name) const;
 
+  // Reset all vector indices
+  void Reset();
+
  private:
   GlobalVectorIndexRegistry() = default;
 

src/server/search/search_family.cc

@@ -1394,7 +1394,7 @@ void SearchFamily::FtSearch(CmdArgList args, const CommandContext& cmd_cntx) {
     if (!vector_index)
       return builder->SendError(string{index_name} + ": no such index");
 
-    docs = vector_index->Search(knn, docs, cmd_cntx, *params);
+    docs = vector_index->Search(knn, docs, search_algo.GetKnnScoreSortOption(), cmd_cntx, *params);
   }
 
   SearchReply(*params, search_algo.GetKnnScoreSortOption(), absl::MakeSpan(docs), builder);