Initial code for adding the SearchPhaseInjectorProcessor interface in…

… Search Pipeline Signed-off-by: Navneet Verma <navneev@amazon.com>
opensearch-project · Apr 24, 2023 · 7945f01 · 7945f01
1 parent 66e49a6
commit 7945f01
Show file tree

Hide file tree

Showing 23 changed files with 377 additions and 45 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -81,6 +81,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x]
 ### Added
 - [Extensions] Moving Extensions APIs to protobuf serialization. ([#6960](https://github.com/opensearch-project/OpenSearch/pull/6960))
+- [SearchPipeline] Initial code for adding the SearchPhaseInjectorProcessor interface in Search Pipeline.([#7283](https://github.com/opensearch-project/OpenSearch/pull/7283))
 
 ### Dependencies
 

diff --git a/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/AbstractSearchAsyncAction.java
@@ -57,6 +57,7 @@
 import org.opensearch.search.internal.InternalSearchResponse;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.internal.ShardSearchRequest;
+import org.opensearch.search.pipeline.SearchPipelineService;
 import org.opensearch.transport.Transport;
 
 import java.util.ArrayDeque;
@@ -116,6 +117,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
     private final boolean throttleConcurrentRequests;
 
     private final List<Releasable> releasables = new ArrayList<>();
+    private final SearchPipelineService searchPipelineService;
 
     AbstractSearchAsyncAction(
         String name,
@@ -134,7 +136,8 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
         SearchTask task,
         SearchPhaseResults<Result> resultConsumer,
         int maxConcurrentRequestsPerNode,
-        SearchResponse.Clusters clusters
+        SearchResponse.Clusters clusters,
+        SearchPipelineService searchPipelineService
     ) {
         super(name);
         final List<SearchShardIterator> toSkipIterators = new ArrayList<>();
@@ -170,6 +173,7 @@ abstract class AbstractSearchAsyncAction<Result extends SearchPhaseResult> exten
         this.indexRoutings = indexRoutings;
         this.results = resultConsumer;
         this.clusters = clusters;
+        this.searchPipelineService = searchPipelineService;
     }
 
     @Override
@@ -696,7 +700,9 @@ private void raisePhaseFailure(SearchPhaseExecutionException exception) {
      * @see #onShardResult(SearchPhaseResult, SearchShardIterator)
      */
     final void onPhaseDone() {  // as a tribute to @kimchy aka. finishHim()
-        executeNextPhase(this, getNextPhase(results, this));
+        final SearchPhase nextPhase = getNextPhase(results, this);
+        searchPipelineService.transformSearchPhase(results, this, this.getSearchPhaseName(), nextPhase.getSearchPhaseName());
+        executeNextPhase(this, nextPhase);
     }
 
     @Override

diff --git a/server/src/main/java/org/opensearch/action/search/ArraySearchPhaseResults.java b/server/src/main/java/org/opensearch/action/search/ArraySearchPhaseResults.java
@@ -66,7 +66,7 @@ boolean hasResult(int shardIndex) {
     }
 
     @Override
-    AtomicArray<Result> getAtomicArray() {
+    public AtomicArray<Result> getAtomicArray() {
         return results;
     }
 }
diff --git a/server/src/main/java/org/opensearch/action/search/CanMatchPreFilterSearchPhase.java b/server/src/main/java/org/opensearch/action/search/CanMatchPreFilterSearchPhase.java
@@ -41,6 +41,7 @@
 import org.opensearch.search.SearchShardTarget;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.search.internal.AliasFilter;
+import org.opensearch.search.pipeline.SearchPipelineService;
 import org.opensearch.search.sort.FieldSortBuilder;
 import org.opensearch.search.sort.MinAndMax;
 import org.opensearch.search.sort.SortOrder;
@@ -90,7 +91,8 @@ final class CanMatchPreFilterSearchPhase extends AbstractSearchAsyncAction<CanMa
         ClusterState clusterState,
         SearchTask task,
         Function<GroupShardsIterator<SearchShardIterator>, SearchPhase> phaseFactory,
-        SearchResponse.Clusters clusters
+        SearchResponse.Clusters clusters,
+        SearchPipelineService searchPipelineService
     ) {
         // We set max concurrent shard requests to the number of shards so no throttling happens for can_match requests
         super(
@@ -110,7 +112,8 @@ final class CanMatchPreFilterSearchPhase extends AbstractSearchAsyncAction<CanMa
             task,
             new CanMatchSearchPhaseResults(shardsIts.size()),
             shardsIts.size(),
-            clusters
+            clusters,
+            searchPipelineService
         );
         this.phaseFactory = phaseFactory;
         this.shardsIts = shardsIts;

diff --git a/server/src/main/java/org/opensearch/action/search/DfsQueryPhase.java b/server/src/main/java/org/opensearch/action/search/DfsQueryPhase.java
@@ -69,7 +69,7 @@ final class DfsQueryPhase extends SearchPhase {
         Function<ArraySearchPhaseResults<SearchPhaseResult>, SearchPhase> nextPhaseFactory,
         SearchPhaseContext context
     ) {
-        super("dfs_query");
+        super(SearchPhaseName.DFS_QUERY.name());
         this.progressListener = context.getTask().getProgressListener();
         this.queryResult = queryResult;
         this.searchResults = searchResults;

diff --git a/server/src/main/java/org/opensearch/action/search/ExpandSearchPhase.java b/server/src/main/java/org/opensearch/action/search/ExpandSearchPhase.java
@@ -62,7 +62,7 @@ final class ExpandSearchPhase extends SearchPhase {
     private final AtomicArray<SearchPhaseResult> queryResults;
 
     ExpandSearchPhase(SearchPhaseContext context, InternalSearchResponse searchResponse, AtomicArray<SearchPhaseResult> queryResults) {
-        super("expand");
+        super(SearchPhaseName.EXPAND.name());
         this.context = context;
         this.searchResponse = searchResponse;
         this.queryResults = queryResults;

diff --git a/server/src/main/java/org/opensearch/action/search/FetchSearchPhase.java b/server/src/main/java/org/opensearch/action/search/FetchSearchPhase.java
@@ -92,7 +92,7 @@ final class FetchSearchPhase extends SearchPhase {
         SearchPhaseContext context,
         BiFunction<InternalSearchResponse, AtomicArray<SearchPhaseResult>, SearchPhase> nextPhaseFactory
     ) {
-        super("fetch");
+        super(SearchPhaseName.FETCH.name());
         if (context.getNumShards() != resultConsumer.getNumShards()) {
             throw new IllegalStateException(
                 "number of shards must match the length of the query results but doesn't:"

diff --git a/server/src/main/java/org/opensearch/action/search/SearchDfsQueryThenFetchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/SearchDfsQueryThenFetchAsyncAction.java
@@ -41,6 +41,7 @@
 import org.opensearch.search.dfs.AggregatedDfs;
 import org.opensearch.search.dfs.DfsSearchResult;
 import org.opensearch.search.internal.AliasFilter;
+import org.opensearch.search.pipeline.SearchPipelineService;
 import org.opensearch.transport.Transport;
 
 import java.util.List;
@@ -76,7 +77,8 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
         final TransportSearchAction.SearchTimeProvider timeProvider,
         final ClusterState clusterState,
         final SearchTask task,
-        SearchResponse.Clusters clusters
+        SearchResponse.Clusters clusters,
+        SearchPipelineService searchPipelineService
     ) {
         super(
             "dfs",
@@ -95,7 +97,8 @@ final class SearchDfsQueryThenFetchAsyncAction extends AbstractSearchAsyncAction
             task,
             new ArraySearchPhaseResults<>(shardsIts.size()),
             request.getMaxConcurrentShardRequests(),
-            clusters
+            clusters,
+            searchPipelineService
         );
         this.queryPhaseResultConsumer = queryPhaseResultConsumer;
         this.searchPhaseController = searchPhaseController;

diff --git a/server/src/main/java/org/opensearch/action/search/SearchPhase.java b/server/src/main/java/org/opensearch/action/search/SearchPhase.java
@@ -41,7 +41,7 @@
  *
  * @opensearch.internal
  */
-abstract class SearchPhase implements CheckedRunnable<IOException> {
+public abstract class SearchPhase implements CheckedRunnable<IOException> {
     private final String name;
 
     protected SearchPhase(String name) {
@@ -54,4 +54,25 @@ protected SearchPhase(String name) {
     public String getName() {
         return name;
     }
+
+    public SearchPhaseName getSearchPhaseName() {
+        return SearchPhaseName.valueOf(name);
+    }
+
+    /**
+     * Enum for different Search Phases in OpenSearch
+     * @opensearch.internal
+     */
+    public enum SearchPhaseName {
+        QUERY("query"),
+        FETCH("fetch"),
+        DFS_QUERY("dfs_query"),
+        EXPAND("expand");
+
+        private final String name;
+
+        SearchPhaseName(final String name) {
+            this.name = name;
+        }
+    }
 }
diff --git a/server/src/main/java/org/opensearch/action/search/SearchPhaseContext.java b/server/src/main/java/org/opensearch/action/search/SearchPhaseContext.java
@@ -50,7 +50,7 @@
  *
  * @opensearch.internal
  */
-interface SearchPhaseContext extends Executor {
+public interface SearchPhaseContext extends Executor {
     // TODO maybe we can make this concrete later - for now we just implement this in the base class for all initial phases
 
     /**

diff --git a/server/src/main/java/org/opensearch/action/search/SearchPhaseResults.java b/server/src/main/java/org/opensearch/action/search/SearchPhaseResults.java
@@ -42,7 +42,7 @@
  *
  * @opensearch.internal
  */
-abstract class SearchPhaseResults<Result extends SearchPhaseResult> {
+public abstract class SearchPhaseResults<Result extends SearchPhaseResult> {
     private final int numShards;
 
     SearchPhaseResults(int numShards) {
@@ -75,7 +75,7 @@ final int getNumShards() {
 
     void consumeShardFailure(int shardIndex) {}
 
-    AtomicArray<Result> getAtomicArray() {
+    public AtomicArray<Result> getAtomicArray() {
         throw new UnsupportedOperationException();
     }
 

diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryThenFetchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/SearchQueryThenFetchAsyncAction.java
@@ -42,6 +42,7 @@
 import org.opensearch.search.internal.AliasFilter;
 import org.opensearch.search.internal.SearchContext;
 import org.opensearch.search.internal.ShardSearchRequest;
+import org.opensearch.search.pipeline.SearchPipelineService;
 import org.opensearch.search.query.QuerySearchResult;
 import org.opensearch.transport.Transport;
 
@@ -81,7 +82,8 @@ class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<SearchPh
         final TransportSearchAction.SearchTimeProvider timeProvider,
         ClusterState clusterState,
         SearchTask task,
-        SearchResponse.Clusters clusters
+        SearchResponse.Clusters clusters,
+        SearchPipelineService searchPipelineService
     ) {
         super(
             "query",
@@ -100,7 +102,8 @@ class SearchQueryThenFetchAsyncAction extends AbstractSearchAsyncAction<SearchPh
             task,
             resultConsumer,
             request.getMaxConcurrentShardRequests(),
-            clusters
+            clusters,
+            searchPipelineService
         );
         this.topDocsSize = SearchPhaseController.getTopDocsSize(request);
         this.trackTotalHitsUpTo = request.resolveTrackTotalHitsUpTo();

diff --git a/server/src/main/java/org/opensearch/action/search/SearchScrollAsyncAction.java b/server/src/main/java/org/opensearch/action/search/SearchScrollAsyncAction.java
@@ -266,7 +266,7 @@ protected SearchPhase sendResponsePhase(
         SearchPhaseController.ReducedQueryPhase queryPhase,
         final AtomicArray<? extends SearchPhaseResult> fetchResults
     ) {
-        return new SearchPhase("fetch") {
+        return new SearchPhase(SearchPhase.SearchPhaseName.FETCH.name()) {
             @Override
             public void run() throws IOException {
                 sendResponse(queryPhase, fetchResults);

diff --git a/server/src/main/java/org/opensearch/action/search/SearchScrollQueryThenFetchAsyncAction.java b/server/src/main/java/org/opensearch/action/search/SearchScrollQueryThenFetchAsyncAction.java
@@ -92,7 +92,7 @@ protected void executeInitialPhase(
 
     @Override
     protected SearchPhase moveToNextPhase(BiFunction<String, String, DiscoveryNode> clusterNodeLookup) {
-        return new SearchPhase("fetch") {
+        return new SearchPhase(SearchPhase.SearchPhaseName.FETCH.name()) {
             @Override
             public void run() {
                 final SearchPhaseController.ReducedQueryPhase reducedQueryPhase = searchPhaseController.reducedScrollQueryPhase(

diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
@@ -345,7 +345,8 @@ public AbstractSearchAsyncAction<? extends SearchPhaseResult> asyncSearchAction(
                     task,
                     new ArraySearchPhaseResults<>(shardsIts.size()),
                     searchRequest.getMaxConcurrentShardRequests(),
-                    clusters
+                    clusters,
+                    searchPipelineService
                 ) {
                     @Override
                     protected void executePhaseOnShard(
@@ -1161,7 +1162,8 @@ public void run() {
                         }
                     };
                 },
-                clusters
+                clusters,
+                searchPipelineService
             );
         } else {
             final QueryPhaseResultConsumer queryResultConsumer = searchPhaseController.newSearchPhaseResults(
@@ -1191,7 +1193,8 @@ public void run() {
                         timeProvider,
                         clusterState,
                         task,
-                        clusters
+                        clusters,
+                        searchPipelineService
                     );
                     break;
                 case QUERY_THEN_FETCH:
@@ -1211,7 +1214,8 @@ public void run() {
                         timeProvider,
                         clusterState,
                         task,
-                        clusters
+                        clusters,
+                        searchPipelineService
                     );
                     break;
                 default: