[Feature] Calcite PPL search result highlighting

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

@@ -45,6 +45,32 @@ public class CalcitePlanContext {
   private static final ThreadLocal<Boolean> legacyPreferredFlag =
       ThreadLocal.withInitial(() -> true);
 
+  /**
+   * Thread-local extra search-source JSON from the PPL request body. Set by PPLService before query
+   * execution and read by AbstractCalciteIndexScan when building the OpenSearch request. The JSON
+   * string is parsed via {@code SearchSourceBuilder.fromXContent()} and selectively merged into the
+   * index scan request.
+   */
+  private static final ThreadLocal<String> extraSearchSource = new ThreadLocal<>();
+
+  public static void setExtraSearchSource(String json) {
+    extraSearchSource.set(json);
+  }
+
+  public static String getExtraSearchSource() {
+    return extraSearchSource.get();
+  }
+
+  public static void clearExtraSearchSource() {
+    extraSearchSource.remove();
+  }
+
+  /** Convenience check: does the extra search source contain a highlight clause? */
+  public static boolean hasHighlightInExtraSearchSource() {
+    String extra = extraSearchSource.get();
+    return extra != null && extra.contains("\"highlight\"");
+  }
+
   @Getter @Setter private boolean isResolvingJoinCondition = false;
   @Getter @Setter private boolean isResolvingSubquery = false;
   @Getter @Setter private boolean inCoalesceFunction = false;

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

@@ -23,6 +23,7 @@
 import static org.opensearch.sql.calcite.utils.PlanUtils.getRelation;
 import static org.opensearch.sql.calcite.utils.PlanUtils.getRexCall;
 import static org.opensearch.sql.calcite.utils.PlanUtils.transformPlanToAttachChild;
+import static org.opensearch.sql.expression.HighlightExpression.HIGHLIGHT_FIELD;
 import static org.opensearch.sql.utils.SystemIndexUtils.DATASOURCES_TABLE_NAME;
 
 import com.google.common.base.Strings;
@@ -407,6 +408,15 @@ public RelNode visitProject(Project node, CalcitePlanContext context) {
       if (!context.isResolvingSubquery()) {
         context.setProjectVisited(true);
       }
+      // When highlight is active, include _highlight in the projection so it survives
+      // through the Calcite pipeline. This matches DSL behavior where _source filtering
+      // does not affect highlights.
+      if (CalcitePlanContext.hasHighlightInExtraSearchSource()) {
+        int hlIndex = currentFields.indexOf(HIGHLIGHT_FIELD);
+        if (hlIndex >= 0) {
+          expandedFields.add(context.relBuilder.field(hlIndex));
+        }
+      }
       context.relBuilder.project(expandedFields);
     }
     return context.relBuilder.peek();

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

@@ -7,6 +7,7 @@
 
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
+import lombok.Setter;
 import org.opensearch.sql.ast.statement.ExplainMode;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.ExecutionEngine;
@@ -22,6 +23,15 @@ public abstract class AbstractPlan {
 
   @Getter protected final QueryType queryType;
 
+  /**
+   * Extra search-source JSON from the PPL request body. Set by PPLService before submitting the
+   * plan to the query manager. The plan carries this across the thread boundary (REST handler
+   * thread → sql-worker thread), and the worker thread sets it as a ThreadLocal before Calcite
+   * planning and execution begin. The JSON is later parsed via {@code
+   * SearchSourceBuilder.fromXContent()} and selectively merged into the index scan request.
+   */
+  @Getter @Setter private String extraSearchSource;
+
   /** Start query execution. */
   public abstract void execute();
 

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

@@ -6,6 +6,7 @@
 package org.opensearch.sql.executor.execution;
 
 import org.opensearch.sql.ast.statement.ExplainMode;
+import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.executor.QueryId;
@@ -34,7 +35,19 @@ public ExplainPlan(
 
   @Override
   public void execute() {
-    plan.explain(explainListener, mode);
+    setExtraSearchSourceThreadLocal();
+    try {
+      plan.explain(explainListener, mode);
+    } finally {
+      CalcitePlanContext.clearExtraSearchSource();
+    }
+  }
+
+  private void setExtraSearchSourceThreadLocal() {
+    String extra = getExtraSearchSource();
+    if (extra != null) {
+      CalcitePlanContext.setExtraSearchSource(extra);
+    }
   }
 
   @Override

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

@@ -10,6 +10,7 @@
 import org.opensearch.sql.ast.statement.ExplainMode;
 import org.opensearch.sql.ast.tree.Paginate;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.executor.QueryId;
@@ -60,10 +61,15 @@ public QueryPlan(
 
   @Override
   public void execute() {
-    if (pageSize.isPresent()) {
-      queryService.execute(new Paginate(pageSize.get(), plan), getQueryType(), listener);
-    } else {
-      queryService.execute(plan, getQueryType(), listener);
+    setExtraSearchSourceThreadLocal();
+    try {
+      if (pageSize.isPresent()) {
+        queryService.execute(new Paginate(pageSize.get(), plan), getQueryType(), listener);
+      } else {
+        queryService.execute(plan, getQueryType(), listener);
+      }
+    } finally {
+      CalcitePlanContext.clearExtraSearchSource();
     }
   }
 
@@ -78,4 +84,11 @@ public void explain(
       queryService.explain(plan, getQueryType(), listener, mode);
     }
   }
+
+  private void setExtraSearchSourceThreadLocal() {
+    String extra = getExtraSearchSource();
+    if (extra != null) {
+      CalcitePlanContext.setExtraSearchSource(extra);
+    }
+  }
 }

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

@@ -23,6 +23,9 @@
 /** Highlight Expression. */
 @Getter
 public class HighlightExpression extends FunctionExpression {
+  /** The field name used to store highlight data on ExprTupleValue rows. */
+  public static final String HIGHLIGHT_FIELD = "_highlight";
+
   private final Expression highlightField;
   private final ExprType type;
 
@@ -46,7 +49,7 @@ public HighlightExpression(Expression highlightField) {
    */
   @Override
   public ExprValue valueOf(Environment<Expression, ExprValue> valueEnv) {
-    String refName = "_highlight";
+    String refName = HIGHLIGHT_FIELD;
     // Not a wilcard expression
     if (this.type == ExprCoreType.ARRAY) {
       refName += "." + StringUtils.unquoteText(getHighlightField().toString());

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

@@ -6,17 +6,21 @@
 package org.opensearch.sql.executor.execution;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.sql.ast.statement.ExplainMode;
+import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.ExecutionEngine;
 import org.opensearch.sql.executor.QueryId;
@@ -56,4 +60,36 @@ public void explainThrowException() {
             });
     assertEquals("explain query can not been explained.", unsupportedExplainException.getMessage());
   }
+
+  @Test
+  public void execute_sets_extra_search_source_threadlocal_for_explain() {
+    String extraSearchSource = "{\"highlight\":{\"fields\":{\"*\":{}}}}";
+    AtomicReference<String> captured = new AtomicReference<>();
+
+    doAnswer(
+            invocation -> {
+              captured.set(CalcitePlanContext.getExtraSearchSource());
+              return null;
+            })
+        .when(queryPlan)
+        .explain(any(), any());
+
+    ExplainPlan explainPlan = new ExplainPlan(queryId, queryType, queryPlan, mode, explainListener);
+    explainPlan.setExtraSearchSource(extraSearchSource);
+    explainPlan.execute();
+
+    assertEquals(extraSearchSource, captured.get());
+  }
+
+  @Test
+  public void execute_clears_extra_search_source_threadlocal_after_explain() {
+    String extraSearchSource = "{\"highlight\":{\"fields\":{\"*\":{}}}}";
+    doNothing().when(queryPlan).explain(any(), any());
+
+    ExplainPlan explainPlan = new ExplainPlan(queryId, queryType, queryPlan, mode, explainListener);
+    explainPlan.setExtraSearchSource(extraSearchSource);
+    explainPlan.execute();
+
+    assertNull(CalcitePlanContext.getExtraSearchSource());
+  }
 }

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

@@ -5,14 +5,18 @@
 
 package org.opensearch.sql.executor.execution;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.commons.lang3.NotImplementedException;
 import org.junit.jupiter.api.DisplayNameGeneration;
 import org.junit.jupiter.api.DisplayNameGenerator;
@@ -22,6 +26,7 @@
 import org.mockito.junit.jupiter.MockitoExtension;
 import org.opensearch.sql.ast.statement.ExplainMode;
 import org.opensearch.sql.ast.tree.UnresolvedPlan;
+import org.opensearch.sql.calcite.CalcitePlanContext;
 import org.opensearch.sql.common.response.ResponseListener;
 import org.opensearch.sql.executor.DefaultExecutionEngine;
 import org.opensearch.sql.executor.ExecutionEngine;
@@ -126,4 +131,35 @@ public void onFailure(Exception e) {
             },
             mode);
   }
+
+  @Test
+  public void execute_sets_extra_search_source_threadlocal() {
+    String extraSearchSource = "{\"highlight\":{\"fields\":{\"*\":{}}}}";
+    AtomicReference<String> captured = new AtomicReference<>();
+
+    doAnswer(
+            invocation -> {
+              captured.set(CalcitePlanContext.getExtraSearchSource());
+              return null;
+            })
+        .when(queryService)
+        .execute(any(), any(), any());
+
+    QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
+    query.setExtraSearchSource(extraSearchSource);
+    query.execute();
+
+    assertEquals(extraSearchSource, captured.get());
+  }
+
+  @Test
+  public void execute_clears_extra_search_source_threadlocal_after_execution() {
+    String extraSearchSource = "{\"highlight\":{\"fields\":{\"*\":{}}}}";
+
+    QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener);
+    query.setExtraSearchSource(extraSearchSource);
+    query.execute();
+
+    assertNull(CalcitePlanContext.getExtraSearchSource());
+  }
 }