Addressing review comments

Signed-off-by: Martin Gaievski <[email protected]>
opensearch-project · Nov 29, 2023 · 19791bb · 19791bb
1 parent b8f794b
commit 19791bb
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java
@@ -53,7 +53,7 @@
 @Log4j2
 public class HybridQueryPhaseSearcher extends QueryPhaseSearcherWrapper {
 
-    final static int MAX_NESTED_SUBQUERY_LIMIT = 20;
+    final static int MAX_NESTED_SUBQUERY_LIMIT = 50;
 
     public HybridQueryPhaseSearcher() {
         super();
@@ -130,7 +130,11 @@ && mightBeWrappedHybridQuery(query)
             && ((BooleanQuery) query).clauses().size() > 0) {
             // extract hybrid query and replace bool with hybrid query
             List<BooleanClause> booleanClauses = ((BooleanQuery) query).clauses();
-            return booleanClauses.stream().findFirst().get().getQuery();
+            Query hybridQuery = booleanClauses.stream().findFirst().get().getQuery();
+            if (!(hybridQuery instanceof HybridQuery)) {
+                throw new IllegalStateException("cannot find hybrid type query in expected location");
+            }
+            return hybridQuery;
         }
         return query;
     }

diff --git a/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java b/src/test/java/org/opensearch/neuralsearch/query/HybridQueryIT.java
@@ -254,7 +254,7 @@ public void testNestedQuery_whenHybridQueryIsWrappedIntoOtherQuery_thenFail() {
     public void testIndexWithNestedFields_whenHybridQuery_thenSuccess() {
         initializeIndexIfNotExist(TEST_MULTI_DOC_INDEX_WITH_NESTED_TYPE_NAME_ONE_SHARD);
 
-        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT);
+        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT3);
         TermQueryBuilder termQuery2Builder = QueryBuilders.termQuery(TEST_TEXT_FIELD_NAME_1, TEST_QUERY_TEXT2);
         HybridQueryBuilder hybridQueryBuilderOnlyTerm = new HybridQueryBuilder();
         hybridQueryBuilderOnlyTerm.add(termQueryBuilder);
@@ -268,13 +268,13 @@ public void testIndexWithNestedFields_whenHybridQuery_thenSuccess() {
             Map.of("search_pipeline", SEARCH_PIPELINE)
         );
 
-        assertEquals(0, getHitCount(searchResponseAsMap));
+        assertEquals(1, getHitCount(searchResponseAsMap));
         assertTrue(getMaxScore(searchResponseAsMap).isPresent());
-        assertEquals(0.0f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION);
+        assertEquals(0.5f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION);
 
         Map<String, Object> total = getTotalHits(searchResponseAsMap);
         assertNotNull(total.get("value"));
-        assertEquals(0, total.get("value"));
+        assertEquals(1, total.get("value"));
         assertNotNull(total.get("relation"));
         assertEquals(RELATION_EQUAL_TO, total.get("relation"));
     }
@@ -303,7 +303,7 @@ public void testIndexWithNestedFields_whenHybridQueryIncludesNested_thenSuccess(
 
         assertEquals(1, getHitCount(searchResponseAsMap));
         assertTrue(getMaxScore(searchResponseAsMap).isPresent());
-        assertTrue(getMaxScore(searchResponseAsMap).get() > 0);
+        assertEquals(0.5f, getMaxScore(searchResponseAsMap).get(), DELTA_FOR_SCORE_ASSERTION);
 
         Map<String, Object> total = getTotalHits(searchResponseAsMap);
         assertNotNull(total.get("value"));