Pass parent filter to inner query in nested query

Pass parent filter to inner query so that inner query can utilize the information Signed-off-by: Heemin Kim <heemin@amazon.com>
opensearch-project · Sep 26, 2023 · c8fe10c · c8fe10c
1 parent 8807d7a
commit c8fe10c
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 0 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
@@ -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);
+    }
 }