Inject QueryCollectorSpecFactory from Plugins to create QueryCollectorContext in QueryPhase

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce a new pull-based ingestion plugin for file-based indexing (for local testing) ([#18591](https://github.com/opensearch-project/OpenSearch/pull/18591))
 - Add support for search pipeline in search and msearch template ([#18564](https://github.com/opensearch-project/OpenSearch/pull/18564))
 - Add BooleanQuery rewrite moving constant-scoring must clauses to filter clauses ([#18510](https://github.com/opensearch-project/OpenSearch/issues/18510))
+- Add functionality for plugins to inject QueryCollectorContext during QueryPhase ([#18637](https://github.com/opensearch-project/OpenSearch/pull/18637))
 - Add support for non-timing info in profiler ([#18460](https://github.com/opensearch-project/OpenSearch/issues/18460))
 
 ### Changed

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

@@ -68,6 +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.QueryCollectorContextSpecFactory;
 import org.opensearch.search.query.QueryPhaseSearcher;
 import org.opensearch.search.rescore.Rescorer;
 import org.opensearch.search.rescore.RescorerBuilder;
@@ -227,6 +228,10 @@ default Optional<ExecutorServiceProvider> getIndexSearcherExecutorProvider() {
         return Optional.empty();
     }
 
+    default List<QueryCollectorContextSpecFactory> getCollectorContextSpecFactories() {
+        return emptyList();
+    }
+
     /**
      * Executor service provider
      */

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

@@ -257,6 +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.QueryCollectorContextSpecRegistry;
 import org.opensearch.search.query.QueryPhase;
 import org.opensearch.search.query.QueryPhaseSearcher;
 import org.opensearch.search.query.QueryPhaseSearcherWrapper;
@@ -350,6 +351,7 @@ public SearchModule(Settings settings, List<SearchPlugin> plugins) {
         indexSearcherExecutorProvider = registerIndexSearcherExecutorProvider(plugins);
         namedWriteables.addAll(SortValue.namedWriteables());
         concurrentSearchDeciderFactories = registerConcurrentSearchDeciderFactories(plugins);
+        registerQueryCollectorContextSpec(plugins);
     }
 
     private Collection<ConcurrentSearchRequestDecider.Factory> registerConcurrentSearchDeciderFactories(List<SearchPlugin> plugins) {
@@ -1297,6 +1299,10 @@ private SearchPlugin.ExecutorServiceProvider registerIndexSearcherExecutorProvid
         return provider;
     }
 
+    private void registerQueryCollectorContextSpec(List<SearchPlugin> plugins) {
+        registerFromPlugin(plugins, SearchPlugin::getCollectorContextSpecFactories, QueryCollectorContextSpecRegistry::registerFactory);
+    }
+
     public FetchPhase getFetchPhase() {
         return new FetchPhase(fetchSubPhases);
     }

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

@@ -0,0 +1,78 @@
+/*
+ * 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.opensearch.common.annotation.ExperimentalApi;
+
+/**
+ * Arguments for {@link QueryCollectorContextSpecRegistry}
+ */
+@ExperimentalApi
+public final class QueryCollectorArguments {
+    private final boolean hasFilterCollector;
+
+    private QueryCollectorArguments(final boolean hasFilterCollector) {
+        this.hasFilterCollector = hasFilterCollector;
+    }
+
+    /**
+     * Whether the query has a filter collector.
+     * @return true if the query has a filter collector, false otherwise
+     */
+    public boolean hasFilterCollector() {
+        return hasFilterCollector;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+
+        QueryCollectorArguments queryCollectorArguments = (QueryCollectorArguments) o;
+        return hasFilterCollector == queryCollectorArguments.hasFilterCollector;
+    }
+
+    @Override
+    public int hashCode() {
+        return Boolean.hashCode(hasFilterCollector);
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public String toString() {
+        return "QueryCollectorArguments[hasFilterCollector=" + hasFilterCollector + "]";
+    }
+
+    /**
+     * Builder for {@link QueryCollectorArguments}
+     */
+    public static class Builder {
+        private boolean hasFilterCollector;
+
+        /**
+         * Set the flag for query has a filter collector.
+         * @param hasFilterCollector true if the query has a filter collector, false otherwise
+         * @return Builder instance
+         */
+        public Builder hasFilterCollector(boolean hasFilterCollector) {
+            this.hasFilterCollector = hasFilterCollector;
+            return this;
+        }
+
+        /**
+         * Build the arguments for the query collector context spec registry.
+         * @return QueryCollectorArguments instance
+         */
+        public QueryCollectorArguments build() {
+            return new QueryCollectorArguments(hasFilterCollector);
+        }
+    }
+}

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

@@ -0,0 +1,50 @@
+/*
+ * 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 org.opensearch.common.annotation.ExperimentalApi;
+
+import java.io.IOException;
+
+/**
+ * interface of QueryCollectorContextSpec
+ */
+@ExperimentalApi
+public interface QueryCollectorContextSpec {
+    /**
+     * Context name for QueryCollectorContext
+     * @return string of context name
+     */
+    String getContextName();
+
+    /**
+     * Create collector
+     * @param in collector
+     * @return collector
+     * @throws IOException
+     */
+    Collector create(Collector in) throws IOException;
+
+    /**
+     * Create collector manager
+     * @param in collector manager
+     * @return collector manager
+     * @throws IOException
+     */
+    CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException;
+
+    /**
+     * Post process query result
+     * @param result query result
+     * @throws IOException
+     */
+    void postProcess(QuerySearchResult result) throws IOException;
+}

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

@@ -0,0 +1,32 @@
+/*
+ * 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.opensearch.common.annotation.ExperimentalApi;
+import org.opensearch.search.internal.SearchContext;
+
+import java.io.IOException;
+import java.util.Optional;
+
+/**
+ *  interface of QueryCollectorContext spec factory
+ */
+@ExperimentalApi
+public interface QueryCollectorContextSpecFactory {
+    /**
+     * @param searchContext context needed to create collector context spec
+     * @param queryCollectorArguments arguments to create collector context spec
+     * @return QueryCollectorContextSpec
+     * @throws IOException
+     */
+    Optional<QueryCollectorContextSpec> createQueryCollectorContextSpec(
+        SearchContext searchContext,
+        QueryCollectorArguments queryCollectorArguments
+    ) throws IOException;
+}

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

@@ -0,0 +1,64 @@
+/*
+ * 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.opensearch.search.internal.SearchContext;
+
+import java.io.IOException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+/**
+ * Registry class to load all collector context spec factories during cluster bootstrapping
+ */
+public final class QueryCollectorContextSpecRegistry {
+    private static final List<QueryCollectorContextSpecFactory> registry = new CopyOnWriteArrayList<>();
+
+    private QueryCollectorContextSpecRegistry() {}
+
+    /**
+     * Get all collector context spec factories
+     * @return list of collector context spec factories
+     */
+    public static List<QueryCollectorContextSpecFactory> getCollectorContextSpecFactories() {
+        return registry;
+    }
+
+    /**
+     * Register factory
+     * @param factory collector context spec factory defined in plugin
+     */
+    public static void registerFactory(QueryCollectorContextSpecFactory factory) {
+        registry.add(factory);
+    }
+
+    /**
+     * Get collector context spec
+     * @param searchContext search context
+     * @param queryCollectorArguments query collector arguments
+     * @return collector context spec
+     * @throws IOException
+     */
+    public static Optional<QueryCollectorContextSpec> getQueryCollectorContextSpec(
+        final SearchContext searchContext,
+        final QueryCollectorArguments queryCollectorArguments
+    ) throws IOException {
+        Iterator<QueryCollectorContextSpecFactory> iterator = registry.iterator();
+        while (iterator.hasNext()) {
+            QueryCollectorContextSpecFactory factory = iterator.next();
+            Optional<QueryCollectorContextSpec> spec = factory.createQueryCollectorContextSpec(searchContext, queryCollectorArguments);
+            if (spec.isEmpty() == false) {
+                return spec;
+            }
+        }
+        return Optional.empty();
+    }
+}

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

@@ -76,6 +76,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.stream.Collectors;
 
@@ -445,9 +446,37 @@
             boolean hasFilterCollector,
             boolean hasTimeout
         ) throws IOException {
+            QueryCollectorContext queryCollectorContext = getQueryCollectorContext(searchContext, hasFilterCollector);
+            return searchWithCollector(searchContext, searcher, query, collectors, queryCollectorContext, hasFilterCollector, hasTimeout);
+        }
+
+        private QueryCollectorContext getQueryCollectorContext(SearchContext searchContext, boolean hasFilterCollector) throws IOException {
             // create the top docs collector last when the other collectors are known
-            final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
-            return searchWithCollector(searchContext, searcher, query, collectors, topDocsFactory, hasFilterCollector, hasTimeout);
+            final Optional<QueryCollectorContext> queryCollectorContextOpt = QueryCollectorContextSpecRegistry.getQueryCollectorContextSpec(
+                searchContext,
+                new QueryCollectorArguments.Builder().hasFilterCollector(hasFilterCollector).build()
+            ).map(queryCollectorContextSpec -> new QueryCollectorContext(queryCollectorContextSpec.getContextName()) {
+                @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);
+                }
+            });
+            if (queryCollectorContextOpt.isPresent()) {
+                return queryCollectorContextOpt.get();
+            } else {
+                return createTopDocsCollectorContext(searchContext, hasFilterCollector);
+            }
         }
 
         protected boolean searchWithCollector(

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

@@ -0,0 +1,49 @@
+/*
+ * 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.opensearch.test.OpenSearchTestCase;
+
+public class QueryCollectorArgumentsTests extends OpenSearchTestCase {
+
+    public void testBuilder() {
+        QueryCollectorArguments args = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
+
+        assertTrue(args.hasFilterCollector());
+    }
+
+    public void testEquals() {
+        QueryCollectorArguments args1 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
+
+        QueryCollectorArguments args2 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
+
+        QueryCollectorArguments args3 = new QueryCollectorArguments.Builder().hasFilterCollector(false).build();
+
+        assertTrue(args1.equals(args2));          // Same values
+        assertFalse(args1.equals(args3));         // Different values
+        assertTrue(args1.equals(args1));          // Same object
+    }
+
+    public void testHashCode() {
+        QueryCollectorArguments args1 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
+
+        QueryCollectorArguments args2 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
+
+        assertEquals(args1.hashCode(), args2.hashCode());
+        assertEquals(args1.hashCode(), args1.hashCode());  // Consistent
+    }
+
+    public void testToString() {
+        QueryCollectorArguments args = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
+
+        String result = args.toString();
+
+        assertEquals("QueryCollectorArguments[hasFilterCollector=true]", result);
+    }
+}