Skip to content

Commit 3f52ac7

Browse files
kderussokderusso
committed
Propagate to
1 parent d2c22a6 commit 3f52ac7

File tree

12 files changed

+186
-28
lines changed

12 files changed

+186
-28
lines changed

‎server/src/main/java/org/elasticsearch/TransportVersions.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ static TransportVersion def(int id) {
209209
public static final TransportVersion INDEX_STATS_AND_METADATA_INCLUDE_PEAK_WRITE_LOAD = def(9_041_0_00);
210210
public static final TransportVersion REPOSITORIES_METADATA_AS_PROJECT_CUSTOM = def(9_042_0_00);
211211
public static final TransportVersion BATCHED_QUERY_PHASE_VERSION = def(9_043_0_00);
212+
public static final TransportVersion RERANK_SNIPPETS = def(9_044_0_00);
212213

213214
/*
214215
* STOP! READ THIS FIRST! No, really,

‎server/src/main/java/org/elasticsearch/action/search/RankFeaturePhase.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
2626
import org.elasticsearch.search.rank.feature.RankFeatureResult;
2727
import org.elasticsearch.search.rank.feature.RankFeatureShardRequest;
28+
import org.elasticsearch.search.rank.feature.Snippets;
2829
import org.elasticsearch.transport.Transport;
2930

3031
import java.util.Arrays;
@@ -172,7 +173,8 @@ public void onFailure(Exception e) {
172173
context.getOriginalIndices(queryResult.getShardIndex()),
173174
queryResult.getContextId(),
174175
queryResult.getShardSearchRequest(),
175-
entry
176+
entry,
177+
rankFeaturePhaseRankCoordinatorContext.snippets()
176178
),
177179
context.getTask(),
178180
listener

‎server/src/main/java/org/elasticsearch/search/rank/context/RankFeaturePhaseRankCoordinatorContext.java‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.apache.lucene.search.ScoreDoc;
1313
import org.elasticsearch.action.ActionListener;
1414
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
15+
import org.elasticsearch.search.rank.feature.Snippets;
1516

1617
import java.util.Arrays;
1718
import java.util.Comparator;
@@ -30,18 +31,24 @@ public abstract class RankFeaturePhaseRankCoordinatorContext {
3031
protected final int from;
3132
protected final int rankWindowSize;
3233
protected final boolean failuresAllowed;
34+
protected final Snippets snippets;
3335

34-
public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize, boolean failuresAllowed) {
36+
public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize, boolean failuresAllowed, Snippets snippets) {
3537
this.size = size < 0 ? DEFAULT_SIZE : size;
3638
this.from = from < 0 ? DEFAULT_FROM : from;
3739
this.rankWindowSize = rankWindowSize;
3840
this.failuresAllowed = failuresAllowed;
41+
this.snippets = snippets;
3942
}
4043

4144
public boolean failuresAllowed() {
4245
return failuresAllowed;
4346
}
4447

48+
public Snippets snippets() {
49+
return snippets;
50+
}
51+
4552
/**
4653
* Computes the updated scores for a list of features (i.e. document-based data). We also pass along an ActionListener
4754
* that should be called with the new scores, and will continue execution to the next phase

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

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,14 @@
1010
package org.elasticsearch.search.rank.feature;
1111

1212
import org.apache.lucene.search.Explanation;
13+
import org.elasticsearch.TransportVersions;
1314
import org.elasticsearch.common.io.stream.StreamInput;
1415
import org.elasticsearch.common.io.stream.StreamOutput;
1516
import org.elasticsearch.search.rank.RankDoc;
1617
import org.elasticsearch.xcontent.XContentBuilder;
1718

1819
import java.io.IOException;
20+
import java.util.List;
1921
import java.util.Objects;
2022

2123
/**
@@ -27,6 +29,7 @@ public class RankFeatureDoc extends RankDoc {
2729

2830
// TODO: update to support more than 1 fields; and not restrict to string data
2931
public String featureData;
32+
public List<String> snippets;
3033

3134
public RankFeatureDoc(int doc, float score, int shardIndex) {
3235
super(doc, score, shardIndex);
@@ -35,6 +38,9 @@ public RankFeatureDoc(int doc, float score, int shardIndex) {
3538
public RankFeatureDoc(StreamInput in) throws IOException {
3639
super(in);
3740
featureData = in.readOptionalString();
41+
if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
42+
snippets = in.readOptionalStringCollectionAsList();
43+
}
3844
}
3945

4046
@Override
@@ -46,20 +52,27 @@ public void featureData(String featureData) {
4652
this.featureData = featureData;
4753
}
4854

55+
public void snippets(List<String> snippets) {
56+
this.snippets = snippets;
57+
}
58+
4959
@Override
5060
protected void doWriteTo(StreamOutput out) throws IOException {
5161
out.writeOptionalString(featureData);
62+
if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
63+
out.writeOptionalStringCollection(snippets);
64+
}
5265
}
5366

5467
@Override
5568
protected boolean doEquals(RankDoc rd) {
5669
RankFeatureDoc other = (RankFeatureDoc) rd;
57-
return Objects.equals(this.featureData, other.featureData);
70+
return Objects.equals(this.featureData, other.featureData) && Objects.equals(this.snippets, other.snippets);
5871
}
5972

6073
@Override
6174
protected int doHashCode() {
62-
return Objects.hashCode(featureData);
75+
return Objects.hash(featureData, snippets);
6376
}
6477

6578
@Override
@@ -70,5 +83,6 @@ public String getWriteableName() {
7083
@Override
7184
protected void doToXContent(XContentBuilder builder, Params params) throws IOException {
7285
builder.field("featureData", featureData);
86+
builder.array("snippets", snippets);
7387
}
7488
}

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

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -57,12 +57,18 @@ public static void prepareForFetch(SearchContext searchContext, RankFeatureShard
5757
new FetchFieldsContext(Collections.singletonList(new FieldAndFormat(rankFeaturePhaseRankShardContext.getField(), null)))
5858
);
5959
try {
60-
SearchHighlightContext searchHighlightContext = new HighlightBuilder().field(field)
61-
.numOfFragments(1) // TODO set num fragments
62-
.preTags("")
63-
.postTags("")
64-
.build(searchContext.getSearchExecutionContext());
65-
searchContext.highlight(searchHighlightContext);
60+
Snippets snippets = request.snippets();
61+
if (snippets != null) {
62+
HighlightBuilder highlightBuilder = new HighlightBuilder().field(field).preTags("").postTags("");
63+
if (snippets.numFragments() != null) {
64+
highlightBuilder.numOfFragments(snippets.numFragments());
65+
}
66+
if (snippets.maxSize() != null) {
67+
highlightBuilder.fragmentSize(snippets.maxSize());
68+
}
69+
SearchHighlightContext searchHighlightContext = highlightBuilder.build(searchContext.getSearchExecutionContext());
70+
searchContext.highlight(searchHighlightContext);
71+
}
6672
} catch (IOException e) {
6773
throw new RuntimeException("Failed to create highlight context", e);
6874
}

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

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,14 @@
99

1010
package org.elasticsearch.search.rank.feature;
1111

12+
import org.elasticsearch.TransportVersions;
1213
import org.elasticsearch.action.IndicesRequest;
1314
import org.elasticsearch.action.OriginalIndices;
1415
import org.elasticsearch.action.search.SearchShardTask;
1516
import org.elasticsearch.action.support.IndicesOptions;
1617
import org.elasticsearch.common.io.stream.StreamInput;
1718
import org.elasticsearch.common.io.stream.StreamOutput;
19+
import org.elasticsearch.core.Nullable;
1820
import org.elasticsearch.search.internal.ShardSearchContextId;
1921
import org.elasticsearch.search.internal.ShardSearchRequest;
2022
import org.elasticsearch.tasks.TaskId;
@@ -38,16 +40,20 @@ public class RankFeatureShardRequest extends TransportRequest implements Indices
3840

3941
private final int[] docIds;
4042

43+
private final Snippets snippets;
44+
4145
public RankFeatureShardRequest(
4246
OriginalIndices originalIndices,
4347
ShardSearchContextId contextId,
4448
ShardSearchRequest shardSearchRequest,
45-
List<Integer> docIds
49+
List<Integer> docIds,
50+
@Nullable Snippets snippets
4651
) {
4752
this.originalIndices = originalIndices;
4853
this.shardSearchRequest = shardSearchRequest;
4954
this.docIds = docIds.stream().flatMapToInt(IntStream::of).toArray();
5055
this.contextId = contextId;
56+
this.snippets = snippets;
5157
}
5258

5359
public RankFeatureShardRequest(StreamInput in) throws IOException {
@@ -56,6 +62,11 @@ public RankFeatureShardRequest(StreamInput in) throws IOException {
5662
shardSearchRequest = in.readOptionalWriteable(ShardSearchRequest::new);
5763
docIds = in.readIntArray();
5864
contextId = in.readOptionalWriteable(ShardSearchContextId::new);
65+
if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
66+
snippets = in.readOptionalWriteable(Snippets::new);
67+
} else {
68+
snippets = null;
69+
}
5970
}
6071

6172
@Override
@@ -65,6 +76,9 @@ public void writeTo(StreamOutput out) throws IOException {
6576
out.writeOptionalWriteable(shardSearchRequest);
6677
out.writeIntArray(docIds);
6778
out.writeOptionalWriteable(contextId);
79+
if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
80+
out.writeOptionalWriteable(snippets);
81+
}
6882
}
6983

7084
@Override
@@ -95,6 +109,10 @@ public ShardSearchContextId contextId() {
95109
return contextId;
96110
}
97111

112+
public Snippets snippets() {
113+
return snippets;
114+
}
115+
98116
@Override
99117
public SearchShardTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
100118
return new SearchShardTask(id, type, action, getDescription(), parentTaskId, headers);

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
/*
2+
* 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".
8+
*/
9+
10+
package org.elasticsearch.search.rank.feature;
11+
12+
import org.elasticsearch.common.io.stream.StreamInput;
13+
import org.elasticsearch.common.io.stream.StreamOutput;
14+
import org.elasticsearch.common.io.stream.Writeable;
15+
16+
import java.io.IOException;
17+
18+
public record Snippets(Integer numFragments, Integer maxSize) implements Writeable {
19+
20+
public Snippets(StreamInput in) throws IOException {
21+
this(in.readOptionalVInt(), in.readOptionalVInt());
22+
}
23+
24+
@Override
25+
public void writeTo(StreamOutput out) throws IOException {
26+
out.writeOptionalVInt(numFragments);
27+
out.writeOptionalVInt(maxSize);
28+
}
29+
}

‎server/src/main/java/org/elasticsearch/search/rank/rerank/RerankingRankFeaturePhaseRankShardContext.java‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,20 @@
1212
import org.apache.logging.log4j.LogManager;
1313
import org.apache.logging.log4j.Logger;
1414
import org.elasticsearch.common.document.DocumentField;
15+
import org.elasticsearch.common.text.Text;
1516
import org.elasticsearch.search.SearchHit;
1617
import org.elasticsearch.search.SearchHits;
18+
import org.elasticsearch.search.fetch.subphase.highlight.HighlightField;
1719
import org.elasticsearch.search.rank.RankShardResult;
1820
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankShardContext;
1921
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
2022
import org.elasticsearch.search.rank.feature.RankFeatureShardResult;
2123

24+
import java.util.ArrayList;
2225
import java.util.Arrays;
26+
import java.util.List;
27+
import java.util.Map;
28+
import java.util.stream.Collectors;
2329

2430
/**
2531
* The {@code ReRankingRankFeaturePhaseRankShardContext} is handles the {@code SearchHits} generated from the {@code RankFeatureShardPhase}
@@ -40,10 +46,19 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
4046
RankFeatureDoc[] rankFeatureDocs = new RankFeatureDoc[hits.getHits().length];
4147
for (int i = 0; i < hits.getHits().length; i++) {
4248
rankFeatureDocs[i] = new RankFeatureDoc(hits.getHits()[i].docId(), hits.getHits()[i].getScore(), shardId);
43-
DocumentField docField = hits.getHits()[i].field(field);
49+
SearchHit hit = hits.getHits()[i];
50+
DocumentField docField = hit.field(field);
4451
if (docField != null) {
4552
rankFeatureDocs[i].featureData(docField.getValue().toString());
4653
}
54+
Map<String, HighlightField> highlightFields = hit.getHighlightFields();
55+
if (highlightFields != null) {
56+
HighlightField highlightField = highlightFields.get(field);
57+
if (highlightField != null) {
58+
List<String> snippets = new ArrayList<>(Arrays.stream(highlightField.fragments()).map(Text::toString).toList());
59+
rankFeatureDocs[i].snippets(snippets);
60+
}
61+
}
4762
}
4863
return new RankFeatureShardResult(rankFeatureDocs);
4964
} catch (Exception ex) {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ public class RandomRankFeaturePhaseRankCoordinatorContext extends RankFeaturePha
2222
private final Integer seed;
2323

2424
public RandomRankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize, Integer seed) {
25-
super(size, from, rankWindowSize, false);
25+
super(size, from, rankWindowSize, false, null);
2626
this.seed = seed;
2727
}
2828

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

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
2424
import org.elasticsearch.search.rank.context.RankFeaturePhaseRankShardContext;
2525
import org.elasticsearch.search.rank.feature.RankFeatureDoc;
26+
import org.elasticsearch.search.rank.feature.Snippets;
2627
import org.elasticsearch.search.rank.rerank.RerankingRankFeaturePhaseRankShardContext;
2728
import org.elasticsearch.xcontent.XContentBuilder;
2829

@@ -54,21 +55,24 @@ public class TextSimilarityRankBuilder extends RankBuilder {
5455
private final String field;
5556
private final Float minScore;
5657
private final boolean failuresAllowed;
58+
private final Snippets snippets;
5759

5860
public TextSimilarityRankBuilder(
5961
String field,
6062
String inferenceId,
6163
String inferenceText,
6264
int rankWindowSize,
6365
Float minScore,
64-
boolean failuresAllowed
66+
boolean failuresAllowed,
67+
Snippets snippets
6568
) {
6669
super(rankWindowSize);
6770
this.inferenceId = inferenceId;
6871
this.inferenceText = inferenceText;
6972
this.field = field;
7073
this.minScore = minScore;
7174
this.failuresAllowed = failuresAllowed;
75+
this.snippets = snippets;
7276
}
7377

7478
public TextSimilarityRankBuilder(StreamInput in) throws IOException {
@@ -83,6 +87,11 @@ public TextSimilarityRankBuilder(StreamInput in) throws IOException {
8387
} else {
8488
this.failuresAllowed = false;
8589
}
90+
if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
91+
this.snippets = in.readOptionalWriteable(Snippets::new);
92+
} else {
93+
this.snippets = null;
94+
}
8695
}
8796

8897
@Override
@@ -105,6 +114,9 @@ public void doWriteTo(StreamOutput out) throws IOException {
105114
if (out.getTransportVersion().onOrAfter(TransportVersions.RERANKER_FAILURES_ALLOWED)) {
106115
out.writeBoolean(failuresAllowed);
107116
}
117+
if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
118+
out.writeOptionalWriteable(snippets);
119+
}
108120
}
109121

110122
@Override
@@ -120,6 +132,9 @@ public void doXContent(XContentBuilder builder, Params params) throws IOExceptio
120132
if (failuresAllowed) {
121133
builder.field(FAILURES_ALLOWED_FIELD.getPreferredName(), true);
122134
}
135+
if (snippets != null) {
136+
137+
}
123138
}
124139

125140
@Override
@@ -179,7 +194,8 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
179194
inferenceId,
180195
inferenceText,
181196
minScore,
182-
failuresAllowed
197+
failuresAllowed,
198+
snippets
183199
);
184200
}
185201

0 commit comments

Comments
 (0)