com.liferay.portal.search.solr.internal.SolrIndexSearcher.java Source code

Java tutorial

Introduction

Here is the source code for com.liferay.portal.search.solr.internal.SolrIndexSearcher.java

Source

/**
 * Copyright (c) 2000-present Liferay, Inc. All rights reserved.
 *
 * This library is free software; you can redistribute it and/or modify it under
 * the terms of the GNU Lesser General Public License as published by the Free
 * Software Foundation; either version 2.1 of the License, or (at your option)
 * any later version.
 *
 * This library is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
 * details.
 */

package com.liferay.portal.search.solr.internal;

import com.liferay.portal.kernel.dao.orm.QueryUtil;
import com.liferay.portal.kernel.dao.search.SearchPaginationUtil;
import com.liferay.portal.kernel.log.Log;
import com.liferay.portal.kernel.log.LogFactoryUtil;
import com.liferay.portal.kernel.search.BaseIndexSearcher;
import com.liferay.portal.kernel.search.Document;
import com.liferay.portal.kernel.search.DocumentImpl;
import com.liferay.portal.kernel.search.Field;
import com.liferay.portal.kernel.search.Hits;
import com.liferay.portal.kernel.search.HitsImpl;
import com.liferay.portal.kernel.search.IndexSearcher;
import com.liferay.portal.kernel.search.Query;
import com.liferay.portal.kernel.search.QueryConfig;
import com.liferay.portal.kernel.search.SearchContext;
import com.liferay.portal.kernel.search.SearchException;
import com.liferay.portal.kernel.search.Sort;
import com.liferay.portal.kernel.search.facet.Facet;
import com.liferay.portal.kernel.search.facet.RangeFacet;
import com.liferay.portal.kernel.search.facet.collector.FacetCollector;
import com.liferay.portal.kernel.search.facet.config.FacetConfiguration;
import com.liferay.portal.kernel.search.filter.FilterTranslator;
import com.liferay.portal.kernel.search.highlight.HighlightUtil;
import com.liferay.portal.kernel.search.query.QueryTranslator;
import com.liferay.portal.kernel.search.suggest.QuerySuggester;
import com.liferay.portal.kernel.util.ArrayUtil;
import com.liferay.portal.kernel.util.GetterUtil;
import com.liferay.portal.kernel.util.ListUtil;
import com.liferay.portal.kernel.util.MapUtil;
import com.liferay.portal.kernel.util.SetUtil;
import com.liferay.portal.kernel.util.StringPool;
import com.liferay.portal.kernel.util.StringUtil;
import com.liferay.portal.kernel.util.Validator;
import com.liferay.portal.search.solr.facet.FacetProcessor;
import com.liferay.portal.search.solr.internal.facet.CompositeFacetProcessor;
import com.liferay.portal.search.solr.internal.facet.SolrFacetFieldCollector;
import com.liferay.portal.search.solr.internal.facet.SolrFacetQueryCollector;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

import org.apache.commons.lang.time.StopWatch;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.SolrQuery.ORDER;
import org.apache.solr.client.solrj.SolrQuery.SortClause;
import org.apache.solr.client.solrj.SolrRequest.METHOD;
import org.apache.solr.client.solrj.SolrServer;
import org.apache.solr.client.solrj.response.FacetField;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.params.FacetParams;

import org.osgi.service.component.annotations.Activate;
import org.osgi.service.component.annotations.Component;
import org.osgi.service.component.annotations.Reference;

/**
 * @author Bruno Farache
 * @author Zsolt Berentey
 * @author Raymond Aug
 */
@Component(immediate = true, property = { "search.engine.impl=Solr",
        "swallow.exception=true" }, service = IndexSearcher.class)
public class SolrIndexSearcher extends BaseIndexSearcher {

    @Override
    public String getQueryString(SearchContext searchContext, Query query) {
        return translateQuery(query, searchContext);
    }

