From bde93ca6fb7735d091995e4fe9a51ab23a32ce43 Mon Sep 17 00:00:00 2001 From: Jialiang Liang Date: Thu, 12 Feb 2026 13:24:41 -0800 Subject: [PATCH] Calcite PPL search result highlighting Signed-off-by: Jialiang Liang --- .../sql/calcite/CalciteRelNodeVisitor.java | 5 ++ .../sql/calcite/utils/PPLHintUtils.java | 37 +++++++++- .../sql/expression/HighlightExpression.java | 5 +- .../executor/OpenSearchExecutionEngine.java | 18 +++++ .../response/OpenSearchResponse.java | 3 +- .../scan/AbstractCalciteIndexScan.java | 11 ++- .../storage/scan/CalciteLogicalIndexScan.java | 22 ++++++ .../scan/OpenSearchIndexEnumerator.java | 32 +++++++++ .../storage/scan/context/PushDownType.java | 4 +- .../sql/protocol/response/QueryResult.java | 47 +++++++++++-- .../format/JdbcResponseFormatter.java | 9 +++ .../format/SimpleJsonResponseFormatter.java | 10 +++ .../protocol/response/QueryResultTest.java | 69 +++++++++++++++++++ .../format/JdbcResponseFormatterTest.java | 50 ++++++++++++++ 14 files changed, 307 insertions(+), 15 deletions(-) diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java index 5825011f653..b3e49d8477b 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java @@ -227,6 +227,11 @@ private RelBuilder scan(RelOptTable tableSchema, CalcitePlanContext context) { public RelNode visitSearch(Search node, CalcitePlanContext context) { // Visit the Relation child to get the scan node.getChild().get(0).accept(this, context); + + // Mark the scan as originating from a search command so that the optimizer + // can scope auto-highlight injection to search queries only. + PPLHintUtils.markSearchCommand(context.relBuilder); + // Create query_string function Function queryStringFunc = AstDSL.function( diff --git a/core/src/main/java/org/opensearch/sql/calcite/utils/PPLHintUtils.java b/core/src/main/java/org/opensearch/sql/calcite/utils/PPLHintUtils.java index 915c45e7083..31d213b32d8 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/utils/PPLHintUtils.java +++ b/core/src/main/java/org/opensearch/sql/calcite/utils/PPLHintUtils.java @@ -6,10 +6,14 @@ package org.opensearch.sql.calcite.utils; import com.google.common.base.Suppliers; +import java.util.List; import java.util.function.Supplier; import lombok.experimental.UtilityClass; +import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.TableScan; import org.apache.calcite.rel.hint.HintStrategyTable; +import org.apache.calcite.rel.hint.Hintable; import org.apache.calcite.rel.hint.RelHint; import org.apache.calcite.rel.logical.LogicalAggregate; import org.apache.calcite.tools.RelBuilder; @@ -19,6 +23,7 @@ public class PPLHintUtils { private static final String HINT_AGG_ARGUMENTS = "AGG_ARGS"; private static final String KEY_IGNORE_NULL_BUCKET = "ignoreNullBucket"; private static final String KEY_HAS_NESTED_AGG_CALL = "hasNestedAggCall"; + public static final String HINT_SEARCH_COMMAND = "SEARCH_COMMAND"; private static final Supplier HINT_STRATEGY_TABLE = Suppliers.memoize( @@ -29,7 +34,7 @@ public class PPLHintUtils { (hint, rel) -> { return rel instanceof LogicalAggregate; }) - // add more here + .hintStrategy(HINT_SEARCH_COMMAND, (hint, rel) -> rel instanceof TableScan) .build()); /** @@ -81,4 +86,34 @@ public static boolean hasNestedAggCall(Aggregate aggregate) { .getOrDefault(KEY_HAS_NESTED_AGG_CALL, "false") .equals("true")); } + + /** + * Mark a scan node as originating from a PPL search command. The scan node may be on top of the + * relBuilder stack directly, or wrapped in a Project (due to alias field wrapping). This hint is + * used to scope auto-highlight injection to search command queries only. + */ + public static void markSearchCommand(RelBuilder relBuilder) { + final RelHint hint = RelHint.builder(HINT_SEARCH_COMMAND).build(); + RelNode top = relBuilder.peek(); + if (top instanceof Hintable) { + // Scan is directly on top of the stack + relBuilder.hints(hint); + } else if (top instanceof org.apache.calcite.rel.core.Project proj) { + RelNode input = proj.getInput(); + if (input instanceof Hintable hintable) { + RelNode hintedInput = hintable.attachHints(List.of(hint)); + RelNode newProject = proj.copy(proj.getTraitSet(), List.of(hintedInput)); + relBuilder.build(); // pop old project + relBuilder.push(newProject); + } + } + if (relBuilder.getCluster().getHintStrategies() == HintStrategyTable.EMPTY) { + relBuilder.getCluster().setHintStrategies(HINT_STRATEGY_TABLE.get()); + } + } + + /** Return true if the scan has the SEARCH_COMMAND hint. */ + public static boolean isSearchCommand(TableScan scan) { + return scan.getHints().stream().anyMatch(hint -> hint.hintName.equals(HINT_SEARCH_COMMAND)); + } } diff --git a/core/src/main/java/org/opensearch/sql/expression/HighlightExpression.java b/core/src/main/java/org/opensearch/sql/expression/HighlightExpression.java index 79cc07f048b..6a6ad43ea49 100644 --- a/core/src/main/java/org/opensearch/sql/expression/HighlightExpression.java +++ b/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 valueEnv) { - String refName = "_highlight"; + String refName = HIGHLIGHT_FIELD; // Not a wilcard expression if (this.type == ExprCoreType.ARRAY) { refName += "." + StringUtils.unquoteText(getHighlightField().toString()); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java index a7eb3ad57be..5e60e448702 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java @@ -53,6 +53,7 @@ import org.opensearch.sql.executor.ExecutionEngine.Schema.Column; import org.opensearch.sql.executor.Explain; import org.opensearch.sql.executor.pagination.PlanSerializer; +import org.opensearch.sql.expression.HighlightExpression; import org.opensearch.sql.expression.function.BuiltinFunctionName; import org.opensearch.sql.expression.function.PPLFuncImpTable; import org.opensearch.sql.monitor.profile.MetricName; @@ -63,6 +64,7 @@ import org.opensearch.sql.opensearch.executor.protector.ExecutionProtector; import org.opensearch.sql.opensearch.functions.DistinctCountApproxAggFunction; import org.opensearch.sql.opensearch.functions.GeoIpFunction; +import org.opensearch.sql.opensearch.storage.scan.OpenSearchIndexEnumerator; import org.opensearch.sql.planner.physical.PhysicalPlan; import org.opensearch.sql.storage.TableScanOperator; import org.opensearch.transport.client.node.NodeClient; @@ -211,6 +213,7 @@ public void execute( client.schedule( () -> { try (PreparedStatement statement = OpenSearchRelRunners.run(context, rel)) { + OpenSearchIndexEnumerator.clearCollectedHighlights(); ProfileMetric metric = QueryProfiling.current().getOrCreateMetric(MetricName.EXECUTE); long execTime = System.nanoTime(); ResultSet result = statement.executeQuery(); @@ -279,6 +282,21 @@ private QueryResponse buildResultSet( values.add(ExprTupleValue.fromExprValueMap(row)); } + // Merge highlight data collected by the enumerator back into ExprTupleValues. + // The Calcite row pipeline only carries schema column values, so highlight metadata + // is collected as a side channel in OpenSearchIndexEnumerator and merged here. + List collectedHighlights = + OpenSearchIndexEnumerator.getAndClearCollectedHighlights(); + for (int i = 0; i < Math.min(values.size(), collectedHighlights.size()); i++) { + ExprValue hl = collectedHighlights.get(i); + if (hl != null) { + Map rowWithHighlight = + new LinkedHashMap<>(ExprValueUtils.getTupleValue(values.get(i))); + rowWithHighlight.put(HighlightExpression.HIGHLIGHT_FIELD, hl); + values.set(i, ExprTupleValue.fromExprValueMap(rowWithHighlight)); + } + } + List columns = new ArrayList<>(metaData.getColumnCount()); for (int i = 1; i <= columnCount; ++i) { String columnName = metaData.getColumnName(i); diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java index 0a47dc64a5e..1422be8ab45 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/response/OpenSearchResponse.java @@ -5,6 +5,7 @@ package org.opensearch.sql.opensearch.response; +import static org.opensearch.sql.expression.HighlightExpression.HIGHLIGHT_FIELD; import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATAFIELD_TYPE_MAP; import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_ID; import static org.opensearch.sql.opensearch.storage.OpenSearchIndex.METADATA_FIELD_INDEX; @@ -200,7 +201,7 @@ private void addHighlightsToBuilder( .map(Text::toString) .collect(Collectors.toList()))); } - builder.put("_highlight", ExprTupleValue.fromExprValueMap(hlBuilder.build())); + builder.put(HIGHLIGHT_FIELD, ExprTupleValue.fromExprValueMap(hlBuilder.build())); } } 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 index 3ab40caee27..e89e5e49d3e 100644 --- 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 + case PROJECT, HIGHLIGHT -> {} case SORT -> dCpu += dRows; case SORT_AGG_METRICS -> { dRows = dRows * .9 / 10; // *.9 because always bucket IS_NOT_NULL @@ -266,6 +266,11 @@ public Map getAliasMapping() { return osIndex.getAliasMapping(); } + @Override + public RelNode withHints(List hintList) { + return buildScan(getCluster(), traitSet, hintList, table, osIndex, schema, pushDownContext); + } + public abstract AbstractCalciteIndexScan copy(); protected List getCollationNames(List collations) { 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 index dbe8306d4b2..5a2ec10fbf1 100644 --- 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 @@ -40,6 +40,7 @@ import org.apache.logging.log4j.Logger; import org.opensearch.search.aggregations.AggregationBuilder; import org.opensearch.search.aggregations.bucket.composite.CompositeAggregationBuilder; +import org.opensearch.search.fetch.subphase.highlight.HighlightBuilder; import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory; import org.opensearch.sql.calcite.utils.PPLHintUtils; import org.opensearch.sql.common.setting.Settings; @@ -158,6 +159,27 @@ public AbstractRelNode pushDownFilter(Filter filter) { (OSRequestBuilderAction) requestBuilder -> requestBuilder.pushDownFilterForCalcite(queryExpression.builder())); + // Auto-inject wildcard highlight for PPL search command result highlighting. + // Only adds highlight when the scan is marked with a SEARCH_COMMAND hint + // (set by CalciteRelNodeVisitor.visitSearch), scoping it to the search command only. + // Uses OSD custom tags so the frontend getHighlightHtml() can convert to . + if (PPLHintUtils.isSearchCommand(this)) { + newScan.pushDownContext.add( + PushDownType.HIGHLIGHT, + "auto_highlight", + (OSRequestBuilderAction) + requestBuilder -> { + if (requestBuilder.getSourceBuilder().highlighter() == null) { + HighlightBuilder highlightBuilder = + new HighlightBuilder() + .field(new HighlightBuilder.Field("*").numOfFragments(0)) + .preTags("@opensearch-dashboards-highlighted-field@") + .postTags("@/opensearch-dashboards-highlighted-field@"); + requestBuilder.getSourceBuilder().highlighter(highlightBuilder); + } + }); + } + // If the query expression is partial, we need to replace the input of the filter with the // partial pushed scan and the filter condition with non-pushed-down conditions. if (queryExpression.isPartial()) { diff --git a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java index 05bd00dcf2c..e4d86ce0fad 100644 --- a/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java +++ b/opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java @@ -5,9 +5,13 @@ package org.opensearch.sql.opensearch.storage.scan; +import static org.opensearch.sql.expression.HighlightExpression.HIGHLIGHT_FIELD; + +import java.util.ArrayList; import java.util.Collections; import java.util.Iterator; import java.util.List; +import java.util.Map; import lombok.EqualsAndHashCode; import lombok.ToString; import org.apache.calcite.linq4j.Enumerator; @@ -27,6 +31,27 @@ */ public class OpenSearchIndexEnumerator implements Enumerator { + /** + * Thread-local collector for highlight data. Since the Calcite row pipeline only carries schema + * column values, highlight metadata from OpenSearch hits is collected here as a side channel. + * After execution, {@link #getAndClearCollectedHighlights()} retrieves the collected data so it + * can be merged back into the ExprTupleValues for the JDBC response. + */ + private static final ThreadLocal> COLLECTED_HIGHLIGHTS = + ThreadLocal.withInitial(ArrayList::new); + + /** Retrieve collected highlights and clear the ThreadLocal. */ + public static List getAndClearCollectedHighlights() { + List result = new ArrayList<>(COLLECTED_HIGHLIGHTS.get()); + COLLECTED_HIGHLIGHTS.get().clear(); + return result; + } + + /** Clear collected highlights (call before starting a new execution). */ + public static void clearCollectedHighlights() { + COLLECTED_HIGHLIGHTS.get().clear(); + } + /** OpenSearch client. */ private final OpenSearchClient client; @@ -111,6 +136,12 @@ public boolean moveNext() { } if (iterator.hasNext()) { current = iterator.next(); + // Collect highlight data as a side channel for the JDBC response. + // The Calcite row (from current()) only carries schema column values, + // so _highlight must be preserved separately. + Map tuple = ExprValueUtils.getTupleValue(current); + ExprValue hl = tuple.get(HIGHLIGHT_FIELD); + COLLECTED_HIGHLIGHTS.get().add(hl != null && !hl.isMissing() ? hl : null); queryCount++; return true; } else { @@ -123,6 +154,7 @@ public void reset() { bgScanner.reset(request); iterator = bgScanner.fetchNextBatch(request).iterator(); queryCount = 0; + COLLECTED_HIGHLIGHTS.get().clear(); } @Override 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 index c763808164d..d6817b8eeab 100644 --- 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/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java index 53badf3950d..9c8119a49dd 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/QueryResult.java @@ -5,11 +5,15 @@ 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; @@ -82,19 +86,48 @@ public Map columnNameTypes() { @Override public Iterator iterator() { - // Any chance to avoid copy for json response generation? return exprValues.stream() .map(ExprValueUtils::getTupleValue) - .map(Map::values) - .map(this::convertExprValuesToValues) + .map( + tuple -> + tuple.entrySet().stream() + .filter(e -> !HIGHLIGHT_FIELD.equals(e.getKey())) + .map(e -> e.getValue().value()) + .toArray(Object[]::new)) .iterator(); } - private String getColumnName(Column column) { - return (column.getAlias() != null) ? column.getAlias() : column.getName(); + /** + * 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. + * + * @return list of highlight maps, one per row + */ + public List> highlights() { + return exprValues.stream() + .map(ExprValueUtils::getTupleValue) + .map( + tuple -> { + ExprValue hl = tuple.get(HIGHLIGHT_FIELD); + if (hl == null || hl.isMissing()) { + return null; + } + Map hlMap = new LinkedHashMap<>(); + for (Map.Entry entry : hl.tupleValue().entrySet()) { + hlMap.put( + entry.getKey(), + entry.getValue().collectionValue().stream() + .map(ExprValue::stringValue) + .collect(Collectors.toList())); + } + return (Map) hlMap; + }) + .collect(Collectors.toList()); } - private Object[] convertExprValuesToValues(Collection exprValues) { - return exprValues.stream().map(ExprValue::value).toArray(Object[]::new); + private String getColumnName(Column column) { + return (column.getAlias() != null) ? column.getAlias() : column.getName(); } } diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java index 8be22af5326..5c696035555 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatter.java @@ -6,6 +6,8 @@ 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; @@ -38,6 +40,12 @@ protected Object buildJsonObject(QueryResult response) { response.getSchema().getColumns().forEach(col -> json.column(fetchColumn(col))); json.datarows(fetchDataRows(response)); + // Populate highlights if present + List> highlights = response.highlights(); + if (highlights.stream().anyMatch(Objects::nonNull)) { + json.highlights(highlights); + } + // Populate other fields json.total(response.size()).size(response.size()).status(200); if (!response.getCursor().equals(Cursor.None)) { @@ -88,6 +96,7 @@ public static class JdbcResponse { private final List schema; private final Object[][] datarows; + private final List> highlights; private final long total; private final long size; private final int status; diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java index ff59ce4cddc..5e57d8aacd6 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java @@ -6,6 +6,8 @@ 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; @@ -54,6 +56,13 @@ public Object buildJsonObject(QueryResult response) { response.columnNameTypes().forEach((name, type) -> json.column(new Column(name, type))); json.datarows(fetchDataRows(response)); + + // Populate highlights if present + List> highlights = response.highlights(); + if (highlights.stream().anyMatch(Objects::nonNull)) { + json.highlights(highlights); + } + formatMetric.set(System.nanoTime() - formatTime); json.profile(QueryProfiling.current().finish()); @@ -79,6 +88,7 @@ public static class JsonResponse { private final List schema; private final Object[][] datarows; + private final List> highlights; private long total; private long size; diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java index fc3402e20a5..f028da0f692 100644 --- a/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java +++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java @@ -7,6 +7,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.fail; import static org.opensearch.sql.data.model.ExprValueUtils.tupleValue; import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; @@ -16,7 +17,11 @@ import com.google.common.collect.ImmutableMap; import java.util.Arrays; import java.util.Collections; +import java.util.List; +import java.util.Map; import org.junit.jupiter.api.Test; +import org.opensearch.sql.data.model.ExprTupleValue; +import org.opensearch.sql.data.model.ExprValueUtils; import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.executor.pagination.Cursor; @@ -106,4 +111,68 @@ void iterate() { i++; } } + + @Test + void iterate_excludes_highlight_from_datarows() { + QueryResult response = + new QueryResult( + schema, + Collections.singletonList( + ExprTupleValue.fromExprValueMap( + ImmutableMap.of( + "name", + ExprValueUtils.stringValue("John"), + "age", + ExprValueUtils.integerValue(20), + "_highlight", + ExprTupleValue.fromExprValueMap( + ImmutableMap.of( + "name", + ExprValueUtils.collectionValue( + ImmutableList.of("John"))))))), + Cursor.None); + + for (Object[] objects : response) { + // datarows should only have schema columns, not _highlight + assertArrayEquals(new Object[] {"John", 20}, objects); + } + } + + @Test + void highlights_returns_highlight_data() { + QueryResult response = + new QueryResult( + schema, + Collections.singletonList( + ExprTupleValue.fromExprValueMap( + ImmutableMap.of( + "name", + ExprValueUtils.stringValue("John"), + "age", + ExprValueUtils.integerValue(20), + "_highlight", + ExprTupleValue.fromExprValueMap( + ImmutableMap.of( + "name", + ExprValueUtils.collectionValue( + ImmutableList.of("John"))))))), + Cursor.None); + + List> highlights = response.highlights(); + assertEquals(1, highlights.size()); + assertEquals(ImmutableMap.of("name", ImmutableList.of("John")), highlights.get(0)); + } + + @Test + void highlights_returns_null_when_no_highlight_data() { + QueryResult response = + new QueryResult( + schema, + Collections.singletonList(tupleValue(ImmutableMap.of("name", "John", "age", 20))), + Cursor.None); + + List> highlights = response.highlights(); + assertEquals(1, highlights.size()); + assertNull(highlights.get(0)); + } } diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java index 16dd1590eea..b8347f32be7 100644 --- a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java +++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java @@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableMap; import com.google.gson.JsonParser; import java.util.Arrays; +import java.util.Collections; import java.util.Map; import org.junit.jupiter.api.DisplayNameGeneration; import org.junit.jupiter.api.DisplayNameGenerator; @@ -207,6 +208,55 @@ void format_server_error_response_due_to_opensearch() { "all shards failed", new IllegalStateException("Execution error")))); } + @Test + void format_response_with_highlights() { + QueryResult response = + new QueryResult( + new Schema( + ImmutableList.of( + new Column("name", null, STRING), new Column("age", null, INTEGER))), + Collections.singletonList( + ExprTupleValue.fromExprValueMap( + ImmutableMap.of( + "name", + stringValue("John"), + "age", + org.opensearch.sql.data.model.ExprValueUtils.integerValue(20), + "_highlight", + ExprTupleValue.fromExprValueMap( + ImmutableMap.of( + "name", + org.opensearch.sql.data.model.ExprValueUtils.collectionValue( + ImmutableList.of("John")))))))); + + assertJsonEquals( + "{" + + "\"schema\":[" + + "{\"name\":\"name\",\"type\":\"keyword\"}," + + "{\"name\":\"age\",\"type\":\"integer\"}" + + "]," + + "\"datarows\":[[\"John\",20]]," + + "\"highlights\":[{\"name\":[\"John\"]}]," + + "\"total\":1," + + "\"size\":1," + + "\"status\":200}", + formatter.format(response)); + } + + @Test + void format_response_without_highlights() { + QueryResult response = + new QueryResult( + new Schema( + ImmutableList.of( + new Column("name", null, STRING), new Column("age", null, INTEGER))), + Collections.singletonList(tupleValue(ImmutableMap.of("name", "John", "age", 20)))); + + // When no highlights, the highlights field should not appear in the JSON + String formatted = formatter.format(response); + assertEquals(false, formatted.contains("\"highlights\"")); + } + private static void assertJsonEquals(String expected, String actual) { assertEquals(JsonParser.parseString(expected), JsonParser.parseString(actual)); }