From e289834e551f692168f7e762da446b58137650da Mon Sep 17 00:00:00 2001
From: tatiana-yan <tatiana.kondakova@gmail.com>
Date: Mon, 28 Sep 2020 18:55:01 +0300
Subject: [PATCH] [search] Do not limit preranker batch size for viewport
 search mode.

---
 search/pre_ranker.cpp                         | 62 +------------------
 search/pre_ranker.hpp                         | 11 +++-
 .../pre_ranker_test.cpp                       |  4 +-
 3 files changed, 12 insertions(+), 65 deletions(-)

diff --git a/search/pre_ranker.cpp b/search/pre_ranker.cpp
index b4c8a7eb1f..cdf9d85863 100644
--- a/search/pre_ranker.cpp
+++ b/search/pre_ranker.cpp
@@ -289,67 +289,7 @@ void PreRanker::FilterForViewportSearch()
   SweepNearbyResults(m_params.m_minDistanceOnMapBetweenResultsX,
                      m_params.m_minDistanceOnMapBetweenResultsY, m_prevEmit, m_results);
 
-  size_t const n = m_results.size();
-
-  if (n <= BatchSize())
-    return;
-
-  size_t const kNumXSlots = 5;
-  size_t const kNumYSlots = 5;
-  size_t const kNumBuckets = kNumXSlots * kNumYSlots;
-  vector<size_t> buckets[kNumBuckets];
-
-  double const sizeX = viewport.SizeX();
-  double const sizeY = viewport.SizeY();
-
-  for (size_t i = 0; i < m_results.size(); ++i)
-  {
-    auto const & p = m_results[i].GetInfo().m_center;
-    int dx = static_cast<int>((p.x - viewport.minX()) / sizeX * kNumXSlots);
-    dx = base::Clamp(dx, 0, static_cast<int>(kNumXSlots) - 1);
-
-    int dy = static_cast<int>((p.y - viewport.minY()) / sizeY * kNumYSlots);
-    dy = base::Clamp(dy, 0, static_cast<int>(kNumYSlots) - 1);
-
-    buckets[dx * kNumYSlots + dy].push_back(i);
-  }
-
-  vector<PreRankerResult> results;
-  double const density = static_cast<double>(BatchSize()) / static_cast<double>(n);
-  for (auto & bucket : buckets)
-  {
-    size_t const m = std::min(static_cast<size_t>(ceil(density * bucket.size())), bucket.size());
-
-    size_t const old =
-        partition(bucket.begin(), bucket.end(),
-                  [this](size_t i) { return m_prevEmit.count(m_results[i].GetId()) != 0; }) -
-        bucket.begin();
-
-    if (m <= old)
-    {
-      for (size_t i : base::RandomSample(old, m, m_rng))
-        results.push_back(m_results[bucket[i]]);
-    }
-    else
-    {
-      for (size_t i = 0; i < old; ++i)
-        results.push_back(m_results[bucket[i]]);
-
-      for (size_t i : base::RandomSample(bucket.size() - old, m - old, m_rng))
-        results.push_back(m_results[bucket[old + i]]);
-    }
-  }
-
-  if (results.size() <= BatchSize())
-  {
-    m_results.swap(results);
-  }
-  else
-  {
-    m_results.clear();
-    for (size_t i : base::RandomSample(results.size(), BatchSize(), m_rng))
-      m_results.push_back(results[i]);
-  }
+  CHECK_LESS_OR_EQUAL(m_results.size(), BatchSize(), ());
 }
 
 void PreRanker::FilterRelaxedResults(bool lastUpdate)
diff --git a/search/pre_ranker.hpp b/search/pre_ranker.hpp
index cfb032c7eb..ec015646f1 100644
--- a/search/pre_ranker.hpp
+++ b/search/pre_ranker.hpp
@@ -11,6 +11,7 @@
 
 #include <algorithm>
 #include <cstddef>
+#include <limits>
 #include <optional>
 #include <random>
 #include <set>
@@ -47,7 +48,9 @@ public:
 
     int m_scale = 0;
 
-    size_t m_batchSize = 100;
+    // Batch size for Everywhere search mode. For viewport search we limit search results number
+    // with SweepNearbyResults.
+    size_t m_everywhereBatchSize = 100;
 
     // The maximum total number of results to be emitted in all batches.
     size_t m_limit = 0;
@@ -87,7 +90,11 @@ public:
   void UpdateResults(bool lastUpdate);
 
   size_t Size() const { return m_results.size() + m_relaxedResults.size(); }
-  size_t BatchSize() const { return m_params.m_batchSize; }
+  size_t BatchSize() const
+  {
+    return m_params.m_viewportSearch ? std::numeric_limits<size_t>::max()
+                                     : m_params.m_everywhereBatchSize;
+  }
   size_t NumSentResults() const { return m_numSentResults; }
   bool HaveFullyMatchedResult() const { return m_haveFullyMatchedResult; }
   size_t Limit() const { return m_params.m_limit; }
diff --git a/search/search_integration_tests/pre_ranker_test.cpp b/search/search_integration_tests/pre_ranker_test.cpp
index 45a034e57d..fb0ba0295f 100644
--- a/search/search_integration_tests/pre_ranker_test.cpp
+++ b/search/search_integration_tests/pre_ranker_test.cpp
@@ -141,9 +141,9 @@ UNIT_CLASS_TEST(PreRankerTest, Smoke)
   params.m_viewport = kViewport;
   params.m_accuratePivotCenter = kPivot;
   params.m_scale = scales::GetUpperScale();
-  params.m_batchSize = kBatchSize;
+  params.m_everywhereBatchSize = kBatchSize;
   params.m_limit = pois.size();
-  params.m_viewportSearch = true;
+  params.m_viewportSearch = false;
   preRanker.Init(params);
 
   vector<double> distances(pois.size());