Skip to content

Commit a75926c

Browse files
mridula-s109mridula-s109
committed
Introduced validations in the code and incorporated in tests - compiling version
1 parent 721b1d1 commit a75926c

File tree

3 files changed

+118
-17
lines changed

3 files changed

+118
-17
lines changed

‎x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRankDoc.java‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package org.elasticsearch.xpack.searchbusinessrules.retriever;
2+
3+
import org.elasticsearch.search.rank.RankDoc;
4+
5+
public class PinnedRankDoc extends RankDoc {
6+
private final boolean isPinned;
7+
private final String pinnedBy;
8+
9+
public PinnedRankDoc(int docId, float score, int shardIndex, boolean isPinned, String pinnedBy) {
10+
super(docId, score, shardIndex);
11+
this.isPinned = isPinned;
12+
this.pinnedBy = pinnedBy;
13+
}
14+
15+
public boolean isPinned() {
16+
return isPinned;
17+
}
18+
19+
public String getPinnedBy() {
20+
return pinnedBy;
21+
}
22+
23+
@Override
24+
public String toString() {
25+
return super.toString() + ", isPinned=" + isPinned + ", pinnedBy=" + pinnedBy;
26+
}
27+
}

‎x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,12 @@ private void validateIdsAndDocs(List<String> ids, List<SpecifiedDocument> docs)
8585
}
8686
}
8787

