Add implementation of multi_match similar to simple_query_string.

Signed-off-by: Yury Fridlyand <[email protected]>
Bit-Quill · Jun 15, 2022 · 88799ad · 88799ad
1 parent 26058b8
commit 88799ad
Show file tree

Hide file tree

Showing 31 changed files with 8,718 additions and 69 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java b/core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java
@@ -6,11 +6,14 @@
 
 package org.opensearch.sql.analysis;
 
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import lombok.Getter;
@@ -32,6 +35,7 @@
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
 import org.opensearch.sql.ast.expression.QualifiedName;
+import org.opensearch.sql.ast.expression.RelevanceFieldList;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
 import org.opensearch.sql.ast.expression.UnresolvedAttribute;
@@ -40,11 +44,13 @@
 import org.opensearch.sql.ast.expression.WindowFunction;
 import org.opensearch.sql.ast.expression.Xor;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
+import org.opensearch.sql.data.model.ExprTupleValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.expression.LiteralExpression;
 import org.opensearch.sql.expression.NamedArgumentExpression;
 import org.opensearch.sql.expression.NamedExpression;
 import org.opensearch.sql.expression.ParseExpression;
@@ -158,6 +164,12 @@ public Expression visitAggregateFunction(AggregateFunction node, AnalysisContext
     }
   }
 
