Add Highlight In SQL

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

@@ -292,6 +292,12 @@ public LogicalPlan visitProject(Project node, AnalysisContext context) {
       child = windowAnalyzer.analyze(expr, context);
     }
 
+    for (UnresolvedExpression expr : node.getProjectList()) {
+      HighlightAnalyzer highlightAnalyzer = new HighlightAnalyzer(expressionAnalyzer, child);
+      child = highlightAnalyzer.analyze(expr, context);
+
+    }
+
     List<NamedExpression> namedExpressions =
         selectExpressionAnalyzer.analyze(node.getProjectList(), context,
             new ExpressionReferenceOptimizer(expressionAnalyzer.getRepository(), child));

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

@@ -11,9 +11,7 @@
 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;
@@ -29,6 +27,7 @@
 import org.opensearch.sql.ast.expression.EqualTo;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.In;
 import org.opensearch.sql.ast.expression.Interval;
 import org.opensearch.sql.ast.expression.Literal;
@@ -44,12 +43,12 @@
 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.HighlightExpression;
 import org.opensearch.sql.expression.LiteralExpression;
 import org.opensearch.sql.expression.NamedArgumentExpression;
 import org.opensearch.sql.expression.NamedExpression;
@@ -191,6 +190,12 @@ public Expression visitWindowFunction(WindowFunction node, AnalysisContext conte
     return expr;
   }
 
+  @Override
+  public Expression visitHighlight(HighlightFunction node, AnalysisContext context) {
+    Expression expr = node.getHighlightField().accept(this, context);
+    return new HighlightExpression(expr);
+  }
+
   @Override
   public Expression visitIn(In node, AnalysisContext context) {
     return visitIn(node.getField(), node.getValueList(), context);

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.analysis;
+
+import lombok.RequiredArgsConstructor;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Alias;
+import org.opensearch.sql.ast.expression.HighlightFunction;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.expression.Expression;
+import org.opensearch.sql.planner.logical.LogicalHighlight;
+import org.opensearch.sql.planner.logical.LogicalPlan;
+
+/**
+ * Analyze the highlight in the {@link AnalysisContext} to construct the {@link
+ * LogicalPlan}.
+ */
+@RequiredArgsConstructor
+public class HighlightAnalyzer extends AbstractNodeVisitor<LogicalPlan, AnalysisContext> {
+  private final ExpressionAnalyzer expressionAnalyzer;
+  private final LogicalPlan child;
+
+  public LogicalPlan analyze(UnresolvedExpression projectItem, AnalysisContext context) {
+    LogicalPlan highlight = projectItem.accept(this, context);
+    return (highlight == null) ? child : highlight;
+  }
+
+  @Override
+  public LogicalPlan visitAlias(Alias node, AnalysisContext context) {
+    if (!(node.getDelegated() instanceof HighlightFunction)) {
+      return null;
+    }
+
+    HighlightFunction unresolved = (HighlightFunction) node.getDelegated();
+    Expression field = expressionAnalyzer.analyze(unresolved.getHighlightField(), context);
+    return new LogicalHighlight(child, field);
+  }
+}

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

@@ -18,6 +18,7 @@
 import org.opensearch.sql.ast.expression.EqualTo;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.In;
 import org.opensearch.sql.ast.expression.Interval;
 import org.opensearch.sql.ast.expression.Let;
@@ -254,4 +255,8 @@ public T visitKmeans(Kmeans node, C context) {
   public T visitAD(AD node, C context) {
     return visitChildren(node, context);
   }
+
+  public T visitHighlight(HighlightFunction node, C context) {
+    return visitChildren(node, context);
+  }
 }

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

@@ -22,6 +22,7 @@
 import org.opensearch.sql.ast.expression.EqualTo;
 import org.opensearch.sql.ast.expression.Field;
 import org.opensearch.sql.ast.expression.Function;
+import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.In;
 import org.opensearch.sql.ast.expression.Interval;
 import org.opensearch.sql.ast.expression.Let;
@@ -261,6 +262,10 @@ public When when(UnresolvedExpression condition, UnresolvedExpression result) {
     return new When(condition, result);
   }
 
+  public UnresolvedExpression highlight(UnresolvedExpression fieldName) {
+    return new HighlightFunction(fieldName);
+  }
+
   public UnresolvedExpression window(UnresolvedExpression function,
                                      List<UnresolvedExpression> partitionByList,
                                      List<Pair<SortOption, UnresolvedExpression>> sortList) {

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

@@ -0,0 +1,34 @@
+/*
+ * 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 Highlight function.
+ */
+@AllArgsConstructor
+@EqualsAndHashCode(callSuper = false)
+@Getter
+@ToString
+public class HighlightFunction extends UnresolvedExpression {
+  private final UnresolvedExpression highlightField;
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitHighlight(this, context);
+  }
+
+  @Override
+  public List<UnresolvedExpression> getChild() {
+    return List.of(highlightField);
+  }
+}

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

@@ -55,6 +55,10 @@ public T visitNamed(NamedExpression node, C context) {
     return node.getDelegated().accept(this, context);
   }
 
+  public T visitHighlight(HighlightExpression node, C context) {
+    return visitNode(node, context);
+  }
+
   public T visitReference(ReferenceExpression node, C context) {
     return visitNode(node, context);
   }

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

@@ -0,0 +1,57 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression;
+
+import java.util.List;
+import lombok.Getter;
+import org.opensearch.sql.common.utils.StringUtils;
+import org.opensearch.sql.data.model.ExprValue;
+import org.opensearch.sql.data.type.ExprCoreType;
+import org.opensearch.sql.data.type.ExprType;
+import org.opensearch.sql.expression.env.Environment;
+import org.opensearch.sql.expression.function.BuiltinFunctionName;
+
+/**
+ * Highlight Expression.
+ */
+@Getter
+public class HighlightExpression extends FunctionExpression {
+  private final Expression highlightField;
+
+  /**
+   * HighlightExpression Constructor.
+   * @param highlightField : Highlight field for expression.
+   */
+  public HighlightExpression(Expression highlightField) {
+    super(BuiltinFunctionName.HIGHLIGHT.getName(), List.of(highlightField));
+    this.highlightField = highlightField;
+  }
+
+  /**
+   * Return collection value matching highlight field.
+   * @param valueEnv : Dataset to parse value from.
+   * @return : collection value of highlight fields.
+   */
+  @Override
+  public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
+    String refName = "_highlight" + "." + StringUtils.unquoteText(getHighlightField().toString());
+    return valueEnv.resolve(DSL.ref(refName, ExprCoreType.STRING));
+  }
+
+  /**
+   * Get type for HighlightExpression.
+   * @return : String type.
+   */
+  @Override
+  public ExprType type() {
+    return ExprCoreType.ARRAY;
+  }
+
+  @Override
+  public <T, C> T accept(ExpressionNodeVisitor<T, C> visitor, C context) {
+    return visitor.visitHighlight(this, context);
+  }
+}

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

