Adds search implementation of context aware segments by shatejas · Pull Request #19558 · opensearch-project/OpenSearch

shatejas · 2025-10-07T22:00:53Z

The implementation aims at pruning search space based on context awareness when context-aware-grouping mapper is presents. It does a best attempt extraction of grouping criteria from the query, if no grouping criteria is found it will search all segements

Dependent on Indexing PR

https://github.com/opensearch-project/OpenSearch/pulls/RS146BIJAY

Related Issues

Resolves #[19093]

Check List

Functionality includes testing.
API changes companion pull request created, if applicable.
Public documentation issue/PR created, if applicable.

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.

Summary by CodeRabbit

Release Notes

New Features
- Introduced context-aware search filtering capabilities that enable segment-level query optimizations based on context criteria.
- Added automatic criteria extraction from queries to support context-aware filtering without manual configuration.
- Enhanced script validation to support stored scripts in context-aware grouping configurations.
Tests
- Expanded test coverage for criteria-based filtering and context-aware query extraction functionality.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

github-actions · 2025-10-07T22:11:52Z

❌ Gradle check result for 0ece83d: 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-10-11T00:05:04Z

❌ Gradle check result for e5db2d9: 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-10-11T00:33:20Z

❌ Gradle check result for 9b887d8: 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?

expani · 2025-10-17T18:08:59Z

server/src/main/java/org/opensearch/search/SearchService.java

        IndexShard shard = indexService.getShard(request.shardId().id());
