Pass parent filter to inner hit query

Signed-off-by: Heemin Kim <heemin@amazon.com>
opensearch-project · May 21, 2024 · 6ee7e12 · 6ee7e12
1 parent 46666c6
commit 6ee7e12
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 10 deletions.
diff --git a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
@@ -401,16 +401,31 @@ 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.getBitsetFilterCache().getBitSetProducer(Queries.newNonNestedFilter());
+            } else {
+                parentFilter = queryShardContext.getBitsetFilterCache().getBitSetProducer(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
@@ -311,6 +311,41 @@ 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();
+        leafInnerHits.setScriptFields(null);
+        leafInnerHits.setHighlightBuilder(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(