Merge pull request #24928 from vespa-engine/balder/-let-arraysize-be-overridden-in-query

Allow override of keep-rank-count via the internal vespa.hitcollector…

6 files changed, 27 insertions, 2 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index 17b95162709..ab8718508b7 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -5350,6 +5350,8 @@
       "public boolean getQueryCache()",
       "public void setRerankCount(int)",
       "public java.lang.Integer getRerankCount()",
+      "public void setKeepRankCount(int)",
+      "public java.lang.Integer getKeepRankCount()",
       "public com.yahoo.prelude.Location getLocation()",
       "public void setLocation(com.yahoo.prelude.Location)",
       "public void setLocation(java.lang.String)",
@@ -5384,6 +5386,7 @@
       "public static final java.lang.String FRESHNESS",
       "public static final java.lang.String QUERYCACHE",
       "public static final java.lang.String RERANKCOUNT",
+      "public static final java.lang.String KEEPRANKCOUNT",
       "public static final java.lang.String MATCH_PHASE",
       "public static final java.lang.String DIVERSITY",
       "public static final java.lang.String SOFTTIMEOUT",
diff --git a/container-search/src/main/java/com/yahoo/search/query/Ranking.java b/container-search/src/main/java/com/yahoo/search/query/Ranking.java
index 279ec3b0ff2..77c1d0aa621 100644
--- a/container-search/src/main/java/com/yahoo/search/query/Ranking.java
+++ b/container-search/src/main/java/com/yahoo/search/query/Ranking.java
@@ -45,6 +45,7 @@ public class Ranking implements Cloneable {
     public static final String FRESHNESS = "freshness";
     public static final String QUERYCACHE = "queryCache";
     public static final String RERANKCOUNT = "rerankCount";
+    public static final String KEEPRANKCOUNT = "keepRankCount";
     public static final String MATCH_PHASE = "matchPhase";
     public static final String DIVERSITY = "diversity";
     public static final String SOFTTIMEOUT = "softtimeout";
@@ -65,6 +66,7 @@ public class Ranking implements Cloneable {
         argumentType.addField(new FieldDescription(FRESHNESS, "string", "datetime"));
         argumentType.addField(new FieldDescription(QUERYCACHE, "boolean"));
         argumentType.addField(new FieldDescription(RERANKCOUNT, "integer"));
+        argumentType.addField(new FieldDescription(KEEPRANKCOUNT, "integer"));
         argumentType.addField(new FieldDescription(MATCH_PHASE,  new QueryProfileFieldType(MatchPhase.getArgumentType()), "matchPhase"));
         argumentType.addField(new FieldDescription(DIVERSITY, new QueryProfileFieldType(Diversity.getArgumentType())));
         argumentType.addField(new FieldDescription(SOFTTIMEOUT, new QueryProfileFieldType(SoftTimeout.getArgumentType())));
@@ -95,6 +97,7 @@ public class Ranking implements Cloneable {
     private boolean queryCache = false;
 
     private Integer rerankCount = null;
+    private Integer keepRankCount = null;
 
     private RankProperties rankProperties = new RankProperties();
 
@@ -161,6 +164,11 @@ public class Ranking implements Cloneable {
     /** Returns the rerank-count that will be used, or null if not set */
     public Integer getRerankCount() { return rerankCount; }
 
+    /** Sets the keep-rank-count that will be used, or null if not set */
+    public void setKeepRankCount(int keepRankCount) { this.keepRankCount = keepRankCount; }
+    /** Returns the keep-rank-count that will be used, or null if not set */
+    public Integer getKeepRankCount() { return keepRankCount; }
+
     /** Returns the location of this query, or null if none */
     public Location getLocation() { return location; }
 
@@ -235,6 +243,8 @@ public class Ranking implements Cloneable {
         prepareNow(freshness);
         if (rerankCount != null)
             rankProperties.put("vespa.hitcollector.heapsize", rerankCount);
+        if (keepRankCount != null)
+            rankProperties.put("vespa.hitcollector.arraysize", keepRankCount);
     }
 
     private void prepareNow(Freshness freshness) {
diff --git a/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java b/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
index 085a6382a13..d5dc8120f29 100644
--- a/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
+++ b/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
@@ -74,6 +74,7 @@ public class QueryProperties extends Properties {
                 if (key.last().equals(Ranking.FRESHNESS)) return ranking.getFreshness();
                 if (key.last().equals(Ranking.QUERYCACHE)) return ranking.getQueryCache();
                 if (key.last().equals(Ranking.RERANKCOUNT)) return ranking.getRerankCount();
+                if (key.last().equals(Ranking.KEEPRANKCOUNT)) return ranking.getKeepRankCount();
                 if (key.last().equals(Ranking.LIST_FEATURES)) return ranking.getListFeatures();
             }
             else if (key.size() >= 3 && key.get(1).equals(Ranking.MATCH_PHASE)) {
@@ -198,6 +199,8 @@ public class QueryProperties extends Properties {
                         ranking.setQueryCache(asBoolean(value, false));
                     else if (key.last().equals(Ranking.RERANKCOUNT))
                         ranking.setRerankCount(asInteger(value, null));
+                    else if (key.last().equals(Ranking.KEEPRANKCOUNT))
+                        ranking.setKeepRankCount(asInteger(value, null));
                     else if (key.last().equals(Ranking.LIST_FEATURES))
                         ranking.setListFeatures(asBoolean(value,false));
                     else
diff --git a/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp b/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp
index 334948bd191..236964c2e6b 100644
--- a/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp
+++ b/searchcore/src/vespa/searchcore/proton/matching/matcher.cpp
@@ -35,6 +35,7 @@ using search::fef::MatchDataLayout;
 using search::fef::MatchData;
 using search::fef::RankSetup;
 using search::fef::indexproperties::hitcollector::HeapSize;
+using search::fef::indexproperties::hitcollector::ArraySize;
 using search::queryeval::Blueprint;
 using search::queryeval::SearchIterator;
 using vespalib::Doom;
@@ -232,8 +233,9 @@ Matcher::match(const SearchRequest &request, vespalib::ThreadBundle &threadBundl
 
         const Properties & rankProperties = request.propertiesMap.rankProperties();
         uint32_t heapSize = HeapSize::lookup(rankProperties, _rankSetup->getHeapSize());
+        uint32_t arraySize = ArraySize::lookup(rankProperties, _rankSetup->getArraySize());
 
-        MatchParams params(searchContext.getDocIdLimit(), heapSize, _rankSetup->getArraySize(),
+        MatchParams params(searchContext.getDocIdLimit(), heapSize, arraySize,
                            _rankSetup->getRankScoreDropLimit(), request.offset, request.maxhits,
                            !_rankSetup->getSecondPhaseRank().empty(), !willNotNeedRanking(request, groupingContext));
 
diff --git a/searchlib/src/vespa/searchlib/fef/indexproperties.cpp b/searchlib/src/vespa/searchlib/fef/indexproperties.cpp
index 70433800469..26a7be005b7 100644
--- a/searchlib/src/vespa/searchlib/fef/indexproperties.cpp
+++ b/searchlib/src/vespa/searchlib/fef/indexproperties.cpp
@@ -587,7 +587,13 @@ const uint32_t ArraySize::DEFAULT_VALUE(10000);
 uint32_t
 ArraySize::lookup(const Properties &props)
 {
-    return lookupUint32(props, NAME, DEFAULT_VALUE);
+    return lookup(props, DEFAULT_VALUE);
+}
+
+uint32_t
+ArraySize::lookup(const Properties &props, uint32_t defaultValue)
+{
+    return lookupUint32(props, NAME, defaultValue);
 }
 
 const vespalib::string EstimatePoint::NAME("vespa.hitcollector.estimatepoint");
diff --git a/searchlib/src/vespa/searchlib/fef/indexproperties.h b/searchlib/src/vespa/searchlib/fef/indexproperties.h
index 3ebfe781400..5ff4ea26bd8 100644
--- a/searchlib/src/vespa/searchlib/fef/indexproperties.h
+++ b/searchlib/src/vespa/searchlib/fef/indexproperties.h
@@ -476,6 +476,7 @@ namespace hitcollector {
         static const vespalib::string NAME;
         static const uint32_t DEFAULT_VALUE;
         static uint32_t lookup(const Properties &props);
+        static uint32_t lookup(const Properties &props, uint32_t defaultValue);
     };
 
     /**