diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java index 2a6fa49a3..5aab7581c 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java @@ -52,6 +52,7 @@ import org.opensearch.core.xcontent.XContentParser; import org.opensearch.index.mapper.TextFieldMapper; import org.opensearch.index.query.MatchAllQueryBuilder; +import org.opensearch.index.query.MatchQueryBuilder; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryBuilders; import org.opensearch.index.query.QueryShardContext; @@ -401,6 +402,80 @@ public void testFromXContent_whenIncorrectFormat_thenFail() { expectThrows(ParsingException.class, () -> HybridQueryBuilder.fromXContent(contentParser2)); } + @SneakyThrows + public void testFromXContent_whenQueriesCountIsGreaterThanFive_thenFail() { + setUpClusterService(); + XContentBuilder xContentBuilder = XContentFactory.jsonBuilder() + .startObject() + .startArray("queries") + .startObject() + .startObject(NeuralQueryBuilder.NAME) + .startObject(VECTOR_FIELD_NAME) + .field(QUERY_TEXT_FIELD.getPreferredName(), QUERY_TEXT) + .field(MODEL_ID_FIELD.getPreferredName(), MODEL_ID) + .field(K_FIELD.getPreferredName(), K) + .field(BOOST_FIELD.getPreferredName(), BOOST) + .endObject() + .endObject() + .endObject() + .startObject() + .startObject(TermQueryBuilder.NAME) + .field(TEXT_FIELD_NAME, TERM_QUERY_TEXT) + .endObject() + .endObject() + .startObject() + .startObject(MatchQueryBuilder.NAME) + .field(QUERY_TEXT_FIELD.getPreferredName(), QUERY_TEXT) + .endObject() + .endObject() + .startObject() + .startObject(TermQueryBuilder.NAME) + .field(MODEL_ID_FIELD.getPreferredName(), MODEL_ID) + .endObject() + .endObject() + .startObject() + .startObject(MatchAllQueryBuilder.NAME) + .endObject() + .endObject() + .startObject() + .startObject(TermQueryBuilder.NAME) + .field(MODEL_ID_FIELD.getPreferredName(), MODEL_ID) + .endObject() + .endObject() + .endArray() + .endObject(); + + NamedXContentRegistry namedXContentRegistry = new NamedXContentRegistry( + List.of( + new NamedXContentRegistry.Entry(QueryBuilder.class, new ParseField(TermQueryBuilder.NAME), TermQueryBuilder::fromXContent), + new NamedXContentRegistry.Entry( + QueryBuilder.class, + new ParseField(NeuralQueryBuilder.NAME), + NeuralQueryBuilder::fromXContent + ), + new NamedXContentRegistry.Entry( + QueryBuilder.class, + new ParseField(HybridQueryBuilder.NAME), + HybridQueryBuilder::fromXContent + ), + new NamedXContentRegistry.Entry( + QueryBuilder.class, + new ParseField(MatchAllQueryBuilder.NAME), + MatchAllQueryBuilder::fromXContent + ), + new NamedXContentRegistry.Entry(QueryBuilder.class, new ParseField(MatchQueryBuilder.NAME), MatchQueryBuilder::fromXContent) + ) + ); + XContentParser contentParser = createParser( + namedXContentRegistry, + xContentBuilder.contentType().xContent(), + BytesReference.bytes(xContentBuilder) + ); + contentParser.nextToken(); + ParsingException exception = expectThrows(ParsingException.class, () -> HybridQueryBuilder.fromXContent(contentParser)); + assertThat(exception.getMessage(), containsString("Number of sub-queries exceeds maximum supported by [hybrid] query")); + } + @SneakyThrows public void testToXContent_whenIncomingJsonIsCorrect_thenSuccessful() { HybridQueryBuilder queryBuilder = new HybridQueryBuilder(); diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java index e20dd5e01..a0d6dcd1a 100644 --- a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java +++ b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java @@ -795,6 +795,49 @@ public void testConcurrentSearchWithMultipleSlices_whenMultipleShardsIndex_thenS } } + @SneakyThrows + public void testPaginationDepth_whenSubqueriesCountIsGreaterThanFive_thenFail() { + initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_NAME); + createSearchPipelineWithResultsPostProcessor(SEARCH_PIPELINE); + TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); + TermQueryBuilder termQuery2Builder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2); + MatchAllQueryBuilder matchAllQueryBuilder = QueryBuilders.matchAllQuery(); + MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT); + MatchQueryBuilder matchQueryBuilder1 = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2); + MatchQueryBuilder matchQueryBuilder2 = QueryBuilders.matchQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3); + HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder(); + hybridQueryBuilder.add(termQueryBuilder); + hybridQueryBuilder.add(termQuery2Builder); + hybridQueryBuilder.add(matchAllQueryBuilder); + hybridQueryBuilder.add(matchQueryBuilder); + hybridQueryBuilder.add(matchQueryBuilder1); + hybridQueryBuilder.add(matchQueryBuilder2); + hybridQueryBuilder.paginationDepth(10); + + ResponseException responseException = assertThrows( + ResponseException.class, + () -> search( + TEST_MULTI_DOC_INDEX_NAME, + hybridQueryBuilder, + null, + 10, + Map.of("search_pipeline", SEARCH_PIPELINE), + null, + null, + null, + false, + null, + 0 + ) + ); + + org.hamcrest.MatcherAssert.assertThat( + responseException.getMessage(), + allOf(containsString("Number of sub-queries exceeds maximum supported by [hybrid] query")) + ); + + } + @SneakyThrows public void testPaginationOnSingleShard_whenConcurrentSearchEnabled_thenSuccessful() { try {