ElasticDocumentSearch.java

/**
 * Copyright (c) 2016 Lemur Consulting Ltd.
 * <p/>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p/>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p/>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package uk.co.flax.biosolr.ontology.search.elasticsearch;

import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.elasticsearch.action.search.SearchRequestBuilder;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.Client;
import org.elasticsearch.index.query.MultiMatchQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptService;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHitField;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.bucket.terms.StringTerms;
import org.elasticsearch.search.aggregations.bucket.terms.Terms;
import org.elasticsearch.search.aggregations.metrics.cardinality.Cardinality;
import org.elasticsearch.search.aggregations.metrics.tophits.TopHits;
import org.elasticsearch.search.aggregations.metrics.tophits.TopHitsBuilder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import uk.co.flax.biosolr.ontology.api.Document;
import uk.co.flax.biosolr.ontology.config.ElasticSearchConfiguration;
import uk.co.flax.biosolr.ontology.search.DocumentSearch;
import uk.co.flax.biosolr.ontology.search.ResultsList;
import uk.co.flax.biosolr.ontology.search.SearchEngineException;

import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * Created by mlp on 17/02/16.
 *
 * @author mlp
 */
public class ElasticDocumentSearch extends ElasticSearchEngine implements DocumentSearch {

	private static final Logger LOGGER = LoggerFactory.getLogger(ElasticDocumentSearch.class);

	private static final String GROUP_FIELD = "study_id";
	private static final String COUNT_AGGREGATION = "numFound";
	private static final String HITS_AGGREGATION = "study";
	private static final String SCORE_AGGREGATION = "topScore";

	private static final String[] DEFAULT_FIELDS = new String[]{
			"title", "first_author", "publication", "efo_uri.label"
	};
	private static final List<String> ANNOTATED_FIELDS = new ArrayList<>();
	static {
		ANNOTATED_FIELDS.add("label");
		ANNOTATED_FIELDS.add("child_labels");
		ANNOTATED_FIELDS.add("parent_labels");
	}

	public ElasticDocumentSearch(Client client, ElasticSearchConfiguration config) {
		super(client, config);
	}

	@Override
	public ResultsList<Document> searchDocuments(String term, int start, int rows,
			List<String> additionalFields, List<String> filters) throws SearchEngineException {
		// Build the query
		MultiMatchQueryBuilder qb = QueryBuilders.multiMatchQuery(term, DEFAULT_FIELDS)
				.minimumShouldMatch("2<25%");
		if (additionalFields != null && additionalFields.size() > 0) {
			List<String> parsedAdditional = parseAdditionalFields(additionalFields);
			parsedAdditional.forEach(qb::field);
		}

		TopHitsBuilder topHitsBuilder = AggregationBuilders.topHits(HITS_AGGREGATION)
				.setFrom(0)
				.setSize(1);

		/* Build the terms aggregation, since we need a result set grouped by study ID.
		 * The "top_score" sub-agg allows us to sort by the top score of the results;
		 * the topHits sub-agg actually pulls back the record data, returning just the first
		 * hit in the aggregation.
		 * Note that we have to get _all_ rows up to and including the last required, annoyingly. */
		AggregationBuilder termsAgg = AggregationBuilders.terms(HITS_AGGREGATION)
				.field(GROUP_FIELD)
				.order(Terms.Order.aggregation(SCORE_AGGREGATION, false))
				.size(start + rows)
				.subAggregation(
						AggregationBuilders.max(SCORE_AGGREGATION)
								.script(new Script("_score", ScriptService.ScriptType.INLINE, "expression", null)))
				.subAggregation(topHitsBuilder);

		// Build the actual search request, including another aggregation to get
		// the number of unique study IDs returned.
		SearchRequestBuilder srb = getClient().prepareSearch(getIndexName())
				.setTypes(getDocumentType())
				.setQuery(qb)
				.setSize(0)
				.addAggregation(termsAgg)
				.addAggregation(AggregationBuilders.cardinality(COUNT_AGGREGATION).field(GROUP_FIELD));
		LOGGER.debug("ES Query: {}", srb.toString());

		SearchResponse response = srb.execute().actionGet();

		// Handle the response
		long total = ((Cardinality)(response.getAggregations().get(COUNT_AGGREGATION))).getValue();
		List<Document> docs;
		if (total == 0) {
			docs = new ArrayList<>();
		} else {
			// Build a map - need to look up annotation data separately.
			// This is because it's not in _source, and the fields() method
			// is not visible for a TopHitsBuilder.
			Map<String, Document> documentMap = new LinkedHashMap<>(rows);
			ObjectMapper mapper = buildObjectMapper();

			int lastIdx = (int)(start + rows <= total ? start + rows : total);
			StringTerms terms = response.getAggregations().get(HITS_AGGREGATION);
			List<Terms.Bucket> termBuckets = terms.getBuckets().subList(start, lastIdx);
			for (Terms.Bucket bucket : termBuckets) {
				TopHits hits = bucket.getAggregations().get(HITS_AGGREGATION);
				SearchHit hit = hits.getHits().getAt(0);
				documentMap.put(hit.getId(), extractDocument(mapper, hit));
			}

			// Populate annotation data for the document
			lookupAnnotationFields(documentMap);

			docs = new ArrayList<>(documentMap.values());
		}

		return new ResultsList<>(docs, start, (start / rows), total);
	}

