Add API for rerank-count

2 files changed, 20 insertions, 2 deletions

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 830a3f4ef81..0aa3f12ca68 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
@@ -39,6 +39,7 @@ public class Ranking implements Cloneable {
     public static final String LIST_FEATURES = "listFeatures";
     public static final String FRESHNESS = "freshness";
     public static final String QUERYCACHE = "queryCache";
+    public static final String RERANKCOUNT = "rerankCount";
     public static final String MATCH_PHASE = "matchPhase";
     public static final String DIVERSITY = "diversity";
     public static final String SOFTTIMEOUT = "softtimeout";
@@ -55,7 +56,8 @@ public class Ranking implements Cloneable {
         argumentType.addField(new FieldDescription(SORTING, "string", "sorting sortspec"));
         argumentType.addField(new FieldDescription(LIST_FEATURES, "string", RANKFEATURES.toString()));
         argumentType.addField(new FieldDescription(FRESHNESS, "string", "datetime"));
-        argumentType.addField(new FieldDescription(QUERYCACHE, "string"));
+        argumentType.addField(new FieldDescription(QUERYCACHE, "boolean"));
+        argumentType.addField(new FieldDescription(RERANKCOUNT, "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())));
@@ -67,7 +69,7 @@ public class Ranking implements Cloneable {
     }
     public static QueryProfileType getArgumentType() { return argumentType; }
 
-    private Query parent;
+    private final Query parent;
 
     /** The location of the query is used for distance ranking */
     private Location location = null;
@@ -85,6 +87,8 @@ public class Ranking implements Cloneable {
 
     private boolean queryCache = false;
 
+    private Integer rerankCount = null;
+
     private RankProperties rankProperties = new RankProperties();
 
     private RankFeatures rankFeatures = new RankFeatures();
@@ -140,6 +144,15 @@ public class Ranking implements Cloneable {
 
     public boolean getQueryCache() { return queryCache; }
 
+    /**
+     * Sets the number of hits for which the second-phase function will be evaluated.
+     * When set, this overrides the setting in the rank profile.
+     */
+    public void setRerankCount(int rerankCount) { this.rerankCount = rerankCount; }
+
+    /** Returns the rerank-count that will be used, or null if not set */
+    public Integer getRerankCount() { return rerankCount; }
+
     /** Returns the location of this query, or null if none */
     public Location getLocation() { return location; }
 
@@ -258,6 +271,8 @@ public class Ranking implements Cloneable {
         matching.prepare(rankProperties);
         softTimeout.prepare(rankProperties);
         prepareNow(freshness);
+        if (rerankCount != null)
+            rankProperties.put("vespa.hitcollector.heapsize", rerankCount);
     }
 
     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 6186f588a21..4c65e8003e5 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
@@ -68,6 +68,7 @@ public class QueryProperties extends Properties {
                 if (key.last().equals(Ranking.SORTING)) return ranking.getSorting();
                 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.LIST_FEATURES)) return ranking.getListFeatures();
             }
             else if (key.size() >= 3 && key.get(1).equals(Ranking.MATCH_PHASE)) {
@@ -187,6 +188,8 @@ public class QueryProperties extends Properties {
                         ranking.setFreshness(asString(value, ""));
                     else if (key.last().equals(Ranking.QUERYCACHE))
                         ranking.setQueryCache(asBoolean(value, false));
+                    else if (key.last().equals(Ranking.RERANKCOUNT))
+                        ranking.setRerankCount(asInteger(value, null));
                     else if (key.last().equals(Ranking.LIST_FEATURES))
                         ranking.setListFeatures(asBoolean(value,false));
                     else