Add NestedField Visit

Signed-off-by: luyuncheng <[email protected]>
opensearch-project · Jun 4, 2024 · e72e9b3 · e72e9b3
1 parent 581fcd2
commit e72e9b3
Show file tree

Hide file tree

Showing 2 changed files with 24 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;
@@ -371,6 +372,14 @@ public void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> innerHit
         }
     }
 
+    @Override
+    public void visit(QueryBuilderVisitor visitor) {
+        visitor.accept(this);
+        if (query != null) {
+            query.visit(visitor.getChildVisitor(BooleanClause.Occur.MUST));
+        }
+    }
+
     /**
      * Context builder for nested inner hits
      *

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;
 
@@ -540,6 +542,19 @@ public void testUpdateMaxDepthSettings() throws Exception {
         });
     }
 
+    public void testVisit() {
+        NestedQueryBuilder queryBuilder = new NestedQueryBuilder(
+            "nested1",
+            new BoolQueryBuilder().must(new TermQueryBuilder("must_field1", "value1"))
+                .filter(new TermQueryBuilder("must_field2", "value2")),
+            ScoreMode.None
+        );
+        List<QueryBuilder> visitorQueries = new ArrayList<>();
+        queryBuilder.visit(createTestVisitor(visitorQueries));
+
+        assertEquals(4, visitorQueries.size());
+    }
+
     void doWithDepth(int depth, ThrowingConsumer<QueryShardContext> test) throws Exception {
         QueryShardContext context = createShardContext();
         int defLimit = context.getIndexSettings().getMaxNestedQueryDepth();