diff --git a/CHANGELOG.md b/CHANGELOG.md index 5991d9593e259..cdb18fd026c9a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,6 +47,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Support retrieving doc values of unsigned long field ([#16543](https://github.com/opensearch-project/OpenSearch/pull/16543)) - Fix rollover alias supports restored searchable snapshot index([#16483](https://github.com/opensearch-project/OpenSearch/pull/16483)) - Fix permissions error on scripted query against remote snapshot ([#16544](https://github.com/opensearch-project/OpenSearch/pull/16544)) +- Add method to create empty query collector context with customizable score mode ([#16660](https://github.com/opensearch-project/OpenSearch/pull/16660)) ### Security 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;