feat(search): Probabilitic cutoffs

Signed-off-by: Vladislav Oleshko <[email protected]>

Author: dranikpg

src/core/search/index_result.h

@@ -39,8 +39,8 @@ class IndexResult {
 
   BorrowedView Borrowed() const;
 
-  // Move out of owned or copy borrowed
-  DocVec Take();
+  // Move out of owned or copy borrowed. Take up to `limit` entries and return original size.
+  std::pair<DocVec, size_t /* full size */> Take(size_t limit = std::numeric_limits<size_t>::max());
 
  private:
   bool IsOwned() const;
@@ -82,20 +82,36 @@ inline IndexResult::BorrowedView IndexResult::Borrowed() const {
   return std::visit(cb, value_);
 }
 
-inline IndexResult::DocVec IndexResult::Take() {
+inline std::pair<IndexResult::DocVec, size_t> IndexResult::Take(size_t limit) {
   if (IsOwned()) {
-    return std::move(std::get<DocVec>(value_));
+    auto& vec = std::get<DocVec>(value_);
+    size_t size = vec.size();
+    return {std::move(vec), size};
   }
 
-  auto cb = [](auto* set) -> DocVec {
+  // Numeric ranges need to be filtered and don't know their size ahead
+  // if (std::holds_alternative<RangeResult>(value_)) {
+  auto cb = [](auto* range) -> std::pair<DocVec, size_t> {
     DocVec out;
-    out.reserve(set->size());
-    for (auto it = set->begin(); it != set->end(); ++it) {
+    out.reserve(range->size());
+    for (auto it = range->begin(); it != range->end(); ++it)
       out.push_back(*it);
-    }
-    return out;
+    size_t total = out.size();
+    return {std::move(out), total};
   };
   return std::visit(cb, Borrowed());
+  //}
+
+  // Generic borrowed results sets don't need to be filtered, so we can tell the result size ahead
+  /*auto cb = [limit](auto* set) -> std::pair<DocVec, size_t> {
+    DocVec out;
+    size_t taken = std::min(limit, set->size());
+    out.reserve(taken);
+    for (auto it = set->begin(); it != set->end() && out.size(); ++it)
+      out.push_back(*it);
+    return {std::move(out), set->size()};
+  };
+  return std::visit(cb, Borrowed());*/
 }
 
 inline bool IndexResult::IsOwned() const {

src/core/search/search.cc

@@ -289,7 +289,7 @@ struct BasicSearch {
 
   // negate -(*subquery*): explicitly compute result complement. Needs further optimizations
   IndexResult Search(const AstNegateNode& node, string_view active_field) {
-    vector<DocId> matched = SearchGeneric(*node.node, active_field).Take();
+    auto matched = SearchGeneric(*node.node, active_field).Take().first;
     vector<DocId> all = indices_->GetAllDocs();
 
     // To negate a result, we have to find the complement of matched to all documents,
@@ -358,7 +358,7 @@ struct BasicSearch {
       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());
+          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`
@@ -405,7 +405,6 @@ struct BasicSearch {
 
     // Top level results don't need to be sorted, because they will be scored, sorted by fields or
     // used by knn
-
     DCHECK(top_level || holds_alternative<AstKnnNode>(node.Variant()) ||
            holds_alternative<AstGeoNode>(node.Variant()) ||
            visit([](auto* set) { return is_sorted(set->begin(), set->end()); }, result.Borrowed()));
@@ -416,16 +415,15 @@ struct BasicSearch {
     return result;
   }
 
-  SearchResult Search(const AstNode& query) {
+  SearchResult Search(const AstNode& query, size_t cuttoff_limit) {
     IndexResult result = SearchGeneric(query, "", true);
 
     // Extract profile if enabled
     optional<AlgorithmProfile> profile =
         profile_builder_ ? make_optional(profile_builder_->Take()) : nullopt;
 
-    auto out = result.Take();
-    const size_t total = out.size();
-    return SearchResult{total, std::move(out), std::move(knn_scores_), std::move(profile),
+    auto [out, total_size] = result.Take(cuttoff_limit);
+    return SearchResult{total_size, std::move(out), std::move(knn_scores_), std::move(profile),
                         std::move(error_)};
   }
 
@@ -654,11 +652,11 @@ bool SearchAlgorithm::Init(string_view query, const QueryParams* params,
   return true;
 }
 
-SearchResult SearchAlgorithm::Search(const FieldIndices* index) const {
+SearchResult SearchAlgorithm::Search(const FieldIndices* index, size_t cuttoff_limit) const {
   auto bs = BasicSearch{index};
   if (profiling_enabled_)
     bs.EnableProfiling();
-  return bs.Search(*query_);
+  return bs.Search(*query_, cuttoff_limit);
 }
 
 optional<KnnScoreSortOption> SearchAlgorithm::GetKnnScoreSortOption() const {

src/core/search/search.h

@@ -197,7 +197,9 @@ class SearchAlgorithm {
   bool Init(std::string_view query, const QueryParams* params,
             const OptionalFilters* filters = nullptr);
 
-  SearchResult Search(const FieldIndices* index) const;
+  // Search on given index with predefined limit for cutting off result ids
+  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;

src/server/search/doc_index.cc

@@ -400,7 +400,9 @@ vector<search::SortableValue> ShardDocIndex::KeepTopKSorted(vector<DocId>* ids,
 SearchResult ShardDocIndex::Search(const OpArgs& op_args, const SearchParams& params,
                                    search::SearchAlgorithm* search_algo) const {
   size_t limit = params.limit_offset + params.limit_total;
-  auto result = search_algo->Search(&*indices_);
+  bool can_cut_off = !params.sort_option && !search_algo->GetKnnScoreSortOption();
+  size_t id_cutoff_limit = can_cut_off ? limit : numeric_limits<size_t>::max();
+  auto result = search_algo->Search(&*indices_, id_cutoff_limit);
   if (!result.error.empty())
     return {facade::ErrorReply(std::move(result.error))};
 
@@ -441,7 +443,10 @@ SearchResult ShardDocIndex::Search(const OpArgs& op_args, const SearchParams& pa
   }
 
   // Cut off unnecessary items
-  result.ids.resize(min(result.ids.size(), limit));
+  size_t serialization_limit = limit;
+  if (!search_algo->GetKnnScoreSortOption())
+    serialization_limit = params.limit_serialization;
+  result.ids.resize(min(result.ids.size(), serialization_limit));
 
   // Serialize documents
   vector<SerializedSearchDoc> out;

src/server/search/doc_index.h

@@ -98,6 +98,7 @@ struct SearchParams {
   size_t limit_offset = 0;
   size_t limit_total = 10;
 
+  size_t limit_serialization = 0;
   /*
   1. If not set -> return all fields
   2. If set but empty -> no fields should be returned

src/server/search/search_family.cc

@@ -959,17 +959,27 @@ void PartialSort(absl::Span<SerializedSearchDoc*> docs, size_t limit, SortOrder
   partial_sort(docs.begin(), docs.begin() + min(limit, docs.size()), docs.end(), cb);
 }
 
-void SearchReply(const SearchParams& params,
+bool SearchReply(const SearchParams& params,
                  std::optional<search::KnnScoreSortOption> knn_sort_option,
                  absl::Span<SearchResult> results, SinkReplyBuilder* builder) {
+  // Count total number of hits
   size_t total_hits = 0;
+  for (auto& shard_results : results)
+    total_hits += shard_results.total_hits;
+
+  // Arrange documents in a stride
   absl::InlinedVector<SerializedSearchDoc*, 5> docs;
   docs.reserve(results.size());
-  for (auto& shard_results : results) {
-    total_hits += shard_results.total_hits;
-    for (auto& doc : shard_results.docs) {
-      docs.push_back(&doc);
+  for (size_t i = 0;; ++i) {
+    bool added = false;
+    for (auto& shard_results : results) {
+      if (i < shard_results.docs.size()) {
+        added = true;
+        docs.push_back(&shard_results.docs[i]);
+      }
     }
+    if (!added)
+      break;
   }
 
   // Reorder and cut KNN results before applying SORT and LIMIT
@@ -995,6 +1005,14 @@ void SearchReply(const SearchParams& params,
     PartialSort(absl::MakeSpan(docs), end, params.sort_option->order,
                 &SerializedSearchDoc::sort_score);
 
+  // Check if we havent' chosen too few documents due to cutoffs
+  size_t left = docs.size() - params.limit_offset;
+  size_t expected = std::min(limit, total_hits - params.limit_offset);
+  if (left < expected)
+    return false;
+
+  // TODO: Check sort correctness
+
   const bool reply_with_ids_only = params.IdsOnly();
   auto* rb = static_cast<RedisReplyBuilder*>(builder);
   RedisReplyBuilder::ArrayScope scope{rb, reply_with_ids_only ? (limit + 1) : (limit * 2 + 1)};
@@ -1011,6 +1029,7 @@ void SearchReply(const SearchParams& params,
 
     SendSerializedDoc(*docs[i], builder);
   }
+  return true;
 }
 
 // Warms up the query parser to avoid first-call slowness
@@ -1279,23 +1298,37 @@ void SearchFamily::FtSearch(CmdArgList args, const CommandContext& cmd_cntx) {
   atomic<bool> index_not_found{false};
   vector<SearchResult> docs(shard_set->size());
 
-  cmd_cntx.tx->ScheduleSingleHop([&](Transaction* t, EngineShard* es) {
-    if (auto* index = es->search_indices()->GetIndex(index_name); index)
-      docs[es->shard_id()] = index->Search(t->GetOpArgs(es), *params, &search_algo);
+  bool succeeded = true;
+  do {
+    if (succeeded)
+      params->limit_serialization =
+          params->limit_offset +
+          std::min(params->limit_total, 2 * params->limit_total / shard_set->size());
     else
-      index_not_found.store(true, memory_order_relaxed);
-    return OpStatus::OK;
-  });
+      params->limit_serialization = params->limit_total + params->limit_offset;
 
-  if (index_not_found.load())
-    return builder->SendError(string{index_name} + ": no such index");
+    cmd_cntx.tx->ScheduleSingleHop([&](Transaction* t, EngineShard* es) {
+      if (auto* index = es->search_indices()->GetIndex(index_name); index)
+        docs[es->shard_id()] = index->Search(t->GetOpArgs(es), *params, &search_algo);
+      else
+        index_not_found.store(true, memory_order_relaxed);
+      return OpStatus::OK;
+    });
 
-  for (const auto& res : docs) {
-    if (res.error)
-      return builder->SendError(*res.error);
-  }
+    if (index_not_found.load())
+      return builder->SendError(string{index_name} + ": no such index");
+
+    for (const auto& res : docs) {
+      if (res.error)
+        return builder->SendError(*res.error);
+    }
 
-  SearchReply(*params, search_algo.GetKnnScoreSortOption(), absl::MakeSpan(docs), builder);
+    bool did_succeed =
+        SearchReply(*params, search_algo.GetKnnScoreSortOption(), absl::MakeSpan(docs), builder);
+    CHECK(did_succeed);
+    DCHECK(succeeded || did_succeed);
+    succeeded = did_succeed;
+  } while (!succeeded);
 }
 
 void SearchFamily::FtProfile(CmdArgList args, const CommandContext& cmd_cntx) {
@@ -1331,6 +1364,7 @@ void SearchFamily::FtProfile(CmdArgList args, const CommandContext& cmd_cntx) {
   std::vector<SearchResult> search_results(shards_count);
   std::vector<absl::Duration> profile_results(shards_count);
 
+  params->limit_serialization = params->limit_offset + params->limit_total;
   cmd_cntx.tx->ScheduleSingleHop([&](Transaction* t, EngineShard* es) {
     auto* index = es->search_indices()->GetIndex(index_name);
     if (!index) {