diff --git a/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java b/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java index 08b048cf682bb..5d1f503d2ac2f 100644 --- a/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java +++ b/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java @@ -51,9 +51,14 @@ import java.io.IOException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.Collectors; import static org.opensearch.search.profile.query.CollectorResult.REASON_SEARCH_MIN_SCORE; import static org.opensearch.search.profile.query.CollectorResult.REASON_SEARCH_MULTI; @@ -67,38 +72,69 @@ */ @PublicApi(since = "1.0.0") public abstract class QueryCollectorContext { - private static final Collector EMPTY_COLLECTOR = new SimpleCollector() { - @Override - public void collect(int doc) {} - @Override - public ScoreMode scoreMode() { - return ScoreMode.COMPLETE_NO_SCORES; - } - }; + private static Collector createEmptyCollector(ScoreMode scoreMode) { + return new SimpleCollector() { + @Override + public void collect(int doc) {} - public static final QueryCollectorContext EMPTY_CONTEXT = new QueryCollectorContext("empty") { + @Override + public ScoreMode scoreMode() { + return scoreMode; + } + }; + } - @Override - Collector create(Collector in) throws IOException { - return EMPTY_COLLECTOR; - } + private static final ReduceableSearchResult EMPTY_RESULT = result -> {}; - @Override - CollectorManager createManager(CollectorManager in) throws IOException { - return new CollectorManager() { - @Override - public Collector newCollector() throws IOException { - return EMPTY_COLLECTOR; - } + private static QueryCollectorContext createEmptyContext(String name, Collector collector) { + return new QueryCollectorContext(name) { + @Override + Collector create(Collector in) { + return collector; + } - @Override - public ReduceableSearchResult reduce(Collection collectors) throws IOException { - return result -> {}; - } - }; - } - }; + @Override + CollectorManager createManager(CollectorManager in) { + return new CollectorManager<>() { + @Override + public Collector newCollector() { + return collector; + } + + @Override + public ReduceableSearchResult reduce(Collection collectors) { + return EMPTY_RESULT; + } + }; + + } + }; + } + + private static final Map CONTEXTS = Arrays.stream(ScoreMode.values()) + .collect( + Collectors.toMap( + Function.identity(), + scoreMode -> createEmptyContext(formatContextName(scoreMode), createEmptyCollector(scoreMode)) + ) + ); + + private static String formatContextName(ScoreMode scoreMode) { + return String.format(Locale.ROOT, "empty_with_score_mode_%s", scoreMode.toString().toLowerCase(Locale.ROOT)); + } + + private static final Collector EMPTY_COLLECTOR = createEmptyCollector(ScoreMode.COMPLETE_NO_SCORES); + public static final QueryCollectorContext EMPTY_CONTEXT = CONTEXTS.get(ScoreMode.COMPLETE_NO_SCORES); + + /** + * Returns the {@link QueryCollectorContext} for the provided {@link ScoreMode} + * + * @param scoreMode The score mode to get the context for + */ + public static QueryCollectorContext getContextForScoreMode(final ScoreMode scoreMode) { + return CONTEXTS.getOrDefault(scoreMode, EMPTY_CONTEXT); + } private String profilerName;