Check for unlikely doc distribution

cbuescher · Jul 5, 2023 · ffa7ea7 · ffa7ea7
1 parent d94461a
commit ffa7ea7
Showing 1 changed file with 24 additions and 2 deletions.
diff --git a/server/src/test/java/org/elasticsearch/search/query/PartialHitCountCollectorTests.java b/server/src/test/java/org/elasticsearch/search/query/PartialHitCountCollectorTests.java
@@ -12,18 +12,22 @@
 import org.apache.lucene.document.Field;
 import org.apache.lucene.document.StringField;
 import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.search.ScoreMode;
 import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.search.Weight;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
+import java.util.List;
 
 public class PartialHitCountCollectorTests extends ESTestCase {
 
@@ -109,7 +113,25 @@ public void testCollectedHitCountEarlyTerminated() throws Exception {
         int totalHitsThreshold = randomInt(2);
         PartialHitCountCollector partialHitCountCollector = new PartialHitCountCollector(totalHitsThreshold);
         searcher.search(query, partialHitCountCollector);
-        assertEquals(totalHitsThreshold, partialHitCountCollector.getTotalHits());
-        assertTrue(partialHitCountCollector.hasEarlyTerminated());
+
+        try {
+            assertEquals(totalHitsThreshold, partialHitCountCollector.getTotalHits());
+            assertTrue(partialHitCountCollector.hasEarlyTerminated());
+        } catch (AssertionError e) {
+            // if we get very unlucky with the document distribution, we can get more than one segment with only one document
+            // in this case early termination might not kick in because we never collect
+            // check if this was the case here and disregard error in that case
+            Weight weight = searcher.createWeight(query, ScoreMode.TOP_DOCS, 1.0f);
+            List<LeafReaderContext> leafContexts = searcher.getLeafContexts();
+            int totalHitsWithoutCollecting = 0;
+            for (LeafReaderContext leafContext : leafContexts) {
+                int count = weight.count(leafContext);
+                if (count != -1) {
+                    totalHitsWithoutCollecting += count;
+                }
+            }
+            assertTrue(totalHitsWithoutCollecting > totalHitsThreshold);
+            assertFalse(partialHitCountCollector.hasEarlyTerminated());
+        }
     }
 }