Peng - drop the AST layer implementation and address the comments.

Signed-off-by: Jialiang Liang <jiallian@amazon.com>

Author: RyanL1997

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

@@ -76,7 +76,6 @@
 import org.opensearch.sql.ast.tree.Flatten;
 import org.opensearch.sql.ast.tree.GraphLookup;
 import org.opensearch.sql.ast.tree.Head;
-import org.opensearch.sql.ast.tree.Highlight;
 import org.opensearch.sql.ast.tree.Join;
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Limit;
@@ -571,11 +570,6 @@ public LogicalPlan visitGraphLookup(GraphLookup node, AnalysisContext context) {
     throw getOnlyForCalciteException("graphlookup");
   }
 
-  @Override
-  public LogicalPlan visitHighlight(Highlight node, AnalysisContext context) {
-    throw getOnlyForCalciteException("highlight");
-  }
-
   /** Build {@link ParseExpression} to context and skip to child nodes. */
   @Override
   public LogicalPlan visitParse(Parse node, AnalysisContext context) {

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

@@ -64,7 +64,6 @@
 import org.opensearch.sql.ast.tree.Flatten;
 import org.opensearch.sql.ast.tree.GraphLookup;
 import org.opensearch.sql.ast.tree.Head;
-import org.opensearch.sql.ast.tree.Highlight;
 import org.opensearch.sql.ast.tree.Join;
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Limit;
@@ -496,8 +495,4 @@ public T visitMvExpand(MvExpand node, C context) {
   public T visitGraphLookup(GraphLookup 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/dsl/AstDSL.java

@@ -60,8 +60,6 @@
 import org.opensearch.sql.ast.tree.FillNull;
 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.HighlightConfig;
 import org.opensearch.sql.ast.tree.Limit;
 import org.opensearch.sql.ast.tree.MinSpanBin;
 import org.opensearch.sql.ast.tree.MvCombine;
@@ -586,14 +584,6 @@ public static Trendline trendline(
     return new Trendline(sortField, Arrays.asList(computations)).attach(input);
   }
 
-  public static Highlight highlight(UnresolvedPlan input, HighlightConfig highlightConfig) {
-    return new Highlight(highlightConfig).attach(input);
-  }
-
-  public static Highlight highlight(UnresolvedPlan input, List<String> highlightFields) {
-    return new Highlight(new HighlightConfig(highlightFields)).attach(input);
-  }
-
   public static AppendPipe appendPipe(UnresolvedPlan input, UnresolvedPlan subquery) {
 
     return new AppendPipe(subquery).attach(input);

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/Highlight.java

@@ -5,18 +5,44 @@
 
 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) and optional 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}.
+ * 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(
-    List<String> fields, List<String> preTags, List<String> postTags, Integer fragmentSize) {
+    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> fields) {
-    this(fields, null, null, null);
+  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/ast/tree/HighlightConfig.java

@@ -98,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/CalcitePlanContext.java

@@ -123,7 +123,6 @@
 import org.opensearch.sql.ast.tree.GraphLookup;
 import org.opensearch.sql.ast.tree.GraphLookup.Direction;
 import org.opensearch.sql.ast.tree.Head;
-import org.opensearch.sql.ast.tree.Highlight;
 import org.opensearch.sql.ast.tree.Join;
 import org.opensearch.sql.ast.tree.Kmeans;
 import org.opensearch.sql.ast.tree.Lookup;
@@ -157,9 +156,9 @@
 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.LogicalHighlight;
 import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit;
 import org.opensearch.sql.calcite.plan.rel.LogicalSystemLimit.SystemLimitType;
 import org.opensearch.sql.calcite.utils.BinUtils;
@@ -229,16 +228,17 @@ public RelNode visitRelation(Relation node, CalcitePlanContext context) {
     context.relBuilder.scan(node.getTableQualifiedName().getParts());
     RelNode scan = context.relBuilder.peek();
 
-    if (scan instanceof AliasFieldsWrappable) {
-      ((AliasFieldsWrappable) scan).wrapProjectForAliasFields(context.relBuilder);
+    // 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
     }
 
-    // Wrap with LogicalHighlight if highlight config is set on context (from API request)
-    if (context.getHighlightConfig() != null) {
-      RelNode input = context.relBuilder.build();
-      LogicalHighlight highlight = LogicalHighlight.create(input, context.getHighlightConfig());
-      context.relBuilder.push(highlight);
-      context.setHighlightConfig(null); // Clear for join safety
+    if (scan instanceof AliasFieldsWrappable) {
+      ((AliasFieldsWrappable) scan).wrapProjectForAliasFields(context.relBuilder);
     }
 
     return context.relBuilder.peek();
@@ -3199,13 +3199,6 @@ static ChartConfig fromArguments(ArgumentMap argMap) {
     }
   }
 
-  @Override
-  public RelNode visitHighlight(Highlight node, CalcitePlanContext context) {
-    context.setHighlightConfig(node.getHighlightConfig());
-    visitChildren(node, context);
-    return context.relBuilder.peek();
-  }
-
   @Override
   public RelNode visitTrendline(Trendline node, CalcitePlanContext context) {
     visitChildren(node, context);

core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.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);
+}