Pass parent filter to inner hit query (opensearch-project#13903)

Signed-off-by: Heemin Kim <heemin@amazon.com> (cherry picked from commit 6c9603a)
heemin32 · May 31, 2024 · dc76830 · dc76830
1 parent c0581a0
commit dc76830
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -142,6 +142,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix tracing context propagation for local transport instrumentation ([#11490](https://github.com/opensearch-project/OpenSearch/pull/11490))
 - Fix parsing of single line comments in `lang-painless` ([#11815](https://github.com/opensearch-project/OpenSearch/issues/11815))
 - Fix typo in API annotation check message ([11836](https://github.com/opensearch-project/OpenSearch/pull/11836))
+- Pass parent filter to inner hit query ([#13903](https://github.com/opensearch-project/OpenSearch/pull/13903))
 
 ### Security
 

diff --git a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
@@ -391,16 +391,32 @@ protected void doBuild(SearchContext parentSearchContext, InnerHitsContext inner
                 }
             }
             String name = innerHitBuilder.getName() != null ? innerHitBuilder.getName() : nestedObjectMapper.fullPath();
-            ObjectMapper parentObjectMapper = queryShardContext.nestedScope().nextLevel(nestedObjectMapper);
-            NestedInnerHitSubContext nestedInnerHits = new NestedInnerHitSubContext(
-                name,
-                parentSearchContext,
-                parentObjectMapper,
-                nestedObjectMapper
-            );
-            setupInnerHitsContext(queryShardContext, nestedInnerHits);
-            queryShardContext.nestedScope().previousLevel();
-            innerHitsContext.addInnerHitDefinition(nestedInnerHits);
+            ObjectMapper parentObjectMapper = queryShardContext.nestedScope().getObjectMapper();
+            BitSetProducer parentFilter;
+            if (parentObjectMapper == null) {
+                parentFilter = queryShardContext.bitsetFilter(Queries.newNonNestedFilter());
+            } else {
+                parentFilter = queryShardContext.bitsetFilter(parentObjectMapper.nestedTypeFilter());
+            }
+            BitSetProducer previousParentFilter = queryShardContext.getParentFilter();
+            try {
+                queryShardContext.setParentFilter(parentFilter);
+                queryShardContext.nestedScope().nextLevel(nestedObjectMapper);
+                try {
+                    NestedInnerHitSubContext nestedInnerHits = new NestedInnerHitSubContext(
+                        name,
+                        parentSearchContext,
+                        parentObjectMapper,
+                        nestedObjectMapper
+                    );
+                    setupInnerHitsContext(queryShardContext, nestedInnerHits);
+                    innerHitsContext.addInnerHitDefinition(nestedInnerHits);
+                } finally {
+                    queryShardContext.nestedScope().previousLevel();
+                }
+            } finally {
+                queryShardContext.setParentFilter(previousParentFilter);
+            }
         }
     }
 

diff --git a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java
@@ -302,6 +302,43 @@ public void testInlineLeafInnerHitsNestedQuery() throws Exception {
         assertThat(innerHitBuilders.get(leafInnerHits.getName()), Matchers.notNullValue());
     }
 
+    public void testParentFilterFromInlineLeafInnerHitsNestedQuery() throws Exception {
+        QueryShardContext queryShardContext = createShardContext();
+        SearchContext searchContext = mock(SearchContext.class);
+        when(searchContext.getQueryShardContext()).thenReturn(queryShardContext);
+
+        MapperService mapperService = mock(MapperService.class);
+        IndexSettings settings = new IndexSettings(newIndexMeta("index", Settings.EMPTY), Settings.EMPTY);
+        when(mapperService.getIndexSettings()).thenReturn(settings);
+        when(searchContext.mapperService()).thenReturn(mapperService);
+
+        InnerHitBuilder leafInnerHits = randomNestedInnerHits();
+        // Set null for values not related with this test case
+        leafInnerHits.setScriptFields(null);
+        leafInnerHits.setHighlightBuilder(null);
+        leafInnerHits.setSorts(null);
+
+        QueryBuilder innerQueryBuilder = spy(new MatchAllQueryBuilder());
+        when(innerQueryBuilder.toQuery(queryShardContext)).thenAnswer(invoke -> {
+            QueryShardContext context = invoke.getArgument(0);
+            if (context.getParentFilter() == null) {
+                throw new Exception("Expect parent filter to be non-null");
+            }
+            return invoke.callRealMethod();
+        });
+        NestedQueryBuilder query = new NestedQueryBuilder("nested1", innerQueryBuilder, ScoreMode.None);
+        query.innerHit(leafInnerHits);
+        final Map<String, InnerHitContextBuilder> innerHitBuilders = new HashMap<>();
+        final InnerHitsContext innerHitsContext = new InnerHitsContext();
+        query.extractInnerHitBuilders(innerHitBuilders);
+        assertThat(innerHitBuilders.size(), Matchers.equalTo(1));
+        assertTrue(innerHitBuilders.containsKey(leafInnerHits.getName()));
+        assertNull(queryShardContext.getParentFilter());
+        innerHitBuilders.get(leafInnerHits.getName()).build(searchContext, innerHitsContext);
+        assertNull(queryShardContext.getParentFilter());
+        verify(innerQueryBuilder).toQuery(queryShardContext);
+    }
+
     public void testInlineLeafInnerHitsNestedQueryViaBoolQuery() {
         InnerHitBuilder leafInnerHits = randomNestedInnerHits();
         NestedQueryBuilder nestedQueryBuilder = new NestedQueryBuilder("path", new MatchAllQueryBuilder(), ScoreMode.None).innerHit(