Rename TopDocsCollectorContext to TopDocsCollectorFactory (#95435)

javanna · web-flow · commit 8e15a1a7ada1 · 2023-04-21T10:00:27.000+02:00
This is to complete the removal of the query collector context abstraction implemented with #95383. The remaining TopDocsCollectorContext is more of a factory than a context object. This commit renames the class and all of its subclasses. It also adds javadocs to its methods to clarify the contract around them.
diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java
@@ -53,7 +53,7 @@
 import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_MULTI;
 import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_POST_FILTER;
 import static org.elasticsearch.search.profile.query.CollectorResult.REASON_SEARCH_TERMINATE_AFTER_COUNT;
-import static org.elasticsearch.search.query.TopDocsCollectorContext.createTopDocsCollectorContext;
+import static org.elasticsearch.search.query.TopDocsCollectorFactory.createTopDocsCollectorFactory;
 
 /**
  * Query phase of a search request, used to run the query and get back from each shard information about the matching documents
@@ -132,7 +132,7 @@ static void executeInternal(SearchContext searchContext) throws QueryPhaseExecut
             }
 
             // create the top docs collector last when the other collectors are known
-            final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(
+            final TopDocsCollectorFactory topDocsFactory = createTopDocsCollectorFactory(
                 searchContext,
                 searchContext.parsedPostFilter() != null || searchContext.minimumScore() != null
             );
diff --git a/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorFactory.java b/server/src/main/java/org/elasticsearch/search/query/TopDocsCollectorFactory.java
@@ -65,21 +65,29 @@
 
 /**
  * Creates and holds the main {@link Collector} that will be used to search and collect top hits.
+ * Once the returned collector has been used to collect hits, allows to enrich the collected top docs to be set to the search context.
  */
-abstract class TopDocsCollectorContext {
+abstract class TopDocsCollectorFactory {
     final String profilerName;
     final DocValueFormat[] sortValueFormats;
 
-    TopDocsCollectorContext(String profilerName, DocValueFormat[] sortValueFormats) {
+    TopDocsCollectorFactory(String profilerName, DocValueFormat[] sortValueFormats) {
         this.profilerName = profilerName;
         this.sortValueFormats = sortValueFormats;
     }
 
+    /**
+     * Returns the collector used to collect top hits, created depending on the incoming request options
+     */
     abstract Collector collector();
 
+    /**
+     * Returns the collected top docs to be set to the {@link QuerySearchResult} within the search context.
+     * To be called after collection, to enrich the top docs and wrap them with our {@link TopDocsAndMaxScore}.
+     */
     abstract TopDocsAndMaxScore topDocsAndMaxScore() throws IOException;
 
-    static class EmptyTopDocsCollectorContext extends TopDocsCollectorContext {
+    static class EmptyTopDocsCollectorFactory extends TopDocsCollectorFactory {
         private final Sort sort;
         private final Collector collector;
         private final Supplier<TotalHits> hitCountSupplier;
@@ -89,7 +97,7 @@ static class EmptyTopDocsCollectorContext extends TopDocsCollectorContext {
          * @param sortAndFormats The sort clause if provided
          * @param trackTotalHitsUpTo The threshold up to which total hit count needs to be tracked
          */
-        private EmptyTopDocsCollectorContext(@Nullable SortAndFormats sortAndFormats, int trackTotalHitsUpTo) {
+        private EmptyTopDocsCollectorFactory(@Nullable SortAndFormats sortAndFormats, int trackTotalHitsUpTo) {
             super(REASON_SEARCH_COUNT, null);
             this.sort = sortAndFormats == null ? null : sortAndFormats.sort;
             if (trackTotalHitsUpTo == SearchContext.TRACK_TOTAL_HITS_DISABLED) {
@@ -130,7 +138,7 @@ TopDocsAndMaxScore topDocsAndMaxScore() {
         }
     }
 
-    static class CollapsingTopDocsCollectorContext extends TopDocsCollectorContext {
+    static class CollapsingTopDocsCollectorFactory extends TopDocsCollectorFactory {
         private final SinglePassGroupingCollector<?> topDocsCollector;
         private final Supplier<Float> maxScoreSupplier;
 
@@ -141,7 +149,7 @@ static class CollapsingTopDocsCollectorContext extends TopDocsCollectorContext {
          * @param numHits The number of collapsed top hits to retrieve.
          * @param trackMaxScore True if max score should be tracked
          */
-        private CollapsingTopDocsCollectorContext(
+        private CollapsingTopDocsCollectorFactory(
             CollapseContext collapseContext,
             @Nullable SortAndFormats sortAndFormats,
             int numHits,
@@ -175,7 +183,7 @@ TopDocsAndMaxScore topDocsAndMaxScore() throws IOException {
         }
     }
 
-    static class SimpleTopDocsCollectorContext extends TopDocsCollectorContext {
+    static class SimpleTopDocsCollectorFactory extends TopDocsCollectorFactory {
 
         private static TopDocsCollector<?> createCollector(
             @Nullable SortAndFormats sortAndFormats,
@@ -207,7 +215,7 @@ private static TopDocsCollector<?> createCollector(
          * @param trackTotalHitsUpTo Threshold up to which total hit count should be tracked
          * @param hasFilterCollector True if the collector chain contains at least one collector that can filter documents out
          */
-        private SimpleTopDocsCollectorContext(
+        private SimpleTopDocsCollectorFactory(
             IndexReader reader,
             Query query,
             @Nullable SortAndFormats sortAndFormats,
@@ -289,11 +297,11 @@ TopDocsAndMaxScore topDocsAndMaxScore() {
         }
     }
 
-    static class ScrollingTopDocsCollectorContext extends SimpleTopDocsCollectorContext {
+    static class ScrollingTopDocsCollectorFactory extends SimpleTopDocsCollectorFactory {
         private final ScrollContext scrollContext;
         private final int numberOfShards;
 
-        private ScrollingTopDocsCollectorContext(
+        private ScrollingTopDocsCollectorFactory(
             IndexReader reader,
             Query query,
             ScrollContext scrollContext,
@@ -403,18 +411,18 @@ static int shortcutTotalHitCount(IndexReader reader, Query query) throws IOExcep
     }
 
     /**
-     * Creates a {@link TopDocsCollectorContext} from the provided <code>searchContext</code>.
+     * Creates a {@link TopDocsCollectorFactory} from the provided <code>searchContext</code>.
      * @param hasFilterCollector True if the collector chain contains at least one collector that can filters document.
      */
-    static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searchContext, boolean hasFilterCollector)
+    static TopDocsCollectorFactory createTopDocsCollectorFactory(SearchContext searchContext, boolean hasFilterCollector)
         throws IOException {
         final IndexReader reader = searchContext.searcher().getIndexReader();
         final Query query = searchContext.rewrittenQuery();
         // top collectors don't like a size of 0
         final int totalNumDocs = Math.max(1, reader.numDocs());
         if (searchContext.size() == 0) {
             // no matter what the value of from is
-            return new EmptyTopDocsCollectorContext(searchContext.sort(), searchContext.trackTotalHitsUpTo());
+            return new EmptyTopDocsCollectorFactory(searchContext.sort(), searchContext.trackTotalHitsUpTo());
         } else if (searchContext.scrollContext() != null) {
             // we can disable the tracking of total hits after the initial scroll query
             // since the total hits is preserved in the scroll context.
@@ -423,7 +431,7 @@ static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searc
                 : SearchContext.TRACK_TOTAL_HITS_ACCURATE;
             // no matter what the value of from is
             int numDocs = Math.min(searchContext.size(), totalNumDocs);
-            return new ScrollingTopDocsCollectorContext(
+            return new ScrollingTopDocsCollectorFactory(
                 reader,
                 query,
                 searchContext.scrollContext(),
@@ -435,9 +443,9 @@ static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searc
                 hasFilterCollector
             );
         } else if (searchContext.collapse() != null) {
-            boolean trackScores = searchContext.sort() == null ? true : searchContext.trackScores();
+            boolean trackScores = searchContext.sort() == null || searchContext.trackScores();
             int numDocs = Math.min(searchContext.from() + searchContext.size(), totalNumDocs);
-            return new CollapsingTopDocsCollectorContext(
+            return new CollapsingTopDocsCollectorFactory(
                 searchContext.collapse(),
                 searchContext.sort(),
                 numDocs,
@@ -453,7 +461,7 @@ static TopDocsCollectorContext createTopDocsCollectorContext(SearchContext searc
                     numDocs = Math.max(numDocs, rescoreContext.getWindowSize());
                 }
             }
-            return new SimpleTopDocsCollectorContext(
+            return new SimpleTopDocsCollectorFactory(
                 reader,
                 query,
                 searchContext.sort(),
diff --git a/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/elasticsearch/search/query/QueryPhaseTests.java
@@ -83,7 +83,7 @@
 import java.util.List;
 import java.util.function.IntUnaryOperator;
 
-import static org.elasticsearch.search.query.TopDocsCollectorContext.hasInfMaxScore;
+import static org.elasticsearch.search.query.TopDocsCollectorFactory.hasInfMaxScore;
 import static org.hamcrest.Matchers.anyOf;
 import static org.hamcrest.Matchers.arrayWithSize;
 import static org.hamcrest.Matchers.equalTo;
@@ -679,15 +679,15 @@ public void testDisableTopScoreCollection() throws Exception {
         context.parsedQuery(new ParsedQuery(q));
         context.setSize(3);
         context.trackTotalHitsUpTo(3);
-        TopDocsCollectorContext topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false);
+        TopDocsCollectorFactory topDocsContext = TopDocsCollectorFactory.createTopDocsCollectorFactory(context, false);
         assertEquals(topDocsContext.collector().scoreMode(), org.apache.lucene.search.ScoreMode.COMPLETE);
         QueryPhase.executeInternal(context);
         assertEquals(5, context.queryResult().topDocs().topDocs.totalHits.value);
         assertEquals(context.queryResult().topDocs().topDocs.totalHits.relation, TotalHits.Relation.EQUAL_TO);
         assertThat(context.queryResult().topDocs().topDocs.scoreDocs.length, equalTo(3));
 
         context.sort(new SortAndFormats(new Sort(new SortField("other", SortField.Type.INT)), new DocValueFormat[] { DocValueFormat.RAW }));
-        topDocsContext = TopDocsCollectorContext.createTopDocsCollectorContext(context, false);
+        topDocsContext = TopDocsCollectorFactory.createTopDocsCollectorFactory(context, false);
         assertEquals(topDocsContext.collector().scoreMode(), org.apache.lucene.search.ScoreMode.TOP_DOCS);
         QueryPhase.executeInternal(context);
         assertEquals(5, context.queryResult().topDocs().topDocs.totalHits.value);
diff --git a/server/src/test/java/org/elasticsearch/search/query/TopDocsCollectorFactoryTests.java b/server/src/test/java/org/elasticsearch/search/query/TopDocsCollectorFactoryTests.java
@@ -25,7 +25,7 @@
 
 import java.io.IOException;
 
-public class TopDocsCollectorContextTests extends ESTestCase {
+public class TopDocsCollectorFactoryTests extends ESTestCase {
 
     public void testShortcutTotalHitCountTextField() throws IOException {
         try (Directory dir = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), dir)) {
@@ -39,7 +39,7 @@ public void testShortcutTotalHitCountTextField() throws IOException {
             iw.commit();
             try (IndexReader reader = iw.getReader()) {
                 final Query testQuery = new FieldExistsQuery("text");
-                int hitCount = TopDocsCollectorContext.shortcutTotalHitCount(reader, testQuery);
+                int hitCount = TopDocsCollectorFactory.shortcutTotalHitCount(reader, testQuery);
                 assertEquals(-1, hitCount);
             }
         }
@@ -59,7 +59,7 @@ public void testShortcutTotalHitCountStringField() throws IOException {
             iw.commit();
             try (IndexReader reader = iw.getReader()) {
                 final Query testQuery = new FieldExistsQuery("string");
-                int hitCount = TopDocsCollectorContext.shortcutTotalHitCount(reader, testQuery);
+                int hitCount = TopDocsCollectorFactory.shortcutTotalHitCount(reader, testQuery);
                 assertEquals(2, hitCount);
             }
         }
@@ -75,7 +75,7 @@ public void testShortcutTotalHitCountNumericField() throws IOException {
             iw.commit();
             try (IndexReader reader = iw.getReader()) {
                 final Query testQuery = new FieldExistsQuery("int");
-                int hitCount = TopDocsCollectorContext.shortcutTotalHitCount(reader, testQuery);
+                int hitCount = TopDocsCollectorFactory.shortcutTotalHitCount(reader, testQuery);
                 assertEquals(1, hitCount);
             }
         }

Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@`
`25`	`25`
`26`	`26`	`import java.io.IOException;`
`27`	`27`
`28`		`-public class TopDocsCollectorContextTests extends ESTestCase {`
	`28`	`+public class TopDocsCollectorFactoryTests extends ESTestCase {`
`29`	`29`
`30`	`30`	`public void testShortcutTotalHitCountTextField() throws IOException {`
`31`	`31`	`try (Directory dir = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), dir)) {`
`@@ -39,7 +39,7 @@ public void testShortcutTotalHitCountTextField() throws IOException {`
`39`	`39`	`iw.commit();`
`40`	`40`	`try (IndexReader reader = iw.getReader()) {`
`41`	`41`	`final Query testQuery = new FieldExistsQuery("text");`
`42`		`- int hitCount = TopDocsCollectorContext.shortcutTotalHitCount(reader, testQuery);`
	`42`	`+ int hitCount = TopDocsCollectorFactory.shortcutTotalHitCount(reader, testQuery);`
`43`	`43`	`assertEquals(-1, hitCount);`
`44`	`44`	`}`
`45`	`45`	`}`
`@@ -59,7 +59,7 @@ public void testShortcutTotalHitCountStringField() throws IOException {`
`59`	`59`	`iw.commit();`
`60`	`60`	`try (IndexReader reader = iw.getReader()) {`
`61`	`61`	`final Query testQuery = new FieldExistsQuery("string");`
`62`		`- int hitCount = TopDocsCollectorContext.shortcutTotalHitCount(reader, testQuery);`
	`62`	`+ int hitCount = TopDocsCollectorFactory.shortcutTotalHitCount(reader, testQuery);`
`63`	`63`	`assertEquals(2, hitCount);`
`64`	`64`	`}`
`65`	`65`	`}`
`@@ -75,7 +75,7 @@ public void testShortcutTotalHitCountNumericField() throws IOException {`
`75`	`75`	`iw.commit();`
`76`	`76`	`try (IndexReader reader = iw.getReader()) {`
`77`	`77`	`final Query testQuery = new FieldExistsQuery("int");`
`78`		`- int hitCount = TopDocsCollectorContext.shortcutTotalHitCount(reader, testQuery);`
	`78`	`+ int hitCount = TopDocsCollectorFactory.shortcutTotalHitCount(reader, testQuery);`
`79`	`79`	`assertEquals(1, hitCount);`
`80`	`80`	`}`
`81`	`81`	`}`