Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 187206b

Browse files
committedFeb 19, 2024·
SQL in search API
Signed-off-by: Vamsi Manohar <reddyvam@amazon.com> SQL in search API Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>
1 parent 84750b3 commit 187206b

11 files changed

+466
-7
lines changed
 

‎server/src/main/java/org/opensearch/action/search/SearchResponseSections.java

+32-1
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@
4040
import org.opensearch.search.SearchExtBuilder;
4141
import org.opensearch.search.SearchHits;
4242
import org.opensearch.search.aggregations.Aggregations;
43+
import org.opensearch.search.externalengine.QueryEngineExtBuilder;
4344
import org.opensearch.search.profile.ProfileShardResult;
4445
import org.opensearch.search.profile.SearchProfileShardResults;
4546
import org.opensearch.search.suggest.Suggest;
@@ -74,6 +75,7 @@ public class SearchResponseSections implements ToXContentFragment {
7475
protected final Boolean terminatedEarly;
7576
protected final int numReducePhases;
7677
protected final List<SearchExtBuilder> searchExtBuilders = new ArrayList<>();
78+
protected final List<QueryEngineExtBuilder> queryEngineExtBuilders = new ArrayList<>();
7779

7880
public SearchResponseSections(
7981
SearchHits hits,
@@ -84,7 +86,7 @@ public SearchResponseSections(
8486
SearchProfileShardResults profileResults,
8587
int numReducePhases
8688
) {
87-
this(hits, aggregations, suggest, timedOut, terminatedEarly, profileResults, numReducePhases, Collections.emptyList());
89+
this(hits, aggregations, suggest, timedOut, terminatedEarly, profileResults, numReducePhases, Collections.emptyList(), Collections.emptyList());
8890
}
8991

9092
public SearchResponseSections(
@@ -107,6 +109,28 @@ public SearchResponseSections(
107109
this.searchExtBuilders.addAll(Objects.requireNonNull(searchExtBuilders, "searchExtBuilders must not be null"));
108110
}
109111

112+
public SearchResponseSections(
113+
SearchHits hits,
114+
Aggregations aggregations,
115+
Suggest suggest,
116+
boolean timedOut,
117+
Boolean terminatedEarly,
118+
SearchProfileShardResults profileResults,
119+
int numReducePhases,
120+
List<SearchExtBuilder> searchExtBuilders,
121+
List<QueryEngineExtBuilder> queryEngineExtBuilders
122+
) {
123+
this.hits = hits;
124+
this.aggregations = aggregations;
125+
this.suggest = suggest;
126+
this.profileResults = profileResults;
127+
this.timedOut = timedOut;
128+
this.terminatedEarly = terminatedEarly;
129+
this.numReducePhases = numReducePhases;
130+
this.searchExtBuilders.addAll(Objects.requireNonNull(searchExtBuilders, "searchExtBuilders must not be null"));
131+
this.queryEngineExtBuilders.addAll(Objects.requireNonNull(queryEngineExtBuilders, "queryEngineExtBuilders must not be null"));
132+
}
133+
110134
public final boolean timedOut() {
111135
return this.timedOut;
112136
}
@@ -166,6 +190,13 @@ public final XContentBuilder toXContent(XContentBuilder builder, Params params)
166190
}
167191
builder.endObject();
168192
}
193+
194+
if(!queryEngineExtBuilders.isEmpty()) {
195+
for (QueryEngineExtBuilder queryEngineExtBuilder: queryEngineExtBuilders) {
196+
queryEngineExtBuilder.toXContent(builder, params);
197+
}
198+
}
199+
169200
return builder;
170201
}
171202

‎server/src/main/java/org/opensearch/action/search/TransportSearchAction.java

+5
Original file line numberDiff line numberDiff line change
@@ -428,6 +428,11 @@ private void executeRequest(
428428
relativeStartNanos,
429429
System::nanoTime
430430
);
431+
// taking over by query engine.
432+
if (!originalSearchRequest.source().queryEngines().isEmpty()) {
433+
originalSearchRequest.source().queryEngines().get(0).executeQuery(originalSearchRequest, originalListener);
434+
return;
435+
}
431436
if (originalSearchRequest.isPhaseTook() == null) {
432437
originalSearchRequest.setPhaseTook(clusterService.getClusterSettings().get(SEARCH_PHASE_TOOK_ENABLED));
433438
}

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

