[Draft] Poc on hybrid collector context

src/main/java/org/opensearch/neuralsearch/plugin/NeuralSearch.java

@@ -15,7 +15,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.function.Supplier;
 
 import org.opensearch.ml.client.MachineLearningNodeClient;
@@ -24,6 +23,7 @@
 import org.opensearch.neuralsearch.highlight.extractor.QueryTextExtractorRegistry;
 import com.google.common.collect.ImmutableList;
 import org.opensearch.action.ActionRequest;
+import org.opensearch.neuralsearch.search.collector.HybridQueryCollectorContextSpecFactory;
 import org.opensearch.neuralsearch.query.NeuralQueryBuilder;
 import org.opensearch.neuralsearch.query.HybridQueryBuilder;
 import org.opensearch.neuralsearch.query.NeuralSparseQueryBuilder;
@@ -38,6 +38,7 @@
 import org.opensearch.neuralsearch.mappingtransformer.SemanticMappingTransformer;
 import org.opensearch.neuralsearch.processor.factory.SemanticFieldProcessorFactory;
 import org.opensearch.plugins.MapperPlugin;
+import org.opensearch.search.query.QueryCollectorContextSpecFactory;
 import org.opensearch.transport.client.Client;
 import org.opensearch.cluster.metadata.IndexNameExpressionResolver;
 import org.opensearch.cluster.node.DiscoveryNodes;
@@ -81,7 +82,6 @@
 import org.opensearch.neuralsearch.processor.rerank.RerankProcessor;
 import org.opensearch.neuralsearch.query.ext.RerankSearchExtBuilder;
 import org.opensearch.neuralsearch.rest.RestNeuralStatsAction;
-import org.opensearch.neuralsearch.search.query.HybridQueryPhaseSearcher;
 import org.opensearch.neuralsearch.transport.NeuralStatsAction;
 import org.opensearch.neuralsearch.transport.NeuralStatsTransportAction;
 import org.opensearch.neuralsearch.util.NeuralSearchClusterUtil;
@@ -100,7 +100,6 @@
 import org.opensearch.search.pipeline.SearchPhaseResultsProcessor;
 import org.opensearch.search.pipeline.SearchRequestProcessor;
 import org.opensearch.search.pipeline.SearchResponseProcessor;
-import org.opensearch.search.query.QueryPhaseSearcher;
 import org.opensearch.threadpool.ExecutorBuilder;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.watcher.ResourceWatcherService;
@@ -235,8 +234,8 @@ public Map<String, Processor.Factory> getProcessors(Processor.Parameters paramet
     }
 
     @Override
-    public Optional<QueryPhaseSearcher> getQueryPhaseSearcher() {
-        return Optional.of(new HybridQueryPhaseSearcher());
+    public List<QueryCollectorContextSpecFactory> getCollectorContextSpecFactories() {
+        return List.of(new HybridQueryCollectorContextSpecFactory());
     }
 
     @Override

src/main/java/org/opensearch/neuralsearch/search/collector/HybridQueryCollectorContextSpec.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.search.collector;
+
+import lombok.extern.log4j.Log4j2;
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.CollectorManager;
+import org.opensearch.common.lucene.search.TopDocsAndMaxScore;
+import org.opensearch.neuralsearch.search.query.HybridCollectorManager;
+import org.opensearch.neuralsearch.search.query.HybridCollectorResultsUtilParams;
+import org.opensearch.neuralsearch.search.query.HybridSearchCollectorResultUtil;
+import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.query.QueryCollectorContextSpec;
+import org.opensearch.search.query.QuerySearchResult;
+import org.opensearch.search.query.ReduceableSearchResult;
+
+import java.io.IOException;
+
+import static org.opensearch.search.profile.query.CollectorResult.REASON_SEARCH_TOP_HITS;
+
+@Log4j2
+public class HybridQueryCollectorContextSpec implements QueryCollectorContextSpec {
+    private HybridCollectorManager collectorManager;
+    private HybridSearchCollector collector;
+    private SearchContext searchContext;
+
+    public HybridQueryCollectorContextSpec(final SearchContext searchContext) throws IOException {
+        this.searchContext = searchContext;
+        this.collectorManager = (HybridCollectorManager) HybridCollectorManager.createHybridCollectorManager(searchContext);
+        this.collector = (HybridSearchCollector) collectorManager.newCollector();
+    }
+
+    @Override
+    public String getContextName() {
+        return REASON_SEARCH_TOP_HITS;
+    }
+
+    @Override
+    public Collector create(Collector in) throws IOException {
+        return collector;
+    }
+
+    @Override
+    public CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException {
+        return collectorManager;
+    }
+
+    @Override
+    public void postProcess(QuerySearchResult result) throws IOException {
+        HybridSearchCollectorResultUtil hybridSearchCollectorResultUtil = new HybridSearchCollectorResultUtil(
+            new HybridCollectorResultsUtilParams.Builder().searchContext(searchContext).build(),
+            collector
+        );
+        TopDocsAndMaxScore topDocsAndMaxScore = hybridSearchCollectorResultUtil.getTopDocsAndAndMaxScore();
+        hybridSearchCollectorResultUtil.reduceCollectorResults(result, topDocsAndMaxScore);
+    }
+}

src/main/java/org/opensearch/neuralsearch/search/collector/HybridQueryCollectorContextSpecFactory.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.opensearch.neuralsearch.search.collector;
+
+import org.apache.lucene.search.Query;
+import org.opensearch.index.query.ParsedQuery;
+import org.opensearch.neuralsearch.query.HybridQuery;
+import org.opensearch.search.internal.SearchContext;
+import org.opensearch.search.query.QueryCollectorArguments;
+import org.opensearch.search.query.QueryCollectorContextSpec;
+import org.opensearch.search.query.QueryCollectorContextSpecFactory;
+
+import java.io.IOException;
+import java.util.Optional;
+
+import static org.opensearch.neuralsearch.util.HybridQueryUtil.extractHybridQuery;
+import static org.opensearch.neuralsearch.util.HybridQueryUtil.isHybridQuery;
+import static org.opensearch.neuralsearch.util.HybridQueryUtil.isWrappedHybridQuery;
+import static org.opensearch.neuralsearch.util.HybridQueryUtil.validateHybridQuery;
+
+public class HybridQueryCollectorContextSpecFactory implements QueryCollectorContextSpecFactory {
+
+    @Override
+    public Optional<QueryCollectorContextSpec> createQueryCollectorContextSpec(
+        SearchContext searchContext,
+        QueryCollectorArguments queryCollectorArguments
+    ) throws IOException {
+        if (isHybridQuery(searchContext.query(), searchContext)) {
+            Query query = extractHybridQuery(searchContext);
+            validateHybridQuery((HybridQuery) query);
+            searchContext.parsedQuery(new ParsedQuery(query));
+            return Optional.of(new HybridQueryCollectorContextSpec(searchContext));
+        } else if (isWrappedHybridQuery(searchContext.query())) {
+            throw new IllegalArgumentException("hybrid query must be a top level query and cannot be wrapped into other queries");
+        }
+        return Optional.empty();
+    }
+}