    @Override
    public Hits search(SearchContext searchContext, Query query) throws SearchException {

        StopWatch stopWatch = new StopWatch();

        stopWatch.start();

        try {
            int total = (int) searchCount(searchContext, query);

            int start = searchContext.getStart();
            int end = searchContext.getEnd();

            if ((start == QueryUtil.ALL_POS) && (end == QueryUtil.ALL_POS)) {
                start = 0;
                end = total;
            }

            int[] startAndEnd = SearchPaginationUtil.calculateStartAndEnd(start, end, total);

            start = startAndEnd[0];
            end = startAndEnd[1];

            Hits hits = doSearchHits(searchContext, query, start, end);

            hits.setStart(stopWatch.getStartTime());

            return hits;
        } catch (Exception e) {
            if (_log.isWarnEnabled()) {
                _log.warn(e, e);
            }

            if (!_swallowException) {
                throw new SearchException(e.getMessage(), e);
            }

            return new HitsImpl();
        } finally {
            if (_log.isInfoEnabled()) {
                stopWatch.stop();

                _log.info("Searching " + query.toString() + " took " + stopWatch.getTime() + " ms");
            }
        }
    }

    public long searchCount(SearchContext searchContext, Query query) throws SearchException {

        StopWatch stopWatch = new StopWatch();

        stopWatch.start();

        try {
            return doSearchCount(searchContext, query);
        } catch (Exception e) {
            if (_log.isWarnEnabled()) {
                _log.warn(e, e);
            }

            if (!_swallowException) {
                throw new SearchException(e.getMessage(), e);
            }

            return 0;
        } finally {
            if (_log.isInfoEnabled()) {
                stopWatch.stop();

                _log.info("Searching " + query.toString() + " took " + stopWatch.getTime() + " ms");
            }
        }
    }

    @Override
    @Reference(target = "(search.engine.impl=Solr)", unbind = "-")
    public void setQuerySuggester(QuerySuggester querySuggester) {
        super.setQuerySuggester(querySuggester);
    }

    public void setSolrServer(SolrServer solrServer) {
        _solrServer = solrServer;
    }

    @Activate
    protected void activate(Map<String, Object> properties) {
        _swallowException = MapUtil.getBoolean(properties, "swallow.exception", true);
    }

    protected void addFacets(SolrQuery solrQuery, SearchContext searchContext) {
        Map<String, Facet> facets = searchContext.getFacets();

        for (Facet facet : facets.values()) {
            if (facet.isStatic()) {
                continue;
            }

            FacetConfiguration facetConfiguration = facet.getFacetConfiguration();

            _facetProcessor.processFacet(solrQuery, facet);

            String facetSort = FacetParams.FACET_SORT_COUNT;

            String order = facetConfiguration.getOrder();

            if (order.equals("OrderValueAsc")) {
                facetSort = FacetParams.FACET_SORT_INDEX;
            }

            solrQuery.add("f." + facetConfiguration.getFieldName() + ".facet.sort", facetSort);
        }

        solrQuery.setFacetLimit(-1);
    }

    protected void addHighlightedField(SolrQuery solrQuery, QueryConfig queryConfig, String fieldName) {

        solrQuery.addHighlightField(fieldName);

        String localizedFieldName = DocumentImpl.getLocalizedName(queryConfig.getLocale(), fieldName);

        solrQuery.addHighlightField(localizedFieldName);
    }

    protected void addHighlights(SolrQuery solrQuery, QueryConfig queryConfig) {
        if (!queryConfig.isHighlightEnabled()) {
            return;
        }

        solrQuery.setHighlight(true);
        solrQuery.setHighlightFragsize(queryConfig.getHighlightFragmentSize());
        solrQuery.setHighlightSimplePost(HighlightUtil.HIGHLIGHT_TAG_CLOSE);
        solrQuery.setHighlightSimplePre(HighlightUtil.HIGHLIGHT_TAG_OPEN);
        solrQuery.setHighlightSnippets(queryConfig.getHighlightSnippetSize());

        for (String highlightFieldName : queryConfig.getHighlightFieldNames()) {
            addHighlightedField(solrQuery, queryConfig, highlightFieldName);
        }

        solrQuery.setHighlightRequireFieldMatch(queryConfig.isHighlightRequireFieldMatch());
    }

    protected void addPagination(SolrQuery solrQuery, int start, int end) {
        solrQuery.setRows(end - start);
        solrQuery.setStart(start);
    }

