#17593 Add Extension Points For Pre and Post Collection of Scores in QueryPhase #18814

atris · 2025-07-22T18:29:15Z

Resolves #17593

Problem

Plugins implementing hybrid queries and neural search currently rely on AggregationProcessor workarounds to inject custom CollectorManager implementations. This approach is fragile and breaks when the aggregation system changes.

Solution

Add extension points directly in QueryPhase around the actual score collection operation:

QueryPhaseExtension interface with beforeScoreCollection() and afterScoreCollection() methods
Extension registration via QueryPhaseSearcher.queryPhaseExtensions()
Execution in DefaultQueryPhaseSearcher with proper error isolation

Implementation

New interface:

@PublicApi(since = "3.0.0")
public interface QueryPhaseExtension {
    void beforeScoreCollection(SearchContext searchContext);
    void afterScoreCollection(SearchContext searchContext);
}

Extension points placed around:
searcher.search(query, collector) // <- actual score collection

Features:
- Extensions execute before and after score collection in try/finally blocks
- Individual extension failures are logged but don't break search
- Default implementations return empty extension lists (backward compatible)
- Zero overhead when no extensions are registered

Usage

public class HybridQueryPhaseSearcher extends DefaultQueryPhaseSearcher {
    @Override
    public List<QueryPhaseExtension> queryPhaseExtensions() {
        return List.of(new HybridQueryExtension(), new NeuralSearchExtension());
    }
}

Testing

Unit tests verify:

Interface functionality
Extension registration
Integration with search execution
Error handling and isolation

No existing functionality affected.

atris · 2025-07-22T18:29:47Z

@getsaurabh02

github-actions · 2025-07-22T18:42:08Z

❌ Gradle check result for da66380: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

owaiskazi19 · 2025-07-22T22:13:26Z

@atris now that we have a way to inject QueryCollectorContext in core, would that solve your need? #18278

martin-gaievski · 2025-07-22T22:50:24Z

server/src/test/java/org/opensearch/search/query/QueryPhaseExtensionTests.java

+        try {
+            searcher.searchWith(searchContext, indexSearcher, query, collectors, false, false);
+        } catch (Exception ignored) {
+        }


we should fail the test in case of exception

martin-gaievski · 2025-07-22T22:54:23Z

server/src/main/java/org/opensearch/search/query/QueryPhaseExtension.java

+ *
+ * @opensearch.api
+ */
+@PublicApi(since = "3.0.0")


this should be 3.2.0 or whatever next version you're aiming for. 3.0.0 has been released some time ago

martin-gaievski · 2025-07-22T23:20:27Z

server/src/main/java/org/opensearch/search/query/QueryPhase.java

