#639: allow metadata fields and score opensearch function

core/src/main/java/org/opensearch/sql/analysis/ExpressionAnalyzer.java

@@ -8,8 +8,7 @@
 
 import static org.opensearch.sql.ast.dsl.AstDSL.and;
 import static org.opensearch.sql.ast.dsl.AstDSL.compare;
-import static org.opensearch.sql.expression.function.BuiltinFunctionName.GTE;
-import static org.opensearch.sql.expression.function.BuiltinFunctionName.LTE;
+import static org.opensearch.sql.ast.expression.QualifiedName.METADATAFIELD_TYPE_MAP;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -31,6 +30,7 @@
 import org.opensearch.sql.ast.expression.Case;
 import org.opensearch.sql.ast.expression.Cast;
 import org.opensearch.sql.ast.expression.Compare;
+import org.opensearch.sql.ast.expression.DataType;
 import org.opensearch.sql.ast.expression.EqualTo;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Function;
@@ -42,6 +42,7 @@
 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.ScoreFunction;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
 import org.opensearch.sql.ast.expression.UnresolvedAttribute;
@@ -51,6 +52,7 @@
 import org.opensearch.sql.ast.expression.Xor;
 import org.opensearch.sql.common.antlr.SyntaxCheckException;
 import org.opensearch.sql.data.model.ExprValueUtils;
+import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
@@ -67,6 +69,7 @@
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
 import org.opensearch.sql.expression.function.FunctionName;
+import org.opensearch.sql.expression.function.OpenSearchFunctions;
 import org.opensearch.sql.expression.parse.ParseExpression;
 import org.opensearch.sql.expression.span.SpanExpression;
 import org.opensearch.sql.expression.window.aggregation.AggregateWindowFunction;
@@ -207,6 +210,79 @@ public Expression visitHighlightFunction(HighlightFunction node, AnalysisContext
     return new HighlightExpression(expr);
   }
 
