Skip to content

Commit add9a9e

Browse files
pmpailispmpailis
authored
Updating RankDocRetrieverBuilderIT for lucene_snapshot branch (#114098)
1 parent 6861f7f commit add9a9e

File tree

1 file changed

+45
-45
lines changed

1 file changed

+45
-45
lines changed

‎server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RankDocRetrieverBuilderIT.java‎

Lines changed: 45 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.elasticsearch.search.sort.FieldSortBuilder;
3737
import org.elasticsearch.search.sort.NestedSortBuilder;
3838
import org.elasticsearch.search.sort.ScoreSortBuilder;
39+
import org.elasticsearch.search.sort.ShardDocSortField;
3940
import org.elasticsearch.search.sort.SortBuilder;
4041
import org.elasticsearch.search.sort.SortOrder;
4142
import org.elasticsearch.test.ESIntegTestCase;
@@ -190,8 +191,10 @@ public void testRankDocsRetrieverBasicWithPagination() {
190191
SearchSourceBuilder source = new SearchSourceBuilder();
191192
StandardRetrieverBuilder standard0 = new StandardRetrieverBuilder();
192193
// this one retrieves docs 1, 4, and 6
193-
standard0.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.queryStringQuery("quick").defaultField(TEXT_FIELD))
194-
.boost(10L);
194+
standard0.queryBuilder = QueryBuilders.boolQuery()
195+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_1")).boost(10L))
196+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_4")).boost(9L))
197+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_6")).boost(8L));
195198
StandardRetrieverBuilder standard1 = new StandardRetrieverBuilder();
196199
// this one retrieves docs 2 and 6 due to prefilter
197200
standard1.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.termsQuery(ID_FIELD, "doc_2", "doc_3", "doc_6")).boost(20L);
@@ -206,8 +209,8 @@ public void testRankDocsRetrieverBasicWithPagination() {
206209
null
207210
);
208211
// the compound retriever here produces a score for a doc based on the percentage of the queries that it was matched on and
209-
// resolves ties based on actual score, rank, and then the doc (we're forcing 1 shard for consistent results)
210-
// so ideal rank would be: 6, 2, 1, 4, 7, 3 and with pagination, we'd just omit the first result
212+
// resolves ties based on actual score, and then the doc (we're forcing 1 shard for consistent results)
213+
// so ideal rank would be: 6, 2, 1, 3, 4, 7 and with pagination, we'd just omit the first result
211214
source.retriever(
212215
new CompoundRetrieverWithRankDocs(
213216
rankWindowSize,
@@ -228,9 +231,9 @@ public void testRankDocsRetrieverBasicWithPagination() {
228231
assertThat(resp.getHits().getTotalHits().relation(), equalTo(TotalHits.Relation.EQUAL_TO));
229232
assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_2"));
230233
assertThat(resp.getHits().getAt(1).getId(), equalTo("doc_1"));
231-
assertThat(resp.getHits().getAt(2).getId(), equalTo("doc_4"));
232-
assertThat(resp.getHits().getAt(3).getId(), equalTo("doc_7"));
233-
assertThat(resp.getHits().getAt(4).getId(), equalTo("doc_3"));
234+
assertThat(resp.getHits().getAt(2).getId(), equalTo("doc_3"));
235+
assertThat(resp.getHits().getAt(3).getId(), equalTo("doc_4"));
236+
assertThat(resp.getHits().getAt(4).getId(), equalTo("doc_7"));
234237
});
235238
}
236239

