Made QueryCollectorContext methods create and postProcess public

Signed-off-by: Martin Gaievski <[email protected]>
opensearch-project · Jan 5, 2024 · 04bce0e · 04bce0e
1 parent ab0f70e
commit 04bce0e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 13 deletions.
diff --git 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
@@ -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
@@ -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