WIP

Author: mkaruza

src/core/search/indices.cc

@@ -11,6 +11,11 @@
 #include <absl/strings/str_split.h>
 
 #include <boost/iterator/function_output_iterator.hpp>
+#include <shared_mutex>
+
+#include "core/search/base.h"
+#include "core/search/vector_utils.h"
+#include "util/fibers/synchronization.h"
 
 #define UNI_ALGO_DISABLE_NFKC_NFKD
 
@@ -492,58 +497,116 @@ bool BaseVectorIndex<T>::Add(T id, const DocumentAccessor& doc, std::string_view
   return true;
 }
 
-template <typename T>
-FlatVectorIndex<T>::FlatVectorIndex(const SchemaField::VectorParams& params,
-                                    PMR_NS::memory_resource* mr)
-    : BaseVectorIndex<T>{params.dim, params.sim}, entries_{mr} {
-  DCHECK(!params.use_hnsw);
-  entries_.reserve(params.capacity * params.dim);
+ShardNoOpVectorIndex::ShardNoOpVectorIndex(const SchemaField::VectorParams& params)
+    : BaseVectorIndex<DocId>{params.dim, params.sim} {
 }
 
-template <typename T>
-void FlatVectorIndex<T>::AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) {
-  DCHECK_LE(id * BaseVectorIndex<T>::dim_, entries_.size());
-  if (id * BaseVectorIndex<T>::dim_ == entries_.size())
-    entries_.resize((id + 1) * BaseVectorIndex<T>::dim_);
+FlatVectorIndex::FlatVectorIndex(const SchemaField::VectorParams& params, ShardId shard_set_size,
+                                 PMR_NS::memory_resource* mr)
+    : BaseVectorIndex<GlobalDocId>{params.dim, params.sim},
+      entries_{mr},
+      shard_vector_locks_(shard_set_size) {
+  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);
+  }
+}
 
-  // TODO: Let get vector write to buf itself
+void FlatVectorIndex::AddVector(GlobalDocId id,
+                                const typename BaseVectorIndex<GlobalDocId>::VectorPtr& vector) {
+  auto shard_id = search::GlobalDocIdShardId(id);
+  auto shard_doc_id = search::GlobalDocIdLocalId(id);
+  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_);
+  }
   if (vector) {
-    memcpy(&entries_[id * BaseVectorIndex<T>::dim_], vector.get(),
-           BaseVectorIndex<T>::dim_ * sizeof(float));
+    memcpy(&entries_[shard_id][shard_doc_id * BaseVectorIndex<GlobalDocId>::dim_], vector.get(),
+           BaseVectorIndex<GlobalDocId>::dim_ * sizeof(float));
   }
 }
 