@@ -192,6 +192,7 @@ public enum BuiltinFunctionName {
   MATCHPHRASE(FunctionName.of("matchphrase")),
   QUERY_STRING(FunctionName.of("query_string")),
   MATCH_BOOL_PREFIX(FunctionName.of("match_bool_prefix")),
+  HIGHLIGHT(FunctionName.of("highlight")),
   MATCH_PHRASE_PREFIX(FunctionName.of("match_phrase_prefix")),
   /**
    * Legacy Relevance Function.

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

@@ -8,19 +8,20 @@
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 import static org.opensearch.sql.data.type.ExprCoreType.STRUCT;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import lombok.experimental.UtilityClass;
+import org.opensearch.sql.ast.dsl.AstDSL;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.expression.Expression;
 import org.opensearch.sql.expression.FunctionExpression;
+import org.opensearch.sql.expression.HighlightExpression;
 import org.opensearch.sql.expression.NamedArgumentExpression;
 import org.opensearch.sql.expression.env.Environment;
 
@@ -50,6 +51,14 @@ public void register(BuiltinFunctionRepository repository) {
     repository.register(match_phrase(BuiltinFunctionName.MATCH_PHRASE));
     repository.register(match_phrase(BuiltinFunctionName.MATCHPHRASE));
     repository.register(match_phrase_prefix());
+    repository.register(highlight());
+  }
+
+  private static FunctionResolver highlight() {
+    FunctionName functionName = BuiltinFunctionName.HIGHLIGHT.getName();
+    FunctionSignature functionSignature = new FunctionSignature(functionName, List.of(STRING));
+    FunctionBuilder functionBuilder = arguments -> new HighlightExpression(arguments.get(0));
+    return new FunctionResolver(functionName, ImmutableMap.of(functionSignature, functionBuilder));
   }
 
   private static FunctionResolver match_bool_prefix() {

core/src/main/java/org/opensearch/sql/planner/logical/LogicalHighlight.java

@@ -0,0 +1,29 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.planner.logical;
+
+import java.util.Collections;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.ToString;
+import org.opensearch.sql.expression.Expression;
+
+@EqualsAndHashCode(callSuper = true)
+@Getter
+@ToString
+public class LogicalHighlight extends LogicalPlan {
+  private final Expression highlightField;
+
+  public LogicalHighlight(LogicalPlan childPlan, Expression field) {
+    super(Collections.singletonList(childPlan));
+    highlightField = field;
+  }
+
+  @Override
+  public <R, C> R accept(LogicalPlanNodeVisitor<R, C> visitor, C context) {
+    return visitor.visitHighlight(this, context);
+  }
+}

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanDSL.java

@@ -62,6 +62,10 @@ public LogicalPlan window(LogicalPlan input,
     return new LogicalWindow(input, windowFunction, windowDefinition);
   }
 
+  public LogicalPlan highlight(LogicalPlan input, Expression field) {
+    return new LogicalHighlight(input, field);
+  }
+
   public static LogicalPlan remove(LogicalPlan input, ReferenceExpression... fields) {
     return new LogicalRemove(input, ImmutableSet.copyOf(fields));
   }

core/src/main/java/org/opensearch/sql/planner/logical/LogicalPlanNodeVisitor.java

@@ -26,6 +26,10 @@ public R visitFilter(LogicalFilter plan, C context) {
     return visitNode(plan, context);
   }
 
+  public R visitHighlight(LogicalHighlight plan, C context) {
+    return visitNode(plan, context);
+  }
+
   public R visitAggregation(LogicalAggregation plan, C context) {
     return visitNode(plan, context);
   }

core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanDSL.java

@@ -105,5 +105,4 @@ public ValuesOperator values(List<LiteralExpression>... values) {
   public static LimitOperator limit(PhysicalPlan input, Integer limit, Integer offset) {
     return new LimitOperator(input, limit, offset);
   }
-
 }

core/src/main/java/org/opensearch/sql/planner/physical/PhysicalPlanNodeVisitor.java

@@ -79,6 +79,4 @@ public R visitMLCommons(PhysicalPlan node, C context) {
   public R visitAD(PhysicalPlan node, C context) {
     return visitNode(node, context);
   }
-
-
 }

core/src/test/java/org/opensearch/sql/analysis/AnalyzerTest.java

@@ -22,6 +22,7 @@
 import static org.opensearch.sql.ast.dsl.AstDSL.qualifiedName;
 import static org.opensearch.sql.ast.dsl.AstDSL.relation;
 import static org.opensearch.sql.ast.dsl.AstDSL.span;
+import static org.opensearch.sql.ast.dsl.AstDSL.stringLiteral;
 import static org.opensearch.sql.ast.tree.Sort.NullOrder;
 import static org.opensearch.sql.ast.tree.Sort.SortOption;
 import static org.opensearch.sql.ast.tree.Sort.SortOption.DEFAULT_ASC;
@@ -45,13 +46,15 @@
 import org.opensearch.sql.ast.dsl.AstDSL;
 import org.opensearch.sql.ast.expression.Argument;
 import org.opensearch.sql.ast.expression.DataType;
+import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.SpanUnit;
 import org.opensearch.sql.ast.tree.AD;
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.RareTopN.CommandType;
 import org.opensearch.sql.exception.SemanticCheckException;
 import org.opensearch.sql.expression.DSL;
+import org.opensearch.sql.expression.HighlightExpression;
 import org.opensearch.sql.expression.config.ExpressionConfig;
 import org.opensearch.sql.expression.window.WindowDefinition;
 import org.opensearch.sql.planner.logical.LogicalAD;
@@ -231,6 +234,22 @@ public void project_source() {
             AstDSL.alias("double_value", AstDSL.field("double_value"))));
   }
 
+  @Test
+  public void project_highlight() {
+    assertAnalyzeEqual(
+        LogicalPlanDSL.project(
+            LogicalPlanDSL.highlight(LogicalPlanDSL.relation("schema"),
+                DSL.literal("fieldA")),
+            DSL.named("highlight(fieldA)", new HighlightExpression(DSL.literal("fieldA")))
+        ),
+        AstDSL.projectWithArg(
+            AstDSL.relation("schema"),
+            AstDSL.defaultFieldsArgs(),
+            AstDSL.alias("highlight(fieldA)", new HighlightFunction(AstDSL.stringLiteral("fieldA")))
+        )
+    );
+  }
+
   @Test
   public void remove_source() {
     assertAnalyzeEqual(