	@Override
	public ResultsList<Document> searchByEfoUri(int start, int rows, String term, String... uris) throws SearchEngineException {
		return null;
	}

	private List<String> parseAdditionalFields(List<String> additional) {
		List<String> parsed;

		if (additional == null || additional.size() == 0) {
			parsed = null;
		} else {
			// Need to add annotation field name to all additional fields
			// Also need to handle hard-coded Solr field names
			parsed = additional.stream()
					.map(add -> add.replaceAll("^efo_uri_(.*)_t$", "$1"))
					.map(add -> getAnnotationField() + "." + add)
					.collect(Collectors.toList());
		}

		return parsed;
	}

	private ObjectMapper buildObjectMapper() {
		ObjectMapper mapper = new ObjectMapper();
		mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
		return mapper;
	}

	private Document extractDocument(ObjectMapper mapper, SearchHit hit) throws SearchEngineException {
		Document doc;

		try {
			doc = mapper.readValue(hit.getSourceAsString(), Document.class);
		} catch (IOException e) {
			LOGGER.error("Error reading document from source: {}", e.getMessage());
			throw new SearchEngineException(e);
		}

		return doc;
	}

	private void lookupAnnotationFields(Map<String, Document> idMap) {
		QueryBuilder qb = QueryBuilders.idsQuery(getDocumentType()).addIds(idMap.keySet());
		SearchRequestBuilder srb = getClient().prepareSearch(getIndexName())
				.addFields("*")
				.setQuery(qb)
				.setSize(idMap.size());
		LOGGER.debug("Annotation field lookup query: {}", srb.toString());

		SearchResponse response = srb.execute().actionGet();
		for (SearchHit hit : response.getHits().getHits()) {
			populateAnnotationFields(hit, idMap.get(hit.getId()));
		}
	}

	private void populateAnnotationFields(SearchHit hit, Document doc) {
		if (doc != null && hit.fields().size() > 0) {
			for (Map.Entry<String, SearchHitField> fieldEntry : hit.fields().entrySet()) {
				if (fieldEntry.getKey().startsWith(getAnnotationField())) {
					String fieldName = fieldEntry.getKey();

					switch (fieldName) {
						case "efo_uri.label":
							doc.setEfoLabels(getStringValues(fieldEntry.getValue().getValues()));
							break;
						case "efo_uri.child_labels":
							doc.setChildLabels(getStringValues(fieldEntry.getValue().getValues()));
							break;
						case "efo_uri.parent_labels":
							doc.setParentLabels(getStringValues(fieldEntry.getValue().getValues()));
							break;
						default:
							String shortName = fieldName.substring("efo_uri.".length());
							if (fieldName.endsWith("_rel_uris")) {
								doc.getRelatedIris().put(shortName, getStringValues(fieldEntry.getValue().getValues()));
							} else if (fieldName.endsWith("_rel_labels")) {
								List<String> labels = getStringValues(fieldEntry.getValue().getValues());
								if (labels != null) {
									doc.getRelatedLabels().put(shortName, labels);
								}
							}
					}
				}
			}
		}
	}

	private List<String> getStringValues(List<Object> fieldValues) {
		List<String> retList;
		if (fieldValues == null || fieldValues.size() == 0) {
			retList = null;
		} else {
			retList = fieldValues.stream().map(Object::toString).collect(Collectors.toList());
		}
		return retList;
	}

}