Skip to content

Commit 703e778

Browse files
vibrantvarunDivyansh Sharma
vibrantvarun
authored and
Divyansh Sharma
committed
Inject QueryCollectorSpecFactory from Plugins to create QueryCollectorContext in QueryPhase (opensearch-project#18637)
* Initial Commit Signed-off-by: Varun Jain <varunudr@amazon.com> * Javadocs Signed-off-by: Varun Jain <varunudr@amazon.com> * Factory design pattern for QueryCollectorContextSpec Signed-off-by: Varun Jain <varunudr@amazon.com> * Fixing Concurrent Segment Search Signed-off-by: Varun Jain <varunudr@amazon.com> * Remove spec Signed-off-by: Varun Jain <varunudr@amazon.com> * Renaming method Signed-off-by: Varun Jain <varunudr@amazon.com> * Javadoc Signed-off-by: Varun Jain <varunudr@amazon.com> * Removing support method from QueryCollectorContextSpecFactory Signed-off-by: Varun Jain <varunudr@amazon.com> * Adding javadoc Signed-off-by: Varun Jain <varunudr@amazon.com> * Add changelog Signed-off-by: Varun Jain <varunudr@amazon.com> * Add experimentalAPI annotation Signed-off-by: Varun Jain <varunudr@amazon.com> * remove redundant check from getQueryCollectorContextSpec method Signed-off-by: Varun Jain <varunudr@amazon.com> * Add unit tests Signed-off-by: Varun Jain <varunudr@amazon.com> * Remove Test annotation Signed-off-by: Varun Jain <varunudr@amazon.com> * Reformatting Signed-off-by: Varun Jain <varunudr@amazon.com> --------- Signed-off-by: Varun Jain <varunudr@amazon.com>
1 parent 6c6ff0f commit 703e778

File tree

10 files changed

+395
-2
lines changed

10 files changed

+395
-2
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
1515
- Introduce a new pull-based ingestion plugin for file-based indexing (for local testing) ([#18591](https://github.com/opensearch-project/OpenSearch/pull/18591))
1616
- Add support for search pipeline in search and msearch template ([#18564](https://github.com/opensearch-project/OpenSearch/pull/18564))
1717
- Add BooleanQuery rewrite moving constant-scoring must clauses to filter clauses ([#18510](https://github.com/opensearch-project/OpenSearch/issues/18510))
18+
- Add functionality for plugins to inject QueryCollectorContext during QueryPhase ([#18637](https://github.com/opensearch-project/OpenSearch/pull/18637))
1819
- Add support for non-timing info in profiler ([#18460](https://github.com/opensearch-project/OpenSearch/issues/18460))
1920

2021
### Changed

‎server/src/main/java/org/opensearch/plugins/SearchPlugin.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@
6868
import org.opensearch.search.deciders.ConcurrentSearchRequestDecider;
6969
import org.opensearch.search.fetch.FetchSubPhase;
7070
import org.opensearch.search.fetch.subphase.highlight.Highlighter;
71+
import org.opensearch.search.query.QueryCollectorContextSpecFactory;
7172
import org.opensearch.search.query.QueryPhaseSearcher;
7273
import org.opensearch.search.rescore.Rescorer;
7374
import org.opensearch.search.rescore.RescorerBuilder;
@@ -227,6 +228,10 @@ default Optional<ExecutorServiceProvider> getIndexSearcherExecutorProvider() {
227228
return Optional.empty();
228229
}
229230

231+
default List<QueryCollectorContextSpecFactory> getCollectorContextSpecFactories() {
232+
return emptyList();
233+
}
234+
230235
/**
231236
* Executor service provider
232237
*/

‎server/src/main/java/org/opensearch/search/SearchModule.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,7 @@
257257
import org.opensearch.search.fetch.subphase.highlight.Highlighter;
258258
import org.opensearch.search.fetch.subphase.highlight.PlainHighlighter;
259259
import org.opensearch.search.fetch.subphase.highlight.UnifiedHighlighter;
260+
import org.opensearch.search.query.QueryCollectorContextSpecRegistry;
260261
import org.opensearch.search.query.QueryPhase;
261262
import org.opensearch.search.query.QueryPhaseSearcher;
262263
import org.opensearch.search.query.QueryPhaseSearcherWrapper;
@@ -350,6 +351,7 @@ public SearchModule(Settings settings, List<SearchPlugin> plugins) {
350351
indexSearcherExecutorProvider = registerIndexSearcherExecutorProvider(plugins);
351352
namedWriteables.addAll(SortValue.namedWriteables());
352353
concurrentSearchDeciderFactories = registerConcurrentSearchDeciderFactories(plugins);
354+
registerQueryCollectorContextSpec(plugins);
353355
}
354356

355357
private Collection<ConcurrentSearchRequestDecider.Factory> registerConcurrentSearchDeciderFactories(List<SearchPlugin> plugins) {
@@ -1297,6 +1299,10 @@ private SearchPlugin.ExecutorServiceProvider registerIndexSearcherExecutorProvid
12971299
return provider;
12981300
}
12991301

1302+
private void registerQueryCollectorContextSpec(List<SearchPlugin> plugins) {
1303+
registerFromPlugin(plugins, SearchPlugin::getCollectorContextSpecFactories, QueryCollectorContextSpecRegistry::registerFactory);
1304+
}
1305+
13001306
public FetchPhase getFetchPhase() {
13011307
return new FetchPhase(fetchSubPhases);
13021308
}

‎server/src/main/java/org/opensearch/search/query/QueryCollectorArguments.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.search.query;
10+
11+
import org.opensearch.common.annotation.ExperimentalApi;
12+
13+
/**
14+
* Arguments for {@link QueryCollectorContextSpecRegistry}
15+
*/
16+
@ExperimentalApi
17+
public final class QueryCollectorArguments {
18+
private final boolean hasFilterCollector;
19+
20+
private QueryCollectorArguments(final boolean hasFilterCollector) {
21+
this.hasFilterCollector = hasFilterCollector;
22+
}
23+
24+
/**
25+
* Whether the query has a filter collector.
26+
* @return true if the query has a filter collector, false otherwise
27+
*/
28+
public boolean hasFilterCollector() {
29+
return hasFilterCollector;
30+
}
31+
32+
@Override
33+
public boolean equals(Object o) {
34+
if (this == o) return true;
35+
if (o == null || getClass() != o.getClass()) return false;
36+
37+
QueryCollectorArguments queryCollectorArguments = (QueryCollectorArguments) o;
38+
return hasFilterCollector == queryCollectorArguments.hasFilterCollector;
39+
}
40+
41+
@Override
42+
public int hashCode() {
43+
return Boolean.hashCode(hasFilterCollector);
44+
}
45+
46+
/**
47+
* {@inheritDoc}
48+
*/
49+
@Override
50+
public String toString() {
51+
return "QueryCollectorArguments[hasFilterCollector=" + hasFilterCollector + "]";
52+
}
53+
54+
/**
55+
* Builder for {@link QueryCollectorArguments}
56+
*/
57+
public static class Builder {
58+
private boolean hasFilterCollector;
59+
60+
/**
61+
* Set the flag for query has a filter collector.
62+
* @param hasFilterCollector true if the query has a filter collector, false otherwise
63+
* @return Builder instance
64+
*/
65+
public Builder hasFilterCollector(boolean hasFilterCollector) {
66+
this.hasFilterCollector = hasFilterCollector;
67+
return this;
68+
}
69+
70+
/**
71+
* Build the arguments for the query collector context spec registry.
72+
* @return QueryCollectorArguments instance
73+
*/
74+
public QueryCollectorArguments build() {
75+
return new QueryCollectorArguments(hasFilterCollector);
76+
}
77+
}
78+
}

‎server/src/main/java/org/opensearch/search/query/QueryCollectorContextSpec.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.search.query;
10+
11+
import org.apache.lucene.search.Collector;
12+
import org.apache.lucene.search.CollectorManager;
13+
import org.opensearch.common.annotation.ExperimentalApi;
14+
15+
import java.io.IOException;
16+
17+
/**
18+
* interface of QueryCollectorContextSpec
19+
*/
20+
@ExperimentalApi
21+
public interface QueryCollectorContextSpec {
22+
/**
23+
* Context name for QueryCollectorContext
24+
* @return string of context name
25+
*/
26+
String getContextName();
27+
28+
/**
29+
* Create collector
30+
* @param in collector
31+
* @return collector
32+
* @throws IOException
33+
*/
34+
Collector create(Collector in) throws IOException;
35+
36+
/**
37+
* Create collector manager
38+
* @param in collector manager
39+
* @return collector manager
40+
* @throws IOException
41+
*/
42+
CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in) throws IOException;
43+
44+
/**
45+
* Post process query result
46+
* @param result query result
47+
* @throws IOException
48+
*/
49+
void postProcess(QuerySearchResult result) throws IOException;
50+
}

‎server/src/main/java/org/opensearch/search/query/QueryCollectorContextSpecFactory.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.search.query;
10+
11+
import org.opensearch.common.annotation.ExperimentalApi;
12+
import org.opensearch.search.internal.SearchContext;
13+
14+
import java.io.IOException;
15+
import java.util.Optional;
16+
17+
/**
18+
* interface of QueryCollectorContext spec factory
19+
*/
20+
@ExperimentalApi
21+
public interface QueryCollectorContextSpecFactory {
22+
/**
23+
* @param searchContext context needed to create collector context spec
24+
* @param queryCollectorArguments arguments to create collector context spec
25+
* @return QueryCollectorContextSpec
26+
* @throws IOException
27+
*/
28+
Optional<QueryCollectorContextSpec> createQueryCollectorContextSpec(
29+
SearchContext searchContext,
30+
QueryCollectorArguments queryCollectorArguments
31+
) throws IOException;
32+
}

‎server/src/main/java/org/opensearch/search/query/QueryCollectorContextSpecRegistry.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.search.query;
10+
11+
import org.opensearch.search.internal.SearchContext;
12+
13+
import java.io.IOException;
14+
import java.util.Iterator;
15+
import java.util.List;
16+
import java.util.Optional;
17+
import java.util.concurrent.CopyOnWriteArrayList;
18+
19+
/**
20+
* Registry class to load all collector context spec factories during cluster bootstrapping
21+
*/
22+
public final class QueryCollectorContextSpecRegistry {
23+
private static final List<QueryCollectorContextSpecFactory> registry = new CopyOnWriteArrayList<>();
24+
25+
private QueryCollectorContextSpecRegistry() {}
26+
27+
/**
28+
* Get all collector context spec factories
29+
* @return list of collector context spec factories
30+
*/
31+
public static List<QueryCollectorContextSpecFactory> getCollectorContextSpecFactories() {
32+
return registry;
33+
}
34+
35+
/**
36+
* Register factory
37+
* @param factory collector context spec factory defined in plugin
38+
*/
39+
public static void registerFactory(QueryCollectorContextSpecFactory factory) {
40+
registry.add(factory);
41+
}
42+
43+
/**
44+
* Get collector context spec
45+
* @param searchContext search context
46+
* @param queryCollectorArguments query collector arguments
47+
* @return collector context spec
48+
* @throws IOException
49+
*/
50+
public static Optional<QueryCollectorContextSpec> getQueryCollectorContextSpec(
51+
final SearchContext searchContext,
52+
final QueryCollectorArguments queryCollectorArguments
53+
) throws IOException {
54+
Iterator<QueryCollectorContextSpecFactory> iterator = registry.iterator();
55+
while (iterator.hasNext()) {
56+
QueryCollectorContextSpecFactory factory = iterator.next();
57+
Optional<QueryCollectorContextSpec> spec = factory.createQueryCollectorContextSpec(searchContext, queryCollectorArguments);
58+
if (spec.isEmpty() == false) {
59+
return spec;
60+
}
61+
}
62+
return Optional.empty();
63+
}
64+
}

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

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
import java.util.List;
7777
import java.util.Map;
7878
import java.util.Objects;
79+
import java.util.Optional;
7980
import java.util.concurrent.ExecutorService;
8081
import java.util.stream.Collectors;
8182

@@ -445,9 +446,37 @@ protected boolean searchWithCollector(
445446
boolean hasFilterCollector,
446447
boolean hasTimeout
447448
) throws IOException {
449+
QueryCollectorContext queryCollectorContext = getQueryCollectorContext(searchContext, hasFilterCollector);
450+
return searchWithCollector(searchContext, searcher, query, collectors, queryCollectorContext, hasFilterCollector, hasTimeout);
451+
}
452+
453+
private QueryCollectorContext getQueryCollectorContext(SearchContext searchContext, boolean hasFilterCollector) throws IOException {
448454
// create the top docs collector last when the other collectors are known
449-
final TopDocsCollectorContext topDocsFactory = createTopDocsCollectorContext(searchContext, hasFilterCollector);
450-
return searchWithCollector(searchContext, searcher, query, collectors, topDocsFactory, hasFilterCollector, hasTimeout);
455+
final Optional<QueryCollectorContext> queryCollectorContextOpt = QueryCollectorContextSpecRegistry.getQueryCollectorContextSpec(
456+
searchContext,
457+
new QueryCollectorArguments.Builder().hasFilterCollector(hasFilterCollector).build()
458+
).map(queryCollectorContextSpec -> new QueryCollectorContext(queryCollectorContextSpec.getContextName()) {
459+
@Override
460+
Collector create(Collector in) throws IOException {
461+
return queryCollectorContextSpec.create(in);
462+
}
463+
464+
@Override
465+
CollectorManager<?, ReduceableSearchResult> createManager(CollectorManager<?, ReduceableSearchResult> in)
466+
throws IOException {
467+
return queryCollectorContextSpec.createManager(in);
468+
}
469+
470+
@Override
471+
void postProcess(QuerySearchResult result) throws IOException {
472+
queryCollectorContextSpec.postProcess(result);
473+
}
474+
});
475+
if (queryCollectorContextOpt.isPresent()) {
476+
return queryCollectorContextOpt.get();
477+
} else {
478+
return createTopDocsCollectorContext(searchContext, hasFilterCollector);
479+
}
451480
}
452481

453482
protected boolean searchWithCollector(

‎server/src/test/java/org/opensearch/search/query/QueryCollectorArgumentsTests.java‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.search.query;
10+
11+
import org.opensearch.test.OpenSearchTestCase;
12+
13+
public class QueryCollectorArgumentsTests extends OpenSearchTestCase {
14+
15+
public void testBuilder() {
16+
QueryCollectorArguments args = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
17+
18+
assertTrue(args.hasFilterCollector());
19+
}
20+
21+
public void testEquals() {
22+
QueryCollectorArguments args1 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
23+
24+
QueryCollectorArguments args2 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
25+
26+
QueryCollectorArguments args3 = new QueryCollectorArguments.Builder().hasFilterCollector(false).build();
27+
28+
assertTrue(args1.equals(args2)); // Same values
29+
assertFalse(args1.equals(args3)); // Different values
30+
assertTrue(args1.equals(args1)); // Same object
31+
}
32+
33+
public void testHashCode() {
34+
QueryCollectorArguments args1 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
35+
36+
QueryCollectorArguments args2 = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
37+
38+
assertEquals(args1.hashCode(), args2.hashCode());
39+
assertEquals(args1.hashCode(), args1.hashCode()); // Consistent
40+
}
41+
42+
public void testToString() {
43+
QueryCollectorArguments args = new QueryCollectorArguments.Builder().hasFilterCollector(true).build();
44+
45+
String result = args.toString();
46+
47+
assertEquals("QueryCollectorArguments[hasFilterCollector=true]", result);
48+
}
49+
}

0 commit comments

Comments
 (0)