+  /**
+   * visitScoreFunction removes the score function from the AST and replaces it with the child
+   * relevance function node. If the optional boost variable is provided, the boost argument
+   * of the relevance function is combined.
+   *
+   * @param node    score function node
+   * @param context analysis context for the query
+   * @return resolved relevance function
+   */
+  public Expression visitScoreFunction(ScoreFunction node, AnalysisContext context) {
+    // if no function argument given, just accept the relevance query and return
+    if (!(node.getFuncArg() instanceof Literal)) {
+      OpenSearchFunctions.OpenSearchFunction relevanceQueryExpr =
+              (OpenSearchFunctions.OpenSearchFunction) node
+                      .getRelevanceQuery().accept(this, context);
+      relevanceQueryExpr.setScoreTracked(true);
+      return relevanceQueryExpr;
+    }
+
+    // note: if an argument exists, and there should only be one, it will be a boost argument
+    Literal boostFunctionArg = (Literal) node.getFuncArg();
+    Double thisBoostValue;
+    if (boostFunctionArg.getType().equals(DataType.DOUBLE)) {
+      thisBoostValue = ((Double) boostFunctionArg.getValue());
+    } else if (boostFunctionArg.getType().equals(DataType.INTEGER)) {
+      thisBoostValue = ((Integer) boostFunctionArg.getValue()).doubleValue();
+    } else {
+      throw new SemanticCheckException(String.format("Expected boost type '%s' but got '%s'",
+          DataType.DOUBLE.name(), boostFunctionArg.getType().name()));
+    }
+
+    // update the existing unresolved expression to add a boost argument if it doesn't exist
+    // OR multiply the existing boost argument
+    Function relevanceQueryUnresolvedExpr = (Function) node.getRelevanceQuery();
+    List<UnresolvedExpression> relevanceFuncArgs = relevanceQueryUnresolvedExpr.getFuncArgs();
+
+    boolean doesFunctionContainBoostArgument = false;
+    List<UnresolvedExpression> updatedFuncArgs = new ArrayList<>();
+    for (UnresolvedExpression expr : relevanceFuncArgs) {
+      String argumentName = ((UnresolvedArgument) expr).getArgName();
+      if (argumentName.equalsIgnoreCase("boost")) {
+        doesFunctionContainBoostArgument = true;
+        Literal boostArgLiteral = (Literal) ((UnresolvedArgument) expr).getValue();
+        Double boostValue =
+            Double.parseDouble((String) boostArgLiteral.getValue()) * thisBoostValue;
+        UnresolvedArgument newBoostArg = new UnresolvedArgument(
+                argumentName,
+                new Literal(boostValue.toString(), DataType.STRING)
+        );
+        updatedFuncArgs.add(newBoostArg);
+      } else {
+        updatedFuncArgs.add(expr);
+      }
+    }
+
+    // since nothing was found, add an argument
+    if (!doesFunctionContainBoostArgument) {
+      UnresolvedArgument newBoostArg = new UnresolvedArgument(
+              "boost", new Literal(Double.toString(thisBoostValue), DataType.STRING));
+      updatedFuncArgs.add(newBoostArg);
+    }
+
+    // create a new function expression with boost argument and resolve it
+    Function updatedRelevanceQueryUnresolvedExpr = new Function(
+            relevanceQueryUnresolvedExpr.getFuncName(),
+            updatedFuncArgs);
+    OpenSearchFunctions.OpenSearchFunction relevanceQueryExpr =
+            (OpenSearchFunctions.OpenSearchFunction) updatedRelevanceQueryUnresolvedExpr
+                    .accept(this, context);
+    relevanceQueryExpr.setScoreTracked(true);
+    return relevanceQueryExpr;
+  }
+
   @Override
   public Expression visitIn(In node, AnalysisContext context) {
     return visitIn(node.getField(), node.getValueList(), context);
@@ -297,6 +373,9 @@ public Expression visitAllFields(AllFields node, AnalysisContext context) {
   @Override
   public Expression visitQualifiedName(QualifiedName node, AnalysisContext context) {
     QualifierAnalyzer qualifierAnalyzer = new QualifierAnalyzer(context);
+    if (node.isMetadataField()) {
+      return visitMetadata(qualifierAnalyzer.unqualified(node), context);
+    }
     return visitIdentifier(qualifierAnalyzer.unqualified(node), context);
   }
 
@@ -313,6 +392,19 @@ public Expression visitUnresolvedArgument(UnresolvedArgument node, AnalysisConte
     return new NamedArgumentExpression(node.getArgName(), node.getValue().accept(this, context));
   }
 
+  /**
+   * If QualifiedName is actually a reserved metadata field, return the expr type associated
+   * with the metadata field.
+   * @param ident   metadata field name
+   * @param context analysis context
+   * @return DSL reference
+   */
+  private Expression visitMetadata(String ident, AnalysisContext context) {
+    ExprCoreType exprCoreType = Optional.ofNullable(METADATAFIELD_TYPE_MAP.get(ident))
+            .orElseThrow(() -> new SemanticCheckException("invalid metadata field"));
+    return DSL.ref(ident, exprCoreType);
+  }
+
   private Expression visitIdentifier(String ident, AnalysisContext context) {
     // ParseExpression will always override ReferenceExpression when ident conflicts
     for (NamedExpression expr : context.getNamedParseExpressions()) {

core/src/main/java/org/opensearch/sql/analysis/ExpressionReferenceOptimizer.java

@@ -19,6 +19,7 @@
 import org.opensearch.sql.expression.conditional.cases.CaseClause;
 import org.opensearch.sql.expression.conditional.cases.WhenClause;
 import org.opensearch.sql.expression.function.BuiltinFunctionRepository;
+import org.opensearch.sql.expression.function.OpenSearchFunctions;
 import org.opensearch.sql.planner.logical.LogicalAggregation;
 import org.opensearch.sql.planner.logical.LogicalPlan;
 import org.opensearch.sql.planner.logical.LogicalPlanNodeVisitor;
@@ -70,8 +71,17 @@ public Expression visitFunction(FunctionExpression node, AnalysisContext context
       final List<Expression> args =
           node.getArguments().stream().map(expr -> expr.accept(this, context))
               .collect(Collectors.toList());
-      return (Expression) repository.compile(context.getFunctionProperties(),
-          node.getFunctionName(), args);
+      Expression optimizedFunctionExpression = (Expression) repository.compile(
+              context.getFunctionProperties(),
+              node.getFunctionName(),
+              args
+      );
+      // Propagate scoreTracked for OpenSearch functions
+      if (optimizedFunctionExpression instanceof OpenSearchFunctions.OpenSearchFunction) {
+        ((OpenSearchFunctions.OpenSearchFunction) optimizedFunctionExpression).setScoreTracked(
+                ((OpenSearchFunctions.OpenSearchFunction)node).isScoreTracked());
+      }
+      return optimizedFunctionExpression;
     }
   }
 

core/src/main/java/org/opensearch/sql/ast/AbstractNodeVisitor.java

@@ -29,6 +29,7 @@
 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.ScoreFunction;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
 import org.opensearch.sql.ast.expression.UnresolvedAttribute;
@@ -278,6 +279,10 @@ public T visitHighlightFunction(HighlightFunction node, C context) {
     return visitChildren(node, context);
   }
 
+  public T visitScoreFunction(ScoreFunction node, C context) {
+    return visitChildren(node, context);
+  }
+
   public T visitStatement(Statement node, C context) {
     return visit(node, context);
   }

core/src/main/java/org/opensearch/sql/ast/dsl/AstDSL.java

@@ -34,6 +34,7 @@
 import org.opensearch.sql.ast.expression.Or;
 import org.opensearch.sql.ast.expression.ParseMethod;
 import org.opensearch.sql.ast.expression.QualifiedName;
+import org.opensearch.sql.ast.expression.ScoreFunction;
 import org.opensearch.sql.ast.expression.Span;
 import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.ast.expression.UnresolvedArgument;
@@ -60,7 +61,6 @@
 import org.opensearch.sql.ast.tree.TableFunction;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
-import org.opensearch.sql.expression.function.BuiltinFunctionName;
 
 /**
  * Class of static methods to create specific node instances.
@@ -140,7 +140,11 @@ public UnresolvedPlan values(List<Literal>... values) {
   }
 
   public static QualifiedName qualifiedName(String... parts) {
-    return new QualifiedName(Arrays.asList(parts));
+    return new QualifiedName(Arrays.asList(parts), Boolean.FALSE);
+  }
+
+  public static QualifiedName qualifiedNameWithMetadata(String... parts) {
+    return new QualifiedName(Arrays.asList(parts), Boolean.TRUE);
   }
 
   public static UnresolvedExpression equalTo(
@@ -285,6 +289,11 @@ public UnresolvedExpression highlight(UnresolvedExpression fieldName,
     return new HighlightFunction(fieldName, arguments);
   }
 
+  public UnresolvedExpression score(UnresolvedExpression relevanceQuery,
+                                    UnresolvedExpression funcArg) {
+    return new ScoreFunction(relevanceQuery, funcArg);
+  }
+
   public UnresolvedExpression window(UnresolvedExpression function,
                                      List<UnresolvedExpression> partitionByList,
                                      List<Pair<SortOption, UnresolvedExpression>> sortList) {

core/src/main/java/org/opensearch/sql/ast/expression/QualifiedName.java

@@ -10,29 +10,58 @@
 import static java.util.stream.Collectors.toList;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.StreamSupport;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.data.type.ExprCoreType;
 
 @Getter
 @EqualsAndHashCode(callSuper = false)
 public class QualifiedName extends UnresolvedExpression {
   private final List<String> parts;
 
+  private final Boolean isMetadataField;
+
   public QualifiedName(String name) {
+    this(name, Boolean.FALSE);
+  }
+
+  public static final String METADATA_FIELD_ID = "_id";
+  public static final String METADATA_FIELD_INDEX = "_index";
+  public static final String METADATA_FIELD_SCORE = "_score";
+  public static final String METADATA_FIELD_MAXSCORE = "_maxscore";
+  public static final String METADATA_FIELD_SORT = "_sort";
+
+  public static final java.util.Map<String, ExprCoreType> METADATAFIELD_TYPE_MAP = ImmutableMap.of(
+      METADATA_FIELD_ID, ExprCoreType.STRING,
+      METADATA_FIELD_INDEX, ExprCoreType.STRING,
+      METADATA_FIELD_SCORE, ExprCoreType.FLOAT,
+      METADATA_FIELD_MAXSCORE, ExprCoreType.FLOAT,
+      METADATA_FIELD_SORT, ExprCoreType.LONG
+  );
+
+  public QualifiedName(String name, Boolean isMetadataField) {
     this.parts = Collections.singletonList(name);
+    this.isMetadataField = isMetadataField;
+  }
+
+  public QualifiedName(Iterable<String> parts) {
+    this(parts, Boolean.FALSE);
   }
 
   /**
    * QualifiedName Constructor.
    */
-  public QualifiedName(Iterable<String> parts) {
+  public QualifiedName(Iterable<String> parts, Boolean isMetadataField) {
+    this.isMetadataField = isMetadataField;
     List<String> partsList = StreamSupport.stream(parts.spliterator(), false).collect(toList());
     if (partsList.isEmpty()) {
       throw new IllegalArgumentException("parts is empty");
@@ -110,4 +139,8 @@ public List<UnresolvedExpression> getChild() {
   public <R, C> R accept(AbstractNodeVisitor<R, C> nodeVisitor, C context) {
     return nodeVisitor.visitQualifiedName(this, context);
   }
+
+  public Boolean isMetadataField() {
+    return this.isMetadataField;
+  }
 }

core/src/main/java/org/opensearch/sql/ast/expression/ScoreFunction.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.expression;
+
+import java.util.List;
+import lombok.AllArgsConstructor;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+
+/**
+ * Expression node of Score function.
+ * Score takes a relevance-search expression as an argument and returns it
+ */
+@AllArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@Getter
+@ToString
+public class ScoreFunction extends UnresolvedExpression {
+  private final UnresolvedExpression relevanceQuery;
+  private final UnresolvedExpression funcArg;
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitScoreFunction(this, context);
+  }
+
+  @Override
+  public List<UnresolvedExpression> getChild() {
+    return List.of(relevanceQuery, funcArg);
+  }
+}

core/src/main/java/org/opensearch/sql/expression/DSL.java

@@ -787,7 +787,19 @@ public static FunctionExpression match_bool_prefix(Expression... args) {
   }
 
   public static FunctionExpression wildcard_query(Expression... args) {
-    return compile(FunctionProperties.None,BuiltinFunctionName.WILDCARD_QUERY, args);
+    return compile(FunctionProperties.None, BuiltinFunctionName.WILDCARD_QUERY, args);
+  }
+
+  public static FunctionExpression score(Expression... args) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.SCORE, args);
+  }
+
+  public static FunctionExpression scorequery(Expression... args) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.SCOREQUERY, args);
+  }
+
+  public static FunctionExpression score_query(Expression... args) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.SCORE_QUERY, args);
   }
 
   public static FunctionExpression now(FunctionProperties functionProperties,

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -112,6 +112,7 @@ public enum BuiltinFunctionName {
   WEEKOFYEAR(FunctionName.of("weekofyear")),
   WEEK_OF_YEAR(FunctionName.of("week_of_year")),
   YEAR(FunctionName.of("year")),
+
   // `now`-like functions
   NOW(FunctionName.of("now")),
   CURDATE(FunctionName.of("curdate")),
@@ -122,6 +123,7 @@ public enum BuiltinFunctionName {
   CURRENT_TIMESTAMP(FunctionName.of("current_timestamp")),
   LOCALTIMESTAMP(FunctionName.of("localtimestamp")),
   SYSDATE(FunctionName.of("sysdate")),
+
   /**
    * Text Functions.
    */
@@ -239,6 +241,10 @@ public enum BuiltinFunctionName {
   MATCH_BOOL_PREFIX(FunctionName.of("match_bool_prefix")),
   HIGHLIGHT(FunctionName.of("highlight")),
   MATCH_PHRASE_PREFIX(FunctionName.of("match_phrase_prefix")),
+  SCORE(FunctionName.of("score")),
+  SCOREQUERY(FunctionName.of("scorequery")),
+  SCORE_QUERY(FunctionName.of("score_query")),
+
   /**
    * Legacy Relevance Function.
    */