Moving getJudgments to abstract class.

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

@@ -35,18 +35,16 @@
 /**
  * A {@link QuerySetRunner} for Amazon OpenSearch.
  */
-public class OpenSearchQuerySetRunner implements QuerySetRunner {
+public class OpenSearchQuerySetRunner extends QuerySetRunner {
 
     private static final Logger LOGGER = LogManager.getLogger(OpenSearchQuerySetRunner.class);
 
-    final Client client;
-
     /**
      * Creates a new query set runner
      * @param client An OpenSearch {@link Client}.
      */
     public OpenSearchQuerySetRunner(final Client client) {
-        this.client = client;
+        super(client);
     }
 
     @Override
@@ -172,48 +170,4 @@ public void onFailure(Exception ex) {
 
     }
 
-    private List<Judgment> getJudgments(final String judgmentsId) throws Exception {
-
-        final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
-        searchSourceBuilder.query(QueryBuilders.matchQuery("_id", judgmentsId));
-        searchSourceBuilder.trackTotalHits(true);
-
-        final SearchRequest getQuerySetSearchRequest = new SearchRequest(SearchQualityEvaluationPlugin.JUDGMENTS_INDEX_NAME);
-        getQuerySetSearchRequest.source(searchSourceBuilder);
-
-        // TODO: Don't use .get()
-        final SearchResponse searchResponse = client.search(getQuerySetSearchRequest).get();
-
-        final List<Judgment> judgments = new ArrayList<>();
-
-        if(searchResponse.getHits().getTotalHits().value == 0) {
-
-            // The judgment_id is probably not valid.
-            // This will return an empty list.
-
-        } else {
-
-            // TODO: Make sure the search gets something back.
-            final Collection<Map<String, Object>> j = (Collection<Map<String, Object>>) searchResponse.getHits().getAt(0).getSourceAsMap().get("judgments");
-
-            for (final Map<String, Object> judgment : j) {
-
-                final String queryId = judgment.get("query_id").toString();
-                final double judgmentValue = Double.parseDouble(judgment.get("judgment").toString());
-                final String query = judgment.get("query").toString();
-                final String document = judgment.get("document").toString();
-
-                final Judgment jobj = new Judgment(queryId, query, document, judgmentValue);
-                LOGGER.info("Judgment: {}", jobj.toJudgmentString());
-
-                judgments.add(jobj);
-
-            }
-
-        }
-
-        return judgments;
-
-    }
-
 }

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

@@ -8,11 +8,34 @@
  */
 package org.opensearch.eval.runners;
 
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.client.Client;
+import org.opensearch.eval.SearchQualityEvaluationPlugin;
+import org.opensearch.eval.judgments.model.Judgment;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.search.builder.SearchSourceBuilder;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+
 /**
- * Interface for query set runners. Classes that implement this interface
+ * Base class for query set runners. Classes that extend this class
  * should be specific to a search engine. See the {@link OpenSearchQuerySetRunner} for an example.
  */
-public interface QuerySetRunner {
+public abstract class QuerySetRunner {
+
+    private static final Logger LOGGER = LogManager.getLogger(QuerySetRunner.class);
+
+    protected final Client client;
+
+    public QuerySetRunner(final Client client) {
+        this.client = client;
+    }
 
     /**
      * Runs the query set.
@@ -24,12 +47,57 @@ public interface QuerySetRunner {
      * @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, final int k) throws Exception;
+    abstract QuerySetRunResult run(String querySetId, final String judgmentsId, final String index, final String idField, final String query, final int k) throws Exception;
 
     /**
      * Saves the query set results to a persistent store, which may be the search engine itself.
      * @param result The {@link QuerySetRunResult results}.
      */
-    void save(QuerySetRunResult result) throws Exception;
+    abstract void save(QuerySetRunResult result) throws Exception;
+
+    public List<Judgment> getJudgments(final String judgmentsId) throws Exception {
+
+        final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
+        searchSourceBuilder.query(QueryBuilders.matchQuery("_id", judgmentsId));
+        searchSourceBuilder.trackTotalHits(true);
+
+        final SearchRequest getQuerySetSearchRequest = new SearchRequest(SearchQualityEvaluationPlugin.JUDGMENTS_INDEX_NAME);
+        getQuerySetSearchRequest.source(searchSourceBuilder);
+
+        // TODO: Don't use .get()
+        final SearchResponse searchResponse = client.search(getQuerySetSearchRequest).get();
+
+        final List<Judgment> judgments = new ArrayList<>();
+
+        if(searchResponse.getHits().getTotalHits().value == 0) {
+
+            // The judgment_id is probably not valid.
+            // This will return an empty list.
+
+        } else {
+
+            // TODO: Make sure the search gets something back.
+            final Collection<Map<String, Object>> j = (Collection<Map<String, Object>>) searchResponse.getHits().getAt(0).getSourceAsMap().get("judgments");
+
+            for (final Map<String, Object> judgment : j) {
+
+                final String queryId = judgment.get("query_id").toString();
+                final double judgmentValue = Double.parseDouble(judgment.get("judgment").toString());
+                final String query = judgment.get("query").toString();
+                final String document = judgment.get("document").toString();
+
+                final Judgment jobj = new Judgment(queryId, query, document, judgmentValue);
+                LOGGER.info("Judgment: {}", jobj.toJudgmentString());
+
+                judgments.add(jobj);
+
+            }
+
+        }
+
+        return judgments;
+
+    }
+
 
 }