+48
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@
6464
import org.opensearch.search.aggregations.pipeline.MovAvgPipelineAggregator;
6565
import org.opensearch.search.aggregations.pipeline.PipelineAggregator;
6666
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
67+
import org.opensearch.search.externalengine.QueryEngineParser;
68+
import org.opensearch.search.externalengine.QueryEngine;
6769
import org.opensearch.search.fetch.FetchSubPhase;
6870
import org.opensearch.search.fetch.subphase.highlight.Highlighter;
6971
import org.opensearch.search.query.QueryPhaseSearcher;
@@ -216,6 +218,10 @@ default Optional<ExecutorServiceProvider> getIndexSearcherExecutorProvider() {
216218
return Optional.empty();
217219
}
218220

221+
default List<QueryEngineSpec<?>> getQueryEnginesSpecs() {
222+
return emptyList();
223+
}
224+
219225
/**
220226
* Executor service provider
221227
*/
@@ -877,4 +883,46 @@ public Map<String, Highlighter> getHighlighters() {
877883
return highlighters;
878884
}
879885
}
886+
887+
/**
888+
* Specification for a {@link SearchExtBuilder} which represents an additional section that can be
889+
* parsed in a search request (within the ext element).
890+
*/
891+
class QueryEngineSpec<T extends QueryEngine> extends SearchExtensionSpec<T, QueryEngineParser<T>> {
892+
/**
893+
* Specification of custom {@link SearchExtBuilder}.
894+
*
895+
* @param name holds the names by which this search ext might be parsed. The {@link ParseField#getPreferredName()} is special as it
896+
* is the name by under which the reader is registered. So it is the name that the search ext should use as its
897+
* {@link NamedWriteable#getWriteableName()} too. It is an error if {@link ParseField#getPreferredName()} conflicts with
898+
* another registered name, including names from other plugins.
899+
* @param reader the reader registered for this search ext's builder. Typically, a reference to a constructor that takes a
900+
* {@link StreamInput}
901+
* @param parser the parser function that reads the search ext builder from xcontent
902+
*/
903+
public QueryEngineSpec(
904+
ParseField name,
905+
Writeable.Reader<? extends T> reader,
906+
QueryEngineParser<T> parser
907+
) {
908+
super(name, reader, parser);
909+
}
910+
911+
/**
912+
* Specification of custom {@link SearchExtBuilder}.
913+
*
914+
* @param name the name by which this search ext might be parsed or deserialized. Make sure that the search ext builder returns this name for
915+
* {@link NamedWriteable#getWriteableName()}. It is an error if this name conflicts with another registered name, including
916+
* names from other plugins.
917+
* @param reader the reader registered for this search ext's builder. Typically, a reference to a constructor that takes a
918+
* {@link StreamInput}
919+
* @param parser the parser function that reads the search ext builder from xcontent
920+
*/
921+
public QueryEngineSpec(String name,
922+
Writeable.Reader<? extends T> reader,
923+
QueryEngineParser<T> parser) {
924+
super(name, reader, parser);
925+
}
926+
927+
}
880928
}

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

