Factory design pattern for QueryCollectorContextSpec

Signed-off-by: Varun Jain <varunudr@amazon.com>

Author: vibrantvarun

server/src/main/java/org/opensearch/plugins/SearchPlugin.java

@@ -68,7 +68,7 @@
 import org.opensearch.search.deciders.ConcurrentSearchRequestDecider;
 import org.opensearch.search.fetch.FetchSubPhase;
 import org.opensearch.search.fetch.subphase.highlight.Highlighter;
-import org.opensearch.search.query.QueryCollectorContextFactory;
+import org.opensearch.search.query.QueryCollectorContextSpecFactory;
 import org.opensearch.search.query.QueryPhaseSearcher;
 import org.opensearch.search.rescore.Rescorer;
 import org.opensearch.search.rescore.RescorerBuilder;
@@ -228,7 +228,7 @@ default Optional<ExecutorServiceProvider> getIndexSearcherExecutorProvider() {
         return Optional.empty();
     }
 
-    default List<FactorySpec<?>> getCollectorContextFactory() {
+    default List<FactorySpec<?>> getCollectorContextSpec() {
         return emptyList();
     }
 
@@ -417,35 +417,35 @@ public QuerySpec(String name, Writeable.Reader<T> reader, QueryParser<T> parser)
       * @param <T>
      */
     class FactorySpec<T extends QueryBuilder> {
-        public String getName() {
-            return name;
-        }
+        // public String getName() {
+        // return name;
+        // }
 
-        public QueryCollectorContextFactory getQueryCollectorContextFactory() {
-            return queryCollectorContextFactory;
+        public QueryCollectorContextSpecFactory getQueryCollectorContextSpec() {
+            return queryCollectorContextSpecFactory;
         }
 
-        private final String name;
+        // private final String name;
+        //
+        // public Class<?> getClassType() {
+        // return classType;
+        // }
 
-        public Class<?> getClassType() {
-            return classType;
-        }
-
-        private final Class<?> classType;
-        private final QueryCollectorContextFactory queryCollectorContextFactory;
+        // private final Class<?> classType;
+        private final QueryCollectorContextSpecFactory queryCollectorContextSpecFactory;
 
         /**
          * Specification of custom {@link Query}.
          *
          * @param name the name by which this query might be parsed or deserialized. Make sure that the query builder returns this name for
          *        {@link NamedWriteable#getWriteableName()}.  It is an error if this name conflicts with another registered name, including
          *        names from other plugins.
-         * @param queryCollectorContextFactory Factory associated with the query builder
+         * @param queryCollectorContextSpecFactory Factory associated with the query builder
          */
-        public FactorySpec(String name, Class<?> classType, QueryCollectorContextFactory queryCollectorContextFactory) {
-            this.name = name;
-            this.classType = classType;
-            this.queryCollectorContextFactory = queryCollectorContextFactory;
+        public FactorySpec(String name, Class<?> classType, QueryCollectorContextSpecFactory queryCollectorContextSpecFactory) {
+            // this.name = name;
+            // this.classType = classType;
+            this.queryCollectorContextSpecFactory = queryCollectorContextSpecFactory;
         }
     }
 

server/src/main/java/org/opensearch/search/SearchModule.java

@@ -257,7 +257,7 @@
 import org.opensearch.search.fetch.subphase.highlight.Highlighter;
 import org.opensearch.search.fetch.subphase.highlight.PlainHighlighter;
 import org.opensearch.search.fetch.subphase.highlight.UnifiedHighlighter;
-import org.opensearch.search.query.QueryCollectorContextFactoryRegistery;
+import org.opensearch.search.query.QueryCollectorContextSpecRegistry;
 import org.opensearch.search.query.QueryPhase;
 import org.opensearch.search.query.QueryPhaseSearcher;
 import org.opensearch.search.query.QueryPhaseSearcherWrapper;
@@ -353,7 +353,7 @@ public SearchModule(Settings settings, List<SearchPlugin> plugins) {
         indexSearcherExecutorProvider = registerIndexSearcherExecutorProvider(plugins);
         namedWriteables.addAll(SortValue.namedWriteables());
         concurrentSearchDeciderFactories = registerConcurrentSearchDeciderFactories(plugins);
-        registerQueryCollectorContextFactory(plugins);
+        registerQueryCollectorContextSpec(plugins);
     }
 
     private Collection<ConcurrentSearchRequestDecider.Factory> registerConcurrentSearchDeciderFactories(List<SearchPlugin> plugins) {
@@ -1301,21 +1301,8 @@ private SearchPlugin.ExecutorServiceProvider registerIndexSearcherExecutorProvid
         return provider;
     }
 
-    private void registerQueryCollectorContextFactory(List<SearchPlugin> plugins) {
-        // QueryCollectorContextFactory queryCollectorContextFactory = null;
-        // for (SearchPlugin plugin : plugins) {
-        // final Optional<QueryCollectorContextFactory> factoryOpt = plugin.getCollectorContextFactory();
-        //
-        // if (factoryOpt.isPresent()) {
-        // queryCollectorContextFactory = factoryOpt.get();
-        // break;
-        // }
-        // }
-        // if (queryCollectorContextFactory == null) {
-        // queryCollectorContextFactory = new TopDocsCollectorContextFactory();
-        // }
-        // return queryCollectorContextFactory;
-        registerFromPlugin(plugins, SearchPlugin::getCollectorContextFactory, QueryCollectorContextFactoryRegistery::registerFactory);
+    private void registerQueryCollectorContextSpec(List<SearchPlugin> plugins) {
+        registerFromPlugin(plugins, SearchPlugin::getCollectorContextSpec, QueryCollectorContextSpecRegistry::registerFactory);
     }
 
     public FetchPhase getFetchPhase() {

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

@@ -0,0 +1,24 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.query;
+
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.CollectorManager;
+
+import java.io.IOException;
+
+public interface QueryCollectorContextSpec {
+    String getProfileName();
+
+    Collector create(Collector in) throws IOException;
+
+    CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException;
+
+    void postProcess(QuerySearchResult result) throws IOException;
+}

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

@@ -8,11 +8,16 @@
 
 package org.opensearch.search.query;
 
+import org.apache.lucene.search.Query;
+import org.opensearch.search.internal.SearchContext;
+
 import java.io.IOException;
 
 /**
  *  QueryCollectorContextFactory
  */
-public interface QueryCollectorContextFactory {
-    QueryCollectorContext createCollectorContext(CollectorContextParams params) throws IOException;
+public interface QueryCollectorContextSpecFactory {
+    boolean supports(Query query);
+
+    QueryCollectorContextSpec create(SearchContext searchContext, boolean hasFilterCollector) throws IOException;
 }

server/src/main/java/org/opensearch/search/query/QueryCollectorContextFactory.java renamed to server/src/main/java/org/opensearch/search/query/QueryCollectorContextSpecFactory.java

@@ -0,0 +1,33 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.query;
+
+import org.apache.lucene.search.Query;
+import org.opensearch.plugins.SearchPlugin;
+
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+/**
+ * QueryCollectorContextFactoryRegistery
+ */
+public class QueryCollectorContextSpecRegistry {
+    // private static final Map<Class<?>, QueryCollectorContextSpec> registry = new ConcurrentHashMap<>();
+
+    private static final List<QueryCollectorContextSpecFactory> registry = new CopyOnWriteArrayList<>();
+
+    public static QueryCollectorContextSpecFactory getFactory(Query query) {
+
+        return registry.stream().filter(entry -> entry.supports(query)).findFirst().orElse(null);
+    }
+
+    public static void registerFactory(SearchPlugin.FactorySpec<?> specs) {
+        registry.add(specs.getQueryCollectorContextSpec());
+    }
+}

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

@@ -83,6 +83,7 @@
 import static org.opensearch.search.query.QueryCollectorContext.createFilteredCollectorContext;
 import static org.opensearch.search.query.QueryCollectorContext.createMinScoreCollectorContext;
 import static org.opensearch.search.query.QueryCollectorContext.createMultiCollectorContext;
+import static org.opensearch.search.query.TopDocsCollectorContext.createTopDocsCollectorContext;
 
 /**
  * Query phase of a search request, used to run the query and get back from each shard information about the matching documents
@@ -445,27 +446,34 @@ protected boolean searchWithCollector(
             boolean hasTimeout
         ) throws IOException {
             // create the top docs collector last when the other collectors are known
-            QueryCollectorContextFactory queryCollectorContextFactory = QueryCollectorContextFactoryRegistery.getFactory(query);
-            CollectorContextParams collectorContextParams = new CollectorContextParams.Builder().withSearchContext(searchContext)
-                .withHasFilterCollector(hasFilterCollector)
-                .build();
-            if (queryCollectorContextFactory == null) {
-                queryCollectorContextFactory = new TopDocsCollectorContextFactory();
+            QueryCollectorContext queryCollectorContext;
+            QueryCollectorContextSpecFactory queryCollectorContextSpecFactory = QueryCollectorContextSpecRegistry.getFactory(query);
+            if (queryCollectorContextSpecFactory == null) {
+                queryCollectorContext = createTopDocsCollectorContext(searchContext, hasFilterCollector);
+            } else {
+                QueryCollectorContextSpec queryCollectorContextSpec = queryCollectorContextSpecFactory.create(
+                    searchContext,
+                    hasFilterCollector
+                );
+                queryCollectorContext = new QueryCollectorContext(queryCollectorContextSpec.getProfileName()) {
+                    @Override
+                    Collector create(Collector in) throws IOException {
+                        return queryCollectorContextSpec.create(in);
+                    }
+
+                    @Override
+                    CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in)
+                        throws IOException {
+                        return queryCollectorContextSpec.createManager(in);
+                    }
+
+                    @Override
+                    void postProcess(QuerySearchResult result) throws IOException {
+                        queryCollectorContextSpec.postProcess(result);
+                    }
+                };
             }
-            final QueryCollectorContext queryCollectorContext = queryCollectorContextFactory.createCollectorContext(collectorContextParams);
-            // final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
-            return searchWithCollector(searchContext, searcher, query, collectors, queryCollectorContext, hasFilterCollector, hasTimeout);
-        }
 
-        protected boolean searchWithCollector(
-            SearchContext searchContext,
-            ContextIndexSearcher searcher,
-            Query query,
-            LinkedList<QueryCollectorContext> collectors,
-            QueryCollectorContext queryCollectorContext,
-            boolean hasFilterCollector,
-            boolean hasTimeout
-        ) throws IOException {
             return QueryPhase.searchWithCollector(
                 searchContext,
                 searcher,
@@ -476,5 +484,25 @@ protected boolean searchWithCollector(
                 hasTimeout
             );
         }
+
+        // protected boolean searchWithCollector(
+        // SearchContext searchContext,
+        // ContextIndexSearcher searcher,
+        // Query query,
+        // LinkedList<QueryCollectorContext> collectors,
+        // QueryCollectorContext queryCollectorContext,
+        // boolean hasFilterCollector,
+        // boolean hasTimeout
+        // ) throws IOException {
+        // return QueryPhase.searchWithCollector(
+        // searchContext,
+        // searcher,
+        // query,
+        // collectors,
+        // queryCollectorContext,
+        // hasFilterCollector,
+        // hasTimeout
+        // );
+        // }
     }
 }