From 021f593be12668e058368e9f8033bf8981590994 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, createManager and
 postProcess public

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

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f7aa083eb2bcb..78b22918a52be 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 ([#11778](https://github.com/opensearch-project/OpenSearch/pull/11778))
 
 ### 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..46ad02428cad1 100644
--- a/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java
+++ b/server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java
@@ -87,9 +87,9 @@ 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;
+    public abstract CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException;
 
     /**
      * Wraps this collector with a profiler
@@ -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,12 +149,12 @@ 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);
             }
 
             @Override
-            CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
+            public CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
                 return new MinimumCollectorManager(in, minScore);
             }
         };
@@ -166,13 +166,13 @@ 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);
             }
 
             @Override
-            CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
+            public CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
                 final Weight filterWeight = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f);
                 return new FilteredCollectorManager(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) {
@@ -244,7 +244,7 @@ protected InternalProfileCollectorManager createWithProfiler(InternalProfileColl
             }
 
             @Override
-            CollectorManager<? extends Collector, ReduceableSearchResult> createManager(
+            public CollectorManager<? extends Collector, ReduceableSearchResult> createManager(
                 CollectorManager<? extends Collector, ReduceableSearchResult> in
             ) throws IOException {
                 final List<CollectorManager<?, ReduceableSearchResult>> managers = new ArrayList<>();
@@ -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<>();
@@ -278,7 +278,7 @@ Collector create(Collector in) {
             }
 
             @Override
-            CollectorManager<? extends Collector, ReduceableSearchResult> createManager(
+            public CollectorManager<? extends Collector, ReduceableSearchResult> createManager(
                 CollectorManager<? extends Collector, ReduceableSearchResult> in
             ) throws IOException {
                 return new EarlyTerminatingCollectorManager<>(in, numHits, true);
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..93e9e4e86bfa2 100644
--- a/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java
+++ b/server/src/main/java/org/opensearch/search/query/TopDocsCollectorContext.java
@@ -170,7 +170,7 @@ private EmptyTopDocsCollectorContext(
         }
 
         @Override
-        CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
+        public CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
             assert in == null;
 
             CollectorManager<?, ReduceableSearchResult> manager = null;
@@ -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,19 +267,19 @@ 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);
         }
 
         @Override
-        CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
+        public CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
             return new CollectorManager<Collector, ReduceableSearchResult>() {
                 @Override
                 public Collector newCollector() throws IOException {
@@ -543,7 +543,7 @@ public ReduceableSearchResult reduce(Collection<Collector> collectors) throws IO
         }
 
         @Override
-        CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
+        public CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
             assert in == null;
             return new SimpleTopDocsCollectorManager();
         }
@@ -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