Skip to content

Commit d4ad759

Browse files
weizijunweizijun
committed
rerank buildRankFeaturePhaseShardContext add SearchContext parameter
TextSimilarityRankBuilder add check field
1 parent 37a3630 commit d4ad759

File tree

11 files changed

+34
-17
lines changed

11 files changed

+34
-17
lines changed

‎server/src/internalClusterTest/java/org/elasticsearch/search/rank/FieldBasedRerankerIT.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.elasticsearch.plugins.Plugin;
2525
import org.elasticsearch.plugins.SearchPlugin;
2626
import org.elasticsearch.search.SearchHits;
27+
import org.elasticsearch.search.internal.SearchContext;
2728
import org.elasticsearch.search.query.QuerySearchResult;
2829
import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
2930
import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
@@ -185,7 +186,7 @@ public ScoreDoc[] rankQueryPhaseResults(
185186
}
186187

187188
@Override
188-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
189+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
189190
return new RankFeaturePhaseRankShardContext(field) {
190191
@Override
191192
public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId) {
@@ -330,7 +331,7 @@ public ScoreDoc[] rankQueryPhaseResults(
330331
}
331332

332333
@Override
333-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
334+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
334335
if (this.throwingRankBuilderType == ThrowingRankBuilderType.THROWING_RANK_FEATURE_PHASE_SHARD_CONTEXT)
335336
return new RankFeaturePhaseRankShardContext(field) {
336337
@Override
@@ -339,7 +340,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
339340
}
340341
};
341342
else {
342-
return super.buildRankFeaturePhaseShardContext();
343+
return super.buildRankFeaturePhaseShardContext(searchContext);
343344
}
344345
}
345346

‎server/src/internalClusterTest/java/org/elasticsearch/search/rank/MockedRequestActionBasedRerankerIT.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import org.elasticsearch.plugins.Plugin;
3434
import org.elasticsearch.plugins.SearchPlugin;
3535
import org.elasticsearch.search.SearchHits;
36+
import org.elasticsearch.search.internal.SearchContext;
3637
import org.elasticsearch.search.query.QuerySearchResult;
3738
import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
3839
import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
@@ -389,7 +390,7 @@ public QueryPhaseRankCoordinatorContext buildQueryPhaseCoordinatorContext(int si
389390
}
390391

391392
@Override
392-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
393+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
393394
return new RerankingRankFeaturePhaseRankShardContext(field);
394395
}
395396

@@ -532,7 +533,7 @@ public ScoreDoc[] rankQueryPhaseResults(
532533
}
533534

