Introduce a collector manager for QueryPhaseCollector (elastic#97410)

javanna · web-flow · commit 4ade73faaed2 · 2023-07-14T09:15:33.000+02:00
The query phase now relies on QueryPhaseCollector to wrap top docs and aggs collector as well as to apply terminate_after, post_filter and min_score. In order to enable inter-segment concurrency, we need to introduce a corresponding collector manager that can create multiple query phase collectors and reduce them to a single resultset.

Note: we are only introducing the collector manager and corresponding tests, but not using it yet in QueryPhase.
diff --git a/docs/changelog/97410.yaml b/docs/changelog/97410.yaml
@@ -0,0 +1,5 @@
+pr: 97410
+summary: Introduce a collector manager for `QueryPhaseCollector`
+area: Search
+type: enhancement
+issues: []
diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhaseCollector.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhaseCollector.java
@@ -24,7 +24,11 @@
 import org.elasticsearch.common.lucene.Lucene;
 
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * Top-level collector used in the query phase to perform top hits collection as well as aggs collection.
@@ -40,21 +44,26 @@
 final class QueryPhaseCollector implements Collector {
     private final Collector aggsCollector;
     private final Collector topDocsCollector;
-    private final int terminateAfter;
+    private final TerminateAfterChecker terminateAfterChecker;
     private final Weight postFilterWeight;
     private final Float minScore;
     private final boolean cacheScores;
-
-    private int numCollected;
     private boolean terminatedAfter = false;
 
     QueryPhaseCollector(Collector topDocsCollector, Weight postFilterWeight, int terminateAfter, Collector aggsCollector, Float minScore) {
+        this(topDocsCollector, postFilterWeight, resolveTerminateAfterChecker(terminateAfter), aggsCollector, minScore);
+    }
+
+    QueryPhaseCollector(
+        Collector topDocsCollector,
+        Weight postFilterWeight,
+        TerminateAfterChecker terminateAfterChecker,
+        Collector aggsCollector,
+        Float minScore
+    ) {
         this.topDocsCollector = Objects.requireNonNull(topDocsCollector);
         this.postFilterWeight = postFilterWeight;
-        if (terminateAfter < 0) {
-            throw new IllegalArgumentException("terminateAfter must be greater than or equal to 0");
-        }
-        this.terminateAfter = terminateAfter;
+        this.terminateAfterChecker = terminateAfterChecker;
         this.aggsCollector = aggsCollector;
         this.minScore = minScore;
         this.cacheScores = aggsCollector != null && topDocsCollector.scoreMode().needsScores() && aggsCollector.scoreMode().needsScores();
@@ -104,30 +113,16 @@ boolean isTerminatedAfter() {
     }
 
     private boolean shouldCollectTopDocs(int doc, Scorable scorer, Bits postFilterBits) throws IOException {
-        if (isDocWithinMinScore(scorer)) {
-            if (doesDocMatchPostFilter(doc, postFilterBits)) {
-                // terminate_after is purposely applied after post_filter, and terminates aggs collection based on number of filtered
-                // top hits that have been collected. Strange feature, but that has been behaviour for a long time.
-                applyTerminateAfter();
-                return true;
-            }
-        }
-        return false;
+        return isDocWithinMinScore(scorer) && (postFilterBits == null || postFilterBits.get(doc));
     }
 
     private boolean isDocWithinMinScore(Scorable scorer) throws IOException {
         return minScore == null || scorer.score() >= minScore;
     }
 
-    private static boolean doesDocMatchPostFilter(int doc, Bits postFilterBits) {
-        return postFilterBits == null || postFilterBits.get(doc);
-    }
-
-    private void applyTerminateAfter() {
-        if (terminateAfter > 0 && numCollected >= terminateAfter) {
-            terminatedAfter = true;
-            throw new CollectionTerminatedException();
-        }
+    private void earlyTerminate() {
+        terminatedAfter = true;
+        throw new CollectionTerminatedException();
     }
 
     private Bits getPostFilterBits(LeafReaderContext context) throws IOException {
@@ -140,12 +135,14 @@ private Bits getPostFilterBits(LeafReaderContext context) throws IOException {
 
     @Override
     public LeafCollector getLeafCollector(LeafReaderContext context) throws IOException {
-        applyTerminateAfter();
+        if (terminateAfterChecker.isThresholdReached()) {
+            earlyTerminate();
+        }
         Bits postFilterBits = getPostFilterBits(context);
 
         if (aggsCollector == null) {
             final LeafCollector topDocsLeafCollector = topDocsCollector.getLeafCollector(context);
-            if (postFilterBits == null && terminateAfter == 0 && minScore == null) {
+            if (postFilterBits == null && terminateAfterChecker == NO_OP_TERMINATE_AFTER_CHECKER && 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;
@@ -182,7 +179,10 @@ public LeafCollector getLeafCollector(LeafReaderContext context) throws IOExcept
         // 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 (because aggs don't support TOP_SCORES).
-        if (aggsLeafCollector == null && postFilterBits == null && terminateAfter == 0 && minScore == null) {
+        if (aggsLeafCollector == null
+            && postFilterBits == null
+            && terminateAfterChecker == NO_OP_TERMINATE_AFTER_CHECKER
+            && minScore == null) {
             // special case for early terminated aggs
             return new FilterLeafCollector(topDocsLeafCollector) {
                 @Override
@@ -213,7 +213,7 @@ private class TopDocsLeafCollector implements LeafCollector {
 
         TopDocsLeafCollector(Bits postFilterBits, LeafCollector topDocsLeafCollector) {
             assert topDocsLeafCollector != null;
-            assert postFilterBits != null || terminateAfter > 0 || minScore != null;
+            assert postFilterBits != null || terminateAfterChecker != NO_OP_TERMINATE_AFTER_CHECKER || minScore != null;
             this.postFilterBits = postFilterBits;
             this.topDocsLeafCollector = topDocsLeafCollector;
         }
@@ -232,7 +232,11 @@ public DocIdSetIterator competitiveIterator() throws IOException {
         @Override
         public void collect(int doc) throws IOException {
             if (shouldCollectTopDocs(doc, scorer, postFilterBits)) {
-                numCollected++;
+                // terminate_after is purposely applied after post_filter, and terminates aggs collection based on number of filtered
+                // top hits that have been collected. Strange feature, but that has been behaviour for a long time.
+                if (terminateAfterChecker.incrementHitCountAndCheckThreshold()) {
+                    earlyTerminate();
+                }
                 topDocsLeafCollector.collect(doc);
             }
         }
@@ -278,7 +282,9 @@ public void collect(int doc) throws IOException {
             if (shouldCollectTopDocs(doc, scorer, postFilterBits)) {
                 // we keep on counting and checking the terminate_after threshold so that we can terminate aggs collection
                 // even if top docs collection early terminated
-                numCollected++;
+                if (terminateAfterChecker.incrementHitCountAndCheckThreshold()) {
+                    earlyTerminate();
+                }
                 if (topDocsLeafCollector != null) {
                     try {
                         topDocsLeafCollector.collect(doc);
@@ -320,4 +326,135 @@ public DocIdSetIterator competitiveIterator() throws IOException {
             return null;
         }
     }
+
+    static CollectorManager createManager(
+        org.apache.lucene.search.CollectorManager<? extends Collector, Void> topDocsCollectorManager,
+        Weight postFilterWeight,
+        int terminateAfter,
+        org.apache.lucene.search.CollectorManager<? extends Collector, Void> aggsCollectorManager,
+        Float minScore
+    ) {
+        return new CollectorManager(
+            topDocsCollectorManager,
+            postFilterWeight,
+            resolveTerminateAfterChecker(terminateAfter),
+            aggsCollectorManager,
+            minScore
+        );
+    }
+
+    private static TerminateAfterChecker resolveTerminateAfterChecker(int terminateAfter) {
+        if (terminateAfter < 0) {
+            throw new IllegalArgumentException("terminateAfter must be greater than or equal to 0");
+        }
+        return terminateAfter == 0 ? NO_OP_TERMINATE_AFTER_CHECKER : new GlobalTerminateAfterChecker(terminateAfter);
+    }
+
+    private abstract static class TerminateAfterChecker {
+        abstract boolean isThresholdReached();
+
+        abstract boolean incrementHitCountAndCheckThreshold();
+    }
+
+    private static final class GlobalTerminateAfterChecker extends TerminateAfterChecker {
+        private final int terminateAfter;
+        private final AtomicInteger numCollected = new AtomicInteger();
+
+        GlobalTerminateAfterChecker(int terminateAfter) {
+            assert terminateAfter > 0;
+            this.terminateAfter = terminateAfter;
+        }
+
+        boolean isThresholdReached() {
+            return numCollected.getAcquire() >= terminateAfter;
+        }
+
+        boolean incrementHitCountAndCheckThreshold() {
+            return numCollected.incrementAndGet() > terminateAfter;
+        }
+    }
+
+    // no needless counting when terminate_after is not set
+    private static final TerminateAfterChecker NO_OP_TERMINATE_AFTER_CHECKER = new TerminateAfterChecker() {
+        @Override
+        boolean isThresholdReached() {
+            return false;
+        }
+
+        @Override
+        boolean incrementHitCountAndCheckThreshold() {
+            return false;
+        }
+    };
+
+    /**
+     * {@link org.apache.lucene.search.CollectorManager} implementation based on {@link QueryPhaseCollector}.
+     * Wraps two {@link org.apache.lucene.search.CollectorManager}s: one required for top docs collection, and another one optional for
+     * aggs collection. Applies terminate_after consistently across the different collectors by sharing an atomic counter of collected docs.
+     */
+    static class CollectorManager implements org.apache.lucene.search.CollectorManager<QueryPhaseCollector, Void> {
+        private final Weight postFilterWeight;
+        private final TerminateAfterChecker terminateAfterChecker;
+        private final Float minScore;
+        private final org.apache.lucene.search.CollectorManager<? extends Collector, Void> topDocsCollectorManager;
+        private final org.apache.lucene.search.CollectorManager<? extends Collector, Void> aggsCollectorManager;
+
+        private boolean terminatedAfter;
+
+        CollectorManager(
+            org.apache.lucene.search.CollectorManager<? extends Collector, Void> topDocsCollectorManager,
+            Weight postFilterWeight,
+            TerminateAfterChecker terminateAfterChecker,
+            org.apache.lucene.search.CollectorManager<? extends Collector, Void> aggsCollectorManager,
+            Float minScore
+        ) {
+            this.topDocsCollectorManager = topDocsCollectorManager;
+            this.postFilterWeight = postFilterWeight;
+            this.terminateAfterChecker = terminateAfterChecker;
+            this.aggsCollectorManager = aggsCollectorManager;
+            this.minScore = minScore;
+        }
+
+        @Override
+        public QueryPhaseCollector newCollector() throws IOException {
+            Collector aggsCollector = aggsCollectorManager == null ? null : aggsCollectorManager.newCollector();
+            return new QueryPhaseCollector(
+                topDocsCollectorManager.newCollector(),
+                postFilterWeight,
+                terminateAfterChecker,
+                aggsCollector,
+                minScore
+            );
+        }
+
+        @Override
+        public Void reduce(Collection<QueryPhaseCollector> collectors) throws IOException {
+            List<Collector> topDocsCollectors = new ArrayList<>();
+            List<Collector> aggsCollectors = new ArrayList<>();
+            for (QueryPhaseCollector collector : collectors) {
+                topDocsCollectors.add(collector.topDocsCollector);
+                aggsCollectors.add(collector.aggsCollector);
+                if (collector.isTerminatedAfter()) {
+                    terminatedAfter = true;
+                }
+            }
+            @SuppressWarnings("unchecked")
+            org.apache.lucene.search.CollectorManager<Collector, Void> topDocsManager = (org.apache.lucene.search.CollectorManager<
+                Collector,
+                Void>) topDocsCollectorManager;
+            topDocsManager.reduce(topDocsCollectors);
+            if (aggsCollectorManager != null) {
+                @SuppressWarnings("unchecked")
+                org.apache.lucene.search.CollectorManager<Collector, Void> aggsManager = (org.apache.lucene.search.CollectorManager<
+                    Collector,
+                    Void>) aggsCollectorManager;
+                aggsManager.reduce(aggsCollectors);
+            }
+            return null;
+        }
+
+        boolean isTerminatedAfter() {
+            return terminatedAfter;
+        }
+    }
 }
diff --git a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseCollectorTests.java b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseCollectorTests.java