-template <typename T>
-void FlatVectorIndex<T>::Remove(T id, const DocumentAccessor& doc, string_view field) {
+void FlatVectorIndex::Remove(GlobalDocId id, const DocumentAccessor& doc, string_view field) {
   // noop
 }
 
-template <typename T> const float* FlatVectorIndex<T>::Get(T doc) const {
-  return &entries_[doc * dim_];
+const float* FlatVectorIndex::Get(GlobalDocId doc) const {
+  ShardId shard_id = search::GlobalDocIdShardId(doc);
+  shared_lock<util::fb2::SharedMutex> lock{shard_vector_locks_[shard_id]};
+  return &entries_[shard_id][search::GlobalDocIdLocalId(doc) * dim_];
+}
+
+std::vector<std::pair<float, GlobalDocId>> FlatVectorIndex::Knn(float* target) const {
+  std::priority_queue<std::pair<float, search::GlobalDocId>> queue;
+
+  for (size_t shard_id = 0; shard_id < entries_.size(); shard_id++) {
+    shared_lock<util::fb2::SharedMutex> lock{shard_vector_locks_[shard_id]};
+    size_t num_vectors = entries_[shard_id].size() / BaseVectorIndex<GlobalDocId>::dim_;
+    for (GlobalDocId 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 = &entries_[shard_id][id * dim_];
+      float dist = VectorDistance(target, vec, dim_, sim_);
+      queue.emplace(dist, CreateGlobalDocId(shard_id, id));
+    }
+  }
+
+  vector<pair<float, search::GlobalDocId>> out(queue.size());
+  size_t idx = out.size();
+  while (!queue.empty()) {
+    out[--idx] = queue.top();
+    queue.pop();
+  }
+
+  return out;
 }
 
-template <typename T> std::vector<T> FlatVectorIndex<T>::GetAllDocsWithNonNullValues() const {
-  std::vector<T> result;
+std::vector<std::pair<float, GlobalDocId>> FlatVectorIndex::Knn(
+    float* target, const std::vector<GlobalDocId>& allowed) const {
+  std::priority_queue<std::pair<float, search::GlobalDocId>> queue;
+
+  for (auto& doc : allowed) {
+    uint16_t shard_id = search::GlobalDocIdShardId(doc);
+    auto shard_doc_id = search::GlobalDocIdLocalId(doc);
+    shared_lock<util::fb2::SharedMutex> lock{shard_vector_locks_[shard_id]};
+    const float* vec = &entries_[shard_id][shard_doc_id * dim_];
+    float dist = VectorDistance(target, vec, dim_, sim_);
+    queue.emplace(dist, doc);
+  }
+
+  vector<pair<float, search::GlobalDocId>> out(queue.size());
+  size_t idx = out.size();
+  while (!queue.empty()) {
+    out[--idx] = queue.top();
+    queue.pop();
+  }
 
-  size_t num_vectors = entries_.size() / BaseVectorIndex<T>::dim_;
-  result.reserve(num_vectors);
+  return out;
+}
 
-  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);
-    bool is_zero_vector = true;
+std::vector<GlobalDocId> FlatVectorIndex::GetAllDocsWithNonNullValues() const {
+  std::vector<GlobalDocId> result;
+  for (size_t shard_id = 0; shard_id < entries_.size(); shard_id++) {
+    shared_lock<util::fb2::SharedMutex> lock{shard_vector_locks_[shard_id]};
+    size_t num_vectors = entries_[shard_id].size() / BaseVectorIndex<GlobalDocId>::dim_;
+    for (GlobalDocId 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 = &entries_[shard_id][id * dim_];
+      bool is_zero_vector = true;
 
-    // TODO: Consider don't use check for zero vector
-    for (size_t i = 0; i < BaseVectorIndex<T>::dim_; ++i) {
-      if (vec[i] != 0.0f) {  // TODO: Consider using a threshold for float comparison
-        is_zero_vector = false;
-        break;
+      // TODO: Consider don't use check for zero vector
+      for (size_t i = 0; i < BaseVectorIndex<GlobalDocId>::dim_; ++i) {
+        if (vec[i] != 0.0f) {  // TODO: Consider using a threshold for float comparison
+          is_zero_vector = false;
+          break;
+        }
       }
-    }
 
-    if (!is_zero_vector) {
-      result.push_back(id);
+      if (!is_zero_vector) {
+        result.push_back(CreateGlobalDocId(shard_id, id));
+      }
     }
   }
 
@@ -552,9 +615,6 @@ template <typename T> std::vector<T> FlatVectorIndex<T>::GetAllDocsWithNonNullVa
   return result;
 }
 
-template struct FlatVectorIndex<DocId>;
-template struct FlatVectorIndex<GlobalDocId>;
-
 template <typename T> struct HnswlibAdapter {
   // Default setting of hnswlib/hnswalg
   constexpr static size_t kDefaultEfRuntime = 10;
@@ -662,44 +722,37 @@ template <typename T> struct HnswlibAdapter {
   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<T>>(params)} {
+HnswVectorIndex::HnswVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource*)
+    : BaseVectorIndex<GlobalDocId>{params.dim, params.sim},
+      adapter_{make_unique<HnswlibAdapter<GlobalDocId>>(params)} {
   DCHECK(params.use_hnsw);
   // TODO: Patch hnsw to use MR
 }
-template <typename T> HnswVectorIndex<T>::~HnswVectorIndex() {
+HnswVectorIndex::~HnswVectorIndex() {
 }
 
-template <typename T>
-void HnswVectorIndex<T>::AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) {
+void HnswVectorIndex::AddVector(GlobalDocId id,
+                                const typename BaseVectorIndex<GlobalDocId>::VectorPtr& vector) {
   if (vector) {
     adapter_->Add(vector.get(), id);
   }
 }
 
-template <typename T>
-std::vector<std::pair<float, T>> HnswVectorIndex<T>::Knn(float* target, size_t k,
-                                                         std::optional<size_t> ef) const {
+std::vector<std::pair<float, GlobalDocId>> HnswVectorIndex::Knn(float* target, size_t k,
+                                                                std::optional<size_t> ef) const {
   return adapter_->Knn(target, k, ef);
 }
 
-template <typename T>
-std::vector<std::pair<float, T>> HnswVectorIndex<T>::Knn(float* target, size_t k,
-                                                         std::optional<size_t> ef,
-                                                         const std::vector<T>& allowed) const {
+std::vector<std::pair<float, GlobalDocId>> HnswVectorIndex::Knn(
+    float* target, size_t k, std::optional<size_t> ef,
+    const std::vector<GlobalDocId>& allowed) const {
   return adapter_->Knn(target, k, ef, allowed);
 }
 
-template <typename T>
-void HnswVectorIndex<T>::Remove(T id, const DocumentAccessor& doc, string_view field) {
+void HnswVectorIndex::Remove(GlobalDocId id, const DocumentAccessor& doc, string_view field) {
   adapter_->Remove(id);
 }
 
-template struct HnswVectorIndex<DocId>;
-template struct HnswVectorIndex<GlobalDocId>;
-
 GeoIndex::GeoIndex(PMR_NS::memory_resource* mr) : rtree_(make_unique<rtree>()) {
 }
 

src/core/search/indices.h

@@ -8,6 +8,9 @@
 #include <absl/container/flat_hash_map.h>
 #include <absl/container/flat_hash_set.h>
 
+//#include "server/search/global_vector_index.h"
+#include "util/fibers/synchronization.h"
+
 // Wrong warning reported when geometry.hpp is loaded
 #ifndef __clang__
 #pragma GCC diagnostic push
@@ -174,56 +177,80 @@ template <typename T> struct BaseVectorIndex : public BaseIndex<T> {
   VectorSimilarity sim_;
 };
 
+// ShardNoOpVectorIndex is used as placeholder as vector index in each shard. It doesn't implement
+// any functionality so adding documents will not have any effect on it. It is used to support
+// as filter when adding fields.
+struct ShardNoOpVectorIndex : public BaseVectorIndex<DocId> {
+  explicit ShardNoOpVectorIndex(const SchemaField::VectorParams& params);
+
+  void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) override {
+    // noop
+  }
+
+  // Return all documents that have vectors in this index
+  std::vector<DocId> GetAllDocsWithNonNullValues() const override {
+    return {};
+  }
+
+ protected:
+  using BaseVectorIndex<DocId>::dim_;
+  void AddVector(DocId id, const typename BaseVectorIndex<DocId>::VectorPtr& vector) override {
+    // noop
+  }
+};
+
 // Index for vector fields.
 // Only supports lookup by id.
-template <typename T> struct FlatVectorIndex : public BaseVectorIndex<T> {
-  FlatVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource* mr);
+struct FlatVectorIndex : public BaseVectorIndex<GlobalDocId> {
+  FlatVectorIndex(const SchemaField::VectorParams& params, ShardId shard_set_size,
+                  PMR_NS::memory_resource* mr);
+
+  void Remove(GlobalDocId id, const DocumentAccessor& doc, std::string_view field) override;
 
-  void Remove(T id, const DocumentAccessor& doc, std::string_view field) override;
+  const float* Get(GlobalDocId doc) const;
 
-  const float* Get(T doc) const;
+  std::vector<std::pair<float, GlobalDocId>> Knn(float* target) const;
+  std::vector<std::pair<float, GlobalDocId>> Knn(float* target,
+                                                 const std::vector<GlobalDocId>& allowed) const;
 
   // Return all documents that have vectors in this index
-  std::vector<T> GetAllDocsWithNonNullValues() const override;
+  std::vector<GlobalDocId> GetAllDocsWithNonNullValues() const override;
 
  protected:
-  using BaseVectorIndex<T>::dim_;
-  void AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) override;
+  using BaseVectorIndex<GlobalDocId>::dim_;
+  void AddVector(GlobalDocId id,
+                 const typename BaseVectorIndex<GlobalDocId>::VectorPtr& vector) override;
 
  private:
-  PMR_NS::vector<float> entries_;
+  PMR_NS::vector<PMR_NS::vector<float>> entries_;
+  mutable std::vector<util::fb2::SharedMutex> shard_vector_locks_;
 };
 
-extern template struct FlatVectorIndex<DocId>;
-extern template struct FlatVectorIndex<GlobalDocId>;
-
 template <typename T> struct HnswlibAdapter;
-
-template <typename T> struct HnswVectorIndex : public BaseVectorIndex<T> {
+struct HnswVectorIndex : public BaseVectorIndex<GlobalDocId> {
   HnswVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource* mr);
   ~HnswVectorIndex();
 
-  void Remove(T id, const DocumentAccessor& doc, std::string_view field) override;
+  void Remove(GlobalDocId id, const DocumentAccessor& doc, std::string_view field) override;
 
-  std::vector<std::pair<float, T>> Knn(float* target, size_t k, std::optional<size_t> ef) const;
-  std::vector<std::pair<float, T>> Knn(float* target, size_t k, std::optional<size_t> ef,
-                                       const std::vector<T>& allowed) const;
+  std::vector<std::pair<float, GlobalDocId>> Knn(float* target, size_t k,
+                                                 std::optional<size_t> ef) const;
+  std::vector<std::pair<float, GlobalDocId>> Knn(float* target, size_t k, std::optional<size_t> ef,
+                                                 const std::vector<GlobalDocId>& allowed) const;
 
   // TODO: Implement if needed
-  std::vector<T> GetAllDocsWithNonNullValues() const override {
-    return std::vector<T>{};
+  std::vector<GlobalDocId> GetAllDocsWithNonNullValues() const override {
+    return std::vector<GlobalDocId>{};
   }
 
  protected:
-  void AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) override;
+  void AddVector(GlobalDocId id,
+                 const typename BaseVectorIndex<GlobalDocId>::VectorPtr& vector) override;
 
  private:
-  std::unique_ptr<HnswlibAdapter<T>> adapter_;
+  std::unique_ptr<HnswlibAdapter<GlobalDocId>> adapter_;
 };
 
-extern template struct HnswVectorIndex<DocId>;
-extern template struct HnswVectorIndex<GlobalDocId>;
-
 struct GeoIndex : public BaseIndex<DocId> {
   using point =
       boost::geometry::model::point<double, 2,