From b1e815d2625a3d0129d7a527f2f2d8182b9dc8df Mon Sep 17 00:00:00 2001 From: Owais Kazi Date: Wed, 25 Sep 2024 18:16:59 -0700 Subject: [PATCH] Added support for msearch API to pass search pipeline name (#15923) * Added support for search pipeline name in multi search API Signed-off-by: Owais * Updated CHANGELOG Signed-off-by: Owais * Pulled search pipeline in MultiSearchRequest and updated test Signed-off-by: Owais * Updated test Signed-off-by: Owais * Updated SearchRequest with search pipeline from source Signed-off-by: Owais * Added tests for parseSearchRequest Signed-off-by: Owais * Guard serialization with version check Signed-off-by: Owais * Updated version and added another test for serialization Signed-off-by: Owais --------- Signed-off-by: Owais --- CHANGELOG.md | 1 + .../action/search/MultiSearchRequest.java | 4 ++ .../rest/action/search/RestSearchAction.java | 2 +- .../search/builder/SearchSourceBuilder.java | 36 +++++++++++- .../action/search/SearchRequestTests.java | 24 ++++++-- .../builder/SearchSourceBuilderTests.java | 21 +++++++ .../pipeline/SearchPipelineServiceTests.java | 58 +++++++++++++++++++ 7 files changed, 139 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0e6c09177c859..2e76a5f0816d8 100644 --- 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 index 5b887b48f696e..f16d7d1e7d6a3 100644 --- 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 index 3a6b45013e892..05465e32631fd 100644 --- 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 index 8a9704b04566f..dd4e4d073cb1b 100644 --- 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 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 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 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 index 40514c526f190..acda1445bacbb 100644 --- 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 index 9697f4cee0d58..da8ccc9e121e0 100644 --- 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 index f5857922fdff2..b52205996f34b 100644 --- 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();