534535
@Override
535-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
536+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
536537
if (this.throwingRankBuilderType == ThrowingRankBuilderType.THROWING_RANK_FEATURE_PHASE_SHARD_CONTEXT)
537538
return new RankFeaturePhaseRankShardContext(field) {
538539
@Override

‎server/src/main/java/org/elasticsearch/search/rank/RankBuilder.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.elasticsearch.features.NodeFeature;
2222
import org.elasticsearch.search.SearchService;
2323
import org.elasticsearch.search.builder.SearchSourceBuilder;
24+
import org.elasticsearch.search.internal.SearchContext;
2425
import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
2526
import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
2627
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
@@ -106,7 +107,7 @@ public int rankWindowSize() {
106107
* Generates a context used to execute the rank feature phase on the shard. This is responsible for retrieving any needed
107108
* feature data, and passing them back to the coordinator through the appropriate {@link RankShardResult}.
108109
*/
109-
public abstract RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext();
110+
public abstract RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext);
110111

111112
/**
112113
* Generates a context used to perform global ranking during the RankFeature phase,

‎server/src/main/java/org/elasticsearch/search/rank/feature/RankFeatureShardPhase.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ public static void processFetch(SearchContext searchContext) {
6868
}
6969

7070
RankFeaturePhaseRankShardContext rankFeaturePhaseRankShardContext = searchContext.request().source().rankBuilder() != null
71-
? searchContext.request().source().rankBuilder().buildRankFeaturePhaseShardContext()
71+
? searchContext.request().source().rankBuilder().buildRankFeaturePhaseShardContext(searchContext)
7272
: null;
7373
if (rankFeaturePhaseRankShardContext != null) {
7474
// TODO: here we populate the profile part of the fetchResult as well
@@ -94,7 +94,7 @@ public static void processFetch(SearchContext searchContext) {
9494

9595
private static RankFeaturePhaseRankShardContext shardContext(SearchContext searchContext) {
9696
return searchContext.request().source() != null && searchContext.request().source().rankBuilder() != null
97-
? searchContext.request().source().rankBuilder().buildRankFeaturePhaseShardContext()
97+
? searchContext.request().source().rankBuilder().buildRankFeaturePhaseShardContext(searchContext)
9898
: null;
9999
}
100100
}

‎server/src/test/java/org/elasticsearch/action/search/RankFeaturePhaseTests.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import org.elasticsearch.search.SearchPhaseResult;
3333
import org.elasticsearch.search.SearchShardTarget;
3434
import org.elasticsearch.search.builder.SearchSourceBuilder;
35+
import org.elasticsearch.search.internal.SearchContext;
3536
import org.elasticsearch.search.internal.ShardSearchContextId;
3637
import org.elasticsearch.search.query.QuerySearchResult;
3738
import org.elasticsearch.search.rank.RankBuilder;
@@ -900,7 +901,7 @@ public QueryPhaseRankCoordinatorContext buildQueryPhaseCoordinatorContext(int si
900901
}
901902

902903
@Override
903-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
904+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
904905
return rankFeaturePhaseRankShardContext;
905906
}
906907

@@ -976,7 +977,7 @@ private void buildRankFeatureResult(
976977
try {
977978
hits = SearchHits.unpooled(searchHits, new TotalHits(totalHits, TotalHits.Relation.EQUAL_TO), maxScore);
978979
// construct the appropriate RankFeatureDoc objects based on the rank builder
979-
RankFeaturePhaseRankShardContext rankFeaturePhaseRankShardContext = shardRankBuilder.buildRankFeaturePhaseShardContext();
980+
RankFeaturePhaseRankShardContext rankFeaturePhaseRankShardContext = shardRankBuilder.buildRankFeaturePhaseShardContext(null);
980981
RankFeatureShardResult rankShardResult = (RankFeatureShardResult) rankFeaturePhaseRankShardContext.buildRankFeatureShardResult(
981982
hits,
982983
shardTarget.getShardId().id()

‎server/src/test/java/org/elasticsearch/search/rank/RankFeatureShardPhaseTests.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -152,7 +152,7 @@ public QueryPhaseRankCoordinatorContext buildQueryPhaseCoordinatorContext(int si
152152
}
153153

154154
@Override
155-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
155+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
156156
return new RankFeaturePhaseRankShardContext(field) {
157157
@Override
158158
public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId) {

‎test/framework/src/main/java/org/elasticsearch/search/rank/TestRankBuilder.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import org.elasticsearch.client.internal.Client;
1717
import org.elasticsearch.common.io.stream.StreamInput;
1818
import org.elasticsearch.common.io.stream.StreamOutput;
19+
import org.elasticsearch.search.internal.SearchContext;
1920
import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
2021
import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
2122
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
@@ -100,7 +101,7 @@ public QueryPhaseRankCoordinatorContext buildQueryPhaseCoordinatorContext(int si
100101
}
101102

102103
@Override
103-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
104+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
104105
throw new UnsupportedOperationException();
105106
}
106107

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/random/RandomRankBuilder.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.elasticsearch.client.internal.Client;
1515
import org.elasticsearch.common.io.stream.StreamInput;
1616
import org.elasticsearch.common.io.stream.StreamOutput;
17+
import org.elasticsearch.search.internal.SearchContext;
1718
import org.elasticsearch.search.rank.RankBuilder;
1819
import org.elasticsearch.search.rank.RankDoc;
1920
import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
@@ -139,7 +140,7 @@ public QueryPhaseRankCoordinatorContext buildQueryPhaseCoordinatorContext(int si
139140
}
140141

141142
@Override
142-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
143+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
143144
return new RerankingRankFeaturePhaseRankShardContext(field);
144145
}
145146

‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankBuilder.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
import org.elasticsearch.client.internal.Client;
1515
import org.elasticsearch.common.io.stream.StreamInput;
1616
import org.elasticsearch.common.io.stream.StreamOutput;
17+
import org.elasticsearch.index.mapper.Mapper;
1718
import org.elasticsearch.license.License;
1819
import org.elasticsearch.license.LicensedFeature;
20+
import org.elasticsearch.search.internal.SearchContext;
1921
import org.elasticsearch.search.rank.RankBuilder;
2022
import org.elasticsearch.search.rank.RankDoc;
2123
import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
@@ -165,7 +167,14 @@ public QueryPhaseRankCoordinatorContext buildQueryPhaseCoordinatorContext(int si
165167
}
166168

167169
@Override
168-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
170+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
171+
172+
// check field in mapping
173+
Mapper mapper = searchContext.indexShard().mapperService().mappingLookup().getMapper(field);
174+
if (mapper == null) {
175+
throw new IllegalArgumentException("field [" + field + "] does not exist in mapping");
176+
}
177+
169178
return new RerankingRankFeaturePhaseRankShardContext(field);
170179
}
171180

‎x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityTestPlugin.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.elasticsearch.plugins.ActionPlugin;
2525
import org.elasticsearch.plugins.Plugin;
2626
import org.elasticsearch.search.SearchHits;
27+
import org.elasticsearch.search.internal.SearchContext;
2728
import org.elasticsearch.search.rank.RankBuilder;
2829
import org.elasticsearch.search.rank.RankShardResult;
2930
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
@@ -193,7 +194,7 @@ public void doWriteTo(StreamOutput out) throws IOException {
193194
}
194195

195196
@Override
196-
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
197+
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext(SearchContext searchContext) {
197198
if (this.throwingRankBuilderType == AbstractRerankerIT.ThrowingRankBuilderType.THROWING_RANK_FEATURE_PHASE_SHARD_CONTEXT)
198199
return new RankFeaturePhaseRankShardContext(field()) {
199200
@Override
@@ -202,7 +203,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
202203
}
203204
};
204205
else {
205-
return super.buildRankFeaturePhaseShardContext();
206+
return super.buildRankFeaturePhaseShardContext(searchContext);
206207
}
207208
}
208209

0 commit comments

Comments
 (0)