Add inner hits support to hybrid query

CHANGELOG.md

@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Optimize parameter parsing in text chunking processor ([#733](https://github.com/opensearch-project/neural-search/pull/733))
 - Use lazy initialization for priority queue of hits and scores to improve latencies by 20% ([#746](https://github.com/opensearch-project/neural-search/pull/746))
 - Optimize max score calculation in the Query Phase of the Hybrid Search ([765](https://github.com/opensearch-project/neural-search/pull/765))
+- Add inner hits support to hybrid query ([#776](https://github.com/opensearch-project/neural-search/pull/776))
 ### Bug Fixes
 - Total hit count fix in Hybrid Query ([756](https://github.com/opensearch-project/neural-search/pull/756))
 - Fix map type validation issue in multiple pipeline processors ([#661](https://github.com/opensearch-project/neural-search/pull/661))

src/main/java/org/opensearch/neuralsearch/query/HybridQueryBuilder.java

@@ -9,6 +9,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Collectors;
 
@@ -23,6 +24,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.query.AbstractQueryBuilder;
+import org.opensearch.index.query.InnerHitContextBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryRewriteContext;
 import org.opensearch.index.query.QueryShardContext;
@@ -310,4 +312,18 @@ public void visit(QueryBuilderVisitor visitor) {
             subQueryBuilder.visit(subVisitor);
         }
     }
+
+    /**
+     * Extracts the inner hits from the query tree
+     * @param innerHits
+     */
+    @Override
+    protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> innerHits) {
+        if (Objects.isNull(queries)) {
+            return;
+        }
+        for (QueryBuilder query : queries) {
+            InnerHitContextBuilder.extractInnerHits(query, innerHits);
+        }
+    }
 }

src/test/java/org/opensearch/neuralsearch/query/HybridQueryBuilderTests.java

@@ -17,6 +17,7 @@
 import static org.opensearch.neuralsearch.query.NeuralQueryBuilder.MODEL_ID_FIELD;
 import static org.opensearch.neuralsearch.query.NeuralQueryBuilder.QUERY_TEXT_FIELD;
 
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -42,7 +43,10 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.mapper.TextFieldMapper;
+import org.opensearch.index.query.BoolQueryBuilder;
+import org.opensearch.index.query.InnerHitContextBuilder;
 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;
@@ -712,13 +716,28 @@ public void testBoost_whenDefaultBoostSet_thenBuildSuccessfully() {
         assertNotNull(hybridQueryBuilder);
     }
 
-    public void testVisit() {
+    public void testVisit_whenMultipleSubqueries_thenSuccessful() {
         HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder().add(new NeuralQueryBuilder()).add(new NeuralSparseQueryBuilder());
         List<QueryBuilder> visitedQueries = new ArrayList<>();
         hybridQueryBuilder.visit(createTestVisitor(visitedQueries));
         assertEquals(3, visitedQueries.size());
     }
 
+    public void testInnerHits_whenSingleSubquery_thenSuccessful() {
+        HybridQueryBuilder hybridQueryBuilderNeuralQuery = new HybridQueryBuilder().add(new NeuralQueryBuilder());
+        Map<String, InnerHitContextBuilder> innerHitsBuilderDefault = new HashMap<>();
+        hybridQueryBuilderNeuralQuery.extractInnerHitBuilders(innerHitsBuilderDefault);
+        assertNotNull(innerHitsBuilderDefault);
+    }
+
+    public void testInnerHits_whenMultipleSubqueries_thenSuccessful() {
+        HybridQueryBuilder hybridQueryBuilder = new HybridQueryBuilder().add(new BoolQueryBuilder())
+            .add(new MatchQueryBuilder(TEXT_FIELD_NAME, TERM_QUERY_TEXT));
+        Map<String, InnerHitContextBuilder> innerHitsBuilder = new HashMap<>();
+        hybridQueryBuilder.extractInnerHitBuilders(innerHitsBuilder);
+        assertNotNull(innerHitsBuilder);
+    }
+
     private Map<String, Object> getInnerMap(Object innerObject, String queryName, String fieldName) {
         if (!(innerObject instanceof Map)) {
             fail("field name does not map to nested object");