88+
private void validateSort(SearchSourceBuilder source) {
89+
if (source.sorts() != null && source.sorts().isEmpty() == false) {
90+
throw new IllegalArgumentException("Pinned retriever only supports sorting by score. Custom sorting is not allowed.");
91+
}
92+
}
93+
8894
public PinnedRetrieverBuilder(List<String> ids, List<SpecifiedDocument> docs, RetrieverBuilder retrieverBuilder, int rankWindowSize) {
8995
super(new ArrayList<>(), rankWindowSize);
9096
validateIdsAndDocs(ids, docs);
@@ -137,6 +143,7 @@ private QueryBuilder createPinnedQuery(QueryBuilder baseQuery) {
137143

138144
@Override
139145
protected SearchSourceBuilder finalizeSourceBuilder(SearchSourceBuilder source) {
146+
validateSort(source);
140147
source.query(createPinnedQuery(source.query()));
141148
return source;
142149
}
@@ -169,7 +176,10 @@ protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, b
169176
RankDoc[] rankDocs = new RankDoc[scoreDocs.length];
170177
for (int i = 0; i < scoreDocs.length; i++) {
171178
ScoreDoc scoreDoc = scoreDocs[i];
172-
rankDocs[i] = new RankDoc(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex);
179+
boolean isPinned = docs.stream().anyMatch(doc -> doc.id().equals(String.valueOf(scoreDoc.doc))) ||
180+
ids.contains(String.valueOf(scoreDoc.doc));
181+
String pinnedBy = docs.stream().anyMatch(doc -> doc.id().equals(String.valueOf(scoreDoc.doc))) ? "docs" : "ids";
182+
rankDocs[i] = new PinnedRankDoc(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex, isPinned, pinnedBy);
173183
rankDocs[i].rank = i + 1;
174184
}
175185
return rankDocs;
@@ -207,7 +217,8 @@ public QueryBuilder topDocsQuery() {
207217

208218
@Override
209219
public QueryBuilder explainQuery() {
210-
return createPinnedQuery(in.explainQuery());
220+
QueryBuilder baseQuery = in.explainQuery();
221+
return createPinnedQuery(baseQuery);
211222
}
212223
}
213224
}

‎x-pack/plugin/search-business-rules/src/test/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilderTests.java‎

Lines changed: 78 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,28 +27,16 @@
2727
import java.io.IOException;
2828
import java.util.ArrayList;
2929
import java.util.List;
30+
import java.util.stream.IntStream;
31+
import java.util.stream.Collectors;
3032

3133
import static org.elasticsearch.search.rank.RankBuilder.DEFAULT_RANK_WINDOW_SIZE;
3234
import static org.hamcrest.Matchers.equalTo;
3335
import static org.hamcrest.Matchers.instanceOf;
36+
import static org.hamcrest.Matchers.containsString;
3437

3538
public class PinnedRetrieverBuilderTests extends AbstractXContentTestCase<PinnedRetrieverBuilder> {
3639

37-
public static PinnedRetrieverBuilder createRandomPinnedRetrieverBuilder() {
38-
boolean useIds = randomBoolean();
39-
boolean useDocs = useIds == false || randomBoolean();
40-
41-
List<String> ids = useIds ? List.of(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLengthBetween(5, 10)) : new ArrayList<>();
42-
List<SpecifiedDocument> docs = useDocs
43-
? List.of(
44-
new SpecifiedDocument(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLengthBetween(5, 10)),
45-
new SpecifiedDocument(randomAlphaOfLengthBetween(5, 10), randomAlphaOfLengthBetween(5, 10))
46-
)
47-
: new ArrayList<>();
48-
49-
return new PinnedRetrieverBuilder(ids, docs, TestRetrieverBuilder.createRandomTestRetrieverBuilder(), randomIntBetween(1, 100));
50-
}
51-
5240
@Override
5341
protected PinnedRetrieverBuilder createTestInstance() {
5442
return createRandomPinnedRetrieverBuilder();
@@ -62,6 +50,26 @@ protected PinnedRetrieverBuilder doParseInstance(XContentParser parser) throws I
6250
);
6351
}
6452

53+
public static PinnedRetrieverBuilder createRandomPinnedRetrieverBuilder() {
54+
boolean useIds = randomBoolean();
55+
int numItems = randomIntBetween(1, 5);
56+
57+
List<String> ids = useIds ?
58+
IntStream.range(0, numItems)
59+
.mapToObj(i -> randomAlphaOfLengthBetween(5, 10))
60+
.collect(Collectors.toList()) :
61+
new ArrayList<>();
62+
List<SpecifiedDocument> docs = useIds ?
63+
new ArrayList<>() :
64+
IntStream.range(0, numItems)
65+
.mapToObj(i -> new SpecifiedDocument(
66+
randomAlphaOfLengthBetween(5, 10),
67+
randomAlphaOfLengthBetween(5, 10)
68+
))
69+
.collect(Collectors.toList());
70+
return new PinnedRetrieverBuilder(ids, docs, TestRetrieverBuilder.createRandomTestRetrieverBuilder(), randomIntBetween(1, 100));
71+
}
72+
6573
@Override
6674
protected boolean supportsUnknownFields() {
6775
return false;
@@ -93,6 +101,61 @@ protected NamedXContentRegistry xContentRegistry() {
93101
return new NamedXContentRegistry(entries);
94102
}
95103

104+
public void testValidation() {
105+
List<String> ids = List.of("id1", "id2");
106+
List<SpecifiedDocument> docs = List.of(
107+
new SpecifiedDocument("id3", "index3"),
108+
new SpecifiedDocument("id4", "index4")
109+
);
110+
111+
IllegalArgumentException e = expectThrows(
112+
IllegalArgumentException.class,
113+
() -> new PinnedRetrieverBuilder(ids, docs, TestRetrieverBuilder.createRandomTestRetrieverBuilder(), 10)
114+
);
115+
assertThat(e.getMessage(), equalTo("Both 'ids' and 'docs' cannot be specified at the same time"));
116+
117+
e = expectThrows(
118+
IllegalArgumentException.class,
119+
() -> new PinnedRetrieverBuilder(new ArrayList<>(), new ArrayList<>(), TestRetrieverBuilder.createRandomTestRetrieverBuilder(), 10)
120+
);
121+
assertThat(e.getMessage(), equalTo("Either 'ids' or 'docs' must be specified"));
122+
123+
assertNotNull(new PinnedRetrieverBuilder(ids, new ArrayList<>(), TestRetrieverBuilder.createRandomTestRetrieverBuilder(), 10));
124+
assertNotNull(new PinnedRetrieverBuilder(new ArrayList<>(), docs, TestRetrieverBuilder.createRandomTestRetrieverBuilder(), 10));
125+
}
126+
127+
public void testValidateSort() {
128+
PinnedRetrieverBuilder builder = createRandomPinnedRetrieverBuilder();
129+
130+
// Test empty sort is allowed
131+
final SearchSourceBuilder emptySource = new SearchSourceBuilder();
132+
builder.finalizeSourceBuilder(emptySource);
133+
134+
// Test score sort is allowed
135+
final SearchSourceBuilder scoreSource = new SearchSourceBuilder();
136+
scoreSource.sort("_score");
137+
builder.finalizeSourceBuilder(scoreSource);
138+
139+
// Test custom sort is not allowed
140+
final SearchSourceBuilder customSortSource = new SearchSourceBuilder();
141+
customSortSource.sort("field");
142+
IllegalArgumentException e = expectThrows(
143+
IllegalArgumentException.class,
144+
() -> builder.finalizeSourceBuilder(customSortSource)
145+
);
146+
assertThat(e.getMessage(), equalTo("Pinned retriever only supports sorting by score. Custom sorting is not allowed."));
147+
148+
// Test multiple sorts including custom sort is not allowed
149+
final SearchSourceBuilder multipleSortsSource = new SearchSourceBuilder();
150+
multipleSortsSource.sort("_score");
151+
multipleSortsSource.sort("field");
152+
e = expectThrows(
153+
IllegalArgumentException.class,
154+
() -> builder.finalizeSourceBuilder(multipleSortsSource)
155+
);
156+
assertThat(e.getMessage(), equalTo("Pinned retriever only supports sorting by score. Custom sorting is not allowed."));
157+
}
158+
96159
public void testParserDefaults() throws IOException {
97160
// Inner retriever content only sent to parser
98161
String json = """

0 commit comments

Comments
 (0)