Fix test failures for termintate_early cases

Signed-off-by: Sandesh Kumar <[email protected]>
opensearch-project · Feb 8, 2024 · 91ccaef · 91ccaef
1 parent 4f1a611
commit 91ccaef
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/.../org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java b/.../org/opensearch/search/aggregations/bucket/terms/GlobalOrdinalsStringTermsAggregator.java
@@ -171,7 +171,8 @@ LeafBucketCollector termDocFreqCollector(
         BiConsumer<Long, Integer> ordCountConsumer
     ) throws IOException {
         if (weight == null || weight.count(ctx) != ctx.reader().maxDoc()) {
-            // Weight not assigned or top-level query does not match all docs in the segment.
+            // weight.count(ctx) == ctx.reader().maxDoc() implies there are no deleted documents and
+            // top-level query matches all docs in the segment
             return null;
         }
 

diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java
@@ -122,6 +122,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static org.opensearch.search.query.TopDocsCollectorContext.hasInfMaxScore;
+import static org.hamcrest.Matchers.allOf;
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
@@ -437,10 +438,16 @@ public void testTerminateAfterEarlyTermination() throws Exception {
             assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L));
             assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(1));
 
+            // Do not expect an exact match when terminate_after is used in conjunction to size = 0 as an optimization introduced by
+            // https://issues.apache.org/jira/browse/LUCENE-10620 can produce a total hit count >= terminated_after, because
+            // TotalHitCountCollector is used in this case as part of Weight#count() optimization
             context.setSize(0);
             QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher);
             assertTrue(context.queryResult().terminatedEarly());
-            assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L));
+            assertThat(
+                context.queryResult().topDocs().topDocs.totalHits.value,
+                allOf(greaterThanOrEqualTo(1L), lessThanOrEqualTo((long) numDocs))
+            );
             assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0));
         }
 
@@ -466,7 +473,10 @@ public void testTerminateAfterEarlyTermination() throws Exception {
             context.parsedQuery(new ParsedQuery(bq));
             QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher);
             assertTrue(context.queryResult().terminatedEarly());
-            assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L));
+            assertThat(
+                context.queryResult().topDocs().topDocs.totalHits.value,
+                allOf(greaterThanOrEqualTo(1L), lessThanOrEqualTo((long) numDocs))
+            );
             assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0));
         }
         {
@@ -486,9 +496,12 @@ public void testTerminateAfterEarlyTermination() throws Exception {
             context.queryCollectorManagers().put(TotalHitCountCollector.class, manager);
             QueryPhase.executeInternal(context.withCleanQueryResult(), queryPhaseSearcher);
             assertTrue(context.queryResult().terminatedEarly());
-            assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(1L));
+            assertThat(
+                context.queryResult().topDocs().topDocs.totalHits.value,
+                allOf(greaterThanOrEqualTo(1L), lessThanOrEqualTo((long) numDocs))
+            );
             assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0));
-            assertThat(manager.getTotalHits(), equalTo(1));
+            assertThat(manager.getTotalHits(), allOf(greaterThanOrEqualTo(1), lessThanOrEqualTo(numDocs)));
         }
 
         // tests with trackTotalHits and terminateAfter
@@ -503,15 +516,18 @@ public void testTerminateAfterEarlyTermination() throws Exception {
             if (trackTotalHits == -1) {
                 assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo(0L));
             } else {
-                assertThat(context.queryResult().topDocs().topDocs.totalHits.value, equalTo((long) Math.min(trackTotalHits, 10)));
+                assertThat(
+                    context.queryResult().topDocs().topDocs.totalHits.value,
+                    allOf(greaterThanOrEqualTo(Math.min(trackTotalHits, 10L)), lessThanOrEqualTo((long) numDocs))
+                );
             }
             assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(0));
             // The concurrent search terminates the collection when the number of hits is reached by each
             // concurrent collector. In this case, in general, the number of results are multiplied by the number of
             // slices (as the unit of concurrency). To address that, we have to use the shared global state,
             // much as HitsThresholdChecker does.
             if (executor == null) {
-                assertThat(manager.getTotalHits(), equalTo(10));
+                assertThat(manager.getTotalHits(), allOf(greaterThanOrEqualTo(Math.min(trackTotalHits, 10)), lessThanOrEqualTo(numDocs)));
             }
         }