@@ -243,8 +246,10 @@ public void testRankDocsRetrieverWithAggs() {
243246
SearchSourceBuilder source = new SearchSourceBuilder();
244247
StandardRetrieverBuilder standard0 = new StandardRetrieverBuilder();
245248
// this one retrieves docs 1, 4, and 6
246-
standard0.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.queryStringQuery("quick").defaultField(TEXT_FIELD))
247-
.boost(10L);
249+
standard0.queryBuilder = QueryBuilders.boolQuery()
250+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_1")).boost(10L))
251+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_4")).boost(9L))
252+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_6")).boost(8L));
248253
StandardRetrieverBuilder standard1 = new StandardRetrieverBuilder();
249254
// this one retrieves docs 2 and 6 due to prefilter
250255
standard1.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.termsQuery(ID_FIELD, "doc_2", "doc_3", "doc_6")).boost(20L);
@@ -268,13 +273,15 @@ public void testRankDocsRetrieverWithAggs() {
268273
)
269274
)
270275
);
276+
source.size(1);
271277
source.aggregation(new TermsAggregationBuilder("topic").field(TOPIC_FIELD));
272278
SearchRequestBuilder req = client().prepareSearch(INDEX).setSource(source);
273279
ElasticsearchAssertions.assertResponse(req, resp -> {
274280
assertNull(resp.pointInTimeId());
275281
assertNotNull(resp.getHits().getTotalHits());
276-
assertThat(resp.getHits().getTotalHits().value(), equalTo(1L));
282+
assertThat(resp.getHits().getTotalHits().value(), equalTo(5L));
277283
assertThat(resp.getHits().getTotalHits().relation(), equalTo(TotalHits.Relation.EQUAL_TO));
284+
assertThat(resp.getHits().getHits().length, equalTo(1));
278285
assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_2"));
279286
assertNotNull(resp.getAggregations());
280287
assertNotNull(resp.getAggregations().get("topic"));
@@ -292,8 +299,10 @@ public void testRankDocsRetrieverWithCollapse() {
292299
SearchSourceBuilder source = new SearchSourceBuilder();
293300
StandardRetrieverBuilder standard0 = new StandardRetrieverBuilder();
294301
// this one retrieves docs 1, 4, and 6
295-
standard0.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.queryStringQuery("quick").defaultField(TEXT_FIELD))
296-
.boost(10L);
302+
standard0.queryBuilder = QueryBuilders.boolQuery()
303+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_1")).boost(10L))
304+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_4")).boost(9L))
305+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_6")).boost(8L));
297306
StandardRetrieverBuilder standard1 = new StandardRetrieverBuilder();
298307
// this one retrieves docs 2 and 6 due to prefilter
299308
standard1.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.termsQuery(ID_FIELD, "doc_2", "doc_3", "doc_6")).boost(20L);
@@ -308,8 +317,8 @@ public void testRankDocsRetrieverWithCollapse() {
308317
null
309318
);
310319
// the compound retriever here produces a score for a doc based on the percentage of the queries that it was matched on and
311-
// resolves ties based on actual score, rank, and then the doc (we're forcing 1 shard for consistent results)
312-
// so ideal rank would be: 6, 2, 1, 4, 7, 3
320+
// resolves ties based on actual score, and then the doc (we're forcing 1 shard for consistent results)
321+
// so ideal rank would be: 6, 2, 1, 3, 4, 7
313322
// with collapsing on topic field we would have 6, 2, 1, 7
314323
source.retriever(
315324
new CompoundRetrieverWithRankDocs(
@@ -339,7 +348,6 @@ public void testRankDocsRetrieverWithCollapse() {
339348
assertThat(resp.getHits().getAt(1).field(TOPIC_FIELD).getValue().toString(), equalTo("astronomy"));
340349
assertThat(resp.getHits().getAt(2).getId(), equalTo("doc_1"));
341350
assertThat(resp.getHits().getAt(2).field(TOPIC_FIELD).getValue().toString(), equalTo("technology"));
342-
assertThat(resp.getHits().getAt(2).getInnerHits().get("a").getHits().length, equalTo(3));
343351
assertThat(resp.getHits().getAt(2).getInnerHits().get("a").getAt(0).getId(), equalTo("doc_4"));
344352
assertThat(resp.getHits().getAt(2).getInnerHits().get("a").getAt(1).getId(), equalTo("doc_3"));
345353
assertThat(resp.getHits().getAt(2).getInnerHits().get("a").getAt(2).getId(), equalTo("doc_1"));
@@ -348,17 +356,15 @@ public void testRankDocsRetrieverWithCollapse() {
348356
});
349357
}
350358

351-
public void testRankDocsRetrieverWithCollapseAndAggs() {
352-
// same as above, but we only want to bring back the top result from each subsearch
353-
// so that would be 1, 2, and 7
354-
// and final rank would be (based on score): 2, 1, 7
355-
// aggs should still account for the same docs as the testRankDocsRetriever test, i.e. all but doc_5
359+
public void testRankDocsRetrieverWithNestedCollapseAndAggs() {
356360
final int rankWindowSize = 10;
357361
SearchSourceBuilder source = new SearchSourceBuilder();
358362
StandardRetrieverBuilder standard0 = new StandardRetrieverBuilder();
359363
// this one retrieves docs 1 and 6 as doc_4 is collapsed to doc_1
360-
standard0.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.queryStringQuery("quick").defaultField(TEXT_FIELD))
361-
.boost(10L);
364+
standard0.queryBuilder = QueryBuilders.boolQuery()
365+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_1")).boost(10L))
366+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_4")).boost(9L))
367+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_6")).boost(8L));
362368
standard0.collapseBuilder = new CollapseBuilder(TOPIC_FIELD).setInnerHits(
363369
new InnerHitBuilder("a").addSort(new FieldSortBuilder(DOC_FIELD).order(SortOrder.DESC)).setSize(10)
364370
);
@@ -376,8 +382,8 @@ public void testRankDocsRetrieverWithCollapseAndAggs() {
376382
null
377383
);
378384
// the compound retriever here produces a score for a doc based on the percentage of the queries that it was matched on and
379-
// resolves ties based on actual score, rank, and then the doc (we're forcing 1 shard for consistent results)
380-
// so ideal rank would be: 6, 2, 1, 4, 7, 3
385+
// resolves ties based on actual score, and then the doc (we're forcing 1 shard for consistent results)
386+
// so ideal rank would be: 6, 2, 1, 3, 4, 7
381387
source.retriever(
382388
new CompoundRetrieverWithRankDocs(
383389
rankWindowSize,
@@ -393,7 +399,7 @@ public void testRankDocsRetrieverWithCollapseAndAggs() {
393399
ElasticsearchAssertions.assertResponse(req, resp -> {
394400
assertNull(resp.pointInTimeId());
395401
assertNotNull(resp.getHits().getTotalHits());
396-
assertThat(resp.getHits().getTotalHits().value(), equalTo(5L));
402+
assertThat(resp.getHits().getTotalHits().value(), equalTo(6L));
397403
assertThat(resp.getHits().getTotalHits().relation(), equalTo(TotalHits.Relation.EQUAL_TO));
398404
assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_6"));
399405
assertNotNull(resp.getAggregations());
@@ -427,8 +433,8 @@ public void testRankDocsRetrieverWithNestedQuery() {
427433
null
428434
);
429435
// the compound retriever here produces a score for a doc based on the percentage of the queries that it was matched on and
430-
// resolves ties based on actual score, rank, and then the doc (we're forcing 1 shard for consistent results)
431-
// so ideal rank would be: 6, 2, 1, 4, 3, 7
436+
// resolves ties based on actual score, and then the doc (we're forcing 1 shard for consistent results)
437+
// so ideal rank would be: 6, 2, 1, 3, 4, 7
432438
source.retriever(
433439
new CompoundRetrieverWithRankDocs(
434440
rankWindowSize,
@@ -460,8 +466,10 @@ public void testRankDocsRetrieverMultipleCompoundRetrievers() {
460466
SearchSourceBuilder source = new SearchSourceBuilder();
461467
StandardRetrieverBuilder standard0 = new StandardRetrieverBuilder();
462468
// this one retrieves docs 1, 4, and 6
463-
standard0.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.queryStringQuery("quick").defaultField(TEXT_FIELD))
464-
.boost(10L);
469+
standard0.queryBuilder = QueryBuilders.boolQuery()
470+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_1")).boost(10L))
471+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_4")).boost(9L))
472+
.should(QueryBuilders.constantScoreQuery(QueryBuilders.idsQuery().addIds("doc_6")).boost(8L));
465473
StandardRetrieverBuilder standard1 = new StandardRetrieverBuilder();
466474
// this one retrieves docs 2 and 6 due to prefilter
467475
standard1.queryBuilder = QueryBuilders.constantScoreQuery(QueryBuilders.termsQuery(ID_FIELD, "doc_2", "doc_3", "doc_6")).boost(20L);
@@ -506,11 +514,11 @@ public void testRankDocsRetrieverMultipleCompoundRetrievers() {
506514
assertThat(resp.getHits().getTotalHits().value(), equalTo(6L));
507515
assertThat(resp.getHits().getTotalHits().relation(), equalTo(TotalHits.Relation.EQUAL_TO));
508516
assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_4"));
509-
assertThat(resp.getHits().getAt(1).getId(), equalTo("doc_6"));
517+
assertThat(resp.getHits().getAt(1).getId(), equalTo("doc_1"));
510518
assertThat(resp.getHits().getAt(2).getId(), equalTo("doc_2"));
511-
assertThat(resp.getHits().getAt(3).getId(), equalTo("doc_1"));
512-
assertThat(resp.getHits().getAt(4).getId(), equalTo("doc_7"));
513-
assertThat(resp.getHits().getAt(5).getId(), equalTo("doc_3"));
519+
assertThat(resp.getHits().getAt(3).getId(), equalTo("doc_3"));
520+
assertThat(resp.getHits().getAt(4).getId(), equalTo("doc_6"));
521+
assertThat(resp.getHits().getAt(5).getId(), equalTo("doc_7"));
514522
});
515523
}
516524

@@ -545,9 +553,9 @@ public void testRankDocsRetrieverDifferentNestedSorting() {
545553
assertThat(resp.getHits().getTotalHits().relation(), equalTo(TotalHits.Relation.EQUAL_TO));
546554
assertThat(resp.getHits().getAt(0).getId(), equalTo("doc_4"));
547555
assertThat(resp.getHits().getAt(1).getId(), equalTo("doc_1"));
548-
assertThat(resp.getHits().getAt(2).getId(), equalTo("doc_7"));
556+
assertThat(resp.getHits().getAt(2).getId(), equalTo("doc_2"));
549557
assertThat(resp.getHits().getAt(3).getId(), equalTo("doc_6"));
550-
assertThat(resp.getHits().getAt(4).getId(), equalTo("doc_2"));
558+
assertThat(resp.getHits().getAt(4).getId(), equalTo("doc_7"));
551559
});
552560
}
553561

@@ -673,22 +681,14 @@ private RankDoc[] getRankDocs(SearchResponse searchResponse) {
673681
for (int i = 0; i < size; i++) {
674682
var hit = searchResponse.getHits().getAt(i);
675683
long sortValue = (long) hit.getRawSortValues()[hit.getRawSortValues().length - 1];
676-
int doc = decodeDoc(sortValue);
677-
int shardRequestIndex = decodeShardRequestIndex(sortValue);
684+
int doc = ShardDocSortField.decodeDoc(sortValue);
685+
int shardRequestIndex = ShardDocSortField.decodeShardRequestIndex(sortValue);
678686
docs[i] = new RankDoc(doc, hit.getScore(), shardRequestIndex);
679687
docs[i].rank = i + 1;
680688
}
681689
return docs;
682690
}
683691

684-
public static int decodeDoc(long value) {
685-
return (int) value;
686-
}
687-
688-
public static int decodeShardRequestIndex(long value) {
689-
return (int) (value >> 32);
690-
}
691-
692692
record RankDocAndHitRatio(RankDoc rankDoc, float hitRatio) {}
693693

694694
/**

0 commit comments

Comments
 (0)