Add a flag in QueryShardContext to differentiate between a normal query and an inner hit query (#16600)

Signed-off-by: Heemin Kim <[email protected]>
(cherry picked from commit c9edb48)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: github-actions[bot]

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Switch from `buildSrc/version.properties` to Gradle version catalog (`gradle/libs.versions.toml`) to enable dependabot to perform automated upgrades on common libs ([#16284](https://github.com/opensearch-project/OpenSearch/pull/16284))
 - Add dynamic setting allowing size > 0 requests to be cached in the request cache ([#16483](https://github.com/opensearch-project/OpenSearch/pull/16483))
 - Make IndexStoreListener a pluggable interface ([#16583](https://github.com/opensearch-project/OpenSearch/pull/16583))
+- Add a flag in QueryShardContext to differentiate inner hit query ([#16600](https://github.com/opensearch-project/OpenSearch/pull/16600))
 
 ### Dependencies
 - Bump `com.google.apis:google-api-services-compute` from v1-rev20240407-2.0.0 to v1-rev20241021-2.0.0 ([#16502](https://github.com/opensearch-project/OpenSearch/pull/16502), [#16548](https://github.com/opensearch-project/OpenSearch/pull/16548))

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

@@ -413,6 +413,7 @@ protected void doBuild(SearchContext parentSearchContext, InnerHitsContext inner
             try {
                 queryShardContext.setParentFilter(parentFilter);
                 queryShardContext.nestedScope().nextLevel(nestedObjectMapper);
+                queryShardContext.setInnerHitQuery(true);
                 try {
                     NestedInnerHitSubContext nestedInnerHits = new NestedInnerHitSubContext(
                         name,
@@ -427,6 +428,7 @@ protected void doBuild(SearchContext parentSearchContext, InnerHitsContext inner
                 }
             } finally {
                 queryShardContext.setParentFilter(previousParentFilter);
+                queryShardContext.setInnerHitQuery(false);
             }
         }
     }

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

@@ -127,6 +127,7 @@ public class QueryShardContext extends QueryRewriteContext {
     private BitSetProducer parentFilter;
     private DerivedFieldResolver derivedFieldResolver;
     private boolean keywordIndexOrDocValuesEnabled;
+    private boolean isInnerHitQuery;
 
     public QueryShardContext(
         int shardId,
@@ -738,4 +739,12 @@ public BitSetProducer getParentFilter() {
     public void setParentFilter(BitSetProducer parentFilter) {
         this.parentFilter = parentFilter;
     }
+
+    public boolean isInnerHitQuery() {
+        return isInnerHitQuery;
+    }
+
+    public void setInnerHitQuery(boolean isInnerHitQuery) {
+        this.isInnerHitQuery = isInnerHitQuery;
+    }
 }

server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java

@@ -335,6 +335,9 @@ public void testParentFilterFromInlineLeafInnerHitsNestedQuery() throws Exceptio
             if (context.getParentFilter() == null) {
                 throw new Exception("Expect parent filter to be non-null");
             }
+            if (context.isInnerHitQuery() == false) {
+                throw new Exception("Expect it to be inner hit query");
+            }
             return invoke.callRealMethod();
         });
         NestedQueryBuilder query = new NestedQueryBuilder("nested1", innerQueryBuilder, ScoreMode.None);
@@ -345,6 +348,7 @@ public void testParentFilterFromInlineLeafInnerHitsNestedQuery() throws Exceptio
         assertThat(innerHitBuilders.size(), Matchers.equalTo(1));
         assertTrue(innerHitBuilders.containsKey(leafInnerHits.getName()));
         assertNull(queryShardContext.getParentFilter());
+        assertFalse(queryShardContext.isInnerHitQuery());
         innerHitBuilders.get(leafInnerHits.getName()).build(searchContext, innerHitsContext);
         assertNull(queryShardContext.getParentFilter());
         verify(innerQueryBuilder).toQuery(queryShardContext);