-        Engine.SearcherSupplier reader = shard.acquireSearcherSupplier();
+        Set<String> groupingCriterias = Collections.emptySet();
+        if (request.source() != null) {


If I understand correctly, this will make criteria extraction happen at every data node.

Since it's only dependent on the QueryBuilder we can compute this once at the Coordinator and pass to every transport search task.

@expani Thats right it happens on every data node. Can you point me to the right code pointer so I can explore this to compute on coordinator node?

https://github.com/opensearch-project/OpenSearch/blob/main/server/src/main/java/org/opensearch/search/SearchService.java#L1519 Maybe during parseSource which will happen at Coordinator

parseSource also executes on data node.

I looked around and there doesn't seem to be a good place to execute on coordinator - From what I understand, iterates through the shards and executes the search on each shard, there are no operations on query builder that happen on co-ordinator in current state

expani · 2025-10-17T18:11:13Z

server/src/main/java/org/opensearch/index/query/WithFilterQueryBuilder.java

+ *
+ * @opensearch.internal
+ */
+public interface WithFilterQueryBuilder {


I don't see this being instantiated anywhere, Can you elaborate on when this will be used ?

There are plugin queries, like knn query, which can have filters, If a plugin query implements this interface - it will be able to leverage pruning in context aware cases

Added documentation around this

The name is a little misleading, consider CriteriaAwareFilterQueryBuilder/FilterAwareQueryBuilder. Also wondering should it also bind with a filter criteria or you want to keep this generic if this finds common use

Also wondering should it also bind with a filter criteria or you want to keep this generic if this finds common use

The intent was to keep it generic, but let me know if you wouldn't want it generic

The name is a little misleading, consider CriteriaAwareFilterQueryBuilder/FilterAwareQueryBuilder

Will go with FilterAwareQueryBuilder

expani · 2025-10-17T18:18:17Z

server/src/main/java/org/opensearch/common/lucene/index/OpenSearchDirectoryReader.java

+     * {@link org.opensearch.indices.IndicesRequestCache} depends on cache key generated in OpenSearchDirectoryReader
+     * This cache is used to be able to keep the cacheKey of OpensearchDirectoryReader consistent in cases of multiple criteria
+     */
+    private Cache<Set<String>, OpenSearchDirectoryReader> criteriaCache;


Based on the extractCriteria implementation, the key Set<String> will be ordered based on how they appear in the query if there are multiple criterias ( BooleanQuery )

I think we need to preserve the order for the cache lookup to work if query is same but order of criteria fields are different.

Sure I will sort the criterias in extractCriteria

expani · 2025-10-17T18:23:25Z

server/src/main/java/org/opensearch/common/lucene/index/OpenSearchDirectoryReader.java

+        // Do not do this warm up if its CriteriaBasedReader as criteriaBasedReaders are already formed
+        if (!(in instanceof CriteriaBasedReader)) {
+            warmUpCriteriaBasedReader(in.leaves());
+            criteriaCache = CacheBuilder.<Set<String>, OpenSearchDirectoryReader>builder().setMaximumWeight(1000).build();


nit: Cluster setting index.context_aware.search.criteria_cache is not used here

Good catch, will add it

I looked around an now I remember that there is no easy way of being able to extract IndexSettings in reader and thats the reason for hard coding the value. I will find a good way to pass the value

github-actions · 2025-10-20T23:06:48Z

❌ Gradle check result for b6eb84e: 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-10-21T00:27:08Z

❌ Gradle check result for 74ac3dc: 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-10-22T20:17:12Z

❌ Gradle check result for ab40f31: 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-10-23T23:46:34Z

❌ Gradle check result for 71b42b4: 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-11-25T16:30:13Z

❌ Gradle check result for 2b4cb4c: 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?

Bukhtawar · 2025-12-03T09:07:46Z

server/src/main/java/org/opensearch/index/engine/Engine.java

+    public SearcherSupplier acquireSearcherSupplier(
+        Function<Searcher, Searcher> wrapper,
+        SearcherScope scope,
+        @Nullable Set<String> criteria


nit: Use Optional instead?

Bukhtawar · 2025-12-03T09:25:11Z

server/src/main/java/org/opensearch/common/lucene/index/OpenSearchDirectoryReader.java

+
+    private OpenSearchDirectoryReader createChildDirectoryReader(final List<LeafReader> leafReaders, Set<String> criteria)
+        throws IOException {
+        final String uniqCacheKey = UUID.nameUUIDFromBytes(String.join(",", criteria).getBytes(StandardCharsets.UTF_8)).toString();


What happens if we have multiple criterion can the set have out of order but same criteria causing multiple keys(unless you are using a SortedSet)

sorting it before building a cachekey

github-actions · 2026-01-24T04:42:10Z

❌ Gradle check result for cf2c41b: 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 · 2026-01-24T18:42:26Z

❌ Gradle check result for bf66db7: 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 · 2026-01-25T00:19:33Z

❌ Gradle check result for 20e445d: 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?

The implementation aims at pruning search space based on context awareness when context-aware-grouping mapper is presents. It does a best attempt extraction of grouping criteria from the query, if no grouping criteria is found it will search all segements Signed-off-by: Tejas Shah <shatejas@amazon.com>

Signed-off-by: Tejas Shah <shatejas@amazon.com>

github-actions · 2026-01-25T03:45:35Z

❌ Gradle check result for f80670a: 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: Tejas Shah <shatejas@amazon.com>

github-actions · 2026-01-25T05:58:33Z

❌ Gradle check result for 51dccbe: 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: Tejas Shah <shatejas@amazon.com>

github-actions · 2026-01-25T16:27:51Z

❌ Gradle check result for 95a175e: 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 · 2026-01-25T22:52:43Z

❌ Gradle check result for ae09e32: 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: Tejas Shah <shatejas@amazon.com>

github-actions · 2026-01-26T01:31:53Z

❌ Gradle check result for 87f267b: 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: Tejas Shah <shatejas@amazon.com>

github-actions · 2026-01-27T02:36:59Z

❕ Gradle check result for 05d8611: UNSTABLE

Please review all flaky tests that succeeded after retry and create an issue if one does not already exist to track the flaky failure.

codecov · 2026-01-27T02:37:27Z

Codecov Report

❌ Patch coverage is 86.20690% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.28%. Comparing base (672039d) to head (e82547b).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
...common/lucene/index/OpenSearchDirectoryReader.java	84.09%	6 Missing and 1 partial ⚠️
...textaware/ContextAwareCriteriaQueryExtraction.java	92.40%	2 Missing and 4 partials ⚠️
.../main/java/org/opensearch/index/engine/Engine.java	50.00%	2 Missing and 2 partials ⚠️
...main/java/org/opensearch/search/SearchService.java	66.66%	1 Missing and 1 partial ⚠️
.../index/mapper/ContextAwareGroupingFieldMapper.java	0.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #19558      +/-   ##
============================================
+ Coverage     73.25%   73.28%   +0.03%     
- Complexity    71979    72066      +87     
============================================
  Files          5796     5797       +1     
  Lines        329287   329463     +176     
  Branches      47419    47456      +37     
============================================
+ Hits         241203   241459     +256     
+ Misses        68759    68654     -105     
- Partials      19325    19350      +25

☔ 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.

Signed-off-by: Tejas Shah <shatejas@amazon.com>

github-actions · 2026-01-28T01:36:14Z

✅ Gradle check result for e82547b: SUCCESS

Bukhtawar

Thanks for the changes, mostly looks good. Will need a second pass, meanwhile can we add integ tests to cover query filters

github-actions bot added discuss Issues intended to help drive brainstorming and decision making enhancement Enhancement or improvement to existing feature or request lucene Search Search query, autocomplete ...etc labels Oct 7, 2025

shatejas force-pushed the search-draft branch from 0ece83d to e5db2d9 Compare October 10, 2025 23:58

shatejas force-pushed the search-draft branch from e5db2d9 to 9b887d8 Compare October 11, 2025 00:15

expani reviewed Oct 17, 2025

View reviewed changes

shatejas force-pushed the search-draft branch from 9b887d8 to b6eb84e Compare October 20, 2025 22:55

shatejas force-pushed the search-draft branch from b6eb84e to 74ac3dc Compare October 21, 2025 00:16

shatejas requested a review from expani October 21, 2025 00:22

shatejas force-pushed the search-draft branch from 74ac3dc to ab40f31 Compare October 22, 2025 20:01

shatejas force-pushed the search-draft branch from ab40f31 to 71b42b4 Compare October 23, 2025 23:38

RS146BIJAY mentioned this pull request Nov 7, 2025

[META] Context Aware Segments #19918

Open

shatejas force-pushed the search-draft branch from 71b42b4 to 2b4cb4c Compare November 25, 2025 15:17

opensearch-ci-bot mentioned this pull request Nov 25, 2025

[AUTOCUT] Gradle Check Flaky Test Report for ClusterShardLimitIT #19916

Open

Bukhtawar reviewed Dec 3, 2025

View reviewed changes

shatejas marked this pull request as ready for review December 7, 2025 07:14

shatejas requested review from anasalkouz, andrross and ashking94 as code owners December 7, 2025 07:14

shatejas force-pushed the search-draft branch 2 times, most recently from 3612534 to cf2c41b Compare January 24, 2026 03:19

shatejas force-pushed the search-draft branch from bf66db7 to d0af31d Compare January 24, 2026 22:55

shatejas added 3 commits January 24, 2026 19:03

Fixes unit test

bcd57bf

Signed-off-by: Tejas Shah <shatejas@amazon.com>

Retrigger CI run

f80670a

Signed-off-by: Tejas Shah <shatejas@amazon.com>

shatejas force-pushed the search-draft branch from 20e445d to f80670a Compare January 25, 2026 03:05

shatejas added 2 commits January 24, 2026 20:39

Adresses flakiness of testCriteriaBasedReaders test

ce1cfc9

Signed-off-by: Tejas Shah <shatejas@amazon.com>

Fixes spotless apply

51dccbe

Signed-off-by: Tejas Shah <shatejas@amazon.com>

Triggering a CI rerun

95a175e

Signed-off-by: Tejas Shah <shatejas@amazon.com>

fixes spotcheck

87f267b

Signed-off-by: Tejas Shah <shatejas@amazon.com>

shatejas force-pushed the search-draft branch from ae09e32 to 87f267b Compare January 26, 2026 00:12

Merge branch 'main' into search-draft

05d8611

Signed-off-by: Tejas Shah <shatejas@amazon.com>

Merge branch 'main' into search-draft

e82547b

Signed-off-by: Tejas Shah <shatejas@amazon.com>

Bukhtawar reviewed Feb 5, 2026

View reviewed changes

Conversation

shatejas commented Oct 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Related Issues

Check List

Summary by CodeRabbit

Release Notes

Uh oh!

github-actions bot commented Oct 7, 2025

Uh oh!

github-actions bot commented Oct 11, 2025

Uh oh!

github-actions bot commented Oct 11, 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!

Bukhtawar Dec 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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 Oct 20, 2025

Uh oh!

github-actions bot commented Oct 21, 2025

Uh oh!

github-actions bot commented Oct 22, 2025

Uh oh!

github-actions bot commented Oct 23, 2025

Uh oh!

github-actions bot commented Nov 25, 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!

github-actions bot commented Jan 24, 2026

Uh oh!

github-actions bot commented Jan 24, 2026

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Jan 25, 2026

Uh oh!

github-actions bot commented Jan 26, 2026

Uh oh!

github-actions bot commented Jan 27, 2026

Uh oh!

codecov bot commented Jan 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

shatejas commented Oct 7, 2025 •

edited

Loading

Bukhtawar Dec 3, 2025 •

edited

Loading

codecov bot commented Jan 27, 2026 •

edited

Loading