Skip to content

Commit cc07332

Browse files
kderussokderusso
committed
Refactor move/rename SnippetConfig
1 parent 1ad29a7 commit cc07332

File tree

7 files changed

+54
-61
lines changed

7 files changed

+54
-61
lines changed

‎server/src/main/java/org/elasticsearch/search/rank/feature/RerankSnippetInput.java‎ renamed to ‎x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/SnippetConfig.java‎

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
11
/*
22
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3-
* or more contributor license agreements. Licensed under the "Elastic License
4-
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5-
* Public License v 1"; you may not use this file except in compliance with, at
6-
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7-
* License v3.0 only", or the "Server Side Public License, v 1".
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
86
*/
97

10-
package org.elasticsearch.search.rank.feature;
8+
package org.elasticsearch.xpack.inference.rank.textsimilarity;
119

1210
import org.elasticsearch.common.io.stream.StreamInput;
1311
import org.elasticsearch.common.io.stream.StreamOutput;
@@ -17,7 +15,7 @@
1715
import java.io.IOException;
1816
import java.util.Objects;
1917

20-
public class RerankSnippetInput implements Writeable {
18+
public class SnippetConfig implements Writeable {
2119

2220
public final Integer numSnippets;
2321
private final String inferenceText;
@@ -26,22 +24,22 @@ public class RerankSnippetInput implements Writeable {
2624

2725
public static final int DEFAULT_NUM_SNIPPETS = 1;
2826

29-
public RerankSnippetInput(StreamInput in) throws IOException {
27+
public SnippetConfig(StreamInput in) throws IOException {
3028
this.numSnippets = in.readOptionalVInt();
3129
this.inferenceText = in.readString();
3230
this.tokenSizeLimit = in.readOptionalVInt();
3331
this.snippetQueryBuilder = in.readOptionalNamedWriteable(QueryBuilder.class);
3432
}
3533

36-
public RerankSnippetInput(Integer numSnippets) {
34+
public SnippetConfig(Integer numSnippets) {
3735
this(numSnippets, null, null);
3836
}
3937

40-
public RerankSnippetInput(Integer numSnippets, String inferenceText, Integer tokenSizeLimit) {
38+
public SnippetConfig(Integer numSnippets, String inferenceText, Integer tokenSizeLimit) {
4139
this(numSnippets, inferenceText, tokenSizeLimit, null);
4240
}
4341

44-
public RerankSnippetInput(Integer numSnippets, String inferenceText, Integer tokenSizeLimit, QueryBuilder snippetQueryBuilder) {
42+
public SnippetConfig(Integer numSnippets, String inferenceText, Integer tokenSizeLimit, QueryBuilder snippetQueryBuilder) {
4543
this.numSnippets = numSnippets;
4644
this.inferenceText = inferenceText;
4745
this.tokenSizeLimit = tokenSizeLimit;
@@ -76,7 +74,7 @@ public QueryBuilder snippetQueryBuilder() {
7674
public boolean equals(Object o) {
7775
if (this == o) return true;
7876
if (o == null || getClass() != o.getClass()) return false;
79-
RerankSnippetInput that = (RerankSnippetInput) o;
77+
SnippetConfig that = (SnippetConfig) o;
8078
return Objects.equals(numSnippets, that.numSnippets)
8179
&& Objects.equals(inferenceText, that.inferenceText)
8280
&& Objects.equals(tokenSizeLimit, that.tokenSizeLimit)

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

Lines changed: 22 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@
2727
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
2828
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankShardContext;
2929
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
30-
import org.elasticsearch.search.rank.feature.RerankSnippetInput;
3130
import org.elasticsearch.xcontent.XContentBuilder;
3231

3332
import java.io.IOException;
@@ -72,7 +71,7 @@ public class TextSimilarityRankBuilder extends RankBuilder {
7271
private final String field;
7372
private final Float minScore;
7473
private final boolean failuresAllowed;
75-
private final RerankSnippetInput rerankSnippetInput;
74+
private final SnippetConfig snippetConfig;
7675

7776
public TextSimilarityRankBuilder(
7877
String field,
@@ -81,15 +80,15 @@ public TextSimilarityRankBuilder(
8180
int rankWindowSize,
8281
Float minScore,
8382
boolean failuresAllowed,
84-
RerankSnippetInput rerankSnippetInput
83+
SnippetConfig snippetConfig
8584
) {
8685
super(rankWindowSize);
8786
this.inferenceId = inferenceId;
8887
this.inferenceText = inferenceText;
8988
this.field = field;
9089
this.minScore = minScore;
9190
this.failuresAllowed = failuresAllowed;
92-
this.rerankSnippetInput = rerankSnippetInput;
91+
this.snippetConfig = snippetConfig;
9392
}
9493

9594
public TextSimilarityRankBuilder(StreamInput in) throws IOException {
@@ -106,9 +105,9 @@ public TextSimilarityRankBuilder(StreamInput in) throws IOException {
106105
this.failuresAllowed = false;
107106
}
108107
if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
109-
this.rerankSnippetInput = in.readOptionalWriteable(RerankSnippetInput::new);
108+
this.snippetConfig = in.readOptionalWriteable(SnippetConfig::new);
110109
} else {
111-
this.rerankSnippetInput = null;
110+
this.snippetConfig = null;
112111
}
113112
}
114113

@@ -134,7 +133,7 @@ public void doWriteTo(StreamOutput out) throws IOException {
134133
out.writeBoolean(failuresAllowed);
135134
}
136135
if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
137-
out.writeOptionalWriteable(rerankSnippetInput);
136+
out.writeOptionalWriteable(snippetConfig);
138137
}
139138
}
140139

@@ -151,16 +150,16 @@ public void doXContent(XContentBuilder builder, Params params) throws IOExceptio
151150
if (failuresAllowed) {
152151
builder.field(FAILURES_ALLOWED_FIELD.getPreferredName(), true);
153152
}
154-
if (rerankSnippetInput != null) {
155-
builder.field(SNIPPETS_FIELD.getPreferredName(), rerankSnippetInput);
153+
if (snippetConfig != null) {
154+
builder.field(SNIPPETS_FIELD.getPreferredName(), snippetConfig);
156155
}
157156
}
158157

159158
@Override
160159
public RankBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOException {
161160
TextSimilarityRankBuilder rewritten = this;
162-
if (rerankSnippetInput != null) {
163-
QueryBuilder snippetQueryBuilder = rerankSnippetInput.snippetQueryBuilder();
161+
if (snippetConfig != null) {
162+
QueryBuilder snippetQueryBuilder = snippetConfig.snippetQueryBuilder();
164163
if (snippetQueryBuilder == null) {
165164
rewritten = new TextSimilarityRankBuilder(
166165
field,
@@ -169,10 +168,10 @@ public RankBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOExc
169168
rankWindowSize(),
170169
minScore,
171170
failuresAllowed,
172-
new RerankSnippetInput(
173-
rerankSnippetInput.numSnippets(),
174-
rerankSnippetInput.inferenceText(),
175-
rerankSnippetInput.tokenSizeLimit(),
171+
new SnippetConfig(
172+
snippetConfig.numSnippets(),
173+
snippetConfig.inferenceText(),
174+
snippetConfig.tokenSizeLimit(),
176175
new MatchQueryBuilder(field, inferenceText)
177176
)
178177
);
@@ -186,10 +185,10 @@ public RankBuilder rewrite(QueryRewriteContext queryRewriteContext) throws IOExc
186185
rankWindowSize(),
187186
minScore,
188187
failuresAllowed,
189-
new RerankSnippetInput(
190-
rerankSnippetInput.numSnippets(),
191-
rerankSnippetInput.inferenceText(),
192-
rerankSnippetInput.tokenSizeLimit(),
188+
new SnippetConfig(
189+
snippetConfig.numSnippets(),
190+
snippetConfig.inferenceText(),
191+
snippetConfig.tokenSizeLimit(),
193192
rewrittenSnippetQueryBuilder
194193
)
195194
);
@@ -244,7 +243,7 @@ public QueryPhaseRankCoordinatorContext buildQueryPhaseCoordinatorContext(int si
244243

245244
@Override
246245
public RankFeaturePhaseRankShardContext buildRankFeaturePhaseShardContext() {
247-
return new TextSimilarityRerankingRankFeaturePhaseRankShardContext(field, rerankSnippetInput);
246+
return new TextSimilarityRerankingRankFeaturePhaseRankShardContext(field, snippetConfig);
248247
}
249248

250249
@Override
@@ -258,9 +257,7 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
258257
inferenceText,
259258
minScore,
260259
failuresAllowed,
261-
rerankSnippetInput != null
262-
? new RerankSnippetInput(rerankSnippetInput.numSnippets, inferenceText, tokenSizeLimit(inferenceId))
263-
: null
260+
snippetConfig != null ? new SnippetConfig(snippetConfig.numSnippets, inferenceText, tokenSizeLimit(inferenceId)) : null
264261
);
265262
}
266263

@@ -303,12 +300,12 @@ protected boolean doEquals(RankBuilder other) {
303300
&& Objects.equals(field, that.field)
304301
&& Objects.equals(minScore, that.minScore)
305302
&& failuresAllowed == that.failuresAllowed
306-
&& Objects.equals(rerankSnippetInput, that.rerankSnippetInput);
303+
&& Objects.equals(snippetConfig, that.snippetConfig);
307304
}
308305

309306
@Override
310307
protected int doHashCode() {
311-
return Objects.hash(inferenceId, inferenceText, field, minScore, failuresAllowed, rerankSnippetInput);
308+
return Objects.hash(inferenceId, inferenceText, field, minScore, failuresAllowed, snippetConfig);
312309
}
313310

314311
@Override

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import org.elasticsearch.inference.TaskType;
1616
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
1717
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
18-
import org.elasticsearch.search.rank.feature.RerankSnippetInput;
1918
import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
2019
import org.elasticsearch.xpack.core.inference.action.InferenceAction;
2120
import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
@@ -41,7 +40,7 @@ public class TextSimilarityRankFeaturePhaseRankCoordinatorContext extends RankFe
4140
protected final String inferenceId;
4241
protected final String inferenceText;
4342
protected final Float minScore;
44-
protected final RerankSnippetInput rerankSnippetInput;
43+
protected final SnippetConfig snippetConfig;
4544

4645
public TextSimilarityRankFeaturePhaseRankCoordinatorContext(
4746
int size,
@@ -52,14 +51,14 @@ public TextSimilarityRankFeaturePhaseRankCoordinatorContext(
5251
String inferenceText,
5352
Float minScore,
5453
boolean failuresAllowed,
55-
@Nullable RerankSnippetInput rerankSnippetInput
54+
@Nullable SnippetConfig snippetConfig
5655
) {
5756
super(size, from, rankWindowSize, failuresAllowed);
5857
this.client = client;
5958
this.inferenceId = inferenceId;
6059
this.inferenceText = inferenceText;
6160
this.minScore = minScore;
62-
this.rerankSnippetInput = rerankSnippetInput;
61+
this.snippetConfig = snippetConfig;
6362
}
6463

6564
@Override
@@ -81,7 +80,7 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
8180
l.onResponse(originalScores);
8281
} else {
8382
final float[] scores;
84-
if (this.rerankSnippetInput != null) {
83+
if (this.snippetConfig != null) {
8584
scores = extractScoresFromRankedSnippets(rankedDocs, featureDocs);
8685
} else {
8786
scores = extractScoresFromRankedDocs(rankedDocs);

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

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515
import org.elasticsearch.license.XPackLicenseState;
1616
import org.elasticsearch.search.builder.SearchSourceBuilder;
1717
import org.elasticsearch.search.rank.RankDoc;
18-
import org.elasticsearch.search.rank.feature.RerankSnippetInput;
1918
import org.elasticsearch.search.retriever.CompoundRetrieverBuilder;
2019
import org.elasticsearch.search.retriever.RetrieverBuilder;
2120
import org.elasticsearch.search.retriever.RetrieverParserContext;
@@ -62,7 +61,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
6261
String field = (String) args[3];
6362
int rankWindowSize = args[4] == null ? DEFAULT_RANK_WINDOW_SIZE : (int) args[4];
6463
boolean failuresAllowed = args[5] != null && (Boolean) args[5];
65-
RerankSnippetInput snippets = (RerankSnippetInput) args[6];
64+
SnippetConfig snippets = (SnippetConfig) args[6];
6665

6766
return new TextSimilarityRankRetrieverBuilder(
6867
retrieverBuilder,
@@ -75,11 +74,14 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
7574
);
7675
});
7776

78-
private static final ConstructingObjectParser<RerankSnippetInput, RetrieverParserContext> SNIPPETS_PARSER =
79-
new ConstructingObjectParser<>(SNIPPETS_FIELD.getPreferredName(), true, args -> {
77+
private static final ConstructingObjectParser<SnippetConfig, RetrieverParserContext> SNIPPETS_PARSER = new ConstructingObjectParser<>(
78+
SNIPPETS_FIELD.getPreferredName(),
79+
true,
80+
args -> {
8081
Integer numSnippets = (Integer) args[0];
81-
return new RerankSnippetInput(numSnippets);
82-
});
82+
return new SnippetConfig(numSnippets);
83+
}
84+
);
8385

8486
static {
8587
PARSER.declareNamedObject(constructorArg(), (p, c, n) -> {
@@ -115,7 +117,7 @@ public static TextSimilarityRankRetrieverBuilder fromXContent(
115117
private final String inferenceText;
116118
private final String field;
117119
private final boolean failuresAllowed;
118-
private final RerankSnippetInput snippets;
120+
private final SnippetConfig snippets;
119121

120122
public TextSimilarityRankRetrieverBuilder(
121123
RetrieverBuilder retrieverBuilder,
@@ -124,7 +126,7 @@ public TextSimilarityRankRetrieverBuilder(
124126
String field,
125127
int rankWindowSize,
126128
boolean failuresAllowed,
127-
RerankSnippetInput snippets
129+
SnippetConfig snippets
128130
) {
129131
super(List.of(RetrieverSource.from(retrieverBuilder)), rankWindowSize);
130132
this.inferenceId = inferenceId;
@@ -144,7 +146,7 @@ public TextSimilarityRankRetrieverBuilder(
144146
boolean failuresAllowed,
145147
String retrieverName,
146148
List<QueryBuilder> preFilterQueryBuilders,
147-
RerankSnippetInput snippets
149+
SnippetConfig snippets
148150
) {
149151
super(retrieverSource, rankWindowSize);
150152
if (retrieverSource.size() != 1) {
@@ -214,7 +216,7 @@ protected SearchSourceBuilder finalizeSourceBuilder(SearchSourceBuilder sourceBu
214216
minScore,
215217
failuresAllowed,
216218
snippets != null
217-
? new RerankSnippetInput(snippets.numSnippets, inferenceText, TextSimilarityRankBuilder.tokenSizeLimit(inferenceId))
219+
? new SnippetConfig(snippets.numSnippets, inferenceText, TextSimilarityRankBuilder.tokenSizeLimit(inferenceId))
218220
: null
219221
)
220222
);

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import org.elasticsearch.search.rank.RankShardResult;
2020
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
2121
import org.elasticsearch.search.rank.feature.RankFeatureShardResult;
22-
import org.elasticsearch.search.rank.feature.RerankSnippetInput;
2322
import org.elasticsearch.search.rank.rerank.RerankingRankFeaturePhaseRankShardContext;
2423
import org.elasticsearch.xcontent.Text;
2524

@@ -28,17 +27,17 @@
2827
import java.util.List;
2928
import java.util.Map;
3029

31-
import static org.elasticsearch.search.rank.feature.RerankSnippetInput.DEFAULT_NUM_SNIPPETS;
30+
import static org.elasticsearch.xpack.inference.rank.textsimilarity.SnippetConfig.DEFAULT_NUM_SNIPPETS;
3231

3332
public class TextSimilarityRerankingRankFeaturePhaseRankShardContext extends RerankingRankFeaturePhaseRankShardContext {
3433

35-
private final RerankSnippetInput snippetRankInput;
34+
private final SnippetConfig snippetRankInput;
3635

3736
// Rough approximation of token size vs. characters in highlight fragments.
3837
// TODO: highlighter should be able to set fragment size by token not length
3938
private static final int TOKEN_SIZE_LIMIT_MULTIPLIER = 5;
4039

41-
public TextSimilarityRerankingRankFeaturePhaseRankShardContext(String field, @Nullable RerankSnippetInput snippetRankInput) {
40+
public TextSimilarityRerankingRankFeaturePhaseRankShardContext(String field, @Nullable SnippetConfig snippetRankInput) {
4241
super(field);
4342
this.snippetRankInput = snippetRankInput;
4443
}

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@
1010
import org.elasticsearch.client.internal.Client;
1111
import org.elasticsearch.inference.TaskType;
1212
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
13-
import org.elasticsearch.search.rank.feature.RerankSnippetInput;
1413
import org.elasticsearch.test.ESTestCase;
1514
import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
1615
import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
@@ -49,7 +48,7 @@ public class TextSimilarityRankFeaturePhaseRankCoordinatorContextTests extends E
4948
"some query",
5049
0.0f,
5150
false,
52-
new RerankSnippetInput(2, "some query", 10)
51+
new SnippetConfig(2, "some query", 10)
5352
);
5453

5554
public void testComputeScores() {

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
import org.elasticsearch.search.rank.RankShardResult;
2929
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
3030
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankShardContext;
31-
import org.elasticsearch.search.rank.feature.RerankSnippetInput;
3231
import org.elasticsearch.search.rank.rerank.AbstractRerankerIT;
3332
import org.elasticsearch.tasks.Task;
3433
import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
@@ -178,9 +177,9 @@ public ThrowingMockRequestActionBasedRankBuilder(
178177
Float minScore,
179178
boolean failuresAllowed,
180179
String throwingType,
181-
RerankSnippetInput rerankSnippetInput
180+
SnippetConfig snippetConfig
182181
) {
183-
super(field, inferenceId, inferenceText, rankWindowSize, minScore, failuresAllowed, rerankSnippetInput);
182+
super(field, inferenceId, inferenceText, rankWindowSize, minScore, failuresAllowed, snippetConfig);
184183
this.throwingRankBuilderType = AbstractRerankerIT.ThrowingRankBuilderType.valueOf(throwingType);
185184
}
186185

0 commit comments

Comments
 (0)