diff --git a/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java b/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java index b3c797f11de6d..79ff7da1a36ba 100644 --- a/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java +++ b/server/src/main/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilder.java @@ -32,6 +32,7 @@ package org.opensearch.index.query.functionscore; +import org.apache.lucene.search.BooleanClause; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.Query; import org.opensearch.common.Nullable; @@ -52,6 +53,7 @@ import org.opensearch.index.query.MatchAllQueryBuilder; import org.opensearch.index.query.MatchNoneQueryBuilder; import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.query.QueryBuilderVisitor; import org.opensearch.index.query.QueryRewriteContext; import org.opensearch.index.query.QueryShardContext; @@ -704,4 +706,12 @@ private static String parseFiltersAndFunctions( } return currentFieldName; } + + @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/functionscore/FunctionScoreQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java index 4e64a1ec03688..78f4163d96d33 100644 --- a/server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java +++ b/server/src/test/java/org/opensearch/index/query/functionscore/FunctionScoreQueryBuilderTests.java @@ -75,6 +75,7 @@ import org.hamcrest.Matcher; import java.io.IOException; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.Collections; @@ -938,4 +939,14 @@ public void testMustRewrite() throws IOException { e = expectThrows(IllegalStateException.class, () -> functionQueryBuilder2.toQuery(context)); assertEquals("Rewrite first", e.getMessage()); } + + public void testVisit() { + TermQueryBuilder termQueryBuilder = new TermQueryBuilder("unmapped_field", "foo"); + FunctionScoreQueryBuilder builder = new FunctionScoreQueryBuilder(termQueryBuilder); + + List visitedQueries = new ArrayList<>(); + builder.visit(createTestVisitor(visitedQueries)); + + assertEquals(2, visitedQueries.size()); + } }