    protected void addSelectedFields(SolrQuery solrQuery, QueryConfig queryConfig) {

        if (queryConfig.isAllFieldsSelected()) {
            return;
        }

        Set<String> selectedFieldNames = SetUtil.fromArray(queryConfig.getSelectedFieldNames());

        if (!selectedFieldNames.contains(Field.UID)) {
            selectedFieldNames.add(Field.UID);
        }

        solrQuery.setFields(selectedFieldNames.toArray(new String[selectedFieldNames.size()]));
    }

    protected void addSnippets(SolrDocument solrDocument, Document document, QueryConfig queryConfig,
            Set<String> queryTerms, Map<String, Map<String, List<String>>> highlights) {

        if (!queryConfig.isHighlightEnabled()) {
            return;
        }

        for (String highlightFieldName : queryConfig.getHighlightFieldNames()) {
            addSnippets(solrDocument, document, queryTerms, highlights, highlightFieldName,
                    queryConfig.getLocale());
        }
    }

    protected void addSnippets(SolrDocument solrDocument, Document document, Set<String> queryTerms,
            Map<String, Map<String, List<String>>> highlights, String fieldName, Locale locale) {

        if (MapUtil.isEmpty(highlights)) {
            return;
        }

        String key = (String) solrDocument.getFieldValue(Field.UID);

        Map<String, List<String>> uidHighlights = highlights.get(key);

        String localizedFieldName = DocumentImpl.getLocalizedName(locale, fieldName);

        String snippetFieldName = localizedFieldName;

        List<String> snippets = uidHighlights.get(localizedFieldName);

        if (snippets == null) {
            snippets = uidHighlights.get(fieldName);

            snippetFieldName = fieldName;
        }

        String snippet = StringPool.BLANK;

        if (ListUtil.isNotEmpty(snippets)) {
            snippet = StringUtil.merge(snippets, StringPool.TRIPLE_PERIOD);

            if (Validator.isNotNull(snippet)) {
                snippet = snippet.concat(StringPool.TRIPLE_PERIOD);
            }
        }

        HighlightUtil.addSnippet(document, queryTerms, snippet, snippetFieldName);
    }

    protected void addSort(SolrQuery solrQuery, Sort[] sorts) {
        if (ArrayUtil.isEmpty(sorts)) {
            return;
        }

        Set<String> sortFieldNames = new HashSet<>(sorts.length);

        for (Sort sort : sorts) {
            if (sort == null) {
                continue;
            }

            String sortFieldName = DocumentImpl.getSortFieldName(sort, "score");

            if (sortFieldNames.contains(sortFieldName)) {
                continue;
            }

            sortFieldNames.add(sortFieldName);

            ORDER order = ORDER.asc;

            if (sort.isReverse() || sortFieldName.equals("score")) {
                order = ORDER.desc;
            }

            solrQuery.addSort(new SortClause(sortFieldName, order));
        }
    }

    protected QueryResponse doSearch(SearchContext searchContext, Query query, int start, int end, boolean count)
            throws Exception {

        QueryConfig queryConfig = query.getQueryConfig();

        SolrQuery solrQuery = new SolrQuery();

        if (!count) {
            addFacets(solrQuery, searchContext);
            addHighlights(solrQuery, queryConfig);
            addPagination(solrQuery, start, end);
            addSelectedFields(solrQuery, queryConfig);
            addSort(solrQuery, searchContext.getSorts());

            solrQuery.setIncludeScore(queryConfig.isScoreEnabled());
        } else {
            solrQuery.setRows(0);
        }

        String queryString = translateQuery(query, searchContext);

        solrQuery.setQuery(queryString);

        String filterQuery = _filterTranslator.translate(query.getPreBooleanFilter(), searchContext);

        solrQuery.setFilterQueries(filterQuery);

        QueryResponse queryResponse = executeSearchRequest(solrQuery);

        if (_log.isInfoEnabled()) {
            _log.info("The search engine processed " + solrQuery.getQuery() + " in "
                    + queryResponse.getElapsedTime() + " ms");
        }

        return queryResponse;
    }

    protected long doSearchCount(SearchContext searchContext, Query query) throws Exception {

        QueryResponse queryResponse = doSearch(searchContext, query, searchContext.getStart(),
                searchContext.getEnd(), true);

        SolrDocumentList solrDocumentList = queryResponse.getResults();

        return solrDocumentList.getNumFound();
    }

