From 484d820c677a3dcdf29e853608829d15b3171086 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Fri, 15 Nov 2024 12:58:35 -0800 Subject: [PATCH 1/4] Add method for getting empty query collector context with particular score mode Signed-off-by: Martin Gaievski --- CHANGELOG-3.0.md | 1 + .../search/query/QueryCollectorContext.java | 90 +++++++++++++------ 2 files changed, 64 insertions(+), 27 deletions(-) diff --git a/CHANGELOG-3.0.md b/CHANGELOG-3.0.md index 48d978bede420..f82a93c59d95b 100644 --- a/CHANGELOG-3.0.md +++ b/CHANGELOG-3.0.md @@ -49,6 +49,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fix 'org.apache.hc.core5.http.ParseException: Invalid protocol version' under JDK 16+ ([#4827](https://github.com/opensearch-project/OpenSearch/pull/4827)) - Fix compression support for h2c protocol ([#4944](https://github.com/opensearch-project/OpenSearch/pull/4944)) - Don't over-allocate in HeapBufferedAsyncEntityConsumer in order to consume the response ([#9993](https://github.com/opensearch-project/OpenSearch/pull/9993)) +- 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; From 7f20f72cdcba5854343f4f361c93b2d2311f1bfc Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Fri, 15 Nov 2024 16:22:04 -0800 Subject: [PATCH 2/4] Added unit tests Signed-off-by: Martin Gaievski --- .../search/query/QueryPhaseTests.java | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java index 84057ab1a1b15..f5b5dc346068e 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java @@ -1197,6 +1197,36 @@ private void createTimeoutCheckerThenWaitThenRun( verifyNoMoreInteractions(mockedSearchContext); } + public void testQueryCollectorContextScoreModes() throws Exception { + // Test default empty context + QueryCollectorContext emptyContext = QueryCollectorContext.EMPTY_CONTEXT; + assertEquals(org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES, emptyContext.create(null).scoreMode()); + + // Test getting context for different score modes + org.apache.lucene.search.ScoreMode[] scoreModes = org.apache.lucene.search.ScoreMode.values(); + for (org.apache.lucene.search.ScoreMode scoreMode : scoreModes) { + QueryCollectorContext context = QueryCollectorContext.getContextForScoreMode(scoreMode); + assertEquals(scoreMode, context.create(null).scoreMode()); + } + + // Test that invalid score mode returns empty context + QueryCollectorContext defaultContext = QueryCollectorContext.getContextForScoreMode(null); + assertEquals(QueryCollectorContext.EMPTY_CONTEXT, defaultContext); + assertEquals(org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES, defaultContext.create(null).scoreMode()); + } + + public void testQueryCollectorContextTopScores() throws Exception { + QueryCollectorContext topScoresContext = QueryCollectorContext.getContextForScoreMode( + org.apache.lucene.search.ScoreMode.TOP_SCORES + ); + + // Verify score mode + assertEquals(org.apache.lucene.search.ScoreMode.TOP_SCORES, topScoresContext.create(null).scoreMode()); + + // Verify it's a different instance than empty context + assertNotEquals(QueryCollectorContext.EMPTY_CONTEXT, topScoresContext); + } + private static class TestSearchContextWithRewriteAndCancellation extends TestSearchContext { private TestSearchContextWithRewriteAndCancellation( From 65bcf14bd970923063416f959ec9e0c2caefbd25 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Mon, 18 Nov 2024 11:33:29 -0800 Subject: [PATCH 3/4] Made empty collector with top_scores mode a public constant Signed-off-by: Martin Gaievski --- CHANGELOG-3.0.md | 2 +- .../search/query/QueryCollectorContext.java | 30 +++++-------------- .../search/query/QueryPhaseTests.java | 24 ++------------- 3 files changed, 11 insertions(+), 45 deletions(-) diff --git a/CHANGELOG-3.0.md b/CHANGELOG-3.0.md index f82a93c59d95b..4319b31cb8766 100644 --- a/CHANGELOG-3.0.md +++ b/CHANGELOG-3.0.md @@ -49,7 +49,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Fix 'org.apache.hc.core5.http.ParseException: Invalid protocol version' under JDK 16+ ([#4827](https://github.com/opensearch-project/OpenSearch/pull/4827)) - Fix compression support for h2c protocol ([#4944](https://github.com/opensearch-project/OpenSearch/pull/4944)) - Don't over-allocate in HeapBufferedAsyncEntityConsumer in order to consume the response ([#9993](https://github.com/opensearch-project/OpenSearch/pull/9993)) -- Add method to create empty query collector context with customizable score mode ([#16660](https://github.com/opensearch-project/OpenSearch/pull/16660)) +- Add empty query collector context constant with TOP_SCORES 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 5d1f503d2ac2f..dea9668a3706c 100644 --- a/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java +++ b/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java @@ -51,14 +51,10 @@ 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; @@ -112,29 +108,19 @@ public ReduceableSearchResult reduce(Collection collectors) { }; } - 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); - } + public static final QueryCollectorContext EMPTY_CONTEXT = createEmptyContext( + formatContextName(ScoreMode.COMPLETE_NO_SCORES), + createEmptyCollector(ScoreMode.COMPLETE_NO_SCORES) + ); + public static final QueryCollectorContext EMPTY_CONTEXT_TOP_SCORES_SCORE_MODE = createEmptyContext( + formatContextName(ScoreMode.TOP_SCORES), + createEmptyCollector(ScoreMode.TOP_SCORES) + ); private String profilerName; diff --git a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java index f5b5dc346068e..846e7b9675dec 100644 --- a/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java +++ b/server/src/test/java/org/opensearch/search/query/QueryPhaseTests.java @@ -1197,28 +1197,8 @@ private void createTimeoutCheckerThenWaitThenRun( verifyNoMoreInteractions(mockedSearchContext); } - public void testQueryCollectorContextScoreModes() throws Exception { - // Test default empty context - QueryCollectorContext emptyContext = QueryCollectorContext.EMPTY_CONTEXT; - assertEquals(org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES, emptyContext.create(null).scoreMode()); - - // Test getting context for different score modes - org.apache.lucene.search.ScoreMode[] scoreModes = org.apache.lucene.search.ScoreMode.values(); - for (org.apache.lucene.search.ScoreMode scoreMode : scoreModes) { - QueryCollectorContext context = QueryCollectorContext.getContextForScoreMode(scoreMode); - assertEquals(scoreMode, context.create(null).scoreMode()); - } - - // Test that invalid score mode returns empty context - QueryCollectorContext defaultContext = QueryCollectorContext.getContextForScoreMode(null); - assertEquals(QueryCollectorContext.EMPTY_CONTEXT, defaultContext); - assertEquals(org.apache.lucene.search.ScoreMode.COMPLETE_NO_SCORES, defaultContext.create(null).scoreMode()); - } - - public void testQueryCollectorContextTopScores() throws Exception { - QueryCollectorContext topScoresContext = QueryCollectorContext.getContextForScoreMode( - org.apache.lucene.search.ScoreMode.TOP_SCORES - ); + public void testEmptyQueryCollectorContextAndDifferentScoreModes() throws Exception { + QueryCollectorContext topScoresContext = QueryCollectorContext.EMPTY_CONTEXT_TOP_SCORES_SCORE_MODE; // Verify score mode assertEquals(org.apache.lucene.search.ScoreMode.TOP_SCORES, topScoresContext.create(null).scoreMode()); From fd75a58808d03d0ebc136a9fdb087c7a151db645 Mon Sep 17 00:00:00 2001 From: Martin Gaievski Date: Mon, 18 Nov 2024 13:54:15 -0800 Subject: [PATCH 4/4] Assume scode mode from collector Signed-off-by: Martin Gaievski --- .../search/query/QueryCollectorContext.java | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) 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 dea9668a3706c..32087af75288d 100644 --- a/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java +++ b/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java @@ -83,7 +83,8 @@ public ScoreMode scoreMode() { private static final ReduceableSearchResult EMPTY_RESULT = result -> {}; - private static QueryCollectorContext createEmptyContext(String name, Collector collector) { + private static QueryCollectorContext createEmptyContext(Collector collector) { + String name = String.format(Locale.ROOT, "empty_with_score_mode_%s", collector.scoreMode().toString().toLowerCase(Locale.ROOT)); return new QueryCollectorContext(name) { @Override Collector create(Collector in) { @@ -108,19 +109,11 @@ public ReduceableSearchResult reduce(Collection collectors) { }; } - 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 = createEmptyContext( - formatContextName(ScoreMode.COMPLETE_NO_SCORES), - createEmptyCollector(ScoreMode.COMPLETE_NO_SCORES) - ); - public static final QueryCollectorContext EMPTY_CONTEXT_TOP_SCORES_SCORE_MODE = createEmptyContext( - formatContextName(ScoreMode.TOP_SCORES), - createEmptyCollector(ScoreMode.TOP_SCORES) - ); + private static final Collector EMPTY_COLLECTOR_TOP_SCORES = createEmptyCollector(ScoreMode.TOP_SCORES); + + public static final QueryCollectorContext EMPTY_CONTEXT = createEmptyContext(EMPTY_COLLECTOR); + public static final QueryCollectorContext EMPTY_CONTEXT_TOP_SCORES_SCORE_MODE = createEmptyContext(EMPTY_COLLECTOR_TOP_SCORES); private String profilerName;