From 58f8cd4810992523d277016374ce8f542b78d109 Mon Sep 17 00:00:00 2001
From: Kathleen DeRusso <kathleen.derusso@elastic.co>
Date: Thu, 31 Oct 2024 13:58:20 -0400
Subject: [PATCH] Add query vector as a param

---
 .../rankeval/TransportRankEvalAction.java     | 22 ++++++++++++-------
 1 file changed, 14 insertions(+), 8 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 acae8f95dea4..cc4abf41e957 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
@@ -234,11 +234,14 @@ public void onResponse(SearchResponse response) {
                         List<RatedRequest> ratedRequests = new ArrayList<>();
                         Map<String, List<RatedDocument>> queryIdToRatedDocumentMap = new HashMap<>();
                         Map<String, String> queryIdToQueryStringMap = new HashMap<>();
+                        Map<String, Map<String, Object>> queryIdToQueryVectorsMap = new HashMap<>();
 
                         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");
+                            @SuppressWarnings("unchecked")
+                            Map<String, Object> queryVector = (Map<String, Object>) source.get("query_vector");
                             String index = (String) source.get("index");
                             String documentId = (String) source.get("document_id");
                             Integer score = (Integer) source.get("score");
@@ -246,17 +249,20 @@ public void onResponse(SearchResponse response) {
                             RatedDocument ratedDocument = new RatedDocument(index, documentId, score);
                             queryIdToRatedDocumentMap.computeIfAbsent(queryId, k -> new ArrayList<>()).add(ratedDocument);
                             queryIdToQueryStringMap.put(queryId, queryString);
+                            queryIdToQueryVectorsMap.put(queryId, queryVector);
                         });
 
                         queryIdToRatedDocumentMap.forEach((queryId, ratedDocuments) -> {
-                            ratedRequests.add(
-                                new RatedRequest(
-                                    queryId,
-                                    ratedDocuments,
-                                    Map.of("query_string", queryIdToQueryStringMap.get(queryId)),
-                                    templateId
-                                )
-                            );
+
+                            Map<String, Object> params = new HashMap<>();
+                            if (queryIdToQueryStringMap.containsKey(queryId)) {
+                                params.put("query_string", queryIdToQueryStringMap.get(queryId));
+                            }
+                            if (queryIdToQueryVectorsMap.containsKey(queryId)) {
+                                params.put("query_vector", queryIdToQueryVectorsMap.get(queryId));
+                            }
+
+                            ratedRequests.add(new RatedRequest(queryId, ratedDocuments, params, templateId));
                         });
 
                         actionListener.onResponse(ratedRequests);