chore(search_family): Minor performance improvements to GetAllResults methods. FIRST PR (#5341)

* fix(search_family): Improve performance of GetAllResults methods

Signed-off-by: Stepan Bagritsevich <[email protected]>

* refactor: address comments

Signed-off-by: Stepan Bagritsevich <[email protected]>

---------

Signed-off-by: Stepan Bagritsevich <[email protected]>

Author: BagritsevichStepan

src/core/search/base.h

@@ -13,6 +13,7 @@
 #include <string_view>
 #include <vector>
 
+#include "absl/container/flat_hash_set.h"
 #include "base/pmr/memory_resource.h"
 #include "core/string_map.h"
 
@@ -80,6 +81,11 @@ struct DocumentAccessor {
   virtual std::optional<StringList> GetTags(std::string_view active_field) const = 0;
 };
 
+// Represents a set of document IDs, used for merging results of inverse indices.
+template <typename Allocator = std::allocator<DocId>>
+using UniqueDocsList = absl::flat_hash_set<DocId, absl::DefaultHashContainerHash<DocId>,
+                                           absl::DefaultHashContainerEq<DocId>, Allocator>;
+
 // Base class for type-specific indices.
 //
 // Queries should be done directly on subclasses with their distinc
@@ -92,9 +98,8 @@ struct BaseIndex {
   virtual void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) = 0;
 
   // Returns documents that have non-null values for this field (used for @field:* queries)
-  virtual std::optional<std::vector<DocId>> GetAllResults() const {
-    return std::nullopt;
-  }
+  // Result must be sorted
+  virtual std::vector<DocId> GetAllDocsWithNonNullValues() const = 0;
 };
 
 // Base class for type-specific sorting indices.

src/core/search/indices.cc

@@ -115,6 +115,24 @@ vector<DocId> NumericIndex::Range(double l, double r) const {
   return out;
 }
 
