Introducing HybridQueryCollectorContextSpec to improve the hybrid search performance #1534

vibrantvarun · 2025-08-29T06:54:07Z

Description

This PR introduces following changes to improve search performance

Implement QueryCollectorContextSpec in Hybrid Query . It will help to directly inject hybrid search collectors in the search.
a. Earlier we relied on HybridQueryAggregationProcessor.preprocess for creating collectorManager.
We also removed empty collector context like DefaultQueryPhaseSearcherWithEmptyContext and ConcurrentQueryPhaseWithEmptyContext from the HybridQueryPhaseSearcher to avoid execution of topDocsCollector. [This was creating additional overhead]
We reduced dependency from HybridQueryAggregationProcessor by removing collectorManager logic from there. It was hacky approach. Now, we only need it to set the size in updateQueryResult during single shard scenario because in that case the fetch phase would run before normalization phase.
We removed creation of filteringWeight from the HybridCollectorManager when post filters are applied. We switched back to the default implementation coming from QueryPhase. Therefore, we no longer need FilteredCollector creation at neural-search end . OpenSearch will take care of it and wrap it in MultiCollector.
We created a dedicated class HybridSearchCollectorResultsUtil for hybrid collector results processing.
It helps to remove multiple collapse, sorting checks.
Class is dedicated for a collector result. So any change specific to result can be done there and it will reflect everywhere.

Related Issues

Resolves
18278

Check List

New functionality includes testing.
New functionality has been documented.
API changes companion pull request created.
Commits are signed per the DCO using --signoff.
Public documentation issue/PR created.

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.

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

vibrantvarun · 2025-08-29T07:17:59Z

Please start reviewing only when it is open for review. Thanks

yuye-aws · 2025-08-29T07:30:28Z

PR description?

vibrantvarun · 2025-08-29T07:48:51Z

@yuye-aws PR is in draft state. Once it is ready for review I will add PR description.

…QueryCollectorContextSpecImpl

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

…QueryCollectorContextSpecImpl

vibrantvarun · 2025-09-02T17:01:38Z

No Integ tests have been added in this PR because it does not contain any new feature. For all hybrid search features, the integ tests are already in place therefore all should pass in this enhancement. Moreover, the BWC tests are also in place.

…QueryCollectorContextSpecImpl

ryanbogan

LGTM

owaiskazi19

Great improvement and coming out of the hack of AggregationProcessor. Did initial review

src/main/java/org/opensearch/neuralsearch/search/collector/HybridQueryCollectorContextSpec.java

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java

owaiskazi19 · 2025-09-02T22:41:22Z

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java

-                filteringWeight,
-                searchContext
-            );
+        return new HybridCollectorManager(


How are we handling concurrent searches?

+1, please put this info into PR description or code level comments

I will add it in code level comments.

owaiskazi19 · 2025-09-02T22:44:18Z

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorResultsUtilParams.java

+        this.searchContext = builder.searchContext;
+        this.topDocsMerger = new TopDocsMerger(searchContext.sort());
+        this.trackTotalHitsUpTo = searchContext.trackTotalHitsUpTo();
+        this.isSortEnabled = searchContext.sort() != null;


Can we have getters here to check if sort is enabled, collapse etc?

There is @Getter added at the top of the class.

...main/java/org/opensearch/neuralsearch/search/query/util/HybridSearchCollectorResultUtil.java

owaiskazi19 · 2025-09-02T22:47:56Z

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java

            searchContext.from(0);
        }

-        Weight filteringWeight = null;


Since we remove filteredWeights and moved back to default post filters, did we verify with this change can we switch back other features of hybrid search to default as well?

Yes. We also made aggregations switch to default. Rest all sorting, collapse, pagination have custom logic specific to Hybrid query.

HybridAggregationProcessor is simply delegating the call to core aggs processor. We can even remove it if we find a way to execute a peice of code which is currently running in its post process. The updateQueryResult method.

HybridQueryPhaseSearcher also can be removed if we find a way to inject a query directly from the plugin. For example when hyrbid query comes with index alias filters then it is wrapped in boolean query. At the plugin end we extract the core hybrid part from the boolean query and inject it in the QueryPhaseSearcher. Already discussed this with @reta as currently there is no way of doing it. Therefore, we have to keep hybridQueryPhaseSearchert.

