From 990ab3954948ea760728d5414d295cecec19cf9f Mon Sep 17 00:00:00 2001
From: Martin Gaievski <gaievski@amazon.com>
Date: Mon, 18 Nov 2024 11:33:29 -0800
Subject: [PATCH] Made empty collector with top_scores mode a public constant

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
---
 .../search/query/QueryCollectorContext.java   | 30 +++++--------------
 .../search/query/QueryPhaseTests.java         | 24 ++-------------
 2 files changed, 10 insertions(+), 44 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 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<Collector> collectors) {
         };
     }
 
-    private static final Map<ScoreMode, QueryCollectorContext> 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());