Add a flag in QueryShardContext to differentiate between a normal que…

…ry and an inner hit query (#16600) (#16620) (cherry picked from commit c9edb48) Signed-off-by: Heemin Kim <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: Michael Froh <[email protected]> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Michael Froh <[email protected]>
opensearch-project · Nov 13, 2024 · 77231fb · 77231fb
1 parent 88c7b20
commit 77231fb
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Increase segrep pressure checkpoint default limit to 30 ([#16577](https://github.com/opensearch-project/OpenSearch/pull/16577/files))
 - 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))
 - Add vertical scaling and SoftReference for snapshot repository data cache ([#16489](https://github.com/opensearch-project/OpenSearch/pull/16489))
 - Add new configuration setting `synonym_analyzer`, to the `synonym` and `synonym_graph` filters, enabling the specification of a custom analyzer for reading the synonym file ([#16488](https://github.com/opensearch-project/OpenSearch/pull/16488)).
 

diff --git 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
@@ -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;
+    }
 }
diff --git 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);