Add query argument in QueryCollectorContextSpecFactory (#19153)

* Add query in QueryCollectorContextSpecFactory

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

* Add javadoc

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

---------

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Author: vibrantvarun

server/src/main/java/org/opensearch/search/query/QueryCollectorContextSpecFactory.java

@@ -8,6 +8,7 @@
 
 package org.opensearch.search.query;
 
+import org.apache.lucene.search.Query;
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.search.internal.SearchContext;
 
@@ -21,12 +22,14 @@
 public interface QueryCollectorContextSpecFactory {
     /**
      * @param searchContext context needed to create collector context spec
+     * @param query required to create collector context spec
      * @param queryCollectorArguments arguments to create collector context spec
      * @return QueryCollectorContextSpec
      * @throws IOException
      */
     Optional<QueryCollectorContextSpec> createQueryCollectorContextSpec(
         SearchContext searchContext,
+        Query query,
         QueryCollectorArguments queryCollectorArguments
     ) throws IOException;
 }

server/src/main/java/org/opensearch/search/query/QueryCollectorContextSpecRegistry.java

@@ -8,6 +8,7 @@
 
 package org.opensearch.search.query;
 
+import org.apache.lucene.search.Query;
 import org.opensearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -43,18 +44,24 @@ public static void registerFactory(QueryCollectorContextSpecFactory factory) {
     /**
      * Get collector context spec
      * @param searchContext search context
+     * @param query required to create collectorContext spec
      * @param queryCollectorArguments query collector arguments
      * @return collector context spec
      * @throws IOException
      */
     public static Optional<QueryCollectorContextSpec> getQueryCollectorContextSpec(
         final SearchContext searchContext,
+        final Query query,
         final QueryCollectorArguments queryCollectorArguments
     ) throws IOException {
         Iterator<QueryCollectorContextSpecFactory> iterator = registry.iterator();
         while (iterator.hasNext()) {
             QueryCollectorContextSpecFactory factory = iterator.next();
-            Optional<QueryCollectorContextSpec> spec = factory.createQueryCollectorContextSpec(searchContext, queryCollectorArguments);
+            Optional<QueryCollectorContextSpec> spec = factory.createQueryCollectorContextSpec(
+                searchContext,
+                query,
+                queryCollectorArguments
+            );
             if (spec.isEmpty() == false) {
                 return spec;
             }

server/src/main/java/org/opensearch/search/query/QueryPhase.java

@@ -446,14 +446,16 @@ protected boolean searchWithCollector(
             boolean hasFilterCollector,
             boolean hasTimeout
         ) throws IOException {
-            QueryCollectorContext queryCollectorContext = getQueryCollectorContext(searchContext, hasFilterCollector);
+            QueryCollectorContext queryCollectorContext = getQueryCollectorContext(searchContext, query, hasFilterCollector);
             return searchWithCollector(searchContext, searcher, query, collectors, queryCollectorContext, hasFilterCollector, hasTimeout);
         }
 
-        private QueryCollectorContext getQueryCollectorContext(SearchContext searchContext, boolean hasFilterCollector) throws IOException {
+        private QueryCollectorContext getQueryCollectorContext(SearchContext searchContext, Query query, boolean hasFilterCollector)
+            throws IOException {
             // create the top docs collector last when the other collectors are known
             final Optional<QueryCollectorContext> queryCollectorContextOpt = QueryCollectorContextSpecRegistry.getQueryCollectorContextSpec(
                 searchContext,
+                query,
                 new QueryCollectorArguments.Builder().hasFilterCollector(hasFilterCollector).build()
             ).map(queryCollectorContextSpec -> new QueryCollectorContext(queryCollectorContextSpec.getContextName()) {
                 @Override

server/src/test/java/org/opensearch/search/query/QueryCollectorContextSpecRegistryTests.java

@@ -49,11 +49,14 @@ public void testGetQueryCollectorContextSpec_WithValidSpec() throws IOException
         QueryCollectorArguments mockArguments = new QueryCollectorArguments.Builder().build();
         // Given
         QueryCollectorContextSpecRegistry.registerFactory(mockFactory1);
-        when(mockFactory1.createQueryCollectorContextSpec(mockSearchContext, mockArguments)).thenReturn(Optional.of(mockSpec));
+        when(mockFactory1.createQueryCollectorContextSpec(mockSearchContext, mockSearchContext.query(), mockArguments)).thenReturn(
+            Optional.of(mockSpec)
+        );
 
         // When
         Optional<QueryCollectorContextSpec> result = QueryCollectorContextSpecRegistry.getQueryCollectorContextSpec(
             mockSearchContext,
+            mockSearchContext.query(),
             mockArguments
         );
 
@@ -69,6 +72,7 @@ public void testGetQueryCollectorContextSpec_NoFactories() throws IOException {
         // When
         Optional<QueryCollectorContextSpec> result = QueryCollectorContextSpecRegistry.getQueryCollectorContextSpec(
             mockSearchContext,
+            mockSearchContext.query(),
             mockArguments
         );