org.elasticsearch.search.controller.SearchPhaseController Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of elasticsearch Show documentation
Show all versions of elasticsearch Show documentation
Elasticsearch subproject :server
/*
* Licensed to Elastic Search and Shay Banon under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. Elastic Search licenses this
* file to you 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
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* 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 org.elasticsearch.search.controller;
import org.apache.lucene.index.Term;
import org.apache.lucene.search.*;
import org.apache.lucene.util.PriorityQueue;
import org.elasticsearch.common.collect.Iterables;
import org.elasticsearch.common.collect.Lists;
import org.elasticsearch.common.collect.Maps;
import org.elasticsearch.common.trove.ExtTIntArrayList;
import org.elasticsearch.common.trove.ExtTObjectIntHasMap;
import org.elasticsearch.search.SearchShardTarget;
import org.elasticsearch.search.dfs.AggregatedDfs;
import org.elasticsearch.search.dfs.DfsSearchResult;
import org.elasticsearch.search.facet.Facet;
import org.elasticsearch.search.facet.internal.InternalFacet;
import org.elasticsearch.search.facet.internal.InternalFacets;
import org.elasticsearch.search.fetch.FetchSearchResult;
import org.elasticsearch.search.fetch.FetchSearchResultProvider;
import org.elasticsearch.search.internal.InternalSearchHit;
import org.elasticsearch.search.internal.InternalSearchHits;
import org.elasticsearch.search.internal.InternalSearchResponse;
import org.elasticsearch.search.query.QuerySearchResult;
import org.elasticsearch.search.query.QuerySearchResultProvider;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;
/**
* @author kimchy (Shay Banon)
*/
public class SearchPhaseController {
private static final ShardDoc[] EMPTY = new ShardDoc[0];
public AggregatedDfs aggregateDfs(Iterable results) {
ExtTObjectIntHasMap dfMap = new ExtTObjectIntHasMap().defaultReturnValue(-1);
long aggMaxDoc = 0;
for (DfsSearchResult result : results) {
for (int i = 0; i < result.freqs().length; i++) {
dfMap.adjustOrPutValue(result.terms()[i], result.freqs()[i], result.freqs()[i]);
}
aggMaxDoc += result.maxDoc();
}
return new AggregatedDfs(dfMap, aggMaxDoc);
}
public ShardDoc[] sortDocs(Collection results) {
if (results.isEmpty()) {
return EMPTY;
}
QuerySearchResultProvider queryResultProvider = Iterables.get(results, 0);
int totalNumDocs = 0;
int queueSize = queryResultProvider.queryResult().from() + queryResultProvider.queryResult().size();
if (queryResultProvider.includeFetch()) {
// if we did both query and fetch on the same go, we have fetched all the docs from each shards already, use them...
queueSize *= results.size();
}
PriorityQueue queue;
if (queryResultProvider.queryResult().topDocs() instanceof TopFieldDocs) {
// sorting, first if the type is a String, chance CUSTOM to STRING so we handle nulls properly (since our CUSTOM String sorting might return null)
TopFieldDocs fieldDocs = (TopFieldDocs) queryResultProvider.queryResult().topDocs();
for (int i = 0; i < fieldDocs.fields.length; i++) {
boolean allValuesAreNull = true;
boolean resolvedField = false;
for (QuerySearchResultProvider resultProvider : results) {
for (ScoreDoc doc : resultProvider.queryResult().topDocs().scoreDocs) {
FieldDoc fDoc = (FieldDoc) doc;
if (fDoc.fields[i] != null) {
allValuesAreNull = false;
if (fDoc.fields[i] instanceof String) {
fieldDocs.fields[i] = new SortField(fieldDocs.fields[i].getField(), SortField.STRING, fieldDocs.fields[i].getReverse());
}
resolvedField = true;
break;
}
}
if (resolvedField) {
break;
}
}
if (!resolvedField && allValuesAreNull && fieldDocs.fields[i].getField() != null) {
// we did not manage to resolve a field (and its not score or doc, which have no field), and all the fields are null (which can only happen for STRING), make it a STRING
fieldDocs.fields[i] = new SortField(fieldDocs.fields[i].getField(), SortField.STRING, fieldDocs.fields[i].getReverse());
}
}
queue = new ShardFieldDocSortedHitQueue(fieldDocs.fields, queueSize);
// we need to accumulate for all and then filter the from
for (QuerySearchResultProvider resultProvider : results) {
QuerySearchResult result = resultProvider.queryResult();
ScoreDoc[] scoreDocs = result.topDocs().scoreDocs;
totalNumDocs += scoreDocs.length;
for (ScoreDoc doc : scoreDocs) {
ShardFieldDoc nodeFieldDoc = new ShardFieldDoc(result.shardTarget(), doc.doc, doc.score, ((FieldDoc) doc).fields);
if (queue.insertWithOverflow(nodeFieldDoc) == nodeFieldDoc) {
// filled the queue, break
break;
}
}
}
} else {
queue = new ScoreDocQueue(queueSize); // we need to accumulate for all and then filter the from
for (QuerySearchResultProvider resultProvider : results) {
QuerySearchResult result = resultProvider.queryResult();
ScoreDoc[] scoreDocs = result.topDocs().scoreDocs;
totalNumDocs += scoreDocs.length;
for (ScoreDoc doc : scoreDocs) {
ShardScoreDoc nodeScoreDoc = new ShardScoreDoc(result.shardTarget(), doc.doc, doc.score);
if (queue.insertWithOverflow(nodeScoreDoc) == nodeScoreDoc) {
// filled the queue, break
break;
}
}
}
}
int resultDocsSize = queryResultProvider.queryResult().size();
if (queryResultProvider.includeFetch()) {
// if we did both query and fetch on the same go, we have fetched all the docs from each shards already, use them...
resultDocsSize *= results.size();
}
if (totalNumDocs < queueSize) {
resultDocsSize = totalNumDocs - queryResultProvider.queryResult().from();
}
if (resultDocsSize <= 0) {
return EMPTY;
}
ShardDoc[] shardDocs = new ShardDoc[resultDocsSize];
for (int i = resultDocsSize - 1; i >= 0; i--) // put docs in array
shardDocs[i] = (ShardDoc) queue.pop();
return shardDocs;
}
public Map docIdsToLoad(ShardDoc[] shardDocs) {
Map result = Maps.newHashMap();
for (ShardDoc shardDoc : shardDocs) {
ExtTIntArrayList list = result.get(shardDoc.shardTarget());
if (list == null) {
list = new ExtTIntArrayList(); // can't be shared!, uses unsafe on it later on
result.put(shardDoc.shardTarget(), list);
}
list.add(shardDoc.docId());
}
return result;
}
public InternalSearchResponse merge(ShardDoc[] sortedDocs, Map queryResults, Map fetchResults) {
boolean sorted = false;
int sortScoreIndex = -1;
QuerySearchResult querySearchResult = Iterables.get(queryResults.values(), 0).queryResult();
if (querySearchResult.topDocs() instanceof TopFieldDocs) {
sorted = true;
TopFieldDocs fieldDocs = (TopFieldDocs) querySearchResult.queryResult().topDocs();
for (int i = 0; i < fieldDocs.fields.length; i++) {
if (fieldDocs.fields[i].getType() == SortField.SCORE) {
sortScoreIndex = i;
}
}
}
// merge facets
InternalFacets facets = null;
if (!queryResults.isEmpty()) {
// we rely on the fact that the order of facets is the same on all query results
QuerySearchResult queryResult = queryResults.values().iterator().next().queryResult();
// we assume the facets are in the same order!
if (queryResult.facets() != null && queryResult.facets().facets() != null && !queryResult.facets().facets().isEmpty()) {
List allFacets = Lists.newArrayList();
for (QuerySearchResultProvider queryResultProvider : queryResults.values()) {
allFacets.addAll(queryResultProvider.queryResult().facets().facets());
}
List mergedFacets = Lists.newArrayList();
for (Facet facet : queryResult.facets().facets()) {
mergedFacets.add(((InternalFacet) facet).aggregate(allFacets));
}
facets = new InternalFacets(mergedFacets);
}
}
// count the total (we use the query result provider here, since we might not get any hits (we scrolled past them))
long totalHits = 0;
float maxScore = Float.NEGATIVE_INFINITY;
for (QuerySearchResultProvider queryResultProvider : queryResults.values()) {
totalHits += queryResultProvider.queryResult().topDocs().totalHits;
if (!Float.isNaN(queryResultProvider.queryResult().topDocs().getMaxScore())) {
maxScore = Math.max(maxScore, queryResultProvider.queryResult().topDocs().getMaxScore());
}
}
if (Float.isInfinite(maxScore)) {
maxScore = Float.NaN;
}
// clean the fetch counter
for (FetchSearchResultProvider fetchSearchResultProvider : fetchResults.values()) {
fetchSearchResultProvider.fetchResult().initCounter();
}
// merge hits
List hits = new ArrayList();
if (!fetchResults.isEmpty()) {
for (ShardDoc shardDoc : sortedDocs) {
FetchSearchResultProvider fetchResultProvider = fetchResults.get(shardDoc.shardTarget());
if (fetchResultProvider == null) {
continue;
}
FetchSearchResult fetchResult = fetchResultProvider.fetchResult();
int index = fetchResult.counterGetAndIncrement();
if (index < fetchResult.hits().internalHits().length) {
InternalSearchHit searchHit = fetchResult.hits().internalHits()[index];
searchHit.score(shardDoc.score());
searchHit.shard(fetchResult.shardTarget());
if (sorted) {
FieldDoc fieldDoc = (FieldDoc) shardDoc;
searchHit.sortValues(fieldDoc.fields);
if (sortScoreIndex != -1) {
searchHit.score(((Number) fieldDoc.fields[sortScoreIndex]).floatValue());
}
}
hits.add(searchHit);
}
}
}
InternalSearchHits searchHits = new InternalSearchHits(hits.toArray(new InternalSearchHit[hits.size()]), totalHits, maxScore);
return new InternalSearchResponse(searchHits, facets);
}
}