Skip to content

Commit

Permalink
Working on converting to a standalone app.
Browse files Browse the repository at this point in the history
  • Loading branch information
jzonthemtn committed Dec 31, 2024
1 parent aa26e3b commit 6c4b241
Show file tree
Hide file tree
Showing 4 changed files with 69 additions and 102 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,16 @@
import org.apache.logging.log4j.Logger;
import org.opensearch.client.json.jackson.JacksonJsonpMapper;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch._types.FieldValue;
import org.opensearch.client.opensearch._types.Refresh;
import org.opensearch.client.opensearch._types.Time;
import org.opensearch.client.opensearch._types.mapping.IntegerNumberProperty;
import org.opensearch.client.opensearch._types.mapping.Property;
import org.opensearch.client.opensearch._types.mapping.TypeMapping;
import org.opensearch.client.opensearch._types.query_dsl.BoolQuery;
import org.opensearch.client.opensearch._types.query_dsl.MatchQuery;
import org.opensearch.client.opensearch._types.query_dsl.Query;
import org.opensearch.client.opensearch._types.query_dsl.TermQuery;
import org.opensearch.client.opensearch.core.BulkRequest;
import org.opensearch.client.opensearch.core.BulkResponse;
import org.opensearch.client.opensearch.core.IndexRequest;
Expand All @@ -43,6 +48,7 @@

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
Expand Down Expand Up @@ -120,6 +126,47 @@ public String indexQuerySet(final QuerySet querySet) throws IOException {

}

@Override
public QuerySet getQuerySet(String querySetId) throws IOException {

final Query query = Query.of(q -> q.term(m -> m.field("_id").value(FieldValue.of(querySetId))));

final SearchResponse<QuerySet> searchResponse = client.search(s -> s.index(Constants.QUERY_SETS_INDEX_NAME).query(query).size(1), QuerySet.class);

// TODO: Handle the query set not being found.

return searchResponse.hits().hits().get(0).source();

}

@Override
public Double getJudgmentValue(final String judgmentsId, final String userQuery, final String documentId) throws Exception {

var boolQuery = BoolQuery.of(bq -> bq
.must(
List.of(
MatchQuery.of(mq -> mq.query(FieldValue.of("judgments_id")).field(judgmentsId)).toQuery(),
MatchQuery.of(mq -> mq.query(FieldValue.of("query")).field(userQuery)).toQuery(),
MatchQuery.of(mq -> mq.query(FieldValue.of("document_id")).field(documentId)).toQuery()
)
)
);

final Query query = Query.of(q -> q.bool(boolQuery));

final SearchResponse<Judgment> searchResponse = client.search(s -> s.index(Constants.JUDGMENTS_INDEX_NAME)
.query(query)
.from(0)
.size(1), Judgment.class);

if(searchResponse.hits().hits().isEmpty()) {
return Double.NaN;
} else {
return searchResponse.hits().hits().get(0).source().getJudgment();
}

}

