Skip to content

Commit d3667ce

Browse files
MaxKsyunzMaxKsyunz
authored andcommitted
Add RelevanceQuery -- a base class for MatchQuery and MatchPhraseQuery.
Signed-off-by: MaxKsyunz <maxk@bitquilltech.com>
1 parent 1fae665 commit d3667ce

File tree

4 files changed

+231
-111
lines changed

4 files changed

+231
-111
lines changed

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MatchPhraseQuery.java

Lines changed: 14 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -23,51 +23,22 @@
2323
/**
2424
* Lucene query that builds a match_phrase query.
2525
*/
26-
public class MatchPhraseQuery extends LuceneQuery {
27-
private final FluentAction analyzer =
28-
(b, v) -> b.analyzer(v.stringValue());
29-
private final FluentAction slop =
30-
(b, v) -> b.slop(Integer.parseInt(v.stringValue()));
31-
private final FluentAction
32-
zeroTermsQuery = (b, v) -> b.zeroTermsQuery(
33-
org.opensearch.index.search.MatchQuery.ZeroTermsQuery.valueOf(v.stringValue()));
34-
35-
interface FluentAction
36-
extends BiFunction<MatchPhraseQueryBuilder, ExprValue, MatchPhraseQueryBuilder> {
26+
public class MatchPhraseQuery extends RelevanceQuery<MatchPhraseQueryBuilder> {
27+
/**
28+
* Default constructor for MatchPhraseQuery configures how RelevanceQuery.build() handles
29+
* named arguments.
30+
*/
31+
public MatchPhraseQuery() {
32+
super(ImmutableMap.<String, QueryBuilderStep<MatchPhraseQueryBuilder>>builder()
33+
.put("analyzer", (b, v) -> b.analyzer(v.stringValue()))
34+
.put("slop", (b, v) -> b.slop(Integer.parseInt(v.stringValue())))
35+
.put("zero_terms_query", (b, v) -> b.zeroTermsQuery(
36+
org.opensearch.index.search.MatchQuery.ZeroTermsQuery.valueOf(v.stringValue())))
37+
.build());
3738
}
3839

39-
ImmutableMap<Object, FluentAction>
40-
argAction =
41-
ImmutableMap.<Object, FluentAction>builder()
42-
.put("analyzer", analyzer)
43-
.put("slop", slop)
44-
.put("zero_terms_query", zeroTermsQuery)
45-
.build();
46-
4740
@Override
48-
public QueryBuilder build(FunctionExpression func) {
49-
List<Expression> arguments = func.getArguments();
50-
if (arguments.size() < 2) {
51-
throw new SemanticCheckException("match_phrase requires at least two parameters");
52-
}
53-
NamedArgumentExpression field = (NamedArgumentExpression) arguments.get(0);
54-
NamedArgumentExpression query = (NamedArgumentExpression) arguments.get(1);
55-
MatchPhraseQueryBuilder queryBuilder = QueryBuilders.matchPhraseQuery(
56-
field.getValue().valueOf(null).stringValue(),
57-
query.getValue().valueOf(null).stringValue());
58-
59-
Iterator<Expression> iterator = arguments.listIterator(2);
60-
while (iterator.hasNext()) {
61-
NamedArgumentExpression arg = (NamedArgumentExpression) iterator.next();
62-
if (!argAction.containsKey(arg.getArgName())) {
63-
throw new SemanticCheckException(String
64-
.format("Parameter %s is invalid for match_phrase function.", arg.getArgName()));
65-
}
66-
(Objects.requireNonNull(
67-
argAction
68-
.get(arg.getArgName())))
69-
.apply(queryBuilder, arg.getValue().valueOf(null));
70-
}
71-
return queryBuilder;
41+
protected MatchPhraseQueryBuilder createQueryBuilder(String field, String query) {
42+
return QueryBuilders.matchPhraseQuery(field, query);
7243
}
7344
}

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/script/filter/lucene/relevance/MatchQuery.java

Lines changed: 30 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -6,79 +6,41 @@
66
package org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance;
77

88
import com.google.common.collect.ImmutableMap;
9-
import java.util.Iterator;
10-
import java.util.function.BiFunction;
9+
import java.util.Map;
1110
import org.opensearch.index.query.MatchQueryBuilder;
1211
import org.opensearch.index.query.Operator;
13-
import org.opensearch.index.query.QueryBuilder;
1412
import org.opensearch.index.query.QueryBuilders;
15-
import org.opensearch.sql.data.model.ExprValue;
16-
import org.opensearch.sql.exception.SemanticCheckException;
17-
import org.opensearch.sql.expression.Expression;
18-
import org.opensearch.sql.expression.FunctionExpression;
19-
import org.opensearch.sql.expression.NamedArgumentExpression;
20-
import org.opensearch.sql.opensearch.storage.script.filter.lucene.LuceneQuery;
2113

22-
public class MatchQuery extends LuceneQuery {
23-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> analyzer =
24-
(b, v) -> b.analyzer(v.stringValue());
25-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> synonymsPhrase =
26-
(b, v) -> b.autoGenerateSynonymsPhraseQuery(Boolean.parseBoolean(v.stringValue()));
27-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> fuzziness =
28-
(b, v) -> b.fuzziness(v.stringValue());
29-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> maxExpansions =
30-
(b, v) -> b.maxExpansions(Integer.parseInt(v.stringValue()));
31-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> prefixLength =
32-
(b, v) -> b.prefixLength(Integer.parseInt(v.stringValue()));
33-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> fuzzyTranspositions =
34-
(b, v) -> b.fuzzyTranspositions(Boolean.parseBoolean(v.stringValue()));
35-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> fuzzyRewrite =
36-
(b, v) -> b.fuzzyRewrite(v.stringValue());
37-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> lenient =
38-
(b, v) -> b.lenient(Boolean.parseBoolean(v.stringValue()));
39-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> operator =
40-
(b, v) -> b.operator(Operator.fromString(v.stringValue()));
41-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> minimumShouldMatch =
42-
(b, v) -> b.minimumShouldMatch(v.stringValue());
43-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> zeroTermsQuery =
44-
(b, v) -> b.zeroTermsQuery(
45-
org.opensearch.index.search.MatchQuery.ZeroTermsQuery.valueOf(v.stringValue()));
46-
private final BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder> boost =
47-
(b, v) -> b.boost(Float.parseFloat(v.stringValue()));
48-
49-
ImmutableMap<Object, Object> argAction = ImmutableMap.builder()
50-
.put("analyzer", analyzer)
51-
.put("auto_generate_synonyms_phrase_query", synonymsPhrase)
52-
.put("fuzziness", fuzziness)
53-
.put("max_expansions", maxExpansions)
54-
.put("prefix_length", prefixLength)
55-
.put("fuzzy_transpositions", fuzzyTranspositions)
56-
.put("fuzzy_rewrite", fuzzyRewrite)
57-
.put("lenient", lenient)
58-
.put("operator", operator)
59-
.put("minimum_should_match", minimumShouldMatch)
60-
.put("zero_terms_query", zeroTermsQuery)
61-
.put("boost", boost)
62-
.build();
14+
/**
15+
* Initializes MatchQueryBuilder from a FunctionExpression.
16+
*/
17+
public class MatchQuery extends RelevanceQuery<MatchQueryBuilder> {
18+
/**
19+
* Default constructor for MatchQuery configures how RelevanceQuery.build() handles
20+
* named arguments.
21+
*/
22+
public MatchQuery() {
23+
super(ImmutableMap.<String, QueryBuilderStep<MatchQueryBuilder>>builder()
24+
.put("analyzer", (b, v) -> b.analyzer(v.stringValue()))
25+
.put("auto_generate_synonyms_phrase_query",
26+
(b, v) -> b.autoGenerateSynonymsPhraseQuery(Boolean.parseBoolean(v.stringValue())))
27+
.put("fuzziness", (b, v) -> b.fuzziness(v.stringValue()))
28+
.put("max_expansions", (b, v) -> b.maxExpansions(Integer.parseInt(v.stringValue())))
29+
.put("prefix_length", (b, v) -> b.prefixLength(Integer.parseInt(v.stringValue())))
30+
.put("fuzzy_transpositions",
31+
(b, v) -> b.fuzzyTranspositions(Boolean.parseBoolean(v.stringValue())))
32+
.put("fuzzy_rewrite", (b, v) -> b.fuzzyRewrite(v.stringValue()))
33+
.put("lenient", (b, v) -> b.lenient(Boolean.parseBoolean(v.stringValue())))
34+
.put("operator", (b, v) -> b.operator(Operator.fromString(v.stringValue())))
35+
.put("minimum_should_match", (b, v) -> b.minimumShouldMatch(v.stringValue()))
36+
.put("zero_terms_query", (b, v) -> b.zeroTermsQuery(
37+
org.opensearch.index.search.MatchQuery.ZeroTermsQuery.valueOf(v.stringValue())))
38+
.put("boost", (b, v) -> b.boost(Float.parseFloat(v.stringValue())))
39+
.build());
40+
}
6341

6442
@Override
65-
public QueryBuilder build(FunctionExpression func) {
66-
Iterator<Expression> iterator = func.getArguments().iterator();
67-
NamedArgumentExpression field = (NamedArgumentExpression) iterator.next();
68-
NamedArgumentExpression query = (NamedArgumentExpression) iterator.next();
69-
MatchQueryBuilder queryBuilder = QueryBuilders.matchQuery(
70-
field.getValue().valueOf(null).stringValue(),
71-
query.getValue().valueOf(null).stringValue());
72-
while (iterator.hasNext()) {
73-
NamedArgumentExpression arg = (NamedArgumentExpression) iterator.next();
74-
if (!argAction.containsKey(arg.getArgName())) {
75-
throw new SemanticCheckException(String
76-
.format("Parameter %s is invalid for match function.", arg.getArgName()));
77-
}
78-
((BiFunction<MatchQueryBuilder, ExprValue, MatchQueryBuilder>) argAction
79-
.get(arg.getArgName()))
80-
.apply(queryBuilder, arg.getValue().valueOf(null));
81-
}
82-
return queryBuilder;
43+
protected MatchQueryBuilder createQueryBuilder(String field, String query) {
44+
return QueryBuilders.matchQuery(field, query);
8345
}
8446
}
Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance;
7+
8+
import java.util.Iterator;
9+
import java.util.List;
10+
import java.util.Map;
11+
import java.util.Objects;
12+
import java.util.function.BiFunction;
13+
import org.opensearch.index.query.MatchPhraseQueryBuilder;
14+
import org.opensearch.index.query.QueryBuilder;
15+
import org.opensearch.index.query.QueryBuilders;
16+
import org.opensearch.sql.data.model.ExprValue;
17+
import org.opensearch.sql.exception.SemanticCheckException;
18+
import org.opensearch.sql.expression.Expression;
19+
import org.opensearch.sql.expression.FunctionExpression;
20+
import org.opensearch.sql.expression.NamedArgumentExpression;
21+
import org.opensearch.sql.opensearch.storage.script.filter.lucene.LuceneQuery;
22+
23+
/**
24+
* Base class for query abstraction that builds a relevance query from function expression.
25+
*/
26+
public abstract class RelevanceQuery<T extends QueryBuilder> extends LuceneQuery {
27+
protected Map<String, QueryBuilderStep<T>> queryBuildActions;
28+
29+
protected RelevanceQuery(Map<String, QueryBuilderStep<T>> actionMap) {
30+
queryBuildActions = actionMap;
31+
}
32+
33+
@Override
34+
public QueryBuilder build(FunctionExpression func) {
35+
List<Expression> arguments = func.getArguments();
36+
if (arguments.size() < 2) {
37+
String queryName = createQueryBuilder("", "").getWriteableName();
38+
throw new SemanticCheckException(
39+
String.format("%s requires at least two parameters", queryName));
40+
}
41+
NamedArgumentExpression field = (NamedArgumentExpression) arguments.get(0);
42+
NamedArgumentExpression query = (NamedArgumentExpression) arguments.get(1);
43+
T queryBuilder = createQueryBuilder(
44+
field.getValue().valueOf(null).stringValue(),
45+
query.getValue().valueOf(null).stringValue());
46+
47+
Iterator<Expression> iterator = arguments.listIterator(2);
48+
while (iterator.hasNext()) {
49+
NamedArgumentExpression arg = (NamedArgumentExpression) iterator.next();
50+
if (!queryBuildActions.containsKey(arg.getArgName())) {
51+
throw new SemanticCheckException(String
52+
.format("Parameter %s is invalid for %s function.", arg.getArgName(), queryBuilder.getWriteableName()));
53+
}
54+
(Objects.requireNonNull(
55+
queryBuildActions
56+
.get(arg.getArgName())))
57+
.apply(queryBuilder, arg.getValue().valueOf(null));
58+
}
59+
return queryBuilder;
60+
}
61+
62+
protected abstract T createQueryBuilder(String field, String query);
63+
64+
/**
65+
* Convenience interface for a function that updates a QueryBuilder
66+
* based on ExprValue.
67+
* @param <T> Concrete query builder
68+
*/
69+
public interface QueryBuilderStep<T extends QueryBuilder> extends
70+
BiFunction<T, ExprValue, T> {
71+
72+
}
73+
}
Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance;
7+
8+
import static org.junit.jupiter.api.Assertions.assertThrows;
9+
import static org.mockito.ArgumentMatchers.any;
10+
import static org.mockito.Mockito.mock;
11+
import static org.mockito.Mockito.times;
12+
import static org.mockito.Mockito.verify;
13+
import static org.mockito.Mockito.when;
14+
import static org.mockito.Mockito.withSettings;
15+
16+
import com.google.common.collect.ImmutableMap;
17+
import java.util.List;
18+
import java.util.stream.Stream;
19+
import org.apache.commons.lang3.NotImplementedException;
20+
import org.junit.jupiter.api.BeforeEach;
21+
import org.junit.jupiter.api.DisplayNameGeneration;
22+
import org.junit.jupiter.api.DisplayNameGenerator;
23+
import org.junit.jupiter.api.Test;
24+
import org.junit.jupiter.params.ParameterizedTest;
25+
import org.junit.jupiter.params.provider.MethodSource;
26+
import org.mockito.Mockito;
27+
import org.opensearch.index.query.QueryBuilder;
28+
import org.opensearch.sql.data.model.ExprStringValue;
29+
import org.opensearch.sql.data.model.ExprValue;
30+
import org.opensearch.sql.data.type.ExprType;
31+
import org.opensearch.sql.exception.SemanticCheckException;
32+
import org.opensearch.sql.expression.Expression;
33+
import org.opensearch.sql.expression.FunctionExpression;
34+
import org.opensearch.sql.expression.LiteralExpression;
35+
import org.opensearch.sql.expression.NamedArgumentExpression;
36+
import org.opensearch.sql.expression.env.Environment;
37+
import org.opensearch.sql.expression.function.FunctionName;
38+
39+
@DisplayNameGeneration(DisplayNameGenerator.ReplaceUnderscores.class)
40+
class RelevanceQueryBuildTest {
41+
42+
public static final NamedArgumentExpression FIELD_ARG = namedArgument("field", "field_A");
43+
public static final NamedArgumentExpression QUERY_ARG = namedArgument("query", "find me");
44+
private RelevanceQuery query;
45+
private QueryBuilder queryBuilder;
46+
47+
@BeforeEach
48+
public void setUp() {
49+
query = mock(RelevanceQuery.class, withSettings().useConstructor(
50+
ImmutableMap.<String, RelevanceQuery.QueryBuilderStep<QueryBuilder>>builder()
51+
.put("boost", (k, v) -> k.boost(Float.parseFloat(v.stringValue()))).build())
52+
.defaultAnswer(Mockito.CALLS_REAL_METHODS));
53+
queryBuilder = mock(QueryBuilder.class);
54+
when(query.createQueryBuilder(any(), any())).thenReturn(queryBuilder);
55+
when(queryBuilder.queryName()).thenReturn("mocked_query");
56+
}
57+
58+
@Test
59+
void first_arg_field_second_arg_query_test() {
60+
query.build(createCall(List.of(FIELD_ARG, QUERY_ARG)));
61+
verify(query, times(1)).createQueryBuilder("field_A", "find me");
62+
}
63+
64+
@Test
65+
void throws_SemanticCheckException_when_wrong_argument_name() {
66+
FunctionExpression expr =
67+
createCall(List.of(FIELD_ARG, QUERY_ARG, namedArgument("wrongArg", "value")));
68+
69+
assertThrows(SemanticCheckException.class, () -> query.build(expr));
70+
}
71+
72+
@Test
73+
void calls_action_when_correct_argument_name() {
74+
FunctionExpression expr =
75+
createCall(List.of(FIELD_ARG, QUERY_ARG, namedArgument("boost", "2.3")));
76+
query.build(expr);
77+
78+
verify(queryBuilder, times(1)).boost(2.3f);
79+
}
80+
81+
@ParameterizedTest
82+
@MethodSource("insufficientArguments")
83+
public void throws_SemanticCheckException_when_no_required_arguments(List<Expression> arguments) {
84+
assertThrows(SemanticCheckException.class, () -> query.build(createCall(arguments)));
85+
}
86+
87+
public static Stream<List<Expression>> insufficientArguments() {
88+
return Stream.of(List.of(),
89+
List.of(namedArgument("field", "field_A")));
90+
}
91+
92+
private static NamedArgumentExpression namedArgument(String field, String fieldValue) {
93+
return new NamedArgumentExpression(field, createLiteral(fieldValue));
94+
}
95+
96+
@Test
97+
private static Expression createLiteral(String value) {
98+
return new LiteralExpression(new ExprStringValue(value));
99+
}
100+
101+
private static FunctionExpression createCall(List<Expression> arguments) {
102+
return new FunctionExpression(new FunctionName("mock_function"), arguments) {
103+
@Override
104+
public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
105+
throw new NotImplementedException("FunctionExpression.valueOf");
106+
}
107+
108+
@Override
109+
public ExprType type() {
110+
throw new NotImplementedException("FunctionExpression.type");
111+
}
112+
};
113+
}
114+
}

0 commit comments

Comments
 (0)