Avoid wrapping leaf collector when possible in QueryPhaseCollector (elastic#97452)

The recently introduced QueryPhaseCollector introduces some overhead
caused by wrapping of the leaf collector that could be avoided.

When either a post_filter, or a min_score, or terminate_after are
provided, we have no other choice than wrapping to filter collection,
but if none of these are used, we can skip wrapping when one of the two
leaf collectors early terminates.

Author: javanna

server/src/main/java/org/elasticsearch/search/query/QueryPhaseCollector.java

@@ -12,6 +12,7 @@
 import org.apache.lucene.search.CollectionTerminatedException;
 import org.apache.lucene.search.Collector;
 import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.FilterLeafCollector;
 import org.apache.lucene.search.FilterScorable;
 import org.apache.lucene.search.LeafCollector;
 import org.apache.lucene.search.Scorable;
@@ -118,7 +119,7 @@ private boolean isDocWithinMinScore(Scorable scorer) throws IOException {
         return minScore == null || scorer.score() >= minScore;
     }
 
-    private boolean doesDocMatchPostFilter(int doc, Bits postFilterBits) {
+    private static boolean doesDocMatchPostFilter(int doc, Bits postFilterBits) {
         return postFilterBits == null || postFilterBits.get(doc);
     }
 
@@ -143,31 +144,36 @@ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOExcept
         Bits postFilterBits = getPostFilterBits(context);
 
         if (aggsCollector == null) {
-            LeafCollector topDocsLeafCollector;
+            LeafCollector tdlc = null;
             try {
-                topDocsLeafCollector = topDocsCollector.getLeafCollector(context);
+                tdlc = topDocsCollector.getLeafCollector(context);
             } catch (@SuppressWarnings("unused") CollectionTerminatedException e) {
                 // TODO we keep on collecting although we have nothing to collect (there is no top docs nor aggs leaf collector).
                 // The reason is only to set the early terminated flag to the QueryResult like some tests expect. This needs fixing.
                 if (terminateAfter == 0) {
                     throw e;
                 }
-                topDocsLeafCollector = null;
+            }
+            final LeafCollector topDocsLeafCollector = tdlc;
+            if (postFilterBits == null && terminateAfter == 0 && minScore == null) {
+                // no need to wrap if we just need to collect unfiltered docs through leaf collector
+                // aggs collector was not originally provided so the overall score mode is that of the top docs collector
+                return topDocsLeafCollector;
             }
             return new TopDocsLeafCollector(postFilterBits, topDocsLeafCollector);
         }
 
-        LeafCollector topDocsLeafCollector;
+        LeafCollector tdlc = null;
         try {
-            topDocsLeafCollector = topDocsCollector.getLeafCollector(context);
+            tdlc = topDocsCollector.getLeafCollector(context);
         } catch (@SuppressWarnings("unused") CollectionTerminatedException e) {
             // top docs collector does not need this segment, but the aggs collector does.
-            topDocsLeafCollector = null;
         }
+        final LeafCollector topDocsLeafCollector = tdlc;
 
-        LeafCollector aggsLeafCollector;
+        LeafCollector alf = null;
         try {
-            aggsLeafCollector = aggsCollector.getLeafCollector(context);
+            alf = aggsCollector.getLeafCollector(context);
         } catch (@SuppressWarnings("unused") CollectionTerminatedException e) {
             // aggs collector does not need this segment, but the top docs collector may.
             if (topDocsLeafCollector == null) {
@@ -177,11 +183,41 @@ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOExcept
                     throw e;
                 }
             }
-            aggsLeafCollector = null;
         }
-        // say that the aggs collector early terminates while the top docs collector does not, we still want to wrap in the same way
+        final LeafCollector aggsLeafCollector = alf;
+
+        if (topDocsLeafCollector == null && terminateAfter == 0 && minScore == null) {
+            // top docs collector early terminated, we can avoid wrapping as long as we don't need to apply terminate_after and min_score.
+            // post_filter does not matter because it's not applied to aggs collection anyways. terminate_after matters only until we
+            // address the different TODOs around needless collection to honour terminate_after after early termination.
+            // aggs don't support skipping low scoring hits, so we can rely on setMinCompetitiveScore being a no-op already.
+            return aggsLeafCollector;
+        }
+
+        // if that the aggs collector early terminates while the top docs collector does not, we still need to wrap the leaf collector
         // to enforce that setMinCompetitiveScore is a no-op. Otherwise we may allow the top docs collector to skip non competitive
-        // hits despite the score mode of the Collector did not allow it.
+        // hits despite the score mode of the Collector did not allow it (because aggs don't support TOP_SCORES).
+        if (aggsLeafCollector == null && postFilterBits == null && terminateAfter == 0 && minScore == null) {
+            // special case for early terminated aggs
+            return new FilterLeafCollector(topDocsLeafCollector) {
+                @Override
+                public void setScorer(Scorable scorer) throws IOException {
+                    super.setScorer(new FilterScorable(scorer) {
+                        @Override
+                        public void setMinCompetitiveScore(float minScore) {
+                            // Ignore calls to setMinCompetitiveScore. The top docs collector may try to skip low
+                            // scoring hits, but the overall score_mode won't allow it because an aggs collector
+                            // was originally provided which never supports TOP_SCORES is not supported for aggs
+                        }
+                    });
+                }
+
+                @Override
+                public DocIdSetIterator competitiveIterator() throws IOException {
+                    return topDocsLeafCollector.competitiveIterator();
+                }
+            };
+        }
         return new CompositeLeafCollector(postFilterBits, topDocsLeafCollector, aggsLeafCollector);
     }
 