+  @Override
+  public Expression visitRelevanceFieldList(RelevanceFieldList node, AnalysisContext context) {
+    return new LiteralExpression(ExprValueUtils.tupleValue(
+        ImmutableMap.copyOf(node.getFieldList())));
+  }
+
   @Override
   public Expression visitFunction(Function node, AnalysisContext context) {
     FunctionName functionName = FunctionName.of(node.getFuncName());

diff --git a/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java b/core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java
@@ -26,6 +26,7 @@
 import org.opensearch.sql.ast.expression.Not;
 import org.opensearch.sql.ast.expression.Or;
 import org.opensearch.sql.ast.expression.QualifiedName;
+import org.opensearch.sql.ast.expression.RelevanceFieldList;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
 import org.opensearch.sql.ast.expression.UnresolvedAttribute;
@@ -110,6 +111,10 @@ public T visitLiteral(Literal node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitRelevanceFieldList(RelevanceFieldList node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitUnresolvedAttribute(UnresolvedAttribute node, C context) {
     return visitChildren(node, context);
   }

diff --git a/core/src/main/java/org/opensearch/sql/ast/expression/RelevanceFieldList.java b/core/src/main/java/org/opensearch/sql/ast/expression/RelevanceFieldList.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.List;
+import java.util.stream.Collectors;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+
+/**
+ * Expression node that includes a list of RelevanceField nodes.
+ */
+@EqualsAndHashCode(callSuper = false)
+@AllArgsConstructor
+public class RelevanceFieldList extends UnresolvedExpression {
+  @Getter
+  private java.util.Map<String, Float> fieldList;
+
+  @Override
+  public List<UnresolvedExpression> getChild() {
+    return List.of();
+  }
+
+  @Override
+  public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
+    return nodeVisitor.visitRelevanceFieldList(this, context);
+  }
+
+  @Override
+  public String toString() {
+    return fieldList
+        .entrySet()
+        .stream()
+        .map(e -> String.format("\"%s\" ^ %s", e.getKey(), e.getValue()))
+        .collect(Collectors.joining(", "));
+  }
+}
diff --git a/core/src/main/java/org/opensearch/sql/data/model/ExprCollectionValue.java b/core/src/main/java/org/opensearch/sql/data/model/ExprCollectionValue.java
@@ -45,7 +45,7 @@ public List<ExprValue> collectionValue() {
   public String toString() {
     return valueList.stream()
         .map(Object::toString)
-        .collect(Collectors.joining(",", "[", "]"));
+        .collect(Collectors.joining(", ", "[", "]"));
   }
 
   @Override

diff --git a/core/src/main/java/org/opensearch/sql/expression/DSL.java b/core/src/main/java/org/opensearch/sql/expression/DSL.java
@@ -658,9 +658,11 @@ public FunctionExpression match_phrase(Expression... args) {
     return compile(BuiltinFunctionName.MATCH_PHRASE, args);
   }
 
+  public FunctionExpression multi_match(Expression... args) {
+    return compile(BuiltinFunctionName.MULTI_MATCH, args);
+  }
+
   private FunctionExpression compile(BuiltinFunctionName bfn, Expression... args) {
     return (FunctionExpression) repository.compile(bfn.getName(), Arrays.asList(args.clone()));
   }
-
-
 }
diff --git a/core/src/main/java/org/opensearch/sql/expression/ExpressionNodeVisitor.java b/core/src/main/java/org/opensearch/sql/expression/ExpressionNodeVisitor.java
@@ -6,7 +6,6 @@
 
 package org.opensearch.sql.expression;
 
-import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.expression.aggregation.Aggregator;
 import org.opensearch.sql.expression.aggregation.NamedAggregator;
 import org.opensearch.sql.expression.conditional.cases.CaseClause;
@@ -93,5 +92,4 @@ public T visitWhen(WhenClause node, C context) {
   public T visitNamedArgument(NamedArgumentExpression node, C context) {
     return visitNode(node, context);
   }
-
 }
diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java
@@ -195,7 +195,8 @@ public enum BuiltinFunctionName {
    */
   QUERY(FunctionName.of("query")),
   MATCH_QUERY(FunctionName.of("match_query")),
-  MATCHQUERY(FunctionName.of("matchquery"));
+  MATCHQUERY(FunctionName.of("matchquery")),
+  MULTI_MATCH(FunctionName.of("multi_match"));
 
   private final FunctionName name;
 

diff --git a/core/src/main/java/org/opensearch/sql/expression/function/OpenSearchFunctions.java b/core/src/main/java/org/opensearch/sql/expression/function/OpenSearchFunctions.java
@@ -6,8 +6,10 @@
 package org.opensearch.sql.expression.function;
 
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
+import static org.opensearch.sql.data.type.ExprCoreType.STRUCT;
 
 import com.google.common.collect.ImmutableMap;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -24,15 +26,17 @@
 @UtilityClass
 public class OpenSearchFunctions {
 
-  public static final int MATCH_MAX_NUM_PARAMETERS = 12;
-  public static final int MATCH_PHRASE_MAX_NUM_PARAMETERS = 3;
+  public static final int MATCH_MAX_NUM_PARAMETERS = 14;
+  public static final int MATCH_PHRASE_MAX_NUM_PARAMETERS = 5;
   public static final int MIN_NUM_PARAMETERS = 2;
+  public static final int MULTI_MATCH_MAX_NUM_PARAMETERS = 17;
 
   /**
    * Add functions specific to OpenSearch to repository.
    */
   public void register(BuiltinFunctionRepository repository) {
     repository.register(match());
+    repository.register(multi_match());
     // Register MATCHPHRASE as MATCH_PHRASE as well for backwards
     // compatibility.
     repository.register(match_phrase(BuiltinFunctionName.MATCH_PHRASE));
@@ -41,31 +45,36 @@ public void register(BuiltinFunctionRepository repository) {
 
   private static FunctionResolver match() {
     FunctionName funcName = BuiltinFunctionName.MATCH.getName();
-    return getRelevanceFunctionResolver(funcName, MATCH_MAX_NUM_PARAMETERS);
+    return getRelevanceFunctionResolver(funcName, MATCH_MAX_NUM_PARAMETERS, STRING);
   }
 
   private static FunctionResolver match_phrase(BuiltinFunctionName matchPhrase) {
     FunctionName funcName = matchPhrase.getName();
-    return getRelevanceFunctionResolver(funcName, MATCH_PHRASE_MAX_NUM_PARAMETERS);
+    return getRelevanceFunctionResolver(funcName, MATCH_PHRASE_MAX_NUM_PARAMETERS, STRING);
+  }
+
+  private static FunctionResolver multi_match() {
+    FunctionName funcName = BuiltinFunctionName.MULTI_MATCH.getName();
+    return getRelevanceFunctionResolver(funcName, MULTI_MATCH_MAX_NUM_PARAMETERS, STRUCT);
   }
 
   private static FunctionResolver getRelevanceFunctionResolver(
-      FunctionName funcName, int maxNumParameters) {
+      FunctionName funcName, int maxNumParameters, ExprCoreType firstArgType) {
     return new FunctionResolver(funcName,
-      getRelevanceFunctionSignatureMap(funcName, maxNumParameters));
+      getRelevanceFunctionSignatureMap(funcName, maxNumParameters, firstArgType));
   }
 
   private static Map<FunctionSignature, FunctionBuilder> getRelevanceFunctionSignatureMap(
-      FunctionName funcName, int numOptionalParameters) {
+      FunctionName funcName, int maxNumParameters, ExprCoreType firstArgType) {
     FunctionBuilder buildFunction = args -> new OpenSearchFunction(funcName, args);
     var signatureMapBuilder = ImmutableMap.<FunctionSignature, FunctionBuilder>builder();
     for (int numParameters = MIN_NUM_PARAMETERS;
-         numParameters <= MIN_NUM_PARAMETERS + numOptionalParameters;
-         numParameters++) {
-      List<ExprType> args = Collections.nCopies(numParameters, STRING);
+         numParameters <= maxNumParameters; numParameters++) {
+      List<ExprType> args = new ArrayList<>(Collections.nCopies(numParameters - 1, STRING));
+      args.add(0, firstArgType);
       signatureMapBuilder.put(new FunctionSignature(funcName, args), buildFunction);
     }
-    return  signatureMapBuilder.build();
+    return signatureMapBuilder.build();
   }
 
   private static class OpenSearchFunction extends FunctionExpression {

diff --git a/core/src/test/java/org/opensearch/sql/analysis/ExpressionAnalyzerTest.java b/core/src/test/java/org/opensearch/sql/analysis/ExpressionAnalyzerTest.java
@@ -10,6 +10,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.opensearch.sql.ast.dsl.AstDSL.field;
+import static org.opensearch.sql.ast.dsl.AstDSL.floatLiteral;
 import static org.opensearch.sql.ast.dsl.AstDSL.function;
 import static org.opensearch.sql.ast.dsl.AstDSL.intLiteral;
 import static org.opensearch.sql.ast.dsl.AstDSL.qualifiedName;
@@ -22,18 +23,23 @@
 import static org.opensearch.sql.data.type.ExprCoreType.STRUCT;
 import static org.opensearch.sql.expression.DSL.ref;
 
+import com.google.common.collect.ImmutableMap;
 import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.opensearch.sql.analysis.symbol.Namespace;
 import org.opensearch.sql.analysis.symbol.Symbol;
 import org.opensearch.sql.ast.dsl.AstDSL;
 import org.opensearch.sql.ast.expression.AllFields;
 import org.opensearch.sql.ast.expression.DataType;
+import org.opensearch.sql.ast.expression.RelevanceFieldList;
 import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
+import org.opensearch.sql.data.model.ExprTupleValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
@@ -359,6 +365,51 @@ void visit_in() {
         () -> analyze(AstDSL.in(field("integer_value"), Collections.emptyList())));
   }
 
+  @Test
+  void multi_match_expression() {
+    assertAnalyzeEqual(
+        dsl.multi_match(
+            dsl.namedArgument("fields", DSL.literal(
+                new ExprTupleValue(new LinkedHashMap<>(ImmutableMap.of(
+                    "field", ExprValueUtils.floatValue(1.F)))))),
+            dsl.namedArgument("query", DSL.literal("sample query"))),
+        AstDSL.function("multi_match",
+            AstDSL.unresolvedArg("fields", new RelevanceFieldList(Map.of(
+                "field", 1.F))),
+            AstDSL.unresolvedArg("query", stringLiteral("sample query"))));
+  }
+
+  @Test
+  void multi_match_expression_with_params() {
+    assertAnalyzeEqual(
+        dsl.multi_match(
+            dsl.namedArgument("fields", DSL.literal(
+                new ExprTupleValue(new LinkedHashMap<>(ImmutableMap.of(
+                    "field", ExprValueUtils.floatValue(1.F)))))),
+            dsl.namedArgument("query", DSL.literal("sample query")),
+            dsl.namedArgument("analyzer", DSL.literal("keyword"))),
+        AstDSL.function("multi_match",
+            AstDSL.unresolvedArg("fields", new RelevanceFieldList(Map.of(
+                "field", 1.F))),
+            AstDSL.unresolvedArg("query", stringLiteral("sample query")),
+            AstDSL.unresolvedArg("analyzer", stringLiteral("keyword"))));
+  }
+
+  @Test
+  void multi_match_expression_two_fields() {
+    assertAnalyzeEqual(
+        dsl.multi_match(
+            dsl.namedArgument("fields", DSL.literal(
+                new ExprTupleValue(new LinkedHashMap<>(ImmutableMap.of(
+                    "field1", ExprValueUtils.floatValue(1.F),
+                    "field2", ExprValueUtils.floatValue(.3F)))))),
+            dsl.namedArgument("query", DSL.literal("sample query"))),
+        AstDSL.function("multi_match",
+            AstDSL.unresolvedArg("fields", new RelevanceFieldList(ImmutableMap.of(
+                "field1", 1.F, "field2", .3F))),
+            AstDSL.unresolvedArg("query", stringLiteral("sample query"))));
+  }
+
   protected Expression analyze(UnresolvedExpression unresolvedExpression) {
     return expressionAnalyzer.analyze(unresolvedExpression, analysisContext);
   }

diff --git a/core/src/test/java/org/opensearch/sql/expression/function/OpenSearchFunctionsTest.java b/core/src/test/java/org/opensearch/sql/expression/function/OpenSearchFunctionsTest.java
@@ -9,18 +9,25 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.opensearch.sql.data.type.ExprCoreType.BOOLEAN;
 
+import com.google.common.collect.ImmutableMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import org.junit.jupiter.api.Test;
+import org.opensearch.sql.data.model.ExprTupleValue;
+import org.opensearch.sql.data.model.ExprValueUtils;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.ExpressionTestBase;
 import org.opensearch.sql.expression.FunctionExpression;
 import org.opensearch.sql.expression.NamedArgumentExpression;
 
 
-
 public class OpenSearchFunctionsTest extends ExpressionTestBase {
   private final NamedArgumentExpression field = new NamedArgumentExpression(
       "field", DSL.literal("message"));
+  private final NamedArgumentExpression fields = new NamedArgumentExpression(
+      "fields", DSL.literal(new ExprTupleValue(new LinkedHashMap<>(ImmutableMap.of(
+          "title", ExprValueUtils.floatValue(1.F),
+          "body", ExprValueUtils.floatValue(.3F))))));
   private final NamedArgumentExpression query = new NamedArgumentExpression(
       "query", DSL.literal("search query"));
   private final NamedArgumentExpression analyzer = new NamedArgumentExpression(
@@ -44,9 +51,9 @@ public class OpenSearchFunctionsTest extends ExpressionTestBase {
   private final NamedArgumentExpression minimumShouldMatch = new NamedArgumentExpression(
       "minimum_should_match", DSL.literal("1"));
   private final NamedArgumentExpression zeroTermsQueryAll = new NamedArgumentExpression(
-          "zero_terms_query", DSL.literal("ALL"));
+      "zero_terms_query", DSL.literal("ALL"));
   private final NamedArgumentExpression zeroTermsQueryNone = new NamedArgumentExpression(
-          "zero_terms_query", DSL.literal("None"));
+      "zero_terms_query", DSL.literal("None"));
   private final NamedArgumentExpression boost = new NamedArgumentExpression(
       "boost", DSL.literal("2.0"));
   private final NamedArgumentExpression slop = new NamedArgumentExpression(
@@ -111,8 +118,8 @@ void match() {
 
     expr = dsl.match(
         field, query, analyzer, autoGenerateSynonymsPhrase, fuzziness, maxExpansions, prefixLength,
-        fuzzyTranspositions, fuzzyRewrite, lenient, operator, minimumShouldMatch, zeroTermsQueryAll,
-        boost);
+        fuzzyTranspositions, fuzzyRewrite, lenient, operator, minimumShouldMatch,
+        zeroTermsQueryNone, boost);
     assertEquals(BOOLEAN, expr.type());
   }
 
@@ -146,4 +153,12 @@ void match_to_string() {
     FunctionExpression expr = dsl.match(field, query);
     assertEquals("match(field=\"message\", query=\"search query\")", expr.toString());
   }
+
+  @Test
+  void multi_match() {
+    FunctionExpression expr = dsl.multi_match(fields, query);
+    assertEquals(String.format("multi_match(fields=%s, query=%s)",
+            fields.getValue().toString(), query.getValue().toString()),
+        expr.toString());
+  }
 }