Pass parent filter to inner query in nested query

heemin32 · heemin32 · commit 99fccbf2075d · 2023-09-26T16:47:29.000-07:00
Pass parent filter to inner query so that inner query can utilize the information

Signed-off-by: Heemin Kim &lt;heemin@amazon.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -112,6 +112,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Force merge with `only_expunge_deletes` honors max segment size ([#10036](https://github.com/opensearch-project/OpenSearch/pull/10036))
 - Add instrumentation in transport service. ([#10042](https://github.com/opensearch-project/OpenSearch/pull/10042))
 - [Tracing Framework] Add support for SpanKind. ([#10122](https://github.com/opensearch-project/OpenSearch/pull/10122))
+- Pass parent filter to inner query in nested query ([#10246](https://github.com/opensearch-project/OpenSearch/pull/10246))
 
 ### Deprecated
 
diff --git a/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/NestedQueryBuilder.java
@@ -318,10 +318,13 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
             parentFilter = context.bitsetFilter(objectMapper.nestedTypeFilter());
         }
 
+        BitSetProducer previousParentFilter = context.getParentFilter();
         try {
+            context.setParentFilter(parentFilter);
             context.nestedScope().nextLevel(nestedObjectMapper);
             innerQuery = this.query.toQuery(context);
         } finally {
+            context.setParentFilter(previousParentFilter);
             context.nestedScope().previousLevel();
         }
 
diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java
@@ -115,6 +115,7 @@ public class QueryShardContext extends QueryRewriteContext {
     private boolean mapUnmappedFieldAsString;
     private NestedScope nestedScope;
     private final ValuesSourceRegistry valuesSourceRegistry;
+    private BitSetProducer parentFilter;
 
     public QueryShardContext(
         int shardId,
@@ -622,4 +623,12 @@ public BitsetFilterCache getBitsetFilterCache() {
     public AggregationUsageService getUsageService() {
         return valuesSourceRegistry.getUsageService();
     }
+
+    public BitSetProducer getParentFilter() {
+        return parentFilter;
+    }
+
+    public void setParentFilter(BitSetProducer parentFilter) {
+        this.parentFilter = parentFilter;
+    }
 }
diff --git a/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/NestedQueryBuilderTests.java
@@ -59,6 +59,8 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 import static org.opensearch.index.IndexSettingsTests.newIndexMeta;
 import static org.opensearch.index.query.InnerHitBuilderTests.randomNestedInnerHits;
 import static org.hamcrest.CoreMatchers.containsString;
@@ -411,4 +413,22 @@ public void testDisallowExpensiveQueries() {
         OpenSearchException e = expectThrows(OpenSearchException.class, () -> queryBuilder.toQuery(queryShardContext));
         assertEquals("[joining] queries cannot be executed when 'search.allow_expensive_queries' is set to false.", e.getMessage());
     }
+
+    public void testSetParentFilterInContext() throws Exception {
+        QueryShardContext queryShardContext = createShardContext();
+        QueryBuilder innerQueryBuilder = spy(new MatchAllQueryBuilderTests().createTestQueryBuilder());
+        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 nqb = new NestedQueryBuilder("nested1", innerQueryBuilder, RandomPicks.randomFrom(random(), ScoreMode.values()));
+
+        assertNull(queryShardContext.getParentFilter());
+        nqb.rewrite(queryShardContext).toQuery(queryShardContext);
+        assertNull(queryShardContext.getParentFilter());
+        verify(innerQueryBuilder).toQuery(queryShardContext);
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -318,10 +318,13 @@ protected Query doToQuery(QueryShardContext context) throws IOException {`
`318`	`318`	`parentFilter = context.bitsetFilter(objectMapper.nestedTypeFilter());`
`319`	`319`	`}`
`320`	`320`
	`321`	`+ BitSetProducer previousParentFilter = context.getParentFilter();`
`321`	`322`	`try {`
	`323`	`+ context.setParentFilter(parentFilter);`
`322`	`324`	`context.nestedScope().nextLevel(nestedObjectMapper);`
`323`	`325`	`innerQuery = this.query.toQuery(context);`
`324`	`326`	`} finally {`
	`327`	`+ context.setParentFilter(previousParentFilter);`
`325`	`328`	`context.nestedScope().previousLevel();`
`326`	`329`	`}`
`327`	`330`