Add IT and Unit test for subqueries count in hybrid clause

Signed-off-by: Varun Jain <[email protected]>
opensearch-project · Oct 28, 2024 · 085710a · 085710a
1 parent 85edf35
commit 085710a
Show file tree

Hide file tree

Showing 2 changed files with 118 additions and 0 deletions.
diff --git 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
@@ -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 {