+14
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,7 @@
114114
import org.opensearch.plugins.SearchPlugin.SignificanceHeuristicSpec;
115115
import org.opensearch.plugins.SearchPlugin.SortSpec;
116116
import org.opensearch.plugins.SearchPlugin.SuggesterSpec;
117+
import org.opensearch.plugins.SearchPlugin.QueryEngineSpec;
117118
import org.opensearch.search.aggregations.AggregationBuilder;
118119
import org.opensearch.search.aggregations.BaseAggregationBuilder;
119120
import org.opensearch.search.aggregations.InternalAggregation;
@@ -240,6 +241,8 @@
240241
import org.opensearch.search.aggregations.pipeline.StatsBucketPipelineAggregationBuilder;
241242
import org.opensearch.search.aggregations.pipeline.SumBucketPipelineAggregationBuilder;
242243
import org.opensearch.search.aggregations.support.ValuesSourceRegistry;
244+
import org.opensearch.search.externalengine.QueryEngine;
245+
import org.opensearch.search.externalengine.SQLQueryEngine;
243246
import org.opensearch.search.fetch.FetchPhase;
244247
import org.opensearch.search.fetch.FetchSubPhase;
245248
import org.opensearch.search.fetch.subphase.ExplainPhase;
@@ -340,6 +343,7 @@ public SearchModule(Settings settings, List<SearchPlugin> plugins) {
340343
registerQueryParsers(plugins);
341344
registerRescorers(plugins);
342345
registerSortParsers(plugins);
346+
registerQueryEngines(plugins);
343347
registerValueFormats();
344348
registerSignificanceHeuristics(plugins);
345349
this.valuesSourceRegistry = registerAggregations(plugins);
@@ -838,6 +842,16 @@ private void registerRescorers(List<SearchPlugin> plugins) {
838842
registerFromPlugin(plugins, SearchPlugin::getRescorers, this::registerRescorer);
839843
}
840844

845+
private void registerQueryEngines(List<SearchPlugin> plugins) {
846+
registerQueryEngine(new SearchPlugin.QueryEngineSpec<>(SQLQueryEngine.NAME, SQLQueryEngine::new, SQLQueryEngine::fromXContent));
847+
registerFromPlugin(plugins, SearchPlugin::getQueryEnginesSpecs, this::registerQueryEngine);
848+
}
849+
850+
private void registerQueryEngine(SearchPlugin.QueryEngineSpec<?> spec) {
851+
namedXContents.add(new NamedXContentRegistry.Entry(QueryEngine.class, spec.getName(), (p, c) -> spec.getParser().fromXContent(p)));
852+
namedWriteables.add(new NamedWriteableRegistry.Entry(QueryEngine.class, spec.getName().getPreferredName(), spec.getReader()));
853+
}
854+
841855
private void registerRescorer(RescorerSpec<?> spec) {
842856
namedXContents.add(new NamedXContentRegistry.Entry(RescorerBuilder.class, spec.getName(), (p, c) -> spec.getParser().apply(p)));
843857
namedWriteables.add(new NamedWriteableRegistry.Entry(RescorerBuilder.class, spec.getName().getPreferredName(), spec.getReader()));

‎server/src/main/java/org/opensearch/search/builder/SearchSourceBuilder.java

+22-5
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
import org.opensearch.search.aggregations.AggregatorFactories;
6161
import org.opensearch.search.aggregations.PipelineAggregationBuilder;
6262
import org.opensearch.search.collapse.CollapseBuilder;
63+
import org.opensearch.search.externalengine.QueryEngine;
6364
import org.opensearch.search.fetch.StoredFieldsContext;
6465
import org.opensearch.search.fetch.subphase.FetchSourceContext;
6566
import org.opensearch.search.fetch.subphase.FieldAndFormat;
@@ -216,6 +217,7 @@ public static HighlightBuilder highlight() {
216217
private PointInTimeBuilder pointInTimeBuilder = null;
217218

218219
private Map<String, Object> searchPipelineSource = null;
220+
private List<QueryEngine> queryEngines = new ArrayList<>();
219221

220222
/**
221223
* Constructs a new search source builder.
@@ -1039,6 +1041,15 @@ public SearchSourceBuilder searchPipelineSource(Map<String, Object> searchPipeli
10391041
return this;
10401042
}
10411043

1044+
public List<QueryEngine> queryEngines() {
1045+
return queryEngines;
1046+
}
1047+
1048+
public List<QueryEngine> queryEngines(List<QueryEngine> queryEngines) {
1049+
this.queryEngines = queryEngines;
1050+
return queryEngines;
1051+
}
1052+
10421053
/**
10431054
* Rewrites this search source builder into its primitive form. e.g. by
10441055
* rewriting the QueryBuilder. If the builder did not change the identity
@@ -1282,11 +1293,17 @@ public void parseXContent(XContentParser parser, boolean checkTrailingTokens) th
12821293
} else if (SEARCH_PIPELINE.match(currentFieldName, parser.getDeprecationHandler())) {
12831294
searchPipelineSource = parser.mapOrdered();
12841295
} else {
1285-
throw new ParsingException(
1286-
parser.getTokenLocation(),
1287-
"Unknown key for a " + token + " in [" + currentFieldName + "].",
1288-
parser.getTokenLocation()
1289-
);
1296+
QueryEngine queryEngine = parser.namedObject(QueryEngine.class, currentFieldName, null);
1297+
if (queryEngine != null) {
1298+
queryEngines.add(queryEngine);
1299+
}
1300+
else {
1301+
throw new ParsingException(
1302+
parser.getTokenLocation(),
1303+
"Unknown key for a " + token + " in [" + currentFieldName + "].",
1304+
parser.getTokenLocation()
1305+
);
1306+
}
12901307
}
12911308
} else if (token == XContentParser.Token.START_ARRAY) {
12921309
if (STORED_FIELDS_FIELD.match(currentFieldName, parser.getDeprecationHandler())) {
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
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.externalengine;
10+
11+
import org.opensearch.action.search.SearchRequest;
12+
import org.opensearch.action.search.SearchResponse;
13+
import org.opensearch.core.action.ActionListener;
14+
import org.opensearch.core.common.io.stream.NamedWriteable;
15+
import org.opensearch.core.xcontent.ToXContentObject;
16+
17+
/**
18+
* QueryEngine abstract interface.
19+
*/
20+
public abstract class QueryEngine implements NamedWriteable, ToXContentObject {
21+
public abstract void executeQuery(SearchRequest searchRequest,
22+
ActionListener<SearchResponse> actionListener);
23+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
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+
/*
10+
* Licensed to Elasticsearch under one or more contributor
11+
* license agreements. See the NOTICE file distributed with
12+
* this work for additional information regarding copyright
13+
* ownership. Elasticsearch licenses this file to you under
14+
* the Apache License, Version 2.0 (the "License"); you may
15+
* not use this file except in compliance with the License.
16+
* You may obtain a copy of the License at
17+
*
18+
* http://www.apache.org/licenses/LICENSE-2.0
19+
*
20+
* Unless required by applicable law or agreed to in writing,
21+
* software distributed under the License is distributed on an
22+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
23+
* KIND, either express or implied. See the License for the
24+
* specific language governing permissions and limitations
25+
* under the License.
26+
*/
27+
28+
/*
29+
* Modifications Copyright OpenSearch Contributors. See
30+
* GitHub history for details.
31+
*/
32+
33+
package org.opensearch.search.externalengine;
34+
35+
import org.opensearch.common.CheckedFunction;
36+
import org.opensearch.common.annotation.PublicApi;
37+
import org.opensearch.core.common.io.stream.NamedWriteable;
38+
import org.opensearch.core.common.io.stream.StreamInput;
39+
import org.opensearch.core.common.io.stream.StreamOutput;
40+
import org.opensearch.core.common.io.stream.Writeable;
41+
import org.opensearch.core.xcontent.ToXContentFragment;
42+
import org.opensearch.plugins.SearchPlugin;
43+
import org.opensearch.plugins.SearchPlugin.SearchExtSpec;
44+
45+
/**
46+
* Intermediate serializable representation of a search ext section. To be subclassed by plugins that support
47+
* a custom section as part of a search request, which will be provided within the ext element.
48+
* Any state needs to be serialized as part of the {@link Writeable#writeTo(StreamOutput)} method and
49+
* read from the incoming stream, usually done adding a constructor that takes {@link StreamInput} as
50+
* an argument.
51+
* <p>
52+
* Registration happens through {@link SearchPlugin#getSearchExts()}, which also needs a {@link CheckedFunction} that's able to parse
53+
* the incoming request from the REST layer into the proper {@link QueryEngineExtBuilder} subclass.
54+
* <p>
55+
* {@link #getWriteableName()} must return the same name as the one used for the registration
56+
* of the {@link SearchExtSpec}.
57+
*
58+
* @see SearchExtSpec
59+
*
60+
* @opensearch.api
61+
*/
62+
@PublicApi(since = "1.0.0")
63+
public abstract class QueryEngineExtBuilder implements NamedWriteable, ToXContentFragment {
64+
65+
public abstract int hashCode();
66+
67+
public abstract boolean equals(Object obj);
68+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
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.externalengine;
10+
11+
import java.io.IOException;
12+
import org.opensearch.core.xcontent.XContentParser;
13+
import org.opensearch.index.query.QueryBuilder;
14+
15+
16+
/**
17+
* Query Engine Parser.
18+
* @param <T> extend QuerEgnien.
19+
*/
20+
@FunctionalInterface
21+
public interface QueryEngineParser<T extends QueryEngine> {
22+
23+
/**
24+
* Creates a new {@link QueryBuilder} from the query held by the
25+
* {@link XContentParser}. The state on the parser contained in this context
26+
* will be changed as a side effect of this method call
27+
*/
28+
T fromXContent(XContentParser parser) throws IOException;
29+
30+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
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.externalengine;
10+
11+
12+
import java.io.IOException;
13+
import java.util.List;
14+
import org.apache.lucene.search.TotalHits;
15+
import org.opensearch.action.search.SearchRequest;
16+
import org.opensearch.action.search.SearchResponse;
17+
import org.opensearch.action.search.ShardSearchFailure;
18+
import org.opensearch.core.ParseField;
19+
import org.opensearch.core.action.ActionListener;
20+
import org.opensearch.core.common.ParsingException;
21+
import org.opensearch.core.common.io.stream.StreamInput;
22+
import org.opensearch.core.common.io.stream.StreamOutput;
23+
import org.opensearch.core.xcontent.XContentBuilder;
24+
import org.opensearch.core.xcontent.XContentParser;
25+
import org.opensearch.search.SearchExtBuilder;
26+
import org.opensearch.search.SearchHit;
27+
import org.opensearch.search.SearchHits;
28+
import org.opensearch.search.internal.InternalSearchResponse;
29+
30+
/**
31+
* SQLQueryEngine.
32+
*/
33+
public class SQLQueryEngine extends QueryEngine {
34+
35+
public static final String NAME = "sql";
36+
37+
public SQLQueryEngine() {
38+
39+
}
40+
public SQLQueryEngine(StreamInput in) {
41+
}
42+
@Override
43+
public String getWriteableName() {
44+
return "sql";
45+
}
46+
47+
@Override
48+
public void writeTo(StreamOutput out) throws IOException {
49+
}
50+
51+
52+
@Override
53+
public void executeQuery(SearchRequest searchRequest,
54+
ActionListener<SearchResponse> actionListener) {
55+
// Creating a minimal response is OK, because SearchResponse self
56+
// is tested elsewhere.
57+
SearchHit[] hits = new SearchHit[] { };
58+
SearchResponse response = new SearchResponse(
59+
new InternalSearchResponse(
60+
new SearchHits(hits, new TotalHits(0, TotalHits.Relation.EQUAL_TO), 0),
61+
null,
62+
null,
63+
null,
64+
false,
65+
null,
66+
1,
67+
List.of(new SQLResponseExternalBuilder("13"))
68+
),
69+
null,
70+
0,
71+
0,
72+
0,
73+
0,
74+
ShardSearchFailure.EMPTY_ARRAY,
75+
SearchResponse.Clusters.EMPTY,
76+
null
77+
);
78+
actionListener.onResponse(response);
79+
}
80+
81+
@Override
82+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
83+
return null;
84+
}
85+
86+
public static QueryEngine fromXContent(XContentParser parser) throws IOException {
87+
XContentParser.Token token;
88+
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
89+
token = parser.nextToken();
90+
}
91+
return new SQLQueryEngine();
92+
}
93+
static class SQLResponseExternalBuilder extends SearchExtBuilder {
94+
95+
static ParseField DUMMY_FIELD = new ParseField("dummy");
96+
97+
protected final String id;
98+
99+
100+
public SQLResponseExternalBuilder(String id) {
101+
this.id = id;
102+
}
103+
104+
public SQLResponseExternalBuilder(StreamInput in) throws IOException {
105+
this.id = in.readString();
106+
}
107+
108+
public String getId() {
109+
return this.id;
110+
}
111+
112+
@Override
113+
public String getWriteableName() {
114+
return DUMMY_FIELD.getPreferredName();
115+
}
116+
117+
@Override
118+
public void writeTo(StreamOutput out) throws IOException {
119+
out.writeString(this.id);
120+
}
121+
122+
@Override
123+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
124+
return builder.field("dummy", id);
125+
}
126+
127+
@Override
128+
public int hashCode() {
129+
return 0;
130+
}
131+
132+
@Override
133+
public boolean equals(Object obj) {
134+
if (obj == null) {
135+
return false;
136+
}
137+
138+
if (!(obj instanceof SQLResponseExternalBuilder)) {
139+
return false;
140+
}
141+
142+
return this.id.equals(((SQLResponseExternalBuilder) obj).getId());
143+
}
144+
145+
public static SQLResponseExternalBuilder parse(XContentParser parser) throws IOException {
146+
String id;
147+
XContentParser.Token token = parser.currentToken();
148+
if (token == XContentParser.Token.VALUE_STRING) {
149+
id = parser.text();
150+
} else {
151+
throw new ParsingException(parser.getTokenLocation(), "Expected a VALUE_STRING but got " + token);
152+
}
153+
if (id == null) {
154+
throw new ParsingException(parser.getTokenLocation(), "no id specified for " + DUMMY_FIELD.getPreferredName());
155+
}
156+
return new SQLResponseExternalBuilder(id);
157+
}
158+
}
159+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
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+
/*
10+
* Licensed to Elasticsearch under one or more contributor
11+
* license agreements. See the NOTICE file distributed with
12+
* this work for additional information regarding copyright
13+
* ownership. Elasticsearch licenses this file to you under
14+
* the Apache License, Version 2.0 (the "License"); you may
15+
* not use this file except in compliance with the License.
16+
* You may obtain a copy of the License at
17+
*
18+
* http://www.apache.org/licenses/LICENSE-2.0
19+
*
20+
* Unless required by applicable law or agreed to in writing,
21+
* software distributed under the License is distributed on an
22+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
23+
* KIND, either express or implied. See the License for the
24+
* specific language governing permissions and limitations
25+
* under the License.
26+
*/
27+
28+
/**
29+
* Search phase that fetches the top hits from the shards after the results of the query phase have been merged. Pluggable by implementing
30+
* {@link org.opensearch.search.fetch.FetchSubPhase} and
31+
* {@link org.opensearch.plugins.SearchPlugin#getFetchSubPhases(org.opensearch.plugins.SearchPlugin.FetchPhaseConstructionContext)}.
32+
*/
33+
/*
34+
* Modifications Copyright OpenSearch Contributors. See
35+
* GitHub history for details.
36+
*/
37+
38+
package org.opensearch.search.externalengine;

‎server/src/main/java/org/opensearch/search/internal/InternalSearchResponse.java

+27-1
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
import org.opensearch.search.SearchExtBuilder;
4343
import org.opensearch.search.SearchHits;
4444
import org.opensearch.search.aggregations.InternalAggregations;
45+
import org.opensearch.search.externalengine.QueryEngineExtBuilder;
4546
import org.opensearch.search.profile.SearchProfileShardResults;
4647
import org.opensearch.search.suggest.Suggest;
4748

@@ -89,6 +90,20 @@ public InternalSearchResponse(
8990
super(hits, aggregations, suggest, timedOut, terminatedEarly, profileResults, numReducePhases, searchExtBuilderList);
9091
}
9192

93+
public InternalSearchResponse(
94+
SearchHits hits,
95+
InternalAggregations aggregations,
96+
Suggest suggest,
97+
SearchProfileShardResults profileResults,
98+
boolean timedOut,
99+
Boolean terminatedEarly,
100+
int numReducePhases,
101+
List<SearchExtBuilder> searchExtBuilderList,
102+
List<QueryEngineExtBuilder> queryEngineExtBuilders
103+
) {
104+
super(hits, aggregations, suggest, timedOut, terminatedEarly, profileResults, numReducePhases, searchExtBuilderList, queryEngineExtBuilders);
105+
}
106+
92107
public InternalSearchResponse(StreamInput in) throws IOException {
93108
super(
94109
new SearchHits(in),
@@ -98,7 +113,8 @@ public InternalSearchResponse(StreamInput in) throws IOException {
98113
in.readOptionalBoolean(),
99114
in.readOptionalWriteable(SearchProfileShardResults::new),
100115
in.readVInt(),
101-
readSearchExtBuildersOnOrAfter(in)
116+
readSearchExtBuildersOnOrAfter(in),
117+
readQueryEngineExtBuildersOnOrAfter(in)
102118
);
103119
}
104120

@@ -112,6 +128,7 @@ public void writeTo(StreamOutput out) throws IOException {
112128
out.writeOptionalWriteable(profileResults);
113129
out.writeVInt(numReducePhases);
114130
writeSearchExtBuildersOnOrAfter(out, searchExtBuilders);
131+
writeQueryEngineExtBuildersOnOrAfter(out, queryEngineExtBuilders);
115132
}
116133

117134
private static List<SearchExtBuilder> readSearchExtBuildersOnOrAfter(StreamInput in) throws IOException {
@@ -123,4 +140,13 @@ private static void writeSearchExtBuildersOnOrAfter(StreamOutput out, List<Searc
123140
out.writeNamedWriteableList(searchExtBuilders);
124141
}
125142
}
143+
private static List<QueryEngineExtBuilder> readQueryEngineExtBuildersOnOrAfter(StreamInput in) throws IOException {
144+
return (in.getVersion().onOrAfter(Version.V_2_12_0)) ? in.readNamedWriteableList(QueryEngineExtBuilder.class) : Collections.emptyList();
145+
}
146+
147+
private static void writeQueryEngineExtBuildersOnOrAfter(StreamOutput out, List<QueryEngineExtBuilder> queryEngineExtBuilders) throws IOException {
148+
if (out.getVersion().onOrAfter(Version.V_2_12_0)) {
149+
out.writeNamedWriteableList(queryEngineExtBuilders);
150+
}
151+
}
126152
}

0 commit comments

Comments
 (0)
Please sign in to comment.