+                try {
+                    extension.beforeScoreCollection(searchContext);
+                } catch (Exception e) {
+                    LOGGER.warn(new ParameterizedMessage("Failed to execute beforeScoreCollection extension [{}]", extension.getClass().getName()), e);


can we make this behavior configurable, and fail execution by default, similar to behavior of ignore_failure flag? This can be part of extension interface

martin-gaievski · 2025-07-22T23:21:25Z

server/src/test/java/org/opensearch/search/query/QueryPhaseExtensionTests.java

+/**
+ * Unit tests for QueryPhaseExtension interface.
+ */
+public class QueryPhaseExtensionTests extends OpenSearchTestCase {


I think we're missing the test case for scenario when extension is throwing exception, if so please add test for such scenario

github-actions · 2025-07-23T15:05:41Z

❌ Gradle check result for 0a0f121: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

github-actions · 2025-07-28T17:02:33Z

❌ Gradle check result for ae4b60c: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

Signed-off-by: Atri Sharma <[email protected]> Signed-off-by: Atri Sharma <[email protected]>

github-actions · 2025-07-28T19:02:13Z

❌ Gradle check result for 51d8de7: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

Signed-off-by: Atri Sharma <[email protected]>

github-actions · 2025-07-28T21:45:19Z

❌ Gradle check result for ddf59ab: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

github-actions · 2025-07-29T07:04:23Z

❌ Gradle check result for ddf59ab: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

atris · 2025-07-29T07:11:43Z

Flaky tests #15813

github-actions · 2025-07-29T08:36:12Z

✅ Gradle check result for ddf59ab: SUCCESS

codecov · 2025-07-29T08:36:47Z

Codecov Report

❌ Patch coverage is 72.50000% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.88%. Comparing base (42a1dd2) to head (ddf59ab).
⚠️ Report is 23 commits behind head on main.

Files with missing lines	Patch %	Lines
...rch/search/query/ConcurrentQueryPhaseSearcher.java	55.00%	8 Missing and 1 partial ⚠️
...n/java/org/opensearch/search/query/QueryPhase.java	89.47%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #18814      +/-   ##
============================================
+ Coverage     72.80%   72.88%   +0.07%     
- Complexity    68564    68570       +6     
============================================
  Files          5567     5567              
  Lines        314844   314949     +105     
  Branches      45675    45690      +15     
============================================
+ Hits         229227   229545     +318     
+ Misses        67028    66770     -258     
- Partials      18589    18634      +45

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

atris · 2025-07-29T09:46:24Z

@owaiskazi19 @martin-gaievski Updated the PR and CI passes. Plsse review

martin-gaievski

Looks good to me, thank you. Please look into failing check for low test coverage, maybe you can came up with more test scenario, the gap between actual and expected procentage is tiny

Signed-off-by: Atri Sharma <[email protected]>

owaiskazi19 · 2025-07-30T17:42:18Z

server/src/test/java/org/opensearch/search/query/QueryPhaseExtensionTests.java

+            beforeContext.set(searchContext);
+            // At this point, query should be available but not results
+            assertNotNull("SearchContext should not be null in beforeScoreCollection", searchContext);
+            assertNotNull("Query should be available in beforeScoreCollection", searchContext.query());


Can we also check

assertNull("Query result should be null", searchContext.queryResult());

owaiskazi19 · 2025-07-30T17:50:24Z

@msfroh can you take another look? Overall LGTM

github-actions · 2025-07-30T18:34:51Z

❌ Gradle check result for 89f4afe: null

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

atris requested review from anasalkouz, andrross, ashking94, Bukhtawar, CEHENKLE, cwperks, dbwiddis, gbbafna, jed326, kotwanikunal, mch2, msfroh, owaiskazi19, reta, Rishikesh1159, sachinpkale, saratvemulapalli, shwetathareja, sohami, VachaShah and a team as code owners July 22, 2025 18:29

github-actions bot added enhancement Enhancement or improvement to existing feature or request Search Search query, autocomplete ...etc labels Jul 22, 2025

martin-gaievski reviewed Jul 22, 2025

View reviewed changes

atris force-pushed the add_querytime_extensions branch from 67ddc71 to ae4b60c Compare July 28, 2025 16:37

Introduce Pre Post Score Collection Hooks to QueryPhase

51d8de7

Signed-off-by: Atri Sharma <[email protected]> Signed-off-by: Atri Sharma <[email protected]>

atris force-pushed the add_querytime_extensions branch from ae4b60c to 51d8de7 Compare July 28, 2025 17:53

Fix review comments and add tests

ddf59ab

Signed-off-by: Atri Sharma <[email protected]>

opensearch-ci-bot mentioned this pull request Jul 28, 2025

[AUTOCUT] Gradle Check Flaky Test Report for FullClusterRestartIT #15813

Open

atris closed this Jul 29, 2025

atris reopened this Jul 29, 2025

atris closed this Jul 29, 2025

atris reopened this Jul 29, 2025

opensearch-ci-bot mentioned this pull request Jul 29, 2025

[AUTOCUT] Gradle Check Flaky Test Report for RemoteRestoreSnapshotIT #14324

Open

martin-gaievski approved these changes Jul 29, 2025

View reviewed changes

Add more tests and have better null handling

89f4afe

Signed-off-by: Atri Sharma <[email protected]>

owaiskazi19 reviewed Jul 30, 2025

View reviewed changes

atris closed this Jul 30, 2025

atris reopened this Jul 30, 2025

atris closed this Jul 30, 2025

atris reopened this Jul 30, 2025

opensearch-ci-bot mentioned this pull request Jul 29, 2025

[AUTOCUT] Gradle Check Flaky Test Report for ResourceAwareTasksTests #14293

Open

atris closed this Jul 30, 2025

atris reopened this Jul 30, 2025

#17593 Add Extension Points For Pre and Post Collection of Scores in QueryPhase #18814

Are you sure you want to change the base?

#17593 Add Extension Points For Pre and Post Collection of Scores in QueryPhase #18814

Conversation

atris commented Jul 22, 2025

Problem

Solution

Implementation

Uh oh!

atris commented Jul 22, 2025

Uh oh!

github-actions bot commented Jul 22, 2025

Uh oh!

owaiskazi19 commented Jul 22, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jul 23, 2025

Uh oh!

github-actions bot commented Jul 28, 2025

Uh oh!

github-actions bot commented Jul 28, 2025

Uh oh!

github-actions bot commented Jul 28, 2025

Uh oh!

github-actions bot commented Jul 29, 2025

Uh oh!

atris commented Jul 29, 2025

Uh oh!

github-actions bot commented Jul 29, 2025

Uh oh!

codecov bot commented Jul 29, 2025

Codecov Report

Uh oh!

atris commented Jul 29, 2025

Uh oh!

martin-gaievski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

owaiskazi19 commented Jul 30, 2025

Uh oh!

github-actions bot commented Jul 30, 2025

Uh oh!

Uh oh!