The estimate is not bounded by [1, K] so it must be capped.

2 files changed, 60 insertions, 2 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/TopKEstimator.java b/container-search/src/main/java/com/yahoo/search/dispatch/TopKEstimator.java
index 8003d9c6744..d3f222a9f3e 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/TopKEstimator.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/TopKEstimator.java
@@ -31,12 +31,12 @@ public class TopKEstimator {
     }
     public int estimateK(int k, int n) {
         return (estimate && n > 1)
-                ? (int)Math.ceil(estimateExactK(k, n, defaultP))
+                ? Math.min(k, (int)Math.ceil(estimateExactK(k, n, defaultP)))
                 : k;
     }
     public int estimateK(int k, int n, double p) {
         return (needEstimate(p) && (n > 1))
-                ? (int)Math.ceil(estimateExactK(k, n, p))
+                ? Math.min(k, (int)Math.ceil(estimateExactK(k, n, p)))
                 : k;
     }
 }
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/TopKEstimatorTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/TopKEstimatorTest.java
index c14e4f984f1..5d978e6b592 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/TopKEstimatorTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/TopKEstimatorTest.java
@@ -25,4 +25,62 @@ public class TopKEstimatorTest {
         assertEquals(44.909040374464155, estimator.estimateExactK(200, 10, 0.99999), 0.0);
         assertEquals(45, estimator.estimateK(200, 10, 0.99999));
     }
+    @Test
+    public void requireHitsAreEstimatedAccordingToPartitionsAndProbabilityForVaryingN_K200() {
+        TopKEstimator estimator = new TopKEstimator(30, 0.99999);
+        assertEquals(200, estimator.estimateExactK(200, 1), 0.0);
+        assertEquals(200, estimator.estimateK(200, 1));
+        assertEquals(137.4727798056239, estimator.estimateExactK(200, 2), 0.0);
+        assertEquals(102.95409291533568, estimator.estimateExactK(200, 3), 0.0);
+        assertEquals(44.909040374464155, estimator.estimateExactK(200, 10), 0.0);
+        assertEquals(28.86025772029091, estimator.estimateExactK(200, 20), 0.0);
+    }
+
+    @Test
+    public void requireHitsAreEstimatedAccordingToPartitionsAndProbabilityForVaryingN_K20() {
+        TopKEstimator estimator = new TopKEstimator(30, 0.99999);
+        assertEquals(20, estimator.estimateExactK(20, 1), 0.0);
+        assertEquals(20, estimator.estimateK(20, 1));
+        assertEquals(21.849933444373328, estimator.estimateExactK(20, 2), 0.0);
+        assertEquals(18.14175840378403, estimator.estimateExactK(20, 3), 0.0);
+        assertEquals(9.87693019124002, estimator.estimateExactK(20, 10), 0.0);
+        assertEquals(6.964137165389415, estimator.estimateExactK(20, 20), 0.0);
+    }
+
+    @Test
+    public void requireHitsAreEstimatedAccordingToPartitionsAndProbabilityForVaryingN_K10_Five9() {
+        TopKEstimator estimator = new TopKEstimator(30, 0.99999);
+        assertEquals(10, estimator.estimateExactK(10, 1), 0.0);
+        assertEquals(10, estimator.estimateK(10, 1));
+        assertEquals(13.379168295125641, estimator.estimateExactK(10, 2), 0.0);
+        assertEquals(11.447448515386741, estimator.estimateExactK(10, 3), 0.0);
+        assertEquals(6.569830753158866, estimator.estimateExactK(10, 10), 0.0);
+        assertEquals(4.717281833573569, estimator.estimateExactK(10, 20), 0.0);
+    }
+
+    @Test
+    public void requireHitsAreEstimatedAccordingToPartitionsAndProbabilityForVaryingN_K10_Four9() {
+        TopKEstimator estimator = new TopKEstimator(30, 0.9999);
+        assertEquals(10, estimator.estimateExactK(10, 1), 0.0);
+        assertEquals(10, estimator.estimateK(10, 1));
+        assertEquals(12.087323848369289, estimator.estimateExactK(10, 2), 0.0);
+        assertEquals(10.230749855131009, estimator.estimateExactK(10, 3), 0.0);
+        assertEquals(5.794676146031378, estimator.estimateExactK(10, 10), 0.0);
+        assertEquals(4.152394782937266, estimator.estimateExactK(10, 20), 0.0);
+    }
+
+    @Test
+    public void requireEstimatesAreRoundeUp() {
+        TopKEstimator estimator = new TopKEstimator(30, 0.9999);
+        assertEquals(5.794676146031378, estimator.estimateExactK(10, 10), 0.0);
+        assertEquals(6, estimator.estimateK(10, 10));
+    }
+
+    @Test
+    public void requireEstimatesAreCappedAtInputK() {
+        TopKEstimator estimator = new TopKEstimator(30, 0.9999);
+        assertEquals(12.087323848369289, estimator.estimateExactK(10, 2), 0.0);
+        assertEquals(10, estimator.estimateK(10, 2));
+    }
+
 }