Update parser - separate parsing function with one and multiple fields.

Signed-off-by: Yury Fridlyand <[email protected]>
Bit-Quill · May 13, 2022 · 6c1d9bb · 6c1d9bb
1 parent dc75ce8
commit 6c1d9bb
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 18 deletions.
diff --git a/sql/src/main/antlr/OpenSearchSQLParser.g4 b/sql/src/main/antlr/OpenSearchSQLParser.g4
@@ -318,12 +318,20 @@ specificFunction
     ;
 
 relevanceFunction
-    : relevanceFunctionName LR_BRACKET
+    : relevanceFunctionType1 | relevanceFunctionType2
+    ;
+
+// Field is a single column
+relevanceFunctionType1
+    : relevanceFunctionType1Name LR_BRACKET
         field=relevanceField COMMA query=relevanceQuery
-        (COMMA relevanceArg)* RR_BRACKET
-    | relevanceFunctionNameEx LR_BRACKET
-        fields=relevanceFields COMMA query=relevanceQuery
-        (COMMA relevanceArg)* RR_BRACKET
+        (COMMA relevanceArg)* RR_BRACKET;
+
+// Field is a list of columns
+relevanceFunctionType2
+    : relevanceFunctionType2Name LR_BRACKET
+        LT_SQR_PRTHS field=relevanceField (COMMA field=relevanceField)* RT_SQR_PRTHS
+        COMMA query=relevanceQuery (COMMA relevanceArg)* RR_BRACKET
     ;
 
 convertedDataType
@@ -385,11 +393,11 @@ flowControlFunctionName
     : IF | IFNULL | NULLIF | ISNULL
     ;
 
-relevanceFunctionName
+relevanceFunctionType1Name
     : MATCH
     ;
 
-relevanceFunctionNameEx
+relevanceFunctionType2Name
     : SIMPLE_QUERY_STRING
     ;
 
@@ -418,11 +426,6 @@ relevanceArgName
 relevanceField
     : relevanceArgValue;
 
-relevanceFields
-// simple_query_string function accepts list of field as well in format '["field1", "field2", ... ]'
-    : (LT_SQR_PRTHS relevanceArgValue (COMMA relevanceArgValue)* RT_SQR_PRTHS)
-    ;
-
 relevanceQuery
     : relevanceArgValue;
 

diff --git a/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java b/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java
@@ -60,6 +60,7 @@
 import org.opensearch.sql.ast.expression.AggregateFunction;
 import org.opensearch.sql.ast.expression.AllFields;
 import org.opensearch.sql.ast.expression.And;
+import org.opensearch.sql.ast.expression.AttributeList;
 import org.opensearch.sql.ast.expression.Case;
 import org.opensearch.sql.ast.expression.Cast;
 import org.opensearch.sql.ast.expression.DataType;
@@ -364,13 +365,26 @@ public UnresolvedExpression visitConvertedDataType(
 
   @Override
   public UnresolvedExpression visitRelevanceFunction(RelevanceFunctionContext ctx) {
+    if (ctx.relevanceFunctionType1() != null) {
+      return new Function(
+          ctx.relevanceFunctionType1().relevanceFunctionType1Name().getText().toLowerCase(),
+          relevanceArgumentsType1(ctx.relevanceFunctionType1()));
+    }
+    else {
+      return new Function(
+          ctx.relevanceFunctionType2().relevanceFunctionType2Name().getText().toLowerCase(),
+          relevanceArgumentsType2(ctx.relevanceFunctionType2()));
+    }
+    /*
+    ctx.relevanceFunctionType1().relevanceFunctionType1Name()
     ParserRuleContext func = ctx.relevanceFunctionNameEx();
     if (func == null) {
       func = ctx.relevanceFunctionName();
     }
     return new Function(
         func.getText().toLowerCase(),
         relevanceArguments(ctx));
+        */
   }
 
   private Function visitFunction(String functionName, FunctionArgsContext args) {
@@ -395,16 +409,12 @@ private QualifiedName visitIdentifiers(List<IdentContext> identifiers) {
     );
   }
 
-  private List<UnresolvedExpression> relevanceArguments(RelevanceFunctionContext ctx) {
+  private List<UnresolvedExpression> relevanceArgumentsType1(OpenSearchSQLParser.RelevanceFunctionType1Context ctx) {
     // all the arguments are defaulted to string values
     // to skip environment resolving and function signature resolving
     ImmutableList.Builder<UnresolvedExpression> builder = ImmutableList.builder();
-    ParserRuleContext field = ctx.field;
-    if (field == null) {
-      field = ctx.fields;
-    }
     builder.add(new UnresolvedArgument("field",
-        new Literal(StringUtils.unquoteText(field.getText()), DataType.STRING)));
+        new Literal(StringUtils.unquoteText(ctx.field.getText()), DataType.STRING)));
     builder.add(new UnresolvedArgument("query",
         new Literal(StringUtils.unquoteText(ctx.query.getText()), DataType.STRING)));
     ctx.relevanceArg().forEach(v -> builder.add(new UnresolvedArgument(
@@ -413,4 +423,24 @@ private List<UnresolvedExpression> relevanceArguments(RelevanceFunctionContext c
     return builder.build();
   }
 
+  private List<UnresolvedExpression> relevanceArgumentsType2(OpenSearchSQLParser.RelevanceFunctionType2Context ctx) {
+    // all the arguments are defaulted to string values
+    // to skip environment resolving and function signature resolving
+    ImmutableList.Builder<UnresolvedExpression> builder = ImmutableList.builder();
+    var fields =
+            ctx.getRuleContexts(OpenSearchSQLParser.RelevanceFieldContext.class);
+    ImmutableList.Builder<UnresolvedExpression> b2 = ImmutableList.builder();
+    for (var i = 0; i < fields.size(); i++) {
+      b2.add(new UnresolvedArgument(String.format("field_%d", i),
+        new Literal(StringUtils.unquoteText(fields.get(i).getText()), DataType.STRING)));
+    }
+    var lst2 = new AttributeList(b2.build());
+    builder.add(new UnresolvedArgument("fields", lst2));
+    builder.add(new UnresolvedArgument("query",
+        new Literal(StringUtils.unquoteText(ctx.query.getText()), DataType.STRING)));
+    ctx.relevanceArg().forEach(v -> builder.add(new UnresolvedArgument(
+        v.relevanceArgName().getText().toLowerCase(), new Literal(StringUtils.unquoteText(
+            v.relevanceArgValue().getText()), DataType.STRING))));
+    return builder.build();
+  }
 }