From 7f16389b44de49ce35e4bb87db19a1c381508ef6 Mon Sep 17 00:00:00 2001
From: Heemin Kim <heemin@amazon.com>
Date: Fri, 8 Nov 2024 11:07:57 -0800
Subject: [PATCH] Add a flag in QueryShardContext to differentiate between a
 normal query and an inner hit query

Signed-off-by: Heemin Kim <heemin@amazon.com>
---
 CHANGELOG.md                                             | 1 +
 .../org/opensearch/index/query/NestedQueryBuilder.java   | 2 ++
 .../org/opensearch/index/query/QueryShardContext.java    | 9 +++++++++
 .../opensearch/index/query/NestedQueryBuilderTests.java  | 4 ++++
 4 files changed, 16 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 60535b2cca895..3c8a4ab3ccb42 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,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.azure:azure-storage-common` from 12.25.1 to 12.27.1 ([#16521](https://github.com/opensearch-project/OpenSearch/pull/16521))
diff --git a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
index 5908882472ce7..ec7e62035a82f 100644
--- a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
+++ b/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);
             }
         }
     }
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
index bccead2b029d0..d717f10b17d9c 100644
--- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
+++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -126,6 +126,7 @@ public class QueryShardContext extends QueryRewriteContext {
     private BitSetProducer parentFilter;
     private DerivedFieldResolver derivedFieldResolver;
     private boolean keywordIndexOrDocValuesEnabled;
+    private boolean isInnerHitQuery;
 
     public QueryShardContext(
         int shardId,
@@ -727,4 +728,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;
+    }
 }
diff --git a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java
index 351011eb1b812..c367d123402d4 100644
--- a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java
+++ b/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);