Introduce top-k-probability and use it to fetch correct proper amount of hits from each partition

8 files changed, 53 insertions, 3 deletions

diff --git a/config-model/src/main/java/com/yahoo/config/model/deploy/TestProperties.java b/config-model/src/main/java/com/yahoo/config/model/deploy/TestProperties.java
index 4c9e9489c63..a772d7c8a1f 100644
--- a/config-model/src/main/java/com/yahoo/config/model/deploy/TestProperties.java
+++ b/config-model/src/main/java/com/yahoo/config/model/deploy/TestProperties.java
@@ -40,6 +40,7 @@ public class TestProperties implements ModelContext.Properties {
     private boolean isFirstTimeDeployment = false;
     private boolean useDedicatedNodeForLogserver = false;
     private boolean useAdaptiveDispatch = false;
+    private double topKProbability = 1.0;
     private double defaultTermwiseLimit = 1.0;
     private Optional<EndpointCertificateSecrets> endpointCertificateSecrets = Optional.empty();
     private AthenzDomain athenzDomain;
@@ -61,6 +62,7 @@ public class TestProperties implements ModelContext.Properties {
     @Override public Optional<EndpointCertificateSecrets> endpointCertificateSecrets() { return endpointCertificateSecrets; }
     @Override public Optional<TlsSecrets> tlsSecrets() { return endpointCertificateSecrets.map(TlsSecrets::new); }
     @Override public double defaultTermwiseLimit() { return defaultTermwiseLimit; }
+    @Override public double defaultTopKProbability() { return topKProbability; }
     @Override public boolean useBucketSpaceMetric() { return true; }
     @Override public Optional<AthenzDomain> athenzDomain() { return Optional.ofNullable(athenzDomain); }
 
@@ -69,6 +71,11 @@ public class TestProperties implements ModelContext.Properties {
         return this;
     }
 
+    public TestProperties setTopKProbability(double probability) {
+        topKProbability = probability;
+        return this;
+    }
+
     public TestProperties setApplicationId(ApplicationId applicationId) {
         this.applicationId = applicationId;
         return this;
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/content/DispatchTuning.java b/config-model/src/main/java/com/yahoo/vespa/model/content/DispatchTuning.java
index 0d15207b6ce..0f9eb5341ab 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/content/DispatchTuning.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/content/DispatchTuning.java
@@ -11,18 +11,25 @@ public class DispatchTuning {
 
     public static final DispatchTuning empty = new DispatchTuning.Builder().build();
 
-    public enum DispatchPolicy { ROUNDROBIN, ADAPTIVE};
+    public enum DispatchPolicy { ROUNDROBIN, ADAPTIVE}
 
     private final Integer maxHitsPerPartition;
     private DispatchPolicy dispatchPolicy;
     private final Double minGroupCoverage;
     private final Double minActiveDocsCoverage;
 
+    public Double getTopkProbability() {
+        return topkProbability;
+    }
+
+    private final Double topkProbability;
+
     private DispatchTuning(Builder builder) {
         maxHitsPerPartition = builder.maxHitsPerPartition;
         dispatchPolicy = builder.dispatchPolicy;
         minGroupCoverage = builder.minGroupCoverage;
         minActiveDocsCoverage = builder.minActiveDocsCoverage;
+        topkProbability = builder.topKProbability;
     }
 
     /** Returns the max number of hits to fetch from each partition, or null to fetch all */
@@ -46,6 +53,7 @@ public class DispatchTuning {
         private DispatchPolicy dispatchPolicy;
         private Double minGroupCoverage;
         private Double minActiveDocsCoverage;
+        private Double topKProbability;
 
         public DispatchTuning build() {
             return new DispatchTuning(this);
@@ -55,6 +63,10 @@ public class DispatchTuning {
             this.maxHitsPerPartition = maxHitsPerPartition;
             return this;
         }
+        public Builder setTopKProbability(Double topKProbability) {
+            this.topKProbability = topKProbability;
+            return this;
+        }
         public Builder setDispatchPolicy(String policy) {
             if (policy != null)
                 dispatchPolicy = toDispatchPolicy(policy);
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/DomTuningDispatchBuilder.java b/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/DomTuningDispatchBuilder.java
index b53d66632a8..d599a1a1aca 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/DomTuningDispatchBuilder.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/content/cluster/DomTuningDispatchBuilder.java
@@ -23,6 +23,7 @@ public class DomTuningDispatchBuilder {
             return builder.build();
         }
         builder.setMaxHitsPerPartition(dispatchElement.childAsInteger("max-hits-per-partition"));
+        builder.setTopKProbability(dispatchElement.childAsDouble("top-k-probability"));
         builder.setDispatchPolicy(dispatchElement.childAsString("dispatch-policy"));
         builder.setMinGroupCoverage(dispatchElement.childAsDouble("min-group-coverage"));
         builder.setMinActiveDocsCoverage(dispatchElement.childAsDouble("min-active-docs-coverage"));
diff --git a/config-model/src/main/java/com/yahoo/vespa/model/search/IndexedSearchCluster.java b/config-model/src/main/java/com/yahoo/vespa/model/search/IndexedSearchCluster.java
index 9746c50450e..56adc227df4 100644
--- a/config-model/src/main/java/com/yahoo/vespa/model/search/IndexedSearchCluster.java
+++ b/config-model/src/main/java/com/yahoo/vespa/model/search/IndexedSearchCluster.java
@@ -53,6 +53,7 @@ public class IndexedSearchCluster extends SearchCluster
     private final DispatchGroup rootDispatch;
     private DispatchSpec dispatchSpec;
     private final boolean useAdaptiveDispatch;
+    private final double defaultTopKProbability;
     private List<SearchNode> searchNodes = new ArrayList<>();
 
     /**
@@ -70,6 +71,7 @@ public class IndexedSearchCluster extends SearchCluster
         unionCfg = new UnionConfiguration(this, documentDbs);
         rootDispatch =  new DispatchGroup(this);
         useAdaptiveDispatch = deployState.getProperties().useAdaptiveDispatch();
+        defaultTopKProbability = deployState.getProperties().defaultTopKProbability();
     }
 
     @Override
@@ -307,7 +309,11 @@ public class IndexedSearchCluster extends SearchCluster
         }
         if (useAdaptiveDispatch)
             builder.distributionPolicy(DistributionPolicy.ADAPTIVE);
-
+        if (tuning.dispatch.getTopkProbability() != null) {
+            builder.topKProbability(tuning.dispatch.getTopkProbability());
+        } else {
+            builder.topKProbability(defaultTopKProbability);
+        }
         if (tuning.dispatch.getMinActiveDocsCoverage() != null)
             builder.minActivedocsPercentage(tuning.dispatch.getMinActiveDocsCoverage());
         if (tuning.dispatch.getMinGroupCoverage() != null)
diff --git a/config-model/src/main/resources/schema/content.rnc b/config-model/src/main/resources/schema/content.rnc
index b1821680b14..481d82ebb4b 100644
--- a/config-model/src/main/resources/schema/content.rnc
+++ b/config-model/src/main/resources/schema/content.rnc
@@ -85,6 +85,7 @@ DispatchTuning = element dispatch {
     element dispatch-policy { string "round-robin" | string "adaptive" | string "random" }? &
     element min-group-coverage { xsd:double }? &
     element min-active-docs-coverage { xsd:double }? &
+    element top-k-probability { xsd:double }? &
     element use-local-node { string "true" | string "false" }?
 }
 
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/content/ContentClusterTest.java b/config-model/src/test/java/com/yahoo/vespa/model/content/ContentClusterTest.java
index b08cc92d20c..4d5df7c1965 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/content/ContentClusterTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/content/ContentClusterTest.java
@@ -933,6 +933,17 @@ public class ContentClusterTest extends ContentBaseTest {
         assertEquals(distributionBits, storDistributormanagerConfig.minsplitcount());
     }
 
+    private void verifyTopKProbabilityPropertiesControl(double topKProbability) {
+        VespaModel model = createEnd2EndOneNode(new TestProperties().setTopKProbability(topKProbability));
+
+        ContentCluster cc = model.getContentClusters().get("storage");
+        DispatchConfig.Builder builder = new DispatchConfig.Builder();
+        cc.getSearch().getConfig(builder);
+
+        DispatchConfig cfg = new DispatchConfig(builder);
+        assertEquals(topKProbability, cfg.topKProbability(), 0.0);
+    }
+
     private void verifyRoundRobinPropertiesControl(boolean useAdaptiveDispatch) {
         VespaModel model = createEnd2EndOneNode(new TestProperties().setUseAdaptiveDispatch(useAdaptiveDispatch));
 
@@ -946,7 +957,6 @@ public class ContentClusterTest extends ContentBaseTest {
         } else {
             assertEquals(DispatchConfig.DistributionPolicy.ROUNDROBIN, cfg.distributionPolicy());
         }
-
     }
 
     @Test
@@ -955,5 +965,12 @@ public class ContentClusterTest extends ContentBaseTest {
         verifyRoundRobinPropertiesControl(true);
     }
 
+    @Test
+    public void default_topKprobability_controlled_by_properties() {
+        verifyTopKProbabilityPropertiesControl(1.0);
+        verifyTopKProbabilityPropertiesControl(0.999);
+        verifyTopKProbabilityPropertiesControl(0.77);
+    }
+
 
 }
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/content/DispatchTuningTest.java b/config-model/src/test/java/com/yahoo/vespa/model/content/DispatchTuningTest.java
index f708d7673e2..8a46aaaa230 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/content/DispatchTuningTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/content/DispatchTuningTest.java
@@ -19,11 +19,13 @@ public class DispatchTuningTest {
                 .setDispatchPolicy("round-robin")
                 .setMinGroupCoverage(7.5)
                 .setMinActiveDocsCoverage(12.5)
+                .setTopKProbability(18.3)
                 .build();
         assertEquals(69, dispatch.getMaxHitsPerPartition().intValue());
         assertEquals(7.5, dispatch.getMinGroupCoverage().doubleValue(), 0.0);
         assertEquals(12.5, dispatch.getMinActiveDocsCoverage().doubleValue(), 0.0);
         assertTrue(DispatchTuning.DispatchPolicy.ROUNDROBIN == dispatch.getDispatchPolicy());
+        assertEquals(18.3, dispatch.getTopkProbability(), 0.0);
     }
     @Test
     public void requireThatRandomDispatchWork() {
@@ -52,6 +54,7 @@ public class DispatchTuningTest {
         assertNull(dispatch.getDispatchPolicy());
         assertNull(dispatch.getMinActiveDocsCoverage());
         assertNull(dispatch.getMinGroupCoverage());
+        assertNull(dispatch.getTopkProbability());
     }
 
 }
diff --git a/config-model/src/test/java/com/yahoo/vespa/model/content/cluster/DomDispatchTuningBuilderTest.java b/config-model/src/test/java/com/yahoo/vespa/model/content/cluster/DomDispatchTuningBuilderTest.java
index 7fa27f74d74..abfb03e41dd 100644
--- a/config-model/src/test/java/com/yahoo/vespa/model/content/cluster/DomDispatchTuningBuilderTest.java
+++ b/config-model/src/test/java/com/yahoo/vespa/model/content/cluster/DomDispatchTuningBuilderTest.java
@@ -47,6 +47,7 @@ public class DomDispatchTuningBuilderTest {
         assertNull(dispatch.getMinGroupCoverage());
         assertNull(dispatch.getMinActiveDocsCoverage());
         assertNull(dispatch.getDispatchPolicy());
+        assertNull(dispatch.getTopkProbability());
     }
 
     @Test
@@ -58,12 +59,14 @@ public class DomDispatchTuningBuilderTest {
                 "      <max-hits-per-partition>69</max-hits-per-partition>" +
                 "      <min-group-coverage>7.5</min-group-coverage>" +
                 "      <min-active-docs-coverage>12.5</min-active-docs-coverage>" +
+                "      <top-k-probability>0.999</top-k-probability>" +
                 "    </dispatch>" +
                 "  </tuning>" +
                 "</content>");
         assertEquals(69, dispatch.getMaxHitsPerPartition().intValue());
         assertEquals(7.5, dispatch.getMinGroupCoverage().doubleValue(), 0.0);
         assertEquals(12.5, dispatch.getMinActiveDocsCoverage().doubleValue(), 0.0);
+        assertEquals(0.999, dispatch.getTopkProbability().doubleValue(), 0.0);
     }
     @Test
     public void requireThatTuningDispatchPolicyRoundRobin() throws Exception {