[Feature] Calcite PPL search result highlighting#5141
[Feature] Calcite PPL search result highlighting#5141RyanL1997 wants to merge 9 commits intoopensearch-project:mainfrom
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds PPL search-result highlighting end-to-end: request parsing, plan-level propagation (ThreadLocal), Calcite schema/plan changes to carry a hidden Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant PPLService
participant Planner as Calcite Planner
participant Worker as Execution Thread
participant CalciteCtx as CalcitePlanContext
participant OpenSearch as OpenSearch Node
participant Enumerator as OpenSearchIndexEnumerator
participant Protocol as QueryResult/Formatter
Client->>PPLService: POST /_plugins/_ppl with highlight
PPLService->>Planner: build plan
PPLService->>Planner: set plan.highlightConfig
Planner->>Worker: submit plan (worker thread)
Worker->>CalciteCtx: setHighlightConfig(plan.highlightConfig)
Worker->>OpenSearch: execute search (HighlightBuilder attached)
OpenSearch-->>Enumerator: search hits with highlight fragments
Enumerator->>Worker: produce rows with hidden _highlight field
Worker->>Protocol: build QueryResult (extract highlights into highlights())
Worker->>CalciteCtx: clearHighlightConfig()
Protocol-->>Client: JSON response with optional "highlights" array
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
📝 WalkthroughAdds PPL search-result highlighting end-to-end: request parsing, plan-level propagation (ThreadLocal), Calcite schema/plan changes to carry a hidden 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
333a24b to
a4d156e
Compare
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (10)
integ-test/src/yamlRestTest/resources/rest-api-spec/test/ppl_highlight.yml (3)
36-41: Add index cleanup to teardown.The teardown disables the Calcite plugin but does not delete the
ppl_highlight_testindex. This can cause test pollution if tests run multiple times or if other tests use the same index name.Proposed fix to add index cleanup
--- teardown: + - do: + indices.delete: + index: ppl_highlight_test + ignore_unavailable: true + - do: query.settings: body: transient: plugins.calcite.enabled : false🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@integ-test/src/yamlRestTest/resources/rest-api-spec/test/ppl_highlight.yml` around lines 36 - 41, The teardown currently only disables the Calcite plugin; add a step to delete the ppl_highlight_test index to avoid test pollution. Inside the existing teardown block (under teardown -> - do:), add a delete-index action targeting the index name ppl_highlight_test (use the same REST/DSL action style as other teardown steps), and make the delete tolerant of missing indices (ignore_unavailable/allow_no_indices or equivalent) so the teardown does not fail if the index is already absent. Ensure this runs before or alongside the plugin disable step so cleanup always occurs.
99-100: Consider adding test case for non-text field behavior.Per the RFC, only text and keyword fields should produce highlight fragments—non-string fields like
code(integer) should not be highlighted. Adding a test that explicitly verifies integer fields are excluded from highlights would strengthen coverage of the documented behavior.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@integ-test/src/yamlRestTest/resources/rest-api-spec/test/ppl_highlight.yml` around lines 99 - 100, Update the ppl_highlight.yml test to explicitly assert that non-text fields (the integer field named "code") are not included in the highlights: after the existing checks for highlights (the existing "highlights" variable and match size assertions), add an assertion that the highlights result does not contain any entry for "code" (e.g., a negative contains/is_false check against highlights.code or a predicate that ensures no fragment keys equal "code"). Reference the "highlights" test output and the field name "code" when adding the new negative assertion so the test verifies integer fields are excluded from highlighting.
62-63: Consider adding assertions for highlight content.The test only verifies that
highlightsis present, but doesn't validate that the highlighted content actually contains the expected tags (e.g.,<em>error</em>) or that the correct fields are highlighted. This makes it possible for a bug to produce malformed highlights while still passing the test.Consider adding a match assertion to verify at least one highlight contains the expected tagged content.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@integ-test/src/yamlRestTest/resources/rest-api-spec/test/ppl_highlight.yml` around lines 62 - 63, Add a match assertion that validates the actual highlight HTML/text content rather than only the presence of "highlights": locate the assertions around "highlights" in ppl_highlight.yml and replace or add to the existing checks (the {"size": 2} and is_true: highlights lines) with a match that asserts at least one highlight string contains the expected tag/term (for example that a highlight value contains "<em>error</em>" or the expected field name wrapped in <em>…</em>), using the same match/assertion syntax used elsewhere in the file so the test fails if highlighted content is missing or malformed.integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLHighlightIT.java (1)
23-165: Consider adding error and boundary condition tests.Per coding guidelines, integration tests should include error condition tests and boundary condition tests. This test class would benefit from:
- Error condition test: Invalid/malformed highlight JSON configuration
- Boundary condition test: Query that returns empty results (no documents match)
- Edge case test: Filter-only query (per RFC, should not produce highlights for purely structured queries)
Suggested additional tests
`@Test` public void testHighlightWithEmptyResults() throws IOException { // Query that matches no documents JSONObject result = executeQueryWithHighlight( "search source=" + TEST_INDEX_ACCOUNT + " \\\"NonExistentTerm12345\\\"", "{\"fields\": {\"*\": {}}, \"pre_tags\": [\"<em>\"], \"post_tags\": [\"</em>\"]}"); assertTrue(result.has("highlights")); JSONArray highlights = result.getJSONArray("highlights"); assertEquals(0, highlights.length()); } `@Test` public void testHighlightWithFilterOnlyQuery() throws IOException { // Filter-only query without full-text search should not produce highlights JSONObject result = executeQueryWithHighlight( "search source=" + TEST_INDEX_ACCOUNT + " | where age > 30", "{\"fields\": {\"*\": {}}, \"pre_tags\": [\"<em>\"], \"post_tags\": [\"</em>\"]}"); // Verify behavior per RFC - no highlights for filter-only queries JSONArray highlights = result.getJSONArray("highlights"); for (int i = 0; i < highlights.length(); i++) { assertTrue("Filter-only query should not produce highlights", highlights.isNull(i)); } }As per coding guidelines: "Include boundary condition tests (min/max values, empty inputs) for all new functions" and "Include error condition tests (invalid inputs, exceptions) for all new functions".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLHighlightIT.java` around lines 23 - 165, Add three new integration tests to CalcitePPLHighlightIT: one that submits malformed highlight JSON via executeQueryWithHighlight (e.g., invalid JSON or missing required fields) and asserts the API returns an error status or error payload; one that runs a full-text search for a non-existent term using executeQueryWithHighlight and asserts "highlights" exists but the JSONArray returned by executeQueryWithHighlight has length 0; and one that executes a filter-only query (e.g., "search source=... | where age > 30") with executeQueryWithHighlight and asserts highlights are null/empty for each hit (use executeQueryWithHighlight and iterate the JSONArray as in testHighlightWithPipedFilter to confirm null entries). Ensure tests use the existing helper methods executeQueryWithHighlight/executeQueryNoHighlight and live in the CalcitePPLHighlightIT class.core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java (1)
89-94: Consider extractingsetHighlightThreadLocal()toAbstractPlanto reduce duplication.This helper method is identical in both
QueryPlanandExplainPlan. While the duplication is minor (4 lines), extracting it toAbstractPlanas a protected method would centralize the logic and ensure both subclasses stay in sync.♻️ Proposed refactor in AbstractPlan
In
AbstractPlan.java:/** * Sets the highlight configuration on the CalcitePlanContext ThreadLocal * for use during Calcite planning and execution. */ protected void setHighlightThreadLocal() { Map<String, Object> config = getHighlightConfig(); if (config != null) { CalcitePlanContext.setHighlightConfig(config); } }Then both
QueryPlanandExplainPlancan simply callsetHighlightThreadLocal()without the private helper.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java` around lines 89 - 94, Extract the duplicated helper into AbstractPlan by moving the setHighlightThreadLocal() logic there: implement a protected setHighlightThreadLocal() in AbstractPlan that calls getHighlightConfig() and, if non-null, invokes CalcitePlanContext.setHighlightConfig(config); then remove the private setHighlightThreadLocal() implementations from QueryPlan and ExplainPlan and replace their usages with a call to the inherited protected setHighlightThreadLocal() method so both subclasses share the same implementation and stay in sync.core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java (1)
156-165: Missing mock setup may cause test instability.This test calls
query.execute()without stubbingqueryService.execute(...). Unlike the first test which usesdoAnswer, this test has no mock behavior defined. With Mockito's default lenient stubbing in@ExtendWith(MockitoExtension.class), the mock returnsnullfor void methods, but this relies on implicit behavior rather than explicit setup.Consider adding explicit stubbing for consistency and clarity:
🔧 Suggested improvement
`@Test` public void execute_clears_highlight_threadlocal_after_execution() { Map<String, Object> highlightConfig = Map.of("fields", Map.of("*", Map.of())); + doAnswer(invocation -> null).when(queryService).execute(any(), any(), any()); QueryPlan query = new QueryPlan(queryId, queryType, plan, queryService, queryListener); query.setHighlightConfig(highlightConfig); query.execute(); assertNull(CalcitePlanContext.getHighlightConfig()); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java` around lines 156 - 165, The test execute_clears_highlight_threadlocal_after_execution calls QueryPlan#execute() but doesn't stub the mock QueryService, which can make the test flaky; add an explicit Mockito stub for queryService.execute(...) (similar to the other test's doAnswer) so that when QueryPlan#execute invokes queryService.execute it returns the expected result or performs the expected side effect, then run the same assertions (ensure you reference QueryPlan, query.execute(), and queryService.execute(...) when adding the stub) to guarantee deterministic behavior.opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java (2)
160-169: Missing highlight cleanup inclose()may leave stale data.The
reset()method clears the collected highlights, butclose()does not. If an enumerator is closed without being reset (e.g., on error paths), stale highlight data from a previous execution could persist in the ThreadLocal and potentially be merged into unrelated subsequent queries on the same thread.🔧 Suggested fix
`@Override` public void close() { iterator = Collections.emptyIterator(); queryCount = 0; + COLLECTED_HIGHLIGHTS.get().clear(); bgScanner.close(); if (request != null) { client.forceCleanup(request); request = null; } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java` around lines 160 - 169, The close() method in OpenSearchIndexEnumerator doesn't clear the per-thread highlight state, which can leave stale highlights in ThreadLocal; update close() (in class OpenSearchIndexEnumerator) to also clear collected highlights by invoking the same cleanup as reset() (or explicitly calling collectedHighlights.remove()/clearCollectedHighlights()) before returning so highlights are not leaked across requests; ensure this runs even when other cleanup (bgScanner.close(), client.forceCleanup(request)) has already executed.
40-53: ThreadLocal not removed after use risks memory leaks in thread pools.The
COLLECTED_HIGHLIGHTSThreadLocal is initialized withwithInitial(ArrayList::new)but never callsremove(). In long-lived thread pools (common in OpenSearch), this keeps the ArrayList instance attached to each thread indefinitely.
getAndClearCollectedHighlights()clears the list contents but the ThreadLocal entry persists. Consider callingCOLLECTED_HIGHLIGHTS.remove()after retrieving and clearing:🔧 Suggested fix
public static List<ExprValue> getAndClearCollectedHighlights() { List<ExprValue> result = new ArrayList<>(COLLECTED_HIGHLIGHTS.get()); - COLLECTED_HIGHLIGHTS.get().clear(); + COLLECTED_HIGHLIGHTS.remove(); return result; } /** Clear collected highlights (call before starting a new execution). */ public static void clearCollectedHighlights() { - COLLECTED_HIGHLIGHTS.get().clear(); + COLLECTED_HIGHLIGHTS.remove(); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java` around lines 40 - 53, The ThreadLocal COLLECTED_HIGHLIGHTS currently only clears the held ArrayList but never removes the ThreadLocal entry, risking memory leaks in pooled threads; update getAndClearCollectedHighlights() to copy contents, clear the list, and then call COLLECTED_HIGHLIGHTS.remove() to detach the entry from the current thread, and change clearCollectedHighlights() to call COLLECTED_HIGHLIGHTS.remove() (instead of just .clear()) so the ThreadLocal is fully removed when starting a new execution; keep references to COLLECTED_HIGHLIGHTS, getAndClearCollectedHighlights, and clearCollectedHighlights to locate the changes.opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java (2)
210-228: Highlights not cleared on exception path.If
statement.executeQuery()orbuildResultSet()throws an exception, the collected highlights fromOpenSearchIndexEnumeratorwill not be cleared (sincegetAndClearCollectedHighlights()is only called inbuildResultSet). This could leave stale data in the ThreadLocal for subsequent queries on the same thread.🔧 Suggested fix
public void execute( RelNode rel, CalcitePlanContext context, ResponseListener<QueryResponse> listener) { 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(); QueryResponse response = buildResultSet(result, rel.getRowType(), context.sysLimit.querySizeLimit()); metric.add(System.nanoTime() - execTime); listener.onResponse(response); } catch (SQLException e) { + OpenSearchIndexEnumerator.clearCollectedHighlights(); throw new RuntimeException(e); } }); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java` around lines 210 - 228, Ensure collected highlights are always cleared on the exception path by moving or adding a call to OpenSearchIndexEnumerator.clearCollectedHighlights() into a finally block around the code that runs the query in OpenSearchExecutionEngine.execute (the block that uses OpenSearchRelRunners.run, statement.executeQuery(), and buildResultSet). Specifically, wrap the try-with-resources/logic such that after ResultSet processing (or if an SQLException/any RuntimeException occurs) you call OpenSearchIndexEnumerator.clearCollectedHighlights() (or getAndClearCollectedHighlights() if that returns and clears) before rethrowing or returning, so no ThreadLocal highlights remain for subsequent queries.
284-297: Silent handling of highlight/row count mismatch may mask bugs.The
Math.min(values.size(), collectedHighlights.size())guard silently ignores mismatches. If the enumerator collected fewer highlights than rows returned (or vice versa), this could indicate a synchronization bug between the Calcite pipeline and the side-channel collection.Consider logging a warning when counts don't match to aid debugging:
🔧 Suggested improvement
List<ExprValue> collectedHighlights = OpenSearchIndexEnumerator.getAndClearCollectedHighlights(); + if (!collectedHighlights.isEmpty() && collectedHighlights.size() != values.size()) { + logger.warn("Highlight count ({}) does not match row count ({})", + collectedHighlights.size(), values.size()); + } for (int i = 0; i < Math.min(values.size(), collectedHighlights.size()); i++) {🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java` around lines 284 - 297, The loop currently uses Math.min(values.size(), collectedHighlights.size()) and silently ignores mismatched counts; update OpenSearchExecutionEngine after calling OpenSearchIndexEnumerator.getAndClearCollectedHighlights() to check if values.size() != collectedHighlights.size() and emit a warning-level log (including both sizes and context like "highlight/row count mismatch") so mismatches are visible for debugging; keep the existing merge logic that sets HighlightExpression.HIGHLIGHT_FIELD into ExprTupleValue.fromExprValueMap(values) but log the discrepancy before the for-loop to aid tracing.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In
`@integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLHighlightIT.java`:
- Around line 56-74: The test testHighlightWithSpecificField currently only
asserts that firstname is not highlighted; update it to also assert that the
address field is present and contains highlight snippets: after obtaining
JSONArray highlights and each non-null JSONObject hl (from
executeQueryWithHighlight result), add a positive assertion that
hl.has("address") is true and that hl.getJSONArray("address").length() > 0 (or
that the address value is non-empty) to ensure the address field is actually
highlighted.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java`:
- Around line 211-258: Rename the two test methods in JdbcResponseFormatterTest:
change format_response_with_highlights() to testFormatResponseWithHighlights()
and format_response_without_highlights() to
testFormatResponseWithoutHighlights(); update any references to these methods
(if used reflectively) so the `@Test-annotated` methods follow the
test<Functionality><Condition> naming convention while keeping their bodies
(uses of QueryResult, formatter.format(response), assertJsonEquals, and
assertEquals checks) unchanged.
- Around line 211-258: Add additional unit tests in JdbcResponseFormatterTest to
cover null/boundary/multi-row/error cases for highlight formatting: create tests
that (1) format an empty QueryResult (no rows) and assert that "highlights" is
absent and total/size are zero, (2) format multiple-row QueryResult where some
rows include a valid "_highlight" and others omit it and assert highlights array
aligns per-row (null or absent entries where appropriate), and (3) pass invalid
or null highlight payloads (e.g., "_highlight" present but value null or wrong
type) to formatter.format and assert graceful handling (no exception and either
omission or null entries) to satisfy NULL/boundary/error/multi-document
coverage; reference existing test helpers and symbols like QueryResult, Schema,
Column, tupleValue, ExprTupleValue, and formatter.format when adding these
cases.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java`:
- Around line 187-219: Two unit test methods in SimpleJsonResponseFormatterTest
do not follow the required naming pattern; rename the methods
formatResponseWithHighlights to testFormatResponseHighlightsPresent and
formatResponseWithoutHighlightsOmitsField to
testFormatResponseHighlightsOmittedWhenAbsent (or the exact pattern
test<Functionality><Condition>) so they match test<Functionality><Condition>;
update any internal references within the SimpleJsonResponseFormatterTest class
(imports, annotations like `@Test` remain unchanged) and ensure the new method
names still run under the existing test framework.
- Around line 187-219: Add unit tests to cover null/boundary/multi-row/error
cases for highlight formatting in SimpleJsonResponseFormatter: create tests that
call SimpleJsonResponseFormatter.format with (1) an empty QueryResult (no rows)
and assert the output omits "highlights" and yields an empty result, (2)
multiple rows where some rows include an "_highlight" ExprTupleValue and others
don't and assert highlights appear only for rows with the field, (3) a row
containing a null "_highlight" value and assert the formatter handles it
gracefully (omits or represents as null per contract), and (4) a row with an
invalid "_highlight" payload type (e.g., non-ExprTupleValue or non-collection
elements) and assert the formatter either throws a documented exception or
sanitizes the output; use the existing helpers QueryResult,
ExprTupleValue.fromExprValueMap, tupleValue, collectionValue and
SimpleJsonResponseFormatter(COMPACT) to construct inputs and assertions.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java`:
- Around line 115-198: Rename the JUnit test methods to follow the
test<Functionality><Condition> pattern: change
iterate_excludes_highlight_from_datarows to
testIterateExcludesHighlightFromDatarows, highlights_returns_highlight_data to
testHighlightsReturnsHighlightData,
highlights_returns_null_when_no_highlight_data to
testHighlightsReturnsNullWhenNoHighlightData, and
highlights_returns_null_when_highlight_is_missing to
testHighlightsReturnsNullWhenHighlightIsMissing; update any references/imports
or test framework annotations if needed so the new method names compile and run.
- Around line 115-198: Add unit tests for QueryResult.highlights() to cover
empty input, multi-row mixed highlights, null/invalid highlight payloads, and
boundary cases: 1) empty rows: construct a QueryResult with an empty list and
assert highlights() returns empty list; 2) multi-row mix: build a QueryResult
with multiple ExprTupleValue rows where some rows include a valid "_highlight",
some have ExprValueUtils.LITERAL_MISSING, some have null/invalid types (e.g.,
non-collection), and assert highlights() returns a list whose entries align
per-document (valid maps for valid highlights, null for missing/invalid); 3)
boundary empty highlight: include a "_highlight" whose field value is an empty
collection and assert it returns an empty map or empty collection entry as per
current contract; and 4) error condition: add a test where "_highlight" is an
unexpected type and assert highlights() yields null for that entry (or throws if
the contract requires), referencing QueryResult.highlights() and QueryResultTest
to locate where to add these cases.
---
Nitpick comments:
In `@core/src/main/java/org/opensearch/sql/executor/execution/QueryPlan.java`:
- Around line 89-94: Extract the duplicated helper into AbstractPlan by moving
the setHighlightThreadLocal() logic there: implement a protected
setHighlightThreadLocal() in AbstractPlan that calls getHighlightConfig() and,
if non-null, invokes CalcitePlanContext.setHighlightConfig(config); then remove
the private setHighlightThreadLocal() implementations from QueryPlan and
ExplainPlan and replace their usages with a call to the inherited protected
setHighlightThreadLocal() method so both subclasses share the same
implementation and stay in sync.
In `@core/src/test/java/org/opensearch/sql/executor/execution/QueryPlanTest.java`:
- Around line 156-165: The test
execute_clears_highlight_threadlocal_after_execution calls QueryPlan#execute()
but doesn't stub the mock QueryService, which can make the test flaky; add an
explicit Mockito stub for queryService.execute(...) (similar to the other test's
doAnswer) so that when QueryPlan#execute invokes queryService.execute it returns
the expected result or performs the expected side effect, then run the same
assertions (ensure you reference QueryPlan, query.execute(), and
queryService.execute(...) when adding the stub) to guarantee deterministic
behavior.
In
`@integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLHighlightIT.java`:
- Around line 23-165: Add three new integration tests to CalcitePPLHighlightIT:
one that submits malformed highlight JSON via executeQueryWithHighlight (e.g.,
invalid JSON or missing required fields) and asserts the API returns an error
status or error payload; one that runs a full-text search for a non-existent
term using executeQueryWithHighlight and asserts "highlights" exists but the
JSONArray returned by executeQueryWithHighlight has length 0; and one that
executes a filter-only query (e.g., "search source=... | where age > 30") with
executeQueryWithHighlight and asserts highlights are null/empty for each hit
(use executeQueryWithHighlight and iterate the JSONArray as in
testHighlightWithPipedFilter to confirm null entries). Ensure tests use the
existing helper methods executeQueryWithHighlight/executeQueryNoHighlight and
live in the CalcitePPLHighlightIT class.
In `@integ-test/src/yamlRestTest/resources/rest-api-spec/test/ppl_highlight.yml`:
- Around line 36-41: The teardown currently only disables the Calcite plugin;
add a step to delete the ppl_highlight_test index to avoid test pollution.
Inside the existing teardown block (under teardown -> - do:), add a delete-index
action targeting the index name ppl_highlight_test (use the same REST/DSL action
style as other teardown steps), and make the delete tolerant of missing indices
(ignore_unavailable/allow_no_indices or equivalent) so the teardown does not
fail if the index is already absent. Ensure this runs before or alongside the
plugin disable step so cleanup always occurs.
- Around line 99-100: Update the ppl_highlight.yml test to explicitly assert
that non-text fields (the integer field named "code") are not included in the
highlights: after the existing checks for highlights (the existing "highlights"
variable and match size assertions), add an assertion that the highlights result
does not contain any entry for "code" (e.g., a negative contains/is_false check
against highlights.code or a predicate that ensures no fragment keys equal
"code"). Reference the "highlights" test output and the field name "code" when
adding the new negative assertion so the test verifies integer fields are
excluded from highlighting.
- Around line 62-63: Add a match assertion that validates the actual highlight
HTML/text content rather than only the presence of "highlights": locate the
assertions around "highlights" in ppl_highlight.yml and replace or add to the
existing checks (the {"size": 2} and is_true: highlights lines) with a match
that asserts at least one highlight string contains the expected tag/term (for
example that a highlight value contains "<em>error</em>" or the expected field
name wrapped in <em>…</em>), using the same match/assertion syntax used
elsewhere in the file so the test fails if highlighted content is missing or
malformed.
In
`@opensearch/src/main/java/org/opensearch/sql/opensearch/executor/OpenSearchExecutionEngine.java`:
- Around line 210-228: Ensure collected highlights are always cleared on the
exception path by moving or adding a call to
OpenSearchIndexEnumerator.clearCollectedHighlights() into a finally block around
the code that runs the query in OpenSearchExecutionEngine.execute (the block
that uses OpenSearchRelRunners.run, statement.executeQuery(), and
buildResultSet). Specifically, wrap the try-with-resources/logic such that after
ResultSet processing (or if an SQLException/any RuntimeException occurs) you
call OpenSearchIndexEnumerator.clearCollectedHighlights() (or
getAndClearCollectedHighlights() if that returns and clears) before rethrowing
or returning, so no ThreadLocal highlights remain for subsequent queries.
- Around line 284-297: The loop currently uses Math.min(values.size(),
collectedHighlights.size()) and silently ignores mismatched counts; update
OpenSearchExecutionEngine after calling
OpenSearchIndexEnumerator.getAndClearCollectedHighlights() to check if
values.size() != collectedHighlights.size() and emit a warning-level log
(including both sizes and context like "highlight/row count mismatch") so
mismatches are visible for debugging; keep the existing merge logic that sets
HighlightExpression.HIGHLIGHT_FIELD into ExprTupleValue.fromExprValueMap(values)
but log the discrepancy before the for-loop to aid tracing.
In
`@opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/OpenSearchIndexEnumerator.java`:
- Around line 160-169: The close() method in OpenSearchIndexEnumerator doesn't
clear the per-thread highlight state, which can leave stale highlights in
ThreadLocal; update close() (in class OpenSearchIndexEnumerator) to also clear
collected highlights by invoking the same cleanup as reset() (or explicitly
calling collectedHighlights.remove()/clearCollectedHighlights()) before
returning so highlights are not leaked across requests; ensure this runs even
when other cleanup (bgScanner.close(), client.forceCleanup(request)) has already
executed.
- Around line 40-53: The ThreadLocal COLLECTED_HIGHLIGHTS currently only clears
the held ArrayList but never removes the ThreadLocal entry, risking memory leaks
in pooled threads; update getAndClearCollectedHighlights() to copy contents,
clear the list, and then call COLLECTED_HIGHLIGHTS.remove() to detach the entry
from the current thread, and change clearCollectedHighlights() to call
COLLECTED_HIGHLIGHTS.remove() (instead of just .clear()) so the ThreadLocal is
fully removed when starting a new execution; keep references to
COLLECTED_HIGHLIGHTS, getAndClearCollectedHighlights, and
clearCollectedHighlights to locate the changes.
| @Test | ||
| public void testHighlightWithSpecificField() throws IOException { | ||
| JSONObject result = | ||
| executeQueryWithHighlight( | ||
| "search source=" + TEST_INDEX_ACCOUNT + " \\\"Street\\\"", | ||
| "{\"fields\": {\"address\": {}}, \"pre_tags\": [\"<em>\"], \"post_tags\":" | ||
| + " [\"</em>\"]}"); | ||
|
|
||
| assertTrue(result.has("highlights")); | ||
| JSONArray highlights = result.getJSONArray("highlights"); | ||
|
|
||
| for (int i = 0; i < highlights.length(); i++) { | ||
| if (!highlights.isNull(i)) { | ||
| JSONObject hl = highlights.getJSONObject(i); | ||
| // Only address field should be highlighted, not other text fields | ||
| assertFalse("Should not highlight firstname field", hl.has("firstname")); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
Strengthen assertion to verify address field IS highlighted.
The test only asserts that firstname is NOT highlighted but doesn't positively verify that the address field IS highlighted. This could pass even if highlighting is completely broken for the target field.
Proposed fix to add positive assertion
for (int i = 0; i < highlights.length(); i++) {
if (!highlights.isNull(i)) {
JSONObject hl = highlights.getJSONObject(i);
// Only address field should be highlighted, not other text fields
assertFalse("Should not highlight firstname field", hl.has("firstname"));
+ assertTrue("Should highlight address field", hl.has("address"));
}
}
+ // Verify at least one non-null highlight was found
+ boolean foundHighlight = false;
+ for (int i = 0; i < highlights.length(); i++) {
+ if (!highlights.isNull(i)) {
+ foundHighlight = true;
+ break;
+ }
+ }
+ assertTrue("Expected at least one highlight entry", foundHighlight);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePPLHighlightIT.java`
around lines 56 - 74, The test testHighlightWithSpecificField currently only
asserts that firstname is not highlighted; update it to also assert that the
address field is present and contains highlight snippets: after obtaining
JSONArray highlights and each non-null JSONObject hl (from
executeQueryWithHighlight result), add a positive assertion that
hl.has("address") is true and that hl.getJSONArray("address").length() > 0 (or
that the address value is non-empty) to ensure the address field is actually
highlighted.
| @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("<em>John</em>")))))))); | ||
|
|
||
| assertJsonEquals( | ||
| "{" | ||
| + "\"schema\":[" | ||
| + "{\"name\":\"name\",\"type\":\"keyword\"}," | ||
| + "{\"name\":\"age\",\"type\":\"integer\"}" | ||
| + "]," | ||
| + "\"datarows\":[[\"John\",20]]," | ||
| + "\"highlights\":[{\"name\":[\"<em>John</em>\"]}]," | ||
| + "\"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\"")); | ||
| } |
There was a problem hiding this comment.
Rename tests to match the required naming pattern.
Use test<Functionality><Condition> (e.g., testFormatResponseWithHighlights, testFormatResponseWithoutHighlights). As per coding guidelines, "Test naming should follow pattern test".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java`
around lines 211 - 258, Rename the two test methods in
JdbcResponseFormatterTest: change format_response_with_highlights() to
testFormatResponseWithHighlights() and format_response_without_highlights() to
testFormatResponseWithoutHighlights(); update any references to these methods
(if used reflectively) so the `@Test-annotated` methods follow the
test<Functionality><Condition> naming convention while keeping their bodies
(uses of QueryResult, formatter.format(response), assertJsonEquals, and
assertEquals checks) unchanged.
Add null/boundary/multi-row/error coverage for highlight formatting.
The new tests only cover a single-row highlight and absence. Add coverage for empty results, multiple rows with mixed highlight presence, and invalid/null highlight payloads to satisfy required per-document test coverage. As per coding guidelines, "NULL input tests must be included for all new functions", "Include boundary condition tests (min/max values, empty inputs) for all new functions", "Include error condition tests (invalid inputs, exceptions) for all new functions", and "Include multi-document tests for per-document operations".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/format/JdbcResponseFormatterTest.java`
around lines 211 - 258, Add additional unit tests in JdbcResponseFormatterTest
to cover null/boundary/multi-row/error cases for highlight formatting: create
tests that (1) format an empty QueryResult (no rows) and assert that
"highlights" is absent and total/size are zero, (2) format multiple-row
QueryResult where some rows include a valid "_highlight" and others omit it and
assert highlights array aligns per-row (null or absent entries where
appropriate), and (3) pass invalid or null highlight payloads (e.g.,
"_highlight" present but value null or wrong type) to formatter.format and
assert graceful handling (no exception and either omission or null entries) to
satisfy NULL/boundary/error/multi-document coverage; reference existing test
helpers and symbols like QueryResult, Schema, Column, tupleValue,
ExprTupleValue, and formatter.format when adding these cases.
| @Test | ||
| void formatResponseWithHighlights() { | ||
| QueryResult response = | ||
| new QueryResult( | ||
| schema, | ||
| Collections.singletonList( | ||
| ExprTupleValue.fromExprValueMap( | ||
| ImmutableMap.of( | ||
| "firstname", | ||
| stringValue("John"), | ||
| "age", | ||
| integerValue(20), | ||
| "_highlight", | ||
| ExprTupleValue.fromExprValueMap( | ||
| ImmutableMap.of( | ||
| "firstname", | ||
| collectionValue(ImmutableList.of("<em>John</em>")))))))); | ||
| SimpleJsonResponseFormatter formatter = new SimpleJsonResponseFormatter(COMPACT); | ||
| String result = formatter.format(response); | ||
| assertTrue(result.contains("\"highlights\"")); | ||
| assertTrue(result.contains("\"firstname\":[\"<em>John</em>\"]")); | ||
| } | ||
|
|
||
| @Test | ||
| void formatResponseWithoutHighlightsOmitsField() { | ||
| QueryResult response = | ||
| new QueryResult( | ||
| schema, | ||
| Collections.singletonList(tupleValue(ImmutableMap.of("firstname", "John", "age", 20)))); | ||
| SimpleJsonResponseFormatter formatter = new SimpleJsonResponseFormatter(COMPACT); | ||
| String result = formatter.format(response); | ||
| assertFalse(result.contains("\"highlights\"")); | ||
| } |
There was a problem hiding this comment.
Rename tests to match the required naming pattern.
Use test<Functionality><Condition> (e.g., testFormatResponseHighlightsPresent, testFormatResponseHighlightsOmittedWhenAbsent). As per coding guidelines, "Test naming should follow pattern test".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java`
around lines 187 - 219, Two unit test methods in SimpleJsonResponseFormatterTest
do not follow the required naming pattern; rename the methods
formatResponseWithHighlights to testFormatResponseHighlightsPresent and
formatResponseWithoutHighlightsOmitsField to
testFormatResponseHighlightsOmittedWhenAbsent (or the exact pattern
test<Functionality><Condition>) so they match test<Functionality><Condition>;
update any internal references within the SimpleJsonResponseFormatterTest class
(imports, annotations like `@Test` remain unchanged) and ensure the new method
names still run under the existing test framework.
Add null/boundary/multi-row/error coverage for highlight formatting.
The new highlight tests only cover a single-row happy path and the “no highlights” case. Add coverage for empty results, multiple rows with mixed highlight presence, and invalid/null highlight payloads to satisfy required test coverage for new per-document behavior. As per coding guidelines, "NULL input tests must be included for all new functions", "Include boundary condition tests (min/max values, empty inputs) for all new functions", "Include error condition tests (invalid inputs, exceptions) for all new functions", and "Include multi-document tests for per-document operations".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java`
around lines 187 - 219, Add unit tests to cover null/boundary/multi-row/error
cases for highlight formatting in SimpleJsonResponseFormatter: create tests that
call SimpleJsonResponseFormatter.format with (1) an empty QueryResult (no rows)
and assert the output omits "highlights" and yields an empty result, (2)
multiple rows where some rows include an "_highlight" ExprTupleValue and others
don't and assert highlights appear only for rows with the field, (3) a row
containing a null "_highlight" value and assert the formatter handles it
gracefully (omits or represents as null per contract), and (4) a row with an
invalid "_highlight" payload type (e.g., non-ExprTupleValue or non-collection
elements) and assert the formatter either throws a documented exception or
sanitizes the output; use the existing helpers QueryResult,
ExprTupleValue.fromExprValueMap, tupleValue, collectionValue and
SimpleJsonResponseFormatter(COMPACT) to construct inputs and assertions.
| @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("<em>John</em>"))))))), | ||
| 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("<em>John</em>"))))))), | ||
| Cursor.None); | ||
|
|
||
| List<Map<String, Object>> highlights = response.highlights(); | ||
| assertEquals(1, highlights.size()); | ||
| assertEquals(ImmutableMap.of("name", ImmutableList.of("<em>John</em>")), 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<Map<String, Object>> highlights = response.highlights(); | ||
| assertEquals(1, highlights.size()); | ||
| assertNull(highlights.get(0)); | ||
| } | ||
|
|
||
| @Test | ||
| void highlights_returns_null_when_highlight_is_missing() { | ||
| QueryResult response = | ||
| new QueryResult( | ||
| schema, | ||
| Collections.singletonList( | ||
| ExprTupleValue.fromExprValueMap( | ||
| ImmutableMap.of( | ||
| "name", | ||
| ExprValueUtils.stringValue("John"), | ||
| "age", | ||
| ExprValueUtils.integerValue(20), | ||
| "_highlight", | ||
| ExprValueUtils.LITERAL_MISSING))), | ||
| Cursor.None); | ||
|
|
||
| List<Map<String, Object>> highlights = response.highlights(); | ||
| assertEquals(1, highlights.size()); | ||
| assertNull(highlights.get(0)); | ||
| } |
There was a problem hiding this comment.
Rename tests to match the required naming pattern.
Use test<Functionality><Condition> (e.g., testIterateExcludesHighlightFromDatarows, testHighlightsReturnsNullWhenMissing). As per coding guidelines, "Test naming should follow pattern test".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java`
around lines 115 - 198, Rename the JUnit test methods to follow the
test<Functionality><Condition> pattern: change
iterate_excludes_highlight_from_datarows to
testIterateExcludesHighlightFromDatarows, highlights_returns_highlight_data to
testHighlightsReturnsHighlightData,
highlights_returns_null_when_no_highlight_data to
testHighlightsReturnsNullWhenNoHighlightData, and
highlights_returns_null_when_highlight_is_missing to
testHighlightsReturnsNullWhenHighlightIsMissing; update any references/imports
or test framework annotations if needed so the new method names compile and run.
Add null/boundary/multi-row/error coverage for highlights().
Current tests cover a single-row happy path and missing highlights, but not empty input, mixed multi-row alignment, or invalid/null highlight payloads. Add those to satisfy required coverage for the new per-document behavior. As per coding guidelines, "NULL input tests must be included for all new functions", "Include boundary condition tests (min/max values, empty inputs) for all new functions", "Include error condition tests (invalid inputs, exceptions) for all new functions", and "Include multi-document tests for per-document operations".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In
`@protocol/src/test/java/org/opensearch/sql/protocol/response/QueryResultTest.java`
around lines 115 - 198, Add unit tests for QueryResult.highlights() to cover
empty input, multi-row mixed highlights, null/invalid highlight payloads, and
boundary cases: 1) empty rows: construct a QueryResult with an empty list and
assert highlights() returns empty list; 2) multi-row mix: build a QueryResult
with multiple ExprTupleValue rows where some rows include a valid "_highlight",
some have ExprValueUtils.LITERAL_MISSING, some have null/invalid types (e.g.,
non-collection), and assert highlights() returns a list whose entries align
per-document (valid maps for valid highlights, null for missing/invalid); 3)
boundary empty highlight: include a "_highlight" whose field value is an empty
collection and assert it returns an empty map or empty collection entry as per
current contract; and 4) error condition: add a test where "_highlight" is an
unexpected type and assert highlights() yields null for that entry (or throws if
the contract requires), referencing QueryResult.highlights() and QueryResultTest
to locate where to add these cases.
| * | ||
| * @return highlight JSONObject from request, or null if not specified | ||
| */ | ||
| public JSONObject getHighlight() { |
There was a problem hiding this comment.
- Add doc to describe
highlighparameters. include, what field can be highlight.
There was a problem hiding this comment.
added at: docs/user/ppl/interfaces/endpoint.md (link)
| // is collected as a side channel in OpenSearchIndexEnumerator and merged here. | ||
| List<ExprValue> collectedHighlights = | ||
| OpenSearchIndexEnumerator.getAndClearCollectedHighlights(); | ||
| for (int i = 0; i < Math.min(values.size(), collectedHighlights.size()); i++) { |
There was a problem hiding this comment.
Highlight-to-row positional misalignment when Calcite operators sit above the index scan. The COLLECTED_HIGHLIGHTS ThreadLocal collects one entry per row the enumerator produces (raw scan order). After Calcite processing, buildResultSet iterates the final ResultSet, which may contain fewer rows in different order if non-pushable operators (filter, sort, dedup, head/tail) were applied above the scan.
There was a problem hiding this comment.
Good catch. I've addressed this by replacing the COLLECTED_HIGHLIGHTS ThreadLocal side-channel entirely. Instead, _highlight is now a hidden column in the Calcite RowType (SqlTypeName.ANY).
How it works: CalciteLogicalIndexScan.buildInitialSchema() appends _highlight to the scan's schema when highlight config is present. OpenSearchResponse.addHighlightsToBuilder() builds the highlight data as an ExprTupleValue and puts it directly into the row tuple. OpenSearchIndexEnumerator.current() carries it as an opaque ExprValue in the Calcite row array.
Because _highlight is now a regular column, Calcite operators (filter, sort, dedup, head/tail) naturally preserve it alongside the row data — no positional misalignment is possible. When a filter removes a row, its highlight goes with it. When sort reorders rows, highlights move with them.
On the output side, OpenSearchExecutionEngine.buildResultSet() extracts _highlight from each row and excludes it from the response schema (hidden from users). QueryResult.highlights() produces a parallel array for the formatters. | fields projections also preserve _highlight automatically via CalciteRelNodeVisitor.visitProject(), matching DSL behavior where _source filtering never affects highlights.
| - "</em>" | ||
|
|
||
| - match: {"size": 2} | ||
| - is_true: highlights |
There was a problem hiding this comment.
assertions only check is_true/is_false, no content validation
| * 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<List<ExprValue>> COLLECTED_HIGHLIGHTS = |
There was a problem hiding this comment.
Is highlighting intended only for single-source search commands, or should it work with joins/subqueries?
COLLECTED_HIGHLIGHTS is a single static ThreadLocal shared across all enumerator instances. If a query involves a join or union of two OpenSearch indices, multiple enumerator instances will append to the same list.
There was a problem hiding this comment.
Yes, Highlighting is single-source only, consistent with DSL where highlighting is inherently single-index per request — there's no DSL equivalent of highlighting across a join.
Regarding the COLLECTED_HIGHLIGHTS ThreadLocal concern — this is no longer applicable. We replaced the ThreadLocal side-channel entirely with a hidden _highlight column approach (see the response on the positional misalignment comment). There is no shared static list anymore.
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java (1)
264-305:⚠️ Potential issue | 🟠 MajorPreserve
_highlightcolumn through projections to maintain downstream availability.
pushDownProjectrebuilds the schema strictly fromselectedColumnsindices. Since_highlightis a hidden column (added bybuildInitialSchemawhen highlight config exists) and is not typically part of user projections, it gets dropped from the new schema. This breaks downstream code that expects_highlightto be present (e.g.,tuple.get(HIGHLIGHT_FIELD)inQueryResultandOpenSearchIndexEnumerator).Re-append
_highlightto the schema when it exists in the current row type but is absent fromselectedColumns, similar to the suggested approach in the review.Suggested fix
- RelDataType newSchema = builder.build(); + // Preserve hidden _highlight column if present in input schema + boolean highlightPresent = + fieldList.stream().anyMatch(f -> HIGHLIGHT_FIELD.equals(f.getName())); + boolean highlightSelected = + selectedColumns.stream() + .anyMatch(i -> HIGHLIGHT_FIELD.equals(fieldList.get(i).getName())); + if (highlightPresent && !highlightSelected) { + builder.add(HIGHLIGHT_FIELD, SqlTypeName.ANY); + } + RelDataType newSchema = builder.build();🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java` around lines 264 - 305, pushDownProject currently rebuilds newSchema from selectedColumns and drops the hidden _highlight field; detect if getRowType() contains HIGHLIGHT_FIELD but newSchema.getFieldNames() does not, and append the _highlight field back into the builder before builder.build(). Concretely, in CalciteLogicalIndexScan.pushDownProject use getRowType().getField(HIGHLIGHT_FIELD) (or scan getRowType().getFieldList() to find the field) to retrieve the RelDataTypeField for HIGHLIGHT_FIELD and call builder.add(...) for it if it exists and its name is missing from selectedColumns/newSchema, so that newSchema and the created CalciteLogicalIndexScan preserve the _highlight column for downstream code that references HIGHLIGHT_FIELD.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In
`@opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java`:
- Around line 264-305: pushDownProject currently rebuilds newSchema from
selectedColumns and drops the hidden _highlight field; detect if getRowType()
contains HIGHLIGHT_FIELD but newSchema.getFieldNames() does not, and append the
_highlight field back into the builder before builder.build(). Concretely, in
CalciteLogicalIndexScan.pushDownProject use
getRowType().getField(HIGHLIGHT_FIELD) (or scan getRowType().getFieldList() to
find the field) to retrieve the RelDataTypeField for HIGHLIGHT_FIELD and call
builder.add(...) for it if it exists and its name is missing from
selectedColumns/newSchema, so that newSchema and the created
CalciteLogicalIndexScan preserve the _highlight column for downstream code that
references HIGHLIGHT_FIELD.
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
Signed-off-by: Jialiang Liang <jiallian@amazon.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java`:
- Around line 411-418: When highlight is enabled
(CalcitePlanContext.getHighlightConfig() != null) the highlight preservation
block in CalciteRelNodeVisitor can add HIGHLIGHT_FIELD twice into expandedFields
(duplicate of expandProjectFields' AllFields expansion); modify the logic around
currentFields/index hlIndex and expandedFields to first scan expandedFields for
an existing RexInputRef that references the same input index (or any existing
projection that already corresponds to HIGHLIGHT_FIELD) and only add
context.relBuilder.field(hlIndex) if no matching RexInputRef exists, so the
projection is not duplicated and downstream tuple.get(HIGHLIGHT_FIELD) continues
to work.
| // 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.getHighlightConfig() != null) { | ||
| int hlIndex = currentFields.indexOf(HIGHLIGHT_FIELD); | ||
| if (hlIndex >= 0) { | ||
| expandedFields.add(context.relBuilder.field(hlIndex)); | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's see the full context of the visitProject method and the code around lines 411-418
cd core/src/main/java/org/opensearch/sql/calcite && wc -l CalciteRelNodeVisitor.javaRepository: opensearch-project/sql
Length of output: 97
🏁 Script executed:
# Read the visitProject method context (expanded range to understand the logic)
cd core/src/main/java/org/opensearch/sql/calcite && sed -n '380,430p' CalciteRelNodeVisitor.javaRepository: opensearch-project/sql
Length of output: 1792
🏁 Script executed:
# Search for how _highlight is used downstream (extraction/handling)
rg "HIGHLIGHT_FIELD|_highlight|getHighlight" -n -A 3 -B 1 --type javaRepository: opensearch-project/sql
Length of output: 50378
🏁 Script executed:
# Check if there are similar guards against duplicate projection in the codebase
rg "expandedFields|alreadyProjected|anyMatch.*RexInputRef" -n --type javaRepository: opensearch-project/sql
Length of output: 1744
🏁 Script executed:
# Search for tests related to highlight projection
rg "highlight.*[Pp]roject|visitProject.*highlight" -n --type javaRepository: opensearch-project/sql
Length of output: 365
🏁 Script executed:
# Check if _highlight can be explicitly selected in user queries
# Look for test cases with explicit _highlight selection
rg "SELECT.*_highlight|project.*_highlight" -n --type java -B 2 -A 2 | head -50Repository: opensearch-project/sql
Length of output: 1538
🏁 Script executed:
# Check expandProjectFields to see if it could include _highlight from user input
cd core/src/main/java/org/opensearch/sql/calcite && sed -n '440,495p' CalciteRelNodeVisitor.javaRepository: opensearch-project/sql
Length of output: 2319
🏁 Script executed:
# Search for how currentFields is populated and whether _highlight would be present
cd core/src/main/java/org/opensearch/sql/calcite && sed -n '393,425p' CalciteRelNodeVisitor.javaRepository: opensearch-project/sql
Length of output: 1328
🏁 Script executed:
# Check if there are tests specifically covering _highlight in projections with other fields
rg "highlight.*project|project.*highlight" -n --type java -A 5 | head -60Repository: opensearch-project/sql
Length of output: 3601
🏁 Script executed:
# Check if _highlight is classified as a metadata field
rg "isMetadataField|METADATA|_highlight" -n --type java -B 2 -A 2 | grep -A 5 "isMetadataField"Repository: opensearch-project/sql
Length of output: 4271
🏁 Script executed:
# Get the isMetadataField implementation
cd core/src/main/java/org/opensearch/sql/calcite && rg "isMetadataField" -A 5 CalciteRelNodeVisitor.javaRepository: opensearch-project/sql
Length of output: 1952
🏁 Script executed:
# Check if _highlight is in METADATAFIELD_TYPE_MAP
rg "METADATAFIELD_TYPE_MAP|_highlight" -n --type java | grep -B 5 -A 5 "METADATAFIELD_TYPE_MAP"Repository: opensearch-project/sql
Length of output: 7296
🏁 Script executed:
# Find where METADATAFIELD_TYPE_MAP is defined
fd "OpenSearchConstants.java" --type f | head -1 | xargs cat | head -100Repository: opensearch-project/sql
Length of output: 1137
🏁 Script executed:
# Verify that _highlight is NOT in the metadata field map and would be included in wildcard projections
cd core/src/main/java/org/opensearch/sql/opensearch/storage && cat -n OpenSearchIndex.java | sed -n '65,75p'Repository: opensearch-project/sql
Length of output: 172
🏁 Script executed:
# Check if there are any existing tests that combine SELECT * with highlight
rg "SELECT \*|AllFields" -n --type java integ-test/src/test/java/org/opensearch/sql/sql/HighlightFunctionIT.javaRepository: opensearch-project/sql
Length of output: 48
🏁 Script executed:
# Find the correct path for OpenSearchIndex.java
fd "OpenSearchIndex.java" --type fRepository: opensearch-project/sql
Length of output: 243
🏁 Script executed:
# Check the METADATAFIELD_TYPE_MAP definition to confirm _highlight is NOT in it
rg "METADATAFIELD_TYPE_MAP\|_highlight" opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java -A 10 -B 2Repository: opensearch-project/sql
Length of output: 48
🏁 Script executed:
# Check if tests verify SELECT * with highlight scenarios
rg "SELECT.*\*.*highlight|highlight.*SELECT" integ-test/ -n --type java -A 3 -B 1Repository: opensearch-project/sql
Length of output: 48
🏁 Script executed:
# Look for expandWildcardPattern to understand if _highlight gets included
rg "expandWildcardPattern|WildcardUtils" -n --type java -A 5 | head -50Repository: opensearch-project/sql
Length of output: 5426
🏁 Script executed:
# Verify the exact scenario: when ALL fields are selected, does _highlight get added to expandedFields?
# Trace through: AllFields case adds all non-metadata fields
# _highlight is in currentFields (added by CalciteLogicalIndexScan when highlight config is set)
# _highlight is NOT in METADATAFIELD_TYPE_MAP, so !isMetadataField returns true
# So _highlight WOULD be added to expandedFields via the AllFields expansion
# Then the highlight logic adds it AGAIN at line 417
# Confirm _highlight is added to schema when highlight is active
rg "buildSchema|HIGHLIGHT_FIELD" opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java -B 3 -A 3Repository: opensearch-project/sql
Length of output: 1710
🏁 Script executed:
# Check if there are integration tests that cover SELECT * with highlight
fd "HighlightFunctionIT.java" --type f -exec cat {} \; | head -150Repository: opensearch-project/sql
Length of output: 253
🏁 Script executed:
# Read the HighlightFunctionIT.java test file
cat integ-test/src/test/java/org/opensearch/sql/sql/HighlightFunctionIT.java | head -200Repository: opensearch-project/sql
Length of output: 7756
🏁 Script executed:
# Let me verify the exact duplicate scenario one more time by checking both code paths:
# Path 1: AllFields expansion (lines 479-482)
# Path 2: Explicit highlight logic (lines 414-417)
# Confirm that expandProjectFields is called BEFORE the highlight logic
cd core/src/main/java/org/opensearch/sql/calcite && sed -n '393,425p' CalciteRelNodeVisitor.java | cat -nRepository: opensearch-project/sql
Length of output: 1559
🏁 Script executed:
# Now trace through a hypothetical SELECT * query with highlight:
# 1. Project node has AllFields as projectList
# 2. visitProject is called
# 3. isSingleAllFieldsProject returns true
# 4. handleAllFieldsProject is called (not visitProject logic at 411-420)
# Check handleAllFieldsProject
cd core/src/main/java/org/opensearch/sql/calcite && sed -n '424,445p' CalciteRelNodeVisitor.javaRepository: opensearch-project/sql
Length of output: 951
🏁 Script executed:
# Check what happens in the non-AllFields-only case
# If someone does: SELECT name, * with highlight
# Then isSingleAllFieldsProject returns FALSE
# Then expandProjectFields is called
# Then highlight logic runs
# Verify: when is expandProjectFields called?
cd core/src/main/java/org/opensearch/sql/calcite && sed -n '395,404p' CalciteRelNodeVisitor.javaRepository: opensearch-project/sql
Length of output: 397
🏁 Script executed:
# Final verification: construct the exact scenario that would trigger the bug
# SELECT name, * with highlight enabled
# Trace:
# 1. Project has [Field("name"), AllFields()]
# 2. isSingleAllFieldsProject returns FALSE (size=2, not 1)
# 3. expandProjectFields processes both:
# - Field "name" → added to expandedFields
# - AllFields → expands to all currentFields (which includes _highlight)
# 4. Then highlight logic adds _highlight AGAIN
# Check if there are tests covering SELECT <field>, * scenario
rg "SELECT.*,.*\*" integ-test/ core/src/test -n --type java | head -20Repository: opensearch-project/sql
Length of output: 3138
Guard against duplicate _highlight projection when using SELECT with mixed explicit fields and wildcards.
When highlight is enabled and a user queries with mixed explicit fields and AllFields (e.g., SELECT name, *), the code can add _highlight twice to the projection:
expandProjectFieldsincludes_highlightvia AllFields expansion (since_highlightis not in METADATAFIELD_TYPE_MAP)- The highlight preservation logic at lines 414-417 adds it again unconditionally
Calcite may rename duplicate columns (e.g., _highlight0, _highlight1), which breaks downstream extraction that uses exact name matching via tuple.get(HIGHLIGHT_FIELD). The suggested guard checking for existing RexInputRef matches prevents this duplication.
🔧 Suggested guard
if (CalcitePlanContext.getHighlightConfig() != null) {
int hlIndex = currentFields.indexOf(HIGHLIGHT_FIELD);
- if (hlIndex >= 0) {
- expandedFields.add(context.relBuilder.field(hlIndex));
+ if (hlIndex >= 0) {
+ boolean alreadyProjected =
+ expandedFields.stream()
+ .anyMatch(
+ rex ->
+ (rex instanceof RexInputRef ref && ref.getIndex() == hlIndex)
+ || (rex instanceof RexCall call
+ && call.getKind() == SqlKind.AS
+ && call.getOperands().get(0) instanceof RexInputRef ref
+ && ref.getIndex() == hlIndex));
+ if (!alreadyProjected) {
+ expandedFields.add(context.relBuilder.field(hlIndex));
+ }
}
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // 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.getHighlightConfig() != null) { | |
| int hlIndex = currentFields.indexOf(HIGHLIGHT_FIELD); | |
| if (hlIndex >= 0) { | |
| expandedFields.add(context.relBuilder.field(hlIndex)); | |
| } | |
| // 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.getHighlightConfig() != null) { | |
| int hlIndex = currentFields.indexOf(HIGHLIGHT_FIELD); | |
| if (hlIndex >= 0) { | |
| boolean alreadyProjected = | |
| expandedFields.stream() | |
| .anyMatch( | |
| rex -> | |
| (rex instanceof RexInputRef ref && ref.getIndex() == hlIndex) | |
| || (rex instanceof RexCall call | |
| && call.getKind() == SqlKind.AS | |
| && call.getOperands().get(0) instanceof RexInputRef ref | |
| && ref.getIndex() == hlIndex)); | |
| if (!alreadyProjected) { | |
| expandedFields.add(context.relBuilder.field(hlIndex)); | |
| } | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java`
around lines 411 - 418, When highlight is enabled
(CalcitePlanContext.getHighlightConfig() != null) the highlight preservation
block in CalciteRelNodeVisitor can add HIGHLIGHT_FIELD twice into expandedFields
(duplicate of expandProjectFields' AllFields expansion); modify the logic around
currentFields/index hlIndex and expandedFields to first scan expandedFields for
an existing RexInputRef that references the same input index (or any existing
projection that already corresponds to HIGHLIGHT_FIELD) and only add
context.relBuilder.field(hlIndex) if no matching RexInputRef exists, so the
projection is not duplicated and downstream tuple.get(HIGHLIGHT_FIELD) continues
to work.
Description.
Related Issues
Check List
--signoffor-s.By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.