@Override
public Collection<UbiQuery> getUbiQueries() throws IOException {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -30,4 +30,21 @@ public abstract class SearchEngine {
public abstract String indexQuerySet(QuerySet querySet) throws IOException;
public abstract Collection<UbiQuery> getUbiQueries() throws IOException;

/**
* Gets a query set from the index.
* @param querySetId The ID of the query set to get.
* @return The query set as a collection of maps of query to frequency
* @throws IOException Thrown if the query set cannot be retrieved.
*/
public abstract QuerySet getQuerySet(String querySetId) throws IOException;

/**
* Get a judgment from the index.
* @param judgmentsId The ID of the judgments to find.
* @param query The user query.
* @param documentId The document ID.
* @return The value of the judgment, or <code>NaN</code> if the judgment cannot be found.
*/
public abstract Double getJudgmentValue(final String judgmentsId, final String query, final String documentId) throws Exception;

}
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,10 @@

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.opensearch.client.opensearch.core.SearchRequest;
import org.opensearch.client.opensearch.core.SearchResponse;
import org.opensearch.eval.Constants;
import org.opensearch.eval.engine.SearchEngine;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

/**
* Base class for query set runners. Classes that extend this class
Expand Down Expand Up @@ -58,99 +53,6 @@ abstract QuerySetRunResult run(String querySetId, final String judgmentsId, fina
*/
abstract void save(QuerySetRunResult result) throws Exception;

/**
* Gets a query set from the index.
* @param querySetId The ID of the query set to get.
* @return The query set as a collection of maps of query to frequency
* @throws Exception Thrown if the query set cannot be retrieved.
*/
public final Collection<Map<String, Long>> getQuerySet(final String querySetId) throws Exception {

// Get the query set.
final SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
sourceBuilder.query(QueryBuilders.matchQuery("_id", querySetId));

// Will be at most one match.
sourceBuilder.from(0);
sourceBuilder.size(1);

final SearchRequest searchRequest = new SearchRequest(Constants.QUERY_SETS_INDEX_NAME).source(sourceBuilder);

// TODO: Don't use .get()
final SearchResponse searchResponse = client.search(searchRequest).get();

if(searchResponse.getHits().getHits().length > 0) {

// The queries from the query set that will be run.
return (Collection<Map<String, Long>>) searchResponse.getHits().getAt(0).getSourceAsMap().get("queries");

} else {

LOGGER.error("Unable to get query set with ID {}", querySetId);

// The query set was not found.
throw new RuntimeException("The query set with ID " + querySetId + " was not found.");

}

}

/**
* Get a judgment from the index.
* @param judgmentsId The ID of the judgments to find.
* @param query The user query.
* @param documentId The document ID.
* @return The value of the judgment, or <code>NaN</code> if the judgment cannot be found.
*/
public Double getJudgmentValue(final String judgmentsId, final String query, final String documentId) throws Exception {

// Find a judgment that matches the judgments_id, query_id, and document_id fields in the index.

final BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
boolQueryBuilder.must(QueryBuilders.termQuery("judgments_id", judgmentsId));
boolQueryBuilder.must(QueryBuilders.termQuery("query", query));
boolQueryBuilder.must(QueryBuilders.termQuery("document_id", documentId));

final SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
searchSourceBuilder.query(boolQueryBuilder);

// Will be a max of 1 result since we are getting the judgments by ID.
searchSourceBuilder.from(0);
searchSourceBuilder.size(1);

// Only include the judgment field in the response.
final String[] includeFields = new String[] {"judgment"};
final String[] excludeFields = new String[] {};
searchSourceBuilder.fetchSource(includeFields, excludeFields);

final SearchRequest searchRequest = new SearchRequest(Constants.JUDGMENTS_INDEX_NAME).source(searchSourceBuilder);

Double judgment = Double.NaN;

final SearchResponse searchResponse = client.search(searchRequest).get();

if (searchResponse.getHits().getHits().length > 0) {

final Map<String, Object> j = searchResponse.getHits().getAt(0).getSourceAsMap();

// LOGGER.debug("Judgment contains a value: {}", j.get("judgment"));

// TODO: Why does this not exist in some cases?
if(j.containsKey("judgment")) {
judgment = (Double) j.get("judgment");
}

} else {

// No judgment for this query/doc pair exists.
judgment = Double.NaN;

}

return judgment;

}

/**
* Gets the judgments for a query / document pairs.
* @param judgmentsId The judgments collection for which the judgment to retrieve belongs.
Expand All @@ -174,7 +76,7 @@ protected RelevanceScores getRelevanceScores(final String judgmentsId, final Str
final String documentId = orderedDocumentIds.get(i);

// Find the judgment value for this combination of query and documentId from the index.
final Double judgmentValue = getJudgmentValue(judgmentsId, query, documentId);
final Double judgmentValue = searchEngine.getJudgmentValue(judgmentsId, query, documentId);

// If a judgment for this query/doc pair is not found, Double.NaN will be returned.
if(!Double.isNaN(judgmentValue)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import org.opensearch.eval.metrics.NdcgSearchMetric;
import org.opensearch.eval.metrics.PrecisionSearchMetric;
import org.opensearch.eval.metrics.SearchMetric;
import org.opensearch.eval.model.data.QuerySet;
import org.opensearch.eval.utils.TimeUtils;

import java.util.ArrayList;
Expand Down Expand Up @@ -48,15 +49,15 @@ public QuerySetRunResult run(final String querySetId, final String judgmentsId,
final String searchPipeline, final String idField, final String query,
final int k, final double threshold) throws Exception {

final Collection<Map<String, Long>> querySet = getQuerySet(querySetId);
LOGGER.info("Found {} queries in query set {}", querySet.size(), querySetId);
final QuerySet querySet = searchEngine.getQuerySet(querySetId);
LOGGER.info("Found {} queries in query set {}", querySet.getQuerySetQueries().size(), querySetId);

try {

// The results of each query.
final List<QueryResult> queryResults = new ArrayList<>();

for (Map<String, Long> queryMap : querySet) {
for (Map<String, Long> queryMap : querySet.getQuerySetQueries()) {

// Loop over each query in the map and run each one.
for (final String userQuery : queryMap.keySet()) {
Expand Down

0 comments on commit 6c4b241

Please sign in to comment.