Allowing for metrics @ k.

opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/SearchQualityEvaluationRestHandler.java

@@ -158,11 +158,16 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             final String judgmentsId = request.param("judgments_id");
             final String index = request.param("index");
             final String idField = request.param("id_field", "_id");
+            final int k = Integer.parseInt(request.param("k", "10"));
 
             if(querySetId == null || judgmentsId == null || index == null) {
                 return restChannel -> restChannel.sendResponse(new BytesRestResponse(RestStatus.BAD_REQUEST, "{\"error\": \"Missing required parameters.\"}"));
             }
 
+            if(k < 1) {
+                return restChannel -> restChannel.sendResponse(new BytesRestResponse(RestStatus.BAD_REQUEST, "{\"error\": \"k must be a positive integer.\"}"));
+            }
+
             if(!request.hasContent()) {
                 return restChannel -> restChannel.sendResponse(new BytesRestResponse(RestStatus.BAD_REQUEST, "{\"error\": \"Missing query in body.\"}"));
             }
@@ -178,7 +183,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             try {
 
                 final OpenSearchQuerySetRunner openSearchQuerySetRunner = new OpenSearchQuerySetRunner(client);
-                final QuerySetRunResult querySetRunResult = openSearchQuerySetRunner.run(querySetId, judgmentsId, index, idField, query);
+                final QuerySetRunResult querySetRunResult = openSearchQuerySetRunner.run(querySetId, judgmentsId, index, idField, query, k);
                 openSearchQuerySetRunner.save(querySetRunResult);
 
             } catch (Exception ex) {

opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/OpenSearchQuerySetRunner.java

@@ -47,7 +47,7 @@ public OpenSearchQuerySetRunner(final Client client) {
     }
 
     @Override
-    public QuerySetRunResult run(final String querySetId, final String judgmentsId, final String index, final String idField, final String query) {
+    public QuerySetRunResult run(final String querySetId, final String judgmentsId, final String index, final String idField, final String query, final int k) {
 
         // Get the query set.
         final SearchSourceBuilder getQuerySetSearchSourceBuilder = new SearchSourceBuilder();
@@ -105,7 +105,7 @@ public void onResponse(final SearchResponse searchResponse) {
 
                             }
 
-                            queryResults.add(new QueryResult(query, orderedDocumentIds));
+                            queryResults.add(new QueryResult(query, orderedDocumentIds, k));
 
                         }
 
@@ -120,7 +120,7 @@ public void onFailure(Exception ex) {
             }
 
             // TODO: Calculate the search metrics given the results and the judgments.
-            final SearchMetrics searchMetrics = new SearchMetrics();
+            final SearchMetrics searchMetrics = new SearchMetrics(k);
 
             return new QuerySetRunResult(queryResults, searchMetrics);
 

opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/QueryResult.java

@@ -19,17 +19,19 @@ public class QueryResult {
     private final List<String> orderedDocumentIds;
 
     // TODO: Calculate these metrics.
-    private final SearchMetrics searchMetrics = new SearchMetrics();
+    private final SearchMetrics searchMetrics;
 
     /**
      * Creates the search results.
      * @param query The query used to generate this result.
      * @param orderedDocumentIds A list of ordered document IDs in the same order as they appeared
      *                           in the query.
+     * @param k The k used for metrics calculation, i.e. DCG@k.
      */
-    public QueryResult(final String query, final List<String> orderedDocumentIds) {
+    public QueryResult(final String query, final List<String> orderedDocumentIds, final int k) {
         this.query = query;
         this.orderedDocumentIds = orderedDocumentIds;
+        this.searchMetrics = new SearchMetrics(k);
     }
 
     /**

opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/QuerySetRunner.java

@@ -21,9 +21,10 @@ public interface QuerySetRunner {
      * @param index The name of the index to run the query sets against.
      * @param idField The field in the index that is used to uniquely identify a document.
      * @param query The query that will be used to run the query set.
+     * @param k The k used for metrics calculation, i.e. DCG@k.
      * @return The query set {@link QuerySetRunResult results} and calculated metrics.
      */
-    QuerySetRunResult run(String querySetId, final String judgmentsId, final String index, final String idField, final String query);
+    QuerySetRunResult run(String querySetId, final String judgmentsId, final String index, final String idField, final String query, final int k);
 
     /**
      * Saves the query set results to a persistent store, which may be the search engine itself.

opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/SearchMetrics.java

@@ -16,9 +16,15 @@
  */
 public class SearchMetrics {
 
-    private double dcg_at_10 = 0.0;
-    private double ndcg_at_10 = 0.0;
-    private double prec_at_10 = 0.0;
+    private final int k;
+
+    private double dcg = 0.0;
+    private double ndcg = 0.0;
+    private double precision = 0.0;
+
+    public SearchMetrics(final int k) {
+        this.k = k;
+    }
 
     /**
      * Gets the metrics as a map for ease of indexing.
@@ -27,36 +33,40 @@ public class SearchMetrics {
     public Map<String, Double> getSearchMetricsAsMap() {
 
         final Map<String, Double> metrics = new HashMap<>();
-        metrics.put("dcg_at_10", dcg_at_10);
-        metrics.put("ndcg_at_10", ndcg_at_10);
-        metrics.put("prec_at_10", prec_at_10);
+        metrics.put("dcg_at_" + k, dcg);
+        metrics.put("ndcg_at_" + k, ndcg);
+        metrics.put("prec_at_" + k, precision);
 
         return metrics;
 
     }
 
-    public double getDcg_at_10() {
-        return dcg_at_10;
+    public int getK() {
+        return k;
+    }
+
+    public double getDcg() {
+        return dcg;
     }
 
-    public void setDcg_at_10(double dcg_at_10) {
-        this.dcg_at_10 = dcg_at_10;
+    public void setDcg(double dcg) {
+        this.dcg = dcg;
     }
 
-    public double getNdcg_at_10() {
-        return ndcg_at_10;
+    public double getNdcg() {
+        return ndcg;
     }
 
-    public void setNdcg_at_10(double ndcg_at_10) {
-        this.ndcg_at_10 = ndcg_at_10;
+    public void setNdcg(double ndcg) {
+        this.ndcg = ndcg;
     }
 
-    public double getPrec_at_10() {
-        return prec_at_10;
+    public double getPrecision() {
+        return precision;
     }
 
-    public void setPrec_at_10(double prec_at_10) {
-        this.prec_at_10 = prec_at_10;
+    public void setPrecision(double precision) {
+        this.precision = precision;
     }
 
 }