    protected Hits doSearchHits(SearchContext searchContext, Query query, int start, int end) throws Exception {

        QueryResponse queryResponse = doSearch(searchContext, query, start, end, false);

        Hits hits = processResponse(queryResponse, searchContext, query);

        return hits;
    }

    protected QueryResponse executeSearchRequest(SolrQuery solrQuery) throws Exception {

        return _solrServer.query(solrQuery, METHOD.POST);
    }

    protected Hits processResponse(QueryResponse queryResponse, SearchContext searchContext, Query query) {

        long startTime = System.currentTimeMillis();

        SolrDocumentList solrDocumentList = queryResponse.getResults();

        updateFacetCollectors(queryResponse, searchContext);

        Hits hits = new HitsImpl();

        List<Document> documents = new ArrayList<>();
        Set<String> queryTerms = new HashSet<>();
        List<Float> scores = new ArrayList<>();

        QueryConfig queryConfig = query.getQueryConfig();
        Map<String, Map<String, List<String>>> highlights = queryResponse.getHighlighting();

        for (SolrDocument solrDocument : solrDocumentList) {
            Document document = processSolrDocument(solrDocument, queryConfig);

            documents.add(document);

            addSnippets(solrDocument, document, queryConfig, queryTerms, highlights);

            float score = GetterUtil.getFloat(String.valueOf(solrDocument.getFieldValue("score")));

            scores.add(score);
        }

        hits.setDocs(documents.toArray(new Document[documents.size()]));
        hits.setLength((int) solrDocumentList.getNumFound());
        hits.setQuery(query);
        hits.setQueryTerms(queryTerms.toArray(new String[queryTerms.size()]));
        hits.setScores(ArrayUtil.toFloatArray(scores));
        hits.setSearchTime(queryResponse.getQTime());
        hits.setStart(startTime);

        return hits;
    }

    protected Document processSolrDocument(SolrDocument solrDocument, QueryConfig queryConfig) {

        Document document = new DocumentImpl();

        Collection<String> fieldNames = solrDocument.getFieldNames();

        for (String fieldName : fieldNames) {
            Collection<Object> fieldValues = solrDocument.getFieldValues(fieldName);

            Field field = new Field(fieldName,
                    ArrayUtil.toStringArray(fieldValues.toArray(new Object[fieldValues.size()])));

            document.add(field);
        }

        populateUID(document, queryConfig);

        return document;
    }

    @Reference(service = CompositeFacetProcessor.class, unbind = "-")
    protected void setFacetProcessor(FacetProcessor<SolrQuery> facetProcessor) {
        _facetProcessor = facetProcessor;
    }

    @Reference(target = "(search.engine.impl=Solr)", unbind = "-")
    protected void setFilterTranslator(FilterTranslator<String> filterTranslator) {

        _filterTranslator = filterTranslator;
    }

    @Reference(target = "(search.engine.impl=Solr)", unbind = "-")
    protected void setQueryTranslator(QueryTranslator<String> queryTranslator) {
        _queryTranslator = queryTranslator;
    }

    protected String translateQuery(Query query, SearchContext searchContext) {
        return _queryTranslator.translate(query, searchContext);
    }

    protected void updateFacetCollectors(QueryResponse queryResponse, SearchContext searchContext) {

        Map<String, Facet> facetsMap = searchContext.getFacets();

        List<FacetField> facetFields = queryResponse.getFacetFields();

        if (ListUtil.isEmpty(facetFields)) {
            return;
        }

        for (FacetField facetField : facetFields) {
            Facet facet = facetsMap.get(facetField.getName());

            FacetCollector facetCollector = null;

            if (facet instanceof RangeFacet) {
                facetCollector = new SolrFacetQueryCollector(facetField.getName(), queryResponse.getFacetQuery());
            } else {
                facetCollector = new SolrFacetFieldCollector(facetField.getName(), facetField);
            }

            facet.setFacetCollector(facetCollector);
        }
    }

    private static final Log _log = LogFactoryUtil.getLog(SolrIndexSearcher.class);

    private FacetProcessor<SolrQuery> _facetProcessor;
    private FilterTranslator<String> _filterTranslator;
    private QueryTranslator<String> _queryTranslator;
    private SolrServer _solrServer;
    private boolean _swallowException;

}