Add getSortedHeapHits

2 files changed, 27 insertions, 21 deletions

diff --git a/searchlib/src/vespa/searchlib/queryeval/hitcollector.cpp b/searchlib/src/vespa/searchlib/queryeval/hitcollector.cpp
index f9196bdaef7..7ed4705b8b8 100644
--- a/searchlib/src/vespa/searchlib/queryeval/hitcollector.cpp
+++ b/searchlib/src/vespa/searchlib/queryeval/hitcollector.cpp
@@ -4,8 +4,7 @@
 #include <vespa/searchlib/common/bitvector.h>
 #include <vespa/searchlib/common/sort.h>
 
-namespace search {
-namespace queryeval {
+namespace search::queryeval {
 
 void
 HitCollector::sortHitsByScore(size_t topn)
@@ -53,16 +52,14 @@ HitCollector::HitCollector(uint32_t numDocs,
       _needReScore(false)
 {
     if (_maxHitsSize > 0) {
-        _collector.reset(new RankedHitCollector(*this));
+        _collector = std::make_unique<RankedHitCollector>(*this);
     } else {
-        _collector.reset(new DocIdCollector<false>(*this));
+        _collector = std::make_unique<DocIdCollector<false>>(*this);
     }
     _hits.reserve(maxHitsSize);
 }
 
-HitCollector::~HitCollector()
-{
-}
+HitCollector::~HitCollector() = default;
 
 void
 HitCollector::RankedHitCollector::collect(uint32_t docId, feature_t score)
@@ -113,7 +110,7 @@ HitCollector::RankedHitCollector::collectAndChangeCollector(uint32_t docId, feat
             hc._docIdVector.push_back(hc._hits[i].first);
         }
         hc._docIdVector.push_back(docId);
-        newCollector.reset(new DocIdCollector<true>(hc));
+        newCollector = std::make_unique<DocIdCollector<true>>(hc);
     } else {
         // start using bit vector
         hc._bitVector = BitVector::create(hc._numDocs);
@@ -123,7 +120,7 @@ HitCollector::RankedHitCollector::collectAndChangeCollector(uint32_t docId, feat
             hc._bitVector->setBit(hc._hits[i].first);
         }
         hc._bitVector->setBit(docId);
-        newCollector.reset(new BitVectorCollector<true>(hc));
+        newCollector = std::make_unique<BitVectorCollector<true>>(hc);
     }
     // treat hit vector as a heap
     std::make_heap(hc._hits.begin(), hc._hits.end(), ScoreComparator());
@@ -168,7 +165,7 @@ HitCollector::DocIdCollector<CollectRankedHit>::collectAndChangeCollector(uint32
     std::vector<uint32_t> emptyVector;
     emptyVector.swap(hc._docIdVector);
     hc._bitVector->setBit(docId);
-    hc._collector.reset(new BitVectorCollector<CollectRankedHit>(hc)); // note - self-destruct.
+    hc._collector = std::make_unique<BitVectorCollector<CollectRankedHit>>(hc); // note - self-destruct.
 }
 
 std::vector<feature_t>
@@ -184,6 +181,20 @@ HitCollector::getSortedHeapScores()
     return scores;
 }
 
+std::vector<HitCollector::Hit>
+HitCollector::getSortedHeapHits()
+{
+    std::vector<Hit> scores;
+    size_t scoresToReturn = std::min(_hits.size(), static_cast<size_t>(_maxReRankHitsSize));
+    scores.reserve(scoresToReturn);
+    sortHitsByScore(scoresToReturn);
+    for (size_t i = 0; i < scoresToReturn; ++i) {
+        scores.push_back(_hits[_scoreOrder[i]]);
+    }
+    return scores;
+}
+
+
 size_t
 HitCollector::reRank(DocumentScorer &scorer)
 {
@@ -205,8 +216,7 @@ HitCollector::reRank(DocumentScorer &scorer, size_t count)
 
     Scores &initScores = _ranges.first;
     Scores &finalScores = _ranges.second;
-    initScores = Scores(_reRankedHits.back().second,
-                        _reRankedHits.front().second);
+    initScores = Scores(_reRankedHits.back().second, _reRankedHits.front().second);
     finalScores = Scores(std::numeric_limits<feature_t>::max(),
                          -std::numeric_limits<feature_t>::max());
 
@@ -335,5 +345,4 @@ HitCollector::getResultSet(HitRank default_value)
     return rs;
 }
 
-} // namespace queryeval
-} // namespace search
+}
diff --git a/searchlib/src/vespa/searchlib/queryeval/hitcollector.h b/searchlib/src/vespa/searchlib/queryeval/hitcollector.h
index 1bf2bc21e95..4f20a5cae83 100644
--- a/searchlib/src/vespa/searchlib/queryeval/hitcollector.h
+++ b/searchlib/src/vespa/searchlib/queryeval/hitcollector.h
@@ -10,16 +10,14 @@
 #include <vespa/vespalib/util/sort.h>
 #include <vespa/fastos/dynamiclibrary.h>
 
-namespace search {
-
-namespace queryeval {
+namespace search::queryeval {
 
 /**
  * This class is used to store all hits found during parallel query evaluation.
  **/
 class HitCollector {
 public:
-    typedef std::pair<uint32_t, feature_t> Hit;
+    using Hit = std::pair<uint32_t, feature_t>;
 
     /**
      * Interface used to calculate the second phase score for the documents being re-ranked.
@@ -173,6 +171,7 @@ public:
      * in the heap. These are the candidates for re-ranking.
      */
     std::vector<feature_t> getSortedHeapScores();
+    std::vector<Hit> getSortedHeapHits();
 
     /**
      * Re-ranks the m (=maxHeapSize) best hits by invoking the score()
@@ -200,6 +199,4 @@ private:
     HitCollector &operator=(const HitCollector &);  // Not implemented
 };
 
-} // namespace queryeval
-} // namespace search
-
+}