Skip to content

Commit

Permalink
Script Score Query Builder Visit Method (#10198)
Browse files Browse the repository at this point in the history
* Script Score Builder Visit Method

Signed-off-by: Varun Jain <[email protected]>

* ScriptScore Query Builder Visit Method

Signed-off-by: Varun Jain <[email protected]>

* gradlew spotlessApply

Signed-off-by: Varun Jain <[email protected]>

---------

Signed-off-by: Varun Jain <[email protected]>
  • Loading branch information
vibrantvarun authored Sep 25, 2023
1 parent f6646fb commit a8969cb
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@

package org.opensearch.index.query.functionscore;

import org.apache.lucene.search.BooleanClause;
import org.apache.lucene.search.Query;
import org.opensearch.OpenSearchException;
import org.opensearch.common.lucene.search.function.ScriptScoreQuery;
Expand All @@ -45,6 +46,7 @@
import org.opensearch.index.query.InnerHitContextBuilder;
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;
import org.opensearch.script.ScoreScript;
Expand Down Expand Up @@ -224,4 +226,12 @@ protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> inner
InnerHitContextBuilder.extractInnerHits(query(), innerHits);
}

@Override
public void visit(QueryBuilderVisitor visitor) {
visitor.accept(this);
if (query != null) {
QueryBuilderVisitor subVisitor = visitor.getChildVisitor(BooleanClause.Occur.MUST);
subVisitor.accept(query);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,9 @@
import org.opensearch.test.AbstractQueryTestCase;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
import static org.hamcrest.CoreMatchers.instanceOf;
Expand Down Expand Up @@ -140,4 +142,11 @@ public void testDisallowExpensiveQueries() {
OpenSearchException e = expectThrows(OpenSearchException.class, () -> queryBuilder.toQuery(queryShardContext));
assertEquals("[script score] queries cannot be executed when 'search.allow_expensive_queries' is set to false.", e.getMessage());
}

public void testVisit() {
ScriptScoreQueryBuilder scriptQueryBuilder = doCreateTestQueryBuilder();
List<QueryBuilder> visitedQueries = new ArrayList<>();
scriptQueryBuilder.visit(createTestVisitor(visitedQueries));
assertEquals(2, visitedQueries.size());
}
}

0 comments on commit a8969cb

Please sign in to comment.