martin-gaievski

Please add more unit tests, at a glance we're missing tests for HybridQueryCollectorContextSpec and HybridQueryCollectorContextSpecFactory.
Also please make sure you executed full set of tests for aggregations using test_aggs param, those are skipped by default when you do integTest gradle command. Example commend ./gradlew ':integTest' -Dtest_aggs=true --tests "org.opensearch.neuralsearch.query.aggregation.*IT"

src/main/java/org/opensearch/neuralsearch/search/query/HybridAggregationProcessor.java

martin-gaievski · 2025-09-03T00:31:55Z

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java

-                filteringWeight,
-                searchContext
-            );
+        return new HybridCollectorManager(


+1, please put this info into PR description or code level comments

vibrantvarun · 2025-09-03T17:53:43Z

@martin-gaievski

Also please make sure you executed full set of tests for aggregations using test_aggs param, those are skipped by default when you do integTest gradle command. Example commend ./gradlew ':integTest' -Dtest_aggs=true --tests "org.opensearch.neuralsearch.query.aggregation.*IT"

 ./gradlew ':integTest' -Dtest_aggs=true --tests "org.opensearch.neuralsearch.query.aggregation.*IT"
WARNING: A restricted method in java.lang.System has been called
WARNING: java.lang.System::load has been called by net.rubygrapefruit.platform.internal.NativeLibraryLoader in an unnamed module (file:/Users/varunudr/.gradle/wrapper/dists/gradle-8.14.3-all/10utluxaxniiv4wxiphsi49nj/gradle-8.14.3/lib/native-platform-0.22-milestone-28.jar)
WARNING: Use --enable-native-access=ALL-UNNAMED to avoid a warning for callers in this module
WARNING: Restricted methods will be blocked in a future release unless native access is enabled

Starting a Gradle Daemon (subsequent builds will be faster)
=======================================
OpenSearch Build Hamster says Hello!
  Gradle Version        : 8.14.3
  OS Info               : Mac OS X 15.6.1 (aarch64)
  JDK Version           : 24 (Amazon Corretto JDK)
  JAVA_HOME             : /Library/Java/JavaVirtualMachines/amazon-corretto-24.jdk/Contents/Home
  Random Testing Seed   : 2F3C5D8DBB9754F6
  Crypto Standard       : any-supported
=======================================

> Task :compileJava
Note: /Users/varunudr/search/neural-search/src/main/java/org/opensearch/neuralsearch/search/query/util/HybridSearchCollectorResultUtil.java uses unchecked or unsafe operations.
Note: Recompile with -Xlint:unchecked for details.

> Task :compileTestJava
Note: Some input files use unchecked or unsafe operations.
Note: Recompile with -Xlint:unchecked for details.

> Task :integTest
WARNING: A terminally deprecated method in sun.misc.Unsafe has been called
WARNING: sun.misc.Unsafe::objectFieldOffset has been called by net.bytebuddy.dynamic.loading.ClassInjector$UsingUnsafe$Dispatcher$CreationAction
WARNING: Please consider reporting this to the maintainers of class net.bytebuddy.dynamic.loading.ClassInjector$UsingUnsafe$Dispatcher$CreationAction
WARNING: sun.misc.Unsafe::objectFieldOffset will be removed in a future release

WARNING: A restricted method in java.lang.System has been called
WARNING: java.lang.System::load has been called by com.sun.jna.Native in an unnamed module (file:/Users/varunudr/.gradle/caches/modules-2/files-2.1/net.java.dev.jna/jna/5.16.0/ebea09f91dc9f7048099f963fb8d6f919f0a4d9c/jna-5.16.0.jar)
WARNING: Use --enable-native-access=ALL-UNNAMED to avoid a warning for callers in this module
WARNING: Restricted methods will be blocked in a future release unless native access is enabled


[Incubating] Problems report is available at: file:///Users/varunudr/search/neural-search/build/reports/problems/problems-report.html

Deprecated Gradle features were used in this build, making it incompatible with Gradle 9.0.

You can use '--warning-mode all' to show the individual deprecation warnings and determine if they come from your own scripts or plugins.

For more on this, please refer to https://docs.gradle.org/8.14.3/userguide/command_line_interface.html#sec:command_line_warnings in the Gradle documentation.

BUILD SUCCESSFUL in 1m 1s
23 actionable tasks: 12 executed, 11 up-to-date

…ctorContextSpecFactory Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

owaiskazi19

Coming along

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorManager.java

owaiskazi19 · 2025-09-03T23:23:33Z

src/main/java/org/opensearch/neuralsearch/search/query/HybridCollectorResultsUtilParams.java

+        this.searchContext = builder.searchContext;
+        this.topDocsMerger = new TopDocsMerger(searchContext.sort());
+        this.trackTotalHitsUpTo = searchContext.trackTotalHitsUpTo();
+        this.isSortEnabled = searchContext.sort() != null;


owaiskazi19 · 2025-09-03T23:24:21Z

src/test/java/org/opensearch/neuralsearch/search/query/HybridAggregationProcessorTests.java

    }

    @SneakyThrows
