Adding Highlight support in PPL. Adding wildcard support for highligh…

…t in SQL and PPL.

Signed-off-by: forestmvey <[email protected]>

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

@@ -10,6 +10,7 @@
 import static org.opensearch.sql.ast.tree.Sort.NullOrder.NULL_LAST;
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.ASC;
 import static org.opensearch.sql.ast.tree.Sort.SortOrder.DESC;
+import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 import static org.opensearch.sql.data.type.ExprCoreType.STRUCT;
 import static org.opensearch.sql.utils.MLCommonsConstants.RCF_ANOMALOUS;
 import static org.opensearch.sql.utils.MLCommonsConstants.RCF_ANOMALY_GRADE;
@@ -31,8 +32,10 @@
 import org.opensearch.sql.analysis.symbol.Namespace;
 import org.opensearch.sql.analysis.symbol.Symbol;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.Alias;
 import org.opensearch.sql.ast.expression.Argument;
 import org.opensearch.sql.ast.expression.Field;
+import org.opensearch.sql.ast.expression.HighlightFunction;
 import org.opensearch.sql.ast.expression.Let;
 import org.opensearch.sql.ast.expression.Literal;
 import org.opensearch.sql.ast.expression.Map;
@@ -43,6 +46,7 @@
 import org.opensearch.sql.ast.tree.Eval;
 import org.opensearch.sql.ast.tree.Filter;
 import org.opensearch.sql.ast.tree.Head;
+import org.opensearch.sql.ast.tree.Highlight;
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.Parse;
@@ -55,6 +59,7 @@
 import org.opensearch.sql.ast.tree.Sort.SortOption;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.ast.tree.Values;
+import org.opensearch.sql.common.utils.StringUtils;
 import org.opensearch.sql.data.model.ExprMissingValue;
 import org.opensearch.sql.data.type.ExprCoreType;
 import org.opensearch.sql.exception.SemanticCheckException;
@@ -71,6 +76,7 @@
 import org.opensearch.sql.planner.logical.LogicalDedupe;
 import org.opensearch.sql.planner.logical.LogicalEval;
 import org.opensearch.sql.planner.logical.LogicalFilter;
+import org.opensearch.sql.planner.logical.LogicalHighlight;
 import org.opensearch.sql.planner.logical.LogicalLimit;
 import org.opensearch.sql.planner.logical.LogicalMLCommons;
 import org.opensearch.sql.planner.logical.LogicalPlan;
@@ -329,6 +335,23 @@ public LogicalPlan visitEval(Eval node, AnalysisContext context) {
     return new LogicalEval(child, expressionsBuilder.build());
   }
 
+  /**
+   * Build {@link LogicalHighlight}.
+   */
+  @Override
+  public LogicalPlan visitHighlight(Highlight node, AnalysisContext context) {
+    LogicalPlan child = node.getChild().get(0).accept(this, context);
+
+    TypeEnvironment env = context.peek();
+    env.define(new Symbol(Namespace.FIELD_NAME,
+        (((Alias) node.getExpression()).getName())), STRING);
+
+    HighlightFunction unresolved = (HighlightFunction) ((Alias)node.getExpression()).getDelegated();
+    Expression field = expressionAnalyzer.analyze(unresolved.getHighlightField(), context);
+    return new LogicalHighlight(child, field);
+  }
+
+
   /**
    * Build {@link ParseExpression} to context and skip to child nodes.
    */

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

@@ -43,12 +43,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.common.utils.StringUtils;
 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;
 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,9 +192,11 @@ public Expression visitWindowFunction(WindowFunction node, AnalysisContext conte
   }
 
   @Override
-  public Expression visitHighlight(HighlightFunction node, AnalysisContext context) {
+  public Expression visitHighlightFunction(HighlightFunction node, AnalysisContext context) {
     Expression expr = node.getHighlightField().accept(this, context);
-    return new HighlightExpression(expr);
+    String highlightStr = "highlight(" + StringUtils.unquoteText(expr.toString()) + ")";
+    return new ReferenceExpression(highlightStr, expr.toString().contains("*")
+        ? ExprCoreType.STRUCT : ExprCoreType.ARRAY);
   }
 
   @Override

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

@@ -40,6 +40,7 @@
 import org.opensearch.sql.ast.tree.Eval;
 import org.opensearch.sql.ast.tree.Filter;
 import org.opensearch.sql.ast.tree.Head;
+import org.opensearch.sql.ast.tree.Highlight;
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.Parse;
@@ -256,7 +257,11 @@ public T visitAD(AD node, C context) {
     return visitChildren(node, context);
   }
 
-  public T visitHighlight(HighlightFunction node, C context) {
+  public T visitHighlightFunction(HighlightFunction node, C context) {
+    return visitChildren(node, context);
+  }
+
+  public T visitHighlight(Highlight node, C context) {
     return visitChildren(node, context);
   }
 }

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

@@ -24,7 +24,7 @@ public class HighlightFunction extends UnresolvedExpression {
 
   @Override
   public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
-    return nodeVisitor.visitHighlight(this, context);
+    return nodeVisitor.visitHighlightFunction(this, context);
   }
 
   @Override

core/src/main/java/org/opensearch/sql/ast/tree/Highlight.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+import lombok.EqualsAndHashCode;
+import lombok.Getter;
+import lombok.RequiredArgsConstructor;
+import lombok.Setter;
+import lombok.ToString;
+import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.expression.UnresolvedExpression;
+
+/**
+ * AST node represent Highlight operation.
+ */
+@Getter
+@Setter
+@ToString
+@EqualsAndHashCode(callSuper = false)
+@RequiredArgsConstructor
+public class Highlight extends UnresolvedPlan {
+  private final UnresolvedExpression expression;
+  private UnresolvedPlan child;
+
+  @Override
+  public Highlight attach(UnresolvedPlan child) {
+    this.child = child;
+    return this;
+  }
+
+  @Override
+  public List<UnresolvedPlan> getChild() {
+    return ImmutableList.of(this.child);
+  }
+
+  @Override
+  public <T, C> T accept(AbstractNodeVisitor<T, C> nodeVisitor, C context) {
+    return nodeVisitor.visitHighlight(this, context);
+  }
+}

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

@@ -55,10 +55,6 @@ 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/function/OpenSearchFunctions.java

@@ -15,13 +15,11 @@
 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;
 
@@ -51,14 +49,6 @@ 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() {