From 54c87856dd8c40b3c6d4af44a559b58124db9fb6 Mon Sep 17 00:00:00 2001
From: Martin Gaievski <gaievski@amazon.com>
Date: Fri, 15 Nov 2024 12:58:35 -0800
Subject: [PATCH] Add method for getting empty query collector context with
 particular score mode

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
---
 .../search/query/QueryCollectorContext.java   | 90 +++++++++++++------
 1 file changed, 63 insertions(+), 27 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 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<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
-            return new CollectorManager<Collector, ReduceableSearchResult>() {
-                @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<Collector> collectors) throws IOException {
-                    return result -> {};
-                }
-            };
-        }
-    };
+            @Override
+            CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) {
+                return new CollectorManager<>() {
+                    @Override
+                    public Collector newCollector() {
+                        return collector;
+                    }
+
+                    @Override
+                    public ReduceableSearchResult reduce(Collection<Collector> collectors) {
+                        return EMPTY_RESULT;
+                    }
+                };
+
+            }
+        };
+    }
+
+    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);
+    }
 
     private String profilerName;