@@ -259,7 +295,7 @@ private class CompositeLeafCollector implements LeafCollector {
 
         @Override
         public void setScorer(Scorable scorer) throws IOException {
-            if (cacheScores) {
+            if (cacheScores && topDocsLeafCollector != null && aggsLeafCollector != null) {
                 scorer = ScoreCachingWrappingScorer.wrap(scorer);
             }
             scorer = new FilterScorable(scorer) {
@@ -300,9 +336,9 @@ public void collect(int doc) throws IOException {
                     }
                 }
             }
-            // min_score is applied to aggs as well as top hits
-            if (isDocWithinMinScore(scorer)) {
-                if (aggsLeafCollector != null) {
+            if (aggsLeafCollector != null) {
+                // min_score is applied to aggs as well as top hits
+                if (isDocWithinMinScore(scorer)) {
                     try {
                         aggsLeafCollector.collect(doc);
                     } catch (@SuppressWarnings("unused") CollectionTerminatedException e) {

server/src/test/java/org/elasticsearch/search/query/QueryPhaseCollectorTests.java

@@ -40,6 +40,7 @@
 import org.apache.lucene.tests.search.DummyTotalHitCountCollector;
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.test.ESTestCase;
+import org.hamcrest.CoreMatchers;
 import org.hamcrest.Matchers;
 
 import java.io.IOException;
@@ -1012,21 +1013,53 @@ public void testCacheScoresIfNecessary() throws IOException {
             QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector(c1, null, 0, c2, null);
             queryPhaseCollector.getLeafCollector(ctx).setScorer(new MinCompetitiveScoreScorable());
         }
+        {
+            // both collectors need scores => caching, but one early terminates
+            Collector c1 = new TerminateAfterCollector(new MockCollector(ScoreMode.COMPLETE, MinCompetitiveScoreScorable.class), 0);
+            Collector c2 = new MockCollector(ScoreMode.COMPLETE, MinCompetitiveScoreScorable.class);
+            QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector(c1, null, 0, c2, null);
+            queryPhaseCollector.getLeafCollector(ctx).setScorer(new MinCompetitiveScoreScorable());
+            queryPhaseCollector = new QueryPhaseCollector(c2, null, 0, c1, null);
+            queryPhaseCollector.getLeafCollector(ctx).setScorer(new MinCompetitiveScoreScorable());
+        }
     }
 
-    public void testCompetitiveIteratorTopDocsOnly() throws IOException {
+    public void testNoWrappingIfUnnecessaryTopDocsOnly() throws IOException {
         MockCollector mockCollector = new MockCollector(randomFrom(ScoreMode.values()));
         QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector(mockCollector, null, 0, null, null);
         LeafReaderContext context = searcher.getLeafContexts().get(0);
         LeafCollector leafCollector = queryPhaseCollector.getLeafCollector(context);
+        assertThat(leafCollector, CoreMatchers.sameInstance(mockCollector));
+    }
+
+    public void testNoWrappingIfUnnecessaryTopDocsEarlyTerminated() throws IOException {
+        TerminateAfterCollector topDocsCollector = new TerminateAfterCollector(new MockCollector(randomFrom(ScoreMode.values())), 0);
+        MockCollector aggsCollector = new MockCollector(randomScoreModeExceptTopScores());
+        QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector(topDocsCollector, null, 0, aggsCollector, null);
+        LeafReaderContext context = searcher.getLeafContexts().get(0);
+        LeafCollector leafCollector = queryPhaseCollector.getLeafCollector(context);
+        assertThat(leafCollector, CoreMatchers.sameInstance(aggsCollector));
+    }
+
+    public void testCompetitiveIteratorNoAggs() throws IOException {
+        // use a post_filter so that we wrap the top docs leaf collector, as this test verifies that
+        // the wrapper calls competitiveIterator when appropriated
+        Weight postFilterWeight = searcher.createWeight(new MatchAllDocsQuery(), ScoreMode.COMPLETE_NO_SCORES, 1.0f);
+        MockCollector mockCollector = new MockCollector(randomFrom(ScoreMode.values()));
+        QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector(mockCollector, postFilterWeight, 0, null, null);
+        LeafReaderContext context = searcher.getLeafContexts().get(0);
+        LeafCollector leafCollector = queryPhaseCollector.getLeafCollector(context);
         leafCollector.competitiveIterator();
         assertTrue(mockCollector.competitiveIteratorCalled);
     }
 
-    public void testCompetitiveIteratorTopDocsOnlyCollectionTerminated() throws IOException {
+    public void testCompetitiveIteratorNoAggsCollectionTerminated() throws IOException {
+        // use a post_filter so that we wrap the top docs leaf collector, as this test verifies that
+        // the wrapper calls competitiveIterator when appropriated
+        Weight postFilterWeight = searcher.createWeight(new MatchAllDocsQuery(), ScoreMode.COMPLETE_NO_SCORES, 1.0f);
         MockCollector mockCollector = new MockCollector(randomFrom(ScoreMode.values()));
         TerminateAfterCollector terminateAfterCollector = new TerminateAfterCollector(mockCollector, 1);
-        QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector(terminateAfterCollector, null, 0, null, null);
+        QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector(terminateAfterCollector, postFilterWeight, 0, null, null);
         LeafReaderContext context = searcher.getLeafContexts().get(0);
         LeafCollector leafCollector = queryPhaseCollector.getLeafCollector(context);
         leafCollector.competitiveIterator();