diff --git a/CHANGELOG.md b/CHANGELOG.md
index fed2bb54783a7..51f904fc69f41 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -191,6 +191,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))
- Capture information for additional query types and aggregation types ([#11582](https://github.com/opensearch-project/OpenSearch/pull/11582))
### 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 in 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 collectors
@@ -149,12 +149,12 @@ static InternalProfileCollector createQueryCollectorWithProfiler(List 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 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> managers = new ArrayList<>();
@@ -267,7 +267,7 @@ static QueryCollectorContext createEarlyTerminationCollectorContext(int numHits)
* can terminate the collection independently of the provided in {@link Collector}.
*/
@Override
- Collector create(Collector in) {
+ public Collector create(Collector in) {
assert collector == null;
List 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() {
@Override
public Collector newCollector() throws IOException {
@@ -543,7 +543,7 @@ public ReduceableSearchResult reduce(Collection 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