review - make highlight embeded into datarows/schemas

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

Author: RyanL1997

docs/user/ppl/interfaces/endpoint.md

@@ -204,7 +204,7 @@ Expected output (trimmed):
 
 ## Highlight
 
-You can add a `highlight` parameter to the PPL request body to enable search-result highlighting. This parameter follows the same semantics as the [OpenSearch highlight API](https://docs.opensearch.org/latest/search-plugins/searching-data/highlight/). When enabled, the response includes a top-level `highlights` array containing matching fragments with the specified tags. Each entry in the `highlights` array corresponds to the row at the same index in `datarows`.
+You can add a `highlight` parameter to the PPL request body to enable search-result highlighting. This parameter follows the same semantics as the [OpenSearch highlight API](https://docs.opensearch.org/latest/search-plugins/searching-data/highlight/). When enabled, the response includes a `_highlight` column in `schema` and `datarows` containing matching fragments with the specified tags. Each `_highlight` value in a datarow is an object whose keys are field names and whose values are arrays of highlight fragments for the corresponding row.
 
 Two formats are supported:
 
@@ -246,25 +246,21 @@ Expected output (trimmed):
   "schema": [
     { "name": "account_number", "type": "bigint" },
     { "name": "firstname", "type": "string" },
-    { "name": "lastname", "type": "string" }
+    { "name": "lastname", "type": "string" },
+    { "name": "_highlight", "type": "struct" }
   ],
   "datarows": [
-    [578, "Holmes", "Mcknight"],
-    [828, "Blanche", "Holmes"],
-    [1, "Amber", "Duke"]
-  ],
-  "highlights": [
-    {
+    [578, "Holmes", "Mcknight", {
       "firstname": ["@opensearch-dashboards-highlighted-field@Holmes@/opensearch-dashboards-highlighted-field@"],
       "firstname.keyword": ["@opensearch-dashboards-highlighted-field@Holmes@/opensearch-dashboards-highlighted-field@"]
-    },
-    {
+    }],
+    [828, "Blanche", "Holmes", {
       "lastname": ["@opensearch-dashboards-highlighted-field@Holmes@/opensearch-dashboards-highlighted-field@"],
       "lastname.keyword": ["@opensearch-dashboards-highlighted-field@Holmes@/opensearch-dashboards-highlighted-field@"]
-    },
-    {
+    }],
+    [1, "Amber", "Duke", {
       "address": ["880 @opensearch-dashboards-highlighted-field@Holmes@/opensearch-dashboards-highlighted-field@ Lane"]
-    }
+    }]
   ],
   "total": 3,
   "size": 3
@@ -292,8 +288,8 @@ Exceeding these limits returns an error.
 
 ### Notes
 
-- Highlighting requires a search term in the PPL statement (e.g. `source=accounts "Holmes"`). Without a search term (e.g. just `source=accounts`), the `highlights` array entries will be empty.
-- The `highlights` array in the response is parallel to `datarows` — each entry contains the highlighted fragments for the corresponding row.
+- Highlighting requires a search term in the PPL statement (e.g. `source=accounts "Holmes"`). Without a search term (e.g. just `source=accounts`), the `_highlight` values in datarows will be empty objects.
+- The `_highlight` column appears in `schema` and `datarows` as a regular column. Each `_highlight` value is an object whose keys are field names and whose values are arrays of highlight fragments.
 - In the simple array format, `["*"]` highlights all fields. Specific field names like `["firstname", "lastname"]` scope highlighting to those fields only.
 - In the object format, each key in the `fields` object is a field name or wildcard. Each value is an object of per-field highlight options. Supported per-field options: `fragment_size`, `number_of_fragments`, `type` (`plain`, `unified`, `fvh`), `pre_tags`, `post_tags`, `require_field_match`, `no_match_size`, `order`. Use `{}` for defaults. Example: `{"title": {"fragment_size": 200}, "body": {"type": "plain"}}`.
-- Highlights may include fields that are not explicitly projected in the `schema`/`datarows`. For example, using `{"*": {}}` highlights all fields that matched the search query, including fields not selected by `| fields`. In the example above, the `address` field appears in `highlights` because it contains a match ("880 Holmes Lane") even though only `account_number`, `firstname`, and `lastname` are projected.
+- Highlights may include fields that are not explicitly projected in the other columns. For example, using `{"*": {}}` highlights all fields that matched the search query, including fields not selected by `| fields`. In the example above, the `address` field appears in `_highlight` because it contains a match ("880 Holmes Lane") even though only `account_number`, `firstname`, and `lastname` are projected as separate columns.

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteHighlightIT.java

@@ -44,13 +44,15 @@ public void testHighlightWildcardWithSearchQuery() throws IOException {
   public void testHighlightContainsMatchingFragments() throws IOException {
     JSONObject result =
         executeQueryWithHighlight("source=" + TEST_INDEX_ACCOUNT + " \"Holmes\"", "[\"*\"]");
-    JSONArray highlights = result.getJSONArray("highlights");
-    assertTrue("highlights array should not be empty", highlights.length() > 0);
+    int hlIndex = getHighlightColumnIndex(result);
+    assertTrue("_highlight column should exist in schema", hlIndex >= 0);
+    JSONArray dataRows = result.getJSONArray("datarows");
+    assertTrue("datarows should not be empty", dataRows.length() > 0);
     // At least one highlight entry should have non-empty data
     boolean foundFragment = false;
-    for (int i = 0; i < highlights.length(); i++) {
-      JSONObject hlEntry = highlights.getJSONObject(i);
-      if (hlEntry.length() > 0) {
+    for (int i = 0; i < dataRows.length(); i++) {
+      JSONObject hlEntry = dataRows.getJSONArray(i).optJSONObject(hlIndex);
+      if (hlEntry != null && hlEntry.length() > 0) {
         foundFragment = true;
         break;
       }
@@ -69,11 +71,11 @@ public void testHighlightOsdObjectFormat() throws IOException {
     assertTrue(dataRows.length() > 0);
     assertHighlightsExist(result);
     // Verify custom tags are applied
-    JSONArray highlights = result.getJSONArray("highlights");
+    int hlIndex = getHighlightColumnIndex(result);
     boolean foundCustomTag = false;
-    for (int i = 0; i < highlights.length(); i++) {
-      String hlStr = highlights.getJSONObject(i).toString();
-      if (hlStr.contains("<b>")) {
+    for (int i = 0; i < dataRows.length(); i++) {
+      JSONObject hlEntry = dataRows.getJSONArray(i).optJSONObject(hlIndex);
+      if (hlEntry != null && hlEntry.toString().contains("<b>")) {
         foundCustomTag = true;
         break;
       }
@@ -93,11 +95,12 @@ public void testHighlightOsdObjectFormatWithDashboardsTags() throws IOException
     assertTrue(dataRows.length() > 0);
     assertHighlightsExist(result);
     // Verify dashboards tags are applied
-    JSONArray highlights = result.getJSONArray("highlights");
+    int hlIndex = getHighlightColumnIndex(result);
     boolean foundDashboardsTag = false;
-    for (int i = 0; i < highlights.length(); i++) {
-      String hlStr = highlights.getJSONObject(i).toString();
-      if (hlStr.contains("@opensearch-dashboards-highlighted-field@")) {
+    for (int i = 0; i < dataRows.length(); i++) {
+      JSONObject hlEntry = dataRows.getJSONArray(i).optJSONObject(hlIndex);
+      if (hlEntry != null
+          && hlEntry.toString().contains("@opensearch-dashboards-highlighted-field@")) {
         foundDashboardsTag = true;
         break;
       }
@@ -162,7 +165,7 @@ public void testHighlightWildcardInSearchText() throws IOException {
     JSONObject result =
         executeQueryWithHighlight("source=" + TEST_INDEX_ACCOUNT + " \"Holm*\"", "[\"*\"]");
     assertTrue("Response should contain datarows", result.has("datarows"));
-    assertTrue("Response should contain highlights array", result.has("highlights"));
+    assertHighlightsExist(result);
   }
 
   @Test
@@ -172,7 +175,7 @@ public void testHighlightMixedFullTextAndStructuredFilter() throws IOException {
             "source=" + TEST_INDEX_ACCOUNT + " \"Holmes\" | where age > 30 | fields firstname, age",
             "[\"*\"]");
     assertTrue("Response should contain datarows", result.has("datarows"));
-    assertTrue("Response should contain highlights array", result.has("highlights"));
+    assertHighlightsExist(result);
   }
 
   @Test
@@ -201,7 +204,7 @@ public void testHighlightWithHead() throws IOException {
     assertTrue("Response should contain datarows", result.has("datarows"));
     JSONArray dataRows = result.getJSONArray("datarows");
     assertTrue("Should return at most 2 rows", dataRows.length() <= 2);
-    assertTrue("Response should contain highlights array", result.has("highlights"));
+    assertHighlightsExist(result);
   }
 
   @Test
@@ -222,13 +225,16 @@ public void testHighlightBooleanOrSearch() throws IOException {
     assertTrue("OR search should return results", dataRows.length() > 0);
     assertHighlightsExist(result);
     // Verify highlights contain fragments for both search terms
-    JSONArray highlights = result.getJSONArray("highlights");
+    int hlIndex = getHighlightColumnIndex(result);
     boolean foundHolmes = false;
     boolean foundBond = false;
-    for (int i = 0; i < highlights.length(); i++) {
-      String hlStr = highlights.getJSONObject(i).toString();
-      if (hlStr.contains("Holmes")) foundHolmes = true;
-      if (hlStr.contains("Bond")) foundBond = true;
+    for (int i = 0; i < dataRows.length(); i++) {
+      JSONObject hlEntry = dataRows.getJSONArray(i).optJSONObject(hlIndex);
+      if (hlEntry != null) {
+        String hlStr = hlEntry.toString();
+        if (hlStr.contains("Holmes")) foundHolmes = true;
+        if (hlStr.contains("Bond")) foundBond = true;
+      }
     }
     assertTrue("Highlights should contain Holmes fragments", foundHolmes);
     assertTrue("Highlights should contain Bond fragments", foundBond);
@@ -243,13 +249,16 @@ public void testHighlightBooleanAndSearch() throws IOException {
     assertTrue("AND search should return results", dataRows.length() > 0);
     assertHighlightsExist(result);
     // Verify highlights contain fragments for both terms
-    JSONArray highlights = result.getJSONArray("highlights");
+    int hlIndex = getHighlightColumnIndex(result);
     boolean foundHolmes = false;
     boolean foundLane = false;
-    for (int i = 0; i < highlights.length(); i++) {
-      String hlStr = highlights.getJSONObject(i).toString();
-      if (hlStr.contains("Holmes")) foundHolmes = true;
-      if (hlStr.contains("Lane")) foundLane = true;
+    for (int i = 0; i < dataRows.length(); i++) {
+      JSONObject hlEntry = dataRows.getJSONArray(i).optJSONObject(hlIndex);
+      if (hlEntry != null) {
+        String hlStr = hlEntry.toString();
+        if (hlStr.contains("Holmes")) foundHolmes = true;
+        if (hlStr.contains("Lane")) foundLane = true;
+      }
     }
     assertTrue("Highlights should contain Holmes fragments", foundHolmes);
     assertTrue("Highlights should contain Lane fragments", foundLane);
@@ -282,10 +291,10 @@ public void testHighlightNoSearchQuery() throws IOException {
   }
 
   @Test
-  public void testWithoutHighlightNoHighlightArray() throws IOException {
-    // Without highlight parameter, highlights array should NOT appear
+  public void testWithoutHighlightNoHighlightColumn() throws IOException {
+    // Without highlight parameter, _highlight column should NOT appear in schema
     JSONObject result = executeQuery("source=" + TEST_INDEX_BANK);
-    assertFalse("Response should NOT contain highlights array", result.has("highlights"));
+    assertTrue("_highlight column should NOT be in schema", getHighlightColumnIndex(result) < 0);
   }
 
   /**
@@ -316,10 +325,26 @@ protected JSONObject executeQueryWithHighlight(String query, String highlightJso
     return jsonify(getResponseBody(response, true));
   }
 
-  /** Assert that the response contains a non-empty highlights array. */
+  /**
+   * Find the index of the _highlight column in the schema array.
+   *
+   * @return the column index, or -1 if not present
+   */
+  private int getHighlightColumnIndex(JSONObject result) {
+    JSONArray schema = result.getJSONArray("schema");
+    for (int i = 0; i < schema.length(); i++) {
+      if ("_highlight".equals(schema.getJSONObject(i).getString("name"))) {
+        return i;
+      }
+    }
+    return -1;
+  }
+
+  /** Assert that the response contains a _highlight column with non-empty highlight data. */
   private void assertHighlightsExist(JSONObject result) {
-    assertTrue("Response should contain highlights array", result.has("highlights"));
-    JSONArray highlights = result.getJSONArray("highlights");
-    assertTrue("Highlights array should not be empty", highlights.length() > 0);
+    int hlIndex = getHighlightColumnIndex(result);
+    assertTrue("Schema should contain _highlight column", hlIndex >= 0);
+    JSONArray dataRows = result.getJSONArray("datarows");
+    assertTrue("datarows should not be empty", dataRows.length() > 0);
   }
 }

protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java

@@ -5,15 +5,11 @@
 
 package org.opensearch.sql.protocol.response;
 
-import static org.opensearch.sql.expression.HighlightExpression.HIGHLIGHT_FIELD;
-
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
-import java.util.List;
 import java.util.Locale;
 import java.util.Map;
-import java.util.stream.Collectors;
 import lombok.Getter;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
@@ -75,7 +71,6 @@ public int size() {
   public Map<String, String> columnNameTypes() {
     Map<String, String> colNameTypes = new LinkedHashMap<>();
     schema.getColumns().stream()
-        .filter(column -> !HIGHLIGHT_FIELD.equals(getColumnName(column)))
         .forEach(
             column ->
                 colNameTypes.put(
@@ -90,41 +85,10 @@ public Iterator<Object[]> iterator() {
         .map(ExprValueUtils::getTupleValue)
         .map(
             tuple ->
-                tuple.entrySet().stream()
-                    .filter(e -> !HIGHLIGHT_FIELD.equals(e.getKey()))
-                    .map(e -> e.getValue().value())
-                    .toArray(Object[]::new))
+                tuple.entrySet().stream().map(e -> e.getValue().value()).toArray(Object[]::new))
         .iterator();
   }
 
-  /**
-   * Extract highlight data from each result row. Each row may contain a {@code _highlight} field
-   * added by {@code OpenSearchResponse.addHighlightsToBuilder()} and preserved through projection.
-   * Returns a list parallel to datarows where each entry is either a map of field name to highlight
-   * fragments, or null if no highlight data exists for that row.
-   */
-  public List<Map<String, Object>> highlights() {
-    return exprValues.stream()
-        .map(ExprValueUtils::getTupleValue)
-        .map(
-            tuple -> {
-              ExprValue hl = tuple.get(HIGHLIGHT_FIELD);
-              if (hl == null || hl.isMissing() || hl.isNull()) {
-                return null;
-              }
-              Map<String, Object> hlMap = new LinkedHashMap<>();
-              for (Map.Entry<String, ExprValue> entry : hl.tupleValue().entrySet()) {
-                hlMap.put(
-                    entry.getKey(),
-                    entry.getValue().collectionValue().stream()
-                        .map(ExprValue::stringValue)
-                        .collect(Collectors.toList()));
-              }
-              return (Map<String, Object>) hlMap;
-            })
-        .collect(Collectors.toList());
-  }
-
   private String getColumnName(Column column) {
     return (column.getAlias() != null) ? column.getAlias() : column.getName();
   }

protocol/src/main/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java

@@ -6,8 +6,6 @@
 package org.opensearch.sql.protocol.response.format;
 
 import java.util.List;
-import java.util.Map;
-import java.util.Objects;
 import lombok.Builder;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
@@ -57,11 +55,6 @@ public Object buildJsonObject(QueryResult response) {
 
     json.datarows(fetchDataRows(response));
 
-    List<Map<String, Object>> highlights = response.highlights();
-    if (highlights.stream().anyMatch(Objects::nonNull)) {
-      json.highlights(highlights);
-    }
-
     formatMetric.set(System.nanoTime() - formatTime);
 
     json.profile(QueryProfiling.current().finish());
@@ -87,7 +80,6 @@ public static class JsonResponse {
     private final List<Column> schema;
 
     private final Object[][] datarows;
-    private final List<Map<String, Object>> highlights;
 
     private long total;
     private long size;