Skip to content

Commit

Permalink
Merge remote-tracking branch 'origin/dev-simple_query_string-#192-imp…
Browse files Browse the repository at this point in the history
…l' into dev-simple_query_string-#192-sql-simple-max_ast
  • Loading branch information
MaxKsyunz committed May 14, 2022
2 parents 730cc7b + 155d90d commit 12cf311
Show file tree
Hide file tree
Showing 6 changed files with 169 additions and 46 deletions.
5 changes: 5 additions & 0 deletions core/src/main/java/org/opensearch/sql/expression/DSL.java
Original file line number Diff line number Diff line change
Expand Up @@ -654,4 +654,9 @@ public FunctionExpression match(Expression... args) {
return (FunctionExpression) repository
.compile(BuiltinFunctionName.MATCH.getName(), Arrays.asList(args.clone()));
}

public FunctionExpression simple_query_string(Expression... args) {
return (FunctionExpression) repository
.compile(BuiltinFunctionName.SIMPLE_QUERY_STRING.getName(), Arrays.asList(args.clone()));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -187,6 +187,7 @@ public enum BuiltinFunctionName {
* Relevance Function.
*/
MATCH(FunctionName.of("match")),
SIMPLE_QUERY_STRING(FunctionName.of("simple_query_string")),

/**
* Legacy Relevance Function.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,11 @@

import static org.opensearch.sql.data.type.ExprCoreType.STRING;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import lombok.ToString;
import lombok.experimental.UtilityClass;
import org.opensearch.sql.data.model.ExprValue;
import org.opensearch.sql.data.type.ExprCoreType;
Expand All @@ -25,54 +25,40 @@
public class OpenSearchFunctions {
public void register(BuiltinFunctionRepository repository) {
repository.register(match());
repository.register(simple_query_string());
}

private static FunctionResolver match() {
FunctionName funcName = BuiltinFunctionName.MATCH.getName();
// At most field, query, and all optional parameters
final int matchMaxNumParameters = 14;
return getRelevanceFunctionResolver(funcName, matchMaxNumParameters);
}

private static FunctionResolver simple_query_string() {
FunctionName funcName = BuiltinFunctionName.SIMPLE_QUERY_STRING.getName();
// At most field, query, and all optional parameters
// TODO 16 ? See org.opensearch.index.query.SimpleQueryStringBuilder.class
final int matchPhraseMaxNumParameters = 12;
return getRelevanceFunctionResolver(funcName, matchPhraseMaxNumParameters);
}

private static FunctionResolver getRelevanceFunctionResolver(
FunctionName funcName, int maxNumParameters) {
return new FunctionResolver(funcName,
ImmutableMap.<FunctionSignature, FunctionBuilder>builder()
.put(new FunctionSignature(funcName, ImmutableList.of(STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList.of(STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList.of(STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING,
STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING,
STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING,
STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING,
STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.put(new FunctionSignature(funcName, ImmutableList
.of(STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING, STRING,
STRING, STRING, STRING, STRING, STRING)),
args -> new OpenSearchFunction(funcName, args))
.build());
getRelevanceFunctionSignatureMap(funcName, maxNumParameters));
}

private static Map<FunctionSignature, FunctionBuilder> getRelevanceFunctionSignatureMap(
FunctionName funcName, int maxNumParameters) {
final int minNumParameters = 2;
FunctionBuilder buildFunction = args -> new OpenSearchFunction(funcName, args);
var signatureMapBuilder = ImmutableMap.<FunctionSignature, FunctionBuilder>builder();
for (int numParameters = minNumParameters; numParameters <= maxNumParameters; numParameters++) {
List<ExprType> args = Collections.nCopies(numParameters, STRING);
signatureMapBuilder.put(new FunctionSignature(funcName, args), buildFunction);
}
return signatureMapBuilder.build();
}

private static class OpenSearchFunction extends FunctionExpression {
Expand Down Expand Up @@ -106,4 +92,4 @@ public String toString() {
return String.format("%s(%s)", functionName, String.join(", ", args));
}
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.opensearch.sql.data.type.ExprCoreType.BOOLEAN;

import org.apache.commons.lang3.NotImplementedException;
import org.junit.jupiter.api.Test;
import org.opensearch.sql.expression.DSL;
import org.opensearch.sql.expression.ExpressionTestBase;
Expand Down Expand Up @@ -121,4 +122,9 @@ void match_to_string() {
FunctionExpression expr = dsl.match(field, query);
assertEquals("match(field=\"message\", query=\"search query\")", expr.toString());
}

@Test
void simple_query_string() {
throw new NotImplementedException();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.opensearch.sql.opensearch.storage.script.filter.lucene.TermQuery;
import org.opensearch.sql.opensearch.storage.script.filter.lucene.WildcardQuery;
import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.MatchQuery;
import org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance.SimpleQueryString;
import org.opensearch.sql.opensearch.storage.serialization.ExpressionSerializer;

@RequiredArgsConstructor
Expand All @@ -55,6 +56,7 @@ public class FilterQueryBuilder extends ExpressionNodeVisitor<QueryBuilder, Obje
.put(BuiltinFunctionName.QUERY.getName(), new MatchQuery())
.put(BuiltinFunctionName.MATCH_QUERY.getName(), new MatchQuery())
.put(BuiltinFunctionName.MATCHQUERY.getName(), new MatchQuery())
.put(BuiltinFunctionName.SIMPLE_QUERY_STRING.getName(), new SimpleQueryString())
.build();

/**
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.sql.opensearch.storage.script.filter.lucene.relevance;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.json.JSONArray;
import org.json.JSONObject;
import org.json.JSONString;
import org.opensearch.index.query.Operator;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryBuilders;
import org.opensearch.index.query.SimpleQueryStringBuilder;
import org.opensearch.index.query.SimpleQueryStringFlag;
import org.opensearch.sql.data.model.ExprValue;
import org.opensearch.sql.exception.SemanticCheckException;
import org.opensearch.sql.expression.Expression;
import org.opensearch.sql.expression.FunctionExpression;
import org.opensearch.sql.expression.NamedArgumentExpression;
import org.opensearch.sql.opensearch.storage.script.filter.lucene.LuceneQuery;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.stream.Collectors;

public class SimpleQueryString extends LuceneQuery {
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> analyzeWildcard =
(b, v) -> b.analyzeWildcard(Boolean.parseBoolean(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> analyzer =
(b, v) -> b.analyzer(v.stringValue());
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> autoGenerateSynonymsPhraseQuery =
(b, v) -> b.autoGenerateSynonymsPhraseQuery(Boolean.parseBoolean(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> defaultOperator =
(b, v) -> b.defaultOperator(Operator.fromString(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> flags =
(b, v) -> b.flags(SimpleQueryStringFlag.valueOf(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> fuzzyMaxExpansions =
(b, v) -> b.fuzzyMaxExpansions(Integer.parseInt(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> fuzzyPrefixLength =
(b, v) -> b.fuzzyPrefixLength(Integer.parseInt(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> fuzzyTranspositions =
(b, v) -> b.fuzzyTranspositions(Boolean.parseBoolean(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> lenient =
(b, v) -> b.lenient(Boolean.parseBoolean(v.stringValue()));
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> minimumShouldMatch =
(b, v) -> b.minimumShouldMatch(v.stringValue());
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> quoteFieldSuffix =
(b, v) -> b.quoteFieldSuffix(v.stringValue());
private final BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder> boost =
(b, v) -> b.boost(Float.parseFloat(v.stringValue()));

ImmutableMap<Object, Object> argAction = ImmutableMap.builder()
.put("analyze_wildcard", analyzeWildcard)
.put("analyzer", analyzer)
.put("auto_generate_synonyms_phrase_query", autoGenerateSynonymsPhraseQuery)
.put("flags", flags)
.put("fuzzy_max_expansions", fuzzyMaxExpansions)
.put("fuzzy_prefix_length", fuzzyPrefixLength)
.put("fuzzy_transpositions", fuzzyTranspositions)
.put("lenient", lenient)
.put("default_operator", defaultOperator)
.put("minimum_should_match", minimumShouldMatch)
.put("quote_field_suffix", quoteFieldSuffix)
.put("boost", boost)
.build();

@Override
public QueryBuilder build(FunctionExpression func) {
Iterator<Expression> iterator = func.getArguments().iterator();
NamedArgumentExpression fields = (NamedArgumentExpression) iterator.next();
NamedArgumentExpression query = (NamedArgumentExpression) iterator.next();
SimpleQueryStringBuilder queryBuilder = QueryBuilders.simpleQueryStringQuery(
query.getValue().valueOf(null).stringValue());
queryBuilder.fields(parseFields(fields.getValue().valueOf(null).stringValue()));
while (iterator.hasNext()) {
NamedArgumentExpression arg = (NamedArgumentExpression) iterator.next();
if (!argAction.containsKey(arg.getArgName())) {
throw new SemanticCheckException(String
.format("Parameter %s is invalid for simple_query_string function.", arg.getArgName()));
}
((BiFunction<SimpleQueryStringBuilder, ExprValue, SimpleQueryStringBuilder>) argAction
.get(arg.getArgName()))
.apply(queryBuilder, arg.getValue().valueOf(null));
}
return queryBuilder;
}

private Map<String, Float> parseFields(String fields) {
try {
// TODO support elements wrapped by single quotes
var arr = new JSONArray(fields);
if (!arr.toList().stream().allMatch(s -> s instanceof String))
throw new Exception("All listed elements should be strings.");

var lst = arr.toList().stream().map(String::valueOf).collect(Collectors.toList());
var res = new HashMap<String, Float>();
//var builder = ImmutableMap.builder();
for (var elem : lst) {
if (!elem.contains("^")) {
res.put(elem, 1F);
continue;
}

var parts = elem.split("\\^");
var weight = Float.parseFloat(parts[parts.length - 1]);
var field = Arrays.stream(parts).limit(parts.length - 1).collect(Collectors.joining("^"));
res.put(field, weight);
}
return res;
}
catch (Exception e) {
throw new SemanticCheckException(String.format(
"%s: Incorrect value '%s' specified for 'fields' argument of 'simple_query_string' function."
+ "The format is: '[\"field1\", \"field2\", ...]'.", e.getMessage(), fields));
}
}
}

0 comments on commit 12cf311

Please sign in to comment.