Added support for msearch API to pass search pipeline name (opensearc…

…h-project#15923) * Added support for search pipeline name in multi search API Signed-off-by: Owais <owaiskazi19@gmail.com> * Updated CHANGELOG Signed-off-by: Owais <owaiskazi19@gmail.com> * Pulled search pipeline in MultiSearchRequest and updated test Signed-off-by: Owais <owaiskazi19@gmail.com> * Updated test Signed-off-by: Owais <owaiskazi19@gmail.com> * Updated SearchRequest with search pipeline from source Signed-off-by: Owais <owaiskazi19@gmail.com> * Added tests for parseSearchRequest Signed-off-by: Owais <owaiskazi19@gmail.com> * Guard serialization with version check Signed-off-by: Owais <owaiskazi19@gmail.com> * Updated version and added another test for serialization Signed-off-by: Owais <owaiskazi19@gmail.com> --------- Signed-off-by: Owais <owaiskazi19@gmail.com>
ruai0511 · Oct 4, 2024 · b1e815d · b1e815d
1 parent 8256b25
commit b1e815d
Show file tree

Hide file tree

Showing 7 changed files with 139 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Implement WithFieldName interface in ValuesSourceAggregationBuilder & FieldSortBuilder ([#15916](https://github.com/opensearch-project/OpenSearch/pull/15916))
 - Add successfulSearchShardIndices in searchRequestContext ([#15967](https://github.com/opensearch-project/OpenSearch/pull/15967))
 - Remove identity-related feature flagged code from the RestController ([#15430](https://github.com/opensearch-project/OpenSearch/pull/15430))
+- Add support for msearch API to pass search pipeline name - ([#15923](https://github.com/opensearch-project/OpenSearch/pull/15923))
 
 ### Dependencies
 - Bump `com.azure:azure-identity` from 1.13.0 to 1.13.2 ([#15578](https://github.com/opensearch-project/OpenSearch/pull/15578))

diff --git a/server/src/main/java/org/opensearch/action/search/MultiSearchRequest.java b/server/src/main/java/org/opensearch/action/search/MultiSearchRequest.java
@@ -310,6 +310,10 @@ public static void readMultiLineFormat(
             ) {
                 consumer.accept(searchRequest, parser);
             }
+
+            if (searchRequest.source() != null && searchRequest.source().pipeline() != null) {
+                searchRequest.pipeline(searchRequest.source().pipeline());
+            }
             // move pointers
             from = nextMarker + 1;
         }

diff --git a/server/src/main/java/org/opensearch/rest/action/search/RestSearchAction.java b/server/src/main/java/org/opensearch/rest/action/search/RestSearchAction.java
@@ -210,7 +210,7 @@ public static void parseSearchRequest(
         searchRequest.routing(request.param("routing"));
         searchRequest.preference(request.param("preference"));
         searchRequest.indicesOptions(IndicesOptions.fromRequest(request, searchRequest.indicesOptions()));
-        searchRequest.pipeline(request.param("search_pipeline"));
+        searchRequest.pipeline(request.param("search_pipeline", searchRequest.source().pipeline()));
 
         checkRestTotalHits(request, searchRequest);
         request.paramAsBoolean(INCLUDE_NAMED_QUERIES_SCORE_PARAM, false);

diff --git a/server/src/main/java/org/opensearch/search/builder/SearchSourceBuilder.java b/server/src/main/java/org/opensearch/search/builder/SearchSourceBuilder.java
@@ -224,6 +224,8 @@ public static HighlightBuilder highlight() {
 
     private Map<String, Object> searchPipelineSource = null;
 
+    private String searchPipeline;
+
     /**
      * Constructs a new search source builder.
      */
@@ -297,6 +299,9 @@ public SearchSourceBuilder(StreamInput in) throws IOException {
                 derivedFields = in.readList(DerivedField::new);
             }
         }
+        if (in.getVersion().onOrAfter(Version.V_3_0_0)) {
+            searchPipeline = in.readOptionalString();
+        }
     }
 
     @Override
@@ -377,6 +382,9 @@ public void writeTo(StreamOutput out) throws IOException {
                 out.writeList(derivedFields);
             }
         }
+        if (out.getVersion().onOrAfter(Version.V_3_0_0)) {
+            out.writeOptionalString(searchPipeline);
+        }
     }
 
     /**
@@ -1111,6 +1119,13 @@ public Map<String, Object> searchPipelineSource() {
         return searchPipelineSource;
     }
 
+    /**
+     * @return a search pipeline name defined within the search source (see {@link org.opensearch.search.pipeline.SearchPipelineService})
+     */
+    public String pipeline() {
+        return searchPipeline;
+    }
+
     /**
      * Define a search pipeline to process this search request and/or its response. See {@link org.opensearch.search.pipeline.SearchPipelineService}.
      */
@@ -1119,6 +1134,14 @@ public SearchSourceBuilder searchPipelineSource(Map<String, Object> searchPipeli
         return this;
     }
 
+    /**
+     * Define a search pipeline name to process this search request and/or its response. See {@link org.opensearch.search.pipeline.SearchPipelineService}.
+     */
+    public SearchSourceBuilder pipeline(String searchPipeline) {
+        this.searchPipeline = searchPipeline;
+        return this;
+    }
+
     /**
      * Rewrites this search source builder into its primitive form. e.g. by
      * rewriting the QueryBuilder. If the builder did not change the identity
@@ -1216,6 +1239,7 @@ private SearchSourceBuilder shallowCopy(
         rewrittenBuilder.pointInTimeBuilder = pointInTimeBuilder;
         rewrittenBuilder.derivedFieldsObject = derivedFieldsObject;
         rewrittenBuilder.derivedFields = derivedFields;
+        rewrittenBuilder.searchPipeline = searchPipeline;
         return rewrittenBuilder;
     }
 
@@ -1283,6 +1307,8 @@ public void parseXContent(XContentParser parser, boolean checkTrailingTokens) th
                     sort(parser.text());
                 } else if (PROFILE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
                     profile = parser.booleanValue();
+                } else if (SEARCH_PIPELINE.match(currentFieldName, parser.getDeprecationHandler())) {
+                    searchPipeline = parser.text();
                 } else {
                     throw new ParsingException(
                         parser.getTokenLocation(),
@@ -1612,6 +1638,10 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
 
         }
 
+        if (searchPipeline != null) {
+            builder.field(SEARCH_PIPELINE.getPreferredName(), searchPipeline);
+        }
+
         return builder;
     }
 
@@ -1889,7 +1919,8 @@ public int hashCode() {
             trackTotalHitsUpTo,
             pointInTimeBuilder,
             derivedFieldsObject,
-            derivedFields
+            derivedFields,
+            searchPipeline
         );
     }
 
@@ -1934,7 +1965,8 @@ public boolean equals(Object obj) {
             && Objects.equals(trackTotalHitsUpTo, other.trackTotalHitsUpTo)
             && Objects.equals(pointInTimeBuilder, other.pointInTimeBuilder)
             && Objects.equals(derivedFieldsObject, other.derivedFieldsObject)
-            && Objects.equals(derivedFields, other.derivedFields);
+            && Objects.equals(derivedFields, other.derivedFields)
+            && Objects.equals(searchPipeline, other.searchPipeline);
     }
 
     @Override

diff --git a/server/src/test/java/org/opensearch/action/search/SearchRequestTests.java b/server/src/test/java/org/opensearch/action/search/SearchRequestTests.java
@@ -42,6 +42,8 @@
 import org.opensearch.geometry.LinearRing;
 import org.opensearch.index.query.GeoShapeQueryBuilder;
 import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.rest.RestRequest;
+import org.opensearch.rest.action.search.RestSearchAction;
 import org.opensearch.search.AbstractSearchTestCase;
 import org.opensearch.search.Scroll;
 import org.opensearch.search.builder.PointInTimeBuilder;
@@ -50,14 +52,18 @@
 import org.opensearch.search.rescore.QueryRescorerBuilder;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.VersionUtils;
+import org.opensearch.test.rest.FakeRestRequest;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.IntConsumer;
 
 import static java.util.Collections.emptyMap;
+import static org.opensearch.action.search.SearchType.DFS_QUERY_THEN_FETCH;
 import static org.opensearch.test.EqualsHashCodeTestUtils.checkEqualsAndHashCode;
 import static org.hamcrest.Matchers.equalTo;
+import static org.mockito.Mockito.mock;
 
 public class SearchRequestTests extends AbstractSearchTestCase {
 
@@ -242,6 +248,19 @@ public void testCopyConstructor() throws IOException {
         assertNotSame(deserializedRequest, searchRequest);
     }
 
+    public void testParseSearchRequestWithUnsupportedSearchType() throws IOException {
+        RestRequest restRequest = new FakeRestRequest();
+        SearchRequest searchRequest = createSearchRequest();
+        IntConsumer setSize = mock(IntConsumer.class);
+        restRequest.params().put("search_type", "query_and_fetch");
+
+        IllegalArgumentException exception = expectThrows(
+            IllegalArgumentException.class,
+            () -> RestSearchAction.parseSearchRequest(searchRequest, restRequest, null, namedWriteableRegistry, setSize)
+        );
+        assertEquals("Unsupported search type [query_and_fetch]", exception.getMessage());
+    }
+
     public void testEqualsAndHashcode() throws IOException {
         checkEqualsAndHashCode(createSearchRequest(), SearchRequest::new, this::mutate);
     }
@@ -268,10 +287,7 @@ private SearchRequest mutate(SearchRequest searchRequest) {
         );
         mutators.add(
             () -> mutation.searchType(
-                randomValueOtherThan(
-                    searchRequest.searchType(),
-                    () -> randomFrom(SearchType.DFS_QUERY_THEN_FETCH, SearchType.QUERY_THEN_FETCH)
-                )
+                randomValueOtherThan(searchRequest.searchType(), () -> randomFrom(DFS_QUERY_THEN_FETCH, SearchType.QUERY_THEN_FETCH))
             )
         );
         mutators.add(() -> mutation.source(randomValueOtherThan(searchRequest.source(), this::createSearchSourceBuilder)));

diff --git a/server/src/test/java/org/opensearch/search/builder/SearchSourceBuilderTests.java b/server/src/test/java/org/opensearch/search/builder/SearchSourceBuilderTests.java
@@ -421,6 +421,27 @@ public void testDerivedFieldsParsingAndSerializationObjectType() throws IOExcept
         }
     }
 
+    public void testSearchPipelineParsingAndSerialization() throws IOException {
+        String restContent = "{ \"query\": { \"match_all\": {} }, \"from\": 0, \"size\": 10, \"search_pipeline\": \"my_pipeline\" }";
+        String expectedContent = "{\"from\":0,\"size\":10,\"query\":{\"match_all\":{\"boost\":1.0}},\"search_pipeline\":\"my_pipeline\"}";
+
+        try (XContentParser parser = createParser(JsonXContent.jsonXContent, restContent)) {
+            SearchSourceBuilder searchSourceBuilder = SearchSourceBuilder.fromXContent(parser);
+            searchSourceBuilder = rewrite(searchSourceBuilder);
+
+            try (BytesStreamOutput output = new BytesStreamOutput()) {
+                searchSourceBuilder.writeTo(output);
+                try (StreamInput in = new NamedWriteableAwareStreamInput(output.bytes().streamInput(), namedWriteableRegistry)) {
+                    SearchSourceBuilder deserializedBuilder = new SearchSourceBuilder(in);
+                    String actualContent = deserializedBuilder.toString();
+                    assertEquals(expectedContent, actualContent);
+                    assertEquals(searchSourceBuilder.hashCode(), deserializedBuilder.hashCode());
+                    assertNotSame(searchSourceBuilder, deserializedBuilder);
+                }
+            }
+        }
+    }
+
     public void testAggsParsing() throws IOException {
         {
             String restContent = "{\n"

diff --git a/server/src/test/java/org/opensearch/search/pipeline/SearchPipelineServiceTests.java b/server/src/test/java/org/opensearch/search/pipeline/SearchPipelineServiceTests.java
@@ -969,6 +969,64 @@ public void testInlinePipeline() throws Exception {
         }
     }
 
+    public void testInlineDefinedPipeline() throws Exception {
+        SearchPipelineService searchPipelineService = createWithProcessors();
+
+        SearchPipelineMetadata metadata = new SearchPipelineMetadata(
+            Map.of(
+                "p1",
+                new PipelineConfiguration(
+                    "p1",
+                    new BytesArray(
+                        "{"
+                            + "\"request_processors\": [{ \"scale_request_size\": { \"scale\" : 2 } }],"
+                            + "\"response_processors\": [{ \"fixed_score\": { \"score\" : 2 } }]"
+                            + "}"
+                    ),
+                    MediaTypeRegistry.JSON
+                )
+            )
+
+        );
+        ClusterState clusterState = ClusterState.builder(new ClusterName("_name")).build();
+        ClusterState previousState = clusterState;
+        clusterState = ClusterState.builder(clusterState)
+            .metadata(Metadata.builder().putCustom(SearchPipelineMetadata.TYPE, metadata))
+            .build();
+        searchPipelineService.applyClusterState(new ClusterChangedEvent("", clusterState, previousState));
+
+        SearchSourceBuilder sourceBuilder = SearchSourceBuilder.searchSource().size(100).pipeline("p1");
+        SearchRequest searchRequest = new SearchRequest().source(sourceBuilder);
+        searchRequest.pipeline(searchRequest.source().pipeline());
+
+        // Verify pipeline
+        PipelinedRequest pipelinedRequest = syncTransformRequest(
+            searchPipelineService.resolvePipeline(searchRequest, indexNameExpressionResolver)
+        );
+        Pipeline pipeline = pipelinedRequest.getPipeline();
+        assertEquals("p1", pipeline.getId());
+        assertEquals(1, pipeline.getSearchRequestProcessors().size());
+        assertEquals(1, pipeline.getSearchResponseProcessors().size());
+
+        // Verify that pipeline transforms request
+        assertEquals(200, pipelinedRequest.source().size());
+
+        int size = 10;
+        SearchHit[] hits = new SearchHit[size];
+        for (int i = 0; i < size; i++) {
+            hits[i] = new SearchHit(i, "doc" + i, Collections.emptyMap(), Collections.emptyMap());
+            hits[i].score(i);
+        }
+        SearchHits searchHits = new SearchHits(hits, new TotalHits(size * 2, TotalHits.Relation.EQUAL_TO), size);
+        SearchResponseSections searchResponseSections = new SearchResponseSections(searchHits, null, null, false, false, null, 0);
+        SearchResponse searchResponse = new SearchResponse(searchResponseSections, null, 1, 1, 0, 10, null, null);
+
+        SearchResponse transformedResponse = syncTransformResponse(pipelinedRequest, searchResponse);
+        for (int i = 0; i < size; i++) {
+            assertEquals(2.0, transformedResponse.getHits().getHits()[i].getScore(), 0.0001);
+        }
+    }
+
     public void testInfo() {
         SearchPipelineService searchPipelineService = createWithProcessors();
         SearchPipelineInfo info = searchPipelineService.info();