fix pushdown and add IT and explain IT

RyanL1997 · RyanL1997 · commit a7c44d82e635 · 2026-03-15T13:49:57.000-07:00
Signed-off-by: Jialiang Liang &lt;jiallian@amazon.com&gt;
diff --git a/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java b/integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteExplainIT.java
@@ -2795,4 +2795,29 @@ public void testNoMvWithEval() throws IOException {
         "Expected explain to contain both CONCAT and ARRAY_JOIN",
         result.toLowerCase().contains("concat") && result.toLowerCase().contains("array_join"));
   }
+
+  @Test
+  public void testHighlightWildcardExplain() throws IOException {
+    String query = "source=" + TEST_INDEX_ACCOUNT + " | highlight *";
+    var result = explainQueryYaml(query);
+    String expected = loadExpectedPlan("explain_highlight_wildcard.yaml");
+    assertYamlEqualsIgnoreId(expected, result);
+  }
+
+  @Test
+  public void testHighlightSingleTermExplain() throws IOException {
+    String query = "source=" + TEST_INDEX_ACCOUNT + " | highlight \\\"Holmes\\\"";
+    var result = explainQueryYaml(query);
+    String expected = loadExpectedPlan("explain_highlight_single_term.yaml");
+    assertYamlEqualsIgnoreId(expected, result);
+  }
+
+  @Test
+  public void testHighlightWithFilterExplain() throws IOException {
+    String query =
+        "source=" + TEST_INDEX_ACCOUNT + " | highlight * | where age > 30 | fields firstname, age";
+    var result = explainQueryYaml(query);
+    String expected = loadExpectedPlan("explain_highlight_with_filter.yaml");
+    assertYamlEqualsIgnoreId(expected, result);
+  }
 }
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_highlight_single_term.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_highlight_single_term.yaml
@@ -0,0 +1,8 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _highlight=[$17])
+        LogicalHighlight(highlightArgs=[[Holmes]])
+          CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[HIGHLIGHT->[Holmes], PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname, _highlight], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["account_number","firstname","address","balance","gender","city","employer","state","age","email","lastname"],"excludes":[]},"highlight":{"fragment_size":2147483647,"fields":{"*":{"highlight_query":{"query_string":{"query":"\"Holmes\"","default_field":"*","fields":[],"type":"best_fields","default_operator":"or","max_determinized_states":10000,"enable_position_increments":true,"fuzziness":"AUTO","fuzzy_prefix_length":0,"fuzzy_max_expansions":50,"phrase_slop":0,"escape":false,"auto_generate_synonyms_phrase_query":true,"fuzzy_transpositions":true,"boost":1.0}}}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_highlight_wildcard.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_highlight_wildcard.yaml
@@ -0,0 +1,8 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _highlight=[$17])
+        LogicalHighlight(highlightArgs=[[*]])
+          CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[HIGHLIGHT->[*], PROJECT->[account_number, firstname, address, balance, gender, city, employer, state, age, email, lastname, _highlight], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","_source":{"includes":["account_number","firstname","address","balance","gender","city","employer","state","age","email","lastname"],"excludes":[]},"highlight":{"fragment_size":2147483647,"fields":{"*":{}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
diff --git a/integ-test/src/test/resources/expectedOutput/calcite/explain_highlight_with_filter.yaml b/integ-test/src/test/resources/expectedOutput/calcite/explain_highlight_with_filter.yaml
@@ -0,0 +1,9 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(firstname=[$1], age=[$8], _highlight=[$17])
+        LogicalFilter(condition=[>($8, 30)])
+          LogicalHighlight(highlightArgs=[[*]])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[HIGHLIGHT->[*], FILTER->>($8, 30), PROJECT->[firstname, age, _highlight], LIMIT->10000], OpenSearchRequestBuilder(sourceBuilder={"from":0,"size":10000,"timeout":"1m","query":{"range":{"age":{"from":30,"to":null,"include_lower":false,"include_upper":true,"boost":1.0}}},"_source":{"includes":["firstname","age"],"excludes":[]},"highlight":{"fragment_size":2147483647,"fields":{"*":{}}}}, requestedTotalSize=10000, pageSize=null, startFrom=0)])
diff --git a/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_highlight_single_term.yaml b/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_highlight_single_term.yaml
@@ -0,0 +1,10 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _highlight=[$17])
+        LogicalHighlight(highlightArgs=[[Holmes]])
+          CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableCalc(expr#0..17=[{inputs}], proj#0..10=[{exprs}], _highlight=[$t17])
+        CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[HIGHLIGHT->[Holmes]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","highlight":{"fragment_size":2147483647,"fields":{"*":{"highlight_query":{"query_string":{"query":"\"Holmes\"","default_field":"*","fields":[],"type":"best_fields","default_operator":"or","max_determinized_states":10000,"enable_position_increments":true,"fuzziness":"AUTO","fuzzy_prefix_length":0,"fuzzy_max_expansions":50,"phrase_slop":0,"escape":false,"auto_generate_synonyms_phrase_query":true,"fuzzy_transpositions":true,"boost":1.0}}}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
diff --git a/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_highlight_wildcard.yaml b/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_highlight_wildcard.yaml
@@ -0,0 +1,10 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(account_number=[$0], firstname=[$1], address=[$2], balance=[$3], gender=[$4], city=[$5], employer=[$6], state=[$7], age=[$8], email=[$9], lastname=[$10], _highlight=[$17])
+        LogicalHighlight(highlightArgs=[[*]])
+          CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableCalc(expr#0..17=[{inputs}], proj#0..10=[{exprs}], _highlight=[$t17])
+        CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[HIGHLIGHT->[*]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","highlight":{"fragment_size":2147483647,"fields":{"*":{}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
diff --git a/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_highlight_with_filter.yaml b/integ-test/src/test/resources/expectedOutput/calcite_no_pushdown/explain_highlight_with_filter.yaml
@@ -0,0 +1,11 @@
+calcite:
+  logical: |
+    LogicalSystemLimit(fetch=[10000], type=[QUERY_SIZE_LIMIT])
+      LogicalProject(firstname=[$1], age=[$8], _highlight=[$17])
+        LogicalFilter(condition=[>($8, 30)])
+          LogicalHighlight(highlightArgs=[[*]])
+            CalciteLogicalIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]])
+  physical: |
+    EnumerableLimit(fetch=[10000])
+      EnumerableCalc(expr#0..17=[{inputs}], expr#18=[30], expr#19=[>($t8, $t18)], firstname=[$t1], age=[$t8], _highlight=[$t17], $condition=[$t19])
+        CalciteEnumerableIndexScan(table=[[OpenSearch, opensearch-sql_test_index_account]], PushDownContext=[[HIGHLIGHT->[*]], OpenSearchRequestBuilder(sourceBuilder={"from":0,"timeout":"1m","highlight":{"fragment_size":2147483647,"fields":{"*":{}}}}, requestedTotalSize=2147483647, pageSize=null, startFrom=0)])
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/AbstractCalciteIndexScan.java
@@ -130,7 +130,7 @@ public double estimateRowCount(RelMetadataQuery mq) {
             (rowCount, operation) ->
                 switch (operation.type()) {
                   case AGGREGATION -> mq.getRowCount((RelNode) operation.digest());
-                  case PROJECT, SORT, SORT_EXPR -> rowCount;
+                  case PROJECT, SORT, SORT_EXPR, HIGHLIGHT -> rowCount;
                   case SORT_AGG_METRICS ->
                       NumberUtil.min(rowCount, osIndex.getQueryBucketSize().doubleValue());
                   // Refer the org.apache.calcite.rel.metadata.RelMdRowCount
@@ -176,8 +176,8 @@ public double estimateRowCount(RelMetadataQuery mq) {
           dRows = mq.getRowCount((RelNode) operation.digest());
           dCpu += dRows * getAggMultiplier(operation);
         }
-        // Ignored Project in cost accumulation, but it will affect the external cost
-        case PROJECT -> {}
+        // Ignored Project and Highlight in cost accumulation, but they affect the external cost
+        case PROJECT, HIGHLIGHT -> {}
         case SORT -> dCpu += dRows;
         case SORT_AGG_METRICS -> {
           dRows = dRows * .9 / 10; // *.9 because always bucket IS_NOT_NULL
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteEnumerableIndexScan.java
@@ -119,7 +119,6 @@ public Result implement(EnumerableRelImplementor implementor, Prefer pref) {
       @Override
       public Enumerator<Object> enumerator() {
         OpenSearchRequestBuilder requestBuilder = pushDownContext.createRequestBuilder();
-        applyHighlightPushDown(requestBuilder, pushDownContext.getHighlightArgs());
         return new OpenSearchIndexEnumerator(
             osIndex.getClient(),
             getRowType().getFieldNames(),
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java
@@ -92,7 +92,13 @@ public RelNode pushDownHighlight(List<String> highlightArgs) {
         HighlightExpression.HIGHLIGHT_FIELD,
         getCluster().getTypeFactory().createSqlType(SqlTypeName.ANY));
     CalciteLogicalIndexScan newScan = copyWithNewSchema(schemaBuilder.build());
-    newScan.getPushDownContext().setHighlightArgs(highlightArgs);
+    newScan
+        .getPushDownContext()
+        .add(
+            PushDownType.HIGHLIGHT,
+            highlightArgs,
+            (OSRequestBuilderAction)
+                requestBuilder -> applyHighlightPushDown(requestBuilder, highlightArgs));
     return newScan;
   }
 
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownContext.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownContext.java
@@ -42,8 +42,6 @@ public class PushDownContext extends AbstractCollection<PushDownOperation> {
   private boolean isRareTopPushed = false;
   private boolean isScriptPushed = false;
 
-  @lombok.Setter private List<String> highlightArgs;
-
   public PushDownContext(OpenSearchIndex osIndex) {
     this.osIndex = osIndex;
   }
@@ -52,7 +50,6 @@ public PushDownContext(OpenSearchIndex osIndex) {
   public PushDownContext clone() {
     PushDownContext newContext = new PushDownContext(osIndex);
     newContext.addAll(this);
-    newContext.highlightArgs = this.highlightArgs;
     return newContext;
   }
 
diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownType.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/context/PushDownType.java
@@ -15,7 +15,7 @@ public enum PushDownType {
   SCRIPT, // script in predicate
   SORT_AGG_METRICS, // convert composite aggregate to terms or multi-terms bucket aggregate
   RARE_TOP, // convert composite aggregate to nested aggregate
-  SORT_EXPR
-  // HIGHLIGHT,
+  SORT_EXPR,
+  HIGHLIGHT
   // NESTED
 }
diff --git a/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLHighlightTest.java b/ppl/src/test/java/org/opensearch/sql/ppl/calcite/CalcitePPLHighlightTest.java
@@ -0,0 +1,91 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.ppl.calcite;
+
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.test.CalciteAssert;
+import org.junit.Test;
+
+public class CalcitePPLHighlightTest extends CalcitePPLAbstractTest {
+
+  public CalcitePPLHighlightTest() {
+    super(CalciteAssert.SchemaSpec.SCOTT_WITH_TEMPORAL);
+  }
+
+  @Test
+  public void testHighlightWildcard() {
+    String ppl = "source=EMP | highlight *";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalHighlight(highlightArgs=[[*]])\n" + "  LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+  }
+
+  @Test
+  public void testHighlightSingleTerm() {
+    String ppl = "source=EMP | highlight \"error\"";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalHighlight(highlightArgs=[[error]])\n"
+            + "  LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+  }
+
+  @Test
+  public void testHighlightMultipleTerms() {
+    String ppl = "source=EMP | highlight \"error\", \"login\"";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalHighlight(highlightArgs=[[error, login]])\n"
+            + "  LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+  }
+
+  @Test
+  public void testHighlightWithFilter() {
+    String ppl = "source=EMP | highlight * | where SAL > 2000";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalFilter(condition=[>($5, 2000)])\n"
+            + "  LogicalHighlight(highlightArgs=[[*]])\n"
+            + "    LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+  }
+
+  @Test
+  public void testHighlightWithFields() {
+    String ppl = "source=EMP | highlight \"error\" | fields ENAME, JOB";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalProject(ENAME=[$1], JOB=[$2], _highlight=[$8])\n"
+            + "  LogicalHighlight(highlightArgs=[[error]])\n"
+            + "    LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+  }
+
+  @Test
+  public void testHighlightWithSort() {
+    String ppl = "source=EMP | highlight * | sort SAL";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalSort(sort0=[$5], dir0=[ASC-nulls-first])\n"
+            + "  LogicalHighlight(highlightArgs=[[*]])\n"
+            + "    LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+  }
+
+  @Test
+  public void testHighlightWithFilterAndFields() {
+    String ppl = "source=EMP | highlight \"error\" | where SAL > 2000 | fields ENAME, SAL";
+    RelNode root = getRelNode(ppl);
+    String expectedLogical =
+        "LogicalProject(ENAME=[$1], SAL=[$5], _highlight=[$8])\n"
+            + "  LogicalFilter(condition=[>($5, 2000)])\n"
+            + "    LogicalHighlight(highlightArgs=[[error]])\n"
+            + "      LogicalTableScan(table=[[scott, EMP]])\n";
+    verifyLogical(root, expectedLogical);
+  }
+}
