Templated BaseIndex class

Author: mkaruza

src/core/search/base.h

@@ -79,16 +79,16 @@ struct DocumentAccessor {
 //
 // Queries should be done directly on subclasses with their distinc
 // query functions. All results for all index types should be sorted.
-struct BaseIndex {
+template <typename T> struct BaseIndex {
   virtual ~BaseIndex() = default;
 
   // Returns true if the document was added / indexed
-  virtual bool Add(DocId id, const DocumentAccessor& doc, std::string_view field) = 0;
-  virtual void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) = 0;
+  virtual bool Add(T id, const DocumentAccessor& doc, std::string_view field) = 0;
+  virtual void Remove(T id, const DocumentAccessor& doc, std::string_view field) = 0;
 
   // Returns documents that have non-null values for this field (used for @field:* queries)
   // Result must be sorted
-  virtual std::vector<DocId> GetAllDocsWithNonNullValues() const = 0;
+  virtual std::vector<T> GetAllDocsWithNonNullValues() const = 0;
 
   /* Called at the end of indexes rebuilding after all initial Add calls are done.
      Some indices may need to finalize internal structures. See RangeTree for example. */
@@ -97,10 +97,9 @@ struct BaseIndex {
 };
 
 // Base class for type-specific sorting indices.
-struct BaseSortIndex : BaseIndex {
-  virtual SortableValue Lookup(DocId doc) const = 0;
-  virtual std::vector<SortableValue> Sort(std::vector<DocId>* ids, size_t limit,
-                                          bool desc) const = 0;
+template <typename T> struct BaseSortIndex : BaseIndex<T> {
+  virtual SortableValue Lookup(T doc) const = 0;
+  virtual std::vector<SortableValue> Sort(std::vector<T>* ids, size_t limit, bool desc) const = 0;
 };
 
 /* Used in iterators of inverse indices.

src/core/search/indices.cc

@@ -473,14 +473,12 @@ absl::flat_hash_set<std::string> TagIndex::Tokenize(std::string_view value) cons
   return NormalizeTags(value, case_sensitive_, separator_);
 }
 
-BaseVectorIndex::BaseVectorIndex(size_t dim, VectorSimilarity sim) : dim_{dim}, sim_{sim} {
+template <typename T>
+BaseVectorIndex<T>::BaseVectorIndex(size_t dim, VectorSimilarity sim) : dim_{dim}, sim_{sim} {
 }
 
-std::pair<size_t /*dim*/, VectorSimilarity> BaseVectorIndex::Info() const {
-  return {dim_, sim_};
-}
-
-bool BaseVectorIndex::Add(DocId id, const DocumentAccessor& doc, std::string_view field) {
+template <typename T>
+bool BaseVectorIndex<T>::Add(T id, const DocumentAccessor& doc, std::string_view field) {
   auto vector = doc.GetVector(field);
   if (!vector)
     return false;
@@ -494,36 +492,40 @@ bool BaseVectorIndex::Add(DocId id, const DocumentAccessor& doc, std::string_vie
   return true;
 }
 
-FlatVectorIndex::FlatVectorIndex(const SchemaField::VectorParams& params,
-                                 PMR_NS::memory_resource* mr)
-    : BaseVectorIndex{params.dim, params.sim}, entries_{mr} {
+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);
 }
 
-void FlatVectorIndex::AddVector(DocId id, const VectorPtr& vector) {
-  DCHECK_LE(id * dim_, entries_.size());
-  if (id * dim_ == entries_.size())
-    entries_.resize((id + 1) * dim_);
+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_);
 
   // TODO: Let get vector write to buf itself
   if (vector) {
-    memcpy(&entries_[id * dim_], vector.get(), dim_ * sizeof(float));
+    memcpy(&entries_[id * BaseVectorIndex<T>::dim_], vector.get(),
+           BaseVectorIndex<T>::dim_ * sizeof(float));
   }
 }
 
-void FlatVectorIndex::Remove(DocId id, const DocumentAccessor& doc, string_view field) {
+template <typename T>
+void FlatVectorIndex<T>::Remove(T id, const DocumentAccessor& doc, string_view field) {
   // noop
 }
 
-const float* FlatVectorIndex::Get(DocId doc) const {
+template <typename T> const float* FlatVectorIndex<T>::Get(T doc) const {
   return &entries_[doc * dim_];
 }
 
-std::vector<DocId> FlatVectorIndex::GetAllDocsWithNonNullValues() const {
+template <typename T> std::vector<T> FlatVectorIndex<T>::GetAllDocsWithNonNullValues() const {
   std::vector<DocId> result;
 
-  size_t num_vectors = entries_.size() / dim_;
+  size_t num_vectors = entries_.size() / BaseVectorIndex<T>::dim_;
   result.reserve(num_vectors);
 
   for (DocId id = 0; id < num_vectors; ++id) {
@@ -533,7 +535,7 @@ std::vector<DocId> FlatVectorIndex::GetAllDocsWithNonNullValues() const {
     bool is_zero_vector = true;
 
     // TODO: Consider don't use check for zero vector
-    for (size_t i = 0; i < dim_; ++i) {
+    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;
@@ -550,6 +552,8 @@ std::vector<DocId> FlatVectorIndex::GetAllDocsWithNonNullValues() const {
   return result;
 }
 
+template struct FlatVectorIndex<DocId>;
+
 struct HnswlibAdapter {
   // Default setting of hnswlib/hnswalg
   constexpr static size_t kDefaultEfRuntime = 10;
@@ -623,35 +627,43 @@ struct HnswlibAdapter {
   hnswlib::HierarchicalNSW<float> world_;
 };
 
-HnswVectorIndex::HnswVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource*)
-    : BaseVectorIndex{params.dim, params.sim}, adapter_{make_unique<HnswlibAdapter>(params)} {
+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)} {
   DCHECK(params.use_hnsw);
   // TODO: Patch hnsw to use MR
 }
-
-HnswVectorIndex::~HnswVectorIndex() {
+template <typename T> HnswVectorIndex<T>::~HnswVectorIndex() {
 }
 
-void HnswVectorIndex::AddVector(DocId id, const VectorPtr& vector) {
+template <typename T>
+void HnswVectorIndex<T>::AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) {
   if (vector) {
     adapter_->Add(vector.get(), id);
   }
 }
 
-std::vector<std::pair<float, DocId>> HnswVectorIndex::Knn(float* target, size_t k,
-                                                          std::optional<size_t> ef) const {
+template <typename T>
+std::vector<std::pair<float, T>> HnswVectorIndex<T>::Knn(float* target, size_t k,
+                                                         std::optional<size_t> ef) const {
   return adapter_->Knn(target, k, ef);
 }
-std::vector<std::pair<float, DocId>> HnswVectorIndex::Knn(float* target, size_t k,
-                                                          std::optional<size_t> ef,
-                                                          const std::vector<DocId>& allowed) const {
+
+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 {
   return adapter_->Knn(target, k, ef, allowed);
 }
 
-void HnswVectorIndex::Remove(DocId id, const DocumentAccessor& doc, string_view field) {
+template <typename T>
+void HnswVectorIndex<T>::Remove(T id, const DocumentAccessor& doc, string_view field) {
   adapter_->Remove(id);
 }
 
+template struct HnswVectorIndex<DocId>;
+
 GeoIndex::GeoIndex(PMR_NS::memory_resource* mr) : rtree_(make_unique<rtree>()) {
 }
 

src/core/search/indices.h

@@ -39,7 +39,7 @@ namespace dfly::search {
 
 // Index for integer fields.
 // Range bounds are queried in logarithmic time, iteration is constant.
-struct NumericIndex : public BaseIndex {
+struct NumericIndex : public BaseIndex<DocId> {
   // Temporary base class for range tree.
   // It is used to use two different range trees depending on the flag use_range_tree.
   // If the flag is true, RangeTree is used, otherwise a simple implementation with btree_set.
@@ -76,7 +76,7 @@ struct NumericIndex : public BaseIndex {
 };
 
 // Base index for string based indices.
-template <typename C> struct BaseStringIndex : public BaseIndex {
+template <typename C> struct BaseStringIndex : public BaseIndex<DocId> {
   using Container = BlockList<C>;
   using VecOrPtr = std::variant<std::vector<DocId>, const Container*>;
 
@@ -157,65 +157,72 @@ struct TagIndex : public BaseStringIndex<SortedVector<DocId>> {
   char separator_;
 };
 
-struct BaseVectorIndex : public BaseIndex {
-  std::pair<size_t /*dim*/, VectorSimilarity> Info() const;
+template <typename T> struct BaseVectorIndex : public BaseIndex<T> {
+  std::pair<size_t /*dim*/, VectorSimilarity> Info() const {
+    return {dim_, sim_};
+  }
 
-  bool Add(DocId id, const DocumentAccessor& doc, std::string_view field) override final;
+  bool Add(T id, const DocumentAccessor& doc, std::string_view field) override final;
 
  protected:
   BaseVectorIndex(size_t dim, VectorSimilarity sim);
 
   using VectorPtr = decltype(std::declval<OwnedFtVector>().first);
-  virtual void AddVector(DocId id, const VectorPtr& vector) = 0;
+  virtual void AddVector(T id, const VectorPtr& vector) = 0;
 
   size_t dim_;
   VectorSimilarity sim_;
 };
 
 // Index for vector fields.
 // Only supports lookup by id.
-struct FlatVectorIndex : public BaseVectorIndex {
+template <typename T> struct FlatVectorIndex : public BaseVectorIndex<T> {
   FlatVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource* mr);
 
-  void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) override;
+  void Remove(T id, const DocumentAccessor& doc, std::string_view field) override;
 
-  const float* Get(DocId doc) const;
+  const float* Get(T doc) const;
 
   // Return all documents that have vectors in this index
-  std::vector<DocId> GetAllDocsWithNonNullValues() const override;
+  std::vector<T> GetAllDocsWithNonNullValues() const override;
 
  protected:
-  void AddVector(DocId id, const VectorPtr& vector) override;
+  using BaseVectorIndex<T>::dim_;
+  void AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) override;
 
  private:
   PMR_NS::vector<float> entries_;
 };
 
+extern template struct FlatVectorIndex<DocId>;
+
 struct HnswlibAdapter;
 
-struct HnswVectorIndex : public BaseVectorIndex {
+template <typename T> struct HnswVectorIndex : public BaseVectorIndex<T> {
   HnswVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource* mr);
   ~HnswVectorIndex();
 
-  void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) override;
+  void Remove(T id, const DocumentAccessor& doc, std::string_view field) override;
 
-  std::vector<std::pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef) const;
-  std::vector<std::pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef,
-                                           const std::vector<DocId>& allowed) const;
+  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;
 
   // TODO: Implement if needed
-  std::vector<DocId> GetAllDocsWithNonNullValues() const override {
-    return std::vector<DocId>{};
+  std::vector<T> GetAllDocsWithNonNullValues() const override {
+    return std::vector<T>{};
   }
 
  protected:
-  void AddVector(DocId id, const VectorPtr& vector) override;
+  void AddVector(T id, const typename BaseVectorIndex<T>::VectorPtr& vector) override;
 
  private:
   std::unique_ptr<HnswlibAdapter> adapter_;
 };
 
-struct GeoIndex : public BaseIndex {
+extern template struct HnswVectorIndex<DocId>;
+
+struct GeoIndex : public BaseIndex<DocId> {
   using point =
       boost::geometry::model::point<double, 2,
                                     boost::geometry::cs::geographic<boost::geometry::degree>>;