diff --git a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/OpenSearchQuerySetRunner.java b/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/OpenSearchQuerySetRunner.java index f6abc38..0a77b26 100644 --- a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/OpenSearchQuerySetRunner.java +++ b/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 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 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> j = (Collection>) searchResponse.getHits().getAt(0).getSourceAsMap().get("judgments"); - - for (final Map 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; - - } - } diff --git a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/QuerySetRunner.java b/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/QuerySetRunner.java index 6851172..a03ddbe 100644 --- a/opensearch-search-quality-evaluation-plugin/src/main/java/org/opensearch/eval/runners/QuerySetRunner.java +++ b/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 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 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> j = (Collection>) searchResponse.getHits().getAt(0).getSourceAsMap().get("judgments"); + + for (final Map 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; + + } + }