Merge pull request #6811 from vespa-engine/lesters/add-dynamic-rerank-count

Add query override of rerank-count

5 files changed, 12 insertions, 8 deletions

diff --git a/searchcore/src/tests/proton/matching/matching_test.cpp b/searchcore/src/tests/proton/matching/matching_test.cpp
index de6a452baf3..7c6779fdc63 100644
--- a/searchcore/src/tests/proton/matching/matching_test.cpp
+++ b/searchcore/src/tests/proton/matching/matching_test.cpp
@@ -53,6 +53,7 @@ using namespace search;
 using search::attribute::test::MockAttributeContext;
 using search::index::schema::DataType;
 using storage::spi::Timestamp;
+using search::fef::indexproperties::hitcollector::HeapSize;
 
 void inject_match_phase_limiting(Properties &setup, const vespalib::string &attribute, size_t max_hits, bool descending)
 {
@@ -287,7 +288,7 @@ struct MyWorld {
         Matcher::SP matcher = createMatcher();
         search::fef::Properties overrides;
         auto mtf = matcher->create_match_tools_factory(*req, searchContext, attributeContext, metaStore, overrides);
-        auto diversity = mtf->createDiversifier();
+        auto diversity = mtf->createDiversifier(HeapSize::lookup(config));
         EXPECT_EQUAL(expectDiverse, static_cast<bool>(diversity));
     }
 
diff --git a/searchcore/src/vespa/searchcore/proton/matching/match_master.cpp b/searchcore/src/vespa/searchcore/proton/matching/match_master.cpp
index d974be1ce3a..4d49e9b5d1b 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/match_master.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/match_master.cpp
@@ -64,7 +64,7 @@ MatchMaster::match(const MatchParams &params,
     fastos::StopWatch query_latency_time;
     query_latency_time.start();
     vespalib::DualMergeDirector mergeDirector(threadBundle.size());
-    MatchLoopCommunicator communicator(threadBundle.size(), params.heapSize, mtf.createDiversifier());
+    MatchLoopCommunicator communicator(threadBundle.size(), params.heapSize, mtf.createDiversifier(params.heapSize));
     TimedMatchLoopCommunicator timedCommunicator(communicator);
     DocidRangeScheduler::UP scheduler = createScheduler(threadBundle.size(), numSearchPartitions, params.numDocs);
 
diff --git a/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp b/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp
index a00a90d7a10..28d56b7e0a2 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/match_tools.cpp
@@ -204,7 +204,7 @@ MatchToolsFactory::createMatchTools() const
 }
 
 std::unique_ptr<IDiversifier>
-MatchToolsFactory::createDiversifier() const
+MatchToolsFactory::createDiversifier(uint32_t heapSize) const
 {
     if ( !_diversityParams.enabled() ) {
         return std::unique_ptr<IDiversifier>();
@@ -214,8 +214,8 @@ MatchToolsFactory::createDiversifier() const
         LOG(warning, "Skipping diversity due to no %s attribute.", _diversityParams.attribute.c_str());
         return std::unique_ptr<IDiversifier>();
     }
-    size_t max_per_group = _rankSetup.getHeapSize()/_diversityParams.min_groups;
-    return DiversityFilter::create(*attr, _rankSetup.getHeapSize(), max_per_group, _diversityParams.min_groups,
+    size_t max_per_group = heapSize/_diversityParams.min_groups;
+    return DiversityFilter::create(*attr, heapSize, max_per_group, _diversityParams.min_groups,
                                    _diversityParams.cutoff_strategy == DiversityParams::CutoffStrategy::STRICT);
 }
 
diff --git a/searchcore/src/vespa/searchcore/proton/matching/match_tools.h b/searchcore/src/vespa/searchcore/proton/matching/match_tools.h
index 8f04eebc50e..0ecf6eb5b78 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/match_tools.h
+++ b/searchcore/src/vespa/searchcore/proton/matching/match_tools.h
@@ -124,7 +124,7 @@ public:
     const MaybeMatchPhaseLimiter &match_limiter() const { return *_match_limiter; }
     MatchTools::UP createMatchTools() const;
     bool should_diversify() const { return _diversityParams.enabled(); }
-    std::unique_ptr<search::queryeval::IDiversifier> createDiversifier() const;
+    std::unique_ptr<search::queryeval::IDiversifier> createDiversifier(uint32_t heapSize) const;
     search::queryeval::Blueprint::HitEstimate estimate() const { return _query.estimate(); }
     bool has_first_phase_rank() const { return !_rankSetup.getFirstPhaseRank().empty(); }
     std::unique_ptr<AttributeOperationTask> createOnMatchTask() const;
diff --git a/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp b/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp
index be0a720f1c1..b32af7e3e5a 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp
@@ -28,6 +28,7 @@ using search::FeatureSet;
 using search::attribute::IAttributeContext;
 using search::fef::MatchDataLayout;
 using search::fef::MatchData;
+using search::fef::indexproperties::hitcollector::HeapSize;
 using search::queryeval::Blueprint;
 using search::queryeval::SearchIterator;
 using vespalib::Doom;
@@ -242,14 +243,16 @@ Matcher::match(const SearchRequest &request, vespalib::ThreadBundle &threadBundl
             return reply;
         }
 
-        MatchParams params(searchContext.getDocIdLimit(), _rankSetup->getHeapSize(), _rankSetup->getArraySize(),
+        const Properties & rankProperties = request.propertiesMap.rankProperties();
+        uint32_t heapSize = HeapSize::lookup(rankProperties, _rankSetup->getHeapSize());
+
+        MatchParams params(searchContext.getDocIdLimit(), heapSize, _rankSetup->getArraySize(),
                            _rankSetup->getRankScoreDropLimit(), request.offset, request.maxhits,
                            !_rankSetup->getSecondPhaseRank().empty(), !willNotNeedRanking(request, groupingContext));
 
         ResultProcessor rp(attrContext, metaStore, sessionMgr, groupingContext, sessionId,
                            request.sortSpec, params.offset, params.hits, request.should_drop_sort_data());
 
-        const Properties & rankProperties = request.propertiesMap.rankProperties();
         size_t numThreadsPerSearch = computeNumThreadsPerSearch(mtf->estimate(), rankProperties);
         LimitedThreadBundleWrapper limitedThreadBundle(threadBundle, numThreadsPerSearch);
         MatchMaster master;