+vector<DocId> NumericIndex::GetAllDocsWithNonNullValues() const {
+  UniqueDocsList<> unique_docs;
+  std::vector<DocId> result;
+
+  unique_docs.reserve(entries_.size());
+  result.reserve(entries_.size());
+
+  for (const auto& [_, doc_id] : entries_) {
+    const auto [__, is_new] = unique_docs.insert(doc_id);
+    if (is_new) {
+      result.push_back(doc_id);
+    }
+  }
+
+  std::sort(result.begin(), result.end());
+  return result;
+}
+
 template <typename C>
 BaseStringIndex<C>::BaseStringIndex(PMR_NS::memory_resource* mr, bool case_sensitive)
     : case_sensitive_{case_sensitive}, entries_{mr} {
@@ -196,6 +214,26 @@ template <typename C> vector<string> BaseStringIndex<C>::GetTerms() const {
   return res;
 }
 
+template <typename C> vector<DocId> BaseStringIndex<C>::GetAllDocsWithNonNullValues() const {
+  UniqueDocsList<> unique_docs;
+  std::vector<DocId> result;
+
+  unique_docs.reserve(entries_.size());
+  result.reserve(entries_.size());
+
+  for (const auto& [_, container] : entries_) {
+    for (const auto& doc_id : container) {
+      auto [_, is_new] = unique_docs.insert(doc_id);
+      if (is_new) {
+        result.push_back(doc_id);
+      }
+    }
+  }
+
+  std::sort(result.begin(), result.end());
+  return result;
+}
+
 template struct BaseStringIndex<CompressedSortedSet>;
 template struct BaseStringIndex<SortedVector>;
 
@@ -264,6 +302,36 @@ const float* FlatVectorIndex::Get(DocId doc) const {
   return &entries_[doc * dim_];
 }
 
+std::vector<DocId> FlatVectorIndex::GetAllDocsWithNonNullValues() const {
+  std::vector<DocId> result;
+
+  size_t num_vectors = entries_.size() / dim_;
+  result.reserve(num_vectors);
+
+  for (DocId 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;
+
+    // TODO: Consider don't use check for zero vector
+    for (size_t i = 0; i < 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);
+    }
+  }
+
+  // Result is already sorted by id, no need to sort again
+  // Also it has no duplicates
+  return result;
+}
+
 struct HnswlibAdapter {
   // Default setting of hnswlib/hnswalg
   constexpr static size_t kDefaultEfRuntime = 10;

src/core/search/indices.h

@@ -35,9 +35,7 @@ struct NumericIndex : public BaseIndex {
 
   std::vector<DocId> Range(double l, double r) const;
 
-  std::optional<std::vector<DocId>> GetAllResults() const override {
-    return Range(-std::numeric_limits<double>::infinity(), std::numeric_limits<double>::infinity());
-  }
+  std::vector<DocId> GetAllDocsWithNonNullValues() const override;
 
  private:
   using Entry = std::pair<double, DocId>;
@@ -62,19 +60,7 @@ template <typename C> struct BaseStringIndex : public BaseIndex {
   // Returns all the terms that appear as keys in the reverse index.
   std::vector<std::string> GetTerms() const;
 
-  std::optional<std::vector<DocId>> GetAllResults() const override {
-    absl::flat_hash_set<DocId> unique_docs;
-
-    for (const auto& [term, container] : entries_) {
-      for (const DocId& id : container) {
-        unique_docs.insert(id);
-      }
-    }
-
-    auto result = std::vector<DocId>(unique_docs.begin(), unique_docs.end());
-    std::sort(result.begin(), result.end());
-    return result;
-  }
+  std::vector<DocId> GetAllDocsWithNonNullValues() const override;
 
  protected:
   using StringList = DocumentAccessor::StringList;
@@ -152,31 +138,7 @@ struct FlatVectorIndex : public BaseVectorIndex {
   const float* Get(DocId doc) const;
 
   // Return all documents that have vectors in this index
-  std::optional<std::vector<DocId>> GetAllResults() const override {
-    std::vector<DocId> result;
-    size_t num_vectors = entries_.size() / dim_;
-    result.reserve(num_vectors);
-
-    for (DocId 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;
-
-      for (size_t i = 0; i < dim_; ++i) {
-        if (vec[i] != 0.0f) {
-          is_zero_vector = false;
-          break;
-        }
-      }
-
-      if (!is_zero_vector) {
-        result.push_back(id);
-      }
-    }
-
-    return result;
-  }
+  std::vector<DocId> GetAllDocsWithNonNullValues() const override;
 
  protected:
   void AddVector(DocId id, const VectorPtr& vector) override;
@@ -187,10 +149,6 @@ struct FlatVectorIndex : public BaseVectorIndex {
 
 struct HnswlibAdapter;
 
-// This index does't have GetAllResults method
-// because it's not possible to get all vectors from the index
-// It depends on the Hnswlib implementation
-// TODO: Consider adding GetAllResults method in the future
 struct HnswVectorIndex : public BaseVectorIndex {
   HnswVectorIndex(const SchemaField::VectorParams& params, PMR_NS::memory_resource* mr);
   ~HnswVectorIndex();
@@ -201,6 +159,11 @@ struct HnswVectorIndex : public BaseVectorIndex {
   std::vector<std::pair<float, DocId>> Knn(float* target, size_t k, std::optional<size_t> ef,
                                            const std::vector<DocId>& allowed) const;
 
+  // TODO: Implement if needed
+  std::vector<DocId> GetAllDocsWithNonNullValues() const override {
+    return std::vector<DocId>{};
+  }
+
  protected:
   void AddVector(DocId id, const VectorPtr& vector) override;
 

src/core/search/search.cc

@@ -184,17 +184,25 @@ struct BasicSearch {
     profile_builder_ = ProfileBuilder{};
   }
 
+  BaseIndex* GetBaseIndex(string_view field) {
+    auto index = indices_->GetIndex(field);
+    if (!index) {
+      error_ = absl::StrCat("Invalid field: ", field);
+      return nullptr;
+    }
+    return index;
+  }
+
   // Get casted sub index by field
   template <typename T> T* GetIndex(string_view field) {
     static_assert(is_base_of_v<BaseIndex, T>);
 
-    auto index = indices_->GetIndex(field);
-    if (!index) {
-      error_ = absl::StrCat("Invalid field: ", field);
+    auto base_index = GetBaseIndex(field);
+    if (!base_index) {
       return nullptr;
     }
 
-    auto* casted_ptr = dynamic_cast<T*>(index);
+    auto* casted_ptr = dynamic_cast<T*>(base_index);
     if (!casted_ptr) {
       error_ = absl::StrCat("Wrong access type for field: ", field);
       return nullptr;
@@ -246,7 +254,7 @@ struct BasicSearch {
   // Efficiently unify multiple sub results with specified logical op
   IndexResult UnifyResults(vector<IndexResult>&& sub_results, LogicOp op) {
     if (sub_results.empty())
-      return vector<DocId>{};
+      return IndexResult{};
 
     // Unifying from smallest to largest is more efficient.
     // AND: the result only shrinks, so starting with the smallest is most optimal.
@@ -283,7 +291,7 @@ struct BasicSearch {
   }
 
   IndexResult Search(monostate, string_view) {
-    return vector<DocId>{};
+    return IndexResult{};
   }
 
   IndexResult Search(const AstStarNode& node, string_view active_field) {
@@ -321,26 +329,12 @@ struct BasicSearch {
     // Try to get a sort index first, as `@field:*` might imply wanting sortable behavior
     BaseSortIndex* sort_index = indices_->GetSortIndex(active_field);
     if (sort_index) {
-      if (auto result = sort_index->GetAllResults()) {
-        return std::move(*result);
-      }
-    }
-
-    // If sort index doesn't exist or doesn't support GetAllResults, try regular index
-    BaseIndex* base_index = indices_->GetIndex(active_field);
-    if (base_index) {
-      if (auto result = base_index->GetAllResults()) {
-        return std::move(*result);
-      }
+      return {sort_index->GetAllDocsWithNonNullValues()};
     }
 
-    // If we get here, neither index could handle the request
-    if (!base_index && !sort_index) {
-      error_ = absl::StrCat("Invalid field: ", active_field);
-    } else {
-      error_ = absl::StrCat("Wrong access type for field: ", active_field);
-    }
-    return IndexResult{};
+    // If sort index doesn't exist try regular index
+    BaseIndex* base_index = GetBaseIndex(active_field);
+    return base_index ? IndexResult{base_index->GetAllDocsWithNonNullValues()} : IndexResult{};
   }
 
   IndexResult Search(const AstPrefixNode& node, string_view active_field) {

src/core/search/sort_indices.cc

@@ -29,7 +29,8 @@ template <typename T> bool SimpleValueSortIndex<T>::ParsedSortValue::IsNullValue
 }
 
 template <typename T>
-SimpleValueSortIndex<T>::SimpleValueSortIndex(PMR_NS::memory_resource* mr) : values_{mr} {
+SimpleValueSortIndex<T>::SimpleValueSortIndex(PMR_NS::memory_resource* mr)
+    : values_{mr}, null_values_(mr) {
 }
 
 template <typename T> SortableValue SimpleValueSortIndex<T>::Lookup(DocId doc) const {
@@ -90,6 +91,23 @@ void SimpleValueSortIndex<T>::Remove(DocId id, const DocumentAccessor& doc,
   values_[id] = T{};
 }
 
+template <typename T>
+std::vector<DocId> SimpleValueSortIndex<T>::GetAllDocsWithNonNullValues() const {
+  std::vector<DocId> result;
+  result.reserve(values_.size());
+
+  auto empty_value = T{};
+  for (DocId id = 0; id < values_.size(); ++id) {
+    if (values_[id] != empty_value) {
+      result.push_back(id);
+    }
+  }
+
+  // Result is already sorted by DocId
+  // Also it has no duplicates
+  return result;
+}
+
 template <typename T> PMR_NS::memory_resource* SimpleValueSortIndex<T>::GetMemRes() const {
   return values_.get_allocator().resource();
 }

src/core/search/sort_indices.h

@@ -42,21 +42,7 @@ template <typename T> struct SimpleValueSortIndex : public BaseSortIndex {
   void Remove(DocId id, const DocumentAccessor& doc, std::string_view field) override;
 
   // Override GetAllResults to return all documents with non-null values
-  std::optional<std::vector<DocId>> GetAllResults() const override {
-    std::vector<DocId> result;
-
-    for (DocId id = 0; id < values_.size(); ++id) {
-      // Check if id is not present in null_values_
-      // Also need to handle deleted elements - in them T should be empty
-      // Different types of T have different "empty" values, but we can check
-      // if this value is the default for the given type
-      if (!null_values_.contains(id) && !(values_[id] == T{})) {
-        result.push_back(id);
-      }
-    }
-
-    return result;
-  }
+  std::vector<DocId> GetAllDocsWithNonNullValues() const override;
 
  protected:
   virtual ParsedSortValue Get(const DocumentAccessor& doc, std::string_view field_value) = 0;
@@ -65,7 +51,7 @@ template <typename T> struct SimpleValueSortIndex : public BaseSortIndex {
 
  private:
   PMR_NS::vector<T> values_;
-  absl::flat_hash_set<DocId> null_values_;
+  UniqueDocsList<PMR_NS::polymorphic_allocator<DocId>> null_values_;
 };
 
 struct NumericSortIndex : public SimpleValueSortIndex<double> {