From 04bce0e21b0cbb70ba0b6864e7bbd99c30c91930 Mon Sep 17 00:00:00 2001
From: Martin Gaievski <gaievski@amazon.com>
Date: Fri, 5 Jan 2024 12:09:08 -0800
Subject: [PATCH] Made QueryCollectorContext methods create and postProcess
 public

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
---
 CHANGELOG.md                                       |  1 +
 .../search/query/QueryCollectorContext.java        | 12 ++++++------
 .../search/query/TopDocsCollectorContext.java      | 14 +++++++-------
 3 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f7aa083eb2bcb..922827a29b538 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -184,6 +184,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Automatically add scheme to discovery.ec2.endpoint ([#11512](https://github.com/opensearch-project/OpenSearch/pull/11512))
 - Restore support for Java 8 for RestClient ([#11562](https://github.com/opensearch-project/OpenSearch/pull/11562))
 - Add deleted doc count in _cat/shards ([#11678](https://github.com/opensearch-project/OpenSearch/pull/11678))
+- Refactor QueryCollectorContext to improve extensibility ([xxx](link))
 
 ### Deprecated
 
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 91762bee2ac08..7ee49269abe42 100644
--- a/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java
+++ b/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java
@@ -87,7 +87,7 @@ public ScoreMode scoreMode() {
      * Creates a collector that delegates documents to the provided <code>in</code> collector.
      * @param in The delegate collector
      */
-    abstract Collector create(Collector in) throws IOException;
+    public abstract Collector create(Collector in) throws IOException;
 
     abstract CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException;
 
@@ -116,7 +116,7 @@ protected InternalProfileCollectorManager createWithProfiler(InternalProfileColl
      *
      * @param result The query search result to populate
      */
-    void postProcess(QuerySearchResult result) throws IOException {}
+    public void postProcess(QuerySearchResult result) throws IOException {}
 
     /**
      * Creates the collector tree from the provided <code>collectors</code>
@@ -149,7 +149,7 @@ static InternalProfileCollector createQueryCollectorWithProfiler(List<QueryColle
     static QueryCollectorContext createMinScoreCollectorContext(float minScore) {
         return new QueryCollectorContext(REASON_SEARCH_MIN_SCORE) {
             @Override
-            Collector create(Collector in) {
+            public Collector create(Collector in) {
                 return new MinimumScoreCollector(in, minScore);
             }
 
@@ -166,7 +166,7 @@ CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, Re
     static QueryCollectorContext createFilteredCollectorContext(IndexSearcher searcher, Query query) {
         return new QueryCollectorContext(REASON_SEARCH_POST_FILTER) {
             @Override
-            Collector create(Collector in) throws IOException {
+            public Collector create(Collector in) throws IOException {
                 final Weight filterWeight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f);
                 return new FilteredCollector(in, filterWeight);
             }
@@ -187,7 +187,7 @@ static QueryCollectorContext createMultiCollectorContext(
     ) {
         return new QueryCollectorContext(REASON_SEARCH_MULTI) {
             @Override
-            Collector create(Collector in) throws IOException {
+            public Collector create(Collector in) throws IOException {
                 List<Collector> subCollectors = new ArrayList<>();
                 subCollectors.add(in);
                 for (CollectorManager<? extends Collector, ReduceableSearchResult> manager : subs) {
@@ -267,7 +267,7 @@ static QueryCollectorContext createEarlyTerminationCollectorContext(int numHits)
              * can terminate the collection independently of the provided <code>in</code> {@link Collector}.
              */
             @Override
-            Collector create(Collector in) {
+            public Collector create(Collector in) {
                 assert collector == null;
 
                 List<Collector> subCollectors = new ArrayList<>();
diff --git a/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java b/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java
index 65d3948c8401e..6d61f7ad76b41 100644
--- a/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java
+++ b/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java
@@ -205,13 +205,13 @@ CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, Re
         }
 
         @Override
-        Collector create(Collector in) {
+        public Collector create(Collector in) {
             assert in == null;
             return collector;
         }
 
         @Override
-        void postProcess(QuerySearchResult result) {
+        public void postProcess(QuerySearchResult result) {
             final TotalHits totalHitCount = hitCountSupplier.get();
             final TopDocs topDocs;
             if (sort != null) {
@@ -267,13 +267,13 @@ private CollapsingTopDocsCollectorContext(
         }
 
         @Override
-        Collector create(Collector in) throws IOException {
+        public Collector create(Collector in) throws IOException {
             assert in == null;
             return collector;
         }
 
         @Override
-        void postProcess(QuerySearchResult result) throws IOException {
+        public void postProcess(QuerySearchResult result) throws IOException {
             final CollapseTopFieldDocs topDocs = topDocsCollector.getTopDocs();
             result.topDocs(new TopDocsAndMaxScore(topDocs, maxScoreSupplier.get()), sortFmt);
         }
@@ -556,7 +556,7 @@ protected ReduceableSearchResult reduceWith(final TopDocs topDocs, final float m
         }
 
         @Override
-        Collector create(Collector in) {
+        public Collector create(Collector in) {
             assert in == null;
             return collector;
         }
@@ -619,7 +619,7 @@ TopDocsAndMaxScore newTopDocs() {
         }
 
         @Override
-        void postProcess(QuerySearchResult result) throws IOException {
+        public void postProcess(QuerySearchResult result) throws IOException {
             final TopDocsAndMaxScore topDocs = newTopDocs();
             result.topDocs(topDocs, sortAndFormats == null ? null : sortAndFormats.formats);
         }
@@ -684,7 +684,7 @@ protected ReduceableSearchResult reduceWith(final TopDocs topDocs, final float m
         }
 
         @Override
-        void postProcess(QuerySearchResult result) throws IOException {
+        public void postProcess(QuerySearchResult result) throws IOException {
             final TopDocsAndMaxScore topDocs = newTopDocs();
             if (scrollContext.totalHits == null) {
                 // first round