From 1e2b047a0e20beb0e118b6f09d3d973d647f699a Mon Sep 17 00:00:00 2001
From: Kathleen DeRusso <kathleen.derusso@elastic.co>
Date: Wed, 30 Oct 2024 11:36:00 -0400
Subject: [PATCH] Add all queries to history list

---
 .../rankeval/TransportRankEvalAction.java     | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java b/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java
index adba5e5e0b79..d828fc70fc37 100644
--- a/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java
+++ b/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java
@@ -48,6 +48,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -224,15 +225,16 @@ public void onResponse(GetResponse response) {
                 // Now that we have the corpus, identify the queries and qrels
                 SearchRequest queryRequest = new SearchRequestBuilder(client).setIndices("search_relevance_evaluation_corpora")
                     .setSource(new SearchSourceBuilder().query(new TermQueryBuilder("dataset_id", corpus.id())))
-                    .setSize(250) // Arbitrary number, but dbpedia has at most 164 judgments per query
+                    .setSize(10000) // Ideally this cap should be the configured max result window, hard coding for now
                     .request();
                 client.search(queryRequest, new ActionListener<>() {
                     @Override
                     public void onResponse(SearchResponse response) {
                         List<RatedRequest> ratedRequests = new ArrayList<>();
-                        List<RatedDocument> ratedDocuments = new ArrayList<>();
+                        Map<String, List<RatedDocument>> queryIdToRatedDocumentMap = new HashMap<>();
+                        Map<String, String> queryIdToQueryStringMap = new HashMap<>();
 
-                        for (SearchHit hit : response.getHits().getHits()) {
+                        Arrays.stream(response.getHits().getHits()).forEach(hit -> {
                             Map<String, Object> source = hit.getSourceAsMap();
                             String queryId = (String) source.get("query_id");
                             String queryString = (String) source.get("query_string");
@@ -241,10 +243,21 @@ public void onResponse(SearchResponse response) {
                             Integer score = (Integer) source.get("score");
 
                             RatedDocument ratedDocument = new RatedDocument(index, documentId, score);
-                            ratedDocuments.add(ratedDocument);
+                            queryIdToRatedDocumentMap.computeIfAbsent(queryId, k -> new ArrayList<>()).add(ratedDocument);
+                            queryIdToQueryStringMap.put(queryId, queryString);
+                        });
+
+                        queryIdToRatedDocumentMap.forEach((queryId, ratedDocuments) -> {
+                            ratedRequests.add(
+                                new RatedRequest(
+                                    queryId,
+                                    ratedDocuments,
+                                    Map.of("query_string", queryIdToQueryStringMap.get(queryId)),
+                                    templateId
+                                )
+                            );
+                        });
 
-                            ratedRequests.add(new RatedRequest(queryId, ratedDocuments, Map.of("query_string", queryString), templateId));
-                        }
                         actionListener.onResponse(ratedRequests);
                     }