-    public void testCollectorManager_whenHybridQueryAndNotConcurrentSearch_thenSuccessful() {


Now that we have rely on HybridQueryCollectorContextSpec to handle concurrent searches, how can we test concurrent searches in hybrid search? Through ITs probably

Actually we rely on ConcurrentQueryPhaseSearcher

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

martin-gaievski

Code looks good, please conduct manual testing for possible regression (like missing search hits) and how exactly this affects response times and throughput

owaiskazi19

Looks good to me. Can we have benchmark numbers before merging this in?

vibrantvarun · 2025-09-04T17:03:24Z

Looks good to me. Can we have benchmark numbers before merging this in?

@owaiskazi19 We are merging this in feature branch. Before merging in main branch I will benchmark the final product.

vibrantvarun · 2025-09-04T17:04:29Z

I will wait for one more maintainer approval. @bzhangam , @heemin32 @junqiu-lei Can one of you take a look as you have most context about this.

Thanks

bzhangam · 2025-09-04T19:58:44Z

Looks good to me. Thanks for the improvement and also reorganize the codes.

…rch performance (opensearch-project#1534) * QueryCollectorContextSpec implemenation for hybrid query Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * HybridQueryCollectContextSpec implemenation Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Optimizations in Hybrid Query Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Reformat files and add javadoc Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add changelog Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Update Changelog Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add HybridSearchCollectorResultUtilTests Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add unit test of HybridQueryCollectorContextSpec and HybridQueryCollectorContextSpecFactory Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add isHybridQuery check in postProcess of HybridAggregationProcessor Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Fix unit tests Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Modify comments to use release branch code Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> --------- Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

…rch performance (#1576) * Introducing HybridQueryCollectorContextSpec to improve the hybrid search performance (#1534) * QueryCollectorContextSpec implemenation for hybrid query Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * HybridQueryCollectContextSpec implemenation Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Optimizations in Hybrid Query Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Reformat files and add javadoc Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add changelog Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Update Changelog Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add HybridSearchCollectorResultUtilTests Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add unit test of HybridQueryCollectorContextSpec and HybridQueryCollectorContextSpecFactory Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Add isHybridQuery check in postProcess of HybridAggregationProcessor Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Fix unit tests Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Modify comments to use release branch code Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> --------- Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> * Optimize HybridAggregationProcessor postProcess Signed-off-by: vibrantvarun <jainvarun4996@gmail.com> --------- Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

vibrantvarun added 4 commits August 20, 2025 13:34

QueryCollectorContextSpec implemenation for hybrid query

638f241

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

HybridQueryCollectContextSpec implemenation

602a5e8

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Optimizations in Hybrid Query

29ce71d

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Merge branch 'main' into QueryCollectorContextSpecImpl

87a6885

vibrantvarun changed the base branch from main to feature/querycollectocontextspec_in_hybrid_query August 29, 2025 06:54

vibrantvarun changed the title ~~Query collector context spec impl~~ Introducing HybridQueryCollectorContext to improve the hybrid search performance Aug 29, 2025

vibrantvarun added 6 commits August 29, 2025 10:14

Merge branch 'feature/querycollectocontextspec_in_hybrid_query' into …

dedcfb8

…QueryCollectorContextSpecImpl

Reformat files and add javadoc

eeb7d0e

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Add changelog

04b9dd7

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Update Changelog

d75e1bb

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Add HybridSearchCollectorResultUtilTests

6e44193

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Merge branch 'feature/querycollectocontextspec_in_hybrid_query' into …

b702c22

…QueryCollectorContextSpecImpl

vibrantvarun marked this pull request as ready for review September 2, 2025 16:16

vibrantvarun requested review from VijayanB, heemin32, jmazanec15, junqiu-lei, martin-gaievski, minalsha, model-collapse, naveentatikonda, navneet1v, sean-zheng-amazon, vamshin, yuye-aws, zane-neo and zhichao-aws as code owners September 2, 2025 16:16

vibrantvarun requested a review from bzhangam as a code owner September 2, 2025 16:16

Merge branch 'feature/querycollectocontextspec_in_hybrid_query' into …

d98e0b9

…QueryCollectorContextSpecImpl

ryanbogan approved these changes Sep 2, 2025

View reviewed changes

owaiskazi19 reviewed Sep 2, 2025

View reviewed changes

martin-gaievski reviewed Sep 3, 2025

View reviewed changes

Add unit test of HybridQueryCollectorContextSpec and HybridQueryColle…

4970328

…ctorContextSpecFactory Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

vibrantvarun requested review from martin-gaievski and owaiskazi19 September 3, 2025 18:54

vibrantvarun added 2 commits September 3, 2025 12:51

Add isHybridQuery check in postProcess of HybridAggregationProcessor

723fbba

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Fix unit tests

6d3d27a

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

owaiskazi19 reviewed Sep 3, 2025

View reviewed changes

vibrantvarun added 2 commits September 4, 2025 08:41

Modify comments to use release branch code

fa778e9

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

Rebasing

e170c53

Signed-off-by: vibrantvarun <jainvarun4996@gmail.com>

vibrantvarun changed the title ~~Introducing HybridQueryCollectorContext to improve the hybrid search performance~~ Introducing HybridQueryCollectorContextSpec to improve the hybrid search performance Sep 4, 2025

martin-gaievski approved these changes Sep 4, 2025

View reviewed changes

vibrantvarun requested a review from owaiskazi19 September 4, 2025 17:02

owaiskazi19 approved these changes Sep 4, 2025

View reviewed changes

bzhangam approved these changes Sep 4, 2025

View reviewed changes

vibrantvarun merged commit f2d02b3 into opensearch-project:feature/querycollectocontextspec_in_hybrid_query Sep 4, 2025
59 checks passed

vibrantvarun mentioned this pull request Sep 23, 2025

Introducing HybridQueryCollectorContextSpec to improve the hybrid search performance #1576

Merged

5 tasks

Introducing HybridQueryCollectorContextSpec to improve the hybrid search performance #1534

Introducing HybridQueryCollectorContextSpec to improve the hybrid search performance #1534

Uh oh!

Conversation

vibrantvarun commented Aug 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issues

Check List

Uh oh!

vibrantvarun commented Aug 29, 2025

Uh oh!

yuye-aws commented Aug 29, 2025

Uh oh!

vibrantvarun commented Aug 29, 2025

Uh oh!

vibrantvarun commented Sep 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ryanbogan left a comment

Choose a reason for hiding this comment

Uh oh!

owaiskazi19 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

martin-gaievski left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vibrantvarun commented Sep 3, 2025

Uh oh!

owaiskazi19 left a comment

Choose a reason for hiding this comment

Uh oh!

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!

martin-gaievski left a comment

Choose a reason for hiding this comment

Uh oh!

owaiskazi19 left a comment

Choose a reason for hiding this comment

Uh oh!

vibrantvarun commented Sep 4, 2025

Uh oh!

vibrantvarun commented Sep 4, 2025

Uh oh!

bzhangam commented Sep 4, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

vibrantvarun commented Aug 29, 2025 •

edited

Loading

vibrantvarun commented Sep 2, 2025 •

edited

Loading