Fix the visit of inner query for NestedQueryBuilder (#14968)

Signed-off-by: zhichao-aws <[email protected]>
opensearch-project · Jul 25, 2024 · b21c13e · b21c13e
1 parent d09642c
commit b21c13e
Show file tree

Hide file tree

Showing 2 changed files with 20 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
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.ReaderUtil;
+import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.MultiCollector;
 import org.apache.lucene.search.Query;
@@ -505,4 +506,12 @@ public TopDocsAndMaxScore topDocs(SearchHit hit) throws IOException {
             }
         }
     }
+
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (query != null) {
+            visitor.getChildVisitor(BooleanClause.Occur.MUST).accept(query);
+        }
+    }
 }
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,8 +59,10 @@
 import org.hamcrest.Matchers;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
@@ -565,4 +567,13 @@ void doWithDepth(int depth, ThrowingConsumer<QueryShardContext> test) throws Exc
                 );
         }
     }
+
+    public void testVisit() {
+        NestedQueryBuilder builder = new NestedQueryBuilder("path", new MatchAllQueryBuilder(), ScoreMode.None);
+
+        List<QueryBuilder> visitedQueries = new ArrayList<>();
+        builder.visit(createTestVisitor(visitedQueries));
+
+        assertEquals(2, visitedQueries.size());
+    }
 }