test + fixes

Author: dranikpg

src/core/search/index_result.h

@@ -90,28 +90,28 @@ inline std::pair<IndexResult::DocVec, size_t> IndexResult::Take(size_t limit) {
   }
 
   // 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(range->size());
-    for (auto it = range->begin(); it != range->end(); ++it)
-      out.push_back(*it);
-    size_t total = out.size();
-    return {std::move(out), total};
-  };
-  return std::visit(cb, Borrowed());
-  //}
+  if (std::holds_alternative<RangeResult>(value_)) {
+    auto cb = [limit](auto* range) -> std::pair<DocVec, size_t> {
+      DocVec out;
+      out.reserve(range->size());
+      for (auto it = range->begin(); it != range->end(); ++it)
+        out.push_back(*it);
+      size_t total = out.size();
+      out.resize(std::min(out.size(), limit));
+      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> {
+  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.reserve(set->size());
+    for (auto it = set->begin(); it != set->end() && out.size() < limit; ++it)
       out.push_back(*it);
     return {std::move(out), set->size()};
   };
-  return std::visit(cb, Borrowed());*/
+  return std::visit(cb, Borrowed());
 }
 
 inline bool IndexResult::IsOwned() const {

src/server/search/search_family.cc

@@ -1007,7 +1007,7 @@ bool SearchReply(const SearchParams& params,
 
   // 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);
+  size_t expected = std::min(params.limit_total, total_hits - params.limit_offset);
   if (left < expected)
     return false;
 
@@ -1303,7 +1303,7 @@ void SearchFamily::FtSearch(CmdArgList args, const CommandContext& cmd_cntx) {
     if (succeeded)
       params->limit_serialization =
           params->limit_offset +
-          std::min(params->limit_total, 2 * params->limit_total / shard_set->size());
+          std::min(params->limit_total, params->limit_total / shard_set->size());
     else
       params->limit_serialization = params->limit_total + params->limit_offset;
 
@@ -1325,9 +1325,9 @@ void SearchFamily::FtSearch(CmdArgList args, const CommandContext& cmd_cntx) {
 
     bool did_succeed =
         SearchReply(*params, search_algo.GetKnnScoreSortOption(), absl::MakeSpan(docs), builder);
-    CHECK(did_succeed);
     DCHECK(succeeded || did_succeed);
     succeeded = did_succeed;
+    cmd_cntx.tx->Refurbish();
   } while (!succeeded);
 }
 

src/server/search/search_family_test.cc

@@ -751,6 +751,24 @@ TEST_F(SearchFamilyTest, ReturnOptionJson) {
               MatchEntry("k1", "name", "dragon"));
 }
 
+TEST_F(SearchFamilyTest, TestUnbalanced) {
+  Run({"ft.create", "i1", "schema", "title", "text"});
+
+  size_t added = 0;
+  for (size_t i = 0; i < 200; i++) {
+    string key = absl::StrCat("k", i);
+    if (Shard(key, shard_set->size()) != 1)
+      continue;
+
+    Run({"hset", key, "name", "olaf"});
+    added++;
+  }
+
+  auto res = Run({"ft.search", "i1", "*", "LIMIT", "0", "15"});
+  EXPECT_THAT(res.GetVec()[0], IntArg(added));
+  EXPECT_EQ(res.GetVec().size(), 1 + 2 * 15);
+}
+
 TEST_F(SearchFamilyTest, TestStopWords) {
   Run({"ft.create", "i1", "STOPWORDS", "3", "red", "green", "blue", "SCHEMA", "title", "TEXT"});