[Feature] PPL Highlight Support (#5234)

Author: RyanL1997

core/src/main/java/org/opensearch/sql/ast/statement/Query.java

@@ -11,6 +11,7 @@
 import lombok.Setter;
 import lombok.ToString;
 import org.opensearch.sql.ast.AbstractNodeVisitor;
+import org.opensearch.sql.ast.tree.HighlightConfig;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.executor.QueryType;
 
@@ -25,6 +26,7 @@ public class Query extends Statement {
   protected final UnresolvedPlan plan;
   protected final int fetchSize;
   private final QueryType queryType;
+  private HighlightConfig highlightConfig;
 
   @Override
   public <R, C> R accept(AbstractNodeVisitor<R, C> visitor, C context) {

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

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ast.tree;
+
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Bundles highlight configuration: field names (or wildcards) with per-field options, and optional
+ * global pre/post tags and fragment size. Supports both the simple array format ({@code ["*"]}) and
+ * the rich OSD object format with {@code pre_tags}, {@code post_tags}, {@code fields}, and {@code
+ * fragment_size}.
+ *
+ * <p>The {@code fields} map keys are field names or wildcards; the values are per-field option maps
+ * that are passed through to the OpenSearch highlight builder (e.g. {@code fragment_size}, {@code
+ * number_of_fragments}, {@code type}).
+ */
+public record HighlightConfig(
+    Map<String, Map<String, Object>> fields,
+    List<String> preTags,
+    List<String> postTags,
+    Integer fragmentSize) {
+
+  /** Convenience constructor for the simple array format (fields only, no tag/size overrides). */
+  public HighlightConfig(List<String> fieldNames) {
+    this(toFieldMap(fieldNames), null, null, null);
+  }
+
+  /** Returns the field names as a list (for display and iteration). */
+  public List<String> fieldNames() {
+    return fields == null ? List.of() : List.copyOf(fields.keySet());
+  }
+
+  private static Map<String, Map<String, Object>> toFieldMap(List<String> fieldNames) {
+    if (fieldNames == null) {
+      return null;
+    }
+    Map<String, Map<String, Object>> map = new LinkedHashMap<>();
+    for (String name : fieldNames) {
+      map.put(name, Map.of());
+    }
+    return map;
+  }
+}

core/src/main/java/org/opensearch/sql/calcite/CalcitePlanContext.java

@@ -22,6 +22,7 @@
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.tools.FrameworkConfig;
 import org.opensearch.sql.ast.expression.UnresolvedExpression;
+import org.opensearch.sql.ast.tree.HighlightConfig;
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper;
 import org.opensearch.sql.calcite.utils.CalciteToolsHelper.OpenSearchRelBuilder;
 import org.opensearch.sql.common.setting.Settings;
@@ -45,6 +46,7 @@ public class CalcitePlanContext {
   private static final ThreadLocal<Boolean> legacyPreferredFlag =
       ThreadLocal.withInitial(() -> true);
 
+  @Getter @Setter private HighlightConfig highlightConfig;
   @Getter @Setter private boolean isResolvingJoinCondition = false;
   @Getter @Setter private boolean isResolvingSubquery = false;
   @Getter @Setter private boolean inCoalesceFunction = false;
@@ -96,6 +98,7 @@ private CalcitePlanContext(CalcitePlanContext parent) {
     this.relBuilder = parent.relBuilder; // Share the same relBuilder
     this.rexBuilder = parent.rexBuilder; // Share the same rexBuilder
     this.functionProperties = parent.functionProperties;
+    this.highlightConfig = parent.highlightConfig;
     this.rexLambdaRefMap = new HashMap<>(); // New map for lambda variables
     this.capturedVariables = new ArrayList<>(); // New list for captured variables
     this.inLambdaContext = true; // Mark that we're inside a lambda

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java

@@ -156,6 +156,7 @@
 import org.opensearch.sql.ast.tree.Values;
 import org.opensearch.sql.ast.tree.Window;
 import org.opensearch.sql.calcite.plan.AliasFieldsWrappable;
+import org.opensearch.sql.calcite.plan.HighlightPushDown;
 import org.opensearch.sql.calcite.plan.OpenSearchConstants;
 import org.opensearch.sql.calcite.plan.rel.LogicalGraphLookup;
 import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit;
@@ -171,6 +172,7 @@
 import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.exception.CalciteUnsupportedException;
 import org.opensearch.sql.exception.SemanticCheckException;
+import org.opensearch.sql.expression.HighlightExpression;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 import org.opensearch.sql.expression.function.PPLFuncImpTable;
 import org.opensearch.sql.expression.parse.RegexCommonUtils;
@@ -225,10 +227,21 @@ public RelNode visitRelation(Relation node, CalcitePlanContext context) {
     }
     context.relBuilder.scan(node.getTableQualifiedName().getParts());
     RelNode scan = context.relBuilder.peek();
+
+    // Eagerly push down highlight config to the scan (highlight is a scan hint, not an operator)
+    if (context.getHighlightConfig() != null && scan instanceof HighlightPushDown) {
+      RelNode newScan = ((HighlightPushDown) scan).pushDownHighlight(context.getHighlightConfig());
+      context.relBuilder.build(); // pop old scan
+      context.relBuilder.push(newScan);
+      scan = newScan;
+      context.setHighlightConfig(null); // consumed
+    }
+
     if (scan instanceof AliasFieldsWrappable) {
-      return ((AliasFieldsWrappable) scan).wrapProjectForAliasFields(context.relBuilder);
+      ((AliasFieldsWrappable) scan).wrapProjectForAliasFields(context.relBuilder);
     }
-    return scan;
+
+    return context.relBuilder.peek();
   }
 
   // This is a tool method to add an existed RelOptTable to builder stack, not used for now
@@ -419,6 +432,12 @@ public RelNode visitProject(Project node, CalcitePlanContext context) {
     List<RexNode> expandedFields =
         expandProjectFields(node.getProjectList(), currentFields, context);
 
+    // Include _highlight in projections when the highlight column is present in the schema
+    int hlIndex = currentFields.indexOf(HighlightExpression.HIGHLIGHT_FIELD);
+    if (hlIndex >= 0) {
+      expandedFields.add(context.relBuilder.field(hlIndex));
+    }
+
     if (node.isExcluded()) {
       validateExclusion(expandedFields, currentFields);
       context.relBuilder.projectExcept(expandedFields);

core/src/main/java/org/opensearch/sql/calcite/plan/HighlightPushDown.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.plan;
+
+import org.apache.calcite.rel.RelNode;
+import org.opensearch.sql.ast.tree.HighlightConfig;
+
+/**
+ * Interface for scan nodes that support highlight pushdown. Highlight is a scan hint (not a
+ * relational operator), so it is pushed down eagerly during plan construction rather than via an
+ * optimizer rule.
+ */
+public interface HighlightPushDown {
+  RelNode pushDownHighlight(HighlightConfig highlightConfig);
+}

core/src/main/java/org/opensearch/sql/executor/QueryService.java

@@ -27,6 +27,7 @@
 import org.opensearch.sql.analysis.AnalysisContext;
 import org.opensearch.sql.analysis.Analyzer;
 import org.opensearch.sql.ast.statement.ExplainMode;
+import org.opensearch.sql.ast.tree.HighlightConfig;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.calcite.CalciteRelNodeVisitor;
@@ -87,8 +88,17 @@ public void execute(
       UnresolvedPlan plan,
       QueryType queryType,
       ResponseListener<ExecutionEngine.QueryResponse> listener) {
+    execute(plan, queryType, null, listener);
+  }
+
+  /** Execute with optional highlight config. */
+  public void execute(
+      UnresolvedPlan plan,
+      QueryType queryType,
+      HighlightConfig highlightConfig,
+      ResponseListener<ExecutionEngine.QueryResponse> listener) {
     if (shouldUseCalcite(queryType)) {
-      executeWithCalcite(plan, queryType, listener);
+      executeWithCalcite(plan, queryType, highlightConfig, listener);
     } else {
       executeWithLegacy(plan, queryType, listener, Optional.empty());
     }
@@ -100,8 +110,18 @@ public void explain(
       QueryType queryType,
       ResponseListener<ExecutionEngine.ExplainResponse> listener,
       ExplainMode mode) {
+    explain(plan, queryType, null, listener, mode);
+  }
+
+  /** Explain with optional highlight config. */
+  public void explain(
+      UnresolvedPlan plan,
+      QueryType queryType,
+      HighlightConfig highlightConfig,
+      ResponseListener<ExecutionEngine.ExplainResponse> listener,
+      ExplainMode mode) {
     if (shouldUseCalcite(queryType)) {
-      explainWithCalcite(plan, queryType, listener, mode);
+      explainWithCalcite(plan, queryType, highlightConfig, listener, mode);
     } else {
       explainWithLegacy(plan, queryType, listener, mode, Optional.empty());
     }
@@ -110,6 +130,7 @@ public void explain(
   public void executeWithCalcite(
       UnresolvedPlan plan,
       QueryType queryType,
+      HighlightConfig highlightConfig,
       ResponseListener<ExecutionEngine.QueryResponse> listener) {
     CalcitePlanContext.run(
         () -> {
@@ -121,6 +142,7 @@ public void executeWithCalcite(
             CalcitePlanContext context =
                 CalcitePlanContext.create(
                     buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
+            context.setHighlightConfig(highlightConfig);
             RelNode relNode = analyze(plan, context);
             RelNode calcitePlan = convertToCalcitePlan(relNode, context);
             analyzeMetric.set(System.nanoTime() - analyzeStart);
@@ -140,6 +162,7 @@ public void executeWithCalcite(
   public void explainWithCalcite(
       UnresolvedPlan plan,
       QueryType queryType,
+      HighlightConfig highlightConfig,
       ResponseListener<ExecutionEngine.ExplainResponse> listener,
       ExplainMode mode) {
     CalcitePlanContext.run(
@@ -149,6 +172,7 @@ public void explainWithCalcite(
             CalcitePlanContext context =
                 CalcitePlanContext.create(
                     buildFrameworkConfig(), SysLimit.fromSettings(settings), queryType);
+            context.setHighlightConfig(highlightConfig);
             context.run(
                 () -> {
                   RelNode relNode = analyze(plan, context);

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java

@@ -8,6 +8,7 @@
 import java.util.Optional;
 import org.apache.commons.lang3.NotImplementedException;
 import org.opensearch.sql.ast.statement.ExplainMode;
+import org.opensearch.sql.ast.tree.HighlightConfig;
 import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
 import org.opensearch.sql.common.response.ResponseListener;
@@ -29,18 +30,32 @@ public class QueryPlan extends AbstractPlan {
 
   protected final Optional<Integer> pageSize;
 
+  protected final HighlightConfig highlightConfig;
+
   /** Constructor. */
   public QueryPlan(
       QueryId queryId,
       QueryType queryType,
       UnresolvedPlan plan,
       QueryService queryService,
       ResponseListener<ExecutionEngine.QueryResponse> listener) {
+    this(queryId, queryType, plan, queryService, listener, null);
+  }
+
+  /** Constructor with highlight config. */
+  public QueryPlan(
+      QueryId queryId,
+      QueryType queryType,
+      UnresolvedPlan plan,
+      QueryService queryService,
+      ResponseListener<ExecutionEngine.QueryResponse> listener,
+      HighlightConfig highlightConfig) {
     super(queryId, queryType);
     this.plan = plan;
     this.queryService = queryService;
     this.listener = listener;
     this.pageSize = Optional.empty();
+    this.highlightConfig = highlightConfig;
   }
 
   /** Constructor with page size. */
@@ -56,14 +71,15 @@ public QueryPlan(
     this.queryService = queryService;
     this.listener = listener;
     this.pageSize = Optional.of(pageSize);
+    this.highlightConfig = null;
   }
 
   @Override
   public void execute() {
     if (pageSize.isPresent()) {
       queryService.execute(new Paginate(pageSize.get(), plan), getQueryType(), listener);
     } else {
-      queryService.execute(plan, getQueryType(), listener);
+      queryService.execute(plan, getQueryType(), highlightConfig, listener);
     }
   }
 
@@ -75,7 +91,7 @@ public void explain(
           new NotImplementedException(
               "`explain` feature for paginated requests is not implemented yet."));
     } else {
-      queryService.explain(plan, getQueryType(), listener, mode);
+      queryService.explain(plan, getQueryType(), highlightConfig, listener, mode);
     }
   }
 }

core/src/main/java/org/opensearch/sql/executor/execution/QueryPlanFactory.java

@@ -122,7 +122,12 @@ public AbstractPlan visitQuery(
       }
     } else {
       return new QueryPlan(
-          QueryId.queryId(), node.getQueryType(), node.getPlan(), queryService, context.getLeft());
+          QueryId.queryId(),
+          node.getQueryType(),
+          node.getPlan(),
+          queryService,
+          context.getLeft(),
+          node.getHighlightConfig());
     }
   }
 

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

@@ -23,6 +23,8 @@
 /** Highlight Expression. */
 @Getter
 public class HighlightExpression extends FunctionExpression {
+  public static final String HIGHLIGHT_FIELD = "_highlight";
+
   private final Expression highlightField;
   private final ExprType type;
 

core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java

@@ -52,15 +52,15 @@ public void execute_no_page_size() {
     QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
     query.execute();
 
-    verify(queryService, times(1)).execute(any(), any(), any());
+    verify(queryService, times(1)).execute(any(), any(), any(), any());
   }
 
   @Test
   public void explain_no_page_size() {
     QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
     query.explain(explainListener, mode);
 
-    verify(queryService, times(1)).explain(plan, queryType, explainListener, mode);
+    verify(queryService, times(1)).explain(plan, queryType, null, explainListener, mode);
   }
 
   @Test