@Override protected SuggestResponse newResponse( SuggestRequest request, AtomicReferenceArray shardsResponses, ClusterState clusterState) { int successfulShards = 0; int failedShards = 0; final Map<String, List<Suggest.Suggestion>> groupedSuggestions = new HashMap<>(); List<ShardOperationFailedException> shardFailures = null; for (int i = 0; i < shardsResponses.length(); i++) { Object shardResponse = shardsResponses.get(i); if (shardResponse == null) { // simply ignore non active shards } else if (shardResponse instanceof BroadcastShardOperationFailedException) { failedShards++; if (shardFailures == null) { shardFailures = newArrayList(); } shardFailures.add( new DefaultShardOperationFailedException( (BroadcastShardOperationFailedException) shardResponse)); } else { Suggest suggest = ((ShardSuggestResponse) shardResponse).getSuggest(); Suggest.group(groupedSuggestions, suggest); successfulShards++; } } return new SuggestResponse( new Suggest(Suggest.reduce(groupedSuggestions)), shardsResponses.length(), successfulShards, failedShards, shardFailures); }
public InternalSearchResponse merge( ScoreDoc[] sortedDocs, AtomicArray<? extends QuerySearchResultProvider> queryResultsArr, AtomicArray<? extends FetchSearchResultProvider> fetchResultsArr) { List<? extends AtomicArray.Entry<? extends QuerySearchResultProvider>> queryResults = queryResultsArr.asList(); List<? extends AtomicArray.Entry<? extends FetchSearchResultProvider>> fetchResults = fetchResultsArr.asList(); if (queryResults.isEmpty()) { return InternalSearchResponse.empty(); } QuerySearchResult firstResult = queryResults.get(0).value.queryResult(); boolean sorted = false; int sortScoreIndex = -1; if (firstResult.topDocs() instanceof TopFieldDocs) { sorted = true; TopFieldDocs fieldDocs = (TopFieldDocs) firstResult.queryResult().topDocs(); for (int i = 0; i < fieldDocs.fields.length; i++) { if (fieldDocs.fields[i].getType() == SortField.Type.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 if (firstResult.facets() != null && firstResult.facets().facets() != null && !firstResult.facets().facets().isEmpty()) { List<Facet> aggregatedFacets = Lists.newArrayList(); List<Facet> namedFacets = Lists.newArrayList(); for (Facet facet : firstResult.facets()) { // aggregate each facet name into a single list, and aggregate it namedFacets.clear(); for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : queryResults) { for (Facet facet1 : entry.value.queryResult().facets()) { if (facet.getName().equals(facet1.getName())) { namedFacets.add(facet1); } } } if (!namedFacets.isEmpty()) { Facet aggregatedFacet = ((InternalFacet) namedFacets.get(0)) .reduce(new InternalFacet.ReduceContext(cacheRecycler, namedFacets)); aggregatedFacets.add(aggregatedFacet); } } facets = new InternalFacets(aggregatedFacets); } } // 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; boolean timedOut = false; Boolean terminatedEarly = null; for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : queryResults) { QuerySearchResult result = entry.value.queryResult(); if (result.searchTimedOut()) { timedOut = true; } if (result.terminatedEarly() != null) { if (terminatedEarly == null) { terminatedEarly = result.terminatedEarly(); } else if (result.terminatedEarly()) { terminatedEarly = true; } } totalHits += result.topDocs().totalHits; if (!Float.isNaN(result.topDocs().getMaxScore())) { maxScore = Math.max(maxScore, result.topDocs().getMaxScore()); } } if (Float.isInfinite(maxScore)) { maxScore = Float.NaN; } // clean the fetch counter for (AtomicArray.Entry<? extends FetchSearchResultProvider> entry : fetchResults) { entry.value.fetchResult().initCounter(); } // merge hits List<InternalSearchHit> hits = new ArrayList<>(); if (!fetchResults.isEmpty()) { for (ScoreDoc shardDoc : sortedDocs) { FetchSearchResultProvider fetchResultProvider = fetchResultsArr.get(shardDoc.shardIndex); 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); } } } // merge suggest results Suggest suggest = null; if (!queryResults.isEmpty()) { final Map<String, List<Suggest.Suggestion>> groupedSuggestions = new HashMap<>(); boolean hasSuggestions = false; for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : queryResults) { Suggest shardResult = entry.value.queryResult().queryResult().suggest(); if (shardResult == null) { continue; } hasSuggestions = true; Suggest.group(groupedSuggestions, shardResult); } suggest = hasSuggestions ? new Suggest(Suggest.Fields.SUGGEST, Suggest.reduce(groupedSuggestions)) : null; } // merge addAggregation InternalAggregations aggregations = null; if (!queryResults.isEmpty()) { if (firstResult.aggregations() != null && firstResult.aggregations().asList() != null) { List<InternalAggregations> aggregationsList = new ArrayList<>(queryResults.size()); for (AtomicArray.Entry<? extends QuerySearchResultProvider> entry : queryResults) { aggregationsList.add((InternalAggregations) entry.value.queryResult().aggregations()); } aggregations = InternalAggregations.reduce( aggregationsList, new ReduceContext(null, bigArrays, scriptService)); } } InternalSearchHits searchHits = new InternalSearchHits( hits.toArray(new InternalSearchHit[hits.size()]), totalHits, maxScore); return new InternalSearchResponse( searchHits, facets, aggregations, suggest, timedOut, terminatedEarly); }
public static Suggest readSuggest(StreamInput in) throws IOException { Suggest result = new Suggest(); result.readFrom(in); return result; }