Optimize innerhits query performance

Signed-off-by: kkewwei <kewei.11@bytedance.com>
Signed-off-by: kkewwei <kkewwei@163.com>

Author: kkewwei

CHANGELOG.md

@@ -72,6 +72,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Allow extended plugins to be optional ([#16909](https://github.com/opensearch-project/OpenSearch/pull/16909))
 - Use the correct type to widen the sort fields when merging top docs ([#16881](https://github.com/opensearch-project/OpenSearch/pull/16881))
 - Limit reader writer separation to remote store enabled clusters [#16760](https://github.com/opensearch-project/OpenSearch/pull/16760)
+- Optimize innerhits query performance [#16937](https://github.com/opensearch-project/OpenSearch/pull/16937)
 
 ### Deprecated
 - Performing update operation with default pipeline or final pipeline is deprecated ([#16712](https://github.com/opensearch-project/OpenSearch/pull/16712))

server/src/main/java/org/opensearch/index/mapper/DocumentMapper.java

@@ -35,9 +35,7 @@
 import org.apache.lucene.document.StoredField;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Query;
-import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.Scorer;
-import org.apache.lucene.search.Weight;
+import org.apache.lucene.util.BitSet;
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.OpenSearchGenerationException;
 import org.opensearch.common.annotation.PublicApi;
@@ -53,6 +51,7 @@
 import org.opensearch.index.analysis.IndexAnalyzers;
 import org.opensearch.index.mapper.MapperService.MergeReason;
 import org.opensearch.index.mapper.MetadataFieldMapper.TypeParser;
+import org.opensearch.index.query.NestedQueryBuilder;
 import org.opensearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -270,25 +269,15 @@ public ParsedDocument createNoopTombstoneDoc(String index, String reason) throws
      * Returns the best nested {@link ObjectMapper} instances that is in the scope of the specified nested docId.
      */
     public ObjectMapper findNestedObjectMapper(int nestedDocId, SearchContext sc, LeafReaderContext context) throws IOException {
+        if (sc instanceof NestedQueryBuilder.NestedInnerHitSubContext) {
+            ObjectMapper objectMapper = ((NestedQueryBuilder.NestedInnerHitSubContext) sc).getChildObjectMapper();
+            assert objectMappers().containsKey(objectMapper.fullPath());
+            assert containSubDocIdWithObjectMapper(nestedDocId, objectMapper, sc, context);
+            return objectMapper;
+        }
         ObjectMapper nestedObjectMapper = null;
         for (ObjectMapper objectMapper : objectMappers().values()) {
-            if (!objectMapper.nested().isNested()) {
-                continue;
-            }
-
-            Query filter = objectMapper.nestedTypeFilter();
-            if (filter == null) {
-                continue;
-            }
-            // We can pass down 'null' as acceptedDocs, because nestedDocId is a doc to be fetched and
-            // therefore is guaranteed to be a live doc.
-            final Weight nestedWeight = filter.createWeight(sc.searcher(), ScoreMode.COMPLETE_NO_SCORES, 1f);
-            Scorer scorer = nestedWeight.scorer(context);
-            if (scorer == null) {
-                continue;
-            }
-
-            if (scorer.iterator().advance(nestedDocId) == nestedDocId) {
+            if (containSubDocIdWithObjectMapper(nestedDocId, objectMapper, sc, context)) {
                 if (nestedObjectMapper == null) {
                     nestedObjectMapper = objectMapper;
                 } else {
@@ -301,6 +290,25 @@ public ObjectMapper findNestedObjectMapper(int nestedDocId, SearchContext sc, Le
         return nestedObjectMapper;
     }
 
+    private boolean containSubDocIdWithObjectMapper(int nestedDocId, ObjectMapper objectMapper, SearchContext sc, LeafReaderContext context)
+        throws IOException {
+        if (!objectMapper.nested().isNested()) {
+            return false;
+        }
+        Query filter = objectMapper.nestedTypeFilter();
+        if (filter == null) {
+            return false;
+        }
+        // We can pass down 'null' as acceptedDocs, because nestedDocId is a doc to be fetched and
+        // therefore is guaranteed to be a live doc.
+        BitSet nestedDocIds = sc.bitsetFilterCache().getBitSetProducer(filter).getBitSet(context);
+        if (nestedDocIds != null && nestedDocIds.get(nestedDocId)) {
+            return true;
+        } else {
+            return false;
+        }
+    }
+
     public DocumentMapper merge(Mapping mapping, MergeReason reason) {
         Mapping merged = this.mapping.merge(mapping, reason);
         return new DocumentMapper(mapperService, merged);

server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java

@@ -438,7 +438,7 @@ protected void doBuild(SearchContext parentSearchContext, InnerHitsContext inner
      *
      * @opensearch.internal
      */
-    static final class NestedInnerHitSubContext extends InnerHitsContext.InnerHitSubContext {
+    public static final class NestedInnerHitSubContext extends InnerHitsContext.InnerHitSubContext {
 
         private final ObjectMapper parentObjectMapper;
         private final ObjectMapper childObjectMapper;
@@ -507,6 +507,10 @@ public TopDocsAndMaxScore topDocs(SearchHit hit) throws IOException {
                 return new TopDocsAndMaxScore(td, maxScore);
             }
         }
+
+        public ObjectMapper getChildObjectMapper() {
+            return childObjectMapper;
+        }
     }
 
     @Override

server/src/main/java/org/opensearch/search/fetch/FetchPhase.java

@@ -38,10 +38,7 @@
 import org.apache.lucene.index.ReaderUtil;
 import org.apache.lucene.search.DocIdSetIterator;
 import org.apache.lucene.search.Query;
-import org.apache.lucene.search.ScoreMode;
-import org.apache.lucene.search.Scorer;
 import org.apache.lucene.search.TotalHits;
-import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.BitSet;
 import org.opensearch.common.CheckedBiConsumer;
 import org.opensearch.common.annotation.PublicApi;
@@ -55,7 +52,6 @@
 import org.opensearch.core.common.text.Text;
 import org.opensearch.core.tasks.TaskCancelledException;
 import org.opensearch.core.xcontent.MediaType;
-import org.opensearch.index.IndexSettings;
 import org.opensearch.index.fieldvisitor.CustomFieldsVisitor;
 import org.opensearch.index.fieldvisitor.FieldsVisitor;
 import org.opensearch.index.mapper.DocumentMapper;
@@ -501,7 +497,6 @@ private SearchHit.NestedIdentity getInternalNestedIdentity(
         ObjectMapper current = nestedObjectMapper;
         String originalName = nestedObjectMapper.name();
         SearchHit.NestedIdentity nestedIdentity = null;
-        final IndexSettings indexSettings = context.getQueryShardContext().getIndexSettings();
         do {
             Query parentFilter;
             nestedParentObjectMapper = current.getParentObjectMapper(mapperService);
@@ -520,14 +515,13 @@ private SearchHit.NestedIdentity getInternalNestedIdentity(
                 current = nestedParentObjectMapper;
                 continue;
             }
-            final Weight childWeight = context.searcher()
-                .createWeight(context.searcher().rewrite(childFilter), ScoreMode.COMPLETE_NO_SCORES, 1f);
-            Scorer childScorer = childWeight.scorer(subReaderContext);
-            if (childScorer == null) {
+            BitSet childIter = context.bitsetFilterCache()
+                .getBitSetProducer(context.searcher().rewrite(childFilter))
+                .getBitSet(subReaderContext);
+            if (childIter == null) {
                 current = nestedParentObjectMapper;
                 continue;
             }
-            DocIdSetIterator childIter = childScorer.iterator();
 
             BitSet parentBits = context.bitsetFilterCache().getBitSetProducer(parentFilter).getBitSet(subReaderContext);
 
@@ -541,8 +535,8 @@ private SearchHit.NestedIdentity getInternalNestedIdentity(
              * that appear before him.
              */
             int previousParent = parentBits.prevSetBit(currentParent);
-            for (int docId = childIter.advance(previousParent + 1); docId < nestedSubDocId
-                && docId != DocIdSetIterator.NO_MORE_DOCS; docId = childIter.nextDoc()) {
+            for (int docId = childIter.nextSetBit(previousParent + 1); docId < nestedSubDocId
+                && docId != DocIdSetIterator.NO_MORE_DOCS; docId = childIter.nextSetBit(docId + 1)) {
                 offset++;
             }
             currentParent = nestedSubDocId;